fix(PP): size the mamba pool per pipeline stage, not per whole model (#33666)

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -24,6 +24,7 @@ from sglang.srt.configs.model_config import (
     is_minimax_sparse,
 )
 from sglang.srt.distributed.parallel_state import get_world_group
+from sglang.srt.distributed.utils import get_pp_indices
 from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
     get_kv_cache_quant_method,
@@ -1817,6 +1818,27 @@ class KVCacheConfigurator:
         server_args = self.server_args
         assert config is not None
 
+        # mamba_cache_per_req covers every mamba layer, but under PP a rank only
+        # allocates its own [start_layer, end_layer) slice. Charge the largest
+        # per-stage share so every rank derives the same pool without a collective.
+        all_mamba_layers = config.mamba2_cache_params.layers
+        if self.ps.pp_size > 1 and all_mamba_layers:
+            max_stage_mamba_layers = max(
+                sum(1 for i in all_mamba_layers if start <= i < end)
+                for start, end in (
+                    get_pp_indices(
+                        self.model_config.num_hidden_layers, rank, self.ps.pp_size
+                    )
+                    for rank in range(self.ps.pp_size)
+                )
+            )
+        else:
+            max_stage_mamba_layers = len(all_mamba_layers)
+        pp_layer_scale = max_stage_mamba_layers / max(len(all_mamba_layers), 1)
+        stage_per_req = int(
+            config.mamba2_cache_params.mamba_cache_per_req * pp_layer_scale
+        )
+
         has_spec_dec = not self.spec_algorithm.is_none()
         # ReplaySSM drops the per-step intermediate_ssm scratch, so its mamba budget
         # no longer reserves the (1 + D/ratio) intermediate factor -- the whole
@@ -1844,6 +1866,7 @@ class KVCacheConfigurator:
             )
         else:
             replayssm_ring_per_req = 0
+        replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale)
         if has_spec_dec:
             assert get_spec().speculative_num_draft_tokens is not None
             assert get_schedule().max_running_requests is not None
@@ -1865,7 +1888,7 @@ class KVCacheConfigurator:
                     get_schedule().max_mamba_cache_size // ratio,
                 )
                 intermediate_size = (
-                    config.mamba2_cache_params.mamba_cache_per_req
+                    stage_per_req
                     * (capped_reqs + 1)
                     * get_spec().speculative_num_draft_tokens
                 )
@@ -1884,15 +1907,15 @@ class KVCacheConfigurator:
             # pool's padding slot). Skipped under replayssm.
             if has_spec_dec and not replayssm_active:
                 intermediate_size = (
-                    config.mamba2_cache_params.mamba_cache_per_req
+                    stage_per_req
                     * (get_schedule().max_mamba_cache_size + 1)
                     * get_spec().speculative_num_draft_tokens
                 )
                 total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30))
         else:
             # Use ratio-based calculation to auto-fit available memory
-            assert config.mamba2_cache_params.mamba_cache_per_req > 0
-            per_req = config.mamba2_cache_params.mamba_cache_per_req
+            assert stage_per_req > 0
+            per_req = stage_per_req
 
             # Solve jointly for max_mamba_cache_size (K), including the pool's
             # +1 padding slot on both buffers (see memory_pool.py):
@@ -1941,7 +1964,7 @@ class KVCacheConfigurator:
                 f"Not enough GPU memory for hybrid (mamba/linear-attention) state cache. "
                 f"Computed max_mamba_cache_size={get_schedule().max_mamba_cache_size} "
                 f"(total_rest_memory={total_rest_memory:.2f} GB, "
-                f"mamba_cache_per_req={config.mamba2_cache_params.mamba_cache_per_req / (1 << 20):.2f} MB). "
+                f"mamba_cache_per_req={stage_per_req / (1 << 20):.2f} MB). "
                 f"Try: (1) reduce --max-running-requests, "
                 f"(2) increase --mem-fraction-static, "
                 f"(3) reduce --speculative-num-draft-tokens, or "
@@ -1953,7 +1976,7 @@ class KVCacheConfigurator:
         # the ring is not allocated).
         mamba_state_memory = (
             (get_schedule().max_mamba_cache_size + 1)
-            * (config.mamba2_cache_params.mamba_cache_per_req + replayssm_ring_per_req)
+            * (stage_per_req + replayssm_ring_per_req)
             / (1 << 30)
         )
         return total_rest_memory - mamba_state_memory

test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py

@@ -232,5 +232,84 @@ class TestMambaDonatedAllocRatio(unittest.TestCase):
         self.assertEqual(len(cache.prefix_nodes), N - 1)
 
 
+class TestPPMambaPoolSizing(unittest.TestCase):
+    """A PP rank only allocates mamba state for its own [start_layer, end_layer)
+    slice, so charging it for the whole model's layers starves the pool. Sizing
+    uses the largest per-stage share, which also keeps every rank on the same
+    pool size (and hence the same max_running_requests / pp_max_micro_batch_size)
+    without a collective."""
+
+    # Kimi-K3 shaped: 93 layers, linear attention everywhere except every 4th and
+    # the last, so the 69 mamba layers split unevenly over 8 stages (9 or 8 each).
+    TOTAL_LAYERS = 93
+    MAMBA_LAYERS = [i for i in range(93) if (i + 1) % 4 != 0 and i <= 90]
+    BUDGET_GB = 8.0
+
+    @classmethod
+    def _pool_size(cls, pp_rank, pp_size):
+        from sglang.srt import runtime_context as rc
+        from sglang.srt.configs.mamba_utils import (
+            Mamba2CacheParams,
+            Mamba2StateDType,
+            Mamba2StateShape,
+        )
+        from sglang.srt.distributed.utils import get_pp_indices
+        from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
+        from sglang.srt.runtime_context import get_schedule
+
+        shape = Mamba2StateShape(
+            conv=[(4096, 3)],
+            temporal=(64, 128, 128),
+            intermediate_size=0,
+            conv_dim=0,
+            ssm_state_size=0,
+            num_heads=0,
+            head_dim=0,
+            state_size=0,
+            conv_kernel=0,
+            num_k_heads_per_tp=8,
+        )
+        params = Mamba2CacheParams(
+            shape=shape,
+            dtype=Mamba2StateDType(conv=torch.bfloat16, temporal=torch.float32),
+            layers=list(cls.MAMBA_LAYERS),
+        )
+        start, end = get_pp_indices(cls.TOTAL_LAYERS, pp_rank, pp_size)
+        fake = SimpleNamespace(
+            mambaish_config=SimpleNamespace(mamba2_cache_params=params),
+            server_args=SimpleNamespace(),
+            spec_algorithm=SimpleNamespace(is_none=lambda: True),
+            layer_info=SimpleNamespace(start_layer=start, end_layer=end),
+            ps=SimpleNamespace(attn_dp_size=1, pp_size=pp_size),
+            hybrid_gdn_config=None,
+            model_config=SimpleNamespace(
+                hf_config=SimpleNamespace(), num_hidden_layers=cls.TOTAL_LAYERS
+            ),
+        )
+        with rc.get_context().override_server_args(
+            disable_radix_cache=False,
+            max_mamba_cache_size=None,
+            max_running_requests=None,
+            mamba_full_memory_ratio=0.5,
+            enable_linear_replayssm_spec=False,
+        ):
+            KVCacheConfigurator._handle_max_mamba_cache(fake, cls.BUDGET_GB)
+            return get_schedule().max_mamba_cache_size
+
+    def test_stage_is_not_charged_for_the_whole_model(self):
+        solo = self._pool_size(0, 1)
+        staged = self._pool_size(0, 8)
+        # The busiest stage holds 9 of the 69 mamba layers, so it should fit
+        # roughly 69/9 more slots than a rank holding all of them. pp_size=1 is
+        # unchanged: that rank does hold every layer.
+        self.assertGreater(staged, solo * 5)
+
+    def test_every_stage_agrees_on_the_pool_size(self):
+        sizes = {self._pool_size(r, 8) for r in range(8)}
+        self.assertEqual(
+            len(sizes), 1, f"per-rank pool sizes diverged: {sorted(sizes)}"
+        )
+
+
 if __name__ == "__main__":
     unittest.main()