[Kernel][SM70] Accelerate Qwen3.8 QSA paged prefill

docs/design/sm70_v100_migration_control.md

@@ -44189,3 +44189,40 @@ Interpretation:
   `47836839b542fb73494caa64adc14cd660e38c535c4a5e67e16d3a763196dac7`
   across repeated runs. A separate FP16 Dense/Indexer candidate screen remains
   benchmark-only and does not alter production dispatch.
+
+## 2026-08-28 Qwen3.8 QSA page4 XQA prefill audit
+
+- The SM70-only route converts each selected four-token QSA block into a
+  virtual Flash-V100 page, sorts physical microblocks for cache locality, and
+  keeps a marked causal tail last. It is restricted to FP16 Hq6/Hkv1/D256,
+  the 2051-token QSA selection layout, compatible contiguous or interleaved
+  KV strides, and at least 4096 query rows. The route remains default-on and
+  retains `VLLM_SM70_QSA_XQA_PAGE4=0` as an operational escape hatch.
+- A production-shaped interleaved-KV V100 A/B at 4096 rows measures the
+  established Triton route at `27.8303 ms` and page4 XQA, including table
+  construction and sort, at `8.84736 ms` (`3.1456x`). Across 6,291,456 FP16
+  outputs, maximum absolute difference is `6.104e-5`, relative L2 is
+  `2.845e-4`, cosine is `0.99999994`, and all outputs are finite.
+- A separate nonmonotonic contiguous-page A/B passes with maximum absolute
+  difference `3.815e-6`, relative L2 `3.645e-4`, and cosine `1.0`. Causal
+  tails of one, two, and three tokens each remain below relative L2 `3.64e-4`
+  with cosine at least `0.99999988`. The 4095-row boundary takes the Triton
+  fallback and is bitwise identical with or without the newly forwarded
+  metadata.
+- The hybrid 784-token scheduler / 16-token kernel geometry is also exercised
+  with a nonmonotonic 128-entry virtual page table after the physical-page
+  correction. Page4 XQA passes at maximum absolute difference `3.815e-6`,
+  relative L2 `3.631e-4`, and cosine `0.99999994`; its CUDA Graph replay is
+  bitwise equal to eager output.
+- Prewarmed CUDA Graph capture succeeds on V100; two replays are bitwise
+  identical to eager page4 XQA with output hash
+  `9b4c76f8420d6e349dc7d552c72d6f0a861332e7e8e8f62459a1c48f0faf278f`.
+  The table kernel now clamps padded or stale query positions to the live
+  request length and admits a partial tail only when the expanded QSA indices
+  contain that exact token, preventing an invalid synthetic tail page.
+- Raising the shared page-ID capacity from 8 to 32 does not reduce the
+  declared two-block V100 occupancy: the page4 padded kernel uses 45,568
+  bytes per CTA (`91,136 < 98,304` bytes for two CTAs), and the pipeline
+  variant uses 47,616 bytes (`95,232 < 98,304`). Existing FP16 page-16 and
+  page-784 XQA-to-scalar smokes pass at relative L2 `3.03e-5` and `4.10e-5`,
+  respectively.

flash-attention-v100/kernel/flash_decode_paged.cu

@@ -38,7 +38,9 @@ constexpr int kThreadsPerBlock = 256;
 constexpr int kWarpsPerBlock = kThreadsPerBlock / kWarpSize;
 constexpr int kXQATCBlockN = 128;
 constexpr int kXQATCStride = 128;
-constexpr int kXQATCPageIdsCapacity = kXQATCBlockN / 16;
+// QSA exposes selected four-token microblocks as virtual paged KV. Keep enough
+// page slots for one full XQA tile at that minimum supported granularity.
+constexpr int kXQATCPageIdsCapacity = kXQATCBlockN / 4;
 constexpr int kXQATC256WideWarpCount = 8;
 constexpr int kXQATC256WideThreads = kXQATC256WideWarpCount * kWarpSize;
 constexpr int kXQATC256WideBlockM = 8;
