[vaibhavjindal]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add opt-in Liger vocab-parallel cross entropy through ops/loss.

Stacked on #8656; this is the CE-only successor to the corresponding part of #5438. It does not depend on the separate RMSNorm contribution.

--cross-entropy-loss-fusion --cross-entropy-fusion-impl liger binds the adapter through the existing shared selector. No model-forward dispatch or normalization changes. Disabled fusion remains native unfused CE, and unselected paths do not import Liger.

The adapter returns FP32 per-token losses, preserves incoming gradient dtype, and uses the supplied TP group. Optional original global argmax preserves MTP prediction semantics, including smallest-index ties. Missing selected dependencies fail through ops._backends.require.

CI: the H100 ops, models and determinism-kernel buckets install liger-kernel==0.8.5 with --no-deps in latest/dev runs on dev images. The shared export check fails before tests on a broken installation. The image's IMAGE_TYPE marker excludes LTS even when GitHub labels its recipe environment dev; runtime installation remains optional.

Performance (Liger v0.8.5)

Measured on H100 80GB with exact release v0.8.5 (1ab779f), PyTorch 2.13.0.2+cu130, Triton 3.7.1 and TE 2.20.2.0+cu130. Times are median synchronized wall latency; lower is better.

Vocab-parallel CE, forward + backward: BF16, sequence 2048, batch 1, global vocabulary 128K. 10 warmups, 3 x 50 samples. TP results use the slowest rank; fresh input clones are outside timing for every backend.

TP	Native unfused (ms)	Native fused (ms)	Liger (ms)	Speedup vs fused	TE (ms)
1	5.340	3.646	2.041	1.79x	0.711
2	2.777	1.925	1.098	1.75x	0.507
4	1.500	1.056	0.737	1.43x	0.404
8	0.871	0.816	0.613	1.33x	0.381

At TP8, peak extra allocated memory is 188.5 MiB native fused / 188.5 MiB Liger / 0.30 MiB TE, excluding the live input. Liger improves native CE latency here; TE remains faster and uses less extra memory.

Full-model ablation: 32-layer Llama2-7B shape, H4096, FFN11008, vocab32K, BF16, batch1, TP8; only this PR's operation changes. Three independent process launches/backend, 5 warmups +20 steps/sequence, forward + backward without optimizer or CUDA graphs. Table latencies are medians of the three process medians; the last column shows paired trial ranges (positive = slower).

Sequence	Native model (ms)	Liger-only model (ms)	Paired latency change
512	83.544	83.000	-3.24% to +0.37%
2048	84.253	84.004	-0.30% to +3.74%
4096	155.667	155.574	-0.09% to -0.03%

The isolated CE win does not establish a material general end-to-end speedup.

Numerics: the separate full-gradient diagnostic has BF16 relative-L2 differences of 1.94-1.98% versus native and approximately 2.45e-06 in FP32. These are not convergence-equivalence claims.

Issue tracking

Related to #5488 and #5438. Retarget to main after #8656 lands, before merging.

Code provenance

No external kernel source copied. The adapter calls the installed Liger API; installation guidance identifies v0.8.5 or a newer compatible version.

Contribution process

• Unit tests, GPT/MTP forward/backward parity, and registered bit-exact kernel replay coverage
• Dedicated functional-test recipe
• Typing and user documentation
• tools/autoformat.sh against pull-request/8656

Validated independently on H100: FP32/BF16, TP=1/2/4, interleaved non-WORLD groups, strided inputs, arbitrary upstream gradients, and native/PyTorch parity.

Latest-release focused regression suite: 495 passed per rank on four GPUs (Liger v0.8.5, locally built from the release tag).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.