SIGN IN SIGN UP

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