[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; they can be taken separately.

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.

Motivation

Bug 1: rows lose softmax mass

Set V to all ones. The output must then be exactly 1.0 everywhere, since softmax weights sum to 1. It is not:

	worst row	rows below 1.0
bf16	0.99996	0%
fp8	0.637	58–83%

Why. P (the softmax weights) is cast to e4m3 before the PV matmul. The smallest number e4m3 can hold is 2**-9. With thousands of keys, most weights are smaller than that, so they become zero. But l_row — the sum used to normalise — is computed before that cast, and still counts them. The output is divided by a sum that includes weights the matmul dropped.

This only shows up when attention is peaked. On near-uniform attention every weight is similar and nothing is lost, so a random-init model will not reveal it.

Bug 2: int32 overflow on the shape

jit_argument._LayoutPlan packs shapes as int32 ("i" * len(shape)); 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, where no single dim is that large — the comment above that branch already warns about this.

Changes

Bug 1 — scale P up before the cast, so weights use e4m3's full range instead of only its top. l_row is scaled by the same factor, so the output is unchanged apart from the tail that no longer underflows. This is free: the scale folds into an FMA that already runs, added as an optional bias argument to _scale_sub_score_pair, so bf16 callers are untouched.

How much scale is safe is set by how large exp2 can get: the lazy path holds the running max until a tile exceeds it by RESCALE_THRESHOLD, leaving log2(448) - THRESHOLD; the eager path rebases every tile and gets all of log2(448). Both paths derive their scale from that.

Doing this surfaced that the fp8 kernel never read DUALWAVE_SWP_LAZY_RESCALE — both of its call sites used lazy_correct_o unconditionally, so setting the flag did nothing on that path. The bf16 kernel has always honoured it; the fp8 one now does too.

Bug 2 — run one kernel launch per batch entry when the total would overflow. Batch entries are independent in attention, and a leading slice of a contiguous tensor is still contiguous, so this needs no copy and no kernel change. Guarded on numel >= 2**31.

fp8 now defaults to the eager rescale. The lazy path cannot lift P far enough to keep the softmax tail out of e4m3's flush-to-zero range, so leaving it on ships a known-wrong default. dualwave_swp_lazy_rescale becomes Optional and resolves to False for fp8, True for bf16 — bf16 is untouched and explicit callers are unaffected. Bug 2's path only runs at sizes that raise today.

At B=2 H=8 D=128 with V=ones, where the exact output is 1.0:

S	k scale	worst row	mean |o-1|	TFLOP/s
1024	1	0.9844 → 0.9844	0.0139 → 0.0105	165.0 → 158.4
1024	200	0.9492 → 0.9883	0.0034 → 0.0014	162.7 → 160.4
12288	1	0.9805 → 0.9805	0.0137 → 0.0044	1733.2 → 1632.1
12288	200	0.7656 → 0.9922	0.0021 → 0.0004	1732.0 → 1611.4

Error drops 3-5x across the board for 1.4-7.0%, most of it at long sequences.

Performance Results (MI355X, gfx950)

Attention forward, B=2 D=128, non-causal.

Configuration	Before	After	Change
fp8, H=64, S=27280	1972 TFLOP/s	1873 TFLOP/s	−5.0%
fp8, H=64, S=61380	2023	1881	−7.0%
fp8, H=64, S=180180	1990	1849	−7.1%
fp8, H=32, S=61380	2024	1918	−5.2%
bf16, H=64, S=27280	1253	1253	unchanged

fp8 over bf16 goes from 1.57–1.60x to 1.47–1.51x. The cost is the eager rescale, not the scaling.

Accuracy on six q/k/v triples captured from a real 80B model. "Floor" is the error from the e4m3 format alone — quantise, dequantise, run in fp32 — the best any fp8 kernel could do.

	Before	After
worst capture	2.44x floor	1.05x floor
range over six	1.06–2.44x	1.00–1.21x

A CUDA fp8 attention kernel reaches 1.05x of the same floor on the same tensors, so 1.05x is the achievable target rather than a number I picked.

Choosing the threshold

The default threshold of 8 leaves a 2x scale, which helps at no cost. Lowering it buys the rest. Measured at B=2 S=12288 H=8 D=128 with a wide score range, V=ones so the exact output is 1.0:

threshold	P scale	worst row	TFLOP/s
8 (current default)	2x	0.766	1641
6	7x	0.910	1632
4	28x	0.988	1548
2	112x	0.996	1498
0 (eager)	448x	0.996	1506

There is no cheap middle ground. By the time the threshold is low enough to fix the row error, it costs what the eager path costs. So the choice is between the free partial fix at the default and the full fix at −8%; there is nothing in between.

Testing

# the two new tests
pytest tests/kernels/test_flash_attn_fwd.py -k "fp8_softmax_normalises or fp8_out_tensor" -v

# the whole forward suite
pytest tests/kernels/test_flash_attn_fwd.py -q

• Unit tests added — test_fp8_softmax_normalises uses the V=ones property above, so it needs no reference implementation, and covers both rescale paths with the bound each can reach; test_fp8_out_tensor_is_filled_and_returned covers the out contract on the split path
• Existing test_flash_attn_fwd.py passes — full file 97 passed cold (FLYDSL_RUNTIME_ENABLE_CACHE=0, 2m49s), with the new default in place
• Performance benchmarks run — table above
• Tested on MI300X — the dense fp8 path is gfx950-only; the test skips elsewhere
• Also checked with causal masking, and over S from 1k to 12k
• 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.