minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[MLX] Size request capacity by attention DP (#32115)

Browse Source 
Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
This commit is contained in:
Xuanyi Li
2026-07-29 18:18:22 -07:00
committed by GitHub
co-authored by R0CKSTAR
parent 1d9c292547
commit 8fbf960980
2 changed files with 148 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
+22 -10
View File
@@ -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()
test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py
+126
View File
@@ -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()