[MLX] Support gpt-oss: sliding-window attention, attention sinks, sm_scale (#30050)

Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/srt/arg_groups/overrides.py

@@ -51,6 +51,7 @@ from sglang.srt.utils.common import (
     is_gfx95_supported,
     is_hip,
     is_mnnvl_fabric_device,
+    is_mps,
     is_musa,
     is_npu,
     is_sm90_supported,
@@ -921,7 +922,8 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict:
             overrides["attention_backend"] = "intel_xpu"
         elif is_hip():
             overrides["attention_backend"] = "aiter"
-        else:
+        elif not is_mps():
+            # No triton on macOS; MPS keeps the platform default.
             overrides["attention_backend"] = "triton"
     if is_xpu():
         # Check for bf16 dtype on Intel XPU. Reads the pristine dtype request,

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

@@ -139,7 +139,23 @@ def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel:
         # AOT kernel currently requires rope_dim == head_dim.
         return MlxAOTRoPEKernel()
 
-    base = float(getattr(rope, "base", 10000.0))
+    # The kernel computes vanilla RoPE from a scalar base. Scaled variants
+    # such as YarnRoPE/Llama3RoPE/SuScaledRoPE expose no ``base`` and bake
+    # their scaling into precomputed ``_freqs`` (plus an ``mscale`` factor
+    # applied outside mx.fast.rope), while linear scaling keeps ``base`` but
+    # sets ``scale != 1`` on nn.RoPE. The kernel has inputs for none of
+    # these, so they must fall back to mx.fast.rope.
+    base = getattr(rope, "base", None)
+    if base is None:
+        return MlxAOTRoPEKernel()
+    if getattr(rope, "_freqs", None) is not None:
+        return MlxAOTRoPEKernel()
+    if float(getattr(rope, "mscale", 1.0)) != 1.0:
+        return MlxAOTRoPEKernel()
+    if float(getattr(rope, "scale", 1.0)) != 1.0:
+        return MlxAOTRoPEKernel()
+    base = float(base)
+
     num_qo_heads = get_num_heads(sample_attn)
     if num_qo_heads is None:
         return MlxAOTRoPEKernel()

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

@@ -1,7 +1,10 @@
 """Cache components for the MLX backend."""
 
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_attention_scale,
+    get_container_window_size,
     get_head_dim,
+    get_layer_window_sizes,
     get_num_heads,
     get_num_kv_heads,
     is_attention_module,
@@ -11,6 +14,7 @@ from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
     AttentionOffsetCache,
     ContiguousAttentionKVCache,
     PoolBackedAttentionKVCache,
+    make_attention_mask,
 )
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
     MlxAttentionKVPool,
@@ -40,12 +44,16 @@ __all__ = [
     "AttentionOffsetCache",
     "ContiguousAttentionKVCache",
     "find_attention_layers",
+    "get_attention_scale",
+    "get_container_window_size",
     "get_head_dim",
     "get_context",
+    "get_layer_window_sizes",
     "get_num_layers",
     "get_num_heads",
     "get_num_kv_heads",
     "is_attention_module",
+    "make_attention_mask",
     "MLXAttentionWrapper",
     "MlxAttentionKVPool",
     "MlxAuxiliaryStateComponent",

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

@@ -4,10 +4,12 @@ from __future__ import annotations
 
 from typing import Any, Iterable
 
-# ``rope`` and ``scale`` are required by MLXAttentionWrapper. Keeping them in
-# the contract also prevents recurrent mixers such as DeltaNet from being
-# mistaken for softmax attention just because they expose projection layers.
-ATTENTION_API_ATTRS = ("q_proj", "k_proj", "v_proj", "o_proj", "rope", "scale")
+# ``rope`` and a softmax scale are required by MLXAttentionWrapper. Keeping
+# them in the contract also prevents recurrent mixers such as DeltaNet from
+# being mistaken for softmax attention just because they expose projections.
+ATTENTION_API_ATTRS = ("q_proj", "k_proj", "v_proj", "o_proj", "rope")
+# Any one of these satisfies the scale requirement (gpt_oss uses ``sm_scale``).
+SCALE_ATTRS = ("scale", "sm_scale")
 NUM_HEAD_ATTRS = ("n_heads", "num_heads", "num_attention_heads")
 NUM_KV_HEAD_ATTRS = ("n_kv_heads", "num_k_heads", "num_kv_heads", "num_key_value_heads")
 SLIDING_ATTENTION_ATTRS = (
@@ -17,6 +19,9 @@ SLIDING_ATTENTION_ATTRS = (
     "use_sliding_window",
     "is_swa",
 )
+# mlx-lm containers name their scalar sliding window either ``window_size``
+# (gpt_oss, gemma4) or ``sliding_window`` (olmo3, llama SWA variants, ...).
+WINDOW_SIZE_ATTRS = ("window_size", "sliding_window")
 
 
 def first_present_attr(module: Any, names: Iterable[str]) -> Any | None:
@@ -50,14 +55,47 @@ def get_head_dim(module: Any) -> int | None:
     return None
 
 
+def get_attention_scale(module: Any) -> float | None:
+    return first_present_attr(module, SCALE_ATTRS)
+
+
 def is_attention_module(module: Any) -> bool:
     return (
         all(hasattr(module, attr) for attr in ATTENTION_API_ATTRS)
+        and any(hasattr(module, attr) for attr in SCALE_ATTRS)
         and get_num_heads(module) is not None
         and get_num_kv_heads(module) is not None
     )
 
 
+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)
+
+
+def get_layer_window_sizes(model: Any) -> dict[int, int | None]:
+    """Per-layer sliding-window sizes from the mlx-lm container convention.
+
+    Containers such as gpt_oss or olmo3 expose ``layer_types`` (one entry
+    per layer, ``"sliding_attention"`` marking windowed layers) plus a
+    scalar window (see ``WINDOW_SIZE_ATTRS``).  Returns
+    ``{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)
+    layer_types = getattr(container, "layer_types", None)
+    window_size = get_container_window_size(model)
+    if not layer_types or window_size is None:
+        return {}
+    return {
+        idx: window_size if layer_type == "sliding_attention" else None
+        for idx, layer_type in enumerate(layer_types)
+    }
+
+
 def uses_sliding_window_attention(*modules: Any) -> bool:
     return any(
         bool(getattr(module, attr, False))

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

@@ -5,6 +5,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 import mlx.core as mx
+from mlx_lm.models.base import create_causal_mask
 
 if TYPE_CHECKING:
     from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
@@ -12,6 +13,23 @@ if TYPE_CHECKING:
     )
 
 
+def make_attention_mask(N, offset, return_array=False, window_size=None):
+    """Mirror mlx_lm ``cache.create_attention_mask`` for cache shims.
+
+    Containers delegate mask creation to ``cache.make_mask`` whenever the
+    cache exposes it, so the shims must honor ``window_size`` (sliding-window
+    layers pass it, including for N == 1) or windowed models silently fall
+    back to full attention.
+    """
+    if window_size is not None:
+        return create_causal_mask(N, offset, window_size=window_size)
+    if N == 1:
+        return None
+    if return_array:
+        return create_causal_mask(N, offset)
+    return "causal"
+
+
 class AttentionOffsetCache:
     """Data-free shim satisfying mlx-lm's cache protocol.
 
@@ -25,8 +43,10 @@ class AttentionOffsetCache:
     def state(self):
         return ()  # Empty — safe for mx.eval unpacking
 
-    def make_mask(self, N, **kwargs):
-        return None if N == 1 else "causal"
+    def make_mask(self, N, return_array=False, window_size=None, **kwargs):
+        return make_attention_mask(
+            N, self.offset, return_array=return_array, window_size=window_size
+        )
 
     def update_and_fetch(self, keys, values):
         raise RuntimeError("AttentionOffsetCache should not store data")
@@ -60,6 +80,11 @@ class ContiguousAttentionKVCache:
         self.offset = 0
         self.max_seq_len = max_seq_len
 
+    def make_mask(self, N, return_array=False, window_size=None, **kwargs):
+        return make_attention_mask(
+            N, self.offset, return_array=return_array, window_size=window_size
+        )
+
     def _allocate(self, keys: mx.array) -> None:
         """Allocate buffers matching the first key tensor's shape."""
         B, n_kv_heads, _, head_dim = keys.shape
@@ -77,9 +102,6 @@ class ContiguousAttentionKVCache:
             return ()
         return (self.keys, self.values)
 
-    def make_mask(self, N, **kwargs):
-        return None if N == 1 else "causal"
-
     def _grow(self, required: int) -> None:
         """Double the buffer until it can hold *required* tokens."""
         new_max = self.max_seq_len
@@ -173,8 +195,10 @@ class PoolBackedAttentionKVCache:
             return (self._full_keys, self._full_values)
         return ()
 
-    def make_mask(self, N, **kwargs):
-        return None if N == 1 else "causal"
+    def make_mask(self, N, return_array=False, window_size=None, **kwargs):
+        return make_attention_mask(
+            N, self.offset, return_array=return_array, window_size=window_size
+        )
 
     def update_and_fetch(
         self, keys: mx.array, values: mx.array

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

@@ -15,6 +15,7 @@ from sglang.srt.hardware_backend.mlx.aot import (
     MlxAOTRoPEContext,
 )
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_attention_scale,
     get_head_dim,
     get_num_heads,
     get_num_kv_heads,
@@ -120,12 +121,29 @@ 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.
     """
 
-    def __init__(self, inner: nn.Module, layer_idx: int):
+    def __init__(
+        self, inner: nn.Module, layer_idx: int, window_size: int | None = None
+    ):
         super().__init__()
         object.__setattr__(self, "_inner", inner)
         object.__setattr__(self, "_layer_idx", layer_idx)
+        object.__setattr__(self, "_window_size", window_size)
+        # Resolved once at patch time (weights are loaded before patching and
+        # the inner module is never swapped afterwards), keeping the decode
+        # hot path free of attribute scans and failing fast on a bad module.
+        scale = get_attention_scale(inner)
+        if scale is None:
+            raise RuntimeError(
+                f"Cannot determine attention scale for {type(inner).__name__}"
+            )
+        object.__setattr__(self, "_scale", scale)
+        object.__setattr__(self, "_sinks", getattr(inner, "sinks", None))
 
     def __call__(self, x: mx.array, mask: Any = None, cache: Any = None) -> mx.array:
         ctx = get_context()
@@ -200,7 +218,17 @@ class MLXAttentionWrapper(nn.Module):
             keys = inner.rope(keys, offset=offsets)
 
         layer_caches = ctx.attention_layer_caches[attention_pool_idx]
-        pad_sizes = ctx.pad_sizes
+        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]
 
         # TODO: replace per-request loop with native batched/ragged
         # attention once mx.fast.scaled_dot_product_attention supports
@@ -212,6 +240,9 @@ class MLXAttentionWrapper(nn.Module):
             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:, :]
 
             pad = pad_sizes[i]
             if pad > 0:
@@ -226,17 +257,34 @@ 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 ctx.needs_padding:
-            mask_bool = ctx.positions[None, :] >= ctx.valid_lens[:, None]
+        if pad_mask is not None:
             attn_mask = mx.where(
-                mask_bool[:, None, None, :],
+                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=inner.scale, mask=attn_mask
+            queries,
+            keys_b,
+            values_b,
+            scale=self._scale,
+            mask=attn_mask,
+            **sink_kwargs,
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1)

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

@@ -1,16 +1,21 @@
 """Model introspection and attention patching."""
 
+import logging
 from typing import Any
 
 import mlx.nn as nn
 
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_container_window_size,
+    get_layer_window_sizes,
     is_attention_module,
 )
 from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
     MLXAttentionWrapper,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def _find_attention_attr(layer: Any) -> str | None:
     """Return the direct child name that satisfies the attention contract."""
@@ -43,6 +48,19 @@ def patch_model_attention(model: Any) -> int:
     is set, so it is always installed and never removed.
     """
     layer_list, attn_attrs = find_attention_layers(model)
+    window_sizes = get_layer_window_sizes(model)
+    if not window_sizes and get_container_window_size(model) is not None:
+        # e.g. gemma3-style containers derive per-layer windows from a
+        # pattern instead of ``layer_types``. Prefill masks (delegated to
+        # the container) honor the window, but batched decode cannot
+        # without a per-layer map, so outputs would diverge past the
+        # window. Surface it instead of silently splitting semantics.
+        logger.warning(
+            "Model %s declares a sliding window but no per-layer "
+            "layer_types map; MLX batched decode will not apply the "
+            "window and long-context output may be incorrect.",
+            type(model).__name__,
+        )
     patched = 0
     for idx, (layer, attn_attr) in enumerate(zip(layer_list, attn_attrs)):
         if attn_attr is None:
@@ -50,7 +68,11 @@ def patch_model_attention(model: Any) -> int:
         attn = getattr(layer, attn_attr)
         if isinstance(attn, MLXAttentionWrapper):
             continue
-        setattr(layer, attn_attr, MLXAttentionWrapper(attn, idx))
+        setattr(
+            layer,
+            attn_attr,
+            MLXAttentionWrapper(attn, idx, window_size=window_sizes.get(idx)),
+        )
         patched += 1
     return patched
 

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

@@ -308,3 +308,12 @@ class MlxModelRunnerStub(ModelRunner):
     def alloc_memory_pool(self, memory_pool_config=None):
         """No-op: MLX manages its own KV cache."""
         pass
+
+    def init_attention_backends(self):
+        """No-op: attention runs inside the MLX runner.
+
+        The backend named by ``server_args.attention_backend`` would never
+        be used, and building one can crash: some backends read real KV
+        buffers in ``__init__``, which this stub never allocates.
+        """
+        self.attn_backend = None

python/sglang/srt/server_args.py

@@ -5359,28 +5359,30 @@ 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).
-
-            supported_backends = [
-                "triton",
-                "trtllm_mha",
-                "fa3",
-                "fa4",
-                "ascend",
-                "intel_amx",
-                "intel_xpu",
-                "aiter",
-            ]
-            prefill_attn_backend, decode_attn_backend = (
-                self._resolved_attention_backends()
-            )
-            assert (
-                prefill_attn_backend in supported_backends
-                and decode_attn_backend in supported_backends
-            ), (
-                f"GptOssForCausalLM requires one of {supported_backends} attention backend, but got the following backends\n"
-                f"- Prefill: {prefill_attn_backend}\n"
-                f"- Decode: {decode_attn_backend}\n"
-            )
+            # 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():
+                supported_backends = [
+                    "triton",
+                    "trtllm_mha",
+                    "fa3",
+                    "fa4",
+                    "ascend",
+                    "intel_amx",
+                    "intel_xpu",
+                    "aiter",
+                ]
+                prefill_attn_backend, decode_attn_backend = (
+                    self._resolved_attention_backends()
+                )
+                assert (
+                    prefill_attn_backend in supported_backends
+                    and decode_attn_backend in supported_backends
+                ), (
+                    f"GptOssForCausalLM requires one of {supported_backends} attention backend, but got the following backends\n"
+                    f"- Prefill: {prefill_attn_backend}\n"
+                    f"- Decode: {decode_attn_backend}\n"
+                )
 
             quant_method = get_quantization_config(hf_config)
             is_mxfp4_quant_format = quant_method == "mxfp4"