[Paged KV] Inline metal kernel. deprecate hf pytorch kernel

[WindChimeRan]

Overview

This PR eliminates the external kernels-community/paged-attention HuggingFace dependency by vendoring the Metal shader files directly and replacing the PyTorch MPS dispatch layer with a nanobind C++ extension that dispatches through MLX's own Metal command encoder.

This removes the MLX↔PyTorch MPS fence synchronization bubble.

Key architectural change:

• original: mlx.array → PyTorch MPS bridge → HF kernel → PyTorch MPS bridge → mlx.array
• this PR: mlx.array → native Metal dispatch → mlx.array

fully zero-copy. single Metal command queue.

FAQ:

• Q: Why mlx.core.fast.metal_kernel Can't Be Used Here? This is much more simpler.

• A: mlx.core.fast.metal_kernel always allocates fresh output arrays and provides inputs as read-only. Paged KV
  cache requires in-place, scatter-write mutation of a persistent buffer (via reshape_and_cache) and zero-copy reads from that same buffer (via paged_attention_v1). If we copied the entire cache in and out on every token, it would negate the whole purpose of paging. The C++ nanobind bridge is the minimum complexity needed to get direct Metal buffer access through MLX's internal command encoder.

• Q: What's the difference between kernel and kernel_v1?

• A: kernel is a drop-in replacement for the HuggingFace kernel, originally vendored from an older version of mistral.rs. This PR's end-to-end deterministic tests confirm no regression from the replacement. kernel_v1 is the latest Metal kernel from the mistral.rs repo — more mature, with preliminary scaffolding for variable-length kernels and gpt-oss sink attention support. That said, neither will persist beyond this PR. Both are slated for deprecation once we introduce first-class variable-length kernel support, which is a prerequisite for continuous batching, chunked prefill, and MQA Scorer speculative decoding.

Test

python -m pytest tests/test_paged_deterministic.py -v -s

5 PASSED. The model output is identical to the hf kernel version before.

Benchmark (M1 Pro 32G RAM, macos 26.1)

Usage

Page KV Cache On

VLLM_METAL_USE_PAGED_ATTENTION=1 VLLM_METAL_MEMORY_FRACTION=0.3 vllm serve Qwen/Qwen3-0.6B --max-model-len 2048

vllm bench serve --backend vllm --model Qwen/Qwen3-0.6B \
    --endpoint /v1/completions \
    --dataset-name sharegpt \
    --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
    --num-prompts 100 \
    --request-rate 10 \
    --max-concurrency 32

baseline: default, Page KV Cache Off, mlx_lm

vllm serve Qwen/Qwen3-0.6B --max-model-len 2048

mlx_lm baseline

============ Serving Benchmark Result ============
Successful requests:                     100
Failed requests:                         0
Maximum request concurrency:             32
Request rate configured (RPS):           10.00
Benchmark duration (s):                  345.33
Total input tokens:                      23260
Total generated tokens:                  22061
Request throughput (req/s):              0.29
Output token throughput (tok/s):         63.88
Peak output token throughput (tok/s):    96.00
Peak concurrent requests:                35.00
Total token throughput (tok/s):          131.24
---------------Time to First Token----------------
Mean TTFT (ms):                          69638.28
Median TTFT (ms):                        80145.51
P99 TTFT (ms):                           111926.69
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          156.87
Median TPOT (ms):                        92.12
P99 TPOT (ms):                           1007.68
---------------Inter-token Latency----------------
Mean ITL (ms):                           92.78
Median ITL (ms):                         68.17
P99 ITL (ms):                            298.20
==================================================

hf kernel main

