The Power-of-2 Wall
Why GPU block-I/O hardware insists on power-of-two extents — and how a single awkward dimension can lock an entire kernel out of the fast path
Two Ways to Move a Tile
Every tiled GPU kernel has to get a 2D block of data from memory into registers. On modern hardware there are two very different ways to do it, and the gap between them is large.
Block I/O (fast)
A single tensor descriptor drives a hardware 2D-block transfer — one wide, coalesced read of the whole tile. On Intel Xe this is the 2D-block load; CUDA calls the analogous idea TMA.
Pointer gather (slow)
Compute a pointer per element (or per row), then load with explicit masks. Correct for any shape, but it leaves the block-transfer hardware idle and is often ~2× slower on memory-bound tiles.
The fast path is the one you want — but it comes with a hardware constraint. This course is about satisfying that constraint for a shape that naively fails it, instead of falling back to the slow path.
The Constraint: Power-of-Two Inner Extent
The catch with block I/O is that the descriptor's inner (contiguous) block extent must be a power of two. This isn't a software style rule — it falls out of how the descriptor is encoded and what the hardware's validator will accept.
A Concrete Place This Bites: Attention
A clean, real example is the query (Q) load inside a fused-attention kernel. The kernel tiles a block of query rows and loads them through a tensor descriptor whose inner extent is (query heads per group) × (head dimension).
When it's a power of two
If both the head count per group and the head dimension are powers of two, the inner extent is too — the descriptor validates and the fast block load fires.
When it isn't
A single non-power-of-two factor — say 7 query heads per group instead of 4 or 8 — makes the inner extent non-pow2 and knocks the whole Q load onto the slow pointer path.
We use that attention Q load as the running example, but the technique is general: pad the awkward axis up to a power of two, read the wider block, and mask the slack. By the end you'll be able to apply it to any block-I/O constraint, not just this one.
Check Your Understanding
What must be a power of two for a block-I/O descriptor to validate?
What is this course's strategy for a shape that fails the power-of-two constraint?
GQA and the Power-of-2 Gate
How grouped-query attention packs query heads into a tile — and why a non-power-of-two group size gets locked out of the fast path
One KV Head, Many Query Heads
In grouped-query attention (GQA), several query heads share one key/value head. The ratio that matters is num_queries_per_kv = num_query_heads / num_kv_heads — call it nq. It is the number of query rows that ride along with each KV head.
Power-of-two group size
e.g. 32 query heads over 8 KV heads → nq = 4. A power of two. This is the common, well-trodden case the fast path was built for.
Non-power-of-two group size
e.g. 28 query heads over 4 KV heads → nq = 7. Not a power of two. Some real models land here, and this single fact is what trips the gate.
How the Kernel Tiles It
The unified-attention kernel processes a tile of BLOCK_M rows at a time. Those rows are a 2D pack: BLOCK_Q query positions × nq query heads per KV group. Each row is one (position, head) pair sharing the tile's KV.
# num_queries_per_kv heads per KV group
BLOCK_M = (16 if nq <= 16
else next_pow2(nq))
BLOCK_Q = BLOCK_M // nq
# row r -> which position, which head
query_pos = base + r // nq
head_in_group = r % nqA BLOCK_M-row tile is laid out as BLOCK_Q positions, each expanded into nq head rows.
For pow2 nq, BLOCK_M / nq divides cleanly and the packing is dense — every row is a real head.
The kernel recovers the position with r // nq and the head with r % nq to index into Q.
This is fine arithmetically for any nq — the trouble starts only when this tile must be loaded via a tensor descriptor.
Two Ways to Load Q
Once the rows are addressed, the kernel must actually fetch Q from memory. There are two code paths, and the choice is gated.
TD path (fast)
A tl.make_tensor_descriptor drives the Xe 2D-block load hardware — one descriptor, one wide hardware-accelerated read of the whole tile.
Pointer path (slow)
A gather over per-row pointers with explicit masks. Correct for any shape, but it leaves the 2D-block I/O hardware idle — the source of the ~2× penalty.
# Q/O take the fast TD path only if
# BOTH nq and head_size are pow2
_is_pow2_nq = (nq & (nq - 1)) == 0
_is_pow2_hs = head_size == \
next_pow2(head_size)
use_td_qo = use_td and _is_pow2_nq \
and _is_pow2_hs_is_pow2_nq is the bit-trick test for "power of two": true for 1, 2, 4, 8, 16…
For nq = 7 it is false, so use_td_qo is false.
Result: a non-pow2 group size silently falls back to the slow pointer path for both the Q load and the O store — even though the hardware is perfectly capable of the load.
The technique in this course removes _is_pow2_nq from this gate. But first: why was it ever there?
Check Your Understanding
For a model with 28 query heads and 4 KV heads, what is num_queries_per_kv?
With the old gate, which path did a non-power-of-two group size take for its Q load?
Why the TD Path Demands a Power of 2
The tensor-descriptor validator has a hard rule about the inner extent — and a flattened nq × HEAD_SIZE axis violates it for nq = 7
The Descriptor Sees a Flattened Axis
The TD Q load doesn't describe a 3D (positions, heads, head_dim) tensor. It flattens the head and head-dim axes into one contiguous inner extent and describes a 2D block: (BLOCK_Q positions, nq × HEAD_SIZE columns).
q_desc = tl.make_tensor_descriptor(
base=q_base,
shape=(q_block_local_len,
nq * HEAD_SIZE),
strides=(query_stride_0, 1),
block_shape=(BLOCK_Q,
nq * HEAD_SIZE_PADDED), # <-- inner
)The inner block_shape dimension is nq × HEAD_SIZE_PADDED — all the heads of a KV group laid end to end.
HEAD_SIZE_PADDED is already next_pow2(head_size), so head-dim padding was a solved problem.
But the head count nq still multiplies in raw. For nq = 4, 4 × 128 = 512 — a power of two. For nq = 7, 7 × 128 = 896 — not a power of two.
The Validator's Rule
The Xe 2D-block I/O lowering requires the flattened inner extent of a tensor descriptor's block to be a power of two. This isn't a Triton style preference — it's what the hardware descriptor encoding and its validator accept.
The validator cares about the extent the descriptor reads, not how many heads are real. If we read a power-of-two number of head-slots and simply ignore the extras, the descriptor is valid — and the load is correct. That is the whole idea behind the fix.
There's Already a Precedent
The kernel already pads one axis for exactly this reason. HEAD_SIZE_PADDED = next_pow2(head_size) rounds an awkward head dimension (say 96 or 80) up to the next power of two, and the descriptor zero-pads the read past the real HEAD_SIZE. A dim_mask then drops the padding columns.
The inner HEAD_SIZE axis is rounded to HEAD_SIZE_PADDED; the descriptor reads zeros past the real extent and dim_mask masks them out. Proven, shipping behavior.
If padding the head dim to a power of two is safe and masked, then padding the head count nq to next_pow2(nq) — and masking the extra head-slots — is the same trick on a second axis.
Check Your Understanding
What must be a power of two for the TD load to validate?
What existing mechanism does the fix take as its precedent?
The Fix: Pad the Head Count, Mask the Slack
Pack BLOCK_M = BLOCK_Q × next_pow2(nq), read NQ_PADDED head-slots through the descriptor, then mask the padding rows in-kernel
Three Moves
Conceptually the change is three coordinated moves that turn the non-pow2 head count into a padded, masked, power-of-two read — small in code, but each move has to line up with the others.
BLOCK_M around NQ_PADDEDDefine NQ_PADDED = next_pow2(nq). Recompute BLOCK_Q = max(BLOCK_M // NQ_PADDED, 1) and BLOCK_M = BLOCK_Q × NQ_PADDED, so each query position now owns NQ_PADDED head-slots, not nq.
NQ_PADDED × HEAD_SIZE_PADDEDThe TD block's inner extent becomes a power of two. The descriptor zero-pads the read past the real nq × HEAD_SIZE — exactly like the head-dim padding already does.
Rows where head_in_group >= nq are pure padding. Fold that into query_mask_1 so they never contribute to softmax and never alias the next KV group.
The New Row Arithmetic
Because each group now spans NQ_PADDED slots instead of nq, every index that used nq to split position from head switches to NQ_PADDED.
# BLOCK_M = BLOCK_Q * NQ_PADDED
NQ_PADDED: tl.constexpr = \
BLOCK_M // BLOCK_Q
head_in_group = offs_m % NQ_PADDED
query_pos = base + offs_m // NQ_PADDED
query_offset_1 = \
kv_head_idx * nq + head_in_groupoffs_m // nq → offs_m // NQ_PADDED for the query position; offs_m % nq → offs_m % NQ_PADDED for the head slot.
For pow2 nq, NQ_PADDED == nq, so every line is bit-for-bit identical to before. A genuine no-op for any pow2 group size.
For nq = 7, NQ_PADDED = 8: each group has 7 real head rows + 1 padding row that must be masked off so it doesn't bleed into the next KV group's output.
The Mask That Makes It Safe
The padding slots carry garbage-or-zero values. Two conditions must both hold for a row to count — the original out-of-range head check and the new padding check.
query_mask_1 = tl.where(
(query_offset_1 < num_query_heads)
& (head_in_group < nq), # NEW
1, 0,
).to(tl.int1)query_offset_1 < num_query_heads — the original guard against running off the end of the head axis.
head_in_group < nq — the new guard. It kills the NQ_PADDED − nq padding slots per group (slot 7 when nq=7).
Without the second clause, slot 7's stale data would alias into kv_head_idx + 1's rows and silently corrupt the output. The mask is the correctness linchpin.
Watch One Group Flow Through
Step through how a single non-pow2 KV group (nq = 7) is loaded, padded, and masked on the new path.
Check Your Understanding
For nq = 7, what is NQ_PADDED?
What would happen if the padding head-slots were not masked?
The Load/Store Asymmetry
Why the TD load can pad but the TD store cannot — and the loop-bound bug the repacking would have introduced
A Padded Read Is Safe. A Padded Write Is Not.
The fix enables the TD path for the Q load on non-pow2 GQA — but deliberately keeps the O store on the pow2-only path. The asymmetry is fundamental, not an oversight.
TD load — can pad
Reading extra padding columns is harmless: the descriptor returns zeros, and the mask discards them downstream. Nothing in memory is touched. Non-pow2 GQA can use it.
TD store — cannot pad
A padded write would emit NQ_PADDED heads' worth of columns — and the padding head would overwrite the adjacent KV group's output. There is no mask on a wide block store. Pow2 only.
The TD store writes the real nq heads contiguously. If it wrote NQ_PADDED heads, slot 7's bytes would land on top of the next group's head 0. A masked read drops slack; a wide write clobbers neighbors. So the store stays pow2-only and non-pow2 falls back to the masked pointer store.
Splitting the Gate in Two
The old code had a single use_td_qo flag controlling both Q load and O store. The fix splits it: a relaxed load gate plus a separate, stricter store gate.
# host: load gate drops _is_pow2_nq
use_td_qo = use_td and _is_pow2_hs
# kernel: store stays pow2-nq only
USE_TD_STORE: tl.constexpr = \
USE_TD_QO and \
(BLOCK_M // BLOCK_Q == nq)use_td_qo loses the _is_pow2_nq term — the TD Q load now fires whenever head size is pow2, regardless of head count.
BLOCK_M // BLOCK_Q == nq is true iff NQ_PADDED == nq, i.e. nq was already a power of two. That is the store's enabling condition.
Pow2 GQA: both gates fire — fast load and fast store, exactly as before. Non-pow2 GQA: fast load, masked pointer store. Strictly better than the old all-or-nothing.
The Loop-Bound Bug Hiding in the Repack
Repacking BLOCK_M quietly broke a helper that computed how far each Q-block reaches. The old expression assumed BLOCK_M = BLOCK_Q × nq; after padding, that overshoots.
# largest in-block query offset
# BEFORE:
- (BLOCK_M - 1) // nq
# AFTER:
+ (BLOCK_Q - 1)The largest in-block query-position offset is just BLOCK_Q − 1. The old form derived it as (BLOCK_M − 1) // nq.
With padding, BLOCK_M = BLOCK_Q × NQ_PADDED > BLOCK_Q × nq, so (BLOCK_M − 1) // nq computes a position offset that is too large — over-counting the sequence prefix and the query-row span.
Computing it from BLOCK_Q directly is correct for both cases: for pow2 GQA the two expressions are identical, so again a clean no-op.
Any code that derived a query-position from BLOCK_M via // nq implicitly assumed dense packing. Once BLOCK_M carries padding slots, those derivations must use BLOCK_Q (positions) or NQ_PADDED (slots) — never the raw nq.
Check Your Understanding
Why does the TD store stay restricted to power-of-2 GQA?
What is the correct largest in-block query-position offset after repacking?
A Companion Fix: Widen the Prefill Tile
The same kernel had a second problem — a Q tile too small to feed the matrix engine — and the cure runs straight into the same power-of-2 rule
The Matrix Engine Is Starving
Getting Q onto the fast load path fixes movement. But prefill is dominated by the two matmuls inside attention — Q@K and P@V — and those dots can still be too small to keep the hardware busy.
Default tile, starved engine
The common GQA≤16 case defaults to BLOCK_M = 16. That makes each Q@K / P@V dot tiny — measured at a flat ~5.7 TFLOPS, against a vendor kernel's 50+.
Wider tile, fed engine
Doubling the Q tile to BLOCK_M = 32 for prefill gives the matrix engine enough work per instruction to approach its throughput — roughly 2× on prefill.
Module 4 was about the Q load — getting bytes in via the fast path. This is about the matmul tile — giving the compute units enough work. They are orthogonal: this fix stands alone and does not depend on the non-pow2 Q-load change.
The Change: Bump BLOCK_M to 32 for Prefill
The fix is a single guarded reassignment of the tile size, placed right after the default BLOCK_M / BLOCK_Q are computed.
_is_pow2_nq = \
(nq & (nq - 1)) == 0
if max_seqlen_q > 1 \
and use_td \
and _is_pow2_nq \
and BLOCK_M < 32:
BLOCK_Q = max(32 // nq, 1)
BLOCK_M = BLOCK_Q * nqmax_seqlen_q > 1 — prefill only. Decode (one query token) is latency/occupancy bound, and a wider tile can regress it.
use_td — only on the tensor-descriptor path, where the wide tile turns into an efficient block load.
_is_pow2_nq — the now-familiar gate. See the next screen for why it's load-bearing here.
BLOCK_M < 32 — only widen tiles that are actually too small; never shrink one that's already ≥ 32.
Why Pow-2 GQA — Again
This fix gates on _is_pow2_nq for the very same reason Module 2's gate existed. Widening BLOCK_M for a non-pow2 group size would produce an invalid tile.
Not every power-of-two wall is worth padding around. For the Q load, padding unlocked a 2× for an otherwise-stuck shape. For tile widening, the pow2 shapes already get the win and the non-pow2 path stays correct on the default — so the simplest safe move is to gate, not pad.
What It Buys
Like the Q-load fix, this was validated against a reference SDPA across the full prefill/decode shape matrix — and is a strict no-op everywhere it shouldn't fire.
For a pow2 group size at a multi-thousand-token prefill, the wider tile feeds the matrix engine enough work to roughly halve kernel time.
Decode (max_seqlen_q == 1) and non-pow2 GQA fail a guard and keep the default tile. Verified byte-identical to baseline.
It's an honest partial win: prefill throughput roughly doubles, but a residual structural gap to the hand-tuned vendor kernel remains a separate item.
Check Your Understanding
Why does widening BLOCK_M to 32 help prefill?
Why is the tile-widening gated to power-of-2 GQA?
Validation & The General Pattern
How to prove a pad-and-mask change is correct, the ~2× it recovers, and why the same trick generalizes far beyond attention
The Payoff: From Slow Path to Fast Path
The whole point of the exercise is to move a non-pow2 shape off the slow pointer gather and onto the hardware block load. On a memory-bound prefill tile, that swap is worth roughly 2×.
On a prefill-sized Q tile the load is memory-bound. The pointer path issues many narrow, partially-masked accesses; the descriptor path issues one wide, coalesced block transfer. Removing that inefficiency is the entire win — no math changed.
Correctness: How to Trust a Mask
A padding-and-masking change is exactly the kind of fix that produces subtly wrong numbers if the mask is off by one. The discipline that earns trust is to diff against a reference across the full matrix of shapes.
2D prefill, 3D decode, long-context, multi-sequence, mixed batches, sliding-window, and scattered block tables. Every case must match a trusted reference — the padding axis interacts with all of them.
Test the non-pow2 shape (the one you changed for) and a pow2 control side by side. Identical-to-baseline output on both confirms the mask drops exactly the padding slots and nothing else.
Because NQ_PADDED == nq for power-of-two shapes, every rewritten line collapses to its original form. Verifying byte-identical output on pow2 shapes proves the change is a true no-op there.
The Pattern Generalizes
Nothing about pad-and-mask is specific to attention. Any time a hardware fast path imposes a shape constraint you don't naturally satisfy, the same three-step recipe applies.
Round the offending dimension up to whatever the hardware accepts — a power of two, a multiple of the tile width, an alignment boundary.
A padded read is safe — the extra lanes come back as zeros. A padded write is only safe if it can't clobber a neighbor; otherwise keep the slow path for the write.
Make sure the padding lanes never influence a result and never alias real data. The mask is the correctness linchpin — get it right and the speedup is free.
float4 load needs a multiple-of-4, aligned extent. Pad the row, read vectorized, mask the tail — identical reasoning, smaller scale.
Final Check
Where does the ~2× speedup come from?
What best confirms the mask is correct?
A power-of-two hardware constraint doesn't have to gate off a whole shape. Pad the awkward axis, read the wider block, mask the slack — the same trick kernels already use on the head dimension, applied one axis over. Reach for it whenever a fast path's shape rule is the only thing standing between you and the hardware.