From cd7517fb17bc26f5e63b348a73cc0dab70be8d48 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:31:53 +0400 Subject: [PATCH] [Bugfix] Fix ReplaySSM ring shape assertion when TP does not divide Mamba2 n_groups (#60704) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/kernels/mamba/test_ssu_dispatch.py | 24 +++++++++++++++++++ .../layers/mamba/mamba_utils.py | 3 ++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 3b3dc2dc6b..e12369ef8a 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -637,3 +637,27 @@ def test_replayssm_physical_ring_shape( (8, expected_ring_len), (2, expected_ring_len, 16), ) + + +def test_replayssm_ring_shape_extends_groups_for_head_shards(): + """A single group under TP is replicated to every head shard, as in + ``mamba2_state_shape``, so each rank's B_cache holds one group.""" + base_shapes = MambaStateShapeCalculator.mamba2_state_shape( + tp_world_size=2, + intermediate_size=256, + n_groups=1, + num_heads=4, + head_dim=64, + state_size=128, + conv_kernel=4, + ) + + shapes = MambaStateShapeCalculator.append_replayssm_ring( + base_shapes, + n_groups=1, + tp_world_size=2, + logical_window=16, + backend=MambaBackendEnum.TRITON, + ) + + assert shapes[4] == (1, 16, 128) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 31e89ffe55..5c423739b7 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -238,13 +238,14 @@ class MambaStateShapeCalculator: """Append the physical ReplaySSM ring shapes. ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; - B_cache uses the un-extended ``n_groups``. + B_cache extends ``n_groups`` for head shards as in ``mamba2_state_shape``. """ ring_buffer_len = logical_window if backend == MambaBackendEnum.FLASHINFER: # FlashInfer keeps the live window and current verify window together. ring_buffer_len += 1 + num_speculative_tokens local_nheads, head_dim, state_size = base_shapes[1] + n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size) local_ngroups = divide(n_groups, tp_world_size) return ( *base_shapes,