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