[Peuqui]

Purpose

The QSA prefill attention of Qwen4Exp takes the native Flash-V100 grouped and XQA page4 routes only on sm_70 devices and only on ranks with six query heads on one KV head, the Flash-Next TP4 rank. Every other rank falls back to the Triton split-k kernel. Triton 3.6 has no tensor-core MMA before SM80, so on SM70 and SM75 its tl.dot runs as FP32 FMA. Flash-Next has 24 query heads on 2 KV heads, so on our TP2xPP2 entry (12/1 per rank) and our PP4 entry (24/2 per stage) no rank reached the native routes, neither the Tesla V100 stages nor the Quadro RTX 8000 stages.

Four commits:

1. Test repair. The page4 planner test imported the top-level flash_attn_v100_cuda name, which can resolve to another tree's build; after any test that loaded the extension through flash_attn_v100 (as the QSA routes do) every planner case failed with "type PreparedPolicy is already registered". It also patched _SM70_QSA_XQA_PAGE4, _SM70_QSA_XQA_PAGE4_MIN_ROWS and _SM70_QSA_GROUPED_PAGE4, which no longer exist since the page4 switches moved into the sm70 sparse policy. It now imports through flash_attn_v100.flash_attn_interface and asserts the policy defaults it relies on.
2. 12/24-head ranks. All query heads of a QSA row share one latent KV head and one selection, so a rank with 6*R heads runs as R consecutive six-head pseudo-rows that repeat the row's selection, position and request. The kernels are unchanged; the shape check accepts multiples of six heads and the minimum-row threshold counts pseudo-rows.
3. Turing. fused_mma.h names the two rows a lane holds in the 16x16x16 FP32 accumulator for both architectures (accumulator_16x16_rows(): Volta HMMA.884 quad pairs r and r + 2, Turing m16n8k8 lane / 4 and lane / 4 + 8); the two row-scaling helpers of the grouped kernel used the Volta formula inline. The planner hash table follows the device's opt-in shared memory: 8192 entries (96 KiB, unchanged on Volta) or 5120 entries (60 KiB) on 64 KiB devices; a group of eight rows holds at most 8 x 513 microblocks, so both tables are complete, and the plan order depends only on category and logical owner. Host checks test the shared memory the kernels need instead of compute capability 7.0, and the new grouped_sparse_page4_device_supported() reports whether the build has a kernel image and enough shared memory for a device. The QSA gate asks it per device (cached); a build without sm_75, or an extension without this query, keeps the device on Triton and logs why. setup.py of Flash-V100 and _volta_cuda_arch_list() in the root setup.py accept 7.5, so FLASH_ATTN_V100_CUDA_ARCH_LIST="7.0;7.5" builds both. Route logs use per-process scope, because in a mixed pipeline the affected stage is often not the first local rank.
4. Several KV heads per rank. The route runs each KV head's block of query heads against that head's slice of the K/V-interleaved cache; the kernels already take block, token and head strides separately, and the page4 view and the microblock ids now use the cache's token stride (num_kv_heads x 256) instead of assuming 256 (_qsa_xqa_page4_page_stride). One KV head per rank is unchanged.

Decode is not affected: MTP decode has 5 rows per request, below the route's minimum of 64 rows, and stays on Triton.

Not a duplicate: gh pr list --repo 1CatAI/1Cat-vLLM --state all --search with "page4 Turing", "grouped page4", "Flash-V100 sm_75", "Flash-V100 Turing", "QSA Turing", "QSA prefill attention", "page4 KV heads", "accumulator layout", "sm75 flash-attention-v100" and "QSA 24 heads", and the same searches over issues, find no PR or issue that admits these rank shapes or Turing to the page4 routes. Related: #828 (Leonccaa) uses the grouped page4 route for QSA DCP prefill attention, a different path; both PRs change the cache page-stride condition of the same gate, which gives one conflict hunk in vllm/models/qwen4_exp/nvidia/ops/qsa.py. #1229 (ours) admits the cuBLAS QSA indexer and exact top-k on Turing; it changes other lines of qsa.py, this PR applies with and without it, and both arms of the serving A/B below had it. #1205 changes the QSA selector and staging in qsa.py without a conflict. git merge-tree against all 122 open PR heads adds conflicts with #828 (qsa.py, see above), #903 (fused_mha.h, flash_decode_paged.cu, fused_mha_api.cpp: new declarations and bindings next to its new split route) and the draft series #1179-#1198 (flash_decode_paged.cu, where the new device query sits next to functions the series renames; #1194-#1198 also qsa.py and test_qsa_ops.py).

