[MLX] Window-bounded SWA KV storage and in-graph sampling (#34166)

Co-authored-by: Siming Deng <siming_deng_stat@163.com>
Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
Co-authored-by: Jiminator <Jiminator@users.noreply.github.com>
Co-authored-by: damahua <damahua@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/arg_groups/overrides.py

@@ -61,6 +61,7 @@ from sglang.srt.utils.common import (
     is_xpu,
     xpu_has_xmx_support,
 )
+from sglang.srt.utils.tensor_bridge import use_mlx
 
 logger = logging.getLogger(__name__)
 
@@ -933,8 +934,12 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict:
             overrides["attention_backend"] = "intel_xpu"
         elif is_hip():
             overrides["attention_backend"] = "aiter"
-        elif not is_mps():
-            # No triton on macOS; MPS keeps the platform default.
+        elif not (is_mps() and use_mlx()):
+            # No triton on macOS, but only the MLX runner can actually serve
+            # gpt-oss there -- it owns attention, so it keeps the platform
+            # default. macOS *without* MLX must still fall through to triton
+            # and fail fast below: torch_native has neither sliding-window nor
+            # attention-sink support, so accepting it would silently mis-serve.
             overrides["attention_backend"] = "triton"
     if is_xpu():
         # Check for bf16 dtype on Intel XPU. Reads the pristine dtype request,

python/sglang/srt/constrained/xgrammar_backend.py

@@ -38,6 +38,7 @@ from sglang.srt.constrained.base_grammar_backend import (
 )
 from sglang.srt.constrained.utils import is_legacy_structural_tag
 from sglang.srt.utils import is_hip
+from sglang.srt.utils.common import is_pin_memory_available
 
 _is_hip = is_hip()
 
@@ -58,14 +59,16 @@ MAX_ROLLBACK_TOKENS = 200
 
 
 def _allocate_token_bitmask(vocab_size: int, batch_size: int) -> torch.Tensor:
-    # Always allocate a pinned bitmask so the later H2D to the device can be a
-    # genuine non_blocking copy (a pageable source silently downgrades it to a
-    # blocking copy).
+    # Allocate a pinned bitmask where pinning exists so the later H2D to the
+    # device can be a genuine non_blocking copy (a pageable source silently
+    # downgrades it to a blocking copy). MPS torch has no pin-memory kernel
+    # and asserts on pin_memory=True; the MLX path consumes the mask on the
+    # CPU anyway.
     return torch.full(
         get_bitmask_shape(batch_size, vocab_size),
         -1,
         dtype=bitmask_dtype,
-        pin_memory=True,
+        pin_memory=is_pin_memory_available(),
     )
 
 
@@ -132,6 +135,12 @@ class XGrammarGrammar(BaseGrammarObject):
             import sgl_kernel_npu  # noqa: F401
 
             torch.ops.npu.apply_token_bitmask(logits, vocab_mask)
+        elif logits.device.type == "cpu":
+            # Used by the MLX backend, which builds its additive mask rows
+            # on the CPU before inserting them into the lazy graph.
+            from xgrammar import apply_token_bitmask_inplace
+
+            apply_token_bitmask_inplace(logits, vocab_mask, backend="cpu")
         else:
             raise RuntimeError(f"Unsupported device: {logits.device.type}")
 

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

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import Any, Callable, Optional
 
 import mlx.core as mx
 
@@ -12,11 +12,6 @@ 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:
@@ -219,7 +214,10 @@ class MlxAOTKernelContext:
         req_ids: list[str],
         req_pool_idx: dict[str, int],
         req_to_token_pool: Any | None,
-        layer_caches: list[list[ContiguousAttentionKVCache]],
+        # Only .offset is read (absolute on every cache kind) and the slot
+        # lookup is layer-agnostic; the wrapper gates the fused pool scatter
+        # to full-attention layers.
+        layer_caches: list[list[Any]],
     ) -> MlxAOTKernelContext:
         """Build optional AOT context for one batched decode step."""
         if not aot_kernels.rope.enabled or kv_pool is None:

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

@@ -14,6 +14,7 @@ from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
     AttentionOffsetCache,
     ContiguousAttentionKVCache,
     PoolBackedAttentionKVCache,
+    WindowedAttentionKVCache,
     make_attention_mask,
 )
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
@@ -64,4 +65,5 @@ __all__ = [
     "PoolBackedAttentionKVCache",
     "set_context",
     "uses_sliding_window_attention",
+    "WindowedAttentionKVCache",
 ]

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

@@ -68,11 +68,15 @@ def is_attention_module(module: Any) -> bool:
     )
 
 
+def language_model_container(model: Any) -> Any:
+    """The mlx-lm container carrying the layer list and model-level config."""
+    root = getattr(model, "language_model", model)
+    return getattr(root, "model", root)
+
+
 def get_container_window_size(model: Any) -> int | None:
     """The container-level scalar sliding window, if the model declares one."""
-    root = getattr(model, "language_model", model)
-    container = getattr(root, "model", root)
-    return first_present_attr(container, WINDOW_SIZE_ATTRS)
+    return first_present_attr(language_model_container(model), WINDOW_SIZE_ATTRS)
 
 
 def get_layer_window_sizes(model: Any) -> dict[int, int | None]:
@@ -84,10 +88,9 @@ def get_layer_window_sizes(model: Any) -> dict[int, int | None]:
     ``{layer_idx: window or None}``, or ``{}`` when the model does not
     follow the convention.
     """
-    root = getattr(model, "language_model", model)
-    container = getattr(root, "model", root)
+    container = language_model_container(model)
     layer_types = getattr(container, "layer_types", None)
-    window_size = get_container_window_size(model)
+    window_size = first_present_attr(container, WINDOW_SIZE_ATTRS)
     if not layer_types or window_size is None:
         return {}
     return {

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

@@ -21,8 +21,19 @@ def make_attention_mask(N, offset, return_array=False, window_size=None):
     layers pass it, including for N == 1) or windowed models silently fall
     back to full attention.
     """
-    if window_size is not None:
+    if window_size is not None and offset + N > window_size:
         return create_causal_mask(N, offset, window_size=window_size)
+    # Either no window, or a window that cannot bind.  The lowest query
+    # position is ``offset``, so ``offset + N <= window_size`` means every
+    # causally visible key is inside the band and the banded mask is
+    # elementwise identical to a plain causal one -- the shortcut mlx_lm's own
+    # RotatingKVCache.make_mask takes, and these shims stand in for exactly
+    # that cache on sliding layers.  Worth the branch because a materialised
+    # mask forces mx.fast.scaled_dot_product_attention off its fused causal
+    # path: ~2x slower per layer, plus an N x (offset + N) allocation.
+    # It also survives a window-bounded store: offset + N <= window_size
+    # implies N + window_size - 1 >= offset + N, so no keys are dropped and
+    # the mask width still matches the returned key length.
     if N == 1:
         return None
     if return_array:
@@ -141,9 +152,130 @@ class ContiguousAttentionKVCache:
         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, :]
+    def get_kv(self, window: int | None = None) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, min(offset, window), head_dim).
+
+        ``window`` keeps only the trailing window a sliding-window layer can
+        attend to.  Slicing here rather than slicing the full history and then
+        slicing again costs one op instead of two per request per layer.
+        """
+        start = 0 if window is None else max(0, self.offset - window)
+        return (
+            self.keys[:, :, start : self.offset, :],
+            self.values[:, :, start : self.offset, :],
+        )
+
+    def reset(self) -> None:
+        """Reset for reuse, keeping allocated buffers."""
+        self.offset = 0
+
+
+class WindowedAttentionKVCache:
+    """Sliding-window attention KV buffer for one request and one layer.
+
+    Holds the trailing ``window`` tokens plus the in-flight chunk, in
+    temporal order, instead of the full sequence.  ``offset`` stays
+    absolute (RoPE positions, decode bookkeeping); the dropped prefix
+    shows up only in the shorter arrays returned by
+    ``update_and_fetch``/``get_kv`` and in the mask offset ``make_mask``
+    clamps to, so mask width always equals returned key length.
+    """
+
+    __slots__ = ("keys", "values", "offset", "window", "_local")
+
+    def __init__(self, window: int):
+        self.window = window
+        self.keys: mx.array | None = None
+        self.values: mx.array | None = None
+        self.offset = 0  # absolute: every token ever written
+        self._local = 0  # tokens currently in the buffer
+
+    @property
+    def state(self):
+        """Arrays for ``mx.eval`` unpacking."""
+        if self.keys is None:
+            return ()
+        return (self.keys, self.values)
+
+    def reset(self) -> None:
+        """Reset for reuse, keeping allocated buffers."""
+        self.offset = 0
+        self._local = 0
+
+    def make_mask(self, N, return_array=False, window_size=None, **kwargs):
+        kept = min(self._local, self.window)
+        if window_size is None and self.offset > kept:
+            raise RuntimeError(
+                "WindowedAttentionKVCache holds only the trailing window and "
+                "cannot serve a full-context attention mask"
+            )
+        # No N == 1 shortcut here, tempting as it looks: mlx_lm's banded mask is
+        # ``linds < rinds + window_size`` (strict), so a window of W admits
+        # exactly W keys.  Once ``kept == window`` this buffer returns W + 1 of
+        # them -- the trailing window plus the token just written -- and the
+        # oldest must still be masked out.
+        return make_attention_mask(
+            N, kept, return_array=return_array, window_size=window_size
+        )
+
+    def _append(self, keys: mx.array, values: mx.array) -> tuple[int, int]:
+        """Append a chunk in place; return the (start, end) span it serves.
+
+        Split out from ``update_and_fetch`` so the decode path can skip
+        building the two return slices, which its caller discards in
+        favour of ``get_kv``.
+        """
+        S = keys.shape[2]
+        kept = min(self._local, self.window)
+        capacity = self.window + max(S, self.window)
+        held = self.keys.shape[2] if self.keys is not None else 0
+        if self._local + S > held or held > capacity:
+            # Compact the trailing window into a right-sized buffer: this
+            # allocates on the first write, drops history when the buffer
+            # fills (amortised O(1) per decode token), and shrinks back to
+            # 2 * window once an oversized prefill chunk is behind us.
+            B, n_kv_heads, _, head_dim = keys.shape
+            new_k = mx.zeros((B, n_kv_heads, capacity, head_dim), dtype=keys.dtype)
+            new_v = mx.zeros((B, n_kv_heads, capacity, head_dim), dtype=keys.dtype)
+            if kept:
+                src = slice(self._local - kept, self._local)
+                new_k[:, :, :kept, :] = self.keys[:, :, src, :]
+                new_v[:, :, :kept, :] = self.values[:, :, src, :]
+            self.keys, self.values, self._local = new_k, new_v, kept
+        start, end = self._local - kept, self._local + S
+        self.keys[:, :, self._local : end, :] = keys
+        self.values[:, :, self._local : end, :] = values
+        self._local = end
+        self.offset += S
+        return start, end
+
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]:
+        """Append a chunk and return the kept trailing window plus the chunk.
+
+        The kept prefix is ``min(local, window)``, matching what
+        ``make_mask`` clamps to earlier in the same forward pass.
+        """
+        start, end = self._append(keys, values)
+        return self.keys[:, :, start:end, :], self.values[:, :, start:end, :]
+
+    def write_token(self, k: mx.array, v: mx.array) -> None:
+        """Write one token. k, v shape: (1, n_kv_heads, 1, head_dim)."""
+        self._append(k, v)
+
+    def get_kv(self, window: int | None = None) -> tuple[mx.array, mx.array]:
+        """Return buffered trailing K/V: (1, n_kv_heads, kept, head_dim).
+
+        ``window`` mirrors :meth:`ContiguousAttentionKVCache.get_kv`, but the
+        slice is buffer-relative: this buffer holds at most ``2 * window``
+        tokens, so the trailing window starts from ``_local``, not ``offset``.
+        """
+        start = 0 if window is None else max(0, self._local - window)
+        return (
+            self.keys[:, :, start : self._local, :],
+            self.values[:, :, start : self._local, :],
+        )
 
 
 class PoolBackedAttentionKVCache:

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

@@ -1,10 +1,10 @@
 """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.
+The pool stores full-attention layers only and stays intentionally
+uniform: every stored layer must share the same KV shape and
+full-context KV semantics.  Sliding-window layers keep window-bounded
+per-request caches instead; heterogeneous KV shapes need per-layer pools.
 
 Slot 0 is reserved as padding (1-based indexing).
 """

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

@@ -22,6 +22,7 @@ from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
 )
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
     ContiguousAttentionKVCache,
+    WindowedAttentionKVCache,
 )
 
 _thread_local = threading.local()
@@ -34,9 +35,15 @@ class BatchedDecodeContext:
 
     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_layer_caches[cache_idx][req_idx], dense over attention layers.
+    # Windowed caches hold only trailing-window KV, so read them through
+    # write_token/get_kv, never by slicing .keys at an absolute offset.
+    attention_layer_caches: list[
+        list[ContiguousAttentionKVCache | WindowedAttentionKVCache]
+    ]
     attention_pool_index_by_layer: dict[int, int] = field(default_factory=dict)
+    # Dense index into the shared pool's buffers; full-attention layers only.
+    full_kv_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
@@ -50,6 +57,11 @@ class BatchedDecodeContext:
     needs_padding: bool = field(init=False)
     pad_sizes: list[int] = field(init=False)
     positions: Optional[mx.array] = field(init=False)
+    # Padding metadata memo, keyed by window size.  It depends only on
+    # ``seq_lens`` and the window, so every layer sharing a window reuses one
+    # entry instead of rebuilding it (gpt-oss decodes 24 attention layers per
+    # step, in two window classes).
+    _padding_by_window: dict = field(init=False, default_factory=dict)
 
     def __post_init__(self) -> None:
         seq_lens = self.seq_lens
@@ -64,6 +76,56 @@ class BatchedDecodeContext:
             self.attention_pool_index_by_layer = {
                 idx: idx for idx in range(len(self.attention_layer_caches))
             }
+        if self.aot.rope is not None and not self.full_kv_pool_index_by_layer:
+            # No silent default here: the fused scatter addresses pool buffers
+            # by full-attention index, so falling back to the cache index would
+            # write the wrong buffer whenever sliding-window layers are
+            # interleaved. A model with a pool always has full layers to index.
+            raise ValueError(
+                "BatchedDecodeContext requires full_kv_pool_index_by_layer "
+                "when the fused AOT RoPE + pool-scatter kernel is active"
+            )
+
+    def decode_padding(
+        self, window: int | None
+    ) -> tuple[list[int], Optional[mx.array]]:
+        """Right-pad sizes and the keep-mask for one decode step.
+
+        Requests are padded to a common KV width so they can be batched into
+        one SDPA call.  Without a window that width is ``max_len``; a
+        sliding-window layer only reads the trailing ``window`` keys, so its
+        width is ``max(min(seq_len + 1, window))`` instead -- which is why the
+        context's full-length metadata cannot be reused for it.
+
+        The mask is boolean (``True`` keeps the key), broadcast-shaped
+        ``(B, 1, 1, width)``, and ``None`` when no request needs padding.
+        Cached per window: all layers in the step share one build.
+        """
+        cached = self._padding_by_window.get(window, None)
+        if cached is not None:
+            return cached
+
+        if window is None:
+            pad_sizes = self.pad_sizes
+            keep = (
+                self.positions[None, :] < self.valid_lens[:, None]
+                if self.needs_padding
+                else None
+            )
+        else:
+            eff_lens = [min(n + 1, window) for n in self.seq_lens]
+            max_eff = max(eff_lens)
+            pad_sizes = [max_eff - n for n in eff_lens]
+            keep = (
+                mx.arange(max_eff)[None, :]
+                < mx.array(eff_lens, dtype=mx.int32)[:, None]
+                if min(eff_lens) < max_eff
+                else None
+            )
+
+        result = (pad_sizes, None if keep is None else keep[:, None, None, :])
+        self._padding_by_window[window] = result
+        return result
 
     @classmethod
     def from_decode(
@@ -77,10 +139,15 @@ class BatchedDecodeContext:
         req_to_token_pool: Any | None,
         attention_layer_indices: list[int] | None = None,
         attention_pool_index_by_layer: dict[int, int] | None = None,
+        full_kv_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])))
+        # One arbitrary attention layer speaks for the whole step: every
+        # attention cache's ``offset`` is the ABSOLUTE sequence position, so
+        # they all agree even though a windowed cache stores far fewer
+        # tokens than that (see the class docstring's read-through rule).
         seq_lens = [
             caches[i][attention_layer_indices[0]].offset for i in range(batch_size)
         ]
@@ -93,6 +160,7 @@ class BatchedDecodeContext:
             seq_lens=seq_lens,
             attention_layer_caches=attention_layer_caches,
             attention_pool_index_by_layer=attention_pool_index_by_layer or {},
+            full_kv_pool_index_by_layer=full_kv_pool_index_by_layer or {},
             aot=MlxAOTKernelContext.from_decode(
                 aot_kernels=aot_kernels,
                 kv_pool=kv_pool,
@@ -122,9 +190,11 @@ class MLXAttentionWrapper(nn.Module):
     When ``BatchedDecodeContext`` is set, performs per-request RoPE,
     cache writes, and batched SDPA.  Otherwise delegates to inner module.
 
-    ``window_size`` marks a sliding-window layer: the pool keeps the full
-    KV history and the wrapper attends to the trailing window only, which
-    is numerically identical to a rotating cache.
+    ``window_size`` marks a sliding-window layer: the wrapper attends to
+    the trailing window of the cached keys only, which is numerically
+    identical to a rotating cache.  Both cache kinds keep KV in temporal
+    order and report absolute offsets, so the same trailing-window slice
+    works whether the cache holds full history or only the window.
     """
 
     def __init__(
@@ -142,8 +212,27 @@ class MLXAttentionWrapper(nn.Module):
             raise RuntimeError(
                 f"Cannot determine attention scale for {type(inner).__name__}"
             )
+        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__}"
+            )
         object.__setattr__(self, "_scale", scale)
