liger-kernel-nightly 0.6.4.dev20260113145602__tar.gz → 0.6.4.dev20260116105204__tar.gz

{liger_kernel_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/PKG-INFO

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

{liger_kernel_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/benchmark/scripts/benchmark_tvd.py

@@ -9,6 +9,9 @@ from utils import parse_benchmark_script_args
 from utils import run_benchmarks
 
 from liger_kernel.transformers.tvd import LigerTVDLoss
+from liger_kernel.utils import infer_device
+
+device = infer_device()
 
 
 class TorchTVDLoss(torch.nn.Module):
@@ -40,8 +43,8 @@ def bench_speed_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
     torch_tvd = TorchTVDLoss(reduction=reduction)
     liger_tvd = LigerTVDLoss(reduction=reduction)
 
-    _input = torch.randn(B * T, V, requires_grad=True, device="cuda").softmax(dim=-1)
-    target = torch.randn(B * T, V, device="cuda").softmax(dim=-1)
+    _input = torch.randn(B * T, V, requires_grad=True, device=device).softmax(dim=-1)
+    target = torch.randn(B * T, V, device=device).softmax(dim=-1)
 
     def fwd():
         if input.kernel_provider == "liger":
@@ -82,8 +85,8 @@ def bench_memory_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput
     V = input.x
     B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"]
 
-    _input = torch.randn(B * T, V, requires_grad=True, device="cuda").softmax(dim=-1)
-    target = torch.randn(B * T, V, device="cuda").softmax(dim=-1)
+    _input = torch.randn(B * T, V, requires_grad=True, device=device).softmax(dim=-1)
+    target = torch.randn(B * T, V, device=device).softmax(dim=-1)
 
     def fwd():
         if input.kernel_provider == "liger":

{liger_kernel_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.4.dev20260113145602"
+version = "0.6.4.dev20260116105204"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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
@@ -247,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,)](
@@ -307,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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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
@@ -436,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,
@@ -516,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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/liger_kernel/transformers/model/gemma3.py

@@ -8,7 +8,6 @@ import torch.nn as nn
 from transformers.cache_utils import Cache
 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
@@ -268,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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/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_nightly-0.6.4.dev20260113145602 → liger_kernel_nightly-0.6.4.dev20260116105204}/src/liger_kernel_nightly.egg-info/PKG-INFO

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