[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
@@ -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:
@@ -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(