[Bugfix][Kernel] Fix fp8 dense attention: softmax normalisation and an int32 dim overflow

[JohnQinAMD]

Summary

Two independent bugs in the dense fp8 attention forward on gfx950, one commit each.

1. Attention weights do not sum to 1. Some rows come back with only 64% of their weight. bf16 is correct on the same inputs.
2. The kernel raises struct.error for large inputs. Any dense fp8 call with B*S*H*D >= 2**31 fails. At D=128, H=64 that is S >= 131072.

Bug 1: rows lose softmax mass

Set V to all ones and the output must be exactly 1.0, since softmax weights sum to 1. bf16 gives 0.99996; fp8 gives 0.637, with 58–83% of rows short.

P is cast to e4m3 before the PV matmul, and e4m3 bottoms out at 2**-9. With thousands of keys most weights fall below that and become zero — but l_row, the normaliser, is summed before the cast and still counts them. Only peaked attention shows this; a random-init model will not.

Fix: scale P up before the cast so it uses e4m3's whole range. l_row scales with it, so nothing changes but the tail that no longer underflows. Free — the scale rides an FMA that already runs, as an optional bias on _scale_sub_score_pair; bf16 never passes one.

The safe lift is log2(448) - RESCALE_THRESHOLD, so fp8 picks its threshold from the KV length: 6 at or below S=4096, 4 above — see below. bf16 keeps 8.

This also surfaced that the fp8 kernel never read DUALWAVE_SWP_LAZY_RESCALE — both call sites used lazy_correct_o unconditionally, so the flag did nothing. It works now.

Bug 2: int32 overflow on the shape

jit_argument._LayoutPlan packs shapes as int32; only strides honour use_32bit_stride. The fp8 path flattens Q/K/V/O to 1-D, so the packed dim is the total element count and it overflows. bf16 passes the natural 4-D shape and is exempt.

Fix: one launch per batch entry when the total would overflow. Entries are independent and a leading slice of a contiguous tensor is still contiguous, so no copy and no kernel change.

Performance and accuracy (MI355X, gfx950)

Eight attention layers built from trained q/k projections — random init hides this bug. V=ones, so the worst row is the least of mean(o) over rows. Worst layer of the eight; throughput relative to main. Bold is what this PR selects.

S	main	thr 8	thr 6	thr 4	eager
1,024	0.934	0.953 / +0.2%	0.957 / −0.1%	0.957 / −1.2%	0.981 / −1.7%
4,096	0.867	0.910 / −1.2%	0.957 / −2.7%	0.957 / −4.6%	0.981 / −6.7%
16,384	0.734	0.797 / −1.3%	0.910 / −1.9%	0.965 / −2.8%	0.984 / −8.0%
32,768	0.664	0.734 / −1.8%	0.879 / −2.4%	0.961 / −3.0%	0.984 / −8.7%
131,072	0.447	0.551 / −1.4%	0.789 / −2.1%	0.906 / −2.4%	0.984 / −10.2%

Scaling P is free; the threshold is what costs. At S=4096 the scaling alone is −1.2% and takes the worst row 0.867 → 0.910; the threshold takes it the rest of the way.

Why not one threshold. The lazy path holds one running max across every tile, so as S grows more tiles sit far below it and lose their tail — 6 falls 0.957 → 0.789 across this range while 4 holds 0.906–0.965. But at and below 4096 the two are equally accurate and 6 is 2 points cheaper. Both curves move with S in opposite directions, so the crossover is a length. S=131,072 is Bug 2's own regime — the second commit exists to let that shape launch at all — which is where the default has to be right.

Why not eager. It is flat at 0.98, but its cost climbs with S (−1.7% → −10.2%) while threshold 4's falls, so it is worst exactly where help is needed most. It stays reachable as flydsl_flash_attn_func(..., dualwave_swp_lazy_rescale=False), which now selects it on fp8 instead of being ignored.

The kernel is not specialised on S, so selecting on it widens the fp8 build cache from one variant to two; the bound is a constant to keep it at two. bf16 is untouched by construction and measured unchanged.

Testing

pytest tests/kernels/test_flash_attn_fwd.py -k fp8 -v
pytest tests/kernels/test_flash_attn_fwd.py -q

