[Bugfix][SM70] Stabilize QSA page4 order under KV relocation

[Leonccaa]

Summary

Fix SM70 QSA page4 attention changing its reduction order when the KV allocator
relocates otherwise identical logical pages. This is a correctness repair,
not an AWQ grouped-decode optimization or a scheduler-policy change.

Upstream target: Leonccaa:fix/sm70-qsa-page4-logical-order to 1CatAI:main.
The reviewed commit is 5fa8a605dab12cc9ee15459d9ac6b88d95c7be3a, based on
fbcef6e2f959e95bbe4ca807931abfa2393546e7; it contains exactly four changed
files. The fork-local review is Leonccaa #7.
The code is unchanged after that review. The numerical follow-up below is
documented for continuity, not included as additional code or runtime switches.

Only two production files change:

• The grouped planner keeps its physical-page union and token masks, but sorts
  by (active-row category, first contributing query, logical page) using
  NVIDIA CUB. Physical addresses become payload rather than ordering keys.
• The single-row XQA planner, including grouped-batch tails, keeps its existing
  GPU sort but changes the key from physical address to logical token position.

No attention arithmetic kernel, route threshold, quantization, weight layout,
KV contents, or public extension signature changes.

Investigation and introduction history

With frozen AWQ / TP4 / MTP0 / FP16 KV / MRv2 / FULL_AND_PIECEWISE, fixed prompt
IDs and batch order, and experimental AWQ grouped decode disabled, repeated
C4 first-token runs first diverged at Layer3 QSA. Q/K/V, selected token IDs and
all effective logical K/V were bitwise identical; physical page tables differed.
The actual step had 185 query rows (184 grouped plus one XQA tail).

All 23 grouped plans retained the same logical page/mask sets, but 17 changed
order. Fixed physical inputs replayed identically. A diagnostic logical-sort
intervention restored Layer3 equality; applying it to all 12 QSA layers restored
all four ranks' 175 observed boundaries and complete logits. Removing sorting
while keeping observation/synchronization brought the difference back.

• 1Cat #378 /
  #382 introduced the SM70 virtual-page4 fork, including physical sorting
  for the single-row path.
• 1Cat #387
  (94ce990ce85abeb12e3948ee1c4f518594bacf25) introduced the grouped physical
  hash union and hash-slot-order emission.
• 1Cat #466
  (186c9e3585b109b88c070603a181dd6826400153) lowered admission from 4096 to
  64 query rows and enabled mixed grouped/XQA tails. It widened exposure; it did
  not introduce the grouped hash planner.
• This is not introduced by 1Cat #477
  or the recent AWQ grouped-decode work.

These are source-history findings, not a GPU bisection of each old release.
The design note includes the full intervention/reversal evidence and limits.

Upstream and duplicate-work check

Original vLLM QSA traverses logical selections before physical addressing and
does not contain this SM70 hash planner. The existing Triton fallback was also
bitwise allocation-invariant in all four captured-rank replay pairs.

Related work does not duplicate this repair:

• 1Cat #394 already stabilizes selection on this tested route; those selected
  indices were equal before this fix.
• vLLM #55122 addresses top-k
  membership/tie/order, not ordering introduced later by a physical planner.
• vLLM #54873 skips unused
  selection entries and tunes sparse-attention launches.
• RFC #55394 uses logical
  ordering in a related GB10 query-union prototype. Its single-request prototype
  is not a ready-made SM70 concurrent replacement. We reuse the principle, not
  the platform-specific prototype.
• Open 1Cat #481 includes separate QSA selection/decode work, not this grouped
  planner repair: its QSA edits concern the selector sidecar and output-gate
  fusion, not logical ownership or page4 ordering.
• Open 1Cat #476 extends batched
  DFlash2 grouped verification and retains the sparse-page4 planner; it does
  not implement this repair. QSA/page4/logical-order searches and the overlapping
  PR diffs were rechecked on 2026-09-04; no direct duplicate was found.

Validation

CUDA 12.8, V100 / SM70:

# In flash-attention-v100, using the project's Python 3.12 environment:
CUDA_HOME=/usr/local/cuda MAX_JOBS=2 TORCH_CUDA_ARCH_LIST=7.0 \
  ../.venv/bin/python setup.py build_ext --inplace -j1

# From the repository root, with the rebuilt extension importable:
.venv/bin/python -m pytest -q tests/kernels/test_sm70_qsa_page4_plan.py \
  tests/models/qwen4_exp/test_qsa_ops.py
