[Bugfix] Fix ReplaySSM ring shape assertion when TP does not divide Mamba2 n_groups (#60704)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

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)

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,