[MLX] Fix Apple Silicon server startup; align MLX tests with upstream (#28770)

Co-authored-by: “toufupi” <“byte2016@outlook.com”>

test/registered/unit/hardware_backend/mlx/test_attention_patching.py

@@ -397,6 +397,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
         scheduler.server_args = SimpleNamespace(
             enable_two_batch_overlap=False,
             cuda_graph_config=None,
+            speculative_algorithm=None,
         )
         scheduler.spec_algorithm = SpeculativeAlgorithm.NONE
         scheduler.req_to_token_pool = ReqToTokenPool(
@@ -423,6 +424,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
         runner = object.__new__(MlxModelRunner)
         runner._req_token_ids = {"r0": [8]}
         runner._decode_step_ct = 0
+        runner._clear_steps = 0
         calls = []
         runner._store_auxiliary_state = lambda req_pool_idx, cache: calls.append(
             (req_pool_idx, cache)
@@ -1180,7 +1182,7 @@ class TestMlxOverlapScheduler(unittest.TestCase):
             model_config=None,
             token_to_kv_pool_allocator=None,
             tree_cache=tree_cache,
-            hisparse_coordinator=None,
+            hisparse_coordinator=SimpleNamespace(request_finished=lambda req: None),
             req_to_token_pool=None,
             decode_offload_manager=None,
             metrics_collector=None,
@@ -1195,34 +1197,65 @@ class TestMlxOverlapScheduler(unittest.TestCase):
             output_streamer=None,
             abort_request=lambda req: None,
         )
+        # Stub out the methods _handle_finish_state_updated_req calls that
+        # are not relevant to this test.  SchedulerBatchResultProcessor is
+        # @dataclass(slots=True, frozen=True), so patches go on the class.
+        noop_stubs = {
+            "_mamba_prefix_cache_update": lambda *a, **k: None,
+            "_maybe_collect_routed_experts": lambda *a, **k: None,
+            "_maybe_collect_indexer_topk": lambda *a, **k: None,
+            "_maybe_collect_customized_info": lambda *a, **k: None,
+        }
+        saved = {
+            name: getattr(SchedulerBatchResultProcessor, name)
+            for name in noop_stubs
+        }
+        for name, value in noop_stubs.items():
+            setattr(SchedulerBatchResultProcessor, name, value)
         req = SimpleNamespace(
             rid="r0",
             finished=lambda: True,
             multimodal_inputs=None,
             session=None,
             return_routed_experts=False,
+            mamba_lazy_is_insert=True,
             time_stats=SimpleNamespace(
                 set_completion_time=lambda: events.append(("completion", "r0"))
             ),
         )
+        batch = SimpleNamespace()
+        result = SimpleNamespace()
+        i = 0
+        logits_output = SimpleNamespace(customized_info=None)
         original_release = batch_result_processor_module.release_kv_cache
         original_get_indexer = batch_result_processor_module.get_global_indexer_capturer
+        original_get_server_args = (
+            batch_result_processor_module.get_global_server_args
+        )
 
-        def fake_release_kv_cache(release_req, tree_cache):
+        def fake_release_kv_cache(release_req, tree_cache, is_insert=False):
             events.append(("release", release_req.rid))
             self.assertIs(tree_cache, processor.tree_cache)
 
         batch_result_processor_module.release_kv_cache = fake_release_kv_cache
         batch_result_processor_module.get_global_indexer_capturer = lambda: None
+        batch_result_processor_module.get_global_server_args = lambda: SimpleNamespace(
+            enable_mamba_extra_buffer_lazy=lambda: False
+        )
         try:
-            SchedulerBatchResultProcessor._handle_finished_req(
-                processor, req, 0, SimpleNamespace(customized_info=None)
+            SchedulerBatchResultProcessor._handle_finish_state_updated_req(
+                processor, req, batch, result, i, logits_output
             )
         finally:
+            for name, original in saved.items():
+                setattr(SchedulerBatchResultProcessor, name, original)
             batch_result_processor_module.release_kv_cache = original_release
             batch_result_processor_module.get_global_indexer_capturer = (
                 original_get_indexer
             )
+            batch_result_processor_module.get_global_server_args = (
+                original_get_server_args
+            )
 
         self.assertEqual(
             events,

test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py

@@ -0,0 +1,79 @@
+"""Guard the MLX stub's ``alloc_memory_pool`` override against drift.
+
+The base ``ModelRunner.alloc_memory_pool`` runs ``_init_pools`` which
+asserts ``is_draft_worker`` (model_runner_kv_cache_mixin.py:409); the
+MLX stub manages its own KV cache via ``MlxAttentionKVPool`` and must
+short-circuit that GPU-allocation path.  If the override is lost, every
+MLX startup crashes inside ``Scheduler.init_target_memory_pool``.
+
+The checks are signature/identity-only and MLX-gated because importing
+the stub pulls in ``mlx.core``.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+_HAS_MLX = importlib.util.find_spec("mlx") is not None
+_SKIP_REASON = "requires mlx"
+
+if _HAS_MLX:
+    from sglang.srt.hardware_backend.mlx.model_runner_stub import MlxModelRunnerStub
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
+class TestMlxRunnerPoolContract(unittest.TestCase):
+    """``MlxModelRunnerStub.alloc_memory_pool`` must override the base."""
+
+    def test_stub_overrides_base_alloc_memory_pool(self):
+        self.assertIn(
+            "alloc_memory_pool",
+            vars(MlxModelRunnerStub),
+            msg=(
+                "MlxModelRunnerStub lost its alloc_memory_pool override. "
+                "Without it the base ModelRunner.alloc_memory_pool runs "
+                "_init_pools, which asserts is_draft_worker "
+                "(model_runner_kv_cache_mixin.py:409) and crashes every "
+                "MLX startup. Re-add the no-op override."
+            ),
+        )
+        self.assertIsNot(
+            MlxModelRunnerStub.alloc_memory_pool,
+            ModelRunner.alloc_memory_pool,
+            msg="alloc_memory_pool must be overridden on the MLX stub, "
+            "not inherited from ModelRunner.",
+        )
+
+    def test_stub_alloc_memory_pool_binds_with_no_args(self):
+        sig = inspect.signature(MlxModelRunnerStub.alloc_memory_pool)
+        try:
+            sig.bind(object())
+        except TypeError as exc:
+            self.fail(
+                "MlxModelRunnerStub.alloc_memory_pool must accept a no-arg "
+                f"call (scheduler default): {exc}"
+            )
+
+    def test_stub_alloc_memory_pool_binds_with_optional_config(self):
+        class _FakeConfig:
+            pass
+
+        sig = inspect.signature(MlxModelRunnerStub.alloc_memory_pool)
+        try:
+            sig.bind(object(), _FakeConfig())
+        except TypeError as exc:
+            self.fail(
+                "MlxModelRunnerStub.alloc_memory_pool must accept an "
+                f"optional MemoryPoolConfig argument: {exc}"
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()