[MLX] Size request capacity by attention DP

[wirybeaver]

Summary

• Size MLX request capacity using ps.attn_dp_size, matching the canonical KV-cache resolver.
• Preserve the full configured request limit on independent pure-DP replicas.
• Partition request and explicit auxiliary-state capacity across attention-DP KV-cache owners.
• Keep infeasible auxiliary-state diagnostics in global CLI units for --max-mamba-cache-size.
• Add focused initialized regressions for pure DP, DP attention, hybrid auxiliary-state sharding, and the global-unit diagnostic.

Motivation

This is a prerequisite for #32101 and the Gemma 4 Apple Silicon MVP in #32102.

The MLX stub previously divided max_running_requests by the system dp_size. That is correct when DP attention is enabled because attn_dp_size == dp_size, but it underallocates pure-DP replicas: with DP attention disabled, attn_dp_size == 1 and every independent replica may need the full configured request limit.

For max_running_requests=8 and dp_size=4:

• pure DP: attn_dp_size=1, so each replica needs 8 slots;
• DP attention: attn_dp_size=4, so each KV-cache owner needs 2 slots.

This follows KVCacheConfigurator.resolve_max_num_reqs(), which partitions the limit by ps.attn_dp_size. For hybrid/linear-attention models, explicit max_mamba_cache_size is also sharded per attention-DP owner before sizing the MLX auxiliary-state pool.

Local Checks

Run from ~/sglang-attn-dp-request-capacity on ebb4e9cb6:

Check	Command	Result
Focused MLX unit	PYTHONPATH=python ~/uvenv/bin/python -m pytest test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py -q	4 passed
Ruff	uvx ruff@0.15.1 check python/sglang/srt/hardware_backend/mlx/model_runner_stub.py test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py	passed
Black	uvx black@26.1.0 --check python/sglang/srt/hardware_backend/mlx/model_runner_stub.py test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py	passed
isort	uvx isort@7.0.0 --check-only python/sglang/srt/hardware_backend/mlx/model_runner_stub.py test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py	passed
Registered tests	PYTHONPATH=python python3 scripts/ci/check_registered_tests.py	passed
py_compile	python3 -m py_compile python/sglang/srt/hardware_backend/mlx/model_runner_stub.py test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py	passed

This PR contains no Gemma 4 model-specific behavior; #32102 can rebase on it after merge.

---

CI States

Latest PR Test (Base): ❌ Run #30464730041
Latest PR Test (Extra): ❌ Run #30464721566

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[yeahdongcn]

cc @noob-se7en (he is working on a redesign of model_runner_stub to avoid this kind of mismatch with the real model runner).

[jlee5814]

Would it be worth having the new cases go through initialize()? The fixture sets stub.dp_size, an attribute production never has, which is the same construction that hid the missing attribute in #30547's suite.

[wirybeaver]

Thanks, updated in 9c7e71c6e and follow-up ebb4e9cb6.

Changes:

• _resolve_max_running_requests() now derives the explicit auxiliary-state cap per attention-DP owner (max_mamba_cache_size // ps.attn_dp_size) and initialize() uses the same value when constructing MlxAuxiliaryStateReqToTokenPool.
• Reworked the new regression to call initialize() and use ParallelState.trivial(..., attn_dp_size=...) instead of setting a synthetic stub.dp_size attribute.
• Added the hybrid attention-DP case with max_running_requests=8, max_mamba_cache_size=16: expected max_running_requests == 1 and auxiliary_state_pool.size == 4.
• Switched the registered unit test to CustomTestCase.
• Updated the infeasible auxiliary-state diagnostic to report the global CLI max_mamba_cache_size, include the derived per-worker cap, and recommend ratio * ps.attn_dp_size; added a regression for global cap 15 recommending 16.

Moved the local checks into the PR summary to keep this thread focused on the review feedback.

[yeahdongcn]

Nit: please preserve the global cap for diagnostics, or convert the recommendation back to global CLI units.

The division here is correct for per-owner allocation, but _resolve_max_running_requests() later interpolates this returned value as max_mamba_cache_size and tells the user to increase --max-mamba-cache-size to at least ratio. That flag is global, while both values in the message are now per-owner.

For attn_dp_size=4 and radix ratio 4:

global cap 4  -> error reports local cap 1 and says increase the flag to at least 4
global cap 15 -> error reports local cap 3 and still says at least 4
global cap 16 -> succeeds with local cap 4 and one request per owner

Following the current advice therefore repeats the same startup failure. Please retain/report the original global value and recommend at least ratio * self.ps.attn_dp_size (16 in this case), or explicitly label the reported cap as per-owner.

[yeahdongcn]

@wirybeaver Could you check if you still want to address this or just drop it? I'm fine with both (will run ci after you reply). Thanks!

[wirybeaver]

Addressed in ebb4e9cb6. I kept the allocation path per attention-DP worker, but changed the startup diagnostic back to global CLI units: it now reports max_mamba_cache_size=<global>, includes the derived per-worker cap, and recommends ratio * ps.attn_dp_size for --max-mamba-cache-size (16 for the attn-DP=4, ratio=4 case). Added test_attention_dp_auxiliary_error_reports_global_cli_units for the 15 -> recommend 16 regression.

Checked with ~/uvenv: focused MLX unit test is 4 passed; ruff, black, isort, registered-test check, py_compile all pass.

[yeahdongcn]

/tag-and-rerun-ci