[MLX] Fix startup crash when reporting preloaded weights (#37035)

This commit is contained in:
Ali Ihsan Nergiz
2026-08-30 17:36:13 -07:00
committed by GitHub
parent bb5e619860
commit 0da6a66856
2 changed files with 25 additions and 4 deletions
@@ -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(