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>

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