Write align-mode GDN state rows in place with a Metal kernel (#634)
## Summary Prefix caching should help when requests share a prefix and add essentially no cost when they do not. On hybrid GDN models, `main` did neither: shared-prefix traffic improved by only 1.03×, while traffic without shared prefixes became 2.24× slower. Qwen3.5+ prefix caching `align` will be the default in the next vllm version (currently it's off by default): https://github.com/vllm-project/vllm/pull/50991 This PR makes it behave like prefix caching: | Model / mode | Shared prefix | No shared prefix | | --- | ---: | ---: | | Qwen3-0.6B, dense reference | **1.32× faster** | 1.02×, effectively free | | Qwen3.5-0.8B, align, `main` | 1.03×, barely faster | **2.24× slower** | | Qwen3.5-0.8B, align, this PR | **1.26× faster** | 0.98×, effectively free | The dense row is the behavioral control: shared prefixes help, while unrelated traffic stays neutral. The PR restores that behavior for the hybrid model. <details> <summary>Benchmark setting</summary> Each cell compares APC enabled with disabled for the same model and workload in the same benchmark session. All runs use 100 prompts, concurrency 16, request rate `inf`, 2304 input tokens, 128 output tokens, `--ignore-eos`, and seed 1234. The shared-prefix workload uses 10 repeated 2048-token prefixes with 256-token suffixes, producing a 61% cache-hit rate; the no-sharing workload uses random prompts. The server runs on an Apple M5 Pro with 64 GB unified memory, macOS 26.2, vllm 0.27.0, mlx 0.32.0, `--max-model-len 4096`, paged attention enabled, and `VLLM_METAL_MEMORY_FRACTION=0.5`. Each cell is one run. Repeat `main` runs for the no-sharing workload were 153.10 s and 156.73 s, so changes within approximately 3% should be treated as noise. </details> ## Roadmap * #584 prototype of align-mode GDN prefix caching * #634 align mode kernel: after this PR, prefix caching becomes actually very useful, if you have enough memory. * future work: fix kv cache double count issue. * then we can mark this feature as support (currently experimental) * future work: extend the support to LFM * not planned: `all` mode prefix caching ## What was wrong & Why we need a `Primitive` metal kernel Under `mamba_cache_mode="align"`, 18 linear layers share 6 physical state pools indexed by scheduler block id. At 1650+ cached blocks, each recurrent pool is about 1.6 GiB. Every state write used MLX assignment. The available MLX 0.32.0 paths behave as follows: | Write path | What `eval_gpu` does | Result for this pool | | --- | --- | --- | | `pool[ids] = rows` | [`Scatter`](https://github.com/ml-explore/mlx/blob/7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247/mlx/backend/metal/indexing.cpp#L247-L256) calls `copy_gpu(pool, out)` before scattering | Full-pool copy when donation fails | | `pool[a:b] = rows` | [`SliceUpdate`](https://github.com/ml-explore/mlx/blob/7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247/mlx/backend/metal/indexing.cpp#L731-L764) calls the same `copy_gpu` path before copying the slice | Same full-pool copy | | `mx.fast.metal_kernel` | [`CustomKernel`](https://github.com/ml-explore/mlx/blob/7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247/mlx/backend/metal/custom_kernel.cpp#L21-L38) creates fresh output storage | Cannot express an aliased pool write | | This PR: MLX `Primitive` | Aliases the pool with `copy_shared_buffer`, then writes selected rows | No full-pool pre-copy, with a graph dependency | For a contiguous pool, `copy_gpu` avoids the copy only by donating the source buffer. [`array::is_donatable()`](https://github.com/ml-explore/mlx/blob/7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247/mlx/array.h#L293-L296) requires both the array descriptor and its data buffer to have `use_count() == 1`. `store_conv_state` and `store_recurrent_state` alias each physical pool across its sibling layers, so donation is unavailable. A row update therefore rewrote up to all 1650+ rows. The cache owned three affected paths: the deferred-state drain, `zero_slots`, and `copy_slots`. Three more writers reached into the pools directly. The hot one was `gdn_lazy.try_conv_decode`, which wrote the conv pool once per linear layer on every pure-decode step. This made a fresh request slower as unrelated requests filled the prefix cache (9.2 ms to 112.5 ms decode ITL). ## What changed - Add `gdn_state_scatter`, an MLX `Primitive` that aliases the pool, writes only selected rows with a 2D/vectorized Metal kernel, and returns the post-write graph handle. - Route all six writers through `write_conv_rows` or `write_recurrent_rows`, which scatter the rows and rebind every sibling layer to that handle. This adds no `mx.eval` barriers. Recurrence, cache placement, and pending-state transitions are unchanged. ## Validation - `tools/hybrid_apc_parity_matrix.py` on Qwen/Qwen3.5-0.8B: 6 runs, 480 comparisons, 0 mismatches. The matrix covers engine defaults and `max_num_batched_tokens=1088` to force aligned chunked splits. - Focused GDN kernel, wrapper, and state-manager tests: 92 passed. ## Build This changes C++ and Metal sources. Run `python -m vllm_metal.metal.build` to refresh prebuilt artifacts, or set `VLLM_METAL_BUILD_FROM_SOURCE=1` for a source build. ## Scope This changes state-write traffic, not align-mode memory usage. The remaining 1.7% no-sharing overhead and the 8.2 ms versus 6.5 ms decode gap are not addressed here. --------- Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>
R
Ranran committed
acce6140320fc90482b9fe80d3f4b9573c171595
Parent: aa6d961
Committed by GitHub <noreply@github.com>
on 8/21/2026, 3:23:54 PM