liger-kernel-nightly 0.4.0.dev20241107052928__py3-none-any.whl → 0.6.3.dev20251121010306__py3-none-any.whl

liger_kernel/ops/jsd.py

@@ -5,6 +5,7 @@ import triton
 import triton.language as tl
 
 from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.utils import infer_device
 
 
 @triton.jit
@@ -18,7 +19,7 @@ def _jsd_kernel(
     dX_ptr,
     dX_stride,
     label_ptr,
-    beta,
+    beta: tl.constexpr,
     n_non_ignore: int,
     ignore_index: tl.constexpr,
     n_cols,
@@ -50,21 +51,49 @@ def _jsd_kernel(
         X = tl.load(X_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32)
         Y = tl.load(Y_ptr + offsets, mask=mask, other=float("-inf")).to(tl.float32)
 
-        Q = tl.exp(X)
-        P = tl.exp(Y)
-        M = beta * P + (1 - beta) * Q
-        log_M = tl.log(M)
+        if beta == 0.0:  # forward KL
+            Y_max = tl.max(Y, axis=0)
+            Y_shifted = Y - Y_max
+            Y_prob = tl.exp(Y_shifted) * tl.exp(Y_max)  # Compensate for the shift
+            loss = Y_prob * (Y - X)
+            dX = -Y_prob
+        elif beta == 1.0:  # reverse KL
+            X_max = tl.max(X, axis=0)
+            X_shifted = X - X_max
+            X_prob = tl.exp(X_shifted) * tl.exp(X_max)  # Compensate for the shift
+            loss = X_prob * (X - Y)
+            dX = loss + X_prob
+        else:
+            max_val = tl.maximum(tl.max(X, axis=0), tl.max(Y, axis=0))
+            X_shifted = X - max_val
+            Y_shifted = Y - max_val
+
+            # Pre-compute exp(max_val) since it's used twice
+            exp_max = tl.exp(max_val)
+
+            # Compute exp terms with compensation
+            Q = tl.exp(X_shifted) * exp_max  # = exp(X)
+            P = tl.exp(Y_shifted) * exp_max  # = exp(Y)
+
+            # Pre-compute common terms
+            beta_P = beta * P
+            one_minus_beta_Q = (1 - beta) * Q
+            M = beta_P + one_minus_beta_Q
+            log_M = tl.log(M)  # No need to compensate as M is already in original scale
+
+            loss = beta_P * Y + one_minus_beta_Q * X - M * log_M
+            dX = one_minus_beta_Q * (X - log_M)
+
+        # Pre-compute scaling factor
+        scale = 1.0 / n_non_ignore
+        loss = loss * scale
+        dX = dX * scale
 
-        loss = beta * P * Y + (1 - beta) * Q * X - M * log_M
-        # reduction == "batchmean"
-        loss = loss / n_non_ignore
         tl.store(loss_ptr + offsets, loss, mask=mask)
-
-        dX = (1 - beta) * Q * (X - log_M) / n_non_ignore
         tl.store(dX_ptr + offsets, dX, mask=mask)
 
 
-MAX_FUSED_SIZE = 65536
+MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536
 
 
 def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label):
@@ -89,9 +118,7 @@ def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label):
         loss_stride=loss.stride(-2),
         dX_ptr=dX,
         dX_stride=dX.stride(-2),
-        label_ptr=(
-            shift_labels if has_label else torch.empty(1, device=_input.device)
-        ),  # dummy ptr if no label
+        label_ptr=(shift_labels if has_label else torch.empty(1, device=_input.device)),  # dummy ptr if no label
         beta=beta,
         n_non_ignore=n_non_ignore,
         ignore_index=ignore_index,
@@ -142,7 +169,7 @@ class LigerJSDFunction(torch.autograd.Function):
             _input (torch.Tensor): predict values with shape (BT, V) in logspace
             target (torch.Tensor): ground truth values with shape (BT, V) in logspace
             shift_labels (Optional[torch.LongTensor]): indicator of next predicted vocab with shape (BT) where each value is in [0, V-1].
-            beta (float): coefficient beta of generalized JSD in the open interval (0, 1)
+            beta (float): coefficient beta of generalized JSD in the interval [0, 1]. It implements forward/reverse KL when beta equals 0 and 1 respectively. Default: `0.5`
             ignore_index (int): the index to ignore. Default: -100
 
         Returns:
