liger-kernel-nightly 0.6.4.dev20260107111351__py3-none-any.whl → 0.6.4.dev20260116023519__py3-none-any.whl

liger_kernel/ops/backends/_ascend/ops/__init__.py

@@ -26,6 +26,9 @@ from liger_kernel.ops.backends._ascend.ops.rope import rope_forward
 from liger_kernel.ops.backends._ascend.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.backends._ascend.ops.swiglu import swiglu_backward
 from liger_kernel.ops.backends._ascend.ops.swiglu import swiglu_forward
+from liger_kernel.ops.backends._ascend.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops.backends._ascend.ops.tvd import tv_distance_forward_triton
+from liger_kernel.ops.backends._ascend.ops.tvd import tvd_backward_triton
 
 __all__ = [
     "LigerGELUMulFunction",
@@ -40,4 +43,7 @@ __all__ = [
     "LigerSiLUMulFunction",
     "swiglu_forward",
     "swiglu_backward",
+    "LigerTVDLossFunction",
+    "tv_distance_forward_triton",
+    "tvd_backward_triton",
 ]

liger_kernel/ops/backends/_ascend/ops/geglu.py

@@ -130,20 +130,26 @@ def geglu_forward(a, b):
     dtype_size = a.element_size()
     # GEGLU forward tiling strategy:
     # - Calculates maximum safe block size based on UB capacity
-    # - Memory analysis:
-    #   * Inputs: a, b
-    #   * Intermediates: a_cubed, tanh_arg, tanh_result, geglu_a
-    #   * Output: c
-    #   * Total: ~7x * BLOCK_SIZE * dtype_size
-    # - Uses memory_multiplier=7.0 * BLOCK_SIZE * dtype_size * 8 bits for safety
+    # - Memory analysis (only buffers that occupy UB, excluding temporary variables):
+    #   * Inputs: a_row (4 bytes, float32), b_row (dtype_size bytes)
+    #   * Output: c_row (dtype_size bytes)
+    #   * Temporary variables (a_cubed, tanh_arg, tanh_result, geglu_a) are optimized to registers
+    #     and don't occupy UB since they are only used once
+    #   * For float16: a_row(4) + b_row(2) + c_row(2) = 8 bytes/element, ratio = 8/2 = 4.0
+    #   * For float32: a_row(4) + b_row(4) + c_row(4) = 12 bytes/element, ratio = 12/4 = 3.0
+    # - Uses memory_multiplier=4.0 (float16) or 3.0 (float32) * BLOCK_SIZE * dtype_size * 8 bits
     # - shapes: ((n_cols,),)
     # - tiling_dims: (0,) means first dimension can be tiled
     # - Returns: ((block_size,),)
     shapes = ((n_cols,),)
+    if dtype_size == 2:
+        memory_multiplier = 4.0
+    else:
+        memory_multiplier = 3.0
     tile_shapes = compute_default_tiling_strategy(
         safety_margin=0.80,
         dtype_size=dtype_size,
-        memory_multiplier=7.0,
+        memory_multiplier=memory_multiplier,
         shapes=shapes,
         tiling_dims=(0,),
     )
@@ -187,18 +193,34 @@ def geglu_backward(a, b, dc):
     dtype_size = dc.element_size()
     # GEGLU backward tiling strategy:
     # - Calculates maximum safe block size based on UB capacity
-    # - Memory analysis:
-    #   * More intermediates for gradient computation compared to forward
-    #   * Total: ~10x * BLOCK_SIZE * dtype_size
-    # - Uses memory_multiplier=10.0 * BLOCK_SIZE * dtype_size * 8 bits for safety
+    # - Memory analysis: Peak memory usage occurs when executing line 103 (term1 calculation)
+    #   At this point, the following buffers simultaneously occupy UB:
+    #   1. dc_row = tl.load(dc + col_offsets, ...)  # dtype_size bytes
+    #   2. a_row = tl.load(a + col_offsets, ...).to(tl.float32)  # 4 bytes (float32)
+    #   3. b_row = tl.load(b + col_offsets, ...)  # dtype_size bytes
+    #   4. tanh_result = tanh(tanh_arg)  # 4 bytes (float32), used in lines 95, 103, 104
+    #   5. geglu_a = 0.5 * a_row * (1 + tanh_result)  # 4 bytes (float32), used in lines 96, 98
+    #   6. db_row = dc_row.cast(tl.float32) * geglu_a  # 4 bytes (float32, computed at line 98, stored at line 109)
+    #   Note: term1 (line 103) is a temporary variable optimized to registers and doesn't occupy UB
+    #   Temporary variables (a_cubed, tanh_arg, term1, tanh_sq, term2) are optimized to registers
+    #   and don't occupy UB since they are only used once
+    #   * For float16: dc_row(2) + a_row(4) + b_row(2) + tanh_result(4) + geglu_a(4) + db_row(4)
+    #     = 20 bytes/element, ratio = 20/2 = 10.0
+    #   * For float32: dc_row(4) + a_row(4) + b_row(4) + tanh_result(4) + geglu_a(4) + db_row(4)
+    #     = 24 bytes/element, ratio = 24/4 = 6.0
+    # - Uses memory_multiplier=10.0 (float16) or 6.0 (float32) * BLOCK_SIZE * dtype_size * 8 bits
     # - shapes: ((n_cols,),)
     # - tiling_dims: (0,) means first dimension can be tiled
     # - Returns: ((block_size,),)
     shapes = ((n_cols,),)
+    if dtype_size == 2:
+        memory_multiplier = 10.0
+    else:
+        memory_multiplier = 6.0
     tile_shapes = compute_default_tiling_strategy(
         safety_margin=0.80,
         dtype_size=dtype_size,
-        memory_multiplier=10.0,
+        memory_multiplier=memory_multiplier,
         shapes=shapes,
         tiling_dims=(0,),
     )

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

@@ -0,0 +1,221 @@
+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
+
+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,
+    reduction: tl.constexpr = "batchmean",
+):
+    thread_id = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for pid in range(thread_id, total_rows, num_threads):
+        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))
+
+    MAX_BATCH_PER_KERNEL = 65535  # The maximum processing capacity of each kernel in npu
+    if BT <= MAX_BATCH_PER_KERNEL:
+        grid = (BT,)
+    else:
+        grid = (MAX_BATCH_PER_KERNEL,)
+
+    out_size = (BT, V) if reduction == "none" else (BT,)
+    output_tensor = torch.zeros(out_size, device=p.device, dtype=torch.float32)
+    grads = torch.empty_like(p)
+
+    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,
+        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

