A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 18:45 UTC
[PERF] [Qwen3.5] Split mixed prefill+decode batches: route decodes to the recurrent kernel (#44700)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Vadim Gimpelson committed
Jun 6, 2026 at 18:13 UTC
fa27d4e9cf3c8d8a5a143f38c346b27c02b2c2e3
5 files changed
+426
-31
tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml
new
+12
@@ -0,0 +1,12 @@
1
+model_name: "nvidia/Qwen3.5-397B-A17B-NVFP4"
2
+accuracy_threshold: 0.88
3
+tolerance: 0.03
4
+num_questions: 1319
5
+num_fewshot: 5
6
+server_args: >-
7
+ --max-model-len 4096
8
+ --data-parallel-size 2
9
+ --enable-expert-parallel
10
+ --max-num-seqs 384
11
+ --spec-method mtp
12
+ --spec-tokens 3
tests/evals/gsm8k/configs/models-qwen35-blackwell.txt
+2
-1
@@ -1,3 +1,4 @@
1
Qwen3.5-35B-A3B-DEP2.yaml
2
Qwen3.5-35B-A3B-FP8-DEP2.yaml
3
-Qwen3.5-397B-A17B-NVFP4-DEP2.yaml
\ No newline at end of file
3
+Qwen3.5-397B-A17B-NVFP4-DEP2.yaml
4
+Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml
\ No newline at end of file
tests/kernels/mamba/test_gdn_forward_core_split.py
new
+296
@@ -0,0 +1,296 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Integration test for the non-spec decode split in
4
+``GatedDeltaNet._forward_core``.
5
+
6
+On a pure non-spec batch that mixes prefills with 1-token decodes, the layer
7
+peels the decodes (the contiguous decode-first front slice) off to
8
+``fused_sigmoid_gating_delta_rule_update`` -- the same recurrent update kernel
9
+the spec-decode path uses -- and runs only the prefill tail through
10
+``chunk_gated_delta_rule``. This must produce the same core-attention output and
11
+the same ssm-state pool update as running *everything* through
12
+``chunk_gated_delta_rule`` (the previous behavior).
13
+
14
+Both paths are exercised through the REAL ``_forward_core``:
15
+
16
+* ``meta_split`` is built by the real ``GDNAttentionMetadataBuilder`` for a
17
+ mixed batch, so ``num_decodes > 0`` triggers the peel (and the builder rebases
18
+ ``chunk_indices``/``chunk_offsets`` to the prefill-only tail).
19
+* ``meta_unified`` is the same metadata with the decodes reclassified as
20
+ prefills and full-batch chunk metadata, which forces ``_forward_core`` through
21
+ the existing chunk-only path on identical inputs (the conv is unified over all
22
+ non-spec tokens in both paths, so it cancels out and only the recurrent split
23
+ is compared).
24
+
25
+The Triton/FLA chunk backend is forced so the prefill-only ``chunk_indices``
26
+must stay consistent with the rebased ``cu_seqlens`` (a stringent, backend
27
+portable check of the split wiring).
28
+"""
29
+
30
+from __future__ import annotations
31
+
32
+import dataclasses
33
+import types
34
+from unittest.mock import patch
35
+
36
+import pytest
37
+import torch
38
+
39
+from vllm.platforms import current_platform
40
+
41
+if not (
42
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
43
+):
44
+ pytest.skip(
45
+ reason="GDN _forward_core split test uses the CuteDSL prefill backend "
46
+ "(requires CUDA SM10x).",
47
+ allow_module_level=True,
48
+ )
49
+
50
+from tests.v1.attention.utils import ( # noqa: E402
51
+ BatchSpec,
52
+ create_common_attn_metadata,
53
+ create_vllm_config,
54
+)
55
+from vllm.config import set_current_vllm_config # noqa: E402
56
+from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402
57
+ prepare_chunk_indices,
58
+ prepare_chunk_offsets,
59
+)
60
+from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE # noqa: E402
61
+from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402
62
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402
63
+ ChunkGatedDeltaRule,
64
+ QwenGatedDeltaNetAttention,
65
+)
66
+from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402
67
+ MambaStateShapeCalculator,
68
+)
69
+from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402
70
+ GDNAttentionMetadataBuilder,
71
+)
72
+from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402
73
+
74
+# Small GDN dims; head_k_dim/head_v_dim=128 keeps the chunk/update kernels happy.
75
+H = 4 # num key heads
76
+HV = 8 # num value heads
77
+K = 128 # head_k_dim
78
+V = 128 # head_v_dim
79
+CONV_KERNEL = 4
80
+KEY_DIM = H * K
81
+VALUE_DIM = HV * V
82
+CONV_DIM = 2 * KEY_DIM + VALUE_DIM
83
+BLOCK_SIZE = 16
84
+PREFIX = "model.layers.0.linear_attn"
85
+
86
+
87
+def _make_vllm_config():
88
+ # A small, ungated GDN model whose config is cached locally; only the config
89
+ # (scheduler/cache/compilation/hf) is used here, never the weights. Inject
90
+ # linear_key_head_dim=128 and request the CuteDSL prefill backend -- the
91
+ # supported GDN chunk kernel on Blackwell (the Triton/FLA chunk kernel is
92
+ # unsupported on SM10x). CuteDSL consumes chunk_indices/chunk_offsets, so
93
+ # this also exercises the prefill-only chunk-metadata wiring.
94
+ cfg = create_vllm_config(
95
+ model_name="Qwen/Qwen3.5-0.8B",
96
+ block_size=BLOCK_SIZE,
97
+ hf_config_override={"linear_key_head_dim": K},
98
+ )
99
+ cfg.additional_config = {"gdn_prefill_backend": "cutedsl"}
100
+ return cfg
101
+
102
+
103
+def _build_layer(
104
+ vllm_config, conv_state, ssm_state, A_log, dt_bias, conv_weight, conv_bias
105
+):
106
+ """A minimal object that runs the real ``_forward_core`` bound to it."""
107
+ layer = types.SimpleNamespace()
108
+ layer.prefix = PREFIX
109
+ layer.enable_packed_recurrent_decode = False
110
+ layer.tp_size = 1
111
+ layer.num_k_heads = H
112
+ layer.num_v_heads = HV
113
+ layer.head_k_dim = K
114
+ layer.head_v_dim = V
115
+ layer.key_dim = KEY_DIM
116
+ layer.value_dim = VALUE_DIM
117
+ layer.activation = "silu"
118
+ layer.A_log = A_log
119
+ layer.dt_bias = dt_bias
120
+ layer.conv1d = types.SimpleNamespace(weight=conv_weight, bias=conv_bias)
121
+ layer.kv_cache = (conv_state, ssm_state)
122
+ with set_current_vllm_config(vllm_config):
123
+ layer.chunk_gated_delta_rule = ChunkGatedDeltaRule()
124
+ for name in (
125
+ "rearrange_mixed_qkv",
126
+ "_forward_core",
127
+ ):
128
+ setattr(
129
+ layer,
130
+ name,
131
+ types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer),
132
+ )
133
+ return layer
134
+
135
+
136
+def _run_forward_core(layer, meta, mixed_qkv, b, a, num_tokens):
137
+ core_attn_out = torch.zeros(
138
+ num_tokens, HV, V, dtype=mixed_qkv.dtype, device=mixed_qkv.device
139
+ )
140
+ ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta})
141
+ with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx):
142
+ layer._forward_core(
143
+ mixed_qkv=mixed_qkv.clone(),
144
+ b=b.clone(),
145
+ a=a.clone(),
146
+ core_attn_out=core_attn_out,
147
+ )
148
+ return core_attn_out
149
+
150
+
151
+@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
152
+@pytest.mark.parametrize("num_decodes,prefill_lens", [(3, [512, 300]), (4, [64, 5])])
153
+@pytest.mark.parametrize("fresh_prefill", [False, True])
154
+def test_forward_core_split_matches_unified(
155
+ state_dtype: torch.dtype,
156
+ num_decodes: int,
157
+ prefill_lens: list[int],
158
+ fresh_prefill: bool,
159
+) -> None:
160
+ torch.manual_seed(0)
161
+ device = torch.device("cuda")
162
+ vllm_config = _make_vllm_config()
163
+
164
+ # Decode-first batch: D 1-token decodes (with context), then the prefills.
165
+ decode_seq_lens = [64] * num_decodes
166
+ prefill_seq_lens = [
167
+ pl if (fresh_prefill and i == 0) else pl + 37
168
+ for i, pl in enumerate(prefill_lens)
169
+ ]
170
+ seq_lens = decode_seq_lens + prefill_seq_lens
171
+ query_lens = [1] * num_decodes + list(prefill_lens)
172
+ batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
173
+
174
+ builder = GDNAttentionMetadataBuilder(
175
+ kv_cache_spec=MambaSpec(
176
+ block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,)
177
+ ),
178
+ layer_names=[PREFIX],
179
+ vllm_config=vllm_config,
180
+ device=device,
181
+ )
182
+ common = create_common_attn_metadata(
183
+ batch, BLOCK_SIZE, device, arange_block_indices=True
184
+ )
185
+ with set_current_vllm_config(vllm_config):
186
+ meta_split = builder.build(common_prefix_len=0, common_attn_metadata=common)
187
+
188
+ assert meta_split.spec_sequence_masks is None
189
+ assert meta_split.num_decodes == num_decodes
190
+ assert meta_split.num_prefills == len(prefill_lens)
191
+ assert meta_split.num_decode_tokens == num_decodes
192
+ assert builder.gdn_prefill_backend == "cutedsl"
193
+
194
+ num_tokens = sum(query_lens)
195
+
196
+ # Full-batch chunk metadata for the unified reference path, built the same
197
+ # way the builder would for a non-split batch (backend-matched).
198
+ cu_full = meta_split.non_spec_query_start_loc
199
+ if builder.gdn_prefill_backend == "cutedsl":
200
+ from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
201
+ prepare_metadata_cutedsl,
202
+ )
203
+
204
+ full_ci, full_co = prepare_metadata_cutedsl(
205
+ cu_full, int(cu_full[-1].item()), FLA_CHUNK_SIZE
206
+ )
207
+ else:
208
+ cu_full_cpu = cu_full.cpu()
209
+ full_ci = prepare_chunk_indices(cu_full_cpu, FLA_CHUNK_SIZE).to(device)
210
+ full_co = prepare_chunk_offsets(cu_full_cpu, FLA_CHUNK_SIZE).to(device)
211
+ meta_unified = dataclasses.replace(
212
+ meta_split,
213
+ num_decodes=0,
214
+ num_decode_tokens=0,
215
+ num_prefills=meta_split.num_decodes + meta_split.num_prefills,
216
+ num_prefill_tokens=(
217
+ meta_split.num_decode_tokens + meta_split.num_prefill_tokens
218
+ ),
219
+ chunk_indices=full_ci,
220
+ chunk_offsets=full_co,
221
+ # Unified path: the chunk kernel processes the full non-spec batch.
222
+ prefill_query_start_loc=meta_split.non_spec_query_start_loc,
223
+ prefill_state_indices=meta_split.non_spec_state_indices_tensor,
224
+ prefill_has_initial_state=meta_split.has_initial_state,
225
+ )
226
+
227
+ # Size the state pools from the indices the builder actually produced.
228
+ pool_size = int(meta_split.non_spec_state_indices_tensor.max().item()) + 1
229
+ conv_state_shape, temporal_state_shape = (
230
+ MambaStateShapeCalculator.gated_delta_net_state_shape(
231
+ 1, H, HV, K, V, CONV_KERNEL, num_spec=0
232
+ )
233
+ )
234
+ conv_state0 = (
235
+ torch.randn(pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device)
236
+ * 0.05
237
+ )
238
+ ssm_state0 = (
239
+ torch.randn(pool_size, *temporal_state_shape, dtype=state_dtype, device=device)
240
+ * 0.05
241
+ )
242
+
243
+ A_log = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
244
+ dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
245
+ conv_weight = (
246
+ torch.randn(CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device) * 0.1
247
+ )
248
+ conv_bias = torch.randn(CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1
249
+
250
+ mixed_qkv = (
251
+ torch.randn(num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device) * 0.1
252
+ )
253
+ a = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1
254
+ b = torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) * 0.1
255
+
256
+ # ---- Split path (real _forward_core, meta_split) ----
257
+ conv_state_split = conv_state0.clone()
258
+ ssm_state_split = ssm_state0.clone()
259
+ layer_split = _build_layer(
260
+ vllm_config,
261
+ conv_state_split,
262
+ ssm_state_split,
263
+ A_log,
264
+ dt_bias,
265
+ conv_weight,
266
+ conv_bias,
267
+ )
268
+ out_split = _run_forward_core(layer_split, meta_split, mixed_qkv, b, a, num_tokens)
269
+
270
+ # ---- Unified path (real _forward_core, meta_unified) ----
271
+ conv_state_unified = conv_state0.clone()
272
+ ssm_state_unified = ssm_state0.clone()
273
+ layer_unified = _build_layer(
274
+ vllm_config,
275
+ conv_state_unified,
276
+ ssm_state_unified,
277
+ A_log,
278
+ dt_bias,
279
+ conv_weight,
280
+ conv_bias,
281
+ )
282
+ out_unified = _run_forward_core(
283
+ layer_unified, meta_unified, mixed_qkv, b, a, num_tokens
284
+ )
285
+
286
+ # Conv is unified in both paths, so the conv-state update must be identical.
287
+ torch.testing.assert_close(conv_state_split, conv_state_unified, atol=0, rtol=0)
288
+
289
+ # Chunk vs. recurrent update accumulate in different orders; mirror the
290
+ # tolerances used by the kernel-level parity test.
291
+ if state_dtype == torch.float32:
292
+ atol = rtol = 2e-2
293
+ else:
294
+ atol = rtol = 6e-2
295
+ torch.testing.assert_close(out_split, out_unified, atol=atol, rtol=rtol)
296
+ 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
@@ -66,7 +66,7 @@ from vllm.utils.torch_utils import (
66
)
67
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
68
69
-# Optional ROCm AITER Triton kernels for the GDN decode fast-path.
69
+# Optional ROCm AITER Triton kernels for the GDN decode path.
70
# Availability is checked centrally via rocm_aiter_ops; the actual function
71
# references are imported here so that they can be called without per-call
72
# import overhead.
@@ -897,8 +897,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
897
projected_states_ba,
898
z,
899
core_attn_out,
900
- fast_kernel=True,
900
layer_name=_encode_layer_name(self.prefix),
901
+ use_aiter=True,
902
)
903
904
self._output_projection(core_attn_out, z, output, num_tokens)
@@ -958,7 +958,6 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
958
b,
959
a,
960
core_attn_out,
961
- fast_kernel=False,
961
layer_name=_encode_layer_name(self.prefix),
962
)
963
@@ -1206,7 +1205,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1205
qkvz/ba layout.
1206
1207
For decode-only (no spec, no prefill) interleaved-GQA layouts,
1209
- dispatches directly to ``_forward_core_decode_fast``. Otherwise unpacks
1208
+ dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks
1209
the packed layout and falls through to ``_forward_core``.
1210
1211
Args:
@@ -1237,7 +1236,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1236
and attn_metadata.num_prefills == 0
1237
and attn_metadata.num_decodes > 0
1238
):
1240
- return self._forward_core_decode_fast(
1239
+ return self._forward_core_decode_aiter(
1240
qkvz=qkvz,
1241
ba=ba,
1242
z_out=z_out,
@@ -1391,6 +1390,15 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1390
mixed_qkv_non_spec = None
1391
1392
query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
1393
+
1394
+ # Split mixed non-spec-decode+prefill to process independently
1395
+ split_non_spec = (
1396
+ spec_sequence_masks is None
1397
+ and attn_metadata.num_prefills > 0
1398
+ and attn_metadata.num_decodes > 0
1399
+ )
1400
+ num_decode_tokens = attn_metadata.num_decode_tokens
1401
+
1402
if attn_metadata.num_prefills > 0:
1403
assert mixed_qkv_non_spec is not None, (
1404
"mixed_qkv_non_spec must be provided for prefill path"
@@ -1402,6 +1410,15 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1410
a_non_spec = a
1411
b_non_spec = b
1412
1413
+ if split_non_spec:
1414
+ conv_output_prefill = mixed_qkv_non_spec[num_decode_tokens:]
1415
+ a_prefill = a_non_spec[num_decode_tokens:]
1416
+ b_prefill = b_non_spec[num_decode_tokens:]
1417
+ else:
1418
+ conv_output_prefill = mixed_qkv_non_spec
1419
+ a_prefill = a_non_spec
1420
+ b_prefill = b_non_spec
1421
+
1422
(
1423
query_non_spec,
1424
key_non_spec,
@@ -1409,9 +1426,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1426
g_non_spec,
1427
beta_non_spec,
1428
) = fused_post_conv_prep(
1412
- conv_output=mixed_qkv_non_spec,
1413
- a=a_non_spec,
1414
- b=b_non_spec,
1429
+ conv_output=conv_output_prefill,
1430
+ a=a_prefill,
1431
+ b=b_prefill,
1432
A_log=self.A_log,
1433
dt_bias=self.dt_bias,
1434
num_k_heads=self.num_k_heads // self.tp_size,
@@ -1459,12 +1476,42 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1476
else:
1477
core_attn_out_spec, last_recurrent_state = None, None
1478
1462
- # 2.2: Process the remaining part
1479
+ # 2.2: Process non-spec-decode part
1480
+ if split_non_spec:
1481
+ query_decode, key_decode, value_decode = self.rearrange_mixed_qkv(
1482
+ mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index]
1483
+ )
1484
+ core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update(
1485
+ A_log=self.A_log,
1486
+ a=a[:num_decode_tokens],
1487
+ b=b[:num_decode_tokens],
1488
+ dt_bias=self.dt_bias,
1489
+ q=query_decode,
1490
+ k=key_decode,
1491
+ v=value_decode,
1492
+ initial_state=ssm_state,
1493
+ inplace_final_state=True,
1494
+ cu_seqlens=non_spec_query_start_loc[ # type: ignore[index]
1495
+ : attn_metadata.num_decodes + 1
1496
+ ],
1497
+ ssm_state_indices=non_spec_state_indices_tensor,
1498
+ use_qk_l2norm_in_kernel=True,
1499
+ )
1500
+ else:
1501
+ core_attn_out_decode = None
1502
+
1503
+ # 2.3: Process the remaining part (prefill chunk, or non-spec decode-only)
1504
if attn_metadata.num_prefills > 0:
1464
- assert non_spec_state_indices_tensor is not None
1465
- initial_state = ssm_state[non_spec_state_indices_tensor].contiguous() # type: ignore[index]
1466
- assert has_initial_state is not None
1467
- initial_state[~has_initial_state, ...] = 0 # type: ignore[operator]
1505
+ # State indices, initial-state mask and cu_seqlens for the chunk
1506
+ # kernel are precomputed by the metadata builder (the prefill tail
1507
+ # when decodes are peeled off, else the full non-spec batch), so they
1508
+ # don't need to be re-derived per layer.
1509
+ prefill_state_indices = attn_metadata.prefill_state_indices
1510
+ prefill_has_initial_state = attn_metadata.prefill_has_initial_state
1511
+ assert prefill_state_indices is not None
1512
+ assert prefill_has_initial_state is not None
1513
+ initial_state = ssm_state[prefill_state_indices]
1514
+ initial_state[~prefill_has_initial_state, ...] = 0
1515
(
1516
core_attn_out_non_spec,
1517
last_recurrent_state,
@@ -1476,15 +1523,20 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1523
beta=beta_non_spec,
1524
initial_state=initial_state,
1525
output_final_state=True,
1479
- cu_seqlens=non_spec_query_start_loc,
1526
+ cu_seqlens=attn_metadata.prefill_query_start_loc,
1527
chunk_indices=attn_metadata.chunk_indices,
1528
chunk_offsets=attn_metadata.chunk_offsets,
1529
use_qk_l2norm_in_kernel=False,
1530
)
1531
# Init cache
1485
- ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
1486
- ssm_state.dtype
1487
- )
1532
+ ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
1533
+
1534
+ if split_non_spec:
1535
+ # Stitch the peeled decode outputs in front of the prefill
1536
+ # outputs (decode-first order).
1537
+ core_attn_out_non_spec = torch.cat(
1538
+ [core_attn_out_decode, core_attn_out_non_spec], dim=1
1539
+ )
1540
elif attn_metadata.num_decodes > 0:
1541
core_attn_out_non_spec, last_recurrent_state = (
1542
fused_sigmoid_gating_delta_rule_update(
@@ -1523,7 +1575,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1575
else:
1576
core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
1577
1526
- def _forward_core_decode_fast(
1578
+ def _forward_core_decode_aiter(
1579
self,
1580
qkvz: torch.Tensor,
1581
ba: torch.Tensor,
@@ -1649,17 +1701,17 @@ def qwen_gdn_attention_core(
1701
b_or_ba: torch.Tensor,
1702
a_or_z_out: torch.Tensor,
1703
core_attn_out: torch.Tensor,
1652
- fast_kernel: bool,
1704
layer_name: LayerNameType,
1705
+ use_aiter: bool = False,
1706
) -> None:
1707
"""Custom op dispatching to _forward_core or _forward_core_rocm.
1708
1709
Handles conv1d + recurrent attention only; input/output projections
1710
are performed by the caller.
1711
1660
- When ``fast_kernel=False`` (standard path):
1712
+ When ``use_aiter=False`` (standard path):
1713
qkv_or_qkvz is [q, k, v], b_or_ba is b, a_or_z_out is a (read-only).
1662
- When ``fast_kernel=True`` (AITER Triton fast path, ROCm only):
1714
+ When ``use_aiter=True`` (AITER Triton path, ROCm only):
1715
qkv_or_qkvz is [q, k, v, z], b_or_ba is [b, a], a_or_z_out is the
1716
z output buffer (mutated in-place).
1717
@@ -1668,7 +1720,7 @@ def qwen_gdn_attention_core(
1720
layer_name = _resolve_layer_name(layer_name)
1721
forward_context: ForwardContext = get_forward_context()
1722
self = forward_context.no_compile_layers[layer_name]
1671
- if fast_kernel:
1723
+ if use_aiter:
1724
self._forward_core_rocm(
1725
qkvz=qkv_or_qkvz,
1726
ba=b_or_ba,
@@ -1689,8 +1741,8 @@ def gdn_attention_core_fake(
1741
b_or_ba: torch.Tensor,
1742
a_or_z_out: torch.Tensor,
1743
core_attn_out: torch.Tensor,
1692
- fast_kernel: bool,
1744
layer_name: LayerNameType,
1745
+ use_aiter: bool = False,
1746
) -> None:
1747
"""Fake implementation for torch.compile."""
1748
return
vllm/v1/attention/backends/gdn_attn.py
+41
-7
@@ -67,6 +67,10 @@ class GDNAttentionMetadata:
67
# Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices)
68
chunk_indices: torch.Tensor | None = None
69
chunk_offsets: torch.Tensor | None = None
70
+ # Chunk-kernel inputs for prefill
71
+ prefill_query_start_loc: torch.Tensor | None = None
72
+ prefill_state_indices: torch.Tensor | None = None
73
+ prefill_has_initial_state: torch.Tensor | None = None
74
75
# The following attributes are for triton implementation of causal_conv1d
76
nums_dict: dict | None = None
@@ -322,19 +326,42 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
326
327
chunk_indices: torch.Tensor | None = None
328
chunk_offsets: torch.Tensor | None = None
329
+ prefill_query_start_loc: torch.Tensor | None = None
330
+ prefill_state_indices: torch.Tensor | None = None
331
+ prefill_has_initial_state: torch.Tensor | None = None
332
if num_prefills > 0:
333
from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
334
335
+ # In a mixed non-spec batch, decodes are peeled off to the recurrent
336
+ # kernel (decode-first front slice), so build chunk metadata from the
337
+ # rebased prefill-only cu_seqlens; otherwise use the full non-spec one.
338
+ # _forward_core keys off the same condition, so they agree.
339
+ if spec_sequence_masks is None and num_decodes > 0:
340
+ assert non_spec_query_start_loc is not None
341
+ assert non_spec_query_start_loc_cpu is not None
342
+ assert non_spec_state_indices_tensor is not None
343
+ prefill_query_start_loc = (
344
+ non_spec_query_start_loc[num_decodes:] - num_decode_tokens
345
+ )
346
+ prefill_query_start_loc_cpu = (
347
+ non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens
348
+ )
349
+ prefill_state_indices = non_spec_state_indices_tensor[num_decodes:]
350
+ else:
351
+ prefill_query_start_loc = non_spec_query_start_loc
352
+ prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu
353
+ prefill_state_indices = non_spec_state_indices_tensor
354
+
355
if self.gdn_prefill_backend == "cutedsl":
356
from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
357
prepare_metadata_cutedsl,
358
)
359
333
- assert non_spec_query_start_loc is not None
334
- assert non_spec_query_start_loc_cpu is not None
335
- total_tokens = int(non_spec_query_start_loc_cpu[-1].item())
360
+ assert prefill_query_start_loc is not None
361
+ assert prefill_query_start_loc_cpu is not None
362
+ total_tokens = int(prefill_query_start_loc_cpu[-1].item())
363
chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
337
- non_spec_query_start_loc,
364
+ prefill_query_start_loc,
365
total_tokens,
366
FLA_CHUNK_SIZE,
367
)
@@ -348,12 +375,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
375
prepare_chunk_offsets,
376
)
377
351
- assert non_spec_query_start_loc_cpu is not None
378
+ assert prefill_query_start_loc_cpu is not None
379
chunk_indices = prepare_chunk_indices(
353
- non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
380
+ prefill_query_start_loc_cpu, FLA_CHUNK_SIZE
381
).to(device=gpu_device, non_blocking=True)
382
chunk_offsets = prepare_chunk_offsets(
356
- non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
383
+ prefill_query_start_loc_cpu, FLA_CHUNK_SIZE
384
).to(device=gpu_device, non_blocking=True)
385
386
if num_prefills > 0:
@@ -367,6 +394,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
394
device=query_start_loc.device,
395
)
396
)
397
+ if spec_sequence_masks is None and num_decodes > 0:
398
+ prefill_has_initial_state = has_initial_state[num_decodes:]
399
+ else:
400
+ prefill_has_initial_state = has_initial_state
401
else:
402
has_initial_state = None
403
@@ -458,6 +489,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
489
has_initial_state=has_initial_state,
490
chunk_indices=chunk_indices,
491
chunk_offsets=chunk_offsets,
492
+ prefill_query_start_loc=prefill_query_start_loc,
493
+ prefill_state_indices=prefill_state_indices,
494
+ prefill_has_initial_state=prefill_has_initial_state,
495
spec_query_start_loc=spec_query_start_loc,
496
non_spec_query_start_loc=non_spec_query_start_loc,
497
spec_state_indices_tensor=spec_state_indices_tensor,