Fix wrong grad_weight/grad_bias for fused linear cross entropy with reduction='none'

[kashif]

Summary

LigerFusedLinearCrossEntropyFunction returns wrong grad_weight and grad_bias when reduction="none" and the upstream gradient is not uniform.

With reduction="none" the loss is a per-token vector, so backward() receives a per-token grad_output of shape (BT,). That is the whole point of the mode (token weighting, masking, per-sample scaling). But element_mul_kernel loads grad_output with a scalar tl.load(grad_output_ptr), so every gradient ends up multiplied by grad_output[0].

There is already a comment in the code acknowledging this:

# Need extra calculations for backward if reduction=='none'. Not supporting reduction='none' now.

The failure is silent. Forward is correct, backward returns plausible-looking numbers that are simply wrong.

Why it cannot be patched in backward

Forward computes the logits gradient per chunk and immediately projects it, collapsing the token dimension:

grad_weight += grad_logits_chunk.t() @ input_chunk   # sum over tokens
grad_bias   += grad_logits_chunk.sum(dim=0)          # sum over tokens

By the time backward runs we hold sum_i (g_i outer x_i). What we need is sum_i go_i * (g_i outer x_i). Those differ by a per-token weight applied inside the sum, and no scalar multiply on the already-summed result can recover it. The information is gone.

grad_input is different. It is per-row and never summed, so scaling it row-wise afterwards is valid.

Fix

For reduction="none", defer the gradient computation to backward, where grad_output is known, and fold it into each row of the logits gradient before projecting. mean and sum keep the existing compute-in-forward fast path, so the common case is unchanged.

The mechanism already existed: use_token_scaling does exactly this per-row scaling before the projection. This reuses it with the upstream gradient.

Cost is one extra chunked GEMM in backward for reduction="none" only. Memory is unchanged, still chunked, still never materializing the full BT x V logits.

Test

Adds test_correctness_reduction_none_per_token_grad_output, which backprops a non-uniform grad_output and compares grad_input, grad_weight and grad_bias against plain autograd.

The existing tests only ever call loss.backward(torch.ones_like(loss)). With uniform ones the scaling is a no-op, which is why this went unnoticed.

Before the fix (errors vs autograd):

	grad_output = ones	grad_output = rand
grad_input	4.8e-07	1.96
grad_weight	7.2e-07	1.49
grad_bias	4.8e-07	0.88

After the fix all three match autograd to 1e-4.

Rest of the suite is unchanged: 125 passed, plus the full test_cross_entropy.py.