liger-kernel-nightly 0.5.10.dev20250624183504__py3-none-any.whl → 0.6.4.dev20251121224847__py3-none-any.whl

liger_kernel/transformers/__init__.py

@@ -5,17 +5,27 @@ 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.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
@@ -23,18 +33,26 @@ 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_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_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
@@ -42,6 +60,11 @@ if TYPE_CHECKING:
     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
@@ -79,18 +102,24 @@ def __getattr__(name: str):
     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_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",
@@ -98,6 +127,13 @@ def __getattr__(name: str):
         "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",
     }
 
     if name in monkey_patch_symbols:
@@ -117,13 +153,23 @@ __all__ = [
     "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
@@ -133,18 +179,24 @@ if _TRANSFORMERS_AVAILABLE:
             "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_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",
@@ -152,5 +204,12 @@ if _TRANSFORMERS_AVAILABLE:
             "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",
         ]
     )

liger_kernel/transformers/cross_entropy.py

@@ -3,6 +3,7 @@ from typing import Optional
 import torch
 
 from liger_kernel.ops.cross_entropy 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/experimental/__init__.py

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

liger_kernel/transformers/functional.py

@@ -1,7 +1,11 @@
+from dataclasses import dataclass
 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
@@ -11,6 +15,7 @@ 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
@@ -20,6 +25,13 @@ from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.tvd 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
 # `weight` and `size_average` are placeholders and not implemented yet
 def liger_cross_entropy(
@@ -34,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,
@@ -45,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(
@@ -63,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,
@@ -76,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(
@@ -253,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)
 

liger_kernel/transformers/fused_add_rms_norm.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.fused_add_rms_norm 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}"
+        )

liger_kernel/transformers/fused_linear_cross_entropy.py

@@ -3,6 +3,7 @@ from typing import Optional
 import torch
 
 from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
+from liger_kernel.transformers.functional import CrossEntropyOutput
 
 
 class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
@@ -15,6 +16,9 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         reduction: str = "mean",
         softcap: Optional[float] = None,
         return_z_loss: bool = False,
+        accum_dtype: Optional[torch.dtype] = None,
+        use_token_scaling: bool = False,
+        return_token_accuracy: bool = False,
     ):
         super().__init__()
         assert (label_smoothing >= 0) and (label_smoothing <= 1), (
@@ -23,7 +27,8 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         assert reduction in {
             "mean",
             "sum",
-        }, f"reduction must be 'mean' or 'sum'. Got: {reduction}"
+            "none",
+        }, f"reduction must be 'mean' or 'sum' or 'none'. Got: {reduction}"
         assert softcap is None or softcap > 0, f"softcap must greater than 0.0 or None. Got: {softcap}"
         self.ce_weight = ce_weight
         self.ignore_index = ignore_index
@@ -32,9 +37,12 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
         self.reduction = reduction
         self.softcap = softcap
         self.return_z_loss = return_z_loss
+        self.accum_dtype = accum_dtype
+        self.use_token_scaling = use_token_scaling
+        self.return_token_accuracy = return_token_accuracy
 
     def forward(self, lin_weight, _input, target, bias=None):
-        loss, z_loss = LigerFusedLinearCrossEntropyFunction.apply(
+        loss, z_loss, token_accuracy = LigerFusedLinearCrossEntropyFunction.apply(
             _input,
             lin_weight,
             target,
@@ -46,7 +54,11 @@ class LigerFusedLinearCrossEntropyLoss(torch.nn.Module):
             self.reduction,
             self.softcap,
             self.return_z_loss,
+            self.accum_dtype,
+            self.use_token_scaling,
+            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/grpo_loss.py

@@ -1,3 +1,6 @@
+import torch
+
+from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
 from liger_kernel.ops.grpo_loss import GrpoLossFunction
 
 
@@ -13,12 +16,20 @@ def triton_grpo_loss(
     eps_low=0.2,
     eps_high=0.4,
     inplace=True,
+    loss_type="dapo",
+    max_completion_length=None,
+    importance_sampling_level="token",
+    reduce=False,
 ):
     assert logits is not None and completion_ids is not None and advantages is not None, (
         "must provide logits、completion_ids and advantages"
     )
+    if importance_sampling_level != "token":
+        raise ValueError(
+            f"Triton GRPO loss only supports token-level importance sampling. Got {importance_sampling_level}."
+        )
 
-    return GrpoLossFunction.apply(
+    per_token_loss, per_token_kl, is_clipped = GrpoLossFunction.apply(
         logits,
         old_logp,
         ref_logp,
@@ -31,6 +42,50 @@ def triton_grpo_loss(
         eps_high,
         inplace,
     )
+    if not reduce:
+        return per_token_loss, per_token_kl, is_clipped
+
+    loss = _reduce_grpo_loss(
+        per_token_loss,
+        completion_mask,
+        loss_type=loss_type,
+        max_completion_length=max_completion_length,
+    )
+
+    metrics = []
+    if beta != 0.0 and per_token_kl is not None:
+        metrics.append(_masked_mean(per_token_kl, completion_mask))
+    metrics.append(_masked_mean(is_clipped.float(), completion_mask))
+    return loss, metrics
+
+
+def _reduce_grpo_loss(per_token_loss, completion_mask, loss_type, max_completion_length):
+    mask = completion_mask
+    if mask is None:
+        mask = torch.ones_like(per_token_loss, dtype=per_token_loss.dtype, device=per_token_loss.device)
+    mask = mask.to(per_token_loss.dtype)
+
+    if loss_type == "grpo":
+        per_seq = (per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)
+        return per_seq.mean()
+    if loss_type == "bnpo":
+        return (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0)
+    if loss_type == "dr_grpo":
+        if max_completion_length is None:
+            raise ValueError("max_completion_length must be provided when using loss_type='dr_grpo'")
+        batch = per_token_loss.shape[0]
+        return (per_token_loss * mask).sum() / (batch * max_completion_length)
+    if loss_type == "dapo":
+        normalizer = LigerFusedLinearPPOBase._compute_dapo_normalizer(mask)
+        return (per_token_loss * mask).sum() / normalizer
+    raise ValueError(f"Unsupported loss_type '{loss_type}' for Triton GRPO loss.")
+
+
+def _masked_mean(values, mask):
+    if mask is None:
+        mask = torch.ones_like(values, dtype=values.dtype, device=values.device)
+    mask = mask.to(values.dtype)
+    return (values * mask).sum() / mask.sum().clamp(min=1.0)
 
 
 # This is a demo how to use grpo_loss in GRPOTrainer. The Trl version must be 0.16

liger_kernel/transformers/llama4_rope.py

@@ -0,0 +1,93 @@
+"""
+Liger Kernel implementation of Llama4 Rotary Position Embedding (RoPE).
+Supports both text and vision RoPE variants with fused operations for optimal performance.
+"""
+
+import torch
+
+from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction
+
+
+def liger_llama4_text_rotary_pos_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Liger-optimized implementation of Llama4 text rotary position embedding.
+
+    This implementation uses a fused Triton kernel for complex multiplication,
+    providing significant performance improvements over the original PyTorch implementation.
+
+    Args:
+        xq (torch.Tensor): Query tensor of shape (batch_size, seq_len, num_heads, head_dim)
+        xk (torch.Tensor): Key tensor of shape (batch_size, seq_len, num_heads, head_dim)
+        freqs_cis (torch.Tensor): Complex frequency tensor from Llama4TextRotaryEmbedding
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Rotated query and key tensors
+    """
+    # Use fused Triton kernel for complex RoPE
+    return LigerLlama4RopeFunction.apply(xq, xk, freqs_cis)
+
+
+def liger_llama4_vision_rotary_pos_emb(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    freqs_ci: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Liger-optimized implementation of Llama4 vision rotary position embedding.
+
+    This implementation uses the same fused Triton kernel as text RoPE,
+    providing performance improvements for vision transformer attention.
+
+    Args:
+        query (torch.Tensor): Query tensor of shape (batch_size, seq_len, num_heads, head_dim)
+        key (torch.Tensor): Key tensor of shape (batch_size, seq_len, num_heads, head_dim)
+        freqs_ci (torch.Tensor): Complex frequency tensor for 2D positions
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Rotated query and key tensors
+    """
+    # Handle broadcasting for vision RoPE
+    if freqs_ci.dim() == 3:
+        try:
+            # Try the regular 3D expansion
+            freqs_ci = freqs_ci.unsqueeze(0).expand(query.shape[0], -1, -1)
+        except RuntimeError as e:
+            if "expand" in str(e) and "4" in str(e):
+                # The tensor is actually 4D internally, handle it differently
+                freqs_ci = freqs_ci.squeeze(1)  # Remove the middle dimension
+                freqs_ci = freqs_ci.unsqueeze(0).expand(query.shape[0], -1, -1)
+            else:
+                raise e
+    elif freqs_ci.dim() == 4:  # (1, seq_len, 1, head_dim//2) - already properly shaped
+        # Squeeze the middle dimension to get (1, seq_len, head_dim//2)
+        freqs_ci = freqs_ci.squeeze(2)
+    elif freqs_ci.dim() == 2:  # (seq_len, head_dim//2) - needs expansion
+        freqs_ci = freqs_ci.unsqueeze(0).expand(query.shape[0], -1, -1)
+    else:
+        raise ValueError(f"Unexpected freqs_ci shape: {freqs_ci.shape}")
+
+    # Use the same fused kernel as text RoPE
+    return LigerLlama4RopeFunction.apply(query, key, freqs_ci)
+
+
+# Note: We only patch the functions, not the classes
+# The original Llama4TextRotaryEmbedding and Llama4VisionRotaryEmbedding classes remain unchanged
+
+
+# Convenience functions for monkey patching
+def apply_liger_llama4_rope_full(modeling_module):
+    """
+    Apply Liger optimizations to Llama4 RoPE functions.
+
+    Args:
+        modeling_module: The transformers modeling module to patch
+    """
+    # Replace the text RoPE function
+    modeling_module.apply_rotary_emb = liger_llama4_text_rotary_pos_emb
+
+    # Replace the vision RoPE function
+    modeling_module.vision_apply_rotary_emb = liger_llama4_vision_rotary_pos_emb