Skip to content

[rmsnorm] Support FP32 weights with FP16/BF16 activations #884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
coderfeli merged 1 commit into from
Jul 24, 2026
Merged

[rmsnorm] Support FP32 weights with FP16/BF16 activations #884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions docs/prebuilt_kernels_guide.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
| Kernel | Builder Function | API Style | Dtypes | Key Feature |
|---|---|---|---|---|
| **LayerNorm** | `build_layernorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Two-pass vectorized normalization |
| **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | LDS-cached 3-pass pipeline |
| **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16; optional fp32 weight | LDS-cached 3-pass pipeline |
| **Softmax** | `build_softmax_module(M, N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Online softmax, adaptive block size |
| **GEMM** | `compile_preshuffle_gemm(...)` | `@flyc.kernel` | fp8, int8, fp16, bf16 | Preshuffle B, ping-pong LDS, MFMA 16x16 |
| **FlashAttention** | `build_flash_attn_func_module(...)` | `@flyc.kernel` | bf16, f16 (any arch); fp8 e4m3fn (gfx950, D=128, dense) | Dual-wave SWP fwd, GQA/MQA, causal, descale ABI |
Expand Down Expand Up @@ -68,14 +68,20 @@ from kernels.norm.rmsnorm_kernel import build_rmsnorm_module
executor = build_rmsnorm_module(N=8192, dtype_str="bf16", store_rstd=False)
```

`build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS)` optionally
writes the per-row reciprocal std (`rstd`) for use by the backward pass.
`build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS,
BLOCK_THREADS=BLOCK_THREADS, weight_dtype_str=None)` optionally writes the
per-row reciprocal std (`rstd`) for use by the backward pass.
`weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally
support FP32 weights.

**Backward:** `build_rmsnorm_bwd_module(N, dtype_str)` builds the fused RMSNorm
backward kernel (grid `(M,)`, one block per row). Kernel signature
**Backward:** `build_rmsnorm_bwd_module(N, dtype_str,
weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`,
one block per row). Kernel signature
`rmsnorm_bwd_kernel(Input, Gamma, DY, Rstd, DX, DWeight)`: it reads the forward
`Rstd`, writes `DX` (input grad) and atomic-adds into `DWeight` (fp32 weight
grad). `eps` is baked into `Rstd` by the forward, so it is not needed here.
The public plain and fused-add training wrappers return `dweight` in the
original weight dtype.

**Configuration Constants:** Same as LayerNorm (BLOCK_THREADS=256, VEC_WIDTH=8, etc.)

Expand Down
41 changes: 37 additions & 4 deletions kernels/norm/rmsnorm_autotune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
"""Two-track RMSNorm autotuning through the normal direct JIT path."""

from flydsl.autotune import Config, autotune
from kernels.norm.rmsnorm_common import BLOCK_THREADS
from kernels.norm.rmsnorm_common import (
BLOCK_THREADS,
resolve_rmsnorm_weight_dtype,
)
from kernels.norm.rmsnorm_kernel import SMALL_N_THRESHOLD, rmsnorm_direct

_SEARCH_CONFIGS = (
Expand All @@ -22,22 +25,51 @@ def _default_config(*_args, **_kwargs):
return Config(BLOCK_THREADS=BLOCK_THREADS)


def _search_configs(input_t, gamma, output, m_in, N, dtype_str="bf16", stream=None):
def _search_configs(
input_t,
gamma,
output,
m_in,
N,
dtype_str="bf16",
weight_dtype_str="bf16",
stream=None,
):
if N <= SMALL_N_THRESHOLD:
return [_default_config()]
return list(_SEARCH_CONFIGS)


_rmsnorm_tuner = autotune(
configs=_search_configs,
key=["m_in", "N", "dtype_str"],
key=["m_in", "N", "dtype_str", "weight_dtype_str"],
default=_default_config,
)(rmsnorm_direct)


def rmsnorm_autotuned(input_t, gamma, output, m_in, dtype_str="bf16", stream=None):
def rmsnorm_autotuned(
input_t,
gamma,
output,
m_in,
dtype_str=None,
stream=None,
weight_dtype_str=None,
):
import torch

from kernels.norm.rmsnorm_common import torch_dtype_to_str

input_dtype_str = torch_dtype_to_str(input_t.dtype)
if dtype_str is not None and dtype_str != input_dtype_str:
raise ValueError(f"dtype_str={dtype_str!r} does not match input dtype {input_dtype_str!r}")
dtype_str = input_dtype_str

gamma_dtype_str = torch_dtype_to_str(gamma.dtype)
if weight_dtype_str is not None and weight_dtype_str != gamma_dtype_str:
raise ValueError(f"weight_dtype_str={weight_dtype_str!r} does not match gamma dtype {gamma_dtype_str!r}")
weight_dtype_str = resolve_rmsnorm_weight_dtype(dtype_str, gamma_dtype_str)

with torch.cuda.device(input_t.device):
launch_stream = torch.cuda.current_stream() if stream is None else stream
return _rmsnorm_tuner(
Expand All @@ -47,5 +79,6 @@ def rmsnorm_autotuned(input_t, gamma, output, m_in, dtype_str="bf16", stream=Non
m_in,
N=int(input_t.shape[-1]),
dtype_str=dtype_str,
weight_dtype_str=weight_dtype_str,
stream=launch_stream,
)
Loading
Loading