@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 31a8a26662

[Qwen3.8-Flash-Next] Improve QSA sparse GQA for prefill and short-ctx decode (#54873)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Kimi <noreply@moonshot.cn>

Thien Tran committed Sep 4, 2026 at 21:08 UTC 31a8a266622781917cef482d01db415cfa3cbd92

6 files changed +366 -107

tests/models/qwen4_exp/test_qsa_reference.py

+65 -26

@@ -754,7 +754,10 @@ def test_qsa_block_expansion_correctness() -> None:

754 sequence_lengths.index_select(0, token_to_req.long()) // 4,

755 ).to(torch.int32)

756

757 - actual = torch.empty((2, 11), device="cuda", dtype=torch.int32)

757 + # Packed layout: one trailing column per row holds the valid-entry count

758 + # (never a token index). Row 0: 1 visible block + 2 tail; row 1: 2 blocks

759 + # + 3 tail.

760 + actual = torch.empty((2, 12), device="cuda", dtype=torch.int32)

761 qsa_indexer_ops.expand_qsa_block_indices(

762 blocks,

763 query_positions,

@@ -771,19 +774,35 @@ def test_qsa_block_expansion_correctness() -> None:

774 token_topk=8,

775 )

776

774 - torch.testing.assert_close(actual, expected)

777 + torch.testing.assert_close(actual[:, :11], expected)

778 + assert actual[:, 11].tolist() == [6, 11]

779

780

781 @requires_qsa_kernels

782 @pytest.mark.parametrize(

779 - ("num_rows", "num_query_heads", "num_kv_heads", "page_size"),

783 + (

784 + "num_rows",

785 + "num_query_heads",

786 + "num_kv_heads",

787 + "page_size",

788 + "use_prefill_config",

789 + "num_requests",

790 + ),

791 [

781 - # Kernel-visible pages with --block-size 256 and hybrid-cache alignment.

782 - pytest.param(1, 24, 2, 1792, id="tp1_split64"),

783 - pytest.param(16, 12, 1, 1792, id="tp2_split32"),

784 - pytest.param(32, 6, 1, 1024, id="tp4_split8"),

785 - pytest.param(257, 6, 1, 1024, id="tp4_split4"),

786 - pytest.param(513, 6, 1, 1024, id="tp4_split1"),

792 + # Production page sizes from hybrid-cache block alignment: 784/800

793 + # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head

794 + # splits are per-rank TP1/TP2/TP4; the largest batch runs both

795 + # use_prefill_config variants.

796 + pytest.param(1, 24, 2, 1600, True, 2, id="tp1_r1"),

797 + pytest.param(16, 12, 1, 1600, True, 3, id="tp2_r16"),

798 + pytest.param(32, 6, 1, 800, True, 5, id="tp4_r32"),

799 + pytest.param(128, 24, 2, 1568, True, 7, id="tp1_r128"),

800 + pytest.param(257, 6, 1, 800, True, 13, id="tp4_r257"),

801 + pytest.param(513, 6, 1, 784, True, 17, id="tp4_r513"),

802 + pytest.param(700, 6, 1, 800, True, 23, id="tp4_r700"),

803 + pytest.param(1024, 24, 2, 1600, True, 33, id="tp1_r1024"),

804 + pytest.param(2048, 24, 2, 1600, True, 63, id="tp1_r2048_prefill"),

805 + pytest.param(2048, 24, 2, 1600, False, 63, id="tp1_r2048_uniform"),

806 ],

807 )

