执行摘要

• 一句话：MLX 后端支持 gpt-oss：滑动窗口注意力、attention sinks、sm_scale
• 推荐动作：值得精读。关注三个设计决策：一是“读取时应用窗口 vs 旋转缓存”的取舍——保留完整 KV 并用 banded mask/尾部截断，在 softmax 下数值等价且不破坏 radix 前缀复用，代价是显存冗余，这是很实用的权衡；二是测试防假阳性设计——刻意要求提示词 >128 token 保证窗口真正参与，否则短提示词下窗口层和全注意力输出相同、测试会空过；三是 AOT 内核门控的 fail-safe 策略——宁可回退到 mx.fast.rope 也绝不用 vanilla 内核算错 scaled RoPE。做平台后端或注意力系统的人可参考其中的契约设计与测试方法论。

功能与动机

PR body 明确指出 gpt-oss 无法在 MLX 后端服务：mlx_lm 的 gpt_oss AttentionBlock 把 softmax scale 命名为 sm_scale 而契约只接受 scale，导致 find_attention_layers 加载即抛 ValueError；MLX 后端完全没有滑动窗口支持，window=128 的层在 prefill/extend 期间静默退化为全注意力（短提示词测试无法暴露）；batched decode 不应用窗口、丢弃 sinks 并直接读 inner.scale；可选的 AOT Metal RoPE 内核会错误接受 YarnRoPE 并算出 base=10000 的 vanilla RoPE；此外 server_args.py 的 GptOss 分支强制 attention_backend=triton，在 macOS 上导致启动即崩溃。

实现拆解

按四层拆解实现：

1. 契约层（attention_contract.py）：把 scale 从必选属性改为 SCALE_ATTRS = ("scale", "sm_scale") 二选一，新增 get_attention_scale()；新增 WINDOW_SIZE_ATTRS = ("window_size", "sliding_window") 兼容 gpt_oss/gemma4 与 olmo3/llama 的命名差异，get_container_window_size() 读取容器级标量窗口，get_layer_window_sizes() 依据 mlx-lm 的 layer_types 惯例把窗口映射到每个具体层（"sliding_attention" 标记窗口层，返回 {layer_idx: window|None}）。is_attention_module() 同步放宽为“任一 scale 属性存在”，避免误拒 gpt_oss 同时继续排除 DeltaNet 等非注意力模块。

2. Mask 层（attention_kv_cache.py）：新增模块级 make_attention_mask(N, offset, return_array, window_size)，直接委托 mlx_lm 自身的 create_causal_mask 生成 banded mask，关键点是 window_size 非 None 时 N==1 也返回真实 mask（此前返回 None 会导致 decode 阶段窗口被静默禁用）；三个 cache shim（AttentionOffsetCache、ContiguousAttentionKVCache、PoolBackedAttentionKVCache）的 make_mask 全部改为委托该函数，修复滑窗层 prefill 退化为全注意力的问题。

3. 解码层（attention_wrapper.py + model_patching.py）：MLXAttentionWrapper.__init__ 新增可选 window_size，并在加载/打补丁时一次性解析并缓存 _scale 与 _sinks（热路径不再扫描属性，scale 缺失则 fail-fast）；_batched_decode 对窗口层按 min(seq_len+1, window) 截断每个请求 KV 的尾部、按窗口长度重建 padding mask（上下文里的 padding 元数据是全量长度的）、仅当模块有 sinks 时才把 sinks= 传给 mx.fast.scaled_dot_product_attention。patch_model_attention 把 get_layer_window_sizes 的结果逐个传给 wrapper，并在“声明了 scalar 窗口但没有 layer_types 映射”时（如 gemma3 风格模型）打 WARNING，避免 decode 与 prefill 语义分裂而无人知晓。

4. 平台接线与安全回退（model_runner_stub.py / server_args.py / arg_groups/overrides.py / aot.py）：MlxModelRunnerStub 覆盖 init_attention_backends() 为 no-op（attn_backend = None），防止基类按 server_args 命名构造真实后端时在 _DummyKVCache 上崩溃；server_args.py 的 gpt-oss 支持后端断言与 arg_groups/overrides.py 的后端强制均改为 elif not is_mps() 平台谓词，MLX 服务时保留平台默认 torch_native，非 Apple Silicon 路径行为完全不变；aot.py 的 _build_rope_kernel 在识别到 base 缺失、预计算 _freqs、mscale != 1 或线性 scale != 1 时一律回退到 mx.fast.rope，避免用 vanilla 内核算错 scaled RoPE。

