[FA4][CuTe DSL] Add head_dim=256 support with exp2 FMA emulation optimization (forward + backward, SM100)

[Johnsonms]

Problem

The SM100 hdim=256 forward kernel uses exp2 (hardware SFU) for softmax in the inner loop. On Blackwell, the SFU pipeline can be a throughput bottleneck relative to the TMEM/MMA pipelines.

This PR depends on https://github.com/Dao-AILab/flash-attention/pull/2412/commits, will start to review aftehr 2412 merged

Solution

Replace a fraction of exp2 SFU calls with a degree-3 FMA polynomial emulation (ex2_emulation_2) to shift work from the XU/SFU pipeline onto the FMA pipeline, improving overall instruction-level parallelism.

Implementation details:

• utils.py provides ex2_emulation_2(x, y) — a packed-pair FP32 degree-3 polynomial using evaluate_polynomial_2 +
  combine_int_frac_ex2 (bit manipulation for integer+fractional recombination)
• In sm100_hd256_2cta_fmha_forward.py, the softmax inner loop processes S score fragments. Each fragment pair is
  emulated based on: k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res
• The last fragment of each tile always uses hardware SFU for numerical stability
• Tuned parameters (freq=4, res=3): because k steps by 2, k % 4 ∈ {0, 2}, so the condition fires 50% of the time →
  50% SFU / 50% FMA emulation

NCU analysis (seqlen=4096, non-causal):

TMEM is the bottleneck post-optimization. Sweeping beyond 50% emulation increases instruction overhead without
additional speedup (benchmarked).

1. Performance (B200, hdim=256, bf16, locked clocks)
2. Forward — non-causal: +2% to +22% across seqlens 1k–128k (sustained +4–6% at seqlen ≥ 4k)
3. Forward — causal: +0.7% to +8.4% across seqlens 1k–128k
4. Backward: no regression (within ±2% noise)
5. No backward change

The exp2 emulation is applied only in the softmax forward path (compute_softmax_p_warp). Backward uses a separate
kernel — no modification needed.

e2e Benchmark

Base (3c015f1):

  PYTHONPATH=/tmp/fa-3c015f1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \
    python benchmarks/benchmark_attn.py \                                                                                                               
    --fwd --bwd --backend fa4 --headdim 256 \                                                                                                           
    --seqlen 1024,2048,4096,8192,16384,32768,65536,98304,131072 \                                                                                       
    --causal both --warmup 5 --rep 20         

Exp2 (e122e67):

  PYTHONPATH=/tmp/fa-e122e67 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \                                                                                 
    python benchmarks/benchmark_attn.py \                                                                                                               
    --fwd --bwd --backend fa4 --headdim 256 \               
    --seqlen 1024,2048,4096,8192,16384,32768,65536,98304,131072 \                                                                                       
    --causal both --warmup 5 --rep 20  

--- Conclusions (3-run avg, locked clocks @ 1965 MHz):

• FWD non-causal: +24% at 1k, +3–8% at ≥8k seqlen; strongest at 96k–128k (+7–8%).
• FWD causal: flat at 1k–32k (within noise); clear gains at 64k–128k (+2–5% 🟢).
• BWD: neutral across all configs (within ±1% noise). BWD causal 64k −3.2% 🔴 is a single-batch outlier inconsistent with the exp2 change scope — noise.