Kernel Bottleneck Tuning

Kernel Bottleneck Tuning (validated on FA4 FP4/FP8 PV, GB300)

Methodology that produced +15% (FP4 PV log-domain quant) and +24% (FP8 PV P-split) on this repo. Follow the order — each step rules out a class of fix before you spend effort on the next.

1. Measure instruction throughputs in isolation

Microbenchmark the suspect instructions with serial dependency chains (8 independent chains/thread, feed output back to input so the loop can't be hoisted; %clock64 with "memory" clobber around the timed region — naive volatile-asm loops get scheduled away). Template: agent_space/bench_cvt_throughput.cu. This gave: ex2 = 32/clk/SM, bf16-cvt = 62, e4m3-cvt = 32, e2m1-cvt = 57, mixed ex2+e4m3 = 54.6 (partial dual-issue).

2. Trace the real pipeline, coarsely

FA4_PROFILE_PIPELINE=1 python3 flash_attn/cute/debug/trace_pipeline.py --pv_mode {bf16,fp8,fp4}. Use the default coarse mode for absolute numbers (per-step period/busy, MMA wait-P). FA4_PROFILE_DETAIL=1 gives the per-phase breakdown but inflates spans ~15-30%: each %clock inline asm is side-effecting and blocks ptxas from interleaving across the boundary. Never tune from detailed/instrumented runs alone.

3. Verify scheduling hypotheses in SASS before reordering source

Don't software-pipeline by hand until you've checked the compiler hasn't already done it: dump PTX (CUTE_DSL_KEEP_PTX=1, strip NULs with tr -d '\000'), ptxas -arch=sm_103a -O3, nvdisasm -c, then compare unit-run interleaving (e.g. MUFU↔F2FP transition counts) between variants. On FA4 the whole softmax step is one fully-unrolled branch-free region (range_constexpr = trace-time unrolled; its unroll= kwarg is silently discarded), so ptxas already interleaves optimally — source reordering and chunk-pipelining measured exactly 0. Corollary: register-resident fragments REQUIRE the unrolled form — a materialized cutlass.range loop makes the index dynamic and spills the fragment to local memory (measured 8.4x slowdown).

4. If scheduling has no headroom: eliminate instructions algebraically

The wins come from algebra, not order. Pattern that paid off — log-domain refactoring: exp2 is monotonic, so max(exp2 s) = exp2(max s), and any multiplicative post-exp scale folds into the exp2 argument as a subtract: exp2(s)/exp2(m) = exp2(s-m). For FA4 FP4 quant this deleted all per-group divisions and the whole post-exp scaling pass (+15%). Look for: divisions that could be exp2 of a negated log, multiply passes foldable into an existing exp/log, reductions movable across a monotonic function to where the data already is.

5. Tune the producer→consumer handoff split

Warp-specialized pipelines hand data over in chunks gated by two mbarriers (first fraction at signal 1, rest at an embedded wait). The split fraction is a tuning knob with hardware-generation-dependent optima: FP8 PV's 1/2 (a B200 tuning) cost 24% on B300 vs 3/4. Safety contract: the producer-side split fraction must be >= the consumer GEMM's pre_mbar_tiles fraction (both derive from mbar_p_split); violating it reads unwritten TMEM (1/4 crashed). Sweep via FA4_FP8_PV_P_SPLIT_NUM/DEN (+ FA4_FP4_PV_P_SPLIT_*, FA4_FP4_PV_TMEM_STORE_REP for chunk count).

6. Validation discipline (every change)

• Bitwise harness with SAVED inputs (agent_space/check_pv_pipeline.py): tensor creation (nvfp4 quantize amax) is NOT reproducible across processes — comparing fresh-tensor runs produces phantom diffs.
• Pure reorderings must be bitwise-identical; algebraic changes must show error vs an FP32 reference identical to baseline (agent_space/check_fp4_log2.py).
• Flake-test ~20 runs against the saved reference: the FP4 baseline had a pre-existing 1-in-20 nondeterminism — a single-run pass can lie, and a "failure" may predate your change (bisect with the env knobs).
• Env-knob variants need a fresh process per setting (constexprs read env at init; the in-memory compile cache won't recompile).