[MLX] Support radix cache (#21509)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>

python/sglang/bench_one_batch.py

@@ -524,10 +524,18 @@ 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):
@@ -535,9 +543,19 @@ class _MlxBenchRunner:
 
     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)

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

@@ -0,0 +1,205 @@
+"""ContiguousKVCache, PoolBackedCache and OffsetCache for MLX backend."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import mlx.core as mx
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool
+
+
+class OffsetCache:
+    """Data-free shim satisfying mlx-lm's cache protocol.
+
+    Provides ``make_mask`` and ``state`` without storing actual K/V.
+    """
+
+    def __init__(self, offset: int = 0):
+        self.offset = offset
+
+    @property
+    def state(self):
+        return ()  # Empty — safe for mx.eval unpacking
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def update_and_fetch(self, keys, values):
+        raise RuntimeError("OffsetCache should not store data")
+
+
+_DEFAULT_MAX_SEQ_LEN = 4096
+
+
+class ContiguousKVCache:
+    """Pre-allocated KV buffer for one request × one layer.
+
+    Shape ``(1, n_kv_heads, max_seq_len, head_dim)``.  Slice assignment
+    instead of ``mx.concatenate``.  Lazy-allocated on first write.
+    """
+
+    __slots__ = ("keys", "values", "offset", "max_seq_len")
+
+    def __init__(
+        self,
+        n_kv_heads: int | None = None,
+        head_dim: int | None = None,
+        max_seq_len: int = _DEFAULT_MAX_SEQ_LEN,
+        dtype: mx.Dtype | None = None,
+    ):
+        if n_kv_heads is not None and head_dim is not None and dtype is not None:
+            self.keys = mx.zeros((1, n_kv_heads, max_seq_len, head_dim), dtype=dtype)
+            self.values = mx.zeros((1, n_kv_heads, max_seq_len, head_dim), dtype=dtype)
+        else:
+            self.keys = None
+            self.values = None
+        self.offset = 0
+        self.max_seq_len = max_seq_len
+
+    def _allocate(self, keys: mx.array) -> None:
+        """Allocate buffers matching the first key tensor's shape."""
+        B, n_kv_heads, _, head_dim = keys.shape
+        self.keys = mx.zeros(
+            (B, n_kv_heads, self.max_seq_len, head_dim), dtype=keys.dtype
+        )
+        self.values = mx.zeros(
+            (B, n_kv_heads, self.max_seq_len, head_dim), dtype=keys.dtype
+        )
+
+    @property
+    def state(self):
+        """Arrays for ``mx.eval`` unpacking."""
+        if self.keys is None:
+            return ()
+        return (self.keys, self.values)
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def _grow(self, required: int) -> None:
+        """Double the buffer until it can hold *required* tokens."""
+        new_max = self.max_seq_len
+        while new_max < required:
+            new_max *= 2
+        B, n_kv_heads, _, head_dim = self.keys.shape
+        new_k = mx.zeros((B, n_kv_heads, new_max, head_dim), dtype=self.keys.dtype)
+        new_v = mx.zeros((B, n_kv_heads, new_max, head_dim), dtype=self.values.dtype)
+        if self.offset > 0:
+            new_k[:, :, : self.offset, :] = self.keys[:, :, : self.offset, :]
+            new_v[:, :, : self.offset, :] = self.values[:, :, : self.offset, :]
+        self.keys = new_k
+        self.values = new_v
+        self.max_seq_len = new_max
+
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]:
+        """Append K/V and return all valid K/V up to current offset."""
+        if self.keys is None:
+            self._allocate(keys)
+        S = keys.shape[2]
+        end = self.offset + S
+        if end > self.max_seq_len:
+            self._grow(end)
+        self.keys[:, :, self.offset : end, :] = keys
+        self.values[:, :, self.offset : end, :] = values
+        self.offset = end
+        return self.keys[:, :, :end, :], self.values[:, :, :end, :]
+
+    def write_token(self, k: mx.array, v: mx.array) -> None:
+        """Write one token. k, v shape: (1, n_kv_heads, 1, head_dim)."""
+        self.keys[:, :, self.offset : self.offset + 1, :] = k
+        self.values[:, :, self.offset : self.offset + 1, :] = v
+        self.offset += 1
+
+    def get_kv(self) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""
+        return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :]
+
+
+class PoolBackedCache:
+    """Lazily gathers cached KV from the shared pool during forward pass.
+
+    Each ``update_and_fetch`` gathers this layer's prefix from the pool
+    on demand, keeping operations in the lazy compute graph.  Convert to
+    ``ContiguousKVCache`` via ``to_contiguous`` after the forward pass.
+    """
+
+    __slots__ = (
+        "_pool",
+        "_layer_idx",
+        "_slots",
+        "offset",
+        "_full_keys",
+        "_full_values",
+        "_new_keys",
+        "_new_values",
+    )
+
+    def __init__(
+        self,
+        pool: MlxKVPool,
+        layer_idx: int,
+        slots: mx.array,
+        prefix_len: int,
+    ):
+        self._pool = pool
+        self._layer_idx = layer_idx
+        self._slots = slots
+        self.offset = prefix_len
+        self._full_keys: mx.array | None = None
+        self._full_values: mx.array | None = None
+        self._new_keys: mx.array | None = None
+        self._new_values: mx.array | None = None
+
+    @property
+    def keys(self) -> mx.array | None:
+        return self._full_keys
+
+    @property
+    def values(self) -> mx.array | None:
+        return self._full_values
+
+    @property
+    def state(self):
+        if self._full_keys is not None:
+            return (self._full_keys, self._full_values)
+        return ()
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]:
+        """Gather cached prefix from pool, concatenate with new K/V."""
+        S = keys.shape[2]
+
+        if self.offset > 0:
+            k_cached, v_cached = self._pool.get_kv(
+                self._layer_idx, self._slots[: self.offset]
+            )
+            # Pool layout (S, n_kv_heads, head_dim) → cache (1, n_kv_heads, S, head_dim)
+            k_cached = k_cached.transpose(1, 0, 2)[None]
+            v_cached = v_cached.transpose(1, 0, 2)[None]
+            k_all = mx.concatenate([k_cached, keys], axis=2)
+            v_all = mx.concatenate([v_cached, values], axis=2)
+        else:
+            k_all = keys
+            v_all = values
+
+        self.offset += S
+        self._full_keys = k_all
+        self._full_values = v_all
+        self._new_keys = keys
+        self._new_values = values
+        return k_all, v_all
+
+    def to_contiguous(self, max_seq_len: int = 4096) -> ContiguousKVCache:
+        """Convert to ContiguousKVCache reusing forward-pass arrays."""
+        cache = ContiguousKVCache(max_seq_len=max_seq_len)
+        if self._full_keys is not None:
+            cache.update_and_fetch(self._full_keys, self._full_values)
+        return cache

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

@@ -0,0 +1,80 @@
+"""Flat KV pool with per-layer buffers of shape (pool_size, n_kv_heads, head_dim).
+
+Slot 0 is reserved as padding (1-based indexing).
+"""
+
+import logging
+
+import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+
+class MlxKVPool:
+    """Pre-allocated KV pool indexed by integer slot IDs."""
+
+    def __init__(
+        self,
+        pool_size: int,
+        num_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        dtype: mx.Dtype = mx.float16,
+    ):
+        self.pool_size = pool_size
+        self.num_layers = num_layers
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.dtype = dtype
+
+        # Per-layer buffers: (pool_size, n_kv_heads, head_dim)
+        self.k_buffer: list[mx.array] = [
+            mx.zeros((pool_size, n_kv_heads, head_dim), dtype=dtype)
+            for _ in range(num_layers)
+        ]
+        self.v_buffer: list[mx.array] = [
+            mx.zeros((pool_size, n_kv_heads, head_dim), dtype=dtype)
+            for _ in range(num_layers)
+        ]
+
+        mem_mb = (pool_size * n_kv_heads * head_dim * 2 * num_layers * dtype.size) / (
+            1024 * 1024
+        )
+        logger.info(
+            f"MlxKVPool: {pool_size} slots × {num_layers} layers "
+            f"× {n_kv_heads} heads × {head_dim} dim, "
+            f"dtype={dtype}, ~{mem_mb:.1f} MB"
+        )
+
+    def set_kv(self, layer_id: int, slots: mx.array, k: mx.array, v: mx.array) -> None:
+        """Scatter K/V into *slots* for one layer."""
+        self.k_buffer[layer_id][slots] = k
+        self.v_buffer[layer_id][slots] = v
+
+    def get_kv(self, layer_id: int, slots: mx.array) -> tuple[mx.array, mx.array]:
+        """Gather K/V from *slots* for one layer."""
+        return self.k_buffer[layer_id][slots], self.v_buffer[layer_id][slots]
+
+    def get_kv_all_layers(self, slots: mx.array) -> tuple[mx.array, mx.array]:
+        """Gather K/V from *slots* across all layers."""
+        k_all = mx.stack([self.k_buffer[i][slots] for i in range(self.num_layers)])
+        v_all = mx.stack([self.v_buffer[i][slots] for i in range(self.num_layers)])
+        return k_all, v_all
+
+    def set_kv_all_layers(
+        self, slots: mx.array, k_all: mx.array, v_all: mx.array
+    ) -> None:
+        """Scatter K/V into *slots* across all layers."""
+        for i in range(self.num_layers):
+            self.set_kv(i, slots, k_all[i], v_all[i])
+
+    def all_buffers(self) -> list[mx.array]:
+        """Return all buffer arrays (for ``mx.eval``)."""
+        return self.k_buffer + self.v_buffer
+
+    def clear(self) -> None:
+        """Zero all buffers."""
+        shape = (self.pool_size, self.n_kv_heads, self.head_dim)
+        for i in range(self.num_layers):
+            self.k_buffer[i] = mx.zeros(shape, dtype=self.dtype)
+            self.v_buffer[i] = mx.zeros(shape, dtype=self.dtype)

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

@@ -0,0 +1,46 @@
+"""Model introspection and attention patching."""
+
+from typing import Any
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
+    MLXAttentionWrapper,
+)
+
+
+def find_attention_layers(model: Any) -> tuple[list[Any], str]:
+    """Find transformer layers and the attention attribute name."""
+    root = getattr(model, "language_model", model)
+    container = getattr(root, "model", root)
+    layer_list = getattr(container, "layers", None) or getattr(root, "layers", [])
+
+    if layer_list:
+        sample = layer_list[0]
+        if hasattr(sample, "self_attn"):
+            return layer_list, "self_attn"
+        if hasattr(sample, "attention"):
+            return layer_list, "attention"
+        raise ValueError(f"No attention attribute in layer type {type(sample)}")
+    return layer_list, "self_attn"
+
+
+def patch_model_attention(model: Any) -> int:
+    """Install MLXAttentionWrapper on all attention layers (idempotent).
+
+    The wrapper delegates to the inner module when no BatchedDecodeContext
+    is set, so it is always installed and never removed.
+    """
+    layer_list, attn_attr = find_attention_layers(model)
+    patched = 0
+    for idx, layer in enumerate(layer_list):
+        attn = getattr(layer, attn_attr)
+        if isinstance(attn, MLXAttentionWrapper):
+            continue
+        setattr(layer, attn_attr, MLXAttentionWrapper(attn, idx))
+        patched += 1
+    return patched
+
+
+def get_num_layers(model: Any) -> int:
+    """Return the number of transformer layers."""
+    layer_list, _ = find_attention_layers(model)
+    return len(layer_list)

python/sglang/srt/hardware_backend/mlx/model_runner.py

@@ -1,96 +1,79 @@
-"""End-to-end MLX model runner for Apple Silicon.
+"""MLX model runner for Apple Silicon.
 
-Runs the entire model within MLX, bypassing PyTorch MPS entirely.
+Slot allocation and radix-trie prefix matching are handled by the
+scheduler (``TokenToKVPoolAllocator`` / ``RadixCache``).  This runner
+reads cached KV from ``MlxKVPool``, runs the forward pass, and writes
+new KV back.  Each request also keeps a ``ContiguousKVCache`` for
+decode-time attention.
 """
 
 import logging
 import time
-from dataclasses import dataclass
 
 import mlx.core as mx
+import psutil
 from mlx_lm import load as mlx_lm_load
-from mlx_lm.models.cache import (
-    BatchKVCache,
-    BatchRotatingKVCache,
-    KVCache,
-    RotatingKVCache,
-    make_prompt_cache,
+
+from sglang.srt.hardware_backend.mlx.kv_cache import (
+    BatchedDecodeContext,
+    ContiguousKVCache,
+    MLXAttentionWrapper,
+    OffsetCache,
+    PoolBackedCache,
+    clear_context,
+    find_attention_layers,
+    get_num_layers,
+    patch_model_attention,
+    set_context,
 )
+from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class MlxRequestState:
-    """Per-request state for MLX inference."""
-
-    token_ids: list[int]
-    cache: list  # List of KVCache per layer
-    generated_tokens: int = 0
-
-
-def _merge_kv_caches(
-    caches_list: list[list],
-) -> list:
-    """Merge multiple per-request caches into batched caches."""
-    if not caches_list:
-        return []
-
-    num_layers = len(caches_list[0])
-    merged = []
-
-    for layer_idx in range(num_layers):
-        layer_caches = [caches[layer_idx] for caches in caches_list]
-        if isinstance(layer_caches[0], KVCache):
-            batch_cache = BatchKVCache.merge(layer_caches)
-        elif isinstance(layer_caches[0], RotatingKVCache):
-            batch_cache = BatchRotatingKVCache.merge(layer_caches)
-        else:
-            raise TypeError(f"Unsupported cache type: {type(layer_caches[0]).__name__}")
-        merged.append(batch_cache)
-
-    return merged
-
-
-def _extract_kv_cache(batch_caches: list, idx: int) -> list:
-    """Extract a single request's cache from batched caches.
-
-    Works with both BatchKVCache (has .extract) and plain KVCache
-    populated with batched data of shape (B, H, L, D).
-    """
-    extracted = []
-    for cache in batch_caches:
-        if hasattr(cache, "extract"):
-            extracted.append(cache.extract(idx))
-        else:
-            # Plain KVCache with batched data — slice along batch dim
-            new_cache = KVCache()
-            new_cache.keys = mx.contiguous(cache.keys[idx : idx + 1])
-            new_cache.values = mx.contiguous(cache.values[idx : idx + 1])
-            new_cache.offset = cache.offset
-            extracted.append(new_cache)
-    return extracted
-
-
 class MlxModelRunner:
-    """Model runner that executes the entire model in MLX.
-
-    This avoids the MPS<->MLX tensor bridge overhead by keeping all
-    computation within MLX.
-    """
+    """MLX model runner with radix-cache prefix sharing."""
 
     def __init__(
         self,
         model_path: str,
         trust_remote_code: bool = False,
+        disable_radix_cache: bool = False,
+        pool_size: int | None = None,
+        mem_fraction_static: float = 0.8,
     ):
         self.model_path = model_path
         self.trust_remote_code = trust_remote_code
         self.model = None
-        self._request_states: dict[str, MlxRequestState] = {}
+        self.disable_radix_cache = disable_radix_cache
+        self._mem_fraction_static = mem_fraction_static
 
         self._load_model()
 
+        # Pin MLX allocations to prevent OS paging
+        device_info = mx.device_info()
+        max_wired = int(device_info.get("max_recommended_working_set_size", 0))
+        if max_wired > 0:
+            mx.set_wired_limit(max_wired)
+            logger.info(f"Wired memory limit set to {max_wired / (1024**3):.1f} GB")
+
+        patch_model_attention(self.model)
+
+        self._num_layers = get_num_layers(self.model)
+        self._max_seq_len = 4096  # doubles on overflow
+
+        self._req_caches: dict[str, list[ContiguousKVCache | PoolBackedCache]] = {}
+        self._req_token_ids: dict[str, list[int]] = {}
+        self._cache_pool: list[list[ContiguousKVCache]] = []  # reusable caches
+
+        self._kv_pool: MlxKVPool | None = None
+        self._req_to_token_pool: ReqToTokenPool | None = None
+        self._req_pool_idx: dict[str, int] = {}
+        self._req_synced_offset: dict[str, int] = {}
+
+        self._pool_size = self._compute_pool_size(pool_size)
+
     @staticmethod
     def _extract_logits(model_output):
         """Extract logits from model output, handling both tuple and direct returns."""
@@ -98,6 +81,29 @@ class MlxModelRunner:
             return model_output[0]
         return model_output
 
+    def _acquire_cache(self) -> list[ContiguousKVCache]:
+        """Get a reusable cache list from the pool, or create a new one."""
+        if self._cache_pool:
+            cache = self._cache_pool.pop()
+            for c in cache:
+                c.offset = 0
+            return cache
+        return [
+            ContiguousKVCache(max_seq_len=self._max_seq_len)
+            for _ in range(self._num_layers)
+        ]
+
+    def _release_cache(self, cache: list[ContiguousKVCache]) -> None:
+        """Return a cache list to the pool for reuse."""
+        self._cache_pool.append(cache)
+
+    @staticmethod
+    def _eval_with_cache(
+        token_result: mx.array, cache: list[ContiguousKVCache | PoolBackedCache]
+    ) -> None:
+        """Evaluate token result and all cache buffers in one mx.eval call."""
+        mx.eval(token_result, *[s for c in cache for s in c.state])
+
     def _load_model(self):
         """Load model using mlx_lm."""
         logger.info(f"Loading MLX model: {self.model_path}")
@@ -107,200 +113,336 @@ class MlxModelRunner:
             self.model_path,
             tokenizer_config={"trust_remote_code": self.trust_remote_code},
         )
+        # Force-evaluate weights so mx.get_active_memory() reflects
+        # actual usage before KV pool sizing.
+        mx.eval(self.model.parameters())
 
         load_time = time.time() - start_time
         logger.info(f"MLX model loaded in {load_time:.2f}s")
 
+    def _get_attn_config(self) -> tuple[int, int, mx.Dtype]:
+        """Return (n_kv_heads, head_dim, dtype) from the model."""
+        layer_list, attn_attr = find_attention_layers(self.model)
+        if not layer_list:
+            raise RuntimeError("Cannot determine attention config: no layers found")
+        sample_attn = getattr(layer_list[0], attn_attr)
+        if isinstance(sample_attn, MLXAttentionWrapper):
+            sample_attn = sample_attn._inner
+        n_kv_heads = sample_attn.n_kv_heads
+        if hasattr(sample_attn, "head_dim"):
+            head_dim = sample_attn.head_dim
+        elif hasattr(sample_attn, "k_proj") and hasattr(sample_attn.k_proj, "weight"):
+            head_dim = sample_attn.k_proj.weight.shape[0] // n_kv_heads
+        else:
+            raise RuntimeError("Cannot determine head_dim from attention module")
+        dtype = mx.float16
+        if hasattr(sample_attn, "k_proj") and hasattr(sample_attn.k_proj, "weight"):
+            dtype = sample_attn.k_proj.weight.dtype
+        return n_kv_heads, head_dim, dtype
+
+    def _compute_pool_size(self, explicit_size: int | None) -> int:
+        """Determine pool slot count (auto-size from available memory if needed)."""
+        if explicit_size is not None:
+            return explicit_size
+        n_kv_heads, head_dim, dtype = self._get_attn_config()
+        num_layers = self._num_layers
+        sys_available = psutil.virtual_memory().available
+        mlx_limit = mx.device_info().get(
+            "max_recommended_working_set_size",
+            mx.device_info().get("memory_size", 0),
+        )
+        mlx_used = mx.get_active_memory()
+        mlx_usable = int(mlx_limit * self._mem_fraction_static)
+        kv_budget = min(
+            max(mlx_usable - mlx_used, 0),
+            int(sys_available * self._mem_fraction_static),
+        )
+        bytes_per_slot = 2 * num_layers * n_kv_heads * head_dim * dtype.size
+        pool_size = max(kv_budget // bytes_per_slot, 256)
+        logger.info(
+            f"Auto-sized KV pool: "
+            f"sys_available={sys_available / (1024**3):.2f} GB, "
+            f"mlx_limit={mlx_limit / (1024**3):.1f} GB, "
+            f"mlx_used={mlx_used / (1024**3):.2f} GB, "
+            f"kv_budget={kv_budget / (1024**3):.2f} GB, "
+            f"bytes_per_slot={bytes_per_slot}, pool_size={pool_size}"
+        )
+        return pool_size
+
+    @property
+    def pool_size(self) -> int:
+        return self._pool_size
+
+    def init_kv_pool(self, req_to_token_pool: ReqToTokenPool) -> None:
+        """Create MlxKVPool (+1 for padding slot 0) and wire scheduler pools."""
+        self._req_to_token_pool = req_to_token_pool
+        if self.disable_radix_cache:
+            return
+        n_kv_heads, head_dim, dtype = self._get_attn_config()
+        # +1 for padding slot 0
+        self._kv_pool = MlxKVPool(
+            pool_size=self._pool_size + 1,
+            num_layers=self._num_layers,
+            n_kv_heads=n_kv_heads,
+            head_dim=head_dim,
+            dtype=dtype,
+        )
+        logger.info(
+            f"KV pool initialized: pool_size={self._pool_size} "
+            f"(buffer size {self._pool_size + 1} incl. padding slot 0), "
+            f"{self._num_layers} layers, {n_kv_heads} kv_heads, {head_dim} head_dim"
+        )
+
     def prefill(
         self,
         req_id: str,
-        token_ids: list[int],
+        new_token_ids: list[int],
+        full_token_ids: list[int],
+        prefix_slot_ids: list[int],
+        new_slot_ids: list[int],
+        req_pool_idx: int,
     ) -> int:
-        """Run prefill for a single request.
+        """Prefill a request.  Returns next_token_id."""
+        num_layers = self._num_layers
+        prefix_len = len(prefix_slot_ids)
 
-        If a request with the same req_id already has state (e.g. from a
-        previous partial prefill), the existing KV cache is reused and only
-        the new tokens are fed through the model.
+        if self.disable_radix_cache:
+            cache = self._acquire_cache()
+            input_ids = mx.array([new_token_ids], dtype=mx.int32)
+            model_output = self.model(input_ids, cache=cache)
+            logits = self._extract_logits(model_output)
+            next_token_mlx = mx.argmax(logits[:, -1, :], axis=-1)
+            self._eval_with_cache(next_token_mlx, cache)
+            next_token = int(next_token_mlx.item())
 
-        Args:
-            req_id: Request identifier
-            token_ids: Input token IDs (full sequence, including any
-                previously prefilled tokens)
+            self._req_token_ids[req_id] = list(full_token_ids) + [next_token]
+            self._req_caches[req_id] = cache
+            self._req_pool_idx[req_id] = req_pool_idx
+            self._req_synced_offset[req_id] = 0
+            return next_token
 
-        Returns:
-            Next token ID (greedy sampled)
-        """
-        existing_state = self._request_states.get(req_id)
-        if existing_state is not None:
-            # Continuation: reuse existing cache, feed only new tokens
-            cached_input_len = (
-                len(existing_state.token_ids) - existing_state.generated_tokens
-            )
-            new_tokens = token_ids[cached_input_len:]
-            cache = existing_state.cache
+        assert self._kv_pool is not None
+
+        new_token_count = len(new_token_ids)
+
+        if prefix_len > 0:
+            slot_ids_mx = mx.array(prefix_slot_ids, dtype=mx.int32)
+            cache = [
+                PoolBackedCache(self._kv_pool, i, slot_ids_mx, prefix_len)
+                for i in range(num_layers)
+            ]
         else:
-            new_tokens = token_ids
-            cache = make_prompt_cache(self.model)
+            cache = self._acquire_cache()
 
-        input_ids = mx.array([new_tokens], dtype=mx.int32)
+        if new_token_count > 0:
+            extend_tokens = new_token_ids
+        else:
+            # Full cache hit — rerun last token to get next-token logits
+            extend_tokens = full_token_ids[-1:]
+            for c in cache:
+                c.offset = max(c.offset - 1, 0)
+
+        input_ids = mx.array([extend_tokens], dtype=mx.int32)
         model_output = self.model(input_ids, cache=cache)
-
         logits = self._extract_logits(model_output)
 
         last_logits = logits[:, -1, :]
         next_token_mlx = mx.argmax(last_logits, axis=-1)
 
-        # Evaluate everything together
-        mx.eval(next_token_mlx, *[c.state for c in cache])
+        # Convert PoolBackedCache → ContiguousKVCache for decode
+        if prefix_len > 0:
+            contiguous_cache = self._acquire_cache()
+            for layer_idx in range(num_layers):
+                pbc = cache[layer_idx]
+                contiguous_cache[layer_idx].update_and_fetch(
+                    pbc._full_keys, pbc._full_values
+                )
+            cache = contiguous_cache
+
+        self._eval_with_cache(next_token_mlx, cache)
         next_token = int(next_token_mlx.item())
 
-        # Store state for future decoding
-        self._request_states[req_id] = MlxRequestState(
-            token_ids=list(token_ids) + [next_token],
-            cache=cache,
-            generated_tokens=1,
-        )
+        if new_slot_ids:
+            self._sync_new_kv_to_pool(cache, prefix_len, new_slot_ids)
+
+        self._req_token_ids[req_id] = list(full_token_ids) + [next_token]
+        self._req_caches[req_id] = cache
+        self._req_pool_idx[req_id] = req_pool_idx
+        self._req_synced_offset[req_id] = prefix_len + len(new_slot_ids)
 
         return next_token
 
-    def prefill_batch(
+    def extend(
         self,
-        req_ids: list[str],
-        token_ids_list: list[list[int]],
-    ) -> list[int]:
-        """Run batched prefill for multiple requests in a single forward pass.
+        req_id: str,
+        new_token_ids: list[int],
+        new_slot_ids: list[int],
+    ) -> int:
+        """Continue prefill for a chunked request.  Returns next_token_id."""
+        assert req_id in self._req_caches, f"extend called for unknown request {req_id}"
 
-        When all sequences have the same length, they are stacked into a single
-        batch tensor for one forward pass.  For variable-length sequences the
-        method falls back to serial prefill.
+        cache = self._req_caches[req_id]
 
-        Args:
-            req_ids: List of request identifiers
-            token_ids_list: List of token ID sequences, one per request
-
-        Returns:
-            List of next token IDs (greedy sampled)
-        """
-        if len(req_ids) == 1:
-            return [self.prefill(req_ids[0], token_ids_list[0])]
-
-        # Check if all sequences have the same length (enables true batching)
-        lengths = [len(tids) for tids in token_ids_list]
-        if len(set(lengths)) != 1:
-            # Variable lengths – fall back to serial prefill
-            return [
-                self.prefill(rid, tids) for rid, tids in zip(req_ids, token_ids_list)
-            ]
-
-        # All same length – use a single set of fresh caches;
-        # they'll be populated with shape (batch_size, ...) on the first forward pass
-        batch_cache = make_prompt_cache(self.model)
-
-        # Stack into (batch_size, seq_len)
-        batched_input = mx.array(
-            [list(tids) for tids in token_ids_list], dtype=mx.int32
-        )
-
-        # Single forward pass
-        model_output = self.model(batched_input, cache=batch_cache)
+        input_ids = mx.array([new_token_ids], dtype=mx.int32)
+        model_output = self.model(input_ids, cache=cache)
         logits = self._extract_logits(model_output)
 
         last_logits = logits[:, -1, :]
-        next_tokens_mlx = mx.argmax(last_logits, axis=-1)
+        next_token_mlx = mx.argmax(last_logits, axis=-1)
+        self._eval_with_cache(next_token_mlx, cache)
+        next_token = int(next_token_mlx.item())
 
-        # Evaluate everything together
-        mx.eval(next_tokens_mlx, *[c.state for c in batch_cache])
-        next_tokens = next_tokens_mlx.tolist()
+        prev_tokens = self._req_token_ids[req_id]
+        if prev_tokens:
+            prev_tokens.pop()  # remove stale intermediate token
+        prev_tokens.extend(new_token_ids)
+        prev_tokens.append(next_token)
 
-        # Extract individual caches and store per-request state
-        for i, req_id in enumerate(req_ids):
-            individual_cache = _extract_kv_cache(batch_cache, i)
-            self._request_states[req_id] = MlxRequestState(
-                token_ids=list(token_ids_list[i]) + [next_tokens[i]],
-                cache=individual_cache,
-                generated_tokens=1,
-            )
+        # Sync new chunk KV to pool immediately
+        if not self.disable_radix_cache and new_slot_ids:
+            synced = self._req_synced_offset[req_id]
+            self._sync_new_kv_to_pool(cache, synced, new_slot_ids)
+            self._req_synced_offset[req_id] = synced + len(new_slot_ids)
 
-        return next_tokens
+        return next_token
+
+    def _sync_new_kv_to_pool(
+        self,
+        cache: list[ContiguousKVCache],
+        cache_start: int,
+        slot_ids: list[int],
+    ) -> None:
+        """Sync KV from contiguous cache to pool at the given slot IDs."""
+        if not slot_ids or self._kv_pool is None:
+            return
+        num_layers = len(cache)
+        end = cache_start + len(slot_ids)
+        slot_ids_mx = mx.array(slot_ids, dtype=mx.int32)
+        # Transpose cache (1, n_kv_heads, S, head_dim) → pool (S, n_kv_heads, head_dim)
+        k_all = mx.stack(
+            [
+                cache[i].keys[0, :, cache_start:end, :].transpose(1, 0, 2)
+                for i in range(num_layers)
+            ]
+        )
+        v_all = mx.stack(
+            [
+                cache[i].values[0, :, cache_start:end, :].transpose(1, 0, 2)
+                for i in range(num_layers)
+            ]
+        )
+        self._kv_pool.set_kv_all_layers(slot_ids_mx, k_all, v_all)
+
+    def _sync_decode_kv_to_pool(self, req_id: str) -> None:
+        """Sync un-flushed decode KV for *req_id* to the shared pool."""
+        if self._kv_pool is None or self._req_to_token_pool is None:
+            return
+        cache = self._req_caches.get(req_id)
+        if cache is None:
+            return
+        current_offset = cache[0].offset
+        synced_offset = self._req_synced_offset.get(req_id, 0)
+        if current_offset <= synced_offset:
+            return
+        req_pool_idx = self._req_pool_idx.get(req_id)
+        if req_pool_idx is None:
+            return
+        # Read slot IDs from scheduler's req_to_token_pool
+        slot_ids = (
+            self._req_to_token_pool.req_to_token[
+                req_pool_idx, synced_offset:current_offset
+            ]
+            .to(dtype=int)
+            .tolist()
+        )
+        self._sync_new_kv_to_pool(cache, synced_offset, slot_ids)
+        self._req_synced_offset[req_id] = current_offset
+
+    def flush_all_decode_kv(self) -> None:
+        """Sync all active requests' un-flushed decode KV to the pool."""
+        if self.disable_radix_cache or self._kv_pool is None:
+            return
+        for req_id in list(self._req_caches.keys()):
+            self._sync_decode_kv_to_pool(req_id)
 
     def decode_batch(
         self,
         req_ids: list[str],
     ) -> list[int]:
-        """Run batched decode for multiple requests.
+        """Decode one token per request."""
+        batch_size = len(req_ids)
+        num_layers = self._num_layers
 
-        Args:
-            req_ids: List of request IDs to decode
+        caches = [self._req_caches[rid] for rid in req_ids]
+        seq_lens = [caches[i][0].offset for i in range(batch_size)]
 
-        Returns:
-            List of next token IDs
-        """
-        if len(req_ids) == 1:
-            return [self._decode_single(req_ids[0])]
+        if batch_size == 1:
+            cache = caches[0]
+            last_token = self._req_token_ids[req_ids[0]][-1]
+            input_ids = mx.array([[last_token]], dtype=mx.int32)
+            model_output = self.model(input_ids, cache=cache)
+            logits = self._extract_logits(model_output)
+            next_tokens_mlx = mx.argmax(logits[:, -1, :], axis=-1)
+            self._eval_with_cache(next_tokens_mlx, cache)
+        else:
+            layer_caches = [
+                [caches[i][layer_idx] for i in range(batch_size)]
+                for layer_idx in range(num_layers)
+            ]
+            ctx = BatchedDecodeContext(
+                batch_size=batch_size,
+                seq_lens=seq_lens,
+                layer_caches=layer_caches,
+            )
+            set_context(ctx)
+            try:
+                max_offset = max(seq_lens)
+                shim_cache = [OffsetCache(offset=max_offset) for _ in range(num_layers)]
+                last_tokens = [self._req_token_ids[rid][-1] for rid in req_ids]
+                batched_input = mx.array(last_tokens, dtype=mx.int32)[:, None]
+                model_output = self.model(batched_input, cache=shim_cache)
+                logits = self._extract_logits(model_output)
+                next_tokens_mlx = mx.argmax(logits[:, -1, :], axis=-1)
 
-        decode_reqs = []
-        for req_id in req_ids:
-            state = self._request_states[req_id]
-            decode_reqs.append((req_id, state))
+                eval_targets = [next_tokens_mlx]
+                for c_list in caches:
+                    for c in c_list:
+                        eval_targets.append(c.keys)
+                        eval_targets.append(c.values)
+                mx.eval(*eval_targets)
+            finally:
+                clear_context()
 
-        return self._batched_decode(decode_reqs)
-
-    def _decode_single(self, req_id: str) -> int:
-        """Decode a single token for one request."""
-        state = self._request_states[req_id]
-        last_token = state.token_ids[-1]
-
-        input_ids = mx.array([[last_token]], dtype=mx.int32)
-        model_output = self.model(input_ids, cache=state.cache)
-
-        logits = self._extract_logits(model_output)
-
-        last_logits = logits[:, -1, :]
-        next_token_mlx = mx.argmax(last_logits, axis=-1)
-
-        mx.eval(next_token_mlx, *[c.state for c in state.cache])
-        next_token = int(next_token_mlx.item())
-
-        state.token_ids.append(next_token)
-        state.generated_tokens += 1
-
-        return next_token
-
-    def _batched_decode(
-        self, decode_reqs: list[tuple[str, MlxRequestState]]
-    ) -> list[int]:
-        """Run a single batched forward pass for multiple decode requests."""
-        last_tokens = [state.token_ids[-1] for _, state in decode_reqs]
-
-        # Merge individual KV caches into batched cache
-        caches_list = [state.cache for _, state in decode_reqs]
-        batch_cache = _merge_kv_caches(caches_list)
-
-        # Create batched input: shape (batch_size, 1)
-        batched_input = mx.array(last_tokens, dtype=mx.int32)[:, None]
-
-        # Single forward pass
-        model_output = self.model(batched_input, cache=batch_cache)
-        logits = self._extract_logits(model_output)
-
-        next_token_logits = logits[:, -1, :]
-        next_tokens_mlx = mx.argmax(next_token_logits, axis=-1)
-
-        mx.eval(next_tokens_mlx, *[c.state for c in batch_cache])
         next_tokens = next_tokens_mlx.tolist()
 
-        # Extract updated caches back to individual requests
-        for i, (_, state) in enumerate(decode_reqs):
-            state.cache = _extract_kv_cache(batch_cache, i)
-            state.token_ids.append(next_tokens[i])
-            state.generated_tokens += 1
+        for i, rid in enumerate(req_ids):
+            self._req_token_ids[rid].append(next_tokens[i])
 
         return next_tokens
 
+    def has_request(self, req_id: str) -> bool:
+        """Check if a request has active state."""
+        return req_id in self._req_caches
+
     def remove_request(self, req_id: str):
-        """Clean up state for a completed request."""
-        self._request_states.pop(req_id, None)
+        """Sync remaining decode KV to pool, then release request state."""
+        if not self.disable_radix_cache:
+            self._sync_decode_kv_to_pool(req_id)
+
+        self._req_token_ids.pop(req_id, None)
+        cache = self._req_caches.pop(req_id, None)
+        if cache is not None:
+            self._release_cache(cache)
+        self._req_pool_idx.pop(req_id, None)
+        self._req_synced_offset.pop(req_id, None)
 
     def clear(self):
         """Clear all request states."""
-        self._request_states.clear()
+        self._req_token_ids.clear()
+        for cache in self._req_caches.values():
+            self._cache_pool.append(cache)
+        self._req_caches.clear()
+        self._req_pool_idx.clear()
+        self._req_synced_offset.clear()
+        if self._kv_pool is not None:
+            self._kv_pool.clear()

python/sglang/srt/hardware_backend/mlx/model_runner_stub.py

@@ -1,11 +1,7 @@
 """Lightweight ModelRunner stub for MLX on Apple Silicon.
 
-Subclasses ModelRunner but overrides both load_model() and initialize()
-to skip PyTorch weight loading entirely.  No GPU memory is consumed:
-the KV cache pool uses a zero-allocation _DummyKVCache, and only
-CPU-side bookkeeping structures (req_to_token_pool,
-token_to_kv_pool_allocator) are created so the SGLang scheduler can
-function.  The actual KV cache is managed by the MLX model runner.
+Skips PyTorch weight loading.  Creates only the CPU-side bookkeeping
+(req_to_token_pool, token_to_kv_pool_allocator) the scheduler needs.
 """
 
 import logging
@@ -78,6 +74,10 @@ class MlxModelRunnerStub(ModelRunner):
     the minimal bookkeeping pools needed by the scheduler are created.
     """
 
+    def __init__(self, *args, mlx_pool_size: int | None = None, **kwargs):
+        self._mlx_pool_size = mlx_pool_size
+        super().__init__(*args, **kwargs)
+
     def load_model(self):
         """Set only the metadata that downstream code needs, without
         loading any PyTorch model weights."""
@@ -128,9 +128,12 @@ class MlxModelRunnerStub(ModelRunner):
         # KV cache dtype
         self.kv_cache_dtype = self.dtype
 
-        # Pool sizing — use context_len as the capacity.
-        # No actual GPU memory is consumed because _DummyKVCache is empty.
-        self.max_total_num_tokens = self.model_config.context_len
+        # Pool sizing — use the MLX runner's auto-sized pool if available,
+        # otherwise fall back to context_len.
+        if self._mlx_pool_size is not None:
+            self.max_total_num_tokens = self._mlx_pool_size
+        else:
+            self.max_total_num_tokens = self.model_config.context_len
         self.max_running_requests = min(
             self.max_total_num_tokens // 2,
             4096,

python/sglang/srt/hardware_backend/mlx/tp_worker.py

@@ -1,13 +1,8 @@
 """MLX-specific TpModelWorker subclass for Apple Silicon.
 
-Overrides the standard TpModelWorker to route forward passes through
-the native MLX model runner, avoiding PyTorch MPS entirely for inference.
-
-PyTorch model weights are never loaded.  A lightweight ModelRunner stub
-(MlxModelRunnerStub) provides only the minimal bookkeeping structures
-(req_to_token_pool, token_to_kv_pool_allocator with a zero-memory
-dummy KV cache) that the SGLang scheduler expects.  The actual KV cache
-is managed internally by the MLX model runner.
+Routes forward passes through the MLX model runner, bypassing PyTorch
+MPS.  A lightweight stub provides scheduler bookkeeping; the actual
+KV data lives in MlxKVPool.
 """
 
 import logging
@@ -33,12 +28,23 @@ class MlxTpModelWorker(TpModelWorker):
     """
 
     def _init_model_runner(self):
-        """Override to use a lightweight ModelRunner that skips weight loading."""
+        """Create MLX runner first (auto-sizes pool), then stub with matching size."""
         from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
         from sglang.srt.hardware_backend.mlx.model_runner_stub import (
             MlxModelRunnerStub,
         )
 
+        logger.info("Initializing MlxModelRunner for end-to-end MLX inference")
+        init_kwargs = dict(
+            model_path=self.server_args.model_path,
+            trust_remote_code=self.server_args.trust_remote_code,
+            disable_radix_cache=self.server_args.disable_radix_cache,
+            mem_fraction_static=self.server_args.mem_fraction_static,
+        )
+        if self.server_args.max_total_tokens is not None:
+            init_kwargs["pool_size"] = self.server_args.max_total_tokens
+        self._mlx_runner = MlxModelRunner(**init_kwargs)
+
         self._model_runner = MlxModelRunnerStub(
             model_config=self.model_config,
             mem_fraction_static=self.server_args.mem_fraction_static,
@@ -56,20 +62,22 @@ class MlxTpModelWorker(TpModelWorker):
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
             memory_pool_config=self.memory_pool_config,
+            mlx_pool_size=self._mlx_runner.pool_size,
         )
 
-        # Initialize the MLX model runner (loads weights via MLX, not PyTorch)
-        logger.info("Initializing MlxModelRunner for end-to-end MLX inference")
-        self._mlx_runner = MlxModelRunner(
-            model_path=self.server_args.model_path,
-            trust_remote_code=self.server_args.trust_remote_code,
-        )
         self._mlx_active_rids: set[str] = set()
+        self._mlx_pool_initialized = False
 
     def get_pad_input_ids_func(self):
         """Override since the stub ModelRunner has no real model."""
         return None
 
+    def _ensure_mlx_pool_initialized(self):
+        """Lazily initialize the MlxKVPool after the stub's pools are ready."""
+        if not self._mlx_pool_initialized:
+            self._mlx_runner.init_kv_pool(self._model_runner.req_to_token_pool)
+            self._mlx_pool_initialized = True
+
     def forward_batch_generation(
         self,
         model_worker_batch: ModelWorkerBatch,
@@ -80,6 +88,7 @@ class MlxTpModelWorker(TpModelWorker):
     ) -> GenerationBatchResult:
         """Override to route through MLX model runner."""
         if model_worker_batch is not None:
+            self._ensure_mlx_pool_initialized()
             return self._forward_batch_generation_mlx(model_worker_batch)
 
         # Fallback to standard path for None batches
@@ -95,11 +104,7 @@ class MlxTpModelWorker(TpModelWorker):
         self,
         model_worker_batch: ModelWorkerBatch,
     ) -> GenerationBatchResult:
-        """Run forward pass through the MLX model runner.
-
-        Bypasses the standard ModelRunner forward+sample and uses native MLX
-        inference for the entire model. Only supports greedy sampling.
-        """
+        """Run forward pass through the MLX model runner (greedy only)."""
         from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 
         forward_mode = model_worker_batch.forward_mode
@@ -111,32 +116,61 @@ class MlxTpModelWorker(TpModelWorker):
                 can_run_cuda_graph=False,
             )
 
-        # Auto-cleanup: remove MLX state for requests no longer in the batch
+        # Auto-cleanup: remove MLX state for requests no longer in the batch.
         current_rids = {req.rid for req in reqs}
-        stale_rids = self._mlx_active_rids - current_rids
-        for rid in stale_rids:
-            self._mlx_runner.remove_request(rid)
-        self._mlx_active_rids = current_rids
+        if forward_mode.is_decode():
+            stale_rids = self._mlx_active_rids - current_rids
+            for rid in stale_rids:
+                self._mlx_runner.remove_request(rid)
+            self._mlx_active_rids = current_rids
+        else:
+            self._mlx_active_rids |= current_rids
 
         next_token_ids_list = []
 
         if forward_mode.is_extend():
-            # Prefill (or MIXED): extract per-request tokens from concatenated input_ids
+            # Ensure pool is up-to-date before PoolBackedCache reads it
+            # for prefix-cached prefills.  Only runs on extend batches.
+            self._mlx_runner.flush_all_decode_kv()
             input_ids_cpu = model_worker_batch.input_ids.cpu().tolist()
+            out_cache_loc_cpu = model_worker_batch.out_cache_loc.cpu().tolist()
             extend_seq_lens = model_worker_batch.extend_seq_lens
-            offset = 0
+
+            offset = 0  # into input_ids_cpu
+            slot_offset = 0  # into out_cache_loc_cpu
             prefill_rids = []
+            extend_rids = []
             decode_rids = []
+
             for i, req in enumerate(reqs):
                 seq_len = extend_seq_lens[i]
                 req_token_ids = input_ids_cpu[offset : offset + seq_len]
+                req_new_slots = out_cache_loc_cpu[slot_offset : slot_offset + seq_len]
                 offset += seq_len
-                if req.rid in self._mlx_runner._request_states:
-                    # MIXED mode: this request already has MLX state, decode it
-                    decode_rids.append(req.rid)
+                slot_offset += seq_len
+
+                if self._mlx_runner.has_request(req.rid):
+                    if seq_len > 1:
+                        # Chunked prefill continuation
+                        next_token = self._mlx_runner.extend(
+                            req.rid, req_token_ids, req_new_slots
+                        )
+                        extend_rids.append((req.rid, next_token))
+                    else:
+                        # MIXED mode: single-token decode
+                        decode_rids.append(req.rid)
                 else:
-                    # Prefill: new request
-                    next_token = self._mlx_runner.prefill(req.rid, req_token_ids)
+                    # New prefill
+                    prefix_slot_ids = req.prefix_indices.tolist()
+                    full_token_ids = list(req.fill_ids)
+                    next_token = self._mlx_runner.prefill(
+                        req_id=req.rid,
+                        new_token_ids=req_token_ids,
+                        full_token_ids=full_token_ids,
+                        prefix_slot_ids=prefix_slot_ids,
+                        new_slot_ids=req_new_slots,
+                        req_pool_idx=req.req_pool_idx,
+                    )
                     prefill_rids.append((req.rid, next_token))
 
             # Batch decode all existing requests at once
@@ -147,16 +181,17 @@ class MlxTpModelWorker(TpModelWorker):
                 decode_map = {}
 
             prefill_map = dict(prefill_rids)
+            extend_map = dict(extend_rids)
 
-            # Reassemble in original request order
             for req in reqs:
                 if req.rid in decode_map:
                     next_token_ids_list.append(decode_map[req.rid])
+                elif req.rid in extend_map:
+                    next_token_ids_list.append(extend_map[req.rid])
                 else:
                     next_token_ids_list.append(prefill_map[req.rid])
 
         elif forward_mode.is_decode():
-            # Decode: batch decode all requests
             req_ids = [req.rid for req in reqs]
             next_token_ids_list = self._mlx_runner.decode_batch(req_ids)