@@ -241,7 +241,7 @@ def compute_default_tiling_strategy(
         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 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.

liger_kernel/ops/fused_add_rms_norm.py

@@ -8,6 +8,7 @@ import triton.language as tl
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import set_large_grf_mode
 from liger_kernel.ops.utils import torch_to_triton_dtype
 from liger_kernel.utils import get_npu_multi_processor_count
 from liger_kernel.utils import is_npu_available
@@ -162,23 +163,21 @@ def _fused_add_rms_norm_backward_kernel(
 
     dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_ptr += row_start * dY_row_stride
-    dX_ptr += row_start * dX_row_stride
-    if has_dS_out:
-        dS_out_ptr += row_start * dS_out_row_stride
-
-    X_ptr += row_start * X_row_stride
-    RSTD_ptr += row_start
-
     W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
     W_row = W_row + offset
 
-    for _ in range(row_start, row_end):
-        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
-        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
+    for row_idx in range(row_start, row_end):
+        dy_base = dY_ptr + row_idx * dY_row_stride
+        dx_base = dX_ptr + row_idx * dX_row_stride
+
+        x_base = X_ptr + row_idx * X_row_stride
+        rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
+
+        dY_row = tl.load(dy_base + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(x_base + col_offsets, mask=mask, other=0.0)
 
         # Get cached rms
-        rstd_row = tl.load(RSTD_ptr)
+        rstd_row = tl.load(rstd_base)
 
         X_row = X_row.to(tl.float32)
 
@@ -195,11 +194,11 @@ def _fused_add_rms_norm_backward_kernel(
         dX_row = rstd_row * m
 
         if has_dS_out:
-            dS_out_row = tl.load(dS_out_ptr + col_offsets, mask=mask, other=0.0)
+            ds_base = dS_out_ptr + row_idx * dS_out_row_stride
+            dS_out_row = tl.load(ds_base + col_offsets, mask=mask, other=0.0)
             dX_row += (rstd_row) * (
                 -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
             ) + dS_out_row
-            dS_out_ptr += dS_out_row_stride
         else:
             dX_row += (rstd_row) * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row)
 
@@ -210,12 +209,7 @@ def _fused_add_rms_norm_backward_kernel(
             # here X_row is already in fp32 (see previous if block)
             dW_row += dY_row * (X_row * rstd_row)
 
-        tl.store(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
-
-        dY_ptr += dY_row_stride
-        dX_ptr += dX_row_stride
-        X_ptr += X_row_stride
-        RSTD_ptr += RSTD_row_stride
+        tl.store(dx_base + col_offsets, dX_row.to(X_dtype), mask=mask)
 
     tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
@@ -254,7 +248,7 @@ def fused_add_rms_norm_forward(X, R, W, eps, offset, casting_mode):
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     # TODO: add _block_fused_add_rms_norm_forward_kernel
     _fused_add_rms_norm_forward_kernel[(n_rows,)](
@@ -314,7 +308,7 @@ def fused_add_rms_norm_backward(dY, dS_out, S, W, RSTD, offset, casting_mode, BL
     # XPU-specific optimization
     kernel_args = {}
     if S.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     # TODO: add _block_fused_add_rms_norm_backward_kernel
     _fused_add_rms_norm_backward_kernel[grid](

liger_kernel/ops/group_norm.py

@@ -6,6 +6,7 @@ import triton.language as tl
 
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.utils import infer_device
 from liger_kernel.utils import is_npu_available
 
 if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
@@ -18,7 +19,10 @@ if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
 else:
     from triton.language.math import rsqrt
 
-MAX_FUSED_SIZE = 65536
+if infer_device() == "npu":
+    MAX_FUSED_SIZE = 16384  # 8192
+else:
+    MAX_FUSED_SIZE = 65536
 
 
 @triton.jit
@@ -78,15 +82,14 @@ def _group_norm_forward_kernel(
     for channel_idx in tl.range(group_idx * channels_per_group, (group_idx + 1) * channels_per_group):
         W = tl.load(W_ptr + channel_idx)
         B = tl.load(B_ptr + channel_idx)
-        for i in range(0, hidden_size_per_channel, BLOCK_SIZE):
+        # Calculate channel offset within the group
+        channel_offset = (channel_idx - group_idx * channels_per_group) * hidden_size_per_channel
+        for i in tl.range(0, hidden_size_per_channel, BLOCK_SIZE):
             hidden_size_offsets = i + block_range
             mask = hidden_size_offsets < hidden_size_per_channel
-            X = tl.load(X_ptr + hidden_size_offsets, mask=mask, other=m)
+            X = tl.load(X_ptr + channel_offset + hidden_size_offsets, mask=mask, other=m)
             Y = (X - m) * rstd * W + B
-            tl.store(Y_ptr + hidden_size_offsets, Y, mask=mask)
-
-        X_ptr += hidden_size_per_channel
-        Y_ptr += hidden_size_per_channel
+            tl.store(Y_ptr + channel_offset + hidden_size_offsets, Y, mask=mask)
 
     tl.store(Mean_ptr + batch_idx * Mean_row_stride + group_idx * Mean_col_stride, m)
     tl.store(RSTD_ptr + batch_idx * RSTD_row_stride + group_idx * RSTD_col_stride, rstd)

liger_kernel/ops/kl_div.py

@@ -21,7 +21,12 @@ def get_num_warps(BLOCK_SIZE):
     return num_warps
 
 
-MAX_FUSED_SIZE = 65536 // 4  # 65536 // 4 or 8 works the best
+if infer_device() == "xpu":
+    MAX_FUSED_SIZE = 8192
+elif infer_device() == "npu":
+    MAX_FUSED_SIZE = 8192
+else:
+    MAX_FUSED_SIZE = 65536 // 4  # 65536 // 4 or 8 works the best
 
 REDUCTION_LITERAL = Literal["none", "sum", "mean", "batchmean"]
 
@@ -116,11 +121,7 @@ def _kldiv_kernel_backward(
 
 def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
     BT, V = y_pred.shape
-    BLOCK_SIZE = (
-        min(8192, triton.next_power_of_2(V))
-        if infer_device() == "xpu"
-        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-    )
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
     num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)
@@ -159,11 +160,7 @@ def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
 
 def kldiv_backward_triton(target, grad_output, new_grads, log_target):
     BT, V = target.shape
-    BLOCK_SIZE = (
-        min(8192, triton.next_power_of_2(V))
-        if infer_device() == "xpu"
-        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
-    )
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
     num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)

liger_kernel/ops/layer_norm.py

@@ -8,6 +8,8 @@ import triton.language as tl
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import set_large_grf_mode
+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") and not is_npu_available():
@@ -124,14 +126,14 @@ def _layer_norm_backward_kernel(
     w = tl.load(W_ptr + cols, mask=mask, other=0.0)
     w_f32 = w.to(tl.float32)
 
-    # Calculate pointers for this specific row
-    row_X_ptr = X_ptr + row_start * stride_x
-    row_DX_ptr = DX_ptr + row_start * stride_dx
-    row_DY_ptr = DY_ptr + row_start * stride_dy
-    row_Mean_ptr = Mean_ptr + row_start
-    row_RSTD_ptr = RSTD_ptr + row_start
+    for row_idx in range(row_start, row_end):
+        # Calculate pointers for this specific row
+        row_X_ptr = X_ptr + row_idx * stride_x
+        row_DX_ptr = DX_ptr + row_idx * stride_dx
+        row_DY_ptr = DY_ptr + row_idx * stride_dy
+        row_Mean_ptr = Mean_ptr + row_idx * stride_mean
+        row_RSTD_ptr = RSTD_ptr + row_idx * stride_rstd
 
-    for _ in range(row_start, row_end):
         # Load data for this row
         x = tl.load(row_X_ptr + cols, mask=mask, other=0.0)
         dy = tl.load(row_DY_ptr + cols, mask=mask, other=0.0)
@@ -160,12 +162,6 @@ def _layer_norm_backward_kernel(
         dW_row += dw
         db_row += db
 
-        row_X_ptr += stride_x
-        row_DX_ptr += stride_dx
-        row_DY_ptr += stride_dy
-        row_Mean_ptr += stride_mean
-        row_RSTD_ptr += stride_rstd
-
     tl.store(DW_ptr + row_block_id * stride_dw + cols, dW_row, mask=mask)
     tl.store(DB_ptr + row_block_id * stride_db + cols, db_row, mask=mask)
 
@@ -204,7 +200,7 @@ def layer_norm_forward(X, W, B, eps):
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     # Launch kernel with one thread block per row for optimal performance
     grid = (n_rows,)
@@ -254,6 +250,8 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":
         sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
+    elif X.device.type == "npu":
+        sm_count = get_npu_multi_processor_count()
 
     # fp32 for numerical stability especially.
     _DW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
@@ -272,7 +270,8 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
     kernel_args = {"num_warps": num_warps}
     # XPU-specific optimization
     if X.device.type == "xpu":
-        kernel_args.update({"grf_mode": "large", "num_warps": 32, "num_stages": 4})
+        kernel_args.update({"num_warps": 32, "num_stages": 4})
+        set_large_grf_mode(kernel_args)
 
     # Launch kernel with one thread block per row for optimal performance
     _layer_norm_backward_kernel[grid](
@@ -301,6 +300,7 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
     DX = DX.view(*shape)
     DW = _DW.sum(dim=0).to(W.dtype)
     DB = _DB.sum(dim=0).to(B.dtype)
+
     return DX, DW, DB
 
 

liger_kernel/ops/poly_norm.py

@@ -7,6 +7,7 @@ import triton.language as tl
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import set_large_grf_mode
 from liger_kernel.utils import get_npu_multi_processor_count
 from liger_kernel.utils import is_npu_available
 
@@ -140,20 +141,19 @@ def _poly_norm_backward_kernel(
     w1 = tl.load(W_ptr + 1).to(tl.float32)
     w2 = tl.load(W_ptr + 2).to(tl.float32)
 
-    dY_ptr += row_start * dY_row_stride
-    dX_ptr += row_start * dX_row_stride
-    X_ptr += row_start * X_row_stride
-    RSTD_ptr += row_start * RSTD_row_stride
+    for row_idx in range(row_start, row_end):
+        dy_base = dY_ptr + row_idx * dY_row_stride
+        x_base = X_ptr + row_idx * X_row_stride
+        dx_base = dX_ptr + row_idx * dX_row_stride
+        rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
 
-    for _ in range(row_start, row_end):
-        # Load input and gradient
-        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
-        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
+        dY_row = tl.load(dy_base + col_offsets, mask=mask, other=0.0).to(tl.float32)
+        X_row = tl.load(x_base + col_offsets, mask=mask, other=0.0).to(tl.float32)
 
         # Load cached rstd values
-        rstd_3 = tl.load(RSTD_ptr + 0).to(tl.float32)
-        rstd_2 = tl.load(RSTD_ptr + 1).to(tl.float32)
-        rstd_1 = tl.load(RSTD_ptr + 2).to(tl.float32)
+        rstd_3 = tl.load(rstd_base + 0).to(tl.float32)
+        rstd_2 = tl.load(rstd_base + 1).to(tl.float32)
+        rstd_1 = tl.load(rstd_base + 2).to(tl.float32)
 
         # Compute powers
         X_pow3 = X_row * X_row * X_row
@@ -190,13 +190,7 @@ def _poly_norm_backward_kernel(
         dX_row = grad_x_3 + grad_x_2 + grad_x_1
 
         # Store gradient
-        tl.store(dX_ptr + col_offsets, dX_row, mask=mask)
-
-        # Update pointers
-        dY_ptr += dY_row_stride
-        dX_ptr += dX_row_stride
-        X_ptr += X_row_stride
-        RSTD_ptr += RSTD_row_stride
+        tl.store(dx_base + col_offsets, dX_row, mask=mask)
 
     # Store accumulated gradients (scalars)
     tl.store(dW_ptr + row_block_id * dW_row_stride + 0, dW0_acc)
@@ -239,7 +233,7 @@ def poly_norm_forward(X, W, B, eps=1e-6):
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     # Launch kernel
     _poly_norm_forward_kernel[(n_rows,)](
@@ -310,7 +304,7 @@ def poly_norm_backward(dY, X, W, RSTD, BLOCK_SIZE, num_warps, in_place):
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     # Launch backward kernel
     _poly_norm_backward_kernel[grid](

liger_kernel/ops/rms_norm.py

@@ -20,6 +20,7 @@ import triton.language as tl
 from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import set_large_grf_mode
 from liger_kernel.ops.utils import torch_to_triton_dtype
 from liger_kernel.utils import get_npu_multi_processor_count
 from liger_kernel.utils import is_npu_available
@@ -70,11 +71,11 @@ def _rms_norm_forward_kernel(
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    Y_ptr += row_idx * Y_row_stride
-    X_ptr += row_idx * X_row_stride
-    RSTD_ptr += row_idx * RSTD_row_stride
+    y_base = Y_ptr + row_idx * Y_row_stride
+    x_base = X_ptr + row_idx * X_row_stride
+    rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
 
-    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
+    X_row = tl.load(x_base + col_offsets, mask=mask, other=0)
     X_row_dtype = X_row.dtype
     if elementwise_affine:
         W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
@@ -99,7 +100,7 @@ def _rms_norm_forward_kernel(
     # We can save time by caching rms with minimal memory overhead
     # because rms is much smaller compared to X_row, as rms is for each row.
     # However, on the computation side, it can save 4 operations (*, sum, /, sqrt).
-    tl.store(RSTD_ptr, rstd)
+    tl.store(rstd_base, rstd)
 
     X_row = X_row * rstd
 
@@ -115,7 +116,7 @@ def _rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_GEMMA:
         Y_row = Y_row.to(X_row_dtype)
 
-    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
+    tl.store(y_base + col_offsets, Y_row, mask=mask)
 
 
 @triton.jit
@@ -155,22 +156,22 @@ def _rms_norm_backward_kernel(
     if elementwise_affine:
         dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_ptr += row_start * dY_row_stride
-    dX_ptr += row_start * dX_row_stride
-
-    X_ptr += row_start * X_row_stride
-    RSTD_ptr += row_start
-
     if elementwise_affine:
         W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
         W_row = W_row + offset
 
-    for _ in range(row_start, row_end):
-        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
-        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
+    for row_idx in range(row_start, row_end):
+        dy_base = dY_ptr + row_idx * dY_row_stride
+        dx_base = dX_ptr + row_idx * dX_row_stride
+
+        x_base = X_ptr + row_idx * X_row_stride
+        rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
+
+        dY_row = tl.load(dy_base + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(x_base + col_offsets, mask=mask, other=0.0)
 
         # Get cached rms
-        rstd_row = tl.load(RSTD_ptr)
+        rstd_row = tl.load(rstd_base)
 
         X_row = X_row.to(tl.float32)
 
@@ -205,12 +206,7 @@ def _rms_norm_backward_kernel(
                 # here X_row is already in fp32 (see previous if block)
                 dW_row += dY_row * (X_row * rstd_row)
 
-        tl.store(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
-
-        dY_ptr += dY_row_stride
-        dX_ptr += dX_row_stride
-        X_ptr += X_row_stride
-        RSTD_ptr += RSTD_row_stride
+        tl.store(dx_base + col_offsets, dX_row.to(X_dtype), mask=mask)
 
     if elementwise_affine:
         tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
@@ -441,7 +437,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
     if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
         _rms_norm_forward_kernel[(n_rows,)](
             Y,
@@ -521,7 +517,7 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     # XPU-specific optimization
     kernel_args = {}
     if X.device.type == "xpu":
-        kernel_args["grf_mode"] = "large"
+        set_large_grf_mode(kernel_args)
 
     if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode:
         _rms_norm_backward_kernel[grid](

liger_kernel/ops/utils.py

@@ -139,3 +139,14 @@ def get_npu_core_count(default: int = 20) -> int:
         return int(props.get("num_vectorcore", default))
     except Exception:
         return default
+
+
+def set_large_grf_mode(kernel_args: dict):
+    """Set large GRF mode for XPU devices."""
+    # On XPU triton installed along with pytorch-xpu will be called `pytorch-triton-xpu`,
+    # triton XPU installed from source will be called `triton`.
+    if compare_version("pytorch-triton-xpu", operator.ge, "3.6.0") or compare_version("triton", operator.ge, "3.6.0"):
+        kernel_args["grf_mode"] = "256"
+    else:
+        # API was changed in https://github.com/intel/intel-xpu-backend-for-triton/pull/5430
+        kernel_args["grf_mode"] = "large"

liger_kernel/transformers/__init__.py

@@ -33,6 +33,7 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel  # noqa: F401
     from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_exaone4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_falcon_h1  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # noqa: F401
@@ -136,6 +137,7 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_smolvlm",
         "apply_liger_kernel_to_hunyuan_v1_dense",
         "apply_liger_kernel_to_hunyuan_v1_moe",
+        "apply_liger_kernel_to_exaone4",
     }
 
     if name in monkey_patch_symbols:
@@ -214,5 +216,6 @@ if _TRANSFORMERS_AVAILABLE:
             "apply_liger_kernel_to_smolvlm",
             "apply_liger_kernel_to_hunyuan_v1_dense",
             "apply_liger_kernel_to_hunyuan_v1_moe",
+            "apply_liger_kernel_to_exaone4",
         ]
     )

liger_kernel/transformers/model/exaone4.py

@@ -0,0 +1,136 @@
+from typing import List
+from typing import Optional
+from typing import Union
+
+import torch
+
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
+from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    **kwargs,
+) -> LigerCausalLMOutputWithPast:
+    r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+    Returns:
+
+    Example:
+
+    ````python
+    >>> from transformers import AutoTokenizer, Exaone4ForCausalLM
+
+    >>> model = Exaone4ForCausalLM.from_pretrained("LGAI-EXAONE/EXAONE-4.0-1.2B")
+    >>> tokenizer = AutoTokenizer.from_pretrained("LGAI-EXAONE/EXAONE-4.0-1.2B")
+
+    >>> prompt = "Hey, are you conscious? Can you talk to me?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+    ```"""
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
+
+    shift_labels = kwargs.pop("shift_labels", None)
+    # Remove output-control parameters that shouldn't be passed to loss functions
+    kwargs.pop("return_dict", None)
+    logits = None
+    loss = None
+    token_accuracy = None
+
+    if skip_logits and labels is None and shift_labels is None:
+        raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+    if skip_logits is None:
+        # By default, if in training mode, don't materialize logits
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    # Compute loss
+    if skip_logits:
+        result = LigerForCausalLMLoss(
+            hidden_states=kept_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            shift_labels=shift_labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+        loss, _, token_accuracy = unpack_cross_entropy_result(result)
+
+    else:
+        logits = self.lm_head(kept_hidden_states)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        output = ((loss,) + output) if loss is not None else output
+        output = output + (token_accuracy,) if token_accuracy is not None else output
+        return output
+
+    # Return custom output class with accuracy field
+    return LigerCausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        token_accuracy=token_accuracy,
+    )

liger_kernel/transformers/model/gemma2.py

@@ -7,7 +7,7 @@ from typing import Union
 import torch
 
 from torch.nn import CrossEntropyLoss
-from transformers.cache_utils import HybridCache
+from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.utils.deprecation import deprecate_kwarg
 
@@ -24,7 +24,7 @@ def lce_forward_deprecated(
     input_ids: torch.LongTensor = None,
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[HybridCache] = None,
+    past_key_values: Optional[Cache] = None,
     inputs_embeds: Optional[torch.FloatTensor] = None,
     labels: Optional[torch.LongTensor] = None,
     use_cache: Optional[bool] = None,
@@ -149,7 +149,7 @@ def lce_forward(
     input_ids: torch.LongTensor = None,
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[HybridCache] = None,
+    past_key_values: Optional[Cache] = None,
     inputs_embeds: Optional[torch.FloatTensor] = None,
     labels: Optional[torch.LongTensor] = None,
     use_cache: Optional[bool] = None,

liger_kernel/transformers/model/gemma3.py

@@ -6,10 +6,8 @@ import torch
 import torch.nn as nn
 
 from transformers.cache_utils import Cache
-from transformers.cache_utils import HybridCache
 from transformers.utils import logging
 
-from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
 from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
 from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast
@@ -23,7 +21,7 @@ def causal_forward(
     input_ids: torch.LongTensor = None,
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[HybridCache] = None,
+    past_key_values: Optional[Cache] = None,
     inputs_embeds: Optional[torch.FloatTensor] = None,
     labels: Optional[torch.LongTensor] = None,
     use_cache: Optional[bool] = None,
@@ -269,8 +267,15 @@ def multimodal_forward(
         shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size)
         shift_labels = shift_labels.view(-1).to(hidden_device)
 
-        lce = LigerFusedLinearCrossEntropyLoss()
-        result = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+        result = LigerForCausalLMLoss(
+            hidden_states=shift_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=shift_labels,
+            hidden_size=self.config.text_config.hidden_size,
+            shift_labels=shift_labels,
+            final_logit_softcapping=getattr(self.config.text_config, "final_logit_softcapping", None),
+            **lm_kwargs,
+        )
         loss, _, token_accuracy = unpack_cross_entropy_result(result)
 
     else:

liger_kernel/transformers/model/loss_utils.py

@@ -1,3 +1,5 @@
+import inspect
+
 from typing import Optional
 from typing import Tuple
 
@@ -71,6 +73,10 @@ def LigerForCausalLMLoss(
     return_token_accuracy: bool = False,
     **kwargs,
 ):
+    # Filter out inapplicable kwargs to liger_fused_linear_cross_entropy
+    applicable_params = inspect.signature(F.liger_fused_linear_cross_entropy).parameters
+    kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
+
     # Skip upcast since intermediate values for the loss are all fp32 in kernel
     if shift_labels is None:
         # Shift so that token < n predict n

liger_kernel/transformers/monkey_patch.py

@@ -2821,6 +2821,83 @@ def apply_liger_kernel_to_hunyuan_v1_moe(
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
 
 
+def apply_liger_kernel_to_exaone4(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace EXAONE4 models.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is True.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+        loaded. Default is None.
+    """
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.exaone4 import modeling_exaone4
+    from transformers.models.exaone4.modeling_exaone4 import Exaone4Model
+
+    from liger_kernel.transformers.model.exaone4 import lce_forward as exaone4_lce_forward
+
+    if rope:
+        modeling_exaone4.apply_rotary_pos_emb = liger_rotary_pos_emb
+
+    if rms_norm:
+        # EXAONE4 requires in_place=False to avoid gradient issues
+        class Exaone4LigerRMSNorm(LigerRMSNorm):
+            def __init__(self, hidden_size, eps=1e-6, **kwargs):
+                super().__init__(hidden_size, eps, **kwargs)
+                self.in_place = False
+
+        modeling_exaone4.Exaone4RMSNorm = Exaone4LigerRMSNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        if model is not None:
+            model.forward = MethodType(exaone4_lce_forward, model)
+        else:
+            modeling_exaone4.Exaone4ForCausalLM.forward = exaone4_lce_forward
+
+    if swiglu:
+        modeling_exaone4.Exaone4MLP = LigerSwiGLUMLP
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+
+        # get the base model from the model instance
+        base_model: Exaone4Model = getattr(model, model.base_model_prefix, model)
+
+        if rms_norm:
+            _patch_rms_norm_module(base_model.norm, in_place=False)
+        for decoder_layer in base_model.layers:
+            if swiglu:
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
+            if rms_norm:
+                _patch_rms_norm_module(decoder_layer.post_attention_layernorm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.post_feedforward_layernorm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.self_attn.q_norm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.self_attn.k_norm, in_place=False)
+
+
 # Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
 MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "gemma": apply_liger_kernel_to_gemma,
@@ -2862,6 +2939,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "smolvlm": apply_liger_kernel_to_smolvlm,
     "hunyuan_v1_dense": apply_liger_kernel_to_hunyuan_v1_dense,
     "hunyuan_v1_moe": apply_liger_kernel_to_hunyuan_v1_moe,
+    "exaone4": apply_liger_kernel_to_exaone4,
 }
 
 

liger_kernel/transformers/tiled_mlp.py

@@ -57,11 +57,7 @@ class LigerTiledGEGLUMLP(nn.Module):
         Returns:
             Output tensor of the same shape as input
         """
-        compute_params = [
-            self.gate_proj.weight,
-            self.up_proj.weight,
-            self.down_proj.weight,
-        ]
+        compute_params = [p for p in self.parameters() if p.requires_grad]
 
         return apply_tiled_mlp(
             fn=self._mlp_forward,
@@ -118,11 +114,7 @@ class LigerTiledSwiGLUMLP(nn.Module):
         Returns:
             Output tensor of the same shape as input
         """
-        compute_params = [
-            self.gate_proj.weight,
-            self.up_proj.weight,
-            self.down_proj.weight,
-        ]
+        compute_params = [p for p in self.parameters() if p.requires_grad]
 
         return apply_tiled_mlp(
             fn=self._mlp_forward,

{liger_kernel_nightly-0.6.4.dev20260107111351.dist-info → liger_kernel_nightly-0.6.4.dev20260116023519.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.4.dev20260107111351
+Version: 0.6.4.dev20260116023519
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.6.4.dev20260107111351.dist-info → liger_kernel_nightly-0.6.4.dev20260116023519.dist-info}/RECORD

@@ -19,42 +19,43 @@ liger_kernel/chunked_loss/simpo_loss.py,sha256=fy2w8KbhMrBv7b1jdIeH3bBFxY52bPQPZ
 liger_kernel/ops/__init__.py,sha256=F3m9qlXbgttykKEBsrMFf1WyK_0H8CKqLuDnFRR-cvc,7237
 liger_kernel/ops/cross_entropy.py,sha256=DnXFRZ9TGN1SnEo8xGBFFPLNQaen8aLVNPJ1em-LbK4,22910
 liger_kernel/ops/dyt.py,sha256=4XmkCCZaPPM8Tl4QHo6vSF2m68jrwsnjucrbyOJvZpM,5628
-liger_kernel/ops/fused_add_rms_norm.py,sha256=lvwrLsKvoAQqS9KatgBkAyy0Xdecado-g0rvXYXaBak,14237
+liger_kernel/ops/fused_add_rms_norm.py,sha256=E4SqFDw13ixd6S3DMhB1HlvtxAfuPL_DiHkgpk3exCI,14174
 liger_kernel/ops/fused_linear_cross_entropy.py,sha256=1gx2qljre9PVc861iknFnNCGC-P35D2w1cc_yMDO9ow,16239
 liger_kernel/ops/fused_linear_jsd.py,sha256=CSoprxb-YcJy-YUKiTcYkxN8sb9h2kdk_iHuncvSV5c,9683
 liger_kernel/ops/fused_neighborhood_attention.py,sha256=vPi5xbnh6wxyZehaqo6Tuilqo2fN5SGDiONjnNmIKqs,35556
 liger_kernel/ops/geglu.py,sha256=-ruMACDsFH1YsAak6BGvZ0ktLGIrBE6yGF0dAyR82UU,4307
-liger_kernel/ops/group_norm.py,sha256=zoy-TcNkYtKGmGhTFJmnyiG_4Es4ZphpqP8jtUSI6-I,10912
+liger_kernel/ops/group_norm.py,sha256=7BqYIP5-HQCdvHKMJlA6jCQoYKZjbtsoD9-eXld5qzk,11133
 liger_kernel/ops/grpo_loss.py,sha256=2SyOujtF9I3xiNo4wFf4s6MeiDotE_qeYfRWgj_bOBE,9573
 liger_kernel/ops/jsd.py,sha256=onHp5T3MbvJaVz5Vup7Ww6EQp_HTaZeayTjJk6FgQMY,7042
-liger_kernel/ops/kl_div.py,sha256=ZjGdDLKWksHT9dZ0xF_TDgAkj5cuMTwwT5tr9E-_24o,8734
-liger_kernel/ops/layer_norm.py,sha256=-4UEyko9eKgBi5LNmfdEU2hTpJOWVnEy5iYjJkMvHmk,10598
+liger_kernel/ops/kl_div.py,sha256=MZZb7eAPMXlydYVV4uL9aTytXFkdQdp-jmiDw9tC0pg,8652
+liger_kernel/ops/layer_norm.py,sha256=D1qPDn0HVHfyOmNHQyMDKv7f_JEnFsFxzHgfq9B4rI8,10696
 liger_kernel/ops/llama4_rope.py,sha256=-aqdZzllklTN8b9--e-TsWY_ntGCN8-tyseT4x0bd8s,8223
 liger_kernel/ops/multi_token_attention.py,sha256=Oz_RXDp-OSS_R_HuGmaETHdAJ7Toda_70OfE7TXMUlY,7645
-liger_kernel/ops/poly_norm.py,sha256=5IdJEZnbbhblkL_X8UhSD4A2CooQbOAZJw8nAekWNs4,11372
+liger_kernel/ops/poly_norm.py,sha256=BBwdOtSzW02W-c-UAN8pzn2vAU-AM3gCsWqZnSE5zf4,11288
 liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU,8481
-liger_kernel/ops/rms_norm.py,sha256=r97gpPmhbKz9qrBjxUEX0XP04aYu4psJeLe3KnhPZyo,21852
+liger_kernel/ops/rms_norm.py,sha256=bd5ZAdiqh2iO7a7FdwWH7woslJEVyPlDKXSoUqDZ3GQ,21874
 liger_kernel/ops/rope.py,sha256=v-7JHRrv-5ImoROkpKfl30WwWI4qTa2tAl7zQeB4ml4,8956
 liger_kernel/ops/softmax.py,sha256=tgORx6MK1IDDtZKqGarj0IPIVjqAIEUXXYPiinhRdtI,5864
 liger_kernel/ops/sparsemax.py,sha256=AeWe1xgkHJFEKWTj2vu_0hj7LztGvjqXAps-QTpCY0U,5087
 liger_kernel/ops/swiglu.py,sha256=D7nd4u_LInwsIRNCDdY77lqnTz8-W5dJrpEAt8zEO_A,3033
 liger_kernel/ops/tiled_mlp.py,sha256=eyMFsFFgHch8a_6R6IYRG24_jqKg5GF_BQUoQuAG8SY,4529
 liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
-liger_kernel/ops/utils.py,sha256=Xu6MJ2-lbp4hSmI0JGImKguKU0KqWnFQDgQwOxSieyc,4360
+liger_kernel/ops/utils.py,sha256=90V8P0ElZeBathDhmIKm_506Nhrsr1ojO0qRl53_Tn0,4909
 liger_kernel/ops/backends/README.md,sha256=ZP59UUqD1WW8LwM5Y-cTpSM-Dtgdp8Wku2mE9kqAc2E,4185
 liger_kernel/ops/backends/__init__.py,sha256=-mgef3cHfDFeL5NbXbq1TI7ngCahE9qqL3aMaHnXvis,629
 liger_kernel/ops/backends/registry.py,sha256=yJa_Sh2FZ__iPCIU8h2nOQbnsFQh1I-_czROLtb1uQM,1637
 liger_kernel/ops/backends/_ascend/__init__.py,sha256=6n0keOX9H-kLadBdVZlx-Ce0ZLVJvLiEfR-9-uxmYUk,221
 liger_kernel/ops/backends/_ascend/ascend-ub-manager-design.md,sha256=FVXHSO1KY4ZFxCAE5r4hOYB2Q8ANyrJZ7WnFJ_GeQOA,19605
-liger_kernel/ops/backends/_ascend/ub_manager.py,sha256=3h7sncZk00veBJS37a01YPt1SLeAxJj5N3lPdv1wXAk,13174
-liger_kernel/ops/backends/_ascend/ops/__init__.py,sha256=R1iS9R0EtmGbrN0cSkIiRtZouVl7ndiPVZJIoEALb7s,1748
-liger_kernel/ops/backends/_ascend/ops/geglu.py,sha256=hs1Cdhw0pbgZFiK1srLuo8DCe8jtnmhjm5SS2vw8-0M,8421
+liger_kernel/ops/backends/_ascend/ub_manager.py,sha256=3Utke2Dwx9huB0Qoch1KU2CXKN3JS5DbP9_JusIbfQU,13174
+liger_kernel/ops/backends/_ascend/ops/__init__.py,sha256=N41VgPn8D_YJpHez1-UEYTtA-JZxpERmAzN7WcDfE2U,2067
+liger_kernel/ops/backends/_ascend/ops/geglu.py,sha256=M3YFE44UREf91PtOvY0X_GZouUxeeDCy3GmXDrvRLQk,10131
 liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py,sha256=pUYcstJ4FuzDTkuhmQaO3U9gcVQoNCpzuwwUdtES5hM,11015
 liger_kernel/ops/backends/_ascend/ops/rope.py,sha256=nOwtm6_eSnzDjl2S-jvGpwHrumAOgWfr5pNg6SL3R2k,10842
 liger_kernel/ops/backends/_ascend/ops/swiglu.py,sha256=yrbEgIgeCZyayMYHCRNq7LntZE9cEemht39_TFPro0k,4682
+liger_kernel/ops/backends/_ascend/ops/tvd.py,sha256=4Q_DXSuVRqummX5dwFT5zOgQpdaWViLbMPjJ3kWy2IE,7745
 liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
 liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
-liger_kernel/transformers/__init__.py,sha256=4sqcDbOZ_JtS9Ag-7oyuhq5jN298GyzjJFu9J-DyyZQ,10872
+liger_kernel/transformers/__init__.py,sha256=h7U1Vxrg5OoqOstBmZMd-0G0LROYleYt_fS-RpvEq84,11057
 liger_kernel/transformers/auto_model.py,sha256=RnJhK8xHamRnnswgRLG_muJE1i6T6LszjK8lC6vonhE,2410
 liger_kernel/transformers/cross_entropy.py,sha256=08H8RxSxGX_52UzrHNnSZ_wWH-uvU8KrRiDmVrkOw14,1996
 liger_kernel/transformers/dyt.py,sha256=Rng-MZQSprnGGWFtpmYKt7MIX26vFUYbq5ruM4MjH-U,719
@@ -71,7 +72,7 @@ liger_kernel/transformers/jsd.py,sha256=_KlOX8YcdONU0tq0bIRDQ5VDBwtywm3Ro-FmlmI0
 liger_kernel/transformers/kl_div.py,sha256=94VR4uuj-2dZCTEnwFksvDi-LporrpB5HgmYtQCZnw0,402
 liger_kernel/transformers/layer_norm.py,sha256=l4nsT_Zj4CdVZOM7F0I0Ox-lmLHyIJzqQvVaF0o0HbI,895
 liger_kernel/transformers/llama4_rope.py,sha256=A_nxcS_KiUCyNeL2FAZX7yUhDsX7krrI9BG49OaN_nM,3627
-liger_kernel/transformers/monkey_patch.py,sha256=ESFIi_7hQMcnUtRLjAMJ9kbzSbwToDhpOfFa6aQ-SrY,135534
+liger_kernel/transformers/monkey_patch.py,sha256=hCFLKviPteLyDTUxjehiUS6k4hEx2GHDEualDhKpEYs,138949
 liger_kernel/transformers/multi_token_attention.py,sha256=LtEjG7qy1-JK-HIPaz8zZ4P08aSZTnj5D635Pa04Onc,1730
 liger_kernel/transformers/poly_norm.py,sha256=T3VdLQHLcCY7KzNzrc6IJRs8SzO8Yc7a0BS_2p6d7Wo,1367
 liger_kernel/transformers/qwen2vl_mrope.py,sha256=0hOBR3j2Yd6xbT4z9BNRKEy1D0eyOUsIW6EmI_3PPNI,1033
@@ -80,16 +81,17 @@ liger_kernel/transformers/rope.py,sha256=-W9aYLa2hMOmmG5yeHcvPsOI5UTc95ylYxUddxk
 liger_kernel/transformers/softmax.py,sha256=VI5QGHYpXSiXckgovEnDGcXwitimsxKB0GX-AT4dAC4,256
 liger_kernel/transformers/sparsemax.py,sha256=Os49bSpPX4pWymsasv_3j20m8GFaI54e03XFPkHiPE0,393
 liger_kernel/transformers/swiglu.py,sha256=LpgikAs9hibAL7G6itygBbOlW9tZe5s4D2IGAKGpbPw,4284
-liger_kernel/transformers/tiled_mlp.py,sha256=gPsz7b0kxpk3mre7o1uGBt-XdNvMUN7IIqnUYIur-T0,4628
+liger_kernel/transformers/tiled_mlp.py,sha256=_Go2bN8huL4I0EHBPXNfpIRaEukl8hiQEEJIwpJST20,4498
 liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
 liger_kernel/transformers/tvd.py,sha256=GYjhtXgS3RTPveOTN2gyK4uBnjs6ii2vkSZRX21QpqA,446
 liger_kernel/transformers/experimental/__init__.py,sha256=oQqk-f32JYgWEP9DJCj6ty6bbJSGrdXsFDQFwGeX6vI,127
 liger_kernel/transformers/experimental/embedding.py,sha256=bjy9hHj--ivy6xEWdiE6qLy9uLyeS4PsBEgl_MdDrng,858
 liger_kernel/transformers/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liger_kernel/transformers/model/exaone4.py,sha256=T5Ef2FnkJ-i8ktRWvBB5GXFOIyJmvMPyGsDFt5awpmE,5802
 liger_kernel/transformers/model/falcon_h1.py,sha256=heUZ4wUt2ATmtBtmv8Rcro3pQl6fV9T0pburjTTW7os,5004
 liger_kernel/transformers/model/gemma.py,sha256=pAri4PYpknsFfkvyo8Ez2NNlqrUDW-KkExUXTGZAcH4,10621
-liger_kernel/transformers/model/gemma2.py,sha256=qa9Ok42vFojVGNmASTH3Ek566Vu507kjd--ZpZDKX9M,12024
-liger_kernel/transformers/model/gemma3.py,sha256=ZUrFCc-pfF8jYHV0HsptBr98hx6p2q9ea0kSzVAoFPo,14966
+liger_kernel/transformers/model/gemma2.py,sha256=KgSpXVi04c8hVFa7dqJtjzVobz6z7BNTvGc1WjoV4nk,12006
+liger_kernel/transformers/model/gemma3.py,sha256=2XPmtpZxR55wccKflIDqf2AwHJdxypUbd62fLuZ8two,15092
 liger_kernel/transformers/model/glm4.py,sha256=bSp22iPIjsli4-c_usUOsyh1Bs2gIK8X6ynS0azseUs,5900
 liger_kernel/transformers/model/glm4v.py,sha256=dd-BQpccDCp1SbIxcJ5rG8xcwYQK3KOv1Tgm9TGnZc4,6594
 liger_kernel/transformers/model/glm4v_moe.py,sha256=zKhMdOOrRhlrvCSFaeVYfddL1ubpY8edEO91TN81n98,7135
@@ -99,7 +101,7 @@ liger_kernel/transformers/model/internvl.py,sha256=OOutracs9qrPHSU7FVYar08yinvGr
 liger_kernel/transformers/model/llama.py,sha256=kqZeONzwTBzudoChlKMzq1w23BtYGbxWZC1l1V__JTw,13410
 liger_kernel/transformers/model/llama4.py,sha256=PfkynGVI0xxMs3EtyYpCgaALI6stu25OIrTIymE-pvg,4853
 liger_kernel/transformers/model/llava.py,sha256=yoADM_BuIEummtTDiwWqjfUjXUMZD78VJzS0TRj5GJ4,15687
-liger_kernel/transformers/model/loss_utils.py,sha256=mAV6NsE1xR2smQMlr_n9afh4ek3BhIfieZdTn1Z-9Fw,2836
+liger_kernel/transformers/model/loss_utils.py,sha256=tNbC94Z4Ca2mlv3MRhnqfpJ7sBc5MZJtt1-mzMMJT1M,3088
 liger_kernel/transformers/model/mistral.py,sha256=OcwOzVDMwwDbVccVPv-AaocznzWwzLT3aRaKK5SMaAg,6030
 liger_kernel/transformers/model/mixtral.py,sha256=YcBDoTEJDgLFJ_RTo180DYGxR8D5Ad9-idumif7kCPE,12130
 liger_kernel/transformers/model/mllama.py,sha256=vAHwCm63sn4kpAY0rDGf_N0HR7KRTBVpBYDVTPOaZTg,12079
@@ -122,9 +124,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=tX0h63aOFe3rNqTmk6JpMf75UPo981yzEa6TghnjS0Q,5370
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/METADATA,sha256=Mzy4eM7hocfOx4KYOI_qKR056hH-RyAOcd99Ju-qY5k,25660
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.6.4.dev20260107111351.dist-info/RECORD,,
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/METADATA,sha256=Ja1hknX3Qd5-8K5-BO7pX4Ln11BgPKgBrYBjf291kzU,25660
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/WHEEL,sha256=WnJ8fYhv8N4SYVK2lLYNI6N0kVATA7b0piVUNvqIIJE,91
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.6.4.dev20260116023519.dist-info/RECORD,,

{liger_kernel_nightly-0.6.4.dev20260107111351.dist-info → liger_kernel_nightly-0.6.4.dev20260116023519.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.3.2)
+Generator: setuptools (75.3.3)
 Root-Is-Purelib: true
 Tag: py3-none-any