liger-kernel-nightly 0.5.10.dev20250528223524__tar.gz → 0.5.10.dev20250531184114__tar.gz

{liger_kernel_nightly-0.5.10.dev20250528223524 → liger_kernel_nightly-0.5.10.dev20250531184114}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.10.dev20250528223524
+Version: 0.5.10.dev20250531184114
 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 → liger_kernel_nightly-0.5.10.dev20250531184114}/README.md

@@ -277,6 +277,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 → liger_kernel_nightly-0.5.10.dev20250531184114}/benchmark/scripts/benchmark_softmax.py

@@ -8,7 +8,7 @@ from utils import _test_memory
 from utils import parse_benchmark_script_args
 from utils import run_benchmarks
 
-from liger_kernel.transformers.softmax import LigerKernelSoftmax
+from liger_kernel.transformers.softmax import LigerSoftmax
 from liger_kernel.utils import infer_device
 
 device = infer_device()
@@ -23,7 +23,7 @@ def bench_speed_softmax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOut
     dtype = extra_benchmark_config["dtype"]
 
     x_shape = (M, N)
-    liger_softmax = LigerKernelSoftmax().to(device).to(dtype)
+    liger_softmax = LigerSoftmax().to(device).to(dtype)
     torch_softmax = torch.nn.Softmax(dim=-1).to(device).to(dtype)
 
     x = torch.randn(x_shape, dtype=dtype, device=device)
@@ -72,7 +72,7 @@ def bench_memory_softmax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
     dtype = extra_benchmark_config.get("dtype", torch.float32)
 
     torch_softmax = torch.nn.Softmax(dim=-1)
-    liger_softmax = LigerKernelSoftmax().to(device).to(dtype)
+    liger_softmax = LigerSoftmax().to(device).to(dtype)
 
     x = torch.randn(shape, device=device, dtype=dtype, requires_grad=True)
 

{liger_kernel_nightly-0.5.10.dev20250528223524 → liger_kernel_nightly-0.5.10.dev20250531184114}/docs/Low-Level-APIs.md

@@ -9,6 +9,8 @@
 | 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`                  |
 
 
@@ -51,6 +53,19 @@ This kernel combines linear transformations with cross-entropy loss calculations
 !!! Example "Try it out"
     You can experiment as shown in this example [here](https://colab.research.google.com/drive/1Z2QtvaIiLm5MWOs7X6ZPS1MN3hcIJFbj?usp=sharing)
 
+### Multi Token Attention
+
+The Multi Token Attention kernel implementation provides and optimized fused implementation of multi-token attention over the implemented Pytorch model baseline. This is a new attention mechanism that can operate on multiple Q and K inputs introduced by Meta Research.
+
+Paper: https://arxiv.org/abs/2504.00927
+
+### Softmax
+
+The Softmax kernel implementation provides an optimized implementation of the softmax operation, which is a fundamental component in neural networks for converting raw scores into probability distributions.
+
+The implementation shows notable speedups compared to the Softmax PyTorch implementation
+
+
 ### Sparsemax
 
 Sparsemax is a sparse alternative to softmax that produces sparse probability distributions. This kernel implements an efficient version of the sparsemax operation that can be used as a drop-in replacement for softmax in attention mechanisms or classification tasks.

{liger_kernel_nightly-0.5.10.dev20250528223524 → liger_kernel_nightly-0.5.10.dev20250531184114}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.10.dev20250528223524"
+version = "0.5.10.dev20250531184114"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.5.10.dev20250528223524 → liger_kernel_nightly-0.5.10.dev20250531184114}/src/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 → liger_kernel_nightly-0.5.10.dev20250531184114}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.5.10.dev20250528223524
+Version: 0.5.10.dev20250531184114
 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 → liger_kernel_nightly-0.5.10.dev20250531184114}/test/transformers/test_softmax.py

@@ -6,7 +6,7 @@ from test.utils import set_seed
 from test.utils import supports_bfloat16
 
 from liger_kernel.transformers.functional import liger_softmax
-from liger_kernel.transformers.softmax import LigerKernelSoftmax
+from liger_kernel.transformers.softmax import LigerSoftmax
 from liger_kernel.utils import infer_device
 
 device = infer_device()
@@ -47,7 +47,7 @@ def test_liger_softmax(shape, dtype, atol, rtol):
 
     torch_softmax = torch.nn.Softmax(dim=-1)
     ref_out = torch_softmax(x1)
-    liger_softmax = LigerKernelSoftmax().to(device).to(dtype)
+    liger_softmax = LigerSoftmax().to(device).to(dtype)
     liger_out = liger_softmax(x2)
 
     assert_verbose_allclose(ref_out, liger_out, atol=atol, rtol=rtol)