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

liger_kernel/ops/fused_linear_jsd.py

@@ -0,0 +1,245 @@
+from typing import Optional
+
+import torch
+import triton
+
+from liger_kernel.ops.jsd import _jsd_kernel
+from liger_kernel.ops.utils import (
+    amp_custom_bwd,
+    amp_custom_fwd,
+    element_mul_kernel,
+    is_hip,
+)
+
+# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
+# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
+# The optimal maximum block size depends on your hardware, your kernel, and your dtype
+MAX_FUSED_SIZE = 65536 // 2
+
+
+def fused_linear_jsd_forward(
+    student_input,
+    student_weight,
+    teacher_input,
+    teacher_weight,
+    shift_labels,
+    jsd_beta,
+    ignore_index,
+    has_label,
+    temperature,
+):
+    device = student_input.device
+    dtype = student_input.dtype
+
+    # inputs have shape: BT x H
+    # materialized activations will have shape: BT x V
+    # the increase in memory = BT x V
+    # reduction can be achieved by partitioning the number of tokens BT into smaller chunks.
+    # for ex: if we were to achieve the same memory consumption as BT x H, then the chunk size should be:
+    # inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor
+    # for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048
+    BT, H = student_input.shape
+    V = student_weight.shape[0]
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+
+    inc_factor = triton.cdiv(V, H)  # (V + H - 1) // H
+    chunk_size = triton.next_power_of_2(
+        triton.cdiv(BT, inc_factor)
+    )  # (BT + inc_factor - 1) // inc_factor
+    num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size
+
+    grad_weight = (
+        torch.zeros_like(student_weight, device=device)
+        if student_weight.requires_grad
+        else None
+    )
+    grad_input = torch.zeros_like(student_input)
+    # we use fp32 for loss accumulator
+    loss_1d = torch.zeros((BT, V), dtype=torch.float32, device=device)
+
+    if has_label:
+        n_non_ignore = (shift_labels != ignore_index).sum().item()
+    else:
+        n_non_ignore = BT
+
+    for chunk_id in range(num_chunks):
+        start_idx = chunk_id * chunk_size
+        end_idx = min((chunk_id + 1) * chunk_size, BT)
+
+        # chunk both inputs, shape: chunk_size x H
+        student_input_chunk = student_input[start_idx:end_idx]
+        teacher_input_chunk = teacher_input[start_idx:end_idx]
+
+        # shape: chunk_size x V
+        # For anything starting from logits to the final JSD loss, we do computation
+        # in FP32 to avoid losing numerical stability.
+        student_logits_chunk = (student_input_chunk @ student_weight.t()).to(
+            torch.float32
+        )
+        teacher_logits_chunk = (teacher_input_chunk @ teacher_weight.t()).to(
+            torch.float32
+        )
+        chunk_n_rows = student_logits_chunk.shape[0]
+
+        # unreduced loss
+        loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size
+        # log-softmax with temperature
+        student_logits_chunk = student_logits_chunk / temperature
+        teacher_logits_chunk = teacher_logits_chunk / temperature
+        student_prob_chunk = torch.log_softmax(student_logits_chunk, dim=-1)
+        teacher_prob_chunk = torch.log_softmax(teacher_logits_chunk, dim=-1)
+
+        # ensure _input and target are contiguous
+        student_prob_chunk = student_prob_chunk.contiguous()
+        teacher_prob_chunk = teacher_prob_chunk.contiguous()
+
+        # Here we calculate the gradient of prob_chunk in place so we can save memory.
+        _jsd_kernel[(chunk_n_rows,)](
+            X_ptr=student_prob_chunk,
+            X_stride=student_prob_chunk.stride(-2),
+            Y_ptr=teacher_prob_chunk,
+            Y_stride=teacher_prob_chunk.stride(-2),
+            loss_ptr=loss_1d_slice,
+            loss_stride=loss_1d_slice.stride(-2),
+            dX_ptr=student_prob_chunk,
+            dX_stride=student_prob_chunk.stride(-2),
+            label_ptr=(
+                shift_labels[start_idx:end_idx]
+                if has_label
+                else torch.empty(1, device=device)
+            ),  # dummy ptr if no label
+            beta=jsd_beta,
+            n_non_ignore=n_non_ignore,
+            ignore_index=ignore_index,
+            n_cols=V,
+            BLOCK_SIZE=BLOCK_SIZE,
+            HAS_LABEL=has_label,
+        )
+        loss_1d[start_idx:end_idx] = loss_1d_slice
+        # gradients of prob_chunk in place, shape: chunk_size x V
+        # gradients of logits_chunk in place, shape: chunk_size x V
+        student_logits_chunk = (
+            student_prob_chunk
+            - torch.softmax(student_logits_chunk, dim=-1)
+            * student_prob_chunk.sum(dim=-1, keepdim=True).broadcast_to(
+                student_prob_chunk.shape
+            )
+        ) / temperature
+        # now we traverse back to grad w.r.t. input to `lm_head` and grad
+        # w.r.t. `lm_head` which should be computed in original dtype
+        student_logits_chunk = student_logits_chunk.to(dtype)
+        grad_input[start_idx:end_idx] = student_logits_chunk @ student_weight
+
+        if grad_weight is not None:
+            grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
+
+    loss = torch.sum(loss_1d)
+    return loss, grad_input, grad_weight
+
+
+def fused_linear_jsd_backward(grad_output, grad_input, grad_weight):
+    # If JSD is the last layer, grad_output is 1.0. Skip the mul to save time
+    if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
+        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
+        BT, H = grad_input.shape
+        n_rows = BT
+        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+
+        element_mul_kernel[(n_rows,)](
+            grad_input,
+            grad_input.stride(-2),
+            grad_output,
+            H,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=32 if not is_hip() else 16,
+        )
+
+        # handle grad_weight
+        if grad_weight is not None:
+            V, H = grad_weight.shape
+            n_rows = V
+
+            element_mul_kernel[(n_rows,)](
+                grad_weight,
+                grad_weight.stride(-2),
+                grad_output,
+                H,
+                BLOCK_SIZE=BLOCK_SIZE,
+                num_warps=32 if not is_hip() else 16,
+            )
+
+    return grad_input, grad_weight
+
+
+class LigerFusedLinearJSDFunction(torch.autograd.Function):
+    """
+    Fusing the last linear layer with generalized JSD
+
+    Handle the forward and backward pass of the final linear layer via JSD by avoiding
+    the materialization of the large logits tensor. Since JSD is the last layer, we can
+    compute the gradient at the forward pass.
+    """
+
+    @staticmethod
+    @amp_custom_fwd
+    def forward(
+        ctx,
+        student_input: torch.Tensor,
+        student_weight: torch.Tensor,
+        teacher_input: torch.Tensor,
+        teacher_weight: torch.Tensor,
+        shift_labels: Optional[torch.Tensor] = None,
+        jsd_beta: float = 0.5,
+        ignore_index: int = -100,
+        temperature: float = 1.0,
+    ):
+        """
+        Args:
+
+            student_input (torch.tensor): input of the last projection layer in student model, with shape (B*T, H), where B is batch size, T is sequence length, H is hidden dimension.
+            student_weight (torch.tensor): the last projection layer in student model, with shape (V, H), where V is vocab size
+            teacher_input (torch.tensor): input of the last projection layer in teacher model, with shape (B*T, H), where B is batch size, T is sequence length, H is hidden dimension.
+            teacher_weight (torch.tensor): the last projection layer in teacher model, with shape (V, H), where V is vocab size
+            shift_labels (Optional[torch.LongTensor]): indicator of next predicted vocab with shape (BT) where each value is in [0, V-1].
+            jsd_beta (float): coefficient beta of generalized JSD in the open interval (0, 1). Default: `0.5`
+            ignore_index (int): the index to ignore. Default: -100
+            temperature (float): temperature in softmax function to control the output probability distribution. Default: `1.0`
+
+        Returns:
+            loss (torch.Tensor): generalized JSD
+        """
+        has_label = False
+        if shift_labels is not None:
+            assert shift_labels.shape == (
+                teacher_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, grad_input, grad_weight = fused_linear_jsd_forward(
+            student_input,
+            student_weight,
+            teacher_input,
+            teacher_weight,
+            shift_labels,
+            jsd_beta,
+            ignore_index,
+            has_label,
+            temperature,
+        )
+        # downcast to dtype and store for backward
+        ctx.save_for_backward(
+            grad_input.detach(),
+            grad_weight.detach() if grad_weight is not None else None,
+        )
+        return loss
+
+    @staticmethod
+    @amp_custom_bwd
+    def backward(ctx, grad_output):
+        (grad_input, grad_weight) = ctx.saved_tensors
+        grad_input, grad_weight = fused_linear_jsd_backward(
+            grad_output, grad_input, grad_weight
+        )
+        return (grad_input, grad_weight, None, None, None, None, None, None)

liger_kernel/ops/geglu.py

@@ -25,7 +25,7 @@ else:
 def _geglu_tanh_forward_kernel(
     a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     a += program_id * stride
@@ -52,7 +52,7 @@ def _geglu_tanh_forward_kernel(
 def _geglu_tanh_backward_kernel(
     dc, a, b, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     dc += program_id * stride

liger_kernel/ops/jsd.py

@@ -0,0 +1,176 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+@triton.jit
+def _jsd_kernel(
+    X_ptr,  # input in logspace, X = log Q
+    X_stride,
+    Y_ptr,  # ground truth in logspace, Y = log P
+    Y_stride,
+    loss_ptr,
+    loss_stride,
+    dX_ptr,
+    dX_stride,
+    label_ptr,
+    beta,
+    n_non_ignore: int,
+    ignore_index: tl.constexpr,
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_LABEL: tl.constexpr,
+):
+    # JSD(P || Q) = (KL(P || M) + KL(Q || M)) / 2, M = (1/2) * (P + Q) = (1/2) * (e ^ Y + e ^ X)
+    #             = sum(P * log P + Q * log Q - 2 * M * log M) / 2
+    #             = sum(e ^ Y * Y + e ^ X * X - 2 * M * log M) / 2
+    # grad_x_i = 0.5 * Q * (X - log_M)
+    pid = tl.program_id(0).to(tl.int64)
+    X_ptr += pid * X_stride
+    dX_ptr += pid * dX_stride
+    Y_ptr += pid * Y_stride
+    loss_ptr += pid * loss_stride
+    label_ptr += pid
+
+    if HAS_LABEL:
+        label = tl.load(label_ptr)
+        if label == ignore_index:
+            for i in range(0, n_cols, BLOCK_SIZE):
+                offsets = i + tl.arange(0, BLOCK_SIZE)
+                tl.store(dX_ptr + offsets, 0.0, mask=offsets < n_cols)
+            return
+
+    for i in range(0, n_cols, BLOCK_SIZE):
+        offsets = i + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_cols
+        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)
+
+        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
+
+
+def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label):
+    BT, V = _input.shape
+    n_rows = BT
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
+    # non reduction loss
+    loss = torch.zeros(_input.shape, dtype=torch.float32, device=_input.device)
+    dX = torch.empty_like(_input)
+
+    if has_label:
+        n_non_ignore = (shift_labels != ignore_index).sum().item()
+    else:
+        n_non_ignore = BT
+
+    _jsd_kernel[(n_rows,)](
+        X_ptr=_input,  # input in logspace, X = log Q
+        X_stride=_input.stride(-2),
+        Y_ptr=target,  # ground truth in logspace, Y = log P
+        Y_stride=target.stride(-2),
+        loss_ptr=loss,
+        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
+        beta=beta,
+        n_non_ignore=n_non_ignore,
+        ignore_index=ignore_index,
+        n_cols=V,
+        BLOCK_SIZE=BLOCK_SIZE,
+        HAS_LABEL=has_label,
+    )
+
+    loss = torch.sum(loss)
+    return loss.to(_input.dtype), dX
+
+
+def jsd_backward(dX, grad_output):
+    # If jsd is the last layer, grad_output is 1.0. Skip the mul to save time
+    if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
+        return dX
+    else:
+        return grad_output * dX
+
+
+class LigerJSDFunction(torch.autograd.Function):
+    r"""
+    This class implements the forward and backward pass for the generalized Jensen-Shannon Divergence.
+    .. math::
+        JSD(\beta)(P || Q)
+            = \beta * KLDiv(P || (\beta * P + (1 - \beta) * Q)) + (1 - \beta) * KLDiv(Q || (\beta * P + (1 - \beta) * Q))
+
+    .. note::
+        As all the other losses in PyTorch, this function expects the first argument,
+        :attr:`_input`, to be the predictions, the output of the student model, in log-space
+        and the second, :attr:`target`, to be the observations, the output of the teacher model, in log-space.
+        This differs from the standard mathematical notation :math:`JSD(P || Q)` where
+        :math:`P` denotes the teacher model and :math:`Q` denotes the student model.
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        _input: torch.Tensor,
+        target: torch.Tensor,
+        shift_labels: Optional[torch.Tensor] = None,
+        beta: float = 0.5,
+        ignore_index: int = -100,
+    ) -> torch.Tensor:
+        """
+        Args:
+            _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)
+            ignore_index (int): the index to ignore. Default: -100
+
+        Returns:
+            loss (torch.Tensor): generalized JSD
+        """
+        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}"
+            shift_labels = shift_labels.contiguous()
+            has_label = True
+
+        loss, dX = jsd_forward(
+            _input, target, shift_labels, beta, ignore_index, has_label
+        )
+        ctx.save_for_backward(dX)
+        return loss
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        (dX,) = ctx.saved_tensors
+        dX = jsd_backward(dX, grad_output)
+        return (
+            dX,
+            None,
+            None,
+            None,
+            None,
+        )

liger_kernel/ops/kl_div.py

@@ -4,13 +4,13 @@ import torch
 import triton
 import triton.language as tl
 
-from liger_kernel.ops.utils import ensure_contiguous
+from liger_kernel.ops.utils import ensure_contiguous, is_hip
 
 
 def get_num_warps(BLOCK_SIZE):
     num_warps = 4
     if BLOCK_SIZE >= 32768:
-        num_warps = 32
+        num_warps = 32 if not is_hip() else 16
     elif BLOCK_SIZE >= 8192:
         num_warps = 16
     elif BLOCK_SIZE >= 2048:

liger_kernel/ops/rms_norm.py

@@ -10,6 +10,7 @@ https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddec
 Modifications made by Yanning Chen, 2024.
 """
 
+import math
 import operator
 
 import torch
@@ -20,6 +21,7 @@ from liger_kernel.ops.utils import (
     calculate_settings,
     compare_version,
     ensure_contiguous,
+    torch_to_triton_dtype,
 )
 
 if compare_version("triton", operator.ge, "3.0.0"):
@@ -84,6 +86,10 @@ def _rms_norm_forward_kernel(
         W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
+    if casting_mode == _CASTING_MODE_NONE:
+        eps = eps.to(X_row_dtype)
+        offset = offset.to(X_row_dtype)
+
     mean_square = tl.sum(X_row * X_row, axis=0) / n_cols
     rstd = rsqrt(mean_square + eps)
 
@@ -100,6 +106,9 @@ def _rms_norm_forward_kernel(
 
     Y_row = X_row * (offset + W_row)
 
+    if casting_mode == _CASTING_MODE_GEMMA:
+        Y_row = Y_row.to(X_row_dtype)
+
     tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
 
 
@@ -109,14 +118,17 @@ def _rms_norm_backward_kernel(
     dY_row_stride,
     X_ptr,
     X_row_stride,
+    X_dtype: tl.constexpr,
     W_ptr,
     W_row_stride,
     RSTD_ptr,
     RSTD_row_stride,
     dW_ptr,
     dW_row_stride,
+    n_rows,
     n_cols,
     offset,
+    rows_per_program: tl.constexpr,
     casting_mode: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
@@ -125,54 +137,60 @@ def _rms_norm_backward_kernel(
     dw = sum(dy * (x / RMS)). summation over BxT dimension
     """
 
-    row_idx = tl.program_id(0)
+    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)
     col_offsets = tl.arange(0, BLOCK_SIZE)
     mask = col_offsets < n_cols
 
-    dY_ptr += row_idx * dY_row_stride
-    X_ptr += row_idx * X_row_stride
-    RSTD_ptr += row_idx * RSTD_row_stride
-    dW_ptr += row_idx * dW_row_stride
+    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0)
-    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
-    original_x_dtype = X_row.dtype
-
-    # Get cached rms
-    rstd_row = tl.load(RSTD_ptr)
+    dY_ptr += row_start * dY_row_stride
+    X_ptr += row_start * X_row_stride
+    RSTD_ptr += row_start
 
+    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
     W_row = W_row + offset
 
-    X_row = X_row.to(tl.float32)
+    for _ in range(row_start, row_end):
+        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
 
-    # Different bacward graphs for different casting modes
-    if casting_mode == _CASTING_MODE_LLAMA:
-        m = (dY_row * W_row).to(tl.float32)
+        # Get cached rms
+        rstd_row = tl.load(RSTD_ptr)
 
-    elif casting_mode == _CASTING_MODE_GEMMA:
-        dY_row, W_row = (
-            dY_row.to(tl.float32),
-            W_row.to(tl.float32),
-        )
+        X_row = X_row.to(tl.float32)
 
-    m = dY_row * W_row
+        # Different bacward graphs for different casting modes
+        if casting_mode == _CASTING_MODE_LLAMA:
+            m = (dY_row * W_row).to(tl.float32)
 
-    dX_row = rstd_row * m
+        elif casting_mode == _CASTING_MODE_GEMMA:
+            dY_row = dY_row.to(tl.float32)
+            m = dY_row * W_row
+        else:
+            m = dY_row * W_row
 
-    dX_row += (rstd_row) * (
-        -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
-    )
+        dX_row = rstd_row * m
 
-    # calculate the gradient of W
-    if casting_mode == _CASTING_MODE_LLAMA:
-        dW_row = dY_row * (X_row * rstd_row).to(original_x_dtype)
-    else:
-        # here X_row is already in fp32 (see previous if block)
-        dW_row = dY_row * (X_row * rstd_row)
+        dX_row += (rstd_row) * (
+            -(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
+        )
+
+        # calculate the gradient of W
+        if casting_mode == _CASTING_MODE_LLAMA:
+            dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
+        else:
+            # here X_row is already in fp32 (see previous if block)
+            dW_row += dY_row * (X_row * rstd_row)
 
-    tl.store(dY_ptr + col_offsets, dX_row, mask=mask)
-    tl.store(dW_ptr + col_offsets, dW_row, mask=mask)
+        tl.store(dY_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
+
+        dY_ptr += dY_row_stride
+        X_ptr += X_row_stride
+        RSTD_ptr += RSTD_row_stride
+
+    tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
 
 _str_to_casting_mode = {
@@ -238,31 +256,38 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
     dim = shape[-1]
     dY = dY.view(-1, dim)
     n_rows, n_cols = dY.shape
-    dW = torch.empty_like(
-        X,
-        dtype=(torch.float32 if casting_mode == _CASTING_MODE_GEMMA.value else W.dtype),
-    )
 
+    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    # fp32 for numerical stability especially.
+    _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+
+    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,)
     # Here we use dY to store the value of dX to save memory
-    _rms_norm_backward_kernel[(n_rows,)](
+    _rms_norm_backward_kernel[grid](
         dY,
         dY.stride(0),
         X,
         X.stride(0),
+        torch_to_triton_dtype[X.dtype],
         W,
         W.stride(0),
         RSTD,
         RSTD.stride(0),
-        dW,
-        dW.stride(0),
+        _dW,
+        _dW.stride(0),
+        n_rows,
         n_cols,
         offset,
+        rows_per_program,
         casting_mode,
         BLOCK_SIZE=BLOCK_SIZE,
         num_warps=num_warps,
     )
     dX = dY.view(*shape)
-    dW = torch.sum(dW, dim=0).to(W.dtype)
+    dW = _dW.sum(dim=0).to(W.dtype)
     return dX, dW
 
 

liger_kernel/ops/swiglu.py

@@ -14,7 +14,7 @@ def silu(x):
 def _swiglu_forward_kernel(
     a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     a_ptr += program_id * stride
@@ -35,7 +35,7 @@ def _swiglu_forward_kernel(
 def _swiglu_backward_kernel(
     dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     dc_ptr += program_id * stride