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 7f1c70393..52ffce041 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -107,6 +107,16 @@ class MlxModelRunnerStub(ModelRunner): # that path working instead of raising AttributeError. prefill_aware_swa = False + @property + def preloaded_weights_bytes(self) -> int: + """Return zero for the base Torch loader accounting hook. + + The native MLX runner materializes weights before sizing its own KV + pool. This stub has no Torch ``ModelLoader`` or Torch-owned weights to + report. + """ + return 0 + @staticmethod def validate_startup_weight_load_mode(server_args) -> None: if server_args.is_startup_weight_load_overlap: 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 index 7db768e98..3af37759b 100644 --- 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 @@ -11,8 +11,11 @@ attention backend named by ``server_args.attention_backend``; MLX never uses one, and some backends read real KV buffers in ``__init__``, which crashes on ``_DummyKVCache``. -The checks are signature/identity-only and MLX-gated because importing -the stub pulls in ``mlx.core``. +The base ``preloaded_weights_bytes`` property reads the Torch model loader. +The MLX stub never creates one because its native runner owns model loading +and KV sizing, so it must report zero for the Torch accounting hook. + +The checks are MLX-gated because importing the stub pulls in ``mlx.core``. """ from __future__ import annotations @@ -22,6 +25,7 @@ import inspect import unittest from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=1, suite="base-a-test-cpu") register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx") @@ -35,8 +39,15 @@ if _HAS_MLX: @unittest.skipUnless(_HAS_MLX, _SKIP_REASON) -class TestMlxRunnerPoolContract(unittest.TestCase): - """``MlxModelRunnerStub.alloc_memory_pool`` must override the base.""" +class TestMlxRunnerPoolContract(CustomTestCase): + """Guard the stub's scheduler-facing ``ModelRunner`` contracts.""" + + def test_stub_reports_no_preloaded_torch_weights(self): + runner = object.__new__(MlxModelRunnerStub) + runner.pre_model_load_memory = 1.0 + self.assertEqual(runner.preloaded_weights_bytes, 0) + runner.account_preloaded_weights(runner.preloaded_weights_bytes) + self.assertEqual(runner.pre_model_load_memory, 1.0) def test_stub_overrides_base_alloc_memory_pool(self): self.assertIn(