[MLX] Support gpt-oss: sliding-window attention, attention sinks, sm_scale

[LarrySimingDeng]

Motivation

Part of #19137 (gpt_oss.py entry).

gpt-oss (GptOssForCausalLM) could not be served on the MLX backend: find_attention_layers raised ValueError at load because mlx_lm's gpt_oss AttentionBlock names its softmax scale sm_scale while the attention contract required scale. Fixing the load alone is not enough — gpt-oss interleaves sliding-window (window=128) and full-attention layers and uses per-head attention sinks, and the MLX backend had no sliding-window support at all:

• The cache shims' make_mask ignored window_size, so sliding-window layers silently degraded to full attention during prefill/extend. Sequences up to 128 tokens produce identical output either way (the window never engages), which makes this failure mode invisible to short-prompt tests.
• Batched decode applied no window, dropped the attention sinks, and read inner.scale directly.
• The opt-in AOT Metal RoPE kernel's gating accepted YarnRoPE (used by gpt-oss), whose base and scaling live in precomputed _freqs with an mscale factor applied outside mx.fast.rope; the kernel would have silently computed vanilla RoPE with base=10000.
• Serving crashed even before the first forward: the GptOss branch in server_args.py forces attention_backend=triton when nothing else matches, and on macOS the scheduler then either crashed constructing TritonAttnBackend against the MLX stub's dummy KV pool or routed write_cache_indices to a triton kernel launch.

Modifications

All behavior changes are in python/sglang/srt/hardware_backend/mlx/, plus two small mps-only carve-outs:

