执行摘要
- 一句话：修复 MLX 后端 canary_manager 缺失和 overlap loop 输入未物化崩溃
- 推荐动作：该 PR 值得阅读，尤其对于理解 SGLang 调度器中硬件后端抽象（FutureMap、canary_manager）的设计和演进。展示了一个小而优雅的修复：通过添加类属性和调整初始化顺序避免重写整个基类流程。同时，与 #26952 的对比体现了不同修复策略的权衡。
功能与动机
用户报告在 macOS 上使用 MLX 后端运行 sglang 服务器时崩溃（Issue #26832）。分析发现两个根本原因：1）MlxModelRunnerStub 的 initialize 轻量实现跳过了 install_canary，导致 canary_manager 属性缺失；2）随着 deferred input materialization 的引入，CUDA 路径调用了 resolve_forward_inputs，但 MLX overlap 循环没有，导致 input_ids 为 None。本 PR 同时修复这两个问题。
实现拆解
-
添加 canary_manager 类属性：在 MlxModelRunnerStub 中添加 canary_manager = None 作为类属性。这样在基类 ModelRunner 的 install_canary 未被调用时，下游代码对 canary_manager is not None 的检查依然安全，不会抛出 AttributeError。
-
调整 init_overlap 中 FutureMap 创建时机：在 scheduler.py 的 init_overlap 中，将 MLX 分支从函数开头移到 FutureMap 创建之后。这样 MLX 后端也能获得一个真实的 FutureMap 实例（通过基类通用的 create_future_map），而非 None。这为下一步在 overlap 循环中使用 resolve_forward_inputs 提供了依赖。
-
在 MLX overlap 循环中调用 resolve_forward_inputs：在 scheduler_mixin.py 的 _launch_fresh 函数内，在调用 async_forward_batch_generation_mlx 之前，先执行 resolve_forward_inputs(batch, self.future_map)，将 CPU staging 中的 tokens 或 FutureMap relay 中的 decode tokens 填充到 batch.input_ids，避免 None 引用。
-
添加回归测试：在 test_attention_patching.py 中新增两个测试：test_mlx_scheduler_init_overlap_keeps_future_map_relay 验证 init_overlap 后 future_map 不为 None 且能正常 stash；test_overlap_loop_materializes_prefill_input_ids 验证 _launch_fresh 执行后 batch.input_ids 已被物化。
关键文件：
python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 init_overlap）: 调整了 init_overlap 中 FutureMap 创建与 MLX 分支的顺序，使 MLX 也拥有 FutureMap，从而支撑 overlap loop 中的 resolve_forward_inputs。python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py（模块 MLX 流水线；类别 source；类型 dependency-wiring；符号 _launch_fresh, event_loop_overlap_mlx）: 在 _launch_fresh 中添加 resolve_forward_inputs 调用，物化 batch.input_ids，防止 forward 时出现 None 引用。python/sglang/srt/hardware_backend/mlx/model_runner_stub.py（模块 MLX 后端；类别 source；类型 data-contract；符号 MlxModelRunnerStub）: 添加 canary_manager = None 作为类属性，防止因跳过基类 initialize 中的 install_canary 而引发 AttributeError。test/registered/unit/hardware_backend/mlx/test_attention_patching.py（模块 MLX 测试；类别 test；类型 test-coverage；符号 test_mlx_scheduler_init_overlap_keeps_future_map_relay, test_overlap_loop_materializes_prefill_input_ids, _StopLoop, fake_forward）: 添加了两个回归测试，直接验证修复点：FutureMap 创建和 overlap 循环输入物化。
关键符号：init_overlap, _launch_fresh, test_mlx_scheduler_init_overlap_keeps_future_map_relay, test_overlap_loop_materializes_prefill_input_ids
关键源码片段
python/sglang/srt/managers/scheduler.py
调整了 init_overlap 中 FutureMap 创建与 MLX 分支的顺序，使 MLX 也拥有 FutureMap，从而支撑 overlap loop 中的 resolve_forward_inputs。
def init_overlap(self):
self.device_module = torch.get_device_module(self.device)
# FutureMap is always-on: input_ids relay used in both modes.
# Workers not on BaseSpecWorker (e.g. FrozenKVMTPWorker) lack the
# override; fall back to target-only so the helper still produces a
# safe decision (no accidental opt-out for unaudited shapes).
if self.draft_worker is not None:
attn_backends = getattr(
self.draft_worker,
"spec_v2_attn_backends",
(self.tp_worker.model_runner.attn_backend,),
)
else:
attn_backends = (self.tp_worker.model_runner.attn_backend,)
needs_cpu_seq_lens = decide_needs_cpu_seq_lens(self.server_args, attn_backends)
self.future_map = self.spec_algorithm.create_future_map(
self.device,
self.req_to_token_pool,
needs_cpu_seq_lens=needs_cpu_seq_lens,
)
# MLX 分支移到 FutureMap 创建之后，这样 MLX 也拥有真实的 FutureMap，
# 供后续 overlap loop 中的 resolve_forward_inputs 使用。
if use_mlx():
self.result_queue: Deque = deque()
return
# forward_stream_ctx / copy_stream are also used by PP (non-overlap)
# via scheduler_pp_mixin; init unconditionally to match main.
self.forward_stream_ctx: CudaStreamContext = self.device_module.stream(
self.forward_stream
)
self.copy_stream: CudaStream = self.device_module.Stream()
self.copy_stream_ctx: CudaStreamContext = self.device_module.stream(
self.copy_stream
)
if not self.enable_overlap:
return
self.batch_record_buf = [None] * 2
self.batch_record_ct = 0
python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py
在 _launch_fresh 中添加 resolve_forward_inputs 调用，物化 batch.input_ids，防止 forward 时出现 None 引用。
def _launch_fresh(batch: "ScheduleBatch") -> MlxPendingJob:
# 物化 batch.input_ids：从 CPU staging（prefill）或 FutureMap relay（decode）
# 中填充 input_ids。随着 deferred input materialization 的引入，
# get_next_batch_to_run 会留下未设置的 input_ids；CUDA 路径调用
# resolve_forward_inputs 处理，MLX overlap 循环也必须这样做，否则
# async_forward_batch_generation_mlx 会解引用 None 的 input_ids。
resolve_forward_inputs(batch, self.future_map)
lazy_tokens, prefills, extends, decode, mode = (
self.tp_worker.async_forward_batch_generation_mlx(batch)
)
return MlxPendingJob(
lazy_tokens=lazy_tokens,
prefills=prefills,
extends=extends,
decode=decode,
mode=mode,
batch_copy=batch.copy(),
schedule_batch=batch,
reqs=list(batch.reqs),
)
python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
添加 canary_manager = None 作为类属性，防止因跳过基类 initialize 中的 install_canary 而引发 AttributeError。
class MlxModelRunnerStub(ModelRunner):
"""ModelRunner that skips PyTorch weight loading and KV cache allocation.
Overrides both load_model() and initialize() so that no PyTorch model
weights are loaded and no large KV cache tensors are allocated. Only
the minimal bookkeeping pools needed by the scheduler are created.
"""
# No KV canary on the MLX path. The base ModelRunner installs it via
# install_canary() in its full initialize(), which this lightweight override
# skips. Downstream consumers (scheduler, cuda graph runner, speculative
# workers) all guard with `canary_manager is not None`, so default to None
# as a class attribute to keep those checks working instead of raising
# AttributeError.
canary_manager = None
def __init__(self, *args, mlx_pool_size: int | None = None, **kwargs):
self._mlx_pool_size = mlx_pool_size
super().__init__(*args, **kwargs)
评论区精华
风险与影响
- 风险：主要风险：修改了核心调度器
init_overlap 的控制流顺序，可能影响非 MLX 后端（如 CUDA）的初始化路径。但变更仅限于将 MLX 分支从函数开头移到 FutureMap 创建之后，并添加了 return 语句，逻辑上不会影响下面的 CUDA 流创建。同时，测试只能在 Apple Silicon 硬件上运行，持续集成可能无法覆盖。另外，添加的类属性 canary_manager = None 需确保所有消费者确实使用 is not None 守卫，没有直接访问该属性的场景。
- 影响：对用户：MLX 后端（Apple Silicon）用户可正常启动服务器并进行推理，不再出现 AttributeError 和崩溃。对系统：影响限于 MLX 后端，CUDA 等其他后端无影响（因 use_mlx() 返回 False 时不进入相关分支）。对团队：提供了 MLX 后端适配 deferred input materialization 的参考实现，也促成了与 #26952 方案的对比讨论。
- 风险标记：核心路径变更, 硬件依赖测试
关联脉络
- PR #26952 [MLX] Fix AttributeError in overlap scheduler on Apple Silicon: 相同问题通过不同方案修复（使 resolve_forward_inputs 接受 future_map=None），本 PR body 中进行了对比说明且保证无冲突。
参与讨论