[Perf] Enable the XCD swizzle and pin the softmax basis for gfx950 bf16 attention

[JohnQinAMD]

Summary

Two changes to the bf16 DUALWAVE_SWP dense path on gfx950: enable the XCD workgroup remap that already existed but was reachable only from the fp8 builder, and pin the softmax basis to the prologue tile max, which in turn lets the softmax scale move out of Q. Together they put FlyDSL 5.4-6.1% above aiter with the accuracy gap against it closed from 25% to 2%.

Motivation

Throughput decayed with sequence length where aiter's stayed flat, and accuracy sat 25% worse than aiter at every length. Each had a single cause, and the two fixes turn out to depend on one another.

Workgroups map to XCDs as linear_id % 8 and the grid is head-fast, so with H % 8 == 0 every workgroup of a head lands on the same XCD and each XCD juggles H/8 concurrent K/V streams against its L2 slice. At B=2 H=32 S=180,180 that is a 30.0% L2 hit rate and 1360 GB of HBM reads against a compulsory 8.9 GB.

Separately, Q is pre-scaled by sm*log2e and narrowed back to bf16 before the QK MFMA. That extra narrowing is the whole of the accuracy deficit: modelled, P and O narrowing alone give 2.350e-03 against aiter's measured 2.296e-03, and adding the Q narrowing gives 2.896e-03 against FlyDSL's measured 2.869e-03.

Changes

• Plumb xcd_swizzle through the bf16 builder and reuse the fp8 dispatch condition (non-causal, H % NUM_XCD_GFX950 == 0, num_q_blocks >= MIN_Q_BLOCKS_XCD_SWIZZLE). Split-K is unreachable here, taking _build_splitk above this site. Auto-dispatched; the remap is bijective so output is bit-identical.
• Add dualwave_swp_fixed_basis, off by default. It pins the basis, which makes the per-tile reduce_max and the online rescale dead code, and moves the scale into sub_m as one FMA where a subtract already stood. One trait drives both: raw logits without a pinned basis costs 6% rather than 0.66%, since four of its five scale multiplies live in the rescale path.
• Promote the block_m literal to DUALWAVE_SWP_BLOCK_M so the dispatch can count q-blocks before traits exist.
• No change to the fp8, varlen, paged or split-K paths.

Performance

MI355X, bf16, D=128, non-causal, B=2 H=32, median of three, aiter measured in the same process on the same tensors. Before means every option off, including the swizzle, which otherwise auto-dispatches on these shapes. Cells are TFLOP/s and relative L2 against float32.

Configuration	Before	After	Improvement
S=61,380	1258 / 2.87e-03	1290 / 2.35e-03	+2.5% speed, 18% less error
S=180,180	1221 / 2.88e-03	1302 / 2.35e-03	+6.6% speed, 18% less error
S=239,580	1196 / 2.87e-03	1303 / 2.35e-03	+8.9% speed, 18% less error
aiter, same shapes	1225-1232 / 2.29e-03	—	this lands 5.4-6.1% above it

Causal gains 17% accuracy on the same terms. LSE improves three orders of magnitude, 1.5e-03 → 1.9e-06, since it reads m directly rather than through a normalisation that divides the error out; that matters for the backward pass.

Testing

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

Disabling the cache is not optional here. This change touches helper code outside
the traced closure, so the JIT cache key does not move with it and a warm cache
serves the previous kernel — CLAUDE.md calls this case out. Four tests fail that
way and pass cold.

• Unit tests added/updated — four new tests in tests/kernels/test_flash_attn_fwd.py: bit-identity of the remap over H=8/16/32/64, fallback when H % 8 != 0 both forced and auto, accuracy improvement for causal and non-causal, and LSE improvement. Full file passes 45/45.
• Performance benchmarks run — table above.
• Tested on MI355X (gfx950). Not run on MI300X: the DUALWAVE_SWP path is gfx950-only and every new test carries @_requires_gfx950.

Dependencies

• No new third-party dependencies added

Breaking Changes

None. The remap is bit-identical, so enabling it by default changes timing only. The pinned basis is not bit-identical and therefore defaults off.

One case is uncovered by it: a later tile whose scores exceed the basis by more than ~109 exp2 units would overflow. Benign data cannot reach that — the measured worst case for the prologue max against the true row max is 5.4 exp2 units — but massive-activation data could, and there is no trip test for it. Adding one was measured and rejected: a trip test needs the per-tile row max, which is exactly what pinning the basis removes, so it gives back more than the option gains (1298 -> 1240 TFLOP/s). Raising the lazy-rescale threshold instead buys 0.15%, since the cost is per-tile rather than on the rare path.

[Copilot]

Pull request overview

This PR wires up the existing gfx950 “XCD head-slow remap” (XCD swizzle) to the bf16/f16 dense DUALWAVE_SWP flash-attention path so large-sequence throughput doesn’t degrade due to poor L2 residency across XCDs.

Changes:

• Plumbs an xcd_swizzle build flag through the dense dualwave builder and exposes it via flydsl_flash_attn_func(..., dualwave_swp_xcd_swizzle=...) (auto/force on/off).
• Adds auto-dispatch logic for enabling the remap on eligible dense shapes (mirroring the existing fp8 behavior).
• Adds gfx950 tests to assert bit-identical output when swizzle is toggled and to validate correctness when heads don’t divide the XCD count.

Reviewed changes

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

File	Description
tests/kernels/test_flash_attn_fwd.py	Adds tests for bit-identical output and correctness around the XCD swizzle option.
kernels/attention/flash_attn_interface.py	Exposes dualwave_swp_xcd_swizzle, computes auto-dispatch condition, and passes the flag into the cached dualwave builder.
kernels/attention/flash_attn_gfx950.py	Adds a private _xcd_swizzle build parameter and forwards it into traits creation.

---

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

[JohnQinAMD]

The single-line form was Black's own output at line-length 120 — it collapsed the multi-line comprehension I originally wrote, and black --check passed on it. So the formatting gate was not at risk.

The readability point stands, though. Replaced the comprehension with an explicit local run(flag) helper, which reads better and is stable under Black.

[JohnQinAMD]

Good catch — the docstring was narrower than the actual guarantee. There are two guards: the dispatch condition auto-selects against H % NUM_XCD_GFX950 != 0, and _init_dualwave_thread_mapping re-checks it independently, which is what makes forcing the flag safe. Docstring now states both.

[JohnQinAMD]

Agreed, and it is exactly the claim the PR description makes, so it should be pinned. The test is now parametrized over dualwave_swp_xcd_swizzle in [None, True], so the forced-True case at H=12 is covered alongside the auto-selected one. Both must still match SDPA.