linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(cutedsl): fuse Gemma RMSNorm weight offset in FP32 (#1462)

Browse Source 
## Summary

The CuTe DSL Gemma paths in RMSNorm and FusedAddRMSNorm add the offset
to BF16/FP16 weights before promoting them to FP32. This introduces
extra rounding in the forward output and input gradient. For example,
BF16 `0.00390625 + 1` rounds to `1.0`, while the FP32 sum is
`1.00390625`.

Pass the weight and offset separately, then add the offset in FP32
inside the forward and backward kernels. This also removes the separate
weight-add kernel and temporary vector.

The fused offset handling follows
[Quack](https://github.com/Dao-AILab/quack/blob/10ce505a7ecf68e46753009bd08dcf9bd6c059ee/quack/rmsnorm.py).

## Testing Done

- Hardware Type: NVIDIA H100 and NVIDIA B200.
- Existing selected suites: **233 passed, 0 failed, 0 skipped on each
GPU, before and after**; 125 deselected. Covers RMSNorm, LayerNorm,
FusedAddRMSNorm and CUDA Graph replay.
- No new pytest cases or changes to existing test tolerances.

The existing BF16 Gemma RMSNorm test at M=128, N=4096 passes on both
revisions. Its maximum absolute error against Triton improves
identically on both GPUs:

| Output | Before | After |
|---|---:|---:|
| Y | 0.0625 | 0 |
| dX | 0.0625 | 0.0078125 |
| dW | 0 | 0 |

Cache-cleared benchmarks show lower forward latency across the seven
measured BF16 RMSNorm shapes: **3.8–31.4% on H100 and 13.9–42.1% on
B200**. Removing the BF16 weight vector saves **8/16/36 KiB** of
forward/full peak allocation at N=4096/8192/18432. Full-call timings
have more variation; details are below.

- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Full-repository and model-convergence suites were not run. No end-to-end
training speed, memory or convergence claim is made.

<details>
<summary>Environment, test command and original pytest output</summary>

Python 3.12.10, PyTorch 2.9.1, CUDA 12.8, Triton 3.5.1, CuTe DSL 4.6.0,
TVM FFI 0.1.13.post3. CUDA-visible memory: H100 79.2 GiB; B200 178.4
GiB.

Run on the base commit `95b01e9` and this branch:

```bash
LIGER_KERNEL_IMPL=cutedsl python -m pytest \
  test/cutedsl/test_rms_norm.py test/cutedsl/test_rms_norm_stream.py \
  test/ops/test_rms_norm.py test/ops/test_layer_norm.py \
  test/ops/test_fused_add_rms_norm.py \
  -k "nvidia-cutedsl or test_rms_norm_parity or test_rms_norm_mixed_weight_dtype or graph" -q
```

Original pytest summary lines, labelled by GPU and revision:

```text
H100 — before
========= 233 passed, 125 deselected, 2 warnings in 116.55s (0:01:56) ==========

H100 — after
========== 233 passed, 125 deselected, 2 warnings in 90.10s (0:01:30) ==========

B200 — before
========= 233 passed, 125 deselected, 2 warnings in 153.33s (0:02:33) ==========

B200 — after
========= 233 passed, 125 deselected, 2 warnings in 102.62s (0:01:42) ==========
```

The error table uses
`test_rms_norm_parity[False-True-1.0-gemma-dtype1-2-64-4096]`. Some
irregular parity shapes exercise the existing inline fallback.

Supplemental diagnostics on each GPU passed 122 numerical cases against
a CuTe FP32-pre-add control and 26 controls for non-Gemma/zero-offset
behavior, offset cache keys and CUDA Graph replay with changed X, W and
dY. Coverage includes BF16/FP16/FP32, mixed FP32 weights, strided
inputs, in-place/out-of-place backward, residual-output gradients and
hidden sizes 512–18432, including boundaries around 4096, 8192 and
16384. These diagnostics are separate from the committed pytest suite.

</details>

<details>
<summary>PyTorch accuracy comparison and minimal reproduction</summary>

Relative L2 error against PyTorch/autograd, using BF16 X/W, M=128,
N=4096, offset=1, seed=42 and weight scale=0.01:

| GPU | Op | Y before → after | dX before → after |
|---|---|---:|---:|
| H100 | RMSNorm | 0.00292 → 2.3e-05 | 0.00288 → 1.22e-05 |
| H100 | FusedAddRMSNorm | 0.00289 → 2.53e-05 | 0.00191 → 0.000253 |
| B200 | RMSNorm | 0.00292 → 2.3e-05 | 0.00288 → 1.21e-05 |
| B200 | FusedAddRMSNorm | 0.00289 → 2.53e-05 | 0.00191 → 0.000256 |

FusedAddRMSNorm includes a residual-output gradient. Its backward uses a
rounded saved residual, which contributes to the remaining difference
from PyTorch. dW is unchanged for fixed inputs and upstream gradients.

Run this unchanged on each revision using the existing Gemma PyTorch
reference:

```python
import torch
from liger_kernel.ops.backends._cutedsl.rms_norm import _LigerRMSNormCuTeDSLFunction
from liger_kernel.testing._op_helpers import pytorch_reference_rms_norm

torch.manual_seed(42)
x = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)
w = (torch.randn(4096, device="cuda") * 0.01).bfloat16().requires_grad_()
dy = torch.randn_like(x)
xr = x.detach().clone().requires_grad_()
wr = w.detach().clone().requires_grad_()
y = _LigerRMSNormCuTeDSLFunction.apply(x, w, 1e-6, 1.0, "gemma", False, None)
yr = pytorch_reference_rms_norm(xr, wr, 1e-6, 1.0, "gemma")
actual = (y, *torch.autograd.grad(y, (x, w), dy))
expected = (yr, *torch.autograd.grad(yr, (xr, wr), dy.clone()))
for name, a, b in zip(("Y", "dX", "dW"), actual, expected):
    delta = a.detach().double() - b.detach().double()
    rel_l2 = delta.norm() / b.detach().double().norm().clamp_min(1e-30)
    print(name, "max_abs=", delta.abs().max().item(), "relative_l2=", rel_l2.item())
```

</details>

<details>
<summary>Benchmark results and measurement details</summary>

Cache-cleared `triton.testing.do_bench`, warmup=5 ms, rep=30 ms, median
of nine shuffled-provider rounds. Compilation and input/dY generation
are excluded. `in_place=False`; full includes `autograd.grad`. An
identical baseline is interleaved as an A/A control. GPU clocks were not
locked; full timings include host enqueue gaps.

Measured 16 configurations × forward/backward/full = 48 combinations per
GPU, including seven BF16 RMSNorm shapes, three FusedAddRMSNorm shapes,
other dtypes and zero-offset controls. Profiling confirms removal of one
weight-add kernel from each affected forward.

BF16 results, before → after (µs):

| GPU | Op | M × N | Forward | Forward + backward |
|---|---|---|---:|---:|
| H100 | RMSNorm | 1 × 4096 | 8.864 → 6.080 | 102.640 → 95.184 |
| H100 | RMSNorm | 128 × 4096 | 9.728 → 6.880 | 99.856 → 87.680 |
| H100 | RMSNorm | 2048 × 4096 | 20.864 → 18.944 | 103.072 → 98.160 |
| H100 | RMSNorm | 8192 × 4096 | 54.048 → 51.968 | 141.472 → 139.488 |
| H100 | RMSNorm | 128 × 8192 | 10.592 → 8.000 | 103.744 → 92.160 |
| H100 | RMSNorm | 2048 × 8192 | 31.808 → 29.504 | 107.216 → 94.672 |
| H100 | RMSNorm | 128 × 18432 | 12.448 → 10.016 | 102.352 → 92.320 |
| H100 | FusedAddRMSNorm | 128 × 4096 | 10.528 → 7.872 | 130.656 →
120.576 |
| H100 | FusedAddRMSNorm | 2048 × 4096 | 30.800 → 28.224 | 130.672 →
119.888 |
| H100 | FusedAddRMSNorm | 2048 × 8192 | 53.440 → 50.720 | 128.032 →
116.304 |
| B200 | RMSNorm | 1 × 4096 | 11.104 → 6.432 | 61.120 → 55.680 |
| B200 | RMSNorm | 128 × 4096 | 11.296 → 8.192 | 63.648 → 58.624 |
| B200 | RMSNorm | 2048 × 4096 | 17.280 → 14.336 | 71.104 → 64.448 |
| B200 | RMSNorm | 8192 × 4096 | 33.664 → 28.992 | 97.440 → 95.232 |
| B200 | RMSNorm | 128 × 8192 | 12.288 → 8.192 | 71.008 → 62.528 |
| B200 | RMSNorm | 2048 × 8192 | 23.360 → 18.496 | 66.256 → 64.528 |
| B200 | RMSNorm | 128 × 18432 | 13.216 → 10.048 | 59.968 → 61.008 |
| B200 | FusedAddRMSNorm | 128 × 4096 | 13.152 → 8.192 | 71.760 → 64.432
|
| B200 | FusedAddRMSNorm | 2048 × 4096 | 22.496 → 18.464 | 85.216 →
77.312 |
| B200 | FusedAddRMSNorm | 2048 × 8192 | 33.568 → 28.704 | 92.096 →
79.008 |

H100 RMSNorm backward-only changes range from −0.8% to +4.7%; B200
backward-only medians are unchanged. Backward incremental allocation and
full-call live allocation are unchanged.

The B200 128×18432 full result has substantial A/A variation, and its
apparent slowdown changed sign on repeat. A separate 21-round
zero-offset repeat measured 53.760 → 56.096 µs (+4.3%), while the
identical-baseline A/A control differed by +5.2%; forward was unchanged
and backward differed by −0.15%. These measurements do not resolve a
stable kernel regression or establish zero overhead.

</details>

Co-authored-by: Arup De <[email protected]>
This commit is contained in:
luca huang
2026-10-06 03:56:30 +00:00
committed by GitHub
co-authored by Arup De
parent 80fb156fce
commit e2039eb8c7
3 changed files with 88 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/backends/_cutedsl/_cute_lib/rmsnorm_fwd.py
+13 -4
View File
@@ -60,9 +60,10 @@ class RMSNorm(ReductionBase):
sum-of-squares reduction.
"""
def __init__(self, dtype: Type[cutlass.Numeric], N: int, is_layernorm: bool = False):
def __init__(self, dtype: Type[cutlass.Numeric], N: int, is_layernorm: bool = False, weight_offset: float = 0.0):
super().__init__(dtype, N, stage=2 if is_layernorm else 1)
self.is_layernorm = is_layernorm
self.weight_offset = weight_offset
self.reload_from = None if N <= (16384 if is_layernorm else 8192) else "smem"
self.delay_w_load = False
@@ -327,7 +328,10 @@ class RMSNorm(ReductionBase):
x_hat = (x - mean) * rstd if const_expr(self.is_layernorm) else x * rstd
y = x_hat
if const_expr(mW is not None):
y *= tXrW.load().to(cute.Float32)
w = tXrW.load().to(cute.Float32)
if const_expr(self.weight_offset != 0.0):
w += Float32(self.weight_offset)
y *= w
if const_expr(mB is not None):
y += tXrB.load().to(cute.Float32)
tXrO.store(y.to(tXrO.element_type))
@@ -355,6 +359,7 @@ def _compile_fwd(
has_mean,
is_layernorm,
per_head,
weight_offset=0.0,
):
"""Compile the RMSNorm/LayerNorm forward kernel for a given dtype mix.
@@ -374,6 +379,7 @@ def _compile_fwd(
has_mean,
is_layernorm,
per_head,
weight_offset,
)
if key in _FWD_COMPILE_CACHE:
return _FWD_COMPILE_CACHE[key]
@@ -391,7 +397,7 @@ def _compile_fwd(
rstd_cute = fake_tensor(Float32, batch_shape) if has_rstd else None
mean_cute = fake_tensor(Float32, batch_shape) if has_mean else None
compiled = cute.compile(
RMSNorm(dtype, N, is_layernorm=is_layernorm),
RMSNorm(dtype, N, is_layernorm=is_layernorm, weight_offset=weight_offset),
x_cute,
weight_cute,
bias_cute,
@@ -419,6 +425,7 @@ def _run_fwd(
residual_out: Optional[Tensor],
eps: float,
is_layernorm: bool,
weight_offset: float = 0.0,
) -> None:
"""Run the compiled kernel. Mirrors ``quack.rmsnorm._rmsnorm_fwd`` but
drops the ``torch.library.custom_op`` registration — we are always
@@ -448,6 +455,7 @@ def _run_fwd(
mean is not None,
is_layernorm,
per_head,
weight_offset,
)(x, weight, bias, residual, out, residual_out, rstd, mean, eps)
@@ -463,6 +471,7 @@ def rmsnorm_fwd(
residual_dtype: Optional[torch.dtype] = None,
eps: float = 1e-6,
store_rstd: bool = False,
weight_offset: float = 0.0,
) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
"""RMSNorm forward. Drop-in replacement for ``quack.rmsnorm.rmsnorm_fwd``.
@@ -479,7 +488,7 @@ def rmsnorm_fwd(
residual_out = torch.empty_like(x, dtype=residual_dtype if residual_dtype is not None else x.dtype)
else:
residual_out = None
_run_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False)
_run_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False, weight_offset)
# residual_out is None if residual is None and residual_dtype == input_dtype
if residual_out is None:
residual_out = x
src/liger_kernel/ops/backends/_cutedsl/fused_add_rms_norm.py
+26 -20
View File
@@ -122,8 +122,8 @@ def _fused_add_rms_norm_cutedsl_forward(
) -> Tuple[Tensor, Tensor, Optional[Tensor], Tensor]:
"""Forward via the inlined ``_cute_lib.rmsnorm_fwd`` with ``residual=R``.
Returns ``(Y, S, W_eff, RSTD)`` where ``S = X + R`` is the summed residual
(saved for the backward), ``W_eff`` is ``weight + offset``, and ``RSTD``
Returns ``(Y, S, W_kernel, RSTD)`` where ``S = X + R`` is the summed residual
(saved for the backward), ``W_kernel`` is the kernel weight (raw for Gemma), and ``RSTD``
is fp32.
"""
shape = X.shape
@@ -163,25 +163,27 @@ def _fused_add_rms_norm_cutedsl_forward(
X_flat = X_in.view(-1, N).contiguous()
R_flat = R.view(-1, N).contiguous()
# Fold offset into the weight host-side.
if offset != 0.0:
w_eff = W + offset
# Gemma adds offset in FP32 inside the forward/backward kernels.
weight_offset = offset if casting_mode == _CASTING_MODE_GEMMA else 0.0
if offset != 0.0 and casting_mode != _CASTING_MODE_GEMMA:
w_kernel = W + offset
else:
w_eff = W
# Only the mixed-dtype gemma route promotes w_eff host-side; the
w_kernel = W
# Only the mixed-dtype gemma route promotes w_kernel host-side; the
# kernel up-casts in-register otherwise (bit-exact).
if _promote_gemma_fp32 and w_eff.dtype != torch.float32:
w_eff = w_eff.to(torch.float32)
w_eff = w_eff.contiguous()
if _promote_gemma_fp32 and w_kernel.dtype != torch.float32:
w_kernel = w_kernel.to(torch.float32)
w_kernel = w_kernel.contiguous()
# rmsnorm_fwd with residual=R fuses the add into the kernel.
# Returns (out, residual_out=S, rstd).
out, S, rstd = _cutedsl_rmsnorm_fwd(
X_flat,
weight=w_eff,
weight=w_kernel,
residual=R_flat,
eps=eps,
store_rstd=True,
weight_offset=weight_offset,
)
# Back-cast if the promote route above widened the output (it is the
@@ -195,17 +197,18 @@ def _fused_add_rms_norm_cutedsl_forward(
if S.dtype != X.dtype:
S = S.to(X.dtype)
return out.view(shape), S.view(shape), w_eff, rstd
return out.view(shape), S.view(shape), w_kernel, rstd
def _fused_add_rms_norm_cutedsl_backward(
dY: Tensor,
dS_out: Optional[Tensor],
S: Tensor,
W_eff: Tensor,
W_kernel: Tensor,
RSTD: Tensor,
in_place: bool,
casting_mode: int,
weight_offset: float = 0.0,
) -> Tuple[Tensor, Tensor, Tensor]:
"""Backward via the inline CuTe DSL rms_norm backward kernel.
@@ -213,7 +216,7 @@ def _fused_add_rms_norm_cutedsl_backward(
residual-add makes ``dX = dR``; an optional ``dS_out`` (gradient flowing
to the residual output) is added host-side after the kernel.
"""
# Reuse the rms_norm backward: it computes dX_rms from (dY, S, W_eff, RSTD).
# Reuse the rms_norm backward: it computes dX_rms from (dY, S, W_kernel, RSTD).
# ``dS_out`` is folded into the kernel epilogue (dx += dS in fp32 before the
# single store round — strictly closer to exact than a post-hoc two-tensor
# bf16 add, and it removes the host pass's full (M,N) read+write trip: the
@@ -234,11 +237,12 @@ def _fused_add_rms_norm_cutedsl_backward(
dX, dW = _rms_norm_cutedsl_backward(
dY.view(-1, N).contiguous() if dY.dim() > 2 else dY,
S.view(-1, N) if S.dim() > 2 else S,
W_eff,
W_kernel,
RSTD,
in_place,
casting_mode,
ds=(dS_out.view(-1, N) if (dS_out is not None and dS_out.dim() > 2) else dS_out),
weight_offset=weight_offset,
)
if dY.dim() > 2:
dX = dX.view(shape)
@@ -251,7 +255,7 @@ def _fused_add_rms_norm_cutedsl_backward(
class _LigerFusedAddRMSNormCuTeDSLFunction(torch.autograd.Function):
"""Autograd wrapper for fused_add_rms_norm via CuTe DSL.
Saves ``S`` (the summed residual), ``W_eff``, and ``RSTD`` for the backward.
Saves ``S`` (the summed residual), ``W_kernel``, and ``RSTD`` for the backward.
"""
@staticmethod
@@ -272,12 +276,13 @@ class _LigerFusedAddRMSNormCuTeDSLFunction(torch.autograd.Function):
else:
casting_mode_int = casting_mode
Y, S, W_eff, RSTD = _fused_add_rms_norm_cutedsl_forward(X, R, W, eps, offset, casting_mode_int)
Y, S, W_kernel, RSTD = _fused_add_rms_norm_cutedsl_forward(X, R, W, eps, offset, casting_mode_int)
ctx.in_place = in_place
ctx.weight_dtype = W.dtype
ctx.casting_mode = casting_mode_int
ctx.save_for_backward(S, W_eff, RSTD)
ctx.weight_offset = offset if casting_mode_int == _CASTING_MODE_GEMMA else 0.0
ctx.save_for_backward(S, W_kernel, RSTD)
return Y, S
@staticmethod
@@ -286,15 +291,16 @@ class _LigerFusedAddRMSNormCuTeDSLFunction(torch.autograd.Function):
if dS_out is not None:
dS_out = _to_local_if_dtensor(dS_out).contiguous()
S, W_eff, RSTD = ctx.saved_tensors
S, W_kernel, RSTD = ctx.saved_tensors
dX, dR, dW = _fused_add_rms_norm_cutedsl_backward(
dY,
dS_out,
S,
W_eff,
W_kernel,
RSTD,
ctx.in_place,
ctx.casting_mode,
weight_offset=ctx.weight_offset,
)
dW = dW.to(ctx.weight_dtype)
src/liger_kernel/ops/backends/_cutedsl/rms_norm.py
+49 -38
View File
@@ -22,7 +22,6 @@ Math (matches the Triton kernel in :mod:`liger_kernel.ops.rms_norm`)::
rstd = 1 / sqrt(mean(x^2) + eps)
y = (x * rstd) * (offset + w)
w_eff = w + offset (host-side, cheap)
wdy = (offset + w) * dy (fp32)
c = mean(wdy * (x * rstd))
dx = rstd * (wdy - (x * rstd) * c)
@@ -43,12 +42,10 @@ does not directly expose; we adapt them as follows:
for API parity, but the math runs through the same code path as
``"llama"`` — slightly more precise than the Triton ``none`` mode.
``offset`` is folded into the weight host-side (``weight + offset``)
before the kernel call — the inline forward has no offset argument, but
this single fused pre-add is far cheaper than the launch itself.
The backward uses the **effective** weight ``w + offset`` (we save it in
ctx). This matches Triton's ``W_row + offset`` line inside the bwd kernel.
For Gemma, the forward and backward add ``offset`` to the loaded FP32
weight in-register. Other casting modes keep their host-side pre-add.
The saved weight and compile-keyed offset reproduce the same effective
weight in the backward without a separate Gemma weight buffer.
Capability
----------
@@ -154,10 +151,12 @@ class _LigerRMSNormCuTeDSLBackward(ReductionBase):
N: int,
has_weight: bool = True,
casting_mode: int = _CASTING_MODE_LLAMA,
weight_offset: float = 0.0,
):
super().__init__(dtype, N, stage=2, reduction_dtype=Float32)
self.has_weight = has_weight
self.casting_mode = casting_mode
self.weight_offset = weight_offset
# Beyond 16K we reload wdy from smem instead of holding the
# fragment live — matches Quack's ``RMSNormBackward.reload_wdy``.
self.reload_wdy = None if N <= 16 * 1024 else "smem"
@@ -188,7 +187,7 @@ class _LigerRMSNormCuTeDSLBackward(ReductionBase):
def __call__(
self,
mX: cute.Tensor, # Input [M, N]
mW: Optional[cute.Tensor], # Effective weight (w + offset) [N,] or None
mW: Optional[cute.Tensor], # Kernel weight [N,] or None; Gemma offset is added in-register
mdO: cute.Tensor, # dY [M, N]
mRstd: cute.Tensor, # RSTD [M,] (fp32)
mdX: cute.Tensor, # dX [M, N]
@@ -382,8 +381,10 @@ class _LigerRMSNormCuTeDSLBackward(ReductionBase):
wdy = dout
if const_expr(mW is not None):
# mW already contains (offset + weight) — folded host-side.
wdy = wdy * tXrW.load().to(Float32)
w = tXrW.load().to(Float32)
if const_expr(self.weight_offset != 0.0):
w += Float32(self.weight_offset)
wdy = wdy * w
if const_expr(self.cluster_n > 1):
cute.arch.mbarrier_wait(mbar_empty_ptr + stage, producer_phase)
@@ -418,7 +419,10 @@ class _LigerRMSNormCuTeDSLBackward(ReductionBase):
dout = tXrdO.load().to(cute.Float32)
wdy = dout
if const_expr(mW is not None):
wdy = wdy * tXrW.load().to(Float32)
w = tXrW.load().to(Float32)
if const_expr(self.weight_offset != 0.0):
w += Float32(self.weight_offset)
wdy = wdy * w
# dx = rstd * (wdy - x_hat * c)
dx = (wdy - x_hat * mean_xhat_wdy) * rstd_val
@@ -518,9 +522,10 @@ def _get_bwd_kernel(
has_weight: bool,
casting_mode: int,
has_ds: bool = False,
weight_offset: float = 0.0,
):
"""Return a compiled backward kernel, building it on first miss."""
key = (x_dtype, weight_dtype, dx_dtype, N, has_weight, casting_mode, has_ds)
key = (x_dtype, weight_dtype, dx_dtype, N, has_weight, casting_mode, has_ds, weight_offset)
if key in _BWD_COMPILE_CACHE:
return _BWD_COMPILE_CACHE[key]
@@ -547,6 +552,7 @@ def _get_bwd_kernel(
N,
has_weight=has_weight,
casting_mode=casting_mode,
weight_offset=weight_offset,
)
compiled = cute.compile(
kernel,
@@ -577,8 +583,8 @@ def _rms_norm_cutedsl_forward(
) -> Tuple[Tensor, Tensor, Optional[Tensor], Tensor]:
"""Forward via the inlined ``_cute_lib.rmsnorm_fwd``.
Returns ``(Y, X_flat, W_eff, RSTD)``. ``W_eff`` is the effective weight
``weight + offset`` (or ``None`` if ``weight is None``), kept around for
Returns ``(Y, X_flat, W_kernel, RSTD)``. ``W_kernel`` is the weight
passed to the kernel (raw weight for Gemma, pre-added otherwise), kept around for
the inline backward. ``RSTD`` is fp32 (the kernel always produces it
in fp32 when ``store_rstd=True``).
"""
@@ -611,28 +617,29 @@ def _rms_norm_cutedsl_forward(
x_in = x
x_flat = x_in.view(-1, N).contiguous()
# Fold offset into the weight host-side — Quack has no offset arg, but
# one extra elementwise add on a (N,) vector is essentially free.
# Gemma adds offset after promoting the loaded weight to FP32 in the
# kernel. Other modes retain their existing host-side rounding.
weight_offset = offset if casting_mode == _CASTING_MODE_GEMMA else 0.0
if weight is not None:
if offset != 0.0:
w_eff = weight + offset
if offset != 0.0 and casting_mode != _CASTING_MODE_GEMMA:
w_kernel = weight + offset
else:
# Avoid copying when offset is the common 0.0 — Quack accepts
# the original weight directly.
w_eff = weight
# Only the mixed-dtype gemma route promotes w_eff host-side;
# Gemma and zero-offset calls use the original weight directly.
w_kernel = weight
# Only the mixed-dtype gemma route promotes w_kernel host-side;
# otherwise the kernel up-casts in-register (bit-exact).
if _promote_gemma_fp32 and w_eff.dtype != torch.float32:
w_eff = w_eff.to(torch.float32)
w_eff = w_eff.contiguous()
if _promote_gemma_fp32 and w_kernel.dtype != torch.float32:
w_kernel = w_kernel.to(torch.float32)
w_kernel = w_kernel.contiguous()
else:
w_eff = None
w_kernel = None
out, _residual_out, rstd = _cutedsl_rmsnorm_fwd(
x_flat,
weight=w_eff,
weight=w_kernel,
eps=eps,
store_rstd=True,
weight_offset=weight_offset,
)
# Back-cast if the promote route above widened the output (it is the
@@ -640,17 +647,18 @@ def _rms_norm_cutedsl_forward(
if out.dtype != x.dtype:
out = out.to(x.dtype)
return out.view(shape), x_flat, w_eff, rstd
return out.view(shape), x_flat, w_kernel, rstd
def _rms_norm_cutedsl_backward(
dy: Tensor,
x_flat: Tensor,
w_eff: Optional[Tensor],
w_kernel: Optional[Tensor],
rstd: Tensor,
in_place: bool,
casting_mode: int,
ds: Optional[Tensor] = None,
weight_offset: float = 0.0,
) -> Tuple[Tensor, Optional[Tensor]]:
"""Backward via the inline CuTe DSL kernel.
@@ -699,7 +707,7 @@ def _rms_norm_cutedsl_backward(
# Saturate the grid: never launch more SMs than rows.
sm_count = min(sm_count, max(M, 1))
has_weight = w_eff is not None
has_weight = w_kernel is not None
dw_partial = torch.empty((sm_count, N), dtype=torch.float32, device=x_flat.device) if has_weight else None
ds_flat = None
@@ -713,19 +721,20 @@ def _rms_norm_cutedsl_backward(
compiled = _get_bwd_kernel(
x_dtype=x_flat.dtype,
weight_dtype=w_eff.dtype if w_eff is not None else None,
weight_dtype=w_kernel.dtype if w_kernel is not None else None,
dx_dtype=dx.dtype,
N=N,
has_weight=has_weight,
casting_mode=casting_mode,
has_ds=ds_flat is not None,
weight_offset=weight_offset,
)
# The environment-stream argument configured in _get_bwd_kernel is supplied
# by TVM FFI from the current PyTorch stream, so callers omit it here.
compiled(
x_flat,
w_eff,
w_kernel,
dy_flat,
rstd,
dx,
@@ -750,7 +759,7 @@ class _LigerRMSNormCuTeDSLFunction(torch.autograd.Function):
"""Autograd wrapper. Saves enough state for the inline-kernel backward.
We save ``x_flat`` (the 2D view passed to the kernel; fp32 when the
mixed-dtype gemma route promoted host-side) and ``w_eff`` (``weight + offset``),
mixed-dtype gemma route promoted host-side) and the kernel weight,
so the backward doesn't redo any of the host-side pre-processing.
"""
@@ -775,15 +784,16 @@ class _LigerRMSNormCuTeDSLFunction(torch.autograd.Function):
else:
casting_mode_int = casting_mode
Y, X_flat, W_eff, RSTD = _rms_norm_cutedsl_forward(X, W, eps, offset, casting_mode_int)
Y, X_flat, W_kernel, RSTD = _rms_norm_cutedsl_forward(X, W, eps, offset, casting_mode_int)
ctx.in_place = in_place
ctx.elementwise_affine = W is not None
ctx.casting_mode = casting_mode_int
ctx.weight_offset = offset if casting_mode_int == _CASTING_MODE_GEMMA else 0.0
# Save the user's original weight dtype so we can cast dW back to it.
ctx.weight_dtype = W.dtype if W is not None else None
if W is not None:
ctx.save_for_backward(X_flat, W_eff, RSTD)
ctx.save_for_backward(X_flat, W_kernel, RSTD)
else:
ctx.save_for_backward(X_flat, RSTD)
return Y
@@ -793,18 +803,19 @@ class _LigerRMSNormCuTeDSLFunction(torch.autograd.Function):
dY = _to_local_if_dtensor(dY).contiguous()
if ctx.elementwise_affine:
X_flat, W_eff, RSTD = ctx.saved_tensors
X_flat, W_kernel, RSTD = ctx.saved_tensors
else:
X_flat, RSTD = ctx.saved_tensors
W_eff = None
W_kernel = None
dX, dW = _rms_norm_cutedsl_backward(
dY,
X_flat,
W_eff,
W_kernel,
RSTD,
ctx.in_place,
ctx.casting_mode,
weight_offset=ctx.weight_offset,
)
if ctx.elementwise_affine and dW is not None:
dW = dW.to(ctx.weight_dtype)