Fix paged-attention KV cache dtype + size accounting (issue #119) (#125)

This PR is:
- To align the Metal paged-attention KV cache dtype with the model's
dtype (fixes batched decode parity for #119).
- To compute KV cache byte sizes via `torch.dtype.itemsize` instead of
allocating temporary tensors.

Notes:
- `tests/test_metal_kernel_paged.py::test_batched_decode_matches` now
passes.
- `tests/test_metal_kernel_paged.py::test_greedy_output_matches` remains
xfailed (tracked in #119). This is a remaining single-request greedy
parity mismatch between the paged-kernel path and the standard path;
fixing it likely requires deeper kernel/offset semantics work, so I'm
keeping it out of this PR to keep scope tight.

Quick manual smoke test:

Terminal 1:
```bash
vllm serve Qwen/Qwen3-0.6B --host 127.0.0.1 --port 8000 --max-model-len 2048
```

Terminal 2 (single request):
```bash
curl -fsS http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"Qwen/Qwen3-0.6B","messages":[{"role":"user","content":"Write a 2-sentence apple story."}],"max_tokens":512,"temperature":0.8}' \
| jq -r '.choices[0].message.content'
```

Terminal 2 (concurrent 4 requests):
```bash
for i in 1 2 3 4; do
  (
    echo "===== req $i ====="
    curl -fsS http://127.0.0.1:8000/v1/chat/completions \
      -H 'Content-Type: application/json' \
      -d "{\"model\":\"Qwen/Qwen3-0.6B\",\"messages\":[{\"role\":\"user\",\"content\":\"Write a 2-sentence apple story (${i}).\"}],\"max_tokens\":256,\"temperature\":0.8}" \
    | jq -r '.choices[0].message.content'
    echo
  ) &
done
wait
```

Related: #119

---------

Signed-off-by: Yuan Lik Xun <lxyuan0420@gmail.com>

Author: LxYuan0420

tests/test_metal_kernel_paged.py

@@ -21,6 +21,8 @@
     import torch
     from mlx_lm import load as mlx_lm_load
     from mlx_lm.models.cache import make_prompt_cache
+
+    from vllm_metal.kv_cache_dtype import infer_kv_cache_dtype_from_model
 except ImportError as exc:
     pytest.skip(
         f"Metal kernel paged attention tests require mlx/torch/mlx_lm: {exc}",
@@ -68,6 +70,21 @@ def _paged_attention_ops_available() -> None:
 # ---------------------------------------------------------------------------
 
 
+def _test_infer_paged_kv_dtype(model) -> torch.dtype:
+    """Test-only helper: choose a float dtype for MPSPagedKVCache.
+
+    This is deliberately local to this test module. Production code uses
+    `vllm_metal.kv_cache_dtype.infer_kv_cache_dtype_from_model()`.
+    """
+    result = infer_kv_cache_dtype_from_model(model)
+    if result.warning is not None:
+        raise AssertionError(
+            "KV cache dtype inference unexpectedly fell back during tests: "
+            f"{result.warning}"
+        )
+    return result.dtype
+
+
 def _greedy_generate_standard(model, token_ids: list[int], max_new: int) -> list[int]:
     """Generate tokens using the standard mlx_lm KVCache path."""
     cache = make_prompt_cache(model)
@@ -109,7 +126,7 @@ def _greedy_generate_metal_kernel(
         head_dim=head_dim,
         num_blocks=num_blocks,
         block_size=BLOCK_SIZE,
-        dtype=torch.float16,
+        dtype=_test_infer_paged_kv_dtype(model),
     )
 
     n_patched = patch_model_attention_metal_kernel(model, mps_cache, BLOCK_SIZE)
@@ -190,9 +207,6 @@ def test_greedy_output_matches(self, qwen3_model):
         )
 
     @pytest.mark.slow
-    @pytest.mark.xfail(
-        reason="Metal paged-attention parity mismatch vs standard path (see #119)"
-    )
     def test_batched_decode_matches(self, qwen3_model):
         """Batched Metal kernel paged decode must match per-request sequential."""
         model, tokenizer = qwen3_model
@@ -225,7 +239,7 @@ def test_batched_decode_matches(self, qwen3_model):
             head_dim=head_dim,
             num_blocks=num_blocks,
             block_size=BLOCK_SIZE,
-            dtype=torch.float16,
+            dtype=_test_infer_paged_kv_dtype(model),
         )
         patch_model_attention_metal_kernel(model, mps_cache, BLOCK_SIZE)
 
@@ -300,7 +314,7 @@ def test_patch_replaces_self_attn(self, qwen3_model):
             head_dim=args.head_dim,
             num_blocks=32,
             block_size=BLOCK_SIZE,
-            dtype=torch.float16,
+            dtype=_test_infer_paged_kv_dtype(model),
         )
         patch_model_attention_metal_kernel(model, mps_cache, BLOCK_SIZE)
 
@@ -323,7 +337,7 @@ def test_fallback_when_no_context(self, qwen3_model):
             head_dim=args.head_dim,
             num_blocks=32,
             block_size=BLOCK_SIZE,
-            dtype=torch.float16,
+            dtype=_test_infer_paged_kv_dtype(model),
         )
         patch_model_attention_metal_kernel(model, mps_cache, BLOCK_SIZE)
 

