liger-kernel-nightly 0.6.4.dev20251208235806__py3-none-any.whl → 0.6.4.dev20251209171241__py3-none-any.whl

liger_kernel/transformers/monkey_patch.py

@@ -35,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
@@ -1754,8 +1753,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
@@ -1829,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

liger_kernel/transformers/rope.py

@@ -1,4 +1,3 @@
-from typing import Optional
 from typing import Tuple
 
 import torch
@@ -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_nightly-0.6.4.dev20251208235806.dist-info → liger_kernel_nightly-0.6.4.dev20251209171241.dist-info}/METADATA

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

{liger_kernel_nightly-0.6.4.dev20251208235806.dist-info → liger_kernel_nightly-0.6.4.dev20251209171241.dist-info}/RECORD

@@ -60,12 +60,12 @@ liger_kernel/transformers/jsd.py,sha256=DGqRnxIZxsvxo0_tbbxX3b-sDbDjC_yKufyRIHCc
 liger_kernel/transformers/kl_div.py,sha256=WLffFbh1EExD2Eb1F7lN11fo9JJC-0751WJjZAF1Fj8,409
 liger_kernel/transformers/layer_norm.py,sha256=c9pk3PEasOKYR0rhe5e5nNrnYKVCEW4VC8S6LpCq9EQ,906
 liger_kernel/transformers/llama4_rope.py,sha256=kS6PSHEwf3dS7hD7C7p8S0geugx2EMCiP0h0F7LsUoY,3639
-liger_kernel/transformers/monkey_patch.py,sha256=0ER5BjQXcIKwgL2e7ji3_DIm1DaJzevzo53aXSe2YJU,135862
+liger_kernel/transformers/monkey_patch.py,sha256=3MtDn6_1lljiloWHFK_GOo9fCO61QwgXh6OAu9KQdAc,135705
 liger_kernel/transformers/multi_token_attention.py,sha256=K3NIY9_5TPgZ4_Rahn0xnkMXxD_fmlJHK4CWGYvGQp0,1752
 liger_kernel/transformers/poly_norm.py,sha256=g5tC75i3qy1_N26ZUP-jfpct7ivQAEdJfIfx8IXzeyE,1377
 liger_kernel/transformers/qwen2vl_mrope.py,sha256=5EwSqrMdsL9MYspeBMXBsNJKvH0MOmRrtJXAJlnnlOI,1047
 liger_kernel/transformers/rms_norm.py,sha256=HwddVqrqS58jE-M2_4NkFGARtCDBhGnkKyjBN9b3FYI,3004
-liger_kernel/transformers/rope.py,sha256=VMlDZI6zss9mLaLcN5XCE_ktmYRwAi_Eh4TIgO6NrIQ,2361
+liger_kernel/transformers/rope.py,sha256=fmSG2h3GVtx9nqPNUOqH8TvGRKaJhsVZJMzsdVIMqPM,2872
 liger_kernel/transformers/softmax.py,sha256=yadlAgE4V2JByMwrDDa2s5SUBp8Jgd57xwnVvAWoBaI,264
 liger_kernel/transformers/sparsemax.py,sha256=0lQA0UEOs4mu8CMruZ3VLhImxQVXJWhPsAKUsYA7vj8,403
 liger_kernel/transformers/swiglu.py,sha256=dRR69wDWSWfdjtnsTECyxQqWVo5QkdXdXm9SpSQ4Jvw,4291
@@ -111,9 +111,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=tX0h63aOFe3rNqTmk6JpMf75UPo981yzEa6TghnjS0Q,5370
 liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
 liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/METADATA,sha256=tjEDcZKe6AONv56kfZB58iuv1zOQJ1BvIjHMVio3A1U,25468
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.6.4.dev20251208235806.dist-info/RECORD,,
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/METADATA,sha256=UgQQVN7vzeOeYXRTEMJucnAgorAIj6QRpv8vBj5fin8,25468
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.6.4.dev20251209171241.dist-info/RECORD,,