@@ -742,17 +744,21 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector(
   const int row = copy_idx / panel_d_stride_uint4;
   const int vec_col = copy_idx % panel_d_stride_uint4;
   const int token_offset = tile_page_offset + kv_tile_start + row;
-  static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 16 || BLOCK_SIZE == 784 ||
-                    BLOCK_SIZE == 800 || BLOCK_SIZE == 1568 ||
-                    BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296,
+  static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 4 || BLOCK_SIZE == 16 ||
+                    BLOCK_SIZE == 784 || BLOCK_SIZE == 800 ||
+                    BLOCK_SIZE == 1568 || BLOCK_SIZE == 1648 ||
+                    BLOCK_SIZE == 3296,
                 "Unsupported paged-KV block-size specialization");
   static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16 ||
                     BLOCK_SIZE == 800 || BLOCK_SIZE == 1568 ||
                     BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296,
                 "The fixed-stride Hkv=1 layout requires a specialized page");
   int logical_block;
   int block_offset;
-  if constexpr (BLOCK_SIZE == 16) {
+  if constexpr (BLOCK_SIZE == 4) {
+    logical_block = token_offset >> 2;
+    block_offset = token_offset & 3;
+  } else if constexpr (BLOCK_SIZE == 16) {
     logical_block = token_offset >> 4;
     block_offset = token_offset & 15;
   } else if constexpr (BLOCK_SIZE == 784) {
@@ -840,7 +846,10 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel(
       const int token_offset = tile_page_offset + kv_tile_start + row;
       int logical_block;
       int block_offset;
-      if constexpr (BLOCK_SIZE == 16) {
+      if constexpr (BLOCK_SIZE == 4) {
+        logical_block = token_offset >> 2;
+        block_offset = token_offset & 3;
+      } else if constexpr (BLOCK_SIZE == 16) {
         logical_block = token_offset >> 4;
         block_offset = token_offset & 15;
       } else if constexpr (BLOCK_SIZE == 784) {
@@ -4456,9 +4465,13 @@ at::Tensor flash_attention_decode_paged_xqa(
       (batch_context_route == XQABatchContextRoute::kDualCta ||
        batch_context_route == XQABatchContextRoute::kDualCtaSplit) &&
       xqa_e5m2_batch_wide_load_enabled();
+  const bool use_qsa_page4 =
+      q.size(0) >= 4096 && q_per_kv == 6 && partition_size == 256 &&
+      k_cache.size(1) == 4 && k_cache.size(2) == 1 &&
+      k_cache.scalar_type() == at::kHalf && block_table.size(1) == 513;
   const bool use_g6_dual_cta =
-      use_g6_p1024_auto || use_g6_p1024_sawtooth || use_mtp5_dual_cta ||
-      use_e5m2_g6_dual_cta ||
+      use_qsa_page4 || use_g6_p1024_auto || use_g6_p1024_sawtooth ||
+      use_mtp5_dual_cta || use_e5m2_g6_dual_cta ||
       batch_context_route == XQABatchContextRoute::kDualCta ||
       batch_context_route == XQABatchContextRoute::kDualCtaSplit ||
       (xqa_g6_dual_cta_enabled() && (use_padded_smem || use_g6_dual_cta_dense));
@@ -4467,6 +4480,8 @@ at::Tensor flash_attention_decode_paged_xqa(
       batch_context_route == XQABatchContextRoute::kDualCtaSplit ||
       (use_g6_dual_cta && xqa_split_reduce_enabled());
   const bool supports_block16_index = use_g6_dual_cta && k_cache.size(1) == 16;
+  const bool use_block4_index =
+      use_g6_dual_cta && partition_size == 256 && k_cache.size(1) == 4;
   const bool supports_block16_contiguous_layout =
       supports_block16_index && k_cache.size(2) == 1 &&
       k_cache.stride(0) == 4096 && k_cache.stride(1) == 256 &&
@@ -4808,6 +4823,13 @@ at::Tensor flash_attention_decode_paged_xqa(
             v_scale, launch_num_partitions, use_split_reduce,
             split_reduce_dim_tile, stream);
       }
+    } else if (use_block4_index) {
+      launch_flash_attention_decode_paged_xqa_tc_256_wide<
+          256, 6, true, kXQATCG6DualCtaThreads, 2, 4, false>(
+          q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits,
+          exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale,
+          launch_num_partitions, use_split_reduce, split_reduce_dim_tile,
+          stream);
     } else if (block16_layout_mode == 2) {
       launch_flash_attention_decode_paged_xqa_tc_256_wide<
           256, 6, true, kXQATCG6DualCtaThreads, 2, 16, true>(

tests/models/qwen4_exp/test_qsa_cache.py

@@ -53,6 +53,8 @@ def test_qsa_forward_splits_local_flash_cache_layout(monkeypatch) -> None:
     logical_indices = torch.zeros(1, 4, dtype=torch.int32)
     block_table = torch.zeros(1, 1, dtype=torch.int32)
     token_to_req = torch.zeros(1, dtype=torch.int32)
+    query_positions = torch.zeros(1, dtype=torch.int64)
+    sequence_lengths = torch.ones(1, dtype=torch.int32)
     captured = {}
 
     def fake_sparse_attention(
@@ -63,13 +65,18 @@ def fake_sparse_attention(
         block_table_arg,
         token_to_req_arg,
         output_arg,
+        *,
+        query_positions,
+        sequence_lengths,
     ):
         captured["key_cache"] = key_cache_arg
         captured["value_cache"] = value_cache_arg
         assert torch.equal(query_arg, query)
         assert torch.equal(logical_indices_arg, logical_indices)
         assert torch.equal(block_table_arg, block_table)
         assert torch.equal(token_to_req_arg, token_to_req)
+        captured["query_positions"] = query_positions
+        captured["sequence_lengths"] = sequence_lengths
         output_arg.fill_(1)
         return output_arg
 
@@ -89,10 +96,14 @@ def fake_sparse_attention(
         SimpleNamespace(num_actual_tokens=1, block_table=block_table),
         output,
         token_to_req,
+        query_positions=query_positions,
+        sequence_lengths=sequence_lengths,
     )
 
     expected_key, expected_value = kv_cache.unbind(1)
     assert torch.equal(captured["key_cache"], expected_key)
     assert torch.equal(captured["value_cache"], expected_value)
+    assert torch.equal(captured["query_positions"], query_positions)
+    assert torch.equal(captured["sequence_lengths"], sequence_lengths)
     assert result is output
     assert torch.equal(output, torch.ones_like(output))

tests/models/qwen4_exp/test_qsa_ops.py

@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import pytest
 import torch
 
 from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops
 from vllm.models.qwen4_exp.nvidia.ops.qsa import (
     _qsa_indexer_cublas_shape_supported,
     _qsa_sparse_launch_profile,
+    _qsa_xqa_page4_shape_supported,
 )
 
 
@@ -36,6 +38,170 @@ def test_qsa_indexer_cublas_accepts_only_exact_single_request_shape():
     assert not _qsa_indexer_cublas_shape_supported(query[:, :3], cache, page_table)
 
 
+def test_qsa_xqa_page4_accepts_only_exact_sm70_prefill_shape():
+    query = torch.empty(8, 6, 256, dtype=torch.float16)
+    key_cache = torch.empty(2, 400, 1, 256, dtype=torch.float16)
+    value_cache = torch.empty_like(key_cache)
+    indices = torch.empty(8, 2051, dtype=torch.int32)
+    page_table = torch.empty(1, 2, dtype=torch.int32)
+    token_to_request = torch.zeros(8, dtype=torch.int32)
+    query_positions = torch.arange(8, dtype=torch.int64)
+    sequence_lengths = torch.full((1,), 8, dtype=torch.int32)
+
+    assert _qsa_xqa_page4_shape_supported(
+        query,
+        key_cache,
+        value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    strided_query = torch.empty(8, 6, 257, dtype=torch.float16)[..., :256]
+    assert _qsa_xqa_page4_shape_supported(
+        strided_query,
+        key_cache,
+        value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    interleaved_cache = torch.empty(2, 2, 400, 1, 256, dtype=torch.float16)
+    interleaved_key_cache, interleaved_value_cache = interleaved_cache.unbind(1)
+    assert _qsa_xqa_page4_shape_supported(
+        query,
+        interleaved_key_cache,
+        interleaved_value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    assert not _qsa_xqa_page4_shape_supported(
+        query.to(torch.bfloat16),
+        key_cache,
+        value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    assert not _qsa_xqa_page4_shape_supported(
+        query,
+        key_cache[:, :398],
+        value_cache[:, :398],
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    assert not _qsa_xqa_page4_shape_supported(
+        query,
+        key_cache,
+        value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions.to(torch.int32),
+        sequence_lengths,
+    )
+
+
+def test_qsa_xqa_page4_route_uses_configured_boundary(monkeypatch):
+    rows = 8
+    query = torch.empty(rows, 6, 256, dtype=torch.float16)
+    key_cache = torch.empty(2, 400, 1, 256, dtype=torch.float16)
+    value_cache = torch.empty_like(key_cache)
+    indices = torch.empty(rows, 2051, dtype=torch.int32)
+    page_table = torch.empty(1, 2, dtype=torch.int32)
+    token_to_request = torch.zeros(rows, dtype=torch.int32)
+    query_positions = torch.arange(rows, dtype=torch.int64)
+    sequence_lengths = torch.full((1,), rows, dtype=torch.int32)
+    monkeypatch.setattr(qsa_ops, "_SM70_QSA_XQA_PAGE4", True)
+    monkeypatch.setattr(qsa_ops, "_SM70_QSA_XQA_PAGE4_MIN_ROWS", rows)
+    monkeypatch.setattr(
+        qsa_ops.current_platform,
+        "is_device_capability",
+        lambda capability: capability == 70,
+    )
+
+    args = (
+        key_cache,
+        value_cache,
+        indices,
+        page_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+    )
+    assert qsa_ops._use_sm70_qsa_xqa_page4(query, *args)
+    assert not qsa_ops._use_sm70_qsa_xqa_page4(
+        query[:-1],
+        key_cache,
+        value_cache,
+        indices[:-1],
+        page_table,
+        token_to_request[:-1],
+        query_positions[:-1],
+        sequence_lengths,
+    )
+    monkeypatch.setattr(qsa_ops, "_SM70_QSA_XQA_PAGE4", False)
+    assert not qsa_ops._use_sm70_qsa_xqa_page4(query, *args)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_qsa_xqa_page4_table_rejects_stale_or_invalid_tail_metadata():
+    indices = torch.full((1, 2051), -1, dtype=torch.int32, device="cuda")
+    indices[:, :2048] = torch.arange(2048, dtype=torch.int32, device="cuda")
+    block_table = torch.tensor([[2, 0, 1]], dtype=torch.int32, device="cuda")
+    token_to_request = torch.zeros(1, dtype=torch.int32, device="cuda")
+    query_positions = torch.tensor([3000], dtype=torch.int64, device="cuda")
+    sequence_lengths = torch.tensor([2049], dtype=torch.int32, device="cuda")
+
+    _, xqa_lengths = qsa_ops._qsa_xqa_page4_block_table(
+        indices,
+        block_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+        num_cache_blocks=3,
+        page_size=784,
+    )
+    assert xqa_lengths.item() == 2048
+
+    invalid_request = torch.full_like(token_to_request, -1)
+    _, invalid_lengths = qsa_ops._qsa_xqa_page4_block_table(
+        indices,
+        block_table,
+        invalid_request,
+        query_positions,
+        sequence_lengths,
+        num_cache_blocks=3,
+        page_size=784,
+    )
+    assert invalid_lengths.item() == 0
+
+    indices[:, 2048] = 2048
+    query_positions.fill_(2048)
+    physical_pages, tail_lengths = qsa_ops._qsa_xqa_page4_block_table(
+        indices,
+        block_table,
+        token_to_request,
+        query_positions,
+        sequence_lengths,
+        num_cache_blocks=3,
+        page_size=784,
+    )
+    assert tail_lengths.item() == 2049
+    assert physical_pages[0, 512].item() == 316
+
+
 def test_qsa_indexer_cublas_does_not_capture_decode_rows(monkeypatch):
     cache = torch.empty(2, 400, 1, 128, dtype=torch.float16)
     page_table = torch.empty(1, 2, dtype=torch.int32)