[JohnQin-llm]

Summary

Restore the XCD workgroup remap on the gfx950 bf16 dense DUALWAVE_SWP path. It shipped in #919 and applied to bf16 automatically; #950 put a default-off trait in front of it while optimising fp8, and bf16 has run without it since 8/3. The remap is bijective, so this is a timing change with bit-identical output.

#919 gated the remap on shape alone. #950 added a trait:

-    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
+    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

XCD_SWIZZLE defaults to False and only the fp8 builder passes it. That PR was scoped to fp8 and never touched the bf16 builder — xcd does not appear in flash_attn_gfx950.py at all — and no test covered the mapping, so nothing failed.

Changes

• Plumb xcd_swizzle through the bf16 builder and reuse the existing dispatch condition (non-causal, H % NUM_XCD_GFX950 == 0, num_q_blocks >= MIN_Q_BLOCKS_XCD_SWIZZLE). Auto-dispatched, with dualwave_swp_xcd_swizzle=True/False to force. Split-K takes the _build_splitk branch above this site and is excluded structurally.
• Promote the block_m literal to DUALWAVE_SWP_BLOCK_M so the dispatch can count q-blocks before traits exist.
• Two tests, which perf(attention): keep a head's Q-blocks on one XCD in the dualwave mapping #919 shipped without: bit-identity of the remap over H=8/16/32/64, and the H % 8 != 0 fallback.
• No change to the fp8, varlen, paged or split-K paths.

Result

Same configuration, median of three, cache cold. "Before" forces the remap off, which is what main does today.

S	before	after	improvement	rel L2 (both)
61,380	1251.1	1263.0	+1.0%	2.8668e-03
180,180	1083.0	1271.0	+17.4%	2.8764e-03
239,580	1062.3	1268.3	+19.4%	2.8665e-03

Long sequences go from 15.0% behind aiter to 1.4% ahead, and land 1.4% short of the 1286.8 measured at #919's merge — that remainder is #930 dropping a fastmath flag, fixed separately in #1035. Short ones barely move: the decay only bites once a head's K/V outgrows the XCD's L2 slice, the mechanism #919 documented (L2 hit 2.36% → 93.62% at S=65536 H=64).

Testing

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

The cache flag 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 — docs/kernel_authoring_guide.md calls this case out.

• Unit tests added — 93 passed cold (3m10s), against a FlyDSL built from this branch's own base.
• Performance benchmarks run — tables above.
• Tested on MI355X (gfx950). Not run on MI300X: the path is gfx950-only and both new tests carry @_requires_gfx950.
• No new third-party dependencies added

Breaking Changes

None. The remap is bijective and the tests assert bit-identical output.

cc @binding7012 @coderfeli — #950 is where the default was introduced; flagging so the fp8 side can confirm nothing there depends on bf16 staying off.

[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.

[JohnQin-llm]

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.

[JohnQin-llm]

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.

[yanguahe]

LGTM

[jhinpan]

Post-merge follow-up from a three-pass review: I found no correctness blocker in the landed code. The remap is bijective over the launched grid, the ineligible guards are preserved, and XCD_SWIZZLE is represented in the specialization/cache keys.

Two durable follow-ups would prevent this feature from silently disappearing again:

1. The tests do not positively prove that the default None path enables the swizzle. The eligible test compares explicit False/True, while None is only exercised with H=12, where it must remain disabled. A disconnected auto-selector would therefore leave both tests green. Please add a host-side selector/dispatch assertion for an eligible shape, including the 63/64 q-block boundary.
2. The XCD comments describe contradictory geometry. With H % 8 == 0, the stock linear workgroup ID already pins head h to XCD h % 8; the remapped IDs span XCDs. The benefit is the smaller concurrently resident K/V working set, not lifetime head-to-XCD affinity. The same H=32 evidence being neutral near 240 q-blocks but strongly positive near 704 also shows that sequence length matters, so the current H/8-only wording should be corrected.

These are test/documentation follow-ups, not a request to revert the merged fix.