Caching less to serve more requests

Qwen3.5 4B with the public DFlash drafter at 15 draft tokens per step has most of them accepted when the output repeats the prompt. On that setup vLLM 0.31 keeps 16 copies of each request’s recurrent state, which take 960 of the 1,415 MiB an 8k-token request uses on a B200. We keep one state and a short log of inputs instead (ReplaySSM), 72 MiB, so the same memory holds 274 requests of 8k tokens instead of 105, and 5.1× as many on Qwen3.5 35B-A3B. At high load this gives two separate results:

• Capacity. At 256 concurrent requests vLLM 0.31 runs 117 and queues the rest, with a median time to first token of 22.5 s. Ours runs 252 and starts streaming in 0.4 s.
• Throughput. From 64 requests, where both fit, ours produces 1.3–1.6× the output tokens per second, because it does not write 16 state copies per step.

Where the memory goes

Qwen3.5 4B has 8 attention layers, with a KV cache that grows with the request, and 24 Gated DeltaNet layers, a form of linear attention that keeps a fixed-size recurrent state per request instead. Speculative decoding has to undo rejected tokens. A KV cache does that by ignoring the rejected entries, but a recurrent state is overwritten by every token. vLLM 0.31 therefore writes the state after every verified position and continues from the accepted one. At 15 draft tokens that is 16 copies per request, in the same memory pool as the KV cache.

A reader can compute this cost for their own model before running anything:

• state copies per request = (draft tokens + 1) × recurrent layers × page size.
• The page size is not the state size. vLLM gives every cache group one page size, set by the largest, and pads each state copy to it. Read it from the KV-cache group sizes in vLLM’s startup log.
• Here (15 + 1) × 24 × 2.5 MiB = 960 MiB, with 2 MiB of fp32 state plus a convolution window per layer, padded to a 2.5 MiB page, against a KV cache of 455 MiB for an 8k-token request, target and drafter together.

Each draft token adds 60 MiB per request. Past about 7 draft tokens the copies outweigh an 8k-token request’s KV cache, and for longer requests their share falls because the KV cache grows and the copies do not.

What we changed

ReplaySSM comes from Dao AI Lab (post), which proposed it for vLLM with NVIDIA (vLLM #49887). Our version started as a port of that pull request. It keeps one checkpoint state per request and a buffer of the per-token inputs since that checkpoint, and computes each layer’s output from the two. Rolling back rejected draft tokens removes their buffer entries, and the full state is rewritten (the “flush”) only when the buffer fills. Neither vLLM 0.31 nor SGLang 0.5.21 serves it for Gated DeltaNet with DFlash drafts.

Upstream measured drafts of up to 7 tokens. At 15 draft tokens the buffer fills every step, so the flush runs every step and its kernel sets the speed. We made two changes:

• We stopped reserving the speculative state copies. Each layer keeps one state plus a 32-slot input buffer, 72 MiB per request across 24 layers.
• The flush applies the step’s accepted tokens as one block of fp32 tensor-core matrix operations, in one launch per layer, instead of one token after another.

Timed in isolation at 16 requests of 16 rows, the state update takes 32.6 µs per layer against 138.3 µs for vLLM 0.31, 0.8 ms against 3.3 ms per step over 24 layers. Our flush works in blocks of 16 rows, and at 17 rows it jumps to 124.8 µs. A deployment fixes its draft length, so the kernel can be sized to it. A general engine has to serve every length.

What it gave

We tested high load with prompts of 3.3k–4.5k tokens and 3,072 output tokens per request (8.7–9.0 of 15 draft tokens accepted, since long copies drift more) and two server lifetimes per arm. Check the prefill budget before reading capacity. In a first test with --max-num-batched-tokens 8192, ours queued 72 of 256 requests with the cache only 44% full, so the prefill budget was holding them back while memory sat free. We raised the budget to 16,384 for both arms:

The two results have different causes:

• Capacity shows as waiting. From 128 requests vLLM 0.31’s cache is 99% full. It holds 117 requests, evicts and recomputes 50–63 per two-minute window and queues the rest, so the median request takes 42.0 s at 256 against our 26.7 s. Ours runs up to 311 at once with the cache at most 78% full and never evicts.
• Throughput comes from the shorter step. At 64 requests both arms fit and ours already gives 1.47× the output tokens per second. Both arms stay at their throughput plateau from 64 requests on. By arithmetic on the state sizes, not a profile, vLLM 0.31 writes about 875 MiB of state per request per step, about 55 GiB per step at 64 requests, and ours writes one copy.

At low load, with 512 output tokens, speed per request was 1.16–1.24× at 1 to 16 requests, and outputs match vLLM 0.31 byte for byte on 99.1–99.4% of requests, inside the 99.1–99.8% that two vLLM lifetimes give. With 3,072 tokens they match on 97.0–97.7%, about 1.5 points below that band, because the fp32 state update runs in a different order. Scored against the source passage on the same prompts, both versions miss 0.246% of words on Qwen3.5 4B, and on Qwen3.5 35B-A3B ours misses 0.025% against 0.028% for vLLM 0.31, so the differences cost no copy fidelity. This check covers copying only, not open-ended text.

Other models

The same build, unchanged, on two more public Qwen3.5 models with their public DFlash drafters, one B200 each:

The 35B model gains 5.1× because of the page size in the formula above. vLLM gives every part of the cache one page size, and the 35B’s drafter needs 4 KiB per token, so each 2.28 MiB state copy is padded to a 5 MiB page. Its 16 copies across 30 layers take 2.4 GiB, 85% of each request’s memory. At 64 concurrent requests of about 7k tokens, vLLM 0.31 already runs only 33 of them on this model.

Where it helps

• Drafts are long and mostly accepted, so there are many copies and each step advances many tokens.
• Requests are short to medium, so the copies are a large share of each request’s memory.
• Enough requests run at once to fill the cache.

We chose 15 draft tokens because the public drafter proposes blocks of 16 and our flush works in 16-row blocks. A longer draft needs the 32-row version, at 74–76 µs per layer against 32.6. Our build runs without prefix caching, which vLLM requires before merging ReplaySSM for Gated DeltaNet, and it runs on a vLLM nightly from before 0.31, which alone is 3–4% faster than 0.31.