liger-kernel 0.6.4__py3-none-any.whl → 0.6.5__py3-none-any.whl

liger_kernel/ops/backends/_ascend/ops/tvd.py

@@ -0,0 +1,223 @@
+from typing import Literal
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.backends._ascend.ub_manager import compute_default_tiling_strategy
+from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import get_npu_core_count
+
+MAX_FUSED_SIZE = 65536 // 4
+
+REDUCTION_LITERAL = Literal["none", "sum", "mean", "batchmean"]
+
+
+@triton.jit
+def _tv_distance_kernel(
+    p_ptr,
+    p_stride,
+    q_ptr,
+    q_stride,
+    loss_ptr,
+    loss_stride,
+    grads_ptr,
+    grads_stride,
+    label_ptr,
+    ignore_index: tl.constexpr,
+    n_cols,  # V
+    total_rows: tl.constexpr,  # BT
+    BLOCK_SIZE: tl.constexpr,
+    HAS_LABEL: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
+    reduction: tl.constexpr = "batchmean",
+):
+    thread_id = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for pid in tl.range(thread_id, total_rows, num_threads, num_stages=NUM_STAGES):
+        p_row_ptr = p_ptr + pid * p_stride
+        q_row_ptr = q_ptr + pid * q_stride
+        loss_row_ptr = loss_ptr + pid * loss_stride
+        grads_row_ptr = grads_ptr + pid * grads_stride
+        label_row_ptr = label_ptr + pid
+
+        base_offsets = tl.arange(0, BLOCK_SIZE)
+
+        should_skip = False
+        if HAS_LABEL:
+            label = tl.load(label_row_ptr)
+            if label == ignore_index:
+                should_skip = True
+
+        if should_skip:
+            for i in range(0, n_cols, BLOCK_SIZE):
+                offsets = i + base_offsets
+                mask = offsets < n_cols
+                tl.store(grads_row_ptr + offsets, 0.0, mask=mask)
+                if reduction == "none":
+                    tl.store(loss_row_ptr + offsets, 0.0, mask=mask)
+        else:
+            loss_sum = 0.0
+            for i in range(0, n_cols, BLOCK_SIZE):
+                offsets = i + base_offsets
+                mask = offsets < n_cols
+
+                p = tl.load(p_row_ptr + offsets, mask=mask, other=0.0)
+                q = tl.load(q_row_ptr + offsets, mask=mask, other=0.0)
+
+                # TVD(P || Q) = 0.5 * |P - Q|
+                tv_loss = 0.5 * tl.abs(p - q)
+                grad_res = tl.where(p > q, 0.5, -0.5)
+
+                tl.store(grads_row_ptr + offsets, grad_res, mask=mask)
+
+                if reduction == "none":
+                    tl.store(loss_row_ptr + offsets, tv_loss, mask=mask)
+                else:
+                    loss_sum += tl.sum(tv_loss, axis=0)
+
+            if reduction != "none":
+                tl.store(loss_row_ptr, loss_sum)
+
+
+def tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label):
+    BT, V = p.shape
+
+    # TVD forward tiling strategy
+    # - In main loop (calculate loss and grad):
+    #   * p: BLOCK_Q elements
+    #   * q: BLOCK_Q elements
+    #   * tv_loss: BLOCK_Q elements
+    #   * grad_res: BLOCK_Q elements
+    #   * loss_sum: BLOCK_Q elements (when reduction != "none")
+    #   * Total: 4 * BLOCK_Q elements or 5 * BLOCK_Q elements when reduction != "none"
+    # - Since loss_sum is not necessarily used in every calculation,
+    # - and considering the consumption of other shared memory and the potential memory consumption of the HAS_LABEL loop.
+    # - Conservative estimate: 5 * BLOCK_Q * dtype_size * 8 bits
+    # - For safety, use: memory_multiplier=5.0 * BLOCK_SIZE * pad_hd * dtype_size * 8 bits
+    # - shapes: ((V,),)
+    # - tiling_dims: (0,) means first dimension of each shape can be tiled
+    # - Returns: ((block_size,),
+    shapes = ((V,),)
+    tile_shapes = compute_default_tiling_strategy(
+        safety_margin=0.80,
+        # In the TVD calculation, many data are implicitly converted to f32, so the size of f32 can be directly used.
+        dtype_size=4,
+        memory_multiplier=5.0,
+        shapes=shapes,
+        tiling_dims=(0,),
+    )
+
+    if tile_shapes is not None and len(tile_shapes) > 0 and len(tile_shapes[0]) > 0:
+        # Strategy returns ((block_size,),)
+        BLOCK_SIZE = tile_shapes[0][0]
+    else:
+        # Fallback to desired block size if no best practice found (no tiling needed)
+        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+
+    num_cores = get_npu_core_count()
+    grid = (min(num_cores, BT),)
+
+    out_size = (BT, V) if reduction == "none" else (BT,)
+
+    # The loss and grid accumulation on BF16 platform of NPU will have precision errors.
+    output_tensor = torch.zeros(out_size, device=p.device, dtype=torch.float32)
+    grads = torch.empty_like(p, dtype=torch.float32)
+
+    n_non_ignore = (shift_labels != ignore_index).sum().item() if has_label else BT
+
+    _tv_distance_kernel[grid](
+        p,
+        p.stride(0),
+        q,
+        q.stride(0),
+        output_tensor,
+        output_tensor.stride(0),
+        grads,
+        grads.stride(0),
+        shift_labels if has_label else torch.empty(1, device=p.device),
+        ignore_index,
+        V,
+        BT,
+        BLOCK_SIZE=BLOCK_SIZE,
+        HAS_LABEL=has_label,
+        NUM_STAGES=3 if BT < 4096 else 4,
+        reduction=reduction,
+    )
+
+    if reduction == "batchmean":
+        return output_tensor.sum() / n_non_ignore, grads / n_non_ignore
+    elif reduction == "sum":
+        return output_tensor.sum(dim=0), grads
+    elif reduction == "mean":
+        return output_tensor.sum() / (n_non_ignore * V), grads / (n_non_ignore * V)
+    else:
+        return output_tensor, grads
+
+
+def tvd_backward_triton(grad_output, grads):
+    # If this is the last layer, grad_output is 1.0. Skip the mul then.
+    if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        return grads
+
+    return grads * grad_output
+
+
+class LigerTVDLossFunction(torch.autograd.Function):
+    """
+    Class implementing the forward and backward pass for the Total Variation Distance Loss using Triton.
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        p: torch.Tensor,
+        q: torch.Tensor,
+        shift_labels: Optional[torch.Tensor] = None,
+        reduction: REDUCTION_LITERAL = "batchmean",
+        ignore_index: int = -100,
+    ) -> torch.Tensor:
+        """A forward pass for the Total Variation Distance Loss.
+
+        Args:
+            ctx: Torch autograd context
+            p (torch.Tensor): A tensor of shape (BT, V) containing the first distribution.
+            q (torch.Tensor): A tensor of shape (BT, V) containing the second distribution.
+            shift_labels (Optional[torch.Tensor]): A tensor of shape (BT,) containing the labels.
+            reduction (REDUCTION_LITERAL, optional): The reduction method to be applied. Defaults to "batchmean".
+            ignore_index (int, optional): The index to ignore during loss calculation. Defaults to -100.
+
+        Returns:
+            torch.Tensor: The computed Total Variation Distance Loss.
+        """
+        has_label = False
+        if shift_labels is not None:
+            assert shift_labels.shape == (p.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
+            shift_labels = shift_labels.contiguous()
+            has_label = True
+
+        loss, grads = tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label)
+        ctx.save_for_backward(grads)
+        return loss
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        """A backward pass for the Total Variation Distance Loss.
+
+        Args:
+            ctx: Torch autograd context
+            grad_output (torch.Tensor): The gradient of the loss with respect to the output.
+
+        Returns:
+            tuple[torch.Tensor, None, None, None, None]: The gradient of the loss with respect to the inputs.
+        """
+        (grads,) = ctx.saved_tensors
+        grads = tvd_backward_triton(grad_output, grads)
+
+        return grads, None, None, None, None

liger_kernel/ops/backends/_ascend/ub_manager.py

@@ -0,0 +1,367 @@
+"""
+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
+
+
+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 get_soc_spec.
+
+        Returns:
+            UB capacity in bits.
+
+        Raises:
+            RuntimeError: If UB capacity cannot be detected and no environment variable is set.
+        """
+        # Check environment variable first (in bits)
+        env_capacity = os.getenv("ASCEND_UB_CAPACITY_BITS")
+        if env_capacity is not None:
+            try:
+                capacity_bits = int(env_capacity)
+                if capacity_bits > 0:
+                    return capacity_bits
+            except ValueError:
+                pass
+
+        # Try to get from get_soc_spec (returns bytes, convert to bits)
+        if is_npu_available():
+            try:
+                from tbe.common.platform import get_soc_spec
+
+                # Query UB size (get_soc_spec returns size in bytes)
+                ub_size_bytes = get_soc_spec("UB_SIZE")
+
+                if ub_size_bytes is None or ub_size_bytes <= 0:
+                    raise ValueError(f"Invalid UB_SIZE from get_soc_spec: {ub_size_bytes}")
+
+                # Convert bytes to bits
+                ub_capacity_bits = ub_size_bytes * 8
+                return ub_capacity_bits
+
+            except ImportError:
+                raise RuntimeError(
+                    "Cannot import tbe.common.platform.get_soc_spec. "
+                    "Please ensure CANN environment variables are sourced "
+                    "(e.g., source /usr/local/Ascend/ascend-toolkit/set_env.sh)"
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to detect UB capacity from get_soc_spec: {e}. "
+                    "Please set ASCEND_UB_CAPACITY_BITS environment variable as fallback."
+                )
+
+        # If NPU is not available, raise error
+        raise RuntimeError(
+            "NPU is not available and UB capacity cannot be detected. "
+            "Please set ASCEND_UB_CAPACITY_BITS environment variable."
+        )
+
+
+# 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, 4.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
@@ -142,13 +143,16 @@ def liger_cross_entropy_kernel(
         block_max = tl.max(X_block)
 
         # Track argmax for accuracy computation
-        if RETURN_TOKEN_ACCURACY and block_max > m:
+        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
-            argmax_idx = tl.min(masked_offsets)
+            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
@@ -288,7 +292,13 @@ def liger_cross_entropy_kernel(
 # 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(