FA4 head_dim=256 B200 Regression Check (2026-05-16) with 2026-04-23

### TL;DR Ran an apples-to-apples FA4 `head_dim=256` benchmark on B200 between commit **27b4eb9** (PR #2412, the original hd=256 add, 2026-04-23) and current **`main` @ `0409f9a`**. On my box the two commits are within ±3% across all 48 measured points (a few small **improvements** at long-seqlen GQA, consistent with #2488 "+3-9% exp2 FMA emu"). **No regression.** However, our absolute numbers are **5–9% below** previously-reported 04.23 numbers from another B200, and other internal reports place hd=256 fwd at +10–28% higher than this box at long seqlen. I'd like to converge on what config produces those higher numbers so the community has a reproducible hd=256-on-B200 baseline. ### Environment - GPU: NVIDIA B200 (SM100), 1755 MHz locked (`nvidia-smi --lock-gpu-clocks=1755,1755`) - Driver: 580.159.03, CUDA 13.1 - torch: 2.11.0a0+eb65b36914.nv26.02 - nvidia-cutlass-dsl: 4.5.0 - flash-attn / FA4 cute: commits below - dtype: bf16, `hdim = hdim_v = 256`, MHA (nheads=8) and GQA (nheads=8, nheads_kv=1, `--gqa-ratio 8`) - num_splits=1, `softmax_scale = hdim**-0.5` - batch sizes auto from `--total-seqlen 32k` → 8/4/2/1/1/1 for 4K/8K/16K/32K/64K/128K - Timing: triton `do_bench` median, `warmup=5 rep=30` ### Commands ```bash # MHA CUDA_VISIBLE_DEVICES=0 python benchmarks/benchmark_attn.py \ --fwd --bwd --headdim 256 \ --seqlen 4096,8192,16384,32768,65536,131072 \ --num-splits 1 --backend fa4 --warmup 5 --rep 30 # GQA (ratio 8) CUDA_VISIBLE_DEVICES=1 python benchmarks/benchmark_attn.py \ --fwd --bwd --headdim 256 \ --seqlen 4096,8192,16384,32768,65536,131072 \ --gqa-ratio 8 --num-splits 1 --backend fa4 --warmup 5 --rep 30 ``` ### Results — TFLOPS (peak BF16 B200 ≈ 2250) #### MHA — Forward | seqlen | bs | causal | 04.23 (27b4eb9) | main (0409f9a) | Δ | |--------|---:|:------:|----------------:|---------------:|------:| | 4K | 8 | F | 1442 | 1442 | +0.0% | | 8K | 4 | F | 1530 | 1548 | +1.2% | | 16K | 2 | F | 1608 | 1608 | +0.0% | | 32K | 1 | F | 1502 | 1492 | −0.7% | | 64K | 1 | F | 1422 | 1435 | +0.9% | | 128K| 1 | F | 1356 | 1387 | +2.3% | | 4K | 8 | T | 1149 | 1149 | +0.0% | | 8K | 4 | T | 1340 | 1340 | +0.0% | | 16K | 2 | T | 1456 | 1457 | +0.1% | | 32K | 1 | T | 1520 | 1497 | −1.5% | | 64K | 1 | T | 1368 | 1364 | −0.3% | | 128K| 1 | T | 1215 | 1224 | +0.7% | #### MHA — Backward | seqlen | bs | causal | 04.23 | main | Δ | |--------|---:|:------:|-----:|-----:|------:| | 4K | 8 | F | 699 | 698 | −0.1% | | 8K | 4 | F | 839 | 830 | −1.1% | | 16K | 2 | F | 873 | 880 | +0.8% | | 32K | 1 | F | 915 | 917 | +0.2% | | 64K | 1 | F | 894 | 890 | −0.4% | | 128K| 1 | F | 908 | 909 | +0.1% | | 4K | 8 | T | 469 | 469 | +0.0% | | 8K | 4 | T | 646 | 646 | +0.0% | | 16K | 2 | T | 758 | 739 | −2.5% | | 32K | 1 | T | 800 | 804 | +0.5% | | 64K | 1 | T | 835 | 858 | +2.8% | | 128K| 1 | T | 847 | 839 | −0.9% | #### GQA (ratio 8) — Forward | seqlen | bs | causal | 04.23 | main | Δ | |--------|---:|:------:|-----:|-----:|------:| | 4K | 8 | F | 1486 | 1487 | +0.1% | | 8K | 4 | F | 1592 | 1553 | −2.4% | | 16K | 2 | F | 1596 | 1646 | +3.1% | | 32K | 1 | F | 1546 | 1510 | −2.3% | | 64K | 1 | F | 1474 | 1476 | +0.1% | | 128K| 1 | F | 1393 | 1450 | +4.1% | | 4K | 8 | T | 1181 | 1180 | −0.1% | | 8K | 4 | T | 1377 | 1376 | −0.1% | | 16K | 2 | T | 1482 | 1490 | +0.5% | | 32K | 1 | T | 1536 | 1548 | +0.8% | | 64K | 1 | T | 1451 | 1408 | −3.0% | | 128K| 1 | T | 1245 | 1382 | **+11.0%** | #### GQA (ratio 8) — Backward | seqlen | bs | causal | 04.23 | main | Δ | |--------|---:|:------:|-----:|-----:|------:| | 4K | 8 | F | 881 | 884 | +0.3% | | 8K | 4 | F | 893 | 901 | +0.9% | | 16K | 2 | F | 891 | 907 | +1.8% | | 32K | 1 | F | 899 | 888 | −1.2% | | 64K | 1 | F | 922 | 919 | −0.3% | | 128K| 1 | F | 925 | 925 | +0.0% | | 4K | 8 | T | 691 | 690 | −0.1% | | 8K | 4 | T | 861 | 855 | −0.7% | | 16K | 2 | T | 844 | 891 | **+5.6%** | | 32K | 1 | T | 851 | 848 | −0.4% | | 64K | 1 | T | 888 | 899 | +1.2% | | 128K| 1 | T | 867 | 865 | −0.2% | ### Observations **No regression** between 27b4eb9 and current main. Two cells show meaningful **improvement** (GQA fwd causal 128K: +11%, GQA bwd causal 16K: +5.6%), the +11% being roughly consistent with #2488's claimed "+3-9% from exp2 FMA emulation".