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

Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>

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"requests: max_mamba_cache_size={global_aux_state_size} "
-                    f"only {aux_state_size // ratio} concurrent requests "
+                    f"backs only {aux_state_size // ratio} concurrent requests "
-                    f"({ratio} slots per request). Increase "
+                    f"per attention-DP worker (per-worker auxiliary-state "
-                    f"--max-mamba-cache-size to at least {ratio}, or leave it "
+                    f"cap={aux_state_size}, {ratio} slots per request). "
-                    f"unset to size the pool from the concurrency cap."
+                    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

@@ -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()