liger-kernel-nightly 0.5.5.dev20250402185702__py3-none-any.whl → 0.6.4.dev20260112233432__py3-none-any.whl

liger_kernel/transformers/__init__.py

@@ -1,30 +1,221 @@
-from liger_kernel.transformers.auto_model import AutoLigerKernelForCausalLM  # noqa: F401
+import importlib
+
+from typing import TYPE_CHECKING
+
+# Always-safe imports (independent of 'transformers')
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.dyt import LigerDyT  # noqa: F401
+from liger_kernel.transformers.fused_add_rms_norm import LigerFusedAddRMSNorm  # noqa: F401
 from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD  # noqa: F401
 from liger_kernel.transformers.geglu import LigerGEGLUMLP  # noqa: F401
 from liger_kernel.transformers.jsd import LigerJSD  # noqa: F401
+from liger_kernel.transformers.kl_div import LigerKLDIVLoss  # noqa: F401
 from liger_kernel.transformers.layer_norm import LigerLayerNorm  # 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_gemma  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2  # 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_llama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_paligemma  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_5_vl  # noqa: F401
-from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
+from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb  # noqa: F401
+from liger_kernel.transformers.llama4_rope import liger_llama4_vision_rotary_pos_emb  # noqa: F401
+from liger_kernel.transformers.multi_token_attention import LigerMultiTokenAttention  # noqa: F401
+from liger_kernel.transformers.poly_norm import LigerPolyNorm  # noqa: F401
 from liger_kernel.transformers.rms_norm import LigerRMSNorm  # noqa: F401
 from liger_kernel.transformers.rope import liger_rotary_pos_emb  # noqa: F401