5. 测试与 CI 配套：新增 test_sliding_window_attention.py（24 个单测：契约接受、窗口解析两个命名、mask 与 mlx_lm 参考的网格比对、banded 语义、N==1 仍保持、batched decode 与手工参考 float-tight 对比、sinks 语义、AOT 门控）与 test_gpt_oss_mlx_correctness.py（黑盒 serving smoke + MlxModelRunner 与未打补丁 mlx_lm 的逐 token 参考等价，提示词强制 >128 token 以真正触发窗口）；按 #30121 的 suite 注册机制挂到 base-a-test-cpu、stage-a-unit-test-mlx、stage-b-e2e-mlx；test_mlx_runner_pool_contract.py 增加对 init_attention_backends 覆盖的防漂移断言。

关键文件：

• python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py（模块 注意力契约；类别 source；类型 core-logic；符号 get_attention_scale, get_container_window_size, get_layer_window_sizes）: 注意力契约的核心改动：scale 属性放宽为 scale/sm_scale 二选一，新增 per-layer 窗口解析（layer_types + window_size/sliding_window），是所有 MLX 注意力模型加载与识别的入口。
• python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py（模块 KV 缓存层；类别 source；类型 core-logic；符号 make_attention_mask, make_mask）: 新增 make_attention_mask 统一三个 cache shim 的 make_mask，修复滑动窗口层 prefill/extend 静默退化为全注意力的问题；N==1 也保持 banded mask 是 decode 正确性的关键。
• python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py（模块 解码包装器；类别 source；类型 core-logic；符号 init, _batched_decode）: MLXAttentionWrapper 是 batched decode 的核心执行单元：新增 window_size 支持、scale/sinks 缓存、窗口截断 + 局部 padding mask，无窗口模型保持原路径。
• python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py（模块 模型打补丁；类别 source；类型 data-contract；符号 patch_model_attention）: 把 per-layer window sizes 接入 MLXAttentionWrapper，并针对“有窗口但无 layer_types 映射”的模型告警，防止 decode 与 prefill 语义分裂。
• python/sglang/srt/hardware_backend/mlx/aot.py（模块 AOT 内核；类别 source；类型 core-logic；符号 _build_rope_kernel）: AOT Metal RoPE 内核门控收紧：拒绝 YarnRoPE 等 scaled 变体回退 mx.fast.rope，避免 gpt-oss 开启自定义内核时算错旋转角。
• python/sglang/srt/hardware_backend/mlx/model_runner_stub.py（模块 运行器桩；类别 source；类型 data-contract；符号 init_attention_backends）: 覆盖 init_attention_backends 为 no-op 并保持 attn_backend = None，防止 server_args 命名的 torch 后端在 MLX stub 的 _DummyKVCache 上构造崩溃。
• python/sglang/srt/server_args.py（模块 服务端参数；类别 source；类型 core-logic）: gpt-oss 的 CUDA 后端断言在 MPS 上会因 attention_backend 未设置而误失败，需用平台谓词收窄；同时配套 arg_groups/overrides.py 的后端强制逻辑。
• test/registered/unit/hardware_backend/mlx/test_sliding_window_attention.py（模块 滑动窗口测试；类别 test；类型 test-coverage；符号 _tiny_gpt_oss_model, TestGptOssAttentionContract, test_gpt_oss_attention_passes_contract, test_get_attention_scale_prefers_scale_over_sm_scale）: 24 个单元测试覆盖契约、mask 与 mlx_lm 参考网格比对、batched decode 与手工参考严格对比、AOT 门控，是本次功能正确性的主要防线。
• test/registered/mlx/models_e2e/test_gpt_oss_mlx_correctness.py（模块 端到端测试；类别 test；类型 test-coverage；符号 _available_gb, TestGptOssMlxCorrectness, test_basic_generation_nonempty, test_simple_arithmetic）: 端到端正确性验证：黑盒 serving smoke + 与未打补丁 mlx_lm 的逐 token 参考等价，提示词强制 >128 token 确保窗口真正参与，是防假阳性的关键设计。
• test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py（模块 桩契约测试；类别 test；类型 test-coverage；符号 test_stub_overrides_base_init_attention_backends, test_stub_init_attention_backends_keeps_attn_backend_none）: 防漂移守护：断言 MlxModelRunnerStub 仍覆盖 init_attention_backends 且 attn_backend 保持 None，防止未来重构丢失该覆盖导致 MLX 启动崩溃。

