diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 70a77f5929..cd01fba49e 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/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 + ) diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index 1f4bf9dc37..feef0afa05 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/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, ) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 6b2c6b292b..5fece32a1b 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/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(