Skip to content

fix(fp8): guard every K stage aliasing the persist-D Q tile area (RFC FC-17) - #382

Merged
DefTruth merged 1 commit into
mainfrom
dev
Oct 14, 2026
Merged

DefTruth merged 1 commit into
mainfrom
dev

Conversation

@DefTruth

Copy link
Copy Markdown
Member

Problem (RFC FC-17)

--cuda-impl fp8 --D 192 returned wrong rows: O_err ≈ 3.1 (fp8 dense noise floor is ~0.02-0.05), allclose=False. The error is confined to the second half of every kBr=128 tile (row % 128 ∈ [64,127]), hits all 32 heads, is non-deterministic across runs (bad-row set Jaccard ≈ 0.38), is independent of every quantization knob, and is absent on the fp16 CUTE path (same D) and on the sm89 fp8 path.

Root cause

cute/fp8/sm_120/persist_d.cuh drops the Q tile from the smem allocation and lets the K stages reuse that area (kPersistQs2r). The hand-off only waited for the consumer Q s2r before the last logical stage's first TMA (k_tile == kStagesK-1) — correct only while the Q tile holds a single K tile, i.e. while kBc == kBr (D ≤ 128).

For D > 128 the launcher shrinks kBc to 64, so kQTileElements/kKTileElements = kBr/kBc = 2: the Q tile holds two K tiles and the prologue's stage 0 (physical slot 1) also lands inside the live Q tile with no guard. Its TMA raced the consumers' Q ldmatrix s2r and the Q area ended up holding whichever data arrived last.

That single defect explains every fingerprint: only the kBc < kBr branch is affected (D = 160/192/224, not just 192); the bad rows sit in the upper half because those rows lose 2/3 of their D range to the overlap while the lower half loses 1/3 (measurably degraded too — lower-half mean 0.0505 / max 0.32 vs the 0.0015 noise floor — just below the 0.3 gross-error threshold); it floats across runs; fp16 is clean because its k_base = q_base + kQTileElements never overlaps; sm89 is clean because kBr == kBc there, so only one slot can land in the Q area.

Fix

Move the q_consumed wait to the first stage that maps into the Q area:

  • kQReuseSlots = kQTileElements / kKTileElements, kQGuardStage = (kPersistQs2r && kQReuseSlots > 1) ? 0 : kStagesK - 1
  • the prologue K loop waits on q_consumed when s == kQGuardStage; the main-loop guard becomes k_tile == kQGuardStage

The two guard sites are disjoint in k_tile (prologue s ∈ [0, kStagesK-2], main loop k_tile ≥ kStagesK-1), so exactly one wait runs; kQReuseSlots == 1 (D ≤ 128) keeps the previous codegen. Two static asserts now pin the invariants the reuse relies on (the Q tile holds a whole number of K tiles; the K region spans the Q tile area, else the V prologue would alias it unguarded).

Verification (RTX PRO 5000, sm_120)

before after
D=192 dense B1H32N8192 (max_abs / mean_abs) 3.121 / 0.0117 0.0188 / 0.00078
D=192 dense, 4 repeat runs non-deterministic bitwise identical
bench CLI --cuda-impl fp8 --D 192 --N 8192 0/6 tasks 6/6 tasks (O_err 0.0132 / 0.0527 / 0.0146 / 0.0166 / 0.0137 / 0.0132)
108-config matrix (N × seed × dense/causal × 6 quant knobs) gross-error rows everywhere 0 gross-error rows; the 7 residual N=512/causal cases match the untouched sm89 reference value for value (fp8 noise floor)
D=64 / D=128 output — bitwise identical to the pre-fix build
perf, 6 CLI tasks @ N=8192 — within ±1%

stages=2 and stages=3 both clean and numerically identical; FFPA_FP8_SM89_FORCE=1 (simulated sm89 at D=192) clean.

Tests / docs

  • tests/test_ffpa_fp8.py: D=192 added to the parity sweep, plus a dedicated test_fp8_persist_d_q_reuse_race. It must use unscaled fixtures — with the usual randn * 0.5 the corruption stays inside tolerance (pre-fix 0.0137 → PASS), so the scaled sweep cannot see it. Red-green confirmed against the pre-fix binary.
  • docs/user_guide/install.md: the public D=192 warning is now a version bound (affected ≤ 0.2.4, fixed on the 0.2.5 dev series) and its scope is corrected (D > 128 window; sm89 not affected).
  • Skill/RFC docs (ffpa-cuda-understand): FC-17 checked off with root cause, evidence and the follow-up.

Known gaps

  • D=160/224 share the branch by construction but could not be measured here (not in the compiled headdim set; adding 224 fails an unrelated sm89 D % 64 == 0 static assert).
  • The two fp8 tests live in a file that is skipped on bf16-only builds, so the shipped wheel's fp8 regression still rests on the bench CLI matrix (recorded as a follow-up).

RFC FC-17: for D>128 the launcher shrinks kBc to 64, so the Q tile area
holds two K stages while only the last one was gated by q_consumed (the
prologue's K TMA then raced the consumer Q s2r and corrupted ~15% of the
upper tile-half rows, non-deterministically). Move the wait to the first
stage that maps into the Q area, plus two static asserts for the layout
invariants it relies on.

D=192 O_err 3.13 -> 0.013, deterministic across runs, D<=128 bitwise
unchanged, perf within +-1%; 108-config matrix has no gross-error rows.
@DefTruth
DefTruth merged commit 333dd75 into main Oct 14, 2026
2 checks passed
@DefTruth
DefTruth deleted the dev branch October 14, 2026 08:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant