[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>
ai-jz committed
Oct 1, 2026 at 07:59 UTC
b538d807ffe0976af7e5d107ccfe77e676b83b70
3 files changed
+83
-18
tests/v1/attention/test_gdn_metadata_builder.py
+48
-8
@@ -17,6 +17,7 @@ from tests.v1.attention.utils import (
17
)
18
from vllm.config import SpeculativeConfig
19
from vllm.config.compilation import CUDAGraphMode
20
+from vllm.v1.attention.backend import CommonAttentionMetadata
21
from vllm.v1.attention.backends.gdn_attn import (
22
GDNAttentionMetadata,
23
GDNAttentionMetadataBuilder,
@@ -383,14 +384,9 @@ def test_cudagraph_capture_batch_stays_decode_only():
384
torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
385
386
386
-@pytest.mark.parametrize("num_spec", [0, 3])
387
-@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
388
-def test_checkpoint_metadata_preserves_non_spec_order(
389
- num_spec, prefix_match_unit, expected_offset
390
-):
391
- """Only eligible non-spec rows get a checkpoint in their reserved page."""
392
- from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
393
-
387
+def _create_checkpoint_builder_and_batch(
388
+ num_spec: int, prefix_match_unit: int = 16
389
+) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]:
390
vllm_config = create_vllm_config(
391
model_name="Qwen/Qwen3.5-0.8B",
392
block_size=BLOCK_SIZE,
@@ -431,7 +427,20 @@ def test_checkpoint_metadata_preserves_non_spec_order(
427
num_decode_draft_tokens_cpu=torch.tensor([2, -1, -1], dtype=torch.int32),
428
num_accepted_tokens=torch.ones(3, dtype=torch.int32),
429
)
430
+ return builder, common, kwargs
431
+
432
+
433
+@pytest.mark.parametrize("num_spec", [0, 3])
434
+@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
435
+def test_checkpoint_metadata_preserves_non_spec_order(
436
+ num_spec, prefix_match_unit, expected_offset
437
+):
438
+ """Only eligible non-spec rows get a checkpoint in their reserved page."""
439
+ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
440
441
+ builder, common, kwargs = _create_checkpoint_builder_and_batch(
442
+ num_spec, prefix_match_unit
443
+ )
444
actual = builder.build(0, common, **kwargs)
445
assert actual.checkpoint is not None
446
offsets = [expected_offset, 0] if num_spec else [0, expected_offset, 0]
@@ -444,3 +453,34 @@ def test_checkpoint_metadata_preserves_non_spec_order(
453
torch.testing.assert_close(
454
actual.checkpoint.state_indices, torch.tensor(slots, dtype=torch.int32)
455
)
456
+
457
+
458
+@pytest.mark.parametrize("num_spec", [0, 3])
459
+def test_update_block_table_regathers_checkpoint(num_spec):
460
+ """Checkpoint pages come from the target group's block table, not the
461
+ source group's (otherwise the target group's checkpoint page is never
462
+ written while the prefix cache treats it as valid)."""
463
+ src, common, kwargs = _create_checkpoint_builder_and_batch(num_spec)
464
+ dst, _, _ = _create_checkpoint_builder_and_batch(num_spec)
465
+ ref, _, _ = _create_checkpoint_builder_and_batch(num_spec)
466
+ other_table = common.block_table_tensor + 100
467
+ other = common.replace(block_table_tensor=other_table)
468
+ dst.mamba_aligned_state_indices = mamba_get_block_table_tensor(
469
+ other_table, other.seq_lens, dst.kv_cache_spec, "align"
470
+ )
471
+
472
+ source = src.build(0, common, **kwargs)
473
+ expected = ref.build(0, other, **kwargs)
474
+ actual = dst.update_block_table(source, other_table, other.slot_mapping)
475
+
476
+ assert source.checkpoint is not None and expected.checkpoint is not None
477
+ assert actual.checkpoint is not None
478
+ torch.testing.assert_close(
479
+ actual.checkpoint.state_indices, expected.checkpoint.state_indices
480
+ )
481
+ torch.testing.assert_close(
482
+ actual.checkpoint.checkpoint_offsets, expected.checkpoint.checkpoint_offsets
483
+ )
484
+ assert not torch.equal(
485
+ actual.checkpoint.state_indices, source.checkpoint.state_indices
486
+ )
vllm/model_executor/layers/mamba/checkpoint.py
+29
-10
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
-from dataclasses import dataclass
5
+from dataclasses import dataclass, replace
6
7
import torch
8
@@ -58,10 +58,33 @@ def compute_mamba_prefill_checkpoints(
58
return offsets, cols
59
60
61
+def _gather_checkpoint_state_indices(
62
+ block_table: torch.Tensor, request_rows: torch.Tensor, block_cols: torch.Tensor
63
+) -> torch.Tensor:
64
+ return torch.where(
65
+ block_cols >= 0, block_table[request_rows, block_cols], NULL_BLOCK_ID
66
+ )
67
+
68
+
69
@dataclass
70
class MambaPrefillCheckpointMetadata:
71
checkpoint_offsets: torch.Tensor
72
state_indices: torch.Tensor
73
+ # Gather indices into the block table, kept so that another KV cache group
74
+ # with the same spec can re-derive its own ``state_indices``.
75
+ request_rows: torch.Tensor | None = None
76
+ block_cols: torch.Tensor | None = None
77
+
78
+ def regather_state_indices(
79
+ self, block_table: torch.Tensor
80
+ ) -> "MambaPrefillCheckpointMetadata":
81
+ assert self.request_rows is not None and self.block_cols is not None
82
+ return replace(
83
+ self,
84
+ state_indices=_gather_checkpoint_state_indices(
85
+ block_table, self.request_rows, self.block_cols
86
+ ),
87
+ )
88
89
90
class MambaPrefillCheckpointBuilder:
@@ -111,17 +134,13 @@ class MambaPrefillCheckpointBuilder:
134
checkpoint_cols_tensor = async_tensor_h2d(
135
checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device
136
)
114
- checkpoint_state_indices = m.block_table_tensor[
115
- request_rows_tensor, checkpoint_cols_tensor
116
- ]
117
- checkpoint_state_indices = torch.where(
118
- checkpoint_cols_tensor >= 0,
119
- checkpoint_state_indices,
120
- NULL_BLOCK_ID,
121
- )
137
return MambaPrefillCheckpointMetadata(
138
checkpoint_offsets_tensor,
124
- checkpoint_state_indices,
139
+ _gather_checkpoint_state_indices(
140
+ m.block_table_tensor, request_rows_tensor, checkpoint_cols_tensor
141
+ ),
142
+ request_rows=request_rows_tensor,
143
+ block_cols=checkpoint_cols_tensor,
144
)
145
146
vllm/v1/attention/backends/gdn_attn.py
+6
@@ -621,6 +621,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
621
"""Re-gather this group's state indices. The other fields are
622
batch-level and stay shared with ``metadata``."""
623
m = metadata
624
+ checkpoint = (
625
+ m.checkpoint.regather_state_indices(blk_table)
626
+ if m.checkpoint is not None
627
+ else None
628
+ )
629
if self.vllm_config.cache_config.mamba_cache_mode == "align":
630
assert self.mamba_aligned_state_indices is not None
631
blk_table = self.mamba_aligned_state_indices
@@ -662,6 +667,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
667
spec_state_indices_tensor=spec_indices,
668
non_spec_state_indices_tensor=non_spec_indices,
669
prefill_state_indices=prefill_indices,
670
+ checkpoint=checkpoint,
671
)
672
673
def build_for_cudagraph_capture(