[Perf] Fuse log-softmax into fused_linear_jsd kernel

[hiwuhgds-pixel]

Summary

fused_linear_jsd computed its forward in two steps: two torch.log_softmax
calls materialized fp32 log-probability tensors, _jsd_kernel consumed them, and
the log-softmax Jacobian was then applied back in PyTorch. This PR replaces that
with a single Triton kernel that derives both log-softmaxes from the logits
in-kernel and applies the chain rule inline.

On an H100 with the default chunking (H=4096, V=128256, bf16): 1.50-2.29x
faster on the full pass and 9.5-44.2% lower peak memory, both growing with
B*T. At num_chunks=1 the same change is 2.37-2.68x faster and 26.3-71.2%
lighter.

Math and dtype contract are unchanged — same expression, same fp32 compute, and
@amp_custom_fwd / @amp_custom_bwd are kept.

Details

The technique is the fused_linear_cross_entropy pattern applied to JSD:

• Online softmax in-kernel instead of materializing log-probs — the same
  algorithm (Milakov & Gimelshein, Algorithm 3) liger_cross_entropy_kernel uses.
  JSD needs 3 passes rather than CE's 2: pass 1 builds m/d for student and
  teacher together, pass 2 accumulates the loss and dX_sum, pass 3 recomputes
  dX and writes the gradient. dX_sum is a cross-column term, so it has to be
  known before any gradient is stored.
• The GEMM keeps the input dtype, and the cast to fp32 happens at tl.load
  rather than on the whole logits tensor after the matmul.
• loss_1d shrinks from (BT, V) to (BT,) — the kernel reduces over V
  itself.
• The student logits buffer is reused in place as the gradient buffer, the
  same trick FLCE uses.
• e^x is emitted as exp2(x * log2 e) to reach the hardware ex2.approx path.

Peak memory drops from three structural changes, per chunk of C rows at vocab
V with s = bytes per element of the input dtype:

	before	after
loss_1d (unchunked, scales with BT)	BT · V · 4	BT · 4
the two log-prob tensors	2 · C · V · 4	—
the two logits tensors	2 · C · V · 4	2 · C · V · s

The PyTorch Jacobian expression also allocated several more C · V temporaries
(softmax, the broadcast product, the subtraction, the dtype cast) that the
kernel no longer needs, so the measured saving exceeds what the table accounts
for. Only loss_1d scales with BT rather than C, which is why the reduction
is largest when chunking is coarsest — at num_chunks=1, −26.3% at BT=1024
rising to −71.2% at BT=8192.

ops/jsd.py is deliberately untouched — see the open questions.

Benchmarks

Hidden size: 4096, Vocab size: 128256, bf16, NVIDIA H100 80GB HBM3. "Liger old"
is the current implementation; the default num_chunks for this shape is 32.
Speed is the median over the benchmark harness's runs, memory is peak allocated.

Speed — full pass (ms, median)

B*T	torch	Liger old	Liger new	new vs old
1024	26.41	129.52	86.26	1.50x
2048	50.11	152.54	88.28	1.73x
4096	98.43	178.41	89.98	1.98x
8192	197.89	227.74	99.31	2.29x

Memory — full pass (MB, peak)

B*T	torch	Liger old	Liger new	reduction
1024	10609	5662	5122	−9.5%
2048	17146	6250	5169	−17.3%
4096	30220	7425	5265	−29.1%
8192	56368	9776	5455	−44.2%

Open questions

1. The fp32 rounding point moved, and #336 is the reason to ask.
#336 ("Fix FusedLinearJSD
precision issue when using AMP") made the path cast logits to fp32 right after the
matmul, so that "all the computation between logit to final JSD loss happen on
FP32", guarded by test_amp. This PR still does all logit→loss computation in
fp32 — but in registers rather than in HBM, so the GEMM output itself now rounds
to the input dtype. chunked_loss/jsd_loss.py already runs F.log_softmax on
unconverted bf16 logits, so the two paths were inconsistent before this change and
are consistent after it. test_amp passes. Still, this narrows the intent of
#336, so it should be a maintainer call rather than mine.

