feat: add Megatron fused linear cross entropy backends

[justinhh4]

Summary

Add a Megatron-compatible fused vocab-parallel linear cross-entropy API with three local-compute backends:

• Triton (default): portable Triton projection, CE, dX, and dW kernels with NCCL tensor-parallel collectives
• CuTile: CuTile local compute with NCCL collectives
• CuTe DSL: persistent SM100 CuTe projection using Liger's existing SM100 GEMM helper, Triton CE, PyTorch/cuBLAS backward GEMMs, and NCCL collectives

Each implementation materializes one local low-precision projection buffer and saves it for backward. Backward recomputes FP32 exponentials from those logits, converts the buffer to dlogits in place, and reuses it for the projection gradients. Predicted-logit and sum-exp statistics share one NCCL SUM reduction, while the low-precision dX reduction runs asynchronously and overlaps dW/dBias computation.

Triton kernels follow Liger's standard flat src/liger_kernel/ops/ layout. The optional backends use the existing LIGER_KERNEL_IMPL=cutile|cutedsl replacement mechanism; this PR adds no new selector or backend registry. CuTe DSL requires exact SM100 hardware and reuses the _sm100_gemm.py infrastructure already on main.

Megatron integration

apply_liger_kernel_to_megatron(fused_linear_cross_entropy=True) installs FLCE through Megatron-Core 0.18's GPTModel._postprocess output-processor hook. Standard labeled GPT forwards with a native BF16/FP16 ColumnParallelLinear output layer route directly from hidden states to per-token loss; inference calls and caller-supplied output processors remain unchanged.

The automatic path fails explicitly for unsupported configurations, including gathered logits, sequence parallelism, gradient-accumulation fusion, deferred embedding wgrad, disabled output dgrad reduction, MTP, MuP output scaling, and separately returned output bias. LigerMegatronFusedLinearCrossEntropy remains available for explicit/custom wiring.

The existing cross_entropy=True patch now also repairs Megatron's import-time consumer bindings in LanguageModule and the tensor_parallel package export, so both fused and unfused CE routes reach Liger even when Megatron was imported before patching.

Correctness

• 254 tests pass under test/megatron
• 31 focused FLCE tests across Triton, CuTile, and CuTe DSL
• TP1 and TP2; BF16 and FP16
• Bias/no-bias, ignore-index, target validation, dispatch, Triton split-K, patch idempotency, and output-processor coverage
• Actual Megatron-Core 0.18 TP2 GPT forward/backward parity for all three backends against the materialized path, including loss, output-weight gradient, and transformer gradient
• Actual Megatron-Core fused and unfused LanguageModule CE routing
• Distributed benchmark correctness gate against materialized linear + vocab-parallel CE
• make checkstyle
• Clean wheel build includes every backend and no nested liger_kernel/ops/triton/ package

B200 benchmark

Full forward + backward, TP4, 16,384 tokens, hidden size 4,096, global vocab 128,256, BF16:

Provider	Forward (ms)	Backward (ms)	Full (ms)	Peak Torch memory (MB)
Megatron-compatible materialized linear + CE	3.9368	6.5434	10.1531	3827.5
Liger Triton	4.0878	6.9272	10.8324	2201.9
Liger CuTile	3.6115	6.8471	10.3739	2201.9
Liger CuTe DSL	4.4299	6.3169	10.1092	2201.9

All three Liger paths reduce measured peak Torch memory by about 42%. CuTile is 2.2% behind the materialized baseline full step and has the fastest forward. CuTe DSL is effectively at baseline full-step speed (0.4% faster in this run). The portable all-Triton backend is 6.7% slower on B200.

Peak memory is measured through the PyTorch allocator. The benchmark discards stale autograd graphs before each memory sample.

Follow-up

The faster NVSHMEM prototype is intentionally excluded: CUDA RDC, NVSHMEM runtime initialization, symmetric-workspace ownership, and packaging need a separate integration surface.