liger-kernel-nightly 0.6.2.dev20251011154427__py3-none-any.whl → 0.6.4.dev20260107111351__py3-none-any.whl

liger_kernel/ops/backends/_ascend/ub_manager.py

@@ -0,0 +1,349 @@
+"""
+Unified Buffer (UB) Manager for Ascend NPU.
+
+This module provides UB capacity detection and tiling strategy computation
+for running Triton kernels on Ascend NPU. It automatically calculates
+optimal block sizes based on UB capacity constraints to prevent UB overflow.
+"""
+
+import os
+
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+import triton
+
+from liger_kernel.utils import is_npu_available
+
+# Default UB capacities for different NPU models (in bits)
+_DEFAULT_UB_CAPACITIES = {
+    "Ascend910B1": 2097152,  # ~256 KB
+    "Ascend910B4": 1572864,  # ~192 KB
+    "default": 2097152,  # ~256 KB
+}
+
+
+def _normalize_tiling_dims(tiling_dim: Union[int, Tuple[int, ...]]) -> set:
+    """
+    Normalize tiling dimension specification to a set of dimension indices.
+
+    Args:
+        tiling_dim: Either an int (single dimension) or tuple of ints (multiple dimensions).
+
+    Returns:
+        Set of dimension indices that can be tiled.
+    """
+    if isinstance(tiling_dim, int):
+        return {tiling_dim}
+    elif isinstance(tiling_dim, tuple):
+        return set(tiling_dim)
+    else:
+        return set()
+
+
+def _default_strategy(
+    ub_capacity_bits: int,
+    safety_margin: float,
+    dtype_size: int,
+    memory_multiplier: float,
+    shapes: Tuple[Tuple[int, ...], ...],
+    tiling_dims: Tuple[Union[int, Tuple[int, ...]], ...],
+) -> Tuple[int, ...]:
+    """
+    Default tiling strategy: calculate maximum safe block size based on UB capacity.
+
+    This is a unified strategy function that works for all kernels by abstracting
+    the memory calculation as: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits
+
+    Args:
+        ub_capacity_bits: UB capacity in bits
+        safety_margin: Safety margin as a float (e.g., 0.80 for 80%)
+        dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32)
+        memory_multiplier: Memory multiplier for estimating peak memory usage
+        shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes.
+            - For ROPE: ((n_q_head, hd), (n_kv_head, hd))
+            - For GEGLU: ((n_cols,),)
+        tiling_dims: Tuple specifying which dimensions can be tiled for each shape.
+            Each element can be:
+            - int: single dimension index (e.g., 0 for first dimension)
+            - tuple of ints: multiple dimensions that can be tiled together
+            - For ROPE: (0, 0) means first dimension of each shape can be tiled
+            - For GEGLU: (0,) means first dimension of the shape can be tiled
+            Length must match len(shapes).
+
+    Returns:
+        Tuple of maximum safe block sizes, one for each shape.
+        Each element is a power of 2.
+
+    Note:
+        For each shape, fixed dimensions (non-tiling) are multiplied together to get unit_param.
+        The final block size is computed in compute_default_tiling_strategy by taking
+        min(desired_block_size, max_safe_block_size) where desired_block_size = triton.next_power_of_2(original_dim).
+    """
+    if not shapes or not tiling_dims:
+        return ()
+
+    # Calculate max_safe_block_size for each tiling dimension
+    max_safe_sizes = []
+
+    for shape, tiling_dim in zip(shapes, tiling_dims):
+        # Normalize tiling_dim to a set of dimension indices
+        tiling_dim_set = _normalize_tiling_dims(tiling_dim)
+
+        # Validate tiling dimensions are within shape bounds
+        if not tiling_dim_set:
+            raise ValueError(
+                f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints."
+            )
+        if any(dim_idx < 0 or dim_idx >= len(shape) for dim_idx in tiling_dim_set):
+            raise ValueError(
+                f"Invalid tiling_dim: {tiling_dim} for shape {shape}. "
+                f"All dimension indices must be in range [0, {len(shape)})."
+            )
+
+        # Calculate unit_param: product of fixed (non-tiling) dimensions
+        unit_param = 1.0
+        for dim_idx, dim_size in enumerate(shape):
+            if dim_idx not in tiling_dim_set:
+                if dim_size <= 0:
+                    # Invalid dimension size, use conservative default
+                    unit_param = 1.0
+                    break
+                unit_param *= float(dim_size)
+
+        # Ensure unit_param is at least 1.0
+        if unit_param <= 0:
+            unit_param = 1.0
+
+        # Calculate maximum safe block size based on UB capacity
+        # Memory: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits
+        SAFE_UB_CAPACITY_BITS = int(ub_capacity_bits * safety_margin)
+
+        # Solve: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 <= SAFE_UB_CAPACITY_BITS
+        # BLOCK_SIZE <= SAFE_UB_CAPACITY_BITS / (memory_multiplier * unit_param * dtype_size * 8)
+        max_block_size = int(SAFE_UB_CAPACITY_BITS // (memory_multiplier * unit_param * dtype_size * 8))
+        max_block_size = max(1, max_block_size)
+
+        # Find largest power of 2 <= max_block_size
+        # Use triton.next_power_of_2(max_block_size + 1) // 2 to get the largest power of 2 <= max_block_size
+        safe_block_size = triton.next_power_of_2(max_block_size + 1) // 2
+        max_safe_sizes.append(safe_block_size)
+
+    return tuple(max_safe_sizes)
+
+
+class UBManager:
+    """
+    Unified Buffer Manager for Ascend NPU.
+
+    Provides UB capacity detection and management for Ascend NPU devices.
+    The UB capacity is used by tiling strategy functions to calculate optimal block sizes.
+    """
+
+    def __init__(self, ub_capacity_bits: Optional[int] = None):
+        """
+        Initialize UB Manager.
+
+        Args:
+            ub_capacity_bits: UB capacity in bits. If None, will be detected automatically.
+        """
+        self._npu_model = self._detect_npu_model()
+        self._ub_capacity_bits = ub_capacity_bits or self._detect_ub_capacity()
+
+    @property
+    def ub_capacity_bits(self) -> int:
+        """Get UB capacity in bits."""
+        return self._ub_capacity_bits
+
+    @property
+    def ub_capacity_bytes(self) -> int:
+        """Get UB capacity in bytes."""
+        return self._ub_capacity_bits // 8
+
+    @property
+    def npu_model(self) -> str:
+        """Get detected NPU model name."""
+        return self._npu_model
+
+    def _detect_npu_model(self) -> str:
+        """Detect NPU model from device properties."""
+        if not is_npu_available():
+            return "unknown"
+
+        try:
+            dev_props = torch.npu.get_device_properties(0)
+            # Try to get model name from device properties
+            return dev_props.name
+        except Exception:
+            pass
+
+        return "default"
+
+    def _detect_ub_capacity(self) -> int:
+        """
+        Detect UB capacity from environment variable or device properties.
+
+        Returns:
+            UB capacity in bits.
+        """
+        # Check environment variable first
+        env_capacity = os.getenv("ASCEND_UB_CAPACITY_BITS")
+        if env_capacity is not None:
+            try:
+                return int(env_capacity)
+            except ValueError:
+                pass
+
+        # Try to get from device properties
+        if is_npu_available():
+            try:
+                dev_props = torch.npu.get_device_properties(0)
+                if hasattr(dev_props, "ub_capacity_bits"):
+                    return dev_props.ub_capacity_bits
+            except Exception:
+                pass
+
+        # Fall back to model-based defaults
+        model = self._npu_model
+        return _DEFAULT_UB_CAPACITIES.get(model, _DEFAULT_UB_CAPACITIES["default"])
+
+
+# Global singleton instance
+_ub_manager: Optional[UBManager] = None
+
+
+def get_ub_manager() -> UBManager:
+    """Get global UB manager instance."""
+    global _ub_manager
+    if _ub_manager is None:
+        _ub_manager = UBManager()
+    return _ub_manager
+
+
+def compute_default_tiling_strategy(
+    safety_margin: float = 0.80,
+    dtype_size: Optional[int] = None,
+    memory_multiplier: Optional[float] = None,
+    shapes: Optional[Tuple[Tuple[int, ...], ...]] = None,
+    tiling_dims: Optional[Tuple[Union[int, Tuple[int, ...]], ...]] = None,
+) -> Optional[Tuple[Tuple[int, ...], ...]]:
+    """
+    Compute tiling strategy using the default strategy function.
+
+    This function directly calls the default strategy and computes the final
+    tiling result. All kernels use the same unified strategy function, so
+    there's no need for kernel_name-based lookup.
+
+    Args:
+        safety_margin: Safety margin as a float (e.g., 0.80 for 80%). Default is 0.80.
+        dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32).
+            Must be provided. If None or <= 0, defaults to 4 (float32).
+        memory_multiplier: Memory multiplier for estimating peak memory usage.
+            - For GEGLU: typically 10.0 for backward, 7.0 for forward
+            - For ROPE: typically 3.0
+            If None, defaults to 10.0 (conservative estimate).
+        shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes.
+            - For ROPE: ((n_q_head, hd), (n_kv_head, hd))
+            - For GEGLU: ((n_cols,),)
+            Can pass original shapes (will handle padding internally) or padded shapes.
+        tiling_dims: Tuple specifying which dimensions can be tiled for each shape.
+            Each element can be:
+            - int: single dimension index (e.g., 0 for first dimension)
+            - tuple of ints: multiple dimensions that can be tiled together
+            - For ROPE: (0, 0) means first dimension of each shape can be tiled
+            - For GEGLU: (0,) means first dimension of the shape can be tiled
+            Length must match len(shapes). Cannot be empty.
+
+    Returns:
+        Tuple of tiled shapes with same structure as input shapes.
+        Tiling dimensions are replaced with computed block sizes (power of 2),
+        while non-tiling dimensions are padded to next power of 2.
+        - For ROPE: ((block_size_q, pad_hd), (block_size_kv, pad_hd))
+        - For GEGLU: ((block_size,),)
+        Returns None if shapes or tiling_dims is None or empty.
+
+    Examples:
+        >>> # ROPE forward
+        >>> strategy = compute_default_tiling_strategy(
+        ...     safety_margin=0.90,
+        ...     dtype_size=4,
+        ...     memory_multiplier=3.0,
+        ...     shapes=((32, 128), (32, 128)),
+        ...     tiling_dims=(0, 0)
+        ... )
+        >>> # Returns: ((block_size_q, 128), (block_size_kv, 128))
+        >>> # GEGLU forward
+        >>> strategy = compute_default_tiling_strategy(
+        ...     safety_margin=0.80,
+        ...     dtype_size=2,
+        ...     memory_multiplier=7.0,
+        ...     shapes=((4096,),),
+        ...     tiling_dims=(0,)
+        ... )
+        >>> # Returns: ((block_size,),)
+    """
+    ub_manager = get_ub_manager()
+
+    if shapes is None or not shapes or tiling_dims is None or not tiling_dims:
+        return None
+
+    if len(shapes) != len(tiling_dims):
+        return None
+
+    if dtype_size is None or dtype_size <= 0:
+        dtype_size = 4  # Default to float32
+
+    if memory_multiplier is None or memory_multiplier <= 0:
+        memory_multiplier = 10.0  # Default conservative estimate
+
+    # Call strategy to get max_safe_block_size for each shape
+    max_supported = _default_strategy(
+        ub_manager.ub_capacity_bits,
+        safety_margin,
+        dtype_size,
+        memory_multiplier,
+        shapes,
+        tiling_dims,
+    )
+
+    if not max_supported or len(max_supported) != len(shapes):
+        return None
+
+    # Build result: same structure as shapes, with tiling dims replaced by computed block sizes
+    result = []
+    for shape, tiling_dim, max_safe in zip(shapes, tiling_dims, max_supported):
+        result_shape = list(shape)
+
+        # Normalize tiling_dim to a set of dimension indices
+        tiling_dim_set = _normalize_tiling_dims(tiling_dim)
+
+        # Validate tiling dimensions are within shape bounds
+        if not tiling_dim_set:
+            raise ValueError(
+                f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints."
+            )
+        if any(dim_idx < 0 or dim_idx >= len(result_shape) for dim_idx in tiling_dim_set):
+            raise ValueError(
+                f"Invalid tiling_dim: {tiling_dim} for shape {shape}. "
+                f"All dimension indices must be in range [0, {len(result_shape)})."
+            )
+
+        # Replace tiling dimensions with computed block sizes
+        # For each tiling dimension, compute: min(desired, max_safe)
+        for dim_idx in tiling_dim_set:
+            original_dim = result_shape[dim_idx]
+            desired = triton.next_power_of_2(original_dim)
+            final_val = min(desired, max_safe)
+            final_val = max(1, final_val)  # Ensure at least 1
+            result_shape[dim_idx] = final_val
+
+        # Pad non-tiling dimensions to next power of 2
+        for dim_idx, dim_size in enumerate(result_shape):
+            if dim_idx not in tiling_dim_set:
+                result_shape[dim_idx] = triton.next_power_of_2(dim_size)
+
+        result.append(tuple(result_shape))
+
+    return tuple(result)

liger_kernel/ops/backends/registry.py

@@ -0,0 +1,61 @@
+"""
+Vendor registry for Liger-Kernel multi-backend support.
+
+This module defines VendorInfo and the registry for vendor registration.
+Each vendor registers itself by calling register_vendor() in its __init__.py.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+# Dynamically get backends package path to avoid hardcoding
+_BACKENDS_PACKAGE = __name__.rsplit(".", 1)[0]  # "liger_kernel.ops.backends"
+
+
+@dataclass
+class VendorInfo:
+    """
+    Information about a chip vendor and its supported device.
+
+    Attributes:
+        vendor: Vendor name (e.g., "ascend", "intel", "nvidia")
+        device: Device type this vendor supports (e.g., "npu", "xpu")
+    """
+
+    vendor: str
+    device: str
+
+    @property
+    def module_path(self) -> str:
+        """Auto-generated module path based on vendor name."""
+        return f"{_BACKENDS_PACKAGE}._{self.vendor}.ops"
+
+
+# Registry mapping device types to their vendor info
+# Vendors register themselves via register_vendor()
+VENDOR_REGISTRY: dict[str, VendorInfo] = {}
+
+
+def register_vendor(vendor_info: VendorInfo) -> None:
+    """
+    Register a vendor's info in the global registry.
+
+    This should be called in each vendor's __init__.py to register itself.
+
+    Args:
+        vendor_info: VendorInfo instance to register
+    """
+    VENDOR_REGISTRY[vendor_info.device] = vendor_info
+
+
+def get_vendor_for_device(device: str) -> Optional[VendorInfo]:
+    """
+    Get the VendorInfo for a given device type.
+
+    Args:
+        device: Device type (e.g., "npu", "xpu")
+
+    Returns:
+        VendorInfo if found, None otherwise
+    """
+    return VENDOR_REGISTRY.get(device)

liger_kernel/ops/cross_entropy.py

@@ -10,8 +10,9 @@ from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import element_mul_kernel
 from liger_kernel.ops.utils import is_hip
 from liger_kernel.utils import infer_device
+from liger_kernel.utils import is_npu_available
 
-if compare_version("triton", operator.ge, "3.0.0"):
+if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
     try:
         # typical import path with dispatch available
         from triton.language.extra.libdevice import tanh
@@ -32,6 +33,8 @@ def liger_cross_entropy_kernel(
     loss_ptr,
     z_loss_ptr,
     loss_stride,
+    token_accuracy_ptr,
+    token_accuracy_stride,
     n_cols,
     n_non_ignore,
     sum_non_ignore_weight,
@@ -42,6 +45,7 @@ def liger_cross_entropy_kernel(
     reduction: tl.constexpr,  # set it as constexpr since reduction is always known at compile time
     softcap,
     RETURN_Z_LOSS: tl.constexpr,
+    RETURN_TOKEN_ACCURACY: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
     HAS_SOFTCAPPING: tl.constexpr,
@@ -60,6 +64,8 @@ def liger_cross_entropy_kernel(
     loss_ptr: Pointer to tensor to store the loss.
     z_loss_ptr: Pointer to tensor to store the z loss. No operation if RETURN_Z_LOSS is 0.
     loss_stride (int): The stride of the loss tensor.
+    token_accuracy_ptr: Pointer to tensor to store the per-token accuracy. No operation if RETURN_TOKEN_ACCURACY is 0.
+    token_accuracy_stride (int): The stride of the token accuracy tensor.
     n_cols (int): The number of columns in the input tensor.
     n_non_ignore (float): The number of non-ignored elements in the batch.
     sum_non_ignore_weight (float): The sum of non-ignored target's weights in the batch.
@@ -69,7 +75,8 @@ def liger_cross_entropy_kernel(
     lse_square_scale (float): The scaler of (logsumexp(_input)) ^ 2 adding to the loss for the stability of training.
     reduction (str): The string for the reduction to apply
     softcap (float): The upper threshold for scaling logits to the range (-softcap, +softcap).
-    RETURN_Z_LOSS (int): The boolean value to decide whether storing z loss to z_loss_ptr or not. It must be 0 or 1.
+    RETURN_Z_LOSS (int): The boolean value to decide whether to store z loss to z_loss_ptr or not. It must be 0 or 1.
+    RETURN_TOKEN_ACCURACY (int): The boolean value to decide whether to store per-token accuracy to token_accuracy_ptr or not. It must be 0 or 1.
     BLOCK_SIZE (int): The block size for Triton operations.
     HAS_WEIGHT (bool): The boolean value to determine whether assigning weight to each of the classes.
     HAS_SOFTCAPPING (bool): The boolean value to determine whether applying soft-capping or not.
@@ -92,11 +99,17 @@ def liger_cross_entropy_kernel(
         for i in range(0, n_cols, BLOCK_SIZE):
             X_offsets = i + tl.arange(0, BLOCK_SIZE)
             tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols)
+        # For ignored tokens, set token accuracy to 0
+        if RETURN_TOKEN_ACCURACY:
+            token_accuracy_ptr += program_id * token_accuracy_stride
+            tl.store(token_accuracy_ptr, 0.0)
         return
 
     loss_ptr += program_id * loss_stride
     if RETURN_Z_LOSS:
         z_loss_ptr += program_id * loss_stride
+    if RETURN_TOKEN_ACCURACY:
+        token_accuracy_ptr += program_id * token_accuracy_stride
 
     if HAS_WEIGHT:
         weight_y = tl.load(weight_ptr + y).cast(tl.float32)
@@ -107,6 +120,7 @@ def liger_cross_entropy_kernel(
     # 3. [Online softmax] first pass: find max + sum
     m = float("-inf")  # m is the max value. use the notation from the paper
     d = 0.0  # d is the sum. use the notation from the paper
+    argmax_idx = 0  # Track the index of the maximum value for token accuracy computation
     ori_X_y = tl.load(X_ptr + y).cast(tl.float32)  # we need to store the original value of X_y for the loss calculation
     if HAS_SOFTCAPPING:
         ori_X_y = softcap * tanh(ori_X_y / softcap)
@@ -127,6 +141,19 @@ def liger_cross_entropy_kernel(
         if HAS_SOFTCAPPING:
             X_block = softcap * tanh(X_block / softcap)
         block_max = tl.max(X_block)
+
+        # Track argmax for accuracy computation
+        if RETURN_TOKEN_ACCURACY:
+            # Find the index of the maximum value in this block
+            is_max_mask = X_block == block_max
+            # Mask out invalid indices with a value larger than n_cols
+            masked_offsets = tl.where(is_max_mask, X_offsets, n_cols)
+            # Get the first (smallest) index where max occurs
+            current_block_argmax_idx = tl.min(masked_offsets)
+
+            is_new_max = block_max > m
+            argmax_idx = tl.where(is_new_max, current_block_argmax_idx, argmax_idx)
+
         if label_smoothing > 0:
             # scale X beforehand to avoid overflow
             if HAS_WEIGHT:
@@ -256,12 +283,22 @@ def liger_cross_entropy_kernel(
     tl.store(loss_ptr, loss)
     if RETURN_Z_LOSS:
         tl.store(z_loss_ptr, z_loss)
+    if RETURN_TOKEN_ACCURACY:
+        # Store 1.0 if prediction is correct, 0.0 otherwise
+        is_correct = 1.0 if argmax_idx == y else 0.0
+        tl.store(token_accuracy_ptr, is_correct)
 
 
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
 # The optimal maximum block size depends on your hardware, your kernel, and your dtype
-MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2  # the best size we found by manually tuning
+# the best size we found by manually tuning on xpu and npu.
+if infer_device() == "xpu":
+    MAX_FUSED_SIZE = 4096
+elif infer_device() == "npu":
+    MAX_FUSED_SIZE = 2048
+else:
+    MAX_FUSED_SIZE = 65536 // 2
 
 
 def cross_entropy_forward(
@@ -274,8 +311,12 @@ def cross_entropy_forward(
     reduction,
     softcap,
     return_z_loss,
+    return_token_accuracy=False,
 ):
     assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}"
+    assert isinstance(return_token_accuracy, bool), (
+        f"return_token_accuracy must be True or False. Got: {return_token_accuracy}"
+    )
 
     BT, V = _input.shape
     n_rows = BT
@@ -285,6 +326,9 @@ def cross_entropy_forward(
     # unreduced loss
     loss_1d = torch.zeros(n_rows, dtype=_input.dtype, device=_input.device)
     z_loss_1d = torch.zeros(n_rows, dtype=_input.dtype, device=_input.device) if return_z_loss else None
+    token_accuracy_1d = (
+        torch.zeros(n_rows, dtype=torch.float32, device=_input.device) if return_token_accuracy else None
+    )
 
     target_mask = target != ignore_index
     n_non_ignore = target_mask.sum().item()
@@ -321,6 +365,10 @@ def cross_entropy_forward(
         loss_ptr=loss_1d,
         z_loss_ptr=z_loss_1d,
         loss_stride=loss_1d.stride(-1),  # always 1
+        token_accuracy_ptr=token_accuracy_1d,
+        token_accuracy_stride=token_accuracy_1d.stride(-1)
+        if return_token_accuracy
+        else 0,  # always 1 if accuracy is enabled
         n_cols=V,
         n_non_ignore=n_non_ignore,
         sum_non_ignore_weight=sum_non_ignore_weight,
@@ -331,6 +379,7 @@ def cross_entropy_forward(
         reduction=reduction,
         softcap=softcap,
         RETURN_Z_LOSS=return_z_loss,
+        RETURN_TOKEN_ACCURACY=return_token_accuracy,
         BLOCK_SIZE=BLOCK_SIZE,
         HAS_WEIGHT=True if weight is not None else False,
         HAS_SOFTCAPPING=True if softcap is not None else False,
@@ -343,11 +392,14 @@ def cross_entropy_forward(
     if reduction == "none":
         loss = loss_1d
         z_loss = z_loss_1d if return_z_loss else None
+        token_accuracy = token_accuracy_1d if return_token_accuracy else None
     else:
         loss = torch.sum(loss_1d)
         z_loss = torch.sum(z_loss_1d) if return_z_loss else None
+        # For accuracy, we compute the mean across all non-ignored tokens
+        token_accuracy = torch.sum(token_accuracy_1d) / n_non_ignore if return_token_accuracy else None
 
-    return loss, z_loss, _input
+    return loss, z_loss, token_accuracy, _input
 
 
 def cross_entropy_backward(_input, grad_output):
@@ -395,6 +447,7 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         reduction: str = "mean",
         softcap: Optional[float] = None,
         return_z_loss: bool = False,
+        return_token_accuracy: bool = False,
     ):
         """
         The forward pass of the Liger Cross Entropy loss.
@@ -409,12 +462,15 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
         reduction (str): The reduction to apply to the output: "none" | "mean | "sum".
         softcap (Optional[float]): The upper threshold for scaling logits to the range (-softcap, +softcap).
-        return_z_loss (bool): When `return_z_loss` is `True`, returns (loss, z_loss) instead of (loss, None). Default: `False`
+        return_z_loss (bool): When `return_z_loss` is `True`, returns (loss, z_loss, token_accuracy) instead of (loss, None, None). Default: `False`
+        return_token_accuracy (bool): When `return_token_accuracy` is `True`, computes and returns per-token accuracy without materializing logits. Default: `False`
 
         Returns:
-        tuple: A tuple with the compouted losses with respect to loss and z loss. The elements are tensors or None.
+        tuple: A tuple with the computed losses and accuracy: (loss, z_loss, token_accuracy). z_loss and token_accuracy are None if not requested.
         """
-        loss, z_loss, _input = cross_entropy_forward(
+        input_requires_grad = _input.requires_grad
+
+        loss, z_loss, token_accuracy, _input = cross_entropy_forward(
             _input,
             target,
             weight,
@@ -424,29 +480,35 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
             reduction,
             softcap,
             return_z_loss,
+            return_token_accuracy,
         )
         # TODO: investigation
         # If we don't detach the _input tensor, the memory will double
         # Not sure why but seems that there will be a time both grad and value exist but in different location
-        ctx.save_for_backward(_input.detach())
+        if input_requires_grad:
+            ctx.save_for_backward(_input.detach())
         ctx.return_z_loss = return_z_loss
+        ctx.return_token_accuracy = return_token_accuracy
 
-        return loss, z_loss
+        return loss, z_loss, token_accuracy
 
     @staticmethod
-    def backward(ctx, grad_output, grad_ouput2):
+    def backward(ctx, grad_output, grad_output2, grad_output3):
         """
         The backward pass of the Liger Cross Entropy loss.
 
         Parameters:
         ctx : The context object with saved tensors.
         grad_output (tensor): The tensor containing the gradient of the loss with respect to the output.
-        grad_output2 (tenosr): No use.
+        grad_output2 (tensor): No use. Gradient for z_loss (not used as z_loss is only for logging).
+        grad_output3 (tensor): No use. Gradient for token_accuracy (not used as token_accuracy is only for metrics).
         Returns:
         tuple: A tuple with the gradients with respect to the inputs. The elements are tensors or None.
         """
         if ctx.return_z_loss:
-            del grad_ouput2  # z_loss is only for logging
+            del grad_output2  # z_loss is only for logging
+        if ctx.return_token_accuracy:
+            del grad_output3  # token_accuracy is only for metrics
 
         (_input,) = ctx.saved_tensors
         _input = cross_entropy_backward(_input, grad_output)
@@ -460,4 +522,5 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
             None,
             None,
             None,
+            None,
         )

liger_kernel/ops/dyt.py

@@ -7,8 +7,10 @@ import triton.language as tl
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.ops.utils import infer_device
+from liger_kernel.utils import get_npu_multi_processor_count
+from liger_kernel.utils import is_npu_available
 
-if compare_version("triton", operator.ge, "3.0.0"):
+if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
     try:
         # typical import path with dispatch available
         from triton.language.extra.libdevice import tanh
@@ -125,7 +127,8 @@ def liger_dyt_bwd(dy, x, alpha, gamma, beta):
         NUM_SMS = torch.cuda.get_device_properties(x.device).multi_processor_count
     elif device == "xpu":
         NUM_SMS = torch.xpu.get_device_properties(x.device).gpu_subslice_count
-
+    elif device == "npu":
+        NUM_SMS = get_npu_multi_processor_count()
     da = torch.zeros(NUM_SMS, triton.cdiv(N, 512), dtype=torch.float32, device=x.device)
     dg = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device)
     db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None