chore: import upstream snapshot with attribution

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

@@ -0,0 +1,240 @@
+"""AOT kernel selection and decode-context helpers for the MLX backend."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import mlx.core as mx
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+        ContiguousAttentionKVCache,
+    )
+
+
+def _load_metal_rope_pool_fused():
+    try:
+        from sgl_kernel import metal
+    except ImportError as exc:
+        raise ImportError(
+            "sgl_kernel.metal is not importable. Install sgl-kernel in the "
+            "active environment before enabling SGLANG_MLX_USE_CUSTOM_ROPE."
+        ) from exc
+
+    import_error = getattr(metal, "_IMPORT_ERROR", None)
+    if getattr(metal, "_metal", None) is None or import_error is not None:
+        reason = f" Reason: {import_error}." if import_error is not None else ""
+        raise ImportError(
+            "sgl_kernel.metal is importable, but the native Metal extension "
+            f"or metallib is not available.{reason} Install the Metal kernels "
+            "with `uv run sgl-kernel/setup_metal.py install` from the SGLang "
+            "repo root in the active environment."
+        ) from import_error
+    return metal.rope_pool_fused
+
+
+@dataclass
+class MlxAOTRoPEKernel:
+    base: float = 0.0
+    config: dict[str, Any] = field(default_factory=dict)
+    rope_pool_fused: Optional[Any] = None
+
+    @property
+    def enabled(self) -> bool:
+        return (
+            self.base > 0.0 and bool(self.config) and self.rope_pool_fused is not None
+        )
+
+
+@dataclass
+class MlxAOTKernelBuildInputs:
+    sample_attn: Any
+    n_kv_heads: int
+    head_dim: int
+
+
+@dataclass(frozen=True)
+class MlxAOTKernelSpec:
+    name: str
+    kernel_attr: str
+    is_enabled: Callable[[], bool]
+    build: Callable[[MlxAOTKernelBuildInputs], Any]
+
+
+@dataclass
+class MlxAOTKernelSet:
+    rope: MlxAOTRoPEKernel = field(default_factory=MlxAOTRoPEKernel)
+    selected_kernel_names: tuple[str, ...] = ()
+
+
+class MlxAOTKernelRegistry:
+    """Registry for optional MLX AOT kernels.
+
+    Each spec owns one kernel field on ``MlxAOTKernelSet``. The registry is the
+    only place that checks kernel enablement policy and model support.
+    """
+
+    def __init__(self, specs: tuple[MlxAOTKernelSpec, ...]):
+        self._specs = specs
+
+    @property
+    def registered_kernel_names(self) -> tuple[str, ...]:
+        return tuple(spec.name for spec in self._specs)
+
+    def build_kernel_set(
+        self,
+        *,
+        sample_attn: Any,
+        n_kv_heads: int,
+        head_dim: int,
+    ) -> MlxAOTKernelSet:
+        inputs = MlxAOTKernelBuildInputs(
+            sample_attn=sample_attn,
+            n_kv_heads=n_kv_heads,
+            head_dim=head_dim,
+        )
+        kernel_set = MlxAOTKernelSet()
+        selected_kernel_names = []
+        for spec in self._specs:
+            if not spec.is_enabled():
+                continue
+            kernel = spec.build(inputs)
+            if getattr(kernel, "enabled", False):
+                if not hasattr(kernel_set, spec.kernel_attr):
+                    raise ValueError(
+                        f"AOT kernel {spec.name} targets unknown kernel-set "
+                        f"attribute {spec.kernel_attr}"
+                    )
+                setattr(kernel_set, spec.kernel_attr, kernel)
+                selected_kernel_names.append(spec.name)
+        kernel_set.selected_kernel_names = tuple(selected_kernel_names)
+        if kernel_set.selected_kernel_names:
+            logger.info(
+                "MLX AOT kernels selected: %s",
+                ", ".join(kernel_set.selected_kernel_names),
+            )
+        return kernel_set
+
+
+def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel:
+    from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+        get_num_heads,
+    )
+
+    sample_attn = getattr(inputs.sample_attn, "_inner", inputs.sample_attn)
+    rope = getattr(sample_attn, "rope", None)
+    if rope is None or getattr(rope, "traditional", False):
+        return MlxAOTRoPEKernel()
+
+    rope_dim = int(getattr(rope, "dims", 0))
+    if rope_dim == 0:
+        return MlxAOTRoPEKernel()
+    if rope_dim != inputs.head_dim:
+        # AOT kernel currently requires rope_dim == head_dim.
+        return MlxAOTRoPEKernel()
+
+    base = float(getattr(rope, "base", 10000.0))
+    num_qo_heads = get_num_heads(sample_attn)
+    if num_qo_heads is None:
+        return MlxAOTRoPEKernel()
+    config = {
+        "head_dim": int(inputs.head_dim),
+        "rope_dim": rope_dim,
+        "num_qo_heads": int(num_qo_heads),
+        "num_kv_heads": int(inputs.n_kv_heads),
+    }
+    try:
+        rope_pool_fused = _load_metal_rope_pool_fused()
+    except Exception as exc:  # noqa: BLE001
+        logger.info(
+            "AOT Metal RoPE kernel not available (%s) - falling back to "
+            "mx.fast.rope.",
+            exc,
+        )
+        return MlxAOTRoPEKernel()
+
+    logger.info(
+        f"AOT Metal RoPE kernel ENABLED: head_dim={inputs.head_dim}, "
+        f"n_heads={config['num_qo_heads']}, n_kv={config['num_kv_heads']}, "
+        f"base={base}"
+    )
+    return MlxAOTRoPEKernel(
+        base=base,
+        config=config,
+        rope_pool_fused=rope_pool_fused,
+    )
+
+
+MLX_AOT_KERNEL_REGISTRY = MlxAOTKernelRegistry(
+    specs=(
+        MlxAOTKernelSpec(
+            name="metal_rope_pool_fused",
+            kernel_attr="rope",
+            is_enabled=lambda: envs.SGLANG_MLX_USE_CUSTOM_ROPE.get(),
+            build=_build_rope_kernel,
+        ),
+    )
+)
+
+
+@dataclass
+class MlxAOTRoPEContext:
+    kernel: MlxAOTRoPEKernel
+    kv_pool: Any
+    new_token_slots: Optional[mx.array] = None
+
+
+@dataclass
+class MlxAOTKernelContext:
+    rope: Optional[MlxAOTRoPEContext] = None
+
+    @classmethod
+    def from_decode(
+        cls,
+        *,
+        aot_kernels: MlxAOTKernelSet,
+        kv_pool: Any | None,
+        req_ids: list[str],
+        req_pool_idx: dict[str, int],
+        req_to_token_pool: Any | None,
+        layer_caches: list[list[ContiguousAttentionKVCache]],
+    ) -> MlxAOTKernelContext:
+        """Build optional AOT context for one batched decode step."""
+        if not aot_kernels.rope.enabled or kv_pool is None:
+            return cls()
+
+        new_token_slots = None
+        if req_to_token_pool is not None:
+            try:
+                slot_ids = []
+                for req_idx, req_id in enumerate(req_ids):
+                    pool_idx = req_pool_idx.get(req_id)
+                    if pool_idx is None:
+                        raise KeyError(req_id)
+                    slot = int(
+                        req_to_token_pool.req_to_token[
+                            pool_idx, layer_caches[0][req_idx].offset
+                        ].item()
+                    )
+                    slot_ids.append(slot)
+                new_token_slots = mx.array(slot_ids, dtype=mx.int32)
+            except Exception as exc:  # noqa: BLE001
+                logger.warning(
+                    "AOT RoPE: failed to resolve new-token slots (%s); "
+                    "falling back to RoPE-only for this decode step",
+                    exc,
+                )
+
+        return cls(
+            rope=MlxAOTRoPEContext(
+                kernel=aot_kernels.rope,
+                kv_pool=kv_pool,
+                new_token_slots=new_token_slots,
+            )
+        )

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)

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

@@ -0,0 +1,210 @@
+"""Lightweight ModelRunner stub for MLX on Apple Silicon.
+
+Skips PyTorch weight loading.  Creates only the CPU-side bookkeeping
+(req_to_token_pool, token_to_kv_pool_allocator) the scheduler needs.
+"""
+
+import logging
+from typing import Tuple
+
+import torch
+
+from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import (
+    MlxAuxiliaryStateReqToTokenPool,
+)
+from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
+from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
+from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class _DummyKVCache(KVCache):
+    """Scheduler-facing KV cache that allocates no GPU memory.
+
+    Satisfies the KVCache interface so that TokenToKVPoolAllocator can be
+    constructed, but every buffer access raises. The MLX backend manages
+    attention KV and auxiliary state internally.
+    """
+
+    def __init__(self, size: int, dtype: torch.dtype, device: str):
+        # Bypass KVCache.__init__ to avoid custom_mem_pool / memory_saver
+        # initialization that may touch CUDA APIs.
+        self.size = size
+        self.page_size = 1
+        self.dtype = dtype
+        self.store_dtype = dtype
+        self.device = device
+        self.layer_num = 0
+        self.start_layer = 0
+        self.end_layer = 0
+        self.mem_usage = 0
+        self.cpu_offloading_chunk_size = 8192
+        self.layer_transfer_counter = None
+        self.enable_custom_mem_pool = False
+        self.custom_mem_pool = None
+
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        raise RuntimeError("_DummyKVCache has no key buffer (MLX manages cache)")
+
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        raise RuntimeError("_DummyKVCache has no value buffer (MLX manages cache)")
+
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise RuntimeError("_DummyKVCache has no kv buffer (MLX manages cache)")
+
+    def set_kv_buffer(self, layer, loc, cache_k, cache_v) -> None:
+        raise RuntimeError("_DummyKVCache cannot set kv buffer (MLX manages cache)")
+
+    def get_kv_size_bytes(self):
+        return 0, 0
+
+
+class _DummyModel:
+    """Minimal stand-in so that `inspect.signature(model.forward)` and
+    `getattr(model, ...)` calls in ModelRunner.__init__ don't crash."""
+
+    @staticmethod
+    def forward():
+        pass
+
+
+class MlxModelRunnerStub(ModelRunner):
+    """ModelRunner that skips PyTorch weight loading and KV cache allocation.
+
+    Overrides both load_model() and initialize() so that no PyTorch model
+    weights are loaded and no large KV cache tensors are allocated.  Only
+    the minimal bookkeeping pools needed by the scheduler are created.
+    """
+
+    # No KV canary on the MLX path. The base ModelRunner installs it via
+    # install_canary() in its full initialize(), which this lightweight override
+    # skips. Downstream consumers (scheduler, cuda graph runner, speculative
+    # workers) all guard with `canary_manager is not None`, so default to None
+    # as a class attribute to keep those checks working instead of raising
+    # AttributeError.
+    canary_manager = None
+
+    # No prefill-aware SWA on the MLX path. The base ModelRunner derives this in
+    # its full initialize() from `model.is_prefill_aware_swa()`, which this
+    # lightweight override skips (and `_DummyModel` does not implement). The
+    # scheduler reads `model_runner.prefill_aware_swa` unconditionally when
+    # admitting a prefill batch, so default to False as a class attribute to keep
+    # that path working instead of raising AttributeError.
+    prefill_aware_swa = False
+
+    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."""
+        logger.info(
+            "MLX stub: skipping PyTorch model weight loading "
+            "(inference runs through MLX)"
+        )
+
+        self.model = _DummyModel()
+
+        self.sliding_window_size = None
+        if (
+            self.model_config.is_hybrid_swa
+            and self.model_config.sliding_window_size is not None
+        ):
+            self.sliding_window_size = self.model_config.sliding_window_size
+        elif self.model_config.attention_chunk_size is not None:
+            self.sliding_window_size = self.model_config.attention_chunk_size
+
+        self.dtype = self.model_config.dtype
+        self.weight_load_mem_usage = 0
+
+    def initialize(self):
+        """Lightweight initialize that skips heavy PyTorch setup.
+
+        Creates minimal req_to_token_pool and token_to_kv_pool_allocator
+        with a dummy KV cache (zero GPU memory) so the scheduler works.
+        """
+        from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=self.server_args.enable_memory_saver
+        )
+
+        # Load model (sets metadata only)
+        self.sampler = None
+        self.load_model()
+
+        # Layer metadata
+        model_num_layers = max(
+            self.model_config.num_hidden_layers,
+            self.model_config.num_attention_layers,
+        )
+        self.start_layer = 0
+        self.end_layer = model_num_layers
+        self.num_effective_layers = model_num_layers
+
+        # KV cache dtype
+        self.kv_cache_dtype = self.dtype
+
+        # 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,
+        )
+        self.is_hybrid_swa = False
+
+        # Create minimal pools
+        if self.mambaish_config is not None:
+            auxiliary_state_size = self.server_args.max_mamba_cache_size
+            if auxiliary_state_size is None:
+                auxiliary_state_size = self.max_running_requests * 4
+            self.req_to_token_pool = MlxAuxiliaryStateReqToTokenPool(
+                size=self.max_running_requests,
+                max_context_len=self.model_config.context_len,
+                device="cpu",
+                enable_memory_saver=False,
+                auxiliary_state_size=auxiliary_state_size,
+            )
+        else:
+            self.req_to_token_pool = ReqToTokenPool(
+                size=self.max_running_requests,
+                max_context_len=self.model_config.context_len,
+                device="cpu",
+                enable_memory_saver=False,
+            )
+
+        dummy_kv = _DummyKVCache(
+            size=self.max_total_num_tokens,
+            dtype=self.kv_cache_dtype,
+            device="cpu",
+        )
+        self.token_to_kv_pool = dummy_kv
+        self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
+            size=self.max_total_num_tokens,
+            dtype=self.kv_cache_dtype,
+            device="cpu",
+            kvcache=dummy_kv,
+            need_sort=False,
+        )
+
+        # No CUDA graphs, no attention backend
+        self.decode_cuda_graph_runner = None
+        self.graph_mem_usage = 0
+        self.attn_backend = None
+
+        logger.info(
+            f"MLX stub: initialized minimal pools "
+            f"(max_total_num_tokens={self.max_total_num_tokens}, "
+            f"max_running_requests={self.max_running_requests}, "
+            f"zero GPU KV cache allocation)"
+        )
+
+    def alloc_memory_pool(self, memory_pool_config=None):
+        """No-op: MLX manages its own KV cache."""
+        pass

python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py

@@ -0,0 +1,573 @@
+"""Path B fusion for SwitchGLU: gate gather_qmv with silu(gate) * x_up epilogue.
+
+Why this exists
+---------------
+The existing `FusedSwitchUpGate` (fused_switch_glu.py) concatenates up_proj
+and gate_proj weights along the output dim and runs one gather_qmm. That saves
+one kernel launch per layer but doubles the matmul's output dim, which pushes
+MLX's quantized GEMV into a worse tile/occupancy config. At bs >= 4 on
+Qwen3-30B-A3B-4bit this is a net regression (~2% slower at bs=32).
+
+Path B keeps up_proj and gate_proj separate (matmul kernels see their natural
+N — no tile regression) and instead fuses the *activation* into the gate
+matmul. Concretely:
+
+    Baseline (3 kernels per MoE layer in the swiglu front-half):
+        x_up   = gather_qmm(x, W_up)
+        x_gate = gather_qmm(x, W_gate)
+        out    = silu(x_gate) * x_up         # 1 compiled kernel via mlx_lm.swiglu
+
+    Path B (2 kernels):
+        x_up = gather_qmm(x, W_up)
+        out  = fused_gate_qmv_silu_mul(x, W_gate, ..., x_up)
+                                              # one custom Metal kernel
+
+This removes one kernel launch per MoE layer with no change in matmul kernel
+shapes. The measured end to end decode impact is within run to run noise (bs=1,
+K=12 interleaved trials on Qwen3-30B-A3B-4bit: on minus off 0.4%, a quarter of
+the noise band), so v1 lands off by default as a correct fusion substrate, not
+a measured speedup.
+
+Scope of v1
+-----------
+Targets the configuration shared by Qwen3-30B-A3B-4bit and Qwen1.5-MoE-A2.7B-4bit:
+- bits=4, mode='affine', group_size=64
+- K (input_dim) divisible by 512  (Qwen3: 2048, Qwen1.5: 2048)
+- N (output_dim) divisible by 8   (Qwen3: 768, Qwen1.5: 1408)
+- Scales/biases dtype matches the input dtype (bf16 or fp16)
+
+Anything outside that falls back to the unfused mlx_lm path.
+
+When the fast `gather_qmv` from #22283 lands, the qmv inner loop here can be
+replaced by a call into it; the epilogue (silu * x_up) doesn't change.
+"""
+
+from __future__ import annotations
+
+import logging
+import weakref
+
+import mlx.core as mx
+import mlx.nn as nn
+
+logger = logging.getLogger(__name__)
+
+# Constants matching MLX's affine_qmv_fast for bits=4, group_size=64.
+# Lifted directly from mlx/include/mlx/backend/metal/kernels/quantized.h
+# (qmv_fast_impl), so the inner-loop layout matches MLX's own gather_qmm in
+# that regime.
+_BITS = 4
+_GROUP_SIZE = 64
+_SIMD_SIZE = 32
+_PACK_FACTOR = 8  # 32 / bits
+_BYTES_PER_PACK = 4  # sizeof(uint32_t)
+_PACKS_PER_THREAD = 2
+_NUM_SIMDGROUPS = 2
+_RESULTS_PER_SIMDGROUP = 4
+_VALUES_PER_THREAD = _PACK_FACTOR * _PACKS_PER_THREAD  # 16
+_BLOCK_SIZE = _VALUES_PER_THREAD * _SIMD_SIZE  # 512
+_ROWS_PER_TG = _NUM_SIMDGROUPS * _RESULTS_PER_SIMDGROUP  # 8
+
+
+# Metal source for the fused kernel.
+# Body only — mx.fast.metal_kernel auto-generates the kernel signature
+# based on input_names / output_names and the template params below.
+_KERNEL_SOURCE = r"""
+    // Mirrors qmv_fast_impl<T, group_size=64, bits=4> from MLX's quantized.h
+    // with a silu(result) * x_up write epilogue.
+    //
+    // Inputs:
+    //   x       [M_tok, K]                — pre-gather activations (T)
+    //   w       [E, N, K * 4 / 32]        — packed 4-bit weights (uint32)
+    //   s       [E, N, K / GROUP_SIZE]    — affine scales (T)
+    //   b       [E, N, K / GROUP_SIZE]    — affine biases (T)
+    //   idx     [M_tok * TOPK]            — expert per (token, topk) pair (uint32)
+    //   x_up    [M_tok * TOPK, N]         — precomputed up output (T)
+    // Output:
+    //   y       [M_tok * TOPK, N]         — silu(gate_qmv(x)) * x_up
+
+    constexpr int BITS = 4;
+    constexpr int GROUP_SIZE = 64;
+    constexpr int SIMD_SIZE = 32;
+    constexpr int PACK_FACTOR = 8;
+    constexpr int BYTES_PER_PACK = 4;
+    constexpr int PACKS_PER_THREAD = 2;
+    constexpr int NUM_SIMDGROUPS = 2;
+    constexpr int RESULTS_PER_SIMDGROUP = 4;
+    constexpr int VALUES_PER_THREAD = PACK_FACTOR * PACKS_PER_THREAD;   // 16
+    constexpr int BLOCK_SIZE = VALUES_PER_THREAD * SIMD_SIZE;            // 512
+    constexpr int SCALE_STEP_PER_THREAD = GROUP_SIZE / VALUES_PER_THREAD; // 4
+
+    // Compile-time problem dims (template params)
+    constexpr int K = IN_VEC_SIZE;
+    constexpr int N = OUT_VEC_SIZE;
+    constexpr int TOPK = TOP_K;
+
+    constexpr int in_vec_size_w = K * BYTES_PER_PACK / PACK_FACTOR;    // K/2 bytes per row
+    constexpr int in_vec_size_g = K / GROUP_SIZE;                       // groups per row
+
+    // tid.x = (token, topk) pair index in M_tok*TOPK
+    // tid.y = output row block index (each block writes 8 rows)
+    uint mt = threadgroup_position_in_grid.x;
+    uint out_row_block = threadgroup_position_in_grid.y;
+    uint simd_gid = simdgroup_index_in_threadgroup;
+    uint simd_lid = thread_index_in_simdgroup;
+
+    uint m = mt / TOPK;                      // index into pre-gather x
+    uint e = idx[mt];                        // expert id
+
+    // Base pointers for this (expert) row block.
+    // Weights / scales / biases live in [E, N, ...] tensors.
+    const device uint8_t* ws = (const device uint8_t*)(w
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_w / BYTES_PER_PACK));
+    const device T* scales_p = s
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_g);
+    const device T* biases_p = b
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_g);
+
+    // out_row indexes the first of RESULTS_PER_SIMDGROUP rows this simdgroup handles.
+    int out_row = out_row_block * NUM_SIMDGROUPS * RESULTS_PER_SIMDGROUP
+                + simd_gid * RESULTS_PER_SIMDGROUP;
+
+    ws += out_row * in_vec_size_w + simd_lid * PACKS_PER_THREAD * BYTES_PER_PACK;
+    scales_p += out_row * in_vec_size_g + simd_lid / SCALE_STEP_PER_THREAD;
+    biases_p += out_row * in_vec_size_g + simd_lid / SCALE_STEP_PER_THREAD;
+
+    // Pre-gather x: shape [M_tok, K]
+    const device T* x_p = x + uint64_t(m) * uint64_t(K)
+                            + uint64_t(simd_lid) * uint64_t(VALUES_PER_THREAD);
+
+    float result[RESULTS_PER_SIMDGROUP] = {0, 0, 0, 0};
+    thread float x_thread[VALUES_PER_THREAD];
+
+    // Outer loop over K in BLOCK_SIZE-wide chunks.
+    // K % BLOCK_SIZE == 0 is required (checked in Python wrapper).
+    for (int k = 0; k < K; k += BLOCK_SIZE) {
+        // --- load_vector for bits=4 ---
+        float sum = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i += 4) {
+            float a0 = float(x_p[i]);
+            float a1 = float(x_p[i + 1]);
+            float a2 = float(x_p[i + 2]);
+            float a3 = float(x_p[i + 3]);
+            sum += a0 + a1 + a2 + a3;
+            x_thread[i]     = a0;
+            x_thread[i + 1] = a1 / 16.0f;
+            x_thread[i + 2] = a2 / 256.0f;
+            x_thread[i + 3] = a3 / 4096.0f;
+        }
+
+        // For each of the 4 output rows this simdgroup is responsible for...
+        for (int row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
+            const device uint16_t* ws_u16 =
+                (const device uint16_t*)(ws + row * in_vec_size_w);
+            float scale_v = float(scales_p[row * in_vec_size_g]);
+            float bias_v  = float(biases_p[row * in_vec_size_g]);
+
+            // --- qdot for bits=4, values_per_thread=16 ---
+            float accum = 0;
+            for (int i = 0; i < VALUES_PER_THREAD / 4; i++) {
+                uint16_t packed = ws_u16[i];
+                accum += (x_thread[4 * i]     * float(packed & 0x000f) +
+                          x_thread[4 * i + 1] * float(packed & 0x00f0) +
+                          x_thread[4 * i + 2] * float(packed & 0x0f00) +
+                          x_thread[4 * i + 3] * float(packed & 0xf000));
+            }
+            result[row] += scale_v * accum + sum * bias_v;
+        }
+
+        ws        += BLOCK_SIZE * BYTES_PER_PACK / PACK_FACTOR;  // += 256 bytes
+        scales_p  += BLOCK_SIZE / GROUP_SIZE;                    // += 8 groups
+        biases_p  += BLOCK_SIZE / GROUP_SIZE;
+        x_p       += BLOCK_SIZE;
+    }
+
+    // Write epilogue: simd-sum across lanes, then silu(gate) * x_up.
+    device T* y_p = y + uint64_t(mt) * uint64_t(N) + uint64_t(out_row);
+    const device T* x_up_p = x_up + uint64_t(mt) * uint64_t(N) + uint64_t(out_row);
+
+    for (int row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
+        float gate_v = simd_sum(result[row]);
+        if (simd_lid == 0) {
+            // silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+            float silu_v = gate_v / (1.0f + metal::precise::exp(-gate_v));
+            y_p[row] = T(silu_v * float(x_up_p[row]));
+        }
+    }
+"""
+
+
+# Build kernel lazily so import-time on non-MLX systems doesn't fail.
+_kernel_cache: dict = {}
+
+# AOT pre-compile cache: tracks which (dtype, K, N, T) tuples have been warmed.
+# MLX specializes the Metal kernel on template args at first dispatch, so the
+# first call per unique tuple pays a ~3ms compile cost. Pre-warming at patch
+# time moves that cost out of the first forward pass.
+_aot_warmed: set = set()
+
+
+def _get_kernel(dtype: mx.Dtype):
+    """Return a compiled mx.fast.metal_kernel for the given input dtype.
+
+    Kept per-dtype because mx.fast.metal_kernel specializes on template args
+    at first call, and we want clean separation between fp16 / bf16 variants.
+    """
+    if dtype not in _kernel_cache:
+        # Strip the `mlx.core.` prefix and any dots from the dtype repr so the
+        # kernel name is a valid C identifier (Metal's host_name attribute and
+        # function-name slot don't accept '.').
+        dtype_tag = str(dtype).replace("mlx.core.", "").replace(".", "_")
+        _kernel_cache[dtype] = mx.fast.metal_kernel(
+            name=f"affine_gather_qmv_silu_mul_4bit_gs64_{dtype_tag}",
+            input_names=["x", "w", "s", "b", "idx", "x_up"],
+            output_names=["y"],
+            source=_KERNEL_SOURCE,
+        )
+    return _kernel_cache[dtype]
+
+
+def fused_gate_qmv_silu_mul(
+    x: mx.array,
+    gate_w: mx.array,
+    gate_s: mx.array,
+    gate_b: mx.array,
+    indices: mx.array,
+    x_up: mx.array,
+) -> mx.array:
+    """Compute ``silu(gather_qmm(x, W_gate)) * x_up`` in one kernel.
+
+    Shapes:
+        x       : (..., 1, 1, K)               input activations (pre-gather)
+        gate_w  : (E, N, K // PACK_FACTOR)     packed 4-bit weights
+        gate_s  : (E, N, K // GROUP_SIZE)      affine scales
+        gate_b  : (E, N, K // GROUP_SIZE)      affine biases
+        indices : (..., T)                     expert per (token, topk) pair
+        x_up    : (..., T, 1, N)               pre-computed up output
+        y       : (..., T, 1, N)               returned
+
+    Numerical contract: equivalent within floating point tolerance to::
+
+        x_gate = mx.gather_qmm(x, gate_w, gate_s, gate_b, rhs_indices=indices,
+                               transpose=True, group_size=GROUP_SIZE, bits=BITS,
+                               mode='affine')
+        y = nn.silu(x_gate) * x_up
+
+    up to floating-point ordering of accumulations (which matches MLX's own
+    qmv_fast_impl exactly).
+    """
+    # Validate the regime this kernel supports. Outside it, the caller should
+    # fall back to the unfused MLX path.
+    K = gate_s.shape[-1] * _GROUP_SIZE
+    N = gate_w.shape[-2]
+    if K % _BLOCK_SIZE != 0:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: K={K} not divisible by {_BLOCK_SIZE}. "
+            f"Use the unfused path."
+        )
+    if N % _ROWS_PER_TG != 0:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: N={N} not divisible by {_ROWS_PER_TG}."
+        )
+    # Sanity: scales/biases dtype must match x dtype for in-kernel float() conversion.
+    if gate_s.dtype != x.dtype or gate_b.dtype != x.dtype:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: dtype mismatch x={x.dtype} "
+            f"s={gate_s.dtype} b={gate_b.dtype}"
+        )
+
+    # Shape handling: x always has K as its last axis. M_tok is the number of
+    # distinct pre-gather tokens (= x.size // K). T is the top_k axis carried
+    # by `indices`. In the unsorted SwitchGLU path x has shape (B, 1, 1, K) and
+    # indices is (B, T); in the sorted path x has shape (B*T, 1, K) and
+    # indices is (B*T, 1) after our reshape. Either way: M_tok * T == idx.size
+    # and M_tok * K == x.size.
+    assert x.shape[-1] == K, f"x last dim {x.shape[-1]} != K={K}"
+    M_tok = x.size // K
+    T = indices.shape[-1]
+    assert (
+        M_tok * T == indices.size
+    ), f"M_tok({M_tok}) * T({T}) != indices.size({indices.size})"
+    x_flat = x.reshape(M_tok, K)
+    idx_flat = indices.reshape(M_tok * T)
+    if idx_flat.dtype != mx.uint32:
+        idx_flat = idx_flat.astype(mx.uint32)
+
+    # x_up has N as its last axis and total size M_tok * T * N. The singleton
+    # rank dims (1 or 2 of them) get folded away by reshape.
+    assert (
+        x_up.shape[-1] == N and x_up.size == M_tok * T * N
+    ), f"x_up shape {x_up.shape} does not match M_tok({M_tok})*T({T})*N({N})"
+    x_up_flat = x_up.reshape(M_tok * T, N)
+
+    kernel = _get_kernel(x.dtype)
+    (y_flat,) = kernel(
+        inputs=[x_flat, gate_w, gate_s, gate_b, idx_flat, x_up_flat],
+        template=[
+            ("T", x.dtype),
+            ("IN_VEC_SIZE", K),
+            ("OUT_VEC_SIZE", N),
+            ("TOP_K", T),
+        ],
+        # grid is in *threads*, not threadgroups: total threads = product of
+        # (grid_x, grid_y, grid_z). One threadgroup processes one (mt, row_block).
+        # Threadgroup is 64 = 2 simdgroups × 32 lanes.
+        grid=(M_tok * T * 64, N // _ROWS_PER_TG, 1),
+        threadgroup=(64, 1, 1),
+        output_shapes=[(M_tok * T, N)],
+        output_dtypes=[x.dtype],
+    )
+
+    # Reshape to x_up's shape, which is exactly what self.activation(x_up,
+    # x_gate) would have returned (silu*mul is shape-preserving).
+    return y_flat.reshape(x_up.shape)
+
+
+def _aot_warm_kernel(switch_mlp, top_k: int) -> None:
+    """Pre-compile the fused kernel for the (dtype, K, N, T) tuples this
+    layer will dispatch at runtime.
+
+    MLX's mx.fast.metal_kernel specializes Metal source on template args
+    at first dispatch (one Metal compile per unique tuple, ~3ms each).
+    Issuing one dummy dispatch per shape moves that compile out of the
+    first forward pass and into model init. The module-level _aot_warmed
+    set means only the first layer per shape actually compiles; the
+    remaining 47 hit the cache and no-op.
+
+    Warms both the unsorted decode shape (T=top_k) and the sorted
+    large-batch shape (T=1) used by the gather-sort path.
+    """
+    gate = switch_mlp.gate_proj
+    K = gate.scales.shape[-1] * _GROUP_SIZE
+    N = gate.weight.shape[-2]
+    dtype = gate.scales.dtype
+    for T in (top_k, 1):
+        key = (dtype, K, N, T)
+        if key in _aot_warmed:
+            continue
+        M_tok = 1
+        x_dummy = mx.zeros((M_tok, 1, 1, K), dtype=dtype)
+        idx_dummy = mx.zeros((M_tok, T), dtype=mx.uint32)
+        x_up_dummy = mx.zeros((M_tok, T, 1, N), dtype=dtype)
+        out = fused_gate_qmv_silu_mul(
+            x_dummy,
+            gate["weight"],
+            gate["scales"],
+            gate.get("biases"),
+            idx_dummy,
+            x_up_dummy,
+        )
+        mx.eval(out)
+        _aot_warmed.add(key)
+        logger.info(
+            "Path B AOT: warmed fused kernel dtype=%s K=%d N=%d T=%d",
+            dtype,
+            K,
+            N,
+            T,
+        )
+
+
+def can_fuse(switch_mlp) -> bool:
+    """Cheap structural check: does this SwitchGLU match the Path B v1 regime?"""
+    try:
+        from mlx_lm.models.switch_layers import (
+            QuantizedSwitchLinear,
+            SwiGLU,
+            SwitchGLU,
+        )
+    except ImportError:
+        return False
+    up = switch_mlp.up_proj
+    gate = switch_mlp.gate_proj
+    if not isinstance(up, QuantizedSwitchLinear) or not isinstance(
+        gate, QuantizedSwitchLinear
+    ):
+        return False
+    # A model that overrides SwitchGLU.__call__ runs a custom forward, but the
+    # patch installs a subclass __call__ that imposes the stock semantics
+    # fused_forward reimplements, silently bypassing the override. Decline when
+    # the forward is not the stock SwitchGLU.__call__ (evaluated at patch time,
+    # before the class swap, so this sees the model's real class).
+    if type(switch_mlp).__call__ is not SwitchGLU.__call__:
+        return False
+    # fused_forward bakes silu into both the kernel and its fallback, and a
+    # swapped activation= leaves __call__ stock, so the check above cannot see
+    # it. Exact type, fail closed: a SwiGLU subclass may change the math.
+    if type(getattr(switch_mlp, "activation", None)) is not SwiGLU:
+        return False
+    # Learned per-expert bias, added after the matmul in
+    # QuantizedSwitchLinear.__call__ as ``x + bias[indices]`` whenever
+    # ``"bias" in self``. This is the affine learned bias, distinct from the
+    # quant ``biases`` (zero-points) the kernel already consumes. The fused
+    # kernel recomputes the gate matmul in-register and has no slot for the
+    # learned bias, so a gate carrying one would silently drop it. up_proj runs
+    # its normal path (its bias, if any, is already in x_up), so only the gate
+    # is at risk: fall back to the unfused path when the gate has a learned bias.
+    if "bias" in gate:
+        return False
+    if up.bits != 4 or up.group_size != 64 or up.mode != "affine":
+        return False
+    if gate.bits != 4 or gate.group_size != 64 or gate.mode != "affine":
+        return False
+    if up.biases is None or gate.biases is None:
+        return False
+    # The kernel reads gate scales/biases as the activation dtype (in-kernel
+    # float() then writes back as the activation dtype), so the runtime
+    # precondition in fused_gate_qmv_silu_mul requires a single shared float
+    # dtype across the quant params. Check it here so a mismatch falls through
+    # to the unfused path instead of raising at forward. The packed weight is
+    # uint32 and intentionally excluded from the dtype contract.
+    param_dtype = gate.scales.dtype
+    if (
+        gate.biases.dtype != param_dtype
+        or up.scales.dtype != param_dtype
+        or up.biases.dtype != param_dtype
+    ):
+        return False
+    # Validate from the gate projection: fused_gate_qmv_silu_mul recomputes the
+    # gate matmul and keys K/N off gate dims, so the gate is the operand that
+    # must satisfy the tiling constraints (equivalent to up only while
+    # up.shape == gate.shape).
+    K = gate.scales.shape[-1] * _GROUP_SIZE
+    N = gate.weight.shape[-2]
+    if K % _BLOCK_SIZE != 0 or N % _ROWS_PER_TG != 0:
+        return False
+    return True
+
+
+_fallback_warned = False
+
+
+def _fused_gate_or_fallback(gate_proj, x, idx, x_up, sorted_indices=False):
+    """silu(gate_qmv(x)) * x_up via the fused kernel; on ValueError fall back to
+    the unfused gate projection. gather_qmm tolerates the activation dtype the
+    fused kernel rejects, which can_fuse cannot pre-check at patch time. Warns
+    once.
+    """
+    # Kernel layout only: the Metal kernel reads T from indices.shape[-1], and
+    # _gather_sort folded top_k into M_tok, so the sorted path needs an explicit
+    # T=1 axis. The fallback must not see it: gather_qmm broadcasts an (M_tok, 1)
+    # index against sorted x's (M_tok,) batch dim into an M_tok x M_tok cross
+    # product.
+    gate_idx = idx.reshape(-1, 1) if sorted_indices else idx
+    gw = gate_proj["weight"]
+    gs = gate_proj["scales"]
+    gb = gate_proj.get("biases")
+    try:
+        return fused_gate_qmv_silu_mul(x, gw, gs, gb, gate_idx, x_up)
+    except ValueError as e:
+        global _fallback_warned
+        if not _fallback_warned:
+            logger.warning(
+                "Path B: fused gate kernel declined inputs (%s); using the "
+                "unfused gate path for this and matching calls.",
+                e,
+            )
+            _fallback_warned = True
+        # Reference expression by construction: the same projection call with
+        # the same flat idx up_proj/down_proj receive.
+        return nn.silu(gate_proj(x, idx, sorted_indices=sorted_indices)) * x_up
+
+
+class FusedSwitchSwiGLU(nn.Module):
+    """SwitchGLU forward with Path B fusion installed.
+
+    Wraps an existing SwitchGLU instance. Reads up_proj / gate_proj weights
+    directly (no concatenation), runs up_proj as usual, then calls the fused
+    kernel for ``silu(gate_qmv(x)) * x_up`` in one shot.
+
+    Replaces ``switch_mlp.__call__`` via patch_switch_glu_with_fused_swiglu;
+    SwitchGLU.up_proj / gate_proj / activation are *not* replaced and remain
+    available for fallback paths (e.g. sorted-indices large-batch case).
+    """
+
+    def __init__(self, switch_mlp):
+        super().__init__()
+        # Weak ref: sw stores the bound fused_forward, so a strong ref here would
+        # cycle (sw -> method -> self -> sw). sw outlives every call into it.
+        self._switch_mlp = weakref.proxy(switch_mlp)
+
+    def fused_forward(self, x, indices):
+        """Same contract as SwitchGLU.__call__ but with fused activation."""
+        from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort
+
+        sw = self._switch_mlp
+        x = mx.expand_dims(x, (-2, -3))
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+
+        x_up = sw.up_proj(x, idx, sorted_indices=do_sort)
+
+        swiglu = _fused_gate_or_fallback(
+            sw.gate_proj, x, idx, x_up, sorted_indices=do_sort
+        )
+
+        out = sw.down_proj(swiglu, idx, sorted_indices=do_sort)
+
+        if do_sort:
+            out = _scatter_unsort(out, inv_order, indices.shape)
+        return out.squeeze(-2)
+
+
+def patch_switch_glu_with_fused_swiglu(model) -> int:
+    """Install Path B on every eligible SwitchGLU in the model.
+
+    Replaces ``switch_mlp.__call__`` with FusedSwitchSwiGLU.fused_forward.
+    Leaves ``up_proj``, ``gate_proj``, ``activation`` in place so the original
+    code path is still reachable for sorted-large-batch (handled internally) and
+    for any callers that bypass the patched __call__.
+
+    Returns number of layers patched.
+    """
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    patched = 0
+    for layer in model.model.layers:
+        mlp = getattr(layer, "mlp", None)
+        if mlp is None:
+            continue
+        sw = getattr(mlp, "switch_mlp", None)
+        if not isinstance(sw, SwitchGLU):
+            continue
+        if not can_fuse(sw):
+            continue
+        # Idempotent: skip if already patched.
+        if getattr(sw, "_path_b_installed", False):
+            continue
+        # AOT: pre-compile the fused kernel for this layer's shapes. top_k
+        # lives on the parent MoE block (e.g. Qwen3MoeSparseMoeBlock.top_k).
+        # If absent, skip warming and fall back to lazy JIT on first dispatch.
+        top_k = getattr(mlp, "top_k", None)
+        if top_k is not None:
+            _aot_warm_kernel(sw, int(top_k))
+        # One-off SwitchGLU subclass rather than rewriting up_proj/gate_proj: the
+        # activation fusion folds silu(gate)*x_up into the gate matmul, which has
+        # to intercept the forward (the projection level can't express it).
+        # can_fuse declines a non-stock __call__, so a customized forward falls
+        # back unpatched. Python resolves __call__ on the type, so swap
+        # sw.__class__ to a subclass; cache it on the exact class (cls.__dict__,
+        # not hasattr which walks the MRO) so a SwitchGLU subclass gets its own
+        # entry instead of being downcast to the base.
+        sw._path_b_call = FusedSwitchSwiGLU(sw).fused_forward
+        cls = type(sw)
+        if "_PathBSubclass" not in cls.__dict__:
+            cls._PathBSubclass = type(
+                f"{cls.__name__}_PathB",
+                (cls,),
+                {"__call__": lambda self, *a, **kw: self._path_b_call(*a, **kw)},
+            )
+        sw.__class__ = cls._PathBSubclass
+        sw._path_b_installed = True
+        patched += 1
+
+    if patched == 0:
+        logger.warning(
+            "patch_switch_glu_with_fused_swiglu: no eligible SwitchGLU found"
+        )
+    else:
+        logger.info(f"patch_switch_glu_with_fused_swiglu: patched {patched} layers")
+    return patched

python/sglang/srt/hardware_backend/mlx/moe/tests/test_fused_swiglu.py

@@ -0,0 +1,464 @@
+"""Numerical equivalence and eligibility tests for the Path B fused swiglu kernel.
+
+Two groups:
+  * Model-based equivalence (``@requires_model``): loads a small MoE model, runs
+    the fused gate_qmv + silu + ×x_up kernel against the unfused reference
+    (``mx.gather_qmm`` + ``nn.silu(gate) * x_up``) on both the unsorted and
+    sorted paths. Gated by SGLANG_MLX_TEST_MODEL so CI hosts without a model
+    cache skip them.
+  * Synthetic eligibility (no model, MLX only): the learned-bias fallback. The
+    fused kernel recomputes the gate matmul and has no slot for the per-expert
+    learned bias QuantizedSwitchLinear adds after the matmul, so ``can_fuse``
+    must exclude a gate carrying one, and the patch must leave such a layer
+    unfused. These run whenever MLX is importable.
+"""
+
+import os
+
+import pytest
+
+mx = pytest.importorskip("mlx.core")
+
+
+# Model-based tests need a real checkpoint; synthetic tests below do not.
+requires_model = pytest.mark.skipif(
+    not os.environ.get("SGLANG_MLX_TEST_MODEL"),
+    reason="Set SGLANG_MLX_TEST_MODEL to a HuggingFace model id to enable",
+)
+
+
+def _max_rel_diff(a, b):
+    diff = mx.abs(a.astype(mx.float32) - b.astype(mx.float32))
+    max_abs = diff.max().item()
+    ref_max = mx.abs(a.astype(mx.float32)).max().item()
+    return max_abs, max_abs / max(ref_max, 1e-9)
+
+
+@requires_model
+def test_fused_gate_qmv_silu_mul_matches_unfused():
+    """Kernel output matches ``nn.silu(gate_qmv) * x_up`` within bf16 ULP."""
+    import mlx.nn as nn
+    from mlx_lm import load
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        can_fuse,
+        fused_gate_qmv_silu_mul,
+    )
+
+    model, _ = load(os.environ["SGLANG_MLX_TEST_MODEL"])
+    sw = model.model.layers[0].mlp.switch_mlp
+    assert can_fuse(sw), "layer 0 not eligible for fused swiglu"
+
+    up = sw.up_proj
+    gate = sw.gate_proj
+    in_dim = up.scales.shape[-1] * up.group_size
+    out_dim = up.weight.shape[-2]
+    num_experts = up.weight.shape[0]
+    dtype = up.scales.dtype
+
+    # Two batch sizes both take the unsorted path (indices.size < 64).
+    for B, TOPK in [(1, 8), (4, 8)]:
+        x = mx.random.normal(shape=(B, 1, 1, in_dim)).astype(dtype)
+        indices = mx.random.randint(0, num_experts, shape=(B, TOPK)).astype(mx.uint32)
+
+        x_up = up(x, indices, sorted_indices=False)
+        x_gate = gate(x, indices, sorted_indices=False)
+        y_ref = nn.silu(x_gate) * x_up
+
+        y_fused = fused_gate_qmv_silu_mul(
+            x, gate["weight"], gate["scales"], gate.get("biases"), indices, x_up
+        )
+        mx.eval(y_ref, y_fused)
+
+        assert y_ref.shape == y_fused.shape
+
+        max_abs, rel = _max_rel_diff(y_ref, y_fused)
+        # 2 % relative covers ~2 bf16 ULPs at typical activation magnitudes;
+        # the kernel's fp32 accumulation order matches MLX's qmv_fast_impl so
+        # most elements should land within 1 ULP.
+        assert rel < 2e-2, f"B={B} TOPK={TOPK}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+@requires_model
+def test_patched_switchglu_matches_unpatched():
+    """Full SwitchGLU forward equivalence on both sorted and unsorted paths."""
+    from mlx_lm import load
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    model, _ = load(os.environ["SGLANG_MLX_TEST_MODEL"])
+    sw = model.model.layers[0].mlp.switch_mlp
+    in_dim = sw.up_proj.scales.shape[-1] * sw.up_proj.group_size
+    num_experts = sw.up_proj.weight.shape[0]
+    dtype = sw.up_proj.scales.dtype
+
+    cases = []
+    # B=2 TOPK=8 -> indices.size=16 < 64 -> unsorted
+    # B=8 TOPK=8 -> indices.size=64 -> sorted
+    for B, TOPK, label in [(2, 8, "unsorted"), (8, 8, "sorted")]:
+        x = mx.random.normal(shape=(B, in_dim)).astype(dtype)
+        indices = mx.random.randint(0, num_experts, shape=(B, TOPK)).astype(mx.uint32)
+        out_ref = sw(x, indices)
+        mx.eval(out_ref)
+        cases.append((label, x, indices, out_ref))
+
+    n_patched = patch_switch_glu_with_fused_swiglu(model)
+    assert n_patched > 0, "no SwitchGLU layers were patched"
+
+    for label, x, indices, out_ref in cases:
+        out_fused = sw(x, indices)
+        mx.eval(out_fused)
+        max_abs, rel = _max_rel_diff(out_ref, out_fused)
+        # 5 % is generous; in practice we see <0.6 % on 48-layer Qwen3-MoE.
+        # The looser bound here absorbs cross-layer ULP propagation through
+        # down_proj's quantized matmul.
+        assert rel < 5e-2, f"full forward {label}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+# Learned-bias fallback (synthetic, no model): a gate with a learned bias must
+# not fuse, since the kernel has no slot for the bias added after the matmul.
+def _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias):
+    """Small quantized SwitchGLU; gate carries a learned bias iff gate_bias.
+
+    in_dim=512 keeps K%512==0 and hidden%8==0, inside the Path B v1 regime, so
+    the bias-free build is genuinely fusion-eligible (the True control).
+    """
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    sw.up_proj = sw.up_proj.to_quantized(group_size=64, bits=4, mode="affine")
+    sw.down_proj = sw.down_proj.to_quantized(group_size=64, bits=4, mode="affine")
+    gate = sw.gate_proj
+    if gate_bias:
+        # Learned per-expert bias (E, N), nonzero so dropping it would change
+        # the result. to_quantized copies it into the QuantizedSwitchLinear.
+        gate.bias = mx.random.normal((n_experts, hidden)) * 0.1
+    sw.gate_proj = gate.to_quantized(group_size=64, bits=4, mode="affine")
+    return sw
+
+
+def test_can_fuse_excludes_learned_gate_bias():
+    """can_fuse: False for a gate with a learned bias, True when bias-free."""
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import can_fuse
+
+    sw_free = _quantized_switch_glu(512, 64, 8, gate_bias=False)
+    sw_bias = _quantized_switch_glu(512, 64, 8, gate_bias=True)
+    assert "bias" not in sw_free.gate_proj
+    assert "bias" in sw_bias.gate_proj
+    assert can_fuse(sw_free) is True, "bias-free gate in regime should fuse"
+    assert can_fuse(sw_bias) is False, "gate with learned bias must fall back"
+
+
+def test_patch_falls_back_on_gate_bias():
+    """Patching a biased-gate SwitchGLU is a no-op; the forward stays bias-correct."""
+    import types
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 8, 4, 2  # 2*4=8 < 64 -> unsorted
+    sw = _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias=True)
+
+    x = mx.random.normal((B, in_dim))
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_before = sw(x, indices)
+    mx.eval(out_before)
+
+    # Minimal model stand-in: the patch walks model.model.layers[*].mlp.switch_mlp.
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+
+    n_patched = patch_switch_glu_with_fused_swiglu(model)
+    assert n_patched == 0, "biased gate must not be patched"
+
+    out_after = sw(x, indices)
+    mx.eval(out_after)
+    d = mx.abs(out_before.astype(mx.float32) - out_after.astype(mx.float32))
+    diff = d.max().item()
+    assert diff == 0.0, f"forward changed after (no-op) patch: max|delta|={diff:.3e}"
+
+
+# Model-free numerical equivalence + non-stock-forward guard: the central
+# correctness check, runs without a model download (skips where Metal is absent).
+def test_fused_matches_unfused_synthetic():
+    """Synthetic quantized gate weights: fused kernel vs the unfused
+    gather_qmm + silu*x_up path, within the kernel's bf16 bound, plus finiteness."""
+    mx.random.seed(0)
+
+    import mlx.nn as nn
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        fused_gate_qmv_silu_mul,
+    )
+
+    # Gate regime: K%512==0, N%8==0, bits=4, group_size=64, affine.
+    E, N, K, TOPK = 4, 16, 512, 2
+    dtype = mx.bfloat16
+    gate_w = (mx.random.normal((E, N, K)) * 0.02).astype(dtype)
+    gwq, gs, gb = mx.quantize(gate_w, group_size=64, bits=4)
+    mx.eval(gwq, gs, gb)
+
+    # Two routing patterns: spread (hi=E) and collisions (many tokens, few experts).
+    for B, hi in [(2, E), (4, max(1, E // 2))]:
+        x = mx.random.normal((B, 1, 1, K)).astype(dtype)
+        idx = mx.random.randint(0, hi, shape=(B, TOPK)).astype(mx.uint32)
+        x_up = mx.random.normal((B, TOPK, 1, N)).astype(dtype)
+
+        x_gate = mx.gather_qmm(
+            x,
+            gwq,
+            gs,
+            gb,
+            rhs_indices=idx,
+            transpose=True,
+            group_size=64,
+            bits=4,
+            mode="affine",
+        )
+        y_ref = nn.silu(x_gate) * x_up
+        y_fused = fused_gate_qmv_silu_mul(x, gwq, gs, gb, idx, x_up)
+        mx.eval(y_ref, y_fused)
+
+        assert y_ref.shape == y_fused.shape
+        # A broken kernel must not leak NaN/Inf into the downstream down_proj matmul.
+        assert bool(
+            mx.all(mx.isfinite(y_fused.astype(mx.float32))).item()
+        ), f"B={B} hi={hi}: non-finite fused output"
+        # Same bf16 bound as the @requires_model kernel test.
+        max_abs, rel = _max_rel_diff(y_ref, y_fused)
+        assert rel < 2e-2, f"B={B} hi={hi}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_can_fuse_declines_nonstock_call():
+    """can_fuse: False when SwitchGLU.__call__ is overridden (the fused subclass
+    would impose stock semantics and silently bypass the override), True for stock."""
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import can_fuse
+
+    # hidden=64 keeps down_proj's input dim divisible by the quant group size.
+    sw_stock = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    assert can_fuse(sw_stock) is True, "stock in-regime SwitchGLU should fuse"
+
+    class _CustomSwitchGLU(SwitchGLU):
+        def __call__(self, x, indices):  # overridden forward
+            return super().__call__(x, indices)
+
+    sw_custom = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    sw_custom.__class__ = _CustomSwitchGLU  # same swap mechanism the patch uses
+    assert can_fuse(sw_custom) is False, "non-stock __call__ must fall back"
+
+
+def test_can_fuse_declines_non_silu_activation():
+    """can_fuse: False for a non SiLU activation (the kernel and the fallback
+    both bake in silu, which would silently replace the module's formula),
+    True for the stock SwiGLU control."""
+    import types
+
+    import mlx.nn as nn
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        can_fuse,
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    # Same build as _quantized_switch_glu, but the activation kwarg is the
+    # subject under test, so construct directly.
+    sw = SwitchGLU(512, 64, 4, activation=nn.gelu, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        proj = getattr(sw, name)
+        setattr(sw, name, proj.to_quantized(group_size=64, bits=4, mode="affine"))
+    assert can_fuse(sw) is False, "non SiLU activation must fall back"
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=4)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert patch_switch_glu_with_fused_swiglu(model) == 0, "gelu module must not patch"
+
+    sw_stock = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    assert can_fuse(sw_stock) is True, "stock SwiGLU activation should fuse"
+
+
+def test_fused_forward_falls_back_on_dtype_mismatch():
+    """A runtime activation dtype the fused kernel rejects but gather_qmm
+    tolerates (bf16 gate params, fp16 activations) must fall back, not crash,
+    and match the unfused forward."""
+    import types
+
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 2  # 2*4=8 < 64 -> unsorted
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        lin = getattr(sw, name)
+        lin.weight = lin.weight.astype(mx.bfloat16)  # bf16 weight -> bf16 scales
+        setattr(sw, name, lin.to_quantized(group_size=64, bits=4, mode="affine"))
+    assert sw.gate_proj.scales.dtype == mx.bfloat16
+
+    # fp16 activations mismatch the bf16 gate params: the fused kernel raises,
+    # the unfused gather_qmm tolerates it.
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+
+    out_ref = sw(x, indices)  # stock forward, unpatched
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert patch_switch_glu_with_fused_swiglu(model) == 1, "layer should patch"
+
+    out_fb = sw(x, indices)  # patched -> kernel raises -> fallback, no crash
+    mx.eval(out_fb)
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    assert rel < 1e-3, f"fallback != unfused: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+# Fallback index contract (synthetic, no model): the fallback must see the
+# untouched flat indices. The sorted path's (M_tok, 1) kernel reshape once
+# leaked into the fallback and broadcast an M_tok x M_tok cross product
+# (PR #26188 review repro).
+def _bf16_quantized_switch_glu(in_dim, hidden, n_experts):
+    """Quantize from bf16 weights so fp16 activations trip the kernel's runtime
+    dtype check while the unfused path tolerates them."""
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        lin = getattr(sw, name)
+        lin.weight = lin.weight.astype(mx.bfloat16)
+        setattr(sw, name, lin.to_quantized(group_size=64, bits=4, mode="affine"))
+    return sw
+
+
+def test_sorted_dtype_mismatch_fallback_matches_reference(monkeypatch):
+    """Reviewer repro: B*T == 64 takes the sorted path, the kernel rejects fp16
+    activations on bf16 params, and the fallback must match the reference in
+    shape and value."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 16  # 16*4 = 64 -> sorted
+    sw = _bf16_quantized_switch_glu(in_dim, hidden, n_experts)
+
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_ref = sw(x, indices)
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    out_fb = sw(x, indices)
+    mx.eval(out_fb)
+    assert out_fb.shape == out_ref.shape
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    # Post fix the fallback runs the same MLX ops as the stock forward, so the
+    # bound only absorbs compiled vs eager elementwise ordering (~1 fp16 ULP).
+    assert bool(
+        mx.allclose(
+            out_fb.astype(mx.float32),
+            out_ref.astype(mx.float32),
+            rtol=2e-3,
+            atol=2e-4,
+        ).item()
+    ), f"sorted fallback != reference: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_unsorted_dtype_mismatch_fallback_matches_reference(monkeypatch):
+    """Sibling guard: same dtype mismatch on the unsorted path (B*T < 64)."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 2  # 2*4 = 8 < 64 -> unsorted
+    sw = _bf16_quantized_switch_glu(in_dim, hidden, n_experts)
+
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_ref = sw(x, indices)
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    out_fb = sw(x, indices)
+    mx.eval(out_fb)
+    assert out_fb.shape == out_ref.shape
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    assert bool(
+        mx.allclose(
+            out_fb.astype(mx.float32),
+            out_ref.astype(mx.float32),
+            rtol=2e-3,
+            atol=2e-4,
+        ).item()
+    ), f"unsorted fallback != reference: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_forced_kernel_rejection_falls_back_correctly(monkeypatch):
+    """Any ValueError from the fused kernel, not just a dtype mismatch, must
+    take the identical fallback: force one via monkeypatch and check both
+    routing paths against the unpatched module."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k = 512, 64, 4, 4
+    sw = _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias=False)
+
+    cases = []
+    # B=2 -> 8 < 64 -> unsorted; B=16 -> 64 -> sorted.
+    for B, label in [(2, "unsorted"), (16, "sorted")]:
+        x = mx.random.normal((B, in_dim))
+        indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+        out_ref = sw(x, indices)
+        mx.eval(out_ref)
+        cases.append((label, x, indices, out_ref))
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    # Patch before installing the raiser: _aot_warm_kernel dispatches the real
+    # kernel at patch time and does not catch ValueError.
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    def raiser(*args, **kwargs):
+        raise ValueError("forced rejection")
+
+    monkeypatch.setattr(fused_swiglu, "fused_gate_qmv_silu_mul", raiser)
+
+    for label, x, indices, out_ref in cases:
+        out_fb = sw(x, indices)
+        mx.eval(out_fb)
+        assert out_fb.shape == out_ref.shape, label
+        max_abs, rel = _max_rel_diff(out_ref, out_fb)
+        assert bool(
+            mx.allclose(
+                out_fb.astype(mx.float32),
+                out_ref.astype(mx.float32),
+                rtol=1e-5,
+                atol=1e-6,
+            ).item()
+        ), f"forced rejection {label}: max_abs={max_abs:.3e} rel={rel:.2%}"

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

@@ -0,0 +1,60 @@
+"""Parent-death watchdog for MLX workers on Apple Silicon.
+
+macOS has no ``PR_SET_PDEATHSIG`` equivalent, so the kernel will not signal a
+worker process when its parent dies; the worker would be reparented to PID 1
+and leak (holding GPU/host memory and ports). This module emulates PDEATHSIG
+with a daemon thread that watches the parent PID via kqueue and SIGKILLs the
+current process once it gets orphaned.
+"""
+
+import os
+import select
+import signal
+import threading
+
+
+def start_parent_death_watcher() -> None:
+    """SIGKILL this process once its current parent exits (macOS only).
+
+    kqueue with an ``EVFILT_PROC`` / ``NOTE_EXIT`` filter is the native,
+    event-driven mechanism on macOS (exposed via ``select.kqueue`` /
+    ``select.kevent``), so the watcher thread blocks until the parent actually
+    exits instead of waking up to poll.
+
+    ``SIGKILL`` is sent from this watcher thread and is uncatchable /
+    unblockable, so it works even when the main thread is stuck inside a
+    blocking native call (e.g. an MLX/Metal ``mx.eval`` / ``.tolist()``).
+    """
+    original_ppid = os.getppid()
+
+    def _watch_parent():
+        kq = select.kqueue()
+        kev = select.kevent(
+            original_ppid,
+            filter=select.KQ_FILTER_PROC,
+            flags=select.KQ_EV_ADD,
+            fflags=select.KQ_NOTE_EXIT,
+        )
+        try:
+            # Register the EVFILT_PROC / NOTE_EXIT watch on the parent PID.
+            kq.control([kev], 0, None)
+        except (ProcessLookupError, OSError):
+            # The parent already exited before we could register the watch
+            # (ESRCH); we are already orphaned.
+            os.kill(os.getpid(), signal.SIGKILL)
+            return
+        # Guard against the race where the parent exits between reading
+        # original_ppid and registering the watch above.
+        if os.getppid() != original_ppid:
+            os.kill(os.getpid(), signal.SIGKILL)
+            return
+        # Block until the parent exits, then terminate ourselves.
+        kq.control(None, 1, None)
+        os.kill(os.getpid(), signal.SIGKILL)
+
+    watcher = threading.Thread(
+        target=_watch_parent,
+        name="parent-death-watcher",
+        daemon=True,
+    )
+    watcher.start()

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

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import gzip
+import json
+import logging
+import os
+import shutil
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+import torch
+
+from sglang.srt.managers.io_struct import ProfileReqOutput
+from sglang.srt.utils.tensor_bridge import use_mlx
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetalCaptureProfiler:
+    label: str
+    trace_path: Path
+    stop_capture: Callable[[], None]
+    standalone: bool
+
+    @classmethod
+    def start_mlx(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            import mlx.core as mx
+
+            mx.metal.start_capture(str(trace_path))
+        except RuntimeError as e:
+            return None, _capture_error("MLX", e)
+
+        return cls._started(
+            label="MLX",
+            trace_path=trace_path,
+            stop_capture=mx.metal.stop_capture,
+            standalone=True,
+        )
+
+    @classmethod
+    def start_mps(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            if not hasattr(torch, "mps") or not hasattr(torch.mps, "profiler"):
+                raise RuntimeError("torch.mps.profiler is not available")
+            context = torch.mps.profiler.metal_capture(str(trace_path))
+            context.__enter__()
+        except RuntimeError as e:
+            return None, _capture_error("MPS", e)
+
+        return cls._started(
+            label="MPS",
+            trace_path=trace_path,
+            stop_capture=lambda: context.__exit__(None, None, None),
+            standalone=False,
+        )
+
+    @classmethod
+    def _started(
+        cls,
+        *,
+        label: str,
+        trace_path: Path,
+        stop_capture: Callable[[], None],
+        standalone: bool,
+    ):
+        profiler = cls(
+            label=label,
+            trace_path=trace_path,
+            stop_capture=stop_capture,
+            standalone=standalone,
+        )
+        logger.info("%s Metal capture started, saving to %s", label, trace_path)
+        return profiler, ProfileReqOutput(success=True, message="Succeeded")
+
+    def stop(self) -> str:
+        self.stop_capture()
+
+        logger.info(
+            "%s Metal capture stopped. Trace saved to: %s",
+            self.label,
+            self.trace_path,
+        )
+        return f" Metal trace: {self.trace_path}"
+
+
+def _capture_error(label: str, error: RuntimeError) -> ProfileReqOutput:
+    return ProfileReqOutput(
+        success=False,
+        message=(
+            f"Failed to start {label} Metal capture: {error}. "
+            "Set MTL_CAPTURE_ENABLED=1 in the server's environment "
+            "before launching to enable GPU trace capture."
+        ),
+    )
+
+
+class MetalTorchProfiler:
+    def __init__(
+        self,
+        *,
+        start_metal_capture: Callable[[Path], tuple[Any, ProfileReqOutput]],
+        torch_profiler: Optional[Any] = None,
+    ):
+        self.start_metal_capture = start_metal_capture
+        self.torch_profiler = torch_profiler
+        self.metal_profiler = None
+
+    def start(self):
+        trace_path = _new_temp_gputrace_path()
+        self.metal_profiler, result = self.start_metal_capture(trace_path)
+        if not result.success:
+            raise RuntimeError(result.message)
+        if self.torch_profiler is not None:
+            try:
+                self.torch_profiler.start()
+            except Exception:
+                self.metal_profiler.stop()
+                raise
+
+    def stop(self):
+        try:
+            if self.torch_profiler is not None:
+                self.torch_profiler.stop()
+        finally:
+            if self.metal_profiler is not None:
+                self.metal_profiler.stop()
+
+    def export_chrome_trace(self, path: str):
+        if self.torch_profiler is not None:
+            self.torch_profiler.export_chrome_trace(path)
+        else:
+            _write_empty_chrome_trace(path)
+
+        if self.metal_profiler is None:
+            return
+
+        final_path = _unique_gputrace_path_for_chrome_trace(path)
+        final_path.parent.mkdir(parents=True, exist_ok=True)
+        if self.metal_profiler.trace_path.exists():
+            shutil.move(str(self.metal_profiler.trace_path), str(final_path))
+            logger.info("Metal trace saved to: %s", final_path)
+
+
+def apply_metal_profiler_patches() -> None:
+    if getattr(torch.profiler.profile, "_sglang_metal_patched", False):
+        return
+
+    original_profile = torch.profiler.profile
+
+    def profile(*args, **kwargs):
+        activities = _get_activities(args, kwargs)
+        if not _has_cuda_activity(activities):
+            return original_profile(*args, **kwargs)
+
+        if use_mlx():
+            return MetalTorchProfiler(
+                start_metal_capture=MetalCaptureProfiler.start_mlx
+            )
+
+        torch_activities = [
+            activity for activity in activities if not _is_cuda_activity(activity)
+        ]
+        torch_profiler = None
+        if torch_activities:
+            patched_args, patched_kwargs = _replace_activities(
+                args, kwargs, torch_activities
+            )
+            torch_profiler = original_profile(*patched_args, **patched_kwargs)
+
+        return MetalTorchProfiler(
+            start_metal_capture=MetalCaptureProfiler.start_mps,
+            torch_profiler=torch_profiler,
+        )
+
+    profile._sglang_metal_patched = True
+    profile._sglang_original_profile = original_profile
+    torch.profiler.profile = profile
+
+
+def _get_activities(args, kwargs):
+    if "activities" in kwargs:
+        return kwargs["activities"]
+    if args:
+        return args[0]
+    return None
+
+
+def _replace_activities(args, kwargs, activities):
+    kwargs = dict(kwargs)
+    if "activities" in kwargs:
+        kwargs["activities"] = activities
+        return args, kwargs
+
+    if args:
+        args = list(args)
+        args[0] = activities
+        return tuple(args), kwargs
+
+    kwargs["activities"] = activities
+    return args, kwargs
+
+
+def _has_cuda_activity(activities) -> bool:
+    if activities is None:
+        return False
+    return any(_is_cuda_activity(activity) for activity in activities)
+
+
+def _is_cuda_activity(activity) -> bool:
+    return activity == torch.profiler.ProfilerActivity.CUDA
+
+
+def _new_temp_gputrace_path() -> Path:
+    output_dir = Path(os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")).expanduser()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(100):
+        candidate = (
+            output_dir / f"sglang-metal-{os.getpid()}-{time.time_ns()}-{i}.gputrace"
+        )
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path in {output_dir}")
+
+
+def _unique_gputrace_path_for_chrome_trace(path: str) -> Path:
+    chrome_path = Path(path).expanduser()
+    name = chrome_path.name
+    if name.endswith(".trace.json.gz"):
+        name = name[: -len(".trace.json.gz")] + ".gputrace"
+    else:
+        name = chrome_path.stem + ".gputrace"
+
+    base = chrome_path.with_name(name)
+    if not base.exists():
+        return base
+
+    stem = base.name[: -len(".gputrace")]
+    for i in range(100):
+        candidate = base.with_name(f"{stem}-{time.time_ns()}-{i}.gputrace")
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path for {base}")
+
+
+def _write_empty_chrome_trace(path: str):
+    trace = {"traceEvents": []}
+    Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+    if str(path).endswith(".gz"):
+        with gzip.open(path, "wt") as f:
+            json.dump(trace, f)
+    else:
+        with open(path, "w") as f:
+            json.dump(trace, f)

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

@@ -0,0 +1,263 @@
+"""MLX overlap scheduling mixin for the SGLang scheduler.
+
+Provides ``event_loop_overlap_mlx``, which pipelines MLX forward
+passes by keeping two in-flight lazy graphs queued on the GPU while
+the scheduler runs its CPU-side bookkeeping on the tokens of the
+older one.  The lazy-graph primitives live in
+``hardware_backend/mlx/tp_worker.py`` and ``model_runner.py``.
+
+Each request's attention KV lives in per-request, per-layer
+``ContiguousAttentionKVCache`` objects that ``MLXAttentionWrapper`` mutates
+in place during the forward pass. Chained decodes reuse the same cache objects:
+step N+1's graph reads step N's lazy writes via MLX's dependency tracking, so
+the GPU runs both steps back-to-back with no idle gap.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import mlx.core as mx
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.overlap_utils import resolve_forward_inputs
+from sglang.srt.utils import DynamicGradMode
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.model_runner import (
+        MlxPendingDecode,
+        MlxPendingExtend,
+        MlxPendingPrefill,
+    )
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+@dataclass
+class MlxPendingJob:
+    """Unfinished MLX work and graphs queued on the GPU.
+
+    Attributes:
+        lazy_tokens: Lazily evaluated token IDs produced by the forward
+            pass.  Unevaluated; calling ``.tolist()`` / ``.item()`` /
+            ``mx.eval`` on it will block until the Metal kernel finishes.
+            ``None`` for idle batches.
+        prefills: MLX prefill state returned by the model worker — one
+            entry per new request in an extend batch.  Used by
+            ``finalize_mlx_result`` to commit per-request caches.  Empty
+            list for pure-decode steps.
+        extends: Chunked-prefill-continuation state, one entry per
+            already-active request whose extend seq_len > 1.  Also empty
+            for pure-decode steps.
+        decode: Decode state covering full-decode mode AND mixed
+            single-token decodes inside an extend batch.  Used as the
+            chaining root by :meth:`async_chained_decode_mlx`.
+        mode: One of ``"decode"``, ``"extend"``, ``"idle"`` describing
+            which forward pass produced this job.  Drives finalise
+            dispatch and whether chaining is safe.
+        batch_copy: Snapshot of the :class:`ScheduleBatch` at launch
+            time.  Decoupled from the live batch so
+            ``process_batch_result`` can update request state without
+            racing against the next scheduling decision.
+        schedule_batch: The full scheduler batch.  Unlike ``batch_copy``,
+            this keeps allocator/cache fields needed when a prefill batch
+            becomes the next running decode batch.
+        reqs: Snapshot of ``batch.reqs`` at launch time.  The overlap
+            loop uses this to check ``req.finished()`` on the previous
+            step's request list without holding a reference to the
+            mutable batch object.
+    """
+
+    lazy_tokens: Optional[mx.array]
+    prefills: list[MlxPendingPrefill]
+    extends: list[MlxPendingExtend]
+    decode: Optional[MlxPendingDecode]
+    mode: str
+    batch_copy: ScheduleBatch
+    schedule_batch: ScheduleBatch
+    reqs: List[Req]
+
+
+class SchedulerMlxOverlapMixin:
+    """Mixin that adds MLX overlap scheduling to :class:`Scheduler`."""
+
+    def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob):
+        # Account for this completed forward step. The standard scheduler does
+        # this inside run_batch(), but the MLX overlap loop bypasses run_batch,
+        # so without this forward_ct never advances on MLX. That stalls the
+        # watchdog liveness counter and, more importantly, breaks step-bounded
+        # profiling: _profile_batch_predicate auto-starts/stops based on
+        # forward_ct, so `--profile-steps` (and the server /start_profile
+        # num_steps path) only takes effect once the counter moves here.
+        self.forward_ct += 1
+        self.profiler_manager._profile_batch_predicate(pending.schedule_batch)
+
+        result = self.tp_worker.finalize_mlx_result(
+            pending.prefills,
+            pending.extends,
+            pending.decode,
+            pending.mode,
+            pending.reqs,
+        )
+        if result.next_token_ids is not None:
+            pending.batch_copy.input_ids = result.next_token_ids
+            pending.schedule_batch.input_ids = result.next_token_ids
+        self.last_batch = pending.schedule_batch
+        self.process_batch_result(pending.batch_copy, result)
+
+    @DynamicGradMode()
+    def event_loop_overlap_mlx(self: Scheduler):
+        """MLX-specific overlap loop modelled on ``mlx_lm.generate.generate_step``.
+
+        At steady state we keep TWO in-flight MLX graphs queued on the
+        GPU:
+
+        * ``pending_curr`` — the step whose tokens we are about to block
+          on and feed into the scheduler's bookkeeping.
+        * ``pending_next`` — the step that was built on top of
+          ``pending_curr``'s still-lazy output tokens via
+          ``async_chained_decode_mlx`` and has already been handed to
+          ``mx.async_eval``.  Because MLX tracks the full dependency
+          graph, the GPU will execute ``pending_next`` back-to-back
+          with ``pending_curr`` — there is no scheduling gap on the
+          device.
+
+        Bookkeeping timeline for a steady-state decode loop:
+
+            iter k:
+              build pending_next  (CPU graph build + mx.async_eval; cheap)
+              block on pending_curr via .tolist() (wait only on curr's tokens)
+              process_batch_result(pending_curr)   <-- GPU is running pending_next
+              pending_curr = pending_next
+
+        The chain is broken (we fall back to a "schedule + launch" step)
+        whenever any of the following holds:
+
+        * ``pending_curr`` is not a pure decode (e.g. prefill/extend).
+        * The waiting queue has new requests that need prefill.
+        * Any req in ``pending_curr`` just finished this iteration, so
+          the composition for ``pending_next`` would need to shrink.
+
+        When the chain breaks mid-flight we still finalise the
+        already-launched ``pending_next`` normally (its tokens are
+        valid for all surviving reqs).  With RadixCache-backed caches
+        (#21509) there is no ``extract_cache`` step: per-request caches
+        are the source of truth and are never merged into a shared
+        batched buffer.
+        """
+        pending_curr: Optional[MlxPendingJob] = None
+        pending_next: Optional[MlxPendingJob] = None
+
+        def _launch_fresh(batch: ScheduleBatch) -> MlxPendingJob:
+            # Materialize batch.input_ids from CPU staging (prefill) or the
+            # FutureMap relay (decode) before the forward. With deferred input
+            # materialization, get_next_batch_to_run leaves input_ids unset; the
+            # CUDA paths call resolve_forward_inputs for this, but the MLX overlap
+            # loop must do it too, otherwise async_forward_batch_generation_mlx
+            # dereferences a None input_ids.
+            resolve_forward_inputs(batch, self.future_map)
+            lazy_tokens, prefills, extends, decode, mode = (
+                self.tp_worker.async_forward_batch_generation_mlx(batch)
+            )
+            return MlxPendingJob(
+                lazy_tokens=lazy_tokens,
+                prefills=prefills,
+                extends=extends,
+                decode=decode,
+                mode=mode,
+                batch_copy=batch.copy(),
+                schedule_batch=batch,
+                reqs=list(batch.reqs),
+            )
+
+        def _launch_chained(prev: MlxPendingJob) -> MlxPendingJob:
+            assert prev.decode is not None
+            lazy_tokens, prefills, extends, decode, mode = (
+                self.tp_worker.async_chained_decode_mlx(prev.decode)
+            )
+            # Composition is identical to prev: reuse a fresh batch copy
+            # of the same underlying ScheduleBatch so process_batch_result
+            # updates the same req objects with the new token.
+            return MlxPendingJob(
+                lazy_tokens=lazy_tokens,
+                prefills=prefills,
+                extends=extends,
+                decode=decode,
+                mode=mode,
+                batch_copy=prev.batch_copy.copy(),
+                schedule_batch=prev.schedule_batch,
+                reqs=prev.reqs,
+            )
+
+        while True:
+            recv_reqs = self.request_receiver.recv_requests()
+            self.process_input_requests(recv_reqs)
+            if self._engine_paused:
+                continue
+
+            # 1. If pending_curr is a pure decode AND no new prefill is waiting,
+            #    build pending_next on top of it NOW — before we block on curr.
+            can_chain = (
+                pending_curr is not None
+                and pending_curr.mode == "decode"
+                and pending_curr.decode is not None
+                and not self.waiting_queue
+            )
+            if can_chain and pending_next is None:
+                # Build + launch the chained step BEFORE we block on
+                # pending_curr — this is the "no idle gap" trick.
+                # GPU now has 2 steps queued.
+                pending_next = _launch_chained(pending_curr)
+                self.result_queue.append(pending_next)
+
+            # 2. Finalize/process on pending_curr's tokens.  (GPU is already
+            #    executing pending_next at this point.)
+            if pending_curr is not None:
+                self._finalize_mlx_pending_job(pending_curr)
+                self.result_queue.popleft()
+                pending_curr = None
+
+            # 3. Decide whether pending_next is still valid (if no reqs finished)
+            #    and promote it.
+            finished_any = any(
+                req.finished() for req in (pending_next.reqs if pending_next else [])
+            )
+            new_prefill_waiting = bool(self.waiting_queue)
+            if (
+                pending_next is not None
+                and not finished_any
+                and not new_prefill_waiting
+            ):
+                pending_curr = pending_next
+                pending_next = None
+                self.cur_batch_for_debug = pending_curr.schedule_batch
+                self.last_batch = pending_curr.schedule_batch
+                if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
+                    self.invariant_checker.self_check_during_busy()
+                continue
+
+            # 4. Chain is broken. Finalise pending_next (if any), then
+            #    schedule fresh.
+            if pending_next is not None:
+                self._finalize_mlx_pending_job(pending_next)
+                self.result_queue.popleft()
+                pending_next = None
+            plan = self.get_next_batch_to_run(
+                running_batch=self.running_batch, last_batch=self.last_batch
+            )
+            self.running_batch = plan.running_batch
+            next_batch = plan.batch_to_run
+            self.cur_batch_for_debug = next_batch
+            if next_batch:
+                pending_curr = _launch_fresh(next_batch)
+                self.result_queue.append(pending_curr)
+            else:
+                self.on_idle()
+
+            self.last_batch = next_batch
+            if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
+                self.invariant_checker.self_check_during_busy()

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

@@ -0,0 +1,528 @@
+"""MLX-specific TpModelWorker subclass for Apple Silicon.
+
+Routes forward passes through the MLX model runner, bypassing PyTorch
+MPS.  A lightweight stub provides scheduler bookkeeping; the actual
+attention KV data lives in MlxAttentionKVPool.
+
+The worker also exposes an async (lazy-eval) surface used by the MLX
+overlap scheduler: ``async_forward_batch_generation_mlx`` launches a
+batch without blocking on the GPU, ``async_chained_decode_mlx`` builds
+the next decode step on top of a still-lazy previous decode, and
+``finalize_mlx_result`` blocks on the lazy outputs and produces a
+normal ``GenerationBatchResult``.
+"""
+
+import logging
+from typing import Optional, Union
+
+import mlx.core as mx
+import torch
+
+from sglang.srt.hardware_backend.mlx.model_runner import (
+    MlxPendingDecode,
+    MlxPendingExtend,
+    MlxPendingPrefill,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+
+logger = logging.getLogger(__name__)
+
+
+class MlxTpModelWorker(TpModelWorker):
+    """A tensor parallel model worker that routes inference through MLX.
+
+    Inherits from TpModelWorker for scheduler integration, but replaces
+    the standard ModelRunner with MlxModelRunnerStub (no PyTorch weights,
+    zero-memory KV cache) and delegates all forward passes to a native
+    MlxModelRunner.
+    """
+
+    def _init_model_runner(self):
+        """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,
+            quantization=self.server_args.quantization,
+        )
+        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,
+            gpu_id=self.gpu_id,
+            tp_rank=self.tp_rank,
+            tp_size=self.tp_size,
+            moe_ep_rank=self.moe_ep_rank,
+            moe_ep_size=self.ep_size,
+            pp_rank=self.pp_rank,
+            pp_size=self.pp_size,
+            nccl_port=self.nccl_port,
+            dp_rank=self.dp_rank,
+            server_args=self.server_args,
+            is_draft_worker=self.is_draft_worker,
+            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,
+        )
+
+        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 MLX cache pools after the stub pools are ready."""
+        if not self._mlx_pool_initialized:
+            self._mlx_runner.init_cache_pools(self._model_runner.req_to_token_pool)
+            self._mlx_pool_initialized = True
+
+    def forward_batch_generation(
+        self,
+        batch: Optional[ScheduleBatch],
+        forward_batch: Optional[ForwardBatch] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+        is_verify: bool = False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
+    ) -> GenerationBatchResult:
+        """Override to route through MLX model runner."""
+        if batch is not None:
+            self._ensure_mlx_pool_initialized()
+            return self._forward_batch_generation_mlx(batch)
+
+        # Fallback to standard path for None batches
+        return super().forward_batch_generation(
+            batch,
+            forward_batch,
+            pp_proxy_tensors,
+            is_verify,
+            skip_attn_backend_init,
+        )
+
+    def _cleanup_stale_rids(self, forward_mode, current_rids: set[str]) -> None:
+        """Remove MLX state for decode-mode requests that dropped out of the batch."""
+        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
+
+    def prepare_for_kv_cache_release(self, req) -> None:
+        """Snapshot MLX auxiliary state at the scheduler's radix insert point."""
+        if self._mlx_runner.has_request(req.rid):
+            self._mlx_runner.store_auxiliary_state_for_request(req.rid)
+            # Prefer the just-snapshotted live auxiliary state for the final
+            # insert. Any older tracked slot is released during component cleanup.
+            req.mamba_last_track_seqlen = None
+
+    def _route_extend_request(self, rid: str, decoding_rids: set[str]) -> str:
+        """Classify a request within an extend / mixed batch.
+
+        Shared by the sync (:meth:`_forward_batch_generation_mlx`) and async
+        (:meth:`_async_extend_batch`) paths so both route identically.
+
+        Returns one of:
+
+        * ``"prefill"``      -- not seen before; start a fresh prefill.
+        * ``"decode"``       -- a genuine single-token decode step mixed into
+          this batch (present in ``batch.decoding_reqs``).
+        * ``"continuation"`` -- a chunked-prefill continuation.  Routing keys on
+          request state, **not** ``seq_len``: a final continuation chunk can be
+          exactly one token, which must still extend.  Routing it as a decode
+          would drop the real token and feed the model its own previous-chunk
+          prediction, silently corrupting the output.
+        """
+        if not self._mlx_runner.has_request(rid):
+            return "prefill"
+        if rid in decoding_rids:
+            return "decode"
+        return "continuation"
+
+    def _forward_batch_generation_mlx(
+        self, batch: ScheduleBatch
+    ) -> GenerationBatchResult:
+        """Run forward pass through the MLX model runner (greedy only)."""
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+        forward_mode = batch.forward_mode
+        reqs = batch.reqs
+
+        if forward_mode.is_idle():
+            return GenerationBatchResult(
+                logits_output=LogitsProcessorOutput(next_token_logits=None),
+                can_run_cuda_graph=False,
+            )
+
+        self._cleanup_stale_rids(forward_mode, {req.rid for req in reqs})
+
+        next_token_ids_list: list[int] = []
+
+        if forward_mode.is_extend():
+            # Ensure pool is up-to-date before pool-backed attention reads it
+            # for prefix-cached prefills.  Only runs on extend batches.
+            self._mlx_runner.flush_all_decode_kv()
+            input_ids_cpu = batch.input_ids.cpu().tolist()
+            out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
+            extend_seq_lens = batch.extend_lens
+
+            offset = 0  # into input_ids_cpu
+            slot_offset = 0  # into out_cache_loc_cpu
+            prefill_rids: list[tuple[str, int]] = []
+            extend_rids: list[tuple[str, int]] = []
+            decode_rids: list[str] = []
+            # Genuine decode steps mixed into this extend batch; see
+            # _route_extend_request.
+            decoding_rids = {r.rid for r in (batch.decoding_reqs or [])}
+
+            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
+                slot_offset += seq_len
+
+                route = self._route_extend_request(req.rid, decoding_rids)
+                if route == "continuation":
+                    next_token = self._mlx_runner.extend(
+                        req.rid, req_token_ids, req_new_slots
+                    )
+                    extend_rids.append((req.rid, next_token))
+                elif route == "decode":
+                    decode_rids.append(req.rid)
+                else:  # "prefill"
+                    prefix_slot_ids = req.prefix_indices.tolist()
+                    full_token_ids = list(req.get_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,
+                        req=req,
+                    )
+                    prefill_rids.append((req.rid, next_token))
+
+            # Batch decode all existing requests at once
+            if decode_rids:
+                decode_results = self._mlx_runner.decode_batch(decode_rids)
+                decode_map = dict(zip(decode_rids, decode_results))
+            else:
+                decode_map = {}
+
+            prefill_map = dict(prefill_rids)
+            extend_map = dict(extend_rids)
+
+            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():
+            req_ids = [req.rid for req in reqs]
+            next_token_ids_list = self._mlx_runner.decode_batch(req_ids)
+
+        else:
+            raise ValueError(
+                f"MLX runner does not support forward mode: {forward_mode}"
+            )
+
+        next_token_ids = torch.tensor(
+            next_token_ids_list, dtype=torch.long, device="cpu"
+        )
+
+        return GenerationBatchResult(
+            logits_output=LogitsProcessorOutput(next_token_logits=None),
+            next_token_ids=next_token_ids,
+            can_run_cuda_graph=False,
+        )
+
+    def async_forward_batch_generation_mlx(self, batch: ScheduleBatch) -> tuple[
+        Union[mx.array, None],
+        list[MlxPendingPrefill],
+        list[MlxPendingExtend],
+        Optional[MlxPendingDecode],
+        str,
+    ]:
+        """Start an async (lazy) forward pass through the MLX model runner.
+
+        Returns ``(lazy_result, prefills, extends, decode, mode)``:
+
+        * ``lazy_result`` — an ``mx.array`` that, when evaluated, forces
+          materialisation of the whole batch's outputs.  ``None`` for
+          idle batches.
+        * ``prefills`` — list of :class:`MlxPendingPrefill` for new
+          requests in an extend batch.
+        * ``extends`` — list of :class:`MlxPendingExtend` for chunked
+          prefill continuations in an extend batch.
+        * ``decode`` — :class:`MlxPendingDecode` for the decode
+          sub-batch (covers full decode mode AND mixed decodes inside
+          an extend batch).
+        * ``mode`` — one of ``"idle"``, ``"decode"``, ``"extend"``.
+
+        The caller must make sure the returned pendings are fed into a
+        subsequent ``mx.async_eval`` or ``.item()`` / ``.tolist()`` call
+        — :meth:`finalize_mlx_result` does that.
+        """
+        self._ensure_mlx_pool_initialized()
+
+        forward_mode = batch.forward_mode
+        reqs = batch.reqs
+
+        if forward_mode.is_idle():
+            return None, [], [], None, "idle"
+
+        self._cleanup_stale_rids(forward_mode, {req.rid for req in reqs})
+
+        if forward_mode.is_decode():
+            req_ids = [req.rid for req in reqs]
+            pending_decode = self._mlx_runner.decode_batch_start(req_ids)
+            mx.async_eval(pending_decode.lazy_tokens)
+            return pending_decode.lazy_tokens, [], [], pending_decode, "decode"
+
+        if forward_mode.is_extend():
+            # TODO (changminbark): Implement per-batch flushing using prefix_slot_ids
+            # Ensure the pool is up-to-date before pool-backed attention
+            # reads it for prefix-cached prefills. Mirror the sync path.
+            self._mlx_runner.flush_all_decode_kv()
+            return self._async_extend_batch(batch)
+
+        raise ValueError(
+            f"MLX async runner does not support forward mode: {forward_mode}"
+        )
+
+    def _async_extend_batch(self, batch: ScheduleBatch) -> tuple[
+        Union[mx.array, None],
+        list[MlxPendingPrefill],
+        list[MlxPendingExtend],
+        Optional[MlxPendingDecode],
+        str,
+    ]:
+        """Launch each request in an EXTEND batch lazily and kick GPU work."""
+        reqs = batch.reqs
+        input_ids_cpu = batch.input_ids.cpu().tolist()
+        out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
+        extend_seq_lens = batch.extend_lens
+
+        offset = 0
+        slot_offset = 0
+        pending_prefills: list[MlxPendingPrefill] = []
+        pending_extends: list[MlxPendingExtend] = []
+        mixed_decode_rids: list[str] = []
+        # Genuine decode steps mixed into this extend batch; see
+        # _route_extend_request.
+        decoding_rids = {r.rid for r in (batch.decoding_reqs or [])}
+
+        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
+            slot_offset += seq_len
+
+            route = self._route_extend_request(req.rid, decoding_rids)
+            if route == "continuation":
+                pending_extends.append(
+                    self._mlx_runner.extend_start(
+                        req_id=req.rid,
+                        new_token_ids=req_token_ids,
+                        new_slot_ids=req_new_slots,
+                    )
+                )
+            elif route == "decode":
+                mixed_decode_rids.append(req.rid)
+            else:  # "prefill"
+                prefix_slot_ids = req.prefix_indices.tolist()
+                full_token_ids = list(req.get_fill_ids())
+                pending_prefills.append(
+                    self._mlx_runner.prefill_start(
+                        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,
+                        req=req,
+                    )
+                )
+
+        pending_mixed_decode: Optional[MlxPendingDecode] = None
+        if mixed_decode_rids:
+            pending_mixed_decode = self._mlx_runner.decode_batch_start(
+                mixed_decode_rids
+            )
+
+        # Stack lazy tokens so the caller has a single handle to evaluate
+        # after CPU scheduling work.  We also hand every cache buffer
+        # (and the decode cache arrays) to mx.async_eval so the GPU
+        # kernel-launch stream sees everything the next step depends on
+        # before we actually block on anything.
+        prefill_ext_tokens: list[mx.array] = [p.lazy_token for p in pending_prefills]
+        prefill_ext_tokens.extend(e.lazy_token for e in pending_extends)
+
+        async_args: list[mx.array] = []
+        if prefill_ext_tokens:
+            lazy_stacked = mx.stack(prefill_ext_tokens, axis=0)
+            async_args.append(lazy_stacked)
+        else:
+            lazy_stacked = None
+
+        for p in pending_prefills:
+            async_args.extend(self._cache_state(p.cache))
+        for e in pending_extends:
+            async_args.extend(self._cache_state(self._mlx_runner._req_caches[e.req_id]))
+        if pending_mixed_decode is not None:
+            async_args.append(pending_mixed_decode.lazy_tokens)
+            for c_list in pending_mixed_decode.caches:
+                async_args.extend(self._cache_state(c_list))
+
+        if async_args:
+            mx.async_eval(*async_args)
+
+        return (
+            lazy_stacked,
+            pending_prefills,
+            pending_extends,
+            pending_mixed_decode,
+            "extend",
+        )
+
+    @staticmethod
+    def _cache_state(cache_list) -> list[mx.array]:
+        """Flatten a per-layer cache list to its ``state`` arrays."""
+        arrays: list[mx.array] = []
+
+        def collect(value):
+            if isinstance(value, mx.array):
+                arrays.append(value)
+            elif value is None:
+                return
+            elif isinstance(value, (list, tuple)):
+                for item in value:
+                    collect(item)
+            elif isinstance(value, dict):
+                for item in value.values():
+                    collect(item)
+
+        for cache in cache_list:
+            collect(getattr(cache, "state", ()))
+        return arrays
+
+    def async_chained_decode_mlx(
+        self,
+        prev_pending: MlxPendingDecode,
+    ) -> tuple[mx.array, list, list, MlxPendingDecode, str]:
+        """Launch a decode step that chains off a still-lazy previous decode.
+
+        This is the "no idle gap" pipelining primitive: build the next
+        decode's compute graph using ``prev_pending.lazy_tokens`` (still
+        unevaluated) as its input ids, hand the combined graph to
+        ``mx.async_eval``, and return.  The GPU runs the new step
+        immediately after ``prev_pending`` with no scheduling gap, while
+        the caller is free to block on ``prev_pending`` and run CPU-side
+        bookkeeping.
+
+        Preconditions (caller must ensure):
+
+        * ``prev_pending`` was produced by a previous decode start
+          (either :meth:`async_forward_batch_generation_mlx` in decode
+          mode or a previous :meth:`async_chained_decode_mlx`).
+        * The batch composition for this step is identical to
+          ``prev_pending`` — same requests, same order.  Composition
+          changes (finished reqs, new prefills) must break the chain.
+        * ``prev_pending`` should be finalised BEFORE the returned
+          pending, so per-request token lists are appended in order.
+
+        Returns a 5-tuple matching
+        :meth:`async_forward_batch_generation_mlx` for the decode case:
+        ``(lazy_tokens, [], [], pending_decode, "decode")``.  The empty
+        prefill/extend lists are always absent for chained decodes.
+        """
+        pending = self._mlx_runner.decode_batch_start_chained(prev_pending)
+        mx.async_eval(pending.lazy_tokens)
+        return pending.lazy_tokens, [], [], pending, "decode"
+
+    def finalize_mlx_result(
+        self,
+        prefills: list[MlxPendingPrefill],
+        extends: list[MlxPendingExtend],
+        decode: Optional[MlxPendingDecode],
+        mode: str,
+        reqs: list,
+    ) -> GenerationBatchResult:
+        """Materialise a lazy MLX result into a :class:`GenerationBatchResult`.
+
+        The blocking wait happens inside ``decode_batch_finalize`` /
+        ``prefill_finalize`` / ``extend_finalize`` via ``.tolist()`` /
+        ``.item()`` on the specific lazy outputs.
+        """
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+        if mode == "idle":
+            return GenerationBatchResult(
+                logits_output=LogitsProcessorOutput(next_token_logits=None),
+                can_run_cuda_graph=False,
+            )
+
+        if mode == "decode":
+            assert decode is not None
+            next_tokens_list = self._mlx_runner.decode_batch_finalize(decode)
+
+        elif mode == "extend":
+            prefill_map: dict[str, int] = {}
+            for pending_p in prefills:
+                prefill_map[pending_p.req_id] = self._mlx_runner.prefill_finalize(
+                    pending_p
+                )
+
+            extend_map: dict[str, int] = {}
+            for pending_e in extends:
+                extend_map[pending_e.req_id] = self._mlx_runner.extend_finalize(
+                    pending_e
+                )
+
+            decode_map: dict[str, int] = {}
+            if decode is not None:
+                mixed_tokens = self._mlx_runner.decode_batch_finalize(decode)
+                decode_map = {
+                    rid: tok for rid, tok in zip(decode.req_ids, mixed_tokens)
+                }
+
+            next_tokens_list = []
+            for req in reqs:
+                if req.rid in decode_map:
+                    next_tokens_list.append(decode_map[req.rid])
+                elif req.rid in extend_map:
+                    next_tokens_list.append(extend_map[req.rid])
+                else:
+                    next_tokens_list.append(prefill_map[req.rid])
+
+        else:
+            raise ValueError(f"Unknown MLX async mode: {mode}")
+
+        next_token_ids = torch.tensor(next_tokens_list, dtype=torch.long, device="cpu")
+        return GenerationBatchResult(
+            logits_output=LogitsProcessorOutput(next_token_logits=None),
+            next_token_ids=next_token_ids,
+            can_run_cuda_graph=False,
+        )