[Qwen3.8-Flash-Next] Fuse main QK-norm/RoPE/gate and KV-cache write into the QSA pre-indexer launch (#57097)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Signed-off-by: Thien <thien@inferact.ai> Co-authored-by: Thien <thien@inferact.ai>
Shuolei Wang committed
Oct 1, 2026 at 14:21 UTC
c9578f10ab3eaf252742960db3fd08f11cf6d25a
6 files changed
+325
-47
.buildkite/test_areas/models_basic.yaml
+1
-1
@@ -178,7 +178,7 @@ steps:
178
- vllm/v1/attention/backends/short_conv_attn.py
179
- tests/models/qwen4_exp/
180
commands:
181
- - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_pre_indexer.py models/qwen4_exp/test_qsa_reference.py
181
+ - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py
182
mirror:
183
amd:
184
label: ":amd: (MI355 DPX) Qwen4 Exp"
tests/models/qwen4_exp/test_qsa_prepare.py
renamed
+87
-6
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Correctness tests for the fused QSA pre-indexer."""
3
+"""Correctness tests for the fused QSA prepare kernel."""
4
5
import pytest
6
import torch
@@ -15,9 +15,7 @@ from vllm.models.qwen4_exp.nvidia.ops.qsa import (
15
qsa_compress_groups_with_ratio,
16
qsa_store_cache_rows,
17
)
18
-from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import (
19
- qsa_pre_indexer,
20
-)
18
+from vllm.models.qwen4_exp.nvidia.ops.qsa_prepare import qsa_prepare
19
from vllm.platforms import current_platform
20
from vllm.triton_utils import HAS_TRITON
21
@@ -36,6 +34,7 @@ ROPE_POS_OFFSET = D
34
RTOL = 1.6e-2
35
ATOL = 1e-2
36
MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8])
37
+MAIN_HQ, MAIN_HK, MAIN_D, MAIN_PAGE = 6, 2, 256, 16
38
39
40
def _make_block_table(block_counts):
@@ -61,6 +60,84 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N
60
assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all())
61
62
63
+def _make_main_inputs(num_tokens: int, fp8_cache: bool) -> dict:
64
+ """Random main-attention arguments with a paged cache in vLLM's layout."""
65
+ num_blocks = num_tokens // MAIN_PAGE + 2
66
+ cache = torch.zeros(
67
+ num_blocks,
68
+ MAIN_HK,
69
+ MAIN_PAGE,
70
+ 2 * MAIN_D,
71
+ dtype=torch.uint8 if fp8_cache else torch.bfloat16,
72
+ device="cuda",
73
+ ).transpose(1, 2)
74
+ slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens]
75
+ slots[-1] = -1
76
+ norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2
77
+ return dict(
78
+ main_qkv=torch.randn(
79
+ num_tokens,
80
+ 2 * (MAIN_HQ + MAIN_HK) * MAIN_D,
81
+ dtype=torch.bfloat16,
82
+ device="cuda",
83
+ ),
84
+ main_q_norm_weight=norm_weights[0],
85
+ main_k_norm_weight=norm_weights[1],
86
+ main_eps=EPS,
87
+ main_kv_cache=cache.view(torch.float8_e4m3fn) if fp8_cache else cache,
88
+ main_slot_mapping=slots,
89
+ main_k_scale=0.5,
90
+ main_v_scale=2.0,
91
+ )
92
+
93
+
94
+def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None:
95
+ """Compare with fused_qk_rmsnorm_rope_gate + reshape_and_cache_flash."""
96
+ from vllm._custom_ops import reshape_and_cache_flash
97
+ from vllm.model_executor.layers.fused_qk_norm_rope import (
98
+ fused_qk_rmsnorm_rope_gate,
99
+ )
100
+
101
+ q_gate, k, v = main["main_qkv"].split(
102
+ [2 * MAIN_HQ * MAIN_D, MAIN_HK * MAIN_D, MAIN_HK * MAIN_D], dim=-1
103
+ )
104
+ q, k, gate = fused_qk_rmsnorm_rope_gate(
105
+ q_gate,
106
+ k,
107
+ main["main_q_norm_weight"],
108
+ main["main_k_norm_weight"],
109
+ rope.cos_sin_cache,
110
+ positions,
111
+ EPS,
112
+ MAIN_HQ,
113
+ MAIN_HK,
114
+ MAIN_D,
115
+ rope.rotary_dim,
116
+ mrope_section=MROPE_SECTION if positions.ndim == 2 else None,
117
+ norm_beta=1.0,
118
+ )
119
+ kv_cache = main["main_kv_cache"]
120
+ fp8_cache = kv_cache.dtype == torch.float8_e4m3fn
121
+ cache = torch.zeros_like(kv_cache.view(torch.uint8) if fp8_cache else kv_cache)
122
+ key_cache, value_cache = cache.split(MAIN_D, dim=-1)
123
+ reshape_and_cache_flash(
124
+ k.view(-1, MAIN_HK, MAIN_D),
125
+ v.view(-1, MAIN_HK, MAIN_D),
126
+ key_cache,
127
+ value_cache,
128
+ main["main_slot_mapping"],
129
+ "fp8" if fp8_cache else "auto",
130
+ torch.tensor(main["main_k_scale"], device="cuda"),
131
+ torch.tensor(main["main_v_scale"], device="cuda"),
132
+ )
133
+ torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL)
134
+ assert torch.equal(gate_out.flatten(1), gate)
135
+ if fp8_cache:
136
+ assert_fp8_within_one_ulp(kv_cache, cache.view(torch.float8_e4m3fn))
137
+ else:
138
+ torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL)
139
+
140
+
141
@requires_qsa_kernels
142
@pytest.mark.usefixtures("default_vllm_config")
143
@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn])
@@ -86,7 +163,7 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N
163
pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"),
164
],
165
)
89
-def test_qsa_fused_pre_indexer_matches_unfused(
166
+def test_qsa_fused_prepare_matches_unfused(
167
indexer_dtype,
168
mrope,
169
is_2d_positions,
@@ -246,7 +323,9 @@ def test_qsa_fused_pre_indexer_matches_unfused(
323
k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
324
325
fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device)
249
- qsa_pre_indexer(
326
+ # The main-attention cache follows the recipe's pairing with the indexer.
327
+ main = _make_main_inputs(num_tokens, indexer_dtype == torch.float8_e4m3fn)
328
+ main_q_out, main_gate_out = qsa_prepare(
329
projected_qk[:, : HQ * D],
330
projected_qk[:, HQ * D :],
331
positions,
@@ -266,7 +345,9 @@ def test_qsa_fused_pre_indexer_matches_unfused(
345
compress_ratio=CR,
346
mrope_section=MROPE_SECTION if mrope else None,
347
rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None,
348
+ **main,
349
)
350
+ _check_main_outputs(main, main_q_out, main_gate_out, rope, positions)
351
352
unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D)
353
unfused_query = gemma_rmsnorm(
tests/models/qwen4_exp/test_qsa_reference.py
+15
-4
@@ -56,16 +56,24 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
56
rope_position_offset=0,
57
),
58
compressed_key_cache=SimpleNamespace(kv_cache=torch.empty(0)),
59
- use_fused_pre_indexer=True,
59
rotary_emb=SimpleNamespace(cos_sin_cache=torch.empty(0)),
60
q_layernorm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6),
61
k_layernorm=SimpleNamespace(weight=torch.ones(1)),
62
compress_ratio=2,
63
)
64
+ attn = SimpleNamespace(
65
+ use_fused_qsa_prepare=True,
66
+ kv_cache=torch.empty(0, 1, 1, 2),
67
+ kv_cache_dtype="auto",
68
+ q_norm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6),
69
+ k_norm=SimpleNamespace(weight=torch.ones(1)),
70
+ _k_scale_float=1.0,
71
+ _v_scale_float=1.0,
72
+ )
73
74
monkeypatch.setattr(
75
indexer_qsa,
68
- "qsa_pre_indexer",
76
+ "qsa_prepare",
77
lambda *args, **kwargs: updates.append((args, kwargs)),
78
)
79
monkeypatch.setattr(
@@ -79,11 +87,14 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
87
lambda *args, **kwargs: selections.append((args, kwargs)),
88
)
89
82
- actual = indexer_qsa.QSAIndexer.forward(
90
+ actual, _ = indexer_qsa.QSAIndexer.forward(
91
indexer,
92
torch.zeros(2, 2),
93
torch.tensor([7, 8]),
94
rows,
95
+ attn=attn,
96
+ qkv=torch.zeros(2, 4),
97
+ slot_mapping=torch.arange(2),
98
)
99
100
assert actual is rows
@@ -512,7 +523,6 @@ def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
523
_metadata=lambda: (raw_metadata, compressed_metadata),
524
skip_topk=True,
525
index_kv_heads=1,
515
- use_fused_pre_indexer=False,
526
index_n_heads=1,
527
index_head_dim=64,
528
indexer_dtype=torch.bfloat16,
@@ -535,6 +545,7 @@ def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
545
torch.cat((torch.ones_like(padded_keys), padded_keys), dim=-1),
546
torch.zeros(8, dtype=torch.long, device=device),
547
torch.full((5, 5), -1, dtype=torch.int32, device=device),
548
+ attn=SimpleNamespace(use_fused_qsa_prepare=False),
549
)
550
torch.testing.assert_close(raw_cache[0, :, 0], keys[[4, 1, 2, 3]])
551
expected_compressed = torch.zeros_like(compressed_cache)
vllm/models/qwen4_exp/nvidia/indexer_qsa.py
+36
-10
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Qwen4Exp weight-free QSA indexer."""
4
5
-from typing import cast
5
+from typing import TYPE_CHECKING, cast
6
7
import torch
8
from torch import nn
@@ -23,7 +23,10 @@ from ..common.qsa_cache import (
23
QSAKeyStateCache,
24
canonical_qsa_rope_positions,
25
)
26
-from .ops.qsa_pre_indexer import qsa_pre_indexer
26
+from .ops.qsa_prepare import qsa_prepare
27
+
28
+if TYPE_CHECKING:
29
+ from .qsa import Qwen4ExpQSAAttention
30
31
32
def apply_qsa_rope(
@@ -235,7 +238,11 @@ class QSAIndexer(nn.Module):
238
projected_qk: torch.Tensor,
239
positions: torch.Tensor,
240
out: torch.Tensor | None = None,
238
- ) -> torch.Tensor:
241
+ *,
242
+ attn: "Qwen4ExpQSAAttention",
243
+ qkv: torch.Tensor | None = None,
244
+ slot_mapping: torch.Tensor | None = None,
245
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
246
"""Update side caches and select token indices from pre-projected Q/K.
247
248
Returns the packed buffer of shape [num_tokens, output_width + 1]:
@@ -243,12 +250,17 @@ class QSAIndexer(nn.Module):
250
request-relative token indices, and the trailing column is the row's
251
valid-entry count (the attention kernel's loop bound, never a token
252
index).
253
+
254
+ With ``attn.use_fused_qsa_prepare``, the same launch also writes
255
+ ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares
256
+ its Q and gate from ``qkv``, returned as the second element (None
257
+ otherwise). ``qkv`` and ``slot_mapping`` are only read in that mode.
258
"""
259
metadata = self._metadata()
260
if metadata is None:
261
# Preserve step-0 indices when later MTP steps reuse the buffer.
262
if self.skip_topk and out is not None:
251
- return out
263
+ return out, None
264
result = torch.full(
265
(projected_qk.shape[0], self.packed_output_width),
266
-1,
@@ -259,8 +271,8 @@ class QSAIndexer(nn.Module):
271
result[:, -1] = 0
272
if out is not None:
273
out.copy_(result)
262
- return out
263
- return result
274
+ return out, None
275
+ return result, None
276
277
from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows
278
from .ops.qsa_indexer import (
@@ -284,14 +296,20 @@ class QSAIndexer(nn.Module):
296
raw_key_state_cache = self.raw_key_cache
297
compressed_key_cache = self.compressed_key_cache.kv_cache
298
287
- if self.use_fused_pre_indexer:
299
+ main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None
300
+ if attn.use_fused_qsa_prepare:
301
+ if qkv is None or slot_mapping is None:
302
+ raise ValueError("fused QSA prepare requires qkv and slot_mapping")
303
q = projected_q.new_empty(
304
num_tokens,
305
self.index_n_heads,
306
self.index_head_dim,
307
dtype=self.indexer_dtype,
308
)
294
- qsa_pre_indexer(
309
+ main_kv_cache = attn.kv_cache.transpose(1, 2)
310
+ if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"):
311
+ main_kv_cache = main_kv_cache.view(torch.float8_e4m3fn)
312
+ main_outputs = qsa_prepare(
313
projected_q,
314
raw_keys,
315
positions,
@@ -315,6 +333,14 @@ class QSAIndexer(nn.Module):
333
if raw_key_state_cache.rope_position_cache is not None
334
else None
335
),
336
+ main_qkv=qkv[:num_tokens],
337
+ main_q_norm_weight=attn.q_norm.weight,
338
+ main_k_norm_weight=attn.k_norm.weight,
339
+ main_eps=attn.q_norm.variance_epsilon,
340
+ main_kv_cache=main_kv_cache,
341
+ main_slot_mapping=slot_mapping[:num_tokens],
342
+ main_k_scale=attn._k_scale_float,
343
+ main_v_scale=attn._v_scale_float,
344
)
345
else:
346
# Unfused reference path
@@ -385,7 +411,7 @@ class QSAIndexer(nn.Module):
411
if self.skip_topk:
412
if out is None:
413
raise RuntimeError("QSA top-k reuse requires an output buffer")
388
- return out
414
+ return out, main_outputs
415
416
if out is None:
417
out = torch.empty(
@@ -448,7 +474,7 @@ class QSAIndexer(nn.Module):
474
self.token_topk,
475
out,
476
)
451
- return out
477
+ return out, main_outputs
478
479
480
__all__ = ["QSAIndexer", "apply_qsa_rope"]
vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py
renamed
+146
-8
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Fused QSA pre-indexer kernel for Qwen4Exp."""
3
+"""Fused QSA prepare kernel for Qwen4Exp."""
4
5
import torch
6
@@ -80,6 +80,27 @@ def _norm_rope(
80
return tl.reshape(result, (TILE_T, TILE_H, D))
81
82
83
+@triton.jit
84
+def _to_dst_dtype(x, dst, scale):
85
+ """Round to BF16 like the unfused path, then scale for an FP8 destination."""
86
+ out_ty = dst.dtype.element_ty
87
+ x = x.to(tl.bfloat16)
88
+ if out_ty == tl.float8e4nv:
89
+ x = x.to(tl.float32) / scale
90
+ return x.to(out_ty)
91
+
92
+
93
+@triton.jit
94
+def _store_rotated(dst, y, o1, o2, scale):
95
+ """Store a normalized head whose first ``2 * len(o1)`` dims are rotated."""
96
+ HALF: tl.constexpr = o1.shape[0]
97
+ dims = tl.arange(0, y.shape[0])
98
+ rot = tl.arange(0, HALF)
99
+ tl.store(dst + dims, _to_dst_dtype(y, dst, scale), mask=dims >= 2 * HALF)
100
+ tl.store(dst + rot, _to_dst_dtype(o1, dst, scale))
101
+ tl.store(dst + HALF + rot, _to_dst_dtype(o2, dst, scale))
102
+
103
+
104
@triton.jit(
105
do_not_specialize=[
106
"num_tokens",
@@ -88,7 +109,7 @@ def _norm_rope(
109
"num_k_work",
110
]
111
)
91
-def _qsa_pre_indexer_kernel(
112
+def _qsa_prepare_kernel(
113
q_ptr,
114
q_stride_token,
115
k_ptr,
@@ -132,8 +153,86 @@ def _qsa_pre_indexer_kernel(
153
CACHE_HAS_ROPE_POS: tl.constexpr,
154
MROPE_H: tl.constexpr,
155
MROPE_W: tl.constexpr,
156
+ main_qkv_ptr,
157
+ main_q_norm_weight_ptr,
158
+ main_k_norm_weight_ptr,
159
+ main_eps,
160
+ main_q_out_ptr,
161
+ main_gate_out_ptr,
162
+ main_cache_ptr,
163
+ main_cache_stride_block,
164
+ main_cache_stride_token,
165
+ main_cache_stride_head,
166
+ main_slots_ptr,
167
+ main_k_scale,
168
+ main_v_scale,
169
+ MAIN_HQ: tl.constexpr,
170
+ MAIN_HK: tl.constexpr,
171
+ MAIN_D: tl.constexpr,
172
+ MAIN_PAGE_SIZE: tl.constexpr,
173
):
174
pid = tl.program_id(0)
175
+ num_index_work = num_k_work + tl.cdiv(num_tokens, TILE_T_Q) * tl.cdiv(HQ, TILE_H_Q)
176
+ if pid >= num_index_work:
177
+ # Main attention: one program per (token, Q or KV head) after the
178
+ # indexer work. RoPE covers the first D // 2 dims, the width of the
179
+ # shared cos/sin table.
180
+ main_pid = pid - num_index_work
181
+ token = main_pid // (MAIN_HQ + MAIN_HK)
182
+ head = main_pid % (MAIN_HQ + MAIN_HK)
183
+ HALF: tl.constexpr = D // 4
184
+ dims = tl.arange(0, MAIN_D)
185
+ rot = tl.arange(0, HALF)
186
+ row = main_qkv_ptr + token * (2 * (MAIN_HQ + MAIN_HK) * MAIN_D)
187
+ is_k = head >= MAIN_HQ
188
+ kv_head = head - MAIN_HQ
189
+ if is_k:
190
+ src = row + (2 * MAIN_HQ + kv_head) * MAIN_D
191
+ weight_ptr = main_k_norm_weight_ptr
192
+ else:
193
+ src = row + 2 * head * MAIN_D
194
+ weight_ptr = main_q_norm_weight_ptr
195
+ x = tl.load(src + dims).to(tl.float32)
196
+ inv_rms = tl.rsqrt(tl.sum(x * x, axis=0) / MAIN_D + main_eps)
197
+ w = tl.load(weight_ptr + dims).to(tl.float32) + 1.0
198
+ y = x * inv_rms * w
199
+ x1 = tl.load(src + rot).to(tl.float32)
200
+ x2 = tl.load(src + HALF + rot).to(tl.float32)
201
+ w1 = tl.load(weight_ptr + rot).to(tl.float32) + 1.0
202
+ w2 = tl.load(weight_ptr + HALF + rot).to(tl.float32) + 1.0
203
+ x1 = (x1 * inv_rms * w1).to(tl.bfloat16).to(tl.float32)
204
+ x2 = (x2 * inv_rms * w2).to(tl.bfloat16).to(tl.float32)
205
+ pos = tl.load(pos_ptr + token * pos_stride_token).to(tl.int64)
206
+ if IS_2D_POSITIONS:
207
+ pos_h = tl.load(pos_ptr + pos_stride_axis + token * pos_stride_token)
208
+ pos_w = tl.load(pos_ptr + 2 * pos_stride_axis + token * pos_stride_token)
209
+ is_h = (rot % 3 == 1) & (rot < 3 * MROPE_H)
210
+ is_w = (rot % 3 == 2) & (rot < 3 * MROPE_W)
211
+ pos = tl.where(
212
+ is_h, pos_h.to(tl.int64), tl.where(is_w, pos_w.to(tl.int64), pos)
213
+ )
214
+ cos = tl.load(cos_sin_ptr + pos * (D // 2) + rot).to(tl.float32)
215
+ sin = tl.load(cos_sin_ptr + pos * (D // 2) + HALF + rot).to(tl.float32)
216
+ o1 = x1 * cos - x2 * sin
217
+ o2 = x2 * cos + x1 * sin
218
+ if is_k:
219
+ slot = tl.load(main_slots_ptr + token).to(tl.int64)
220
+ if slot >= 0:
221
+ dst = (
222
+ main_cache_ptr
223
+ + (slot // MAIN_PAGE_SIZE) * main_cache_stride_block
224
+ + (slot % MAIN_PAGE_SIZE) * main_cache_stride_token
225
+ + kv_head * main_cache_stride_head
226
+ )
227
+ _store_rotated(dst, y, o1, o2, main_k_scale)
228
+ v = tl.load(src + MAIN_HK * MAIN_D + dims)
229
+ tl.store(dst + MAIN_D + dims, _to_dst_dtype(v, dst, main_v_scale))
230
+ else:
231
+ out = (token * MAIN_HQ + head) * MAIN_D
232
+ _store_rotated(main_q_out_ptr + out, y, o1, o2, None)
233
+ gate = tl.load(src + MAIN_D + dims)
234
+ tl.store(main_gate_out_ptr + out + dims, gate)
235
+ return
236
# K work occupies the first programs; the remaining programs tile Q. This
237
# keeps both paths in one launch while leaving their register shapes
238
# independent.
@@ -400,7 +499,7 @@ def _qsa_pre_indexer_kernel(
499
tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot)
500
501
403
-def qsa_pre_indexer(
502
+def qsa_prepare(
503
q: torch.Tensor,
504
k: torch.Tensor,
505
positions: torch.Tensor,
@@ -421,11 +520,28 @@ def qsa_pre_indexer(
520
compress_ratio: int,
521
mrope_section: tuple[int, int, int] | None,
522
rope_pos_offset: int | None,
424
-) -> None:
425
- """Normalize Q, compress K, then update the circular raw state."""
523
+ main_qkv: torch.Tensor,
524
+ main_q_norm_weight: torch.Tensor,
525
+ main_k_norm_weight: torch.Tensor,
526
+ main_eps: float,
527
+ main_kv_cache: torch.Tensor,
528
+ main_slot_mapping: torch.Tensor,
529
+ main_k_scale: float,
530
+ main_v_scale: float,
531
+) -> tuple[torch.Tensor, torch.Tensor]:
532
+ """Normalize Q, compress K, then update the circular raw state.
533
+
534
+ Also prepares the main attention (QK-norm/RoPE, gate copy, K/V cache write)
535
+ and returns its Q and gate.
536
+ """
537
num_tokens = q.shape[0]
538
+ main_head_dim = main_kv_cache.shape[-1] // 2
539
+ num_main_kv_heads = main_kv_cache.shape[2]
540
+ num_main_q_heads = main_qkv.shape[1] // (2 * main_head_dim) - num_main_kv_heads
541
+ main_q_out = main_qkv.new_empty(num_tokens, num_main_q_heads, main_head_dim)
542
+ main_gate_out = torch.empty_like(main_q_out)
543
if num_tokens == 0:
428
- return
544
+ return main_q_out, main_gate_out
545
num_q_heads, head_dim = q_out.shape[1:]
546
assert cos_sin_cache.shape[-1] * 2 == head_dim
547
assert q.shape == (num_tokens, num_q_heads * head_dim)
@@ -450,6 +566,9 @@ def qsa_pre_indexer(
566
pos_stride_axis, pos_stride_token = 0, positions.stride(0)
567
section = mrope_section if mrope_section is not None else (0, 0, 0)
568
assert len(section) == 3
569
+ qkv_width = 2 * (num_main_q_heads + num_main_kv_heads) * main_head_dim
570
+ assert main_qkv.shape == (num_tokens, qkv_width) and main_qkv.is_contiguous()
571
+ assert main_slot_mapping.shape == (num_tokens,)
572
573
if num_tokens <= 4096:
574
TILE_T_Q, TILE_H_Q = 2, 2
@@ -457,7 +576,8 @@ def qsa_pre_indexer(
576
TILE_T_Q, TILE_H_Q = 2, 4
577
num_k_work = k_work_metadata.shape[0]
578
num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q)
460
- _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)](
579
+ num_main_work = num_tokens * (num_main_q_heads + num_main_kv_heads)
580
+ _qsa_prepare_kernel[(num_k_work + num_q_work + num_main_work,)](
581
q,
582
q.stride(0),
583
k,
@@ -501,8 +621,26 @@ def qsa_pre_indexer(
621
CACHE_HAS_ROPE_POS=cache_has_rope_pos,
622
MROPE_H=section[1],
623
MROPE_W=section[2],
624
+ main_qkv_ptr=main_qkv,
625
+ main_q_norm_weight_ptr=main_q_norm_weight,
626
+ main_k_norm_weight_ptr=main_k_norm_weight,
627
+ main_eps=main_eps,
628
+ main_q_out_ptr=main_q_out,
629
+ main_gate_out_ptr=main_gate_out,
630
+ main_cache_ptr=main_kv_cache,
631
+ main_cache_stride_block=main_kv_cache.stride(0),
632
+ main_cache_stride_token=main_kv_cache.stride(1),
633
+ main_cache_stride_head=main_kv_cache.stride(2),
634
+ main_slots_ptr=main_slot_mapping,
635
+ main_k_scale=main_k_scale,
636
+ main_v_scale=main_v_scale,
637
+ MAIN_HQ=num_main_q_heads,
638
+ MAIN_HK=num_main_kv_heads,
639
+ MAIN_D=main_head_dim,
640
+ MAIN_PAGE_SIZE=main_kv_cache.shape[1],
641
num_warps=1,
642
)
643
+ return main_q_out, main_gate_out
644
645
508
-__all__ = ["qsa_pre_indexer"]
646
+__all__ = ["qsa_prepare"]
vllm/models/qwen4_exp/nvidia/qsa.py
+40
-18
@@ -407,6 +407,12 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
407
quant_config=quant_config,
408
prefix=f"{prefix}.indexer",
409
)
410
+ # One launch does the indexer prepare, the main QK-norm/RoPE/gate and
411
+ # the main K/V cache write (see QSAIndexer.forward); otherwise all of
412
+ # them take the separate kernels.
413
+ self.use_fused_qsa_prepare = (
414
+ self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer
415
+ )
416
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
417
# PACKED selection buffer: the trailing column holds each row's
418
# valid-entry count (written by the expand kernel) — never a token
@@ -446,12 +452,15 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
452
self,
453
projected_qk: torch.Tensor,
454
positions: torch.Tensor,
449
- query: torch.Tensor,
450
- key: torch.Tensor,
451
- value: torch.Tensor,
455
+ query: torch.Tensor | None,
456
+ key: torch.Tensor | None,
457
+ value: torch.Tensor | None,
458
output: torch.Tensor,
453
- output_gate: torch.Tensor,
459
+ output_gate: torch.Tensor | None,
460
+ qkv: torch.Tensor,
461
) -> None:
462
+ # query/key/value/output_gate are None when the fused prepare runs
463
+ # inside the indexer launch.
464
metadata = get_forward_context().attn_metadata
465
if isinstance(metadata, list):
466
metadata = metadata[0]
@@ -469,21 +478,29 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
478
)
479
if side_metadata.num_actual_tokens != num_tokens:
480
raise RuntimeError("QSA main and side metadata token counts disagree")
472
- selected = self.indexer(
481
+ selected, main_outputs = self.indexer(
482
projected_qk,
483
positions,
484
self.topk_indices_buffer[:num_tokens],
485
+ attn=self,
486
+ qkv=qkv,
487
+ slot_mapping=main_metadata.slot_mapping,
488
)
489
if selected.shape != (num_tokens, self.indexer.packed_output_width):
490
raise RuntimeError("QSA indexer returned an invalid selection shape")
491
impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
480
- impl.do_kv_cache_update(
481
- self,
482
- key,
483
- value,
484
- self.kv_cache,
485
- main_metadata.slot_mapping,
486
- )
492
+ if main_outputs is None:
493
+ assert key is not None and value is not None
494
+ impl.do_kv_cache_update(
495
+ self,
496
+ key,
497
+ value,
498
+ self.kv_cache,
499
+ main_metadata.slot_mapping,
500
+ )
501
+ else:
502
+ query, output_gate = main_outputs
503
+ assert query is not None and output_gate is not None
504
impl.forward_qsa(
505
self,
506
query,
@@ -503,13 +520,17 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
520
hidden_states: torch.Tensor,
521
) -> torch.Tensor:
522
qkv, _ = self.qkv_proj(hidden_states)
506
- q, k, v, gate = self._project_qkv_gate(qkv, positions)
507
- assert gate is not None
523
num_tokens = hidden_states.shape[0]
509
- query = q.view(num_tokens, self.num_heads, self.head_dim)
510
- key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
511
- value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
512
- attn_output = torch.empty_like(query)
524
+ if not self.use_fused_qsa_prepare:
525
+ q, k, v, gate = self._project_qkv_gate(qkv, positions)
526
+ assert gate is not None
527
+ query = q.view(num_tokens, self.num_heads, self.head_dim)
528
+ key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
529
+ value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
530
+ else:
531
+ # Norm/RoPE/gate and the K/V cache write happen inside _run_qsa.
532
+ query = key = value = gate = None
533
+ attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim)
534
# Keep the index projection outside the eager break.
535
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
536
self._run_qsa(
@@ -520,6 +541,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
541
value,
542
attn_output,
543
gate,
544
+ qkv,
545
)
546
flat_output = attn_output.view(num_tokens, -1)
547
output, _ = self.o_proj(flat_output)