关键符号：get_attention_scale, get_container_window_size, get_layer_window_sizes, make_attention_mask, MLXAttentionWrapper.init, MLXAttentionWrapper._batched_decode, patch_model_attention, _build_rope_kernel, MlxModelRunnerStub.init_attention_backends

关键源码片段

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

注意力契约的核心改动：scale 属性放宽为 scale/sm_scale 二选一，新增 per-layer 窗口解析（layer_types + window_size/sliding_window），是所有 MLX 注意力模型加载与识别的入口。

# 注意力契约：mlx-lm 各模型对 softmax scale 的命名不统一，
# gpt_oss 使用 sm_scale，而大多数模型使用 scale。
SCALE_ATTRS = ("scale", "sm_scale")
​
# mlx-lm 容器对滑动窗口标量的命名有两种：window_size（gpt_oss、gemma4）
# 与 sliding_window（olmo3、llama SWA 变体）。
WINDOW_SIZE_ATTRS = ("window_size", "sliding_window")
​
​
def get_attention_scale(module: Any) -> float | None:
    # 任选其一即可满足契约；first_present_attr 按元组顺序取第一个存在属性。
    return first_present_attr(module, SCALE_ATTRS)
​
​
def get_container_window_size(model: Any) -> int | None:
    # 逐层定位容器：language_model -> model，读取容器级 scalar 窗口。
    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]:
    """按 mlx-lm 容器惯例解析每层窗口大小。
​
    gpt_oss / olmo3 等容器暴露 layer_types（每层一个条目，
    "sliding_attention" 标记窗口层）配一个 scalar 窗口。
    返回 {layer_idx: window 或 None}；不遵循该惯例时返回空 dict。
    """
    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)
    }

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

MLXAttentionWrapper 是 batched decode 的核心执行单元：新增 window_size 支持、scale/sinks 缓存、窗口截断 + 局部 padding mask，无窗口模型保持原路径。

# 滑动窗口层：pool 保留完整 KV 历史，但最新 token 只关注尾部 window 个 key。
# 对 softmax 注意力而言，这与旋转缓存数值等价，同时不破坏 radix 前缀复用。
window = self._window_size
if window is None:
    # 无窗口模型走原路径，pad 元数据直接复用上下文里的全量版本。
    pad_sizes = ctx.pad_sizes
else:
    # 上下文共享的 padding 元数据按全量长度计算，
    # 窗口层必须按 min(seq + 1, window) 重建有效长度与 pad。
    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]
​
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:
        # 只保留尾部 window 个 key/value，效果等价于 banded mask。
        k_all = k_all[:, :, -window:, :]
        v_all = v_all[:, :, -window:, :]
​
    pad = pad_sizes[i]
    if pad > 0:
        # 补齐对齐 batch 长度，pad 位置后续会被 mask 屏蔽。
        k_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=k_all.dtype)
        v_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=v_all.dtype)
        k_all = mx.concatenate([k_all, k_pad], axis=2)
        v_all = mx.concatenate([v_all, v_pad], axis=2)

评论区精华

Review 中几个关键交锋：

• server_args 平台特殊化方向：yeahdongcn 反对在 server_args.py 引入 _mlx_serving = is_mps() and use_mlx() 特例，认为让通用验证路径理解 MLX 细节是错误方向。作者回应：纯 revert 会破坏启动，因为 MPS 上 attention_backend 在断言执行前仍是 None（torch_native 默认值稍后才填充）；最终采用 elif not is_mps() 平台谓词同时收束 overrides 与断言，既保留非 MPS 的快速失败，又不污染 MLX 路径。
• 热路径缓存与 fail-fast：gemini-code-assist 建议把 get_attention_scale 与 getattr(inner, "sinks", None) 从每次 decode 的热路径移到 __init__ 一次性解析，既省掉每次属性扫描，又能在加载期尽早暴露坏模块。作者采纳，_scale 与 _sinks 均在 patch 时缓存。
• AOT RoPE 门控需复核：yeahdongcn 对 _build_rope_kernel 拒绝 scaled RoPE 的判定（_freqs、mscale、scale）@adityavaid 复核，后续无异议。
• 测试专用 env 注册：yeahdongcn 建议 SGLANG_MLX_TEST_* 只在测试文件里读 os.environ，不要进入全局 Envs 注册表；作者同意并移除。

• 注释风格：yeahdongcn 指出 model_runner_stub.py 注释“读起来像 AI 生成的”要求精简，作者已修剪。

  • server_args.py 平台特殊化的设计方向 (design): 作者说明纯 revert 会导致 MPS 上 attention_backend 为 None 时断言先行失败，最终保留断言并用 if not is_mps() 平台谓词同时收束 overrides 与 server_args。
  • attention scale 与 sinks 的热路径缓存 (performance): 作者采纳：_scale 与 _sinks 在 wrapper 构造时解析并缓存，scale 缺失直接抛 RuntimeError。
  • AOT RoPE 门控的正确性复核 (correctness): 复核后无异议，门控逻辑随 PR 合入；默认路径不受影响（SGLANG_MLX_USE_CUSTOM_ROPE 默认关闭）。
  • 测试专用环境变量是否进入全局注册表 (design): 作者同意：“Sure, it makes sense”，注册被移除，e2e 测试直接读 os.environ。
  • stub 注释风格 (style): 作者回复 “Done, trimmed”，注释已精简。

风险与影响

• 风险：
  • 窗口层 KV 冗余存储：设计上窗口层在 pool 里保留完整 KV 历史、读取时截断，长序列下会持续占用不被再读的 KV slot，PR 自述将 per-layer windowed pool 留作 TODO。极端长连接 + 多层滑窗会放大内存占用。
  • gemma3 风格模型语义分裂：model_patching.py 只对“没有 layer_types 但有 scalar window”的模型告警，decode 实际不应用窗口；若未来接入该类模型，>128 token 的输出会静默错误，目前仅靠日志提示。
  • AOT 门控依赖 mlx 内部属性：_build_rope_kernel 检查 _freqs、mscale、scale 等非公共属性，mlx 升级可能改变命名或语义，导致门控错误放行（算错）或误拒绝（性能回退）；需要依赖版本测试守护。
  • 注意力契约放宽的波及面：ATTENTION_API_ATTRS 把 scale 移出必选集合，is_attention_module 改为任一 scale 属性，所有 MLX 模型加载与 find_attention_layers 都会经过新判定；若未来某模型把 sm_scale 用于非 softmax 语义，可能被误判为注意力模块。
  • CI 覆盖缺口：新增测试全部依赖 mlx/mlx_lm，当前 CI runner 全部跳过，真实回归只能依赖 Apple Silicon 本机或门控的 macOS lane，存在覆盖盲区。
  • 影响：对 Apple Silicon 用户，gpt-oss 系列（如 gpt-oss-20b-MXFP4-Q8）现在可直接通过 SGLANG_USE_MLX=1 服务，且贪心解码与未打补丁 mlx_lm 逐 token 一致；已用 >128 token 提示词验证窗口真正参与计算。对既有 MLX 模型，无窗口 decode 路径逻辑保持字节级不变，qwen MoE 参考等价与服务 smoke 均通过。对非 MLX 平台（CUDA/ROCm/XPU/CPU），改动仅用 if not is_mps() 收窄了 gpt-oss 后端断言与 override，行为完全不变。对团队而言，这是 Apple 支持 roadmap（#19137）的关键里程碑，为后续 olmo3、gemma4、llama SWA 变体等窗口模型铺平了契约与基础设施。
  • 风险标记：窗口层 KV 全量存储未裁剪, gemma3 风格模型 decode 不应用窗口（仅告警）, AOT RoPE 门控依赖 mlx 内部属性, 非 Apple Silicon CI 全部跳过, 注意力契约放宽影响所有 MLX 模型加载

关联脉络

• PR #19137 [Roadmap] Apple Device Support (2026 Q2): 本 PR 直接承接 roadmap 的 gpt_oss 条目，PR body 开头即标注 Part of #19137，是 Apple Silicon 支持路线图上的关键功能补齐。
• PR #29440 qwen MoE MLX correctness tests structure: PR body 明确说明正确性测试结构遵循 qwen MoE MLX 测试（#29440），包括 serving smoke 与参考等价的组织方式。
• PR #30121 MLX CI suite registration mechanism: commit 说明滑动窗口单测与 gpt-oss e2e 测试按 #30121 的 run_suite.py suite 注册机制挂到 stage-a-unit-test-mlx / stage-b-e2e-mlx。
• PR #29217 forward_ct accounting change: commit 2f6407b 对齐 FakeOverlapScheduler stub 与 #29217 的 forward_ct 记账方式，属于本 PR 处理 main 分支冲突时的配套调整。