liger-kernel-nightly 0.5.10.dev20250611191801__py3-none-any.whl → 0.6.4.dev20260112233432__py3-none-any.whl

liger_kernel/ops/rms_norm.py

@@ -21,8 +21,10 @@ from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.ops.utils import compare_version
 from liger_kernel.ops.utils import ensure_contiguous
 from liger_kernel.ops.utils import torch_to_triton_dtype
+from liger_kernel.utils import get_npu_multi_processor_count
+from liger_kernel.utils import is_npu_available
 
-if compare_version("triton", operator.ge, "3.0.0"):
+if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available():
     try:
         # typical import path with dispatch available
         from triton.language.extra.libdevice import rsqrt
@@ -52,6 +54,7 @@ def _rms_norm_forward_kernel(
     eps,
     offset,
     casting_mode: tl.constexpr,  # constexpr so the `if` blocks can be optimized out
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -63,17 +66,18 @@ def _rms_norm_forward_kernel(
     3. https://arxiv.org/pdf/1910.07467
     """
 
-    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
-    RSTD_ptr += row_idx * RSTD_row_stride
+    y_base = Y_ptr + row_idx * Y_row_stride
+    x_base = X_ptr + row_idx * X_row_stride
+    rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
 
-    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
+    X_row = tl.load(x_base + col_offsets, mask=mask, other=0)
     X_row_dtype = X_row.dtype
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
+    if elementwise_affine:
+        W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
 
     # On Llama, only rstd is computed on fp32
     if casting_mode == _CASTING_MODE_LLAMA:
@@ -81,7 +85,8 @@ def _rms_norm_forward_kernel(
 
     # Gemma computes everything on fp32, and then casts back the output to the original dtype
     if casting_mode == _CASTING_MODE_GEMMA:
-        W_row = W_row.to(tl.float32)
+        if elementwise_affine:
+            W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
     if casting_mode == _CASTING_MODE_NONE:
@@ -94,7 +99,7 @@ def _rms_norm_forward_kernel(
     # We can save time by caching rms with minimal memory overhead
     # because rms is much smaller compared to X_row, as rms is for each row.
     # However, on the computation side, it can save 4 operations (*, sum, /, sqrt).
-    tl.store(RSTD_ptr, rstd)
+    tl.store(rstd_base, rstd)
 
     X_row = X_row * rstd
 
@@ -102,12 +107,15 @@ def _rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_LLAMA:
         X_row = X_row.to(X_row_dtype)
 
-    Y_row = X_row * (offset + W_row)
+    if elementwise_affine:
+        Y_row = X_row * (offset + W_row)
+    else:
+        Y_row = X_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)
+    tl.store(y_base + col_offsets, Y_row, mask=mask)
 
 
 @triton.jit
@@ -128,8 +136,9 @@ def _rms_norm_backward_kernel(
     n_rows,
     n_cols,
     offset,
-    rows_per_program: tl.constexpr,
+    rows_per_program,
     casting_mode: tl.constexpr,
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -137,61 +146,69 @@ def _rms_norm_backward_kernel(
     dw = sum(dy * (x / RMS)). summation over BxT dimension
     """
 
-    row_block_id = tl.program_id(0)
+    row_block_id = tl.program_id(0).to(tl.int64)
     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
 
-    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    if elementwise_affine:
+        dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    dY_ptr += row_start * dY_row_stride
-    dX_ptr += row_start * dX_row_stride
+    if elementwise_affine:
+        W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
+        W_row = W_row + offset
 
-    X_ptr += row_start * X_row_stride
-    RSTD_ptr += row_start
+    for row_idx in range(row_start, row_end):
+        dy_base = dY_ptr + row_idx * dY_row_stride
+        dx_base = dX_ptr + row_idx * dX_row_stride
 
-    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
-    W_row = W_row + offset
+        x_base = X_ptr + row_idx * X_row_stride
+        rstd_base = RSTD_ptr + row_idx * RSTD_row_stride
 
-    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)
+        dY_row = tl.load(dy_base + col_offsets, mask=mask, other=0.0)
+        X_row = tl.load(x_base + col_offsets, mask=mask, other=0.0)
 
         # Get cached rms
-        rstd_row = tl.load(RSTD_ptr)
+        rstd_row = tl.load(rstd_base)
 
         X_row = X_row.to(tl.float32)
 
         # Different bacward graphs for different casting modes
         if casting_mode == _CASTING_MODE_LLAMA:
-            m = (dY_row * W_row).to(tl.float32)
+            if elementwise_affine:
+                m = (dY_row * W_row).to(tl.float32)
+            else:
+                m = dY_row.to(tl.float32)
 
         elif casting_mode == _CASTING_MODE_GEMMA:
             dY_row = dY_row.to(tl.float32)
-            m = dY_row * W_row
+            if elementwise_affine:
+                m = dY_row * W_row
+            else:
+                m = dY_row
         else:
-            m = dY_row * W_row
+            if elementwise_affine:
+                m = dY_row * W_row
+            else:
+                m = dY_row
 
         dX_row = rstd_row * m
 
         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)
