From f583a359489f017de60fcd63eb826b7e741f8258 Mon Sep 17 00:00:00 2001 From: kimseunghyun Date: Sun, 11 Oct 2026 05:46:16 +0900 Subject: [PATCH 1/4] [Perf][GDN] Gather and scatter prefill SSM states without Python indexing Before the chunk kernel, each GDN layer gathered the prefill states with ssm_state[idx], zeroed the rows without initial state with a boolean __setitem__, and scattered the final states back with ssm_state[idx] = x. The two Python advanced-indexing calls and the per-layer mask inversion are host work in every layer of a mixed step. Use index_select, masked_fill_ and index_copy_ instead. The metadata builder now computes the inverted mask once per step, shaped to broadcast over the [num_prefills, HV, V, K] state. The values are bit-identical: the state indices are non-negative and unique per prefill. Test: tests/kernels/mamba/test_gdn_forward_core_mixed.py compares the state that the chunk kernel receives and the pool after the step with the indexing reference, bit for bit (CUDA). Co-authored-by: Claude Opus 5.5 Signed-off-by: kimseunghyun --- .../mamba/test_gdn_forward_core_mixed.py | 241 ++++++++++++++++++ .../mamba/test_gdn_forward_core_split.py | 1 + .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 12 +- vllm/v1/attention/backends/gdn_attn.py | 6 + 4 files changed, 255 insertions(+), 5 deletions(-) create mode 100644 tests/kernels/mamba/test_gdn_forward_core_mixed.py diff --git a/tests/kernels/mamba/test_gdn_forward_core_mixed.py b/tests/kernels/mamba/test_gdn_forward_core_mixed.py new file mode 100644 index 000000000000..0fe7ae15e301 --- /dev/null +++ b/tests/kernels/mamba/test_gdn_forward_core_mixed.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Parity tests for the host-side tensor plumbing of the GDN mixed step. + +``GatedDeltaNet._forward_core`` gathers, zeroes and scatters the prefill SSM +states around ``chunk_gated_delta_rule``. These tests run the real +``_forward_core`` on batches built by the real ``GDNAttentionMetadataBuilder`` +and require bit-identical outputs and states against the reference op sequence. +""" + +from __future__ import annotations + +import types +from unittest.mock import patch + +import pytest +import torch + +from vllm.platforms import current_platform + +if not (current_platform.is_cuda() and current_platform.has_device_capability(80)): + pytest.skip( + reason="GDN _forward_core runs Triton kernels that need CUDA 8.0+.", + allow_module_level=True, + ) + +from tests.v1.attention.utils import ( # noqa: E402 + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig, set_current_vllm_config # noqa: E402 +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402 + ChunkGatedDeltaRule, + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402 + MambaStateShapeCalculator, +) +from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 + +NUM_SPEC = 3 +SPEC_TOKENS = NUM_SPEC + 1 +H = 2 +HV = 4 +K = 128 +V = 128 +CONV_KERNEL = 4 +CONV_DIM = 2 * H * K + HV * V +Z_DIM = HV * V +BLOCK_SIZE = 16 +PREFIX = "model.layers.0.linear_attn" + +# (seq_lens, query_lens, draft_tokens); -1 marks a request without drafts. +BATCHES = { + "prefills": ([96, 300, 70], [96, 64, 5], [-1, -1, -1]), + "spec-decodes-then-prefills": ( + [128, 90, 64, 200], + [SPEC_TOKENS, 2, 64, 37], + [NUM_SPEC, 1, -1, -1], + ), +} + + +def _make_vllm_config(): + config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + hf_config_override={"linear_key_head_dim": K}, + ) + config.cache_config.mamba_cache_mode = "none" + config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=NUM_SPEC + ) + return config + + +def _build_metadata(vllm_config, batch: BatchSpec, draft_tokens: list[int]): + device = torch.device("cuda") + builder = GDNAttentionMetadataBuilder( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=NUM_SPEC, + ), + layer_names=[PREFIX], + vllm_config=vllm_config, + device=device, + ) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ) + # Keep the real states off NULL_BLOCK_ID. + common.block_table_tensor.add_(1) + with set_current_vllm_config(vllm_config): + return builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.ones( + batch.batch_size, dtype=torch.int32, device=device + ), + num_decode_draft_tokens_cpu=torch.tensor(draft_tokens, dtype=torch.int32), + ) + + +def _build_layer(vllm_config, kv_cache, weights): + layer = types.SimpleNamespace( + prefix=PREFIX, + enable_packed_recurrent_decode=False, + tp_size=1, + num_k_heads=H, + num_v_heads=HV, + head_k_dim=K, + head_v_dim=V, + key_dim=H * K, + value_dim=HV * V, + activation="silu", + A_log=weights["A_log"], + dt_bias=weights["dt_bias"], + conv1d=types.SimpleNamespace(weight=weights["conv_weight"], bias=None), + kv_cache=kv_cache, + ) + with set_current_vllm_config(vllm_config): + layer.chunk_gated_delta_rule = ChunkGatedDeltaRule() + for name in ("rearrange_mixed_qkv", "_forward_core"): + setattr( + layer, + name, + types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer), + ) + return layer + + +def _make_inputs(metadata: GDNAttentionMetadata, num_tokens: int): + """Random states, weights and a strided mixed_qkv, as split from qkvz.""" + device = torch.device("cuda") + state_indices = [ + indices + for indices in ( + metadata.spec_state_indices_tensor, + metadata.non_spec_state_indices_tensor, + ) + if indices is not None and indices.numel() > 0 + ] + pool_size = max(int(indices.max().item()) for indices in state_indices) + 1 + conv_state_shape, ssm_state_shape = ( + MambaStateShapeCalculator.gated_delta_net_state_shape( + 1, H, HV, K, V, CONV_KERNEL, NUM_SPEC + ) + ) + conv_state = 0.05 * torch.randn( + pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device + ) + ssm_state = 0.01 * torch.randn( + pool_size, *ssm_state_shape, dtype=torch.float32, device=device + ) + weights = { + "A_log": 0.1 * torch.randn(HV, dtype=torch.float32, device=device), + "dt_bias": 0.1 * torch.randn(HV, dtype=torch.float32, device=device), + "conv_weight": 0.1 + * torch.randn(CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device), + } + mixed_qkvz = 0.1 * torch.randn( + num_tokens, CONV_DIM + Z_DIM, dtype=torch.bfloat16, device=device + ) + b = 0.1 * torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) + a = 0.1 * torch.randn_like(b) + return (conv_state, ssm_state), weights, mixed_qkvz, b, a + + +def _run_forward_core(layer, metadata, mixed_qkvz, b, a, num_tokens): + core_attn_out = torch.zeros( + num_tokens, HV, V, dtype=mixed_qkvz.dtype, device=mixed_qkvz.device + ) + context = types.SimpleNamespace(attn_metadata={PREFIX: metadata}) + with patch.object( + qwen_gdn_linear_attn, "get_forward_context", return_value=context + ): + layer._forward_core( + mixed_qkv=mixed_qkvz.clone()[:, :CONV_DIM], + b=b.clone(), + a=a.clone(), + core_attn_out=core_attn_out, + ) + return core_attn_out + + +@pytest.mark.parametrize("batch_name", BATCHES) +@torch.inference_mode() +def test_prefill_state_gather_and_scatter_match_indexing(batch_name: str) -> None: + """The chunk kernel sees ``ssm_state[idx]`` with fresh rows zeroed, and + ``ssm_state[idx] = final_state`` is the only change it makes to the pool.""" + torch.manual_seed(0) + seq_lens, query_lens, draft_tokens = BATCHES[batch_name] + vllm_config = _make_vllm_config() + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + metadata = _build_metadata(vllm_config, batch, draft_tokens) + assert metadata.num_prefills > 0 + prefill_state_indices = metadata.prefill_state_indices + prefill_has_initial_state = metadata.prefill_has_initial_state + assert prefill_state_indices is not None + assert prefill_has_initial_state is not None + assert not bool(prefill_has_initial_state.all()) + + num_tokens = batch.compute_num_tokens() + kv_cache, weights, mixed_qkvz, b, a = _make_inputs(metadata, num_tokens) + ssm_state = kv_cache[1] + # A fresh prefill must not read the stale state of its slot. + fresh_slots = prefill_state_indices[~prefill_has_initial_state].long() + ssm_state[fresh_slots] = float("nan") + layer = _build_layer(vllm_config, kv_cache, weights) + + chunk_calls = [] + chunk_rule = layer.chunk_gated_delta_rule + + def recording_chunk_rule(**kwargs): + pool_before = ssm_state.clone() + initial_state = kwargs["initial_state"].clone() + output, final_state = chunk_rule(**kwargs) + chunk_calls.append((pool_before, initial_state, final_state)) + return output, final_state + + layer.chunk_gated_delta_rule = recording_chunk_rule + _run_forward_core(layer, metadata, mixed_qkvz, b, a, num_tokens) + + assert len(chunk_calls) == 1 + pool_before, initial_state, final_state = chunk_calls[0] + expected_initial_state = pool_before[prefill_state_indices] + expected_initial_state[~prefill_has_initial_state, ...] = 0 + torch.testing.assert_close( + initial_state, expected_initial_state, atol=0, rtol=0, equal_nan=True + ) + expected_pool = pool_before.clone() + expected_pool[prefill_state_indices] = final_state.to(expected_pool.dtype) + torch.testing.assert_close(ssm_state, expected_pool, atol=0, rtol=0) diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py index 8a99e223be3b..4b87e91ebf13 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_split.py +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -227,6 +227,7 @@ def test_forward_core_split_matches_unified( prefill_query_start_loc=meta_split.non_spec_query_start_loc, prefill_state_indices=meta_split.non_spec_state_indices_tensor, prefill_has_initial_state=meta_split.has_initial_state, + prefill_no_initial_state_mask=~meta_split.has_initial_state.view(-1, 1, 1, 1), ) # Size the state pools from the indices the builder actually produced. diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1a3f8183c4f1..ac9c2a881ab7 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1737,11 +1737,11 @@ def _forward_core( # when decodes are peeled off, else the full non-spec batch), so they # don't need to be re-derived per layer. prefill_state_indices = attn_metadata.prefill_state_indices - prefill_has_initial_state = attn_metadata.prefill_has_initial_state + no_initial_state_mask = attn_metadata.prefill_no_initial_state_mask assert prefill_state_indices is not None - assert prefill_has_initial_state is not None - initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 + assert no_initial_state_mask is not None + initial_state = ssm_state.index_select(0, prefill_state_indices) + initial_state.masked_fill_(no_initial_state_mask, 0) plan = attn_metadata.flashinfer_checkpoint_plan if plan is None: ( @@ -1796,7 +1796,9 @@ def _forward_core( recurrent_rows=plan.rows, ) # Init cache - ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) + ssm_state.index_copy_( + 0, prefill_state_indices, last_recurrent_state.to(ssm_state.dtype) + ) if split_non_spec: # Stitch the peeled decode outputs in front of the prefill diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 16a2c39176b2..38b8a5fa76d5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -123,6 +123,9 @@ class GDNAttentionMetadata: prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + # ~prefill_has_initial_state as [num_prefills, 1, 1, 1], to zero the + # gathered [num_prefills, HV, V, K] initial states in one masked_fill_. + prefill_no_initial_state_mask: torch.Tensor | None = None aiter_prefill_metadata: object | None = None # The following attributes are for triton implementation of causal_conv1d @@ -513,6 +516,7 @@ def build( # type: ignore[override] prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + prefill_no_initial_state_mask: torch.Tensor | None = None aiter_prefill_metadata: object | None = None if num_prefills > 0: # In a mixed non-spec batch, decodes are peeled off to the recurrent @@ -569,6 +573,7 @@ def build( # type: ignore[override] prefill_has_initial_state = has_initial_state[num_decodes:] else: prefill_has_initial_state = has_initial_state + prefill_no_initial_state_mask = ~prefill_has_initial_state.view(-1, 1, 1, 1) else: has_initial_state = None @@ -683,6 +688,7 @@ def build( # type: ignore[override] prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, prefill_has_initial_state=prefill_has_initial_state, + prefill_no_initial_state_mask=prefill_no_initial_state_mask, aiter_prefill_metadata=aiter_prefill_metadata, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, From 34631f90ffe26a17bc9b610c80337fd71f37b373 Mon Sep 17 00:00:00 2001 From: kimseunghyun Date: Sun, 11 Oct 2026 05:50:37 +0900 Subject: [PATCH 2/4] [Perf][GDN] Slice instead of gather when spec-decode rows lead the batch In a mixed step with spec decodes, each GDN layer split the tokens with four index_select calls (mixed_qkv, a, b for spec; mixed_qkv for non-spec), two more for the non-spec a and b, and merged the outputs with two index_copy_ calls. The batch reorder puts decode rows first, so the spec-decode rows usually lead the batch. The metadata builder now records this on the CPU (spec_tokens_first). Then the spec tokens are [0, num_spec_decode_tokens) and the non-spec tokens follow, so the layer uses slices and two copy_ calls. Other batches keep the gather path. The slices have the same strides as the non-spec path, which passes mixed_qkv to the same kernels directly. The values are bit-identical. Tests: - tests/v1/attention/test_gdn_metadata_builder.py: the flag is True exactly when the gather indices are the two ranges (CPU). - tests/kernels/mamba/test_gdn_forward_core_mixed.py: the slice path and the gather path give bit-identical outputs, conv states and SSM states (CUDA). Co-authored-by: Claude Opus 5.5 Signed-off-by: kimseunghyun --- .../mamba/test_gdn_forward_core_mixed.py | 33 +++++++++++++++-- .../v1/attention/test_gdn_metadata_builder.py | 36 +++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 28 +++++++++++++-- vllm/v1/attention/backends/gdn_attn.py | 9 +++++ 4 files changed, 101 insertions(+), 5 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_forward_core_mixed.py b/tests/kernels/mamba/test_gdn_forward_core_mixed.py index 0fe7ae15e301..8317f0fe690a 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_mixed.py +++ b/tests/kernels/mamba/test_gdn_forward_core_mixed.py @@ -3,13 +3,15 @@ """Parity tests for the host-side tensor plumbing of the GDN mixed step. ``GatedDeltaNet._forward_core`` gathers, zeroes and scatters the prefill SSM -states around ``chunk_gated_delta_rule``. These tests run the real -``_forward_core`` on batches built by the real ``GDNAttentionMetadataBuilder`` -and require bit-identical outputs and states against the reference op sequence. +states around ``chunk_gated_delta_rule``, and splits a batch with spec decodes +into spec and non-spec tokens. These tests run the real ``_forward_core`` on +batches built by the real ``GDNAttentionMetadataBuilder`` and require +bit-identical outputs and states against the reference op sequence. """ from __future__ import annotations +import dataclasses import types from unittest.mock import patch @@ -239,3 +241,28 @@ def recording_chunk_rule(**kwargs): expected_pool = pool_before.clone() expected_pool[prefill_state_indices] = final_state.to(expected_pool.dtype) torch.testing.assert_close(ssm_state, expected_pool, atol=0, rtol=0) + + +@torch.inference_mode() +def test_spec_token_slices_match_gathers() -> None: + """Spec-first batches slice the spec and non-spec tokens; the result is + bit-identical to the index_select / index_copy_ path.""" + torch.manual_seed(0) + seq_lens, query_lens, draft_tokens = BATCHES["spec-decodes-then-prefills"] + vllm_config = _make_vllm_config() + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + metadata = _build_metadata(vllm_config, batch, draft_tokens) + assert metadata.spec_tokens_first + gather_metadata = dataclasses.replace(metadata, spec_tokens_first=False) + + num_tokens = batch.compute_num_tokens() + kv_cache, weights, mixed_qkvz, b, a = _make_inputs(metadata, num_tokens) + results = [] + for step_metadata in (metadata, gather_metadata): + step_kv_cache = tuple(state.clone() for state in kv_cache) + layer = _build_layer(vllm_config, step_kv_cache, weights) + out = _run_forward_core(layer, step_metadata, mixed_qkvz, b, a, num_tokens) + results.append((out, *step_kv_cache)) + + for sliced, gathered in zip(*results): + torch.testing.assert_close(sliced, gathered, atol=0, rtol=0) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 937bf2f30d7d..0cbba18aaa7d 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -277,6 +277,42 @@ def test_has_initial_state_after_reclassification(): assert meta.has_initial_state[0].item() is True +@pytest.mark.parametrize( + "seq_lens,query_lens,draft_tokens,expected", + [ + pytest.param([20, 30, 100], [3, 2, 50], [2, 1, -1], True, id="spec-first"), + pytest.param( + [20, 65, 100], [3, 1, 50], [2, -1, -1], True, id="spec-first-reclassified" + ), + pytest.param([100, 20], [50, 3], [-1, 2], False, id="prefill-first"), + pytest.param( + [20, 100, 30], [3, 50, 3], [2, -1, 2], False, id="spec-interleaved" + ), + ], +) +def test_spec_tokens_first_matches_token_indices( + seq_lens: list[int], + query_lens: list[int], + draft_tokens: list[int], + expected: bool, +): + """spec_tokens_first holds exactly when the gather indices are two ranges.""" + builder = _create_gdn_builder(num_speculative_tokens=2) + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + meta = _build(builder, batch, num_decode_draft_tokens=draft_tokens) + + assert meta.num_prefills > 0 and meta.num_spec_decodes > 0 + assert meta.spec_token_indx is not None + assert meta.non_spec_token_indx is not None + num_spec_tokens = meta.num_spec_decode_tokens + num_tokens = batch.compute_num_tokens() + is_two_ranges = meta.spec_token_indx.tolist() == list( + range(num_spec_tokens) + ) and meta.non_spec_token_indx.tolist() == list(range(num_spec_tokens, num_tokens)) + assert meta.spec_tokens_first is expected + assert is_two_ranges is expected + + def test_full_cudagraph_spec_metadata_uses_request_count(): """FULL cudagraph token padding must not pad request-indexed metadata.""" num_speculative_tokens = 3 diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index ac9c2a881ab7..fc22ff5eae5e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1551,12 +1551,30 @@ def _forward_core( self.conv1d.weight.size(0), self.conv1d.weight.size(2) ) + # When the spec-decode rows lead the batch, the spec and non-spec + # tokens are two contiguous ranges, so slices replace the gathers. + spec_tokens_first = ( + spec_sequence_masks is not None and attn_metadata.spec_tokens_first + ) + num_spec_tokens = attn_metadata.num_spec_decode_tokens + non_spec_tokens = slice( + num_spec_tokens, + num_spec_tokens + + attn_metadata.num_prefill_tokens + + attn_metadata.num_decode_tokens, + ) + if spec_sequence_masks is not None: if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: mixed_qkv_spec = mixed_qkv a_spec = a b_spec = b mixed_qkv_non_spec = None + elif spec_tokens_first: + mixed_qkv_spec = mixed_qkv[:num_spec_tokens] + a_spec = a[:num_spec_tokens] + b_spec = b[:num_spec_tokens] + mixed_qkv_non_spec = mixed_qkv[non_spec_tokens] else: mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) a_spec = a.index_select(0, spec_token_indx) @@ -1633,7 +1651,10 @@ def _forward_core( assert mixed_qkv_non_spec is not None, ( "mixed_qkv_non_spec must be provided for prefill path" ) - if spec_sequence_masks is not None: + if spec_tokens_first: + a_non_spec = a[non_spec_tokens] + b_non_spec = b[non_spec_tokens] + elif spec_sequence_masks is not None: a_non_spec = a.index_select(0, non_spec_token_indx) b_non_spec = b.index_select(0, non_spec_token_indx) else: @@ -1830,7 +1851,10 @@ def _forward_core( core_attn_out_non_spec, last_recurrent_state = None, None # 3. Merge core attention output - if spec_sequence_masks is not None and core_attn_out_non_spec is not None: + if spec_tokens_first and core_attn_out_non_spec is not None: + core_attn_out[:num_spec_tokens].copy_(core_attn_out_spec.squeeze(0)) + core_attn_out[non_spec_tokens].copy_(core_attn_out_non_spec.squeeze(0)) + elif spec_sequence_masks is not None and core_attn_out_non_spec is not None: core_attn_out.index_copy_(0, spec_token_indx, core_attn_out_spec.squeeze(0)) core_attn_out.index_copy_( 0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 38b8a5fa76d5..52bbdc50fdd5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -112,6 +112,10 @@ class GDNAttentionMetadata: spec_sequence_masks_cpu: torch.Tensor | None = None # shape: [batch,] spec_token_indx: torch.Tensor | None = None non_spec_token_indx: torch.Tensor | None = None + # True when a batch with spec decodes and prefills puts every spec-decode + # row first: the spec tokens are then the leading num_spec_decode_tokens + # tokens and the non-spec tokens follow them. + spec_tokens_first: bool = False num_accepted_tokens: torch.Tensor | None = None # shape: [batch,] uniform_spec_sequence_length: int | None = None # None for ragged batches @@ -338,6 +342,7 @@ def build( # type: ignore[override] ) uniform_spec_sequence_length = None + spec_tokens_first = False spec_sequence_masks_cpu: torch.Tensor | None = None if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: spec_sequence_masks = None @@ -458,6 +463,9 @@ def build( # type: ignore[override] non_spec_query_start_loc = None non_spec_query_start_loc_cpu = None else: + spec_tokens_first = bool( + spec_sequence_masks_cpu[:num_spec_decodes].all() + ) spec_token_masks = torch.repeat_interleave( spec_sequence_masks, query_lens, @@ -698,6 +706,7 @@ def build( # type: ignore[override] spec_sequence_masks_cpu=spec_sequence_masks_cpu, spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, + spec_tokens_first=spec_tokens_first, num_accepted_tokens=num_accepted_tokens, uniform_spec_sequence_length=uniform_spec_sequence_length, nums_dict=nums_dict, From aedc212d33d74acf8a71ce5a0fb42ceea55ef0db Mon Sep 17 00:00:00 2001 From: kimseunghyun Date: Sun, 11 Oct 2026 05:52:03 +0900 Subject: [PATCH 3/4] [Perf][GDN] Pass strided q/k/v views to the recurrent update rearrange_mixed_qkv copies q, k and v into one contiguous buffer: one split, three reshape copies, one cat, three slices and three views. In _forward_core, every caller passes the result to fused_sigmoid_gating_delta_rule_update, which calls .contiguous() on q, k and v itself. Add a contiguous=False mode that returns the three strided views, and use it at the three recurrent-update call sites (spec decodes, peeled decodes, decode-only batches). This removes about seven host calls and one copy kernel per call site. The other callers (CPU ops, OLMo) keep the contiguous buffer. The values are bit-identical. Test: tests/kernels/mamba/test_gdn_forward_core_mixed.py compares the strided views with the contiguous copies for spec, mixed spec, mixed non-spec and decode-only batches, bit for bit (CUDA). Co-authored-by: Claude Opus 5.5 Signed-off-by: kimseunghyun --- .../mamba/test_gdn_forward_core_mixed.py | 41 +++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 23 ++++++++--- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_forward_core_mixed.py b/tests/kernels/mamba/test_gdn_forward_core_mixed.py index 8317f0fe690a..036a57aad600 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_mixed.py +++ b/tests/kernels/mamba/test_gdn_forward_core_mixed.py @@ -12,6 +12,7 @@ from __future__ import annotations import dataclasses +import functools import types from unittest.mock import patch @@ -67,6 +68,13 @@ [NUM_SPEC, 1, -1, -1], ), } +# Batches whose recurrent-update q/k/v come from rearrange_mixed_qkv. +RECURRENT_BATCHES = { + "spec-decodes-then-prefills": BATCHES["spec-decodes-then-prefills"], + "spec-decodes": ([128, 90], [SPEC_TOKENS, 2], [NUM_SPEC, 1]), + "decodes-then-prefills": ([128, 90, 200], [1, 1, 37], [-1, -1, -1]), + "decodes": ([128, 90], [1, 1], [-1, -1]), +} def _make_vllm_config(): @@ -266,3 +274,36 @@ def test_spec_token_slices_match_gathers() -> None: for sliced, gathered in zip(*results): torch.testing.assert_close(sliced, gathered, atol=0, rtol=0) + + +def _rearrange_mixed_qkv_copies(layer, mixed_qkv, *, contiguous): + return QwenGatedDeltaNetAttention.rearrange_mixed_qkv(layer, mixed_qkv) + + +@pytest.mark.parametrize("batch_name", RECURRENT_BATCHES) +@torch.inference_mode() +def test_strided_qkv_views_match_contiguous_copies(batch_name: str) -> None: + """The recurrent update gives bit-identical results for the strided q/k/v + views and for the contiguous q/k/v copies of rearrange_mixed_qkv.""" + torch.manual_seed(0) + seq_lens, query_lens, draft_tokens = RECURRENT_BATCHES[batch_name] + vllm_config = _make_vllm_config() + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + metadata = _build_metadata(vllm_config, batch, draft_tokens) + + num_tokens = batch.compute_num_tokens() + kv_cache, weights, mixed_qkvz, b, a = _make_inputs(metadata, num_tokens) + results = [] + for contiguous in (False, True): + step_kv_cache = tuple(state.clone() for state in kv_cache) + layer = _build_layer(vllm_config, step_kv_cache, weights) + if contiguous: + # Ignore the call site's contiguous=False: copy as before. + layer.rearrange_mixed_qkv = functools.partial( + _rearrange_mixed_qkv_copies, layer + ) + out = _run_forward_core(layer, metadata, mixed_qkvz, b, a, num_tokens) + results.append((out, *step_kv_cache)) + + for strided, copied in zip(*results): + torch.testing.assert_close(strided, copied, atol=0, rtol=0) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index fc22ff5eae5e..5a94aaad5ab4 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -981,14 +981,17 @@ def prepare_gdn_attention_core_inputs( return mixed_qkv_out, z_out, b_out, a_out - def rearrange_mixed_qkv(self, mixed_qkv): - """Split packed qkv into contiguous (1, seq, heads, dim) tensors. + def rearrange_mixed_qkv(self, mixed_qkv, *, contiguous: bool = True): + """Split packed qkv into (1, seq, heads, dim) tensors. The original code used ``rearrange(x, "l (h d) -> 1 l h d", d=...)`` followed by ``.contiguous()`` on each tensor. This version flattens all three splits into a single buffer via ``torch.cat`` so that torch.compile emits one Triton copy kernel instead of three separate contiguous() calls. + + With ``contiguous=False``, return strided views of ``mixed_qkv`` + instead. Use it only for a consumer that accepts strided input. """ if mixed_qkv is None: return None, None, None @@ -999,6 +1002,12 @@ def rearrange_mixed_qkv(self, mixed_qkv): v_dim = self.value_dim // self.tp_size query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + if not contiguous: + return ( + query.view(1, seq_len, q_dim // self.head_k_dim, self.head_k_dim), + key.view(1, seq_len, k_dim // self.head_k_dim, self.head_k_dim), + value.view(1, seq_len, v_dim // self.head_v_dim, self.head_v_dim), + ) fused = torch.cat( [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 @@ -1637,7 +1646,10 @@ def _forward_core( else: mixed_qkv_non_spec = None - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + # The recurrent update below makes its own contiguous q/k/v copies. + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv( + mixed_qkv_spec, contiguous=False + ) # Split mixed non-spec-decode+prefill to process independently split_non_spec = ( @@ -1695,7 +1707,7 @@ def _forward_core( beta_non_spec = beta_non_spec.unsqueeze(0) else: query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( - mixed_qkv_non_spec + mixed_qkv_non_spec, contiguous=False ) g_non_spec = None beta_non_spec = None @@ -1730,7 +1742,8 @@ def _forward_core( # 2.2: Process non-spec-decode part if split_non_spec: query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( - mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] + mixed_qkv_non_spec[:num_decode_tokens], # type: ignore[index] + contiguous=False, ) core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, From 8286a881217f52494e7f9d43e6ab425e1ad189c0 Mon Sep 17 00:00:00 2001 From: kimseunghyun Date: Sun, 11 Oct 2026 09:16:44 +0900 Subject: [PATCH 4/4] [Bugfix][GDN] Use int64 prefill state indices for index_copy_ index_copy_ only accepts int64 indices, but the GDN state indices are int32, so every prefill step raised "index_copy_(): Expected a long tensor for index, but got Int" (found by the CUDA tests on A100). The metadata builder (build and update_block_table) now also stores prefill_state_indices_int64, cast once per step, and the layer uses it for index_select and index_copy_. prefill_state_indices keeps its dtype for its other users. Test: tests/v1/attention/test_gdn_metadata_builder.py checks the int64 copy against prefill_state_indices and after update_block_table (CPU). Co-authored-by: Claude Opus 5.5 Signed-off-by: kimseunghyun --- .../mamba/test_gdn_forward_core_split.py | 3 +++ .../v1/attention/test_gdn_metadata_builder.py | 22 +++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 3 ++- vllm/v1/attention/backends/gdn_attn.py | 9 ++++++++ 4 files changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py index 4b87e91ebf13..60d1163eecf7 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_split.py +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -226,6 +226,9 @@ def test_forward_core_split_matches_unified( # Unified path: the chunk kernel processes the full non-spec batch. prefill_query_start_loc=meta_split.non_spec_query_start_loc, prefill_state_indices=meta_split.non_spec_state_indices_tensor, + prefill_state_indices_int64=meta_split.non_spec_state_indices_tensor.to( + torch.int64 + ), prefill_has_initial_state=meta_split.has_initial_state, prefill_no_initial_state_mask=~meta_split.has_initial_state.view(-1, 1, 1, 1), ) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 0cbba18aaa7d..b05dca543dc1 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -246,6 +246,7 @@ def test_update_block_table_matches_build( "spec_state_indices_tensor", "non_spec_state_indices_tensor", "prefill_state_indices", + "prefill_state_indices_int64", ) source_indices = [getattr(source, f) for f in fields] source_indices = [t if t is None else t.clone() for t in source_indices] @@ -277,6 +278,27 @@ def test_has_initial_state_after_reclassification(): assert meta.has_initial_state[0].item() is True +@pytest.mark.parametrize( + "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() +) +def test_prefill_state_indices_int64_matches_prefill_state_indices( + test_case: GDNBuildTestCase, +): + """The int64 copy (for index_copy_) holds the same prefill state indices.""" + builder = _create_gdn_builder(test_case.num_speculative_tokens) + batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens) + meta = _build(builder, batch, test_case.num_decode_draft_tokens) + + if meta.prefill_state_indices is None: + assert meta.prefill_state_indices_int64 is None + return + assert meta.prefill_state_indices_int64 is not None + assert meta.prefill_state_indices_int64.dtype == torch.int64 + assert torch.equal( + meta.prefill_state_indices_int64, meta.prefill_state_indices.to(torch.int64) + ) + + @pytest.mark.parametrize( "seq_lens,query_lens,draft_tokens,expected", [ diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 5a94aaad5ab4..32474880040f 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1770,7 +1770,8 @@ def _forward_core( # kernel are precomputed by the metadata builder (the prefill tail # when decodes are peeled off, else the full non-spec batch), so they # don't need to be re-derived per layer. - prefill_state_indices = attn_metadata.prefill_state_indices + # index_copy_ needs int64 indices; the builder casts them once. + prefill_state_indices = attn_metadata.prefill_state_indices_int64 no_initial_state_mask = attn_metadata.prefill_no_initial_state_mask assert prefill_state_indices is not None assert no_initial_state_mask is not None diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 52bbdc50fdd5..dd91a19e6572 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -126,6 +126,8 @@ class GDNAttentionMetadata: # Chunk-kernel inputs for prefill prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None + # prefill_state_indices as int64, for index_copy_ into the SSM state. + prefill_state_indices_int64: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None # ~prefill_has_initial_state as [num_prefills, 1, 1, 1], to zero the # gathered [num_prefills, HV, V, K] initial states in one masked_fill_. @@ -523,6 +525,7 @@ def build( # type: ignore[override] chunk_offsets: torch.Tensor | None = None prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None + prefill_state_indices_int64: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None prefill_no_initial_state_mask: torch.Tensor | None = None aiter_prefill_metadata: object | None = None @@ -546,6 +549,8 @@ def build( # type: ignore[override] prefill_query_start_loc = non_spec_query_start_loc prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu prefill_state_indices = non_spec_state_indices_tensor + assert prefill_state_indices is not None + prefill_state_indices_int64 = prefill_state_indices.to(torch.int64) if self.gdn_prefill_backend == "aiter_flydsl": # AITER carries its own reusable varlen metadata and has no use @@ -695,6 +700,7 @@ def build( # type: ignore[override] chunk_offsets=chunk_offsets, prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, + prefill_state_indices_int64=prefill_state_indices_int64, prefill_has_initial_state=prefill_has_initial_state, prefill_no_initial_state_mask=prefill_no_initial_state_mask, aiter_prefill_metadata=aiter_prefill_metadata, @@ -797,6 +803,9 @@ def update_block_table( spec_state_indices_tensor=spec_indices, non_spec_state_indices_tensor=non_spec_indices, prefill_state_indices=prefill_indices, + prefill_state_indices_int64=( + None if prefill_indices is None else prefill_indices.to(torch.int64) + ), checkpoint=checkpoint, )