minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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()