808 def test_qsa_sparse_paged_attention_correctness(

@@ -791,10 +810,11 @@ def test_qsa_sparse_paged_attention_correctness(

810 num_query_heads: int,

811 num_kv_heads: int,

812 page_size: int,

813 + use_prefill_config: bool,

814 + num_requests: int,

815 ) -> None:

816 torch.manual_seed(2)

817 head_dim = 256

797 - num_requests = 2

818 num_selected_pages = 64

819 # Keep the newest page outside the synthetic top-k as causal headroom.

820 num_pages_per_request = num_selected_pages + 1

@@ -822,13 +842,22 @@ def test_qsa_sparse_paged_attention_correctness(

842 rows_per_request = math.ceil(num_rows / num_requests)

843 row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32)

844 token_to_req = row_indices // rows_per_request

825 - request_row_counts = torch.tensor(

826 - [rows_per_request, num_rows - rows_per_request],

845 + # Uniform row split; the last request takes the remainder (possibly 0).

846 + request_row_counts = torch.full(

847 + (num_requests,), rows_per_request, device="cuda", dtype=torch.int32

848 + )

849 + request_row_counts[-1] = num_rows - rows_per_request * (num_requests - 1)

850 +

851 + # Mix context lengths: every third request is short-context, attending

852 + # to only its first few pages; the rest fill their cache.

853 + context_lengths = torch.full(

854 + (num_requests,),

855 + num_pages_per_request * page_size - 1,

856 device="cuda",

857 dtype=torch.int32,

858 )

830 -

831 - context_length = num_pages_per_request * page_size - 1

859 + short_requests = torch.arange(num_requests, device="cuda") % 3 == 1

860 + context_lengths[short_requests] = request_row_counts[short_requests] + 8

861 block_topk = indexer_budget // indexer_compress_ratio

862 compressed_blocks_per_page = page_size // indexer_compress_ratio

863 selection = torch.arange(block_topk, device="cuda")

@@ -842,17 +871,19 @@ def test_qsa_sparse_paged_attention_correctness(

871 )

872 rows_within_request = row_indices % rows_per_request

873 query_positions = (

845 - context_length - request_row_counts[token_to_req.long()] + rows_within_request

874 + context_lengths[token_to_req.long()]

875 + - request_row_counts[token_to_req.long()]

876 + + rows_within_request

877 ).to(torch.int64)

847 - sequence_lengths = torch.full(

848 - (num_requests,), context_length, device="cuda", dtype=torch.int32

849 - )

878 + sequence_lengths = context_lengths

879 visible_blocks = torch.minimum(

880 (query_positions + 1) // indexer_compress_ratio,

881 sequence_lengths.index_select(0, token_to_req.long()) // indexer_compress_ratio,

882 ).to(torch.int32)

883 + # +1: the packed trailing column holds each row's valid-entry count

884 + # (never a token index); the reference reads only the selection region.

885 logical_indices = torch.empty(

855 - (num_rows, selection_width), device="cuda", dtype=torch.int32

886 + (num_rows, selection_width + 1), device="cuda", dtype=torch.int32

887 )

888 qsa_indexer_ops.expand_qsa_block_indices(

889 block_indices,

@@ -862,7 +893,7 @@ def test_qsa_sparse_paged_attention_correctness(

893 indexer_budget,

894 logical_indices,

895 )

865 - assert logical_indices.shape == (num_rows, selection_width)

896 + assert logical_indices.shape == (num_rows, selection_width + 1)

897 scale = q.shape[-1] ** -0.5

898

899 actual = qsa_ops.qsa_sparse_paged_attention(

@@ -872,12 +903,13 @@ def test_qsa_sparse_paged_attention_correctness(

903 logical_indices,

904 block_table,

905 token_to_req,

906 + use_prefill_config=use_prefill_config,

907 )

908 expected = _qsa_sparse_paged_attention_reference(

909 q,

910 k_cache,

911 v_cache,

880 - logical_indices,

912 + logical_indices[:, :selection_width],

913 block_table,

914 token_to_req,

915 scale,

@@ -948,8 +980,10 @@ def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int)

980 block_indices[prefill_slice],

981 max_seq_len=sequence_lengths.max().item(),

982 )

983 + # +1: the packed trailing count column (never a token index; excluded

984 + # from the comparison).

985 actual = torch.empty(

952 - (rows, token_topk + compress_ratio - 1), device="cuda", dtype=torch.int32

986 + (rows, token_topk + compress_ratio), device="cuda", dtype=torch.int32

987 )

988 qsa_indexer_ops.expand_qsa_block_indices(

989 block_indices,

@@ -977,7 +1011,10 @@ def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int)

1011 token_topk,

1012 )

1013

980 - torch.testing.assert_close(actual.sort().values, expected.sort().values)

1014 + torch.testing.assert_close(

1015 + actual[:, : token_topk + compress_ratio - 1].sort().values,

1016 + expected.sort().values,

1017 + )

1018

1019

1020 @requires_qsa_kernels

@@ -1001,7 +1038,7 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) ->

1038 block_indices=block_indices,

1039 max_seq_len=64, # clamps to the page-table capacity

1040 )

1004 - selected = torch.empty((2, 2051), device="cuda", dtype=torch.int32)

1041 + selected = torch.empty((2, 2052), device="cuda", dtype=torch.int32)

1042 qsa_indexer_ops.expand_qsa_block_indices(

1043 block_indices,

1044 query_positions,

@@ -1013,8 +1050,10 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) ->

1050

