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

liger_kernel/ops/llama4_rope.py

@@ -0,0 +1,225 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def _prepare_freqs(freqs_cis: torch.Tensor, seq_len: int, head_dim_half: int):
+    # Split or unpack complex frequencies into real and imag parts
+    if freqs_cis.is_complex():
+        freqs_real = freqs_cis.real
+        freqs_imag = freqs_cis.imag
+    else:
+        # Already split: last dim should be 2*head_dim_half
+        if freqs_cis.shape[-1] == 2 * head_dim_half:
+            freqs_real = freqs_cis[..., :head_dim_half]
+            freqs_imag = freqs_cis[..., head_dim_half:]
+        else:
+            raise ValueError(
+                f"Unexpected freqs_cis shape for non-complex input: {freqs_cis.shape}, expected last dim = {2 * head_dim_half}"
+            )
+
+    # Canonicalize to shape (seq_len, head_dim_half):
+    # 1) Ensure the last dimension is head_dim_half
+    if freqs_real.shape[-1] != head_dim_half:
+        raise ValueError(f"Unexpected last dim for freqs: {freqs_real.shape[-1]} (expected {head_dim_half})")
+    # 2) Flatten all leading dims to a single row dimension
+    freqs_real = freqs_real.reshape(-1, head_dim_half)
+    freqs_imag = freqs_imag.reshape(-1, head_dim_half)
+    # 3) If we have fewer rows than seq_len, allow broadcasting when single row
+    if freqs_real.shape[0] < seq_len:
+        if freqs_real.shape[0] == 1:
+            freqs_real = freqs_real.expand(seq_len, -1)
+            freqs_imag = freqs_imag.expand(seq_len, -1)
+        else:
+            raise ValueError(f"Insufficient rows in freqs: {freqs_real.shape[0]} < seq_len={seq_len}")
+    # 4) If we have more rows than seq_len (e.g., batch present), take the first seq_len rows
+    elif freqs_real.shape[0] > seq_len:
+        freqs_real = freqs_real[:seq_len]
+        freqs_imag = freqs_imag[:seq_len]
+
+    return freqs_real, freqs_imag
+
+
+def _maybe_to_dtype(t: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    return t if t.dtype == dtype else t.to(dtype)
+
+
+def _maybe_contiguous(t: torch.Tensor) -> torch.Tensor:
+    return t if t.is_contiguous() else t.contiguous()
+
+
+def _cast_and_contiguous(q, k, freqs_real, freqs_imag):
+    # Choose compute dtype: use fp32 only when inputs are fp32; otherwise keep input dtype for performance
+    compute_dtype = torch.float32 if q.dtype == torch.float32 else q.dtype
+
+    # Make sure q/k share the same dtype before casting to compute dtype
+    if k.dtype != q.dtype:
+        k = k.to(q.dtype)
+
+    q = _maybe_contiguous(_maybe_to_dtype(q, compute_dtype))
+    k = _maybe_contiguous(_maybe_to_dtype(k, compute_dtype))
+    freqs_real = _maybe_contiguous(_maybe_to_dtype(freqs_real, compute_dtype))
+    freqs_imag = _maybe_contiguous(_maybe_to_dtype(freqs_imag, compute_dtype))
+    return q, k, freqs_real, freqs_imag
+
+
+@triton.jit
+def _llama4_rope_kernel(
+    q_ptr,
+    k_ptr,
+    freqs_real_ptr,
+    freqs_imag_ptr,
+    q_row_stride,
+    k_row_stride,
+    q_head_stride,
+    k_head_stride,
+    freqs_row_stride,
+    seq_len,
+    batch_size,
+    imag_sign,
+    head_dim_half: tl.constexpr,
+    n_q_heads: tl.constexpr,
+    n_k_heads: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    H100-optimized RoPE kernel with improved parallelization across heads and dimensions.
+    Grid: (batch*seq, head)
+    """
+    # 2D grid
+    pid_bs = tl.program_id(0)  # over batch*seq
+    pid_h = tl.program_id(1)  # over heads
+
+    batch_idx = pid_bs // seq_len
+    seq_idx = pid_bs % seq_len
+
+    # Bounds check
+    if batch_idx >= batch_size or seq_idx >= seq_len:
+        return
+
+    # Base pointers for this (batch, seq) position
+    base_offset = batch_idx * seq_len + seq_idx
+    q_base = q_ptr + base_offset * q_row_stride
+    k_base = k_ptr + base_offset * k_row_stride
+
+    # Tiling over dim/2
+    for d_start in tl.static_range(0, head_dim_half, BLOCK_SIZE):
+        d_indices = d_start + tl.arange(0, BLOCK_SIZE)
+        mask_d = d_indices < head_dim_half
+
+        # Load frequencies once per tile (freqs layout: [seq_len, head_dim_half])
+        freq_idx = d_indices
+        freqs_real = tl.load(freqs_real_ptr + seq_idx * freqs_row_stride + freq_idx, mask=mask_d, other=0.0)
+        freqs_imag = tl.load(freqs_imag_ptr + seq_idx * freqs_row_stride + freq_idx, mask=mask_d, other=0.0)
+        freqs_imag = freqs_imag * imag_sign
+
+        # Process one query head per program in pid_h
+        if pid_h < n_q_heads:
+            q_head_ptr = q_base + pid_h * q_head_stride
+            q_real = tl.load(q_head_ptr + d_indices * 2, mask=mask_d, other=0.0)
+            q_imag = tl.load(q_head_ptr + d_indices * 2 + 1, mask=mask_d, other=0.0)
+
+            # Complex multiply with FMAs: (a+ib)*(c+i d) = (a*c - b*d) + i(a*d + b*c)
+            new_q_real = tl.math.fma(q_real, freqs_real, -(q_imag * freqs_imag))
+            new_q_imag = tl.math.fma(q_real, freqs_imag, q_imag * freqs_real)
+
+            tl.store(q_head_ptr + d_indices * 2, new_q_real, mask=mask_d)
+            tl.store(q_head_ptr + d_indices * 2 + 1, new_q_imag, mask=mask_d)
+
+        # Process one key head per program in pid_h
+        if pid_h < n_k_heads:
+            k_head_ptr = k_base + pid_h * k_head_stride
+            k_real = tl.load(k_head_ptr + d_indices * 2, mask=mask_d, other=0.0)
+            k_imag = tl.load(k_head_ptr + d_indices * 2 + 1, mask=mask_d, other=0.0)
+
+            new_k_real = tl.math.fma(k_real, freqs_real, -(k_imag * freqs_imag))
+            new_k_imag = tl.math.fma(k_real, freqs_imag, k_imag * freqs_real)
+
+            tl.store(k_head_ptr + d_indices * 2, new_k_real, mask=mask_d)
+            tl.store(k_head_ptr + d_indices * 2 + 1, new_k_imag, mask=mask_d)
+
+
+def _select_kernel_meta(head_dim_half: int):
+    # Heuristic tuning for block size and num_warps
+    if head_dim_half >= 256:
+        return 128, 8
+    if head_dim_half >= 96:
+        return 128, 4
+    if head_dim_half >= 48:
+        return 64, 4
+    if head_dim_half >= 24:
+        return 32, 2
+    return 16, 2
+
+
+def llama4_rope_forward(q, k, freqs_cis, BLOCK_SIZE: int = None, imag_sign: float = 1.0):
+    # Save original dtype for casting back
+    original_dtype = q.dtype
+
+    batch_size, seq_len, n_q_heads, head_dim = q.shape
+    _, _, n_k_heads, _ = k.shape
+    head_dim_half = head_dim // 2
+
+    # Prepare frequencies
+    freqs_real, freqs_imag = _prepare_freqs(freqs_cis, seq_len, head_dim_half)
+
+    # Cast to appropriate dtype and make contiguous only when needed
+    q, k, freqs_real, freqs_imag = _cast_and_contiguous(q, k, freqs_real, freqs_imag)
+
+    # H100-optimized meta-params
+    if BLOCK_SIZE is None:
+        BLOCK_SIZE, num_warps = _select_kernel_meta(head_dim_half)
+    else:
+        # Provide a default num_warps if caller pins BLOCK_SIZE
+        _, num_warps = _select_kernel_meta(head_dim_half)
+
+    # 2D grid: one program per (batch, seq, head)
+    n_heads_max = max(n_q_heads, n_k_heads)
+    grid = (batch_size * seq_len, n_heads_max)
+
+    # Launch kernel
+    _llama4_rope_kernel[grid](
+        q,
+        k,
+        freqs_real,
+        freqs_imag,
+        q.stride(1),
+        k.stride(1),
+        q.stride(2),
+        k.stride(2),
+        freqs_real.stride(0),
+        seq_len,
+        batch_size,
+        imag_sign,
+        head_dim_half,
+        n_q_heads,
+        n_k_heads,
+        BLOCK_SIZE,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+
+    # Cast back to original dtype only if it differs from compute dtype
+    if q.dtype != original_dtype:
+        q = q.to(original_dtype)
+    if k.dtype != original_dtype:
+        k = k.to(original_dtype)
+
+    return q, k
+
+
+class LigerLlama4RopeFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, freqs_cis, BLOCK_SIZE: int = None):
+        q_out, k_out = llama4_rope_forward(q, k, freqs_cis, BLOCK_SIZE, imag_sign=1.0)
+        ctx.save_for_backward(freqs_cis.detach() if isinstance(freqs_cis, torch.Tensor) else freqs_cis)
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        return q_out, k_out
+
+    @staticmethod
+    def backward(ctx, dq, dk):
+        (freqs_cis,) = ctx.saved_tensors
+        BLOCK_SIZE = getattr(ctx, "BLOCK_SIZE", None)
+        # Use imag_sign=-1.0 for conjugate without materializing a new tensor
+        dq_out, dk_out = llama4_rope_forward(dq, dk, freqs_cis, BLOCK_SIZE, imag_sign=-1.0)
+        return dq_out, dk_out, None

liger_kernel/ops/poly_norm.py

@@ -0,0 +1,386 @@
+import operator
+
+import torch
+import triton
+import triton.language as tl
+
+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:
+        from triton.language.extra.libdevice import rsqrt
+    except ModuleNotFoundError:
+        from triton.language.extra.cuda.libdevice import rsqrt
+else:
+    from triton.language.math import rsqrt
+
+
+@triton.jit
+def _poly_norm_forward_kernel(
+    Y_ptr,
+    Y_row_stride,
+    X_ptr,
+    X_row_stride,
+    W_ptr,  # weight: [3] for [w0, w1, w2]
+    B_ptr,  # bias: scalar
+    RSTD_ptr,  # cache rstd for backward: shape (n_rows, 3)
+    RSTD_row_stride,
+    n_cols,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    PolyNorm formula:
+        y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+        where norm(u) = u / sqrt(mean(u²) + ε)
+
+    Reference:
+    1. https://github.com/BryceZhuo/PolyCom/
+    2. https://arxiv.org/pdf/2411.03884
+
+    Cache rstd values for backward pass
+    """
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    # Load pointers
+    Y_ptr += row_idx * Y_row_stride
+    X_ptr += row_idx * X_row_stride
+    RSTD_ptr += row_idx * RSTD_row_stride
+
+    # Load input row
+    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
+
+    # Load weights and bias
+    w0 = tl.load(W_ptr + 0)
+    w1 = tl.load(W_ptr + 1)
+    w2 = tl.load(W_ptr + 2)
+    b = tl.load(B_ptr)
+
+    # Compute x³, x², x
+    X_pow3 = X_row * X_row * X_row
+    X_pow2 = X_row * X_row
+    X_pow1 = X_row
+
+    # Compute norm(x³): norm(u) = u * rsqrt(mean(u²) + eps)
+    mean_square_3 = tl.sum(X_pow3 * X_pow3, axis=0) / n_cols
+    rstd_3 = rsqrt(mean_square_3 + eps)
+    norm_x3 = X_pow3 * rstd_3
+
+    # Compute norm(x²)
+    mean_square_2 = tl.sum(X_pow2 * X_pow2, axis=0) / n_cols
+    rstd_2 = rsqrt(mean_square_2 + eps)
+    norm_x2 = X_pow2 * rstd_2
+
+    # Compute norm(x)
+    mean_square_1 = tl.sum(X_pow1 * X_pow1, axis=0) / n_cols
+    rstd_1 = rsqrt(mean_square_1 + eps)
+    norm_x1 = X_pow1 * rstd_1
+
+    # Cache rstd values for backward
+    tl.store(RSTD_ptr + 0, rstd_3)
+    tl.store(RSTD_ptr + 1, rstd_2)
+    tl.store(RSTD_ptr + 2, rstd_1)
+
+    # Compute output: y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+    Y_row = w0 * norm_x3 + w1 * norm_x2 + w2 * norm_x1 + b
+
+    # Store output
+    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
+
+
+@triton.jit
+def _poly_norm_backward_kernel(
+    dY_ptr,
+    dY_row_stride,
+    dX_ptr,
+    dX_row_stride,
+    X_ptr,
+    X_row_stride,
+    W_ptr,
+    RSTD_ptr,
+    RSTD_row_stride,
+    dW_ptr,  # shape: (n_programs, 3)
+    dW_row_stride,
+    dB_ptr,  # shape: (n_programs,)
+    n_rows,
+    n_cols,
+    rows_per_program: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    PolyNorm Backward Kernel Gradient:
+        ∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
+
+    where:
+        - D_p = RMS(x^p) = 1/rstd_p
+        - S_p = sum(grad * x^p) over the row
+        - d = n_cols
+        - p ∈ {3, 2, 1}
+    """
+    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
+
+    # Initialize accumulators for weight and bias gradients (scalars)
+    dW0_acc = 0.0
+    dW1_acc = 0.0
+    dW2_acc = 0.0
+    dB_acc = 0.0
+
+    # Load weights
+    w0 = tl.load(W_ptr + 0).to(tl.float32)
+    w1 = tl.load(W_ptr + 1).to(tl.float32)
+    w2 = tl.load(W_ptr + 2).to(tl.float32)
+
+    dY_ptr += row_start * dY_row_stride
+    dX_ptr += row_start * dX_row_stride
+    X_ptr += row_start * X_row_stride
+    RSTD_ptr += row_start * RSTD_row_stride
+
+    for _ in range(row_start, row_end):
+        # Load input and gradient
+        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
+        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
+
+        # Load cached rstd values
+        rstd_3 = tl.load(RSTD_ptr + 0).to(tl.float32)
+        rstd_2 = tl.load(RSTD_ptr + 1).to(tl.float32)
+        rstd_1 = tl.load(RSTD_ptr + 2).to(tl.float32)
+
+        # Compute powers
+        X_pow3 = X_row * X_row * X_row
+        X_pow2 = X_row * X_row
+        X_pow1 = X_row
+
+        # Accumulate bias gradient: dB = sum(dY)
+        dB_acc += tl.sum(dY_row, axis=0)
+
+        # Compute gradient w.r.t. input using closed-form formula
+        # For p=3: ∂L/∂x from w0 * norm(x³)
+        S_3 = tl.sum(dY_row * X_pow3, axis=0)  # scalar
+        grad_x_3 = w0 * (
+            3.0 * X_pow2 * rstd_3 * dY_row
+            - (3.0 / n_cols) * X_row * X_row * X_row * X_row * X_row * (rstd_3 * rstd_3 * rstd_3) * S_3
+        )
+
+        # For p=2: ∂L/∂x from w1 * norm(x²)
+        S_2 = tl.sum(dY_row * X_pow2, axis=0)  # scalar
+        grad_x_2 = w1 * (
+            2.0 * X_row * rstd_2 * dY_row - (2.0 / n_cols) * X_row * X_row * X_row * (rstd_2 * rstd_2 * rstd_2) * S_2
+        )
+
+        # For p=1: ∂L/∂x from w2 * norm(x)
+        S_1 = tl.sum(dY_row * X_pow1, axis=0)  # scalar
+        grad_x_1 = w2 * (1.0 * rstd_1 * dY_row - (1.0 / n_cols) * X_row * (rstd_1 * rstd_1 * rstd_1) * S_1)
+
+        # Accumulate weight gradients using closed-form: dW_p = rstd_p * S_p
+        dW0_acc += rstd_3 * S_3
+        dW1_acc += rstd_2 * S_2
+        dW2_acc += rstd_1 * S_1
+
+        # Total gradient
+        dX_row = grad_x_3 + grad_x_2 + grad_x_1
+
+        # Store gradient
+        tl.store(dX_ptr + col_offsets, dX_row, mask=mask)
+
+        # Update pointers
+        dY_ptr += dY_row_stride
+        dX_ptr += dX_row_stride
+        X_ptr += X_row_stride
+        RSTD_ptr += RSTD_row_stride
+
+    # Store accumulated gradients (scalars)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 0, dW0_acc)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 1, dW1_acc)
+    tl.store(dW_ptr + row_block_id * dW_row_stride + 2, dW2_acc)
+    tl.store(dB_ptr + row_block_id, dB_acc)
+
+
+def poly_norm_forward(X, W, B, eps=1e-6):
+    """
+    PolyNorm Forward Pass
+
+    Args:
+        X: input tensor of shape (*, H) where H is hidden dimension
+        W: weight tensor of shape (3,) for [w0, w1, w2]
+        B: bias scalar tensor
+        eps: epsilon for numerical stability
+
+    Returns:
+        Y: output tensor of same shape as X
+        X: reshaped input (for backward)
+        RSTD: cached rstd values (for backward)
+        BLOCK_SIZE: block size used
+        num_warps: number of warps used
+    """
+    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)
+
+    # RSTD is to cache rstd for each row
+    Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
+    RSTD = torch.empty((n_rows, 3), dtype=torch.float32, device=X.device)
+
+    # Check constraints
+    assert W.shape[0] == 3, "Weight tensor must have shape (3,)"
+    assert B.numel() == 1, "Bias must be a scalar"
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
+    # Launch kernel
+    _poly_norm_forward_kernel[(n_rows,)](
+        Y,
+        Y.stride(0),
+        X,
+        X.stride(0),
+        W,
+        B,
+        RSTD,
+        RSTD.stride(0),
+        n_cols,
+        eps,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        **kernel_args,
+    )
+
+    return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps
+
+
+def poly_norm_backward(dY, X, W, RSTD, BLOCK_SIZE, num_warps, in_place):
+    """
+    PolyNorm Backward Pass
+
+    Args:
+        dY: gradient of output
+        X: input tensor (already reshaped to 2D)
+        W: weight tensor
+        RSTD: cached rstd values from forward
+        BLOCK_SIZE: block size from forward
+        num_warps: number of warps from forward
+        in_place: whether to in-place modify dY to store dX (saves memory)
+
+    Returns:
+        dX: gradient w.r.t. input
+        dW: gradient w.r.t. weight
+        dB: gradient w.r.t. bias
+    """
+    shape = dY.shape
+    dim = shape[-1]
+    dY = dY.view(-1, dim)
+    n_rows, n_cols = dY.shape
+
+    # Get number of SMs for parallelization
+    import math
+
+    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 or reuse gradients
+    if in_place is True:
+        dX = dY
+    else:
+        dX = torch.zeros_like(dY)
+
+    _dW = torch.empty((sm_count, 3), dtype=torch.float32, device=W.device)
+    _dB = torch.empty((sm_count,), dtype=torch.float32, device=W.device)
+
+    rows_per_program = math.ceil(n_rows / sm_count)
+    grid = (sm_count,)
+
+    # XPU-specific optimization
+    kernel_args = {}
+    if X.device.type == "xpu":
+        kernel_args["grf_mode"] = "large"
+
+    # Launch backward kernel
+    _poly_norm_backward_kernel[grid](
+        dY,
+        dY.stride(0),
+        dX,
+        dX.stride(0),
+        X,
+        X.stride(0),
+        W,
+        RSTD,
+        RSTD.stride(0),
+        _dW,
+        _dW.stride(0),
+        _dB,
+        n_rows,
+        n_cols,
+        rows_per_program,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+        **kernel_args,
+    )
+
+    # Reduce gradients across SMs
+    dX = dX.view(*shape)
+    dW = _dW.sum(dim=0).to(W.dtype)
+    dB = _dB.sum().to(W.dtype)
+
+    return dX, dW, dB
+
+
+class LigerPolyNormFunction(torch.autograd.Function):
+    """
+    PolyNorm Function with forward and backward pass
+
+    PolyNorm formula:
+        y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+        where norm(u) = u / sqrt(mean(u²) + ε)
+
+    Backward uses closed-form gradient:
+        ∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
+    """
+
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, X, W, B, eps=1e-6, in_place=True):
+        """
+        Args:
+            X: input tensor of shape (B, T, H) or (BxT, H)
+            W: weight tensor of shape (3,) for [w0, w1, w2]
+            B: bias scalar
+            eps: epsilon for numerical stability
+            in_place: whether to in-place modify grad_output in backward (saves memory)
+
+        Returns:
+            Y: output tensor of same shape as X
+        """
+        Y, X, RSTD, BLOCK_SIZE, num_warps = poly_norm_forward(X, W, B, eps)
+        ctx.BLOCK_SIZE = BLOCK_SIZE
+        ctx.num_warps = num_warps
+        ctx.in_place = in_place
+        ctx.save_for_backward(X, W, RSTD)
+        return Y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, grad_output):
+        """
+        Args:
+            grad_output: gradient of output
+
+        Returns:
+            dX, dW, dB: gradients w.r.t. X, W, B
+        """
+        X, W, RSTD = ctx.saved_tensors
+        dX, dW, dB = poly_norm_backward(grad_output, X, W, RSTD, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place)
+        return dX, dW, dB, None, None