[JohnQin-llm]

What

_init_dualwave_thread_mapping took block_idx.x/y directly as (head, q_block). The dualwave launch grid is (NUM_HEADS_Q, num_q_blocks, ...) — head-fast — and gfx950 assigns contiguous workgroup ids to an XCD, each with a private ~4 MB L2. So one head's query blocks were spread across all 8 XCDs, and every XCD re-streamed that head's K/V from HBM.

This re-derives (head, q_block) from the linear workgroup id with the head as the slow axis, keeping a head's blocks on one XCD so its K/V stays resident in that XCD's L2 across the concurrent wavefronts. One file, +24/−2.

This is Swizzled Head-first Mapping from arXiv:2511.02132, applied to the gfx950 dualwave kernels.

Mechanism — measured, not assumed

rocprofv3 --pmc TCC_HIT_sum TCC_MISS_sum, one dispatch at S=65536, H=64, D=128, bf16 non-causal:

	L2 hit	hits	misses
before	2.36%	2.073e+08	8.593e+09
after	93.62%	8.239e+09	5.611e+08

15.3× fewer L2 misses, consistent with the 80–97% L2 hit rates 1 reports on MI300X. The reuse is the concurrent-wavefront mechanism described in 2, not "loaded once" — a head's K/V is well over 4 MB at these sequence lengths.

Performance

MI355X gfx950, H=64 D=128 bf16 non-causal batch=1, median of 30 timed iterations, each build in its own JIT cache dir, part at PCI bus 0x85, this branch vs its merge-base:

seqlen	before	after	delta	aiter v3 ASM
16384	1200.9	1318.0	+9.8%	1252.1
65536	1101.7	1329.2	+20.7%	1255.6
131072	1081.9	1326.9	+22.6%	1223.7

Moves the dense non-causal path from losing to aiter's closed ASM at every size to beating it by 5–8%.

Reproducing absolutes: this workload runs into the 1400 W package cap, so sustained clock — and throughput — varies by part. Across all 8 parts in one node the same measurement spans 1204–1328 TFLOP/s after the change (10–11% spread), while the delta is stable at +16.3% … +22.6% (median +20%) at 65536–131072. Compare before/after on one part; don't compare absolutes across parts.

Test steps and results

1. Repo test suite — 9 configs, before and after

python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 128 --dtype bf16 --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 128 --dtype bf16 --causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 64  --dtype bf16 --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 128 --dtype fp16 --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 128 --dtype fp8  --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --head_dim 128 --dtype bf16 --no-causal --num_kv_splits 2
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 20 --head_dim 128 --dtype bf16 --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 1 --seq_len 4096 --num_heads 64 --num_kv_heads 8 --head_dim 128 --dtype bf16 --no-causal
python3 tests/kernels/test_flash_attn_fwd.py --batch 4 --seq_len 2048 --num_heads 64 --head_dim 128 --dtype bf16 --no-causal

config	base	this branch
bf16 dense H=64 S=4096	PASS	PASS
bf16 causal	PASS	PASS
bf16 D=64	PASS	PASS
fp16 dense	PASS	PASS
fp8 dense (shares this mapping via DualwaveFp8KernelContext)	PASS	PASS
split-K = 2 (must take the fallback)	PASS	PASS
num_heads = 20 (% 8 != 0, must take the fallback)	PASS	PASS
GQA kv_heads=8	PASS	PASS
batch=4	PASS	PASS

2. Bit-exactness

The map is a bijection over [0, NUM_HEADS_Q * num_q_blocks), so every (head, q_block) pair is computed exactly once and output is bit-identical. Verified over 7 shapes — D=64/128, causal, batch=2, and both fallback cases:

case	bit-identical	max abs diff
bfl_h64_s4096 / s8192	yes	0.0
causal_h64	yes	0.0
d64_h64	yes	0.0
batch2_h64	yes	0.0
fallback_h60 / fallback_h4	yes	0.0

0 differing elements in every case. Each was independently checked against an fp32 SDPA reference (worst 4.6e-3 relative to peak — the expected bf16 error, unchanged by this commit), so a fault identical in both trees would still have been caught.

3. Reproducing the perf numbers

# one build per variant in its OWN cache dir -- the mapping is chosen at trace
# time, so sharing a cache dir between variants can return the wrong kernel
FLYDSL_RUNTIME_CACHE_DIR=/tmp/fc/$VARIANT python3 your_bench.py --M 65536

Pin the part by PCI bus (torch.cuda.get_device_properties(0).pci_bus_id) — note ROCR_VISIBLE_DEVICES does not enumerate in the same order as rocm-smi, and given the 10–11% part spread that mismatch is easy to misread as a regression.

Guards

Applied only when not SPLITK and NUM_HEADS_Q % 8 == 0, falling back to the direct mapping otherwise. Split-K adds a third grid axis whose blocks share a K/V range this 2-D linearization would not preserve.

The condition is compile-time from traits rather than an env var, deliberately: the mapping is chosen at trace time, so an env-gated version would not participate in the JIT cache key and two builds with identical traits but different mappings could alias in the disk cache.

References

1. M. Choudhary, K. Sangaiah, S. Singh, M. Osama, L. Wu Wills, G. Dasika. Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects. arXiv:2511.02132. Introduces Swizzled Head-first Mapping; reports up to 50% higher performance and 80–97% L2 hit on MI300X.
2. Y. Zhu, Y. Pan, C. Ding. Sawtooth Wavefront Reordering: Enhanced CuTile FlashAttention on NVIDIA GB10. arXiv:2601.16032. Concurrent-wavefront L2 reuse; the mechanism by which a working set larger than L2 still benefits.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR updates the dualwave flash-attention workgroup-to-(head, q_block) mapping on gfx950 so that a single head’s query blocks are assigned contiguously in linear workgroup-id space (making head the slow axis). This is intended to improve K/V cache residency per XCD by avoiding scattering one head’s workgroups across all XCDs.

Changes:

• Add a gfx950-specific constant (NUM_XCD_GFX950 = 8) with explanatory comment.
• In _init_dualwave_thread_mapping, conditionally remap (head, q_block) from the linearized 2D workgroup id when not SPLITK and NUM_HEADS_Q % 8 == 0; otherwise keep the original direct (block_idx.x, block_idx.y) mapping.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderfeli]

clean comments?

[JohnQin-llm]

Done — cut from 13 lines to 4, and the constant's from 2 to 1.

You were right that it was out of place: this file has 64 comment blocks, and before this change 43 were 1 line and 18 were 2 (median 1, previous max 4). Mine was the only one over 4. It now reads:

# Swizzled Head-first Mapping (arXiv:2511.02132): the grid is head-fast, so one
# head's q-blocks scatter across all XCDs and each re-streams its K/V. Re-derive
# (head, q_block) with head as the slow axis to keep them on one XCD. Bijective,
# so output is bit-identical; split-K's third grid axis would not survive it.

Kept only what isn't inferable from the code — the citation, and why split-K is excluded. The mechanism narrative, the L2 counters and the perf tables live in the commit message and PR description instead.

Only comment lines changed: the AST of the file is byte-identical to the revision I tested (ast.dump hash 98d4320861bbd0bd before and after), so the test and benchmark results in the PR description still apply unchanged.

[JohnQin-llm]