+        if elementwise_affine:
+            # 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(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
+        tl.store(dx_base + col_offsets, dX_row.to(X_dtype), mask=mask)
 
-        dY_ptr += dY_row_stride
-        dX_ptr += dX_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)
+    if elementwise_affine:
+        tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
 
 
 @triton.jit
@@ -209,6 +226,7 @@ def _block_rms_norm_forward_kernel(
     eps,
     offset,
     casting_mode: tl.constexpr,  # constexpr so the `if` blocks can be optimized out
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     BLOCK_ROW: tl.constexpr,
 ):
@@ -232,7 +250,8 @@ def _block_rms_norm_forward_kernel(
         other=0,
     )
     X_row_dtype = X_row.dtype
-    W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0)
+    if elementwise_affine:
+        W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0)
 
     # On Llama, only rstd is computed on fp32
     if casting_mode == _CASTING_MODE_LLAMA:
@@ -240,7 +259,8 @@ def _block_rms_norm_forward_kernel(
 
     # Gemma computes everything on fp32, and then casts back the output to the original dtype
     if casting_mode == _CASTING_MODE_GEMMA:
-        W_row = W_row.to(tl.float32)
+        if elementwise_affine:
+            W_row = W_row.to(tl.float32)
         X_row = X_row.to(tl.float32)
 
     if casting_mode == _CASTING_MODE_NONE:
@@ -261,7 +281,10 @@ def _block_rms_norm_forward_kernel(
     if casting_mode == _CASTING_MODE_LLAMA:
         X_row = X_row.to(X_row_dtype)
 
-    Y_row = X_row * (offset + W_row)[None, :]
+    if elementwise_affine:
+        Y_row = X_row * (offset + W_row)[None, :]
+    else:
+        Y_row = X_row
 
     if casting_mode == _CASTING_MODE_GEMMA:
         Y_row = Y_row.to(X_row_dtype)
@@ -291,8 +314,8 @@ def _block_rms_norm_backward_kernel(
     n_rows,
     n_cols,
     offset,
-    rows_per_program: tl.constexpr,
     casting_mode: tl.constexpr,
+    elementwise_affine: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     BLOCK_ROW: tl.constexpr,
 ):
@@ -307,10 +330,11 @@ def _block_rms_norm_backward_kernel(
     col_offsets = tl.arange(0, BLOCK_SIZE)
     col_mask = col_offsets < n_cols
 
-    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    if elementwise_affine:
+        dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
 
-    W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)
-    W_row = W_row + offset
+        W_row = tl.load(W_ptr + col_offsets, mask=col_mask, other=0.0)
+        W_row = W_row + offset
 
     for start in range(pid * BLOCK_ROW, n_rows, NUM_SMS * BLOCK_ROW):
         row_idx = start + tl.arange(0, BLOCK_ROW)
@@ -333,13 +357,22 @@ def _block_rms_norm_backward_kernel(
 
         # Different bacward graphs for different casting modes
         if casting_mode == _CASTING_MODE_LLAMA:
-            m = (dY_row * W_row[None, :]).to(tl.float32)
+            if elementwise_affine:
+                m = (dY_row * W_row[None, :]).to(tl.float32)
+            else:
+                m = dY_row.to(tl.float32)
 
         elif casting_mode == _CASTING_MODE_GEMMA:
             dY_row = dY_row.to(tl.float32)
-            m = dY_row * W_row[None, :]
+            if elementwise_affine:
+                m = dY_row * W_row[None, :]
+            else:
+                m = dY_row
         else:
-            m = dY_row * W_row[None, :]
+            if elementwise_affine:
+                m = dY_row * W_row[None, :]
+            else:
+                m = dY_row
 
         dX_row = rstd_row[:, None] * m
 
@@ -347,12 +380,13 @@ def _block_rms_norm_backward_kernel(
             -(1 / n_cols) * (rstd_row * rstd_row * tl.sum(m * X_row, axis=1))[:, None] * X_row
         )
 
-        # calculate the gradient of W
-        if casting_mode == _CASTING_MODE_LLAMA:
-            dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]).to(X_dtype), 0)
-        else:
-            # here X_row is already in fp32 (see previous if block)
-            dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]), 0)
+        if elementwise_affine:
+            if casting_mode == _CASTING_MODE_LLAMA:
+                # TODO(tcc): use tl.sum(..., dtype=tl.float32) once we upgrade to triton>=3.3.0
+                dW_row += tl.sum((dY_row * (X_row * rstd_row[:, None]).to(X_dtype)).to(tl.float32), 0)
+            else:
+                # here X_row is already in fp32 (see previous if block)
+                dW_row += tl.sum(dY_row * (X_row * rstd_row[:, None]), 0)
 
         tl.store(
             dX_ptr + row_idx[:, None] * dX_row_stride + col_offsets[None, :],
@@ -360,7 +394,8 @@ def _block_rms_norm_backward_kernel(
             mask=row_mask[:, None] & col_mask[None, :],
         )
 
-    tl.store(dW_ptr + pid * dW_row_stride + col_offsets, dW_row, mask=col_mask)
+    if elementwise_affine:
+        tl.store(dW_ptr + pid * dW_row_stride + col_offsets, dW_row, mask=col_mask)
 
 
 _str_to_casting_mode = {
@@ -389,8 +424,14 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
     rstd_dtype = torch.float32 if casting_mode in (_CASTING_MODE_LLAMA.value, _CASTING_MODE_GEMMA.value) else X.dtype
     RSTD = torch.empty(n_rows, dtype=rstd_dtype, device=X.device)
 
-    # Check constraints.
-    assert X.shape[1] == W.shape[0], "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
+    if W is not None:
+        # Check constraints.
+        assert X.shape[1] == W.shape[0], (
+            "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
+        )
+        elementwise_affine = True
+    else:
+        elementwise_affine = False
 
     # XPU-specific optimization
     kernel_args = {}
@@ -403,13 +444,14 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             X,
             X.stride(0),
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             n_cols,
             eps,
             offset,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -423,7 +465,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             X,
             X.stride(0),
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             n_rows,
@@ -431,6 +473,7 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode):
             eps,
             offset,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -449,9 +492,16 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
         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
+    elif X.device.type == "npu":
+        sm_count = get_npu_multi_processor_count()
 
-    # fp32 for numerical stability especially.
-    _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+    if W is not None:
+        # fp32 for numerical stability especially.
+        _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+        elementwise_affine = True
+    else:
+        _dW = None
+        elementwise_affine = False
 
     if n_cols > BLOCK_SIZE:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
@@ -478,16 +528,17 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
             X.stride(0),
             torch_to_triton_dtype[X.dtype],
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             _dW,
-            _dW.stride(0),
+            _dW.stride(0) if elementwise_affine else 0,
             n_rows,
             n_cols,
             offset,
             rows_per_program,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
@@ -504,22 +555,26 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp
             X.stride(0),
             torch_to_triton_dtype[X.dtype],
             W,
-            W.stride(0),
+            W.stride(0) if elementwise_affine else 0,
             RSTD,
             RSTD.stride(0),
             _dW,
-            _dW.stride(0),
+            _dW.stride(0) if elementwise_affine else 0,
             n_rows,
             n_cols,
             offset,
-            rows_per_program,
             casting_mode,
+            elementwise_affine=elementwise_affine,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=num_warps,
             **kernel_args,  # XPU-specific optimization
         )
     dX = dX.view(*shape)
