[Bugfix][RecoverSSM] Correct aligned state indices at exact block boundaries (#59962)

Signed-off-by: apex-mochen <2756823972@qq.com>
Co-authored-by: Codex <noreply@openai.com>

tests/models/kimi_k3/test_kda.py

@@ -692,8 +692,10 @@ def test_kda_spec_decode_correctness(
         pytest.param(True, False, None, True, id="aligned"),
     ],
 )
+@pytest.mark.parametrize("accepted", [(2, 8), (3, 4), (5, 8)])
 @torch.inference_mode()
 def test_kda_recoverssm_verify_and_group_commit(
+    accepted: tuple[int, int],
     monkeypatch: pytest.MonkeyPatch,
     lower_bound: float | None,
     use_request_indices: bool,
@@ -738,7 +740,6 @@ def test_kda_recoverssm_verify_and_group_commit(
     state_indices = torch.tensor(
         [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE
     )
-    accepted = [2, 8]
     if use_request_indices:
         global_num_accepted = torch.tensor(
             [0, accepted[0], 0, accepted[1]],
@@ -891,7 +892,9 @@ def test_kda_recoverssm_verify_and_group_commit(
             if align_mode:
                 assert block_table is not None
                 row = request_indices[seq_idx] if use_request_indices else seq_idx
-                final_block = block_table[row, (4 + commit_len) // 8]
+                final_block = block_table[
+                    row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len]
+                ]
             committed_states[final_block] = committed_state.transpose(-1, -2)
             if align_mode and 4 + commit_len >= 8:
                 _, boundary_state = naive_recurrent_kda(
@@ -944,7 +947,7 @@ def test_kda_recoverssm_verify_and_group_commit(
             if align_mode:
                 assert block_table is not None
                 row = request_indices[seq_idx] if use_request_indices else seq_idx
-                block = block_table[row, (4 + commit_len) // 8]
+                block = block_table[row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len]]
             source_block = state_indices[seq_idx] if align_mode else block
             if conv_state_dim_first:
                 actual_conv = layer.kv_cache[0][block, :, :history_len]

tests/v1/worker/test_mamba_hybrid_model_state.py

@@ -387,7 +387,13 @@ def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None:
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
-def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None:
+@pytest.mark.parametrize(
+    ("computed", "sampled", "expected_column"),
+    [(4, 3, 0), (4, 4, 0), (4, 5, 1), (8, 8, 1), (8, 9, 2)],
+)
+def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias(
+    computed: int, sampled: int, expected_column: int
+) -> None:
     state = object.__new__(MambaHybridModelState)
     state._align_mode = True
     state._mamba_ctx = None
@@ -400,11 +406,13 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -
     metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata(
         num_spec_decodes=1,
         request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"),
-        num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"),
+        num_computed_tokens=torch.tensor(
+            [6, computed], dtype=torch.int32, device="cuda"
+        ),
         block_size=8,
         block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"),
     )
-    num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda")
+    num_sampled = torch.tensor([2, sampled], dtype=torch.int32, device="cuda")
     idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda")
     group = SimpleNamespace(layer_names=["layer"])
 
@@ -412,7 +420,7 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -
 
     state.postprocess_state(idx_mapping, num_sampled)
 
-    expected_state_indices = [-1, 1, -1, -1, -1]
+    expected_state_indices = [-1, expected_column, -1, -1, -1]
     assert state._mamba_state_idx_gpu.tolist() == expected_state_indices
     expected_accepted = [9, 1, 9, 2, 9]
     assert state.num_accepted_tokens_gpu.tolist() == expected_accepted

vllm/models/kimi_k3/nvidia/ops/recoverssm.py

@@ -284,7 +284,8 @@ def _prepare_commit_plan_kernel(
         )
         final_num_computed = num_computed + commit_len
         final_state_col = tl.minimum(
-            final_num_computed // mamba_block_size, block_table_width - 1
+            tl.maximum(final_num_computed - 1, 0) // mamba_block_size,
+            block_table_width - 1,
         )
         final_state_idx = tl.load(
             block_table_ptr

vllm/v1/worker/gpu/model_states/recoverssm.py

@@ -94,7 +94,7 @@ def _postprocess_recoverssm_align_kernel(
     tl.store(
         state_idx_ptr + req_state_idx,
         tl.minimum(
-            (num_computed + num_sampled) // MAMBA_BLOCK_SIZE,
+            tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE,
             BLOCK_TABLE_WIDTH - 1,
         ),
     )