1051 assert selected[0, :2].tolist() == [0, 1]

1052 assert selected[1, :3].tolist() == [0, 1, 2]

1016 - assert torch.all(selected[0, 2:] == -1)

1017 - assert torch.all(selected[1, 3:] == -1)

1053 + assert torch.all(selected[0, 2:2051] == -1)

1054 + assert torch.all(selected[1, 3:2051] == -1)

1055 + # The packed trailing column holds each row's valid-entry count.

1056 + assert selected[:, 2051].tolist() == [2, 3]

1057

1058

1059 @requires_qsa_kernels

vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py

+47 -24

@@ -5,6 +5,8 @@

5 import sys

6 from typing import TYPE_CHECKING, cast

8 +import torch

9 +

10 from vllm.logger import init_logger

12 if TYPE_CHECKING:

@@ -15,37 +17,41 @@ logger = init_logger(__name__)

19 def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None:

18 - """Warm every reachable QSA decode-query-length specialization."""

20 + """Warm every reachable QSA specialization: indexer decode-query-length

21 + profiles plus the sparse attention split-K/merge configs."""

23 qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa")

21 - if qsa_module is None:

24 + attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa")

25 + if qsa_module is None or attn_module is None:

26 return

23 - indexer = next(

24 - (

25 - layer

26 - for layer in worker.get_model().modules()

27 - if isinstance(layer, qsa_module.QSAIndexer)

28 - ),

29 - None,

30 - )

31 - if indexer is None:

27 + indexer = None

28 + owner = None

29 + for layer in worker.get_model().modules():

30 + if indexer is None and isinstance(layer, qsa_module.QSAIndexer):

31 + indexer = layer

32 + elif owner is None and isinstance(layer, attn_module.Qwen4ExpQSAAttention):

33 + owner = layer

34 + if indexer is None or owner is None:

35 return

37 runner = worker.model_runner

35 - prefix = indexer.compressed_key_cache.prefix

36 - group_id = next(

37 - i

38 - for i, group in enumerate(runner.kv_cache_config.kv_cache_groups)

39 - if prefix in group.layer_names

40 - )

41 - if worker.use_v2_model_runner:

42 - runner_v2 = cast("GPUModelRunnerV2", runner)

43 - block_table = runner_v2.block_tables.input_block_tables[group_id]

44 - max_decode_query_len = runner_v2.decode_query_len

45 - else:

