liger-kernel 0.5.10__py3-none-any.whl → 0.6.1__py3-none-any.whl

liger_kernel/ops/geglu.py

@@ -40,7 +40,7 @@ def _geglu_tanh_forward_kernel(a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE
     tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
     tanh_result = tanh(tanh_arg)
     geglu_a = 0.5 * a_row * (1 + tanh_result)
-    c_row = geglu_a * b_row
+    c_row = geglu_a.cast(b_row.dtype) * b_row
     tl.store(c + col_offsets, c_row, mask=mask)
 
 

liger_kernel/ops/layer_norm.py

@@ -1,4 +1,3 @@
-import math
 import operator
 
 import torch
@@ -43,30 +42,45 @@ def _layer_norm_forward_kernel(
     https://arxiv.org/abs/1607.06450
     https://github.com/karpathy/llm.c/blob/master/doc/layernorm/layernorm.md
     """
-    row_idx = tl.program_id(0)
+    row_idx = tl.program_id(0).to(tl.int64)
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    Y_ptr += row_idx * Y_row_stride
-    X_ptr += row_idx * X_row_stride
-    Mean_ptr += row_idx * Mean_row_stride
-    RSTD_ptr += row_idx * RSTD_row_stride
-
-    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
-    B_row = tl.load(B_ptr + col_offsets, mask=mask, other=0)
-
-    mean = tl.sum(X_row, axis=0) / n_cols
-    Xmm = tl.where(mask, X_row - mean, 0)
-    var = tl.sum(Xmm * Xmm, axis=0) / n_cols
+    # Pre-load weights and bias in fp32 to avoid repeated conversions
+    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
+    B_row = tl.load(B_ptr + col_offsets, mask=mask, other=0.0)
+    W_f32 = W_row.to(tl.float32)
+    B_f32 = B_row.to(tl.float32)
+
+    # Calculate pointers for this row
+    row_X_ptr = X_ptr + row_idx * X_row_stride
+    row_Y_ptr = Y_ptr + row_idx * Y_row_stride
+    row_Mean_ptr = Mean_ptr + row_idx * Mean_row_stride
+    row_RSTD_ptr = RSTD_ptr + row_idx * RSTD_row_stride
+
+    # Load input data and convert to fp32 for numerical stability
+    X_row = tl.load(row_X_ptr + col_offsets, mask=mask, other=0.0)
+    X_f32 = X_row.to(tl.float32)
+
+    # Compute statistics in fp32 for numerical stability
+    n_cols_f32 = n_cols.to(tl.float32)
+    mean = tl.sum(X_f32, axis=0) / n_cols_f32
+    X_centered = X_f32 - mean
+    # Apply mask to variance calculation to exclude contributions from masked elements
+    X_centered_masked = tl.where(mask, X_centered, 0.0)
+    var = tl.sum(X_centered_masked * X_centered_masked, axis=0) / n_cols_f32
     rstd = rsqrt(var + eps)
 
-    tl.store(Mean_ptr, mean)
-    tl.store(RSTD_ptr, rstd)
+    # Store statistics (convert back to original dtype only once)
+    tl.store(row_Mean_ptr, mean.to(X_row.dtype))
+    tl.store(row_RSTD_ptr, rstd.to(X_row.dtype))
 
-    Y_row = Xmm * rstd * W_row + B_row
+    # Fused normalization and affine transformation
+    # Y = (X - mean) * rstd * W + B = X_centered * rstd * W + B
+    Y_f32 = X_centered * rstd * W_f32 + B_f32
 
-    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
+    # Store output (single conversion back to original dtype)
+    tl.store(row_Y_ptr + col_offsets, Y_f32.to(X_row.dtype), mask=mask)
 
 
 @triton.jit
@@ -81,73 +95,87 @@ def _layer_norm_backward_kernel(
     DY_ptr,  # pointer to output grad, shape (n_rows, n_cols)
     stride_x,  # stride of each row in input
     stride_dx,  # stride of each row in input grad
-    stride_dw,  # stride of each row in weights grad
-    stride_db,  # stride of each row in bias grad
     stride_dy,  # stride of each row in output grad
-    n_rows,
     n_cols,
-    rows_per_program: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     dtype: tl.constexpr,
+    atomic_dtype: tl.constexpr,
 ):
     """
     References:
     https://arxiv.org/abs/1607.06450
     https://github.com/karpathy/llm.c/blob/master/doc/layernorm/layernorm.md
-    https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
-    https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/layer_norm.py
     """
-    row_block_id = tl.program_id(0)
-    row_start = row_block_id * rows_per_program
-    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
+    row_idx = tl.program_id(0).to(tl.int64)
     cols = tl.arange(0, BLOCK_SIZE)
     mask = cols < n_cols
 
-    dw_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
-    db_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
-
-    X_ptr += row_start * stride_x
-    Mean_ptr += row_start
-    RSTD_ptr += row_start
-    DX_ptr += row_start * stride_dx
-    DY_ptr += row_start * stride_dy
-
-    for _ in range(row_start, row_end):
-        x = tl.load(X_ptr + cols, mask=mask, other=0.0)
-        w = tl.load(W_ptr + cols, mask=mask, other=0.0)
-        dy = tl.load(DY_ptr + cols, mask=mask, other=0.0)
-        mean = tl.load(Mean_ptr)
-        rstd = tl.load(RSTD_ptr)
-
-        x_hat = (x - mean) * rstd
-        wdy = w * dy
-        c1 = tl.sum(x_hat * wdy, axis=0) / n_cols
-        c2 = tl.sum(wdy, axis=0) / n_cols
-        dx = (wdy - (x_hat * c1 + c2)) * rstd
-        tl.store(DX_ptr + cols, dx.to(dtype), mask=mask)
-
-        dw_row += dy * x_hat
-        db_row += dy
-
-        X_ptr += stride_x
-        Mean_ptr += 1
-        RSTD_ptr += 1
-        DX_ptr += stride_dx
-        DY_ptr += stride_dy
-
-    tl.store(DW_ptr + row_block_id * stride_dw + cols, dw_row.to(dtype), mask=mask)
-    tl.store(DB_ptr + row_block_id * stride_db + cols, db_row.to(dtype), mask=mask)
+    # Pre-load weights once (same optimization as forward pass)
+    w = tl.load(W_ptr + cols, mask=mask, other=0.0)
+    w_f32 = w.to(tl.float32)
+    n_cols_f32 = n_cols.to(tl.float32)
+
+    # Calculate pointers for this specific row
+    row_X_ptr = X_ptr + row_idx * stride_x
+    row_DX_ptr = DX_ptr + row_idx * stride_dx
+    row_DY_ptr = DY_ptr + row_idx * stride_dy
+    row_Mean_ptr = Mean_ptr + row_idx
+    row_RSTD_ptr = RSTD_ptr + row_idx
+
+    # Load data for this row
+    x = tl.load(row_X_ptr + cols, mask=mask, other=0.0)
+    dy = tl.load(row_DY_ptr + cols, mask=mask, other=0.0)
+    mean = tl.load(row_Mean_ptr)
+    rstd = tl.load(row_RSTD_ptr)
+
+    # Convert to fp32 for numerical stability
+    x_f32 = x.to(tl.float32)
+    dy_f32 = dy.to(tl.float32)
+    mean_f32 = mean.to(tl.float32)
+    rstd_f32 = rstd.to(tl.float32)
+
+    # Compute backward pass for this row
+    x_hat = (x_f32 - mean_f32) * rstd_f32
+    wdy = w_f32 * dy_f32
+    c1 = tl.sum(x_hat * wdy, axis=0) / n_cols_f32
+    c2 = tl.sum(wdy, axis=0) / n_cols_f32
+    dx = (wdy - (x_hat * c1 + c2)) * rstd_f32
+
+    # Store input gradient
+    tl.store(row_DX_ptr + cols, dx.to(dtype), mask=mask)
+
+    # Accumulate weight and bias gradients using atomic operations
+    dw = dy_f32 * x_hat
+    db = dy_f32
+    tl.atomic_add(DW_ptr + cols, dw.to(atomic_dtype), mask=mask)
+    tl.atomic_add(DB_ptr + cols, db.to(atomic_dtype), mask=mask)
 
 
 def layer_norm_forward(X, W, B, eps):
+    """
+    Args:
+        X: Input tensor of shape (..., hidden_size)
+        W: Weight tensor of shape (hidden_size,)
+        B: Bias tensor of shape (hidden_size,)
+        eps: Small constant for numerical stability
+
+    Returns:
+        Tuple of (output, input, mean, rstd, block_size, num_warps)
+    """
     shape = X.shape
     dim = shape[-1]
     X = X.view(-1, dim)
     n_rows, n_cols = X.shape
+
+    # Calculate optimal block size and warp configuration
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    # Allocate output tensors
     Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
     Mean = torch.empty(n_rows, dtype=X.dtype, device=X.device)
     RSTD = torch.empty(n_rows, dtype=X.dtype, device=X.device)
+
+    # Validate input dimensions
     if X.shape[1] != W.shape[0]:
         raise ValueError(
             f"Incompatible dimensions: input feature size (X.shape[1]={X.shape[1]}) "
@@ -159,7 +187,9 @@ def layer_norm_forward(X, W, B, eps):
     if X.device.type == "xpu":
         kernel_args["grf_mode"] = "large"
 
-    _layer_norm_forward_kernel[(n_rows,)](
+    # Launch kernel with one thread block per row for optimal performance
+    grid = (n_rows,)
+    _layer_norm_forward_kernel[grid](
         Y,
         Y.stride(0),
         X,
@@ -176,35 +206,43 @@ def layer_norm_forward(X, W, B, eps):
         eps,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
-        **kernel_args,  # XPU-specific optimization
+        **kernel_args,
     )
+
     return Y.view(*shape), X, Mean, RSTD, BLOCK_SIZE, num_warps
 
 
 def layer_norm_backward(dY, X, W, B, Mean, RSTD):
+    """
+    Args:
+        dY: Gradient of output
+        X: Input tensor
+        W: Weight tensor
+        B: Bias tensor
+        Mean: Pre-computed mean
+        RSTD: Pre-computed reciprocal standard deviation
+
+    Returns:
+        Tuple of (input_grad, weight_grad, bias_grad)
+    """
     shape = dY.shape
     dim = shape[-1]
     dY = dY.view(-1, dim)
     n_rows, n_cols = dY.shape
 
-    sm_count = 1
-    if X.device.type == "cuda":
-        sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
-    elif X.device.type == "xpu":
-        sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
-
+    # Allocate gradient tensors
     DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
-    _DW = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
-    _DB = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
+    # Use float32 for weight/bias gradients if bfloat16 (due to atomic_add limitation)
+    grad_dtype = torch.float32 if W.dtype == torch.bfloat16 else W.dtype
+    DW = torch.zeros(n_cols, dtype=grad_dtype, device=W.device)
+    DB = torch.zeros(n_cols, dtype=grad_dtype, device=W.device)
 
+    # Calculate optimal block size and warp configuration
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
     if n_cols > BLOCK_SIZE:
-        raise RuntimeError(
-            f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}. Consider using a smaller feature dimension."
-        )
+        raise RuntimeError(f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}.")
 
-    rows_per_program = math.ceil(n_rows / sm_count)
-    grid = (sm_count,)
+    # Determine dtype for triton operations
     triton_dtype = (
         tl.float32
         if X.dtype == torch.float32
@@ -212,41 +250,40 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
         if X.dtype == torch.bfloat16
         else tl.float16
         if X.dtype == torch.float16
-        else tl.float32  # fallback to float32 for other types
+        else tl.float32  # fallback
     )
 
+    # Use float32 for atomic operations if bfloat16 is not supported
+    atomic_dtype = tl.float32 if triton_dtype == tl.bfloat16 else triton_dtype
+
+    kernel_args = {"num_warps": num_warps}
     # XPU-specific optimization
-    kernel_args = {}
     if X.device.type == "xpu":
         kernel_args.update({"grf_mode": "large", "num_warps": 32, "num_stages": 4})
 
+    # Launch kernel with one thread block per row for optimal performance
+    grid = (n_rows,)
     _layer_norm_backward_kernel[grid](
         X,
         W,
         Mean,
         RSTD,
         DX,
-        _DW,
-        _DB,
+        DW,
+        DB,
         dY,
         X.stride(0),
         DX.stride(0),
-        _DW.stride(0),
-        _DB.stride(0),
         dY.stride(0),
-        n_rows,
         n_cols,
-        rows_per_program,
         BLOCK_SIZE=BLOCK_SIZE,
         dtype=triton_dtype,
-        **kernel_args,  # XPU-specific optimization
+        atomic_dtype=atomic_dtype,
+        **kernel_args,
     )
 
-    DW = _DW.sum(dim=0).to(W.dtype)
-    DB = _DB.sum(dim=0).to(W.dtype)
-
     DX = DX.view(*shape)
-    return DX, DW, DB
+    return DX, DW.to(W.dtype), DB.to(W.dtype)
 
 
 class LigerLayerNormFunction(torch.autograd.Function):

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)