liger-kernel-nightly 0.5.9.dev20250515034325__py3-none-any.whl → 0.5.9.dev20250515065336__py3-none-any.whl

liger_kernel/ops/sparsemax.py

@@ -0,0 +1,167 @@
+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 _sparsemax_forward_kernel(
+    x_ptr,
+    x_stride_row,
+    sorted_x_ptr,
+    sorted_x_stride_row,
+    o_ptr,
+    o_stride_row,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+    num_warps: tl.constexpr,
+):
+    pid_row = tl.program_id(0)
+    ptr_x_data_row = x_ptr + pid_row * x_stride_row
+    ptr_sorted_x_data_row = sorted_x_ptr + pid_row * sorted_x_stride_row
+    ptr_output_row = o_ptr + pid_row * o_stride_row
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < n_cols
+
+    z_sorted_block = tl.load(
+        ptr_sorted_x_data_row + offs,
+        mask=mask,
+        other=-float("inf"),
+        cache_modifier=".ca",
+    ).to(tl.float32)
+
+    z_valid = tl.where(mask, z_sorted_block, 0.0)
+    cssv = tl.cumsum(z_valid, 0)
+
+    r = (offs + 1).to(tl.float32)
+    safe_r = tl.where(mask, r, 1.0)
+
+    t_vec = (cssv - 1.0) / safe_r
+
+    support = (z_sorted_block > t_vec) & mask
+
+    k_int = tl.sum(support.to(tl.int32), 0)
+    k_clamped_int = tl.maximum(k_int, 1)
+    k = k_clamped_int.to(tl.float32)
+
+    s = tl.sum(tl.where(support, z_sorted_block, 0.0), 0)
+
+    tau = (s - 1.0) / k
+
+    x_block = tl.load(
+        ptr_x_data_row + offs,
+        mask=mask,
+        other=0.0,
+        cache_modifier=".ca",
+    ).to(tl.float32)
+
+    y = tl.maximum(x_block - tau, 0.0)
+
+    tl.store(
+        ptr_output_row + offs,
+        y.to(ptr_output_row.dtype.element_ty),
+        mask=mask,
+        cache_modifier=".cs",
+    )
+
+
+@triton.jit
+def _sparsemax_backward_kernel(
+    o_ptr, go_ptr, gi_ptr, stride, n_cols, BLOCK_SIZE: tl.constexpr, num_warps: tl.constexpr
+):
+    row = tl.program_id(0)
+    o_row = o_ptr + row * stride
+    go_row = go_ptr + row * stride
+    gi_row = gi_ptr + row * stride
+
+    offs = tl.arange(0, BLOCK_SIZE)
+
+    supp_cnt = tl.zeros((), tl.float32)
+    go_sum = tl.zeros((), tl.float32)
+
+    for i in tl.range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
+        offs_iter = i * BLOCK_SIZE + offs
+        mask_iter = offs_iter < n_cols
+        o_val = tl.load(o_row + offs_iter, mask=mask_iter, other=0.0, cache_modifier=".ca").to(tl.float32)
+        go_val = tl.load(go_row + offs_iter, mask=mask_iter, other=0.0).to(tl.float32)
+        supp = o_val > 0.0
+        go_sum += tl.sum(tl.where(supp, go_val, 0.0))
+        supp_cnt += tl.sum(supp.to(tl.float32))
+
+    for i in tl.range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
+        offs_iter = i * BLOCK_SIZE + offs
+        mask_iter = offs_iter < n_cols
+        o_val = tl.load(o_row + offs_iter, mask=mask_iter, other=0.0, cache_modifier=".ca").to(tl.float32)
+        go_val = tl.load(go_row + offs_iter, mask=mask_iter, other=0.0).to(tl.float32)
+        supp = o_val > 0.0
+        gi_val = tl.where(
+            supp,
+            go_val - tl.cast(go_sum / tl.maximum(supp_cnt, 1e-6), gi_row.dtype.element_ty).to(tl.float32),
+            0.0,
+        )
+        tl.store(gi_row + offs_iter, gi_val.to(gi_row.dtype.element_ty), mask=mask_iter, cache_modifier=".wb")
+
+
+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,
+        )
+
+        ctx.save_for_backward(out_flat)
+        return out_flat.view_as(x_sw).transpose(dim, -1)
+
+    @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

