liger-kernel 0.1.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

liger_kernel/ops/kl_div.py

@@ -0,0 +1,258 @@
+from typing import Literal
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+def get_num_warps(BLOCK_SIZE):
+    num_warps = 4
+    if BLOCK_SIZE >= 32768:
+        num_warps = 32
+    elif BLOCK_SIZE >= 8192:
+        num_warps = 16
+    elif BLOCK_SIZE >= 2048:
+        num_warps = 8
+
+    return num_warps
+
+
+MAX_FUSED_SIZE = 65536 // 4  # 65536 // 4 or 8 works the best
+
+REDUCTION_LITERAL = Literal["none", "sum", "mean", "batchmean"]
+
+_REDUCTION_MODE_NONE = tl.constexpr(0)
+_REDUCTION_MODE_SUM = tl.constexpr(1)
+_REDUCTION_MODE_MEAN = tl.constexpr(2)
+_REDUCTION_MODE_BATCHMEAN = tl.constexpr(3)
+
+_str_to_reduction_mode = {
+    "none": _REDUCTION_MODE_NONE.value,
+    "sum": _REDUCTION_MODE_SUM.value,
+    "mean": _REDUCTION_MODE_MEAN.value,
+    "batchmean": _REDUCTION_MODE_BATCHMEAN.value,
+}
+
+
+@triton.jit
+def _kldiv_kernel_forward(
+    y_ptr,  # [B, S], prediction ptr, the kernel expects the prediction in log-space
+    y_stride,  # int, prediction stride
+    gt_ptr,  # [B, S], ground truth ptr
+    gt_stride,  # int, ground truth stride
+    loss_ptr,  # [B] or [B, S] if reduction == _REDUCTION_MODE_NONE, output ptr
+    loss_stride,  # int, output stride
+    n_cols,  # int, number of columns in the input tensor
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+    log_target: tl.constexpr = False,
+    reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    y_ptr += pid * y_stride
+    gt_ptr += pid * gt_stride
+    loss_ptr += pid * loss_stride
+
+    base_offsets = tl.arange(0, BLOCK_SIZE)
+
+    loss_sum = 0.0
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + base_offsets
+        mask = offsets < n_cols
+        y = tl.load(y_ptr + offsets, mask=mask, other=0.0)
+        y_true = tl.load(gt_ptr + offsets, mask=mask, other=0.0)
+
+        # KL(y_true || y) = y_true * (log(y_true) - log(y))
+        # We compute KL(y_true || y) with y in the log-space
+        if not log_target:
+            loss = y_true * (tl.log(tl.maximum(y_true, eps)) - y)
+        else:
+            loss = tl.exp(y_true) * (y_true - y)
+
+        if reduction == _REDUCTION_MODE_NONE:
+            tl.store(loss_ptr + offsets, loss, mask=mask)
+        else:
+            loss_sum += tl.sum(loss, axis=0)
+
+    if reduction != _REDUCTION_MODE_NONE:
+        tl.store(loss_ptr, loss_sum)
+
+
+@triton.jit
+def _kldiv_kernel_backward(
+    target_ptr,
+    target_stride,
+    new_grads_ptr,
+    new_grads_stride,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+    log_target: tl.constexpr = False,
+):
+    pid = tl.program_id(0).to(tl.int64)
+
+    target_ptr += pid * target_stride
+    new_grads_ptr += pid * new_grads_stride
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_cols
+
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_cols
+
+        target = tl.load(target_ptr + offsets, mask=mask, other=0.0)
+
+        if not log_target:
+            res = target * -1
+        else:
+            res = -tl.exp(target)
+
+        tl.store(new_grads_ptr + offsets, res, mask=mask)
+
+
+def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
+    BT, V = y_pred.shape
+
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    num_warps = get_num_warps(BLOCK_SIZE)
+
+    grid = (BT,)
+    reduction = _str_to_reduction_mode[reduction]
+
+    out_size = (BT, V) if reduction == _REDUCTION_MODE_NONE.value else (BT,)
+    output_tensor = torch.zeros(out_size, device=y_pred.device, dtype=torch.float32)
+
+    _kldiv_kernel_forward[grid](
+        y_pred,
+        y_pred.stride(0),
+        y_true,
+        y_true.stride(0),
+        output_tensor,
+        output_tensor.stride(0),
+        V,
+        eps=eps,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        log_target=log_target,
+        reduction=reduction,
+    )
+
+    # calculated according to the reduction mode same as in Pytorch. In the later versions, `mean` will be changed to the same behavior as `batchmean`
+    # https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html
+    # https://github.com/pytorch/pytorch/blob/d7b57c4d63edb42e1deeeba9497fcb5f1f748ff2/torch/nn/functional.py#L3372
+    if reduction == _REDUCTION_MODE_BATCHMEAN.value:
+        return output_tensor.sum() / BT
+    elif reduction == _REDUCTION_MODE_SUM.value:
+        return output_tensor.sum(dim=0)
+    elif reduction == _REDUCTION_MODE_MEAN.value:
+        return output_tensor.sum() / (BT * V)
+    else:
+        return output_tensor
+
+
+def kldiv_backward_triton(target, grad_output, new_grads, log_target):
+    BT, V = target.shape
+
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    num_warps = get_num_warps(BLOCK_SIZE)
+
+    grid = (BT,)
+
+    # We store the gradients in-place in the input tensor
+    _kldiv_kernel_backward[grid](
+        target,
+        target.stride(0),
+        new_grads,
+        new_grads.stride(0),
+        V,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        log_target=log_target,
+    )
+
+    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul then.
+    if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        return new_grads
+
+    return new_grads * grad_output
+
+
+class LigerKLDivLossFunction(torch.autograd.Function):
+    """
+    Class implementing the forward and backward pass for the KL Divergence Loss using Triton, as defined by the following formula:
+    ```python
+    if log_target:
+        loss = target * (target.log() - input)
+    else:
+        loss = target.exp() * (target - input)
+    ```,
+    then the loss is reduced according to the `reduction` parameter.
+    as defined in the PyTorch documentation: https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        y_pred: torch.Tensor,
+        y_true: torch.Tensor,
+        reduction: REDUCTION_LITERAL = "batchmean",
+        log_target: bool = False,
+        eps: float = 1e-10,
+    ) -> torch.Tensor:
+        """A forward pass for the KL Divergence Loss.
+
+        Args:
+            ctx: Torch autograd context
+            y_pred (torch.Tensor): A tensor of shape (BT, V) containing the predicted values, expected to be log-probabilities.
+            y_true (torch.Tensor): A tensor of shape (BT, V) containing the target values, expected to be either probabilities or log-probabilities, depending on the value of `log_target`.
+            reduction (REDUCTION_LITERAL, optional): Reduction to be used. Defaults to "batchmean".
+            log_target (bool, optional): If set to true, expects the ground truth to already be log-probabilities. Defaults to False.
+            eps: (float, optional): A small value to avoid division by zero. Defaults to 1e-10.
+
+        Returns:
+            torch.Tensor: The computed KL Divergence Loss, with shape (BT, V) if `reduction` is "none", else a scalar.
+        """
+        ctx.save_for_backward(y_true)
+        ctx.reduction = reduction
+        ctx.log_target = log_target
+        return kldiv_forward_triton(
+            y_pred, y_true, log_target=log_target, reduction=reduction, eps=eps
+        )
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        """A backward pass for the KL Divergence Loss.
+
+        Args:
+            ctx: Torch autograd context
+            grad_output (torch.Tensor): The gradient of the loss with respect to the output.
+
+        Returns:
+            tuple[torch.Tensor, None, None, None, None]: The gradient of the loss with respect to the inputs and None for the other arguments of the forward method.
+        """
+        (y_true,) = ctx.saved_tensors
+
+        new_grads = torch.empty_like(y_true)
+
+        derivative = kldiv_backward_triton(
+            y_true, grad_output, new_grads, ctx.log_target
+        )
+
+        if ctx.reduction == "batchmean":
+            derivative = derivative / y_true.shape[0]
+        elif ctx.reduction == "sum" or ctx.reduction == "none":
+            pass
+        elif ctx.reduction == "mean":
+            derivative = derivative / (y_true.shape[0] * y_true.shape[1])
+
+        return (
+            derivative,
+            None,
+            None,
+            None,
+            None,
+        )

liger_kernel/ops/layer_norm.py

@@ -0,0 +1,236 @@
+import math
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import (
+    calculate_settings,
+    compare_version,
+    ensure_contiguous,
+)
+
+if compare_version("triton", operator.ge, "3.0.0"):
+    try:
+        # typical import path with dispatch available
+        from triton.language.extra.libdevice import rsqrt
+    except ModuleNotFoundError:
+        # for working with NGC containers
+        from triton.language.extra.cuda.libdevice import rsqrt
+else:
+    from triton.language.math import rsqrt
+
+
+@triton.jit
+def _layer_norm_forward_kernel(
+    Y_ptr,  # pointer to output, shape (n_rows, n_cols)
+    Y_row_stride,  # stride of each row in output
+    X_ptr,  # pointer to input, shape (n_rows, n_cols)
+    X_row_stride,  # stride of each row in input
+    W_ptr,  # pointer to weights, shape (n_cols,)
+    W_row_stride,  # stride of each row in weights
+    B_ptr,  # pointer to bias, shape (n_cols,)
+    B_row_stride,  # stride of each row in bias
+    Mean_ptr,  # pointer to mean, shape (n_rows,)
+    Mean_row_stride,  # stride of each row in mean
+    RSTD_ptr,  # pointer to rstd, shape (n_rows,)
+    RSTD_row_stride,  # stride of each row in rstd
+    n_cols,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    References:
+    https://arxiv.org/abs/1607.06450
+    https://github.com/karpathy/llm.c/blob/master/doc/layernorm/layernorm.md
+    """
+    row_idx = tl.program_id(0)
+    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
+    var = tl.sum((X_row - mean) * (X_row - mean), axis=0) / n_cols
+    rstd = rsqrt(var + eps)
+
+    tl.store(Mean_ptr, mean)
+    tl.store(RSTD_ptr, rstd)
+
+    Y_row = (X_row - mean) * rstd * W_row + B_row
+
+    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
+
+
+@triton.jit
+def _layer_norm_backward_kernel(
+    X_ptr,  # pointer to input, shape (n_rows, n_cols)
+    W_ptr,  # pointer to weights, shape (n_cols,)
+    Mean_ptr,  # pointer to mean, shape (n_rows,)
+    RSTD_ptr,  # pointer to rstd, shape (n_rows,)
+    DX_ptr,  # pointer to input grad, shape (n_rows, n_cols)
+    DW_ptr,  # pointer to weights grad, shape (n_cols,)
+    DB_ptr,  # pointer to bias grad, shape (n_cols,)
+    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,
+):
+    """
+    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)
+    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)
+
+
+def layer_norm_forward(X, W, B, eps):
+    shape = X.shape
+    dim = shape[-1]
+    X = X.view(-1, dim)
+    n_rows, n_cols = X.shape
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    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)
+    assert (
+        X.shape[1] == W.shape[0]
+    ), f"Incompatible hidden size dimension between input tensor with shape[1] = {X.shape[1]} and weight tensor with shape[0] = {W.shape[0]}"
+
+    _layer_norm_forward_kernel[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        W,
+        W.stride(0),
+        B,
+        B.stride(0),
+        Mean,
+        Mean.stride(0),
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        eps,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return Y.view(*shape), X, Mean, RSTD, BLOCK_SIZE, num_warps
+
+
+def layer_norm_backward(dY, X, W, B, Mean, RSTD):
+    shape = dY.shape
+    dim = shape[-1]
+    dY = dY.view(-1, dim)
+    n_rows, n_cols = dY.shape
+
+    DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
+    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    _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)
+
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    if n_cols > BLOCK_SIZE:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
+    rows_per_program = math.ceil(n_rows / sm_count)
+    grid = (sm_count,)
+    triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16
+    _layer_norm_backward_kernel[grid](
+        X,
+        W,
+        Mean,
+        RSTD,
+        DX,
+        _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,
+    )
+
+    DW = _DW.sum(dim=0).to(W.dtype)
+    DB = _DB.sum(dim=0).to(W.dtype)
+
+    DX = DX.view(*shape)
+    return DX, DW, DB
+
+
+class LigerLayerNormFunction(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, X, W, B, eps):
+        Y, X, Mean, RSTD, BLOCK_SIZE, num_warps = layer_norm_forward(X, W, B, eps)
+        ctx.save_for_backward(X, W, B, Mean, RSTD)
+        return Y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dY):
+        X, W, B, Mean, RSTD = ctx.saved_tensors
+        DX, DW, DB = layer_norm_backward(dY, X, W, B, Mean, RSTD)
+        return DX, DW, DB, None