[MLX] Support radix cache

python/sglang/bench_one_batch.py

@@ -524,20 +524,38 @@ class _MlxBenchRunner:
     def __init__(self, model_runner, server_args):
         from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
 
-        self.mlx_runner = MlxModelRunner(
+        # Radix cache requires the scheduler's allocator/trie; disable in
+        # standalone bench mode where no scheduler is present.
+        init_kwargs = dict(
             model_path=server_args.model_path,
             trust_remote_code=server_args.trust_remote_code,
+            disable_radix_cache=True,
+            mem_fraction_static=server_args.mem_fraction_static,
         )
+        if server_args.max_total_tokens is not None:
+            init_kwargs["pool_size"] = server_args.max_total_tokens
+        self.mlx_runner = MlxModelRunner(**init_kwargs)
+        self.mlx_runner.init_kv_pool(req_to_token_pool=None)
         self.fake_torch_runner = model_runner
 
     def clear(self):
         self.mlx_runner.clear()
 
     def extend(self, reqs):
         req_ids = [str(req.rid) for req in reqs]
-        token_ids_list = [[int(t) for t in req.fill_ids] for req in reqs]
-        next_token_ids = self.mlx_runner.prefill_batch(req_ids, token_ids_list)
-        return torch.tensor(next_token_ids), None, req_ids
+        results = []
+        for rid, req in zip(req_ids, reqs):
+            token_ids = [int(t) for t in req.fill_ids]
+            next_token = self.mlx_runner.prefill(
+                req_id=rid,
+                new_token_ids=token_ids,
+                full_token_ids=token_ids,
+                prefix_slot_ids=[],
+                new_slot_ids=[],
+                req_pool_idx=0,
+            )
+            results.append(next_token)
+        return torch.tensor(results), None, req_ids
 
     def decode(self, next_token_ids, req_ids):
         next_token_ids = self.mlx_runner.decode_batch(req_ids)

python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py

@@ -0,0 +1,35 @@
+"""KV cache components for the MLX backend."""
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
+    BatchedDecodeContext,
+    MLXAttentionWrapper,
+    clear_context,
+    get_context,
+    set_context,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import (
+    ContiguousKVCache,
+    OffsetCache,
+    PoolBackedCache,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool
+from sglang.srt.hardware_backend.mlx.kv_cache.model_patching import (
+    find_attention_layers,
+    get_num_layers,
+    patch_model_attention,
+)
+
+__all__ = [
+    "BatchedDecodeContext",
+    "clear_context",
+    "ContiguousKVCache",
+    "find_attention_layers",
+    "get_context",
+    "get_num_layers",
+    "MLXAttentionWrapper",
+    "MlxKVPool",
+    "OffsetCache",
+    "patch_model_attention",
+    "PoolBackedCache",
+    "set_context",
+]

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py

@@ -0,0 +1,134 @@
+"""Batched decode attention wrapper for MLX backend."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import ContiguousKVCache
+
+_thread_local = threading.local()
+
+
+@dataclass
+class BatchedDecodeContext:
+    """Context set before batched decode, read by attention wrappers."""
+
+    batch_size: int
+    seq_lens: list[int]  # per-request token count before the new token
+    # layer_caches[layer_idx][req_idx] = ContiguousKVCache
+    layer_caches: list[list[ContiguousKVCache]]
+
+
+def set_context(ctx: Optional[BatchedDecodeContext]) -> None:
+    _thread_local.batched_ctx = ctx
+
+
+def get_context() -> Optional[BatchedDecodeContext]:
+    return getattr(_thread_local, "batched_ctx", None)
+
+
+def clear_context() -> None:
+    _thread_local.batched_ctx = None
+
+
+class MLXAttentionWrapper(nn.Module):
+    """Wraps an mlx-lm Attention for batched decode (BS>1).
+
+    When ``BatchedDecodeContext`` is set, performs per-request RoPE,
+    cache writes, and batched SDPA.  Otherwise delegates to inner module.
+    """
+
+    def __init__(self, inner: nn.Module, layer_idx: int):
+        super().__init__()
+        object.__setattr__(self, "_inner", inner)
+        object.__setattr__(self, "_layer_idx", layer_idx)
+
+    def __call__(self, x: mx.array, mask: Any = None, cache: Any = None) -> mx.array:
+        ctx = get_context()
+        if ctx is None:
+            return self._inner(x, mask=mask, cache=cache)
+        return self._batched_decode(x, ctx)
+
+    def _batched_decode(self, x: mx.array, ctx: BatchedDecodeContext) -> mx.array:
+        inner = self._inner
+        layer_idx = self._layer_idx
+        B = ctx.batch_size
+
+        queries = inner.q_proj(x)
+        keys = inner.k_proj(x)
+        values = inner.v_proj(x)
+
+        head_dim = queries.shape[-1] // inner.n_heads
+        queries = queries.reshape(B, 1, inner.n_heads, head_dim)
+        keys = keys.reshape(B, 1, inner.n_kv_heads, head_dim)
+        values = values.reshape(B, 1, inner.n_kv_heads, head_dim)
+
+        if hasattr(inner, "q_norm"):
+            queries = inner.q_norm(queries)
+        if hasattr(inner, "k_norm"):
+            keys = inner.k_norm(keys)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.transpose(0, 2, 1, 3)
+
+        # Vectorized RoPE with per-batch offsets
+        offsets = mx.array(ctx.seq_lens, dtype=mx.int32)
+        queries = inner.rope(queries, offset=offsets)
+        keys = inner.rope(keys, offset=offsets)
+
+        layer_caches = ctx.layer_caches[layer_idx]
+        max_len = max(ctx.seq_lens) + 1
+
+        # TODO: replace per-request loop with native batched/ragged
+        # attention once mx.fast.scaled_dot_product_attention supports
+        # variable-length sequences.
+        all_k = []
+        all_v = []
+
+        for i in range(B):
+            layer_caches[i].write_token(keys[i : i + 1], values[i : i + 1])
+
+            k_all, v_all = layer_caches[i].get_kv()
+            curr_len = layer_caches[i].offset
+
+            if curr_len < max_len:
+                pad = max_len - curr_len
+                k_pad = mx.zeros(
+                    (1, inner.n_kv_heads, pad, head_dim), dtype=k_all.dtype
+                )
+                v_pad = mx.zeros(
+                    (1, inner.n_kv_heads, pad, head_dim), dtype=v_all.dtype
+                )
+                k_all = mx.concatenate([k_all, k_pad], axis=2)
+                v_all = mx.concatenate([v_all, v_pad], axis=2)
+
+            all_k.append(k_all)
+            all_v.append(v_all)
+
+        keys_b = mx.concatenate(all_k, axis=0)
+        values_b = mx.concatenate(all_v, axis=0)
+
+        attn_mask = None
+        seq_lens_plus1 = [s + 1 for s in ctx.seq_lens]
+        if min(seq_lens_plus1) < max_len:
+            positions = mx.arange(max_len)
+            valid_lens = mx.array(seq_lens_plus1, dtype=mx.int32)
+            mask_bool = positions[None, :] >= valid_lens[:, None]
+            attn_mask = mx.where(
+                mask_bool[:, None, None, :],
+                mx.array(mx.finfo(queries.dtype).min, dtype=queries.dtype),
+                mx.array(0.0, dtype=queries.dtype),
+            )
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys_b, values_b, scale=inner.scale, mask=attn_mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1)
+        return inner.o_proj(output)