liger_kernel/transformers/functional.py

@@ -12,6 +12,7 @@ from liger_kernel.ops.layer_norm import LigerLayerNormFunction
 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.sparsemax import LigerSparsemaxFunction
 from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from liger_kernel.ops.tvd import LigerTVDLossFunction
 
@@ -159,6 +160,13 @@ def liger_kl_div(
     )
 
 
+def liger_sparsemax(
+    input,
+    dim: int = -1,
+):
+    return LigerSparsemaxFunction.apply(input, dim)
+
+
 def liger_tvd(
     input,
     target,

liger_kernel/transformers/sparsemax.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.sparsemax import LigerSparsemaxFunction
+
+
+class LigerSparsemax(nn.Module):
+    def __init__(self, dim: int = -1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return LigerSparsemaxFunction.apply(x, self.dim)
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}"

{liger_kernel_nightly-0.5.9.dev20250515034325.dist-info → liger_kernel_nightly-0.5.9.dev20250515065336.dist-info}/METADATA

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

{liger_kernel_nightly-0.5.9.dev20250515034325.dist-info → liger_kernel_nightly-0.5.9.dev20250515065336.dist-info}/RECORD

@@ -28,6 +28,7 @@ liger_kernel/ops/layer_norm.py,sha256=vWCyOm-F2GMAilB-ozJcFeUQQLCJoTE_uiXq-_0uYu
 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/swiglu.py,sha256=KmgMjaJQnbLLgZn2nEpbwHU_xpnYRweCyrLQSVvM1vA,3015
 liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
 liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
@@ -37,7 +38,7 @@ liger_kernel/transformers/__init__.py,sha256=0KX0rxyy0E_uNWVE0PSTzEVzKqc5KdFHtvd
 liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
 liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
 liger_kernel/transformers/dyt.py,sha256=QMqqc14pkE0WhpRZvapfnNAun-6C0C_tHExL2ZJuCUA,648
-liger_kernel/transformers/functional.py,sha256=4h9Pdx_iINBqfv2Zod_c27qOpYXDDwbdVgatQ9_XBmI,5089
+liger_kernel/transformers/functional.py,sha256=2YBfvtdU1GRZuRpJhHgJXeGYa1RvmO6-qQvrKQrLJK4,5259
 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
@@ -50,6 +51,7 @@ liger_kernel/transformers/monkey_patch.py,sha256=k8WIkx_f3ObG6TjhIiN_4KeOABurB2W
 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/sparsemax.py,sha256=0lQA0UEOs4mu8CMruZ3VLhImxQVXJWhPsAKUsYA7vj8,403
 liger_kernel/transformers/swiglu.py,sha256=LZ8YeLIdv2k46JleZMjzubGk98smt6t780kSgcVLsQk,3454
 liger_kernel/transformers/trainer_integration.py,sha256=W3ON51O5GkyzNJsItz0y5rKx-uy2f2cFfveZpqbUdhw,123
 liger_kernel/transformers/tvd.py,sha256=XrRfyJIqN6HFxXk8MYyFVZM1OLz3mtSbRZvWfZ_JerQ,450
@@ -77,9 +79,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
 liger_kernel/transformers/trainer/orpo_trainer.py,sha256=pdekW7l6Qg_aqa5SYKYlSWUF8m3lkOFvFLcIMEHrz9s,8338
 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.9.dev20250515034325.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
-liger_kernel_nightly-0.5.9.dev20250515034325.dist-info/METADATA,sha256=4us0oTXgw5IL_J6GV3yJjrjr_XwkxXt0wVokauQAYIU,23874
-liger_kernel_nightly-0.5.9.dev20250515034325.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
-liger_kernel_nightly-0.5.9.dev20250515034325.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-liger_kernel_nightly-0.5.9.dev20250515034325.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
-liger_kernel_nightly-0.5.9.dev20250515034325.dist-info/RECORD,,
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/METADATA,sha256=IK7MV888DLovn85_Xto_NFKgXq4SILvZB7HDXyeP2uc,23874
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel_nightly-0.5.9.dev20250515065336.dist-info/RECORD,,