[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>

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

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,
+    )

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

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",
 ]

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](

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(