liger-kernel 0.6.4__py3-none-any.whl → 0.6.5__py3-none-any.whl

liger_kernel/transformers/monkey_patch.py

@@ -20,6 +20,7 @@ from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forwa
 from liger_kernel.transformers.model.gemma import lce_forward_deprecated as gemma_lce_forward_deprecated
 from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_forward
 from liger_kernel.transformers.model.gemma2 import lce_forward_deprecated as gemma2_lce_forward_deprected
+from liger_kernel.transformers.model.gpt_oss import lce_forward as gpt_oss_lce_forward
 from liger_kernel.transformers.model.llama import lce_forward as llama_lce_forward
 from liger_kernel.transformers.model.llama import lce_forward_deprecated as llama_lce_forward_deprecated
 from liger_kernel.transformers.model.llava import lce_forward as llava_lce_forward
@@ -34,8 +35,7 @@ from liger_kernel.transformers.model.smollm3 import lce_forward as smollm3_lce_f
 from liger_kernel.transformers.qwen2vl_mrope import liger_multimodal_rotary_pos_emb
 from liger_kernel.transformers.rms_norm import LigerRMSNorm
 from liger_kernel.transformers.rope import liger_rotary_pos_emb
-from liger_kernel.transformers.rope import liger_rotary_pos_emb_with_cast
-from liger_kernel.transformers.rope import liger_rotary_pos_emb_with_cast_and_leading_batch
+from liger_kernel.transformers.rope import liger_rotary_pos_emb_vision
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
@@ -430,7 +430,7 @@ def apply_liger_kernel_to_llava(
                     f"These parameters are not supported by {text_model_name}. Enter the remaining {list(text_kwargs.keys())} except for {list(remain_params)}\n"
                     f"Parameters accepted by {text_model_name}: {list(accept_params.keys())}"
                 )
-            text_kwargs["model"] = model.language_model
+            text_kwargs["model"] = model.model.language_model
             text_liger_fn(**text_kwargs)
         elif text_model_name not in MODEL_TYPE_TO_APPLY_LIGER_FN:
             logger.warning(f"{text_model_name} is not supported by Liger kernel.")
@@ -445,7 +445,7 @@ def apply_liger_kernel_to_llava(
                     f"These parameters are not supported by {vision_model_name}. Enter the remaining {list(vision_kwargs.keys())} except for {list(remain_params)}\n"
                     f"Parameters accepted by {vision_model_name}: {list(accept_params.keys())}"
                 )
-            vision_kwargs["model"] = model.vision_tower
+            vision_kwargs["model"] = model.model.vision_tower
             vision_liger_fn(**vision_kwargs)
         elif vision_model_name not in MODEL_TYPE_TO_APPLY_LIGER_FN:
             logger.warning(f"{vision_model_name} is not supported by Liger kernel.")
