perf: fuse mHC coefficient backward epilogue

[heiheiha798]

Addresses heiheiha798#8.

Engineering record: heiheiha798#9.

Summary

LigerMHCCoeffsFunction.backward currently leaves Triton after the Sinkhorn backward and executes a fixed eager PyTorch coefficient-gradient epilogue before the fused MM/RMS backward. For the usual HC=4, that epilogue operates on only 24 columns but launches 29 pointwise/copy/reduction kernels on B300.

This change replaces only that epilogue with:

• one Triton producer for the pre/post sigmoid derivatives, grad_mix, and FP32 partial grad_b/alpha gradients;
• one deterministic reducer for the 24 grad_b columns and three scalar alpha gradients;
• reference-gradient coverage for pre-only, post-only, and residual-only losses;
• an exact-source comparison tool under benchmark/tools/ that loads clean baseline/candidate worktrees in separate processes and alternates providers.

The comparison helper intentionally lives outside benchmark/scripts/benchmark_*.py, so the documented make run-benchmarks wildcard continues to contain only shared-framework benchmark scripts. Sinkhorn backward, MM/RMS backward, pre/post-residual kernels, forward behavior, public APIs, and the Ascend backend are unchanged. There is no autotune or architecture dispatch.

Testing Done

• Hardware Type: NVIDIA B300 SXM6 AC (SM103)
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

Additional recorded validation:

• Slurm job 4953: full test/transformers/test_mhc.py, 54 passed in 16.11s.
• Coverage includes FP16, BF16, allowed FP32, HC={2,4,8}, odd/ragged token and channel counts, tmax={1,4,20}, and pre-only/post-only/residual-only non-unit upstream gradients. All six gradients are compared against mhc_coeffs_ref with repository tolerances.
• Slurm job 4950: focused compile/correctness smoke, 40 passed.
• Login-node make checkstyle at PR HEAD: passed, 371 files formatted.
• Clean exact-HEAD job 4966: the tracked benchmark/tools/compare_mhc_coeff_epilogue_commits.py loaded baseline 124fb8a2 and candidate 4302c8c from detached clean worktrees, alternated both providers for two rounds, and completed backward/full CUDA-event timing plus JSONL/readable summaries.
• Job 4959 is retained as a failed pre-worker harness attempt. Job 4960 passed after its one-line coordinator fix while the tool was still uncommitted. Review follow-up job 4966 supersedes both for immutable-HEAD traceability.

The full repository test suite and convergence suite were not run, so their template boxes remain unchecked.

Commits and Environment

• Baseline: 124fb8a2a443c1118f7caeff0203aea78424eb2b
• Kernel/test candidate measured by job 4953: 189bf55c4d53e8fbb5ff4e1ae96ae19431d3152b
• PR HEAD: 4302c8c4497983fc8990ac819e32656e4dbd3e79
• GPU: NVIDIA B300 SXM6 AC, SM103, 148 SMs
• Driver 580.126.09; CUDA 13.0; PyTorch 2.13.0+cu130; Triton 3.7.1; Nsight Compute 2025.3.1

Job 4953 verified both detached worktrees, exact commits, import paths, index, and worktree cleanliness before GPU work. Each provider ran in a separate exact-source Python process in one B300 allocation. The two later commits add and then relocate only the comparison tool; kernel/test source remains the measured 189bf55 content.

Performance Method

Primary matrix: BF16, B=4, HC=4, C={1024,2048,4096,8192}, T={128,512,2048}, tmax=20, direct coefficient backward and directly measured forward-plus-backward. Controls cover FP16, HC={2,8}, and tmax=8.

For every comparison:

• baseline/candidate processes use identical seeds, immutable inputs, and independent non-unit upstream gradients;
• providers alternate for seven rounds in one allocation;
• each round uses 10 symmetric warmups and 100 CUDA-event measurements;
• cache clearing occurs before the start event;
• backward reuses an equivalent retained forward graph; full directly measures forward plus backward and is never formed by adding independent quantiles;
• reported p20/median/p80 are quantiles across the seven round medians.