• kv_cache/attention_contract.py: the scale requirement accepts any of ("scale", "sm_scale") (SCALE_ATTRS, get_attention_scale()). New get_layer_window_sizes() reads the mlx-lm container convention — layer_types (with "sliding_attention" marking windowed layers) plus a scalar window named either window_size (gpt_oss, gemma4) or sliding_window (olmo3, llama SWA variants).
• kv_cache/attention_kv_cache.py: module-level make_attention_mask() mirrors mlx_lm cache.create_attention_mask exactly — window_size produces a banded mask (including for N == 1) via mlx_lm's own create_causal_mask; all three cache shims delegate to it. Previously they returned None/"causal" unconditionally.
• kv_cache/attention_wrapper.py: MLXAttentionWrapper takes a window_size. Batched decode truncates each request's KV to the trailing min(window, seq_len) tokens and rebuilds the padding mask for the windowed lengths (the context's shared padding metadata is full-length); attention sinks are passed through to mx.fast.scaled_dot_product_attention when the module has them; the scale is resolved via get_attention_scale. Models without a window take the exact previous path.
• kv_cache/model_patching.py: wires per-layer window sizes into the wrappers; warns when a model declares a scalar window without a layer_types map (gemma3-style pattern models), where batched decode cannot apply the window yet.
• aot.py: _build_rope_kernel now rejects scaled RoPE variants — missing base, precomputed _freqs, mscale != 1, or a linear scale != 1 fall back to mx.fast.rope instead of computing wrong rotations when the opt-in kernel is enabled.
• model_runner_stub.py: overrides init_attention_backends() as a no-op (attn_backend = None). MLX performs attention itself; the base implementation constructs whatever backend server_args names, and model-specific defaults can force one whose __init__ reads real KV buffers (gpt-oss → triton → crash on the stub's _DummyKVCache).
• server_args.py: the GptOss attention-backend forcing and supported-backends assert are skipped only when serving through MLX on Apple Silicon (is_mps() and use_mlx()); the platform default (torch_native) then applies, exactly as for other MLX-served models. The non-MLX macOS torch path keeps failing fast as before, and CUDA/ROCm/XPU/CPU behavior is unchanged.
• environ.py: the MLX correctness-test hooks (SGLANG_MLX_TEST_MODEL / _MEM_FRACTION / _MIN_FREE_GB) are registered as EnvField descriptors per the env-var conventions.

Design: sliding-window layers keep the full KV history in the pool and the window is applied at read time — banded mask in prefill, trailing-window truncation in decode. For softmax attention this is numerically identical to a rotating cache (masked positions do not contribute), and it keeps radix prefix reuse and chunked prefill working unchanged. The cost is that sliding layers store KV they will not read again; a per-layer windowed pool is left as a TODO.

Test

• Unit (test/registered/unit/hardware_backend/mlx/test_sliding_window_attention.py, 24 tests, base-a-test-cpu, skipped without mlx): contract acceptance for a tiny random-weight gpt_oss; get_attention_scale / get_layer_window_sizes (both window attr namings, empty-convention fallback, the window-without-layer_types warning); shim make_mask pinned against mlx_lm's own KVCache.make_mask across an N/offset/window/return_array grid, plus explicit banded-semantics checks (window includes self, N==1 stays banded); _batched_decode against a hand-built decode reference (full untruncated KV + mlx_lm's banded mask vs the wrapper's truncation + local padding mask) for ragged batches crossing the window, all-past-window unequal lengths, single request, and a full-attention layer; sdpa sinks= semantics vs manual softmax-with-sink-column; AOT RoPE gating (vanilla accepted; YarnRoPE, missing base, _freqs, mscale, linear scale rejected). The stub override is drift-guarded in test_mlx_runner_pool_contract.py.
• Correctness (test/registered/mlx/models_e2e/test_gpt_oss_mlx_correctness.py, structure follows the qwen MoE MLX tests from [MLX] Add correctness tests for qwen2_moe and qwen3_moe #29440): black-box serving smoke including a >128-token prompt, and token-for-token equivalence of MlxModelRunner greedy decoding against raw, unpatched mlx_lm greedy generation. Prompts are asserted to be >128 tokens (below that the window never engages and the test would pass vacuously) and ≤2048 (past mlx_lm's prefill chunking the RotatingKVCache reference trims differently). Registered on base-a-test-cpu; skips wherever mlx is absent (all current CI runners) and runs for real on Apple Silicon.
• On device (Apple Silicon 24 GB, mlx-community/gpt-oss-20b-MXFP4-Q8, mlx 0.31.2 / mlx_lm 0.31.3):
  • Token-for-token equivalence: SGLang MLX greedy output matches unpatched mlx_lm exactly (64-token horizon, two >128-token prompts), and batched decode matches solo decode per request.
  • Serving smoke: sglang serve with SGLANG_USE_MLX=1 answers chat completions correctly, including the >128-token sliding-window prompt (3/3 tests).
  • Vanilla sanity: mlx_lm generate on the same checkpoint: 88.8 tok/s decode, 12.2 GB peak.
  • Regressions: qwen1.5-MoE reference equivalence (token-exact) and qwen3-30B serving smoke pass unchanged; the MLX unit-test modules that fail (test_attention_patching::TestMlxOverlapScheduler on forward_ct, test_metal_profiler) fail identically on current main without this PR.

Accuracy Tests

Covered by the token-for-token reference equivalence above: greedy decoding through the SGLang MLX backend is exactly the unpatched mlx_lm output for gpt-oss (sliding-window + sinks engaged) and stays exact for the qwen MoE regressions.

Speed Tests and Profiling

No perf-sensitive path changes for existing models: the no-window decode path is byte-identical logic, and windowed models did not run before. Windowed decode attends over at most window keys per sliding layer, which is strictly less work than the full-attention fallback it replaces.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

cc @yeahdongcn

---

CI States

Latest PR Test (Base): ⏳ Run #31347160622
Latest PR Test (Extra): ❌ Run #31347160551

[gemini-code-assist]

Code Review

This pull request adds comprehensive support for sliding-window attention models (such as gpt-oss) on the MLX backend, introducing banded attention masks, trailing-window KV truncation, and support for attention sinks. It also updates the AOT RoPE kernel gating to reject scaled RoPE variants and adds extensive end-to-end and unit tests. The reviewer feedback suggests optimizing the hot path in MLXAttentionWrapper by resolving and caching the attention scale and sinks during initialization rather than on every decode step, which also enables failing fast if the scale cannot be determined.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[LarrySimingDeng]

Re-reviewed the branch end to end and re-ran the full local matrix on the final commit (18644c5): sliding-window / pool-contract / runner-init / scheduler-mixin unit suites (34 tests), gpt-oss serving smoke, gpt-oss token-for-token reference equivalence vs unpatched mlx_lm, Qwen1.5-MoE reference equivalence, and Qwen3-30B serving smoke — all green on Apple Silicon (24 GB). One hardening was amended in: the opt-in AOT RoPE gate now also rejects nn.RoPE with linear rope scaling (scale != 1), which the vanilla-RoPE Metal kernel cannot represent; pinned by a new gating unit test. Default-path behavior is unchanged (SGLANG_MLX_USE_CUSTOM_ROPE remains off by default).

[LarrySimingDeng]

Resolved the conflict with main. The config resolution refactor moved the GPT OSS attention backend selection into arg_groups/overrides.py, so the MLX skip now lives in _gpt_oss_overrides and the assertion guard stays in server_args.py. MLX unit tests pass locally.

[yeahdongcn]

@adityavaid Could you take a look? Thanks!

[yeahdongcn]

I would revert the server_args.py part of this change. The new _mlx_serving = is_mps() and use_mlx() special case makes the generic GPT-OSS validation path understand an MLX runtime detail, which is the wrong direction for this file.

[LarrySimingDeng]

Dropped use_mlx() here, agreed that's the wrong direction for this file. A plain revert breaks startup though: with the overrides change, nothing sets attention_backend on MPS before this assert runs (the torch_native default is only filled later, in _handle_attention_backend_compatibility), so every gpt-oss launch on Apple Silicon would fail the assert with None backends. Kept the assert and scoped it with the same platform predicate as the overrides side instead: if not is_mps():.

[LarrySimingDeng]

@yeahdongcn The branch is now up to date with the latest main, and the tests are registered with the #30121 suite mechanism: test_sliding_window_attention under stage-a-unit-test-mlx, test_gpt_oss_mlx_correctness under stage-b-e2e-mlx. All six points from the first review round were addressed earlier (replies in the threads above), and the MLX stage-a unit suite passes locally on Apple Silicon. Ready for another look when you have time.

[yeahdongcn]

/tag-and-rerun-ci