46 - block_table = runner.input_batch.block_table[group_id].get_device_tensor(

38 +

39 + def block_table_for(prefix: str) -> torch.Tensor:

40 + group_id = next(

41 + i

42 + for i, group in enumerate(runner.kv_cache_config.kv_cache_groups)

43 + if prefix in group.layer_names

44 + )

45 + if worker.use_v2_model_runner:

46 + runner_v2 = cast("GPUModelRunnerV2", runner)

47 + return runner_v2.block_tables.input_block_tables[group_id]

48 + return runner.input_batch.block_table[group_id].get_device_tensor(

49 runner.max_num_reqs

50 )

51 +

52 + if worker.use_v2_model_runner:

53 + max_decode_query_len = cast("GPUModelRunnerV2", runner).decode_query_len

54 + else:

55 max_decode_query_len = runner.uniform_decode_query_len

57 from vllm.models.qwen4_exp.nvidia.ops.qsa_indexer import (

@@ -56,7 +62,7 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None:

62 assert k_cache.numel()

63 profiles = warmup_qsa_mqa_paged_decode(

64 k_cache,

59 - block_table,

65 + block_table_for(indexer.compressed_key_cache.prefix),

66 num_heads=indexer.index_n_heads,

67 head_dim=indexer.index_head_dim,

68 max_decode_query_len=max_decode_query_len,

@@ -64,3 +70,20 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None:

70 max_num_batched_tokens=runner.max_num_tokens,

71 )

72 logger.info("Warmed up Qwen4Exp QSA decode kernels: %s.", profiles)

73 +

74 + from vllm.models.qwen4_exp.nvidia.ops.qsa import (

75 + warmup_qsa_sparse_paged_attention,

76 + )

77 +

78 + kv_cache = owner.kv_cache

79 + assert kv_cache.numel()

80 + attention_profiles = warmup_qsa_sparse_paged_attention(

81 + kv_cache,

82 + block_table_for(owner.layer_name),

83 + num_query_heads=owner.num_heads,

84 + selection_width=indexer.output_width,

85 + )

86 + logger.info(

87 + "Warmed up Qwen4Exp QSA sparse attention kernels: %s.",

88 + attention_profiles,

89 + )

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

+24 -4

@@ -167,8 +167,19 @@ class QSAIndexer(nn.Module):

167

168 @property

169 def output_width(self) -> int:

170 + """Selection (index) columns per row."""

171 return self.token_topk + self.compress_ratio - 1

172

173 + @property

174 + def packed_output_width(self) -> int:

175 + """Packed selection-buffer width: output_width + 1.

176 +

177 + The trailing column holds each row's valid-entry count (written by

178 + the expand kernel) — never a token index; the sparse attention

179 + kernel reads it as its tile-loop bound.

180 + """

181 + return self.output_width + 1

182 +

183 def _metadata(

184 self,

185 ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:

@@ -211,7 +222,14 @@ class QSAIndexer(nn.Module):

222 positions: torch.Tensor,

223 out: torch.Tensor | None = None,

224 ) -> torch.Tensor:

214 - """Return fixed-width request-relative token indices padded with ``-1``."""

225 + """Select each query row's token indices.

226 +

227 + Returns the packed buffer of shape [num_tokens, output_width + 1]:

228 + the leading ``output_width`` columns are ``-1``-padded

229 + request-relative token indices, and the trailing column is the row's

230 + valid-entry count (the attention kernel's loop bound, never a token

231 + index).

232 + """

233

234 metadata = self._metadata()

235 if metadata is None:

@@ -219,11 +237,13 @@ class QSAIndexer(nn.Module):

237 if self.skip_topk and out is not None:

238 return out

239 result = torch.full(

222 - (hidden_states.shape[0], self.output_width),

240 + (hidden_states.shape[0], self.packed_output_width),

241 -1,

242 dtype=torch.int32,

243 device=hidden_states.device,

244 )

245 + # Inert rows carry a zero valid count (empty loop bound), not -1.

246 + result[:, -1] = 0

247 if out is not None:

248 out.copy_(result)

249 return out

@@ -357,11 +377,11 @@ class QSAIndexer(nn.Module):

377 if out is None:

378 out = torch.empty(

379 num_tokens,

360 - self.output_width,

380 + self.packed_output_width,

381 dtype=torch.int32,

382 device=q.device,

383 )

364 - elif out.shape != (num_tokens, self.output_width):

384 + elif out.shape != (num_tokens, self.packed_output_width):

385 raise ValueError("QSA selection output has an invalid shape")

386

387 num_decode_tokens = compressed_metadata.num_decode_tokens

vllm/models/qwen4_exp/nvidia/ops/qsa.py

+202 -44

@@ -6,10 +6,14 @@ from __future__ import annotations

7 import torch

9 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

10 + TritonWarmupTensor,

11 + triton_scalar_specialization_rep,

12 +)

13 from vllm.triton_utils import HAS_TRITON, tl, triton

12 -@triton.jit

16 +@triton.jit(do_not_specialize=["num_rows", "num_requests"])

17 def _qsa_sparse_paged_gqa_splitk_kernel(

18 q_ptr,

19 k_cache_ptr,

@@ -52,6 +56,12 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

56 request = tl.load(token_to_req_ptr + row)

57 safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)

59 + # The packed selection buffer carries one TRAILING COUNT COLUMN per row

60 + # (column TOPK of a TOPK+1-wide buffer): the row's valid-entry count,

61 + # written by the expand kernel. It is never a token index — the tile loop

62 + # and the index load below only ever cover columns [0, TOPK).

63 + valid_count = tl.load(indices_ptr + row * stride_indices_row + TOPK)

64 +

65 head_offsets = tl.arange(0, BLOCK_M)

66 dim_offsets = tl.arange(0, HEAD_DIM)

67 column_offsets = tl.arange(0, BLOCK_N)

@@ -70,10 +80,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

80 accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)

81 softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634

73 - # Dynamic bounds avoid padded main-loop iterations for uneven splits.

74 - split_tile_start = split_id * NUM_TILES // NUM_SPLITS

75 - split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS

76 - for tile in range(split_tile_start, split_tile_end):

83 + tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N))

84 +

85 + for tile in range(split_id, tile_end, NUM_SPLITS):

86 columns = tile * BLOCK_N + column_offsets

87 logical_token = tl.load(

88 indices_ptr + row * stride_indices_row + columns,

@@ -156,29 +165,23 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

165 max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)),

166 -float("inf"),

167 )

168 + partial_row = (split_id * num_rows + row).to(tl.int64)

