Defer the GDN recurrent decode state scatter (#620)
## Summary
Keep each GDN recurrent decode update in its compact buffer and feed it
directly into the next decode step when the request slot order is
unchanged.
Previously, every decode layer copied that update back into the stable
state pool. The pool is still updated before any operation that needs
it: slot-order changes, prefill or fallback execution, materialization,
release, and align-mode step planning.
```
Old:
kernel output
→ scatter into stable pool
→ clear consumed pending input
New:
kernel output
→ discard the consumed pending input
→ register the output as the new authoritative pending state
```
## Why
For Qwen3.5-0.8B, one recurrent-state slab is 1 MiB per request per GDN
layer (`16 × 128 × 128 × fp32`), across 18 GDN layers. Avoiding the
extra per-token scatter reduces state traffic without changing the
recurrent computation.
## Performance
Serving benchmark on M5 Pro, `Qwen/Qwen3.5-0.8B`, paged path,
`mamba_cache_mode=none`. Sonnet 1024-token input / 128-token output, 100
prompts, request rate 10, `--max-num-seqs 32`. Two independent passes
per arm; values below are the mean of both.
Concurrency 32:
| metric | main | this PR | change |
|---|---:|---:|---:|
| Output token throughput (tok/s) | 398.9 | 431.7 | +8.2% |
| Total token throughput (tok/s) | 3635 | 3908 | +7.5% |
| Mean TPOT (ms) | 66.6 | 59.6 | -10.4% |
| P99 TPOT (ms) | 83.0 | 74.0 | -10.8% |
| Benchmark duration (s) | 31.3 | 29.2 | -6.9% |
Concurrency 8:
| metric | main | this PR | change |
|---|---:|---:|---:|
| Output token throughput (tok/s) | 266.6 | 286.9 | +7.6% |
| Total token throughput (tok/s) | 2428 | 2610 | +7.5% |
| Mean TPOT (ms) | 25.9 | 23.6 | -8.9% |
| P99 TPOT (ms) | 30.3 | 29.3 | -3.4% |
| Benchmark duration (s) | 46.9 | 43.7 | -6.9% |
TTFT is unaffected. Prefill cannot reach this path, which requires
`total_tokens == num_requests`. Median TTFT differs by 0.9% at
concurrency 32; at concurrency 8 the two `main` passes alone span
374-504 ms, bracketing this PR's 489 ms.
The end-to-end gain is smaller than the decode-only delta because this
workload is prefill-heavy at 1024 in / 128 out. Isolating decode on the
same machine gives roughly 1.4x at batch 32.
Align mode still publishes pending state at the start of every step, so
it does not receive the full steady-state benefit.
<details>
<summary>Raw output, concurrency 32</summary>
**main**
```
============ Serving Benchmark Result ============
Successful requests: 100
Maximum request concurrency: 32
Benchmark duration (s): 31.36
Total input tokens: 101402
Total generated tokens: 12478
Request throughput (req/s): 3.19
Output token throughput (tok/s): 397.85
Total token throughput (tok/s): 3630.95
---------------Time to First Token----------------
Mean TTFT (ms): 610.44
Median TTFT (ms): 623.06
P99 TTFT (ms): 1076.93
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 66.67
Median TPOT (ms): 69.86
P99 TPOT (ms): 90.93
---------------Inter-token Latency----------------
Mean ITL (ms): 66.32
Median ITL (ms): 43.87
P99 ITL (ms): 280.05
==================================================
```
**this PR**
```
============ Serving Benchmark Result ============
Successful requests: 100
Maximum request concurrency: 32
Benchmark duration (s): 29.17
Total input tokens: 101402
Total generated tokens: 12656
Request throughput (req/s): 3.43
Output token throughput (tok/s): 433.89
Total token throughput (tok/s): 3910.30
---------------Time to First Token----------------
Mean TTFT (ms): 642.88
Median TTFT (ms): 622.21
P99 TTFT (ms): 1047.37
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 59.18
Median TPOT (ms): 63.29
P99 TPOT (ms): 78.32
---------------Inter-token Latency----------------
Mean ITL (ms): 59.09
Median ITL (ms): 35.39
P99 ITL (ms): 285.98
==================================================
```
</details>
<details>
<summary>Benchmark configuration</summary>
- Hardware: Apple M5 Pro, 64 GB
- Model: `Qwen/Qwen3.5-0.8B` (18 GDN layers + 6 SDPA layers)
- Server: `VLLM_METAL_USE_PAGED_ATTENTION=1`,
`VLLM_METAL_MEMORY_FRACTION=0.5`, `--max-model-len 2048 --max-num-seqs
32`
- Client: `vllm bench serve --dataset-name sonnet --sonnet-input-len
1024 --sonnet-output-len 128 --num-prompts 100 --request-rate 10
--max-concurrency {8,32} --seed 0`
- Arms: `main` at `a8b7e75`, this PR at `8a7caaa`
</details>
## Validation
- `tests/test_gdn_lazy_kernels.py`: 37 passed, including C++ parity and
pending-state transitions.
- Paged output matched MLX inline-cache output with lazy GDN kernels
enabled and disabled.
- Uneven batched generation matched batch-1 goldens at `max_num_seqs=8`
and `32`.
- Hybrid prefix-caching parity passed both 40-case configurations with
zero mismatches and confirmed an aligned restore.
`VLLM_METAL_GDN_LAZY_KERNELS=0` continues to bypass this path.
---------
Signed-off-by: Ranran <ranranhaoranzhang@gmail.com> R
Ranran committed
150dd3292bd940f2eac1b3442ece21355d8ebf19
Parent: a8b7e75
Committed by GitHub <noreply@github.com>
on 8/17/2026, 8:13:23 AM