[NPU] optimize fused_linear_cross_entropy

[sunyi0505]

Summary

Move gradient computation out of forward, add a plain CE fast path, and optionally reuse saved logits to improve Ascend training performance while preserving correct grad_output handling.

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[sunyi0505]

Main Modifications

Introduce plain fast path for Ascend CE:

Add a new kernel: liger_cross_entropy_forward_kernel_plain. A lightweight kernel is used for vanilla cross-entropy cases without ce_weight, softcap, label smoothing, z loss, token metrics, or token scaling.

Remove upfront gradient computation in forward pass

The original implementation precomputed and cached grad_input / grad_weight / grad_bias during the forward pass. The new implementation only computes loss and essential statistics, and offloads all actual gradient calculation to LigerFusedLinearCrossEntropyFunction.backward().

Reimplement backward pass with chunked computation

Reuse cached forward logits if available; otherwise recompute input @ weight.T in the backward pass. The CE backward kernel produces grad_logits, and GEMM is then applied to derive grad_input, grad_weight, and grad_bias.

Add logits storage limit via _get_logits_save_limit_bytes()

Under the plain fast path, logits are cached when their size is within device memory budget to avoid redundant GEMM in the backward pass. The limit is defined as min(4GiB, total_memory / 8) with a hard minimum of 1GiB.

Adjust token scaling logic

The legacy implementation gathered valid targets separately. The new logic adopts safe_targets combined with masking, and caches the complete scaling_factors_full for backward usage.

Remove the element_mul_kernel path

The old backward logic multiplied precomputed forward gradients by the upstream grad_output. The new implementation consumes grad_output directly inside the CE backward kernel.

Motivation

The core objective is to optimize the performance and gradient semantics of fused linear cross entropy on Ascend.

The legacy implementation placed heavy gradient computation into the forward pass, making the forward stage undertake backward workload. It also failed to fully utilize the actual grad_output passed from autograd, especially for reduction="none" scenarios where upstream gradients are vector-valued. By moving gradient computation back to the backward pass, CE gradients are computed strictly based on the real incoming grad_output, yielding cleaner and more consistent semantics.

Furthermore, vanilla cross-entropy is a high-frequency path in language model training. The newly added plain fast path leverages larger block size, fewer kernel branches, and reduced Python/kernel launch overhead. Logits are reused within a controlled memory budget to lower backward recomputation cost. Overall, this change achieves an Ascend-optimized trade-off between memory footprint and GEMM / kernel launch efficiency.

forward:

backward:

full:

no_grad_forward:

ut:

[Tcc0403]

Do you mean reading after for loop might hurt the performance?

[sunyi0505]

Load x[y] directly outside the scan loop. Avoid embedding selection logic into the block loop to reduce undefined behavior risks and register pressure on Ascend, thereby improving performance.

[Tcc0403]

This fallback makes sense but shouldn't we inform user to use cross entropy instead of fused lienar cross entropy?

[sunyi0505]

Added logging in fused_linear_cross_entropy_forward:if the caller already has logits available, it is recommended to use the standalone cross-entropy operator directly.The warning is triggered only once per process to avoid redundant alerts at every training step.

[Tcc0403]

lgtm