perf: remove device to host syncs in cross entropy forward and backward

[kashif]

What

The fused linear cross entropy (and plain cross entropy) forward calls .item() on the non-ignore token count before launching the Triton kernel, and the backward guards the grad scaling with torch.equal(grad_output, 1.0). Both pull a value off the GPU, so each one forces a device to host sync. In a training profile these show up as aten::item, _local_scalar_dense and StreamSynchronize blocks sitting on the critical path inside LigerFusedLinearCrossEntropyFunction.forward.

This PR removes both syncs.

Forward. Keep the non-ignore count and the weighted sums as 0-D device tensors instead of Python scalars, and pass them into the kernel as pointers. The kernel loads and casts them internally, so the mean reduction and the numerics are unchanged.

Backward. Move the "skip the multiply when grad_output is 1.0" decision into element_mul_kernel. It already loads grad_output, so it now returns early on-device when the value is exactly 1.0. This is bit identical because x * 1.0 == x. The wrappers launch the kernel unconditionally and drop the torch.equal check. The reduction='none' path still branches on grad_output.ndim, which is a host side shape check and does not sync.

Why it matters

Dropping the syncs lets the CPU keep queuing work instead of stalling on the GPU. That helps compute and communication overlap under FSDP2 and DeepSpeed, and it unblocks CUDA and HIP graph capture, since a single .item() breaks capture. The math is identical, so this is a pure latency and capture win.

Benchmark

Measured on the pi052 CE path (LigerFusedLinearCrossEntropyLoss, forward plus backward) on an AMD Radeon 890M, ROCm 7.2, with BT=2048 H=2048 V=131072.

metric (per step)	original	this PR
device to host sync CPU cost	4411.5 us	2.3 us
blocking StreamSynchronize	2	0
wall time	3825 ms	3826 ms

Wall time does not move on this iGPU because it is heavily compute bound, so a few ms of CPU stall hides behind the GPU work. The payoff shows on a fast GPU where the CE compute is quick and the stall was blocking the rest of the model from launching, which is the setting the original profile came from.

Testing

test_cross_entropy.py and test_fused_linear_cross_entropy.py pass (306 passed). The only failures are the pre existing test_amp cases, which fail the same way on main from an unrelated AMP addmm dtype issue. test_float32_internal was updated to pass the count as a 0-D tensor to match the new kernel signature.

Heads up for reviewers: this changes the internal signature of liger_cross_entropy_kernel. The n_non_ignore, sum_non_ignore_weight and weight_sum arguments are now pointers to 0-D tensors rather than Python scalars, so any code calling that kernel directly must pass tensors.

🤖 Generated with Claude Code

[kashif]

cc @vaibhavjindal if you can kindly have a look thanks!