+from liger_kernel.transformers.softmax import LigerSoftmax  # noqa: F401
+from liger_kernel.transformers.sparsemax import LigerSparsemax  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP  # noqa: F401
+from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP  # noqa: F401
+from liger_kernel.transformers.tiled_mlp import LigerTiledGEGLUMLP  # noqa: F401
+from liger_kernel.transformers.tiled_mlp import LigerTiledSwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.tvd import LigerTVDLoss  # noqa: F401
+
+# Static-only imports for IDEs and type checkers
+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
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text  # noqa: F401
+    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
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_internvl  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama4  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_paligemma  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_5_vl  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2_vl  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_moe  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_next  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_vl  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen3_vl_moe  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_smollm3  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_smolvlm  # noqa: F401
+
+
+# Check if 'transformers' is installed
+try:
+    import transformers  # noqa: F401
+
+    _TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    _TRANSFORMERS_AVAILABLE = False
+
+
+def is_transformers_available() -> bool:
+    """
+    Returns True if the 'transformers' package is available.
+    Useful for conditional logic in downstream code.
+    """
+    return _TRANSFORMERS_AVAILABLE
+
+
+def __getattr__(name: str):
+    """
+    Handles lazy access to transformer-dependent attributes.
+    If 'transformers' is not installed, raises a user-friendly ImportError.
+    """
+    if not _TRANSFORMERS_AVAILABLE:
+        raise ImportError(
+            f"The attribute '{name}' requires the 'transformers' library, which is not installed.\n"
+            f"Please install it with `pip install transformers` to use this functionality."
+        )
+
+    if name == "AutoLigerKernelForCausalLM":
+        module = importlib.import_module("liger_kernel.transformers.auto_model")
+        return getattr(module, name)
+
+    monkey_patch_symbols = {
+        "_apply_liger_kernel",
+        "_apply_liger_kernel_to_instance",
+        "apply_liger_kernel_to_falcon_h1",
+        "apply_liger_kernel_to_gemma",
+        "apply_liger_kernel_to_gemma2",
+        "apply_liger_kernel_to_gemma3",
+        "apply_liger_kernel_to_gemma3_text",
+        "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",
+        "apply_liger_kernel_to_llava",
+        "apply_liger_kernel_to_llama4",
+        "apply_liger_kernel_to_mistral",
+        "apply_liger_kernel_to_mixtral",
+        "apply_liger_kernel_to_mllama",
+        "apply_liger_kernel_to_olmo2",
+        "apply_liger_kernel_to_olmo3",
+        "apply_liger_kernel_to_paligemma",
+        "apply_liger_kernel_to_phi3",
+        "apply_liger_kernel_to_qwen2",
+        "apply_liger_kernel_to_qwen2_5_vl",
+        "apply_liger_kernel_to_qwen2_vl",
+        "apply_liger_kernel_to_qwen3",
+        "apply_liger_kernel_to_qwen3_moe",
+        "apply_liger_kernel_to_qwen3_next",
+        "apply_liger_kernel_to_qwen3_vl",
+        "apply_liger_kernel_to_qwen3_vl_moe",
+        "apply_liger_kernel_to_smollm3",
+        "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:
+        module = importlib.import_module("liger_kernel.transformers.monkey_patch")
+        return getattr(module, name)
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+# Shared symbols in all environments
+__all__ = [
+    "is_transformers_available",
+    "LigerCrossEntropyLoss",
+    "LigerDyT",
+    "LigerFusedLinearCrossEntropyLoss",
+    "LigerFusedLinearJSD",
+    "LigerGEGLUMLP",
+    "LigerJSD",
+    "LigerLayerNorm",
+    "LigerFusedAddRMSNorm",
+    "LigerPolyNorm",
+    "LigerRMSNorm",
+    "liger_rotary_pos_emb",
+    "liger_llama4_text_rotary_pos_emb",
+    "liger_llama4_vision_rotary_pos_emb",
+    "LigerBlockSparseTop2MLP",
+    "LigerPhi3SwiGLUMLP",
+    "LigerQwen3MoeSwiGLUMLP",
+    "LigerSwiGLUMLP",
+    "LigerTiledGEGLUMLP",
+    "LigerTiledSwiGLUMLP",
+    "LigerTVDLoss",
+    "LigerKLDIVLoss",
+    "LigerMultiTokenAttention",
+    "LigerSoftmax",
+    "LigerSparsemax",
+]
+
+# Add transformer-dependent symbols only if available
+if _TRANSFORMERS_AVAILABLE:
+    __all__.extend(
+        [
+            "AutoLigerKernelForCausalLM",
+            "_apply_liger_kernel",
+            "_apply_liger_kernel_to_instance",
+            "apply_liger_kernel_to_falcon_h1",
+            "apply_liger_kernel_to_gemma",
+            "apply_liger_kernel_to_gemma2",
+            "apply_liger_kernel_to_gemma3",
+            "apply_liger_kernel_to_gemma3_text",
+            "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",
+            "apply_liger_kernel_to_llava",
+            "apply_liger_kernel_to_llama4",
+            "apply_liger_kernel_to_mistral",
+            "apply_liger_kernel_to_mixtral",
+            "apply_liger_kernel_to_mllama",
+            "apply_liger_kernel_to_olmo2",
+            "apply_liger_kernel_to_olmo3",
+            "apply_liger_kernel_to_paligemma",
+            "apply_liger_kernel_to_phi3",
+            "apply_liger_kernel_to_qwen2",
+            "apply_liger_kernel_to_qwen2_5_vl",
+            "apply_liger_kernel_to_qwen2_vl",
+            "apply_liger_kernel_to_qwen3",
+            "apply_liger_kernel_to_qwen3_moe",
+            "apply_liger_kernel_to_qwen3_next",
+            "apply_liger_kernel_to_qwen3_vl",
+            "apply_liger_kernel_to_qwen3_vl_moe",
+            "apply_liger_kernel_to_smollm3",
+            "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/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,8 @@ 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
 
 
 class LigerCrossEntropyLoss(torch.nn.Module):
@@ -15,6 +16,7 @@ class LigerCrossEntropyLoss(torch.nn.Module):
         reduction: str = "mean",
         softcap: Optional[float] = None,
         return_z_loss: bool = False,
+        return_token_accuracy: bool = False,
     ):
         super().__init__()
         assert (label_smoothing >= 0) and (label_smoothing <= 1), (
@@ -33,9 +35,10 @@ class LigerCrossEntropyLoss(torch.nn.Module):
         self.reduction = reduction
         self.softcap = softcap
         self.return_z_loss = return_z_loss
+        self.return_token_accuracy = return_token_accuracy
 
     def forward(self, _input: torch.Tensor, target: torch.Tensor):
-        loss, z_loss = LigerCrossEntropyFunction.apply(
+        loss, z_loss, token_accuracy = LigerCrossEntropyFunction.apply(
             _input,
             target,
             self.weight,
@@ -45,7 +48,9 @@ class LigerCrossEntropyLoss(torch.nn.Module):
             self.reduction,
             self.softcap,
             self.return_z_loss,
+            self.return_token_accuracy,
         )
-        if not self.return_z_loss:
+        if not self.return_z_loss and not self.return_token_accuracy:
             return loss
-        return loss, z_loss
+
+        return CrossEntropyOutput(loss=loss, z_loss=z_loss, token_accuracy=token_accuracy)

liger_kernel/transformers/dyt.py

@@ -1,20 +1,22 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.dyt import LigerDyTFunction
+from liger_kernel.ops import LigerDyTFunction
 
 
 class LigerDyT(nn.Module):
-    def __init__(self, hidden_size, init_alpha=0.5):
+    def __init__(self, hidden_size, beta=True, init_alpha=0.5):
         super().__init__()
         self.hidden_size = hidden_size
         self.init_alpha = init_alpha
         self.alpha = nn.Parameter(torch.ones(1) * init_alpha)
         self.gamma = nn.Parameter(torch.ones(hidden_size))
-        self.beta = nn.Parameter(torch.zeros(hidden_size))
+        self.beta = None
+        if beta:
+            self.beta = nn.Parameter(torch.zeros(hidden_size))
 
     def forward(self, x):
         return LigerDyTFunction.apply(x, self.alpha, self.gamma, self.beta)
 
     def extra_repr(self):
-        return f"{self.hidden_size}, init_alpha={self.init_alpha}"
+        return f"{self.hidden_size}, init_alpha={self.init_alpha}, beta={self.beta}"

liger_kernel/transformers/experimental/__init__.py

@@ -0,0 +1,5 @@
+from liger_kernel.transformers.experimental.embedding import LigerEmbedding  # noqa: F401
+
+__all__ = [
+    "LigerEmbedding",
+]

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/fsdp.py

@@ -0,0 +1,55 @@
+from typing import Any
+from typing import Callable
+
+from torch.distributed.fsdp import FullyShardedDataParallel
+
+
+class _FSDPForwardRedirection:
+    """
+    Modified based on
+    https://github.com/Lightning-AI/pytorch-lightning/blob/d3f9c83d6efa4f1def36aa6c199600946cdb9117/src/lightning/pytorch/strategies/strategy.py#L601-L648
+    Redirect a method call through FullyShardedDataParallel.forward so that the FSDP module's root pre-forward and
+    post-forward can be properly executed around the method call.
+    This is needed in cases where we call a submodule of a FSDP module. For instance, when we want to call only
+    the `LlamaModel` part out of a FSDP-wrapped `LlamaForCausalLM` to get the hidden states without involving
+    GPU-memory-heavy `lm_head` and cross entropy computation, doing this directly (i.e. `model.model.forward()`)
+    will not work because the first `nn.Embedding` layer is not independently wrapped as a FSDP module (because of
+    the transformer-based wrapping policy), and not calling it through FSDP root module forward will not all-gather
+    its parameter, thus resulting in "RuntimeError: 'weight' must be 2-D" error. Similarly, if we want to call just
+    the `lm_head` part of a model, we need this trick too to properly get its params all-gathered.
+    """
+
+    def __call__(
+        self,
+        wrapper_module: FullyShardedDataParallel,
+        method: Callable,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """Reroutes a method call through the `wrapper_module`'s `forward` method.
+        Args:
+            wrapper_module: The module that has `original_module` wrapped.
+            original_module: The module that was wrapped inside `wrapper_module`.
+            method_name: The name of the method that should be called on the `original_module` after inputs get
+                redirected through the `wrapper_module`'s `forward` method.
+            *args: The positional arguments to the method `method_name`. They will get passed to a patched
+                `forward` method instead.
+            **kwargs: The keyword arguments to the method `method_name`. They will get passed to a patched
+                `forward` method instead.
+        """
+        assert isinstance(wrapper_module, FullyShardedDataParallel)
+        original_module = wrapper_module._fsdp_wrapped_module
+        original_forward = original_module.forward
+
+        def wrapped_forward(*_args: Any, **_kwargs: Any) -> Any:
+            # Unpatch ourselves immediately before calling the method `method_name`
+            # because itself may want to call the real `forward`
+            original_module.forward = original_forward  # type: ignore[method-assign]
+            # Call the actual method e.g. `.training_step(...)`
+            out = method(*_args, **_kwargs)
+            return out
+
+        # Patch the original_module's forward so we can redirect the arguments back to the real method
+        original_module.forward = wrapped_forward  # type: ignore[method-assign]
+        wrapper_output = wrapper_module(*args, **kwargs)
+        return wrapper_output

liger_kernel/transformers/functional.py

@@ -1,19 +1,35 @@
+from dataclasses import dataclass
 from typing import Optional
 
-from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
-from liger_kernel.ops.dyt import LigerDyTFunction
-from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
-from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
-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.qwen2vl_mrope import LigerQwen2VLMRopeFunction
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
-from liger_kernel.ops.rope import LigerRopeFunction
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+import torch
+
+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
+class CrossEntropyOutput:
+    loss: torch.Tensor
+    z_loss: Optional[torch.Tensor] = None
+    token_accuracy: Optional[torch.Tensor] = None
 
 
 # conform to the function signature in https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
@@ -30,8 +46,9 @@ def liger_cross_entropy(
     lse_square_scale: float = 0.0,
     softcap: Optional[float] = None,
     return_z_loss: bool = False,
+    return_token_accuracy: bool = False,
 ):
-    loss, z_loss = LigerCrossEntropyFunction.apply(
+    loss, z_loss, token_accuracy = LigerCrossEntropyFunction.apply(
         input,
         target,
         weight,
@@ -41,10 +58,13 @@ def liger_cross_entropy(
         reduction,
         softcap,
         return_z_loss,
+        return_token_accuracy,
     )
-    if not return_z_loss:
+
+    if not return_z_loss and not return_token_accuracy:
         return loss
-    return loss, z_loss
+
+    return CrossEntropyOutput(loss=loss, z_loss=z_loss, token_accuracy=token_accuracy)
 
 
 def liger_fused_linear_cross_entropy(
@@ -59,8 +79,11 @@ def liger_fused_linear_cross_entropy(
     reduction: str = "mean",
     softcap: Optional[float] = None,
     return_z_loss: bool = False,
+    accum_dtype=None,
+    use_token_scaling: bool = False,
+    return_token_accuracy: bool = False,
 ):
-    loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
+    loss, z_loss, token_accuracy = LigerFusedLinearCrossEntropyFunction.apply(
         input,
         weight,
         target,
@@ -72,10 +95,15 @@ def liger_fused_linear_cross_entropy(
         reduction,
         softcap,
         return_z_loss,
+        accum_dtype,
+        use_token_scaling,
+        return_token_accuracy,
     )
-    if not return_z_loss:
+
+    if not return_z_loss and not return_token_accuracy:
         return loss
-    return loss, z_loss
+
+    return CrossEntropyOutput(loss=loss, z_loss=z_loss, token_accuracy=token_accuracy)
 
 
 def liger_fused_linear_jsd(
@@ -159,6 +187,68 @@ def liger_kl_div(
     )
 
 
+def liger_sparsemax(
+    input,
+    dim: int = -1,
+):
+    return LigerSparsemaxFunction.apply(input, dim)
+
+
+def liger_multi_token_attention(
+    scores,
+    weight,
+    bias=None,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    sparse: bool = False,
+):
+    """
+    Functional interface for multi-token attention.
+
+    Args:
+        scores: Input tensor of shape (B, C_in, L, L)
+        weight: Convolution weight tensor of shape (C_out, C_in // groups, K, K)
+        bias: Optional bias tensor of shape (C_out,)
+        stride: Stride for the convolution (default: 1)
+        padding: Padding for the convolution (default: 0)
+        dilation: Dilation factor for the convolution (default: 1)
+        groups: Number of groups for the convolution (default: 1)
+        sparse: Specifies if input tensors are expected to be sparse (default: False)
+    Returns:
+        Output tensor after applying multi-token attention.
+    """
+    return LigerMultiTokenAttentionFunction.apply(scores, weight, bias, stride, padding, dilation, groups, sparse)
+
+
+def liger_fused_neighborhood_attention(
+    query,
+    key,
+    value,
+    kernel_size: int = 7,
+    dilation: int = 1,
+    scale: float = None,
+):
+    """
+    Liger fused neighborhood attention.
+
+    paper: https://arxiv.org/pdf/2504.16922
+
+    Args:
+        query: Query tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        key: Key tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        value: Value tensor of shape [batch_size, num_heads, seq_len, head_dim]
+        kernel_size: Size of the neighborhood window (default: 7)
+        dilation: Dilation factor for the neighborhood (default: 1)
+        scale: Scaling factor for attention scores (default: rsqrt(head_dim))
+
+    Returns:
+        Output tensor of shape [batch_size, num_heads, seq_len, head_dim]
+    """
+    return LigerFusedNeighborhoodAttentionFunction.apply(query, key, value, kernel_size, dilation, scale)
+
+
 def liger_tvd(
     input,
     target,
@@ -187,6 +277,14 @@ def liger_rms_norm(X, W, eps, offset: float = 0.0, casting_mode: str = "llama",
     return LigerRMSNormFunction.apply(X, W, eps, offset, casting_mode, in_place)
 
 
+def liger_poly_norm(X, W, B, eps=1e-6, in_place=True):
+    return LigerPolyNormFunction.apply(X, W, B, eps, in_place)
+
+
+def liger_fused_add_rms_norm(X, R, W, eps, offset: float = 0.0, casting_mode: str = "llama", in_place: bool = True):
+    return LigerFusedAddRMSNormFunction.apply(X, R, W, eps, offset, casting_mode, in_place)
+
+
 def liger_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)
 
@@ -195,5 +293,9 @@ def liger_swiglu(a, b):
     return LigerSiLUMulFunction.apply(a, b)
 
 
+def liger_softmax(x):
+    return LigerSoftmaxFunction.apply(x)
+
+
 def liger_dyt(x, alpha, gamma, beta):
     return LigerDyTFunction.apply(x, alpha, gamma, beta)

liger_kernel/transformers/fused_add_rms_norm.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerFusedAddRMSNormFunction
+
+
+class LigerFusedAddRMSNorm(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        eps=1e-6,
+        offset=0.0,
+        casting_mode="llama",
+        init_fn="ones",
+        in_place=False,
+    ):
+        super().__init__()
+        assert init_fn in [
+            "ones",
+            "zeros",
+        ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
+        self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
+        self.variance_epsilon, self.offset, self.casting_mode, self.in_place = (eps, offset, casting_mode, in_place)
+
+    def forward(self, hidden_states, residual):
+        return LigerFusedAddRMSNormFunction.apply(
+            hidden_states,
+            residual,
+            self.weight,
+            self.variance_epsilon,
+            self.offset,
+            self.casting_mode,
+            self.in_place,
+        )
+
+    def extra_repr(self):
+        return (
+            f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
+        )