pre-commit run --files flash-attention-v100/kernel/flash_decode_paged.cu \
  vllm/models/qwen4_exp/nvidia/ops/qsa.py \
  tests/kernels/test_sm70_qsa_page4_plan.py \
  docs/design/sm70_qsa_page4_allocation_invariance.md

The actual isolated build used
/evidence/test-venv/bin/python setup.py build_ext --build-lib /build/runtime --build-temp /build/temp -j1.
The actual test command was
/evidence/test-venv/bin/python -m pytest /tests/test_sm70_qsa_page4_plan.py /tests/test_qsa_ops.py --confcutdir=/tests -q -p no:cacheprovider --junitxml=/output/pytest.xml,
using a frozen compatible runtime; the commands above are source-tree equivalents.
Result: 26 new regressions + 14 existing QSA tests = 40 passed. All applicable
pre-commit hooks passed. A fresh pre-submission single-V100 rerun of the exact
reviewed sources and binary also passed all 40 tests (10.25 seconds, no skips);
changed-file pre-commit checks passed again. This rerun did not reload the full
model or repeat the performance matrix. Tests cover exact CPU-reference plans, aliasing,
collisions, invalid/empty groups, wide unions, selection permutations, page
sizes 4/16/32, CUDA Graph relocation, FP16/E4M3 contiguous/interleaved KV, and
the 185-row grouped-plus-XQA route.

Negative control: before the fix, 17 of the initial 20 cases failed; only the
three empty-group controls passed. Later added tests are not counted in that
initial run. Captured replay changed from 1625/1852/1515/1866 different output
elements on ranks 0/1/2/3 (max absolute difference 0.00048828125) to zero on
every rank.

Single-V100 operator input	Old page4	Fixed page4	Existing Triton
Captured 185-row step, allocation A	0.3518 ms	0.3615 ms (+2.8%)	2.6552 ms
Captured 185-row step, allocation B	0.3600 ms	0.3635 ms (+1.0%)	2.4049 ms
Synthetic 1024 rows	1.513 ms	1.581 ms (+4.5%)	14.371 ms
Synthetic 8192 rows	9.114 ms	9.782 ms (+7.3%)	99.359 ms

These are CUDA-event medians (3 warmups, 20 samples), FP16 KV, Hq/Hkv/D=6/1/256.
The synthetic rows use four 64K requests with highly shared selections; the
185-row inputs are captured from the short diagnostic batch. These are not E2E tok/s or a
generalized speedup claim. Planner-only 8192-row latency increases from 0.661 to
1.325 ms. The measured cost favors this narrow repair over full Triton fallback.

The planner now uses 96 KiB dynamic shared memory and 128 registers/thread,
with zero reported stack/local bytes. No grouped global workspace is added.
Single-row sorting uses int64 rather than int32 temporary keys, so total
temporary metadata memory is not unchanged.

Natural-EOS full-model sanity: AWQ / TP4 / MTP0 / FP16 KV, frozen prompts,
fixed enqueue/free-slot order, no model tracing or attention intervention.
C1/C4/C8 each ran twice: all 26 outputs stopped naturally before the 96-token
limit, reported finite logprobs, passed basic answer checks, and matched token
IDs exactly between repeats of the same shape. Four worker hashes confirmed
the repaired sources. KV capacity remained 386,392 tokens. Actual C4/C8 route
logs showed 184+1 / 368+2 grouped+XQA rows. The runtime retains the existing
AWQ wrapper admission repair; experimental AWQ grouped decode stayed disabled.

Cross-shape/position generation differences remained in that sanity run: one open-ended prompt
returned 61 tokens at C1, 70 at C4, and 61/71 at its two C8 positions. Each
position repeated exactly within that engine. This is not broad batch invariance
or full quality acceptance. The subsequent, separate investigation is recorded
below; it does not replace or retroactively upgrade this sanity result.

Numerical follow-up: findings and decision to defer

After the allocation-order fix, we investigated the remaining cross-batch and
cross-position output differences. Decision: retain this narrow QSA repair,
but defer further precision/invariance changes. No HC, NCCL, LM-head, AWQ,
or scheduler change from these diagnostic experiments is part of this PR.

What was isolated

• HC projection, independent of AWQ experts: replay of the unquantized FP16
  HC down/up projections showed sensitivity to GEMM batch shape and low-precision
  reduction. This is a numerical computation issue, not a wrong AWQ dispatch
  or a repeated instance of the QSA physical-page ordering defect.
• FP16 all-reduce: identical per-rank input rows at different packed positions
  could reduce differently with the tested default NCCL schedule. Replay of the
  actual Layer0 GDN collective reproduced 43,161 different output elements for
  duplicate requests. This identifies a reduction-order effect, not corrupted
  inputs. The test runtime was Torch 2.10.0+cu128 / NCCL 2.27.5 on four V100s.
• Near-tied next-token candidates: the remaining English example differed
  only at token 28, and versus +. All natural English outputs had 81 tokens;
  the other 80 token IDs matched exactly, and both answers computed 156.
  At the compared C8 positions the candidate probabilities were approximately
  49.98%/49.98% versus 49.59%/50.37%. For eight captured candidate pairs,
  CPU FP64 dot products using the actual FP16 hidden states and LM-head weights,
  followed by FP16 rounding, matched the actual outputs exactly (error at most
  0.5 FP16 ULP). Rounding can turn a slight preference into a tie, then argmax
  token-ID tie-breaking changes the wording. No incorrect LM-head rounding
  was found in this case. This is a conditional last-layer reference, not
  an end-to-end FP64 oracle or a blanket claim that every residual error is benign.

Diagnostic controls and their limits

Within one loaded engine, disabling low-precision reduction only during HC
prefill projections, together with FP32 prefill all-reduce, aligned one Chinese
prompt's full 70-token natural-EOS output across C1/C4/C8 positions. Either
control alone did not. Reverting the controls restored the same engine's
baseline logits and natural outputs. However, the English near-tie remained,
and logit tensors were not globally equal: the full prompt set did not pass
cross-batch token invariance. None of these controls changed decode arithmetic.

A separate real-input NCCL replay showed that NCCL_ALGO=allreduce:tree alone
removed the tested position/shape differences, while one channel alone did not.
The cost was material: the 8192-row FP16 collective increased from 5.572 ms
to 8.478 ms (+52.2%). Tree + Simple + one channel measured 7.907 ms (+41.9%).
These are rank0 CUDA-event medians for a single uncaptured collective, two
warmups and ten samples, not whole-model latency or decode tok/s. There is no
whole-model acceptance of a changed NCCL configuration.

Repeating the old call/warmup sequence recovered the prior Chinese sequences,
but English near-tie choices still differed between processes. Same-engine,
fixed-shape repeats were stable in these runs. The complete source of
cross-process numerical drift remains not fully attributed.

Why no further fix is proposed now

The observed English wording change did not change the task answer, and the
last-layer replay found normal finite-precision rounding rather than faulty
arithmetic. Imposing cross-batch bitwise equality would be a distinct product
contract, with broader GEMM/communication work and measured communication cost.
We therefore leave precision and NCCL defaults unchanged, do not promote the
diagnostic FP32 controls, and do not treat C1's particular wording as the sole
correct output. This is an explicit DEFERRED / NO DEFAULT CHANGE decision,
not a declaration that all numerical questions are solved.

Revisit separately if a matched workload demonstrates answer-quality regression,
reference error beyond the accepted numerical contract, or an explicit
deterministic-output requirement. Any follow-up should keep common-prefix
numerical checks, natural-EOS task quality, and performance costs separate;
it should not expand this allocation-order patch.

Contract and limits

Allocation invariance assumes fixed query grouping, logical selections,
visibility and page-alias relationships. It does not promise bitwise equality
across different batch shapes, different attention implementations or formats.
It does not by itself explain every previous free-running generation difference
or fix the independent AWQ W13 numerical difference. No production deployment
or broad NVFP4/long-context quality acceptance is claimed.

AI assistance: OpenAI Codex assisted with investigation, implementation, test
execution and this description. Leon confirmed human review of all four changed
files and explicitly authorized upstream submission after this description was
updated. The listed test runs were executed with Codex assistance; no independent
maintainer review or upstream CI pass is claimed.

[yangzhuxinyzx]

本轮已阅读逻辑排序修复、原始定位和 V100 验证记录。其主要收益是固定 query grouping 下对 KV 物理重定位保持稳定，代价是已测 page4 算子约 +1% 到 +7.3%，并不代表同等 E2E 损失。项目不要求 greedy/batch bitwise identity，因此后续评估应把已修复的分配顺序影响、实际答案质量和速度成本分开；不应据此推进附带的 FP32 HC/NCCL Tree 诊断方案。本轮优先合入了启动、调度与 AWQ 确定收益项；本 PR 保持 Open，列入下一轮注意力专项复核。