From 8fbf960980022bdad198d9f66deff8e52ae8826c Mon Sep 17 00:00:00 2001 From: Xuanyi Li Date: Wed, 29 Jul 2026 18:18:22 -0700 Subject: [PATCH] [MLX] Size request capacity by attention DP (#32115) Co-authored-by: R0CKSTAR --- .../hardware_backend/mlx/model_runner_stub.py | 32 +++-- .../mlx/test_attn_dp_request_capacity.py | 126 ++++++++++++++++++ 2 files changed, 148 insertions(+), 10 deletions(-) create mode 100644 test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py 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 84f58eaeb..920efb0d0 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -148,12 +148,20 @@ class MlxModelRunnerStub(ModelRunner): return 1 return MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO + def _explicit_aux_state_size_per_worker(self) -> int | None: + """Return the explicit auxiliary-state cap for this attention-DP owner.""" + aux_state_size = self.server_args.max_mamba_cache_size + if aux_state_size is None: + return None + return aux_state_size // self.ps.attn_dp_size + def _resolve_max_running_requests(self) -> int: """Concurrency cap handed to the scheduler. Honors ``--max-running-requests``, mirroring the base runner's clamp (``model_runner_kv_cache_mixin._resolve_max_num_reqs``): the requested - value is split per dp worker and capped by the KV pool capacity. When + value is split across attention-DP KV-cache owners and capped by the KV + pool capacity. Pure-DP replicas retain the full per-replica limit. When the flag is unset, fall back to a capacity-based default. On hybrid / linear-attention models the concurrency is additionally @@ -170,10 +178,10 @@ class MlxModelRunnerStub(ModelRunner): requested_per_worker = None resolved = min(capacity_cap, 4096) else: - requested_per_worker = requested // self.dp_size + requested_per_worker = requested // self.ps.attn_dp_size resolved = min(requested_per_worker, capacity_cap) - aux_state_size = self.server_args.max_mamba_cache_size + aux_state_size = self._explicit_aux_state_size_per_worker() if ( mambaish_config(self.model_config) is not None and aux_state_size is not None @@ -181,19 +189,23 @@ class MlxModelRunnerStub(ModelRunner): ratio = self._aux_state_slots_per_request() resolved = min(resolved, aux_state_size // ratio) if resolved <= 0: + global_aux_state_size = self.server_args.max_mamba_cache_size + min_global_aux_state_size = ratio * self.ps.attn_dp_size raise RuntimeError( f"MLX auxiliary-state cache is too small to serve any " - f"requests: max_mamba_cache_size={aux_state_size} backs " - f"only {aux_state_size // ratio} concurrent requests " - f"({ratio} slots per request). Increase " - f"--max-mamba-cache-size to at least {ratio}, or leave it " - f"unset to size the pool from the concurrency cap." + f"requests: max_mamba_cache_size={global_aux_state_size} " + f"backs only {aux_state_size // ratio} concurrent requests " + f"per attention-DP worker (per-worker auxiliary-state " + f"cap={aux_state_size}, {ratio} slots per request). " + f"Increase --max-mamba-cache-size to at least " + f"{min_global_aux_state_size}, or leave it unset to size " + f"the pool from the concurrency cap." ) if requested_per_worker is not None and resolved < requested_per_worker: logger.warning( "max_running_requests was reduced from the requested %d to %d " - "(per dp worker) due to the available KV cache or " + "(per attention-DP worker) due to the available KV cache or " "auxiliary-state capacity.", requested_per_worker, resolved, @@ -241,7 +253,7 @@ class MlxModelRunnerStub(ModelRunner): # Create minimal pools if mambaish_config(self.model_config) is not None: - auxiliary_state_size = self.server_args.max_mamba_cache_size + auxiliary_state_size = self._explicit_aux_state_size_per_worker() if auxiliary_state_size is None: auxiliary_state_size = ( self.max_running_requests * self._aux_state_slots_per_request() diff --git a/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py new file mode 100644 index 000000000..3db8a7921 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py @@ -0,0 +1,126 @@ +"""Request-capacity accounting for pure DP and DP attention on MLX. + +``max_running_requests`` is partitioned only when attention DP partitions a +logical batch across multiple KV-cache owners. Pure data-parallel replicas own +independent schedulers and caches, so each replica retains the full configured +limit even when the system DP size is greater than one. +""" + +from __future__ import annotations + +import importlib.util +import unittest +from types import SimpleNamespace +from unittest import mock + +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") + +_HAS_MLX = importlib.util.find_spec("mlx") is not None +_SKIP_REASON = "requires mlx" + +if _HAS_MLX: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.hardware_backend.mlx.model_runner_stub import ( + MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO as RATIO, + ) + from sglang.srt.hardware_backend.mlx.model_runner_stub import ( + MlxModelRunnerStub, + ) + + +def _arch(*, hybrid: bool): + return mock.patch( + "sglang.srt.hardware_backend.mlx.model_runner_stub.mambaish_config", + return_value=object() if hybrid else None, + ) + + +def _stub_for_initialize( + *, + dp_size: int, + attn_dp_size: int, + max_running_requests: int = 8, + max_mamba_cache_size: int | None = None, + pool_size: int = 64, +): + stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub) + stub._mlx_pool_size = pool_size + stub.device = "cpu" + stub.ps = ParallelState.trivial(dp_size=dp_size, attn_dp_size=attn_dp_size) + stub.server_args = SimpleNamespace( + enable_memory_saver=False, + max_running_requests=max_running_requests, + max_mamba_cache_size=max_mamba_cache_size, + disable_radix_cache=False, + ) + stub.model_config = SimpleNamespace( + is_hybrid_swa=False, + sliding_window_size=None, + attention_chunk_size=None, + dtype="float16", + num_hidden_layers=1, + num_attention_layers=1, + context_len=64, + use_ngram_embedding=False, + ) + return stub + + +def _initialize_stub(stub, *, hybrid: bool = False): + with _arch(hybrid=hybrid): + stub.initialize() + return stub + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestAttentionDpRequestCapacity(CustomTestCase): + def test_pure_dp_replica_retains_full_request_limit(self): + stub = _initialize_stub( + _stub_for_initialize(dp_size=4, attn_dp_size=1), + ) + self.assertEqual(stub.max_running_requests, 8) + self.assertEqual(stub.req_to_token_pool.size, 8) + + def test_attention_dp_partitions_request_limit(self): + stub = _initialize_stub( + _stub_for_initialize(dp_size=4, attn_dp_size=4), + ) + self.assertEqual(stub.max_running_requests, 2) + self.assertEqual(stub.req_to_token_pool.size, 2) + + def test_attention_dp_partitions_explicit_auxiliary_state_limit(self): + stub = _initialize_stub( + _stub_for_initialize( + dp_size=4, + attn_dp_size=4, + max_running_requests=8, + max_mamba_cache_size=4 * RATIO, + ), + hybrid=True, + ) + + self.assertEqual(stub.max_running_requests, 1) + self.assertEqual(stub.req_to_token_pool.size, 1) + self.assertEqual(stub.req_to_token_pool.auxiliary_state_pool.size, RATIO) + + def test_attention_dp_auxiliary_error_reports_global_cli_units(self): + stub = _stub_for_initialize( + dp_size=4, + attn_dp_size=4, + max_running_requests=8, + max_mamba_cache_size=4 * RATIO - 1, + ) + + with self.assertRaisesRegex( + RuntimeError, + "max_mamba_cache_size=15.*per-worker auxiliary-state cap=3.*" "at least 16", + ): + _initialize_stub(stub, hybrid=True) + + +if __name__ == "__main__": + unittest.main()