Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects #

§1 TL;DR #

Swizzled Head-first Mapping — a ~15-line Triton WG-ID remapping that confines all workgroups of one attention head (the "Attention Compute Cluster") to a single XCD on AMD MI300X, achieving up to 50% forward-pass speedup and sustaining 80–97% L2 cache hit rates versus ~1% for NUMA-unaware baselines.

---

§2 Q1 / Q2 / Q3 #

Q1 痛点 #

Modern AI GPUs have evolved from monolithic dies with unified L2 cache to multi-chiplet architectures (AMD MI300X: 8 XCDs, each with private 4 MB L2). The default workgroup scheduler uses chunked round-robin (chunk=1), spreading workgroups that share the same K/V tensors across different XCDs. This fragments L2 cache: each XCD loads its own copy of shared data from HBM, collapsing aggregate L2 hit rate to as low as ~1% for MHA with many heads and long sequences.

Paper's Figure 1, verbatim (caption: "Evolution of GPU architectures toward disaggregated memory hierarchies").

The progression from (a) unified L2 to (b) dual-die to (c) quad-die chiplet designs makes NUMA effects increasingly severe. AMD MI300X exposes these effects to software; NVIDIA Blackwell hides them via hardware coherency.

Q2 方法 #

Core insight: In FA2, all workgroups within one attention head read the same K and V tensors. Define an Attention Compute Cluster (ACC) as the set of WGs sharing K/V data (= one head for MHA, one KV-group for GQA). The algorithm remaps linear WG IDs so that an entire ACC is assigned to a single XCD before advancing to the next head.

The mapping is:

head_offset = ((wid_per_batch % NUM_XCD) * heads_per_xcd
               + wid_per_batch // (NUM_XCD * blocks_per_head))
block_offset = (wid_per_batch % chunk_size) // NUM_XCD

This ensures: (1) each XCD services one ACC at a time; (2) all blocks within that ACC share a single L2 for K/V; (3) zero cross-XCD redundant HBM fetches for shared tensors.

核心技术壁垒: The non-obvious insight is that head-first iteration alone (Naive Head-first) is insufficient — without spatial swizzling, round-robin dispatch still stripes blocks of one head across all 8 XCDs. The swizzle must simultaneously control both iteration order (head-first) and physical placement (all blocks → one XCD). Getting this wrong in either dimension collapses the benefit.

Q3 结果 #

Metric	Swizzled Head-first	Best baseline
MHA forward speedup (H=128, 128K)	1.0× (reference)	0.64–0.70× (Block-first)
L2 hit rate (MHA, 128K)	90–97%	~1% (Naive Block-first)
GQA forward (Llama-3 family)	1.0×	0.90–0.95× (Naive Head-first)
DeepSeek-V3 prefill (H=128, D=56)	1.0×	<0.65× (Block-first at 128K)
Backward pass (128K)	1.0×	0.91× (Naive approaches)

---

§3 架構 / 方法图 #

Paper's Figure 10, verbatim (caption: "Swizzled Head-first mapping (eight qheads, 128 row blocks, four XCDs): Head-first iteration with spatial swizzling confines each attention head to a single XCD").

Each color = one attention head. All blocks of a given head are packed into a single XCD column. Contrast with Naive Block-first where colors (heads) are interleaved across all XCDs, or Naive Head-first where blocks of the same head are striped across XCDs via round-robin dispatch.

Paper's Figure 4, verbatim (caption: "FlashAttention2 Tiled Compute Partitioning Across Workgroups for a Single Attention Head").

This illustrates why co-location matters: pid0–pid3 all read the entire K^T and V matrices. If they land on different XCDs, each loads K/V independently from HBM; if co-located, they share a single L2-cached copy.

Paper's Figure 2, verbatim (caption: "Impact of workgroup scheduling on cache reuse in the multi-die chiplet architecture of AMD MI300X").

When WGs sharing data land on different XCDs (Die 0 vs Die 1), their private L2 caches cannot share entries, and both must independently fetch from HBM through LLC — doubling bandwidth consumption for the same data.

Method as before/after diff #

Aspect	Before (Naive / Swizzled Block-first)	After (Swizzled Head-first)
Iteration order	Block-first (all heads per block) or Head-first (but unswizzled)	Head-first
XCD assignment	Round-robin chunk=1 by linear WG ID	Swizzled: all WGs of one head → one XCD
ACCs per XCD simultaneously	Multiple (fragments L2)	One at a time
K/V copies in aggregate L2	Replicated across 8 XCDs	Single copy per XCD, fully reused
L2 hit rate (MHA, H=128, 128K)	~1% (Block-first) / 40–60% (Naive Head-first)	90–97%

---

§4 作者证明 #

无形式化作者证明 — 仅实证

The paper provides no convergence theorem, complexity bound, or formal proof. The argument is entirely empirical: measuring L2 hit rates and kernel throughput across a parameter sweep. The "proof" is architectural reasoning + hardware counter measurements.

Notation table #

Symbol	Meaning
$Q, K, V \in \mathbb{R}^{N \times d}$	Query, key, value matrices per head
$N$ (N_CTX)	Sequence length
$d$ (HEAD_DIM)	Per-head dimension
$H_Q, H_K$	Number of query / KV heads
ACC	Attention Compute Cluster: WGs sharing K/V
XCD	Accelerator Complex Die (chiplet on MI300X)
BLOCK_M	Row-block tile size (128)
WG / pid	Workgroup / program ID

Equation physical meaning #

$$S = QK^{\top}, \quad P = \text{softmax}\!\left(\frac{S}{\sqrt{d}}\right), \quad O = PV$$

Standard scaled dot-product attention. The relevant property for this work: within one head, all $Q$-row-blocks share the same $K$ and $V$ — establishing the data-sharing pattern that defines an ACC.

6 minimum checks #

1. Claim: up to 50% speedup over Block-first. Check: Figure 12 shows Naive Block-first at 0.64–0.70× for H=128, N=128K; 1/0.67 ≈ 1.49× confirming ~50%.
2. Claim: L2 hit rate 80–97%. Check: Figure 13 directly plots hit rates; Swizzled Head-first sustains >80% across all tested configs and reaches 97% at favorable points.
3. Claim: Naive Block-first collapses to ~1% L2 hit rate. Check: Figure 13 shows Block-first at <5% for H=128, N≥32K. Consistent with reasoning — 128 heads across 8 XCDs means 16 concurrent ACCs per XCD, each evicting the others' K/V from 4 MB L2.
4. Claim: works for both MHA and GQA. Check: §4.4 Figure 14 shows GQA results; Swizzled Head-first matches or exceeds all baselines. For GQA with 8 KV groups = 8 XCDs, Swizzled Block-first ties (expected — it's the degenerate case where groups perfectly tile XCDs).
5. Claim: backward pass gains ~10%. Check: Figure 16 shows other approaches at 0.91–0.94× of Swizzled Head-first. The paper attributes the ceiling to scalar operation bottlenecks but doesn't precisely isolate them.
6. Claim: minimal code change (~15 lines). Check: Figure 11 shows the complete swizzle logic in 10 lines of Triton. The claim holds — no algorithmic change to FA2, only WG ID remapping.

What formal guarantee would be desirable #

An analytical model predicting L2 hit rate as a function of $(H_Q, N, d, \text{L2\_size}, \text{NUM\_XCD})$ would let practitioners predict when Swizzled Head-first is necessary versus when NUMA effects are negligible (e.g., few heads, short sequences). The paper leaves this to empirical observation.

---

§5 实验与数据 #

Hardware platform #

AMD MI300X: 8 XCDs × 38 CUs = 304 CUs, 4 MB L2/XCD (32 MB total), 192 GB HBM3 @ 5.3 TB/s. Implementation in Triton; profiling via ROCProfiler v3 hardware counters.

Parameter sweep #

Parameter	Values
Batch size	1, 2, 4, 8
Sequence length	2K, 4K, 8K, 16K, 32K, 64K, 128K
$H_Q$	8, 16, 32, 64, 128
HEAD_DIM	128
BLOCK_M	128

Key result: MHA forward performance #



Paper's Figure 12, verbatim (caption: "MHA Performance relative to Swizzled Head-first baseline across varying batch sizes (1-8) and sequence lengths (8K-128K)").

The performance gap widens monotonically with $H_Q$ and $N_{\text{CTX}}$. At $H_Q = 8$ the advantage is modest (<10%) because even naive scheduling only fragments 8 ACCs across 8 XCDs. At $H_Q = 128$, each XCD juggles 16 concurrent ACCs under naive dispatch, causing catastrophic L2 thrashing.

Key result: L2 cache hit rate #



Paper's Figure 13, verbatim (caption: "L2 Cache hit rates for MHA across varying batch sizes (1-8) and sequence lengths (2K-128K)").

This is the most diagnostic figure in the paper. It directly visualizes the mechanism: Swizzled Head-first keeps K/V in L2 by serving one ACC at a time, while Block-first forces each XCD to context-switch between ACCs faster than L2 can retain their working sets.

GQA results (Llama-3 family) #

For GQA with 8 KV groups = 8 XCDs, both Swizzled Block-first and Swizzled Head-first achieve near-identical performance. Naive Block-first still degrades significantly (to 0.75× at H_Q=128, 128K). Naive Head-first shows instability at 0.90–0.95× for high N/batch.

DeepSeek-V3 case study #

MHA with $H_Q = 128$, $H_K = 128$, $d = 56$. This is the worst case for NUMA-unaware scheduling: 128 independent heads (no KV sharing). Block-first drops below 0.65× at 128K. Swizzled Head-first is strictly optimal.

Backward pass #

Gains capped at ~10% (other approaches at 0.91–0.94×). The paper hypothesizes scalar operation bottlenecks in the backward kernel limit the cache-utilization benefit. This is the one regime where the optimization underdelivers relative to the forward pass.

Training recipe & scale #

Not applicable — this is a kernel scheduling algorithm, not a training recipe. No training is involved.

---

§6 论证链 #

Step	Claim	Evidence	Logic
1	Modern chiplet GPUs have private-per-XCD L2 caches creating NUMA effects	MI300X architecture spec: 8 XCDs × 4 MB private L2	Architectural fact
2	FA2 workgroups within one head all read the same K/V (forming an ACC)	FA2 algorithm definition (Fig 4): each WG reads full K^T and V	Algorithmic property
3	Default round-robin scheduling (chunk=1) distributes one ACC's WGs across all XCDs	Fig 2 + driver behavior documentation	Observed hardware behavior
4	Cross-XCD ACC distribution forces each XCD to independently load K/V → L2 thrashing	8 copies of K/V in 8 L2s vs. 1 copy in 1 L2; L2 capacity saturated by concurrent ACCs	Memory hierarchy reasoning
5	Swizzled Head-first remaps WG IDs so all blocks of one head → one XCD	Fig 10/11: head_offset calculation assigns consecutive heads to consecutive XCDs; blocks within head stay on same XCD	Code logic
6	Co-located ACC enables K/V L2 reuse: one load amortized across all WGs in that head	Measured L2 hit 90–97% (Fig 13) vs ~1% baseline	Hardware counter measurement
7	Higher L2 hit rate → fewer HBM accesses → higher effective bandwidth → up to 50% speedup	Fig 12: performance tracks L2 hit rate improvements	Bandwidth-bound kernel argument

---

§7 实现 cross-reference #

AMD AITER repository #

The Swizzled Block-first variant (not Swizzled Head-first) is deployed in AMD's AITER repository (AMD, 2025a) for the FA2 backward pass kernel. The paper's proposed Swizzled Head-first mapping is described as a novel extension not yet merged into a public production repository at time of writing.

Triton implementation #

The complete swizzle logic is shown in Figure 11 (§3.3). Key code:

wid = tl.program_id(0)
wid_per_batch = wid // BATCH
heads_per_xcd = NUM_Q_HEADS // NUM_XCD
blocks_per_head = (SEQLEN_Q + BLOCK_M - 1) // BLOCK_M
chunk_size = NUM_XCD * blocks_per_head

head_offset = ((wid_per_batch % NUM_XCD) * heads_per_xcd
               + wid_per_batch // (NUM_XCD * blocks_per_head))
block_offset = (wid_per_batch % chunk_size) // NUM_XCD
batch_offset = (wid // (blocks_per_head * NUM_Q_HEADS)) % BATCH

The existing chiplet-aware swizzle baseline (Figure 3, §2.2) in Triton:

@triton.jit()
def swizzle_chiplet(wgid, grid, NUM_XCD: tl_constexpr):
    wgids_per_xcd = grid // NUM_XCD
    xcd = wgid % NUM_XCD
    local_wgid = wgid // NUM_XCD
    new_wgid = xcd * wgids_per_xcd + local_wgid
    return new_wgid

核心技术壁垒 (detailed) #

The technical barrier is not the code itself (~15 lines) but the architectural insight required to arrive at it: understanding that (a) ACC is the correct unit of locality for attention (not tile, not batch), (b) head-first iteration must be coupled with spatial swizzling (either alone is insufficient), and (c) the driver's round-robin policy is chunk=1 and mutable, so the kernel must own its placement. Without profiling infrastructure (ROCProfiler L2 counters) and architectural knowledge of XCD-private L2, one would never identify the ~1% hit rate failure mode or know that the fix is so simple.

关键実装细节 #

1. NUM_Q_HEADS must be divisible by NUM_XCD for the integer division heads_per_xcd = NUM_Q_HEADS // NUM_XCD to distribute evenly. When this doesn't hold (e.g., H_Q=12, NUM_XCD=8), the paper doesn't discuss the fallback — likely requires padding or partial-XCD assignment.
1. Backward pass scalar bottleneck: The backward kernel's lower gains (~10% vs 50%) suggest that once L2 hit rate is high, the kernel becomes compute-bound on scalar operations (online softmax correction, gradient accumulation). This means the optimization has diminishing returns for compute-dense kernel variants.

Reproducibility #

• Training code: N/A (no model training)
• Kernel code: Not publicly released as standalone; the swizzle logic is fully specified in Figures 3 and 11
• Closest open reference: Triton's default FA kernel (triton-lang/triton) + AITER (ROCm/aiter) for the Block-first variant
• Community reproduction: No known independent replication at time of writing
• Production usage: AITER deploys Swizzled Block-first for backward pass; Swizzled Head-first proposed for forward pass but deployment status unclear