Selected job 4953 results (milliseconds, p20 / median / p80):

Shape	Mode	Baseline	Candidate	Speedup
BF16 C1024 T128	backward	0.352333 / 0.353600 / 0.354861	0.168710 / 0.180752 / 0.184003	1.956x
BF16 C1024 T512	full	0.530493 / 0.531008 / 0.533990	0.318253 / 0.345856 / 0.365459	1.535x
BF16 C4096 T512	backward	0.380845 / 0.380928 / 0.381005	0.255888 / 0.255952 / 0.256010	1.488x
BF16 C4096 T512	full	0.618496 / 0.618528 / 0.618637	0.492614 / 0.493472 / 0.493517	1.253x
BF16 C8192 T2048	backward	1.901504 / 1.901584 / 1.901613	1.761162 / 1.761200 / 1.761296	1.080x
BF16 C8192 T2048	full	2.396307 / 2.397104 / 2.397190	2.251699 / 2.251712 / 2.251725	1.065x
FP16 C2048 T512	backward	0.351424 / 0.354304 / 0.355942	0.173453 / 0.187520 / 0.188416	1.889x
BF16 HC8 C2048 T512	full	1.014986 / 1.015872 / 1.015936	0.883910 / 0.884208 / 0.884669	1.149x

All 36 declared backward/full comparisons passed the no-regression gate. The minimum ratio was 1.065x; 22 primary comparisons improved by at least 10%.

Profile Evidence

Nsight Compute profiled the same exact-source BF16 B=4, HC=4, C=4096, T=512, tmax=20 backward:

• whole backward: 33 launches / 582.240 us baseline to 6 launches / 378.560 us candidate;
• targeted coefficient epilogue: 29 launches / 214.688 us baseline to 2 launches / 14.400 us candidate;
• candidate producer: 9.024 us, 56 registers/thread, grid 64;
• candidate reducer: 5.376 us, 18 registers/thread, grid 27;
• representative peak incremental allocation: 69,994,496 to 69,796,864 bytes for backward and 73,017,344 to 72,819,712 bytes for full.

NCU timings are intrusive supporting evidence; the CUDA-event table above is the performance result.

Reproduction

Create clean detached worktrees, then run the tracked tool inside a B300 allocation:

git worktree add --detach /tmp/liger-mhc-baseline 124fb8a2a443c1118f7caeff0203aea78424eb2b
git worktree add --detach /tmp/liger-mhc-candidate 4302c8c4497983fc8990ac819e32656e4dbd3e79

.venv/bin/python benchmark/tools/compare_mhc_coeff_epilogue_commits.py \
  --baseline-worktree /tmp/liger-mhc-baseline \
  --candidate-worktree /tmp/liger-mhc-candidate \
  --dtype bfloat16 \
  --batch 4 \
  --tokens 128,512,2048 \
  --streams 4 \
  --channels 1024,2048,4096,8192 \
  --tmax 20 \
  --modes backward,full \
  --rounds 7 \
  --warmup 10 \
  --repetitions 100 \
  --output /tmp/mhc-coeff-epilogue.jsonl

The tool writes machine-readable per-round and summary JSONL to --output, and writes a compact human-readable comparison table to stderr. FP16 and alternate-HC/tmax controls can be reproduced by changing the explicit CLI values.

Risks and Limitations

• grad_b and alpha reductions use a deterministic two-stage order rather than eager PyTorch's reduction tree; tests intentionally use repository tolerances instead of claiming bit identity for those reductions.
• Performance evidence is from one B300 system. The implementation is type-consistent and does not add an SM103-only dispatch, but no performance claim is made for other accelerators.
• The optimization is beneficial over the declared matrix, but the gain narrows at the largest C8192/T2048 point to 1.065x full; this is reported rather than hidden.