[Bugfix][MRV2] Keep GDN prefill checkpoint metadata local to each cache group (#59536)

Signed-off-by: Jingqiao Zhang <ai-jz@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jingqiao Zhang <ai-jz@users.noreply.github.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

tests/v1/attention/test_gdn_metadata_builder.py

@@ -17,6 +17,7 @@ from tests.v1.attention.utils import (
 )
 from vllm.config import SpeculativeConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionMetadata,
     GDNAttentionMetadataBuilder,
@@ -383,14 +384,9 @@ def test_cudagraph_capture_batch_stays_decode_only():
     torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
 
 
-@pytest.mark.parametrize("num_spec", [0, 3])
-@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
-def test_checkpoint_metadata_preserves_non_spec_order(
-    num_spec, prefix_match_unit, expected_offset
-):
-    """Only eligible non-spec rows get a checkpoint in their reserved page."""
-    from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
-
+def _create_checkpoint_builder_and_batch(
+    num_spec: int, prefix_match_unit: int = 16
+) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
         block_size=BLOCK_SIZE,
@@ -431,7 +427,20 @@ def test_checkpoint_metadata_preserves_non_spec_order(
             num_decode_draft_tokens_cpu=torch.tensor([2, -1, -1], dtype=torch.int32),
             num_accepted_tokens=torch.ones(3, dtype=torch.int32),
         )
+    return builder, common, kwargs
 
+
+@pytest.mark.parametrize("num_spec", [0, 3])
+@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
+def test_checkpoint_metadata_preserves_non_spec_order(
+    num_spec, prefix_match_unit, expected_offset
+):
+    """Only eligible non-spec rows get a checkpoint in their reserved page."""
+    from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+    builder, common, kwargs = _create_checkpoint_builder_and_batch(
+        num_spec, prefix_match_unit
+    )
     actual = builder.build(0, common, **kwargs)
     assert actual.checkpoint is not None
     offsets = [expected_offset, 0] if num_spec else [0, expected_offset, 0]
@@ -444,3 +453,34 @@ def test_checkpoint_metadata_preserves_non_spec_order(
     torch.testing.assert_close(
         actual.checkpoint.state_indices, torch.tensor(slots, dtype=torch.int32)
     )
+
+
+@pytest.mark.parametrize("num_spec", [0, 3])
+def test_update_block_table_regathers_checkpoint(num_spec):
+    """Checkpoint pages come from the target group's block table, not the
+    source group's (otherwise the target group's checkpoint page is never
+    written while the prefix cache treats it as valid)."""
+    src, common, kwargs = _create_checkpoint_builder_and_batch(num_spec)
+    dst, _, _ = _create_checkpoint_builder_and_batch(num_spec)
+    ref, _, _ = _create_checkpoint_builder_and_batch(num_spec)
+    other_table = common.block_table_tensor + 100
+    other = common.replace(block_table_tensor=other_table)
+    dst.mamba_aligned_state_indices = mamba_get_block_table_tensor(
+        other_table, other.seq_lens, dst.kv_cache_spec, "align"
+    )
+
+    source = src.build(0, common, **kwargs)
+    expected = ref.build(0, other, **kwargs)
+    actual = dst.update_block_table(source, other_table, other.slot_mapping)
+
+    assert source.checkpoint is not None and expected.checkpoint is not None
+    assert actual.checkpoint is not None
+    torch.testing.assert_close(
+        actual.checkpoint.state_indices, expected.checkpoint.state_indices
+    )
+    torch.testing.assert_close(
+        actual.checkpoint.checkpoint_offsets, expected.checkpoint.checkpoint_offsets
+    )
+    assert not torch.equal(
+        actual.checkpoint.state_indices, source.checkpoint.state_indices
+    )

vllm/model_executor/layers/mamba/checkpoint.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 import torch
 
@@ -58,10 +58,33 @@ def compute_mamba_prefill_checkpoints(
     return offsets, cols
 
 
+def _gather_checkpoint_state_indices(
+    block_table: torch.Tensor, request_rows: torch.Tensor, block_cols: torch.Tensor
+) -> torch.Tensor:
+    return torch.where(
+        block_cols >= 0, block_table[request_rows, block_cols], NULL_BLOCK_ID
+    )
+
+
 @dataclass
 class MambaPrefillCheckpointMetadata:
     checkpoint_offsets: torch.Tensor
     state_indices: torch.Tensor
+    # Gather indices into the block table, kept so that another KV cache group
+    # with the same spec can re-derive its own ``state_indices``.
+    request_rows: torch.Tensor | None = None
+    block_cols: torch.Tensor | None = None
+
+    def regather_state_indices(
+        self, block_table: torch.Tensor
+    ) -> "MambaPrefillCheckpointMetadata":
+        assert self.request_rows is not None and self.block_cols is not None
+        return replace(
+            self,
+            state_indices=_gather_checkpoint_state_indices(
+                block_table, self.request_rows, self.block_cols
+            ),
+        )
 
 
 class MambaPrefillCheckpointBuilder:
@@ -111,17 +134,13 @@ class MambaPrefillCheckpointBuilder:
         checkpoint_cols_tensor = async_tensor_h2d(
             checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device
         )
-        checkpoint_state_indices = m.block_table_tensor[
-            request_rows_tensor, checkpoint_cols_tensor
-        ]
-        checkpoint_state_indices = torch.where(
-            checkpoint_cols_tensor >= 0,
-            checkpoint_state_indices,
-            NULL_BLOCK_ID,
-        )
         return MambaPrefillCheckpointMetadata(
             checkpoint_offsets_tensor,
-            checkpoint_state_indices,
+            _gather_checkpoint_state_indices(
+                m.block_table_tensor, request_rows_tensor, checkpoint_cols_tensor
+            ),
+            request_rows=request_rows_tensor,
+            block_cols=checkpoint_cols_tensor,
         )
 
 

vllm/v1/attention/backends/gdn_attn.py

@@ -621,6 +621,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         """Re-gather this group's state indices. The other fields are
         batch-level and stay shared with ``metadata``."""
         m = metadata
+        checkpoint = (
+            m.checkpoint.regather_state_indices(blk_table)
+            if m.checkpoint is not None
+            else None
+        )
         if self.vllm_config.cache_config.mamba_cache_mode == "align":
             assert self.mamba_aligned_state_indices is not None
             blk_table = self.mamba_aligned_state_indices
@@ -662,6 +667,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
             spec_state_indices_tensor=spec_indices,
             non_spec_state_indices_tensor=non_spec_indices,
             prefill_state_indices=prefill_indices,
+            checkpoint=checkpoint,
         )
 
     def build_for_cudagraph_capture(