@@ -150,15 +177,13 @@ class LigerJSDFunction(torch.autograd.Function):
         """
         has_label = False
         if shift_labels is not None:
-            assert shift_labels.shape == (
-                _input.shape[0],
-            ), f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            assert shift_labels.shape == (_input.shape[0],), (
+                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
+            )
             shift_labels = shift_labels.contiguous()
             has_label = True
 
-        loss, dX = jsd_forward(
-            _input, target, shift_labels, beta, ignore_index, has_label
-        )
+        loss, dX = jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label)
         ctx.save_for_backward(dX)
         return loss
 

liger_kernel/ops/kl_div.py

@@ -4,7 +4,9 @@ import torch
 import triton
 import triton.language as tl
 
-from liger_kernel.ops.utils import ensure_contiguous, is_hip
+from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import is_hip
+from liger_kernel.utils import infer_device
 
 
 def get_num_warps(BLOCK_SIZE):
@@ -23,10 +25,10 @@ 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)
+_REDUCTION_MODE_NONE: tl.constexpr = tl.constexpr(0)
+_REDUCTION_MODE_SUM: tl.constexpr = tl.constexpr(1)
+_REDUCTION_MODE_MEAN: tl.constexpr = tl.constexpr(2)
+_REDUCTION_MODE_BATCHMEAN: tl.constexpr = tl.constexpr(3)
 
 _str_to_reduction_mode = {
     "none": _REDUCTION_MODE_NONE.value,
@@ -114,9 +116,12 @@ def _kldiv_kernel_backward(
 
 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)
+    BLOCK_SIZE = (
+        min(8192, triton.next_power_of_2(V))
+        if infer_device() == "xpu"
+        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    )
+    num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)
     reduction = _str_to_reduction_mode[reduction]
@@ -154,9 +159,12 @@ def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps):  # [BT, V]
 
 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)
+    BLOCK_SIZE = (
+        min(8192, triton.next_power_of_2(V))
+        if infer_device() == "xpu"
+        else min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    )
+    num_warps = 32 if infer_device() == "xpu" else get_num_warps(BLOCK_SIZE)
 
     grid = (BT,)
 
@@ -184,9 +192,9 @@ 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)
+    else:
+        loss = target * (target.log() - 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
@@ -218,9 +226,7 @@ class LigerKLDivLossFunction(torch.autograd.Function):
         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
-        )
+        return kldiv_forward_triton(y_pred, y_true, log_target=log_target, reduction=reduction, eps=eps)
 
     @staticmethod
     @ensure_contiguous
@@ -238,9 +244,7 @@ class LigerKLDivLossFunction(torch.autograd.Function):
 
         new_grads = torch.empty_like(y_true)
 
-        derivative = kldiv_backward_triton(
-            y_true, grad_output, new_grads, ctx.log_target
-        )
+        derivative = kldiv_backward_triton(y_true, grad_output, new_grads, ctx.log_target)
 
         if ctx.reduction == "batchmean":
             derivative = derivative / y_true.shape[0]

liger_kernel/ops/layer_norm.py

@@ -1,15 +1,12 @@
-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,
-)
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import compare_version
+from liger_kernel.ops.utils import ensure_contiguous
 
 if compare_version("triton", operator.ge, "3.0.0"):
     try:
@@ -45,29 +42,44 @@ 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
-    var = tl.sum((X_row - mean) * (X_row - mean), 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
+    mean = tl.sum(X_f32, axis=0) / n_cols
+    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
     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 = (X_row - mean) * 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
@@ -82,78 +94,100 @@ 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)
+
+    # 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
+    c2 = tl.sum(wdy, axis=0) / n_cols
+    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)
-    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,)](
+    # Validate input dimensions
+    if X.shape[1] != W.shape[0]:
+        raise ValueError(
+            f"Incompatible dimensions: input feature size (X.shape[1]={X.shape[1]}) "
+            f"must match weight size (W.shape[0]={W.shape[0]})"
+        )
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
+    # Launch kernel with one thread block per row for optimal performance
+    grid = (n_rows,)
+    _layer_norm_forward_kernel[grid](
         Y,
         Y.stride(0),
         X,
@@ -170,54 +204,84 @@ def layer_norm_forward(X, W, B, eps):
         eps,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
+        **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
 
+    # Allocate gradient tensors
     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)
+    # 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("This layer norm doesn't support feature dim >= 64KB.")
+        raise RuntimeError(f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}.")
+
+    # Determine dtype for triton operations
+    triton_dtype = (
+        tl.float32
+        if X.dtype == torch.float32
+        else tl.bfloat16
+        if X.dtype == torch.bfloat16
+        else tl.float16
+        if X.dtype == torch.float16
+        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
 
-    rows_per_program = math.ceil(n_rows / sm_count)
-    grid = (sm_count,)
-    triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16
+    kernel_args = {"num_warps": num_warps}
+    # XPU-specific optimization
+    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,
+        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):