perf(attention): keep a head's Q-blocks on one XCD in the dualwave mapping - #919
perf(attention): keep a head's Q-blocks on one XCD in the dualwave mapping#919JohnQinAMD wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
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 whennot SPLITKandNUM_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.
| # The map is a bijection over [0, NUM_HEADS_Q * num_q_blocks), so every | ||
| # (head, q_block) pair is still computed exactly once and the result is | ||
| # bit-identical. Split-K adds a third grid axis whose blocks share the | ||
| # same K/V range, so it keeps the direct mapping. |
There was a problem hiding this comment.
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.
460f351 to
177caf5
Compare
…pping Implements Swizzled Head-first Mapping (Choudhary et al., "Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects", arXiv:2511.02132) for the gfx950 dualwave kernels. The dualwave launch grid is (NUM_HEADS_Q, num_q_blocks, ...) -- head-fast -- and `_init_dualwave_thread_mapping` took block_idx.x/y directly as (head, q_block). gfx950 has 8 XCDs each with a private ~4 MB L2, and the hardware assigns contiguous workgroup ids to an XCD, so one head's query blocks were spread across all 8 XCDs and every XCD re-streamed that head's K/V from HBM. Re-derive (head, q_block) from the linear workgroup id with the head as the slow axis, so a head's blocks stay on one XCD and its K/V stays resident in that XCD's L2 across the concurrent wavefronts. The reuse is the concurrent-wavefront mechanism, not "loaded once" -- a head's K/V is well over the 4 MB L2 at these sequence lengths. Mechanism confirmed by hardware counters (rocprofv3, TCC_HIT_sum / TCC_MISS_sum, one dispatch at S=65536 H=64 D=128 bf16 non-causal): before: L2 hit 2.36% (hits 2.073e+08, misses 8.593e+09) after : L2 hit 93.62% (hits 8.239e+09, misses 5.611e+08) a 15.3x reduction in L2 misses, consistent with the 80-97% L2 hit rates reported in the paper on MI300X. Measured on this branch vs its merge-base, MI355X (gfx950), H=64 D=128 bf16 non-causal batch=1, median of 30 timed iterations, each build in its own JIT cache dir, on the part at PCI bus 0x85: 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 This moves the dense non-causal path from losing to aiter's closed ASM at every size to beating it by 5-8%. Note on reproducing the absolute numbers: this workload runs into the 1400 W package cap, so sustained clocks -- and throughput -- vary 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% to +22.6% (median +20%) at 65536-131072. Compare before/after on one part; do not compare absolutes across parts. Correctness: * tests/kernels/test_flash_attn_fwd.py passes on 9 configurations before and after -- bf16/fp16/fp8 dense, causal, D=64, GQA kv=8, batch=4, and both fallback paths (num_kv_splits=2 and num_heads=20). fp8 matters because DualwaveFp8KernelContext shares this mapping function. * 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: 0 differing elements across 7 shapes (D=64/128, causal, batch=2, and both fallback cases), each also checked against an fp32 SDPA reference (worst 4.6e-3 relative to peak, the expected bf16 error, unchanged here). Guarded to non-split-K with 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 that this 2-D linearization would not preserve. The condition is compile-time from traits rather than an env var, so the selected mapping is a function of the build key and cannot alias in the JIT disk cache.
177caf5 to
2149ef8
Compare
|
Re-validated against current
Two notes for reviewers:
🤖 Generated with Claude Code |
Measured on MI355X, B=1 S=16384 H=64 D=128 bf16 causal: the swizzle costs 7.4% there (1067 -> 988 TFLOP/s), turning a win into a regression. Cause: under a causal mask, q-block i does work proportional to i+1. The stock head-fast mapping gives each consecutive run of workgroups 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 workgroup ids, so the XCDs finish at different times. Guarded to `not traits.CAUSAL`. Causal returns to the stock mapping and its stock performance (1066 vs 1067 TFLOP/s), while non-causal keeps the full gain (1328 vs 1331). Output stays bit-identical across all 7 validation shapes.
|
Pushed Measured on MI355X, B=1 S=16384 H=64 D=128 bf16, PCI bus
Cause. Under a causal mask, q-block The guard is 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 🤖 Generated with Claude Code |
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
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
What
_init_dualwave_thread_mappingtookblock_idx.x/ydirectly 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: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: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
DualwaveFp8KernelContext)% 8 != 0, must take the fallback)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: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
Pin the part by PCI bus (
torch.cuda.get_device_properties(0).pci_bus_id) — noteROCR_VISIBLE_DEVICESdoes not enumerate in the same order asrocm-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
🤖 Generated with Claude Code