-    dW = _dW.sum(dim=0).to(W.dtype)
+
+    if elementwise_affine:
+        dW = _dW.sum(dim=0).to(W.dtype)
+    else:
+        dW = None
 
     return dX, dW
 
@@ -553,6 +608,13 @@ class LigerRMSNormFunction(torch.autograd.Function):
         X: (B, T, H) or (BxT, H)
         W: (H,)
         """
+        if isinstance(X, torch.distributed.tensor.DTensor):
+            # Input tensor is output of a tensor parallel module and
+            # needs to be gathered to a local tensor to compute
+            # RMSE layer norm on each TP worker.
+            # TODO: support CP.
+            X = X.full_tensor()
+
         Y, X, RSTD, BLOCK_SIZE, num_warps, casting_mode = rms_norm_forward(X, W, eps, offset, casting_mode, row_mode)
         ctx.offset = offset
         ctx.casting_mode = casting_mode
@@ -560,7 +622,11 @@ class LigerRMSNormFunction(torch.autograd.Function):
         ctx.row_mode = row_mode
         ctx.BLOCK_SIZE = BLOCK_SIZE
         ctx.num_warps = num_warps
-        ctx.save_for_backward(X, W, RSTD)
+        ctx.elementwise_affine = W is not None
+        if W is not None:
+            ctx.save_for_backward(X, W, RSTD)
+        else:
+            ctx.save_for_backward(X, RSTD)
         return Y
 
     @staticmethod
@@ -569,7 +635,18 @@ class LigerRMSNormFunction(torch.autograd.Function):
         """
         Y: (B, T, H) or (BxT, H)
         """
-        X, W, RSTD = ctx.saved_tensors
+        if ctx.elementwise_affine:
+            X, W, RSTD = ctx.saved_tensors
+        else:
+            X, RSTD = ctx.saved_tensors
+            W = None
+
+        if isinstance(dY, torch.distributed.tensor.DTensor):
+            # Gradients are output of a tensor parallel module and
+            # needs to be gathered to a local tensor for computing RMSE layer.
+            # TODO: support CP.
+            dY = dY.full_tensor()
+
         dX, dW = rms_norm_backward(
             dY, X, W, RSTD, ctx.offset, ctx.casting_mode, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place, ctx.row_mode
         )

liger_kernel/ops/rope.py

@@ -32,7 +32,7 @@ def _triton_rope(
 
     # cos size: (1, seq_len, head_dim) or (bsz, seq_len, head_dim)
     # stride: (seq_len * head_dim, head_dim, 1)
-    pid = tl.program_id(0)
+    pid = tl.program_id(0).to(tl.int64)
 
     # locate start address
     q_ptr = q_ptr + pid * q_row_stride

liger_kernel/ops/swiglu.py

@@ -26,7 +26,7 @@ def _swiglu_forward_kernel(a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BL
     # sigmoid requires type float32
     a_row = tl.load(a_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
     b_row = tl.load(b_ptr + col_offsets, mask=mask, other=0)
-    c_row = silu(a_row) * b_row
+    c_row = silu(a_row).cast(b_row.dtype) * b_row
     tl.store(c_ptr + col_offsets, c_row, mask=mask)
 
 

liger_kernel/ops/tiled_mlp.py

@@ -0,0 +1,136 @@
+import math
+
+from typing import Callable
+from typing import List
+from typing import Optional
+
+import torch
+
+from liger_kernel.ops.utils import ensure_contiguous
+
+
+class LigerTiledMLPFunction(torch.autograd.Function):
+    """
+    Based on DeepSpeed's TiledMLP:
+    https://github.com/deepspeedai/DeepSpeed/blob/v0.18.2/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L838
+
+    Perform a tiled MLP computation to massively reduce memory usage needed to compute MLP
+    when using very long sequence lengths.
+
+    This module re-computes `forward` in the `backward`. So the `forward` occurs twice each iteration.
+    And if you're using activation checkpointing it then occurs thrice.
+
+    Args:
+        fn: the function to call on sharded inputs (e.g., mlp.forward)
+        mlp_module: the MLP nn.Module object
+        x: the input to MLP.forward (hidden_states)
+        shards: how many shards to use
+        compute_params: a list of weights engaged in the compute
+
+    Returns:
+        the computed hidden_states
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        fn: Callable,
+        mlp_module: torch.nn.Module,
+        x: torch.Tensor,
+        shards: int,
+        compute_params: Optional[List[torch.nn.Parameter]] = None,
+    ) -> torch.Tensor:
+        ctx.fn = fn
+        ctx.mlp_module = mlp_module
+        ctx.shards = shards
+        ctx.save_for_backward(x)
+
+        # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
+        x_shards = list(torch.chunk(x, chunks=shards, dim=-2))
+        with torch.no_grad():
+            output_shards = [fn(mlp_module, x_shard) for x_shard in x_shards]
+        output_unsharded = torch.cat(output_shards, dim=-2)
+
+        return output_unsharded
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, *grads) -> tuple:
+        fn = ctx.fn
+        (x,) = ctx.saved_tensors
+        mlp_module = ctx.mlp_module
+        shards = ctx.shards
+
+        x_requires_grad = x.requires_grad
+        x = x.detach()
+        # detach() unsets x.requires_grad, so restore it
+        x.requires_grad_(x_requires_grad)
+
+        # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
+        hidden_size = x.shape[-1]
+        x_shape_orig = x.shape
+
+        # flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
+        x = x.view(-1, hidden_size)
+        incoming_grad = grads[0].view(-1, hidden_size)
+        x_grad = torch.zeros_like(x)
+
+        x_shards = list(torch.chunk(x, chunks=shards, dim=0))
+
+        for i, x_shard in enumerate(x_shards):
+            x_shard.requires_grad_(x_requires_grad)
+
+            # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step
+            shard_step = x_shards[i].shape[0]
+            shard_offset = i * x_shards[0].shape[0]
+
+            x_shard.grad = x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
+            incoming_grad_shard = incoming_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
+
+            with torch.enable_grad():
+                output = fn(mlp_module, x_shard)
+            torch.autograd.backward(output, incoming_grad_shard)
+
+        # unflatten
+        x_grad = x_grad.view(x_shape_orig)
+
+        return (None, None, x_grad, None, None)
+
+
+def apply_tiled_mlp(
+    fn: Callable,
+    mlp_module: torch.nn.Module,
+    x: torch.Tensor,
+    num_shards: Optional[int] = None,
+    compute_params: Optional[List[torch.nn.Parameter]] = None,
+) -> torch.Tensor:
+    """
+    Apply tiled MLP computation for memory efficiency.
+
+    Args:
+        fn: the function to call on sharded inputs (e.g., lambda module, x: module(x))
+        mlp_module: the MLP nn.Module object
+        x: the input tensor with shape [bs, seqlen, hidden_size] or [seqlen, hidden_size]
+        num_shards: number of shards to use. If None, automatically calculated as ceil(seqlen / hidden_size)
+        compute_params: list of parameters for DeepSpeed ZeRO optimization
+
+    Returns:
+        output tensor with the same shape as input
+    """
+    if num_shards is None:
+        # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size]
+        hidden_size = x.shape[-1]
+        seqlen = x.shape[-2]
+        num_shards = math.ceil(seqlen / hidden_size)
+
+    # Ensure num_shards is at least 1
+    num_shards = max(1, num_shards)
+
+    return LigerTiledMLPFunction.apply(
+        fn,
+        mlp_module,
+        x,
+        num_shards,
+        compute_params,
+    )

liger_kernel/ops/utils.py

@@ -78,6 +78,8 @@ def get_amp_custom_fwd_bwd() -> Callable:
             functools.partial(torch.amp.custom_fwd, device_type=device),
             functools.partial(torch.amp.custom_bwd, device_type=device),
         )
+    if hasattr(torch, "npu") and getattr(torch.npu, "amp", None) is not None:
+        return torch.npu.amp.custom_fwd, torch.npu.amp.custom_bwd
     return torch.cuda.amp.custom_fwd, torch.cuda.amp.custom_bwd
 
 
@@ -125,3 +127,15 @@ def element_mul_kernel(
         X_offsets = i + tl.arange(0, BLOCK_SIZE)
         X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols)
         tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols)
+
+
+def get_npu_core_count(default: int = 20) -> int:
+    """Return NPU vector core count.
+    Fallback to `default` if Triton runtime or NPU device is unavailable.
+    """
+    try:
+        utils = triton.runtime.driver.active.utils
+        props = utils.get_device_properties(0)
+        return int(props.get("num_vectorcore", default))
+    except Exception:
+        return default