diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 7c009cf75..6724e8bba 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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, diff --git a/python/sglang/srt/hardware_backend/mlx/aot.py b/python/sglang/srt/hardware_backend/mlx/aot.py index e04b7651f..af9784e99 100644 --- a/python/sglang/srt/hardware_backend/mlx/aot.py +++ b/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() diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py index 89c83dc94..985db1e7b 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py +++ b/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", diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py index 80f989f1c..d96b790c6 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py +++ b/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)) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py index 8d9ccef09..0317c85ab 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py index 13f731f4f..38dcdb2b8 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py +++ b/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) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py index ed41317fc..585c700fa 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 95081ef81..ddfb1e1fc 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3fa0eb22d..0086ef228 100644 --- a/python/sglang/srt/server_args.py +++ b/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" diff --git a/test/registered/mlx/models_e2e/test_gpt_oss_mlx_correctness.py b/test/registered/mlx/models_e2e/test_gpt_oss_mlx_correctness.py new file mode 100644 index 000000000..a03dfeea1 --- /dev/null +++ b/test/registered/mlx/models_e2e/test_gpt_oss_mlx_correctness.py @@ -0,0 +1,371 @@ +"""Correctness tests for gpt-oss served on the SGLang MLX backend. + +gpt-oss interleaves sliding-window (window=128) and full-attention layers and +uses per-head attention sinks, so it exercises the MLX backend's +sliding-window path end to end. Two guards: + +1. ``TestGptOssMlxCorrectness`` — black-box serving smoke against a running + server, including a >128-token prompt so the sliding window actually + engages. +2. ``TestGptOssMlxReferenceCorrectness`` — token-for-token equivalence of + ``MlxModelRunner`` greedy decoding against raw, unpatched mlx_lm greedy + generation. SGLang keeps full KV and applies banded masks / + trailing-window truncation, while vanilla mlx_lm uses RotatingKVCache + for sliding layers — mathematically equivalent, so tokens must match + exactly. + +Both follow the structure of the qwen MoE MLX correctness tests +(PR #29440). + +Prompt length matters for both: sequences up to 128 tokens never engage the +window (banded and causal masks coincide), so a short-prompt test passes even +if window handling is completely broken. Prompts here are >128 tokens. They +also stay well under 2048 tokens: past mlx_lm's prefill chunk size the +RotatingKVCache reference trims differently and exact token equality no +longer holds by construction. + +MLX-gated like its siblings: registered on the CPU suite but skipped wherever +``mlx`` is absent (all current CI runners); runs for real only on Apple +Silicon. The default 20B model needs ~11 GB of weights — override with +``SGLANG_MLX_TEST_MODEL`` (e.g. a local download of +``mlx-community/gpt-oss-20b-MXFP4-Q8``). +""" + +from __future__ import annotations + +import gc +import importlib.util +import os +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, + try_cached_model, +) + +# Registered on the CPU suite but skipped wherever mlx is absent; runs for real +# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the +# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch. +register_cpu_ci(est_time=1, suite="base-a-test-cpu") +register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) +_SKIP_REASON = "requires mlx + mlx_lm (Apple Silicon only)" + +MODEL_PATH = os.environ.get( + "SGLANG_MLX_TEST_MODEL", "mlx-community/gpt-oss-20b-MXFP4-Q8" +) +MEM_FRACTION_STATIC = os.environ.get("SGLANG_MLX_TEST_MEM_FRACTION", "0.9") +# Skip (do NOT crash) unless this much system memory is free; an MLX Metal +# OOM is uncatchable and can reboot the machine. ~12 GB suits the default +# 20B MXFP4-Q8 repo (11 GB of weights). +MIN_FREE_GB = float(os.environ.get("SGLANG_MLX_TEST_MIN_FREE_GB", "12")) + +# Filler that pushes every prompt past 128 tokens (the gpt-oss sliding +# window) while staying far below 2048. The question at the end keeps greedy +# answers short and deterministic. +_NUMBER_LIST = "The following is a list of numbers: " + ", ".join( + str(i) for i in range(1, 121) +) +LONG_PROMPTS = [ + _NUMBER_LIST + ". Which number comes right after 57? Answer briefly.", + _NUMBER_LIST + ". What is the sum of the first three numbers? Answer briefly.", +] +MAX_NEW_TOKENS = 64 # equivalence horizon; analysis-channel tokens count too +BATCH_HORIZON = 24 # fixed step count for the batching-isolation test + + +def _available_gb(): + try: + import psutil + + return psutil.virtual_memory().available / 1024**3 + except Exception: + return None # psutil absent -> skip the pre-flight check + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestGptOssMlxCorrectness(CustomTestCase): + @classmethod + def setUpClass(cls): + avail = _available_gb() + if avail is not None and avail < MIN_FREE_GB: + raise unittest.SkipTest( + f"insufficient free memory: {avail:.1f} GB < {MIN_FREE_GB} GB " + f"needed to safely serve {MODEL_PATH} " + f"(override SGLANG_MLX_TEST_MIN_FREE_GB)" + ) + + cls.model = try_cached_model(MODEL_PATH) + cls.base_url = DEFAULT_URL_FOR_TEST + + env = os.environ.copy() + env["SGLANG_USE_MLX"] = "1" + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp-size", + "1", + "--disable-radix-cache", + "--disable-cuda-graph", + "--mem-fraction-static", + MEM_FRACTION_STATIC, + "--max-running-requests", + "1", + "--context-length", + "2048", + ], + env=env, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process is not None: + kill_process_tree(cls.process.pid) + + def _chat(self, messages, max_tokens=64, temperature=0): + resp = requests.post( + f"{self.base_url}/v1/chat/completions", + json={ + "model": MODEL_PATH, + "messages": messages, + "temperature": temperature, + "max_tokens": max_tokens, + }, + timeout=300, + ) + resp.raise_for_status() + return resp.json()["choices"][0]["message"]["content"].strip() + + def test_basic_generation_nonempty(self): + text = self._chat( + [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": "Say hello briefly."}, + ], + max_tokens=32, + ) + self.assertIsInstance(text, str) + self.assertGreater(len(text), 0) + + def test_simple_arithmetic(self): + text = self._chat( + [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": "What is 2+2? Reply with just the number."}, + ], + ) + self.assertIn("4", text) + + def test_long_prompt_engages_sliding_window(self): + # >128 prompt tokens: prefill and decode both run with the sliding + # window engaged on half the layers. The needle sits near the end of + # the prompt, inside the window of the final positions. + text = self._chat( + [ + {"role": "system", "content": "You are a concise assistant."}, + { + "role": "user", + "content": ( + _NUMBER_LIST + ". The secret word is BLUEBERRY. " + "What is the secret word? Answer briefly." + ), + }, + ], + ) + self.assertIn("BLUEBERRY", text.upper()) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestGptOssMlxReferenceCorrectness(CustomTestCase): + @classmethod + def setUpClass(cls): + import mlx.core as mx + from mlx_lm import load + + avail = _available_gb() + if avail is not None and avail < MIN_FREE_GB: + raise unittest.SkipTest( + f"insufficient free memory: {avail:.1f} GB < {MIN_FREE_GB} GB needed " + f"to safely load {MODEL_PATH} (override SGLANG_MLX_TEST_MIN_FREE_GB)" + ) + + model_path = try_cached_model(MODEL_PATH) + + # --- Phase 1: reference tokens from UNPATCHED mlx_lm (one copy resident) --- + try: + ref_model, cls.tokenizer = load( + model_path, tokenizer_config={"trust_remote_code": True} + ) + except Exception as exc: # not cached / offline / bad path + raise unittest.SkipTest(f"could not load {MODEL_PATH}: {exc}") + + eos = getattr(cls.tokenizer, "eos_token_ids", None) or { + cls.tokenizer.eos_token_id + } + cls.eos_ids = set(eos) + + cls.cases = [] # (prompt, prompt_ids, reference_token_ids) + for prompt in LONG_PROMPTS: + prompt_ids = list( + cls.tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], add_generation_prompt=True + ) + ) + # The whole point of this test: the sliding window only engages + # past 128 tokens, and the RotatingKVCache reference only stays + # trim-free below mlx_lm's prefill chunking threshold. + assert 128 < len(prompt_ids) <= 2048, ( + f"prompt must be >128 and <=2048 tokens to exercise the " + f"sliding window, got {len(prompt_ids)}" + ) + ref_ids = cls._reference_greedy( + ref_model, cls.tokenizer, prompt_ids, MAX_NEW_TOKENS + ) + cls.cases.append((prompt, prompt_ids, ref_ids)) + + # --- Release the reference BEFORE building the runner (cap peak at 1x) --- + del ref_model + gc.collect() + mx.clear_cache() + active_gb = mx.get_active_memory() / 1024**3 + if active_gb > 2.0: + raise unittest.SkipTest( + f"reference model not released (active={active_gb:.1f} GB); " + "skipping to avoid a double-resident OOM" + ) + + # --- Phase 2: SGLang runner (one copy resident) --- + from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner + + cls.runner = MlxModelRunner( + model_path=model_path, + trust_remote_code=True, + disable_radix_cache=True, # per-request contiguous caches; no big pool + mem_fraction_static=float(MEM_FRACTION_STATIC), + ) + cls.runner.init_cache_pools(req_to_token_pool=None) + + @classmethod + def tearDownClass(cls): + runner = getattr(cls, "runner", None) + if runner is not None: + runner.clear() + cls.runner = None + gc.collect() + try: + import mlx.core as mx + + mx.clear_cache() + except Exception: + pass + + # --- helpers ---------------------------------------------------------- + + @staticmethod + def _reference_greedy(model, tokenizer, prompt_ids, max_new): + """Ground-truth token ids from raw, unpatched mlx_lm greedy generation.""" + import mlx.core as mx + from mlx_lm import stream_generate + from mlx_lm.sample_utils import make_sampler + + sampler = make_sampler(temp=0.0) # greedy / argmax + out = [] + for resp in stream_generate( + model, tokenizer, mx.array(prompt_ids), max_tokens=max_new, sampler=sampler + ): + out.append(int(resp.token)) + return out + + def _prefill(self, rid, prompt_ids): + return int( + self.runner.prefill( + req_id=rid, + new_token_ids=list(prompt_ids), + full_token_ids=list(prompt_ids), + prefix_slot_ids=[], + new_slot_ids=[], + req_pool_idx=0, + ) + ) + + def _decode(self, rids): + return [int(t) for t in self.runner.decode_batch(rids)] + + def _sglang_greedy(self, rid, prompt_ids, max_new): + """SGLang MLX greedy generation, stopping at EOS like the reference.""" + tok = self._prefill(rid, prompt_ids) + out = [tok] + while len(out) < max_new and tok not in self.eos_ids: + tok = self._decode([rid])[0] + out.append(tok) + self.runner.remove_request(rid) + return out + + def _diff_msg(self, prompt, ref, sgl): + horizon = min(len(ref), len(sgl)) + first = next((j for j in range(horizon) if ref[j] != sgl[j]), horizon) + return ( + f"\nprompt: {prompt[:80]!r}..." + f"\n first divergence @ index {first} (len ref={len(ref)} sgl={len(sgl)})" + f"\n ref text: {self.tokenizer.decode(ref)!r}" + f"\n sgl text: {self.tokenizer.decode(sgl)!r}" + ) + + # --- tests ------------------------------------------------------------ + + def test_greedy_matches_reference_exact(self): + """SGLang MLX greedy output == unpatched mlx_lm greedy output, token-for-token.""" + for i, (prompt, prompt_ids, ref) in enumerate(self.cases): + sgl = self._sglang_greedy(f"ref-{i}", prompt_ids, MAX_NEW_TOKENS) + self.assertEqual(sgl, ref, self._diff_msg(prompt, ref, sgl)) + + def test_batched_decode_matches_solo(self): + """A request's tokens are identical whether decoded alone or in a batch. + + Pins slot/cache isolation for the sliding-window decode path: the + per-request trailing-window truncation and locally rebuilt padding + mask must not let one request's state bleed into another's. + """ + ids_list = [prompt_ids for (_, prompt_ids, _) in self.cases] + + # Solo: prefill, decode a fixed horizon, remove -- one request at a time. + solo = [] + for i, ids in enumerate(ids_list): + seq = [self._prefill(f"solo-{i}", ids)] + for _ in range(BATCH_HORIZON - 1): + seq.append(self._decode([f"solo-{i}"])[0]) + self.runner.remove_request(f"solo-{i}") + solo.append(seq) + + # Batched: prefill all, then advance them together in one decode_batch. + rids = [f"batch-{i}" for i in range(len(ids_list))] + batched = [[self._prefill(rid, ids)] for rid, ids in zip(rids, ids_list)] + for _ in range(BATCH_HORIZON - 1): + for j, t in enumerate(self._decode(rids)): + batched[j].append(t) + for rid in rids: + self.runner.remove_request(rid) + + for i, (prompt, _, _) in enumerate(self.cases): + self.assertEqual( + batched[i], solo[i], self._diff_msg(prompt, solo[i], batched[i]) + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3) diff --git a/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py b/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py index 8bb073a33..7db768e98 100644 --- a/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py +++ b/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py @@ -1,4 +1,4 @@ -"""Guard the MLX stub's ``alloc_memory_pool`` override against drift. +"""Guard the MLX stub's ``ModelRunner`` overrides against drift. The base ``ModelRunner.alloc_memory_pool`` runs ``_init_pools`` which asserts ``is_draft_worker`` (model_runner_kv_cache_mixin.py:409); the @@ -6,6 +6,11 @@ MLX stub manages its own KV cache via ``MlxAttentionKVPool`` and must short-circuit that GPU-allocation path. If the override is lost, every MLX startup crashes inside ``Scheduler.init_target_memory_pool``. +Similarly, the base ``init_attention_backends`` constructs the torch +attention backend named by ``server_args.attention_backend``; MLX never +uses one, and some backends read real KV buffers in ``__init__``, which +crashes on ``_DummyKVCache``. + The checks are signature/identity-only and MLX-gated because importing the stub pulls in ``mlx.core``. """ @@ -75,6 +80,31 @@ class TestMlxRunnerPoolContract(unittest.TestCase): f"optional MemoryPoolConfig argument: {exc}" ) + def test_stub_overrides_base_init_attention_backends(self): + self.assertIn( + "init_attention_backends", + vars(MlxModelRunnerStub), + msg=( + "MlxModelRunnerStub lost its init_attention_backends " + "override. The base implementation constructs the backend " + "named by server_args.attention_backend; some backends " + "read real KV buffers in __init__, which crashes on " + "_DummyKVCache. MLX never uses a torch attention backend " + "— re-add the override that keeps attn_backend = None." + ), + ) + self.assertIsNot( + MlxModelRunnerStub.init_attention_backends, + ModelRunner.init_attention_backends, + msg="init_attention_backends must be overridden on the MLX " + "stub, not inherited from ModelRunner.", + ) + + def test_stub_init_attention_backends_keeps_attn_backend_none(self): + runner = object.__new__(MlxModelRunnerStub) + MlxModelRunnerStub.init_attention_backends(runner) + self.assertIsNone(runner.attn_backend) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/hardware_backend/mlx/test_sliding_window_attention.py b/test/registered/unit/hardware_backend/mlx/test_sliding_window_attention.py new file mode 100644 index 000000000..36e05bac3 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_sliding_window_attention.py @@ -0,0 +1,487 @@ +"""Unit tests for MLX sliding-window attention support (gpt-oss style models). + +gpt-oss interleaves sliding-window and full-attention layers, names its +softmax scale ``sm_scale``, and adds per-head attention sinks. These tests +pin the three seams that make such models work on the MLX backend: + +1. The attention contract accepts ``sm_scale`` and exposes per-layer window + sizes read from the mlx-lm container convention (``layer_types`` + + ``window_size``). +2. The cache shims' ``make_mask`` mirrors mlx_lm's + ``cache.create_attention_mask`` exactly — in particular ``window_size`` + must produce a banded mask (including for N == 1) instead of being + silently dropped, or sliding-window layers degrade to full attention. +3. ``MLXAttentionWrapper._batched_decode`` applies the window by truncating + each request's KV to the trailing window, passes ``sinks`` through, and + uses the contract scale helper. + +The AOT RoPE kernel gating is also pinned: YarnRoPE (used by gpt-oss) bakes +its base and scaling into precomputed ``_freqs`` plus an ``mscale`` factor, +so the vanilla-RoPE Metal kernel must reject it rather than silently compute +with base=10000 and no yarn scaling. +""" + +from __future__ import annotations + +import importlib.util +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") +register_mlx_ci(est_time=6, suite="stage-a-unit-test-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm") is not None +) +_SKIP_REASON = "requires mlx + mlx_lm" + +if _HAS_MLX: + import mlx.core as mx + from mlx_lm.models import gpt_oss + from mlx_lm.models.base import create_causal_mask + from mlx_lm.models.cache import KVCache + + import sglang.srt.hardware_backend.mlx.aot as mlx_aot + from sglang.srt.hardware_backend.mlx.kv_cache import ( + AttentionOffsetCache, + BatchedDecodeContext, + ContiguousAttentionKVCache, + MLXAttentionWrapper, + PoolBackedAttentionKVCache, + find_attention_layers, + get_attention_scale, + get_layer_window_sizes, + is_attention_module, + make_attention_mask, + patch_model_attention, + ) + +TINY_WINDOW = 8 + + +def _tiny_gpt_oss_model(): + """Randomly initialized 4-layer gpt_oss with alternating sliding/full layers.""" + args = gpt_oss.ModelArgs( + num_hidden_layers=4, + num_local_experts=8, + num_experts_per_tok=2, + vocab_size=128, + hidden_size=64, + intermediate_size=64, + head_dim=16, + num_attention_heads=4, + num_key_value_heads=2, + sliding_window=TINY_WINDOW, + rope_theta=150000, + rope_scaling={ + "rope_type": "yarn", + "factor": 32.0, + "beta_fast": 32.0, + "beta_slow": 1.0, + "original_max_position_embeddings": 4096, + "truncate": False, + }, + ) + return gpt_oss.Model(args) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestGptOssAttentionContract(CustomTestCase): + def test_gpt_oss_attention_passes_contract(self): + model = _tiny_gpt_oss_model() + attn = model.model.layers[0].self_attn + + self.assertFalse(hasattr(attn, "scale")) + self.assertTrue(hasattr(attn, "sm_scale")) + self.assertTrue(is_attention_module(attn)) + + layers, attrs = find_attention_layers(model) + self.assertEqual(len(layers), 4) + self.assertEqual(attrs, ["self_attn"] * 4) + + def test_module_without_any_scale_attr_fails_contract(self): + attn = _tiny_gpt_oss_model().model.layers[0].self_attn + scaleless = SimpleNamespace( + q_proj=attn.q_proj, + k_proj=attn.k_proj, + v_proj=attn.v_proj, + o_proj=attn.o_proj, + rope=attn.rope, + num_attention_heads=4, + num_key_value_heads=2, + ) + self.assertFalse(is_attention_module(scaleless)) + + def test_get_attention_scale_prefers_scale_over_sm_scale(self): + self.assertEqual(get_attention_scale(SimpleNamespace(scale=0.5)), 0.5) + self.assertEqual(get_attention_scale(SimpleNamespace(sm_scale=0.25)), 0.25) + self.assertEqual( + get_attention_scale(SimpleNamespace(scale=0.5, sm_scale=0.25)), 0.5 + ) + self.assertIsNone(get_attention_scale(SimpleNamespace())) + + def test_get_layer_window_sizes_reads_gpt_oss_container(self): + windows = get_layer_window_sizes(_tiny_gpt_oss_model()) + self.assertEqual(windows, {0: TINY_WINDOW, 1: None, 2: TINY_WINDOW, 3: None}) + + def test_get_layer_window_sizes_reads_sliding_window_alias(self): + # olmo3/llama-style containers name the scalar ``sliding_window``. + model = SimpleNamespace( + model=SimpleNamespace( + layer_types=["sliding_attention", "full_attention"], + sliding_window=16, + ) + ) + self.assertEqual(get_layer_window_sizes(model), {0: 16, 1: None}) + + def test_get_layer_window_sizes_defaults_to_empty(self): + self.assertEqual(get_layer_window_sizes(SimpleNamespace()), {}) + no_window = SimpleNamespace( + model=SimpleNamespace(layer_types=["sliding_attention"]) + ) + self.assertEqual(get_layer_window_sizes(no_window), {}) + + def test_patch_warns_when_window_declared_but_unmapped(self): + # A container that declares a scalar window without layer_types + # (gemma3-style pattern models): prefill masks honor the window but + # batched decode cannot; the mismatch must be surfaced. + model = _tiny_gpt_oss_model() + model.model.layer_types = [] + with self.assertLogs( + "sglang.srt.hardware_backend.mlx.kv_cache.model_patching", + level="WARNING", + ) as logs: + patch_model_attention(model) + self.assertTrue(any("sliding window" in msg for msg in logs.output)) + wrappers = [layer.self_attn for layer in model.model.layers] + self.assertEqual([w._window_size for w in wrappers], [None] * 4) + + def test_patch_model_attention_assigns_window_sizes(self): + model = _tiny_gpt_oss_model() + self.assertEqual(patch_model_attention(model), 4) + wrappers = [layer.self_attn for layer in model.model.layers] + self.assertTrue(all(isinstance(w, MLXAttentionWrapper) for w in wrappers)) + self.assertEqual( + [w._window_size for w in wrappers], + [TINY_WINDOW, None, TINY_WINDOW, None], + ) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestShimMakeMask(CustomTestCase): + """The shims must return exactly what mlx_lm's own KVCache.make_mask returns.""" + + def _shims(self, offset): + contig = ContiguousAttentionKVCache( + n_kv_heads=1, head_dim=2, max_seq_len=64, dtype=mx.float32 + ) + contig.offset = offset + pool_backed = PoolBackedAttentionKVCache( + pool=None, layer_idx=0, slots=None, prefix_len=offset + ) + return (AttentionOffsetCache(offset=offset), contig, pool_backed) + + def _assert_same_mask(self, got, ref, msg): + if ref is None or isinstance(ref, str): + self.assertEqual(got, ref, msg) + else: + self.assertTrue( + isinstance(got, mx.array) and mx.array_equal(got, ref).item(), + msg, + ) + + def test_shims_match_mlx_lm_reference(self): + cases = [ + (N, offset, window, return_array) + for N in (1, 4) + for offset in (0, 3, 9) + for window in (None, 4) + for return_array in (False, True) + ] + for N, offset, window, return_array in cases: + reference = KVCache() + reference.offset = offset + ref = reference.make_mask(N, return_array=return_array, window_size=window) + for shim in self._shims(offset): + got = shim.make_mask(N, return_array=return_array, window_size=window) + self._assert_same_mask( + got, + ref, + f"{type(shim).__name__} mismatch for N={N} offset={offset} " + f"window={window} return_array={return_array}", + ) + + def test_windowed_mask_is_banded_including_self(self): + # Query at absolute position 6 with W=4 may attend to keys 3..6 + # (j in [i - W + 1, i], the window includes the query itself). + mask = make_attention_mask(1, 6, window_size=4) + self.assertEqual(mask.shape, (1, 7)) + self.assertEqual( + [bool(v) for v in mask[0]], + [False, False, False, True, True, True, True], + ) + + def test_single_token_windowed_mask_is_not_none(self): + # N == 1 must still produce a banded mask when a window is set — + # returning None here silently disables the window during decode. + mask = make_attention_mask(1, 200, window_size=128) + self.assertIsNotNone(mask) + self.assertEqual(mask.shape, (1, 201)) + self.assertEqual(mx.sum(mask).item(), 128) + + def test_prefill_mask_bands_each_query_row(self): + # N=5 rows starting at offset 7, W=4: row i allows [i+4, i+7]. + offset, N, window = 7, 5, 4 + mask = make_attention_mask(N, offset, window_size=window) + self.assertEqual(mask.shape, (N, offset + N)) + for i in range(N): + allowed = {j for j in range(offset + N) if bool(mask[i, j])} + expected = set(range(offset + i - window + 1, offset + i + 1)) + self.assertEqual(allowed, expected, f"row {i}") + + def test_defaults_without_window_are_unchanged(self): + self.assertIsNone(make_attention_mask(1, 5)) + self.assertEqual(make_attention_mask(4, 5), "causal") + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestBatchedDecodeSlidingWindow(CustomTestCase): + """Batched decode must match a hand-built decode-step reference. + + The reference recomputes the decode step from the inner module's own + projections and RoPE, attending over the full (untruncated) KV with + mlx_lm's own ``create_causal_mask`` band. The wrapper instead truncates + KV to the trailing window and pads ragged requests — mathematically + identical, so the comparison is float-tight. + + The reference deliberately mirrors the wrapper's tensor shapes + (projections and RoPE batched over B, not per request, and no + full-sequence forward): MLX matmul/SDPA kernels pick different code + paths per input shape, and e.g. a (2, 1, H) vs (1, 1, H) linear alone + differs by ~1e-3 on this tiny model — far above the tolerance that + makes this test able to catch real bugs. + """ + + HIDDEN = 64 + N_KV_HEADS = 2 + HEAD_DIM = 16 + + def _prefill_cache(self, attn, x, window): + cache = ContiguousAttentionKVCache( + n_kv_heads=self.N_KV_HEADS, + head_dim=self.HEAD_DIM, + max_seq_len=32, + dtype=mx.float32, + ) + prefix = x[:, :-1, :] + attn( + prefix, + make_attention_mask(prefix.shape[1], 0, window_size=window), + cache=cache, + ) + return cache + + def _project_last_tokens(self, attn, xs, offsets): + B, D = len(xs), self.HEAD_DIM + x_last = mx.concatenate([x[:, -1:, :] for x in xs], axis=0) + q = attn.q_proj(x_last).reshape(B, 1, -1, D).transpose(0, 2, 1, 3) + k = attn.k_proj(x_last).reshape(B, 1, -1, D).transpose(0, 2, 1, 3) + v = attn.v_proj(x_last).reshape(B, 1, -1, D).transpose(0, 2, 1, 3) + off = mx.array(offsets, dtype=mx.int32) + return attn.rope(q, offset=off), attn.rope(k, offset=off), v + + def _reference_decode(self, attn, xs, caches, window): + """Full-KV banded-mask decode; must be called before the wrapper + writes the decode token into the shared caches.""" + offsets = [x.shape[1] - 1 for x in xs] + q, k_new, v_new = self._project_last_tokens(attn, xs, offsets) + outs = [] + for i, cache in enumerate(caches): + k_prefix, v_prefix = cache.get_kv() + k = mx.concatenate([k_prefix, k_new[i : i + 1]], axis=2) + v = mx.concatenate([v_prefix, v_new[i : i + 1]], axis=2) + mask = ( + create_causal_mask(1, offsets[i], window_size=window) + if window is not None + else None + ) + outs.append( + mx.fast.scaled_dot_product_attention( + q[i : i + 1], + k, + v, + scale=attn.sm_scale, + mask=mask, + sinks=attn.sinks, + ) + ) + out = mx.concatenate(outs, axis=0) + out = out.transpose(0, 2, 1, 3).reshape(len(xs), 1, -1) + return attn.o_proj(out) + + def _wrapper_decode(self, attn, window, xs, caches): + wrapper = MLXAttentionWrapper(attn, layer_idx=0, window_size=window) + ctx = BatchedDecodeContext( + batch_size=len(xs), + seq_lens=[x.shape[1] - 1 for x in xs], + attention_layer_caches=[caches], + ) + x_last = mx.concatenate([x[:, -1:, :] for x in xs], axis=0) + out = wrapper._batched_decode(x_last, ctx) + mx.eval(out) + return out + + def _assert_matches_reference(self, attn, window, lens): + mx.random.seed(0) + xs = [mx.random.normal((1, L, self.HIDDEN)) for L in lens] + caches = [self._prefill_cache(attn, x, window) for x in xs] + ref = self._reference_decode(attn, xs, caches, window) + got = self._wrapper_decode(attn, window, xs, caches) + for i in range(len(xs)): + diff = mx.abs(got[i : i + 1] - ref[i : i + 1]).max().item() + self.assertLess( + diff, + 1e-5, + f"request {i} (len={lens[i]}, window={window}) diverges " + f"from the manual decode reference by {diff}", + ) + + def test_sliding_layer_ragged_batch_matches_reference(self): + # Request 0 crosses the window (12 > 8), request 1 stays inside (5 < 8): + # covers trailing-window truncation and the local padding mask at once. + model = _tiny_gpt_oss_model() + attn = model.model.layers[0].self_attn + self._assert_matches_reference(attn, TINY_WINDOW, lens=[12, 5]) + + def test_sliding_layer_all_past_window_matches_reference(self): + # Both requests exceed the window with unequal true lengths: the + # shared context reports padding but the windowed lengths are all + # equal, so the correct local pad is zero. Reusing the full-length + # ctx metadata here would inject spurious padding. + model = _tiny_gpt_oss_model() + attn = model.model.layers[0].self_attn + self._assert_matches_reference(attn, TINY_WINDOW, lens=[12, 10]) + + def test_sliding_layer_single_request_matches_reference(self): + # B=1 with truncation: the windowed no-padding branch (mask stays None). + model = _tiny_gpt_oss_model() + attn = model.model.layers[0].self_attn + self._assert_matches_reference(attn, TINY_WINDOW, lens=[12]) + + def test_full_attention_layer_matches_reference(self): + # Full-attention gpt_oss layer: sinks + sm_scale on the unwindowed path. + model = _tiny_gpt_oss_model() + attn = model.model.layers[1].self_attn + self._assert_matches_reference(attn, None, lens=[12, 5]) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestSdpaSinksSemantics(CustomTestCase): + def test_sdpa_sinks_match_manual_softmax_with_sink_column(self): + # Pin mx.fast.scaled_dot_product_attention(sinks=...) to the reference + # semantics gpt-oss relies on: append one per-head sink logit to the + # softmax and drop its probability column afterwards. + mx.random.seed(0) + B, H, Lq, Lk, D = 1, 4, 5, 9, 16 + scale = D**-0.5 + q = mx.random.normal((B, H, Lq, D)) + k = mx.random.normal((B, H, Lk, D)) + v = mx.random.normal((B, H, Lk, D)) + sinks = mx.random.normal((H,)) + offset, window = Lk - Lq, 4 + + rinds = mx.arange(Lk) + linds = mx.arange(offset, offset + Lq) + mask = (linds[:, None] >= rinds[None]) & (linds[:, None] < rinds[None] + window) + + out_fast = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, sinks=sinks + ) + + scores = (q * scale) @ k.transpose(0, 1, 3, 2) + scores = mx.where(mask, scores, mx.finfo(mx.float32).min) + sink_col = mx.broadcast_to(sinks[None, :, None, None], (B, H, Lq, 1)) + probs = mx.softmax(mx.concatenate([scores, sink_col], axis=-1), axis=-1) + out_manual = probs[..., :-1] @ v + + diff = mx.abs(out_fast - out_manual).max().item() + self.assertLess(diff, 1e-6) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestAotRopeKernelGating(CustomTestCase): + """The vanilla-RoPE Metal kernel must reject scaled RoPE variants.""" + + def _build_kernel(self, attn, head_dim=2, n_kv_heads=1): + original_loader = mlx_aot._load_metal_rope_pool_fused + mlx_aot._load_metal_rope_pool_fused = lambda: object() + try: + return mlx_aot._build_rope_kernel( + mlx_aot.MlxAOTKernelBuildInputs( + sample_attn=attn, + n_kv_heads=n_kv_heads, + head_dim=head_dim, + ) + ) + finally: + mlx_aot._load_metal_rope_pool_fused = original_loader + + def test_vanilla_rope_is_accepted(self): + attn = SimpleNamespace( + n_heads=2, + rope=SimpleNamespace(dims=2, traditional=False, base=10000.0), + ) + self.assertTrue(self._build_kernel(attn).enabled) + + def test_gpt_oss_yarn_rope_is_rejected(self): + # YarnRoPE has no ``base`` (it is baked into ``_freqs``) and applies + # mscale outside mx.fast.rope; the kernel would silently compute + # vanilla RoPE with base=10000. + attn = _tiny_gpt_oss_model().model.layers[0].self_attn + kernel = self._build_kernel(attn, head_dim=attn.head_dim, n_kv_heads=2) + self.assertFalse(kernel.enabled) + + def test_missing_base_is_rejected(self): + attn = SimpleNamespace( + n_heads=2, rope=SimpleNamespace(dims=2, traditional=False) + ) + self.assertFalse(self._build_kernel(attn).enabled) + + def test_precomputed_freqs_are_rejected(self): + attn = SimpleNamespace( + n_heads=2, + rope=SimpleNamespace( + dims=2, traditional=False, base=10000.0, _freqs=mx.ones(1) + ), + ) + self.assertFalse(self._build_kernel(attn).enabled) + + def test_nontrivial_mscale_is_rejected(self): + attn = SimpleNamespace( + n_heads=2, + rope=SimpleNamespace(dims=2, traditional=False, base=10000.0, mscale=1.5), + ) + self.assertFalse(self._build_kernel(attn).enabled) + + def test_linear_scale_is_rejected(self): + # rope_scaling type "linear" yields nn.RoPE(..., scale=1/factor); the + # kernel computes unscaled positions and must fall back, while the + # nn.RoPE default scale of exactly 1.0 stays accepted. + def attn(scale): + return SimpleNamespace( + n_heads=2, + rope=SimpleNamespace( + dims=2, traditional=False, base=10000.0, scale=scale + ), + ) + + self.assertFalse(self._build_kernel(attn(0.25)).enabled) + self.assertTrue(self._build_kernel(attn(1.0)).enabled) + + +if __name__ == "__main__": + unittest.main()