Clamp kv_stage to avoid SMEM overflow for small head_dims on SM100

[Johnsonms]

Summary

Fixes #2591.

flash_fwd_sm100.py:335 computes kv_stage without accounting for per-stage overhead (mbar_*, sScale, pipeline counters). For head_dim_padded=16, the heuristic can select 52+ stages, exceeding the sm_100a shared memory limit by ~1 KB and causing kernel launch failure:

cudaErrorInvalidValue: launch shared memory exceeds current GPU arch sm_100a allowed
Allocated: 233472 bytes. Max: 232448 bytes.

This affects small head dimensions (head_dim ∈ [8,16]) with non-trivial sequence lengths where q_stage=2.

Fix

Clamp kv_stage to 32.

This is sufficient to avoid the SMEM overflow and does not impact performance in practice, since attention pipelining saturates well before 32 stages. Existing head_dim_padded >= 64 configurations are unchanged because they already use ≤24 stages.

Test plan

• Added test_flash_attn_small_head_dim covering:

  • d ∈ {8,16,32}
  • causal ∈ {True, False}
  • seqlen ∈ {128, 2048}

• Verified original [FA4, CuTe, SM100] _flash_attn_fwd fails at launch for head_dim < 32 on B200 #2591 repro now passes.

• Verified on B200 (sm_100a): all new tests pass with max error vs Torch SDPA ≤ 0.0078.

• Confirmed no behavior change for head_dim_padded >= 64.

Notes

An alternative fix would incorporate per-stage overhead directly into the heuristic:

(BUDGET - smem_qo) // (smem_kv + PER_STAGE_OVERHEAD)

but that reduces stage count across all configurations. The clamp is a smaller, lower-risk fix for this bug.

_validate_head_dims remains unchanged: small head dimensions are valid on sm_100/110; only the staging heuristic was incorrect.

[Johnsonms]

Cc: @rm-wu for review

[rm-wu]

LGTM! I verified again on B200 and all head_dim in 8, 16, 32, 64, 96, 128 pass with this PR using the repro script in #2591 :

python: 3.12.3
torch:  2.12.0+cu130
device: NVIDIA B200 sm_100  (capability=(10, 0))
Calling flash_attn.cute._flash_attn_fwd with varying head_dim:
(batch=1, seqlen=2048, num_heads=2, bf16, causal=True)
head_dim=  8: OK    out.shape=(1, 2048, 2, 8)  lse.shape=(1, 2, 2048)
head_dim= 16: OK    out.shape=(1, 2048, 2, 16)  lse.shape=(1, 2, 2048)
head_dim= 32: OK    out.shape=(1, 2048, 2, 32)  lse.shape=(1, 2, 2048)
head_dim= 64: OK    out.shape=(1, 2048, 2, 64)  lse.shape=(1, 2, 2048)
head_dim= 96: OK    out.shape=(1, 2048, 2, 96)  lse.shape=(1, 2, 2048)
head_dim=128: OK    out.shape=(1, 2048, 2, 128)  lse.shape=(1, 2, 2048)

[drisspg]

do we actually want to be using 32 kv stages?

[Johnsonms]

Good catch, @drisspg
32 is essentially "the next round number above 26" — picked to be surgical to the broken case. The unbounded formula only exceeds 32 at head_dim_padded ∈ {8, 16}:

hd_padded	q=1	q=2
8 → 16	108	104
16	108	52
32	26	24
64	12	10
96	7	5
128	5	3

(1CTA path; 2CTA is gated to hd_padded ∈ {128, 192} at interface.py:572 and its max value there is 10, also below 32.)

So min(..., 32) only fires at hd_padded ∈ {8, 16}; anything from 27 to ~50 would be equally surgical. Dropping below 26 starts perturbing kernel staging for hd_padded ∈ {32, 64} (clamp=8 would give 3× fewer stages at hd=32) — and we only have perf data at hd=16, not those.

At hd=16 itself, swept clamp on B200 (batch=4, nheads=16, hd=16, bf16, causal=False):

clamp	sl=4096 ms	TFLOPS	sl=16384 ms	TFLOPS
2	0.3778	181.9	1.4705	186.9
4	0.3818	180.0	1.4858	185.0
8	0.3819	180.0	1.4860	185.0
16	0.3824	179.7	1.4874	184.8
32	0.3830	179.4	1.4883	184.7

Everything is within ~1% — clamp value doesn't measurably matter at hd=16 either, so keeping 32 just avoids changing kernel staging anywhere outside the broken case.