AI assistance was used for this change. I reviewed every line and ran the tests and measurements below.

Test Plan

pytest tests/kernels/attention/test_sm70_qsa_page4_head_groups.py tests/kernels/attention/test_sm70_qsa_grouped_page4.py tests/kernels/test_sm70_qsa_page4_plan.py tests/models/qwen4_exp/test_qsa_ops.py tests/models/qwen4_exp/test_qsa_xqa_page4_workspace.py tests/models/qwen4_exp/test_qsa_xqa_page4_workspace_cpu.py tests/models/qwen4_exp/test_qsa_dcp_attention.py tests/models/qwen4_exp/test_qsa_dcp_local_width.py tests/models/qwen4_exp/test_host_kv.py tests/kernels/test_qsa_dense_decode.py tests/models/qwen4_exp/test_e4m3_mtp_gpu.py tests/models/qwen4_exp/test_e4m3_mtp_capture_gpu.py tests/models/qwen4_exp/test_qsa_e4m3.py
pre-commit run --from-ref origin/main --to-ref HEAD
pre-commit run mypy-3.10 --hook-stage manual --from-ref origin/main --to-ref HEAD

The tests ran on main 6f3557c plus this PR (Flash-V100 built from this tree with TORCH_CUDA_ARCH_LIST="7.0;7.5") and on main with its own build (Flash-V100 sm_70), once on a Tesla V100 and once on a Quadro RTX 8000.

Kernel benchmark: one 1,616-row prefill chunk with the real top-k rows, positions, requests and block table dumped from a Flash-Next prefill at 24,240 and 150,288 tokens of context, 13 QSA layers each; random FP16 K/V in the serving layout [blocks, 2, 16, kv_heads, 256]; both routes through qsa_sparse_paged_attention; median of 20 calls.

Serving A/B on our fork (main 6f3557c plus our commits): fork commit eb203d0 (this change) against its parent ee261cc, Flash-Next PP4 (stages on two Quadro RTX 8000 and two Tesla V100) and TP2xPP2 (one TP2 pair per GPU type), torch.compile with CUDA graphs, MTP with 4 speculative tokens. Boots fix cold, base, fix, base, fix; the cold boot is not counted. The base arm runs the production tree, the fix arm the same entries with PYTHONPATH on the fork tree, its own compile cache and its Flash-V100 build for sm_70 and sm_75. Per boot: four ~29k-token prompts (the first long request after a boot is not counted), a decode benchmark (prose, code, edit), quality prompts, and prompts of 184,134 and 249,950 tokens.

Path coverage

Path / provider	Format	Key conditions (shape, TP/EP, graph, capability)	Expected selection / fallback	Evidence (source / static / executed / measured)	Unverified scope
Grouped/XQA page4 prefill, 6 query heads on 1 KV head (TP4 rank)	FP16 and FP8 E4M3 KV	sm_70; now also sm_75 with a 7.5 build	native page4 route	executed (tests on V100 and RTX 8000)	TP4 serving not run
Same, 12 on 1 (TP2 rank)	FP16 KV	sm_70 and sm_75, at least 64 rows	native page4 route, previously Triton	executed + measured (kernel benchmark, TP2xPP2 serving on all four ranks)
Same, 24 on 2 (TP1 pipeline stage)	FP16 KV	sm_70 and sm_75, at least 64 rows	native page4 route per KV head, previously Triton	executed + measured (kernel benchmark, PP4 serving on all four stages)
Same, 24 on 1	FP16 KV	test shape only	native page4 route	executed (tests)	no model with this rank shape run
Device without a kernel image or enough shared memory, or extension without the device query	any	grouped_sparse_page4_device_supported() false or missing	Triton split-k, warning logged once per process	source	not run
QSA DCP prefill, QSA decode	any	DCP path; fewer than 64 rows	unchanged	static

Acceleration and benchmark contract (required for performance changes)

