From 9d0e6a2df41623105423738c767f1320686ba436 Mon Sep 17 00:00:00 2001 From: Lijuan Tang Date: Wed, 3 Jun 2026 09:03:46 -0700 Subject: [PATCH] fix(mlx): set canary_manager and materialize overlap-loop inputs on Apple Silicon (#26882) Signed-off-by: Xiaodong Ye Signed-off-by: LijuanTang94 Co-authored-by: Xiaodong Ye --- .../hardware_backend/mlx/model_runner_stub.py | 8 ++ .../hardware_backend/mlx/scheduler_mixin.py | 8 ++ python/sglang/srt/managers/scheduler.py | 33 ++++---- .../mlx/test_attention_patching.py | 80 +++++++++++++++++++ 4 files changed, 113 insertions(+), 16 deletions(-) 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 17f01e93e..adcba8a85 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -77,6 +77,14 @@ class MlxModelRunnerStub(ModelRunner): 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) diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 673aafda2..ae84db044 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, List, Optional import mlx.core as mx from sglang.srt.environ import envs +from sglang.srt.managers.overlap_utils import resolve_forward_inputs from sglang.srt.utils import DynamicGradMode logger = logging.getLogger(__name__) @@ -142,6 +143,13 @@ class SchedulerMlxOverlapMixin: pending_next: Optional[MlxPendingJob] = None def _launch_fresh(batch: "ScheduleBatch") -> MlxPendingJob: + # Materialize batch.input_ids from CPU staging (prefill) or the + # FutureMap relay (decode) before the forward. With deferred input + # materialization, get_next_batch_to_run leaves input_ids unset; the + # CUDA paths call resolve_forward_inputs for this, but the MLX overlap + # loop must do it too, otherwise async_forward_batch_generation_mlx + # dereferences a 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) ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index d1a7d2c9e..ac316b716 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1167,22 +1167,6 @@ class Scheduler( def init_overlap(self): self.device_module = torch.get_device_module(self.device) - if use_mlx(): - # MLX: no CUDA streams / FutureMap. - self.future_map = None - 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 - ) - # 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 @@ -1202,6 +1186,23 @@ class Scheduler( needs_cpu_seq_lens=needs_cpu_seq_lens, ) + if use_mlx(): + # MLX uses its own overlap loop and does not create CUDA streams, + # but the normal non-overlap scheduler path still relays decode + # input IDs through FutureMap. + 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 diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index e82b1e8dc..6f3546226 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -4,6 +4,7 @@ from __future__ import annotations import importlib.util import unittest +from collections import deque from types import SimpleNamespace from sglang.test.ci.ci_register import register_cpu_ci @@ -381,6 +382,43 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): self.assertEqual(calls, [(1, [[7]], ["r0"])]) self.assertEqual(pending.lazy_tokens.tolist(), [8]) + def test_mlx_scheduler_init_overlap_keeps_future_map_relay(self): + from sglang.srt.managers import scheduler as scheduler_module + from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + scheduler = object.__new__(Scheduler) + scheduler.device = "cpu" + scheduler.draft_worker = None + scheduler.tp_worker = SimpleNamespace( + model_runner=SimpleNamespace(attn_backend=None) + ) + scheduler.server_args = SimpleNamespace( + enable_two_batch_overlap=False, + disable_piecewise_cuda_graph=True, + ) + scheduler.spec_algorithm = SpeculativeAlgorithm.NONE + scheduler.req_to_token_pool = ReqToTokenPool( + size=4, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + ) + scheduler.enable_overlap = False + + original_use_mlx = scheduler_module.use_mlx + scheduler_module.use_mlx = lambda: True + try: + Scheduler.init_overlap(scheduler) + finally: + scheduler_module.use_mlx = original_use_mlx + + self.assertIsNotNone(scheduler.future_map) + indices = torch.tensor([1], dtype=torch.int64) + scheduler.future_map.stash(indices, torch.tensor([7], dtype=torch.int64)) + self.assertEqual(int(scheduler.future_map.output_tokens_buf[1].item()), 7) + def test_decode_finalize_does_not_snapshot_auxiliary_state(self): runner = object.__new__(MlxModelRunner) runner._req_token_ids = {"r0": [8]} @@ -1082,6 +1120,48 @@ class TestMlxOverlapScheduler(unittest.TestCase): self.assertIs(scheduler.processed_batch, batch_copy) self.assertIs(scheduler.processed_result, scheduler.tp_worker.result) + def test_overlap_loop_materializes_prefill_input_ids(self): + # Regression: the MLX overlap loop must materialize batch.input_ids + # (deferred input materialization) before launching the forward. + # Without resolve_forward_inputs in _launch_fresh, input_ids stays + # None and async_forward_batch_generation_mlx dereferences a None. + class _StopLoop(Exception): + pass + + captured = {} + + def fake_forward(batch): + captured["input_ids"] = batch.input_ids + raise _StopLoop + + scheduler = SchedulerMlxOverlapMixin.__new__(SchedulerMlxOverlapMixin) + scheduler.request_receiver = SimpleNamespace(recv_requests=lambda: []) + scheduler.process_input_requests = lambda recv_reqs: None + scheduler._engine_paused = False + scheduler.waiting_queue = [] + scheduler.result_queue = deque() + scheduler.future_map = SimpleNamespace() + scheduler.cur_batch = None + scheduler.last_batch = None + scheduler.tp_worker = SimpleNamespace( + async_forward_batch_generation_mlx=fake_forward + ) + + batch = SimpleNamespace( + prefill_input_ids_cpu=torch.tensor([1, 2, 3], dtype=torch.int64), + input_ids=None, + mix_running_indices=None, + is_spec_v2=False, + device="cpu", + ) + scheduler.get_next_batch_to_run = lambda: batch + + with self.assertRaises(_StopLoop): + scheduler.event_loop_overlap_mlx() + + self.assertIsNotNone(captured["input_ids"]) + self.assertTrue(torch.equal(captured["input_ids"], torch.tensor([1, 2, 3]))) + def test_finished_request_snapshots_before_release(self): events = [] tree_cache = object()