diff --git a/README.md b/README.md index 02a50264f..59732b852 100644 --- a/README.md +++ b/README.md @@ -297,7 +297,7 @@ loss.backward() | Qwen2, Qwen2.5, & QwQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen2-VL, & QVQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl` | RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen2.5-VL | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_5_vl` | RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | -| Qwen3 | `liger_kernel.transformers.apply_liger_kernel_to_qwen3` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | +| Qwen3 | `liger_kernel.transformers.apply_liger_kernel_to_qwen3` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy, optional fused QK-Norm + RoPE | | Qwen3 MoE | `liger_kernel.transformers.apply_liger_kernel_to_qwen3_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen3.5 | `liger_kernel.transformers.apply_liger_kernel_to_qwen3_5` | RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | | Qwen3.5 MoE (Text) & (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_qwen3_5_moe` | RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy | diff --git a/benchmark/scripts/profile_qk_norm_rope.py b/benchmark/scripts/profile_qk_norm_rope.py new file mode 100644 index 000000000..469a187a2 --- /dev/null +++ b/benchmark/scripts/profile_qk_norm_rope.py @@ -0,0 +1,247 @@ +"""Profile the fused QK-Norm + RoPE kernel against the unfused PyTorch baseline. + +The baseline is the *exact* sequence the ``qwen3_attention_forward`` monkeypatch +replaces (see ``modeling_qwen3.Qwen3Attention.forward``): + + q = q_norm(q.view(B, T, n_qh, hd)).transpose(1, 2) # per-head RMSNorm + transpose + k = k_norm(k.view(B, T, n_kh, hd)).transpose(1, 2) + q, k = apply_rotary_pos_emb(q, k, cos, sin) # RoPE + +vs. the single fused Triton kernel: + + q, k = liger_qk_norm_rope(q, k, wq, wk, cos, sin, eps) + +We report forward / backward / full latency (via ``triton.testing.do_bench``), +peak activation memory, and a ``torch.profiler`` CUDA-kernel breakdown so the +"before vs after" kernel behaviour is visible. + +Shapes are the real Qwen3 dense attention configs (head_dim=128, GQA, eps=1e-6). + +Usage:: + + python benchmark/scripts/profile_qk_norm_rope.py # all configs, bf16 + python benchmark/scripts/profile_qk_norm_rope.py --dtype float32 + python benchmark/scripts/profile_qk_norm_rope.py --seq-len 8192 --bsz 1 + python benchmark/scripts/profile_qk_norm_rope.py --trace # dump chrome traces +""" + +import argparse +import gc + +import torch +import triton + +from torch.profiler import ProfilerActivity +from torch.profiler import profile + +from liger_kernel.transformers.qk_norm_rope import liger_qk_norm_rope +from liger_kernel.utils import infer_device + +device = infer_device() + +# (name, n_q_head, n_kv_head, head_dim) -- real Qwen3 dense configs +QWEN3_CONFIGS = [ + ("qwen3_0.6b", 16, 8, 128), + ("qwen3_4b/8b", 32, 8, 128), + ("qwen3_14b", 40, 8, 128), + ("qwen3_32b", 64, 8, 128), +] + +EPS = 1e-6 # Qwen3 rms_norm_eps + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def rms_norm_ref(x, weight, eps): + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + return weight * x.to(input_dtype) + + +def baseline_forward(q, k, wq, wk, cos, sin, eps): + """Unfused reference == what the monkeypatch replaces (q_norm/k_norm + RoPE).""" + qn = rms_norm_ref(q, wq, eps).transpose(1, 2) + kn = rms_norm_ref(k, wk, eps).transpose(1, 2) + return apply_rotary_pos_emb(qn, kn, cos, sin) + + +def fused_forward(q, k, wq, wk, cos, sin, eps): + return liger_qk_norm_rope(q, k, wq, wk, cos, sin, eps) + + +def make_cos_sin(seq_len, head_dim, dtype): + pos = torch.arange(seq_len, device=device, dtype=torch.float32) + inv_freq = 1.0 / (1000000 ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) + freqs = torch.outer(pos, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().unsqueeze(0).to(dtype), emb.sin().unsqueeze(0).to(dtype) + + +def make_inputs(bsz, seq_len, n_qh, n_kh, hd, dtype, requires_grad): + q = torch.randn(bsz, seq_len, n_qh, hd, device=device, dtype=dtype, requires_grad=requires_grad) + k = torch.randn(bsz, seq_len, n_kh, hd, device=device, dtype=dtype, requires_grad=requires_grad) + wq = (1.0 + 0.1 * torch.randn(hd, device=device, dtype=dtype)).requires_grad_(requires_grad) + wk = (1.0 + 0.1 * torch.randn(hd, device=device, dtype=dtype)).requires_grad_(requires_grad) + cos, sin = make_cos_sin(seq_len, hd, dtype) + return q, k, wq, wk, cos, sin + + +def _bench(fn): + """Median / p20 / p80 ms via triton.testing.do_bench.""" + ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8], grad_to_none=None) + return ms + + +def bench_mode(fwd, inputs, mode): + q, k, wq, wk, cos, sin = inputs + + if mode == "forward": + return _bench(lambda: fwd(q, k, wq, wk, cos, sin, EPS)) + + # produce a fixed upstream grad + oq, ok = fwd(q, k, wq, wk, cos, sin, EPS) + gq = torch.randn_like(oq) + gk = torch.randn_like(ok) + + if mode == "backward": + + def run(): + for t in (q, k, wq, wk): + if t.grad is not None: + t.grad = None + torch.autograd.backward((oq, ok), (gq, gk), retain_graph=True) + + return _bench(run) + + if mode == "full": + + def run(): + for t in (q, k, wq, wk): + if t.grad is not None: + t.grad = None + o1, o2 = fwd(q, k, wq, wk, cos, sin, EPS) + torch.autograd.backward((o1, o2), (gq, gk)) + + return _bench(run) + + raise ValueError(mode) + + +def peak_memory_mb(fwd, inputs): + q, k, wq, wk, cos, sin = inputs + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + o1, o2 = fwd(q, k, wq, wk, cos, sin, EPS) + torch.autograd.backward((o1, o2), (torch.randn_like(o1), torch.randn_like(o2))) + return torch.cuda.max_memory_allocated() / (1024**2) + + +def run_latency(bsz, seq_len, dtype): + print(f"\n{'=' * 96}") + print(f"Latency (ms, median) | bsz={bsz} seq_len={seq_len} dtype={dtype} device={torch.cuda.get_device_name()}") + print(f"{'=' * 96}") + header = ( + f"{'config':<12} {'n_qh/n_kh/hd':<14} " + f"{'fwd base':>9} {'fwd fus':>9} {'fwd x':>6} " + f"{'full base':>10} {'full fus':>10} {'full x':>7} " + f"{'mem base':>9} {'mem fus':>9} {'mem save':>9}" + ) + print(header) + print("-" * len(header)) + + for name, n_qh, n_kh, hd in QWEN3_CONFIGS: + inputs = make_inputs(bsz, seq_len, n_qh, n_kh, hd, dtype, requires_grad=True) + + fwd_b = bench_mode(baseline_forward, inputs, "forward") + fwd_f = bench_mode(fused_forward, inputs, "forward") + full_b = bench_mode(baseline_forward, inputs, "full") + full_f = bench_mode(fused_forward, inputs, "full") + mem_b = peak_memory_mb(baseline_forward, inputs) + mem_f = peak_memory_mb(fused_forward, inputs) + + print( + f"{name:<12} {f'{n_qh}/{n_kh}/{hd}':<14} " + f"{fwd_b:>9.4f} {fwd_f:>9.4f} {fwd_b / fwd_f:>5.2f}x " + f"{full_b:>10.4f} {full_f:>10.4f} {full_b / full_f:>6.2f}x " + f"{mem_b:>8.1f}M {mem_f:>8.1f}M {(1 - mem_f / mem_b) * 100:>7.1f}%" + ) + + del inputs + gc.collect() + torch.cuda.empty_cache() + + +def run_profiler(bsz, seq_len, dtype, trace): + """torch.profiler kernel breakdown for one representative config (8B).""" + name, n_qh, n_kh, hd = QWEN3_CONFIGS[1] + print(f"\n{'=' * 96}") + print(f"torch.profiler CUDA kernels | {name} bsz={bsz} seq_len={seq_len} dtype={dtype}") + print(f"{'=' * 96}") + + for tag, fwd in (("BEFORE (unfused baseline)", baseline_forward), ("AFTER (fused kernel)", fused_forward)): + inputs = make_inputs(bsz, seq_len, n_qh, n_kh, hd, dtype, requires_grad=True) + q, k, wq, wk, cos, sin = inputs + + # warmup + for _ in range(5): + o1, o2 = fwd(q, k, wq, wk, cos, sin, EPS) + torch.autograd.backward((o1, o2), (torch.randn_like(o1), torch.randn_like(o2))) + torch.cuda.synchronize() + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(20): + o1, o2 = fwd(q, k, wq, wk, cos, sin, EPS) + torch.autograd.backward((o1, o2), (torch.randn_like(o1), torch.randn_like(o2))) + torch.cuda.synchronize() + + print(f"\n--- {tag} ---") + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=12)) + + if trace: + fname = f"trace_qk_norm_rope_{'baseline' if fwd is baseline_forward else 'fused'}.json" + prof.export_chrome_trace(fname) + print(f"chrome trace -> {fname}") + + del inputs + gc.collect() + torch.cuda.empty_cache() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--bsz", type=int, default=4) + parser.add_argument("--seq-len", type=int, default=2048) + parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") + parser.add_argument("--profiler", action="store_true", help="run torch.profiler kernel breakdown") + parser.add_argument("--trace", action="store_true", help="export chrome traces (implies --profiler)") + parser.add_argument("--no-latency", action="store_true") + args = parser.parse_args() + + dtype = {"bfloat16": torch.bfloat16, "float32": torch.float32}[args.dtype] + + if not torch.cuda.is_available(): + raise SystemExit("CUDA required for profiling.") + + if not args.no_latency: + run_latency(args.bsz, args.seq_len, dtype) + + if args.profiler or args.trace: + run_profiler(args.bsz, args.seq_len, dtype, args.trace) + + +if __name__ == "__main__": + main() diff --git a/src/liger_kernel/ops/__init__.py b/src/liger_kernel/ops/__init__.py index f6676c07e..4d5b40d9d 100644 --- a/src/liger_kernel/ops/__init__.py +++ b/src/liger_kernel/ops/__init__.py @@ -76,6 +76,9 @@ from liger_kernel.ops.poly_norm import LigerPolyNormFunction # noqa: F401 from liger_kernel.ops.poly_norm import poly_norm_backward # noqa: F401 from liger_kernel.ops.poly_norm import poly_norm_forward # noqa: F401 +from liger_kernel.ops.qk_norm_rope import LigerQkNormRopeFunction # noqa: F401 +from liger_kernel.ops.qk_norm_rope import qk_norm_rope_backward # noqa: F401 +from liger_kernel.ops.qk_norm_rope import qk_norm_rope_forward # noqa: F401 from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction # noqa: F401 from liger_kernel.ops.relu_squared import LigerReLUSquaredFunction # noqa: F401 from liger_kernel.ops.relu_squared import relu_squared_backward # noqa: F401 diff --git a/src/liger_kernel/ops/qk_norm_rope.py b/src/liger_kernel/ops/qk_norm_rope.py new file mode 100644 index 000000000..5a2069e98 --- /dev/null +++ b/src/liger_kernel/ops/qk_norm_rope.py @@ -0,0 +1,444 @@ +"""Fused QK-Norm + RoPE Triton kernel. + +Several recent architectures (e.g. Qwen3) apply a per-head RMSNorm to the query +and key projections *before* rotary positional embedding: + + q = q_proj(x).view(B, T, n_qh, hd) + q = q_norm(q).transpose(1, 2) # RMSNorm over the head_dim + k = k_proj(x).view(B, T, n_kh, hd) + k = k_norm(k).transpose(1, 2) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + +Both RMSNorm and RoPE are memory-bound elementwise/reduction ops that only touch +Q and K. Running them as separate kernels round-trips the normalized Q/K through +HBM, and the ``.transpose(1, 2)`` in between forces the RoPE kernel to materialize +a ``.contiguous()`` copy. This module fuses the whole ``RMSNorm -> RoPE`` chain +into a single Triton kernel that reads Q/K once and writes them once, absorbing +the transpose as a pure stride operation. + +The RMSNorm follows the "llama" casting convention (reduction + rstd in fp32, +matching ``Qwen3RMSNorm``). The RoPE follows the HuggingFace Llama/Qwen half- +rotation layout (first half / second half), identical to ``LigerRopeFunction``. +""" + +import operator + +import torch +import triton +import triton.language as tl + +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 _qk_norm_rope_forward_kernel( + q_ptr, + q_row_stride, + k_ptr, + k_row_stride, + oq_ptr, + oq_row_stride, + ok_ptr, + ok_row_stride, + wq_ptr, + wk_ptr, + cos_ptr, + cos_row_stride, + sin_ptr, + sin_row_stride, + rstd_q_ptr, + rstd_k_ptr, + seq_len, + eps, + n_qh: tl.constexpr, + n_kh: tl.constexpr, + hd: tl.constexpr, + cos_bs: tl.constexpr, + pad_n_qh: tl.constexpr, + pad_n_kh: tl.constexpr, + pad_hd_half: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """One program instance == one token. + + q layout: (bsz, seq_len, n_qh, hd) contiguous, so the per-token row stride is + ``n_qh * hd``. We load the left / right halves of every head separately (as + in the RoPE kernel) and reduce over the full head_dim for the RMSNorm. + """ + pid = tl.program_id(0).to(tl.int64) + + q_ptr = q_ptr + pid * q_row_stride + k_ptr = k_ptr + pid * k_row_stride + oq_ptr = oq_ptr + pid * oq_row_stride + ok_ptr = ok_ptr + pid * ok_row_stride + + # ---- locate cos/sin for this token (only the left half is needed) ---- + batch_idx = pid // seq_len + seq_idx = pid % seq_len + cos_ptr = cos_ptr + tl.where( + cos_bs == 1, + seq_idx * cos_row_stride, + batch_idx * (seq_len * cos_row_stride) + seq_idx * cos_row_stride, + ) + sin_ptr = sin_ptr + tl.where( + cos_bs == 1, + seq_idx * sin_row_stride, + batch_idx * (seq_len * sin_row_stride) + seq_idx * sin_row_stride, + ) + half_cols = tl.arange(0, pad_hd_half) + half_mask = half_cols < (hd // 2) + cos_row = tl.load(cos_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + sin_row = tl.load(sin_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + + # ---- RMSNorm weight (loaded once, broadcast over heads) ---- + wq1 = tl.load(wq_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wq2 = tl.load(wq_ptr + (hd // 2) + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wk1 = tl.load(wk_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wk2 = tl.load(wk_ptr + (hd // 2) + half_cols, mask=half_mask, other=0.0).to(tl.float32) + + # ================= Q ================= + q_heads = tl.arange(0, pad_n_qh) + q1_off = q_heads[:, None] * hd + half_cols[None, :] + q2_off = q1_off + (hd // 2) + q_mask = (q_heads[:, None] < n_qh) & half_mask[None, :] + q1 = tl.load(q_ptr + q1_off, mask=q_mask, other=0.0).to(tl.float32) + q2 = tl.load(q_ptr + q2_off, mask=q_mask, other=0.0).to(tl.float32) + + # RMSNorm over the full head_dim (llama casting: reduce in fp32) + ms_q = (tl.sum(q1 * q1, axis=1) + tl.sum(q2 * q2, axis=1)) / hd + rstd_q = rsqrt(ms_q + eps) # (pad_n_qh,) + tl.store(rstd_q_ptr + pid * n_qh + q_heads, rstd_q, mask=q_heads < n_qh) + + qn1 = q1 * rstd_q[:, None] * wq1[None, :] + qn2 = q2 * rstd_q[:, None] * wq2[None, :] + + # RoPE: out1 = n1*cos - n2*sin ; out2 = n2*cos + n1*sin + oq1 = qn1 * cos_row[None, :] - qn2 * sin_row[None, :] + oq2 = qn2 * cos_row[None, :] + qn1 * sin_row[None, :] + tl.store(oq_ptr + q1_off, oq1.to(OUT_DTYPE), mask=q_mask) + tl.store(oq_ptr + q2_off, oq2.to(OUT_DTYPE), mask=q_mask) + + # ================= K ================= + k_heads = tl.arange(0, pad_n_kh) + k1_off = k_heads[:, None] * hd + half_cols[None, :] + k2_off = k1_off + (hd // 2) + k_mask = (k_heads[:, None] < n_kh) & half_mask[None, :] + k1 = tl.load(k_ptr + k1_off, mask=k_mask, other=0.0).to(tl.float32) + k2 = tl.load(k_ptr + k2_off, mask=k_mask, other=0.0).to(tl.float32) + + ms_k = (tl.sum(k1 * k1, axis=1) + tl.sum(k2 * k2, axis=1)) / hd + rstd_k = rsqrt(ms_k + eps) + tl.store(rstd_k_ptr + pid * n_kh + k_heads, rstd_k, mask=k_heads < n_kh) + + kn1 = k1 * rstd_k[:, None] * wk1[None, :] + kn2 = k2 * rstd_k[:, None] * wk2[None, :] + + ok1 = kn1 * cos_row[None, :] - kn2 * sin_row[None, :] + ok2 = kn2 * cos_row[None, :] + kn1 * sin_row[None, :] + tl.store(ok_ptr + k1_off, ok1.to(OUT_DTYPE), mask=k_mask) + tl.store(ok_ptr + k2_off, ok2.to(OUT_DTYPE), mask=k_mask) + + +@triton.jit +def _qk_norm_rope_backward_kernel( + doq_ptr, + doq_row_stride, + dok_ptr, + dok_row_stride, + q_ptr, + q_row_stride, + k_ptr, + k_row_stride, + dq_ptr, + dq_row_stride, + dk_ptr, + dk_row_stride, + wq_ptr, + wk_ptr, + cos_ptr, + cos_row_stride, + sin_ptr, + sin_row_stride, + rstd_q_ptr, + rstd_k_ptr, + dwq_ptr, + dwk_ptr, + dw_row_stride, + n_rows, + seq_len, + n_qh: tl.constexpr, + n_kh: tl.constexpr, + hd: tl.constexpr, + cos_bs: tl.constexpr, + pad_n_qh: tl.constexpr, + pad_n_kh: tl.constexpr, + pad_hd_half: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """Grid-strided over tokens; each program accumulates a partial dWq/dWk. + + ``dwq_ptr`` / ``dwk_ptr`` point to per-program partial buffers of shape + ``(num_programs, hd)`` that the host reduces to the final weight gradients. + """ + pid = tl.program_id(0).to(tl.int64) + num_programs = tl.num_programs(0) + + half_cols = tl.arange(0, pad_hd_half) + half_mask = half_cols < (hd // 2) + + wq1 = tl.load(wq_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wq2 = tl.load(wq_ptr + (hd // 2) + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wk1 = tl.load(wk_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + wk2 = tl.load(wk_ptr + (hd // 2) + half_cols, mask=half_mask, other=0.0).to(tl.float32) + + dwq1_acc = tl.zeros((pad_hd_half,), dtype=tl.float32) + dwq2_acc = tl.zeros((pad_hd_half,), dtype=tl.float32) + dwk1_acc = tl.zeros((pad_hd_half,), dtype=tl.float32) + dwk2_acc = tl.zeros((pad_hd_half,), dtype=tl.float32) + + q_heads = tl.arange(0, pad_n_qh) + k_heads = tl.arange(0, pad_n_kh) + q1_off = q_heads[:, None] * hd + half_cols[None, :] + q2_off = q1_off + (hd // 2) + k1_off = k_heads[:, None] * hd + half_cols[None, :] + k2_off = k1_off + (hd // 2) + q_mask = (q_heads[:, None] < n_qh) & half_mask[None, :] + k_mask = (k_heads[:, None] < n_kh) & half_mask[None, :] + + for token in range(pid, n_rows, num_programs): + batch_idx = token // seq_len + seq_idx = token % seq_len + c_ptr = cos_ptr + tl.where( + cos_bs == 1, + seq_idx * cos_row_stride, + batch_idx * (seq_len * cos_row_stride) + seq_idx * cos_row_stride, + ) + s_ptr = sin_ptr + tl.where( + cos_bs == 1, + seq_idx * sin_row_stride, + batch_idx * (seq_len * sin_row_stride) + seq_idx * sin_row_stride, + ) + cos_row = tl.load(c_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + sin_row = tl.load(s_ptr + half_cols, mask=half_mask, other=0.0).to(tl.float32) + + # ---------------- Q ---------------- + doq1 = tl.load(doq_ptr + token * doq_row_stride + q1_off, mask=q_mask, other=0.0).to(tl.float32) + doq2 = tl.load(doq_ptr + token * doq_row_stride + q2_off, mask=q_mask, other=0.0).to(tl.float32) + # RoPE backward -> grad wrt normed value (dY) + dqn1 = doq1 * cos_row[None, :] + doq2 * sin_row[None, :] + dqn2 = doq2 * cos_row[None, :] - doq1 * sin_row[None, :] + + xq1 = tl.load(q_ptr + token * q_row_stride + q1_off, mask=q_mask, other=0.0).to(tl.float32) + xq2 = tl.load(q_ptr + token * q_row_stride + q2_off, mask=q_mask, other=0.0).to(tl.float32) + rstd_q = tl.load(rstd_q_ptr + token * n_qh + q_heads, mask=q_heads < n_qh, other=0.0) + + # RMSNorm backward (llama): dx = rstd*(m - (1/hd)*rstd^2*sum(m*x)*x), m = dY*w + mq1 = dqn1 * wq1[None, :] + mq2 = dqn2 * wq2[None, :] + sum_mx_q = tl.sum(mq1 * xq1, axis=1) + tl.sum(mq2 * xq2, axis=1) + coef_q = (1.0 / hd) * rstd_q * rstd_q * sum_mx_q # (pad_n_qh,) + dxq1 = rstd_q[:, None] * (mq1 - coef_q[:, None] * xq1) + dxq2 = rstd_q[:, None] * (mq2 - coef_q[:, None] * xq2) + tl.store(dq_ptr + token * dq_row_stride + q1_off, dxq1.to(OUT_DTYPE), mask=q_mask) + tl.store(dq_ptr + token * dq_row_stride + q2_off, dxq2.to(OUT_DTYPE), mask=q_mask) + + # dW += dY * (x * rstd), reduce over heads (token accumulation happens in the loop) + dwq1_acc += tl.sum(dqn1 * (xq1 * rstd_q[:, None]), axis=0) + dwq2_acc += tl.sum(dqn2 * (xq2 * rstd_q[:, None]), axis=0) + + # ---------------- K ---------------- + dok1 = tl.load(dok_ptr + token * dok_row_stride + k1_off, mask=k_mask, other=0.0).to(tl.float32) + dok2 = tl.load(dok_ptr + token * dok_row_stride + k2_off, mask=k_mask, other=0.0).to(tl.float32) + dkn1 = dok1 * cos_row[None, :] + dok2 * sin_row[None, :] + dkn2 = dok2 * cos_row[None, :] - dok1 * sin_row[None, :] + + xk1 = tl.load(k_ptr + token * k_row_stride + k1_off, mask=k_mask, other=0.0).to(tl.float32) + xk2 = tl.load(k_ptr + token * k_row_stride + k2_off, mask=k_mask, other=0.0).to(tl.float32) + rstd_k = tl.load(rstd_k_ptr + token * n_kh + k_heads, mask=k_heads < n_kh, other=0.0) + + mk1 = dkn1 * wk1[None, :] + mk2 = dkn2 * wk2[None, :] + sum_mx_k = tl.sum(mk1 * xk1, axis=1) + tl.sum(mk2 * xk2, axis=1) + coef_k = (1.0 / hd) * rstd_k * rstd_k * sum_mx_k + dxk1 = rstd_k[:, None] * (mk1 - coef_k[:, None] * xk1) + dxk2 = rstd_k[:, None] * (mk2 - coef_k[:, None] * xk2) + tl.store(dk_ptr + token * dk_row_stride + k1_off, dxk1.to(OUT_DTYPE), mask=k_mask) + tl.store(dk_ptr + token * dk_row_stride + k2_off, dxk2.to(OUT_DTYPE), mask=k_mask) + + dwk1_acc += tl.sum(dkn1 * (xk1 * rstd_k[:, None]), axis=0) + dwk2_acc += tl.sum(dkn2 * (xk2 * rstd_k[:, None]), axis=0) + + tl.store(dwq_ptr + pid * dw_row_stride + half_cols, dwq1_acc, mask=half_mask) + tl.store(dwq_ptr + pid * dw_row_stride + (hd // 2) + half_cols, dwq2_acc, mask=half_mask) + tl.store(dwk_ptr + pid * dw_row_stride + half_cols, dwk1_acc, mask=half_mask) + tl.store(dwk_ptr + pid * dw_row_stride + (hd // 2) + half_cols, dwk2_acc, mask=half_mask) + + +def _num_warps(pad_n_qh, pad_n_kh, pad_hd_half): + block = max(pad_n_qh, pad_n_kh) * pad_hd_half + if block >= 8192: + return 16 + if block >= 2048: + return 8 + return 4 + + +def qk_norm_rope_forward(q, k, q_weight, k_weight, cos, sin, eps): + # q: (bsz, seq_len, n_qh, hd) ; k: (bsz, seq_len, n_kh, hd) (pre-transpose layout) + bsz, seq_len, n_qh, hd = q.shape + n_kh = k.shape[2] + assert hd % 2 == 0, "head_dim must be even for RoPE" + + q = q.contiguous() + k = k.contiguous() + cos = cos.contiguous() + sin = sin.contiguous() + + oq = torch.empty_like(q) + ok = torch.empty_like(k) + rstd_q = torch.empty((bsz * seq_len, n_qh), dtype=torch.float32, device=q.device) + rstd_k = torch.empty((bsz * seq_len, n_kh), dtype=torch.float32, device=q.device) + + pad_n_qh = triton.next_power_of_2(n_qh) + pad_n_kh = triton.next_power_of_2(n_kh) + pad_hd_half = triton.next_power_of_2(hd // 2) + n_rows = bsz * seq_len + + out_dtype = {torch.float32: tl.float32, torch.float16: tl.float16, torch.bfloat16: tl.bfloat16}[q.dtype] + + _qk_norm_rope_forward_kernel[(n_rows,)]( + q, + q.stride(1), + k, + k.stride(1), + oq, + oq.stride(1), + ok, + ok.stride(1), + q_weight, + k_weight, + cos, + cos.stride(-2), + sin, + sin.stride(-2), + rstd_q, + rstd_k, + seq_len, + eps, + n_qh, + n_kh, + hd, + cos.shape[0], + pad_n_qh, + pad_n_kh, + pad_hd_half, + out_dtype, + num_warps=_num_warps(pad_n_qh, pad_n_kh, pad_hd_half), + ) + # absorb the transpose: return (bsz, n_head, seq_len, hd) as a strided view + return oq.transpose(1, 2), ok.transpose(1, 2), q, k, rstd_q, rstd_k + + +def qk_norm_rope_backward(doq, dok, q, k, q_weight, k_weight, cos, sin, rstd_q, rstd_k, eps): + # doq/dok arrive as (bsz, n_head, seq_len, hd) transposed views -> back to (bsz, seq_len, n_head, hd) + doq = doq.transpose(1, 2).contiguous() + dok = dok.transpose(1, 2).contiguous() + + bsz, seq_len, n_qh, hd = q.shape + n_kh = k.shape[2] + n_rows = bsz * seq_len + + pad_n_qh = triton.next_power_of_2(n_qh) + pad_n_kh = triton.next_power_of_2(n_kh) + pad_hd_half = triton.next_power_of_2(hd // 2) + + sm_count = 1 + if q.device.type == "cuda": + sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count + n_programs = min(sm_count, n_rows) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + _dwq = torch.zeros((n_programs, hd), dtype=torch.float32, device=q.device) + _dwk = torch.zeros((n_programs, hd), dtype=torch.float32, device=q.device) + + out_dtype = {torch.float32: tl.float32, torch.float16: tl.float16, torch.bfloat16: tl.bfloat16}[q.dtype] + + _qk_norm_rope_backward_kernel[(n_programs,)]( + doq, + doq.stride(1), + dok, + dok.stride(1), + q, + q.stride(1), + k, + k.stride(1), + dq, + dq.stride(1), + dk, + dk.stride(1), + q_weight, + k_weight, + cos, + cos.stride(-2), + sin, + sin.stride(-2), + rstd_q, + rstd_k, + _dwq, + _dwk, + _dwq.stride(0), + n_rows, + seq_len, + n_qh, + n_kh, + hd, + cos.shape[0], + pad_n_qh, + pad_n_kh, + pad_hd_half, + out_dtype, + num_warps=_num_warps(pad_n_qh, pad_n_kh, pad_hd_half), + ) + + dq_weight = _dwq.sum(dim=0).to(q_weight.dtype) + dk_weight = _dwk.sum(dim=0).to(k_weight.dtype) + return dq, dk, dq_weight, dk_weight + + +class LigerQkNormRopeFunction(torch.autograd.Function): + """Fused per-head RMSNorm(Q/K) + RoPE. + + Inputs (matching the pre-transpose projection layout used by Qwen3): + q: (bsz, seq_len, n_q_head, head_dim) + k: (bsz, seq_len, n_kv_head, head_dim) + q_weight, k_weight: (head_dim,) + cos, sin: (1, seq_len, head_dim) or (bsz, seq_len, head_dim) + Returns: + q, k: (bsz, n_head, seq_len, head_dim) (transposed, ready for attention) + """ + + @staticmethod + @ensure_contiguous + def forward(ctx, q, k, q_weight, k_weight, cos, sin, eps): + oq, ok, q_saved, k_saved, rstd_q, rstd_k = qk_norm_rope_forward(q, k, q_weight, k_weight, cos, sin, eps) + ctx.eps = eps + ctx.save_for_backward(q_saved, k_saved, q_weight, k_weight, cos, sin, rstd_q, rstd_k) + return oq, ok + + @staticmethod + def backward(ctx, doq, dok): + q, k, q_weight, k_weight, cos, sin, rstd_q, rstd_k = ctx.saved_tensors + dq, dk, dq_weight, dk_weight = qk_norm_rope_backward( + doq, dok, q, k, q_weight, k_weight, cos, sin, rstd_q, rstd_k, ctx.eps + ) + return dq, dk, dq_weight, dk_weight, None, None, None diff --git a/src/liger_kernel/transformers/functional.py b/src/liger_kernel/transformers/functional.py index 9bb13d462..3ad0f4355 100644 --- a/src/liger_kernel/transformers/functional.py +++ b/src/liger_kernel/transformers/functional.py @@ -21,6 +21,7 @@ from liger_kernel.ops import LigerModulatedRMSNormFunction from liger_kernel.ops import LigerMultiTokenAttentionFunction from liger_kernel.ops import LigerPolyNormFunction +from liger_kernel.ops import LigerQkNormRopeFunction from liger_kernel.ops import LigerQwen2VLMRopeFunction from liger_kernel.ops import LigerReLUSquaredFunction from liger_kernel.ops import LigerRMSNormFunction @@ -288,6 +289,10 @@ def liger_qwen2vl_mrope(q, k, cos, sin, mrope_section, unsqueeze_dim=1): return LigerQwen2VLMRopeFunction.apply(q, k, cos, sin, mrope_section, unsqueeze_dim) +def liger_qk_norm_rope(q, k, q_weight, k_weight, cos, sin, eps): + return LigerQkNormRopeFunction.apply(q, k, q_weight, k_weight, cos, sin, eps) + + def liger_relu_squared(x): return LigerReLUSquaredFunction.apply(x) diff --git a/src/liger_kernel/transformers/model/qwen3_attention.py b/src/liger_kernel/transformers/model/qwen3_attention.py new file mode 100644 index 000000000..bbe4b8299 --- /dev/null +++ b/src/liger_kernel/transformers/model/qwen3_attention.py @@ -0,0 +1,74 @@ +from typing import Callable +from typing import Optional + +import torch + +from transformers.cache_utils import Cache +from transformers.models.qwen3.modeling_qwen3 import ALL_ATTENTION_FUNCTIONS +from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward + +from liger_kernel.transformers.qk_norm_rope import liger_qk_norm_rope + + +def qwen3_attention_forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + **kwargs, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Drop-in replacement for ``Qwen3Attention.forward`` using the fused + QK-Norm + RoPE Triton kernel. + + The reference implementation does (modeling_qwen3.py): + + query_states = self.q_norm(self.q_proj(h).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(h).view(hidden_shape)).transpose(1, 2) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + We fuse ``q_norm``/``k_norm`` + ``apply_rotary_pos_emb`` into a single kernel. + The RMSNorm weights live on ``self.q_norm``/``self.k_norm`` and act on the + ``head_dim`` axis, which is exactly the layout the fused kernel expects. + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + # pre-transpose layout: (bsz, seq_len, n_head, head_dim) + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = liger_qk_norm_rope( + query_states, + key_states, + self.q_norm.weight, + self.k_norm.weight, + cos, + sin, + self.q_norm.variance_epsilon, + ) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 4d33d7e41..91e351f8b 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -1668,10 +1668,18 @@ def apply_liger_kernel_to_qwen3( fused_linear_cross_entropy: bool = True, rms_norm: bool = True, swiglu: bool = True, + qk_norm_rope: bool = False, model: PreTrainedModel = None, ) -> None: """ Apply Liger kernels to replace original implementation in HuggingFace Qwen3 models. + + Args: + qk_norm_rope (bool): Whether to fuse the per-head ``q_norm``/``k_norm`` + RMSNorm with the rotary positional embedding into a single Triton + kernel by replacing ``Qwen3Attention.forward``. When enabled it + supersedes the standalone ``rope`` patch for Q/K (the fused kernel + applies RoPE internally). Default is False. """ assert not (cross_entropy and fused_linear_cross_entropy), ( "cross_entropy and fused_linear_cross_entropy cannot both be True." @@ -1681,8 +1689,9 @@ def apply_liger_kernel_to_qwen3( from transformers.models.qwen3.modeling_qwen3 import Qwen3Model from liger_kernel.transformers.model.qwen3 import lce_forward as qwen3_lce_forward + from liger_kernel.transformers.model.qwen3_attention import qwen3_attention_forward - if rope: + if rope and not qk_norm_rope: modeling_qwen3.apply_rotary_pos_emb = liger_rotary_pos_emb if rms_norm: @@ -1702,6 +1711,9 @@ def apply_liger_kernel_to_qwen3( if swiglu: modeling_qwen3.Qwen3MLP = LigerSwiGLUMLP + if qk_norm_rope: + modeling_qwen3.Qwen3Attention.forward = qwen3_attention_forward + if model is not None: # The model instance already exists, so we need to additionally patch the # instance variables that reference already-instantiated modules @@ -1717,6 +1729,8 @@ def apply_liger_kernel_to_qwen3( if rms_norm: _patch_rms_norm_module(decoder_layer.input_layernorm) _patch_rms_norm_module(decoder_layer.post_attention_layernorm) + if qk_norm_rope: + _bind_method_to_module(decoder_layer.self_attn, "forward", qwen3_attention_forward) def apply_liger_kernel_to_qwen3_moe( diff --git a/src/liger_kernel/transformers/qk_norm_rope.py b/src/liger_kernel/transformers/qk_norm_rope.py new file mode 100644 index 000000000..db6e43ee7 --- /dev/null +++ b/src/liger_kernel/transformers/qk_norm_rope.py @@ -0,0 +1,38 @@ +from typing import Tuple + +import torch + +from liger_kernel.ops.qk_norm_rope import LigerQkNormRopeFunction + + +def liger_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused per-head RMSNorm(Q/K) followed by rotary positional embedding. + + This fuses the ``q_norm``/``k_norm`` + ``apply_rotary_pos_emb`` sequence used by + models such as Qwen3. ``q`` and ``k`` are expected in the *pre-transpose* + projection layout, i.e. the output of ``proj(x).view(bsz, seq_len, n_head, + head_dim)`` before ``.transpose(1, 2)``. The returned tensors are already + transposed to ``(bsz, n_head, seq_len, head_dim)`` so they can be fed directly + into the attention interface. + + Args: + q: query states, shape ``(bsz, seq_len, n_q_head, head_dim)``. + k: key states, shape ``(bsz, seq_len, n_kv_head, head_dim)``. + q_weight: RMSNorm weight for the query, shape ``(head_dim,)``. + k_weight: RMSNorm weight for the key, shape ``(head_dim,)``. + cos: cosine table, shape ``(1, seq_len, head_dim)`` or ``(bsz, seq_len, head_dim)``. + sin: sine table, same shape as ``cos``. + eps: RMSNorm epsilon. + + Returns: + Tuple of query and key tensors, each ``(bsz, n_head, seq_len, head_dim)``. + """ + return LigerQkNormRopeFunction.apply(q, k, q_weight, k_weight, cos, sin, eps) diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index 25099f7be..423e2b15d 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -38,6 +38,7 @@ from liger_kernel.transformers.model.qwen2 import lce_forward as qwen2_lce_forward from liger_kernel.transformers.model.qwen3_5 import lce_forward as qwen3_5_lce_forward from liger_kernel.transformers.model.qwen3_5 import lce_forward_for_multimodal as qwen3_5_lce_forward_for_multimodal +from liger_kernel.transformers.model.qwen3_attention import qwen3_attention_forward from liger_kernel.transformers.model.qwen3_next import lce_forward as qwen3_next_lce_forward from liger_kernel.transformers.model.smollm3 import lce_forward as smolllm3_lce_forward from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN @@ -2192,6 +2193,20 @@ def test_apply_liger_kernel_to_instance_for_qwen3(): pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}") +@pytest.mark.skipif(not is_qwen3_available(), reason="qwen3 module not available") +def test_qwen3_qk_norm_rope_hook_applied(): + # Ensure any monkey patching is cleaned up for subsequent tests + with patch("transformers.models.qwen3.modeling_qwen3") as modeling_mod: + setattr(modeling_mod, "apply_rotary_pos_emb", object()) + modeling_mod.Qwen3Attention = MagicMock() + modeling_mod.Qwen3Attention.forward = object() + + _apply_liger_kernel("qwen3", qk_norm_rope=True) + + assert modeling_mod.Qwen3Attention.forward is qwen3_attention_forward + assert modeling_mod.apply_rotary_pos_emb is not monkey_patch.liger_rotary_pos_emb + + @pytest.mark.skipif(not is_qwen3_available(), reason="qwen3 module not available") def test_apply_liger_kernel_to_instance_for_qwen3_moe(): # Ensure any monkey patching is cleaned up for subsequent tests diff --git a/test/transformers/test_qk_norm_rope.py b/test/transformers/test_qk_norm_rope.py new file mode 100644 index 000000000..b0623bbb8 --- /dev/null +++ b/test/transformers/test_qk_norm_rope.py @@ -0,0 +1,322 @@ +"""Correctness tests for the fused QK-Norm + RoPE Triton kernel. + +The reference stacks a per-head RMSNorm (over ``head_dim``) with the HuggingFace +Llama/Qwen rotary embedding, exactly as ``Qwen3Attention.forward`` does:: + + q = q_norm(q_proj(x).view(B, T, n_qh, hd)).transpose(1, 2) + k = k_norm(k_proj(x).view(B, T, n_kh, hd)).transpose(1, 2) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + +Correctness methodology +----------------------- +A fused kernel that keeps the RMSNorm->RoPE chain in fp32 is *not* bit-identical +to running RMSNorm and RoPE as two separate ops -- in low precision it is +actually more accurate (it rounds once instead of round-tripping the normalized +Q/K through bf16). A fixed ``atol`` on ``fused vs naive`` therefore either has +to be loose enough to be meaningless or it flags precision, not bugs. + +Instead we compute a fully-fp32 "gold" result and assert the fused kernel is +**no worse than the naive same-dtype reference** relative to that gold (plus a +small absolute floor for the fp32 case where the naive path *is* the gold). +This encodes the real property we care about -- "fusing must not lose accuracy" +-- rather than a magic threshold. +""" + +import pytest +import torch + +from liger_kernel.ops.qk_norm_rope import LigerQkNormRopeFunction +from liger_kernel.transformers.qk_norm_rope import liger_qk_norm_rope + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def rms_norm_ref(x, weight, eps): + # llama casting: reduce/rstd in fp32, weight multiply back in input dtype + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + return weight * x.to(input_dtype) + + +def ref_forward(q, k, wq, wk, cos, sin, eps): + # q, k: (B, T, n_head, hd) + q = rms_norm_ref(q, wq, eps).transpose(1, 2) + k = rms_norm_ref(k, wk, eps).transpose(1, 2) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + return q, k + + +def make_cos_sin(bsz, seq_len, head_dim, device, dtype, batched=False, base=10000): + pos = torch.arange(seq_len, device=device, dtype=torch.float32) + inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) + freqs = torch.outer(pos, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + if batched: + cos = cos.unsqueeze(0).expand(bsz, -1, -1) + sin = sin.unsqueeze(0).expand(bsz, -1, -1) + else: + cos = cos.unsqueeze(0) + sin = sin.unsqueeze(0) + return cos.to(dtype), sin.to(dtype) + + +def _max_abs_err(a, b): + return (a.float() - b.float()).abs().max().item() + + +def _assert_not_worse(fused, ref, gold, name, slack=2.0, floor=0.0): + """Assert the fused result is no worse than the naive reference vs fp32 gold. + + ``slack`` absorbs reduction-order differences; ``floor`` is an absolute + allowance for the fp32 case, where the naive reference *is* the gold so its + error is exactly zero. + """ + fused_err = _max_abs_err(fused, gold) + ref_err = _max_abs_err(ref, gold) + assert fused_err <= ref_err * slack + floor, ( + f"{name}: fused err {fused_err:.3e} > ref err {ref_err:.3e} * {slack} + {floor:.3e}" + ) + + +def _run_gold_ref_fused(q, k, wq, wk, cos, sin, eps): + """Run fp32-gold, naive-same-dtype, and fused paths with a shared upstream grad. + + Returns three dicts keyed by q/k/dq/dk/dwq/dwk. + """ + dtype = q.dtype + cos32, sin32 = cos.float(), sin.float() + + def leaves(dt): + return ( + q.clone().to(dt).requires_grad_(True), + k.clone().to(dt).requires_grad_(True), + wq.clone().to(dt).requires_grad_(True), + wk.clone().to(dt).requires_grad_(True), + ) + + # ---- fp32 gold ---- + gq, gk, gwq, gwk = leaves(torch.float32) + gold_q, gold_k = ref_forward(gq, gk, gwq, gwk, cos32, sin32, eps) + grad_q = torch.randn_like(gold_q) + grad_k = torch.randn_like(gold_k) + ((gold_q * grad_q).sum() + (gold_k * grad_k).sum()).backward() + gold = dict(q=gold_q, k=gold_k, dq=gq.grad, dk=gk.grad, dwq=gwq.grad, dwk=gwk.grad) + + # ---- naive same-dtype reference ---- + rq, rk, rwq, rwk = leaves(dtype) + ref_q, ref_k = ref_forward(rq, rk, rwq, rwk, cos, sin, eps) + ((ref_q * grad_q.to(dtype)).sum() + (ref_k * grad_k.to(dtype)).sum()).backward() + ref = dict(q=ref_q, k=ref_k, dq=rq.grad, dk=rk.grad, dwq=rwq.grad, dwk=rwk.grad) + + # ---- fused kernel ---- + fq, fk, fwq, fwk = leaves(dtype) + fus_q, fus_k = liger_qk_norm_rope(fq, fk, fwq, fwk, cos, sin, eps) + ((fus_q * grad_q.to(dtype)).sum() + (fus_k * grad_k.to(dtype)).sum()).backward() + fus = dict(q=fus_q, k=fus_k, dq=fq.grad, dk=fk.grad, dwq=fwq.grad, dwk=fwk.grad) + + return gold, ref, fus + + +def _check_all_not_worse(gold, ref, fus, act_floor, wgrad_floor): + # activations and input grads + for key in ("q", "k", "dq", "dk"): + _assert_not_worse(fus[key], ref[key], gold[key], key, slack=2.0, floor=act_floor) + # weight grads accumulate over all tokens/heads -> large magnitude, so a + # bigger absolute floor for the fp32 reduction-order term. + for key in ("dwq", "dwk"): + _assert_not_worse(fus[key], ref[key], gold[key], key, slack=2.0, floor=wgrad_floor) + + +# Real Qwen3 dense attention shapes (n_q_head, n_kv_head, head_dim), taken from +# the published HuggingFace ``config.json`` files. Qwen3 decouples ``head_dim`` +# (fixed at 128) from ``hidden_size / num_attention_heads`` and uses GQA with 8 +# KV heads across the whole dense family; ``rms_norm_eps`` is 1e-6 everywhere. +QWEN3_REAL_CONFIGS = [ + pytest.param(16, 8, 128, id="qwen3_0.6b"), + pytest.param(32, 8, 128, id="qwen3_4b_8b"), + pytest.param(40, 8, 128, id="qwen3_14b"), + pytest.param(64, 8, 128, id="qwen3_32b"), +] + +# fp32 leaves the naive path == gold, so allow a small absolute reduction-order +# floor; bf16 relies purely on "fused must not be worse than naive". +DTYPE_FLOORS = [ + pytest.param(torch.float32, 1e-4, 1e-2, id="fp32"), # (dtype, act_floor, wgrad_floor) + pytest.param(torch.bfloat16, 0.0, 0.0, id="bf16"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("n_qh, n_kh, hd", QWEN3_REAL_CONFIGS) +@pytest.mark.parametrize( + "bsz, seq_len", + [ + (1, 4096), # single long-context sequence + (4, 2048), # typical training micro-batch + ], +) +@pytest.mark.parametrize("dtype, act_floor, wgrad_floor", DTYPE_FLOORS) +def test_qk_norm_rope_qwen3_real_configs(n_qh, n_kh, hd, bsz, seq_len, dtype, act_floor, wgrad_floor): + """Correctness on real Qwen3 dense attention shapes and training seq lengths. + + Uses the production ``rms_norm_eps`` (1e-6), ``rope_theta`` (1e6) and the GQA + head layout that the fused ``qwen3_attention_forward`` monkeypatch actually + feeds the kernel. Asserts the fused kernel is no worse than the naive + same-dtype path measured against a fully-fp32 gold reference. + """ + device = "cuda" + eps = 1e-6 # Qwen3 rms_norm_eps + torch.manual_seed(0) + + q = torch.randn(bsz, seq_len, n_qh, hd, device=device, dtype=dtype) + k = torch.randn(bsz, seq_len, n_kh, hd, device=device, dtype=dtype) + # RMSNorm weights initialize to ones in Qwen3; jitter around 1.0 to exercise + # a realistic-but-non-trivial per-channel scale. + wq = 1.0 + 0.1 * torch.randn(hd, device=device, dtype=dtype) + wk = 1.0 + 0.1 * torch.randn(hd, device=device, dtype=dtype) + cos, sin = make_cos_sin(bsz, seq_len, hd, device, dtype, batched=False, base=1000000) + + gold, ref, fus = _run_gold_ref_fused(q, k, wq, wk, cos, sin, eps) + + # shape / layout + assert fus["q"].shape == (bsz, n_qh, seq_len, hd) + assert fus["k"].shape == (bsz, n_kh, seq_len, hd) + + _check_all_not_worse(gold, ref, fus, act_floor, wgrad_floor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("bsz, seq_len", [(1, 128), (2, 200), (4, 51)]) +@pytest.mark.parametrize("n_qh, n_kh, hd", [(32, 8, 128), (16, 16, 64), (28, 4, 128)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("batched_cos", [False, True]) +def test_qk_norm_rope_correctness(bsz, seq_len, n_qh, n_kh, hd, dtype, batched_cos): + """Small-shape sweep (odd seq lens, MHA/GQA mixes, batched vs shared cos/sin). + + Same "no worse than naive vs fp32 gold" methodology as the real-config test. + """ + device = "cuda" + eps = 1e-6 + torch.manual_seed(0) + act_floor = 1e-4 if dtype == torch.float32 else 0.0 + wgrad_floor = 1e-2 if dtype == torch.float32 else 0.0 + + q = torch.randn(bsz, seq_len, n_qh, hd, device=device, dtype=dtype) + k = torch.randn(bsz, seq_len, n_kh, hd, device=device, dtype=dtype) + wq = torch.randn(hd, device=device, dtype=dtype) + wk = torch.randn(hd, device=device, dtype=dtype) + cos, sin = make_cos_sin(bsz, seq_len, hd, device, dtype, batched=batched_cos) + + gold, ref, fus = _run_gold_ref_fused(q, k, wq, wk, cos, sin, eps) + + assert fus["q"].shape == (bsz, n_qh, seq_len, hd) + assert fus["k"].shape == (bsz, n_kh, seq_len, hd) + + _check_all_not_worse(gold, ref, fus, act_floor, wgrad_floor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_qk_norm_rope_patch_matches_reference(): + """Sanity-check the fp32 layout/values of the standalone function.""" + device = "cuda" + eps = 1e-6 + bsz, seq_len, n_qh, n_kh, hd = 2, 64, 8, 2, 128 + dtype = torch.float32 + torch.manual_seed(0) + + q = torch.randn(bsz, seq_len, n_qh, hd, device=device, dtype=dtype) + k = torch.randn(bsz, seq_len, n_kh, hd, device=device, dtype=dtype) + wq = torch.randn(hd, device=device, dtype=dtype) + wk = torch.randn(hd, device=device, dtype=dtype) + cos, sin = make_cos_sin(bsz, seq_len, hd, device, dtype, batched=False) + + rq, rk = ref_forward(q, k, wq, wk, cos, sin, eps) + fq, fk = LigerQkNormRopeFunction.apply(q, k, wq, wk, cos, sin, eps) + + assert torch.allclose(rq, fq, atol=1e-4, rtol=1e-4) + assert torch.allclose(rk, fk, atol=1e-4, rtol=1e-4) + # returned tensors must be laid out (bsz, n_head, seq_len, head_dim) + assert fq.shape == (bsz, n_qh, seq_len, hd) + assert fk.shape == (bsz, n_kh, seq_len, hd) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_qwen3_attention_forward_patch_end_to_end(dtype): + """End-to-end: the monkeypatched ``Qwen3Attention.forward`` must be no worse + than the stock HuggingFace attention output (both measured against an fp32 + gold) on a real (scaled-down) Qwen3 config. + + We build an actual ``Qwen3Attention`` layer so the fused kernel is exercised + through ``self.q_norm``/``self.k_norm``/``apply_rotary_pos_emb`` exactly as it + is in a live model. + """ + import types + + from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention + from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + + from liger_kernel.transformers.model.qwen3_attention import qwen3_attention_forward + + device = "cuda" + torch.manual_seed(0) + + bsz, seq_len = 2, 512 + # Real Qwen3 head layout (GQA 32/8, head_dim=128), hidden trimmed so the + # projection weights stay small; head_dim / eps / theta match production. + config = Qwen3Config( + hidden_size=32 * 128, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + rms_norm_eps=1e-6, + max_position_embeddings=40960, + rope_theta=1000000, + attention_dropout=0.0, + _attn_implementation="eager", + ) + attn = Qwen3Attention(config, layer_idx=0).to(device=device).eval() + rotary = Qwen3RotaryEmbedding(config, device=device) + + hidden = torch.randn(bsz, seq_len, config.hidden_size, device=device) + pos_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + + def run(module, h, cos, sin): + with torch.no_grad(): + out, _ = module(h, position_embeddings=(cos, sin), attention_mask=None) + return out + + # fp32 gold: stock forward in fp32 + cos32, sin32 = rotary(hidden, pos_ids) + gold = run(attn, hidden, cos32, sin32) + + # cast the whole layer + inputs to the target dtype + attn = attn.to(dtype=dtype) + h = hidden.to(dtype) + cos, sin = cos32.to(dtype), sin32.to(dtype) + + ref = run(attn, h, cos, sin) # stock forward, target dtype + + attn.forward = types.MethodType(qwen3_attention_forward, attn) + fused = run(attn, h, cos, sin) # fused forward, target dtype + + assert fused.shape == gold.shape + floor = 1e-4 if dtype == torch.float32 else 0.0 + _assert_not_worse(fused, ref, gold, "attn_out", slack=2.0, floor=floor)