liger-kernel-nightly 0.5.10.dev20250524022630__py3-none-any.whl → 0.5.10.dev20250526154149__py3-none-any.whl

liger_kernel/ops/multi_token_attention.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from torch.nn.modules.utils import _pair
+
+from liger_kernel.ops.softmax import _softmax_forward
+from liger_kernel.ops.sparsemax import _sparsemax_backward
+from liger_kernel.ops.sparsemax import _sparsemax_forward
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+@triton.jit
+def _mask_fwd_kernel(
+    scores_ptr,
+    out_ptr,
+    stride_b,
+    stride_m,
+    stride_n,
+    L,
+    mask_val: tl.constexpr,
+    BLOCK: tl.constexpr,
+    num_warps: tl.constexpr,
+):
+    row_block = tl.program_id(0)
+    col_block = tl.program_id(1)
+    batch_id = tl.program_id(2)
+
+    row_idx = row_block * BLOCK + tl.arange(0, BLOCK)
+    col_idx = col_block * BLOCK + tl.arange(0, BLOCK)
+    in_bounds = (row_idx[:, None] < L) & (col_idx[None, :] < L)
+
+    base = scores_ptr + batch_id * stride_b
+    offs = row_idx[:, None] * stride_m + col_idx[None, :] * stride_n
+    future = col_idx[None, :] > row_idx[:, None]
+    mask_load = in_bounds & ~future
+    out = tl.load(base + offs, mask=mask_load, other=mask_val, cache_modifier=".ca")
+    tl.store(out_ptr + batch_id * stride_b + offs, out, mask=in_bounds, cache_modifier=".cs")
+
+
+@triton.jit
+def _mask_bwd_kernel(
+    grad_in_ptr, out_ptr, stride_b, stride_m, stride_n, L, BLOCK: tl.constexpr, num_warps: tl.constexpr
+):
+    row_block = tl.program_id(0)
+    col_block = tl.program_id(1)
+    batch_id = tl.program_id(2)
+
+    row_idx = row_block * BLOCK + tl.arange(0, BLOCK)
+    col_idx = col_block * BLOCK + tl.arange(0, BLOCK)
+    in_bounds = (row_idx[:, None] < L) & (col_idx[None, :] < L)
+
+    base = grad_in_ptr + batch_id * stride_b
+    offs = row_idx[:, None] * stride_m + col_idx[None, :] * stride_n
+    grad_vals = tl.load(base + offs, mask=in_bounds, other=0.0, cache_modifier=".ca")
+
+    future = col_idx[None, :] > row_idx[:, None]
+    zero = tl.zeros(grad_vals.shape, dtype=grad_vals.dtype)
+    out = tl.where(future, zero, grad_vals)
+
+    tl.store(out_ptr + batch_id * stride_b + offs, out, mask=in_bounds, cache_modifier=".wb")
+
+
+def _mask_inf_forward(scores: torch.Tensor) -> torch.Tensor:
+    *batch, L, _ = scores.shape
+    N = int(torch.prod(torch.tensor(batch))) if batch else 1
+    scores_f = scores.view(N, L, L)
+    out = torch.empty_like(scores_f)
+
+    sb, sm, sn = scores_f.stride(0), scores_f.stride(1), scores_f.stride(2)
+    BLOCK_SIZE, num_warps = calculate_settings(L)
+    grid = (triton.cdiv(L, BLOCK_SIZE), triton.cdiv(L, BLOCK_SIZE), N)
+    _mask_fwd_kernel[grid](scores_f, out, sb, sm, sn, L, mask_val=-1e9, BLOCK=BLOCK_SIZE, num_warps=num_warps)
+    return out.view(*batch, L, L)
+
+
+def _mask_inf_backward(grad: torch.Tensor) -> torch.Tensor:
+    *batch, L, _ = grad.shape
+    N = int(torch.prod(torch.tensor(batch))) if batch else 1
+    grad_f = grad.view(N, L, L)
+    out = torch.empty_like(grad_f)
+
+    sb, sm, sn = grad_f.stride(0), grad_f.stride(1), grad_f.stride(2)
+    BLOCK_SIZE, num_warps = calculate_settings(L)
+    grid = (triton.cdiv(L, BLOCK_SIZE), triton.cdiv(L, BLOCK_SIZE), N)
+    _mask_bwd_kernel[grid](grad_f, out, sb, sm, sn, L, BLOCK=BLOCK_SIZE, num_warps=num_warps)
+    return out.view(*batch, L, L)
+
+
+def _mask_zero_forward(scores: torch.Tensor) -> torch.Tensor:
+    *batch, L, _ = scores.shape
+    N = int(torch.prod(torch.tensor(batch))) if batch else 1
+    scores_f = scores.view(N, L, L)
+    out = torch.empty_like(scores_f)
+
+    sb, sm, sn = scores_f.stride(0), scores_f.stride(1), scores_f.stride(2)
+    BLOCK_SIZE, num_warps = calculate_settings(L)
+    grid = (triton.cdiv(L, BLOCK_SIZE), triton.cdiv(L, BLOCK_SIZE), N)
+    _mask_fwd_kernel[grid](scores_f, out, sb, sm, sn, L, mask_val=0.0, BLOCK=BLOCK_SIZE, num_warps=num_warps)
+    return out.view(*batch, L, L)
+
+
+def _mask_zero_backward(grad: torch.Tensor) -> torch.Tensor:
+    *batch, L, _ = grad.shape
+    N = int(torch.prod(torch.tensor(batch))) if batch else 1
+    grad_f = grad.view(N, L, L)
+    out = torch.empty_like(grad_f)
+
+    sb, sm, sn = grad_f.stride(0), grad_f.stride(1), grad_f.stride(2)
+    BLOCK_SIZE, num_warps = calculate_settings(L)
+    grid = (triton.cdiv(L, BLOCK_SIZE), triton.cdiv(L, BLOCK_SIZE), N)
+    _mask_bwd_kernel[grid](grad_f, out, sb, sm, sn, L, BLOCK=BLOCK_SIZE, num_warps=num_warps)
+    return out.view(*batch, L, L)
+
+
+class LigerMultiTokenAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, scores, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, sparse=False):
+        scores_inf = _mask_inf_forward(scores)
+
+        out_flat_sparse = None
+        activation_output = None
+
+        ctx.sparse = sparse
+
+        if sparse:
+            if scores_inf.dtype != torch.float32:
+                raise RuntimeError("Liger sparse multi-token attention currently only supports fp32 input scores")
+            probs_sparse, out_flat_sparse = _sparsemax_forward(scores_inf, dim=-1)
+            activation_output = probs_sparse
+            ctx.save_for_backward(scores_inf, activation_output, out_flat_sparse, weight, bias)
+            ctx.out_flat_sparse_saved = True
+        else:
+            probs_softmax, _, _, _ = _softmax_forward(scores_inf)
+            activation_output = probs_softmax
+            ctx.save_for_backward(scores_inf, activation_output, weight, bias)
+            ctx.out_flat_sparse_saved = False
+
+        out_conv = F.conv2d(
+            activation_output,
+            weight,
+            bias,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+
+        out = _mask_zero_forward(out_conv)
+
+        ctx.stride = _pair(stride)
+        ctx.padding = _pair(padding)
+        ctx.dilation = _pair(dilation)
+        ctx.groups = groups
+        ctx.dim = -1
+
+        return out
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_out):
+        if ctx.out_flat_sparse_saved:
+            scores_inf, activation_output, out_flat_sparse, weight, bias = ctx.saved_tensors
+        else:
+            scores_inf, activation_output, weight, bias = ctx.saved_tensors
+            out_flat_sparse = None
+
+        use_sparsemax = ctx.sparse
+        dim = ctx.dim
+        stride, padding, dilation, groups = (ctx.stride, ctx.padding, ctx.dilation, ctx.groups)
+
+        grad_conv = _mask_zero_backward(grad_out)
+
+        grad_probs = F.conv_transpose2d(
+            grad_conv, weight, None, stride=stride, padding=padding, dilation=dilation, groups=groups
+        )
+
+        grad_weight = torch.nn.grad.conv2d_weight(
+            input=activation_output,
+            weight_size=weight.shape,
+            grad_output=grad_conv,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+        grad_bias = None
+        if bias is not None:
+            grad_bias = grad_conv.sum(dim=(0, 2, 3))
+
+        grad_scores_inf = None
+        if use_sparsemax:
+            if not ctx.out_flat_sparse_saved or out_flat_sparse is None:
+                raise RuntimeError("Internal error: Sparse flag is set but sparse tensor was not saved.")
+            grad_scores_inf = _sparsemax_backward(grad_probs, out_flat_sparse, dim=dim)
+        else:
+            grad_probs_cont = grad_probs
+            probs_cont = activation_output
+            dot = (grad_probs_cont * probs_cont).sum(dim=-1, keepdim=True)
+            grad_scores_inf = probs_cont * (grad_probs_cont - dot)
+
+        grad_scores = _mask_inf_backward(grad_scores_inf)
+
+        return (grad_scores, grad_weight, grad_bias, None, None, None, None, None)

liger_kernel/ops/softmax.py

@@ -0,0 +1,201 @@
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+@triton.jit
+def _softmax_single_block_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < n_cols
+
+    x = tl.load(X_ptr + row_id * X_row_stride + offs, mask=mask, other=-float("inf"), cache_modifier=".ca")
+    m = tl.max(x, axis=0)
+    e = tl.exp(x - m)
+    d = tl.sum(e, axis=0)
+    y = e / d
+    tl.store(Y_ptr + row_id * Y_row_stride + offs, y, mask=mask, cache_modifier=".cs")
+
+
+@triton.jit
+def _softmax_multi_block_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+
+    m = tl.float32(-float("inf"))
+    d = tl.float32(0.0)
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        xblk = tl.load(X_ptr + row_id * X_row_stride + idx, mask=mask, other=-float("inf"), cache_modifier=".ca")
+        blk_max = tl.max(xblk, axis=0)
+        new_m = tl.max(m, blk_max)
+        d = d * tl.exp(m - new_m) + tl.sum(tl.exp(xblk - new_m), axis=0)
+        m = new_m
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        xblk = tl.load(X_ptr + row_id * X_row_stride + idx, mask=mask, other=-float("inf"), cache_modifier=".ca")
+        yblk = tl.exp(xblk - m) / d
+        tl.store(Y_ptr + row_id * Y_row_stride + idx, yblk, mask=mask, cache_modifier=".cs")
+
+
+@triton.jit
+def _softmax_single_block_backward_kernel(
+    dy_ptr,
+    dy_stride,
+    y_ptr,
+    y_stride,
+    dx_ptr,
+    dx_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < n_cols
+
+    dy = tl.load(dy_ptr + row_id * dy_stride + offs, mask=mask, other=0.0)
+    y = tl.load(y_ptr + row_id * y_stride + offs, mask=mask, other=0.0, cache_modifier=".ca")
+    dot = tl.sum(dy * y, axis=0)
+    dx = y * (dy - dot)
+    tl.store(dx_ptr + row_id * dx_stride + offs, dx, mask=mask, cache_modifier=".wb")
+
+
+@triton.jit
+def _softmax_multi_block_backward_kernel(
+    dy_ptr,
+    dy_stride,
+    y_ptr,
+    y_stride,
+    dx_ptr,
+    dx_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_id = tl.program_id(0)
+    offs = tl.arange(0, BLOCK_SIZE)
+    acc = tl.float32(0.0)
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        dy_blk = tl.load(dy_ptr + row_id * dy_stride + idx, mask=mask, other=0.0)
+        y_blk = tl.load(y_ptr + row_id * y_stride + idx, mask=mask, other=0.0, cache_modifier=".ca")
+        acc += tl.sum(dy_blk * y_blk, axis=0)
+
+    for start in tl.range(0, n_cols, BLOCK_SIZE):
+        idx = start + offs
+        mask = idx < n_cols
+        dy_blk = tl.load(dy_ptr + row_id * dy_stride + idx, mask=mask, other=0.0)
+        y_blk = tl.load(y_ptr + row_id * y_stride + idx, mask=mask, other=0.0, cache_modifier=".ca")
+        dx_blk = y_blk * (dy_blk - acc)
+        tl.store(dx_ptr + row_id * dx_stride + idx, dx_blk, mask=mask, cache_modifier=".wb")
+
+
+def _softmax_forward(x: torch.Tensor) -> Tuple[torch.Tensor, int, int, bool]:
+    *batch, n_cols = x.shape
+    x2d = x.contiguous().view(-1, n_cols)
+    n_rows = x2d.shape[0]
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    y2d = torch.empty_like(x2d)
+
+    if n_cols <= BLOCK_SIZE:
+        _softmax_single_block_forward_kernel[(n_rows,)](
+            y2d, y2d.stride(0), x2d, x2d.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
+        )
+        multi_block_launch = False
+    else:
+        _softmax_multi_block_forward_kernel[(n_rows,)](
+            y2d, y2d.stride(0), x2d, x2d.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
+        )
+        multi_block_launch = True
+
+    return y2d.view(*batch, n_cols), BLOCK_SIZE, num_warps, multi_block_launch
+
+
+def _softmax_backward(
+    dy: torch.Tensor,
+    y: torch.Tensor,
+    BLOCK_SIZE: int,
+    num_warps: int,
+    multi_block_launch: bool,
+) -> torch.Tensor:
+    *batch, n_cols = dy.shape
+    dy2d = dy.contiguous().view(-1, n_cols)
+    y2d = y.contiguous().view(-1, n_cols)
+    n_rows = dy2d.shape[0]
+    dx2d = torch.empty_like(dy2d)
+
+    if not multi_block_launch and n_cols <= BLOCK_SIZE:
+        _softmax_single_block_backward_kernel[(n_rows,)](
+            dy2d,
+            dy2d.stride(0),
+            y2d,
+            y2d.stride(0),
+            dx2d,
+            dx2d.stride(0),
+            n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+    else:
+        _softmax_multi_block_backward_kernel[(n_rows,)](
+            dy2d,
+            dy2d.stride(0),
+            y2d,
+            y2d.stride(0),
+            dx2d,
+            dx2d.stride(0),
+            n_cols,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=num_warps,
+        )
+
+    return dx2d.view(*batch, n_cols)
+
+
+class LigerSoftmaxFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, input_: torch.Tensor):
+        y, BLOCK_SIZE, num_warps, multi_block_launch = _softmax_forward(input_)
+        ctx.save_for_backward(y)
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        ctx.num_warps = num_warps
+        ctx.multi_block_launch = multi_block_launch
+        return y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        (y,) = ctx.saved_tensors
+        dx = _softmax_backward(
+            grad_output,
+            y,
+            ctx.BLOCK_SIZE,
+            ctx.num_warps,
+            ctx.multi_block_launch,
+        )
+        return dx

liger_kernel/ops/sparsemax.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 import triton
 import triton.language as tl
@@ -105,63 +107,73 @@ def _sparsemax_backward_kernel(
         tl.store(gi_row + offs_iter, gi_val.to(gi_row.dtype.element_ty), mask=mask_iter, cache_modifier=".wb")
 
 
+def _sparsemax_forward(x: torch.Tensor, dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    if dim < 0:
+        dim += x.dim()
+    x_sw = x.transpose(dim, -1).contiguous()
+    n_cols = x_sw.size(-1)
+    n_rows = x_sw.numel() // n_cols
+    x_flat = x_sw.view(n_rows, n_cols)
+    x_sorted_flat = torch.sort(x_flat.float(), dim=-1, descending=True).values
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    out_flat = torch.empty_like(x_flat)
+    grid = (n_rows,)
+    _sparsemax_forward_kernel[grid](
+        x_flat,
+        x_flat.stride(0),
+        x_sorted_flat,
+        x_sorted_flat.stride(0),
+        out_flat,
+        out_flat.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    y = out_flat.view_as(x_sw).transpose(dim, -1)
+    return y, out_flat
+
+
+def _sparsemax_backward(
+    grad_out: torch.Tensor,
+    out_flat: torch.Tensor,
+    dim: int,
+) -> torch.Tensor:
+    grad_sw = grad_out.transpose(dim, -1).contiguous()
+    n_cols = grad_sw.size(-1)
+    n_rows = grad_sw.numel() // n_cols
+    go_flat = grad_sw.view(n_rows, n_cols)
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    dx_flat = torch.empty_like(go_flat)
+    grid = (n_rows,)
+    _sparsemax_backward_kernel[grid](
+        out_flat,
+        go_flat,
+        dx_flat,
+        out_flat.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    dx = dx_flat.view_as(grad_sw).transpose(dim, -1)
+    return dx
+
+
 class LigerSparsemaxFunction(torch.autograd.Function):
     @staticmethod
     @ensure_contiguous
     def forward(ctx, x: torch.Tensor, dim: int):
-        if dim < 0:
-            dim += x.dim()
-        ctx.dim = dim
-
-        x_sw = x.transpose(dim, -1).contiguous()
-        n_cols = x_sw.size(-1)
-        n_rows = x_sw.numel() // n_cols
-        x_flat = x_sw.view(n_rows, n_cols)
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-        out_flat = torch.empty_like(x_flat)
-        grid = (n_rows,)
-
-        x_sorted_flat = torch.sort(x_flat.float(), dim=-1, descending=True).values
-
-        _sparsemax_forward_kernel[grid](
-            x_flat,
-            x_flat.stride(0),
-            x_sorted_flat,
-            x_sorted_flat.stride(0),
-            out_flat,
-            out_flat.stride(0),
-            n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
+        y, out_flat = _sparsemax_forward(x, dim)
         ctx.save_for_backward(out_flat)
-        return out_flat.view_as(x_sw).transpose(dim, -1)
+        ctx.dim = dim
+        return y
 
     @staticmethod
     @ensure_contiguous
     def backward(ctx, grad_out: torch.Tensor):
         (out_flat,) = ctx.saved_tensors
-        dim = ctx.dim
-
-        go_sw = grad_out.transpose(dim, -1).contiguous()
-        n_cols = go_sw.size(-1)
-        n_rows = go_sw.numel() // n_cols
-        go_flat = go_sw.view(n_rows, n_cols)
-
-        BLOCK_SIZE, num_warps = calculate_settings(n_cols)
-        gi_flat = torch.empty_like(go_flat)
-        grid = (n_rows,)
-
-        _sparsemax_backward_kernel[grid](
-            out_flat,
-            go_flat,
-            gi_flat,
-            out_flat.stride(0),
-            n_cols,
-            BLOCK_SIZE=BLOCK_SIZE,
-            num_warps=num_warps,
-        )
-
-        return gi_flat.view_as(go_sw).transpose(dim, -1), None
+        dx = _sparsemax_backward(grad_out, out_flat, ctx.dim)
+        return dx, None

liger_kernel/transformers/functional.py

@@ -9,9 +9,11 @@ from liger_kernel.ops.group_norm import LigerGroupNormFunction
 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.qwen2vl_mrope import LigerQwen2VLMRopeFunction
 from liger_kernel.ops.rms_norm import LigerRMSNormFunction
 from liger_kernel.ops.rope import LigerRopeFunction
+from liger_kernel.ops.softmax import LigerSoftmaxFunction
 from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
 from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.tvd import LigerTVDLossFunction
@@ -167,6 +169,34 @@ def liger_sparsemax(
     return LigerSparsemaxFunction.apply(input, dim)
 
 
+def liger_multi_token_attention(
+    scores,
+    weight,
+    bias=None,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    sparse: bool = False,
+):
+    """
+    Functional interface for multi-token attention.
+
+    Args:
+        scores: Input tensor of shape (B, C_in, L, L)
+        weight: Convolution weight tensor of shape (C_out, C_in // groups, K, K)
+        bias: Optional bias tensor of shape (C_out,)
+        stride: Stride for the convolution (default: 1)
+        padding: Padding for the convolution (default: 0)
+        dilation: Dilation factor for the convolution (default: 1)
+        groups: Number of groups for the convolution (default: 1)
+        sparse: Specifies if input tensors are expected to be sparse (default: False)
+    Returns:
+        Output tensor after applying multi-token attention.
+    """
+    return LigerMultiTokenAttentionFunction.apply(scores, weight, bias, stride, padding, dilation, groups, sparse)
+
+
 def liger_tvd(
     input,
     target,
@@ -203,5 +233,9 @@ def liger_swiglu(a, b):
     return LigerSiLUMulFunction.apply(a, b)
 
 
+def liger_softmax(x):
+    return LigerSoftmaxFunction.apply(x)
+
+
 def liger_dyt(x, alpha, gamma, beta):
     return LigerDyTFunction.apply(x, alpha, gamma, beta)

liger_kernel/transformers/multi_token_attention.py

@@ -0,0 +1,64 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from torch.nn.modules.utils import _pair
+
+from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
+
+
+class LigerMultiTokenAttention(nn.Module):
+    """
+    Multi-Token Attention:
+        out = mask_{0}(conv2d(softmax(mask_{-\inf}(scores))))
+
+    Reference: https://arxiv.org/pdf/2504.00927
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        sparse: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.padding = _pair(padding)
+        self.dilation = _pair(dilation)
+        self.groups = groups
+        self.sparse = sparse
+
+        self.weight = nn.Parameter(torch.empty(out_channels, in_channels // groups, *self.kernel_size))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_channels))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, scores: torch.Tensor) -> torch.Tensor:
+        return LigerMultiTokenAttentionFunction.apply(
+            scores,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.sparse,
+        )

liger_kernel/transformers/softmax.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.softmax import LigerSoftmaxFunction
+
+
+class LigerKernelSoftmax(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        return LigerSoftmaxFunction.apply(x)

{liger_kernel_nightly-0.5.10.dev20250524022630.dist-info → liger_kernel_nightly-0.5.10.dev20250526154149.dist-info}/METADATA

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

{liger_kernel_nightly-0.5.10.dev20250524022630.dist-info → liger_kernel_nightly-0.5.10.dev20250526154149.dist-info}/RECORD

@@ -26,10 +26,12 @@ liger_kernel/ops/grpo_loss.py,sha256=anRnv7k1-AV3pCC6_TqP0GMg78YYUfRAJrbpx6PVhl0
 liger_kernel/ops/jsd.py,sha256=onHp5T3MbvJaVz5Vup7Ww6EQp_HTaZeayTjJk6FgQMY,7042
 liger_kernel/ops/kl_div.py,sha256=ZjGdDLKWksHT9dZ0xF_TDgAkj5cuMTwwT5tr9E-_24o,8734
 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=PP27OIBmV9By63i13jot9ylDowW0nuxY_JFIkaPLgL4,12078
 liger_kernel/ops/rope.py,sha256=ofmBOkUpZZO-Q8Z5B_LOFYYLD-YT-8WnJ4vGOrDYouI,8943
-liger_kernel/ops/sparsemax.py,sha256=t7JWIyzq1piikXUufayFzsfkzVaCYU-hXPuMs7839pk,4850
+liger_kernel/ops/softmax.py,sha256=tgORx6MK1IDDtZKqGarj0IPIVjqAIEUXXYPiinhRdtI,5864
+liger_kernel/ops/sparsemax.py,sha256=AeWe1xgkHJFEKWTj2vu_0hj7LztGvjqXAps-QTpCY0U,5087
 liger_kernel/ops/swiglu.py,sha256=KmgMjaJQnbLLgZn2nEpbwHU_xpnYRweCyrLQSVvM1vA,3015
 liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
 liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
@@ -40,7 +42,7 @@ liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawX
 liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
 liger_kernel/transformers/dyt.py,sha256=i-4GPaMrl-jab9TVI5qN0-H9qycn_mCbV82ozU4nbmU,723
 liger_kernel/transformers/fsdp.py,sha256=CUiyjTmjkjY7pLXQv8ly9rnzgXw6529csd9pvtJNMYc,3096
-liger_kernel/transformers/functional.py,sha256=2YBfvtdU1GRZuRpJhHgJXeGYa1RvmO6-qQvrKQrLJK4,5259
+liger_kernel/transformers/functional.py,sha256=QmnAFpRgIbp9Rzlfp8QibwiEbf5BUcANxfY68an7o8c,6444
 liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=O8Sg5BT81nTaY9fSGoOY9dOD9ekibwwiuXhdUHaxntQ,1742
 liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
 liger_kernel/transformers/geglu.py,sha256=mrgqzIUVd6lN7fkDKLkw5YaESDxDtFgbot430WwPVOQ,1107
@@ -51,9 +53,11 @@ 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/monkey_patch.py,sha256=DKv5-4KyXLiVhAJ9WVFv1I1i1DzjaudTrhqx6EVYViU,74505
+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=GqCEJuGt0YdqqlMcToE0Wp4A8YFquDa4UUSyH2uFW2A,1191
 liger_kernel/transformers/rope.py,sha256=ZTrTORSAyfcFIKjk6XEeYmk4ROH7xXED9L4g2NFntlE,999
+liger_kernel/transformers/softmax.py,sha256=u7bFo35-cjaAm9of6-DLzmkaNFELOM-9AgyrcvUPifw,270
 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
@@ -82,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.dev20250524022630.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.10.dev20250524022630.dist-info/METADATA,sha256=kjNZA78siLFmGihrwvXrUIiLnTNaJoZglOkbtihcynk,24113
-liger_kernel_nightly-0.5.10.dev20250524022630.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.10.dev20250524022630.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.10.dev20250524022630.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.10.dev20250524022630.dist-info/RECORD,,
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/METADATA,sha256=0CXMJx6ef3SurofjUlAWwMaj-prwFf-xg_nMo-n6UPE,24113
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.10.dev20250526154149.dist-info/RECORD,,