From 2fc557254b3aaf539e80266e52a6d1e1f8da9980 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:10:43 -0700 Subject: [PATCH] fix(PP): size the mamba pool per pipeline stage, not per whole model (#33666) --- .../srt/mem_cache/kv_cache_configurator.py | 35 ++++++-- .../test_mamba_donated_alloc_ratio.py | 79 +++++++++++++++++++ 2 files changed, 108 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 221c73dcc..de9efbf20 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py index 4f3937562..9b108af7a 100644 --- a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py +++ b/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()