Fused RoPE + RMSNorm#
Track 07 - GPU · Notebook 05 · Runtime: ≈10 min on any GPU
Prerequisites:
07_gpu/01(GPU architecture tour),07_gpu/02(Triton softmax).Papers:
Su et al. 2021, RoFormer: Enhanced Transformer with Rotary Position Embedding (2104.09864).
Zhang & Sennrich 2019, Root Mean Square Layer Normalization (1910.07467).
What#
Two operations every modern transformer runs before attention: RMSNorm (a variance-only normalisation) and RoPE (a position-dependent rotation of the Q/K channels). In eager PyTorch they’re 5-8 ops each; fused into one Triton kernel per op they’re single-pass memory-bound kernels at ≥90% of peak bandwidth.
We implement:
A Triton RMSNorm kernel matching
torch.nn.functional.rms_normsemantics.A RoPE kernel that rotates channel pairs by position-dependent angles.
Numerical correctness checks vs reference PyTorch implementations.
Both are memory-bound (one read, one write) so fusion’s main win is halving the number of HBM trips.
from llm_systems_cookbook.nb import bootstrap
import math
import time
import torch
IS_CUDA = torch.cuda.is_available()
s = bootstrap("07_gpu_05_fused_rope_rmsnorm")
RMSNorm#
where g is a learned per-channel gain. One pass through the tensor:
compute the per-row RMS, divide, multiply by gain.
def rmsnorm_torch(x: torch.Tensor, g: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
return (x / rms) * g
# Smoke-test against the formula: a constant row with magnitude c has
# RMS c, so output should be g.
c = 3.0
x_test = torch.full((2, 8), c)
g_test = torch.linspace(0.5, 1.5, 8)
out = rmsnorm_torch(x_test, g_test)
print(f"rmsnorm(const*{c}, g) = {out[0].tolist()}")
s.check(
"rmsnorm_of_constant_equals_gain",
lambda: torch.allclose(out[0], g_test, atol=1e-4),
msg=f"out = {out[0].tolist()}",
)
triton = None
tl = None
rmsnorm_triton = None
if IS_CUDA:
try:
import triton
import triton.language as tl
@triton.jit
def _rmsnorm_kernel(
out_ptr, x_ptr, g_ptr,
x_stride, out_stride,
n_cols,
eps,
BLOCK_SIZE: tl.constexpr,
):
row = tl.program_id(0)
col = tl.arange(0, BLOCK_SIZE)
mask = col < n_cols
x = tl.load(x_ptr + row * x_stride + col, mask=mask, other=0.0).to(tl.float32)
g = tl.load(g_ptr + col, mask=mask, other=0.0).to(tl.float32)
rms = tl.sqrt(tl.sum(x * x, axis=0) / n_cols + eps)
y = (x / rms) * g
tl.store(out_ptr + row * out_stride + col, y.to(tl.float16), mask=mask)
def rmsnorm_triton(x: torch.Tensor, g: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
assert x.is_cuda and g.is_cuda
M, N = x.shape
BLOCK = triton.next_power_of_2(N)
out = torch.empty_like(x)
_rmsnorm_kernel[(M,)](
out, x, g,
x.stride(0), out.stride(0),
N,
eps,
BLOCK_SIZE=BLOCK,
num_warps=4 if BLOCK >= 1024 else 2,
)
return out
print("Triton RMSNorm compiled")
except Exception as e: # noqa: BLE001
print(f"Triton unavailable: {type(e).__name__}: {e}")
rmsnorm_triton = None
if IS_CUDA and rmsnorm_triton is not None:
torch.manual_seed(0)
x = torch.randn((256, 512), device="cuda", dtype=torch.float16)
g = torch.randn(512, device="cuda", dtype=torch.float16)
y_ref = rmsnorm_torch(x.float(), g.float()).half()
y_ours = rmsnorm_triton(x, g)
err = (y_ref - y_ours).abs().max().item()
print(f"max abs err vs torch reference = {err:.3e}")
s.check(
"triton_rmsnorm_matches_torch_1e_minus_2",
lambda: err < 1e-2,
msg=f"max abs err = {err:.3e}",
)
else:
s.skip("triton_rmsnorm_matches_torch_1e_minus_2", "no CUDA / Triton")
RoPE#
RoPE rotates channel pairs (x[2i], x[2i+1]) by position-dependent
angles. For head dim D and position p:
base = 10000 in original RoPE, 500000 in Llama-3. A reference
implementation precomputes cos/sin tables of shape (T, D/2) and
broadcasts.
def rope_torch(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
'''Apply RoPE to the last dim of x, assuming x shape (..., T, D) and
cos/sin shape (T, D/2).'''
D = x.shape[-1]
assert D % 2 == 0
x1 = x[..., 0::2]
x2 = x[..., 1::2]
y1 = x1 * cos - x2 * sin
y2 = x1 * sin + x2 * cos
return torch.stack((y1, y2), dim=-1).flatten(-2)
def build_rope_tables(seq_len: int, head_dim: int, base: float = 10000.0,
device: torch.device | str = "cpu",
dtype: torch.dtype = torch.float32) -> tuple[torch.Tensor, torch.Tensor]:
positions = torch.arange(seq_len, device=device, dtype=dtype)
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim))
angles = positions[:, None] * inv_freq[None, :]
return torch.cos(angles), torch.sin(angles)
# Invariant: RoPE preserves the norm of every pair of channels, so the
# overall L2 norm per position is preserved up to fp rounding.
T = 32
D = 64
cos_t, sin_t = build_rope_tables(T, D)
x = torch.randn(2, T, D)
y = rope_torch(x, cos_t, sin_t)
norm_err = (x.pow(2).sum(dim=-1) - y.pow(2).sum(dim=-1)).abs().max().item()
print(f"rope preserves L2 norm, max err = {norm_err:.2e}")
s.check(
"rope_preserves_norm_per_position",
lambda: norm_err < 1e-4,
msg=f"max err = {norm_err:.2e}",
)
rope_triton = None
if IS_CUDA:
try:
@triton.jit
def _rope_kernel(
out_ptr, x_ptr, cos_ptr, sin_ptr,
stride_bt, stride_bd,
T, D,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
t = pid // 1 # one program per (batch, t) pair collapsed
# Flatten (batch, t) dims outside the kernel; each pid is one row.
row = pid
col = tl.arange(0, BLOCK_SIZE)
mask_even = (col * 2) < D
mask_odd = (col * 2 + 1) < D
x_even = tl.load(x_ptr + row * stride_bt + col * 2, mask=mask_even, other=0.0).to(tl.float32)
x_odd = tl.load(x_ptr + row * stride_bt + col * 2 + 1, mask=mask_odd, other=0.0).to(tl.float32)
# We need the t index relative to the sequence length; the caller passes
# pos_in_seq via stride math - for simplicity the caller prepares a
# contiguous view of shape (M=B*T, D) and passes T.
pos = row % T
c = tl.load(cos_ptr + pos * (D // 2) + col, mask=col < (D // 2), other=0.0).to(tl.float32)
sv = tl.load(sin_ptr + pos * (D // 2) + col, mask=col < (D // 2), other=0.0).to(tl.float32)
y_even = x_even * c - x_odd * sv
y_odd = x_even * sv + x_odd * c
tl.store(out_ptr + row * stride_bt + col * 2, y_even.to(tl.float16), mask=mask_even)
tl.store(out_ptr + row * stride_bt + col * 2 + 1, y_odd.to(tl.float16), mask=mask_odd)
def rope_triton(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
'''Apply RoPE. x shape (B, T, D), cos/sin shape (T, D/2).'''
B, T, D = x.shape
assert D % 2 == 0
x_flat = x.contiguous().view(B * T, D)
out = torch.empty_like(x_flat)
BLOCK = triton.next_power_of_2(D // 2)
_rope_kernel[(B * T,)](
out, x_flat, cos.contiguous(), sin.contiguous(),
x_flat.stride(0), x_flat.stride(1),
T, D,
BLOCK_SIZE=BLOCK,
)
return out.view(B, T, D)
print("Triton RoPE compiled")
except Exception as e: # noqa: BLE001
print(f"Triton RoPE compile failed: {type(e).__name__}: {e}")
rope_triton = None
if IS_CUDA and rope_triton is not None:
B, T, D = 2, 64, 128
torch.manual_seed(0)
x = torch.randn((B, T, D), device="cuda", dtype=torch.float16)
cos_d, sin_d = build_rope_tables(T, D, device="cuda", dtype=torch.float16)
y_ref = rope_torch(x.float(), cos_d.float(), sin_d.float()).half()
y_ours = rope_triton(x, cos_d, sin_d)
err = (y_ref - y_ours).abs().max().item()
print(f"max abs err vs torch reference = {err:.3e}")
s.check(
"triton_rope_matches_torch_1e_minus_2",
lambda: err < 1e-2,
msg=f"err = {err:.3e}",
)
else:
s.skip("triton_rope_matches_torch_1e_minus_2", "no CUDA / Triton")
Unfused vs fused latency#
Both ops are memory-bound, so the win is proportional to the number of HBM trips saved. Measure eager PyTorch against the fused Triton kernels across a few shapes and plot latency side by side. The gap widens on wider rows because the per-launch overhead gets amortised over more useful work.
import matplotlib.pyplot as plt
if IS_CUDA and rmsnorm_triton is not None and rope_triton is not None:
SHAPES = [(512, 512), (1024, 1024), (2048, 2048), (4096, 4096)]
def _bench(fn, *args, n_iter=30):
for _ in range(3):
fn(*args)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(n_iter):
fn(*args)
torch.cuda.synchronize()
return (time.perf_counter() - t0) * 1000.0 / n_iter # ms/iter
rms_eager, rms_fused, rope_eager, rope_fused = [], [], [], []
for M, N in SHAPES:
x = torch.randn((M, N), device="cuda", dtype=torch.float16)
g = torch.randn(N, device="cuda", dtype=torch.float16)
rms_eager.append(_bench(rmsnorm_torch, x, g))
rms_fused.append(_bench(rmsnorm_triton, x, g))
xr = torch.randn((1, M, N), device="cuda", dtype=torch.float16)
cos_d, sin_d = build_rope_tables(M, N, device="cuda", dtype=torch.float16)
rope_eager.append(_bench(rope_torch, xr, cos_d, sin_d))
rope_fused.append(_bench(rope_triton, xr, cos_d, sin_d))
xs = list(range(len(SHAPES)))
labels = [f"{m}x{n}" for m, n in SHAPES]
w = 0.38
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))
for ax, eager, fused, title in [(ax1, rms_eager, rms_fused, "RMSNorm"),
(ax2, rope_eager, rope_fused, "RoPE")]:
ax.bar([i - w/2 for i in xs], eager, w, label="eager", color="tab:gray")
ax.bar([i + w/2 for i in xs], fused, w, label="triton", color="tab:blue")
ax.set_xticks(xs); ax.set_xticklabels(labels, rotation=20)
ax.set_ylabel("ms / iter"); ax.set_title(title); ax.legend()
ax.grid(True, axis="y", alpha=0.3)
fig.tight_layout(); plt.show()
else:
print("skipped - no CUDA / Triton.")
Exercises#
Fuse RMSNorm and RoPE into a single kernel for Q and K together (the layout before attention). The fused kernel does one read/write of Q and K instead of four.
Add the Llama-3 base (500000) and compare the angle schedule.
Implement the backward pass. RMSNorm’s gradient requires the RMS value again; RoPE’s inverse is just the transpose of the same rotation.
References#
HuggingFace’s
transformers/models/llama/modeling_llama.pyfor the reference RoPE implementation.Triton tutorials,
05-layer-norm.py- essentially the same kernel shape as our RMSNorm.
s.summary()
s.save()