============ Serving Benchmark Result ============
Successful requests:                     100
Failed requests:                         0
Maximum request concurrency:             32
Request rate configured (RPS):           10.00
Benchmark duration (s):                  336.21
Total input tokens:                      23260
Total generated tokens:                  22061
Request throughput (req/s):              0.30
Output token throughput (tok/s):         65.62
Peak output token throughput (tok/s):    128.00
Peak concurrent requests:                35.00
Total token throughput (tok/s):          134.80
---------------Time to First Token----------------
Mean TTFT (ms):                          1861.20
Median TTFT (ms):                        1019.14
P99 TTFT (ms):                           6470.41
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          365.51
Median TPOT (ms):                        349.28
P99 TPOT (ms):                           844.43
---------------Inter-token Latency----------------
Mean ITL (ms):                           340.26
Median ITL (ms):                         312.74
P99 ITL (ms):                            811.23
==================================================

This PR (Inline paged_kernel)

============ Serving Benchmark Result ============
Successful requests:                     100
Failed requests:                         0
Maximum request concurrency:             32
Request rate configured (RPS):           10.00
Benchmark duration (s):                  92.14
Total input tokens:                      23260
Total generated tokens:                  22061
Request throughput (req/s):              1.09
Output token throughput (tok/s):         239.43
Peak output token throughput (tok/s):    358.00
Peak concurrent requests:                35.00
Total token throughput (tok/s):          491.87
---------------Time to First Token----------------
Mean TTFT (ms):                          844.34
Median TTFT (ms):                        434.30
P99 TTFT (ms):                           2657.34
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          120.13
Median TPOT (ms):                        107.88
P99 TPOT (ms):                           522.26
---------------Inter-token Latency----------------
Mean ITL (ms):                           99.98
Median ITL (ms):                         86.99
P99 ITL (ms):                            471.89
==================================================

[WindChimeRan]

@LxYuan0420 request for review.

• I'm not very familiar with the cpp + JIT stuff. I hope it works on your machine out-of-the-box.
• Ideally, in the future, maybe we can release the compiled metal kernel binary: install everything in one line together with the wheel Bump vllm to 0.16.0 and install from wheel #134

[ricky-chaoju]

It works on my machine, but I had to manually install nanobind first. Until pre-compiled binaries are bundled in the wheel, should nanobind be added to the pyproject.toml dependencies?

[LxYuan0420]

@LxYuan0420 request for review.

• I'm not very familiar with the cpp + JIT stuff. I hope it works on your machine out-of-the-box.
• Ideally, in the future, maybe we can release the compiled metal kernel binary: install everything in one line together with the wheel Bump vllm to 0.16.0 and install from wheel #134

Thanks for the context. Totally understandable.

Given this PR introduces a new native C++/JIT execution path, I would like a bit more time to complete review and local validation before final approval.

Could you please share before/after perf numbers (main vs this PR, same machine/config)?

The 5-pass deterministic test is helpful for correctness smoke testing, but I would also like performance evidence for this backend swap.

[WindChimeRan]

@ricky-chaoju Thanks for testing! pyproject.toml updated.

[WindChimeRan]

@LxYuan0420

Benchmarks are now in the PR description: TL;DR ~3.6x throughput improvement over the HF kernel path.

Worth noting that the scope here is fairly contained: it's a backend swap behind the existing VLLM_METAL_USE_PAGED_ATTENTION=1 toggle (off by default), and the deterministic tests confirm identical output.

The kernels themselves are temporary. They'll be replaced once we have first-class variable-length support (might be my next PR).

Happy to address any concerns, but wanted to flag that this is low-risk to land and iterate on.

[LxYuan0420]

@LxYuan0420

Benchmarks are now in the PR description: TL;DR ~3.6x throughput improvement over the HF kernel path.

Worth noting that the scope here is fairly contained: it's a backend swap behind the existing VLLM_METAL_USE_PAGED_ATTENTION=1 toggle (off by default), and the deterministic tests confirm identical output.

The kernels themselves are temporary. They'll be replaced once we have first-class variable-length support (might be my next PR).

Happy to address any concerns, but wanted to flag that this is low-risk to land and iterate on.

This is great.

[LxYuan0420]

Thanks for driving this forward. A few minor changes are needed before we merge this

[WindChimeRan]

@LxYuan0420 Fixed. Will cleanup the kernel files and try var len flashatt in the next PR.

[LxYuan0420]

LGTM.

Thanks for the quick update.