tests/test_prefix_cache.py

@@ -302,6 +302,7 @@ def test_rotating_kvcache_merge_handles_prefill_exceeding_max_size(self) -> None
 
         assert isinstance(merged[0], mr.BatchRotatingKVCache)
         assert isinstance(extracted_req0, mr.RotatingKVCache)
+        assert isinstance(extracted_req1, mr.RotatingKVCache)
         assert extracted_req0.offset == cache_req0.offset
         assert extracted_req1.offset == cache_req1.offset
 

vllm_metal/kv_cache_dtype.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+"""KV cache dtype inference and policy.
+
+The Metal paged-attention backend stores *activation* K/V tensors in an
+MPS-backed cache. Those tensors must be floating point. Some models may have
+quantized *weights* (e.g. int8), so we must not derive the KV cache dtype from
+weights without enforcing a float-only policy.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from vllm_metal.paged_attention_common import find_layers_and_attr
+from vllm_metal.pytorch_backend.tensor_bridge import MLX_TO_TORCH_DTYPE
+
+DEFAULT_KV_CACHE_DTYPE = torch.float16
+ALLOWED_KV_CACHE_DTYPES: frozenset[torch.dtype] = frozenset(
+    {
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    }
+)
+
+
+@dataclass(frozen=True)
+class KvCacheDtypeInference:
+    """Result of inferring the KV cache dtype from a model."""
+
+    dtype: torch.dtype
+    warning: str | None = None
+
+
+def infer_kv_cache_dtype_from_model(
+    model: Any, *, default: torch.dtype = DEFAULT_KV_CACHE_DTYPE
+) -> KvCacheDtypeInference:
+    """Infer a float KV-cache dtype from an MLX(-LM/-VLM) model.
+
+    Policy:
+    - If we can map the model's attention weight dtype to torch and it's a
+      supported float dtype, use it.
+    - Otherwise, fall back to *default* and provide a warning string the caller
+      may log.
+    """
+    try:
+        layers, attn_attr = find_layers_and_attr(model)
+        if not layers:
+            raise ValueError("model has no transformer layers")
+
+        attn = getattr(layers[0], attn_attr)
+        # If the model is already patched, unwrap to the real attention module.
+        attn = getattr(attn, "_inner", attn)
+
+        mlx_dtype = attn.q_proj.weight.dtype
+    except (AttributeError, IndexError, TypeError, ValueError) as exc:
+        return KvCacheDtypeInference(
+            dtype=default,
+            warning=f"Cannot infer KV cache dtype from model ({exc}); using {default}",
+        )
+
+    torch_dtype = MLX_TO_TORCH_DTYPE.get(mlx_dtype)
+    if torch_dtype is None:
+        return KvCacheDtypeInference(
+            dtype=default,
+            warning=f"Unsupported MLX dtype for KV cache ({mlx_dtype!r}); using {default}",
+        )
+
+    if torch_dtype not in ALLOWED_KV_CACHE_DTYPES:
+        return KvCacheDtypeInference(
+            dtype=default,
+            warning=(
+                f"Model weight dtype {mlx_dtype!r} maps to non-float torch dtype "
+                f"{torch_dtype}; using {default} for KV cache instead"
+            ),
+        )
+
+    return KvCacheDtypeInference(dtype=torch_dtype)

vllm_metal/metal_kernel_backend/paged_attention.py

@@ -72,7 +72,7 @@
 from vllm_metal.metal_kernel_backend.kernel_loader import get_paged_attention_ops
 from vllm_metal.paged_attention_common import (
     PagedAttentionContext,
-    _find_layers_and_attr,
+    find_layers_and_attr,
     get_context,
 )
 from vllm_metal.pytorch_backend.tensor_bridge import mlx_to_torch, torch_to_mlx
@@ -327,7 +327,7 @@ def patch_model_attention_metal_kernel(
 
     Returns the number of patched layers.
     """
-    layer_list, attn_attr = _find_layers_and_attr(model)
+    layer_list, attn_attr = find_layers_and_attr(model)
     patched = 0
 
     for layer_idx, layer in enumerate(layer_list):

vllm_metal/paged_attention_common.py

