[zheliuyu]

Summary

Align the Ascend backend with the contracts the transformer tests already encode.

• RMSNorm and FusedAddRMSNorm cast scalar eps and offset with tl.cast, and read the tensor dtype from the pointer. torch.compile specializes those scalars to constexpr, which has no .to(), and a host dtype constexpr is emitted as a bare name on NPU.
• Ascend cross entropy accepts loss_dtype and in_place. Per-token loss is accumulated in fp32 and cast back to the input dtype unless loss_dtype=torch.float32. A second backward is rejected after an in-place scale. Plain reduction="none" scales the existing Triton kernel's stored gradient once, so the result matches that kernel bitwise. The masked tail of a weighted, label-smoothed row is dropped before the sum.
• RoPE writes in place when the BSND view needs no copy.

The CUDA kernels are unchanged. Backends other than nvidia-triton and ascend-triton still reject loss_dtype and in_place.

Details

src/liger_kernel/transformers/cross_entropy.py only adds ascend-triton to the existing allowlist. That check runs only when loss_dtype is set or in_place is true.

Plain reduction="none" backward calls liger_kernel.ops.cross_entropy.cross_entropy_forward and multiplies its stored gradient by the upstream row scale, then rounds once. Mean and sum stay on the Ascend backward kernel. Recomputing softmax in the Ascend kernel differs by about one ulp from the stored gradient, which fails torch.equal.

Testing Done

• Hardware Type: Ascend910, CANN 9.1.0, torch 2.10.0, triton-ascend 3.6.0
• targeted pytest on the previously failing cases
• ruff check and ruff format --check on the six touched files
• make test
• make test-convergence

On one idle NPU:

• Compiled RMSNorm forward and backward, compiled FusedAddRMSNorm static and dynamic, fp16 and bf16, plus seqlen-1 RoPE: 9 passed.
• Previously failing cross entropy groups (loss_dtype, per-row grad, matches_torch, in-place second backward, retain_graph) and Qwen3.5 vision RoPE: 104 passed.
• Eager correctness for RMSNorm, FusedAddRMSNorm, cross entropy, and seqlen-1 RoPE: 60 passed, 16 skipped, 2 xpassed.