[Spec Decode] Chunk cold draft-KV ingest to bound first-propose dispatch stall (vllm-project#482 direction 3) (#643)

## Problem

CC: @ricky-chaoju

Issue #482 measures two stacked costs in the draft-model proposer. This
PR addresses the tail end of Problem 1: when a request's draft KV cache
is cold, the first propose ingests the **entire prompt** into the draft
model's KV in one tiled prefill forward, stalling the engine for the
full prompt length (4.59 s measured at an 8k prefix on an M4 Pro).

That single forward has two costs:

1. **Dispatch stall** — one GPU dispatch sized to the whole prompt. The
engine blocks until it completes.
2. **Per-forward logits allocation** — prefill logits are materialized
for every position, so the forward's logits allocation scales as
`draft_vocab × prompt length` (~1.24 G floats at 8k for Qwen3-0.6B).

Context: direction 1 (#630, scheduler-managed draft cache) and direction
2 (#636, skip re-ingesting accepted draft tokens) are already merged.
This is the issue's third fix direction: a mitigation for the cold path
that remains.

## Solution

Chunk the cold ingest into fixed-size forwards in
`_ingest_and_draft_first`, behind a new env knob
`VLLM_METAL_SPEC_INGEST_CHUNK` (default `1024`; set `0` for the old
single-forward behavior):

- **Bounded dispatches.** No single forward is larger than the chunk
(≤1024 tokens), instead of one prompt-length forward.
- **Bounded per-forward logits.** Each ingest forward's logits
allocation scales with the chunk (`draft_vocab × 1024`), not the prompt
(`draft_vocab × 8191`) — 8× lower at 8k. This does not move the engine's
measured peak memory: the cold step's high-water mark is set before the
ingest forwards run, identically in both arms (measured, see benchmark).
- **Identical results.** Chunk starts map to draft position
`draft_seq_len + start` (`vllm_metal/v1/draft_model_proposer.py:649`),
so a chunk after an existing prefix lands in the correct KV slot. The
final row of each plan's last chunk is the last ingested token, whose
logits predict the plan's first draft token — bit-identical to the
single-forward path, verified byte-for-byte in the benchmark below and
by an e2e test.
- `1024` is a multiple of the block size (16).

**What this is not:** this is a dispatch-shape mitigation, not a
throughput or peak-memory optimization. The total ingest work is
unchanged, and the engine still completes the whole ingest synchronously
before the first draft token exists. Measured cold-request latency is
flat to slightly worse (+2–6 %), and measured engine peak memory is
unchanged. Chunking only changes the *shape* of the execution — a
sequence of bounded dispatches instead of one giant one.

## Benchmark

Apple M5, 32 GB, macOS 26.6.2. Draft == target == `Qwen/Qwen3-0.6B`
(greedy, 100 % acceptance — the #482 setup), `num_speculative_tokens=3`,
96 generated tokens, one engine per process, prefix caching on,
`VLLM_METAL_MEMORY_FRACTION=0.4`, `VLLM_ENABLE_V1_MULTIPROCESSING=0`.

- `gen1_cold` — fresh request, cold draft KV (full-prompt ingest)
- `gen2_resubmit` — identical prompt, draft-cache reuse (16-token
ingest)

| prefix | chunking | gen1_cold wall | gen1_cold TPOT | gen2_resubmit
wall | gen2 TPOT | cold extra vs resubmit | lossless |
| -----: | ----------- | -------------: | -------------: |
-----------------: | --------: | ---------------------: |
---------------------- |
| 8192 | off (`0`) | 6.085 s | 63.4 ms | 2.847 s | 29.7 ms | 3.24 s |
pass (both) |
| 8192 | on (`1024`) | 6.206 s | 64.6 ms | 2.787 s | 29.0 ms | 3.42 s |
pass (both) |
| 2048 | off (`0`) | 2.327 s | 24.2 ms | 1.797 s | 18.7 ms | 0.53 s |
gen1 pass, gen2 flake |
| 2048 | on (`1024`) | 2.470 s | 25.7 ms | 1.749 s | 18.2 ms | 0.72 s |
gen1 pass, gen2 flake |

What chunking changes structurally (and what it does not):

| property | unchunked | chunked (`1024`) |
| ----------------------------- | ------------------ |
----------------------------------- |
| cold-ingest dispatch | 1 × 8191 tokens | 8 × ≤1024 tokens |
| per-forward logits allocation | vocab × 8191 | vocab × 1024 (8× lower
at 8k) |
| measured engine peak, gen1 | 13.29 GiB | 13.29 GiB |

Reading the numbers:

- **Outputs are unchanged.** Both arms produce identical tokens at both
prefixes, cold and resubmit.
- **Cold first-request cost** rises slightly: +2 % at 8k, +6 % at 2k
(~120–190 ms absolute) — the price of splitting one forward into chunks.
At 2k there are only two chunks, so fixed per-dispatch overhead is a
larger share of the total.
- **Steady state is untouched**, within cross-process noise (29.0 vs
29.7 ms/TPOT @8k) — the knob only fires on the cold-ingest path.
- **Engine peak memory is unchanged.** `mx.reset_peak_memory()` before
the cold generate and `mx.get_peak_memory()` after (multiprocessing off,
`--skip-warm`, 8k): 13.29 GiB with chunking on and off. Isolating the
first propose with the same reset/read pair gives 12.85 GiB in both
arms. The high-water mark is reached before the ingest forwards begin —
KV-block/step setup for the full prompt allocates ~2.3 GiB over the
post-init floor either way — so the draft's per-forward logits (the only
allocation chunking shrinks) never raise the engine's peak in either
arm. Chunking bounds the dispatch and the per-forward logits buffer; it
does not reduce the engine's peak memory on this setup.
- `gen2_resubmit` at 2048 diverges from its reference at token 23 in
**both** arms identically — a pre-existing reference-process flake (the
run logs warn about seed determinism), not a chunking effect. At 8192
everything is lossless.

## Tests

- `test_draft_cache_policy_gaps.py` — chunk-size-0 single-forward
equivalence; chunk larger than prompt; exact-multiple chunking;
mixed-length plans sharing rounds.
- `test_draft_model_proposer.py` — per-chunk forward counting; small
ingests ignoring the knob (decode path); token-identity e2e between
chunked and unchunked cold ingests. The position-encoding stub's
`VOCAB_SIZE` (64) exceeds the largest tested position (31), so `%
VOCAB_SIZE` cannot fold a dropped chunk offset back onto the correct
token: a mutation test that drops `+ start` fails three of these tests,
and all pass with it.
- `tools/benchmark/draft_resubmit_benchmark.py` — records per-generate
and first-propose peak memory (`peak_mem_gb`, `first_propose_peak_gb`).
- `scripts/lint.sh` (shellcheck, ruff, mypy) passes.

## Configuration

`VLLM_METAL_SPEC_INGEST_CHUNK` — max tokens of cold draft KV ingested
per forward. Default `1024`; `0` restores the single-forward behavior. A
multiple of the block size (16) is recommended.

---------

Signed-off-by: Mihal Dimo <mihal@kakao.com>