From 31a8a266622781917cef482d01db415cfa3cbd92 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 4 Sep 2026 21:08:02 +0800 Subject: [PATCH] [Qwen3.8-Flash-Next] Improve QSA sparse GQA for prefill and short-ctx decode (#54873) Signed-off-by: Thien Tran Co-authored-by: Kimi --- tests/models/qwen4_exp/test_qsa_reference.py | 91 +++++-- .../warmup/qwen4_exp_qsa_warmup.py | 71 +++-- vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 28 +- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 244 +++++++++++++++--- .../qwen4_exp/nvidia/ops/qsa_indexer.py | 14 +- vllm/models/qwen4_exp/nvidia/qsa.py | 23 +- 6 files changed, 365 insertions(+), 106 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 01dcb20f40..75ade4b95f 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -754,7 +754,10 @@ def test_qsa_block_expansion_correctness() -> None: sequence_lengths.index_select(0, token_to_req.long()) // 4, ).to(torch.int32) - actual = torch.empty((2, 11), device="cuda", dtype=torch.int32) + # Packed layout: one trailing column per row holds the valid-entry count + # (never a token index). Row 0: 1 visible block + 2 tail; row 1: 2 blocks + # + 3 tail. + actual = torch.empty((2, 12), device="cuda", dtype=torch.int32) qsa_indexer_ops.expand_qsa_block_indices( blocks, query_positions, @@ -771,19 +774,35 @@ def test_qsa_block_expansion_correctness() -> None: token_topk=8, ) - torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual[:, :11], expected) + assert actual[:, 11].tolist() == [6, 11] @requires_qsa_kernels @pytest.mark.parametrize( - ("num_rows", "num_query_heads", "num_kv_heads", "page_size"), + ( + "num_rows", + "num_query_heads", + "num_kv_heads", + "page_size", + "use_prefill_config", + "num_requests", + ), [ - # Kernel-visible pages with --block-size 256 and hybrid-cache alignment. - pytest.param(1, 24, 2, 1792, id="tp1_split64"), - pytest.param(16, 12, 1, 1792, id="tp2_split32"), - pytest.param(32, 6, 1, 1024, id="tp4_split8"), - pytest.param(257, 6, 1, 1024, id="tp4_split4"), - pytest.param(513, 6, 1, 1024, id="tp4_split1"), + # Production page sizes from hybrid-cache block alignment: 784/800 + # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head + # splits are per-rank TP1/TP2/TP4; the largest batch runs both + # use_prefill_config variants. + pytest.param(1, 24, 2, 1600, True, 2, id="tp1_r1"), + pytest.param(16, 12, 1, 1600, True, 3, id="tp2_r16"), + pytest.param(32, 6, 1, 800, True, 5, id="tp4_r32"), + pytest.param(128, 24, 2, 1568, True, 7, id="tp1_r128"), + pytest.param(257, 6, 1, 800, True, 13, id="tp4_r257"), + pytest.param(513, 6, 1, 784, True, 17, id="tp4_r513"), + pytest.param(700, 6, 1, 800, True, 23, id="tp4_r700"), + pytest.param(1024, 24, 2, 1600, True, 33, id="tp1_r1024"), + pytest.param(2048, 24, 2, 1600, True, 63, id="tp1_r2048_prefill"), + pytest.param(2048, 24, 2, 1600, False, 63, id="tp1_r2048_uniform"), ], ) def test_qsa_sparse_paged_attention_correctness( @@ -791,10 +810,11 @@ def test_qsa_sparse_paged_attention_correctness( num_query_heads: int, num_kv_heads: int, page_size: int, + use_prefill_config: bool, + num_requests: int, ) -> None: torch.manual_seed(2) head_dim = 256 - num_requests = 2 num_selected_pages = 64 # Keep the newest page outside the synthetic top-k as causal headroom. num_pages_per_request = num_selected_pages + 1 @@ -822,13 +842,22 @@ def test_qsa_sparse_paged_attention_correctness( rows_per_request = math.ceil(num_rows / num_requests) row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32) token_to_req = row_indices // rows_per_request - request_row_counts = torch.tensor( - [rows_per_request, num_rows - rows_per_request], + # Uniform row split; the last request takes the remainder (possibly 0). + request_row_counts = torch.full( + (num_requests,), rows_per_request, device="cuda", dtype=torch.int32 + ) + request_row_counts[-1] = num_rows - rows_per_request * (num_requests - 1) + + # Mix context lengths: every third request is short-context, attending + # to only its first few pages; the rest fill their cache. + context_lengths = torch.full( + (num_requests,), + num_pages_per_request * page_size - 1, device="cuda", dtype=torch.int32, ) - - context_length = num_pages_per_request * page_size - 1 + short_requests = torch.arange(num_requests, device="cuda") % 3 == 1 + context_lengths[short_requests] = request_row_counts[short_requests] + 8 block_topk = indexer_budget // indexer_compress_ratio compressed_blocks_per_page = page_size // indexer_compress_ratio selection = torch.arange(block_topk, device="cuda") @@ -842,17 +871,19 @@ def test_qsa_sparse_paged_attention_correctness( ) rows_within_request = row_indices % rows_per_request query_positions = ( - context_length - request_row_counts[token_to_req.long()] + rows_within_request + context_lengths[token_to_req.long()] + - request_row_counts[token_to_req.long()] + + rows_within_request ).to(torch.int64) - sequence_lengths = torch.full( - (num_requests,), context_length, device="cuda", dtype=torch.int32 - ) + sequence_lengths = context_lengths visible_blocks = torch.minimum( (query_positions + 1) // indexer_compress_ratio, sequence_lengths.index_select(0, token_to_req.long()) // indexer_compress_ratio, ).to(torch.int32) + # +1: the packed trailing column holds each row's valid-entry count + # (never a token index); the reference reads only the selection region. logical_indices = torch.empty( - (num_rows, selection_width), device="cuda", dtype=torch.int32 + (num_rows, selection_width + 1), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.expand_qsa_block_indices( block_indices, @@ -862,7 +893,7 @@ def test_qsa_sparse_paged_attention_correctness( indexer_budget, logical_indices, ) - assert logical_indices.shape == (num_rows, selection_width) + assert logical_indices.shape == (num_rows, selection_width + 1) scale = q.shape[-1] ** -0.5 actual = qsa_ops.qsa_sparse_paged_attention( @@ -872,12 +903,13 @@ def test_qsa_sparse_paged_attention_correctness( logical_indices, block_table, token_to_req, + use_prefill_config=use_prefill_config, ) expected = _qsa_sparse_paged_attention_reference( q, k_cache, v_cache, - logical_indices, + logical_indices[:, :selection_width], block_table, token_to_req, scale, @@ -948,8 +980,10 @@ def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int) block_indices[prefill_slice], max_seq_len=sequence_lengths.max().item(), ) + # +1: the packed trailing count column (never a token index; excluded + # from the comparison). actual = torch.empty( - (rows, token_topk + compress_ratio - 1), device="cuda", dtype=torch.int32 + (rows, token_topk + compress_ratio), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.expand_qsa_block_indices( block_indices, @@ -977,7 +1011,10 @@ def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int) token_topk, ) - torch.testing.assert_close(actual.sort().values, expected.sort().values) + torch.testing.assert_close( + actual[:, : token_topk + compress_ratio - 1].sort().values, + expected.sort().values, + ) @requires_qsa_kernels @@ -1001,7 +1038,7 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> block_indices=block_indices, max_seq_len=64, # clamps to the page-table capacity ) - selected = torch.empty((2, 2051), device="cuda", dtype=torch.int32) + selected = torch.empty((2, 2052), device="cuda", dtype=torch.int32) qsa_indexer_ops.expand_qsa_block_indices( block_indices, query_positions, @@ -1013,8 +1050,10 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> assert selected[0, :2].tolist() == [0, 1] assert selected[1, :3].tolist() == [0, 1, 2] - assert torch.all(selected[0, 2:] == -1) - assert torch.all(selected[1, 3:] == -1) + assert torch.all(selected[0, 2:2051] == -1) + assert torch.all(selected[1, 3:2051] == -1) + # The packed trailing column holds each row's valid-entry count. + assert selected[:, 2051].tolist() == [2, 3] @requires_qsa_kernels diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py index 7f2c5a20b6..3cbc3bf2e0 100644 --- a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py +++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py @@ -5,6 +5,8 @@ import sys from typing import TYPE_CHECKING, cast +import torch + from vllm.logger import init_logger if TYPE_CHECKING: @@ -15,37 +17,41 @@ logger = init_logger(__name__) def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: - """Warm every reachable QSA decode-query-length specialization.""" + """Warm every reachable QSA specialization: indexer decode-query-length + profiles plus the sparse attention split-K/merge configs.""" qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa") - if qsa_module is None: + attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa") + if qsa_module is None or attn_module is None: return - indexer = next( - ( - layer - for layer in worker.get_model().modules() - if isinstance(layer, qsa_module.QSAIndexer) - ), - None, - ) - if indexer is None: + indexer = None + owner = None + for layer in worker.get_model().modules(): + if indexer is None and isinstance(layer, qsa_module.QSAIndexer): + indexer = layer + elif owner is None and isinstance(layer, attn_module.Qwen4ExpQSAAttention): + owner = layer + if indexer is None or owner is None: return runner = worker.model_runner - prefix = indexer.compressed_key_cache.prefix - group_id = next( - i - for i, group in enumerate(runner.kv_cache_config.kv_cache_groups) - if prefix in group.layer_names - ) - if worker.use_v2_model_runner: - runner_v2 = cast("GPUModelRunnerV2", runner) - block_table = runner_v2.block_tables.input_block_tables[group_id] - max_decode_query_len = runner_v2.decode_query_len - else: - block_table = runner.input_batch.block_table[group_id].get_device_tensor( + + def block_table_for(prefix: str) -> torch.Tensor: + group_id = next( + i + for i, group in enumerate(runner.kv_cache_config.kv_cache_groups) + if prefix in group.layer_names + ) + if worker.use_v2_model_runner: + runner_v2 = cast("GPUModelRunnerV2", runner) + return runner_v2.block_tables.input_block_tables[group_id] + return runner.input_batch.block_table[group_id].get_device_tensor( runner.max_num_reqs ) + + if worker.use_v2_model_runner: + max_decode_query_len = cast("GPUModelRunnerV2", runner).decode_query_len + else: max_decode_query_len = runner.uniform_decode_query_len from vllm.models.qwen4_exp.nvidia.ops.qsa_indexer import ( @@ -56,7 +62,7 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: assert k_cache.numel() profiles = warmup_qsa_mqa_paged_decode( k_cache, - block_table, + block_table_for(indexer.compressed_key_cache.prefix), num_heads=indexer.index_n_heads, head_dim=indexer.index_head_dim, max_decode_query_len=max_decode_query_len, @@ -64,3 +70,20 @@ def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: max_num_batched_tokens=runner.max_num_tokens, ) logger.info("Warmed up Qwen4Exp QSA decode kernels: %s.", profiles) + + from vllm.models.qwen4_exp.nvidia.ops.qsa import ( + warmup_qsa_sparse_paged_attention, + ) + + kv_cache = owner.kv_cache + assert kv_cache.numel() + attention_profiles = warmup_qsa_sparse_paged_attention( + kv_cache, + block_table_for(owner.layer_name), + num_query_heads=owner.num_heads, + selection_width=indexer.output_width, + ) + logger.info( + "Warmed up Qwen4Exp QSA sparse attention kernels: %s.", + attention_profiles, + ) diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index 11c787b4e5..982e15a34a 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -167,8 +167,19 @@ class QSAIndexer(nn.Module): @property def output_width(self) -> int: + """Selection (index) columns per row.""" return self.token_topk + self.compress_ratio - 1 + @property + def packed_output_width(self) -> int: + """Packed selection-buffer width: output_width + 1. + + The trailing column holds each row's valid-entry count (written by + the expand kernel) — never a token index; the sparse attention + kernel reads it as its tile-loop bound. + """ + return self.output_width + 1 + def _metadata( self, ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None: @@ -211,7 +222,14 @@ class QSAIndexer(nn.Module): positions: torch.Tensor, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Return fixed-width request-relative token indices padded with ``-1``.""" + """Select each query row's token indices. + + Returns the packed buffer of shape [num_tokens, output_width + 1]: + the leading ``output_width`` columns are ``-1``-padded + request-relative token indices, and the trailing column is the row's + valid-entry count (the attention kernel's loop bound, never a token + index). + """ metadata = self._metadata() if metadata is None: @@ -219,11 +237,13 @@ class QSAIndexer(nn.Module): if self.skip_topk and out is not None: return out result = torch.full( - (hidden_states.shape[0], self.output_width), + (hidden_states.shape[0], self.packed_output_width), -1, dtype=torch.int32, device=hidden_states.device, ) + # Inert rows carry a zero valid count (empty loop bound), not -1. + result[:, -1] = 0 if out is not None: out.copy_(result) return out @@ -357,11 +377,11 @@ class QSAIndexer(nn.Module): if out is None: out = torch.empty( num_tokens, - self.output_width, + self.packed_output_width, dtype=torch.int32, device=q.device, ) - elif out.shape != (num_tokens, self.output_width): + elif out.shape != (num_tokens, self.packed_output_width): raise ValueError("QSA selection output has an invalid shape") num_decode_tokens = compressed_metadata.num_decode_tokens diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 312160210c..709acbde00 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -6,10 +6,14 @@ from __future__ import annotations import torch +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, + triton_scalar_specialization_rep, +) from vllm.triton_utils import HAS_TRITON, tl, triton -@triton.jit +@triton.jit(do_not_specialize=["num_rows", "num_requests"]) def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, k_cache_ptr, @@ -52,6 +56,12 @@ def _qsa_sparse_paged_gqa_splitk_kernel( request = tl.load(token_to_req_ptr + row) safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + # The packed selection buffer carries one TRAILING COUNT COLUMN per row + # (column TOPK of a TOPK+1-wide buffer): the row's valid-entry count, + # written by the expand kernel. It is never a token index — the tile loop + # and the index load below only ever cover columns [0, TOPK). + valid_count = tl.load(indices_ptr + row * stride_indices_row + TOPK) + head_offsets = tl.arange(0, BLOCK_M) dim_offsets = tl.arange(0, HEAD_DIM) column_offsets = tl.arange(0, BLOCK_N) @@ -70,10 +80,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel( accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 - # Dynamic bounds avoid padded main-loop iterations for uneven splits. - split_tile_start = split_id * NUM_TILES // NUM_SPLITS - split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS - for tile in range(split_tile_start, split_tile_end): + tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N)) + + for tile in range(split_id, tile_end, NUM_SPLITS): columns = tile * BLOCK_N + column_offsets logical_token = tl.load( indices_ptr + row * stride_indices_row + columns, @@ -156,29 +165,23 @@ def _qsa_sparse_paged_gqa_splitk_kernel( max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), -float("inf"), ) + partial_row = (split_id * num_rows + row).to(tl.int64) tl.store( partial_output_ptr - + ( - (split_id * num_rows + row) * NUM_QUERY_HEADS - + first_head - + head_offsets[:, None] - ) + + (partial_row * NUM_QUERY_HEADS + first_head + head_offsets[:, None]) * HEAD_DIM + dim_offsets[None, :], normalized_output, mask=output_mask, ) tl.store( - partial_lse_ptr - + (split_id * num_rows + row) * NUM_QUERY_HEADS - + first_head - + head_offsets, + partial_lse_ptr + partial_row * NUM_QUERY_HEADS + first_head + head_offsets, partial_lse, mask=head_offsets < GROUP_SIZE, ) -@triton.jit +@triton.jit(do_not_specialize=["num_rows"]) def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, @@ -206,10 +209,10 @@ def _qsa_merge_splitk_kernel( shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf")) weights = tl.math.exp2(shifted) denominator = tl.sum(weights, axis=0) + split_rows = split_offsets.to(tl.int64) * num_rows + row partial_output = tl.load( partial_output_ptr - + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) - * HEAD_DIM + + (split_rows[:, None] * NUM_QUERY_HEADS + head) * HEAD_DIM + dim_offsets[None, :], mask=split_mask[:, None], other=0.0, @@ -408,6 +411,41 @@ def _compress_qsa_groups_kernel( ) +def _select_config( + num_rows: int, num_kv_heads: int, use_prefill_config: bool, num_columns: int +) -> tuple[int, int, int, int]: + """Select (block_n, num_warps, num_tiles, num_splits) for the kernel. + + Tuned on GB300 for the Qwen3.8-Flash-Next TP1/TP2/TP4 shapes, keyed on + base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on + use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the + uniform decode/verify length). + """ + base_programs = num_rows * num_kv_heads + if base_programs > 2048: + BLOCK_N, target_splits, num_warps = ( + (32, 1, 1) if use_prefill_config else (64, 1, 2) + ) + elif base_programs <= 24: + BLOCK_N, target_splits, num_warps = 32, 64, 4 + elif base_programs <= 32: + BLOCK_N, target_splits, num_warps = 32, 16, 1 + elif base_programs <= 64: + BLOCK_N, target_splits, num_warps = 32, 8, 1 + elif base_programs <= 128: + BLOCK_N, target_splits, num_warps = 32, 4, 1 + elif base_programs <= 256: + BLOCK_N, target_splits, num_warps = 32, 8, 1 + elif base_programs <= 512: + BLOCK_N, target_splits, num_warps = 64, 4, 2 + else: + BLOCK_N, target_splits, num_warps = 64, 1, 2 + num_tiles = triton.cdiv(num_columns, BLOCK_N) + # Never more splits than tiles, never empty. + num_splits = min(target_splits, num_tiles) + return BLOCK_N, num_warps, num_tiles, num_splits + + def qsa_sparse_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -415,12 +453,17 @@ def qsa_sparse_paged_attention( logical_indices: torch.Tensor, block_table: torch.Tensor, token_to_req: torch.Tensor, + use_prefill_config: bool, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA directly over paged BF16 K/V caches. - if not q.is_cuda or not HAS_TRITON: - raise RuntimeError("paged QSA sparse attention requires CUDA and Triton") + logical_indices is the PACKED selection buffer: [rows, selection_width + 1] + with the trailing column holding each row's valid-entry count (written by + the expand kernel; never a token index). The kernel reads it as the + tile-loop bound. use_prefill_config only steers the top of the config table; see + _select_config. + """ if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: raise ValueError("QSA sparse attention received invalid Q/K/V shapes") if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]: @@ -429,8 +472,10 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention metadata has invalid shapes") if not all(k_cache.shape[:3]) or not all(block_table.shape): raise ValueError("QSA sparse attention cache and block table must be nonempty") - if logical_indices.shape[1] <= 0: - raise ValueError("QSA sparse attention requires a positive selection width") + if logical_indices.shape[1] < 2: + raise ValueError( + "QSA packed indices need selection columns plus the count column" + ) if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] @@ -444,6 +489,7 @@ def qsa_sparse_paged_attention( assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 assert logical_indices.stride(1) == block_table.stride(1) == 1 assert token_to_req.stride(0) == 1 + if out is None: out = torch.empty_like(q) if out.shape != q.shape: @@ -455,26 +501,10 @@ def qsa_sparse_paged_attention( group_size = q.shape[1] // k_cache.shape[2] block_m = triton.next_power_of_2(group_size) - base_programs = q.shape[0] * k_cache.shape[2] - small_profile_limit = 8 if block_m <= 8 else 4 - - # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes. - # Narrow tiles favor decode; wide tiles improve throughput for prefill. - if base_programs <= small_profile_limit: - block_n, target_splits, partial_warps = 16, 64, 4 - elif base_programs < 32: - block_n, target_splits, partial_warps = 16, 32, 4 - elif base_programs <= 256: - block_n, target_splits, partial_warps = 64, 8, 2 - elif base_programs <= 512: - block_n, target_splits, partial_warps = 64, 4, 2 - else: - block_n, target_splits, partial_warps = 64, 1, 2 - - num_tiles = triton.cdiv(logical_indices.shape[1], block_n) - # Avoid empty splits when the selection width is smaller than the profile. - max_useful_splits = 1 << (num_tiles.bit_length() - 1) - num_splits = min(max_useful_splits, target_splits) + selection_width = logical_indices.shape[1] - 1 # trailing column is the count + block_n, partial_warps, num_tiles, num_splits = _select_config( + q.shape[0], k_cache.shape[2], use_prefill_config, selection_width + ) # Split=1 writes output directly and compiles out all workspace accesses. if num_splits == 1: @@ -518,7 +548,7 @@ def qsa_sparse_paged_attention( q.shape[0], k_cache.shape[0], block_table.shape[0], - TOPK=logical_indices.shape[1], + TOPK=selection_width, PAGE_SIZE=k_cache.shape[1], PAGE_TABLE_WIDTH=block_table.shape[1], GROUP_SIZE=group_size, @@ -551,6 +581,133 @@ def qsa_sparse_paged_attention( return out +def warmup_qsa_sparse_paged_attention( + kv_cache: torch.Tensor, + block_table: torch.Tensor, + *, + num_query_heads: int, + selection_width: int, +) -> tuple[tuple[int, int, int], ...]: + """Compile every production-reachable split-K/merge specialization.""" + + head_dim = kv_cache.shape[-1] // 2 + key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) + num_kv_heads = key_cache.shape[2] + group_size = num_query_heads // num_kv_heads + block_m = triton.next_power_of_2(group_size) + + # Every config the dispatch can pick for this group size. + profiles = { + _select_config(num_rows, num_kv_heads, use_prefill_config, selection_width) + for num_rows in range(1, 8193) + for use_prefill_config in (False, True) + } + + # Scalars constant per deployment get their real values (their divisibility + # specialization is wanted); the batch-varying ones are do_not_specialize'd + # on the kernels, so any value here compiles the only variant. + num_rows = 16 + num_requests = 16 + q_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) + k_cache_ptr = TritonWarmupTensor( + key_cache.dtype, + shape=tuple(key_cache.shape), + strides=tuple(key_cache.stride()), + ) + v_cache_ptr = TritonWarmupTensor( + value_cache.dtype, + shape=tuple(value_cache.shape), + strides=tuple(value_cache.stride()), + ) + # +1: the packed buffer's trailing count column. + indices_ptr = TritonWarmupTensor(torch.int32, shape=(num_rows, selection_width + 1)) + block_table_ptr = TritonWarmupTensor( + block_table.dtype, + shape=tuple(block_table.shape), + strides=tuple(block_table.stride()), + ) + token_to_req_ptr = TritonWarmupTensor(torch.int32) + output_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) + head_stride = head_dim + row_stride = num_query_heads * head_dim + num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0]) + + warmed = [] + for block_n, warps, num_tiles, num_splits in sorted(profiles): + if num_splits == 1: + partial_output_ptr = output_ptr + partial_lse_ptr = output_ptr + else: + partial_output_ptr = TritonWarmupTensor( + torch.float32, + shape=(num_splits, num_rows, num_query_heads, head_dim), + ) + partial_lse_ptr = TritonWarmupTensor( + torch.float32, shape=(num_splits, num_rows, num_query_heads) + ) + _qsa_sparse_paged_gqa_splitk_kernel.warmup( + q_ptr, + k_cache_ptr, + v_cache_ptr, + indices_ptr, + block_table_ptr, + token_to_req_ptr, + partial_output_ptr, + partial_lse_ptr, + output_ptr, + row_stride, + head_stride, + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + selection_width + 1, + block_table.stride(0), + row_stride, + head_stride, + num_rows, + num_cache_blocks, + num_requests, + TOPK=selection_width, + PAGE_SIZE=key_cache.shape[1], + PAGE_TABLE_WIDTH=block_table.shape[1], + GROUP_SIZE=group_size, + HEAD_DIM=head_dim, + NUM_QUERY_HEADS=num_query_heads, + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=warps, + num_stages=2, + grid=(num_rows, num_kv_heads, num_splits), + ) + if num_splits > 1: + _qsa_merge_splitk_kernel.warmup( + partial_output_ptr, + partial_lse_ptr, + output_ptr, + row_stride, + head_stride, + num_rows, + HEAD_DIM=head_dim, + NUM_QUERY_HEADS=num_query_heads, + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + num_warps=2, + num_stages=1, + grid=(num_rows, num_query_heads), + ) + warmed.append((block_n, num_splits, warps)) + return tuple(warmed) + + def qsa_store_cache_rows( cache: torch.Tensor, slot_mapping: torch.Tensor, @@ -708,4 +865,5 @@ __all__ = [ "qsa_compress_groups_with_ratio", "qsa_sparse_paged_attention", "qsa_store_cache_rows", + "warmup_qsa_sparse_paged_attention", ] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 1f0abae468..7efcc9077b 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -262,6 +262,15 @@ def _expand_qsa_indices_kernel( tl.where(valid, token, -1), mask=columns < OUTPUT_WIDTH, ) + # TRAILING COUNT COLUMN: the packed buffer is OUTPUT_WIDTH+1 wide; the + # last column holds this row's valid-entry count (expanded blocks plus + # causal tail). It is NOT a token index — the sparse attention kernel + # reads it as the row's tile-loop bound. + if tl.program_id(1) == 0: + tl.store( + output_ptr + row * stride_output_row + OUTPUT_WIDTH * stride_output_column, + expanded_count + tail_count, + ) def _decode_tiles_per_program(num_requests: int, columns: int) -> int: @@ -433,7 +442,10 @@ def expand_qsa_block_indices( output_width = token_topk + compress_ratio - 1 assert block_indices.shape == (query_positions.numel(), block_topk) assert visible_blocks.shape == query_positions.shape - assert out.shape == (block_indices.shape[0], output_width) + # +1: the packed buffer's trailing column holds each row's valid-entry + # count (never a token index); the index region below only writes + # columns [0, output_width). + assert out.shape == (block_indices.shape[0], output_width + 1) column_block = 256 grid = (block_indices.shape[0], triton.cdiv(output_width, column_block)) _expand_qsa_indices_kernel[grid]( diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 064339dc53..2e135149af 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -30,7 +30,6 @@ from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, _resolve_layer_name, - canonicalize_singleton_dim_strides, direct_register_custom_op, kv_cache_dtype_str_to_dtype, ) @@ -124,6 +123,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): attn_metadata: FlashAttentionMetadata, output: torch.Tensor, token_to_req: torch.Tensor, + use_prefill_config: bool, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: @@ -146,8 +146,6 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): logical_indices = topk_buffer[:num_tokens] token_to_req = token_to_req[:num_tokens] key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) - key_cache = canonicalize_singleton_dim_strides(key_cache) - value_cache = canonicalize_singleton_dim_strides(value_cache) if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16: raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V") @@ -160,6 +158,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): logical_indices, attn_metadata.block_table, token_to_req, + use_prefill_config, output[:num_tokens], ) return output @@ -209,6 +208,11 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): if self.total_num_heads % tp_size: raise ValueError("QSA attention heads must be divisible by TP size") self.num_heads = self.total_num_heads // tp_size + # Decode/verify batches have at most 1 + num_spec query tokens per + # request; use_prefill_config (max_query_len > this) steers the + # config table. Shorter batches take the decode profile — harmless, + # the difference is tile-shape tuning, not correctness. + self._max_decode_query_len = 1 + vllm_config.num_speculative_tokens self.total_num_kv_heads = int(config.num_key_value_heads) if self.total_num_kv_heads >= tp_size: if self.total_num_kv_heads % tp_size: @@ -310,11 +314,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): prefix=f"{prefix}.indexer", ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + # PACKED selection buffer: the trailing column holds each row's + # valid-entry count (written by the expand kernel) — never a token + # index; the sparse attention kernel reads it as its loop bound. + # MTP skip_topk steps reuse rows frozen from step 0; the count is + # a row column, so compaction/reuse keep it paired with the content. self.register_buffer( "topk_indices_buffer", torch.empty( max_tokens, - self.indexer.output_width, + self.indexer.packed_output_width, dtype=torch.int32, ), persistent=False, @@ -369,10 +378,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): positions, self.topk_indices_buffer[:num_tokens], ) - if selected.shape != ( - num_tokens, - self.indexer.output_width, - ): + if selected.shape != (num_tokens, self.indexer.packed_output_width): raise RuntimeError("QSA indexer returned an invalid selection shape") impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) impl.do_kv_cache_update( @@ -391,6 +397,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): main_metadata, output, token_to_req=side_metadata.token_to_req, + use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len, ) def forward(