Re-validated against current main (9a5c08e, which now includes #905/#915/#916), and there is an additional result worth recording: this change also speeds up FP8 attention.

M=65536, H=64, D=128, non-causal, batch 1, median of 12 after 8 warmup, PCI bus 0x85, error vs an fp32 reference:

	TFLOP/s	delta	rel err
bf16, main	1099	—	2.39e-03
bf16, this branch	1331	+21.1%	2.39e-03
fp8 per-tensor, main	1764	—	3.87e-02
fp8 per-tensor, this branch	1863	+5.6%	3.87e-02

_init_dualwave_thread_mapping is shared with DualwaveFp8KernelContext, so the L2-locality fix carries to the FP8 path as well. The gain is smaller there — FP8 halves the K/V footprint, so it thrashes L2 less to begin with — but it is real and free, and the FP8 output is unchanged (3.87e-02 both sides, same as the bit-exactness result for bf16).

Two notes for reviewers:

• The branch is still based on 05d785e; it no longer cherry-picks cleanly onto main because [ROCDL] Add arch-aware universal s_waitcnt #915 touched the same function. I can rebase on request — the change itself is unaffected, it is a context conflict only.
• The +21.1% on bf16 reproduces the original claim against the current baseline, so the gain is not something upstream has since obtained by another route.

🤖 Generated with Claude Code

[JohnQin-llm]

Pushed aeb6d54: the swizzle is now restricted to non-causal. Testing the causal path found that this change was a regression there, which the original PR did not cover.

Measured on MI355X, B=1 S=16384 H=64 D=128 bf16, PCI bus 0x85:

	causal	non-causal (S=65536)
stock mapping	1067	1099
swizzle, unguarded (previous revision)	988 (−7.4%)	1331 (+21.1%)
swizzle, guarded to non-causal (aeb6d54)	1066 (neutral)	1328 (+20.8%)

Cause. Under a causal mask, q-block i does work proportional to i+1. The stock head-fast mapping gives each consecutive run of workgroup ids a uniform amount of work — one q-block index across all heads. Making q_block the fast axis instead clusters steadily-increasing work into consecutive ids, so XCDs finish at markedly different times and the tail dominates. The L2-locality win is real in both cases, but under causal it is outweighed by the load imbalance it creates.

The guard is not traits.CAUSAL, alongside the existing not traits.SPLITK and NUM_HEADS_Q % 8 == 0. Causal falls back to the stock mapping and recovers stock performance exactly (1066 vs 1067); non-causal keeps the full gain.

Output remains bit-identical across all 7 validation shapes (D=64/128, causal, batch=2, and both fallback paths), re-run after the change.

A balanced causal mapping — pairing q-block i with n-1-i so each workgroup gets constant work — would in principle let causal keep the locality benefit too. I have not implemented or measured that, and it is a larger change than this PR should carry; flagging it as a possible follow-up.

🤖 Generated with Claude Code

[coderfeli]

@yanguahe @jhinpan

[jhinpan]

Blocking: this is not the cited Swizzled Head-first mapping under AMD's documented SPX round-robin workgroup scheduling.

For this launch grid, the PR defines wid = block_idx.x + block_idx.y * H, then maps head = wid // Q and q_block = wid % Q. A fixed logical head therefore consumes consecutive physical WIDs (head * Q + q_block). Since SPX assigns WIDs round-robin across the 8 XCDs (XCD = wid % 8), any head with Q >= 8 is spread across all 8 XCDs. This is Naive Head-first ordering: it can improve temporal/concurrent-wave locality, so the reported counters and speedup may be real, but it does not keep a head on one XCD as the title, comment, and PR description claim.

The cited paper's actual swizzle preserves the physical WID's XCD residue, e.g. for MHA when H % 8 == 0:

heads_per_xcd = H // 8
head = (wid % 8) * heads_per_xcd + wid // (8 * Q)
q_block = (wid % (8 * Q)) // 8

Then all WIDs for one head have the same wid % 8.

There is also a concrete GQA concern. For Hq=64/Hkv=8, the existing direct mapping already gives wid % 8 == h_idx % 8 == h_kv_idx, keeping all Q heads that share K/V on one XCD. The proposed mapping spreads every logical Q head—and consequently every KV group—across all 8 XCDs. Correctness/bit-exactness tests cannot detect this locality regression, and the PR provides no GQA performance or L2 comparison.

Please implement the modulo/stride spatial swizzle for MHA, and either preserve the current GQA mapping or add a KV-head-aware GQA swizzle. Please also add an enumerative mapping test asserting the number of XCDs used per head/KV group, plus before/after GQA counters or performance results.

References:

• Cited mapping and formula: https://arxiv.org/html/2511.02132v1
• ROCm/AITER spatial mapping description (wid % NUM_XCDS): ROCm/aiter#3936
• AMD SPX scheduling behavior: https://rocm.blogs.amd.com/software-tools-optimization/compute-memory-modes/README.html