liger-kernel-nightly 0.5.10.dev20250528223524__py3-none-any.whl → 0.5.10.dev20250601024230__py3-none-any.whl

liger_kernel/ops/rms_norm.py

@@ -193,6 +193,7 @@ def _rms_norm_backward_kernel(
 
     tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
+
 @triton.jit
 def _block_rms_norm_forward_kernel(
     Y_ptr,
@@ -225,8 +226,11 @@ def _block_rms_norm_forward_kernel(
     row_mask = row_idx < n_rows
     col_mask = col_offsets < n_cols
 
-
-    X_row = tl.load(X_ptr + row_idx[:, None] * X_row_stride + col_offsets[None, :], mask=row_mask[:, None] & col_mask[None, :] , other=0)
+    X_row = tl.load(
+        X_ptr + row_idx[:, None] * X_row_stride + col_offsets[None, :],
+        mask=row_mask[:, None] & col_mask[None, :],
+        other=0,
+    )
     X_row_dtype = X_row.dtype
     W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0)
 
@@ -262,7 +266,12 @@ def _block_rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_GEMMA:
         Y_row = Y_row.to(X_row_dtype)
 
-    tl.store(Y_ptr + row_idx[:, None] * Y_row_stride + col_offsets[None, :], Y_row, mask=row_mask[:, None] & col_mask[None, :])
+    tl.store(
+        Y_ptr + row_idx[:, None] * Y_row_stride + col_offsets[None, :],
+        Y_row,
+        mask=row_mask[:, None] & col_mask[None, :],
+    )
+
 
 @triton.jit
 def _block_rms_norm_backward_kernel(
@@ -306,8 +315,16 @@ def _block_rms_norm_backward_kernel(
     for start in range(pid * BLOCK_ROW, n_rows, NUM_SMS * BLOCK_ROW):
         row_idx = start + tl.arange(0, BLOCK_ROW)
         row_mask = row_idx < n_rows
-        dY_row = tl.load(dY_ptr + row_idx[:, None] * dY_row_stride + col_offsets[None, :], mask=row_mask[:, None] & col_mask[None, :], other=0.0)
-        X_row = tl.load(X_ptr + row_idx[:, None] * X_row_stride + col_offsets[None, :], mask=row_mask[:, None] & col_mask[None, :], other=0.0)
+        dY_row = tl.load(
+            dY_ptr + row_idx[:, None] * dY_row_stride + col_offsets[None, :],
+            mask=row_mask[:, None] & col_mask[None, :],
+            other=0.0,
+        )
+        X_row = tl.load(
+            X_ptr + row_idx[:, None] * X_row_stride + col_offsets[None, :],
+            mask=row_mask[:, None] & col_mask[None, :],
+            other=0.0,
+        )
 
         # Get cached rms
         rstd_row = tl.load(RSTD_ptr + row_idx * RSTD_row_stride, row_mask)
@@ -326,7 +343,9 @@ def _block_rms_norm_backward_kernel(
 
         dX_row = rstd_row[:, None] * m
 
-        dX_row += (rstd_row[:, None]) * (-(1 / n_cols) * (rstd_row * rstd_row * tl.sum(m * X_row, axis=1))[:, None] * X_row)
+        dX_row += (rstd_row[:, None]) * (
+            -(1 / n_cols) * (rstd_row * rstd_row * tl.sum(m * X_row, axis=1))[:, None] * X_row
+        )
 
         # calculate the gradient of W
         if casting_mode == _CASTING_MODE_LLAMA:
@@ -335,8 +354,11 @@ def _block_rms_norm_backward_kernel(
             # here X_row is already in fp32 (see previous if block)
             dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]), 0)
 
-        tl.store(dX_ptr + row_idx[:, None] * dX_row_stride + col_offsets[None, :], dX_row, mask=row_mask[:, None] & col_mask[None, :])
-
+        tl.store(
+            dX_ptr + row_idx[:, None] * dX_row_stride + col_offsets[None, :],
+            dX_row,
+            mask=row_mask[:, None] & col_mask[None, :],
+        )
 
     tl.store(dW_ptr + pid * dW_row_stride + col_offsets, dW_row, mask=col_mask)
 
@@ -549,15 +571,6 @@ class LigerRMSNormFunction(torch.autograd.Function):
         """
         X, W, RSTD = ctx.saved_tensors
         dX, dW = rms_norm_backward(
-            dY,
-            X,
-            W,
-            RSTD,
-            ctx.offset,
-            ctx.casting_mode,
-            ctx.BLOCK_SIZE,
-            ctx.num_warps,
-            ctx.in_place,
-            ctx.row_mode
+            dY, X, W, RSTD, ctx.offset, ctx.casting_mode, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place, ctx.row_mode
         )
         return dX, dW, None, None, None, None, None

liger_kernel/transformers/monkey_patch.py

@@ -776,7 +776,8 @@ def apply_liger_kernel_to_gemma3_text(
 
     from transformers.models.gemma3 import modeling_gemma3
     from transformers.models.gemma3.modeling_gemma3 import Gemma3DecoderLayer
-    from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM, Gemma3TextModel
+    from transformers.models.gemma3.modeling_gemma3 import Gemma3ForCausalLM
+    from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel
 
     from liger_kernel.transformers.gema3_rms import LigerRMSNormForGemma3
     from liger_kernel.transformers.model.gemma3 import causal_forward

liger_kernel/transformers/rms_norm.py

@@ -37,7 +37,7 @@ class LigerRMSNorm(nn.Module):
             self.offset,
             self.casting_mode,
             self.in_place,
-            self.row_mode
+            self.row_mode,
         )
 
     def extra_repr(self):

liger_kernel/transformers/softmax.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from liger_kernel.ops.softmax import LigerSoftmaxFunction
 
 
-class LigerKernelSoftmax(nn.Module):
+class LigerSoftmax(nn.Module):
     def __init__(self):
         super().__init__()
 

{liger_kernel_nightly-0.5.10.dev20250528223524.dist-info → liger_kernel_nightly-0.5.10.dev20250601024230.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.10.dev20250528223524
+Version: 0.5.10.dev20250601024230
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -325,6 +325,8 @@ loss.backward()
 | GeGLU                           | `liger_kernel.transformers.LigerGEGLUMLP`                   |
 | CrossEntropy                    | `liger_kernel.transformers.LigerCrossEntropyLoss`           |
 | Fused Linear CrossEntropy       | `liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss`|
+| Multi Token Attention           | `liger_kernel.transformers.LigerMultiTokenAttention`        |
+| Softmax                         | `liger_kernel.transformers.LigerSoftmax`                    |
 | Sparsemax                       | `liger_kernel.transformers.LigerSparsemax`                  |
 
 

{liger_kernel_nightly-0.5.10.dev20250528223524.dist-info → liger_kernel_nightly-0.5.10.dev20250601024230.dist-info}/RECORD

@@ -28,7 +28,7 @@ liger_kernel/ops/kl_div.py,sha256=ZjGdDLKWksHT9dZ0xF_TDgAkj5cuMTwwT5tr9E-_24o,87
 liger_kernel/ops/layer_norm.py,sha256=vWCyOm-F2GMAilB-ozJcFeUQQLCJoTE_uiXq-_0uYuI,8356
 liger_kernel/ops/multi_token_attention.py,sha256=Oz_RXDp-OSS_R_HuGmaETHdAJ7Toda_70OfE7TXMUlY,7645
 liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU,8481
-liger_kernel/ops/rms_norm.py,sha256=IDj_V3hwo6tm3FijVbRh6ebUj2A3591MNkMer_gncdM,18749
+liger_kernel/ops/rms_norm.py,sha256=-rcgHwWCxlA-Syec2XhdW4jfOeCDt2r7qwjslgXFYDU,18865
 liger_kernel/ops/rope.py,sha256=ofmBOkUpZZO-Q8Z5B_LOFYYLD-YT-8WnJ4vGOrDYouI,8943
 liger_kernel/ops/softmax.py,sha256=tgORx6MK1IDDtZKqGarj0IPIVjqAIEUXXYPiinhRdtI,5864
 liger_kernel/ops/sparsemax.py,sha256=AeWe1xgkHJFEKWTj2vu_0hj7LztGvjqXAps-QTpCY0U,5087
@@ -52,12 +52,12 @@ liger_kernel/transformers/grpo_loss.py,sha256=uAkUNKSnUGEOqa82L9w2e6AI1kcmG8K45-
 liger_kernel/transformers/jsd.py,sha256=DGqRnxIZxsvxo0_tbbxX3b-sDbDjC_yKufyRIHCcScY,2979
 liger_kernel/transformers/kl_div.py,sha256=WLffFbh1EExD2Eb1F7lN11fo9JJC-0751WJjZAF1Fj8,409
 liger_kernel/transformers/layer_norm.py,sha256=c9pk3PEasOKYR0rhe5e5nNrnYKVCEW4VC8S6LpCq9EQ,906
-liger_kernel/transformers/monkey_patch.py,sha256=a0CXSC8BwZg3vok-ns0udZLUOBkegGQgPDod3H8ilP4,74610
+liger_kernel/transformers/monkey_patch.py,sha256=A91QWjMG7d7302lx-Djjxd_VgwBhYwxAYa1davBFCjU,74668
 liger_kernel/transformers/multi_token_attention.py,sha256=l9VDICK0dfmifUDW668hGscP8AHq2rYcM2oGUa3baRQ,1751
 liger_kernel/transformers/qwen2vl_mrope.py,sha256=5EwSqrMdsL9MYspeBMXBsNJKvH0MOmRrtJXAJlnnlOI,1047
-liger_kernel/transformers/rms_norm.py,sha256=srMS4jdkMCjY4Yqj9jjsy_IkY8KlHdTPLOx4069ZACA,1277
+liger_kernel/transformers/rms_norm.py,sha256=QimExM27kYoAnaZqxb_8mBaUcd72-X01DviJ1dQd55I,1278
 liger_kernel/transformers/rope.py,sha256=ZTrTORSAyfcFIKjk6XEeYmk4ROH7xXED9L4g2NFntlE,999
-liger_kernel/transformers/softmax.py,sha256=u7bFo35-cjaAm9of6-DLzmkaNFELOM-9AgyrcvUPifw,270
+liger_kernel/transformers/softmax.py,sha256=yadlAgE4V2JByMwrDDa2s5SUBp8Jgd57xwnVvAWoBaI,264
 liger_kernel/transformers/sparsemax.py,sha256=0lQA0UEOs4mu8CMruZ3VLhImxQVXJWhPsAKUsYA7vj8,403
 liger_kernel/transformers/swiglu.py,sha256=LZ8YeLIdv2k46JleZMjzubGk98smt6t780kSgcVLsQk,3454
 liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
@@ -86,9 +86,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.5.10.dev20250528223524.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.10.dev20250528223524.dist-info/METADATA,sha256=XqzBAk8PxwjhEYwf_3Xw0sssbGSM3IWW9z3NWlsZ7ZU,24113
-liger_kernel_nightly-0.5.10.dev20250528223524.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.10.dev20250528223524.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.10.dev20250528223524.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.10.dev20250528223524.dist-info/RECORD,,
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/METADATA,sha256=p4YDg6nRS2Zh3pCFi_dj1Yl7DtEi5U3bciMTtrcY-1U,24309
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.10.dev20250601024230.dist-info/RECORD,,