169 tl.store(

170 partial_output_ptr

161 - + (

162 - (split_id * num_rows + row) * NUM_QUERY_HEADS

163 - + first_head

164 - + head_offsets[:, None]

165 - )

171 + + (partial_row * NUM_QUERY_HEADS + first_head + head_offsets[:, None])

172 * HEAD_DIM

173 + dim_offsets[None, :],

174 normalized_output,

175 mask=output_mask,

176 )

177 tl.store(

172 - partial_lse_ptr

173 - + (split_id * num_rows + row) * NUM_QUERY_HEADS

174 - + first_head

175 - + head_offsets,

178 + partial_lse_ptr + partial_row * NUM_QUERY_HEADS + first_head + head_offsets,

179 partial_lse,

180 mask=head_offsets < GROUP_SIZE,

181 )

182

183

181 -@triton.jit

184 +@triton.jit(do_not_specialize=["num_rows"])

185 def _qsa_merge_splitk_kernel(

186 partial_output_ptr,

187 partial_lse_ptr,

@@ -206,10 +209,10 @@ def _qsa_merge_splitk_kernel(

209 shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf"))

210 weights = tl.math.exp2(shifted)

211 denominator = tl.sum(weights, axis=0)

212 + split_rows = split_offsets.to(tl.int64) * num_rows + row

213 partial_output = tl.load(

214 partial_output_ptr

211 - + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head)

212 - * HEAD_DIM

215 + + (split_rows[:, None] * NUM_QUERY_HEADS + head) * HEAD_DIM

216 + dim_offsets[None, :],

217 mask=split_mask[:, None],

218 other=0.0,

@@ -408,6 +411,41 @@ def _compress_qsa_groups_kernel(

411 )

412

413

414 +def _select_config(

415 + num_rows: int, num_kv_heads: int, use_prefill_config: bool, num_columns: int

416 +) -> tuple[int, int, int, int]:

417 + """Select (block_n, num_warps, num_tiles, num_splits) for the kernel.

418 +

419 + Tuned on GB300 for the Qwen3.8-Flash-Next TP1/TP2/TP4 shapes, keyed on

420 + base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on

421 + use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the

422 + uniform decode/verify length).

423 + """

424 + base_programs = num_rows * num_kv_heads

425 + if base_programs > 2048:

426 + BLOCK_N, target_splits, num_warps = (

427 + (32, 1, 1) if use_prefill_config else (64, 1, 2)

428 + )

429 + elif base_programs <= 24:

430 + BLOCK_N, target_splits, num_warps = 32, 64, 4

431 + elif base_programs <= 32:

432 + BLOCK_N, target_splits, num_warps = 32, 16, 1

433 + elif base_programs <= 64:

434 + BLOCK_N, target_splits, num_warps = 32, 8, 1

435 + elif base_programs <= 128:

436 + BLOCK_N, target_splits, num_warps = 32, 4, 1

437 + elif base_programs <= 256:

438 + BLOCK_N, target_splits, num_warps = 32, 8, 1

439 + elif base_programs <= 512:

440 + BLOCK_N, target_splits, num_warps = 64, 4, 2

441 + else:

442 + BLOCK_N, target_splits, num_warps = 64, 1, 2

443 + num_tiles = triton.cdiv(num_columns, BLOCK_N)

444 + # Never more splits than tiles, never empty.

445 + num_splits = min(target_splits, num_tiles)

446 + return BLOCK_N, num_warps, num_tiles, num_splits

447 +

448 +

449 def qsa_sparse_paged_attention(

450 q: torch.Tensor,

451 k_cache: torch.Tensor,

@@ -415,12 +453,17 @@ def qsa_sparse_paged_attention(

453 logical_indices: torch.Tensor,

454 block_table: torch.Tensor,

455 token_to_req: torch.Tensor,

456 + use_prefill_config: bool,

457 out: torch.Tensor | None = None,

458 ) -> torch.Tensor:

420 - """Run sparse GQA directly over paged BF16 K/V caches."""

421 -

422 - if not q.is_cuda or not HAS_TRITON:

423 - raise RuntimeError("paged QSA sparse attention requires CUDA and Triton")

459 + """Run sparse GQA directly over paged BF16 K/V caches.

460 +

461 + logical_indices is the PACKED selection buffer: [rows, selection_width + 1]

462 + with the trailing column holding each row's valid-entry count (written by

463 + the expand kernel; never a token index). The kernel reads it as the

464 + tile-loop bound. use_prefill_config only steers the top of the config table; see

465 + _select_config.

466 + """

467 if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape:

468 raise ValueError("QSA sparse attention received invalid Q/K/V shapes")

469 if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]:

@@ -429,8 +472,10 @@ def qsa_sparse_paged_attention(

472 raise ValueError("QSA sparse attention metadata has invalid shapes")

473 if not all(k_cache.shape[:3]) or not all(block_table.shape):

474 raise ValueError("QSA sparse attention cache and block table must be nonempty")

432 - if logical_indices.shape[1] <= 0:

433 - raise ValueError("QSA sparse attention requires a positive selection width")

475 + if logical_indices.shape[1] < 2:

476 + raise ValueError(

477 + "QSA packed indices need selection columns plus the count column"

478 + )

479 if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]:

480 raise ValueError("QSA sparse attention requires valid grouped-query heads")

481 head_dim = q.shape[2]

@@ -444,6 +489,7 @@ def qsa_sparse_paged_attention(

489 assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1

490 assert logical_indices.stride(1) == block_table.stride(1) == 1

491 assert token_to_req.stride(0) == 1

492 +

493 if out is None:

494 out = torch.empty_like(q)

495 if out.shape != q.shape:

@@ -455,26 +501,10 @@ def qsa_sparse_paged_attention(

501

502 group_size = q.shape[1] // k_cache.shape[2]

503 block_m = triton.next_power_of_2(group_size)

458 - base_programs = q.shape[0] * k_cache.shape[2]

459 - small_profile_limit = 8 if block_m <= 8 else 4

460 -

461 - # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes.

462 - # Narrow tiles favor decode; wide tiles improve throughput for prefill.

463 - if base_programs <= small_profile_limit:

464 - block_n, target_splits, partial_warps = 16, 64, 4

465 - elif base_programs < 32:

466 - block_n, target_splits, partial_warps = 16, 32, 4

467 - elif base_programs <= 256:

468 - block_n, target_splits, partial_warps = 64, 8, 2

469 - elif base_programs <= 512:

470 - block_n, target_splits, partial_warps = 64, 4, 2

471 - else:

472 - block_n, target_splits, partial_warps = 64, 1, 2

473 -

474 - num_tiles = triton.cdiv(logical_indices.shape[1], block_n)

475 - # Avoid empty splits when the selection width is smaller than the profile.

476 - max_useful_splits = 1 << (num_tiles.bit_length() - 1)

477 - num_splits = min(max_useful_splits, target_splits)

504 + selection_width = logical_indices.shape[1] - 1 # trailing column is the count

505 + block_n, partial_warps, num_tiles, num_splits = _select_config(

506 + q.shape[0], k_cache.shape[2], use_prefill_config, selection_width

507 + )

508

509 # Split=1 writes output directly and compiles out all workspace accesses.

510 if num_splits == 1:

@@ -518,7 +548,7 @@ def qsa_sparse_paged_attention(

548 q.shape[0],

549 k_cache.shape[0],

550 block_table.shape[0],

521 - TOPK=logical_indices.shape[1],

551 + TOPK=selection_width,

552 PAGE_SIZE=k_cache.shape[1],

553 PAGE_TABLE_WIDTH=block_table.shape[1],

554 GROUP_SIZE=group_size,

@@ -551,6 +581,133 @@ def qsa_sparse_paged_attention(

581 return out

582

583

584 +def warmup_qsa_sparse_paged_attention(

585 + kv_cache: torch.Tensor,

586 + block_table: torch.Tensor,

587 + *,

588 + num_query_heads: int,

589 + selection_width: int,

590 +) -> tuple[tuple[int, int, int], ...]:

591 + """Compile every production-reachable split-K/merge specialization."""

592 +

593 + head_dim = kv_cache.shape[-1] // 2

594 + key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1)

595 + num_kv_heads = key_cache.shape[2]

596 + group_size = num_query_heads // num_kv_heads

597 + block_m = triton.next_power_of_2(group_size)

598 +

599 + # Every config the dispatch can pick for this group size.

600 + profiles = {

601 + _select_config(num_rows, num_kv_heads, use_prefill_config, selection_width)

602 + for num_rows in range(1, 8193)

603 + for use_prefill_config in (False, True)

604 + }

605 +

606 + # Scalars constant per deployment get their real values (their divisibility

607 + # specialization is wanted); the batch-varying ones are do_not_specialize'd

608 + # on the kernels, so any value here compiles the only variant.

609 + num_rows = 16

610 + num_requests = 16

611 + q_ptr = TritonWarmupTensor(

612 + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)

613 + )

614 + k_cache_ptr = TritonWarmupTensor(

615 + key_cache.dtype,

616 + shape=tuple(key_cache.shape),

617 + strides=tuple(key_cache.stride()),

618 + )

619 + v_cache_ptr = TritonWarmupTensor(

620 + value_cache.dtype,

621 + shape=tuple(value_cache.shape),

622 + strides=tuple(value_cache.stride()),

623 + )

624 + # +1: the packed buffer's trailing count column.

625 + indices_ptr = TritonWarmupTensor(torch.int32, shape=(num_rows, selection_width + 1))

626 + block_table_ptr = TritonWarmupTensor(

627 + block_table.dtype,

628 + shape=tuple(block_table.shape),

629 + strides=tuple(block_table.stride()),

630 + )

631 + token_to_req_ptr = TritonWarmupTensor(torch.int32)

632 + output_ptr = TritonWarmupTensor(

633 + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)

634 + )

635 + head_stride = head_dim

636 + row_stride = num_query_heads * head_dim

637 + num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0])

638 +

639 + warmed = []

640 + for block_n, warps, num_tiles, num_splits in sorted(profiles):

641 + if num_splits == 1:

642 + partial_output_ptr = output_ptr

643 + partial_lse_ptr = output_ptr

644 + else:

645 + partial_output_ptr = TritonWarmupTensor(

646 + torch.float32,

647 + shape=(num_splits, num_rows, num_query_heads, head_dim),

648 + )

649 + partial_lse_ptr = TritonWarmupTensor(

650 + torch.float32, shape=(num_splits, num_rows, num_query_heads)

651 + )

652 + _qsa_sparse_paged_gqa_splitk_kernel.warmup(

653 + q_ptr,

654 + k_cache_ptr,

655 + v_cache_ptr,

656 + indices_ptr,

657 + block_table_ptr,

658 + token_to_req_ptr,

659 + partial_output_ptr,

660 + partial_lse_ptr,

661 + output_ptr,

662 + row_stride,

663 + head_stride,

664 + key_cache.stride(0),

665 + key_cache.stride(1),

666 + key_cache.stride(2),

667 + value_cache.stride(0),

668 + value_cache.stride(1),

669 + value_cache.stride(2),

670 + selection_width + 1,

671 + block_table.stride(0),

672 + row_stride,

673 + head_stride,

674 + num_rows,

675 + num_cache_blocks,

676 + num_requests,

677 + TOPK=selection_width,

678 + PAGE_SIZE=key_cache.shape[1],

679 + PAGE_TABLE_WIDTH=block_table.shape[1],

680 + GROUP_SIZE=group_size,

681 + HEAD_DIM=head_dim,

682 + NUM_QUERY_HEADS=num_query_heads,

683 + NUM_SPLITS=num_splits,

684 + NUM_TILES=num_tiles,

685 + BLOCK_M=block_m,

686 + BLOCK_N=block_n,

687 + num_warps=warps,

688 + num_stages=2,

689 + grid=(num_rows, num_kv_heads, num_splits),

690 + )

691 + if num_splits > 1:

692 + _qsa_merge_splitk_kernel.warmup(

693 + partial_output_ptr,

694 + partial_lse_ptr,

695 + output_ptr,

696 + row_stride,

697 + head_stride,

698 + num_rows,

699 + HEAD_DIM=head_dim,

700 + NUM_QUERY_HEADS=num_query_heads,

701 + NUM_SPLITS=num_splits,

702 + BLOCK_SPLITS=triton.next_power_of_2(num_splits),

703 + num_warps=2,

704 + num_stages=1,

705 + grid=(num_rows, num_query_heads),

706 + )

707 + warmed.append((block_n, num_splits, warps))

708 + return tuple(warmed)

709 +

710 +

711 def qsa_store_cache_rows(

712 cache: torch.Tensor,

713 slot_mapping: torch.Tensor,

@@ -708,4 +865,5 @@ __all__ = [

865 "qsa_compress_groups_with_ratio",

866 "qsa_sparse_paged_attention",

867 "qsa_store_cache_rows",

868 + "warmup_qsa_sparse_paged_attention",

869 ]

vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py

+13 -1

@@ -262,6 +262,15 @@ def _expand_qsa_indices_kernel(

262 tl.where(valid, token, -1),

263 mask=columns < OUTPUT_WIDTH,

264 )

265 + # TRAILING COUNT COLUMN: the packed buffer is OUTPUT_WIDTH+1 wide; the

266 + # last column holds this row's valid-entry count (expanded blocks plus

267 + # causal tail). It is NOT a token index — the sparse attention kernel

268 + # reads it as the row's tile-loop bound.

269 + if tl.program_id(1) == 0:

270 + tl.store(

271 + output_ptr + row * stride_output_row + OUTPUT_WIDTH * stride_output_column,

272 + expanded_count + tail_count,

273 + )

274

275

276 def _decode_tiles_per_program(num_requests: int, columns: int) -> int:

@@ -433,7 +442,10 @@ def expand_qsa_block_indices(

442 output_width = token_topk + compress_ratio - 1

443 assert block_indices.shape == (query_positions.numel(), block_topk)

444 assert visible_blocks.shape == query_positions.shape

436 - assert out.shape == (block_indices.shape[0], output_width)

445 + # +1: the packed buffer's trailing column holds each row's valid-entry

446 + # count (never a token index); the index region below only writes

447 + # columns [0, output_width).

448 + assert out.shape == (block_indices.shape[0], output_width + 1)

449 column_block = 256

450 grid = (block_indices.shape[0], triton.cdiv(output_width, column_block))

451 _expand_qsa_indices_kernel[grid](

vllm/models/qwen4_exp/nvidia/qsa.py

+15 -8

@@ -30,7 +30,6 @@ from vllm.utils.torch_utils import (

30 LayerNameType,

31 _encode_layer_name,

32 _resolve_layer_name,

33 - canonicalize_singleton_dim_strides,

33 direct_register_custom_op,

34 kv_cache_dtype_str_to_dtype,

35 )

@@ -124,6 +123,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

123 attn_metadata: FlashAttentionMetadata,

124 output: torch.Tensor,

125 token_to_req: torch.Tensor,

126 + use_prefill_config: bool,

127 output_scale: torch.Tensor | None = None,

128 output_block_scale: torch.Tensor | None = None,

129 ) -> torch.Tensor:

@@ -146,8 +146,6 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

146 logical_indices = topk_buffer[:num_tokens]

147 token_to_req = token_to_req[:num_tokens]

148 key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

149 - key_cache = canonicalize_singleton_dim_strides(key_cache)

150 - value_cache = canonicalize_singleton_dim_strides(value_cache)

149 if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:

150 raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")

151

@@ -160,6 +158,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

158 logical_indices,

159 attn_metadata.block_table,

160 token_to_req,

161 + use_prefill_config,

162 output[:num_tokens],

163 )

164 return output

@@ -209,6 +208,11 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

208 if self.total_num_heads % tp_size:

209 raise ValueError("QSA attention heads must be divisible by TP size")

210 self.num_heads = self.total_num_heads // tp_size

211 + # Decode/verify batches have at most 1 + num_spec query tokens per

212 + # request; use_prefill_config (max_query_len > this) steers the

213 + # config table. Shorter batches take the decode profile — harmless,

214 + # the difference is tile-shape tuning, not correctness.

215 + self._max_decode_query_len = 1 + vllm_config.num_speculative_tokens

216 self.total_num_kv_heads = int(config.num_key_value_heads)

217 if self.total_num_kv_heads >= tp_size:

218 if self.total_num_kv_heads % tp_size:

@@ -310,11 +314,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

314 prefix=f"{prefix}.indexer",

315 )

316 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens

317 + # PACKED selection buffer: the trailing column holds each row's

318 + # valid-entry count (written by the expand kernel) — never a token

319 + # index; the sparse attention kernel reads it as its loop bound.

320 + # MTP skip_topk steps reuse rows frozen from step 0; the count is

321 + # a row column, so compaction/reuse keep it paired with the content.

322 self.register_buffer(

323 "topk_indices_buffer",

324 torch.empty(

325 max_tokens,

317 - self.indexer.output_width,

326 + self.indexer.packed_output_width,

327 dtype=torch.int32,

328 ),

329 persistent=False,

@@ -369,10 +378,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

378 positions,

379 self.topk_indices_buffer[:num_tokens],

380 )

372 - if selected.shape != (

373 - num_tokens,

374 - self.indexer.output_width,

375 - ):

381 + if selected.shape != (num_tokens, self.indexer.packed_output_width):

382 raise RuntimeError("QSA indexer returned an invalid selection shape")

383 impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)

384 impl.do_kv_cache_update(

@@ -391,6 +397,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

397 main_metadata,

398 output,

399 token_to_req=side_metadata.token_to_req,

400 + use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len,

401 )

402

403 def forward(