[MLX] Fix startup crash when reporting preloaded weights (#37035)
This commit is contained in:
@@ -107,6 +107,16 @@ class MlxModelRunnerStub(ModelRunner):
|
|||||||
# that path working instead of raising AttributeError.
|
# that path working instead of raising AttributeError.
|
||||||
prefill_aware_swa = False
|
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
|
@staticmethod
|
||||||
def validate_startup_weight_load_mode(server_args) -> None:
|
def validate_startup_weight_load_mode(server_args) -> None:
|
||||||
if server_args.is_startup_weight_load_overlap:
|
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
|
uses one, and some backends read real KV buffers in ``__init__``, which
|
||||||
crashes on ``_DummyKVCache``.
|
crashes on ``_DummyKVCache``.
|
||||||
|
|
||||||
The checks are signature/identity-only and MLX-gated because importing
|
The base ``preloaded_weights_bytes`` property reads the Torch model loader.
|
||||||
the stub pulls in ``mlx.core``.
|
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
|
from __future__ import annotations
|
||||||
@@ -22,6 +25,7 @@ import inspect
|
|||||||
import unittest
|
import unittest
|
||||||
|
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
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_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||||
register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx")
|
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)
|
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||||
class TestMlxRunnerPoolContract(unittest.TestCase):
|
class TestMlxRunnerPoolContract(CustomTestCase):
|
||||||
"""``MlxModelRunnerStub.alloc_memory_pool`` must override the base."""
|
"""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):
|
def test_stub_overrides_base_alloc_memory_pool(self):
|
||||||
self.assertIn(
|
self.assertIn(
|
||||||
|
|||||||
Reference in New Issue
Block a user