• Default enabled or opt-in: enabled where admitted, through the existing sm70 sparse policy switches qsa_xqa_page4 and qsa_grouped_page4 (unchanged defaults); Turing needs Flash-V100 built with 7.5.
• Required CLI options and environment switches: none at runtime; build with FLASH_ATTN_V100_CUDA_ARCH_LIST="7.0;7.5" (or a TORCH_CUDA_ARCH_LIST that contains 7.5) for Turing.
• KV cache dtype used for the benchmark: auto (FP16, the model dtype); the tests also cover FP8 E4M3.
• Wheel SHA or source commit: our fork, fix arm eb203d0, base arm ee261cc; this PR is the same change on main (identical diff apart from [Perf][Qwen4Exp] Admit the cuBLAS QSA indexer and exact top-k on Turing #1229's lines in qsa.py and its tests).
• PYTHONPATH, source overlays, or external native libraries used: fix arm with PYTHONPATH on the fork tree and its own Flash-V100 build; nothing else.
• User-entry route-hit, speed, and output-quality evidence: see Test Result.
• Existing kernel/backend registry and capability-rejection reasons: the page4 gate in ops/qsa.py; new per-device capability query grouped_sparse_page4_device_supported(), rejection logged with the device index.
• Per-engine configuration and deprecated-variable compatibility: no new fields or variables.
• New/changed environment metadata and generated reference check: none.
• Route snapshot command and intentional changed rows: none; tools/sm70_route_snapshot.py has no QSA category.
• SM70 registered variables / unregistered reads / model-parameter restrictions, before -> after: unchanged.

Test Result

• Tests on a Tesla V100: this PR 156 passed, 72 failed; main 109 passed, 72 failed, 38 errors (the planner test, fixed by the first commit). On a Quadro RTX 8000: this PR 128 passed, 72 failed, 28 skipped; main 79 passed, 72 failed, 68 skipped (the page4 tests skip on Turing there). No test fails with this PR that passes on main. The 72 failures are tests/models/qwen4_exp/test_host_kv.py and identical on both trees.
• pre-commit and mypy-3.10 pass.
• Kernel benchmark, ms per call (less is better), ranges over the 26 dumps:

GPU	Rank shape (query / KV heads)	Triton split-k	Native page4	Speedup
Quadro RTX 8000	12 / 1	38.89-39.00	6.20-8.25	4.72-6.28x
Quadro RTX 8000	24 / 2	75.96-76.20	12.26-16.37	4.64-6.20x
Tesla V100	12 / 1	22.62-23.43	7.10-9.38	2.41-3.19x
Tesla V100	24 / 2	44.52-44.63	14.06-18.66	2.39-3.17x

Max |native - Triton| 1.22e-4 to 2.44e-4 on every dump (FP16 output).

• Route hit: in all three fix boots of each topology every worker logged "Using SM70 grouped QSA Flash-V100 page4 prefill route" (PP4: Worker_PP0 and PP1 on the Quadro RTX 8000, PP2 and PP3 on the Tesla V100; TP2xPP2: all four ranks); no base boot did.
• Serving, time to first token for the ~29k-token prompts (29,204 tokens on average; mean of three prompts per boot; less is better): TP2xPP2 base 14.83 / 14.80 s (1,969 / 1,973 tok/s) against 11.80 s in both fix boots (2,475 tok/s), -20 %; PP4 base 10.33 / 10.20 s (2,826 / 2,863 tok/s) against 7.67 / 7.70 s (3,809 / 3,793 tok/s), -25 %. On TP2xPP2 the gain comes from the Quadro RTX 8000 ranks: in an earlier A/B of the second commit alone, where only the V100 ranks of that entry reached the native route, the same prompts took 14.7 s against 14.8 s.
• Decode is unchanged: edit task PP4 54 ms per step = 92.3-92.6 tok/s base against 92.2-92.8 tok/s, TP2xPP2 45 ms = 109.8-110.2 tok/s in both arms. The 184k and 250k prefills of this A/B were dominated by stalls of the PLE disk tier in both arms (addressed separately), so we make no claim for them.
• Outputs: the head-group test compares the serving layout for 6/1, 12/1, 24/2 and 24/1 against Triton and an FP32 reference (atol 2e-3) on both GPUs. On PP4 the greedy texts are word-identical in all five boots of both arms. The long-context answers are the same in every boot of an arm but differ between the arms, as expected when the attention kernel changes; on TP2xPP2 the greedy texts of the fix arm also differ between its boots. This A/B still had the stale resident rows fixed in [Bugfix][Qwen4Exp] Decide the PLE cascade merge when the layer is built #1234; with [Bugfix][Qwen4Exp] Decide the PLE cascade merge when the layer is built #1234 on top (our fork commit 6bbcd40, which contains this change), three boots per topology give byte-identical greedy texts and long-context answers. Greedy texts, long-context answers and quality answers were read by hand; both arms pass the same quality prompts.

🤖 Generated with Claude Code

https://claude.ai/code/session_01YaG8VjXn5eN3Aq7eZrA3dp