From 5fa8a605dab12cc9ee15459d9ac6b88d95c7be3a Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:46:31 -0700 Subject: [PATCH] [Bugfix][SM70] Stabilize QSA page4 order under KV relocation Keep physical deduplication and masks while ordering grouped and XQA plans by logical ownership. Add allocation-invariance regressions and document the causal investigation, upstream alternatives, and measured cost. Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- .../sm70_qsa_page4_allocation_invariance.md | 201 +++++++++++ .../kernel/flash_decode_paged.cu | 102 ++++-- tests/kernels/test_sm70_qsa_page4_plan.py | 311 ++++++++++++++++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 15 +- 4 files changed, 597 insertions(+), 32 deletions(-) create mode 100644 docs/design/sm70_qsa_page4_allocation_invariance.md create mode 100644 tests/kernels/test_sm70_qsa_page4_plan.py diff --git a/docs/design/sm70_qsa_page4_allocation_invariance.md b/docs/design/sm70_qsa_page4_allocation_invariance.md new file mode 100644 index 0000000000..a6e9b35008 --- /dev/null +++ b/docs/design/sm70_qsa_page4_allocation_invariance.md @@ -0,0 +1,201 @@ +# SM70 QSA page4 allocation invariance + +## Defect and contract + +The selected logical tokens and their K/V values can remain identical while +the physical KV page allocation changes between requests. The grouped page4 +planner formerly emitted pages in physical hash-slot order, within active-row +categories. Changing allocation therefore changed the online-softmax and +tensor-core reduction order. This can perturb FP16 attention outputs and +downstream logits without an incorrect selected set or incorrect K/V values. + +The fix guarantees a stable logical plan for fixed query grouping, logical +selections, visibility, and page-alias relationships. Physical addresses are +payload, not ordering keys. It does **not** promise bitwise equality across +different batch shapes, changed prefix-sharing relationships, different +attention implementations, or quantization formats. + +## Implementation history + +| Change | Role in this path | +|---|---| +| [#378](https://github.com/1CatAI/1Cat-vLLM/pull/378), integrated with [#382](https://github.com/1CatAI/1Cat-vLLM/pull/382) | Introduced the SM70 Flash-V100 virtual-page4 path; the single-row path sorted physical microblock IDs for locality. | +| [#387](https://github.com/1CatAI/1Cat-vLLM/pull/387), commit `94ce990ce85abeb12e3948ee1c4f518594bacf25` | Added eight-query K/V sharing, physical-page hash deduplication, mask merging, category packing, and hash-slot-order emission. | +| [#466](https://github.com/1CatAI/1Cat-vLLM/pull/466), commit `186c9e3585b109b88c070603a181dd6826400153` | Lowered the default page4 admission from 4096 to 64 actual query rows and admitted grouped prefixes with XQA tails. It widened exposure, rather than introducing the hash planner. | + +These are source-history findings, not a GPU regression bisection of every +historical release. This defect is independent of AWQ grouped **decode**: +the observed first difference arose during QSA grouped **prefill**, with the +experimental AWQ grouped-decode gate disabled throughout the diagnosis. + +## Causal investigation + +The diagnostic contract used one frozen AWQ checkpoint/runtime, four V100s, +TP4, MTP0, FP16 activations and KV, MRv2, FULL_AND_PIECEWISE graphs, +8192 batched tokens, prefix caching off, asynchronous scheduling off, +fixed prompt token IDs, fixed enqueue order, and fixed request-slot order. +The four prompts contained 42/41/58/44 tokens. Each diagnostic request generated +only its first token, with temperature 0, `min_tokens=0`, and `ignore_eos=false`. +This was numerical diagnosis, not an output-quality or throughput benchmark. + +1. In repeated C4 batches, all four ranks had identical inputs through the first + three layers. The first different boundary was Layer3 QSA (zero-based layer + numbering). Attention output max absolute difference was `0.0009765625`; + full-logit max absolute difference was `0.01171875`, with unchanged argmax. +2. Finer observations found identical Q/K/V, logical positions, request mapping, + selected token indices, and all effective K/V read in logical order. + Physical block tables differed. All 23 query groups retained the same + logical page/mask sets, but 17 had a different order. QSA core max absolute + difference was `0.00048828125` on each rank. +3. Holding physical allocation fixed, 24 extra attention replays were bitwise + identical. The measured symptom was not fixed-input kernel randomness. +4. Diagnostic CPU sorting of only Layer3's plan restored equality there and + moved the first difference to Layer7, the next QSA layer. +5. Applying the same control to all 12 QSA layers made all four ranks' 175 + observed boundaries and complete logits bitwise identical. +6. Removing sorting while retaining observation and CPU synchronization brought + the Layer3 difference back. Extra synchronization alone did not explain the + result. + +The actual C4 step had 185 query rows: 184 grouped rows plus one XQA tail. +The short C1 control had 44 rows and did not enter page4. Long C1 prefill can +still enter the path; concurrency labels are not route evidence. + +Instrumentation can affect compilation boundaries. Causal attribution rests on +same-instrumentation intervention/reversal and local fixed-input replays, not +on interpreting diagnostic timings as performance. The earlier freely batched +70-versus-61-token answer divergence and the independent AWQ W13 operator +microdifference are not claimed to be completely explained by this experiment. + +## Upstream comparison and duplicate-work check + +Original vLLM QSA iterates logical selection positions, then uses the block +table for addressing. It does not use this SM70 physical-hash union planner. +Related upstream work must not be conflated with this defect: + +- [1Cat #394](https://github.com/1CatAI/1Cat-vLLM/pull/394) already provides + exact lexicographic QSA top-k selection on the tested SM70 contract. + Those selections were bitwise equal during this investigation. +- [vLLM #55122](https://github.com/vllm-project/vllm/pull/55122) addresses + generic `persistent_topk` membership/tie/order nondeterminism, not ordering + subsequently introduced by the page4 planner. +- [vLLM #54873](https://github.com/vllm-project/vllm/pull/54873) skips unused + sparse-attention selection entries and tunes launch profiles. It does not + fix the 1Cat planner. +- [vLLM RFC #55394](https://github.com/vllm-project/vllm/issues/55394) proposes + a related query-tile union. Its prototype sorts logical blocks before + physical mapping. That principle is useful here; the GB10 single-request + prototype is not a ready-made SM70 concurrent replacement and is not ported. + +At the 2026-09-04 duplicate check, PR #55122 was open, PR #54873 was merged, and +RFC #55394 was open. No direct repair of this planner was found. The existing +Triton fallback remains a correctness/performance control; upstream use does +not by itself establish V100 performance or cross-batch invariance. + +## Narrow repair + +The grouped planner retains physical-page deduplication and OR-merged token +masks. Each entry also records its smallest logical owner: +`(first contributing query within the group, logical four-token block)`. +An atomic minimum makes shared-page ownership independent of insertion order. +CUB block radix sort orders entries by active-row category and logical owner. +The original category packing, eight-page padding, and attention kernel remain. +This uses CUB's existing sorting primitive, not a new sorting algorithm. + +The single-row XQA path, including non-grouped tails, also needs repair. Its +existing GPU `torch.sort` now sorts packed logical keys carrying physical IDs +as payload. The causal partial page remains after complete pages and invalid +slots remain last. Only integer planning changes; attention arithmetic, +weights, quantization, scheduler policy, and route thresholds are unchanged. + +The grouped hash and owner arrays occupy 96 KiB of shared memory. Once entries +are held in registers, that storage is reused for CUB sorting and category +scans. There is no added global-memory grouped workspace or weight/KV copy. +The single-row sorting keys grow from int32 to int64, so temporary metadata +memory is not claimed to be unchanged. Resource use and latency require GPU +measurement; absence of a global grouped allocation does not imply zero cost. +The tested CUDA 12.8 SM70 binary reports 128 registers per planner thread, +zero local memory, and zero stack bytes. Dynamic shared memory is 96 KiB per +CTA; the resource dump's zero static-shared value does not include it. + +## Validation + +The regression is `tests/kernels/test_sm70_qsa_page4_plan.py`. It checks exact +logical-reference plans, shared physical pages, collisions, invalid rows, +all-empty groups, wide unions, selection permutations, page sizes 4/16/32, +graph replay with relocated inputs, contiguous/interleaved FP16 and E4M3 KV, and a +185-row grouped-plus-XQA-tail batch. + +```bash +.venv/bin/python -m pytest -q tests/kernels/test_sm70_qsa_page4_plan.py +.venv/bin/python -m pytest -q tests/models/qwen4_exp/test_qsa_ops.py +``` + +Before this repair, the initial 20-case regression produced 17 failures and +three passes (the all-empty controls). Both attention relocation checks failed +at the bitwise-output assertion. The later single-row and 185-row integration +tests were added separately and must not be counted as part of that initial run. + +The candidate builds with CUDA 12.8 / SM70. On one V100, all 26 new regression +cases and all 14 existing QSA-ops tests pass (40 total). All applicable +pre-commit hooks pass. Captured Layer3 inputs from all four TP ranks reproduce +the baseline relocation difference and become bitwise equal after the repair: + +| Rank | Baseline different output elements | Baseline max absolute difference | Repaired different elements | +|---|---:|---:|---:| +| 0 | 1625 | 0.00048828125 | 0 | +| 1 | 1852 | 0.00048828125 | 0 | +| 2 | 1515 | 0.00048828125 | 0 | +| 3 | 1866 | 0.00048828125 | 0 | + +The existing Triton fallback is also bitwise allocation-invariant for these +four captured pairs. This is direct replay evidence, in addition to the +upstream source inspection above. + +### Bounded operator timings + +CUDA-event medians, three warmups and 20 samples per call, one V100, FP16 KV, +Hq/Hkv/D = 6/1/256, page size 16. The synthetic cases have four 64K requests, +512 shared selected logical blocks per query, randomized physical allocation, +and the stated total query-row count. They favor K/V sharing and are **not** +a server concurrency contract or full-model throughput measurement. + +| Input | Old page4 (ms) | Repaired page4 (ms) | Change | Existing Triton (ms) | +|---|---:|---:|---:|---:| +| Captured 185-row step, allocation A | 0.3518 | 0.3615 | +2.8% | 2.6552 | +| Captured 185-row step, allocation B | 0.3600 | 0.3635 | +1.0% | 2.4049 | +| Synthetic 1024 rows | 1.5130 | 1.5811 | +4.5% | 14.3713 | +| Synthetic 8192 rows | 9.1136 | 9.7823 | +7.3% | 99.3587 | + +The Triton column comes from the candidate probe. Baseline probe Triton +medians were 2.6563/2.6588/14.4645/99.3930 ms respectively; short host-driven +operator calls show noise and these single-session numbers are not confidence +intervals. For 8192 rows, planner-only time grows from 0.6610 to 1.3251 ms. +The cost is real; retaining page4 with logical ordering is still preferable +to the measured full Triton fallback on this V100 contract. Synthetic +reference errors versus Triton remain small (repaired maximum absolute +difference 0.0001220703125 in both cases), not bitwise cross-kernel equality. + +### Natural-EOS full-model sanity + +A bounded AWQ / TP4 / MTP0 / FP16-KV run used the same frozen compatible +runtime as the diagnosis, including its existing AWQ wrapper admission repair +and disabled experimental AWQ grouped-decode gate. No model trace, sampler +replacement, or attention intervention was installed. Prompt IDs, enqueue +order, and free request-slot order were fixed; each C1/C4/C8 shape ran twice +with temperature 0, `min_tokens=0`, `ignore_eos=false`, and a 96-token limit. + +All 26 request outputs stopped naturally, had finite reported logprobs, and +passed basic answer checks. Token IDs matched exactly between the two runs +of each shape. Worker provenance confirmed the repaired binary and Python +source on all four ranks. Actual page4 route logs reported 184 grouped + 1 +XQA row for C4 and 368 grouped + 2 XQA rows for C8. KV capacity was 386,392 +tokens, unchanged from the prior same-configuration run. + +Cross-shape/position differences **remain**: the same open-ended Chinese +prompt produced 61 tokens at C1, 70 at C4, and 61/71 at its two C8 positions. +Each position reproduced its own sequence on the repeat. This run validates +short same-shape repeatability, not cross-batch or cross-position invariance; +the remaining divergence has not been localized by this model sanity check. +It also does not establish long-context quality, NVFP4 model acceptance, E2E +performance, or a resolution of every prior generation divergence. diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 7384538f75..aa43664a5a 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -14,6 +14,7 @@ #include #include #include +#include #include "fp8_kv_utils.cuh" #include "fused_mma.h" @@ -3659,10 +3660,23 @@ constexpr int kGroupedSparseQueries = 8; constexpr int kGroupedSparsePlannerThreads = 512; constexpr int kGroupedSparseHashCapacity = 8192; constexpr unsigned long long kGroupedSparseEmptyEntry = 0x00000000ffffffffULL; +constexpr int kGroupedSparseItemsPerThread = + kGroupedSparseHashCapacity / kGroupedSparsePlannerThreads; +using GroupedSparseSort = + cub::BlockRadixSort; +// Hash entries plus logical owners exactly fit Volta's 96 KiB opt-in limit. +// After loading both into registers, reuse this storage for sorting and scans. +constexpr size_t kGroupedSparsePlannerSharedMemory = + kGroupedSparseHashCapacity * + (sizeof(unsigned long long) + sizeof(uint32_t)); +static_assert(sizeof(GroupedSparseSort::TempStorage) <= + kGroupedSparsePlannerSharedMemory); __device__ __forceinline__ void grouped_sparse_hash_insert( - unsigned long long* __restrict__ hash_table, const int physical_microblock, - const uint32_t token_mask) { + unsigned long long* __restrict__ hash_table, + uint32_t* __restrict__ logical_owners, const int physical_microblock, + const uint32_t token_mask, const int query, const int logical_token) { if (physical_microblock < 0 || token_mask == 0) { return; } @@ -3675,13 +3689,17 @@ __device__ __forceinline__ void grouped_sparse_hash_insert( for (int probe = 0; probe < kGroupedSparseHashCapacity; ++probe) { const unsigned long long old = atomicCAS(hash_table + slot, kGroupedSparseEmptyEntry, desired); - if (old == kGroupedSparseEmptyEntry) { - return; - } - if (static_cast(old) == - static_cast(physical_microblock)) { + if (old == kGroupedSparseEmptyEntry || + static_cast(old) == + static_cast(physical_microblock)) { atomicOr(hash_table + slot, static_cast(token_mask) << 32); + // Nonnegative int32 tokens use at most 29 bits after division by four. + // The first contributing query owns shared pages, independent of request + // slot IDs, physical allocation, insertion order, and hash collisions. + const uint32_t owner = (static_cast(query) << 29) | + (static_cast(logical_token) >> 2); + atomicMin(logical_owners + slot, owner); return; } slot = (slot + 1) & (kGroupedSparseHashCapacity - 1); @@ -3741,15 +3759,13 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla const int physical_page_stride, const int num_cache_blocks) { const int group_idx = blockIdx.x; const int tid = threadIdx.x; - __shared__ int category_counts[8]; - __shared__ int category_offsets[8]; - __shared__ int category_cursors[8]; - __shared__ int - warp_category_prefix[(kGroupedSparsePlannerThreads / kWarpSize) * 8]; extern __shared__ unsigned long long hash_table[]; + auto* logical_owners = + reinterpret_cast(hash_table + kGroupedSparseHashCapacity); for (int slot = tid; slot < kGroupedSparseHashCapacity; slot += kGroupedSparsePlannerThreads) { hash_table[slot] = kGroupedSparseEmptyEntry; + logical_owners[slot] = UINT_MAX; } __syncthreads(); @@ -3793,8 +3809,9 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla first_token, request_idx, request_block_table, request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); - grouped_sparse_hash_insert(hash_table, physical_microblock, - 0xFu << (query * 4)); + grouped_sparse_hash_insert(hash_table, logical_owners, + physical_microblock, 0xFu << (query * 4), + query, first_token); } else { #pragma unroll for (int token_offset = 0; token_offset < 4; ++token_offset) { @@ -3804,8 +3821,9 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla token, request_idx, request_block_table, request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); - grouped_sparse_hash_insert(hash_table, physical_microblock, - 1u << (query * 4 + (token & 3))); + grouped_sparse_hash_insert( + hash_table, logical_owners, physical_microblock, + 1u << (query * 4 + (token & 3)), query, token); } } } @@ -3842,12 +3860,47 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); const uint32_t tail_mask = ((1u << tail_count) - 1) << (query * 4); - grouped_sparse_hash_insert(hash_table, physical_microblock, tail_mask); + grouped_sparse_hash_insert(hash_table, logical_owners, + physical_microblock, tail_mask, query, + selected_tail_token); } } } __syncthreads(); + // Keep the existing physical-page union and masks, but never let physical + // hash slots determine the attention reduction order. Category is primary + // to preserve active-tile packing; the logical owner orders each category. + unsigned long long entries[kGroupedSparseItemsPerThread]; + unsigned long long sort_keys[kGroupedSparseItemsPerThread]; +#pragma unroll + for (int item = 0; item < kGroupedSparseItemsPerThread; ++item) { + const int slot = tid * kGroupedSparseItemsPerThread + item; + const unsigned long long entry = hash_table[slot]; + entries[item] = entry; + const int category = + grouped_sparse_active_m_tiles(static_cast(entry >> 32)); + sort_keys[item] = static_cast(entry) == 0xffffffffu + ? ULLONG_MAX + : (static_cast(category) << 32) | + logical_owners[slot]; + } + __syncthreads(); + auto& sort_storage = + *reinterpret_cast(hash_table); + // Three category bits, 32 owner bits, and one bit separating empty slots. + GroupedSparseSort(sort_storage).Sort(sort_keys, entries, 0, 36); + __syncthreads(); +#pragma unroll + for (int item = 0; item < kGroupedSparseItemsPerThread; ++item) { + hash_table[tid * kGroupedSparseItemsPerThread + item] = entries[item]; + } + auto* category_counts = reinterpret_cast(logical_owners); + int* category_offsets = category_counts + 8; + int* category_cursors = category_offsets + 8; + int* warp_category_prefix = category_cursors + 8; + __syncthreads(); + if (tid < 8) { category_counts[tid] = 0; category_offsets[tid] = 0; @@ -4006,17 +4059,16 @@ at::Tensor flash_attention_grouped_sparse_page4_plan( TORCH_CHECK(properties->major == 7 && properties->minor == 0, "grouped sparse page4 planner supports SM70 only"); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - constexpr size_t kPlannerSharedMemory = - kGroupedSparseHashCapacity * sizeof(unsigned long long); - const cudaError_t smem_status = cudaFuncSetAttribute( - grouped_sparse_page4_plan_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, kPlannerSharedMemory); + const cudaError_t smem_status = + cudaFuncSetAttribute(grouped_sparse_page4_plan_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + kGroupedSparsePlannerSharedMemory); TORCH_CHECK(smem_status == cudaSuccess, "Failed to set grouped sparse page4 planner shared memory: ", cudaGetErrorString(smem_status)); - grouped_sparse_page4_plan_kernel<<(num_groups), - kGroupedSparsePlannerThreads, - kPlannerSharedMemory, stream>>>( + grouped_sparse_page4_plan_kernel<<< + static_cast(num_groups), kGroupedSparsePlannerThreads, + kGroupedSparsePlannerSharedMemory, stream>>>( logical_indices.data_ptr(), block_table.data_ptr(), token_to_req.data_ptr(), query_positions.data_ptr(), sequence_lengths.data_ptr(), output_blocks.data_ptr(), diff --git a/tests/kernels/test_sm70_qsa_page4_plan.py b/tests/kernels/test_sm70_qsa_page4_plan.py new file mode 100644 index 0000000000..08367faa6f --- /dev/null +++ b/tests/kernels/test_sm70_qsa_page4_plan.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Allocation-invariance contract for the SM70 grouped QSA planner. + +Keep logical selections, query grouping, and physical-page aliasing fixed. +Relocating KV pages must preserve masks, category padding, and the logical +reduction order. This is not a cross-batch-shape invariance contract. +""" + +import pytest +import torch + +WIDTH = 2051 +OUTPUT_WIDTH = 4160 + + +@pytest.fixture +def extension(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("requires a V100 / SM70 GPU") + return pytest.importorskip("flash_attn_v100_cuda") + + +def make_case(kind="mixed", page_size=16, permute_selection=False): + generator = torch.Generator().manual_seed(387466) + if kind == "wide": + requests = torch.arange(8, dtype=torch.int32) + lengths = torch.full((8,), 8191, dtype=torch.int32) + positions = lengths.to(torch.int64) - 1 + else: + requests = torch.tensor([0, 0, 1, 1, 2, 3, 3, 3], dtype=torch.int32) + lengths = torch.tensor([61, 57, 59, 63], dtype=torch.int32) + positions = torch.tensor([16, 38, 19, 40, 36, 25, 32, 62]) + pages = (int(lengths.max()) + page_size - 1) // page_size + table = torch.arange(len(lengths) * pages, dtype=torch.int32).view(-1, pages) + if kind == "shared": + table[:, 0] = table[0, 0] + indices = torch.full((8, WIDTH), -1, dtype=torch.int32) + for row, request in enumerate(requests.tolist()): + visible = min(int(positions[row]) + 1, int(lengths[request])) + blocks = torch.randperm(visible // 4, generator=generator)[:512].sort().values + if permute_selection: + blocks = blocks.flip(0) + full = (blocks[:, None] * 4 + torch.arange(4)).flatten() + indices[row, : len(full)] = full.to(torch.int32) + tail = torch.arange(visible // 4 * 4, visible, dtype=torch.int32) + indices[row, len(full) : len(full) + len(tail)] = tail + if kind == "invalid": + requests[1], requests[4] = -1, len(lengths) + indices[6] = -1 + if kind == "empty": + indices.fill_(-1) + return indices, table, requests, positions, lengths + + +def relocate(case, layout): + indices, table, requests, positions, lengths = case + count = int(table.max()) + 1 + if layout == "collision": + # PAGE16, interleaved K/V: stride=8, so 1024 cache pages collide + # in the 8192-entry hash. No K/V allocation is needed for plan tests. + mapping = torch.arange(count, dtype=torch.int32) * 2048 + 7 + elif layout == "shuffled": + mapping = torch.randperm(count, generator=torch.Generator().manual_seed(466)) + mapping = mapping.to(torch.int32) * 3 + 11 + else: + mapping = torch.arange(count, dtype=torch.int32) + return (indices, mapping[table.long()], requests, positions, lengths), mapping + + +def category(mask): + result = 0 + for query in range(8): + if mask & (15 << (query * 4)): + result |= 1 << (query * 6 // 16) + result |= 1 << ((query * 6 + 5) // 16) + return result + + +def reference(case, page_size, physical_stride): + indices, table, requests, positions, lengths = case + entries: dict[int, tuple[int, int]] = {} + for query, request in enumerate(requests.tolist()): + if not 0 <= request < len(lengths): + continue + visible = min(max(int(positions[query]) + 1, 0), int(lengths[request])) + count = min(visible // 4, 512) * 4 + visible % 4 + for token in indices[query, :count].tolist(): + if not 0 <= token < visible: + continue + physical = int(table[request, token // page_size]) + physical = physical * physical_stride + token % page_size // 4 + owner = (query << 29) | (token // 4) + old_mask, old_owner = entries.get(physical, (0, (1 << 32) - 1)) + entries[physical] = ( + old_mask | (1 << (query * 4 + token % 4)), + min(old_owner, owner), + ) + pages: list[int] = [] + masks: list[int] = [] + for group_category in range(1, 8): + bucket = sorted( + (owner, physical, mask) + for physical, (mask, owner) in entries.items() + if category(mask) == group_category + ) + pages.extend(physical for _, physical, _ in bucket) + masks.extend(mask for _, _, mask in bucket) + padding = -len(bucket) % 8 + pages.extend([0] * padding) + masks.extend([0] * padding) + return pages, masks + + +def run_plan(extension, case, page_size, physical_stride, num_cache_blocks=None): + device_case = tuple(t.cuda() for t in case) + if num_cache_blocks is None: + num_cache_blocks = int(case[1].max()) + 1 + pages = torch.full((1, OUTPUT_WIDTH), -17, dtype=torch.int32, device="cuda") + masks = torch.zeros_like(pages, dtype=torch.uint32) + lengths = torch.empty(1, dtype=torch.int32, device="cuda") + + def launch(): + extension.grouped_sparse_page4_plan_fwd( + *device_case, + pages, + masks, + lengths, + page_size, + physical_stride, + num_cache_blocks, + ) + + launch() + return pages, masks, lengths, launch, device_case + + +@pytest.mark.parametrize("kind", ["mixed", "shared", "invalid", "empty", "wide"]) +@pytest.mark.parametrize("layout", ["compact", "shuffled", "collision"]) +def test_plan_matches_logical_reference(extension, kind, layout): + case, _ = relocate(make_case(kind), layout) + expected_pages, expected_masks = reference(case, 16, 8) + pages, masks, lengths, launch, _ = run_plan(extension, case, 16, 8) + assert int(lengths[0]) == len(expected_pages) * 4 + count = len(expected_pages) + assert pages[0, :count].cpu().tolist() == expected_pages + assert masks[0, :count].cpu().tolist() == expected_masks + first = (pages.clone(), masks.clone(), lengths.clone()) + for _ in range(3): + launch() + for before, after in zip(first, (pages, masks, lengths)): + assert torch.equal(before, after) + + +@pytest.mark.parametrize("page_size", [4, 16, 32]) +def test_plan_selection_order_and_graph_relocation(extension, page_size): + case = make_case("shared", page_size) + changed_case = make_case("shared", page_size, permute_selection=True) + changed_case, _ = relocate(changed_case, "shuffled") + stride = page_size // 4 + pages, masks, lengths, launch, device_case = run_plan( + extension, case, page_size, stride, int(changed_case[1].max()) + 1 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + for source, target in zip(changed_case, device_case): + target.copy_(source) + expected_pages, expected_masks = reference(changed_case, page_size, stride) + for _ in range(3): + graph.replay() + count = len(expected_pages) + assert int(lengths[0]) == count * 4 + assert pages[0, :count].cpu().tolist() == expected_pages + assert masks[0, :count].cpu().tolist() == expected_masks + + +@pytest.mark.parametrize("interleaved", [False, True]) +@pytest.mark.parametrize("kv_dtype", ["auto", "fp8_e4m3"]) +def test_attention_is_bitwise_invariant_to_physical_relocation( + extension, interleaved, kv_dtype +): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + case = make_case("shared") + generator = torch.Generator(device="cuda").manual_seed(387) + count = int(case[1].max()) + 1 + kv = torch.randn( + count, 2, 16, 1, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + if kv_dtype == "fp8_e4m3": + kv = kv.to(torch.float8_e4m3fn).view(torch.uint8) + query = torch.randn( + 8, 6, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + outputs = [] + for layout in ("compact", "shuffled"): + remapped, mapping = relocate(case, layout) + cache = torch.zeros( + int(mapping.max()) + 1, 2, 16, 1, 256, dtype=kv.dtype, device="cuda" + ) + cache[mapping.long().cuda()] = kv + key, value = cache[:, 0], cache[:, 1] + if not interleaved: + key, value = key.contiguous(), value.contiguous() + stride = key.stride(0) // (4 * 256) + pages, masks, lengths, _, _ = run_plan(extension, remapped, 16, stride) + physical_k, physical_v = qsa._qsa_xqa_page4_physical_kv(query, key, value) + out = torch.empty_like(query) + lse = torch.empty((8, 6), dtype=torch.float32, device="cuda") + extension.grouped_sparse_page4_fwd( + query, + physical_k, + physical_v, + out, + pages, + masks, + lengths, + lse, + 256**-0.5, + kv_dtype, + 0.125, + 0.25, + ) + outputs.append(out.clone()) + assert torch.equal(*outputs) + + +@pytest.mark.parametrize("permute_selection", [False, True]) +def test_single_row_page4_table_uses_logical_order(extension, permute_selection): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + case, _ = relocate(make_case(permute_selection=permute_selection), "shuffled") + indices, table, requests, positions, lengths = case + pages, seq_lens = qsa._qsa_xqa_page4_block_table( + *(tensor.cuda() for tensor in case), int(table.max()) + 1, 16, 8 + ) + for row, request in enumerate(requests.tolist()): + visible = min(int(positions[row]) + 1, int(lengths[request])) + tokens = sorted(indices[row, :visible:4].tolist()) + expected = [int(table[request, t // 16]) * 8 + t % 16 // 4 for t in tokens] + assert pages[row, : len(expected)].cpu().tolist() == expected + assert int(seq_lens[row]) == visible + + +@pytest.mark.parametrize("kv_dtype", ["auto", "fp8_e4m3"]) +def test_mixed_grouped_and_xqa_tail_is_allocation_invariant( + extension, monkeypatch, kv_dtype +): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + monkeypatch.setattr(qsa, "_SM70_QSA_XQA_PAGE4", True) + monkeypatch.setattr(qsa, "_SM70_QSA_XQA_PAGE4_MIN_ROWS", 64) + monkeypatch.setattr(qsa, "_SM70_QSA_GROUPED_PAGE4", True) + calls = [] + grouped = qsa._qsa_sparse_paged_attention_sm70_grouped_page4 + tail = qsa._qsa_sparse_paged_attention_sm70_xqa_page4_batch + + def record_grouped(query, *args): + calls.append(("grouped", query.shape[0])) + return grouped(query, *args) + + def record_tail(query, *args): + calls.append(("tail", query.shape[0])) + return tail(query, *args) + + monkeypatch.setattr( + qsa, "_qsa_sparse_paged_attention_sm70_grouped_page4", record_grouped + ) + monkeypatch.setattr( + qsa, "_qsa_sparse_paged_attention_sm70_xqa_page4_batch", record_tail + ) + case = make_case("shared") + generator = torch.Generator(device="cuda").manual_seed(185) + count = int(case[1].max()) + 1 + kv = torch.randn( + count, 2, 16, 1, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + if kv_dtype == "fp8_e4m3": + kv = kv.to(torch.float8_e4m3fn).view(torch.uint8) + query = torch.randn( + 185, 6, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + outputs = [] + for layout in ("compact", "shuffled"): + remapped, mapping = relocate(case, layout) + indices, table, requests, positions, lengths = [t.cuda() for t in remapped] + indices = indices.repeat(24, 1)[:185].contiguous() + requests = requests.repeat(24)[:185].contiguous() + positions = positions.repeat(24)[:185].contiguous() + cache = torch.zeros( + int(mapping.max()) + 1, 2, 16, 1, 256, dtype=kv.dtype, device="cuda" + ) + cache[mapping.long().cuda()] = kv + outputs.append( + qsa.qsa_sparse_paged_attention( + query, + cache[:, 0], + cache[:, 1], + indices, + table, + requests, + query_positions=positions, + sequence_lengths=lengths, + kv_cache_dtype=kv_dtype, + k_scale=0.125, + v_scale=0.25, + ).clone() + ) + assert torch.equal(*outputs) + assert calls == [("grouped", 184), ("tail", 1)] * 2 diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b8b370ba18..b5fc967931 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -426,7 +426,6 @@ def _qsa_xqa_page4_table_kernel( 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) @@ -491,13 +490,16 @@ def _qsa_xqa_page4_table_kernel( physical_microblock = ( tl.maximum(physical_page, 0) * PHYSICAL_PAGE_STRIDE + page_offset // 4 ) + # Sort by logical token, not allocator-dependent physical page ID. Keep + # the partial causal page after all complete pages and invalid slots last. + logical_key = safe_token.to(tl.int64) << 31 encoded = tl.where( valid & is_complete, - physical_microblock, + logical_key | physical_microblock.to(tl.int64), tl.where( valid & is_tail, - physical_microblock + TAIL_MARKER, - 2147483647, + (1 << 62) | logical_key | physical_microblock.to(tl.int64), + 9223372036854775807, ), ) tl.store( @@ -1464,7 +1466,7 @@ def _qsa_xqa_page4_block_table( rows = logical_indices.shape[0] encoded_pages = torch.empty( (rows, _SM70_QSA_XQA_PAGE4_PAGES), - dtype=torch.int32, + dtype=torch.int64, device=logical_indices.device, ) xqa_sequence_lengths = torch.empty( @@ -1490,14 +1492,13 @@ def _qsa_xqa_page4_block_table( 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, - ) + ).to(torch.int32) return physical_pages, xqa_sequence_lengths