From ba759b85836fbca326567af8e06f0d434476fba7 Mon Sep 17 00:00:00 2001
From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Date: Fri, 28 Aug 2026 12:13:44 +0800
Subject: [PATCH 1/5] [WIP][Kernel][SM70] Prototype QSA page4 XQA prefill

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
---
 .../kernel/flash_decode_paged.cu              |  38 +-
 tests/models/qwen4_exp/test_qsa_ops.py        |  66 +++
 vllm/models/qwen4_exp/nvidia/ops/qsa.py       | 397 ++++++++++++++++++
 vllm/models/qwen4_exp/nvidia/qsa.py           |  10 +
 4 files changed, 503 insertions(+), 8 deletions(-)

diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu
index 6d92630656..224653d046 100644
--- a/flash-attention-v100/kernel/flash_decode_paged.cu
+++ b/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,9 +744,10 @@ __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 ||
@@ -752,7 +755,10 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector(
                 "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>(
diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py
index ebc3f3273d..3ab64fd88d 100644
--- a/tests/models/qwen4_exp/test_qsa_ops.py
+++ b/tests/models/qwen4_exp/test_qsa_ops.py
@@ -7,6 +7,7 @@
 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 +37,71 @@ 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,
+    )
+
+
 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)
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
index d62a16385d..47378e3d14 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
@@ -27,6 +27,17 @@
 _SM70_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS = int(
     os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS", str(1024**2))
 )
+_SM70_QSA_XQA_PAGE4 = os.getenv("VLLM_SM70_QSA_XQA_PAGE4", "1") == "1"
+_SM70_QSA_XQA_PAGE4_MIN_ROWS = int(
+    os.getenv("VLLM_SM70_QSA_XQA_PAGE4_MIN_ROWS", "4096")
+)
+_SM70_QSA_XQA_PAGE4_PARTITION = 256
+_SM70_QSA_XQA_PAGE4_PAGES = 513
+_SM70_QSA_XQA_PAGE4_MARKER = 1 << 30
+_SM70_QSA_XQA_PAGE4_WORKSPACES: dict[
+    tuple[int, int, int, int, int],
+    tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+] = {}
 
 
 @triton.jit
@@ -321,6 +332,97 @@ def _expand_qsa_indices_kernel(
     )
 
 
