[BugFix] Keep GDN graph inputs stable across padded speculative replay

tests/v1/attention/test_gdn_capture_transition.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Keep graph-captured speculative metadata valid when requests are padded."""
+
+from types import SimpleNamespace as NS
+
+import pytest
+import torch
+
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+from vllm.v1.kv_cache_interface import MambaSpec
+
+
+def _common(query_lens, padded_tokens=32):
+    starts = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(0).to(torch.int32)
+    seq_lens = torch.tensor([64 if n else 0 for n in query_lens], dtype=torch.int32)
+    return CommonAttentionMetadata(
+        query_start_loc=starts,
+        query_start_loc_cpu=starts.clone(),
+        seq_lens=seq_lens,
+        seq_lens_cpu_upper_bound=seq_lens.clone(),
+        num_reqs=len(query_lens),
+        num_actual_tokens=padded_tokens,
+        max_query_len=max(query_lens),
+        max_seq_len=64,
+        block_table_tensor=torch.arange(len(query_lens) * 4, dtype=torch.int32).reshape(
+            len(query_lens), 4
+        ),
+        slot_mapping=torch.arange(padded_tokens, dtype=torch.int64),
+        is_prefilling=torch.zeros(len(query_lens), dtype=torch.bool),
+        causal=True,
+    )
+
+
+@pytest.mark.parametrize("fastpath", [False, True])
+@pytest.mark.parametrize("active_reqs", [1, 5, 8])
+@pytest.mark.parametrize("alias_accepted", [False, True])
+@pytest.mark.parametrize("consumer", ["build", "update_block_table"])
+def test_padded_replay_updates_captured_spec_buffers(
+    monkeypatch, fastpath, active_reqs, consumer, alias_accepted
+):
+    monkeypatch.setenv("VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH", str(int(fastpath)))
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn._resolve_gdn_prefill_backend",
+        lambda config: ("triton", "triton"),
+    )
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.gdn_attn.async_tensor_h2d",
+        lambda data, dtype=None, device="cpu", **kwargs: torch.as_tensor(
+            data, dtype=dtype, device=device
+        ),
+    )
+    config = NS(
+        compilation_config=NS(
+            cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
+            max_cudagraph_capture_size=32,
+        ),
+        speculative_config=NS(num_speculative_tokens=3, parallel_drafting=False),
+        scheduler_config=NS(max_num_seqs=8),
+        parallel_config=NS(decode_context_parallel_size=1),
+        cache_config=NS(mamba_cache_mode="align"),
+    )
+    builder = GDNAttentionMetadataBuilder(
+        MambaSpec(
+            block_size=16,
+            shapes=((16, 64),),
+            dtypes=(torch.float16,),
+            mamba_cache_mode="align",
+        ),
+        ["layer.0"],
+        config,
+        torch.device("cpu"),
+    )
+    builder.mamba_aligned_state_indices = torch.arange(32, dtype=torch.int32).reshape(
+        8, 4
+    )
+    builder.mamba_spec_accepted_tokens = torch.ones(8, dtype=torch.int32)
+    captured = builder.build_for_cudagraph_capture(_common([4] * 8))
+    other = GDNAttentionMetadataBuilder(
+        builder.kv_cache_spec, ["layer.1"], config, torch.device("cpu")
+    )
+    other.mamba_aligned_state_indices = builder.mamba_aligned_state_indices.clone() + 64
+    other.mamba_spec_accepted_tokens = builder.mamba_spec_accepted_tokens
+    captured_other = other.build_for_cudagraph_capture(_common([4] * 8))
+    fields = [
+        "spec_state_indices_tensor",
+        "spec_query_start_loc",
+        "spec_sequence_masks",
+        "num_accepted_tokens",
+    ]
+    captured = captured_other if consumer == "update_block_table" else captured
+    owner = other if consumer == "update_block_table" else builder
+    pointers = {name: getattr(captured, name).data_ptr() for name in fields}
+    for count in (active_reqs, 8, active_reqs):
+        for current in (builder, other):
+            current.mamba_aligned_state_indices.copy_(
+                torch.arange(32, dtype=torch.int32).reshape(8, 4) + 100
+            )
+            current.mamba_aligned_state_indices[count:].fill_(NULL_BLOCK_ID)
+        expected_accepted = torch.tensor([1, 2, 1, 3, 2, 1, 1, 1], dtype=torch.int32)
+        expected_accepted[count:] = 1
+        accepted = expected_accepted.clone()
+        if alias_accepted:
+            builder.mamba_spec_accepted_tokens.copy_(accepted)
+            accepted = builder.mamba_spec_accepted_tokens
+        runtime_common = _common([4] * count + [0] * (8 - count))
+        replay = builder.build(
+            0,
+            runtime_common,
+            accepted,
+            torch.tensor([3] * count + [-1] * (8 - count), dtype=torch.int32),
+        )
+        if consumer == "update_block_table":
+            replay = other.update_block_table(
+                replay, runtime_common.block_table_tensor, None
+            )
+        for name in fields:
+            assert getattr(replay, name).data_ptr() == pointers[name], name
+            torch.testing.assert_close(getattr(captured, name), getattr(replay, name))
+        torch.testing.assert_close(captured.num_accepted_tokens, expected_accepted)
+        torch.testing.assert_close(
+            captured.spec_query_start_loc, runtime_common.query_start_loc
+        )
+        torch.testing.assert_close(
+            captured.spec_sequence_masks, torch.arange(8) < count
+        )
+        torch.testing.assert_close(
+            captured.spec_state_indices_tensor, owner.mamba_aligned_state_indices
+        )

vllm/v1/attention/backends/gdn_attn.py

@@ -198,9 +198,6 @@ def __init__(
         self._decode_state_indices_source: torch.Tensor | None = None
         self._decode_state_indices_view: torch.Tensor | None = None
         self._reuse_spec_decode_inputs = envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH
-        self._uniform_spec_masks = torch.ones(
-            self.decode_cudagraph_max_bs, dtype=torch.bool, device=device
-        )
         self._uniform_spec_masks_cpu = torch.ones(
             self.decode_cudagraph_max_bs, dtype=torch.bool
         )
@@ -247,12 +244,35 @@ def _get_spec_state_indices_view(self, num_reqs: int) -> torch.Tensor:
             self._spec_state_indices_view = source[:num_reqs, : self.num_spec + 1]
         return self._spec_state_indices_view
 
+    def _graph_accepted_tokens(self) -> torch.Tensor:
+        # Uniform and padded replay must update the tensor captured by the graph.
+        if (
+            self._reuse_spec_decode_inputs
+            and self.mamba_spec_accepted_tokens is not None
+        ):
+            return self.mamba_spec_accepted_tokens
+        return self.num_accepted_tokens
+
+    def _store_uniform_spec_inputs(self, num_reqs: int, num_tokens: int) -> None:
+        # Use the fallback destinations so padding cannot leave captured inputs stale.
+        self.spec_state_indices_tensor[:num_reqs].copy_(
+            self._get_spec_state_indices_view(num_reqs), non_blocking=True
+        )
+        self.spec_query_start_loc[: num_reqs + 1].copy_(
+            self._uniform_spec_query_start[: num_reqs + 1], non_blocking=True
+        )
+        self.spec_sequence_masks[:num_reqs].fill_(True)
+        self.spec_token_indx[:num_tokens].copy_(
+            self._uniform_spec_tokens[:num_tokens], non_blocking=True
+        )
+
     def _build_uniform_spec_decode(
         self, m: CommonAttentionMetadata, num_accepted_tokens: torch.Tensor
     ) -> GDNAttentionMetadata:
         num_reqs = m.num_reqs
         assert self.mamba_spec_accepted_tokens is not None
-        accepted = self.mamba_spec_accepted_tokens[:num_reqs]
+        self._store_uniform_spec_inputs(num_reqs, m.num_actual_tokens)
+        accepted = self._graph_accepted_tokens()[:num_reqs]
         accepted.copy_(num_accepted_tokens[:num_reqs], non_blocking=True)
         return GDNAttentionMetadata(
             num_prefills=0,
@@ -262,12 +282,12 @@ def _build_uniform_spec_decode(
             num_spec_decodes=num_reqs,
             num_spec_decode_tokens=m.num_actual_tokens,
             num_actual_tokens=m.num_actual_tokens,
-            spec_query_start_loc=self._uniform_spec_query_start[: num_reqs + 1],
-            spec_state_indices_tensor=self._get_spec_state_indices_view(num_reqs),
-            spec_sequence_masks=self._uniform_spec_masks[:num_reqs],
+            spec_query_start_loc=self.spec_query_start_loc[: num_reqs + 1],
+            spec_state_indices_tensor=self.spec_state_indices_tensor[:num_reqs],
+            spec_sequence_masks=self.spec_sequence_masks[:num_reqs],
             spec_sequence_masks_cpu=self._uniform_spec_masks_cpu[:num_reqs],
-            spec_token_indx=self._uniform_spec_tokens[: m.num_actual_tokens],
-            non_spec_token_indx=self._uniform_spec_tokens[:0],
+            spec_token_indx=self.spec_token_indx[: m.num_actual_tokens],
+            non_spec_token_indx=self.non_spec_token_indx[:0],
             num_accepted_tokens=accepted,
             num_reqs=num_reqs,
             seq_lens=m.seq_lens,
@@ -683,10 +703,11 @@ def build(  # type: ignore[override]
             spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
             spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens)
 
-            self.num_accepted_tokens[:num_spec_decodes].copy_(
+            accepted_buffer = self._graph_accepted_tokens()
+            accepted_buffer[:num_spec_decodes].copy_(
                 num_accepted_tokens, non_blocking=True
             )
-            num_accepted_tokens = self.num_accepted_tokens[:batch_size]
+            num_accepted_tokens = accepted_buffer[:batch_size]
             num_accepted_tokens[num_spec_decodes:].fill_(1)
 
         if (
@@ -760,9 +781,18 @@ def update_block_table(
             and self.mamba_spec_accepted_tokens is not None
         ):
             updated = copy(metadata)
-            updated.spec_state_indices_tensor = self._get_spec_state_indices_view(
-                metadata.num_reqs
+            self._store_uniform_spec_inputs(
+                metadata.num_reqs, metadata.num_actual_tokens
             )
+            updated.spec_state_indices_tensor = self.spec_state_indices_tensor[
+                : metadata.num_reqs
+            ]
+            updated.spec_query_start_loc = self.spec_query_start_loc[
+                : metadata.num_reqs + 1
+            ]
+            updated.spec_sequence_masks = self.spec_sequence_masks[: metadata.num_reqs]
+            updated.spec_token_indx = self.spec_token_indx[: metadata.num_actual_tokens]
+            updated.non_spec_token_indx = self.non_spec_token_indx[:0]
             accepted = self.mamba_spec_accepted_tokens[: metadata.num_reqs]
             assert metadata.num_accepted_tokens is not None
             if accepted.data_ptr() != metadata.num_accepted_tokens.data_ptr():
@@ -874,10 +904,11 @@ def update_block_table(
             )
             spec_query_start_loc = self.spec_query_start_loc[: metadata.num_reqs + 1]
 
-            self.num_accepted_tokens[: metadata.num_reqs].copy_(
+            accepted_buffer = self._graph_accepted_tokens()
+            accepted_buffer[: metadata.num_reqs].copy_(
                 num_accepted_tokens[: metadata.num_reqs], non_blocking=True
             )
-            num_accepted_tokens = self.num_accepted_tokens[: metadata.num_reqs]
+            num_accepted_tokens = accepted_buffer[: metadata.num_reqs]
 
         if (
             self.use_full_cuda_graph