[Bugfix][Kernel] Fix NaN from the lazy rescale on a wide score range

[JohnQinAMD]

Summary

The dense bf16 DUALWAVE_SWP forward returns NaN in its default configuration once the logit spread widens. At B=2 S=8192 H=32 D=128 with K scaled 32x, 9.2% of the output is NaN; at 64x, all of it. dualwave_swp_lazy_rescale=False is finite on the same inputs, and so is aiter. Not a regression — a 7/31 checkout reproduces it identically.

Repro

B, S, H, D = 2, 8192, 32, 128
q = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
k = (torch.randn_like(q).float() * 32).to(torch.bfloat16)
v = torch.randn_like(q)

flydsl_flash_attn_func(q, k, v, causal=False)                                       # 9.2% NaN
flydsl_flash_attn_func(q, k, v, causal=False, dualwave_swp_lazy_rescale=False)      # finite

Root cause

The rescale branch is wave-uniform. below is per-lane, but the ballot collapses it into one decision for the whole wave:

below = fx.Float32(m_diff) <= c_eight_f
ballot = rocdl.ballot(T.i64, as_mlir_value(below))
all_below = arith.cmpi(eq, ballot, _read_exec_i64())

So one lane past DUALWAVE_SWP_RESCALE_THRESHOLD sends every lane into _lazy_rescale_o_rescale, including lanes whose tile max sits below their running max. That body took the tile max unconditionally:

corr = rocdl.exp2(T.f32, as_mlir_value(m_row - m_tile_max))
...
out.append(_anchor_scalar_f32(m_tile_max))

For a lane that did not trigger the branch, m_tile_max < m_row, so corr = exp2(positive) and it scales v_o and l_row up. With a wide spread that reaches inf, and the final v_o / l_row is inf/inf = NaN. The eager rescale_o never had this problem — it takes m_new = maxnumf(m_row, m_tile_max) first.

Evidence: LSE comes back inf on 51,757 of the 52,218 NaN rows (438 NaN, 23 finite), while good rows sit in [80.8, 217.8]. The row statistics break first; the accumulator is downstream of them.

The fix

Take the max, exactly as the eager path does. For a lane that did trigger the branch this is m_tile_max and nothing changes; for a lane dragged along, corr becomes exp2(0) = 1 and its accumulators are left alone.

Result

rel L2 against fp32, MI355X, cache cold (FLYDSL_RUNTIME_ENABLE_CACHE=0):

K scale	before	after	eager (reference)
8	5.468e-03	5.404e-03	5.298e-03
16	7.291e-03	7.209e-03	7.159e-03
24	NaN, 0.2%	8.667e-03	8.637e-03
32	NaN, 9.2%	9.923e-03	9.903e-03
48	NaN, 91.1%	1.207e-02	1.205e-02
64	NaN, 100%	1.391e-02	1.390e-02

The two paths are mathematically the same and now agree to within 0.2%. The two rows that never produced NaN also improve — those lanes were being scaled wrongly all along, just not far enough to overflow.

Throughput is unchanged. 1252.0 / 1077.3 / 1053.9 TFLOP/s before against 1251.0 / 1076.1 / 1053.7 after, at S=61,380 / 180,180 / 239,580, two rounds each. The extra maxnumf sits on the rescale branch only.

The fp8 helper, second commit

DualwaveFp8SoftmaxHelper has the same wave-uniform branch and produces NaN too, though only at a much wider range — its threshold compares (m_tile_max - m_row) * logit_scale, so the same 8 is worth ~91 raw log2 units. At S=12288 with V=ones: 20.0M NaN elements at k*5000, all 25.2M at k*20000.

It cannot take the max. Rebasing to a lower running max is what keeps P large against the base, and pinning it costs accuracy at ranges that work today — mean |o-1| goes 0.0001 → 0.0088 at k*1000. What overflows is the missing bound, not the rebase, so cap the drop at 2**16 instead:

k scale	before	bound 8	bound 16	bound 32
1000	0.0001, 0 NaN	0.0021	0.0001	0.0001
5000	0.0000, 20.0M NaN	0.0025	0.0019	0.0001, 887k NaN
20000	nan, 25.2M NaN	—	0.0006	0.0006, 2.4k NaN

16 leaves every working range bit-for-bit where it is and takes the NaN to zero; 32 starts leaking again. Verified on main and on top of #1020, whose own fp8 tests pass unchanged (11 passed) and whose threshold table is unaffected — k*200 measures 0.0021 either way.

Testing

FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest tests/kernels/test_flash_attn_fwd.py -q

• Unit test added — test_lazy_rescale_survives_a_wide_score_range over K scales 8/32/64, failing on main at 32 and 64. Every existing case uses near-uniform attention, where the running max barely moves and this branch is never reached.
• Full file: 90 passed cold (2m50s), against a FlyDSL built from this branch's base.
• Performance benchmarks run — table above.
• Tested on MI355X (gfx950). The DUALWAVE_SWP path is gfx950-only; the new test carries @_requires_gfx950.
• No new third-party dependencies added

Breaking Changes

None. Output changes only where it was NaN, or where a lane was being scaled by a factor it should never have received.

[Copilot]

Pull request overview

Fixes NaNs in the gfx950-only dense bf16 DUALWAVE_SWP flash-attention forward path when using the lazy online rescale under wide logit ranges by aligning the lazy rescale math with the eager rescale’s max-tracking behavior. Adds a targeted regression test that widens the score range via K scaling and asserts finiteness + accuracy parity versus the eager path.

Changes:

• Update the lazy rescale branch to use m_new = max(m_row, m_tile_max) when computing the correction factor and when anchoring the updated max.
• Add a gfx950-gated test that reproduces the wide-score scenario (via K scaling) and checks lazy rescale produces no NaNs and is not worse than the eager path.

Reviewed changes

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

File	Description
kernels/attention/flash_attn_utils.py	Fix lazy rescale correction/anchoring to prevent overflow → NaN under wave-uniform rescale branching.
tests/kernels/test_flash_attn_fwd.py	Add gfx950 regression test covering wide-score lazy rescale stability and accuracy relative to eager.

---

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

[JohnQinAMD]

Good catch — the assertion covers any non-finite value, so the message now reports NaN and inf separately, and the eager baseline is asserted finite before it is used as the reference. Pushed.