[hd256] Improve forward kernel with exp2 FMA emulation (3% to 9% performance gain)

[Johnsonms]

Change

Replace a fraction of hardware exp2 (SFU) instructions with a polynomial FMA emulation (ex2_emulation_2) in the softmax P-tile computation. The key insight: SM100's SFU throughput is a bottleneck for hdim=256 due to the large tile size. By substituting 3 out of every 4 exp2 calls (ex2_emu_freq=4, ex2_emu_res=3) with packed FMA polynomial approximation, we shift pressure onto the underutilized FMA pipeline.

Additionally, the P write-slot acquisition is moved earlier to overlap any pipeline stall with the exp2 compute.

Kernel-only change; no API change. Backward is untouched.

Validation (this PR vs current origin/main)

[origin/main] : #2487

B200, bf16, hdim=256, MHA (32:32) and GQA (32:2), 3-run means, locked clocks @ 1755 MHz, seqlens 4k..128k.

FWD delta vs origin/main (TFLOPS mean, 3 runs each)

seqlen	causal	MHA 32:32	GQA 32:2
4k	F	-0.2%	+0.7%
8k	F	+0.4%	+0.3%
16k	F	+0.7%	-1.0%
32k	F	+2.3%	-0.3%
64k	F	+5.4%	+5.2%
128k	F	+7.4%	+7.3%
4k	T	+0.3%	+0.9%
8k	T	+0.8%	+0.9%
16k	T	+0.8%	+1.1%
32k	T	+5.2%	-0.4%
64k	T	+1.1%	+1.1%
128k	T	+3.4%	+2.6%

• 19 of 24 cells positive; 4 slightly negative, all within the batch-quantization noise band that origin/main itself already showed in our 3-run regression sweep.
• Peak gain +7.4% (MHA 128k non-causal) — exactly where softmax SFU pressure is worst, consistent with the theory above.
• Averages: MHA fwd +2.3%, GQA fwd +1.5%.

Correctness smoke

pytest tests/cute/test_flash_attn.py::test_flash_attn_output -k "256-False-0-0.0-False-False" on B200:
78 passed, 78 skipped, 0 failed — identical pass/skip count to origin/main.

Co-authered by @Johnsonms and @wangsiyu

[tridao]

I'm a bit surprised that ex2 emulation helps here because the ratio of ex2 vs MMA is lower for large headdim, so ex2 emulation is less necessary.
Did you try 50:50 ratio between ex2 and ex2_emulation?

[Johnsonms]

This

I'm a bit surprised that ex2 emulation helps here because the ratio of ex2 vs MMA is lower for large headdim, so ex2 emulation is less necessary. Did you try 50:50 ratio between ex2 and ex2_emulation?

On the 50:50 question
The original PR values (freq=4, res=3) actually produce exactly 50:50 hw/emu for hd256 （TBH unintended 😊) — since the inner loop steps k by 2, k % 4 only takes values {0, 2}, and res=3 triggers emulation for k%4 ≥ 1, which is every other step.

Using tune_ex2_emu tuned with locked frequency and found freq=14/res=6 (~43% emulation) slightly outperforms that as following.

freq=4/res=3

freq=14/res=6

• Config C (freq=14/res=6) clearly outperforms baseline Config A, especially on long sequences (32k+) with +4% to +9% gains. Short sequences (4k–8k) also improve by +1% to +3%, with almost no regressions.
• Config B (freq=4/res=3) is unstable; causal results are mostly negative, suggesting the 50:50 emulation ratio hurts the causal path.
• The C vs. B gap is mainly driven by causal performance: Config C is ~+12% faster than B on causal 32k.

On why it helps at all
I checked the key NCU metrics:

1. In the baseline config A, FMA pipeline utilization is only ~24.5%, suggesting spare FMA capacity.
2. Emulation shifts ~28% of SFU work onto FMA, increasing FMA utilization from ~24.5% to ~36% and better balancing the two pipelines.
3. SM throughput improves by ~4% (94% → 97–98%), which is consistent with the +2–8% benchmark gains.

So the benefit is likely not simply that “SFU is the bottleneck,” but that FMA had unused capacity. Moving part of the work from SFU to FMA gives a modest but measurable improvement.

Thanks @tridao for leading me to explore this more deeply.

[tridao]

This is great, thanks @Johnsonms