From fa27d4e9cf3c8d8a5a143f38c346b27c02b2c2e3 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Sat, 6 Jun 2026 18:13:50 +0400 Subject: [PATCH] [PERF] [Qwen3.5] Split mixed prefill+decode batches: route decodes to the recurrent kernel (#44700) Signed-off-by: Vadim Gimpelson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml | 12 + .../gsm8k/configs/models-qwen35-blackwell.txt | 3 +- .../mamba/test_gdn_forward_core_split.py | 296 ++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 98 ++++-- vllm/v1/attention/backends/gdn_attn.py | 48 ++- 5 files changed, 426 insertions(+), 31 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml create mode 100644 tests/kernels/mamba/test_gdn_forward_core_split.py diff --git a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml new file mode 100644 index 00000000000..d247515a0f0 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml @@ -0,0 +1,12 @@ +model_name: "nvidia/Qwen3.5-397B-A17B-NVFP4" +accuracy_threshold: 0.88 +tolerance: 0.03 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --max-model-len 4096 + --data-parallel-size 2 + --enable-expert-parallel + --max-num-seqs 384 + --spec-method mtp + --spec-tokens 3 diff --git a/tests/evals/gsm8k/configs/models-qwen35-blackwell.txt b/tests/evals/gsm8k/configs/models-qwen35-blackwell.txt index 908ada3a22c..aef2b9dbe65 100644 --- a/tests/evals/gsm8k/configs/models-qwen35-blackwell.txt +++ b/tests/evals/gsm8k/configs/models-qwen35-blackwell.txt @@ -1,3 +1,4 @@ Qwen3.5-35B-A3B-DEP2.yaml Qwen3.5-35B-A3B-FP8-DEP2.yaml -Qwen3.5-397B-A17B-NVFP4-DEP2.yaml \ No newline at end of file +Qwen3.5-397B-A17B-NVFP4-DEP2.yaml +Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml \ No newline at end of file diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py new file mode 100644 index 00000000000..f2bfc30abdb --- /dev/null +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration test for the non-spec decode split in +``GatedDeltaNet._forward_core``. + +On a pure non-spec batch that mixes prefills with 1-token decodes, the layer +peels the decodes (the contiguous decode-first front slice) off to +``fused_sigmoid_gating_delta_rule_update`` -- the same recurrent update kernel +the spec-decode path uses -- and runs only the prefill tail through +``chunk_gated_delta_rule``. This must produce the same core-attention output and +the same ssm-state pool update as running *everything* through +``chunk_gated_delta_rule`` (the previous behavior). + +Both paths are exercised through the REAL ``_forward_core``: + +* ``meta_split`` is built by the real ``GDNAttentionMetadataBuilder`` for a + mixed batch, so ``num_decodes > 0`` triggers the peel (and the builder rebases + ``chunk_indices``/``chunk_offsets`` to the prefill-only tail). +* ``meta_unified`` is the same metadata with the decodes reclassified as + prefills and full-batch chunk metadata, which forces ``_forward_core`` through + the existing chunk-only path on identical inputs (the conv is unified over all + non-spec tokens in both paths, so it cancels out and only the recurrent split + is compared). + +The Triton/FLA chunk backend is forced so the prefill-only ``chunk_indices`` +must stay consistent with the rebased ``cu_seqlens`` (a stringent, backend +portable check of the split wiring). +""" + +from __future__ import annotations + +import dataclasses +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.is_device_capability_family(100) +): + pytest.skip( + reason="GDN _forward_core split test uses the CuteDSL prefill backend " + "(requires CUDA SM10x).", + allow_module_level=True, + ) + +from tests.v1.attention.utils import ( # noqa: E402 + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import set_current_vllm_config # noqa: E402 +from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402 + prepare_chunk_indices, + prepare_chunk_offsets, +) +from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE # 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 + GDNAttentionMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 + +# Small GDN dims; head_k_dim/head_v_dim=128 keeps the chunk/update kernels happy. +H = 4 # num key heads +HV = 8 # num value heads +K = 128 # head_k_dim +V = 128 # head_v_dim +CONV_KERNEL = 4 +KEY_DIM = H * K +VALUE_DIM = HV * V +CONV_DIM = 2 * KEY_DIM + VALUE_DIM +BLOCK_SIZE = 16 +PREFIX = "model.layers.0.linear_attn" + + +def _make_vllm_config(): + # A small, ungated GDN model whose config is cached locally; only the config + # (scheduler/cache/compilation/hf) is used here, never the weights. Inject + # linear_key_head_dim=128 and request the CuteDSL prefill backend -- the + # supported GDN chunk kernel on Blackwell (the Triton/FLA chunk kernel is + # unsupported on SM10x). CuteDSL consumes chunk_indices/chunk_offsets, so + # this also exercises the prefill-only chunk-metadata wiring. + cfg = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + hf_config_override={"linear_key_head_dim": K}, + ) + cfg.additional_config = {"gdn_prefill_backend": "cutedsl"} + return cfg + + +def _build_layer( + vllm_config, conv_state, ssm_state, A_log, dt_bias, conv_weight, conv_bias +): + """A minimal object that runs the real ``_forward_core`` bound to it.""" + layer = types.SimpleNamespace() + layer.prefix = PREFIX + layer.enable_packed_recurrent_decode = False + layer.tp_size = 1 + layer.num_k_heads = H + layer.num_v_heads = HV + layer.head_k_dim = K + layer.head_v_dim = V + layer.key_dim = KEY_DIM + layer.value_dim = VALUE_DIM + layer.activation = "silu" + layer.A_log = A_log + layer.dt_bias = dt_bias + layer.conv1d = types.SimpleNamespace(weight=conv_weight, bias=conv_bias) + layer.kv_cache = (conv_state, ssm_state) + 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 _run_forward_core(layer, meta, mixed_qkv, b, a, num_tokens): + core_attn_out = torch.zeros( + num_tokens, HV, V, dtype=mixed_qkv.dtype, device=mixed_qkv.device + ) + ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta}) + with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx): + layer._forward_core( + mixed_qkv=mixed_qkv.clone(), + b=b.clone(), + a=a.clone(), + core_attn_out=core_attn_out, + ) + return core_attn_out + + +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_decodes,prefill_lens", [(3, [512, 300]), (4, [64, 5])]) +@pytest.mark.parametrize("fresh_prefill", [False, True]) +def test_forward_core_split_matches_unified( + state_dtype: torch.dtype, + num_decodes: int, + prefill_lens: list[int], + fresh_prefill: bool, +) -> None: + torch.manual_seed(0) + device = torch.device("cuda") + vllm_config = _make_vllm_config() + + # Decode-first batch: D 1-token decodes (with context), then the prefills. + decode_seq_lens = [64] * num_decodes + prefill_seq_lens = [ + pl if (fresh_prefill and i == 0) else pl + 37 + for i, pl in enumerate(prefill_lens) + ] + seq_lens = decode_seq_lens + prefill_seq_lens + query_lens = [1] * num_decodes + list(prefill_lens) + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + + builder = GDNAttentionMetadataBuilder( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,) + ), + layer_names=[PREFIX], + vllm_config=vllm_config, + device=device, + ) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ) + with set_current_vllm_config(vllm_config): + meta_split = builder.build(common_prefix_len=0, common_attn_metadata=common) + + assert meta_split.spec_sequence_masks is None + assert meta_split.num_decodes == num_decodes + assert meta_split.num_prefills == len(prefill_lens) + assert meta_split.num_decode_tokens == num_decodes + assert builder.gdn_prefill_backend == "cutedsl" + + num_tokens = sum(query_lens) + + # Full-batch chunk metadata for the unified reference path, built the same + # way the builder would for a non-split batch (backend-matched). + cu_full = meta_split.non_spec_query_start_loc + if builder.gdn_prefill_backend == "cutedsl": + from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( + prepare_metadata_cutedsl, + ) + + full_ci, full_co = prepare_metadata_cutedsl( + cu_full, int(cu_full[-1].item()), FLA_CHUNK_SIZE + ) + else: + cu_full_cpu = cu_full.cpu() + full_ci = prepare_chunk_indices(cu_full_cpu, FLA_CHUNK_SIZE).to(device) + full_co = prepare_chunk_offsets(cu_full_cpu, FLA_CHUNK_SIZE).to(device) + meta_unified = dataclasses.replace( + meta_split, + num_decodes=0, + num_decode_tokens=0, + num_prefills=meta_split.num_decodes + meta_split.num_prefills, + num_prefill_tokens=( + meta_split.num_decode_tokens + meta_split.num_prefill_tokens + ), + chunk_indices=full_ci, + chunk_offsets=full_co, + # 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_has_initial_state=meta_split.has_initial_state, + ) + + # Size the state pools from the indices the builder actually produced. + pool_size = int(meta_split.non_spec_state_indices_tensor.max().item()) + 1 + conv_state_shape, temporal_state_shape = ( + MambaStateShapeCalculator.gated_delta_net_state_shape( + 1, H, HV, K, V, CONV_KERNEL, num_spec=0 + ) + ) + conv_state0 = ( + torch.randn(pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device) + * 0.05 + ) + ssm_state0 = ( + torch.randn(pool_size, *temporal_state_shape, dtype=state_dtype, device=device) + * 0.05 + ) + + A_log = torch.randn(HV, dtype=torch.float32, device=device) * 0.1 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1 + conv_weight = ( + torch.randn(CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device) * 0.1 + ) + conv_bias = torch.randn(CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1 + + mixed_qkv = ( + torch.randn(num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1 + ) + a = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1 + b = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1 + + # ---- Split path (real _forward_core, meta_split) ---- + conv_state_split = conv_state0.clone() + ssm_state_split = ssm_state0.clone() + layer_split = _build_layer( + vllm_config, + conv_state_split, + ssm_state_split, + A_log, + dt_bias, + conv_weight, + conv_bias, + ) + out_split = _run_forward_core(layer_split, meta_split, mixed_qkv, b, a, num_tokens) + + # ---- Unified path (real _forward_core, meta_unified) ---- + conv_state_unified = conv_state0.clone() + ssm_state_unified = ssm_state0.clone() + layer_unified = _build_layer( + vllm_config, + conv_state_unified, + ssm_state_unified, + A_log, + dt_bias, + conv_weight, + conv_bias, + ) + out_unified = _run_forward_core( + layer_unified, meta_unified, mixed_qkv, b, a, num_tokens + ) + + # Conv is unified in both paths, so the conv-state update must be identical. + torch.testing.assert_close(conv_state_split, conv_state_unified, atol=0, rtol=0) + + # Chunk vs. recurrent update accumulate in different orders; mirror the + # tolerances used by the kernel-level parity test. + if state_dtype == torch.float32: + atol = rtol = 2e-2 + else: + atol = rtol = 6e-2 + torch.testing.assert_close(out_split, out_unified, atol=atol, rtol=rtol) + torch.testing.assert_close(ssm_state_split, ssm_state_unified, atol=atol, rtol=rtol) 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 7a0d50c74e3..a7a609eba9b 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 @@ -66,7 +66,7 @@ from vllm.utils.torch_utils import ( ) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata -# Optional ROCm AITER Triton kernels for the GDN decode fast-path. +# Optional ROCm AITER Triton kernels for the GDN decode path. # Availability is checked centrally via rocm_aiter_ops; the actual function # references are imported here so that they can be called without per-call # import overhead. @@ -897,8 +897,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): projected_states_ba, z, core_attn_out, - fast_kernel=True, layer_name=_encode_layer_name(self.prefix), + use_aiter=True, ) self._output_projection(core_attn_out, z, output, num_tokens) @@ -958,7 +958,6 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): b, a, core_attn_out, - fast_kernel=False, layer_name=_encode_layer_name(self.prefix), ) @@ -1206,7 +1205,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): qkvz/ba layout. For decode-only (no spec, no prefill) interleaved-GQA layouts, - dispatches directly to ``_forward_core_decode_fast``. Otherwise unpacks + dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks the packed layout and falls through to ``_forward_core``. Args: @@ -1237,7 +1236,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): and attn_metadata.num_prefills == 0 and attn_metadata.num_decodes > 0 ): - return self._forward_core_decode_fast( + return self._forward_core_decode_aiter( qkvz=qkvz, ba=ba, z_out=z_out, @@ -1391,6 +1390,15 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): mixed_qkv_non_spec = None query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + + # Split mixed non-spec-decode+prefill to process independently + split_non_spec = ( + spec_sequence_masks is None + and attn_metadata.num_prefills > 0 + and attn_metadata.num_decodes > 0 + ) + num_decode_tokens = attn_metadata.num_decode_tokens + if attn_metadata.num_prefills > 0: assert mixed_qkv_non_spec is not None, ( "mixed_qkv_non_spec must be provided for prefill path" @@ -1402,6 +1410,15 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): a_non_spec = a b_non_spec = b + if split_non_spec: + conv_output_prefill = mixed_qkv_non_spec[num_decode_tokens:] + a_prefill = a_non_spec[num_decode_tokens:] + b_prefill = b_non_spec[num_decode_tokens:] + else: + conv_output_prefill = mixed_qkv_non_spec + a_prefill = a_non_spec + b_prefill = b_non_spec + ( query_non_spec, key_non_spec, @@ -1409,9 +1426,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): g_non_spec, beta_non_spec, ) = fused_post_conv_prep( - conv_output=mixed_qkv_non_spec, - a=a_non_spec, - b=b_non_spec, + conv_output=conv_output_prefill, + a=a_prefill, + b=b_prefill, A_log=self.A_log, dt_bias=self.dt_bias, num_k_heads=self.num_k_heads // self.tp_size, @@ -1459,12 +1476,42 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): else: core_attn_out_spec, last_recurrent_state = None, None - # 2.2: Process the remaining part + # 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] + ) + core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[:num_decode_tokens], + b=b[:num_decode_tokens], + dt_bias=self.dt_bias, + q=query_decode, + k=key_decode, + v=value_decode, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + 1 + ], + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + ) + else: + core_attn_out_decode = None + + # 2.3: Process the remaining part (prefill chunk, or non-spec decode-only) if attn_metadata.num_prefills > 0: - assert non_spec_state_indices_tensor is not None - initial_state = ssm_state[non_spec_state_indices_tensor].contiguous() # type: ignore[index] - assert has_initial_state is not None - initial_state[~has_initial_state, ...] = 0 # type: ignore[operator] + # State indices, initial-state mask and cu_seqlens for the chunk + # 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 + prefill_has_initial_state = attn_metadata.prefill_has_initial_state + 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 ( core_attn_out_non_spec, last_recurrent_state, @@ -1476,15 +1523,20 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): beta=beta_non_spec, initial_state=initial_state, output_final_state=True, - cu_seqlens=non_spec_query_start_loc, + cu_seqlens=attn_metadata.prefill_query_start_loc, chunk_indices=attn_metadata.chunk_indices, chunk_offsets=attn_metadata.chunk_offsets, use_qk_l2norm_in_kernel=False, ) # Init cache - ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to( - ssm_state.dtype - ) + ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) + + if split_non_spec: + # Stitch the peeled decode outputs in front of the prefill + # outputs (decode-first order). + core_attn_out_non_spec = torch.cat( + [core_attn_out_decode, core_attn_out_non_spec], dim=1 + ) elif attn_metadata.num_decodes > 0: core_attn_out_non_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( @@ -1523,7 +1575,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): else: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) - def _forward_core_decode_fast( + def _forward_core_decode_aiter( self, qkvz: torch.Tensor, ba: torch.Tensor, @@ -1649,17 +1701,17 @@ def qwen_gdn_attention_core( b_or_ba: torch.Tensor, a_or_z_out: torch.Tensor, core_attn_out: torch.Tensor, - fast_kernel: bool, layer_name: LayerNameType, + use_aiter: bool = False, ) -> None: """Custom op dispatching to _forward_core or _forward_core_rocm. Handles conv1d + recurrent attention only; input/output projections are performed by the caller. - When ``fast_kernel=False`` (standard path): + When ``use_aiter=False`` (standard path): qkv_or_qkvz is [q, k, v], b_or_ba is b, a_or_z_out is a (read-only). - When ``fast_kernel=True`` (AITER Triton fast path, ROCm only): + When ``use_aiter=True`` (AITER Triton path, ROCm only): qkv_or_qkvz is [q, k, v, z], b_or_ba is [b, a], a_or_z_out is the z output buffer (mutated in-place). @@ -1668,7 +1720,7 @@ def qwen_gdn_attention_core( layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] - if fast_kernel: + if use_aiter: self._forward_core_rocm( qkvz=qkv_or_qkvz, ba=b_or_ba, @@ -1689,8 +1741,8 @@ def gdn_attention_core_fake( b_or_ba: torch.Tensor, a_or_z_out: torch.Tensor, core_attn_out: torch.Tensor, - fast_kernel: bool, layer_name: LayerNameType, + use_aiter: bool = False, ) -> None: """Fake implementation for torch.compile.""" return diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 2c0ff984b41..9323c5d8a46 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -67,6 +67,10 @@ class GDNAttentionMetadata: # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None + # Chunk-kernel inputs for prefill + prefill_query_start_loc: torch.Tensor | None = None + prefill_state_indices: torch.Tensor | None = None + prefill_has_initial_state: torch.Tensor | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -322,19 +326,42 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None + prefill_query_start_loc: torch.Tensor | None = None + prefill_state_indices: torch.Tensor | None = None + prefill_has_initial_state: torch.Tensor | None = None if num_prefills > 0: from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE + # In a mixed non-spec batch, decodes are peeled off to the recurrent + # kernel (decode-first front slice), so build chunk metadata from the + # rebased prefill-only cu_seqlens; otherwise use the full non-spec one. + # _forward_core keys off the same condition, so they agree. + if spec_sequence_masks is None and num_decodes > 0: + assert non_spec_query_start_loc is not None + assert non_spec_query_start_loc_cpu is not None + assert non_spec_state_indices_tensor is not None + prefill_query_start_loc = ( + non_spec_query_start_loc[num_decodes:] - num_decode_tokens + ) + prefill_query_start_loc_cpu = ( + non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens + ) + prefill_state_indices = non_spec_state_indices_tensor[num_decodes:] + else: + 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 + if self.gdn_prefill_backend == "cutedsl": from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( prepare_metadata_cutedsl, ) - assert non_spec_query_start_loc is not None - assert non_spec_query_start_loc_cpu is not None - total_tokens = int(non_spec_query_start_loc_cpu[-1].item()) + assert prefill_query_start_loc is not None + assert prefill_query_start_loc_cpu is not None + total_tokens = int(prefill_query_start_loc_cpu[-1].item()) chunk_indices, chunk_offsets = prepare_metadata_cutedsl( - non_spec_query_start_loc, + prefill_query_start_loc, total_tokens, FLA_CHUNK_SIZE, ) @@ -348,12 +375,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] prepare_chunk_offsets, ) - assert non_spec_query_start_loc_cpu is not None + assert prefill_query_start_loc_cpu is not None chunk_indices = prepare_chunk_indices( - non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE + prefill_query_start_loc_cpu, FLA_CHUNK_SIZE ).to(device=gpu_device, non_blocking=True) chunk_offsets = prepare_chunk_offsets( - non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE + prefill_query_start_loc_cpu, FLA_CHUNK_SIZE ).to(device=gpu_device, non_blocking=True) if num_prefills > 0: @@ -367,6 +394,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] device=query_start_loc.device, ) ) + if spec_sequence_masks is None and num_decodes > 0: + prefill_has_initial_state = has_initial_state[num_decodes:] + else: + prefill_has_initial_state = has_initial_state else: has_initial_state = None @@ -458,6 +489,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] has_initial_state=has_initial_state, chunk_indices=chunk_indices, chunk_offsets=chunk_offsets, + prefill_query_start_loc=prefill_query_start_loc, + prefill_state_indices=prefill_state_indices, + prefill_has_initial_state=prefill_has_initial_state, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor,