Add paged KV cache with HF Metal kernel for kv cache read/write by-reference decode (#92)

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
```

### Benchmark

Apple M2 Max 36GB, ShareGPT, num-prompts 100, request-rate 10,
max-concurrency 32



<img width="2225" height="766" alt="bench_comparison"
src="https://github.com/user-attachments/assets/c66c4847-f522-4cec-947c-ef5321c36a0b"
/>



* TTFT: better
* output troughput: better
* Mean ITL: worse

output equivelence (paged kv cache vs mlx_lm), both:
* Total input tokens:23260
* Total generated tokens:22061

Memory allocation

* mlx_lm use `auto` to only use just enough memory
* paged kv cache use VLLM_METAL_MEMORY_FRACTION, allocate as much memory
as possible.

Paged KV cache trades higher memory usage for better concurrency, making
it systemwide faster than mlx_lm. Whether it's also faster at the kernel
level is unclear, but advanced features like continuous batching and
chunked prefilling are infeasible to support with mlx_lm alone.

### PR Summary


<details> 
<summary> Patch the mlx models with paged attention kernel. </summary>

- mlx_lm requires contiguous kv cache, and this PR use paged kv cache
(not contiguous).
- Paged kv cache is a prerequisite of future continuous batching and
real chunked prefilling.
- Integrates the
https://huggingface.co/kernels-community/paged-attention Metal shader
for paged KV cache on Apple Silicon (This can be replaced by mlx native
page attention or other better kernels in the future)
- Patches existing mlx_lm model attention layers at runtime with a
wrapper that routes to the external Metal kernel for cache read/write,
while keeping MLX for projections and other layers.
- Prefill: standard MLX causal SDPA, then writes K/V to MPS paged cache
via reshape_and_cache
- Decode: zero-copy attention via paged_attention_v1 — reads K/V
directly from block tables on GPU, eliminating the O(seq_len)
gather/copy per layer per step
  - Falls back to original mlx_lm attention when the env var is not set

</details>

<details> 
<summary> Implement the model runner <--> vllm scheduler contract, so
they are aligned. </summary>

- for chunked prefilling 0:n-1: sample-then-drop the last token
- for chunked prefilling n: sample-and-keep the last token
- for decoding: generate 1 new token

</details>


### Known Limitation & Planned Future PRs:


* **[High Priority]** when setting too small
`VLLM_METAL_MEMORY_FRACTION=0.1`. hit `RuntimeError: Not enough free
blocks: need 21, have 0` . This is because: All kv blocks have been
consumed, while the prefilling/decoding have not been finished, then
deadlock. Need to implement vllm-metal's paged kv cache preemption to
align with vllm scheduler contract.
* **[High Priority]** torch_to_mlx in the tensor bridge may not be true
zero copy.
* [Bug] not working with HuggingFaceTB/SmolVLM-Instruct,
#114 might be the fix
* [Feature] re-enable prefix caching under paged kv cache. Prior
version: #80
* [Medium Priority] real chunked prefilling. This PR's implementation is
wasteful. Expected: chunk n prefill read the 0:n-1 kv cache, and only
prefill n. Actual: prefill all 0:n kv cache each time. The time
complexity is quadratic in terms of the number of chunks. Why? just to
satisfy vllm scheduler.
* [Medium Priority] real continuous batching to align with upstream
vllm. This requires var len prefilling & decoding operating on
`[total_num_token, *]` instead of the current `[batch, seq, *]`.
* [Refactor] #97
* [Refactor] five separate forward paths (_prefill_single,
_prefill_single_request_paged .etc) that share the same pre/post
processing. Maybe we can merged duplicate codes, but that's for the sake
of aesthetics.
* [Doc] Readme architecture figure is no longer accurate. 

* ~~[Testing] Need to test on macos 14/15, metal 3.2. It is expected to
work.~~ It works.

### FAQ



<details>
<summary>Why hack the paged KV cache as a global variable?</summary>

The model's `__call__` signature is `(input_ids, cache=...)` — and
`mlx_lm`'s call requires contiguous tensors with no additional
parameters. There's no way to pass `slot_mapping`, `block_tables`, or
any other per-forward metadata down to the attention layers. This design
is inspired by [nano-vllm](https://github.com/GeeeekExplorer/nano-vllm).

</details>

<details>
<summary>Why use this attention kernel?</summary>

This kernel supports variable-length prefilling and decoding, so
attention can be computed over `[total_tokens, *]` instead of `[batch,
seq, *]`. This is essential for supporting real continuous batching in
the future.

</details>

<details>
<summary>Each call to the attention kernel triggers an MLX ↔ Torch round
trip?</summary>

Yes, but it's okay as long as the MLX-to-Torch round trip is implemented
in zero-copy mode. Besides, this kernel can be replaced by better ones
in the future if they become available. The most important thing is to
get the whole system working end-to-end (chunked prefill, continuous
batching) first; then we can swap in better modules later.

</details>

<details>
<summary>What would a future paged attention kernel look like?</summary>

It would need to support variable sequence lengths, `slot_map`, etc. —
similar to `flash_attn_varlen_func` and `flash_attn_with_kvcache` from
FlashAttention. The difficulties are:

1. HuggingFace kernel libraries only expose PyTorch bindings, which
require type conversion from our MLX tensors.
2. As far as I understand, a proper FlashAttention-style implementation
would need to be written directly in Metal, not in MLX.

</details>



### Acknowledgement: 
Early prototype #71

---------

Signed-off-by: ran <hzz5361@psu.edu>

Author: WindChimeRan

README.md

@@ -9,7 +9,7 @@ vLLM Metal is a plugin that enables vLLM to run on Apple Silicon Macs using MLX
 - **MLX-accelerated inference**: faster than PyTorch MPS on Apple Silicon
 - **Unified memory**: True zero-copy operations leveraging Apple Silicon's unified memory architecture
 - **vLLM compatibility**: Full integration with vLLM's engine, scheduler, and OpenAI-compatible API
-- **Paged attention**: Efficient KV cache management for long sequences
+- **Paged attention** *(experimental)*: Efficient KV cache management for long sequences — opt-in via `VLLM_METAL_USE_PAGED_ATTENTION=1` (requires `pip install 'vllm-metal[paged]'`); default path uses MLX-managed KV cache
 - **GQA support**: Grouped-Query Attention for efficient inference
 
 ## Requirements
@@ -78,6 +78,7 @@ Environment variables for customization:
 | `VLLM_METAL_USE_MLX` | `1` | Use MLX for compute (1=yes, 0=no) |
 | `VLLM_MLX_DEVICE` | `gpu` | MLX device (`gpu` or `cpu`) |
 | `VLLM_METAL_BLOCK_SIZE` | `16` | KV cache block size |
+| `VLLM_METAL_USE_PAGED_ATTENTION` | `0` | Enable experimental paged KV cache (requires `pip install 'vllm-metal[paged]'`) |
 | `VLLM_METAL_DEBUG` | `0` | Enable debug logging |
 | `VLLM_USE_MODELSCOPE` | `False` | Set True to change model registry to <https://www.modelscope.cn/> |
 | `VLLM_METAL_MODELSCOPE_CACHE` | None | Specify the absolute path of the local model |

pyproject.toml

@@ -41,6 +41,10 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+paged = [
+    # Paged attention Metal kernel (opt-in, experimental)
+    "kernels>=0.4.5; platform_system == 'Darwin' and platform_machine == 'arm64'",
+]
 vllm = ["vllm>=0.14.0"]
 stt = [
     # Speech-to-text audio processing (Whisper models)
@@ -54,7 +58,7 @@ dev = [
     "mypy>=1.19.1",
 ]
 all = [
-    "vllm-metal[vllm,stt,dev]",
+    "vllm-metal[vllm,paged,stt,dev]",
 ]
 
 [project.urls]

scripts/lib.sh

@@ -49,7 +49,7 @@ ensure_venv() {
 # Install dev dependencies
 install_dev_deps() {
   section "Installing dependencies"
-  uv pip install -e ".[dev]"
+  uv pip install -e ".[dev,paged]"
 }
 
 # Full development environment setup

src/lib.rs

@@ -21,7 +21,7 @@ pub struct BlockAllocator {
     /// Free blocks stored in a deque for O(1) operations
     free_blocks: VecDeque<usize>,
     /// Mapping from sequence ID to allocated blocks
-    sequence_blocks: HashMap<i64, Vec<usize>>,
+    sequence_blocks: HashMap<String, Vec<usize>>,
     /// Total number of blocks
     num_blocks: usize,
 }
@@ -53,7 +53,7 @@ impl BlockAllocator {
     ///
     /// # Raises
     /// RuntimeError if not enough free blocks
-    pub fn allocate_blocks(&mut self, seq_id: i64, num_blocks: usize) -> PyResult<Vec<usize>> {
+    pub fn allocate_blocks(&mut self, seq_id: String, num_blocks: usize) -> PyResult<Vec<usize>> {
         if self.free_blocks.len() < num_blocks {
             return Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
                 "Not enough free blocks: need {}, have {}",
@@ -83,7 +83,7 @@ impl BlockAllocator {
     ///
     /// # Arguments
     /// * `seq_id` - Sequence identifier
-    pub fn free_sequence(&mut self, seq_id: i64) {
+    pub fn free_sequence(&mut self, seq_id: String) {
         if let Some(blocks) = self.sequence_blocks.remove(&seq_id) {
             // Return blocks to the free pool
             for block_idx in blocks {
@@ -99,7 +99,7 @@ impl BlockAllocator {
     ///
     /// # Returns
     /// List of block indices for the sequence
-    pub fn get_sequence_blocks(&self, seq_id: i64) -> Vec<usize> {
+    pub fn get_sequence_blocks(&self, seq_id: String) -> Vec<usize> {
         self.sequence_blocks
             .get(&seq_id)
             .cloned()
@@ -119,7 +119,7 @@ impl BlockAllocator {
     }
 
     /// Check if sequence has blocks allocated.
-    pub fn has_sequence(&self, seq_id: i64) -> bool {
+    pub fn has_sequence(&self, seq_id: String) -> bool {
         self.sequence_blocks.contains_key(&seq_id)
     }
 
@@ -130,7 +130,7 @@ impl BlockAllocator {
     }
 
     /// Get all sequence blocks as a dictionary.
-    pub fn get_all_sequence_blocks(&self) -> HashMap<i64, Vec<usize>> {
+    pub fn get_all_sequence_blocks(&self) -> HashMap<String, Vec<usize>> {
         self.sequence_blocks.clone()
     }
 }
@@ -280,10 +280,10 @@ pub fn compute_kv_block_indices(
 /// Batch compute block indices for multiple sequences.
 #[pyfunction]
 pub fn batch_compute_kv_indices(
-    sequence_blocks: HashMap<i64, Vec<usize>>,
-    seq_lens: HashMap<i64, usize>,
+    sequence_blocks: HashMap<String, Vec<usize>>,
+    seq_lens: HashMap<String, usize>,
     block_size: usize,
-) -> HashMap<i64, Vec<(usize, usize, usize)>> {
+) -> HashMap<String, Vec<(usize, usize, usize)>> {
     let mut result = HashMap::with_capacity(sequence_blocks.len());
 
     for (seq_id, blocks) in sequence_blocks {

tests/test_kernel_loader.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for kernel_loader: OS-aware revision pinning for Metal compatibility.
+
+Verifies that:
+- macOS 16+ uses the latest HF kernel (default revision)
+- macOS 15 and earlier pins to the Nov 2025 compat revision (Metal 3.2)
+- Both revisions actually load and expose the expected ops
+
+Run with:
+    python -m pytest tests/test_kernel_loader.py -v -s
+"""
+
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+pytest.importorskip("kernels")
+
+# ---------------------------------------------------------------------------
+# Unit tests (no network, no GPU)
+# ---------------------------------------------------------------------------
+
+
+class TestNeedsCompatRevision:
+    """Test _needs_compat_revision() with mocked macOS versions."""
+
+    @pytest.mark.parametrize(
+        "ver, expected",
+        [
+            ("15.7.4", True),  # macOS 15 — needs compat
+            ("14.5", True),  # macOS 14 — needs compat
+            ("26.3", False),  # macOS 26 — modern
+            ("", False),  # empty — safe default
+        ],
+    )
+    def test_version_check(self, ver, expected):
+        from vllm_metal.metal_kernel_backend.kernel_loader import _needs_compat_revision
+
+        with mock.patch("platform.mac_ver", return_value=(ver, ("", "", ""), "")):
+            assert _needs_compat_revision() is expected
+
+
+class TestGetKernelRevisionSelection:
+    """Test that get_paged_attention_ops passes the right revision to get_kernel."""
+
+    def _reset_kernel_cache(self):
+        import vllm_metal.metal_kernel_backend.kernel_loader as kl
+
+        kl._kernel = None
+
+    def test_macos_15_uses_compat_revision(self):
+        self._reset_kernel_cache()
+        with (
+            mock.patch("platform.mac_ver", return_value=("15.7.4", ("", "", ""), "")),
+            mock.patch("kernels.get_kernel", return_value=mock.MagicMock()) as mk,
+        ):
+            from vllm_metal.metal_kernel_backend.kernel_loader import (
+                _MACOS15_COMPAT_REVISION,
+                get_paged_attention_ops,
+            )
+
+            get_paged_attention_ops()
+            mk.assert_called_once_with(
+                "kernels-community/paged-attention",
+                revision=_MACOS15_COMPAT_REVISION,
+            )
+        self._reset_kernel_cache()
+
+    def test_macos_26_uses_latest(self):
+        self._reset_kernel_cache()
+        with (
+            mock.patch("platform.mac_ver", return_value=("26.3", ("", "", ""), "")),
+            mock.patch("kernels.get_kernel", return_value=mock.MagicMock()) as mk,
+        ):
+            from vllm_metal.metal_kernel_backend.kernel_loader import (
+                get_paged_attention_ops,
+            )
+
+            get_paged_attention_ops()
+            mk.assert_called_once_with(
+                "kernels-community/paged-attention",
+                revision=None,
+            )
+        self._reset_kernel_cache()
+
+
+# ---------------------------------------------------------------------------
+# Integration tests (require network + MPS)
+# ---------------------------------------------------------------------------
+
+
+def _mps_available() -> bool:
+    try:
+        import torch
+
+        return torch.backends.mps.is_available()
+    except Exception:
+        return False
+
+
+@pytest.mark.skipif(not _mps_available(), reason="MPS not available")
+class TestKernelLoadsForReal:
+    """Actually load the kernel from HuggingFace and verify ops exist."""
+
+    _EXPECTED_OPS = {"reshape_and_cache", "paged_attention_v1"}
+
+    def test_latest_revision_loads(self):
+        from kernels import get_kernel
+
+        kernel = get_kernel("kernels-community/paged-attention")
+        ops = set(dir(kernel))
+        assert self._EXPECTED_OPS <= ops, f"Missing ops: {self._EXPECTED_OPS - ops}"
+
+    def test_compat_revision_loads(self):
+        from kernels import get_kernel
+
+        from vllm_metal.metal_kernel_backend.kernel_loader import (
+            _MACOS15_COMPAT_REVISION,
+        )
+
+        kernel = get_kernel(
+            "kernels-community/paged-attention",
+            revision=_MACOS15_COMPAT_REVISION,
+        )
+        ops = set(dir(kernel))
+        assert self._EXPECTED_OPS <= ops, f"Missing ops: {self._EXPECTED_OPS - ops}"