SIGN IN SIGN UP

Remove per-layer GDN prefill materialization barriers (#632)

## Summary

<img width="2856" height="1661" alt="bench_pr632"
src="https://github.com/user-attachments/assets/ed58eb42-5367-405d-ade0-e97f7dc2b53e"
/>

Multi-request prefill on separate-projection GDN models synchronously
evaluates convolution and recurrent state after every GDN layer. This
change removes those per-layer evaluations and lets the existing
end-of-forward submission evaluate the complete graph.

The GDN recurrence, cache placement, pending-state transitions, and
fallback paths are unchanged. Combined-QKVZ models already use the fully
lazy path.

## Benchmark

Qwen3.5-0.8B on M5 Pro, Sonnet dataset, 100 prompts, concurrency 16,
prefix caching off. Each result is the mean of six runs per branch
across three alternating fresh-server sessions; the first run of each
session was discarded.

| Input / output | Duration: main | Duration: PR | Δ duration | Mean
TTFT: main | Mean TTFT: PR | Δ TTFT |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| 1024 / 128 | 30.15 ± 2.01 s | 27.85 ± 1.05 s | **−7.6%** | 807.0 ±
54.6 ms | 702.5 ± 34.8 ms | **−13.0%** |
| 1024 / 16 | 14.30 ± 0.58 s | 12.44 ± 0.16 s | **−13.0%** | 789.0 ±
37.3 ms | 685.3 ± 9.1 ms | **−13.1%** |

The prefill-heavy case improves more, consistent with removing 36
`mx.eval` barriers from each prefill-containing step: two per GDN layer
across 18 layers. All runs processed the same 101,402 input tokens and
fixed output-token counts.

<details>
<summary>Benchmark configuration</summary>

- Base: `bd32be8`
- Hardware: Apple M5 Pro, 64 GB unified memory, macOS 26.6 (25G72)
- Versions: vllm 0.27.0, mlx 0.32.0, mlx-lm 0.31.3
- Server: `VLLM_METAL_USE_PAGED_ATTENTION=1
VLLM_METAL_MEMORY_FRACTION=0.5 vllm serve Qwen/Qwen3.5-0.8B
--max-model-len 2048 --no-enable-prefix-caching`
- Client: `vllm bench serve --dataset-name sonnet --sonnet-input-len
1024 --sonnet-output-len {128,16} --num-prompts 100 --request-rate inf
--max-concurrency 16 --ignore-eos --temperature 0`
- Scheduler: chunked prefill on, `max_num_batched_tokens=2048`

</details>

## Memory trade-off

Peak prefill graph growth increases from 202 MiB to 432 MiB. This was
measured with MLX 0.32.0, 24 layers, four 512-token requests, fp32
state, and the real recurrent prefill kernel. End-of-forward evaluation
still bounds the graph.

## Validation

A deterministic Qwen3.5-0.8B A/B probe produced identical token IDs for
all five prompts across all four main/PR pairings; same-branch controls
were also identical. Instrumentation confirmed that all 18 GDN layers
took the affected multi-request prefill path (`prefill_multi=18`,
`prefill_single=0`, `max_nreq=5`). An independent GMRID_v3 evaluation
(20 stratified examples across 8 classes, two rounds per branch)
produced byte-identical main/PR completion text for every example at
both 0-shot and 5-shot.

| Revision | 0-shot weighted F1 / EM | 5-shot weighted F1 / EM |
| --- | ---: | ---: |
| main | 0.3900 / 0.4500 | 0.4071 / 0.4500 |
| PR | 0.3900 / 0.4500 | 0.4071 / 0.4500 |

Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>
R
Ranran committed
67100ba77780dec48adeb569724efaf8fe928b19
Parent: 4d48fbc
Committed by GitHub <noreply@github.com> on 8/20/2026, 12:49:32 PM