karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[PERF] [Qwen3.5] Split mixed prefill+decode batches: route decodes to the recurrent kernel (#44700)

Browse Source 
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Vadim Gimpelson
2026-06-06 22:13:50 +08:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 67d3792d99
commit fa27d4e9cf
5 changed files with 426 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml
+12
View File
@@ -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
tests/evals/gsm8k/configs/models-qwen35-blackwell.txt
+2 -1
View File
@@ -1,3 +1,4 @@
Qwen3.5-35B-A3B-DEP2.yaml
Qwen3.5-35B-A3B-FP8-DEP2.yaml
Qwen3.5-397B-A17B-NVFP4-DEP2.yaml
Qwen3.5-397B-A17B-NVFP4-DEP2.yaml
Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml
tests/kernels/mamba/test_gdn_forward_core_split.py
+296
View File
@@ -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)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+75 -23
View File
@@ -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
vllm/v1/attention/backends/gdn_attn.py
+41 -7
View File
@@ -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,