-        object.__setattr__(self, "_sinks", getattr(inner, "sinks", None))
+        object.__setattr__(self, "_n_heads", n_heads)
+        object.__setattr__(self, "_n_kv_heads", n_kv_heads)
+        # None for modules that expose head_dim only through a projection
+        # shape; _batched_decode falls back to the runtime K shape.
+        object.__setattr__(self, "_head_dim", get_head_dim(inner))
+        object.__setattr__(self, "_has_q_norm", hasattr(inner, "q_norm"))
+        object.__setattr__(self, "_has_k_norm", hasattr(inner, "k_norm"))
+        # Only pass sinks when the module has them: the kwarg requires a
+        # recent mlx and must not constrain models without sinks.
+        sinks = getattr(inner, "sinks", None)
+        object.__setattr__(self, "_sinks", sinks)
+        object.__setattr__(
+            self, "_sink_kwargs", {} if sinks is None else {"sinks": sinks}
+        )
 
     def __call__(self, x: mx.array, mask: Any = None, cache: Any = None) -> mx.array:
         ctx = get_context()
@@ -155,18 +244,14 @@ class MLXAttentionWrapper(nn.Module):
         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__}"
-            )
+        n_heads = self._n_heads
+        n_kv_heads = self._n_kv_heads
 
         q_proj_output = inner.q_proj(x)
         keys = inner.k_proj(x)
         values = inner.v_proj(x)
 
-        head_dim = get_head_dim(inner)
+        head_dim = self._head_dim
         if head_dim is None:
             head_dim = keys.shape[-1] // n_kv_heads
 
@@ -188,9 +273,9 @@ class MLXAttentionWrapper(nn.Module):
         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"):
+        if self._has_q_norm:
             queries = inner.q_norm(queries)
-        if hasattr(inner, "k_norm"):
+        if self._has_k_norm:
             keys = inner.k_norm(keys)
 
         queries = queries.transpose(0, 2, 1, 3)
@@ -199,36 +284,32 @@ class MLXAttentionWrapper(nn.Module):
 
         # Vectorized RoPE with per-batch offsets (cached on the context).
         offsets = ctx.offsets
-        attention_pool_idx = ctx.attention_pool_index_by_layer[layer_idx]
+        cache_idx = ctx.attention_pool_index_by_layer[layer_idx]
+        window = self._window_size
 
-        if ctx.aot.rope is not None:
-            # AOT path: real .metallib RoPE + fused KV pool scatter.
+        if ctx.aot.rope is not None and window is None:
+            # AOT path: real .metallib RoPE + fused scatter into this layer's
+            # pool buffer.
             queries, keys = self._rope_custom_aot(
                 queries,
                 keys,
                 values,
                 offsets,
-                attention_pool_idx,
+                ctx.full_kv_pool_index_by_layer[layer_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).
+            # Fallback: MLX's built-in mx.fast.rope. Used when the AOT kernel
+            # isn't built, the model uses an unsupported RoPE variant, or the
+            # layer is sliding-window (windowed KV never enters the pool).
             queries = inner.rope(queries, offset=offsets)
             keys = inner.rope(keys, offset=offsets)
 
-        layer_caches = ctx.attention_layer_caches[attention_pool_idx]
-        window = self._window_size
-        if window is None:
-            pad_sizes = ctx.pad_sizes
-        else:
-            # Sliding-window layer: the cache keeps the full history but the
-            # newest token only attends to the trailing ``window`` keys.  The
-            # padding metadata shared on the context is full-length, so it is
-            # rebuilt locally for the windowed lengths.
-            eff_lens = [min(n + 1, window) for n in ctx.seq_lens]
-            max_eff = max(eff_lens)
-            pad_sizes = [max_eff - n for n in eff_lens]
+        layer_caches = ctx.attention_layer_caches[cache_idx]
+        # A sliding-window layer reads only the trailing ``window`` keys, so its
+        # padded width differs from the unwindowed one.  Both are memoised on
+        # the context and shared by every layer of their kind.
+        pad_sizes, attn_mask = ctx.decode_padding(window)
 
         # TODO: replace per-request loop with native batched/ragged
         # attention once mx.fast.scaled_dot_product_attention supports
@@ -239,10 +320,7 @@ class MLXAttentionWrapper(nn.Module):
         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()
-            if window is not None and k_all.shape[2] > window:
-                k_all = k_all[:, :, -window:, :]
-                v_all = v_all[:, :, -window:, :]
+            k_all, v_all = layer_caches[i].get_kv(window)
 
             pad = pad_sizes[i]
             if pad > 0:
@@ -257,34 +335,13 @@ class MLXAttentionWrapper(nn.Module):
         keys_b = mx.concatenate(all_k, axis=0)
         values_b = mx.concatenate(all_v, axis=0)
 
-        pad_mask = None
-        if window is None:
-            if ctx.needs_padding:
-                pad_mask = ctx.positions[None, :] >= ctx.valid_lens[:, None]
-        elif max(pad_sizes) > 0:
-            eff = mx.array(eff_lens, dtype=mx.int32)
-            pad_mask = mx.arange(max_eff)[None, :] >= eff[:, None]
-
-        attn_mask = None
-        if pad_mask is not None:
-            attn_mask = mx.where(
-                pad_mask[:, None, None, :],
-                mx.array(mx.finfo(queries.dtype).min, dtype=queries.dtype),
-                mx.array(0.0, dtype=queries.dtype),
-            )
-
-        # Only pass sinks when the module has them: the kwarg requires a
-        # recent mlx and must not constrain models without sinks.
-        sink_kwargs = {}
-        if self._sinks is not None:
-            sink_kwargs["sinks"] = self._sinks
         output = mx.fast.scaled_dot_product_attention(
             queries,
             keys_b,
             values_b,
             scale=self._scale,
             mask=attn_mask,
-            **sink_kwargs,
+            **self._sink_kwargs,
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1)
@@ -298,14 +355,14 @@ class MLXAttentionWrapper(nn.Module):
         keys: mx.array,
         values: mx.array,
         positions: mx.array,
-        attention_pool_idx: int,
+        full_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``.
+        rotated K + (untouched) V into ``kv_pool`` buffer ``full_pool_idx``
+        at ``new_token_slots``.
 
         If ``new_token_slots`` is None, slot=-1 sentinel is used (no pool
         write, RoPE-only mode). Returns rotated (queries, keys) in the
@@ -322,8 +379,8 @@ class MLXAttentionWrapper(nn.Module):
         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]
+        k_pool = rope_ctx.kv_pool.k_buffer[full_pool_idx]
+        v_pool = rope_ctx.kv_pool.v_buffer[full_pool_idx]
 
         q_rot, k_rot, k_pool_new, v_pool_new = rope_ctx.kernel.rope_pool_fused(
             q_flat,
@@ -339,8 +396,8 @@ class MLXAttentionWrapper(nn.Module):
             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
+        rope_ctx.kv_pool.k_buffer[full_pool_idx] = k_pool_new
+        rope_ctx.kv_pool.v_buffer[full_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/layout.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Sequence
 
 
@@ -10,9 +10,13 @@ from typing import Any, Sequence
 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.
+    Full-attention layers store softmax-attention KV in the shared MLX
+    attention KV pool; sliding-window layers keep per-request windowed KV
+    only.  Auxiliary layers keep native ``mlx-lm`` cache state.
+
+    ``attention_pool_index_by_layer`` indexes per-request cache arrays
+    (dense over all attention layers); ``full_kv_pool_index_by_layer``
+    indexes the shared pool's buffers (dense over full-attention layers).
     """
 
     layers: tuple[Any, ...]
@@ -20,12 +24,37 @@ class MlxModelCacheLayout:
     attention_layer_indices: tuple[int, ...]
     auxiliary_layer_indices: tuple[int, ...]
     attention_pool_index_by_layer: dict[int, int]
+    # Per-layer sliding window (None or absent = full attention).
+    layer_window_sizes: dict[int, int | None] = field(default_factory=dict)
+    full_attention_layer_indices: tuple[int, ...] = field(init=False)
+    swa_attention_layer_indices: tuple[int, ...] = field(init=False)
+    full_kv_pool_index_by_layer: dict[int, int] = field(init=False)
+
+    def __post_init__(self) -> None:
+        full_indices = tuple(
+            idx
+            for idx in self.attention_layer_indices
+            if self.layer_window_sizes.get(idx) is None
+        )
+        swa_indices = tuple(
+            idx
+            for idx in self.attention_layer_indices
+            if self.layer_window_sizes.get(idx) is not None
+        )
+        object.__setattr__(self, "full_attention_layer_indices", full_indices)
+        object.__setattr__(self, "swa_attention_layer_indices", swa_indices)
+        object.__setattr__(
+            self,
+            "full_kv_pool_index_by_layer",
+            {layer_idx: pool_idx for pool_idx, layer_idx in enumerate(full_indices)},
+        )
 
     @classmethod
     def from_attention_discovery(
         cls,
         layers: Sequence[Any],
         attention_attrs: Sequence[str | None],
+        layer_window_sizes: dict[int, int | None] | None = None,
     ) -> MlxModelCacheLayout:
         if len(layers) != len(attention_attrs):
             raise ValueError(
@@ -50,6 +79,7 @@ class MlxModelCacheLayout:
             attention_layer_indices=attention_layer_indices,
             auxiliary_layer_indices=auxiliary_layer_indices,
             attention_pool_index_by_layer=attention_pool_index_by_layer,
+            layer_window_sizes=dict(layer_window_sizes or {}),
         )
 
     @property
@@ -60,6 +90,14 @@ class MlxModelCacheLayout:
     def num_attention_layers(self) -> int:
         return len(self.attention_layer_indices)
 
+    @property
+    def num_full_attention_layers(self) -> int:
+        return len(self.full_attention_layer_indices)
+
+    @property
+    def has_sliding_window_layers(self) -> bool:
+        return bool(self.swa_attention_layer_indices)
+
     @property
     def has_auxiliary_state(self) -> bool:
         return bool(self.auxiliary_layer_indices)
@@ -70,12 +108,22 @@ class MlxModelCacheLayout:
             raise RuntimeError("MLX model has no supported attention layers")
         return self.attention_layer_indices[0]
 
+    def window_size(self, layer_idx: int) -> int | None:
+        """Sliding window of *layer_idx*, or None for full attention."""
+        return self.layer_window_sizes.get(layer_idx)
+
     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 full_kv_pool_index(self, layer_idx: int) -> int:
+        try:
+            return self.full_kv_pool_index_by_layer[layer_idx]
+        except KeyError as exc:
+            raise KeyError(f"Layer {layer_idx} is not a full-attention layer") from exc
+
     def attention_attr(self, layer_idx: int) -> str:
         attr = self.attention_attrs[layer_idx]
         if attr is None:

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

@@ -5,8 +5,15 @@ scheduler (``TokenToKVPoolAllocator`` / ``RadixCache``).  This runner
 reads cached attention KV from ``MlxAttentionKVPool``, restores any
 native auxiliary layer state, runs the forward pass, and writes the new
 cache state back.  Each request keeps model-shaped cache entries:
-attention layers use ``ContiguousAttentionKVCache`` and auxiliary layers
-use native ``mlx-lm`` cache objects.
+full-attention layers use ``ContiguousAttentionKVCache``, sliding-window
+layers use a fixed-size ``WindowedAttentionKVCache`` on both KV paths,
+and auxiliary layers use native ``mlx-lm`` cache objects.
+
+The shared pool stores full-attention layers only, so no cross-request
+SWA prefix KV exists: a radix prefix hit on a sliding-window model
+recomputes the whole prefix (a trailing-band rebuild is inexact because
+window receptive fields chain backwards through layers).  The
+scheduler's slot bookkeeping is untouched.
 
 The module also exposes a lazy-eval (`*_start` / `*_finalize`) surface
 used by the MLX overlap scheduler to pipeline CPU bookkeeping with
@@ -22,6 +29,7 @@ from dataclasses import dataclass
 from typing import Any
 
 import mlx.core as mx
+import numpy as np
 import psutil
 from mlx.utils import tree_flatten
 from mlx_lm import load as mlx_lm_load
@@ -40,14 +48,29 @@ from sglang.srt.hardware_backend.mlx.kv_cache import (
     MLXAttentionWrapper,
     MlxModelCacheLayout,
     PoolBackedAttentionKVCache,
+    WindowedAttentionKVCache,
     clear_context,
     find_attention_layers,
     get_head_dim,
+    get_layer_window_sizes,
     get_num_kv_heads,
     patch_model_attention,
     set_context,
     uses_sliding_window_attention,
 )
+from sglang.srt.hardware_backend.mlx.sampling import (
+    GREEDY_PARAMS,
+    MlxLazyLogprobs,
+    MlxLogprobSpec,
+    MlxSamplingParams,
+    MlxStepLogprobs,
+    all_greedy,
+    compute_logprobs,
+    lazy_logprob_arrays,
+    sample_tokens,
+    sanitize_logits,
+    scale_by_temperature,
+)
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.runtime_context import (
     mamba_cache_chunk_size,
@@ -72,6 +95,7 @@ class MlxPendingPrefill:
     full_token_ids: list[int]
     req_pool_idx: int
     synced_offset: int
+    lazy_logprobs: MlxLazyLogprobs | None = None
 
 
 @dataclass
@@ -85,9 +109,11 @@ class MlxPendingExtend:
     """
 
     lazy_token: mx.array
+    cache: list[Any]
     req_id: str
     new_token_ids: list[int]
     new_synced_offset: int
+    lazy_logprobs: MlxLazyLogprobs | None = None
 
 
 @dataclass
@@ -104,6 +130,13 @@ class MlxPendingDecode:
     lazy_tokens: mx.array
     req_ids: list[str]
     caches: list[list[Any]]
+    lazy_logprobs: MlxLazyLogprobs | None = None
+    # Carried so a chained step recomputes the same logprob request.
+    logprob_spec: MlxLogprobSpec | None = None
+    # Carried so chained steps keep applying static logit_bias rows.
+    # Never holds a grammar mask on the chained path: grammar batches are
+    # not chain_safe, so their pendings never become a chain root.
+    edit_rows: mx.array | None = None
 
 
 _MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = {
@@ -117,6 +150,14 @@ _MLX_KV_FLOAT_DTYPES = {mx.float16, mx.bfloat16, mx.float32}
 class MlxModelRunner:
     """MLX model runner with radix-cache prefix sharing."""
 
+    # Class defaults cover unit tests that build runners via object.__new__
+    # without running __init__/_load_model, which set the real values.
+    # ``_trunk`` is the headless trunk, resolved in _load_model.
+    _trunk = None
+    _enable_sampling = False
+    _sanitize_nan = False
+    _deterministic_seeding = False
+
     def __init__(
         self,
         model_path: str,
@@ -125,12 +166,27 @@ class MlxModelRunner:
         pool_size: int | None = None,
         mem_fraction_static: float = 0.8,
         quantization: str | None = None,
+        enable_sampling: bool = False,
+        sampling_rng_seed: int = 0,
+        deterministic_seeding: bool = False,
     ):
         self.model_path = model_path
         self.trust_remote_code = trust_remote_code
         self.model = None
         self.disable_radix_cache = disable_radix_cache
         self._mem_fraction_static = mem_fraction_static
+        self._enable_sampling = enable_sampling
+        # --enable-deterministic-inference parity: seed every row (the
+        # sampling module's DEFAULT_SAMPLING_SEED when the request has no
+        # sampling_seed), like the pytorch backend.
+        self._deterministic_seeding = deterministic_seeding
+        self._sanitize_nan = envs.SGLANG_SANITIZE_NAN_LOGITS.get()
+        # RNG state for unseeded sampled rows; split at every sampling
+        # graph build, so runs are reproducible given the same seed and
+        # request schedule.
+        self._rng_key: mx.array | None = (
+            mx.random.key(sampling_rng_seed) if enable_sampling else None
+        )
         # Counter used to trigger periodic mx.clear_cache() calls.
         self._decode_step_ct: int = 0
         self._clear_steps = envs.SGLANG_MLX_CLEAR_CACHE_STEPS.get()
@@ -155,6 +211,8 @@ class MlxModelRunner:
         self._cache_layout = MlxModelCacheLayout.from_attention_discovery(
             layer_list,
             attn_attrs,
+            # Per-layer sliding windows (container convention, e.g. gpt-oss).
+            layer_window_sizes=get_layer_window_sizes(self.model),
         )
         if self._cache_layout.num_attention_layers == 0:
             raise RuntimeError("MLX model has no supported attention layers")
@@ -164,6 +222,16 @@ class MlxModelRunner:
             raise RuntimeError(
                 "MLX models with auxiliary cache state require model.make_cache()."
             )
+        if (
+            self._cache_layout.has_auxiliary_state
+            and self._cache_layout.has_sliding_window_layers
+        ):
+            # Auxiliary-state restore assumes a prefix hit runs only the new
+            # tokens; an SWA prefix hit recomputes the prefix on a fresh cache.
+            raise NotImplementedError(
+                "MLX runner does not support models with both auxiliary "
+                "cache state and sliding-window attention layers."
+            )
         if self._cache_layout.has_auxiliary_state:
             self._model_embed, self._model_norm, self._model_lm_head = (
                 self._extract_model_components()
@@ -172,7 +240,9 @@ class MlxModelRunner:
 
         self._req_caches: dict[str, list[Any]] = {}
         self._req_token_ids: dict[str, list[int]] = {}
-        self._cache_pool: list[list[Any]] = []  # reusable full-attention caches
+        self._req_sampling: dict[str, MlxSamplingParams] = {}
+        # Reusable cache lists, for models without auxiliary layer state.
+        self._cache_pool: list[list[Any]] = []
 
         self._attention_kv_pool: MlxAttentionKVPool | None = None
         self._req_to_token_pool: ReqToTokenPool | None = None
@@ -206,7 +276,12 @@ class MlxModelRunner:
         """Create a model-shaped cache list with attention KV adapters."""
         cache = self._new_cache_skeleton()
         for layer_idx in self._cache_layout.attention_layer_indices:
-            cache[layer_idx] = ContiguousAttentionKVCache(max_seq_len=self._max_seq_len)
+            window = self._cache_layout.window_size(layer_idx)
+            cache[layer_idx] = (
+                WindowedAttentionKVCache(window)
+                if window is not None
+                else ContiguousAttentionKVCache(max_seq_len=self._max_seq_len)
+            )
         return cache
 
     def _acquire_cache(self) -> list[Any]:
@@ -214,7 +289,7 @@ class MlxModelRunner:
         if not self._cache_layout.has_auxiliary_state and self._cache_pool:
             cache = self._cache_pool.pop()
             for c in cache:
-                c.offset = 0
+                c.reset()
             return cache
         return self._new_native_cache()
 
@@ -335,13 +410,18 @@ class MlxModelRunner:
     def _cache_with_pool_backed_attention(
         self, prefix_slot_ids: list[int], prefix_len: int
     ) -> list[Any]:
+        """Build a prefill cache list gathering *prefix_len* pool tokens.
+
+        Only reachable without sliding-window layers (SWA prefix hits
+        recompute instead), so every attention layer here is pool-backed.
+        """
         assert self._attention_kv_pool is not None
         slot_ids_mx = mx.array(prefix_slot_ids, dtype=mx.int32)
         cache = self._new_cache_skeleton()
         for layer_idx in self._cache_layout.attention_layer_indices:
             cache[layer_idx] = PoolBackedAttentionKVCache(
                 self._attention_kv_pool,
-                self._cache_layout.attention_pool_index(layer_idx),
+                self._cache_layout.full_kv_pool_index(layer_idx),
                 slot_ids_mx,
                 prefix_len,
             )
@@ -379,22 +459,15 @@ class MlxModelRunner:
         return arrays
 
     @staticmethod
-    def _eval_with_cache(token_result: mx.array, cache: list[Any]) -> None:
-        """Evaluate token result and all cache buffers in one mx.eval call."""
-        mx.eval(
-            token_result,
-            *[s for c in cache for s in MlxModelRunner._cache_arrays(c)],
-        )
+    def cache_state_arrays(caches: list[list[Any]]) -> list[mx.array]:
+        """Flatten per-request cache lists (``caches[req][layer]``) to arrays.
 
-    @staticmethod
-    def _cache_state_arrays(pending_caches: list[list[Any]]) -> list[mx.array]:
-        """Flatten pending decode cache state list into an array list.
-
-        Safe to hand to ``mx.async_eval``.
+        Pass ``[cache]`` for a single request.  Safe to hand to
+        ``mx.eval`` / ``mx.async_eval``.
         """
         return [
             s
-            for cache_list in pending_caches
+            for cache_list in caches
             for cache in cache_list
             for s in MlxModelRunner._cache_arrays(cache)
         ]
@@ -465,6 +538,17 @@ class MlxModelRunner:
         load_time = time.time() - start_time
         logger.info(f"MLX model loaded in {load_time:.2f}s")
 
+        # mlx-lm models expose the headless trunk as ``Model.model``; without
+        # it, non-final chunked-prefill chunks cannot skip the logit head.
+        trunk = getattr(self.model, "model", None)
+        self._trunk = trunk if callable(trunk) else None
+        if self._trunk is None:
+            logger.info(
+                "Model %s exposes no headless trunk (`.model`); non-final "
+                "chunked-prefill chunks will compute full vocab logits.",
+                type(self.model).__name__,
+            )
+
         # Optional: Path B fusion — keep up_proj/gate_proj weights separate
         # (no matmul-kernel tile regression) but fuse the swiglu activation
         # into the gate matmul via a custom Metal kernel. Activated by
@@ -494,11 +578,13 @@ class MlxModelRunner:
     ) -> tuple[int, int, mx.Dtype]:
         layer = self._cache_layout.layers[layer_idx]
         sample_attn = self._attention_module_for_layer(layer_idx)
-        if uses_sliding_window_attention(layer, sample_attn):
+        unsized_window = self._cache_layout.window_size(layer_idx) is None
+        if unsized_window and uses_sliding_window_attention(layer, sample_attn):
             raise NotImplementedError(
-                "MLX radix attention KV pool does not support sliding-window "
-                f"attention yet at layer {layer_idx}. Sliding-window KV needs "
-                "per-layer/window-aware pools."
+                f"Attention layer {layer_idx} declares sliding-window "
+                "attention but the model exposes no per-layer window map "
+                "(container `layer_types` plus a scalar window), so the MLX "
+                "KV cache cannot bound its sliding-window KV."
             )
         n_kv_heads = get_num_kv_heads(sample_attn)
         if n_kv_heads is None:
@@ -528,7 +614,11 @@ class MlxModelRunner:
         return n_kv_heads, head_dim, dtype
 
     def _get_attn_config(self) -> tuple[int, int, mx.Dtype]:
-        """Return the uniform attention KV config used by the shared MLX pool."""
+        """Return the uniform KV config shared by every attention layer.
+
+        Sizes the shared pool and the AOT kernels; sliding-window layers
+        must match the same shape because they share the decode kernels.
+        """
         if self._cache_layout.num_attention_layers == 0:
             raise RuntimeError(
                 "Cannot determine attention config: no attention module found"
@@ -539,12 +629,11 @@ class MlxModelRunner:
             config = self._attention_kv_config_for_layer(layer_idx)
             if config != first_config:
                 raise NotImplementedError(
-                    "MLX radix attention KV pool requires uniform softmax-attention "
+                    "MLX attention KV caching requires uniform softmax-attention "
                     "KV shape across layers. "
                     f"Layer {first_layer_idx} has {first_config}, "
                     f"but layer {layer_idx} has {config}. "
-                    "Heterogeneous attention KV or sliding-window KV needs "
-                    "per-layer pools."
+                    "Heterogeneous attention KV needs per-layer pools."
                 )
         return first_config
 
@@ -553,7 +642,13 @@ class MlxModelRunner:
         if explicit_size is not None:
             return explicit_size
         n_kv_heads, head_dim, dtype = self._get_attn_config()
-        num_layers = self._cache_layout.num_attention_layers
+        # Only full-attention layers occupy pool slots. All-SWA models have no
+        # pool at all and fall back to the all-layer formula purely to keep the
+        # scheduler's token budget finite.
+        num_layers = (
+            self._cache_layout.num_full_attention_layers
+            or self._cache_layout.num_attention_layers
+        )
         sys_available = psutil.virtual_memory().available
         mlx_limit = mx.device_info().get(
             "max_recommended_working_set_size",
@@ -598,15 +693,44 @@ class MlxModelRunner:
         )
 
     def init_cache_pools(self, req_to_token_pool: ReqToTokenPool | None) -> None:
-        """Create attention KV pool (+1 for padding slot 0)."""
+        """Create the full-attention KV pool (+1 for padding slot 0)."""
         self._req_to_token_pool = req_to_token_pool
         if self.disable_radix_cache:
             return
+        num_pool_layers = self._cache_layout.num_full_attention_layers
+        if self._cache_layout.has_sliding_window_layers:
+            # The pool exists to serve radix prefix hits, and an SWA prefix hit
+            # recomputes the prefix instead of gathering it (see prefill_start),
+            # so on any SWA model the pool has no reader: its sole consumer is
+            # PoolBackedAttentionKVCache, reachable only when
+            # trusted_prefix_len > 0, which requires no SWA layers. Allocating
+            # it anyway would burn the whole auto-sized KV budget
+            # (_compute_pool_size fills mem_fraction_static) on a buffer that is
+            # only ever written. Skipping it also disables the fused AOT
+            # RoPE + pool-scatter kernel for the full layers, whose scatter half
+            # is dead work here; that kernel is opt-in
+            # (SGLANG_MLX_USE_CUSTOM_ROPE, default off), so the default path
+            # loses nothing.
+            #
+            # Un-gate this together with the window-aware shared SWA pool that
+            # restores fast prefix hits: the seams it needs (the layout
+            # partition, the full-pool index, layer-type dispatch) are already
+            # in place.
+            logger.info(
+                "Model has %d sliding-window attention layers; skipping the "
+                "shared attention KV pool (an SWA prefix hit recomputes the "
+                "prefix, so the pool would never be read). Per-request "
+                "windowed caches only.",
+                len(self._cache_layout.swa_attention_layer_indices),
+            )
+            return
+        if num_pool_layers == 0:
+            return
         n_kv_heads, head_dim, dtype = self._get_attn_config()
         # +1 for padding slot 0
         self._attention_kv_pool = MlxAttentionKVPool(
             pool_size=self._pool_size + 1,
-            num_layers=self._cache_layout.num_attention_layers,
+            num_layers=num_pool_layers,
             n_kv_heads=n_kv_heads,
             head_dim=head_dim,
             dtype=dtype,
@@ -614,7 +738,9 @@ class MlxModelRunner:
         logger.info(
             f"Attention KV pool initialized: pool_size={self._pool_size} "
             f"(buffer size {self._pool_size + 1} incl. padding slot 0), "
-            f"{self._cache_layout.num_attention_layers} attention layers, "
+            f"{num_pool_layers} full-attention layers "
+            f"({len(self._cache_layout.swa_attention_layer_indices)} "
+            "sliding-window layers stay per-request), "
             f"{n_kv_heads} kv_heads, {head_dim} head_dim"
         )
 
@@ -627,8 +753,14 @@ class MlxModelRunner:
         new_slot_ids: list[int],
         req_pool_idx: int,
         req: Any | None = None,
+        needs_logits: bool = True,
     ) -> int:
-        """Prefill a request.  Returns next_token_id."""
+        """Prefill a request.  Returns next_token_id.
+
+        One-shot convenience wrapper around ``prefill_start`` /
+        ``prefill_finalize``; logit edits and logprobs are only available
+        through that lazy surface.
+        """
         pending = self.prefill_start(
             req_id=req_id,
             new_token_ids=new_token_ids,
@@ -637,8 +769,9 @@ class MlxModelRunner:
             new_slot_ids=new_slot_ids,
             req_pool_idx=req_pool_idx,
             req=req,
+            needs_logits=needs_logits,
         )
-        self._eval_with_cache(pending.lazy_token, pending.cache)
+        self.eval_pending(pending)
         return self.prefill_finalize(pending)
 
     def extend(
@@ -646,10 +779,16 @@ class MlxModelRunner:
         req_id: str,
         new_token_ids: list[int],
         new_slot_ids: list[int],
+        needs_logits: bool = True,
     ) -> int:
-        """Continue prefill for a chunked request.  Returns next_token_id."""
-        pending = self.extend_start(req_id, new_token_ids, new_slot_ids)
-        self._eval_with_cache(pending.lazy_token, self._req_caches[req_id])
+        """Continue prefill for a chunked request.  Returns next_token_id.
+
+        One-shot convenience wrapper; see :meth:`prefill`.
+        """
+        pending = self.extend_start(
+            req_id, new_token_ids, new_slot_ids, needs_logits=needs_logits
+        )
+        self.eval_pending(pending)
         return self.extend_finalize(pending)
 
     def _sync_new_kv_to_pool(
@@ -658,9 +797,15 @@ class MlxModelRunner:
         cache_start: int,
         slot_ids: list[int],
     ) -> None:
-        """Sync attention KV from contiguous cache to pool at the given slots."""
+        """Sync full-attention KV from contiguous caches to the pool slots.
+
+        Sliding-window layers are skipped: they keep no pool KV, and their
+        buffers are window-local so the absolute slicing below would not
+        apply to them anyway.
+        """
         if not slot_ids or self._attention_kv_pool is None:
             return
+        full_layer_indices = self._cache_layout.full_attention_layer_indices
         end = cache_start + len(slot_ids)
         slot_ids_mx = mx.array(slot_ids, dtype=mx.int32)
         # TODO: Standardize ContiguousAttentionKVCache size to avoid transpose
@@ -668,13 +813,13 @@ class MlxModelRunner:
         k_all = mx.stack(
             [
                 cache[layer_idx].keys[0, :, cache_start:end, :].transpose(1, 0, 2)
-                for layer_idx in self._cache_layout.attention_layer_indices
+                for layer_idx in full_layer_indices
             ]
         )
         v_all = mx.stack(
             [
                 cache[layer_idx].values[0, :, cache_start:end, :].transpose(1, 0, 2)
-                for layer_idx in self._cache_layout.attention_layer_indices
+                for layer_idx in full_layer_indices
             ]
         )
         self._attention_kv_pool.set_kv_all_layers(slot_ids_mx, k_all, v_all)
@@ -711,17 +856,13 @@ class MlxModelRunner:
         for req_id in list(self._req_caches.keys()):
             self._sync_decode_kv_to_pool(req_id)
 
-    def decode_batch(
-        self,
-        req_ids: list[str],
-    ) -> list[int]:
-        """Decode one token per request."""
+    def decode_batch(self, req_ids: list[str]) -> list[int]:
+        """Decode one token per request.
+
+        One-shot convenience wrapper; see :meth:`prefill`.
+        """
         pending = self.decode_batch_start(req_ids)
-        # Evaluate lazy_tokens together with every affected cache buffer so
-        # the attention write-then-read ordering is materialised in one
-        # kernel submission.
-        cache_arrays = self._cache_state_arrays(pending.caches)
-        mx.eval(pending.lazy_tokens, *cache_arrays)
+        self.eval_pending(pending)
         return self.decode_batch_finalize(pending)
 
     def prefill_start(
@@ -733,6 +874,9 @@ class MlxModelRunner:
         new_slot_ids: list[int],
         req_pool_idx: int,
         req: Any | None = None,
+        needs_logits: bool = True,
+        logit_edit_row: mx.array | None = None,
+        logprob_spec: MlxLogprobSpec | None = None,
     ) -> MlxPendingPrefill:
         """Queue a prefill forward pass without evaluating.
 
@@ -740,17 +884,28 @@ class MlxModelRunner:
         next-token ``mx.array`` plus everything needed to commit the
         request in :meth:`prefill_finalize`.  The caller drives the GPU
         by handing ``lazy_token`` (and cache state) to ``mx.async_eval``.
+
+        ``needs_logits=False`` marks the first chunk of a chunked prompt
+        (its next-token output is discarded); see :meth:`extend_start`.
         """
         prefix_len = len(prefix_slot_ids)
         if req is not None:
             req.mamba_last_track_seqlen = None
+        if self._enable_sampling:
+            self._req_sampling[req_id] = (
+                MlxSamplingParams.from_req(
+                    req, deterministic_seeding=self._deterministic_seeding
+                )
+                if req is not None
+                else GREEDY_PARAMS
+            )
 
         if self.disable_radix_cache:
             cache = self._acquire_cache()
             input_ids = mx.array([new_token_ids], dtype=mx.int32)
-            model_output = self.model(input_ids, cache=cache)
-            logits = self._extract_logits(model_output)
-            lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
+            lazy_token, lazy_logprobs = self._forward_lazy_token(
+                input_ids, cache, needs_logits, req_id, logit_edit_row, logprob_spec
+            )
             return MlxPendingPrefill(
                 lazy_token=lazy_token,
                 cache=cache,
@@ -758,9 +913,17 @@ class MlxModelRunner:
                 full_token_ids=list(full_token_ids),
                 req_pool_idx=req_pool_idx,
                 synced_offset=0,
+                lazy_logprobs=lazy_logprobs,
             )
 
-        assert self._attention_kv_pool is not None
+        # A pool is required only where one can actually be read: a model with
+        # full-attention layers and no sliding-window layers.  init_cache_pools
+        # skips it otherwise, and the gather path below is unreachable then.
+        assert (
+            self._attention_kv_pool is not None
+            or self._cache_layout.num_full_attention_layers == 0
+            or self._cache_layout.has_sliding_window_layers
+        )
 
         new_token_count = len(new_token_ids)
         track_len = self._select_auxiliary_state_track_len(
@@ -770,8 +933,22 @@ class MlxModelRunner:
             req=req,
         )
 
-        if prefix_len > 0:
-            cache = self._cache_with_pool_backed_attention(prefix_slot_ids, prefix_len)
+        # Sliding-window layers keep no pool KV, so a prefix hit has nothing to
+        # gather and re-runs the prefix. A trailing-band rebuild would not be
+        # exact: each rebuilt position needs its own window of exact hidden
+        # states, and that dependency chains back through every layer. Only the
+        # run is clamped -- slot ids and synced offsets stay unclamped.
+        if prefix_len > 0 and self._cache_layout.has_sliding_window_layers:
+            trusted_prefix_len = 0
+            run_token_ids = list(full_token_ids[:prefix_len]) + new_token_ids
+        else:
+            trusted_prefix_len = prefix_len
+            run_token_ids = new_token_ids
+
+        if trusted_prefix_len > 0:
+            cache = self._cache_with_pool_backed_attention(
+                prefix_slot_ids[:trusted_prefix_len], trusted_prefix_len
+            )
             pool_backed_attention = True
             restored_auxiliary_state = (
                 not self._cache_layout.has_auxiliary_state
@@ -787,9 +964,9 @@ class MlxModelRunner:
                 # allocated attention KV below.
                 cache = self._acquire_cache()
                 input_ids = mx.array([full_token_ids or new_token_ids], dtype=mx.int32)
-                model_output = self.model(input_ids, cache=cache)
-                logits = self._extract_logits(model_output)
-                lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
+                lazy_token, lazy_logprobs = self._forward_lazy_token(
+                    input_ids, cache, needs_logits, req_id, logit_edit_row, logprob_spec
+                )
                 if new_slot_ids:
                     self._sync_new_kv_to_pool(cache, prefix_len, new_slot_ids)
                 return MlxPendingPrefill(
@@ -799,39 +976,45 @@ class MlxModelRunner:
                     full_token_ids=list(full_token_ids),
                     req_pool_idx=req_pool_idx,
                     synced_offset=prefix_len + len(new_slot_ids),
+                    lazy_logprobs=lazy_logprobs,
                 )
         else:
             cache = self._acquire_cache()
             pool_backed_attention = False
 
-        if new_token_count > 0:
+        if run_token_ids:
             track_new_count = track_len - prefix_len if track_len is not None else None
             if track_new_count is not None and 0 < track_new_count < new_token_count:
+                # aux + SWA is rejected at init, so run_token_ids is
+                # new_token_ids on this branch.
                 input_ids = mx.array([new_token_ids[:track_new_count]], dtype=mx.int32)
-                self.model(input_ids, cache=cache)
+                # Cache side effects only — this intermediate forward's
+                # output is never read, so skip the head when possible.
+                if self._trunk_forward(input_ids, cache) is None:
+                    self.model(input_ids, cache=cache)
                 self._store_tracked_auxiliary_state(req, cache, track_len)
                 if pool_backed_attention:
                     cache = self._materialize_pool_backed_attention(cache)
                     pool_backed_attention = False
                 extend_tokens = new_token_ids[track_new_count:]
             else:
-                extend_tokens = new_token_ids
+                extend_tokens = run_token_ids
         else:
-            # Full cache hit - rerun last token to get next-token logits
+            # Full cache hit - rerun last token to get next-token logits.
+            # Unreachable with SWA layers: a prefix rebuild always leaves run
+            # tokens whose final logits already predict the next token.
             extend_tokens = full_token_ids[-1:]
             for c in cache:
                 c.offset = max(c.offset - 1, 0)
 
         input_ids = mx.array([extend_tokens], dtype=mx.int32)
-        model_output = self.model(input_ids, cache=cache)
-        logits = self._extract_logits(model_output)
+        lazy_token, lazy_logprobs = self._forward_lazy_token(
+            input_ids, cache, needs_logits, req_id, logit_edit_row, logprob_spec
+        )
 
         if track_len is not None and track_len == prefix_len + new_token_count:
             self._store_tracked_auxiliary_state(req, cache, track_len)
 
-        last_logits = logits[:, -1, :]
-        lazy_token = mx.argmax(last_logits, axis=-1)
-
         # Convert pool-backed attention KV to contiguous attention KV for decode.
         # This appends a lazy slice-assign onto the forward graph; the
         # arrays get materialised when the caller evaluates lazy_token.
@@ -848,6 +1031,7 @@ class MlxModelRunner:
             full_token_ids=list(full_token_ids),
             req_pool_idx=req_pool_idx,
             synced_offset=prefix_len + len(new_slot_ids),
+            lazy_logprobs=lazy_logprobs,
         )
 
     def prefill_finalize(self, pending: MlxPendingPrefill) -> int:
@@ -872,8 +1056,16 @@ class MlxModelRunner:
         req_id: str,
         new_token_ids: list[int],
         new_slot_ids: list[int],
+        needs_logits: bool = True,
+        logit_edit_row: mx.array | None = None,
+        logprob_spec: MlxLogprobSpec | None = None,
     ) -> MlxPendingExtend:
-        """Queue chunked-prefill continuation without evaluating."""
+        """Queue chunked-prefill continuation without evaluating.
+
+        ``needs_logits=False`` marks a non-final chunk whose next-token
+        output the scheduler discards; the logit head is skipped when the
+        model exposes a headless trunk.
+        """
         assert (
             req_id in self._req_caches
         ), f"extend_start called for unknown request {req_id}"
@@ -881,9 +1073,9 @@ class MlxModelRunner:
         cache = self._req_caches[req_id]
 
         input_ids = mx.array([new_token_ids], dtype=mx.int32)
-        model_output = self.model(input_ids, cache=cache)
-        logits = self._extract_logits(model_output)
-        lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
+        lazy_token, lazy_logprobs = self._forward_lazy_token(
+            input_ids, cache, needs_logits, req_id, logit_edit_row, logprob_spec
+        )
 
         if not self.disable_radix_cache and new_slot_ids:
             synced = self._req_synced_offset[req_id]
@@ -894,9 +1086,11 @@ class MlxModelRunner:
 
         return MlxPendingExtend(
             lazy_token=lazy_token,
+            cache=cache,
             req_id=req_id,
             new_token_ids=list(new_token_ids),
             new_synced_offset=new_synced_offset,
+            lazy_logprobs=lazy_logprobs,
         )
 
     def extend_finalize(self, pending: MlxPendingExtend) -> int:
@@ -916,6 +1110,201 @@ class MlxModelRunner:
         )
         return next_token
 
+    def _trunk_forward(self, input_ids: mx.array, cache: list[Any]) -> mx.array | None:
+        """Run the model WITHOUT its logit head, for cache side effects only.
+
+        Non-final chunked-prefill chunks discard their next-token output
+        (``extend_finalize`` pops it), yet the full model call still computes
+        vocab-sized float32 logits for every chunk position — for a 200K-vocab
+        model that is ~100x the useful head work and the largest transient
+        allocation in the process.  ``self._trunk`` is resolved once at load;
+        returns None when the model exposes no headless trunk (caller falls
+        back to the full forward).
+        """
+        if self._trunk is None:
+            return None
+        return self._trunk(input_ids, cache=cache)
+
+    def _forward_lazy_token(
+        self,
+        input_ids: mx.array,
+        cache: list[Any],
+        needs_logits: bool,
+        req_id: str,
+        logit_edit_row: mx.array | None = None,
+        logprob_spec: MlxLogprobSpec | None = None,
+    ) -> tuple[mx.array, MlxLazyLogprobs | None]:
+        """Forward one chunk, returning (lazy next-token, lazy logprobs).
+
+        Skips the logit head for discarded-output chunks when possible.
+        """
+        if not needs_logits:
+            hidden = self._trunk_forward(input_ids, cache)
+            if hidden is not None:
+                return self._dummy_next_token(hidden), None
+            # Headless trunk unavailable: run the full model, but keep the
+            # discarded token on argmax — sampling here would consume RNG
+            # state and make final output depend on prefill chunking.
+            model_output = self.model(input_ids, cache=cache)
+            logits = self._extract_logits(model_output)
+            return mx.argmax(logits[:, -1, :], axis=-1), None
+        model_output = self.model(input_ids, cache=cache)
+        logits = self._extract_logits(model_output)
+        edits = logit_edit_row[None, :] if logit_edit_row is not None else None
+        return self._select_tokens_with_logprobs(
+            logits[:, -1, :], [req_id], [cache], edits, logprob_spec
+        )
+
+    def _select_tokens_with_logprobs(
+        self,
+        last_logits: mx.array,
+        req_ids: list[str],
+        caches: list[list[Any]],
+        edit_rows: mx.array | None = None,
+        logprob_spec: MlxLogprobSpec | None = None,
+    ) -> tuple[mx.array, MlxLazyLogprobs | None]:
+        """Pick one token per row of ``last_logits`` — lazily, inside the graph.
+
+        Greedy behavior (sampling disabled, or every row greedy with no
+        logit edits) is exactly the pre-sampling ``mx.argmax`` and consumes
+        no RNG state.  ``edit_rows`` is the worker's pre-combined additive
+        [B, vocab] array (grammar mask + logit_bias), applied before token
+        selection and logprobs, mirroring the CUDA
+        ``ModelRunner._preprocess_logits`` order.  Positions for seeded rows
+        come from the attention cache offsets, which the just-built forward
+        has already advanced past the token being sampled — the same
+        ``seq_len - 1`` the pytorch path feeds its sampler.  They are
+        build-time Python ints, so this is chained-decode safe.
+        """
+        if not self._enable_sampling:
+            return mx.argmax(last_logits, axis=-1), None
+        params = [self._req_sampling[rid] for rid in req_ids]
+        edited = self._edited_logits(last_logits, edit_rows)
+        greedy = all_greedy(params)
+        # Built once and shared: sampling and logprobs both start from
+        # logits/temperature, and MLX does not CSE the two identical graphs.
+        # Stays None when neither needs it (the greedy, no-logprob path).
+        scaled = (
+            scale_by_temperature(edited, params)
+            if not greedy or logprob_spec is not None
+            else None
+        )
+        if greedy:
+            tokens = mx.argmax(edited, axis=-1)
+        else:
+            positions = [self._first_attention_cache(c).offset - 1 for c in caches]
+            self._rng_key, key = mx.random.split(self._rng_key)
+            tokens = sample_tokens(
+                last_logits=edited,
+                params=params,
+                positions=positions,
+                key=key,
+                scaled=scaled,
+            )
+        lazy_logprobs = (
+            compute_logprobs(
+                last_logits=edited,
+                params=params,
+                tokens=tokens,
+                spec=logprob_spec,
+                scaled=scaled,
+            )
+            if logprob_spec is not None
+            else None
+        )
+        return tokens, lazy_logprobs
+
+    def _edited_logits(
+        self, last_logits: mx.array, edit_rows: mx.array | None
+    ) -> mx.array:
+        """Apply the additive logit edits and env-gated NaN sanitization."""
+        edited = last_logits
+        if edit_rows is not None:
+            # The edit rows are sized from SamplingBatchInfo.vocab_size while
+            # these logits come from the model's lm_head.  A model whose head
+            # is padded past the tokenizer vocabulary would otherwise fail as
+            # an opaque broadcast error deep in the lazy graph.
+            if edit_rows.shape[-1] != last_logits.shape[-1]:
+                raise RuntimeError(
+                    "Logit edit rows do not match the model's vocabulary: "
+                    f"sampling_info.vocab_size={edit_rows.shape[-1]} vs "
+                    f"lm_head width {last_logits.shape[-1]}"
+                )
+            edited = edited.astype(mx.float32) + edit_rows
+        if self._sanitize_nan:
+            edited = sanitize_logits(edited.astype(mx.float32))
+        return edited
+
+    def _run_logits_hook(self, last_logits: mx.array, logits_hook) -> mx.array:
+        """Materialize logits and let the worker edit them on the CPU.
+
+        Used for custom logit processors (arbitrary torch callables) — the
+        one edit that cannot be expressed lazily.  Synchronizes the graph;
+        callers gate this to fresh, pure-decode launches, so the chained
+        overlap pipeline never pays for it.
+        """
+        logits32 = last_logits.astype(mx.float32)
+        mx.eval(logits32)
+        return mx.array(logits_hook(np.array(logits32)))
+
+    def collect_logprobs(
+        self, lazy_logprobs: MlxLazyLogprobs | None
+    ) -> MlxStepLogprobs | None:
+        """Materialize one step's lazy logprob arrays into Python lists."""
+        if lazy_logprobs is None:
+            return None
+        mx.eval(*lazy_logprob_arrays(lazy_logprobs))
+        spec = lazy_logprobs.spec
+        if lazy_logprobs.top_val is not None:
+            all_val = lazy_logprobs.top_val.tolist()
+            all_idx = lazy_logprobs.top_idx.tolist()
+            top_val = [all_val[i][:k] for i, k in enumerate(spec.top_ks)]
+            top_idx = [all_idx[i][:k] for i, k in enumerate(spec.top_ks)]
+        else:
+            top_val = [[] for _ in spec.top_ks]
+            top_idx = [[] for _ in spec.top_ks]
+        return MlxStepLogprobs(
+            chosen=lazy_logprobs.chosen.tolist(),
+            top_val=top_val,
+            top_idx=top_idx,
+            token_ids_val=[
+                a.tolist() if a is not None else [] for a in lazy_logprobs.token_ids_val
+            ],
+            token_ids_idx=[list(ids) if ids else [] for ids in spec.token_ids],
+        )
+
+    def eval_pending(
+        self, pending: MlxPendingPrefill | MlxPendingExtend | MlxPendingDecode
+    ) -> None:
+        """Materialize a queued forward: token(s), cache writes and logprobs.
+
+        One ``mx.eval`` for the whole pending, so the attention
+        write-then-read ordering is materialised in a single kernel
+        submission.  Prefill and extend carry one request's per-layer
+        cache; a decode carries one cache list per request.
+        """
+        if isinstance(pending, MlxPendingDecode):
+            tokens, caches = pending.lazy_tokens, pending.caches
+        else:
+            tokens, caches = pending.lazy_token, [pending.cache]
+        mx.eval(
+            tokens,
+            *self.cache_state_arrays(caches),
+            *lazy_logprob_arrays(pending.lazy_logprobs),
+        )
+
+    @staticmethod
+    def _dummy_next_token(hidden: mx.array) -> mx.array:
+        """Graph-connected placeholder token for a skipped-head chunk.
+
+        Value is always 0 (a valid vocab id); it is appended and then popped
+        as the "stale intermediate token" by the next chunk's finalize.
+        Deriving it from ``hidden`` keeps the trunk in the lazy graph handed
+        to ``mx.eval``/``mx.async_eval`` (cache arrays are also evaluated
+        explicitly by both call paths).
+        """
+        return (hidden[:, -1, 0] * 0).astype(mx.int32)
+
     def _extract_model_components(self):
         """Cache embedding, norm, and lm_head for layer-by-layer hybrid forward."""
         root = getattr(self.model, "language_model", self.model)
@@ -936,7 +1325,7 @@ class MlxModelRunner:
         batched_input: mx.array,
         req_ids: list[str],
     ) -> mx.array:
-        """Layer-by-layer hybrid decode for attention plus auxiliary state.
+        """Layer-by-layer hybrid decode; returns [B, vocab] last-token logits.
 
         Attention layers run with batched hidden states via
         ``BatchedDecodeContext``. Auxiliary layers run batched when their
@@ -971,7 +1360,7 @@ class MlxModelRunner:
 
         hidden_states = self._model_norm(hidden_states)
         logits = self._extract_logits(self._model_lm_head(hidden_states))
-        return mx.argmax(logits[:, -1, :], axis=-1)
+        return logits[:, -1, :]
 
     def _decode_auxiliary_layer(
         self,
@@ -1107,15 +1496,15 @@ class MlxModelRunner:
         caches: list[list[Any]],
         input_ids_by_request: list[mx.array],
     ) -> mx.array:
-        lazy_token_list = []
+        lazy_logits_list = []
         for input_ids, cache in zip(input_ids_by_request, caches):
             model_output = self.model(input_ids, cache=cache)
             logits = self._extract_logits(model_output)
-            lazy_token_list.append(mx.argmax(logits[:, -1, :], axis=-1))
+            lazy_logits_list.append(logits[:, -1, :])
         return (
-            lazy_token_list[0]
-            if len(lazy_token_list) == 1
-            else mx.concatenate(lazy_token_list, axis=0)
+            lazy_logits_list[0]
+            if len(lazy_logits_list) == 1
+            else mx.concatenate(lazy_logits_list, axis=0)
         )
 
     def _decode_with_batched_attention(
@@ -1135,7 +1524,7 @@ class MlxModelRunner:
             ]
             model_output = self.model(batched_input, cache=shim_cache)
             logits = self._extract_logits(model_output)
-            return mx.argmax(logits[:, -1, :], axis=-1)
+            return logits[:, -1, :]
         finally:
             clear_context()
 
@@ -1156,9 +1545,16 @@ class MlxModelRunner:
             attention_pool_index_by_layer=(
                 self._cache_layout.attention_pool_index_by_layer
             ),
+            full_kv_pool_index_by_layer=self._cache_layout.full_kv_pool_index_by_layer,
         )
 
-    def decode_batch_start(self, req_ids: list[str]) -> MlxPendingDecode:
+    def decode_batch_start(
+        self,
+        req_ids: list[str],
+        edit_rows: mx.array | None = None,
+        logprob_spec: MlxLogprobSpec | None = None,
+        logits_hook=None,
+    ) -> MlxPendingDecode:
         """Queue a decode forward pass without evaluating.
 
         The caller is responsible for calling ``mx.async_eval`` on the
@@ -1170,18 +1566,31 @@ class MlxModelRunner:
         batched_input = mx.array(last_tokens, dtype=mx.int32)[:, None]
 
         if self._cache_layout.has_auxiliary_state:
-            lazy_tokens = self._decode_with_hybrid_batching(
+            last_logits = self._decode_with_hybrid_batching(
                 caches, batched_input, list(req_ids)
             )
         else:
-            lazy_tokens = self._decode_with_batched_attention(
+            last_logits = self._decode_with_batched_attention(
                 caches, batched_input, list(req_ids)
             )
 
+        if logits_hook is not None:
+            # CUDA edit order: grammar mask + logit_bias first, custom
+            # processors second, sanitization last (inside selection).
+            if edit_rows is not None:
+                last_logits = last_logits.astype(mx.float32) + edit_rows
+                edit_rows = None
+            last_logits = self._run_logits_hook(last_logits, logits_hook)
+        lazy_tokens, lazy_logprobs = self._select_tokens_with_logprobs(
+            last_logits, list(req_ids), caches, edit_rows, logprob_spec
+        )
         return MlxPendingDecode(
             lazy_tokens=lazy_tokens,
             req_ids=list(req_ids),
             caches=caches,
+            lazy_logprobs=lazy_logprobs,
+            logprob_spec=logprob_spec,
+            edit_rows=edit_rows,
         )
 
     def decode_batch_start_chained(
@@ -1215,18 +1624,24 @@ class MlxModelRunner:
         # be written at in step N+1 (and equivalently the RoPE offset).
         batched_input = prev.lazy_tokens[:, None]
         if self._cache_layout.has_auxiliary_state:
-            lazy_tokens = self._decode_with_hybrid_batching(
+            last_logits = self._decode_with_hybrid_batching(
                 caches, batched_input, prev.req_ids
             )
         else:
-            lazy_tokens = self._decode_with_batched_attention(
+            last_logits = self._decode_with_batched_attention(
                 caches, batched_input, prev.req_ids
             )
 
+        lazy_tokens, lazy_logprobs = self._select_tokens_with_logprobs(
+            last_logits, prev.req_ids, caches, prev.edit_rows, prev.logprob_spec
+        )
         return MlxPendingDecode(
             lazy_tokens=lazy_tokens,
             req_ids=prev.req_ids,
             caches=caches,
+            lazy_logprobs=lazy_logprobs,
+            logprob_spec=prev.logprob_spec,
+            edit_rows=prev.edit_rows,
         )
 
     def decode_batch_finalize(
@@ -1266,6 +1681,7 @@ class MlxModelRunner:
             self._sync_decode_kv_to_pool(req_id)
 
         self._req_token_ids.pop(req_id, None)
+        self._req_sampling.pop(req_id, None)
         cache = self._req_caches.pop(req_id, None)
         if cache is not None:
             self._release_cache(cache)
@@ -1275,6 +1691,7 @@ class MlxModelRunner:
     def clear(self):
         """Clear all request states."""
         self._req_token_ids.clear()
+        self._req_sampling.clear()
         for cache in self._req_caches.values():
             self._release_cache(cache)
         self._req_caches.clear()

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

@@ -0,0 +1,491 @@
+"""MLX-native in-graph sampling for the MLX backend.
+
+Token selection (temperature / top-k / top-p / min-p / per-request seed)
+built entirely from ``mx`` ops, so it lives inside the same lazy graph as
+the forward pass.  This is what lets sampling coexist with the overlap
+scheduler: ``decode_batch_start_chained`` feeds step N's still-unevaluated
+sampled tokens as step N+1's input ids, exactly as it does for greedy
+argmax, and the GPU runs both steps back-to-back with no host sync.
+
+An earlier proposal (#25804) bridged MLX logits to the CPU pytorch
+``Sampler`` instead.  That design forces a host sync in the middle of the
+graph-build window, which is precisely what the overlap scheduler exists
+to avoid, so token selection is rebuilt from ``mx`` ops here.
+
+Semantics mirror the sglang pytorch sampling backend
+(``top_k_top_p_min_p_sampling_from_probs_torch`` /
+``multinomial_with_seed`` in ``sglang/srt/layers/sampler.py``):
+
+* ``probs = softmax(logits / temperature)`` per row.
+* Descending sort, then zero out rank >= top_k, cumulative-prob mass
+  beyond top_p (the top token is always kept), and probs below
+  ``max_prob * min_p``.
+* Multinomial sampling via the Gumbel-max identity:
+  ``argmax(log(weights) + gumbel_noise)`` over the masked, unnormalized
+  weights is distributed identically to ``torch.multinomial(weights)``
+  (normalization only shifts ``log`` by a per-row constant).
+* When every row asks for a small enough ``top_k``
+  (:data:`MAX_BOUNDED_TOP_K`), everything after the sort runs on the
+  ``[B, K]`` candidates instead of ``[B, vocab]``.  Weights are zero
+  outside those K, so their ``log`` is ``-inf`` and the full-vocab
+  argmax could never have picked them — same token, a fraction of the
+  work.  Any row without a bounded ``top_k`` sends the batch back to the
+  full-vocab chain.
+* Rows with ``sampling_seed`` set use deterministic Gumbel noise derived
+  from the same MurmurHash3 formula as the CUDA kernel
+  (``sglang/kernels/ops/sampling/murmur_hash.py``): hash(seed, position,
+  token_id) -> uniform -> ``-log(-log(u))``.  Seeded noise is keyed on
+  the token id in every branch (the full-vocab chain scatters the masked
+  weights back through the sort; the bounded chain hashes the candidate
+  ids), so a seeded row's token never depends on whether a batchmate
+  triggered top-k/top-p/min-p filtering or on which chain ran.
+* Greedy rows (``top_k == 1`` after sglang normalization, which rewrites
+  ``temperature < eps`` to ``temperature=1, top_k=1``) short-circuit to
+  ``argmax`` and consume no randomness.
+
+Seeds follow the same gate as every other backend: ``sampling_seed`` is
+consumed only under ``--enable-deterministic-inference``, which then
+seeds every row (:data:`DEFAULT_SAMPLING_SEED` for requests that did not
+ask for one).  See :meth:`MlxSamplingParams.from_req`.
+
+Known deviations from the pytorch backend (not bugs):
+
+* Seeded determinism is MLX-local: noise math runs in float32 (Metal has
+  no float64) and tie order follows MLX's sort, so the same seed on a
+  CUDA backend may pick a different token from the same distribution.
+* Unseeded rows draw their Gumbel noise in whichever space the chain is
+  running (candidate or vocab), so the bounded top-K path consumes the
+  RNG differently from the full-vocab one.  Seeded rows are unaffected —
+  they hash the token id — so ``--enable-deterministic-inference`` is
+  bit-for-bit identical either way.
+* Penalties (frequency/presence/repetition) are not applied on the MLX
+  path (warned once per process).  #25804 skips them as well.
+* Custom logit processors run on pure-decode steps only: the first
+  generated token and decode steps mixed into an extend batch are not
+  processed (``apply_custom_logit_processor`` requires logits rows to
+  match the full ``sampling_info``).  Same scope as #25804, which only
+  hooked the pure-decode path at all.
+* Logprob output covers the sampled token, top-k, and requested token
+  ids for every generated token; prompt/input logprobs
+  (``logprob_start_len``) are not computed.
+
+Logit edits (grammar vocab masks, ``logit_bias``) arrive as a
+pre-combined additive [B, vocab] array built by the worker at graph
+launch — grammar FSM state is always current at a fresh launch because
+the previous token was finalized before scheduling, so the mask is known
+at build time and the graph stays lazy.  Grammar/custom-processor
+batches must not CHAIN (the mask for step N+1 needs token N
+materialized); the scheduler breaks the chain for them.  NaN/inf
+sanitization mirrors ``sanitize_nan_logits`` and is gated on the same
+``SGLANG_SANITIZE_NAN_LOGITS`` env var.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+# Seed given to rows without an explicit ``sampling_seed`` when
+# --enable-deterministic-inference is on.  Mirrors the literal in
+# ``SamplingBatchInfo.from_schedule_batch``.
+DEFAULT_SAMPLING_SEED = 42
+
+# Largest ``top_k`` that still takes the bounded candidate path in
+# :func:`sample_tokens`.  Past this the [B, K] chain stops being
+# meaningfully cheaper than the [B, vocab] one it replaces.
+MAX_BOUNDED_TOP_K = 1024
+
+_warned_ignored_penalties = False
+
+
+@dataclass(frozen=True)
+class MlxSamplingParams:
+    """Per-request sampling parameters, frozen at prefill registration."""
+
+    temperature: float
+    top_k: int
+    top_p: float
+    min_p: float
+    seed: int | None
+
+    @classmethod
+    def from_req(
+        cls, req: Any, deterministic_seeding: bool = False
+    ) -> MlxSamplingParams:
+        sp = req.sampling_params
+        global _warned_ignored_penalties
+        if not _warned_ignored_penalties and (
+            sp.frequency_penalty != 0.0
+            or sp.presence_penalty != 0.0
+            or sp.repetition_penalty != 1.0
+        ):
+            _warned_ignored_penalties = True
+            logger.warning(
+                "MLX sampling ignores frequency/presence/repetition penalties; "
+                "a request specified them. (Warning logged once.)"
+            )
+        # Seed contract, identical to every other backend: SamplingBatchInfo
+        # populates sampling_seed only under --enable-deterministic-inference,
+        # and then seeds every row (default DEFAULT_SAMPLING_SEED).  Outside
+        # that flag a per-request seed is ignored, so it is ignored here too.
+        seed = None
+        if deterministic_seeding:
+            seed = (
+                sp.sampling_seed
+                if sp.sampling_seed is not None
+                else DEFAULT_SAMPLING_SEED
+            )
+        return cls(
+            temperature=sp.temperature,
+            top_k=sp.top_k,
+            top_p=sp.top_p,
+            min_p=sp.min_p,
+            seed=seed,
+        )
+
+    @property
+    def is_greedy(self) -> bool:
+        return self.top_k == 1
+
+
+GREEDY_PARAMS = MlxSamplingParams(
+    temperature=1.0, top_k=1, top_p=1.0, min_p=0.0, seed=None
+)
+
+
+@dataclass(frozen=True)
+class MlxLogprobSpec:
+    """Per-row logprob request for one step.
+
+    Mirrors the CUDA ``OutputLogprobProcessor`` inputs: ``top_ks[i]`` is
+    the row's ``top_logprobs_num`` (0 = none) and ``token_ids[i]`` the
+    row's requested token ids (None = none).
+    """
+
+    top_ks: tuple[int, ...]
+    token_ids: tuple[tuple[int, ...] | None, ...]
+
+
+@dataclass
+class MlxLazyLogprobs:
+    """Lazy logprob arrays for one step; materialized at finalize."""
+
+    chosen: mx.array  # [B]
+    top_val: mx.array | None  # [B, max_k]
+    top_idx: mx.array | None  # [B, max_k]
+    token_ids_val: list[mx.array | None]  # per row
+    spec: MlxLogprobSpec
+
+
+@dataclass
+class MlxStepLogprobs:
+    """Materialized per-step logprobs, per row, cut to the request shape."""
+
+    chosen: list[float]
+    top_val: list[list[float]]
+    top_idx: list[list[int]]
+    token_ids_val: list[list[float]]
+    token_ids_idx: list[list[int]]
+
+
+def lazy_logprob_arrays(lazy_logprobs: MlxLazyLogprobs | None) -> list[mx.array]:
+    """The mx arrays of a lazy logprob bundle, for eval/async_eval calls."""
+    if lazy_logprobs is None:
+        return []
+    arrays = [lazy_logprobs.chosen]
+    if lazy_logprobs.top_val is not None:
+        arrays += [lazy_logprobs.top_val, lazy_logprobs.top_idx]
+    arrays += [a for a in lazy_logprobs.token_ids_val if a is not None]
+    return arrays
+
+
+def all_greedy(params: list[MlxSamplingParams]) -> bool:
+    return all(p.is_greedy for p in params)
+
+
+def sanitize_logits(logits: mx.array) -> mx.array:
+    """Lazy analogue of sglang's ``sanitize_nan_logits``: NaN -> -1e30,
+    +-inf -> +-1e30 (not dtype extremes — temperature division would
+    overflow those back to inf)."""
+    return mx.clip(mx.where(mx.isnan(logits), -1e30, logits), -1e30, 1e30)
+
+
+def scale_by_temperature(
+    last_logits: mx.array, params: list[MlxSamplingParams]
+) -> mx.array:
+    """``logits / temperature`` per row, in float32.
+
+    Both :func:`sample_tokens` and :func:`compute_logprobs` start here.  MLX
+    builds eager graphs and does not eliminate common subexpressions, so a
+    step that samples *and* reports logprobs would otherwise pay two
+    full-vocab divisions; callers needing both compute this once and pass it
+    to each.
+    """
+    temps = mx.array([p.temperature for p in params], dtype=mx.float32)[:, None]
+    return last_logits.astype(mx.float32) / temps
+
+
+def compute_logprobs(
+    last_logits: mx.array,
+    params: list[MlxSamplingParams],
+    tokens: mx.array,
+    spec: MlxLogprobSpec,
+    scaled: mx.array | None = None,
+) -> MlxLazyLogprobs:
+    """Lazy log-probabilities of this step's distribution.
+
+    Matches the pytorch sampler: the distribution is
+    ``log_softmax(edited_logits / temperature)`` — after grammar-mask /
+    logit_bias / sanitization, before top-k/top-p/min-p filtering (the
+    filters affect which token is drawn, not the reported logprobs).
+
+    ``scaled`` optionally supplies :func:`scale_by_temperature`'s result when
+    the caller already built it for :func:`sample_tokens`.
+    """
+    if scaled is None:
+        scaled = scale_by_temperature(last_logits, params)
+    logp = scaled - mx.logsumexp(scaled, axis=-1, keepdims=True)
+
+    chosen = mx.take_along_axis(logp, tokens[:, None], axis=-1).squeeze(-1)
+
+    max_k = max(spec.top_ks) if spec.top_ks else 0
+    if max_k > 0:
+        top_idx = mx.argsort(-logp, axis=-1)[:, :max_k]
+        top_val = mx.take_along_axis(logp, top_idx, axis=-1)
+    else:
+        top_idx = None
+        top_val = None
+
+    token_ids_val: list[mx.array | None] = [
+        logp[row, mx.array(ids)] if ids else None
+        for row, ids in enumerate(spec.token_ids)
+    ]
+    return MlxLazyLogprobs(
+        chosen=chosen,
+        top_val=top_val,
+        top_idx=top_idx,
+        token_ids_val=token_ids_val,
+        spec=spec,
+    )
+
+
+def sample_tokens(
+    last_logits: mx.array,
+    params: list[MlxSamplingParams],
+    positions: list[int],
+    key: mx.array,
+    scaled: mx.array | None = None,
+) -> mx.array:
+    """Select one token per row of ``last_logits`` ([B, vocab], lazy ok).
+
+    Pure ``mx`` ops — the result stays inside the lazy graph.  ``positions``
+    are the absolute sequence positions of the tokens being sampled (only
+    consumed for seeded rows).  Callers should shortcut to ``mx.argmax``
+    when ``all_greedy(params)`` — this function assumes at least one row
+    samples.  ``scaled`` optionally supplies
+    :func:`scale_by_temperature`'s result when the caller already built it
+    for :func:`compute_logprobs`.
+    """
+    batch_size, vocab_size = last_logits.shape
+    logits32 = last_logits.astype(mx.float32)
+    if scaled is None:
+        scaled = scale_by_temperature(logits32, params)
+
+    filtering = any(
+        not p.is_greedy and (p.top_k < vocab_size or p.top_p < 1.0 or p.min_p > 0.0)
+        for p in params
+    )
+    # Token ids the Gumbel-max runs over: None means "the whole vocabulary",
+    # otherwise a [B, K] candidate array that the argmax result indexes into.
+    candidates: mx.array | None = None
+    if filtering:
+        probs = mx.softmax(scaled, axis=-1)
+        width = _candidate_width(params, vocab_size)
+        sorted_idx = mx.argsort(-probs, axis=-1)[:, :width]
+        p_sort = mx.take_along_axis(probs, sorted_idx, axis=-1)
+        ranks = mx.arange(width, dtype=mx.int32)[None, :]
+        # SamplingParams normalizes top_k=-1 to TOP_K_ALL and temperature
+        # below eps to (temperature=1, top_k=1), so min(top_k, width)
+        # is always >= 1: rank 0 survives and log(weights) is never all -inf.
+        top_ks = mx.array([min(p.top_k, width) for p in params], dtype=mx.int32)[
+            :, None
+        ]
+        top_ps = mx.array([p.top_p for p in params], dtype=mx.float32)[:, None]
+        min_ps = mx.array([p.min_p for p in params], dtype=mx.float32)[:, None]
+        cum = mx.cumsum(p_sort, axis=-1)
+        masked_out = (
+            (ranks >= top_ks)
+            | ((cum - p_sort) > top_ps)
+            | (p_sort < p_sort[:, :1] * min_ps)
+        )
+        w_sort = mx.where(masked_out, 0.0, p_sort)
+        if width < vocab_size:
+            # Bounded top-K: every surviving rank is inside the K
+            # candidates, so run the rest of the chain in candidate space
+            # and map the winning rank back to its vocab id at the end.
+            log_weights = mx.log(w_sort)
+            candidates = sorted_idx
+        else:
+            # Scatter the masked weights back to vocab order: noise must be
+            # applied in vocab-id space in every branch, or a seeded row's
+            # token would change when a batchmate happens to need filtering.
+            weights = mx.put_along_axis(
+                mx.zeros_like(probs), sorted_idx, w_sort, axis=-1
+            )
+            log_weights = mx.log(weights)
+    else:
+        # Nothing is masked, so the weights are the plain softmax and
+        # ``log(softmax(scaled)) == scaled - logsumexp(scaled)``: a per-row
+        # constant offset, which the argmax below is invariant to.  Feeding
+        # the scaled logits straight to the Gumbel-max drops two full-vocab
+        # passes (softmax, log) on the common temperature-only batch.
+        log_weights = scaled
+
+    noise = _gumbel_noise(
+        params=params,
+        positions=positions,
+        shape=log_weights.shape,
+        key=key,
+        columns=candidates,
+    )
+    # Gumbel-max over the UNNORMALIZED masked weights: normalization would
+    # only shift log(w) by a per-row constant, which argmax is invariant to.
+    # That is also why seed + min_p is well-defined here, and why the
+    # pytorch backend's `assert sampling_seed is None` under min-p (and its
+    # TODO at layers/sampler.py "probs_sort should be re-normalized for the
+    # use of multinomial_with_seed") has no analogue on this path.
+    sampled = mx.argmax(log_weights + noise, axis=-1)
+    if candidates is not None:
+        sampled = mx.take_along_axis(candidates, sampled[:, None], axis=-1).squeeze(-1)
+
+    greedy = [p.is_greedy for p in params]
+    if not any(greedy):
+        return sampled
+    # A batch that mixes greedy rows in still runs them through the sampled
+    # path above (the row exists either way); overwrite those rows with the
+    # unnoised argmax, which is what makes greedy rows consume no randomness.
+    return mx.where(mx.array(greedy), mx.argmax(logits32, axis=-1), sampled)
+
+
+def _candidate_width(params: list[MlxSamplingParams], vocab_size: int) -> int:
+    """Rank cut-off the filtered chain can run on, or ``vocab_size``.
+
+    Every rank at or beyond a row's ``top_k`` is masked to weight 0, whose
+    ``log`` is ``-inf``, so the Gumbel-max can never select it.  When every
+    row's ``top_k`` is small, the whole chain after the sort — gather,
+    cumsum, mask, log, noise, argmax — can therefore run on ``[B, K]``
+    instead of ``[B, vocab]`` and still pick exactly the same token.
+
+    Falls back to the full vocabulary as soon as one row wants more
+    candidates than :data:`MAX_BOUNDED_TOP_K` (or no top-k at all, which
+    ``SamplingParams`` spells as ``top_k = TOP_K_ALL``).
+    """
+    largest_top_k = max(p.top_k for p in params)
+    if largest_top_k > MAX_BOUNDED_TOP_K or largest_top_k >= vocab_size:
+        return vocab_size
+    return largest_top_k
+
+
+def _gumbel_noise(
+    params: list[MlxSamplingParams],
+    positions: list[int],
+    shape: tuple[int, int],
+    key: mx.array,
+    columns: mx.array | None = None,
+) -> mx.array:
+    """Gumbel noise shaped like the weights: hashed if seeded, RNG otherwise.
+
+    ``columns`` is the [B, K] token ids each weight column stands for on the
+    bounded top-K path; ``None`` means column j is token id j.  Seeded rows
+    hash the token id, so their noise — and therefore their token — is the
+    same either way.
+    """
+    seeded_rows = [p.seed is not None for p in params]
+    if not any(seeded_rows):
+        return mx.random.gumbel(shape=shape, key=key)
+
+    hashed = _murmur_hash32(
+        seeds=[p.seed if p.seed is not None else 0 for p in params],
+        positions=positions,
+        vocab_size=shape[1],
+        columns=columns,
+    )
+    u = hashed.astype(mx.float32) / float(0xFFFFFFFF)
+    # REQUIRED, not cosmetic: uint32(0xFFFFFFFF) rounds UP to 2**32 in
+    # float32, so the quotient can land just above 1.0 and make
+    # log(-log(u)) NaN; and an exact 1.0 gives -log(-log(1)) = +inf, which
+    # would deterministically force that token.  Clamp both ends to the
+    # nearest representable interior values.
+    u = mx.clip(u, 2.0**-32, 1.0 - 2.0**-24)
+    hash_noise = -mx.log(-mx.log(u))
+
+    if all(seeded_rows):
+        return hash_noise
+    random_noise = mx.random.gumbel(shape=shape, key=key)
+    return mx.where(mx.array(seeded_rows)[:, None], hash_noise, random_noise)
+
+
+def _murmur3_mix_py(h: int, k: int) -> int:
+    """One MurmurHash3 block mix on Python ints (exact 32-bit semantics)."""
+    k = (k * 0xCC9E2D51) & 0xFFFFFFFF
+    k = ((k << 15) | (k >> 17)) & 0xFFFFFFFF
+    k = (k * 0x1B873593) & 0xFFFFFFFF
+    h ^= k
+    h = ((h << 13) | (h >> 19)) & 0xFFFFFFFF
+    h = (h * 5 + 0xE6546B64) & 0xFFFFFFFF
+    return h
+
+
+def _murmur_hash32(
+    seeds: list[int],
+    positions: list[int],
+    vocab_size: int,
+    columns: mx.array | None = None,
+) -> mx.array:
+    """Port of ``murmur_hash32`` (Triton) to mx ops: [B, V] uint32.
+
+    Blocks mixed in kernel order: seed_low, seed_high, position, column.
+    The first three are per-row scalars, so they are folded exactly on the
+    CPU with Python ints; only the column block and finalization run as
+    vectorized uint32 ops (verified to wrap like the Triton kernel).
+
+    ``columns`` hashes an explicit [B, K] set of token ids instead of every
+    id in ``[0, vocab_size)`` — same value per (row, token id) either way,
+    which is what keeps a seeded row's token identical on the bounded
+    top-K path.
+    """
+    row_states = []
+    for seed, pos in zip(seeds, positions):
+        seed &= 0xFFFFFFFFFFFFFFFF
+        h = _murmur3_mix_py(0, seed & 0xFFFFFFFF)
+        h = _murmur3_mix_py(h, (seed >> 32) & 0xFFFFFFFF)
+        h = _murmur3_mix_py(h, pos & 0xFFFFFFFF)
+        row_states.append(h)
+
+    h = mx.array(row_states, dtype=mx.uint32)[:, None]
+    if columns is None:
+        k = mx.arange(vocab_size, dtype=mx.uint32)[None, :]
+    else:
+        k = columns.astype(mx.uint32)
+
+    # murmur3_mix on [B, V]
+    k = k * mx.array(0xCC9E2D51, dtype=mx.uint32)
+    k = (k << 15) | (k >> 17)
+    k = k * mx.array(0x1B873593, dtype=mx.uint32)
+    h = h ^ k
+    h = (h << 13) | (h >> 19)
+    h = h * mx.array(5, dtype=mx.uint32) + mx.array(0xE6546B64, dtype=mx.uint32)
+
+    # finalize: len = 16 bytes (seed + pos + col), then fmix32
+    h = h ^ mx.array(16, dtype=mx.uint32)
+    h = h ^ (h >> 16)
+    h = h * mx.array(0x85EBCA6B, dtype=mx.uint32)
+    h = h ^ (h >> 13)
+    h = h * mx.array(0xC2B2AE35, dtype=mx.uint32)
+    h = h ^ (h >> 16)
+    return h

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

@@ -17,23 +17,20 @@ from __future__ import annotations
 
 import logging
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 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.runtime_context import get_device
 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.hardware_backend.mlx.tp_worker import MlxLaunch
     from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
     from sglang.srt.managers.scheduler import Scheduler
 
@@ -43,23 +40,10 @@ 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.
+        launch: The :class:`MlxLaunch` this job is waiting on — the lazy
+            token handle plus the prefill / extend / decode pendings the
+            forward produced, and the mode that 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
@@ -73,19 +57,41 @@ class MlxPendingJob:
             mutable batch object.
     """
 
-    lazy_tokens: Optional[mx.array]
-    prefills: list[MlxPendingPrefill]
-    extends: list[MlxPendingExtend]
-    decode: Optional[MlxPendingDecode]
-    mode: str
+    launch: MlxLaunch
     batch_copy: ScheduleBatch
     schedule_batch: ScheduleBatch
     reqs: List[Req]
+    # False when the batch needs per-step CPU logit state (grammar vocab
+    # masks / custom logit processors under --mlx-enable-sampling): step
+    # N+1's mask needs token N materialized, so such batches must launch
+    # fresh every step instead of chaining.
+    chain_safe: bool = True
+    # Captured at launch when batch.return_logprob, exactly like
+    # Scheduler.run_batch does for the CUDA paths (the live values mutate
+    # before output processing).
+    extend_input_len_per_req: Optional[List[int]] = None
+    extend_logprob_start_len_per_req: Optional[List[int]] = None
 
 
 class SchedulerMlxOverlapMixin:
     """Mixin that adds MLX overlap scheduling to :class:`Scheduler`."""
 
+    def _mlx_batch_chain_safe(self: Scheduler, batch: ScheduleBatch) -> bool:
+        """False when per-step CPU logit state forbids chained decode.
+
+        Grammar vocab masks and custom logit processors depend on the
+        previous token being materialized; a chained step is built before
+        that, so those batches launch fresh every step.
+        """
+        if not get_device().mlx_enable_sampling:
+            return True
+        sampling_info = batch.sampling_info
+        if sampling_info is None:
+            return True
+        # batch.has_grammar, not sampling_info.grammars: the latter is only
+        # populated at forward launch (see _build_logit_edit_rows).
+        return not (batch.has_grammar or sampling_info.has_custom_logit_processor)
+
     def _prepare_mlx_launch(self: Scheduler, batch: ScheduleBatch):
         """Stamp scheduler bookkeeping before an MLX forward is launched."""
         # Match run_batch's launch boundary. In particular, the profiler
@@ -98,12 +104,10 @@ class SchedulerMlxOverlapMixin:
         self.profiler_manager._profile_batch_predicate(batch)
 
     def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob):
-        result = self.tp_worker.finalize_mlx_result(
-            pending.prefills,
-            pending.extends,
-            pending.decode,
-            pending.mode,
-            pending.reqs,
+        result = self.tp_worker.finalize_mlx_result(pending.launch, pending.reqs)
+        result.extend_input_len_per_req = pending.extend_input_len_per_req
+        result.extend_logprob_start_len_per_req = (
+            pending.extend_logprob_start_len_per_req
         )
         if result.next_token_ids is not None:
             pending.batch_copy.input_ids = result.next_token_ids
@@ -163,29 +167,31 @@ class SchedulerMlxOverlapMixin:
             # loop must do it too, otherwise async_forward_batch_generation_mlx
             # dereferences a None input_ids.
             resolve_forward_inputs(batch, self.future_map)
-            # run_batch stamps launch_ts on every scheduler-built forward; the
-            # MLX overlap loop bypasses run_batch, and process_batch_result ->
-            # _record_step_counters subtracts launch_ts unconditionally for
-            # prefill/decode batches. ScheduleBatch.copy() below carries the
-            # stamp to process_batch_result.
-            lazy_tokens, prefills, extends, decode, mode = (
-                self.tp_worker.async_forward_batch_generation_mlx(batch)
-            )
+            launch = self.tp_worker.async_forward_batch_generation_mlx(batch)
+            extend_input_len_per_req = None
+            extend_logprob_start_len_per_req = None
+            if batch.return_logprob:
+                # Mirror Scheduler.run_batch's launch-time copy.
+                extend_input_len_per_req = [
+                    req.extend_range.length if req.extend_range is not None else 0
+                    for req in batch.reqs
+                ]
+                extend_logprob_start_len_per_req = batch.extend_logprob_start_lens
             return MlxPendingJob(
-                lazy_tokens=lazy_tokens,
-                prefills=prefills,
-                extends=extends,
-                decode=decode,
-                mode=mode,
+                launch=launch,
                 batch_copy=batch.copy(),
                 schedule_batch=batch,
                 reqs=list(batch.reqs),
+                chain_safe=self._mlx_batch_chain_safe(batch),
+                extend_input_len_per_req=extend_input_len_per_req,
+                extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
             )
 
         def _launch_chained(prev: MlxPendingJob) -> MlxPendingJob:
-            assert prev.decode is not None
-            # Composition is identical to prev: reuse a fresh batch copy
-            # of the same underlying ScheduleBatch so process_batch_result
+            assert prev.launch.decode is not None
+            # Composition is identical to prev: every scheduler-side field
+            # carries over, and only a fresh batch copy of the same
+            # underlying ScheduleBatch is needed so process_batch_result
             # updates the same req objects with the new token.
             batch_copy = prev.batch_copy.copy()
             self._prepare_mlx_launch(batch_copy)
@@ -193,18 +199,10 @@ class SchedulerMlxOverlapMixin:
             # chain breaks, prepare_for_decode() may run SWA maintenance
             # before the next fresh launch gets a chance to re-stamp it.
             prev.schedule_batch.forward_iter = batch_copy.forward_iter
-            lazy_tokens, prefills, extends, decode, mode = (
-                self.tp_worker.async_chained_decode_mlx(prev.decode)
-            )
-            return MlxPendingJob(
-                lazy_tokens=lazy_tokens,
-                prefills=prefills,
-                extends=extends,
-                decode=decode,
-                mode=mode,
+            return replace(
+                prev,
+                launch=self.tp_worker.async_chained_decode_mlx(prev.launch.decode),
                 batch_copy=batch_copy,
-                schedule_batch=prev.schedule_batch,
-                reqs=prev.reqs,
             )
 
         while True:
@@ -224,8 +222,9 @@ class SchedulerMlxOverlapMixin:
             #    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 pending_curr.launch.mode == "decode"
+                and pending_curr.launch.decode is not None
+                and pending_curr.chain_safe
                 and not self.waiting_queue
             )
             if can_chain and pending_next is None:

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

@@ -13,7 +13,8 @@ normal ``GenerationBatchResult``.
 """
 
 import logging
-from typing import Optional, Union
+from dataclasses import dataclass
+from typing import Optional
 
 import mlx.core as mx
 import torch
@@ -23,6 +24,11 @@ from sglang.srt.hardware_backend.mlx.model_runner import (
     MlxPendingExtend,
     MlxPendingPrefill,
 )
+from sglang.srt.hardware_backend.mlx.sampling import (
+    MlxLogprobSpec,
+    MlxStepLogprobs,
+    lazy_logprob_arrays,
+)
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.managers.utils import GenerationBatchResult
@@ -31,11 +37,45 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     PPProxyTensors,
 )
-from sglang.srt.runtime_context import get_memory, get_model, get_schedule
+from sglang.srt.runtime_context import (
+    get_device,
+    get_exec,
+    get_memory,
+    get_model,
+    get_schedule,
+)
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class MlxLaunch:
+    """One lazily launched MLX forward pass: its handle and its pending work.
+
+    Produced by :meth:`MlxTpModelWorker.async_forward_batch_generation_mlx`
+    and :meth:`MlxTpModelWorker.async_chained_decode_mlx`, consumed by
+    :meth:`MlxTpModelWorker.finalize_mlx_result`.
+
+    Attributes:
+        lazy_tokens: an ``mx.array`` that, when evaluated, forces
+            materialisation of the whole batch's outputs.  ``None`` for
+            idle batches.
+        prefills: one :class:`MlxPendingPrefill` per new request in an
+            extend batch; empty for pure-decode steps.
+        extends: one :class:`MlxPendingExtend` per chunked-prefill
+            continuation; also empty for pure-decode steps.
+        decode: the :class:`MlxPendingDecode` covering full decode mode
+            AND mixed single-token decodes inside an extend batch.
+        mode: one of ``"idle"``, ``"decode"``, ``"extend"``.
+    """
+
+    lazy_tokens: Optional[mx.array]
+    prefills: list[MlxPendingPrefill]
+    extends: list[MlxPendingExtend]
+    decode: Optional[MlxPendingDecode]
+    mode: str
+
+
 class MlxTpModelWorker(TpModelWorker):
     """A tensor parallel model worker that routes inference through MLX.
 
@@ -59,6 +99,11 @@ class MlxTpModelWorker(TpModelWorker):
             disable_radix_cache=get_memory().disable_radix_cache,
             mem_fraction_static=get_schedule().mem_fraction_static,
             quantization=get_model().quantization,
+            enable_sampling=get_device().mlx_enable_sampling,
+            sampling_rng_seed=get_device().random_seed,
+            deterministic_seeding=(
+                get_exec().deterministic.enable_deterministic_inference
+            ),
         )
         if get_schedule().max_total_tokens is not None:
             init_kwargs["pool_size"] = get_schedule().max_total_tokens
@@ -137,8 +182,8 @@ class MlxTpModelWorker(TpModelWorker):
     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.
+        Called once per request from :meth:`_async_extend_batch`, which both
+        the overlap loop and the synchronous entry point launch through.
 
         Returns one of:
 
@@ -157,134 +202,220 @@ class MlxTpModelWorker(TpModelWorker):
             return "decode"
         return "continuation"
 
+    @staticmethod
+    def _chunk_needs_logits(req) -> bool:
+        """False iff this extend chunk is a non-final chunked-prefill chunk.
+
+        The scheduler truncates a chunked request's extend range below the
+        tokens it already knows about; such a chunk's next-token output is
+        discarded (the runner pops it as the stale intermediate token), so
+        the runner may skip the logit head for it.
+        """
+        if req.extend_range is None:
+            return True
+        return req.extend_range.end >= len(req.full_untruncated_fill_ids)
+
+    @staticmethod
+    def _sampling_active(batch: ScheduleBatch) -> bool:
+        return get_device().mlx_enable_sampling and batch.sampling_info is not None
+
+    def _build_logit_edit_rows(
+        self, batch: ScheduleBatch
+    ) -> dict[str, mx.array] | None:
+        """Pre-combine grammar vocab masks and logit_bias into one additive
+        [vocab] float32 row per request, ready to enter the lazy graph.
+
+        Grammar FSM state is current at every fresh launch — the previous
+        token was finalized before this batch was scheduled — so the mask
+        is knowable at graph-build time with no device sync.  The
+        scheduler never chains grammar batches
+        (:attr:`MlxPendingJob.chain_safe`), so a chained step never needs
+        a stale mask.  Mask application reuses the grammar backend's own
+        ``apply_vocab_mask`` on a zeros tensor, which keeps this
+        backend-agnostic (xgrammar / llguidance / outlines).
+        """
+        if not self._sampling_active(batch):
+            return None
+        sinfo = batch.sampling_info
+        # Mirror ForwardBatch.init_new's grammars population — the MLX paths
+        # never build a ForwardBatch, so without this the list stays None
+        # even when requests carry live grammar objects.
+        sinfo.grammars = (
+            [req.grammar for req in batch.reqs] if batch.has_grammar else None
+        )
+        has_grammar = bool(sinfo.grammars)
+        if not has_grammar and sinfo.logit_bias is None:
+            return None
+        if not has_grammar:
+            # logit_bias alone is already the dense [B, vocab] additive row we
+            # want; converting it directly skips a second [B, vocab] float32
+            # allocation and an add on every step (~6 MB of churn per step at
+            # vocab 200k, batch 8).  Not mutated below, so no clone is needed.
+            combined = sinfo.logit_bias.to(device="cpu", dtype=torch.float32)
+        else:
+            combined = torch.zeros(
+                len(batch.reqs), sinfo.vocab_size, dtype=torch.float32
+            )
+            sinfo.update_regex_vocab_mask()
+            if sinfo.grammar_mask is not None:
+                grammar_mask = sinfo.grammar_mask
+                grammar_mask.grammar.apply_vocab_mask(
+                    logits=combined,
+                    vocab_mask=grammar_mask.vocab_mask.to("cpu"),
+                )
+                # Release promptly; mirrors the VRAM-leak note in the CUDA
+                # ModelRunner._preprocess_logits.
+                sinfo.grammar_mask = None
+            if sinfo.logit_bias is not None:
+                combined += sinfo.logit_bias.to("cpu")
+        rows = mx.array(combined.numpy())
+        return {req.rid: rows[i] for i, req in enumerate(batch.reqs)}
+
+    def _logprob_rows(
+        self, batch: ScheduleBatch
+    ) -> dict[str, tuple[int, tuple[int, ...] | None]] | None:
+        """Per-request (top_logprobs_num, token_ids) for logprob output."""
+        if not self._sampling_active(batch) or not batch.return_logprob:
+            return None
+        tops = batch.top_logprobs_nums or [0] * len(batch.reqs)
+        tids = batch.token_ids_logprobs or [None] * len(batch.reqs)
+        rows = {}
+        for req, top_k, token_ids in zip(batch.reqs, tops, tids):
+            if req.return_logprob:
+                rows[req.rid] = (
+                    int(top_k or 0),
+                    tuple(token_ids) if token_ids else None,
+                )
+        return rows or None
+
+    @staticmethod
+    def _logprob_spec_for(
+        rows: dict[str, tuple[int, tuple[int, ...] | None]] | None,
+        rids: list[str],
+    ) -> MlxLogprobSpec | None:
+        if rows is None or not any(rid in rows for rid in rids):
+            return None
+        return MlxLogprobSpec(
+            top_ks=tuple(rows.get(rid, (0, None))[0] for rid in rids),
+            token_ids=tuple(rows.get(rid, (0, None))[1] for rid in rids),
+        )
+
+    def _custom_logits_hook(self, batch: ScheduleBatch):
+        """CPU edit hook for custom logit processors, or None.
+
+        Only built for fresh pure-decode launches; the runner materializes
+        the logits for the hook, so these batches never chain.
+        """
+        if not (
+            self._sampling_active(batch)
+            and batch.sampling_info.has_custom_logit_processor
+        ):
+            return None
+        sinfo = batch.sampling_info
+
+        def hook(logits_np):
+            from sglang.srt.layers.sampler import apply_custom_logit_processor
+
+            # torch.from_numpy shares memory with logits_np, so the
+            # processors' in-place edits land in the returned array.
+            logits_t = torch.from_numpy(logits_np)
+            apply_custom_logit_processor(logits_t, sinfo)
+            return logits_np
+
+        return hook
+
+    @staticmethod
+    def _assemble_logprob_output(step_rows: dict[str, tuple], reqs: list):
+        """Batch-ordered LogitsProcessorOutput from per-request logprob rows.
+
+        Field shapes follow what ``move_logprobs_to_cpu`` and
+        ``add_logprob_return_values`` consume: tensors for values the
+        scheduler ``.tolist()``s, plain lists for token-id indices.
+        """
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+        chosen, top_val, top_idx, tid_val, tid_idx = [], [], [], [], []
+        for req in reqs:
+            row = step_rows.get(req.rid)
+            if row is None:
+                row = (0.0, [], [], [], [])
+            chosen.append(row[0])
+            top_val.append(torch.tensor(row[1], dtype=torch.float32))
+            top_idx.append(torch.tensor(row[2], dtype=torch.long))
+            tid_val.append(torch.tensor(row[3], dtype=torch.float32))
+            tid_idx.append(list(row[4]))
+        return LogitsProcessorOutput(
+            next_token_logits=None,
+            next_token_logprobs=torch.tensor(chosen, dtype=torch.float32),
+            next_token_top_logprobs_val=top_val,
+            next_token_top_logprobs_idx=top_idx,
+            next_token_token_ids_logprobs_val=tid_val,
+            next_token_token_ids_logprobs_idx=tid_idx,
+        )
+
+    @staticmethod
+    def _step_logprob_rows(
+        step: Optional[MlxStepLogprobs], rids: list[str]
+    ) -> dict[str, tuple]:
+        """Split a step's batch logprobs into per-request rows."""
+        if step is None:
+            return {}
+        return {
+            rid: (
+                step.chosen[i],
+                step.top_val[i],
+                step.top_idx[i],
+                step.token_ids_val[i],
+                step.token_ids_idx[i],
+            )
+            for i, rid in enumerate(rids)
+        }
+
+    def _collect_step_logprobs(
+        self,
+        step_rows: dict[str, tuple],
+        lazy_logprobs,
+        rids: list[str],
+    ) -> None:
+        """Materialize one pending's lazy logprobs into ``step_rows``."""
+        step = self._mlx_runner.collect_logprobs(lazy_logprobs)
+        step_rows.update(self._step_logprob_rows(step, rids))
+
     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
+        """Run one forward pass through the MLX model runner, synchronously.
 
-        forward_mode = batch.forward_mode
-        reqs = batch.reqs
+        Reachable only under ``--disable-overlap-schedule``: the default MLX
+        loop drives :meth:`async_forward_batch_generation_mlx` /
+        :meth:`finalize_mlx_result` directly and never calls ``run_batch``.
+        Launching and finalising back-to-back IS the synchronous path — the
+        lazy graph is built exactly the same way, then blocked on
+        immediately — so routing, logit edits, logprob collection and
+        chunk-head skipping have one implementation rather than two that
+        must be kept in step.  It is also strictly cheaper than evaluating
+        each request as it is queued: one ``mx.async_eval`` covers the whole
+        batch.
+        """
+        launch = self.async_forward_batch_generation_mlx(batch)
+        return self.finalize_mlx_result(launch, batch.reqs)
 
-        if forward_mode.is_idle():
-            return GenerationBatchResult(
-                logits_output=LogitsProcessorOutput(next_token_logits=None),
-                can_run_cuda_graph=False,
-            )
+    @staticmethod
+    def _stacked_edit_rows(
+        edit_rows: dict[str, mx.array] | None, req_ids: list[str]
+    ) -> Optional[mx.array]:
+        """Stack the per-request additive edit rows for a decode sub-batch."""
+        if not edit_rows:
+            return None
+        return mx.stack([edit_rows[rid] for rid in req_ids])
 
-        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,
-    ]:
+    def async_forward_batch_generation_mlx(self, batch: ScheduleBatch) -> MlxLaunch:
         """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.
+        See :class:`MlxLaunch` for the returned fields.  The caller must
+        make sure the launch's pendings are fed into a subsequent
+        ``mx.async_eval`` or ``.item()`` / ``.tolist()`` call —
+        :meth:`finalize_mlx_result` does that.
         """
         self._ensure_mlx_pool_initialized()
 
@@ -292,15 +423,33 @@ class MlxTpModelWorker(TpModelWorker):
         reqs = batch.reqs
 
         if forward_mode.is_idle():
-            return None, [], [], None, "idle"
+            return MlxLaunch(
+                lazy_tokens=None, prefills=[], extends=[], decode=None, mode="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"
+            pending_decode = self._mlx_runner.decode_batch_start(
+                req_ids,
+                edit_rows=self._stacked_edit_rows(
+                    self._build_logit_edit_rows(batch), req_ids
+                ),
+                logprob_spec=self._logprob_spec_for(self._logprob_rows(batch), req_ids),
+                logits_hook=self._custom_logits_hook(batch),
+            )
+            mx.async_eval(
+                pending_decode.lazy_tokens,
+                *lazy_logprob_arrays(pending_decode.lazy_logprobs),
+            )
+            return MlxLaunch(
+                lazy_tokens=pending_decode.lazy_tokens,
+                prefills=[],
+                extends=[],
+                decode=pending_decode,
+                mode="decode",
+            )
 
         if forward_mode.is_extend():
             # TODO (changminbark): Implement per-batch flushing using prefix_slot_ids
@@ -313,18 +462,14 @@ class MlxTpModelWorker(TpModelWorker):
             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,
-    ]:
+    def _async_extend_batch(self, batch: ScheduleBatch) -> MlxLaunch:
         """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
+        edit_rows = self._build_logit_edit_rows(batch)
+        logprob_rows = self._logprob_rows(batch)
 
         offset = 0
         slot_offset = 0
@@ -349,6 +494,9 @@ class MlxTpModelWorker(TpModelWorker):
                         req_id=req.rid,
                         new_token_ids=req_token_ids,
                         new_slot_ids=req_new_slots,
+                        needs_logits=self._chunk_needs_logits(req),
+                        logit_edit_row=edit_rows[req.rid] if edit_rows else None,
+                        logprob_spec=self._logprob_spec_for(logprob_rows, [req.rid]),
                     )
                 )
             elif route == "decode":
@@ -365,13 +513,18 @@ class MlxTpModelWorker(TpModelWorker):
                         new_slot_ids=req_new_slots,
                         req_pool_idx=req.req_pool_idx,
                         req=req,
+                        needs_logits=self._chunk_needs_logits(req),
+                        logit_edit_row=edit_rows[req.rid] if edit_rows else None,
+                        logprob_spec=self._logprob_spec_for(logprob_rows, [req.rid]),
                     )
                 )
 
         pending_mixed_decode: Optional[MlxPendingDecode] = None
         if mixed_decode_rids:
             pending_mixed_decode = self._mlx_runner.decode_batch_start(
-                mixed_decode_rids
+                mixed_decode_rids,
+                edit_rows=self._stacked_edit_rows(edit_rows, mixed_decode_rids),
+                logprob_spec=self._logprob_spec_for(logprob_rows, mixed_decode_rids),
             )
 
         # Stack lazy tokens so the caller has a single handle to evaluate
@@ -389,51 +542,28 @@ class MlxTpModelWorker(TpModelWorker):
         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]))
+        for pending in (*pending_prefills, *pending_extends):
+            async_args.extend(self._mlx_runner.cache_state_arrays([pending.cache]))
+            async_args.extend(lazy_logprob_arrays(pending.lazy_logprobs))
         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))
+            async_args.extend(lazy_logprob_arrays(pending_mixed_decode.lazy_logprobs))
+            async_args.extend(
+                self._mlx_runner.cache_state_arrays(pending_mixed_decode.caches)
+            )
 
         if async_args:
             mx.async_eval(*async_args)
 
-        return (
-            lazy_stacked,
-            pending_prefills,
-            pending_extends,
-            pending_mixed_decode,
-            "extend",
+        return MlxLaunch(
+            lazy_tokens=lazy_stacked,
+            prefills=pending_prefills,
+            extends=pending_extends,
+            decode=pending_mixed_decode,
+            mode="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]:
+    def async_chained_decode_mlx(self, prev_pending: MlxPendingDecode) -> MlxLaunch:
         """Launch a decode step that chains off a still-lazy previous decode.
 
         This is the "no idle gap" pipelining primitive: build the next
@@ -455,22 +585,21 @@ class MlxTpModelWorker(TpModelWorker):
         * ``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.
+        Returns an :class:`MlxLaunch` in ``"decode"`` mode; its prefill
+        and extend lists are always empty for a chained decode.
         """
         pending = self._mlx_runner.decode_batch_start_chained(prev_pending)
         mx.async_eval(pending.lazy_tokens)
-        return pending.lazy_tokens, [], [], pending, "decode"
+        return MlxLaunch(
+            lazy_tokens=pending.lazy_tokens,
+            prefills=[],
+            extends=[],
+            decode=pending,
+            mode="decode",
+        )
 
     def finalize_mlx_result(
-        self,
-        prefills: list[MlxPendingPrefill],
-        extends: list[MlxPendingExtend],
-        decode: Optional[MlxPendingDecode],
-        mode: str,
-        reqs: list,
+        self, launch: MlxLaunch, reqs: list
     ) -> GenerationBatchResult:
         """Materialise a lazy MLX result into a :class:`GenerationBatchResult`.
 
@@ -480,28 +609,40 @@ class MlxTpModelWorker(TpModelWorker):
         """
         from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 
-        if mode == "idle":
+        decode = launch.decode
+        if launch.mode == "idle":
             return GenerationBatchResult(
                 logits_output=LogitsProcessorOutput(next_token_logits=None),
                 can_run_cuda_graph=False,
             )
 
-        if mode == "decode":
+        step_logprob_rows: dict[str, tuple] = {}
+
+        if launch.mode == "decode":
             assert decode is not None
             next_tokens_list = self._mlx_runner.decode_batch_finalize(decode)
+            self._collect_step_logprobs(
+                step_logprob_rows, decode.lazy_logprobs, decode.req_ids
+            )
 
-        elif mode == "extend":
+        elif launch.mode == "extend":
             prefill_map: dict[str, int] = {}
-            for pending_p in prefills:
+            for pending_p in launch.prefills:
                 prefill_map[pending_p.req_id] = self._mlx_runner.prefill_finalize(
                     pending_p
                 )
+                self._collect_step_logprobs(
+                    step_logprob_rows, pending_p.lazy_logprobs, [pending_p.req_id]
+                )
 
             extend_map: dict[str, int] = {}
-            for pending_e in extends:
+            for pending_e in launch.extends:
                 extend_map[pending_e.req_id] = self._mlx_runner.extend_finalize(
                     pending_e
                 )
+                self._collect_step_logprobs(
+                    step_logprob_rows, pending_e.lazy_logprobs, [pending_e.req_id]
+                )
 
             decode_map: dict[str, int] = {}
             if decode is not None:
@@ -509,6 +650,9 @@ class MlxTpModelWorker(TpModelWorker):
                 decode_map = {
                     rid: tok for rid, tok in zip(decode.req_ids, mixed_tokens)
                 }
+                self._collect_step_logprobs(
+                    step_logprob_rows, decode.lazy_logprobs, decode.req_ids
+                )
 
             next_tokens_list = []
             for req in reqs:
@@ -520,11 +664,16 @@ class MlxTpModelWorker(TpModelWorker):
                     next_tokens_list.append(prefill_map[req.rid])
 
         else:
-            raise ValueError(f"Unknown MLX async mode: {mode}")
+            raise ValueError(f"Unknown MLX async mode: {launch.mode}")
 
         next_token_ids = torch.tensor(next_tokens_list, dtype=torch.long, device="cpu")
+        logits_output = (
+            self._assemble_logprob_output(step_logprob_rows, reqs)
+            if step_logprob_rows
+            else LogitsProcessorOutput(next_token_logits=None)
+        )
         return GenerationBatchResult(
-            logits_output=LogitsProcessorOutput(next_token_logits=None),
+            logits_output=logits_output,
             next_token_ids=next_token_ids,
             can_run_cuda_graph=False,
         )

python/sglang/srt/managers/scheduler.py

@@ -276,7 +276,12 @@ from sglang.srt.observability.trace import process_tracing_init, trace_set_threa
 from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.platforms import current_platform
 from sglang.srt.plugins import load_plugins
-from sglang.srt.runtime_context import get_context, get_parallel, publish
+from sglang.srt.runtime_context import (
+    get_context,
+    get_device,
+    get_parallel,
+    publish,
+)
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
 from sglang.srt.server_args import PortArgs, ServerArgs
@@ -2605,6 +2610,24 @@ class Scheduler(
             self._add_request_to_queue(req)
             return
 
+        if (
+            get_device().mlx_enable_sampling
+            and req.return_logprob
+            and 0 <= req.logprob_start_len < len(req.origin_input_ids)
+        ):
+            # The MLX sampling path computes output logprobs only; the
+            # prefill result carries no input_token_logprobs, so letting
+            # this through would crash output processing.
+            error_msg = (
+                "Prompt input logprobs (logprob_start_len) are not supported "
+                "on the MLX sampling path; omit logprob_start_len to get "
+                "output logprobs."
+            )
+            req.logprob_start_len = -1
+            req.set_finish_with_abort(error_msg)
+            self._add_request_to_queue(req)
+            return
+
         if recv_req.return_routed_experts:
             error_msg = None
             if recv_req.routed_experts_start_len < 0:

python/sglang/srt/managers/scheduler_components/logprob_result_processor.py

@@ -327,7 +327,13 @@ class SchedulerLogprobResultProcessor:
         else:
             self._initialize_empty_logprob_containers(req)
 
-        if req.logprob.top_logprobs_num > 0:
+        if (
+            req.logprob.top_logprobs_num > 0
+            and output.next_token_top_logprobs_val is not None
+        ):
+            # Guarded like next_token_logprobs above: a backend may leave the
+            # top-logprob fields unset even for a request that asked for them
+            # (indexing None raises TypeError).
             req.logprob.output_top_logprobs_val.append(
                 output.next_token_top_logprobs_val[i]
             )

python/sglang/srt/server_args.py

@@ -1198,6 +1198,25 @@ class ServerArgs:
         NS("device"),
     ] = 1
     random_seed: A[Optional[int], "The random seed.", NS("device")] = None
+    mlx_enable_sampling: A[
+        bool,
+        (
+            "MLX backend only: sample decode tokens (temperature / top-k / "
+            "top-p / min-p) instead of greedy argmax. Sampling runs inside "
+            "the lazy MLX graph, so it works with the overlap scheduler; "
+            "first tokens from prefill/extend are sampled too. Greedy "
+            "requests keep exact argmax behavior. Also enables on the MLX "
+            "path: grammar vocab masks and custom logit processors (these "
+            "break decode chaining per step; custom processors run on "
+            "pure-decode steps only), logit_bias, output logprobs (sampled "
+            "token / top-k / token_ids; prompt input logprobs are not "
+            "computed), NaN sanitization (SGLANG_SANITIZE_NAN_LOGITS), and "
+            "per-request sampling_seed under "
+            "--enable-deterministic-inference (deterministic within MLX "
+            "only). Penalties are not applied."
+        ),
+        NS("device"),
+    ] = False
     watchdog_timeout: A[
         float,
         "Set watchdog timeout in seconds. If a forward batch takes longer than this, the server will crash to prevent hanging.",
@@ -5359,9 +5378,13 @@ class ServerArgs:
         elif model_arch in ["GptOssForCausalLM"]:
             # Attention backend selection + XPU dtype validation moved to the
             # override registry (arg_groups/overrides.py: _gpt_oss_overrides).
-            # None of these backends exist on MPS; attention_backend is still
-            # unset there at this point (the torch_native default fills later).
-            if not is_mps():
+            # None of these backends exist on MPS, and under MLX attention
+            # runs inside the MLX runner, so attention_backend is still unset
+            # at this point (the torch_native default fills later). macOS
+            # *without* MLX is not exempt: it has no runner of its own, so it
+            # must still be held to the supported-backend list instead of
+            # silently landing on torch_native (no SWA, no sinks).
+            if not (is_mps() and use_mlx()):
                 supported_backends = [
                     "triton",
                     "trtllm_mha",