[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.
|
||||
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(
|
||||
|
||||
Reference in New Issue
Block a user