From d8b1c09b4c58606ac6864aa0ba0665273254aba2 Mon Sep 17 00:00:00 2001 From: logprobz <321553542+logprobz@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:02:00 +0000 Subject: [PATCH] Fix GDN graph inputs across padded speculative replay --- .../attention/test_gdn_capture_transition.py | 132 ++++++++++++++++++ vllm/v1/attention/backends/gdn_attn.py | 61 ++++++-- 2 files changed, 178 insertions(+), 15 deletions(-) create mode 100644 tests/v1/attention/test_gdn_capture_transition.py diff --git a/tests/v1/attention/test_gdn_capture_transition.py b/tests/v1/attention/test_gdn_capture_transition.py new file mode 100644 index 000000000000..2e80fdfa6b2c --- /dev/null +++ b/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 + ) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index d28a3d9e5f6b..628b985e7f8a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/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