[NPU] Fix Ascend JSD UT failures and register LayerNorm dispatch (#1452)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Unit tests for JSD / FusedLinearJSD on Ascend failed, mainly due to
signature misalignment with the NVIDIA frontend, and invalid kernel
launches when all samples are ignored. The ascend-triton dispatcher
registration for LayerNorm is also added so that the NPU can
automatically select the UB-aware implementation.

JSD / FusedLinearJSD signature alignment: NVIDIA's LigerJSDFunction /
LigerFusedLinearJSDFunction accepts jsd_impl / jsd_mode (apply(...,
None, None) is used in tests). The Ascend kernel is self-contained and
does not use the inner dispatcher, so it simply accepts and ignores
these two parameters, and the number of backward return values is
supplemented accordingly; supports_inner_impl_dispatch is set to False
for JSD.

Short-circuit for all-ignored case: When n_non_ignore == 0, skip kernel
launch and directly return zero loss and zero gradients. This covers
cases such as test_correctness_all_ignored to prevent division-by-zero /
invalid grid.

LayerNorm registration: Add ops/backends/_ascend/layer_norm.py and
register it in declare_op_locations("layer_norm") within functional.py.
Otherwise, auto-selection only detects the NVIDIA-triton implementation
for CUDA and cannot use _ascend.ops.layer_norm.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
pytest tests/transformers
<img width="1255" height="383" alt="aa9edfdb912726cb88d8905e1b40f870"
src="https://github.com/user-attachments/assets/db596c84-8a94-480c-82a9-2a98351b4510"
/>

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: Atlas 800T A3(X86)
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

src/liger_kernel/functional.py

@@ -47,6 +47,7 @@ declare_op_locations(
         "liger_kernel.ops.backends._triton.layer_norm",
         "liger_kernel.ops.backends._cutedsl.layer_norm",
         "liger_kernel.ops.backends._cutile.layer_norm",
+        "liger_kernel.ops.backends._ascend.layer_norm",
     ),
 )
 

src/liger_kernel/ops/backends/_ascend/layer_norm.py

@@ -0,0 +1,42 @@
+"""Ascend Triton backend registration for ``layer_norm``.
+
+``LigerLayerNorm`` and ``functional.layer_norm`` dispatch through
+:func:`liger_kernel.backends.dispatch`. Without this adapter, auto-select only
+sees ``nvidia-triton`` (the CUDA-oriented kernel in ``ops.layer_norm``) and
+never the UB-aware NPU kernels in ``ops.backends._ascend.ops.layer_norm``.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from liger_kernel.backends import register_op
+from liger_kernel.ops.backends._ascend.dispatch_common import ASCEND_TRITON_CAPABILITY
+from liger_kernel.ops.backends._ascend.dispatch_common import ASCEND_TRITON_RANK
+from liger_kernel.ops.backends._ascend.dispatch_common import DEFAULT_TOLERANCES
+from liger_kernel.ops.backends._ascend.ops.layer_norm import LigerLayerNormFunction
+
+
+@register_op(
+    "layer_norm",
+    impl_name="ascend-triton",
+    capability=ASCEND_TRITON_CAPABILITY,
+    modes=("default",),
+    default_mode="default",
+    preference_rank=ASCEND_TRITON_RANK,
+    tolerances=DEFAULT_TOLERANCES,
+    notes="NPU-tuned LayerNorm (UB-aware fused 2D / row / tiled). Preferred on Ascend.",
+)
+def layer_norm_ascend(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    *,
+    mode: Optional[str] = None,
+) -> torch.Tensor:
+    if mode not in (None, "default"):
+        raise ValueError(f"Ascend layer_norm has only mode='default'; got mode={mode!r}.")
+    return LigerLayerNormFunction.apply(x, weight, bias, eps)

src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py

@@ -79,6 +79,15 @@ def fused_linear_jsd_forward(
     else:
         n_non_ignore = BT
 
+    if n_non_ignore == 0:
+        return (
+            torch.tensor(0.0, device=device, dtype=dtype),
+            grad_input,
+            grad_weight.to(student_weight.dtype)
+            if grad_weight is not None and accum_dtype is not None
+            else grad_weight,
+        )
+
     num_cores = get_npu_core_count()
 
     for chunk_id in range(num_chunks):
@@ -215,6 +224,8 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
         ignore_index: int = -100,
         temperature: float = 1.0,
         accum_dtype: Optional[torch.dtype] = None,
+        jsd_impl=None,
+        jsd_mode=None,
     ):
         """
         Args:
@@ -227,10 +238,12 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
             jsd_beta (float): coefficient beta of generalized JSD in the interval [0, 1]. It implements forward/reverse KL when beta equals 0 and 1 respectively. Default: `0.5`
             ignore_index (int): the index to ignore. Default: -100
             temperature (float): temperature in softmax function to control the output probability distribution. Default: `1.0`
+            jsd_impl / jsd_mode: accepted for NVIDIA JSD signature parity; unused on Ascend.
 
         Returns:
             loss (torch.Tensor): generalized JSD
         """
+        _ = (jsd_impl, jsd_mode)
         has_label = False
         if shift_labels is not None:
             assert shift_labels.shape == (teacher_input.shape[0],), (
@@ -263,4 +276,4 @@ class LigerFusedLinearJSDFunction(torch.autograd.Function):
     def backward(ctx, grad_output):
         (grad_input, grad_weight) = ctx.saved_tensors
         grad_input, grad_weight = fused_linear_jsd_backward(grad_output, grad_input, grad_weight)
-        return (grad_input, grad_weight, None, None, None, None, None, None, None)
+        return (grad_input, grad_weight, None, None, None, None, None, None, None, None, None)

src/liger_kernel/ops/backends/_ascend/ops/jsd.py

@@ -132,6 +132,9 @@ def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label):
     else:
         n_non_ignore = BT
 
+    if n_non_ignore == 0:
+        return torch.tensor(0.0, device=_input.device, dtype=_input.dtype), torch.zeros_like(_input)
+
     # Use NPU core count for grid size
     num_cores = get_npu_core_count()
     grid_size = min(num_cores, n_rows)
@@ -191,6 +194,8 @@ class LigerJSDFunction(torch.autograd.Function):
         shift_labels: Optional[torch.Tensor] = None,
         beta: float = 0.5,
         ignore_index: int = -100,
+        jsd_impl=None,
+        jsd_mode=None,
     ) -> torch.Tensor:
         """
         Args:
@@ -199,10 +204,12 @@ class LigerJSDFunction(torch.autograd.Function):
             shift_labels (Optional[torch.LongTensor]): indicator of next predicted vocab with shape (BT) where each value is in [0, V-1].
             beta (float): coefficient beta of generalized JSD in the interval [0, 1]. It implements forward/reverse KL when beta equals 0 and 1 respectively. Default: `0.5`
             ignore_index (int): the index to ignore. Default: -100
+            jsd_impl / jsd_mode: accepted for NVIDIA JSD signature parity; unused on Ascend.
 
         Returns:
             loss (torch.Tensor): generalized JSD
         """
+        _ = (jsd_impl, jsd_mode)
         has_label = False
         if shift_labels is not None:
             assert shift_labels.shape == (_input.shape[0],), (
@@ -226,4 +233,6 @@ class LigerJSDFunction(torch.autograd.Function):
             None,
             None,
             None,
+            None,
+            None,
         )