liger-kernel-nightly 0.6.4.dev20251202054858__py3-none-any.whl → 0.6.4.dev20260107181130__py3-none-any.whl

liger_kernel/ops/rms_norm.py

@@ -54,6 +54,7 @@ def _rms_norm_forward_kernel(
     eps,
     offset,
     casting_mode: tl.constexpr,  # constexpr so the `if` blocks can be optimized out
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -75,7 +76,8 @@ def _rms_norm_forward_kernel(
 
     X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
     X_row_dtype = X_row.dtype
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
+    if elementwise_affine:
+        W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
 
     # On Llama, only rstd is computed on fp32
     if casting_mode == _CASTING_MODE_LLAMA:
@@ -83,7 +85,8 @@ def _rms_norm_forward_kernel(
 
     # Gemma computes everything on fp32, and then casts back the output to the original dtype
     if casting_mode == _CASTING_MODE_GEMMA:
-        W_row = W_row.to(tl.float32)
+        if elementwise_affine:
+            W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
     if casting_mode == _CASTING_MODE_NONE:
@@ -104,7 +107,10 @@ def _rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_LLAMA:
         X_row = X_row.to(X_row_dtype)
 
-    Y_row = X_row * (offset + W_row)
+    if elementwise_affine:
+        Y_row = X_row * (offset + W_row)
+    else:
+        Y_row = X_row
 
     if casting_mode == _CASTING_MODE_GEMMA:
         Y_row = Y_row.to(X_row_dtype)
@@ -130,8 +136,9 @@ def _rms_norm_backward_kernel(
     n_rows,
     n_cols,
     offset,
-    rows_per_program: tl.constexpr,
+    rows_per_program,
     casting_mode: tl.constexpr,
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -145,7 +152,8 @@ def _rms_norm_backward_kernel(
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    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
@@ -153,8 +161,9 @@ def _rms_norm_backward_kernel(
     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
+    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)
@@ -167,24 +176,34 @@ def _rms_norm_backward_kernel(
 
         # Different bacward graphs for different casting modes
         if casting_mode == _CASTING_MODE_LLAMA:
-            m = (dY_row * W_row).to(tl.float32)
+            if elementwise_affine:
+                m = (dY_row * W_row).to(tl.float32)
+            else:
+                m = dY_row.to(tl.float32)
 
         elif casting_mode == _CASTING_MODE_GEMMA:
             dY_row = dY_row.to(tl.float32)
-            m = dY_row * W_row
+            if elementwise_affine:
+                m = dY_row * W_row
+            else:
+                m = dY_row
         else:
-            m = dY_row * W_row
+            if elementwise_affine:
+                m = dY_row * W_row
+            else:
+                m = dY_row
 
         dX_row = rstd_row * m
 
         dX_row += (rstd_row) * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row)
 
-        # calculate the gradient of W
-        if casting_mode == _CASTING_MODE_LLAMA:
-            dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
-        else:
-            # here X_row is already in fp32 (see previous if block)
-            dW_row += dY_row * (X_row * rstd_row)
+        if elementwise_affine:
+            # calculate the gradient of W
+            if casting_mode == _CASTING_MODE_LLAMA:
+                dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
+            else:
+                # 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)
 
@@ -193,7 +212,8 @@ def _rms_norm_backward_kernel(
         X_ptr += X_row_stride
         RSTD_ptr += RSTD_row_stride
 
-    tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
+    if elementwise_affine:
+        tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
 
 @triton.jit
@@ -211,6 +231,7 @@ def _block_rms_norm_forward_kernel(
     eps,
     offset,
     casting_mode: tl.constexpr,  # constexpr so the `if` blocks can be optimized out
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     BLOCK_ROW: tl.constexpr,
 ):
@@ -234,7 +255,8 @@ def _block_rms_norm_forward_kernel(
         other=0,
     )
     X_row_dtype = X_row.dtype
-    W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0)
+    if elementwise_affine:
+        W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0)
 
     # On Llama, only rstd is computed on fp32
     if casting_mode == _CASTING_MODE_LLAMA:
@@ -242,7 +264,8 @@ def _block_rms_norm_forward_kernel(
 
     # Gemma computes everything on fp32, and then casts back the output to the original dtype
     if casting_mode == _CASTING_MODE_GEMMA:
-        W_row = W_row.to(tl.float32)
+        if elementwise_affine:
+            W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
     if casting_mode == _CASTING_MODE_NONE:
@@ -263,7 +286,10 @@ def _block_rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_LLAMA:
         X_row = X_row.to(X_row_dtype)
 
-    Y_row = X_row * (offset + W_row)[None, :]
+    if elementwise_affine:
+        Y_row = X_row * (offset + W_row)[None, :]
+    else:
+        Y_row = X_row
 
     if casting_mode == _CASTING_MODE_GEMMA:
         Y_row = Y_row.to(X_row_dtype)
@@ -293,8 +319,8 @@ def _block_rms_norm_backward_kernel(
     n_rows,
     n_cols,
     offset,
-    rows_per_program: tl.constexpr,
     casting_mode: tl.constexpr,
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     BLOCK_ROW: tl.constexpr,
 ):
@@ -309,10 +335,11 @@ def _block_rms_norm_backward_kernel(
     col_offsets = tl.arange(0, BLOCK_SIZE)
     col_mask = col_offsets < n_cols
 
-    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    if elementwise_affine:
+        dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)
-    W_row = W_row + offset
+        W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)
+        W_row = W_row + offset
 
     for start in range(pid * BLOCK_ROW, n_rows, NUM_SMS * BLOCK_ROW):
         row_idx = start + tl.arange(0, BLOCK_ROW)
@@ -335,13 +362,22 @@ def _block_rms_norm_backward_kernel(
 
         # Different bacward graphs for different casting modes
         if casting_mode == _CASTING_MODE_LLAMA:
-            m = (dY_row * W_row[None, :]).to(tl.float32)
+            if elementwise_affine:
+                m = (dY_row * W_row[None, :]).to(tl.float32)
+            else:
+                m = dY_row.to(tl.float32)
 
         elif casting_mode == _CASTING_MODE_GEMMA:
             dY_row = dY_row.to(tl.float32)
-            m = dY_row * W_row[None, :]
+            if elementwise_affine:
+                m = dY_row * W_row[None, :]
+            else:
+                m = dY_row
         else:
-            m = dY_row * W_row[None, :]
+            if elementwise_affine:
+                m = dY_row * W_row[None, :]
+            else:
+                m = dY_row
 
         dX_row = rstd_row[:, None] * m
 
@@ -349,13 +385,13 @@ def _block_rms_norm_backward_kernel(
             -(1 / n_cols) * (rstd_row * rstd_row * tl.sum(m * X_row, axis=1))[:, None] * X_row
         )
 
-        # calculate the gradient of W
-        if casting_mode == _CASTING_MODE_LLAMA:
-			# TODO(tcc): use tl.sum(..., dtype=tl.float32) once we upgrade to triton>=3.3.0
-            dW_row += tl.sum((dY_row * (X_row * rstd_row[:, None]).to(X_dtype)).to(tl.float32), 0)
-        else:
-            # here X_row is already in fp32 (see previous if block)
-            dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]), 0)
+        if elementwise_affine:
+            if casting_mode == _CASTING_MODE_LLAMA:
+                # TODO(tcc): use tl.sum(..., dtype=tl.float32) once we upgrade to triton>=3.3.0
+                dW_row += tl.sum((dY_row * (X_row * rstd_row[:, None]).to(X_dtype)).to(tl.float32), 0)
+            else:
+                # here X_row is already in fp32 (see previous if block)
+                dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]), 0)
 
         tl.store(
             dX_ptr + row_idx[:, None] * dX_row_stride + col_offsets[None, :],
@@ -363,7 +399,8 @@ def _block_rms_norm_backward_kernel(
             mask=row_mask[:, None] & col_mask[None, :],
         )
 
-    tl.store(dW_ptr + pid * dW_row_stride + col_offsets, dW_row, mask=col_mask)
+    if elementwise_affine:
+        tl.store(dW_ptr + pid * dW_row_stride + col_offsets, dW_row, mask=col_mask)
 
 
 _str_to_casting_mode = {
@@ -392,8 +429,14 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
     rstd_dtype = torch.float32 if casting_mode in (_CASTING_MODE_LLAMA.value, _CASTING_MODE_GEMMA.value) else X.dtype
     RSTD = torch.empty(n_rows, dtype=rstd_dtype, device=X.device)
 
-    # Check constraints.
-    assert X.shape[1] == W.shape[0], "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
+    if W is not None:
+        # Check constraints.
+        assert X.shape[1] == W.shape[0], (
+            "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
+        )
+        elementwise_affine = True
+    else:
+        elementwise_affine = False
 
     # XPU-specific optimization
     kernel_args = {}
@@ -406,13 +449,14 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             X,
             X.stride(0),
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             n_cols,
             eps,
             offset,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -426,7 +470,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             X,
             X.stride(0),
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             n_rows,
@@ -434,6 +478,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             eps,
             offset,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -455,8 +500,13 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     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)
+    if W is not None:
+        # fp32 for numerical stability especially.
+        _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+        elementwise_affine = True
+    else:
+        _dW = None
+        elementwise_affine = False
 
     if n_cols > BLOCK_SIZE:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
@@ -483,16 +533,17 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
             X.stride(0),
             torch_to_triton_dtype[X.dtype],
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             _dW,
-            _dW.stride(0),
+            _dW.stride(0) if elementwise_affine else 0,
             n_rows,
             n_cols,
             offset,
             rows_per_program,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -509,22 +560,26 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
             X.stride(0),
             torch_to_triton_dtype[X.dtype],
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             _dW,
-            _dW.stride(0),
+            _dW.stride(0) if elementwise_affine else 0,
             n_rows,
             n_cols,
             offset,
-            rows_per_program,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
         )
     dX = dX.view(*shape)
