perf(flce): write the weight-gradient product directly into grad_weight

[ananthv26]

Summary

The default FLCE path (accum_dtype=None) accumulates the weight gradient per chunk as

grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()

This materializes the V×H product twice (a bf16 mm output, then a ~2.0 GB fp32 copy at
llama-3-8B shapes) and re-reads/rewrites grad_weight: several GB of avoidable HBM
traffic per chunk, O(V·H) regardless of token count. It made default FLCE slower than
the HF linear+CrossEntropyLoss baseline at all measured sizes and heavier below
BT=8192.

Changes (Triton kernel and chunk geometry untouched):

• write the product straight into grad_weight via mm/addmm out= on both accumulation paths
• first chunk overwrites (mm, beta=0) instead of accumulating, dropping one
  read-modify-write pass over grad_weight
• allocate grad_weight/grad_input with empty_like; the loop writes every element
• older torch / non-CUDA / mixed-dtype cases keep the previous fallback

Numerics are equal or better: cuBLAS accumulates the product and beta*C in fp32 and
rounds once, where the old path rounded the product to bf16 before adding.

Overlaps with #1324, which also routes same-dtype fp16/bf16/fp32 accumulation through
addmm out=. This PR additionally skips the first-chunk read-modify-write, drops the
zero-fills, and keeps the mixed-dtype fallback. Happy to consolidate the two.

Details

Measured on current main (2798d08) vs this branch, H100 80GB HBM3, torch
2.13.0+cu130 / triton 3.7.1, llama-3-8B shapes (V=128256, H=4096), bf16,
reduction=mean, full pass (fwd+bwd). Standalone script equivalent to
benchmark/scripts/benchmark_fused_linear_cross_entropy.py (same model definitions
and providers), median of triton.testing.do_bench; each config run twice, runs
agreed within ~1–3%.

Default liger provider (accum_dtype=None):

BT	time before → after	peak mem before → after	HF torch (time / mem)
1024	13.21 → 5.86 ms (−56%)	5215 → 2335 MB (−55%)	4.93 ms / 2335 MB
2048	16.89 → 9.75 ms (−42%)	5357 → 2601 MB (−51%)	9.74 ms / 2602 MB
4096	26.39 → 20.18 ms (−24%)	5640 → 3136 MB (−44%)	20.54 ms / 4104 MB
8192	44.47 → 37.34 ms (−16%)	6204 → 4200 MB (−32%)	40.02 ms / 7142 MB

After the change the default path matches or beats the HF baseline on latency at
BT ≥ 2048 (previously slower at every size) and uses less peak memory everywhere,
41% less at BT=8192. At BT=1024 it remains ~1 ms slower on latency at equal memory.
The 2.0–2.9 GB memory drop matches the eliminated V×H fp32 product
(128256×4096×4 B ≈ 2.0 GB).

accum_dtype=torch.float32 (already on the addmm(out_dtype=) path) gains 7–16%
latency from the beta=0 first chunk and the dropped zero-fills, memory flat; the
BT=8192 row is within noise.

The dispatch layer added in #1416 is unaffected: the default (no LIGER_KERNEL_IMPL)
path and the _triton backend adapter both import
LigerFusedLinearCrossEntropyFunction from the patched module;
ascend/cute/cutedsl/cutile backends are opt-in and untouched.

Developed with AI assistance (Cognition's Devin: profiling, patch, and benchmark
runs); human-reviewed.

Testing Done

On this branch (current main 2798d08 + this commit), H100 80GB:

• make test (full suite): passed

• python -m pytest test/transformers/test_fused_linear_cross_entropy.py: 141 passed
  (AMP, accum_dtype, ce_weight, softcap, z-loss, label smoothing, reduction="none")

• python -m pytest test/convergence/bf16/test_mini_models.py -k llama3: passed

• make checkstyle: clean

• Hardware Type: H100-80G-HBM3

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence (llama3 bf16 subset run; full suite not run)