@@ -615,8 +615,8 @@ def apply_liger_kernel_to_mllama(
         # instance variables that reference already-instantiated modules
 
         if isinstance(model, MllamaForConditionalGeneration):
-            language_model: MllamaForCausalLM = model.language_model
-            vision_model: MllamaVisionModel = model.vision_model
+            language_model: MllamaForCausalLM = model.model.language_model
+            vision_model: MllamaVisionModel = model.model.vision_model
             if isinstance(language_model, MllamaForCausalLM):
                 text_model: MllamaTextModel = language_model.model
             else:
@@ -1118,8 +1118,8 @@ def apply_liger_kernel_to_gemma3(
         # instance variables that reference already-instantiated modules
 
         if isinstance(model, Gemma3ForConditionalGeneration):
-            if isinstance(model.vision_tower, SiglipVisionModel):
-                vision_tower = model.vision_tower
+            if isinstance(model.model.vision_tower, SiglipVisionModel):
+                vision_tower = model.model.vision_tower
 
                 _patch_layer_norm_module(vision_tower.vision_model.post_layernorm)
 
@@ -1132,7 +1132,7 @@ def apply_liger_kernel_to_gemma3(
                 raise TypeError("The vision tower must be SiglipVisionModel")
 
             if rms_norm:
-                _patch_rms_norm_module_for_gemma3(model.multi_modal_projector.mm_soft_emb_norm)
+                _patch_rms_norm_module_for_gemma3(model.model.multi_modal_projector.mm_soft_emb_norm)
 
             apply_liger_kernel_to_gemma3_text(
                 rope=rope,
@@ -1140,7 +1140,7 @@ def apply_liger_kernel_to_gemma3(
                 fused_linear_cross_entropy=False,
                 rms_norm=rms_norm,
                 geglu=geglu,
-                model=model.language_model,
+                model=model.model.language_model,
             )
 
         else:
@@ -1228,7 +1228,7 @@ def apply_liger_kernel_to_paligemma(
         if not isinstance(model, PaliGemmaForConditionalGeneration):
             raise TypeError("model have to be of type PaliGemmaForConditionalGeneration")
 
-        vision_tower: SiglipVisionModel = model.vision_tower
+        vision_tower: SiglipVisionModel = model.model.vision_tower
 
         _patch_layer_norm_module(vision_tower.vision_model.post_layernorm)
 
@@ -1238,7 +1238,7 @@ def apply_liger_kernel_to_paligemma(
                 _patch_layer_norm_module(layer.layer_norm1)
                 _patch_layer_norm_module(layer.layer_norm2)
 
-        language_model = model.language_model
+        language_model = model.model.language_model
 
         if isinstance(language_model, (GemmaForCausalLM, GemmaModel)):
             apply_liger_kernel_to_gemma(
@@ -1459,6 +1459,79 @@ def apply_liger_kernel_to_qwen3_moe(
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
 
 
+def apply_liger_kernel_to_gpt_oss(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = False,  # Set to False by default since GPT-OSS has custom expert implementation
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace GPT-OSS models.
+    NOTE: GPT-OSS is supported in transformers >= 4.55.0
+    NOTE: SwiGLU patching is disabled by default for GPT-OSS as it uses a custom expert
+          implementation with clamping and MXFP4 quantization.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is True.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is False.
+            Note: GPT-OSS uses a custom expert implementation, so SwiGLU patching is disabled by default.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+            loaded. Default is None.
+    """
+    if version.parse(transformers.__version__) < version.parse("4.55.0"):
+        logger.warning("GPT-OSS support requires transformers >= 4.55.0")
+        return
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.gpt_oss import modeling_gpt_oss
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssModel
+
+    if rope:
+        modeling_gpt_oss.apply_rotary_pos_emb = liger_rotary_pos_emb
+
+    if rms_norm:
+        modeling_gpt_oss.GptOssRMSNorm = LigerRMSNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        if model is not None:
+            model.forward = MethodType(gpt_oss_lce_forward, model)
+        else:
+            modeling_gpt_oss.GptOssForCausalLM.forward = gpt_oss_lce_forward
+
+    # Note: SwiGLU patching is not implemented for GPT-OSS due to custom expert implementation
+    # with clamping (swiglu_limit=7.0) and MXFP4 quantization
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+
+        # get the base model from the model instance
+        base_model: GptOssModel = getattr(model, model.base_model_prefix, model)
+
+        if rms_norm:
+            _patch_rms_norm_module(base_model.norm)
+        for decoder_layer in base_model.layers:
+            if rms_norm:
+                _patch_rms_norm_module(decoder_layer.input_layernorm)
+                _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
+
+
 def apply_liger_kernel_to_qwen2_vl(
     rope: bool = True,
     cross_entropy: bool = False,
@@ -1520,11 +1593,10 @@ def apply_liger_kernel_to_qwen2_vl(
     if model is not None:
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
-
-        if isinstance(model, (Qwen2VLForConditionalGeneration, Qwen2VLModel)):
-            # Note: language_model and visual properties can be accessed throught conditional class for BC.
-            # Not sure if it is subject to changes in the future.
-            # Reference: https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1698
+        if isinstance(model, Qwen2VLForConditionalGeneration):
+            text_model: Qwen2VLTextModel = model.model.language_model
+            vision_model: Qwen2VisionTransformerPretrainedModel = model.model.visual
+        elif isinstance(model, Qwen2VLModel):
             text_model: Qwen2VLTextModel = model.language_model
             vision_model: Qwen2VisionTransformerPretrainedModel = model.visual
         elif isinstance(model, Qwen2VLTextModel):
@@ -1611,11 +1683,10 @@ def apply_liger_kernel_to_qwen2_5_vl(
     if model is not None:
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
-
-        if isinstance(model, (Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLModel)):
-            # Note: language_model and visual properties can be accessed throught conditional class for BC.
-            # Not sure if it is subject to changes in the future.
-            # Reference: https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L1823
+        if isinstance(model, Qwen2_5_VLForConditionalGeneration):
+            text_model: Qwen2_5_VLTextModel = model.model.language_model
+            vision_model: Qwen2_5_VisionTransformerPretrainedModel = model.model.visual
+        elif isinstance(model, Qwen2_5_VLModel):
             text_model: Qwen2_5_VLTextModel = model.language_model
             vision_model: Qwen2_5_VisionTransformerPretrainedModel = model.visual
         elif isinstance(model, Qwen2_5_VLTextModel):
@@ -1629,7 +1700,7 @@ def apply_liger_kernel_to_qwen2_5_vl(
 
         if vision_model is not None:
             # Patch Qwen2_5_VisionTransformerPretrainedModel
-            for vision_block in model.visual.blocks:
+            for vision_block in vision_model.blocks:
                 if rms_norm:
                     _patch_rms_norm_module(vision_block.norm1)
                     _patch_rms_norm_module(vision_block.norm2)
@@ -1680,8 +1751,8 @@ def apply_liger_kernel_to_qwen3_vl(
     from liger_kernel.transformers.model.qwen3_vl import lce_forward as qwen3_vl_lce_forward
 
     if rope:
-        modeling_qwen3_vl.apply_rotary_pos_emb = liger_rotary_pos_emb_with_cast
-        modeling_qwen3_vl.apply_rotary_pos_emb_vision = liger_rotary_pos_emb_with_cast_and_leading_batch
+        modeling_qwen3_vl.apply_rotary_pos_emb = liger_rotary_pos_emb
+        modeling_qwen3_vl.apply_rotary_pos_emb_vision = liger_rotary_pos_emb_vision
 
     if rms_norm:
         modeling_qwen3_vl.Qwen3VLTextRMSNorm = LigerRMSNorm
@@ -1698,7 +1769,9 @@ def apply_liger_kernel_to_qwen3_vl(
             modeling_qwen3_vl.Qwen3VLForConditionalGeneration.forward = qwen3_vl_lce_forward
 
     if model is not None and rms_norm:
-        if isinstance(model, (Qwen3VLForConditionalGeneration, Qwen3VLModel)):
+        if isinstance(model, Qwen3VLForConditionalGeneration):
+            text_model: Qwen3VLTextModel = model.model.language_model
+        elif isinstance(model, Qwen3VLModel):
             text_model: Qwen3VLTextModel = model.language_model
         elif isinstance(model, Qwen3VLTextModel):
             text_model = model
@@ -1755,8 +1828,8 @@ def apply_liger_kernel_to_qwen3_vl_moe(
     from liger_kernel.transformers.model.qwen3_vl_moe import lce_forward as qwen3_vl_moe_lce_forward
 
     if rope:
-        modeling_qwen3_vl_moe.apply_rotary_pos_emb = liger_rotary_pos_emb_with_cast
-        modeling_qwen3_vl_moe.apply_rotary_pos_emb_vision = liger_rotary_pos_emb_with_cast_and_leading_batch
+        modeling_qwen3_vl_moe.apply_rotary_pos_emb = liger_rotary_pos_emb
+        modeling_qwen3_vl_moe.apply_rotary_pos_emb_vision = liger_rotary_pos_emb_vision
 
     if rms_norm:
         modeling_qwen3_vl_moe.Qwen3VLMoeTextRMSNorm = LigerRMSNorm
@@ -1773,7 +1846,9 @@ def apply_liger_kernel_to_qwen3_vl_moe(
             modeling_qwen3_vl_moe.Qwen3VLMoeForConditionalGeneration.forward = qwen3_vl_moe_lce_forward
 
     if model is not None and rms_norm:
-        if isinstance(model, (Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeModel)):
+        if isinstance(model, Qwen3VLMoeForConditionalGeneration):
+            text_model: Qwen3VLMoeTextModel = model.model.language_model
+        elif isinstance(model, Qwen3VLMoeModel):
             text_model: Qwen3VLMoeTextModel = model.language_model
         elif isinstance(model, Qwen3VLMoeTextModel):
             text_model = model
@@ -2118,10 +2193,10 @@ def apply_liger_kernel_to_glm4v(
     if model is not None:
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
-        if isinstance(model, (Glm4vForConditionalGeneration, Glm4vModel)):
-            # Note: language_model and visual properties can be accessed throught conditional class for BC.
-            # Not sure if it is subject to changes in the future.
-            # Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm4v/modeling_glm4v.py#L1305
+        if isinstance(model, Glm4vForConditionalGeneration):
+            text_model: Glm4vTextModel = model.model.language_model
+            vision_model: Glm4vVisionModel = model.model.visual
+        elif isinstance(model, Glm4vModel):
             text_model: Glm4vTextModel = model.language_model
             vision_model: Glm4vVisionModel = model.visual
         elif isinstance(model, Glm4vTextModel):
@@ -2208,10 +2283,11 @@ def apply_liger_kernel_to_glm4v_moe(
     if model is not None:
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
-        if isinstance(model, (Glm4vMoeForConditionalGeneration, Glm4vMoeModel)):
-            # Note: language_model and visual properties can be accessed throught conditional class for BC.
-            # Not sure if it is subject to changes in the future.
-            # Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py#L337
+        if isinstance(model, Glm4vMoeForConditionalGeneration):
+            text_model: Glm4vMoeTextModel = model.model.language_model
+            vision_model: Glm4vMoeVisionModel = model.model.visual
+            Glm4vMoeTextMoE = modeling_glm4v_moe.Glm4vMoeTextMoE
+        elif isinstance(model, Glm4vMoeModel):
             text_model: Glm4vMoeTextModel = model.language_model
             vision_model: Glm4vMoeVisionModel = model.visual
             Glm4vMoeTextMoE = modeling_glm4v_moe.Glm4vMoeTextMoE
@@ -2314,8 +2390,10 @@ def apply_liger_kernel_to_internvl(
     if model is not None:
         # The model instance already exists, so we need to additionally patch the
         # instance variables that reference already-instantiated modules
-        if isinstance(model, (InternVLForConditionalGeneration, InternVLModel)):
-            # NOTE: language_model and visual properties can be accessed throught conditional class.
+        if isinstance(model, InternVLForConditionalGeneration):
+            text_model = model.model.language_model
+            vision_model: InternVLVisionModel = model.model.vision_tower
+        elif isinstance(model, InternVLModel):
             text_model = model.language_model
             vision_model: InternVLVisionModel = model.vision_tower
         else:
@@ -2743,6 +2821,83 @@ def apply_liger_kernel_to_hunyuan_v1_moe(
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
 
 
+def apply_liger_kernel_to_exaone4(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace EXAONE4 models.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is True.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+        loaded. Default is None.
+    """
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.exaone4 import modeling_exaone4
+    from transformers.models.exaone4.modeling_exaone4 import Exaone4Model
+
+    from liger_kernel.transformers.model.exaone4 import lce_forward as exaone4_lce_forward
+
+    if rope:
+        modeling_exaone4.apply_rotary_pos_emb = liger_rotary_pos_emb
+
+    if rms_norm:
+        # EXAONE4 requires in_place=False to avoid gradient issues
+        class Exaone4LigerRMSNorm(LigerRMSNorm):
+            def __init__(self, hidden_size, eps=1e-6, **kwargs):
+                super().__init__(hidden_size, eps, **kwargs)
+                self.in_place = False
+
+        modeling_exaone4.Exaone4RMSNorm = Exaone4LigerRMSNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        if model is not None:
+            model.forward = MethodType(exaone4_lce_forward, model)
+        else:
+            modeling_exaone4.Exaone4ForCausalLM.forward = exaone4_lce_forward
+
+    if swiglu:
+        modeling_exaone4.Exaone4MLP = LigerSwiGLUMLP
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+
+        # get the base model from the model instance
+        base_model: Exaone4Model = getattr(model, model.base_model_prefix, model)
+
+        if rms_norm:
+            _patch_rms_norm_module(base_model.norm, in_place=False)
+        for decoder_layer in base_model.layers:
+            if swiglu:
+                _bind_method_to_module(decoder_layer.mlp, "forward", LigerSwiGLUMLP.forward)
+            if rms_norm:
+                _patch_rms_norm_module(decoder_layer.post_attention_layernorm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.post_feedforward_layernorm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.self_attn.q_norm, in_place=False)
+                _patch_rms_norm_module(decoder_layer.self_attn.k_norm, in_place=False)
+
+
 # Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
 MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "gemma": apply_liger_kernel_to_gemma,
@@ -2752,6 +2907,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "glm4": apply_liger_kernel_to_glm4,
     "glm4v": apply_liger_kernel_to_glm4v,
     "glm4v_moe": apply_liger_kernel_to_glm4v_moe,
+    "gpt_oss": apply_liger_kernel_to_gpt_oss,
     "internvl": apply_liger_kernel_to_internvl,
     "llama": apply_liger_kernel_to_llama,
     "llama4_text": apply_liger_kernel_to_llama4,
@@ -2783,6 +2939,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "smolvlm": apply_liger_kernel_to_smolvlm,
     "hunyuan_v1_dense": apply_liger_kernel_to_hunyuan_v1_dense,
     "hunyuan_v1_moe": apply_liger_kernel_to_hunyuan_v1_moe,
+    "exaone4": apply_liger_kernel_to_exaone4,
 }
 
 

liger_kernel/transformers/multi_token_attention.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 
 from torch.nn.modules.utils import _pair
 
-from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
+from liger_kernel.ops import LigerMultiTokenAttentionFunction
 
 
 class LigerMultiTokenAttention(nn.Module):

liger_kernel/transformers/poly_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.poly_norm import LigerPolyNormFunction
+from liger_kernel.ops import LigerPolyNormFunction
 
 
 class LigerPolyNorm(nn.Module):

liger_kernel/transformers/qwen2vl_mrope.py

@@ -1,4 +1,4 @@
-from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
+from liger_kernel.ops import LigerQwen2VLMRopeFunction
 
 
 def liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):

liger_kernel/transformers/rms_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops import LigerRMSNormFunction
 
 
 class LigerRMSNorm(nn.Module):
@@ -14,13 +14,18 @@ class LigerRMSNorm(nn.Module):
         init_fn="ones",
         in_place=True,
         row_mode=None,
+        elementwise_affine=True,
     ):
         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.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
+        else:
+            self.register_parameter("weight", None)
         self.variance_epsilon, self.offset, self.casting_mode, self.in_place, self.row_mode = (
             eps,
             offset,
@@ -41,7 +46,7 @@ class LigerRMSNorm(nn.Module):
         )
 
     def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}, row_mode={self.row_mode}"
+        return f"weight_shape={tuple(self.weight.shape) if self.weight is not None else None}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}, row_mode={self.row_mode}"
 
 
 class LigerRMSNormForGemma(LigerRMSNorm):

liger_kernel/transformers/rope.py

@@ -1,9 +1,8 @@
-from typing import Optional
 from typing import Tuple
 
 import torch
 
-from liger_kernel.ops.rope import LigerRopeFunction
+from liger_kernel.ops import LigerRopeFunction
 
 
 def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
@@ -25,39 +24,41 @@ def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)
 
 
-def liger_rotary_pos_emb_with_cast(
+def liger_rotary_pos_emb_vision(
     q: torch.Tensor,
     k: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-    position_ids: Optional[torch.Tensor] = None,
-    unsqueeze_dim: int = 1,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Modified version of liger_rotary_pos_emb for qwen3_vl's apply_rotary_pos_emb_vision function.
+    Manually tranposed the input and output to match the expected shape for liger_rotary_pos_emb.
+    Reference: https://https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L116
+
+    Args:
+        q (torch.Tensor): The query tensor of shape (seq_length, num_heads, head_dim),
+        with stride (num_heads * head_dim, head_dim, 1).
+        k (torch.Tensor): The query tensor of shape (seq_length, num_heads, head_dim),
+        with stride (num_heads * head_dim, head_dim, 1). Same as q.
+        cos (torch.Tensor): The cosine tensor of shape (seq_length, head_dim).
+        sin (torch.Tensor): The sine tensor of shape (seq_length, head_dim).
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The query and key tensors with the same shape and stride as inputs.
+    """
     orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
 
-    q32 = q.to(torch.float32)
-    k32 = k.to(torch.float32)
+    # tranpose to (1, num_heads, seq_length, head_dim) and cast to float32 to match liger_rotary_pos_emb input shape
+    # also unsqueeze for batch dim
+    q32 = q.to(torch.float32).unsqueeze(0).transpose(1, 2)
+    k32 = k.to(torch.float32).unsqueeze(0).transpose(1, 2)
     cos32 = cos.to(torch.float32)
     sin32 = sin.to(torch.float32)
 
-    q_out, k_out = liger_rotary_pos_emb(q32, k32, cos32, sin32, position_ids=position_ids, unsqueeze_dim=unsqueeze_dim)
-    return q_out.to(orig_q_dtype), k_out.to(orig_k_dtype)
-
-
-def liger_rotary_pos_emb_with_cast_and_leading_batch(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    position_ids: Optional[torch.Tensor] = None,
-    unsqueeze_dim: int = 1,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
-
-    q32 = q.to(torch.float32).unsqueeze(0)
-    k32 = k.to(torch.float32).unsqueeze(0)
-    cos32 = cos.to(torch.float32).unsqueeze(0)
-    sin32 = sin.to(torch.float32).unsqueeze(0)
+    q_out, k_out = liger_rotary_pos_emb(q32, k32, cos32, sin32)
 
-    q_out, k_out = liger_rotary_pos_emb(q32, k32, cos32, sin32, position_ids=position_ids, unsqueeze_dim=unsqueeze_dim)
-    return q_out.to(orig_q_dtype).squeeze(0), k_out.to(orig_k_dtype).squeeze(0)
+    # transpose back to (seq_length, num_heads, head_dim) and cast back to original dtype
+    # also squeeze out batch dim
+    q_out = q_out.transpose(1, 2).squeeze(0).to(orig_q_dtype)
+    k_out = k_out.transpose(1, 2).squeeze(0).to(orig_k_dtype)
+    return q_out, k_out

liger_kernel/transformers/softmax.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.softmax import LigerSoftmaxFunction
+from liger_kernel.ops import LigerSoftmaxFunction
 
 
 class LigerSoftmax(nn.Module):

liger_kernel/transformers/sparsemax.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
+from liger_kernel.ops import LigerSparsemaxFunction
 
 
 class LigerSparsemax(nn.Module):

liger_kernel/transformers/swiglu.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
 
 
 class LigerSwiGLUMLP(nn.Module):

liger_kernel/transformers/tiled_mlp.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 import torch.nn as nn
 
-from liger_kernel.ops.geglu import LigerGELUMulFunction
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-from liger_kernel.ops.tiled_mlp import apply_tiled_mlp
+from liger_kernel.ops import LigerGELUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
+from liger_kernel.ops import apply_tiled_mlp
 
 
 class LigerTiledGEGLUMLP(nn.Module):
@@ -57,11 +57,7 @@ class LigerTiledGEGLUMLP(nn.Module):
         Returns:
             Output tensor of the same shape as input
         """
-        compute_params = [
-            self.gate_proj.weight,
-            self.up_proj.weight,
-            self.down_proj.weight,
-        ]
+        compute_params = [p for p in self.parameters() if p.requires_grad]
 
         return apply_tiled_mlp(
             fn=self._mlp_forward,
@@ -118,11 +114,7 @@ class LigerTiledSwiGLUMLP(nn.Module):
         Returns:
             Output tensor of the same shape as input
         """
-        compute_params = [
-            self.gate_proj.weight,
-            self.up_proj.weight,
-            self.down_proj.weight,
-        ]
+        compute_params = [p for p in self.parameters() if p.requires_grad]
 
         return apply_tiled_mlp(
             fn=self._mlp_forward,

liger_kernel/transformers/tvd.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops import LigerTVDLossFunction
 
 
 class LigerTVDLoss(nn.Module):