+@triton.jit
+def _qsa_xqa_page4_table_kernel(
+    indices_ptr,
+    block_table_ptr,
+    token_to_req_ptr,
+    query_positions_ptr,
+    sequence_lengths_ptr,
+    encoded_pages_ptr,
+    xqa_sequence_lengths_ptr,
+    stride_indices_row,
+    stride_table_req,
+    stride_encoded_row,
+    rows,
+    num_cache_blocks,
+    num_requests,
+    PAGE_SIZE: tl.constexpr,
+    PAGE_TABLE_WIDTH: tl.constexpr,
+    COMPLETE_PAGES: tl.constexpr,
+    OUTPUT_PAGES: tl.constexpr,
+    BLOCK_PAGES: tl.constexpr,
+    PHYSICAL_PAGE_STRIDE: tl.constexpr,
+    TAIL_MARKER: tl.constexpr,
+) -> None:
+    row = tl.program_id(0)
+    slots = tl.arange(0, BLOCK_PAGES)
+    request = tl.load(token_to_req_ptr + row)
+    safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
+    query_position = tl.load(query_positions_ptr + row)
+    sequence_length = tl.load(
+        sequence_lengths_ptr + safe_request,
+        mask=(request >= 0) & (request < num_requests),
+        other=0,
+    )
+    visible_tokens = query_position + 1
+    complete_pages = tl.minimum(
+        tl.minimum(visible_tokens // 4, sequence_length // 4),
+        COMPLETE_PAGES,
+    )
+    tail_count = visible_tokens - (visible_tokens // 4) * 4
+    is_complete = slots < complete_pages
+    is_tail = (slots == complete_pages) & (tail_count > 0)
+    selected_token = tl.load(
+        indices_ptr + row * stride_indices_row + slots * 4,
+        mask=(row < rows) & is_complete,
+        other=-1,
+    )
+    tail_token = (visible_tokens // 4) * 4
+    logical_token = tl.where(is_tail, tail_token, selected_token)
+    safe_token = tl.maximum(logical_token, 0)
+    logical_page = safe_token // PAGE_SIZE
+    page_offset = safe_token - logical_page * PAGE_SIZE
+    valid = (
+        (row < rows)
+        & (request >= 0)
+        & (request < num_requests)
+        & (logical_token >= 0)
+        & (logical_page < PAGE_TABLE_WIDTH)
+        & (is_complete | is_tail)
+    )
+    physical_page = tl.load(
+        block_table_ptr
+        + safe_request * stride_table_req
+        + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1),
+        mask=valid,
+        other=-1,
+    )
+    valid &= (physical_page >= 0) & (physical_page < num_cache_blocks)
+    physical_microblock = (
+        tl.maximum(physical_page, 0) * PHYSICAL_PAGE_STRIDE + page_offset // 4
+    )
+    encoded = tl.where(
+        valid & is_complete,
+        physical_microblock,
+        tl.where(
+            valid & is_tail,
+            physical_microblock + TAIL_MARKER,
+            2147483647,
+        ),
+    )
+    tl.store(
+        encoded_pages_ptr + row * stride_encoded_row + slots,
+        encoded,
+        mask=(row < rows) & (slots < OUTPUT_PAGES),
+    )
+    tl.store(
+        xqa_sequence_lengths_ptr + row,
+        complete_pages * 4 + tail_count,
+        mask=row < rows,
+    )
+
+
 @triton.jit
 def _qsa_sparse_paged_gqa_splitk_kernel(
     q_ptr,
@@ -1134,6 +1236,274 @@ def qsa_select_paged_tokens(
     return out
 
 
+def _qsa_xqa_page4_shape_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor | None,
+    sequence_lengths: torch.Tensor | None,
+) -> bool:
+    return (
+        query_positions is not None
+        and sequence_lengths is not None
+        and q.dtype == torch.float16
+        and k_cache.dtype == v_cache.dtype == torch.float16
+        and q.device
+        == k_cache.device
+        == v_cache.device
+        == logical_indices.device
+        == block_table.device
+        == token_to_req.device
+        == query_positions.device
+        == sequence_lengths.device
+        and q.ndim == 3
+        and q.shape[1:] == (6, 256)
+        and q.stride(2) == 1
+        and k_cache.ndim == 4
+        and v_cache.shape == k_cache.shape
+        and k_cache.shape[2:] == (1, 256)
+        and k_cache.shape[1] % 4 == 0
+        and k_cache.stride(3) == v_cache.stride(3) == 1
+        and k_cache.stride(1) == v_cache.stride(1) == 256
+        and k_cache.stride(0) == v_cache.stride(0)
+        and k_cache.stride(0) in (k_cache.shape[1] * 256, 2 * k_cache.shape[1] * 256)
+        and logical_indices.shape == (q.shape[0], 2051)
+        and logical_indices.dtype == torch.int32
+        and logical_indices.stride(1) == 1
+        and block_table.ndim == 2
+        and block_table.dtype == torch.int32
+        and block_table.stride(1) == 1
+        and token_to_req.shape == (q.shape[0],)
+        and token_to_req.dtype == torch.int32
+        and token_to_req.stride(0) == 1
+        and query_positions.shape == (q.shape[0],)
+        and query_positions.dtype == torch.int64
+        and query_positions.stride(0) == 1
+        and sequence_lengths.shape == (block_table.shape[0],)
+        and sequence_lengths.dtype == torch.int32
+        and sequence_lengths.stride(0) == 1
+        and k_cache.shape[0] * (k_cache.stride(0) // (4 * 256))
+        < _SM70_QSA_XQA_PAGE4_MARKER
+    )
+
+
+def _use_sm70_qsa_xqa_page4(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor | None,
+    sequence_lengths: torch.Tensor | None,
+) -> bool:
+    return (
+        _SM70_QSA_XQA_PAGE4
+        and current_platform.is_device_capability(70)
+        and q.shape[0] >= _SM70_QSA_XQA_PAGE4_MIN_ROWS
+        and _qsa_xqa_page4_shape_supported(
+            q,
+            k_cache,
+            v_cache,
+            logical_indices,
+            block_table,
+            token_to_req,
+            query_positions,
+            sequence_lengths,
+        )
+    )
+
+
+def _qsa_xqa_page4_block_table(
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    num_cache_blocks: int,
+    page_size: int,
+    physical_page_stride: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if physical_page_stride is None:
+        physical_page_stride = page_size // 4
+    rows = logical_indices.shape[0]
+    encoded_pages = torch.empty(
+        (rows, _SM70_QSA_XQA_PAGE4_PAGES),
+        dtype=torch.int32,
+        device=logical_indices.device,
+    )
+    xqa_sequence_lengths = torch.empty(
+        (rows,), dtype=torch.int32, device=logical_indices.device
+    )
+    _qsa_xqa_page4_table_kernel[(rows,)](
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+        encoded_pages,
+        xqa_sequence_lengths,
+        logical_indices.stride(0),
+        block_table.stride(0),
+        encoded_pages.stride(0),
+        rows,
+        num_cache_blocks,
+        block_table.shape[0],
+        PAGE_SIZE=page_size,
+        PAGE_TABLE_WIDTH=block_table.shape[1],
+        COMPLETE_PAGES=2048 // 4,
+        OUTPUT_PAGES=_SM70_QSA_XQA_PAGE4_PAGES,
+        BLOCK_PAGES=1024,
+        PHYSICAL_PAGE_STRIDE=physical_page_stride,
+        TAIL_MARKER=_SM70_QSA_XQA_PAGE4_MARKER,
+        num_warps=4,
+    )
+    sorted_pages = torch.sort(encoded_pages, dim=1).values
+    physical_pages = torch.bitwise_and(
+        sorted_pages,
+        _SM70_QSA_XQA_PAGE4_MARKER - 1,
+    )
+    return physical_pages, xqa_sequence_lengths
+
+
+def _qsa_xqa_page4_workspace(
+    q: torch.Tensor,
+    num_partitions: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    device_index = q.device.index if q.device.index is not None else -1
+    stream_id = int(torch.cuda.current_stream(q.device).cuda_stream)
+    key = (device_index, stream_id, q.shape[1], q.shape[2], num_partitions)
+    workspace = _SM70_QSA_XQA_PAGE4_WORKSPACES.get(key)
+    rows = q.shape[0]
+    if workspace is None or workspace[0] < rows:
+        capacity = 1 << (rows - 1).bit_length()
+        temporary_output = torch.empty(
+            (capacity, q.shape[1], num_partitions, q.shape[2]),
+            dtype=torch.float16,
+            device=q.device,
+        )
+        max_logits = torch.empty(
+            (capacity, q.shape[1], num_partitions),
+            dtype=torch.float32,
+            device=q.device,
+        )
+        exp_sums = torch.empty_like(max_logits)
+        active_num_partitions = torch.tensor(
+            [num_partitions], dtype=torch.int32, device=q.device
+        )
+        workspace = (
+            capacity,
+            temporary_output,
+            max_logits,
+            exp_sums,
+            active_num_partitions,
+        )
+        _SM70_QSA_XQA_PAGE4_WORKSPACES[key] = workspace
+    _, temporary_output, max_logits, exp_sums, active_num_partitions = workspace
+    return (
+        temporary_output[:rows],
+        max_logits[:rows],
+        exp_sums[:rows],
+        active_num_partitions,
+    )
+
+
+def _qsa_sparse_paged_attention_sm70_xqa_page4(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor | None:
+    try:
+        from flash_attn_v100.flash_attn_interface import flash_attn_v100_cuda
+    except ImportError:
+        logger.warning_once(
+            "SM70 QSA page4 XQA route is unavailable because Flash-V100 "
+            "could not be imported; using Triton sparse attention."
+        )
+        return None
+    if not hasattr(flash_attn_v100_cuda, "decode_paged_xqa_fwd"):
+        logger.warning_once(
+            "SM70 QSA page4 XQA route is unavailable in this Flash-V100 build; "
+            "using Triton sparse attention."
+        )
+        return None
+
+    virtual_block_table, xqa_sequence_lengths = _qsa_xqa_page4_block_table(
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+        k_cache.shape[0],
+        k_cache.shape[1],
+        k_cache.stride(0) // (4 * q.shape[2]),
+    )
+    num_partitions = math.ceil(logical_indices.shape[1] / _SM70_QSA_XQA_PAGE4_PARTITION)
+    temporary_output, max_logits, exp_sums, active_num_partitions = (
+        _qsa_xqa_page4_workspace(q, num_partitions)
+    )
+    microblock_stride = 4 * q.shape[2]
+    if k_cache.stride(0) == k_cache.shape[1] * q.shape[2]:
+        microblocks_per_cache_block = k_cache.shape[1] // 4
+        physical_k_cache = k_cache.view(
+            k_cache.shape[0] * microblocks_per_cache_block,
+            4,
+            1,
+            q.shape[2],
+        )
+        physical_v_cache = v_cache.view_as(physical_k_cache)
+    else:
+        # The local FlashAttention ABI interleaves K and V inside every
+        # physical cache block. The virtual page IDs carry that doubled block
+        # stride, while this narrow view exposes a four-token page stride to
+        # XQA. No cache data is copied or rearranged.
+        physical_shape = (k_cache.shape[0], 4, 1, q.shape[2])
+        physical_strides = (
+            microblock_stride,
+            q.shape[2],
+            q.shape[2],
+            1,
+        )
+        physical_k_cache = k_cache.as_strided(physical_shape, physical_strides)
+        physical_v_cache = v_cache.as_strided(physical_shape, physical_strides)
+    flash_attn_v100_cuda.decode_paged_xqa_fwd(
+        q,
+        physical_k_cache,
+        physical_v_cache,
+        out,
+        virtual_block_table,
+        xqa_sequence_lengths,
+        temporary_output,
+        max_logits,
+        exp_sums,
+        active_num_partitions,
+        q.shape[2] ** -0.5,
+        _SM70_QSA_XQA_PAGE4_PARTITION,
+        num_partitions,
+        "auto",
+        1.0,
+        1.0,
+        -1,
+        -1,
+        0,
+    )
+    logger.info_once(
+        "Using SM70 QSA Flash-V100 XQA page4 prefill route (rows=%d, partitions=%d).",
+        q.shape[0],
+        num_partitions,
+    )
+    return out
+
+
 def qsa_sparse_paged_attention(
     q: torch.Tensor,
     k_cache: torch.Tensor,
@@ -1142,6 +1512,8 @@ def qsa_sparse_paged_attention(
     block_table: torch.Tensor,
     token_to_req: torch.Tensor,
     out: torch.Tensor | None = None,
+    query_positions: torch.Tensor | None = None,
+    sequence_lengths: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Run sparse GQA directly over paged FP16/BF16 K/V caches."""
 
@@ -1180,6 +1552,31 @@ def qsa_sparse_paged_attention(
     if not q.shape[0]:
         return out
 
+    if _use_sm70_qsa_xqa_page4(
+        q,
+        k_cache,
+        v_cache,
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+    ):
+        assert query_positions is not None and sequence_lengths is not None
+        xqa_output = _qsa_sparse_paged_attention_sm70_xqa_page4(
+            q,
+            k_cache,
+            v_cache,
+            logical_indices,
+            block_table,
+            token_to_req,
+            query_positions,
+            sequence_lengths,
+            out,
+        )
+        if xqa_output is not None:
+            return xqa_output
+
     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]
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
index 25e9ab967d..08a58b43a5 100644
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
@@ -126,6 +126,8 @@ def forward_qsa(
         attn_metadata: FlashAttentionMetadata,
         output: torch.Tensor,
         token_to_req: torch.Tensor,
+        query_positions: torch.Tensor | None = None,
+        sequence_lengths: torch.Tensor | None = None,
         output_scale: torch.Tensor | None = None,
         output_block_scale: torch.Tensor | None = None,
     ) -> torch.Tensor:
@@ -160,6 +162,11 @@ def forward_qsa(
 
         from .ops.qsa import qsa_sparse_paged_attention
 
+        qsa_metadata: dict[str, torch.Tensor] = {}
+        if query_positions is not None:
+            qsa_metadata["query_positions"] = query_positions[:num_tokens]
+        if sequence_lengths is not None:
+            qsa_metadata["sequence_lengths"] = sequence_lengths
         qsa_sparse_paged_attention(
             query[:num_tokens],
             key_cache,
@@ -168,6 +175,7 @@ def forward_qsa(
             attn_metadata.block_table,
             token_to_req,
             output[:num_tokens],
+            **qsa_metadata,
         )
         return output
 
@@ -424,6 +432,8 @@ def _run_qsa(
             main_metadata,
             output,
             token_to_req=side_metadata.token_to_req,
+            query_positions=side_metadata.logical_positions,
+            sequence_lengths=side_metadata.seq_lens,
         )
 
     def forward(

From f0028febea2b48d3370bfd4ff99dc1c151b056bb Mon Sep 17 00:00:00 2001
From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Date: Fri, 28 Aug 2026 13:07:02 +0800
Subject: [PATCH 2/5] [Kernel][SM70] Start #378 latest-main repair

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

From 59efa8fe175dafab08e7ede85a92a5b48a936120 Mon Sep 17 00:00:00 2001
From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Date: Fri, 28 Aug 2026 12:13:44 +0800
Subject: [PATCH 3/5] [WIP][Kernel][SM70] Prototype QSA page4 XQA prefill

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
---
 .../kernel/flash_decode_paged.cu              |  38 +-
 tests/models/qwen4_exp/test_qsa_ops.py        |  66 +++
 vllm/models/qwen4_exp/nvidia/ops/qsa.py       | 397 ++++++++++++++++++
 vllm/models/qwen4_exp/nvidia/qsa.py           |  10 +
 4 files changed, 503 insertions(+), 8 deletions(-)

diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu
index 6d92630656..224653d046 100644
--- a/flash-attention-v100/kernel/flash_decode_paged.cu
+++ b/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,9 +744,10 @@ __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 ||
@@ -752,7 +755,10 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector(
                 "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>(
diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py
index ebc3f3273d..3ab64fd88d 100644
--- a/tests/models/qwen4_exp/test_qsa_ops.py
+++ b/tests/models/qwen4_exp/test_qsa_ops.py
@@ -7,6 +7,7 @@
 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 +37,71 @@ 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,
+    )
+
+
 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)
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
index d62a16385d..47378e3d14 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
@@ -27,6 +27,17 @@
 _SM70_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS = int(
     os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS", str(1024**2))
 )
+_SM70_QSA_XQA_PAGE4 = os.getenv("VLLM_SM70_QSA_XQA_PAGE4", "1") == "1"
+_SM70_QSA_XQA_PAGE4_MIN_ROWS = int(
+    os.getenv("VLLM_SM70_QSA_XQA_PAGE4_MIN_ROWS", "4096")
+)
+_SM70_QSA_XQA_PAGE4_PARTITION = 256
+_SM70_QSA_XQA_PAGE4_PAGES = 513
+_SM70_QSA_XQA_PAGE4_MARKER = 1 << 30
+_SM70_QSA_XQA_PAGE4_WORKSPACES: dict[
+    tuple[int, int, int, int, int],
+    tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+] = {}
 
 
 @triton.jit
@@ -321,6 +332,97 @@ def _expand_qsa_indices_kernel(
     )
 
 
+@triton.jit
+def _qsa_xqa_page4_table_kernel(
+    indices_ptr,
+    block_table_ptr,
+    token_to_req_ptr,
+    query_positions_ptr,
+    sequence_lengths_ptr,
+    encoded_pages_ptr,
+    xqa_sequence_lengths_ptr,
+    stride_indices_row,
+    stride_table_req,
+    stride_encoded_row,
+    rows,
+    num_cache_blocks,
+    num_requests,
+    PAGE_SIZE: tl.constexpr,
+    PAGE_TABLE_WIDTH: tl.constexpr,
+    COMPLETE_PAGES: tl.constexpr,
+    OUTPUT_PAGES: tl.constexpr,
+    BLOCK_PAGES: tl.constexpr,
+    PHYSICAL_PAGE_STRIDE: tl.constexpr,
+    TAIL_MARKER: tl.constexpr,
+) -> None:
+    row = tl.program_id(0)
+    slots = tl.arange(0, BLOCK_PAGES)
+    request = tl.load(token_to_req_ptr + row)
+    safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
+    query_position = tl.load(query_positions_ptr + row)
+    sequence_length = tl.load(
+        sequence_lengths_ptr + safe_request,
+        mask=(request >= 0) & (request < num_requests),
+        other=0,
+    )
+    visible_tokens = query_position + 1
+    complete_pages = tl.minimum(
+        tl.minimum(visible_tokens // 4, sequence_length // 4),
+        COMPLETE_PAGES,
+    )
+    tail_count = visible_tokens - (visible_tokens // 4) * 4
+    is_complete = slots < complete_pages
+    is_tail = (slots == complete_pages) & (tail_count > 0)
+    selected_token = tl.load(
+        indices_ptr + row * stride_indices_row + slots * 4,
+        mask=(row < rows) & is_complete,
+        other=-1,
+    )
+    tail_token = (visible_tokens // 4) * 4
+    logical_token = tl.where(is_tail, tail_token, selected_token)
+    safe_token = tl.maximum(logical_token, 0)
+    logical_page = safe_token // PAGE_SIZE
+    page_offset = safe_token - logical_page * PAGE_SIZE
+    valid = (
+        (row < rows)
+        & (request >= 0)
+        & (request < num_requests)
+        & (logical_token >= 0)
+        & (logical_page < PAGE_TABLE_WIDTH)
+        & (is_complete | is_tail)
+    )
+    physical_page = tl.load(
+        block_table_ptr
+        + safe_request * stride_table_req
+        + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1),
+        mask=valid,
+        other=-1,
+    )
+    valid &= (physical_page >= 0) & (physical_page < num_cache_blocks)
+    physical_microblock = (
+        tl.maximum(physical_page, 0) * PHYSICAL_PAGE_STRIDE + page_offset // 4
+    )
+    encoded = tl.where(
+        valid & is_complete,
+        physical_microblock,
+        tl.where(
+            valid & is_tail,
+            physical_microblock + TAIL_MARKER,
+            2147483647,
+        ),
+    )
+    tl.store(
+        encoded_pages_ptr + row * stride_encoded_row + slots,
+        encoded,
+        mask=(row < rows) & (slots < OUTPUT_PAGES),
+    )
+    tl.store(
+        xqa_sequence_lengths_ptr + row,
+        complete_pages * 4 + tail_count,
+        mask=row < rows,
+    )
+
+
 @triton.jit
 def _qsa_sparse_paged_gqa_splitk_kernel(
     q_ptr,
@@ -1134,6 +1236,274 @@ def qsa_select_paged_tokens(
     return out
 
 
+def _qsa_xqa_page4_shape_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor | None,
+    sequence_lengths: torch.Tensor | None,
+) -> bool:
+    return (
+        query_positions is not None
+        and sequence_lengths is not None
+        and q.dtype == torch.float16
+        and k_cache.dtype == v_cache.dtype == torch.float16
+        and q.device
+        == k_cache.device
+        == v_cache.device
+        == logical_indices.device
+        == block_table.device
+        == token_to_req.device
+        == query_positions.device
+        == sequence_lengths.device
+        and q.ndim == 3
+        and q.shape[1:] == (6, 256)
+        and q.stride(2) == 1
+        and k_cache.ndim == 4
+        and v_cache.shape == k_cache.shape
+        and k_cache.shape[2:] == (1, 256)
+        and k_cache.shape[1] % 4 == 0
+        and k_cache.stride(3) == v_cache.stride(3) == 1
+        and k_cache.stride(1) == v_cache.stride(1) == 256
+        and k_cache.stride(0) == v_cache.stride(0)
+        and k_cache.stride(0) in (k_cache.shape[1] * 256, 2 * k_cache.shape[1] * 256)
+        and logical_indices.shape == (q.shape[0], 2051)
+        and logical_indices.dtype == torch.int32
+        and logical_indices.stride(1) == 1
+        and block_table.ndim == 2
+        and block_table.dtype == torch.int32
+        and block_table.stride(1) == 1
+        and token_to_req.shape == (q.shape[0],)
+        and token_to_req.dtype == torch.int32
+        and token_to_req.stride(0) == 1
+        and query_positions.shape == (q.shape[0],)
+        and query_positions.dtype == torch.int64
+        and query_positions.stride(0) == 1
+        and sequence_lengths.shape == (block_table.shape[0],)
+        and sequence_lengths.dtype == torch.int32
+        and sequence_lengths.stride(0) == 1
+        and k_cache.shape[0] * (k_cache.stride(0) // (4 * 256))
+        < _SM70_QSA_XQA_PAGE4_MARKER
+    )
+
+
+def _use_sm70_qsa_xqa_page4(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor | None,
+    sequence_lengths: torch.Tensor | None,
+) -> bool:
+    return (
+        _SM70_QSA_XQA_PAGE4
+        and current_platform.is_device_capability(70)
+        and q.shape[0] >= _SM70_QSA_XQA_PAGE4_MIN_ROWS
+        and _qsa_xqa_page4_shape_supported(
+            q,
+            k_cache,
+            v_cache,
+            logical_indices,
+            block_table,
+            token_to_req,
+            query_positions,
+            sequence_lengths,
+        )
+    )
+
+
+def _qsa_xqa_page4_block_table(
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    num_cache_blocks: int,
+    page_size: int,
+    physical_page_stride: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if physical_page_stride is None:
+        physical_page_stride = page_size // 4
+    rows = logical_indices.shape[0]
+    encoded_pages = torch.empty(
+        (rows, _SM70_QSA_XQA_PAGE4_PAGES),
+        dtype=torch.int32,
+        device=logical_indices.device,
+    )
+    xqa_sequence_lengths = torch.empty(
+        (rows,), dtype=torch.int32, device=logical_indices.device
+    )
+    _qsa_xqa_page4_table_kernel[(rows,)](
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+        encoded_pages,
+        xqa_sequence_lengths,
+        logical_indices.stride(0),
+        block_table.stride(0),
+        encoded_pages.stride(0),
+        rows,
+        num_cache_blocks,
+        block_table.shape[0],
+        PAGE_SIZE=page_size,
+        PAGE_TABLE_WIDTH=block_table.shape[1],
+        COMPLETE_PAGES=2048 // 4,
+        OUTPUT_PAGES=_SM70_QSA_XQA_PAGE4_PAGES,
+        BLOCK_PAGES=1024,
+        PHYSICAL_PAGE_STRIDE=physical_page_stride,
+        TAIL_MARKER=_SM70_QSA_XQA_PAGE4_MARKER,
+        num_warps=4,
+    )
+    sorted_pages = torch.sort(encoded_pages, dim=1).values
+    physical_pages = torch.bitwise_and(
+        sorted_pages,
+        _SM70_QSA_XQA_PAGE4_MARKER - 1,
+    )
+    return physical_pages, xqa_sequence_lengths
+
+
+def _qsa_xqa_page4_workspace(
+    q: torch.Tensor,
+    num_partitions: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    device_index = q.device.index if q.device.index is not None else -1
+    stream_id = int(torch.cuda.current_stream(q.device).cuda_stream)
+    key = (device_index, stream_id, q.shape[1], q.shape[2], num_partitions)
+    workspace = _SM70_QSA_XQA_PAGE4_WORKSPACES.get(key)
+    rows = q.shape[0]
+    if workspace is None or workspace[0] < rows:
+        capacity = 1 << (rows - 1).bit_length()
+        temporary_output = torch.empty(
+            (capacity, q.shape[1], num_partitions, q.shape[2]),
+            dtype=torch.float16,
+            device=q.device,
+        )
+        max_logits = torch.empty(
+            (capacity, q.shape[1], num_partitions),
+            dtype=torch.float32,
+            device=q.device,
+        )
+        exp_sums = torch.empty_like(max_logits)
+        active_num_partitions = torch.tensor(
+            [num_partitions], dtype=torch.int32, device=q.device
+        )
+        workspace = (
+            capacity,
+            temporary_output,
+            max_logits,
+            exp_sums,
+            active_num_partitions,
+        )
+        _SM70_QSA_XQA_PAGE4_WORKSPACES[key] = workspace
+    _, temporary_output, max_logits, exp_sums, active_num_partitions = workspace
+    return (
+        temporary_output[:rows],
+        max_logits[:rows],
+        exp_sums[:rows],
+        active_num_partitions,
+    )
+
+
+def _qsa_sparse_paged_attention_sm70_xqa_page4(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    logical_indices: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_req: torch.Tensor,
+    query_positions: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor | None:
+    try:
+        from flash_attn_v100.flash_attn_interface import flash_attn_v100_cuda
+    except ImportError:
+        logger.warning_once(
+            "SM70 QSA page4 XQA route is unavailable because Flash-V100 "
+            "could not be imported; using Triton sparse attention."
+        )
+        return None
+    if not hasattr(flash_attn_v100_cuda, "decode_paged_xqa_fwd"):
+        logger.warning_once(
+            "SM70 QSA page4 XQA route is unavailable in this Flash-V100 build; "
+            "using Triton sparse attention."
+        )
+        return None
+
+    virtual_block_table, xqa_sequence_lengths = _qsa_xqa_page4_block_table(
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+        k_cache.shape[0],
+        k_cache.shape[1],
+        k_cache.stride(0) // (4 * q.shape[2]),
+    )
+    num_partitions = math.ceil(logical_indices.shape[1] / _SM70_QSA_XQA_PAGE4_PARTITION)
+    temporary_output, max_logits, exp_sums, active_num_partitions = (
+        _qsa_xqa_page4_workspace(q, num_partitions)
+    )
+    microblock_stride = 4 * q.shape[2]
+    if k_cache.stride(0) == k_cache.shape[1] * q.shape[2]:
+        microblocks_per_cache_block = k_cache.shape[1] // 4
+        physical_k_cache = k_cache.view(
+            k_cache.shape[0] * microblocks_per_cache_block,
+            4,
+            1,
+            q.shape[2],
+        )
+        physical_v_cache = v_cache.view_as(physical_k_cache)
+    else:
+        # The local FlashAttention ABI interleaves K and V inside every
+        # physical cache block. The virtual page IDs carry that doubled block
+        # stride, while this narrow view exposes a four-token page stride to
+        # XQA. No cache data is copied or rearranged.
+        physical_shape = (k_cache.shape[0], 4, 1, q.shape[2])
+        physical_strides = (
+            microblock_stride,
+            q.shape[2],
+            q.shape[2],
+            1,
+        )
+        physical_k_cache = k_cache.as_strided(physical_shape, physical_strides)
+        physical_v_cache = v_cache.as_strided(physical_shape, physical_strides)
+    flash_attn_v100_cuda.decode_paged_xqa_fwd(
+        q,
+        physical_k_cache,
+        physical_v_cache,
+        out,
+        virtual_block_table,
+        xqa_sequence_lengths,
+        temporary_output,
+        max_logits,
+        exp_sums,
+        active_num_partitions,
+        q.shape[2] ** -0.5,
+        _SM70_QSA_XQA_PAGE4_PARTITION,
+        num_partitions,
+        "auto",
+        1.0,
+        1.0,
+        -1,
+        -1,
+        0,
+    )
+    logger.info_once(
+        "Using SM70 QSA Flash-V100 XQA page4 prefill route (rows=%d, partitions=%d).",
+        q.shape[0],
+        num_partitions,
+    )
+    return out
+
+
 def qsa_sparse_paged_attention(
     q: torch.Tensor,
     k_cache: torch.Tensor,
@@ -1142,6 +1512,8 @@ def qsa_sparse_paged_attention(
     block_table: torch.Tensor,
     token_to_req: torch.Tensor,
     out: torch.Tensor | None = None,
+    query_positions: torch.Tensor | None = None,
+    sequence_lengths: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Run sparse GQA directly over paged FP16/BF16 K/V caches."""
 
@@ -1180,6 +1552,31 @@ def qsa_sparse_paged_attention(
     if not q.shape[0]:
         return out
 
+    if _use_sm70_qsa_xqa_page4(
+        q,
+        k_cache,
+        v_cache,
+        logical_indices,
+        block_table,
+        token_to_req,
+        query_positions,
+        sequence_lengths,
+    ):
+        assert query_positions is not None and sequence_lengths is not None
+        xqa_output = _qsa_sparse_paged_attention_sm70_xqa_page4(
+            q,
+            k_cache,
+            v_cache,
+            logical_indices,
+            block_table,
+            token_to_req,
+            query_positions,
+            sequence_lengths,
+            out,
+        )
+        if xqa_output is not None:
+            return xqa_output
+
     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]
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
index 25e9ab967d..08a58b43a5 100644
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
@@ -126,6 +126,8 @@ def forward_qsa(
         attn_metadata: FlashAttentionMetadata,
         output: torch.Tensor,
         token_to_req: torch.Tensor,
+        query_positions: torch.Tensor | None = None,
+        sequence_lengths: torch.Tensor | None = None,
         output_scale: torch.Tensor | None = None,
         output_block_scale: torch.Tensor | None = None,
     ) -> torch.Tensor:
@@ -160,6 +162,11 @@ def forward_qsa(
 
         from .ops.qsa import qsa_sparse_paged_attention
 
+        qsa_metadata: dict[str, torch.Tensor] = {}
+        if query_positions is not None:
+            qsa_metadata["query_positions"] = query_positions[:num_tokens]
+        if sequence_lengths is not None:
+            qsa_metadata["sequence_lengths"] = sequence_lengths
         qsa_sparse_paged_attention(
             query[:num_tokens],
             key_cache,
@@ -168,6 +175,7 @@ def forward_qsa(
             attn_metadata.block_table,
             token_to_req,
             output[:num_tokens],
+            **qsa_metadata,
         )
         return output
 
@@ -424,6 +432,8 @@ def _run_qsa(
             main_metadata,
             output,
             token_to_req=side_metadata.token_to_req,
+            query_positions=side_metadata.logical_positions,
+            sequence_lengths=side_metadata.seq_lens,
         )
 
     def forward(

From 62c0832ce73fd33e725d559ab7089ac24eda5842 Mon Sep 17 00:00:00 2001
From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Date: Fri, 28 Aug 2026 13:13:56 +0800
Subject: [PATCH 4/5] [Kernel][SM70] Harden QSA page4 XQA prefill

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
---
 docs/design/sm70_v100_migration_control.md |  32 +++++++
 tests/models/qwen4_exp/test_qsa_cache.py   |  11 +++
 tests/models/qwen4_exp/test_qsa_ops.py     | 100 +++++++++++++++++++++
 vllm/models/qwen4_exp/nvidia/ops/qsa.py    |  30 +++++--
 4 files changed, 166 insertions(+), 7 deletions(-)

diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md
index b3e6d4d9c6..70abc7d2bf 100644
--- a/docs/design/sm70_v100_migration_control.md
+++ b/docs/design/sm70_v100_migration_control.md
@@ -44189,3 +44189,35 @@ 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.
+- 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.
diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py
index d3cfea8d5f..2879118dec 100644
--- a/tests/models/qwen4_exp/test_qsa_cache.py
+++ b/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,6 +65,9 @@ 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
@@ -70,6 +75,8 @@ def fake_sparse_attention(
         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))
diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py
index 3ab64fd88d..3f521b3c15 100644
--- a/tests/models/qwen4_exp/test_qsa_ops.py
+++ b/tests/models/qwen4_exp/test_qsa_ops.py
@@ -1,6 +1,7 @@
 # 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
@@ -100,6 +101,105 @@ def test_qsa_xqa_page4_accepts_only_exact_sm70_prefill_shape():
         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):
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
index 47378e3d14..ce64cc09da 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
@@ -358,36 +358,52 @@ def _qsa_xqa_page4_table_kernel(
     row = tl.program_id(0)
     slots = tl.arange(0, BLOCK_PAGES)
     request = tl.load(token_to_req_ptr + row)
+    request_is_valid = (request >= 0) & (request < num_requests)
     safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
     query_position = tl.load(query_positions_ptr + row)
     sequence_length = tl.load(
         sequence_lengths_ptr + safe_request,
-        mask=(request >= 0) & (request < num_requests),
+        mask=request_is_valid,
         other=0,
     )
-    visible_tokens = query_position + 1
+    # Padded graph rows use position -1. Clamp malformed or stale positions to
+    # the request's live sequence so they cannot expose a synthetic tail page.
+    visible_tokens = tl.minimum(
+        tl.maximum(query_position + 1, 0),
+        sequence_length,
+    )
     complete_pages = tl.minimum(
         tl.minimum(visible_tokens // 4, sequence_length // 4),
         COMPLETE_PAGES,
     )
     tail_count = visible_tokens - (visible_tokens // 4) * 4
     is_complete = slots < complete_pages
-    is_tail = (slots == complete_pages) & (tail_count > 0)
     selected_token = tl.load(
         indices_ptr + row * stride_indices_row + slots * 4,
         mask=(row < rows) & is_complete,
         other=-1,
     )
     tail_token = (visible_tokens // 4) * 4
-    logical_token = tl.where(is_tail, tail_token, selected_token)
+    selected_tail_token = tl.load(
+        indices_ptr + row * stride_indices_row + complete_pages * 4,
+        mask=(row < rows) & (tail_count > 0),
+        other=-1,
+    )
+    tail_is_valid = (
+        (tail_count > 0)
+        & (selected_tail_token == tail_token)
+        & (selected_tail_token < sequence_length)
+    )
+    is_tail = (slots == complete_pages) & tail_is_valid
+    logical_token = tl.where(is_tail, selected_tail_token, selected_token)
     safe_token = tl.maximum(logical_token, 0)
     logical_page = safe_token // PAGE_SIZE
     page_offset = safe_token - logical_page * PAGE_SIZE
     valid = (
         (row < rows)
-        & (request >= 0)
-        & (request < num_requests)
+        & request_is_valid
         & (logical_token >= 0)
+        & (logical_token < sequence_length)
         & (logical_page < PAGE_TABLE_WIDTH)
         & (is_complete | is_tail)
     )
@@ -418,7 +434,7 @@ def _qsa_xqa_page4_table_kernel(
     )
     tl.store(
         xqa_sequence_lengths_ptr + row,
-        complete_pages * 4 + tail_count,
+        complete_pages * 4 + tl.where(tail_is_valid, tail_count, 0),
         mask=row < rows,
     )
 

From 8cd30655d13d4d91d83b329924850180a0b72461 Mon Sep 17 00:00:00 2001
From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Date: Fri, 28 Aug 2026 13:21:47 +0800
Subject: [PATCH 5/5] [Doc][SM70] Record page16 QSA XQA integration

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
---
 docs/design/sm70_v100_migration_control.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md
index 70abc7d2bf..7692adfb43 100644
--- a/docs/design/sm70_v100_migration_control.md
+++ b/docs/design/sm70_v100_migration_control.md
@@ -44209,6 +44209,11 @@ Interpretation:
   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`.