Chunk the batch to keep FlashAttention's backward in its index range

[priorphil]

Generated by Claude Code.

What

A FlashAttention backward fails with CUDA error: an illegal memory access was encountered once

batch * heads * round_up(seq_q, 128) * max(head_dim, 32) > 2**31

The forward is fine at the same shapes, and the memory-efficient backend is fine at all of them.

_torch_sdpa already chunks the batch to keep batch * heads under the CUDA grid limit. This adds a second bound to that same loop, on the element count above.

Both roundings matter, and each one on its own makes the naive reading wrong:

• a query length that is already a multiple of 128 reaches exactly 2**31 and passes, so the largest shape in the table below is a passing one,
• a head dim of 16 fails at half the raw element count, because it is counted as 32.

Reproducing

Needs a CUDA device with ~24 GiB free.

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

# 400 * 4 * 10496 * 128 = 2_149_580_800 > 2**31
q = torch.randn(400, 4, 10_496, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(400, 4, 128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
v = torch.randn_like(k).requires_grad_(True)

with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, k, v)  # forward is fine
out.sum().backward()                               # illegal memory access

Swapping FLASH_ATTENTION for EFFICIENT_ATTENTION makes it pass.

The small-head-dim case is worth reproducing too, since it is the one a check on the raw element count would miss. (800, 8, 10496, 16) is 1,074,790,400 elements, half the bound, and fails.

Where the boundary is

Measured on an RTX PRO 6000, torch 2.9.0+cu128, bfloat16. One process per shape, since a fault leaves the CUDA context unusable. "counted" is batch * heads * round_up(seq_q, 128) * max(head_dim, 32).

batch	heads	seq_q	head_dim	counted	/ 2**31	backward
8	4	65536	128	268,435,456	0.125	ok
200	8	10496	16	537,395,200	0.250	ok
400	8	10496	16	1,074,790,400	0.500	ok
400	8	10496	32	1,074,790,400	0.500	ok
400	4	10368	128	2,123,366,400	0.989	ok
799	8	10496	32	2,146,893,824	0.99972	ok
512	4	8192	128	2,147,483,648	1.000	ok
800	8	10496	16	2,149,580,800	1.001	fails
800	8	10496	32	2,149,580,800	1.001	fails
400	8	10496	64	2,149,580,800	1.001	fails
400	4	10496	128	2,149,580,800	1.001	fails
800	4	5248	128	2,149,580,800	1.001	fails
200	4	20992	128	2,149,580,800	1.001	fails
1600	8	10496	16	4,299,161,600	2.002	fails
3200	8	10496	16	8,598,323,200	4.004	fails

The 799 and 512 rows sit either side of the boundary within 0.03%, and the 400 / 16 versus 800 / 16 pair fixes the head-dim floor at 32 rather than 64.

Caveats

• The constant is a measured boundary. It is consistent with a signed 32-bit element index over a workspace padded to the kernel's block sizes, but I have not confirmed that in the FlashAttention source, so if a maintainer knows the actual bound that would be a better value than mine.
• One GPU, one torch version, bf16, head dims 16/32/64/128.
• If a single batch entry is on its own outside the bound, chunking the batch cannot help and the guard leaves the call as it is. That needs heads * round_up(seq_q, 128) * max(head_dim, 32) > 2**31, which no shape here reaches.

Tests

tests/test_architectures/test_shared/test_scaled_dot_product_attention.py:

• four tests that drive _torch_sdpa on meta tensors and assert the chunking: over the bound it splits and every call lands inside it, the head-dim-16 case over the bound also splits (its raw count looks safe, so this is the one that catches a check which ignores the padding), inside the bound it stays a single call, and a non-FlashAttention backend is not split. These need no GPU and no memory.
• one slow test that runs a failing shape end to end, skipped unless a CUDA device has ~24 GiB free.

pytest tests/test_architectures/test_shared tests/test_architectures/test_attention_backends.py tests/test_architectures/test_chunked_evaluate.py passes (30 passed, 9 skipped).

Verified on the GPU: the slow test fails unpatched and passes patched. A larger model that previously failed at these shapes now completes, both at the shape where it first failed and at one 1.34x beyond it.