Repository navigation
Use 32-row flex attention blocks for head_dim > 256 - #49444
qgallouedec wants to merge 8 commits into
Conversation
The default flex block sizes fault at head_dim 512 on Blackwell (misaligned address), which breaks attn_implementation="flex_attention" for Gemma 4, and are 1.7x slower than 32-row blocks on Hopper.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
There was a problem hiding this comment.
I wasn't sure about adding this new file. The thing is, there's no existing flex attention test file to put it in, and the per-model test_flex_attention_with_grads forces head_dim 16, so it never reaches this path.
Also the test only catches the bug on Blackwell: it fails on main on B300 and passes with this PR. Our CI runners have no Blackwell, so there it passes either way.
Happy to drop the file or move the test if you prefer.
CI recapDashboard: View test results in Grafana |
Rocketknight1
left a comment
Rocketknight1
left a comment
There was a problem hiding this comment.
I sent an agent to test this on my local hardware and it failed in fp32. Torch's own fp32 default above 256 is (32, 16, 3, 4), so fp32 inference at head_dim 512 went from working on main to Required: 131072 Hardware limit:101376 on my local 4090
fp32 support probably isn't mandatory, so maybe this is fine, but it suggests that the fix might have more hardware-dependent issues! In particular, query.device.type == "cuda" is also true for AMD cards I believe, so they definitely need testing as well before we start playing with kernel settings that can overflow buffers like this
…hich have the shared memory for them
|
Thanks, good catch. Fixed in 2aca365: the 32-row blocks now only apply on Hopper and datacenter Blackwell (compute capability 9.x and 10.x, CUDA builds), where I measured them. Compiled on H100 at head_dim 512, their forward kernel needs 196 KB of shared memory in fp32 and 98 KB in bf16, against 99 KB on a 4090, so consumer GPUs and ROCm keep torch's default blocks, as on main. |
Gemma 4's global layers use head_dim 512, where flex attention's default blocks are slow on H100 and, on torch <= 2.13, crash on B300 with
CUDA error: misaligned address(pytorch/pytorch#200348; nightly avoids it since pytorch/pytorch#200725, which prunes those blocks on sm_10x only). 32-row blocks run on both and match fp32. Re-measured on H100 with torch 2.13 today: 193.0 -> 114.8 ms, 1.7x.This passes
kernel_options={"BLOCK_M": 32, "BLOCK_N": 32, "num_stages": 2, "num_warps": 4}when the caller gives none and head_dim > 256, on Hopper and datacenter Blackwell GPUs (in fp32 these blocks need ~200 KB of shared memory). Other GPUs keep the default blocks.Global-layer fwd+bwd (H32, KV4, D512), ms:
Gemma-4-12B SFT on one B300, packed: 3.10k tok/s with the best option before (flex for the sliding layers, eager for the global ones) -> 6.18k with
flex_attentionand this PR. Per-token log-probs of a packed batch stay as close to fp32 as eager (mean |Δlogp| 0.082 vs 0.090).New test: head_dim 512 fwd+bwd against an fp32 SDPA reference. Passes here, fails on main on B300.