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 56af1b6d0..681adec25 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -196,3 +196,7 @@ class MlxModelRunnerStub(ModelRunner): f"max_running_requests={self.max_running_requests}, " f"zero GPU KV cache allocation)" ) + + def alloc_memory_pool(self, memory_pool_config=None): + """No-op: MLX manages its own KV cache.""" + pass 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 eac480e1d..d890564de 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/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, diff --git a/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py b/test/registered/unit/hardware_backend/mlx/test_mlx_runner_pool_contract.py new file mode 100644 index 000000000..e73162631 --- /dev/null +++ b/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() \ No newline at end of file