• Unit tests added — test_fp8_softmax_normalises needs no reference implementation, and its lazy bound pins the per-length choice: pinned at 6 it reaches 0.09 and at 8 it reaches 0.23, both past the 0.06 allowed
• Existing test_flash_attn_fwd.py passes — 117 passed cold (FLYDSL_RUNTIME_ENABLE_CACHE=0)
• Performance benchmarks run — table above
• Tested on MI300X — the dense fp8 path is gfx950-only; the test skips elsewhere
• No new third-party dependencies added

Breaking Changes

None.

[Copilot]

Pull request overview

Fixes two independent correctness issues in dense fp8 flash-attention: (1) softmax normalization error caused by fp8 underflow in the PV path, and (2) a launch-time crash from int32 shape packing overflow when flattening very large tensors. The changes are localized to fp8 dense attention utilities, the Python interface wrapper, and the gfx950 fp8 kernel path.

Changes:

• Add an optional log2-domain bias to the fused scale*(score-row_max) helper to lift softmax probabilities before fp8 casting (avoiding tail flush-to-zero mass loss under fp8 e4m3).
• Add a dense-fp8 guard to split very-large B*S*H*D launches into per-batch-entry launches to avoid struct.pack int32 shape overflow.
• Honor DUALWAVE_SWP_LAZY_RESCALE by routing through a helper that selects lazy vs eager rescale correction.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
kernels/attention/flash_attn_utils.py	Adds optional bias support in softmax pre-processing to prevent fp8 PV underflow from breaking normalization.
kernels/attention/flash_attn_interface.py	Adds a size-triggered per-batch split path for dense fp8 to avoid int32 ABI shape packing overflow.
kernels/attention/flash_attn_fp8_gfx950.py	Ensures the lazy-rescale flag is actually respected by selecting the appropriate online-rescale correction path.

Suppressed comments (1)

kernels/attention/flash_attn_utils.py:434

• The _scale_sub_score_pair docstring still describes the return as scale * (v_s - row_max_raw), but the function now also supports an optional bias term (log2-domain shift). Updating the docstring will prevent future callers from missing the bias behavior when reading this helper.

    Returns ``scale * (v_s - row_max_raw)`` per element via a single FMA
    (``fma(s, scale, -scale*row_max_raw)``), so the fp8 QK MMA can emit raw

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jhinpan]

Requesting changes for two correctness regressions in the new fp8 batch-split path.

The softmax headroom change and the signed-int32 boundary hold up in local validation, and the exact head passes the full attention file (106/106). The split path, however, (1) silently discards unsupported bias / alibi_slopes / sink arguments before validation and (2) can copy side-stream outputs on the current stream before the kernels finish when out=None. Both were independently reproduced; the inline comments include minimal fixes and regression-test specifications.

One smaller durability item: the library's new fp8 default is eager, but test_fp8_softmax_normalises always passes the flag explicitly and the CLI/benchmark configuration still forces dualwave_swp_lazy_rescale=True. Please let the CLI inherit None and add one test that omits the keyword and matches explicit eager. This is secondary to the two blockers above.

[jhinpan]

The two split-path correctness blockers are fixed. I re-ran the original modifier and delayed-side-stream reproductions, the new targeted cases, and the full attention file: 17/17 targeted and 110/110 overall passed on gfx950.

One requested item remains: the PR changes fp8's library default to eager, but the first-party CLI/benchmark path still forces the old lazy value, and no test can detect a default regression. See the inline comment. Requesting changes only for that remaining default/coverage inconsistency; the implementation fixes themselves hold up.

[jhinpan]

LGTM. The split-path modifier validation and explicit-stream fix remain correct, and the final default/CLI gap is now closed: the config and forwarding sites carry None, the library resolves per dtype, both CLI overrides remain available, and the omitted-keyword fp8 test matches explicit eager bit for bit.

Exact-head gfx950 validation: 13/13 focused cases and the full attention file 111/111 passed. GPU CI jobs are still pending at review time; the completed style, docs, signal, and Navi checks are green.

[jhinpan]

Re-reviewed exact head ef65bdb after the rebase. The eight-commit patch matches the previously approved series; the modifier validation, side-stream output fix, and CLI/default coverage remain intact. On a checkout-matched dd83733 runtime, all 18 focused fp8 regressions and the full attention file (116/116) pass on an isolated gfx950. The current check-signal failure is its pre-check polling timeout, not a code-test failure.

[binding7012]

LGTM