fix(mlx): set canary_manager and materialize overlap-loop inputs on Apple Silicon (#26882)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Signed-off-by: LijuanTang94 <tang.lij@northeastern.edu>
Co-authored-by: Xiaodong Ye <yeahdongcn@gmail.com>

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)

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)
             )

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
 

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()