-    dW = _dW.sum(dim=0).to(W.dtype)
+
+    if elementwise_affine:
+        dW = _dW.sum(dim=0).to(W.dtype)
+    else:
+        dW = None
 
     return dX, dW
 
@@ -558,6 +613,13 @@ class LigerRMSNormFunction(torch.autograd.Function):
         X: (B, T, H) or (BxT, H)
         W: (H,)
         """
+        if isinstance(X, torch.distributed.tensor.DTensor):
+            # Input tensor is output of a tensor parallel module and
+            # needs to be gathered to a local tensor to compute
+            # RMSE layer norm on each TP worker.
+            # TODO: support CP.
+            X = X.full_tensor()
+
         Y, X, RSTD, BLOCK_SIZE, num_warps, casting_mode = rms_norm_forward(X, W, eps, offset, casting_mode, row_mode)
         ctx.offset = offset
         ctx.casting_mode = casting_mode
@@ -565,7 +627,11 @@ class LigerRMSNormFunction(torch.autograd.Function):
         ctx.row_mode = row_mode
         ctx.BLOCK_SIZE = BLOCK_SIZE
         ctx.num_warps = num_warps
-        ctx.save_for_backward(X, W, RSTD)
+        ctx.elementwise_affine = W is not None
+        if W is not None:
+            ctx.save_for_backward(X, W, RSTD)
+        else:
+            ctx.save_for_backward(X, RSTD)
         return Y
 
     @staticmethod
@@ -574,7 +640,18 @@ class LigerRMSNormFunction(torch.autograd.Function):
         """
         Y: (B, T, H) or (BxT, H)
         """
-        X, W, RSTD = ctx.saved_tensors
+        if ctx.elementwise_affine:
+            X, W, RSTD = ctx.saved_tensors
+        else:
+            X, RSTD = ctx.saved_tensors
+            W = None
+
+        if isinstance(dY, torch.distributed.tensor.DTensor):
+            # Gradients are output of a tensor parallel module and
+            # needs to be gathered to a local tensor for computing RMSE layer.
+            # TODO: support CP.
+            dY = dY.full_tensor()
+
         dX, dW = rms_norm_backward(
             dY, X, W, RSTD, ctx.offset, ctx.casting_mode, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place, ctx.row_mode
         )

liger_kernel/ops/utils.py

@@ -127,3 +127,15 @@ def element_mul_kernel(
         X_offsets = i + tl.arange(0, BLOCK_SIZE)
         X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols)
         tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols)
+
+
+def get_npu_core_count(default: int = 20) -> int:
+    """Return NPU vector core count.
+    Fallback to `default` if Triton runtime or NPU device is unavailable.
+    """
+    try:
+        utils = triton.runtime.driver.active.utils
+        props = utils.get_device_properties(0)
+        return int(props.get("num_vectorcore", default))
+    except Exception:
+        return default

liger_kernel/transformers/__init__.py

@@ -41,6 +41,7 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v_moe  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gpt_oss  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_dense  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_moe  # noqa: F401
@@ -110,6 +111,7 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_glm4",
         "apply_liger_kernel_to_glm4v",
         "apply_liger_kernel_to_glm4v_moe",
+        "apply_liger_kernel_to_gpt_oss",
         "apply_liger_kernel_to_granite",
         "apply_liger_kernel_to_internvl",
         "apply_liger_kernel_to_llama",
@@ -187,6 +189,7 @@ if _TRANSFORMERS_AVAILABLE:
             "apply_liger_kernel_to_glm4",
             "apply_liger_kernel_to_glm4v",
             "apply_liger_kernel_to_glm4v_moe",
+            "apply_liger_kernel_to_gpt_oss",
             "apply_liger_kernel_to_granite",
             "apply_liger_kernel_to_internvl",
             "apply_liger_kernel_to_llama",

liger_kernel/transformers/auto_model.py

@@ -1,4 +1,5 @@
 import inspect
+import logging
 
 from transformers import AutoConfig
 from transformers import AutoModelForCausalLM
@@ -6,6 +7,8 @@ from transformers import AutoModelForCausalLM
 from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
 from liger_kernel.transformers.monkey_patch import _apply_liger_kernel
 
+logger = logging.getLogger(__name__)
+
 
 def _get_model_config(model_dir, **model_init_kwargs):
     config = AutoConfig.from_pretrained(model_dir, **model_init_kwargs)
@@ -36,3 +39,21 @@ class AutoLigerKernelForCausalLM(AutoModelForCausalLM):
         applicable_kwargs = {key: value for key, value in kwargs.items() if key not in apply_fn_signature.parameters}
 
         return super().from_pretrained(pretrained_model_name_or_path, *model_args, **applicable_kwargs)
+
+    @classmethod
+    def from_config(cls, config, **kwargs):
+        model_type = getattr(config, "model_type", None)
+        if not model_type:
+            logger.info("Model type could not be determined from model config. No Liger kernels will be applied.")
+            return
+        model_type = config.model_type
+
+        _apply_liger_kernel(model_type, **kwargs)
+
+        # Filter out kwargs that were passed to the apply_liger_* function, which will cause
+        # model initialization errors otherwise
+        apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type]
+        apply_fn_signature = inspect.signature(apply_fn)
+        applicable_kwargs = {key: value for key, value in kwargs.items() if key not in apply_fn_signature.parameters}
+
+        return super().from_config(config, **applicable_kwargs)

liger_kernel/transformers/cross_entropy.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
+from liger_kernel.ops import LigerCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
 
 

liger_kernel/transformers/dyt.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.dyt import LigerDyTFunction
+from liger_kernel.ops import LigerDyTFunction
 
 
 class LigerDyT(nn.Module):

liger_kernel/transformers/experimental/embedding.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.experimental.embedding import LigerEmbeddingFunction
+from liger_kernel.ops import LigerEmbeddingFunction
 
 
 class LigerEmbedding(nn.Module):

liger_kernel/transformers/functional.py

@@ -3,26 +3,26 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
-from liger_kernel.ops.dyt import LigerDyTFunction
-from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
-from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
-from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
-from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
-from liger_kernel.ops.geglu import LigerGELUMulFunction
-from liger_kernel.ops.group_norm import LigerGroupNormFunction
-from liger_kernel.ops.jsd import LigerJSDFunction
-from liger_kernel.ops.kl_div import LigerKLDivLossFunction
-from liger_kernel.ops.layer_norm import LigerLayerNormFunction
-from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
-from liger_kernel.ops.poly_norm import LigerPolyNormFunction
-from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
-from liger_kernel.ops.rope import LigerRopeFunction
-from liger_kernel.ops.softmax import LigerSoftmaxFunction
-from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops import LigerCrossEntropyFunction
+from liger_kernel.ops import LigerDyTFunction
+from liger_kernel.ops import LigerFusedAddRMSNormFunction
+from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
+from liger_kernel.ops import LigerFusedLinearJSDFunction
+from liger_kernel.ops import LigerFusedNeighborhoodAttentionFunction
+from liger_kernel.ops import LigerGELUMulFunction
+from liger_kernel.ops import LigerGroupNormFunction
+from liger_kernel.ops import LigerJSDFunction
+from liger_kernel.ops import LigerKLDivLossFunction
+from liger_kernel.ops import LigerLayerNormFunction
+from liger_kernel.ops import LigerMultiTokenAttentionFunction
+from liger_kernel.ops import LigerPolyNormFunction
+from liger_kernel.ops import LigerQwen2VLMRopeFunction
+from liger_kernel.ops import LigerRMSNormFunction
+from liger_kernel.ops import LigerRopeFunction
+from liger_kernel.ops import LigerSiLUMulFunction
+from liger_kernel.ops import LigerSoftmaxFunction
+from liger_kernel.ops import LigerSparsemaxFunction
+from liger_kernel.ops import LigerTVDLossFunction
 
 
 @dataclass

liger_kernel/transformers/fused_add_rms_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
+from liger_kernel.ops import LigerFusedAddRMSNormFunction
 
 
 class LigerFusedAddRMSNorm(nn.Module):

liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
 
 

liger_kernel/transformers/fused_linear_jsd.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
+from liger_kernel.ops import LigerFusedLinearJSDFunction
 
 
 class LigerFusedLinearJSD(torch.nn.Module):

liger_kernel/transformers/fused_neighborhood_attention.py

@@ -5,7 +5,7 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
+from liger_kernel.ops import LigerFusedNeighborhoodAttentionFunction
 
 
 class LigerFusedNeighborhoodAttention(nn.Module):

liger_kernel/transformers/geglu.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.geglu import LigerGELUMulFunction
+from liger_kernel.ops import LigerGELUMulFunction
 
 
 class LigerGEGLUMLP(nn.Module):

liger_kernel/transformers/group_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.group_norm import LigerGroupNormFunction
+from liger_kernel.ops import LigerGroupNormFunction
 
 
 class LigerGroupNorm(nn.Module):

liger_kernel/transformers/grpo_loss.py

@@ -1,7 +1,7 @@
 import torch
 
 from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
-from liger_kernel.ops.grpo_loss import GrpoLossFunction
+from liger_kernel.ops import GrpoLossFunction
 
 
 def triton_grpo_loss(

liger_kernel/transformers/jsd.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import torch
 
-from liger_kernel.ops.jsd import LigerJSDFunction
+from liger_kernel.ops import LigerJSDFunction
 
 
 class LigerJSD(torch.nn.Module):

liger_kernel/transformers/kl_div.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.kl_div import LigerKLDivLossFunction
+from liger_kernel.ops import LigerKLDivLossFunction
 
 
 class LigerKLDIVLoss(nn.KLDivLoss):

liger_kernel/transformers/layer_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.layer_norm import LigerLayerNormFunction
+from liger_kernel.ops import LigerLayerNormFunction
 
 
 class LigerLayerNorm(nn.Module):

liger_kernel/transformers/llama4_rope.py

@@ -5,7 +5,7 @@ Supports both text and vision RoPE variants with fused operations for optimal pe
 
 import torch
 
-from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction
+from liger_kernel.ops import LigerLlama4RopeFunction
 
 
 def liger_llama4_text_rotary_pos_emb(