2. Should temperature scale the loss by T² rather than T?
Hinton et al. 2015 multiply the soft-target objective by T², because soft-target
gradients scale as 1/T² and the factor keeps gradient magnitude stable as T
changes. Liger divides by T once, everywhere: here, and in
chunked_loss/fused_linear_distillation.py. There is no temperature ** 2 in the
repo and no mention of temperature in docs/, so users currently have to know to
pre-scale themselves. This PR keeps the existing behaviour — changing semantics
does not belong in a perf PR. Worth deciding: fold T² into the kernel, or document
that the caller owns it.

3. Why the kernel is here and not in ops/jsd.py.
_jsd_kernel takes log-probabilities and is public API — LigerJSDFunction is
exported from ops/__init__.py, reached via transformers/jsd.py and
functional.py, and its docstring states the log-space contract. The cutile and
Ascend backends carry the same signature. Rewriting it to take logits would break
all of that, so the fused kernel lives in fused_linear_jsd.py and jsd.py is
unmodified. The cost is two code paths for the same JSD formula. Dedupe or accept
is a maintainer decision.

4. Aside: the chunking heuristic looks mis-tuned for JSD.
While benchmarking I swept num_chunks from 1 to 32. Speed scales far more
strongly with chunk size than FLCE's heuristic assumes. torch does not chunk, so
its column is the same series throughout and serves as the fixed reference:

Speed — full pass (ms, median), new kernel

B*T	torch	nc=1	nc=8	nc=32 (current default)	nc=1 vs default
1024	26.41	9.69	24.25	86.26	8.90x
2048	50.11	16.60	26.19	88.28	5.32x
4096	98.43	30.08	39.81	89.98	2.99x
8192	197.89	60.62	65.38	99.31	1.64x

Memory — full pass (MB, peak), new kernel

B*T	torch	nc=1	nc=8	nc=32 (current default)
1024	10609	5607	5169	5122
2048	17146	6140	5263	5169
4096	30220	7206	5453	5265
8192	56368	9338	5831	5455

Against torch the new kernel is 2.73x / 3.02x / 3.27x / 3.26x faster at nc=1,
but 0.31x / 0.57x / 1.09x / 1.99x of torch's speed at the default — so with the
current heuristic it is slower than plain torch below BT≈4096, and only the
coarse-chunk configuration beats torch everywhere. Part of the reason is
structural: the heuristic pins num_chunks at cdiv(V, H) = 32 for every BT in
this sweep — the chunk size is grown to absorb a larger BT instead, so even at
BT=1024 the sequence is still split into 32 pieces, fragmenting work that would
run faster as a single pass. At B*T=1024, nc=1 buys 8.9x the speed for 9.5%
more memory, and still sits at roughly half of torch's 10609 MB.

inc_factor = cdiv(V, H) is inherited from FLCE, but JSD holds more live buffers
per chunk (two distributions plus M, against CE's one), so the formula does not
transfer and currently splits too finely. Notably chunked_loss/jsd_loss.py
already uses a flat chunk_size = 1024 instead of a formula, so there is
precedent for JSD not following FLCE here. I have left the heuristic alone — it is
independent of this kernel and wants its own PR, but the effect is larger than
this one.

Testing Done

test_correctness gained one shape, (1, 4, 64, 40960). Every pre-existing case
had V ≤ 4096, so with BLOCK_SIZE = min(32768, next_pow2(V)) they all ran
single-block. The old kernel did not care — blocks were independent — but the new
one carries m/d rescaling and dX_sum across blocks, which is the path
production V=128256 takes. V=40960 also leaves a partial trailing block (8192 of
32768 lanes valid), covering the mask × multi-block interaction. 58 → 66 cases.

make test: 3914 passed, 942 skipped, 14 xfailed of 4870 collected. Within that
run, test_fused_linear_jsd.py 66/66 and test_jsd.py 67/67 — jsd.py is
unmodified but no longer imported from here, so it is worth confirming. The eight
new cases cover both dtypes across all four beta branches (0.0 forward KL, 0.1 and
0.5 generalized, 1.0 reverse KL).

make checkstyle: clean, no diff.

• Hardware Type: H100-80G-HBM3 (also verified on A100)
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence