FA4 Forward Kernel Optimization Report

Author: Zhuobin Huang @ mars-compute | Date:: March 9, 2026

Target: FlashAttention-4 Forward Kernel on Blackwell (sm100a)

Tool: mars-compute

Baseline: commit 8451d4e | Optimized: commit 6017fa7

Files Modified:

• flash_attn/cute/flash_fwd_sm100.py
• flash_attn/cute/softmax.py

1. Executive Summary

Through profiling-guided optimization of the FA4 forward kernel's softmax path, mars-compute achieved an overall +1.3% average TFLOPS improvement across all tested configurations (MHA/GQA, causal/non-causal, bf16/fp16, sequence lengths 1k–32k on Blackwell). The optimization reduces hardware exp2 unit pressure by tuning the software emulation parameters in the softmax warp group. The change is pure Python (CuTe DSL), requires no C++ rebuild, and introduces zero accuracy regression for bf16/fp16 workloads.

Overall Summary

* Causal GQA bf16 has one outlier at seq=8192 (-2.0%, likely measurement noise); excluding it, the average is +0.9%.

2. Profiling & Root Cause Analysis

2.1 Baseline Establishment

Initial baseline measurement at batch=8, seqlen=8192, nheads=4, headdim=128, bf16:

• Baseline: ~1541 TFLOPS

2.2 Range Profiling

Range profiling revealed the kernel's overall SM utilization and identified that the tensor pipe (MMA) was not fully utilized:

• Tensor pipe utilization: 76.88%

2.3 Intra-Kernel Tracing

Trace sentinels were inserted into the CuTe DSL source to measure per-warp-group cycle counts. Results revealed the critical bottleneck:

Root cause: The MMA warp completes each PV gemm iteration in ~1712 ns but must wait ~3091 ns for softmax to produce the next P tile. The softmax path is bottlenecked on the hardware exp2 unit — all 8 softmax warps compete for the same exp2 functional unit, creating a pipeline stall.

2.4 SASS Analysis

Binary analysis confirmed dense MUFU.EX2 (multi-function unit exp2) instructions in the softmax hot loop, validating the exp2 pressure hypothesis.

3. Optimization Strategy

The FA4 kernel already includes an exp2 emulation mechanism (ex2_emulation_2) that replaces hardware exp2 with FMA-based polynomial approximation for a fraction of elements. The key insight: the default parameters are too conservative — they don't offload enough work from the exp2 unit to the FMA pipeline.

3.1 Changes Made

Note: commit 6017fa7 restored poly_degree back to 3 (degree-3 polynomial), retaining only the ex2_emu_freq and ex2_emu_start_frg changes as the active optimizations.

3.2 Why This Works

1. ex2_emu_freq 16→10: With freq=16, only every 16th pair of elements uses software emulation (6.25%). With freq=10, every 10th pair is emulated (10%), shifting more work from the exp2 unit to the FMA pipeline. The FMA units have much higher throughput and are underutilized during softmax.
2. ex2_emu_start_frg 1→0: Previously, the first fragment (32 elements) always used hardware exp2. Starting emulation from fragment 0 ensures the exp2 unit pressure relief applies uniformly across all fragments.

3.3 Failed Optimization Attempts

4. Results After Optimization

4.1 MHA 16 Heads (32k tokens, headdim=128)

Non-Causal MHA bf16

Avg: +1.6% | Peak: +2.9% (8k)

Non-Causal MHA fp16

Avg: +2.0% | Peak: +3.1% (32k)

Causal MHA bf16

Avg: +0.8% | Peak: +2.1% (8k)

Causal MHA fp16

Avg: +0.6% | Peak: +1.7% (32k)

4.2 GQA 40Q/8KV (32k tokens, headdim=128)

Non-Causal GQA bf16

Avg: +2.4% | Peak: +2.9% (16k)

Non-Causal GQA fp16

Avg: +1.9% | Peak: +2.9% (16k)

Causal GQA bf16

Avg: +0.3% | Peak: +1.3% (4k)

* The seq=8192 data point shows a -2.0% dip which is likely measurement noise (thermal variance or background activity). Excluding that outlier, the average is +0.9%.

Causal GQA fp16

Avg: +0.8% | Peak: +2.1% (32k)

5. Key Findings

1. Positive gains across nearly all configurations: mars-compute delivers consistent improvements across MHA/GQA, causal/non-causal, bf16/fp16, and sequence lengths 1k–32k, with only one outlier data point.
2. Non-causal configurations benefit more than causal: Non-causal MHA/GQA sees +1.6–2.4% average gains, while causal sees +0.3–0.8%. Causal masking skips ~50% of tiles, reducing the number of softmax iterations where the optimization applies.
3. Gains scale with sequence length: Longer sequences have more softmax iterations per CTA, providing more opportunities for the exp2 emulation to reduce bottleneck pressure. Peak gains typically occur at 8k–32k.
4. GQA benefits comparably to MHA: Despite the different compute-to-memory ratio (5:1 Q-to-KV head ratio), the softmax bottleneck is equally present and equally relieved.
5. bf16 slightly faster than fp16 in absolute TFLOPS: bf16's wider exponent range reduces the frequency of rescale corrections in online softmax, leading to fewer wasted cycles.