@@ -103,7 +103,7 @@ def make_mask(
 # ---------------------------------------------------------------------------
 
 
-def _find_layers_and_attr(model: Any) -> tuple[list[Any], str]:
+def find_layers_and_attr(model: Any) -> tuple[list[Any], str]:
     """Find transformer layers and the attention attribute name.
 
     Returns (layer_list, attn_attr_name) where each layer has

vllm_metal/v1/model_runner.py

@@ -47,6 +47,7 @@
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_metal.config import get_config
+from vllm_metal.kv_cache_dtype import infer_kv_cache_dtype_from_model
 from vllm_metal.paged_attention_common import (
     OffsetCache,
     clear_context,
@@ -617,6 +618,7 @@ def __init__(
         self._paged_kv_cache: Any = None  # MPSPagedKVCache, set by worker
         self._paged_block_size: int = 0
         self._paged_request_seq_lens: dict[str, int] = {}  # req_id → seq_len
+        self.kv_cache_dtype: torch.dtype | None = None
 
     def _is_vlm_model(self) -> bool:
         """Check if the model is a vision-language model (VLM).
@@ -650,6 +652,7 @@ def load_model(self) -> None:
                 )
                 self._extract_model_args()
                 self._resolve_model_dims()
+                self._initialize_kv_cache_dtype()
                 return
 
         # Load model using appropriate backend
@@ -673,9 +676,20 @@ def load_model(self) -> None:
 
         self._extract_model_args()
         self._resolve_model_dims()
+        self._initialize_kv_cache_dtype()
         load_time = time.time() - start_time
         logger.info(f"Model loaded in {load_time:.2f}s: {model_name}")
 
+    def _initialize_kv_cache_dtype(self) -> None:
+        """Infer and store the KV cache dtype for this runner."""
+        if self.model is None:
+            raise RuntimeError("Model not loaded")
+
+        paged_kv_dtype = infer_kv_cache_dtype_from_model(self.model)
+        if paged_kv_dtype.warning:
+            logger.warning("%s", paged_kv_dtype.warning)
+        self.kv_cache_dtype = paged_kv_dtype.dtype
+
     def _extract_model_args(self) -> None:
         """Extract model configuration from loaded model.
 
@@ -782,6 +796,8 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
             Dictionary mapping attention layer names to KV cache specs
         """
         block_size = self.metal_config.block_size
+        if self.kv_cache_dtype is None:
+            raise RuntimeError("KV cache dtype not initialized; load_model() first")
 
         # Create a spec for each layer
         specs: dict[str, KVCacheSpec] = {}
@@ -791,7 +807,7 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
                 block_size=block_size,
                 num_kv_heads=self.num_kv_heads,
                 head_size=self.head_dim,
-                dtype=torch.float16,
+                dtype=self.kv_cache_dtype,
             )
 
         return specs
@@ -817,7 +833,9 @@ def get_cache_block_size_bytes(self) -> int:
 
         # Each block stores key and value for all layers
         # Block memory = 2 * num_layers * block_size * num_kv_heads * head_dim * dtype_size
-        dtype_size = 2  # float16
+        if self.kv_cache_dtype is None:
+            raise RuntimeError("KV cache dtype not initialized; load_model() first")
+        dtype_size = self.kv_cache_dtype.itemsize
         return (
             2
             * self.num_layers

vllm_metal/v1/worker.py

@@ -146,7 +146,6 @@ def _setup_paged_attention(self) -> None:
         max_model_len.
         """
         import psutil
-        import torch
 
         from vllm_metal.metal_kernel_backend.cache import MPSPagedKVCache
         from vllm_metal.metal_kernel_backend.paged_attention import (
@@ -242,13 +241,15 @@ def _setup_paged_attention(self) -> None:
         )
 
         # --- Create cache and patch model ---
+        if runner.kv_cache_dtype is None:
+            raise RuntimeError("KV cache dtype not initialized; runner.load_model()")
         mps_kv_cache = MPSPagedKVCache(
             num_layers=runner.num_layers,
             num_kv_heads=runner.num_kv_heads,
             head_dim=runner.head_dim,
             num_blocks=num_blocks,
             block_size=block_size,
-            dtype=torch.float16,
+            dtype=runner.kv_cache_dtype,
         )
 
         n_patched = patch_model_attention_metal_kernel(
@@ -295,15 +296,18 @@ def _get_model_memory_usage(self) -> int:
         return 0
 
     def _one_sequence_kv_bytes(self) -> int:
-        """Bytes for one max-length sequence of KV cache (K + V, float16)."""
+        """Bytes for one max-length sequence of KV cache (K + V)."""
         runner = self.model_runner
+        dtype_size = (
+            runner.kv_cache_dtype.itemsize if runner.kv_cache_dtype is not None else 2
+        )
         return (
             2  # K and V
             * runner.num_layers
             * self.model_config.max_model_len
             * runner.num_kv_heads
             * runner.head_dim
-            * 2  # float16
+            * dtype_size
         )
 
     def determine_available_memory(self) -> int: