chore: import upstream snapshot with attribution

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

@@ -0,0 +1,59 @@
+"""Cache components for the MLX backend."""
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_head_dim,
+    get_num_heads,
+    get_num_kv_heads,
+    is_attention_module,
+    uses_sliding_window_attention,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+    AttentionOffsetCache,
+    ContiguousAttentionKVCache,
+    PoolBackedAttentionKVCache,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
+    MlxAttentionKVPool,
+)
+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.auxiliary_state import (
+    MlxAuxiliaryStateComponent,
+    MlxAuxiliaryStatePool,
+    MlxAuxiliaryStateReqToTokenPool,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.layout import MlxModelCacheLayout
+from sglang.srt.hardware_backend.mlx.kv_cache.model_patching import (
+    find_attention_layers,
+    get_num_layers,
+    patch_model_attention,
+)
+
+__all__ = [
+    "BatchedDecodeContext",
+    "clear_context",
+    "AttentionOffsetCache",
+    "ContiguousAttentionKVCache",
+    "find_attention_layers",
+    "get_head_dim",
+    "get_context",
+    "get_num_layers",
+    "get_num_heads",
+    "get_num_kv_heads",
+    "is_attention_module",
+    "MLXAttentionWrapper",
+    "MlxAttentionKVPool",
+    "MlxAuxiliaryStateComponent",
+    "MlxAuxiliaryStatePool",
+    "MlxAuxiliaryStateReqToTokenPool",
+    "MlxModelCacheLayout",
+    "patch_model_attention",
+    "PoolBackedAttentionKVCache",
+    "set_context",
+    "uses_sliding_window_attention",
+]

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

@@ -0,0 +1,66 @@
+"""Attention helpers based on duck typing for the MLX backend."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+# ``rope`` and ``scale`` are required by MLXAttentionWrapper. Keeping them in
+# the contract also prevents recurrent mixers such as DeltaNet from being
+# mistaken for softmax attention just because they expose projection layers.
+ATTENTION_API_ATTRS = ("q_proj", "k_proj", "v_proj", "o_proj", "rope", "scale")
+NUM_HEAD_ATTRS = ("n_heads", "num_heads", "num_attention_heads")
+NUM_KV_HEAD_ATTRS = ("n_kv_heads", "num_k_heads", "num_kv_heads", "num_key_value_heads")
+SLIDING_ATTENTION_ATTRS = (
+    "is_sliding",
+    "use_sliding",
+    "is_sliding_window",
+    "use_sliding_window",
+    "is_swa",
+)
+
+
+def first_present_attr(module: Any, names: Iterable[str]) -> Any | None:
+    """Return the first present attribute value without treating 0 as absent."""
+    for name in names:
+        if hasattr(module, name):
+            return getattr(module, name)
+    return None
+
+
+def get_num_heads(module: Any) -> int | None:
+    return first_present_attr(module, NUM_HEAD_ATTRS)
+
+
+def get_num_kv_heads(module: Any) -> int | None:
+    return first_present_attr(module, NUM_KV_HEAD_ATTRS)
+
+
+def get_head_dim(module: Any) -> int | None:
+    head_dim = first_present_attr(module, ("head_dim",))
+    if head_dim is not None:
+        return head_dim
+
+    n_kv_heads = get_num_kv_heads(module)
+    if n_kv_heads is None:
+        return None
+    if hasattr(module, "hidden_size") and hasattr(module, "num_k_heads"):
+        return module.hidden_size // module.num_k_heads
+    if hasattr(module, "k_proj") and hasattr(module.k_proj, "weight"):
+        return module.k_proj.weight.shape[0] // n_kv_heads
+    return None
+
+
+def is_attention_module(module: Any) -> bool:
+    return (
+        all(hasattr(module, attr) for attr in ATTENTION_API_ATTRS)
+        and get_num_heads(module) is not None
+        and get_num_kv_heads(module) is not None
+    )
+
+
+def uses_sliding_window_attention(*modules: Any) -> bool:
+    return any(
+        bool(getattr(module, attr, False))
+        for module in modules
+        for attr in SLIDING_ATTENTION_ATTRS
+    )

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

@@ -0,0 +1,210 @@
+"""Attention KV cache adapters for the 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.attention_kv_pool import (
+        MlxAttentionKVPool,
+    )
+
+
+class AttentionOffsetCache:
+    """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("AttentionOffsetCache should not store data")
+
+
+_DEFAULT_MAX_SEQ_LEN = 4096
+
+
+class ContiguousAttentionKVCache:
+    """Pre-allocated attention KV buffer for one request and 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)."""
+        end = self.offset + 1
+        if end > self.max_seq_len:
+            self._grow(end)
+        self.keys[:, :, self.offset : end, :] = k
+        self.values[:, :, self.offset : end, :] = v
+        self.offset = end
+
+    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 PoolBackedAttentionKVCache:
+    """Lazily gathers cached attention KV from the shared pool during forward.
+
+    Each ``update_and_fetch`` gathers this layer's prefix from the pool
+    on demand, keeping operations in the lazy compute graph.  Convert to
+    ``ContiguousAttentionKVCache`` 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: MlxAttentionKVPool,
+        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) -> ContiguousAttentionKVCache:
+        """Convert to contiguous attention KV reusing forward-pass arrays."""
+        cache = ContiguousAttentionKVCache(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/attention_kv_pool.py

@@ -0,0 +1,86 @@
+"""Flat attention KV pool for the MLX backend.
+
+Each layer buffer has shape ``(pool_size, n_kv_heads, head_dim)``.
+This v1 pool is intentionally uniform: every wrapped softmax-attention
+layer must share the same KV shape and full-context KV semantics.
+Heterogeneous KV shapes and sliding-window KV need per-layer/window-aware
+pools before they can use MLX radix reuse.
+
+Slot 0 is reserved as padding (1-based indexing).
+"""
+
+import logging
+
+import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+
+class MlxAttentionKVPool:
+    """Pre-allocated attention 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-attention-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"MlxAttentionKVPool: {pool_size} slots x {num_layers} layers "
+            f"x {n_kv_heads} heads x {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/attention_wrapper.py

@@ -0,0 +1,298 @@
+"""Batched decode attention wrapper for MLX backend."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from sglang.srt.hardware_backend.mlx.aot import (
+    MlxAOTKernelContext,
+    MlxAOTKernelSet,
+    MlxAOTRoPEContext,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_head_dim,
+    get_num_heads,
+    get_num_kv_heads,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+    ContiguousAttentionKVCache,
+)
+
+_thread_local = threading.local()
+
+
+# TODO: Move from threading to multiprocessing or asyncio
+@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
+    # attention_layer_caches[attention_pool_idx][req_idx] = ContiguousAttentionKVCache
+    attention_layer_caches: list[list[ContiguousAttentionKVCache]]
+    attention_pool_index_by_layer: dict[int, int] = field(default_factory=dict)
+
+    # Optional AOT kernel state. Keep kernel-specific fields out of the regular
+    # MLX decode path so future AOT kernels can be added without growing this
+    # context one field at a time.
+    aot: MlxAOTKernelContext = field(default_factory=MlxAOTKernelContext)
+
+    # Derived tensors/metadata, shared across all layers in one forward pass.
+    offsets: mx.array = field(init=False)
+    max_len: int = field(init=False)
+    valid_lens: mx.array = field(init=False)
+    needs_padding: bool = field(init=False)
+    pad_sizes: list[int] = field(init=False)
+    positions: Optional[mx.array] = field(init=False)
+
+    def __post_init__(self) -> None:
+        seq_lens = self.seq_lens
+        max_seq_len = max(seq_lens)
+        self.offsets = mx.array(seq_lens, dtype=mx.int32)
+        self.max_len = max_seq_len + 1
+        self.valid_lens = self.offsets + 1
+        self.needs_padding = min(seq_lens) < max_seq_len
+        self.pad_sizes = [max_seq_len - s for s in seq_lens]
+        self.positions = mx.arange(self.max_len) if self.needs_padding else None
+        if not self.attention_pool_index_by_layer:
+            self.attention_pool_index_by_layer = {
+                idx: idx for idx in range(len(self.attention_layer_caches))
+            }
+
+    @classmethod
+    def from_decode(
+        cls,
+        *,
+        caches: list[list[Any]],
+        req_ids: list[str],
+        aot_kernels: MlxAOTKernelSet,
+        kv_pool: Any | None,
+        req_pool_idx: dict[str, int],
+        req_to_token_pool: Any | None,
+        attention_layer_indices: list[int] | None = None,
+        attention_pool_index_by_layer: dict[int, int] | None = None,
+    ) -> BatchedDecodeContext:
+        batch_size = len(req_ids)
+        if attention_layer_indices is None:
+            attention_layer_indices = list(range(len(caches[0])))
+        seq_lens = [
+            caches[i][attention_layer_indices[0]].offset for i in range(batch_size)
+        ]
+        attention_layer_caches = [
+            [caches[i][layer_idx] for i in range(batch_size)]
+            for layer_idx in attention_layer_indices
+        ]
+        return cls(
+            batch_size=batch_size,
+            seq_lens=seq_lens,
+            attention_layer_caches=attention_layer_caches,
+            attention_pool_index_by_layer=attention_pool_index_by_layer or {},
+            aot=MlxAOTKernelContext.from_decode(
+                aot_kernels=aot_kernels,
+                kv_pool=kv_pool,
+                req_ids=req_ids,
+                req_pool_idx=req_pool_idx,
+                req_to_token_pool=req_to_token_pool,
+                layer_caches=attention_layer_caches,
+            ),
+        )
+
+
+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
+        n_heads = get_num_heads(inner)
+        n_kv_heads = get_num_kv_heads(inner)
+        if n_heads is None or n_kv_heads is None:
+            raise RuntimeError(
+                f"Cannot determine attention head counts for {type(inner).__name__}"
+            )
+
+        q_proj_output = inner.q_proj(x)
+        keys = inner.k_proj(x)
+        values = inner.v_proj(x)
+
+        head_dim = get_head_dim(inner)
+        if head_dim is None:
+            head_dim = keys.shape[-1] // n_kv_heads
+
+        q_width = n_heads * head_dim
+        gate = None
+        if q_proj_output.shape[-1] == q_width:
+            queries = q_proj_output.reshape(B, 1, n_heads, head_dim)
+        elif q_proj_output.shape[-1] == 2 * q_width:
+            queries, gate = mx.split(
+                q_proj_output.reshape(B, 1, n_heads, 2 * head_dim), 2, axis=-1
+            )
+            gate = gate.reshape(B, 1, q_width)
+        else:
+            raise RuntimeError(
+                f"Unexpected q_proj output shape {q_proj_output.shape} for "
+                f"{type(inner).__name__}"
+            )
+
+        keys = keys.reshape(B, 1, n_kv_heads, head_dim)
+        values = values.reshape(B, 1, 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 (cached on the context).
+        offsets = ctx.offsets
+        attention_pool_idx = ctx.attention_pool_index_by_layer[layer_idx]
+
+        if ctx.aot.rope is not None:
+            # AOT path: real .metallib RoPE + fused KV pool scatter.
+            queries, keys = self._rope_custom_aot(
+                queries,
+                keys,
+                values,
+                offsets,
+                attention_pool_idx,
+                ctx.aot.rope,
+            )
+        else:
+            # Fallback: MLX's built-in mx.fast.rope (used when the AOT kernel
+            # isn't built or the model uses an unsupported RoPE variant).
+            queries = inner.rope(queries, offset=offsets)
+            keys = inner.rope(keys, offset=offsets)
+
+        layer_caches = ctx.attention_layer_caches[attention_pool_idx]
+        pad_sizes = ctx.pad_sizes
+
+        # 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()
+
+            pad = pad_sizes[i]
+            if pad > 0:
+                k_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=k_all.dtype)
+                v_pad = mx.zeros((1, 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
+        if ctx.needs_padding:
+            mask_bool = ctx.positions[None, :] >= ctx.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)
+        if gate is not None:
+            output = output * mx.sigmoid(gate)
+        return inner.o_proj(output)
+
+    @staticmethod
+    def _rope_custom_aot(
+        queries: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        positions: mx.array,
+        attention_pool_idx: int,
+        rope_ctx: MlxAOTRoPEContext,
+    ) -> tuple[mx.array, mx.array]:
+        """AOT path: rotate Q/K and scatter K/V into the shared pool.
+
+        The kernel call does RoPE on Q/K and scatters
+        rotated K + (untouched) V into ``kv_pool`` at ``new_token_slots``
+        for ``layer_idx``.
+
+        If ``new_token_slots`` is None, slot=-1 sentinel is used (no pool
+        write, RoPE-only mode). Returns rotated (queries, keys) in the
+        original 4-D attention layout. ``values`` is unchanged by RoPE.
+        """
+        # (B, n_heads, 1, head_dim) -> (B, n_heads, head_dim) for kernel
+        q_flat = queries[:, :, 0, :]
+        k_flat = keys[:, :, 0, :]
+        v_flat = values[:, :, 0, :]
+        B = q_flat.shape[0]
+
+        if rope_ctx.new_token_slots is None:
+            slots = mx.full((B,), -1, dtype=mx.int32)
+        else:
+            slots = rope_ctx.new_token_slots.astype(mx.int32)
+
+        k_pool = rope_ctx.kv_pool.k_buffer[attention_pool_idx]
+        v_pool = rope_ctx.kv_pool.v_buffer[attention_pool_idx]
+
+        q_rot, k_rot, k_pool_new, v_pool_new = rope_ctx.kernel.rope_pool_fused(
+            q_flat,
+            k_flat,
+            v_flat,
+            positions,
+            slots,
+            k_pool,
+            v_pool,
+            head_dim=rope_ctx.kernel.config["head_dim"],
+            num_qo_heads=rope_ctx.kernel.config["num_qo_heads"],
+            num_kv_heads=rope_ctx.kernel.config["num_kv_heads"],
+            rope_base=rope_ctx.kernel.base,
+        )
+        # Rebind pool buffers (zero-copy donation result).
+        rope_ctx.kv_pool.k_buffer[attention_pool_idx] = k_pool_new
+        rope_ctx.kv_pool.v_buffer[attention_pool_idx] = v_pool_new
+
+        # (B, n_heads, head_dim) -> (B, n_heads, 1, head_dim) for SDPA path
+        return q_rot[:, :, None, :], k_rot[:, :, None, :]

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

@@ -0,0 +1,390 @@
+"""MLX auxiliary-state snapshots for unified radix cache.
+
+Hybrid MLX models may include non-softmax-attention layers whose native
+``mlx-lm`` cache state cannot be reconstructed from the attention KV pool.
+The global scheduler exposes that state through its existing MAMBA component
+contract, so this MLX adapter keeps those scheduler-facing field names while
+storing model-agnostic native cache snapshots.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+import mlx.core as mx
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import EvictParams, InsertResult
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.unified_cache_components.mamba_component import (
+    MambaComponent,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import TreeComponent
+
+_CACHE_ATTRS = ("offset", "lengths", "left_padding")
+_MISSING = object()
+
+
+def _clone_tree(value: Any) -> Any:
+    if isinstance(value, mx.array):
+        return mx.array(value)
+    if isinstance(value, list):
+        return [_clone_tree(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_clone_tree(item) for item in value)
+    if isinstance(value, dict):
+        return {key: _clone_tree(item) for key, item in value.items()}
+    return value
+
+
+def _arrays_in_tree(value: Any) -> list[mx.array]:
+    arrays: list[mx.array] = []
+
+    def collect(item: Any) -> None:
+        if isinstance(item, mx.array):
+            arrays.append(item)
+        elif isinstance(item, (list, tuple)):
+            for child in item:
+                collect(child)
+        elif isinstance(item, dict):
+            for child in item.values():
+                collect(child)
+
+    collect(value)
+    return arrays
+
+
+@dataclass
+class _CacheSnapshot:
+    state: Any
+    meta_state: Any = _MISSING
+    attrs: dict[str, Any] | None = None
+
+
+def _snapshot_cache(cache: Any) -> _CacheSnapshot:
+    state = _clone_tree(getattr(cache, "state", ()))
+    meta_state = (
+        _clone_tree(cache.meta_state) if hasattr(cache, "meta_state") else _MISSING
+    )
+    attrs = {
+        name: _clone_tree(getattr(cache, name))
+        for name in _CACHE_ATTRS
+        if hasattr(cache, name)
+    }
+    arrays = _arrays_in_tree((state, meta_state, attrs))
+    if arrays:
+        mx.eval(*arrays)
+    return _CacheSnapshot(state=state, meta_state=meta_state, attrs=attrs)
+
+
+def _restore_cache(cache: Any, snapshot: _CacheSnapshot) -> None:
+    cache.state = _clone_tree(snapshot.state)
+    if snapshot.meta_state is not _MISSING and hasattr(cache, "meta_state"):
+        cache.meta_state = _clone_tree(snapshot.meta_state)
+    for name, value in (snapshot.attrs or {}).items():
+        setattr(cache, name, _clone_tree(value))
+
+
+class MlxAuxiliaryStatePool:
+    """Index-addressable snapshots of native MLX auxiliary cache state."""
+
+    def __init__(self, size: int, device: str):
+        self.size = size
+        self.device = device
+        self.mamba_cache = None
+        self.mem_usage = 0
+        self._snapshots: dict[int, dict[int, _CacheSnapshot]] = {}
+        self.clear()
+
+    def _tensor(self, indices: Any) -> torch.Tensor:
+        return torch.as_tensor(indices, dtype=torch.int64, device=self.device).view(-1)
+
+    def _index(self, index: Any) -> int:
+        flat = self._tensor(index)
+        assert flat.numel() == 1
+        return int(flat.item())
+
+    def available_size(self) -> int:
+        return int(self.free_slots.numel())
+
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        if need_size > self.available_size():
+            return None
+        slots = self.free_slots[:need_size].clone()
+        self.free_slots = self.free_slots[need_size:]
+        for slot in slots.tolist():
+            self._snapshots.pop(int(slot), None)
+        return slots
+
+    def free(self, indices: Any) -> None:
+        if indices is None:
+            return
+        indices = self._tensor(indices)
+        if indices.numel() == 0:
+            return
+        for slot in indices.tolist():
+            self._snapshots.pop(int(slot), None)
+        self.free_slots = torch.cat([self.free_slots, indices])
+
+    def clear(self) -> None:
+        self.free_slots = torch.arange(
+            1, self.size + 1, dtype=torch.int64, device=self.device
+        )
+        self._snapshots.clear()
+
+    def copy_from(self, src: Any, dst: Any) -> None:
+        src_indices = self._tensor(src)
+        dst_indices = self._tensor(dst)
+        assert src_indices.numel() == dst_indices.numel()
+        for src_idx, dst_idx in zip(src_indices.tolist(), dst_indices.tolist()):
+            snapshot = self._snapshots.get(int(src_idx))
+            if snapshot is None:
+                self._snapshots.pop(int(dst_idx), None)
+            else:
+                self._snapshots[int(dst_idx)] = {
+                    layer_idx: _CacheSnapshot(
+                        state=_clone_tree(cache_snapshot.state),
+                        meta_state=_clone_tree(cache_snapshot.meta_state),
+                        attrs=_clone_tree(cache_snapshot.attrs),
+                    )
+                    for layer_idx, cache_snapshot in snapshot.items()
+                }
+
+    def fork_from(self, src: Any) -> Optional[torch.Tensor]:
+        src_indices = self._tensor(src)
+        dst = self.alloc(src_indices.numel())
+        if dst is None:
+            return None
+        self.copy_from(src_indices, dst)
+        return dst
+
+    def store_cache(
+        self,
+        index: Any,
+        cache: list[Any],
+        layer_indices: Iterable[int],
+    ) -> None:
+        self._snapshots[self._index(index)] = {
+            layer_idx: _snapshot_cache(cache[layer_idx]) for layer_idx in layer_indices
+        }
+
+    def restore_cache(
+        self,
+        index: Any,
+        cache: list[Any],
+        layer_indices: Iterable[int] | None = None,
+    ) -> bool:
+        snapshot = self._snapshots.get(self._index(index))
+        if snapshot is None:
+            return False
+        selected_layers = set(layer_indices) if layer_indices is not None else None
+        for layer_idx, cache_snapshot in snapshot.items():
+            if selected_layers is not None and layer_idx not in selected_layers:
+                continue
+            _restore_cache(cache[layer_idx], cache_snapshot)
+        return True
+
+    def has_snapshot(self, index: Any) -> bool:
+        return self._index(index) in self._snapshots
+
+
+class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
+    """Req-to-token pool with MLX auxiliary-state slot bookkeeping."""
+
+    def __init__(
+        self,
+        *,
+        size: int,
+        max_context_len: int,
+        device: str,
+        enable_memory_saver: bool,
+        auxiliary_state_size: int,
+    ):
+        super().__init__(
+            size=size,
+            max_context_len=max_context_len,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+        )
+        self.mamba_pool = MlxAuxiliaryStatePool(
+            size=auxiliary_state_size,
+            device=device,
+        )
+        # The unified radix base MAMBA component still reads ``mamba_pool``.
+        # Keep the MLX-owned name beside it so local code can avoid model-
+        # specific terminology.
+        self.auxiliary_state_pool = self.mamba_pool
+        self.enable_mamba_extra_buffer = False
+        self.req_index_to_auxiliary_state_index_mapping = torch.zeros(
+            self._alloc_size, dtype=torch.int32, device=device
+        )
+
+    def alloc(self, reqs):
+        select_index = super().alloc(reqs)
+        if select_index is None:
+            return None
+
+        auxiliary_state_indices = []
+        for req in reqs:
+            if getattr(req, "mamba_pool_idx", None) is not None:
+                mid = req.mamba_pool_idx
+            else:
+                allocated = self.auxiliary_state_pool.alloc(1)
+                assert allocated is not None, "Not enough MLX auxiliary state slots"
+                mid = allocated[0]
+                req.mamba_pool_idx = mid
+            auxiliary_state_indices.append(mid.to(dtype=torch.int32))
+        self.req_index_to_auxiliary_state_index_mapping[select_index] = torch.stack(
+            auxiliary_state_indices
+        )
+        return select_index
+
+    def get_auxiliary_state_indices(self, req_indices) -> torch.Tensor:
+        return self.req_index_to_auxiliary_state_index_mapping[req_indices]
+
+    def get_mamba_indices(self, req_indices) -> torch.Tensor:
+        return self.get_auxiliary_state_indices(req_indices)
+
+    def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int:
+        return 0
+
+    def free_mamba_cache(self, req, mamba_ping_pong_track_buffer_to_keep=None):
+        if getattr(req, "mamba_pool_idx", None) is not None:
+            self.auxiliary_state_pool.free(req.mamba_pool_idx.unsqueeze(0))
+            req.mamba_pool_idx = None
+        track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+        if track_buffer is not None:
+            if mamba_ping_pong_track_buffer_to_keep is None:
+                self.auxiliary_state_pool.free(track_buffer)
+            req.mamba_ping_pong_track_buffer = None
+            req.mamba_next_track_idx = None
+
+    def free_auxiliary_state_cache(self, req, track_buffer_to_keep=None):
+        self.free_mamba_cache(
+            req,
+            mamba_ping_pong_track_buffer_to_keep=track_buffer_to_keep,
+        )
+
+    def free(self, req):
+        super().free(req)
+
+    def clear(self):
+        super().clear()
+        self.auxiliary_state_pool.clear()
+        self.req_index_to_auxiliary_state_index_mapping.zero_()
+
+
+class MlxAuxiliaryStateComponent(MambaComponent):
+    """Unified radix component for MLX native auxiliary-state snapshots."""
+
+    def __init__(self, cache, params):
+        if params.enable_mamba_extra_buffer:
+            raise NotImplementedError(
+                "MLX auxiliary-state radix cache does not support "
+                "enable_mamba_extra_buffer yet."
+            )
+        pool = getattr(cache.req_to_token_pool, "auxiliary_state_pool", None)
+        if not isinstance(pool, MlxAuxiliaryStatePool):
+            raise TypeError(
+                "MlxAuxiliaryStateComponent requires MlxAuxiliaryStatePool, "
+                f"got {type(pool)}"
+            )
+        TreeComponent.__init__(self, cache, params)
+        self.enable_mamba_extra_buffer = False
+        self._mamba_pool_host = None
+
+    @staticmethod
+    def _tracked_value(req) -> tuple[object | None, bool]:
+        track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+        track_len = getattr(req, "mamba_last_track_seqlen", None)
+        if track_buffer is not None and track_len is not None:
+            return track_buffer[0].unsqueeze(-1).clone(), True
+        if getattr(req, "mamba_pool_idx", None) is None:
+            return None, False
+        return req.mamba_pool_idx.unsqueeze(-1).clone(), False
+
+    def prepare_for_caching_req(
+        self,
+        req,
+        insert_params,
+        token_ids_len: int,
+        is_finished: bool,
+    ) -> int | None:
+        cache_len = getattr(req, "mamba_last_track_seqlen", None)
+        auxiliary_value, uses_track_slot = self._tracked_value(req)
+        setattr(insert_params, "mlx_auxiliary_state_uses_track_slot", uses_track_slot)
+
+        if auxiliary_value is None:
+            return 0 if is_finished else None
+
+        if cache_len is None:
+            cache_len = token_ids_len
+        if is_finished:
+            insert_params.mamba_value = auxiliary_value
+        else:
+            source_value = auxiliary_value
+            forked_value = self.cache.req_to_token_pool.auxiliary_state_pool.fork_from(
+                source_value
+            )
+            if forked_value is None:
+                self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+                forked_value = (
+                    self.cache.req_to_token_pool.auxiliary_state_pool.fork_from(
+                        source_value
+                    )
+                )
+                assert forked_value is not None, "Can not alloc MLX auxiliary cache"
+            insert_params.mamba_value = forked_value
+        return cache_len
+
+    def cleanup_after_caching_req(
+        self,
+        req,
+        is_finished: bool,
+        insert_result: InsertResult | None = None,
+        insert_params=None,
+    ) -> None:
+        if not is_finished:
+            if (
+                insert_params is not None
+                and insert_params.mamba_value is not None
+                and (insert_result is None or insert_result.mamba_exist)
+            ):
+                self.cache.req_to_token_pool.auxiliary_state_pool.free(
+                    insert_params.mamba_value
+                )
+            if bool(
+                getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False)
+            ):
+                track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+                if track_buffer is not None:
+                    self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
+                req.mamba_ping_pong_track_buffer = None
+                req.mamba_next_track_idx = None
+            req.mamba_last_track_seqlen = None
+            return
+
+        auxiliary_value_exists = (
+            insert_result.mamba_exist if insert_result is not None else True
+        )
+        uses_track_slot = bool(
+            getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False)
+        )
+        if uses_track_slot:
+            keep_track_slot = not auxiliary_value_exists
+            self.cache.req_to_token_pool.free_auxiliary_state_cache(
+                req,
+                track_buffer_to_keep=0 if keep_track_slot else None,
+            )
+        elif auxiliary_value_exists:
+            self.cache.req_to_token_pool.free_auxiliary_state_cache(req)
+        else:
+            # The radix tree now owns the live auxiliary-state slot.
+            track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+            if track_buffer is not None:
+                self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
+                req.mamba_ping_pong_track_buffer = None
+                req.mamba_next_track_idx = None
+            req.mamba_pool_idx = None
+        req.mamba_last_track_seqlen = None

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

@@ -0,0 +1,93 @@
+"""Model cache layout helpers for the MLX backend."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+
+@dataclass(frozen=True)
+class MlxModelCacheLayout:
+    """Map model layers to MLX cache storage components.
+
+    Attention layers store softmax-attention KV in the MLX attention KV pool.
+    Auxiliary layers keep native ``mlx-lm`` cache state and are snapshotted by
+    the MLX auxiliary-state component.
+    """
+
+    layers: tuple[Any, ...]
+    attention_attrs: tuple[str | None, ...]
+    attention_layer_indices: tuple[int, ...]
+    auxiliary_layer_indices: tuple[int, ...]
+    attention_pool_index_by_layer: dict[int, int]
+
+    @classmethod
+    def from_attention_discovery(
+        cls,
+        layers: Sequence[Any],
+        attention_attrs: Sequence[str | None],
+    ) -> MlxModelCacheLayout:
+        if len(layers) != len(attention_attrs):
+            raise ValueError(
+                "Layer count and attention attribute count differ: "
+                f"{len(layers)} != {len(attention_attrs)}"
+            )
+
+        attention_layer_indices = tuple(
+            idx for idx, attr in enumerate(attention_attrs) if attr is not None
+        )
+        auxiliary_layer_indices = tuple(
+            idx for idx, attr in enumerate(attention_attrs) if attr is None
+        )
+        attention_pool_index_by_layer = {
+            layer_idx: pool_idx
+            for pool_idx, layer_idx in enumerate(attention_layer_indices)
+        }
+
+        return cls(
+            layers=tuple(layers),
+            attention_attrs=tuple(attention_attrs),
+            attention_layer_indices=attention_layer_indices,
+            auxiliary_layer_indices=auxiliary_layer_indices,
+            attention_pool_index_by_layer=attention_pool_index_by_layer,
+        )
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layers)
+
+    @property
+    def num_attention_layers(self) -> int:
+        return len(self.attention_layer_indices)
+
+    @property
+    def has_auxiliary_state(self) -> bool:
+        return bool(self.auxiliary_layer_indices)
+
+    @property
+    def first_attention_layer_index(self) -> int:
+        if not self.attention_layer_indices:
+            raise RuntimeError("MLX model has no supported attention layers")
+        return self.attention_layer_indices[0]
+
+    def attention_pool_index(self, layer_idx: int) -> int:
+        try:
+            return self.attention_pool_index_by_layer[layer_idx]
+        except KeyError as exc:
+            raise KeyError(f"Layer {layer_idx} is not an attention layer") from exc
+
+    def attention_attr(self, layer_idx: int) -> str:
+        attr = self.attention_attrs[layer_idx]
+        if attr is None:
+            raise KeyError(f"Layer {layer_idx} is not an attention layer")
+        return attr
+
+    def attention_layer_caches(
+        self,
+        caches_by_request: list[list[Any]],
+    ) -> list[list[Any]]:
+        """Return layer-major attention caches for batched decode."""
+        return [
+            [request_cache[layer_idx] for request_cache in caches_by_request]
+            for layer_idx in self.attention_layer_indices
+        ]

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

@@ -0,0 +1,61 @@
+"""Model introspection and attention patching."""
+
+from typing import Any
+
+import mlx.nn as nn
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    is_attention_module,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
+    MLXAttentionWrapper,
+)
+
+
+def _find_attention_attr(layer: Any) -> str | None:
+    """Return the direct child name that satisfies the attention contract."""
+    if not isinstance(layer, nn.Module):
+        raise TypeError(f"Expected mlx.nn.Module layer, got {type(layer)}")
+    for name, module in layer.children().items():
+        if isinstance(module, MLXAttentionWrapper) or is_attention_module(module):
+            return name
+    return None
+
+
+def find_attention_layers(model: Any) -> tuple[list[Any], list[str | None]]:
+    """Find transformer layers and per-layer attention attribute names."""
+    root = getattr(model, "language_model", model)
+    container = getattr(root, "model", root)
+    layer_list = getattr(container, "layers", None) or getattr(root, "layers", [])
+
+    if layer_list:
+        attn_attrs = [_find_attention_attr(layer) for layer in layer_list]
+        if any(attr is not None for attr in attn_attrs):
+            return layer_list, attn_attrs
+        raise ValueError(f"No attention attribute in layer type {type(layer_list[0])}")
+    return layer_list, []
+
+
+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_attrs = find_attention_layers(model)
+    patched = 0
+    for idx, (layer, attn_attr) in enumerate(zip(layer_list, attn_attrs)):
+        if attn_attr is None:
+            continue
+        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)