From d313f43d3a2b53b8dd33d7b2e84da60d4f408f79 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 28 May 2026 22:39:06 +0800 Subject: [PATCH 01/17] rmsnorm tests: reorganize variants, use helpers to reduce duplication, simplify numeric checks --- tests/kernels/benchmark_common.py | 8 +- tests/kernels/test_rmsnorm.py | 978 +++++++++++++++--------------- 2 files changed, 482 insertions(+), 504 deletions(-) diff --git a/tests/kernels/benchmark_common.py b/tests/kernels/benchmark_common.py index 03ffc4b9e..13691b615 100644 --- a/tests/kernels/benchmark_common.py +++ b/tests/kernels/benchmark_common.py @@ -41,10 +41,10 @@ class PerfRow: aiter_gpu_us: Optional[float] @property - def speedup_aiter_vs_flydsl(self) -> Optional[float]: + def speedup_flydsl_vs_aiter(self) -> Optional[float]: if self.flydsl_gpu_us is None or self.aiter_gpu_us is None: return None - return self.flydsl_gpu_us / self.aiter_gpu_us + return self.aiter_gpu_us / self.flydsl_gpu_us def _fmt_us(x: Optional[float]) -> str: @@ -57,7 +57,7 @@ def print_perf_table(rows: List[PerfRow]) -> None: print("=" * 100) print(f"{'op':10s} {'shape':18s} {'dtype':6s} {'FlyDSL(gpu us)':>14s} {'AIter(gpu us)':>14s} {'speedup':>10s}") for r in rows: - sp = r.speedup_aiter_vs_flydsl + sp = r.speedup_flydsl_vs_aiter sp_s = "-" if sp is None else f"{sp:,.2f}x" print( f"{r.op:10s} {r.shape:18s} {r.dtype:6s} {_fmt_us(r.flydsl_gpu_us):>14s} {_fmt_us(r.aiter_gpu_us):>14s} {sp_s:>10s}" @@ -946,7 +946,7 @@ def main() -> None: print("=" * 100) print(f"{'op':10s} {'shape':18s} {'dtype':6s} {'FlyDSL(gpu us)':>14s} {'torch(gpu us)':>14s} {'speedup':>10s}") for r in wmma_rows: - sp = r.speedup_aiter_vs_flydsl + sp = r.speedup_flydsl_vs_aiter sp_s = "-" if sp is None else f"{sp:,.2f}x" print( f"{r.op:10s} {r.shape:18s} {r.dtype:6s} {_fmt_us(r.flydsl_gpu_us):>14s} {_fmt_us(r.aiter_gpu_us):>14s} {sp_s:>10s}" diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 04eae3c92..e0a7f6407 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -53,6 +53,40 @@ BENCH_ITERS = 100 +def _torch_dtype(dtype: str): + if dtype == "f32": + return DTYPE_FP32 + if dtype == "f16": + return DTYPE_FP16 + if dtype == "bf16": + return DTYPE_BF16 + raise ValueError(f"unsupported dtype: {dtype}") + + +def _get_rmsnorm_configs(): + shapes_env = os.environ.get("ROCDSL_RMSNORM_SHAPES", "").strip() + if shapes_env: + configs = [] + for part in shapes_env.split(";"): + p = part.strip() + if not p: + continue + m_s, n_s, dt = [x.strip() for x in p.split(",")] + configs.append((int(m_s), int(n_s), dt)) + else: + # Prefer N multiples of 2048 to exercise the fast path. + configs = [ + (64, 256, "f32"), # Aligned + (128, 1024, "f32"), # Aligned + (32, 128, "f16"), # Aligned + (64, 2000, "f32"), # Unaligned (tail handling) + (16, 512, "bf16"), # BF16 + (1024, 8192, "bf16"), # BF16 + (32768, 8192, "bf16"), + ] + return configs + + def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting RMSNorm (M={M}, N={N}, dtype={dtype})") @@ -61,6 +95,7 @@ def run_test(M: int, N: int, dtype: str = "f32"): except Exception as e: print(f"[FAIL] Compile failed for (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None + torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) @@ -143,166 +178,6 @@ def kernel_launch(): return ok, flydsl_gpu_us -def test_all(): - print("=" * 80) - print("Running RMSNorm Tests") - print("=" * 80) - - shapes_env = os.environ.get("ROCDSL_RMSNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - # Prefer N multiples of 2048 to exercise the fast path. - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] - - failures = 0 - for M, N, dtype in configs: - ok, flydsl_gpu_us = run_test(M, N, dtype) - if not ok: - failures += 1 - - if do_compare: - import torch - - aiter_us = None - if maybe_enable_aiter(): - try: - from aiter.ops.triton.rmsnorm import rms_norm as aiter_rms_norm - - x = torch.randn( - (M, N), - device="cuda", - dtype=DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32), - ) - w = torch.rand((N,), device="cuda", dtype=x.dtype) - - def run_aiter(): - aiter_rms_norm(x, w, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter rmsnorm gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter rmsnorm skipped: {type(e).__name__}: {e!r}") - - perf_rows.append( - PerfRow(op="rmsnorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us) - ) - - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") - else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - # Ensure a non-zero exit code on failure for shell wrappers. - if failures != 0: - raise SystemExit(1) - - -def _torch_dtype(dtype: str): - if dtype == "f32": - return DTYPE_FP32 - if dtype == "f16": - return DTYPE_FP16 - if dtype == "bf16": - return DTYPE_BF16 - raise ValueError(f"unsupported dtype: {dtype}") - - -def _get_rmsnorm_configs(): - shapes_env = os.environ.get("ROCDSL_RMSNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - return configs - - # Prefer N multiples of 2048 to exercise the fast path. - return [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] - - -def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None): - x = input_dev.to(DTYPE_FP32) - gamma = gamma_dev.to(DTYPE_FP32) - expected = (x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + EPS)) * gamma - if xscale_dev is not None: - expected = expected * xscale_dev.to(DTYPE_FP32) - - yscale = expected.abs().amax(dim=1) / 127.0 - yscale = torch.where(yscale == 0, torch.ones_like(yscale), yscale) - q = torch.clamp(torch.trunc(expected / yscale.unsqueeze(1)), -127, 127).to(torch.int8) - return expected, q, yscale - - -def _bench_aiter_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): - mode = "smoothquant" if is_smooth else "dynamicquant" - torch_dtype = _torch_dtype(dtype) - - try: - if is_smooth: - from aiter.ops.triton.normalization.rmsnorm import ( - rmsnorm2d_fwd_with_smoothquant as aiter_rmsnorm_quant, - ) - else: - from aiter.ops.triton.normalization.rmsnorm import ( - rmsnorm2d_fwd_with_dynamicquant as aiter_rmsnorm_quant, - ) - except Exception as e: - print(f"[Perf] AIter rmsnorm {mode} skipped: {type(e).__name__}: {e!r}") - return None - - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - y = torch.empty((M, N), dtype=torch.int8, device="cuda") - yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") - - if is_smooth: - xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() - - def run_aiter(): - aiter_rmsnorm_quant(y, x, xscale, yscale, w, EPS) - - else: - - def run_aiter(): - aiter_rmsnorm_quant(y, x, yscale, w, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter rmsnorm {mode} gpu: {aiter_us:.1f} us") - return aiter_us - - def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting RMSNorm {mode} (M={M}, N={N}, dtype={dtype})") @@ -315,6 +190,7 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): except Exception as e: print(f"[FAIL] Compile failed for {mode} (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") return False, None + torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) @@ -336,17 +212,7 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): xscale_dev = None if is_smooth: xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() - dequant_tol = 0.25 if is_smooth else 0.2 - scale_tol = 1e-2 if is_smooth else 5e-3 - - # PyTorch Reference: - # RMS(x) = sqrt(mean(x^2) + eps) ; RMSNorm(x) = x / RMS(x) * gamma - # Quant path additionally computes per-row yscale and int8 output from the fp32 reference. - expected, q_ref, yscale_ref = _reference_rmsnorm_quant( - input_dev, - gamma_dev, - xscale_dev=xscale_dev, - ) + scale_tol = 1e-3 print("Launching kernel...") stream = torch.cuda.current_stream() @@ -381,30 +247,30 @@ def kernel_launch(): if flydsl_gpu_us is not None: print(f"[Perf] FlyDSL rmsnorm {mode} gpu: {flydsl_gpu_us:.1f} us") + # PyTorch Reference: + # RMS(x) = sqrt(mean(x^2) + eps) ; RMSNorm(x) = x / RMS(x) * gamma + # Quant path additionally computes per-row yscale and int8 output from the fp32 reference. + q_ref, yscale_ref = _reference_rmsnorm_quant( + input_dev, + gamma_dev, + xscale_dev=xscale_dev, + ) q_out = output_dev.to(torch.int16) q_expected = q_ref.to(torch.int16) yscale_out = yscale_dev.cpu() yscale_expected = yscale_ref.cpu() - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) - error = (output_ref - expected).abs().max().item() - scale_diff = (yscale_out - yscale_expected).abs().max().item() - quant_diff = (q_out - q_expected).abs().max().item() + quant_error = (q_out - q_expected).abs().max().item() + scale_error = (yscale_out - yscale_expected).abs().max().item() - print(f"Max dequant error: {error:.2e} (tol={dequant_tol})") - print(f"Max scale diff: {scale_diff:.2e} (tol={scale_tol})") - print(f"Max quant diff: {quant_diff}") + print(f"Max quant diff: {quant_error}") + print(f"Max scale diff: {scale_error:.2e} (tol={scale_tol})") - ok = error < dequant_tol and scale_diff < scale_tol and quant_diff <= 1 + ok = quant_error <= 1 and scale_error < scale_tol if ok: print("PASSED") - ok = True else: print("FAILED") - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) print("First row Quant Expected:") print(q_expected[0, :8]) print("First row Quant Actual:") @@ -413,168 +279,49 @@ def kernel_launch(): print(yscale_expected[:5]) print("First few YScale Actual:") print(yscale_out[:5]) - ok = False return ok, flydsl_gpu_us -def test_rmsnorm_dynamicquant(): - print("=" * 80) - print("Running RMSNorm DynamicQuant Tests") - print("=" * 80) - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] +def run_fused_add_test(M: int, N: int, dtype: str): + print(f"\nTesting FusedAdd RMSNorm (M={M}, N={N}, dtype={dtype})") - failures = 0 - for M, N, dtype in _get_rmsnorm_configs(): - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False) - if not ok: - failures += 1 + try: + launch_fn = build_fused_add_rmsnorm_module(M, N, dtype) + except Exception as e: + print(f"[FAIL] Compile failed for fused_add rmsnorm (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") + return False, None - if do_compare: - aiter_us = None - if maybe_enable_aiter(): - aiter_us = _bench_aiter_rmsnorm_quant(M, N, dtype, is_smooth=False) - perf_rows.append( - PerfRow( - op="rmsnorm_dq", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, - ) - ) + torch.manual_seed(42) + input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") + if dtype == "f32": + input_dev = input_t.contiguous() + residual_in_dev = residual_t.contiguous() + gamma_dev = gamma_t.contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) + residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) + atol = 1e-4 + elif dtype == "f16": + input_dev = input_t.to(DTYPE_FP16).contiguous() + residual_in_dev = residual_t.to(DTYPE_FP16).contiguous() + gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) + residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) + atol = 1e-2 + elif dtype == "bf16": + input_dev = input_t.to(DTYPE_BF16).contiguous() + residual_in_dev = residual_t.to(DTYPE_BF16).contiguous() + gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) + residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) + atol = 2e-2 else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - # Ensure a non-zero exit code on failure for shell wrappers. - if failures != 0: - raise SystemExit(1) - + raise ValueError(f"unsupported dtype: {dtype}") -def test_rmsnorm_smoothquant(): - print("=" * 80) - print("Running RMSNorm SmoothQuant Tests") - print("=" * 80) - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] - failures = 0 - - for M, N, dtype in _get_rmsnorm_configs(): - ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True) - if not ok: - failures += 1 - - if do_compare: - aiter_us = None - if maybe_enable_aiter(): - aiter_us = _bench_aiter_rmsnorm_quant(M, N, dtype, is_smooth=True) - perf_rows.append( - PerfRow( - op="rmsnorm_sq", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, - ) - ) - - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") - else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - # Ensure a non-zero exit code on failure for shell wrappers. - if failures != 0: - raise SystemExit(1) - - -def _reference_fused_add_rmsnorm(input_dev, residual_in_dev, gamma_dev): - added = input_dev + residual_in_dev - added_fp32 = added.to(DTYPE_FP32) - gamma = gamma_dev.to(DTYPE_FP32) - expected = (added_fp32 / torch.sqrt((added_fp32 * added_fp32).mean(dim=1, keepdim=True) + EPS)) * gamma - return added_fp32, expected - - -def _bench_aiter_fused_add_rmsnorm(M: int, N: int, dtype: str): - torch_dtype = _torch_dtype(dtype) - - try: - from aiter.ops.triton.normalization.rmsnorm import ( - rmsnorm2d_fwd_with_add as aiter_fused_add_rmsnorm, - ) - except Exception as e: - print(f"[Perf] AIter fused_add rmsnorm skipped: {type(e).__name__}: {e!r}") - return None - - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual_in = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - out = torch.empty((M, N), device="cuda", dtype=torch_dtype) - residual_out = torch.empty((M, N), device="cuda", dtype=torch_dtype) - - def run_aiter(): - aiter_fused_add_rmsnorm(out, x, residual_in, residual_out, w, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter fused_add rmsnorm gpu: {aiter_us:.1f} us") - return aiter_us - - -def run_fused_add_test(M: int, N: int, dtype: str): - print(f"\nTesting FusedAdd RMSNorm (M={M}, N={N}, dtype={dtype})") - - try: - launch_fn = build_fused_add_rmsnorm_module(M, N, dtype) - except Exception as e: - print(f"[FAIL] Compile failed for fused_add rmsnorm (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") - return False, None - - torch.manual_seed(42) - input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - - if dtype == "f32": - input_dev = input_t.contiguous() - residual_in_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - output_atol = 1e-4 - residual_atol = 1e-4 - elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_in_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - output_atol = 1e-2 - residual_atol = 1e-2 - elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_in_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - output_atol = 2e-2 - residual_atol = 2e-2 - else: - raise ValueError(f"unsupported dtype: {dtype}") - - print("Launching kernel...") - stream = torch.cuda.current_stream() + print("Launching kernel...") + stream = torch.cuda.current_stream() def kernel_launch(): launch_fn( @@ -607,6 +354,8 @@ def kernel_launch(): if flydsl_gpu_us is not None: print(f"[Perf] FlyDSL fused_add rmsnorm gpu: {flydsl_gpu_us:.1f} us") + # PyTorch Reference: + # RMS(x) = sqrt(mean(x^2) + eps) ; RMSNorm(x) = x / RMS(x) * gamma residual_expected, output_expected = _reference_fused_add_rmsnorm( input_dev, residual_in_dev, @@ -618,121 +367,25 @@ def kernel_launch(): residual_error = (residual_out_ref - residual_expected).abs().max().item() output_error = (output_ref - output_expected).abs().max().item() - print(f"Max residual error: {residual_error:.2e} (atol={residual_atol})") - print(f"Max output error: {output_error:.2e} (atol={output_atol})") + print(f"Max residual error: {residual_error:.2e} (atol={atol})") + print(f"Max output error: {output_error:.2e} (atol={atol})") - ok = residual_error < residual_atol and output_error < output_atol + ok = residual_error < atol and output_error < atol if ok: print("PASSED") else: print("FAILED") + print("First row Residual Expected:") + print(residual_expected[0, :5]) + print("First row Residual Actual:") + print(residual_out_ref[0, :5]) + print("First row Output Expected:") + print(output_expected[0, :5]) + print("First row Output Actual:") + print(output_ref[0, :5]) return ok, flydsl_gpu_us -def test_rmsnorm_fused_add(): - print("=" * 80) - print("Running FusedAdd RMSNorm Tests") - print("=" * 80) - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] - failures = 0 - - for M, N, dtype in _get_rmsnorm_configs(): - ok, flydsl_gpu_us = run_fused_add_test(M, N, dtype) - if not ok: - failures += 1 - - if do_compare: - aiter_us = None - if maybe_enable_aiter(): - aiter_us = _bench_aiter_fused_add_rmsnorm(M, N, dtype) - perf_rows.append( - PerfRow( - op="rmsnorm_add", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, - ) - ) - - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") - else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - # Ensure a non-zero exit code on failure for shell wrappers. - if failures != 0: - raise SystemExit(1) - - -def _reference_fused_add_rmsnorm_quant( - input_dev, - residual_in_dev, - gamma_dev, - *, - xscale_dev=None, -): - added = input_dev + residual_in_dev - residual_expected = added.to(DTYPE_FP32) - expected, q, yscale = _reference_rmsnorm_quant( - added, - gamma_dev, - xscale_dev=xscale_dev, - ) - return residual_expected, expected, q, yscale - - -def _bench_aiter_fused_add_rmsnorm_quant( - M: int, - N: int, - dtype: str, - *, - is_smooth: bool, -): - mode = "smoothquant" if is_smooth else "dynamicquant" - torch_dtype = _torch_dtype(dtype) - - try: - if is_smooth: - from aiter.ops.triton.normalization.rmsnorm import ( - rmsnorm2d_fwd_with_add_smoothquant as aiter_fused_add_rmsnorm_quant, - ) - else: - from aiter.ops.triton.normalization.rmsnorm import ( - rmsnorm2d_fwd_with_add_dynamicquant as aiter_fused_add_rmsnorm_quant, - ) - except Exception as e: - print(f"[Perf] AIter fused_add rmsnorm {mode} skipped: {type(e).__name__}: {e!r}") - return None - - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual_in = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - y = torch.empty((M, N), dtype=torch.int8, device="cuda") - residual_out = torch.empty((M, N), device="cuda", dtype=torch_dtype) - yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") - - if is_smooth: - xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() - - def run_aiter(): - aiter_fused_add_rmsnorm_quant(y, x, residual_in, residual_out, xscale, yscale, w, EPS) - - else: - - def run_aiter(): - aiter_fused_add_rmsnorm_quant(y, x, residual_in, residual_out, yscale, w, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter fused_add rmsnorm {mode} gpu: {aiter_us:.1f} us") - return aiter_us - - def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): mode = "smoothquant" if is_smooth else "dynamicquant" print(f"\nTesting FusedAdd RMSNorm {mode} (M={M}, N={N}, dtype={dtype})") @@ -780,15 +433,7 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): xscale_dev = None if is_smooth: xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() - dequant_tol = 0.25 if is_smooth else 0.2 - scale_tol = 1e-2 if is_smooth else 5e-3 - - residual_expected, expected, q_ref, yscale_ref = _reference_fused_add_rmsnorm_quant( - input_dev, - residual_in_dev, - gamma_dev, - xscale_dev=xscale_dev, - ) + scale_tol = 1e-3 print("Launching kernel...") stream = torch.cuda.current_stream() @@ -840,24 +485,29 @@ def kernel_launch(): if flydsl_gpu_us is not None: print(f"[Perf] FlyDSL fused_add rmsnorm {mode} gpu: {flydsl_gpu_us:.1f} us") + # PyTorch Reference: + # RMS(x) = sqrt(mean(x^2) + eps) ; RMSNorm(x) = x / RMS(x) * gamma + residual_expected, q_ref, yscale_ref = _reference_fused_add_rmsnorm_quant( + input_dev, + residual_in_dev, + gamma_dev, + xscale_dev=xscale_dev, + ) residual_out_ref = residual_out_dev.to(DTYPE_FP32) - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) q_out = output_dev.to(torch.int16) q_expected = q_ref.to(torch.int16) yscale_out = yscale_dev.cpu() yscale_expected = yscale_ref.cpu() residual_error = (residual_out_ref - residual_expected).abs().max().item() - dequant_error = (output_ref - expected).abs().max().item() - scale_diff = (yscale_out - yscale_expected).abs().max().item() - quant_diff = (q_out - q_expected).abs().max().item() + scale_error = (yscale_out - yscale_expected).abs().max().item() + quant_error = (q_out - q_expected).abs().max().item() print(f"Max residual error: {residual_error:.2e} (tol={residual_atol})") - print(f"Max dequant error: {dequant_error:.2e} (tol={dequant_tol})") - print(f"Max scale diff: {scale_diff:.2e} (tol={scale_tol})") - print(f"Max quant diff: {quant_diff}") + print(f"Max scale error: {scale_error:.2e} (tol={scale_tol})") + print(f"Max quant error: {quant_error}") - ok = residual_error < residual_atol and dequant_error < dequant_tol and scale_diff < scale_tol and quant_diff <= 1 + ok = residual_error < residual_atol and scale_error < scale_tol and quant_error <= 1 if ok: print("PASSED") else: @@ -866,10 +516,6 @@ def kernel_launch(): print(residual_expected[0, :5]) print("First row Residual Actual:") print(residual_out_ref[0, :5]) - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) print("First row Quant Expected:") print(q_expected[0, :8]) print("First row Quant Actual:") @@ -881,56 +527,383 @@ def kernel_launch(): return ok, flydsl_gpu_us -def test_rmsnorm_fused_add_dynamicquant(): - print("=" * 80) - print("Running FusedAdd RMSNorm DynamicQuant Tests") - print("=" * 80) - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] - failures = 0 +def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None): + x = input_dev.to(DTYPE_FP32) + gamma = gamma_dev.to(DTYPE_FP32) + normalized = (x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + EPS)) * gamma + if xscale_dev is not None: + normalized = normalized * xscale_dev.to(DTYPE_FP32) - for M, N, dtype in _get_rmsnorm_configs(): - ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=False) - if not ok: - failures += 1 + yscale = normalized.abs().amax(dim=1) / 127.0 + yscale = torch.where(yscale == 0, torch.ones_like(yscale), yscale) + q = torch.clamp(torch.trunc(normalized / yscale.unsqueeze(1)), -127, 127).to(torch.int8) + return q, yscale - if do_compare: - aiter_us = None - if maybe_enable_aiter(): - aiter_us = _bench_aiter_fused_add_rmsnorm_quant(M, N, dtype, is_smooth=False) - perf_rows.append( - PerfRow( - op="rmsnorm_add_dq", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, - ) - ) - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") - else: +def _reference_fused_add_rmsnorm(input_dev, residual_in_dev, gamma_dev): + added = input_dev + residual_in_dev + added_fp32 = added.to(DTYPE_FP32) + gamma = gamma_dev.to(DTYPE_FP32) + expected = (added_fp32 / torch.sqrt((added_fp32 * added_fp32).mean(dim=1, keepdim=True) + EPS)) * gamma + return added_fp32, expected + + +def _reference_fused_add_rmsnorm_quant( + input_dev, + residual_in_dev, + gamma_dev, + *, + xscale_dev=None, +): + added = input_dev + residual_in_dev + residual_expected = added.to(DTYPE_FP32) + q, yscale = _reference_rmsnorm_quant( + added, + gamma_dev, + xscale_dev=xscale_dev, + ) + return residual_expected, q, yscale + + +def _bench_aiter_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + torch_dtype = _torch_dtype(dtype) + + try: + if is_smooth: + from aiter.ops.triton.normalization.rmsnorm import ( + rmsnorm2d_fwd_with_smoothquant as aiter_rmsnorm_quant, + ) + else: + from aiter.ops.triton.normalization.rmsnorm import ( + rmsnorm2d_fwd_with_dynamicquant as aiter_rmsnorm_quant, + ) + except Exception as e: + print(f"[Perf] AIter rmsnorm {mode} skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + y = torch.empty((M, N), dtype=torch.int8, device="cuda") + yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") + + if is_smooth: + xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + + def run_aiter(): + aiter_rmsnorm_quant(y, x, xscale, yscale, w, EPS) + + else: + + def run_aiter(): + aiter_rmsnorm_quant(y, x, yscale, w, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter rmsnorm {mode} gpu: {aiter_us:.1f} us") + return aiter_us + + +def _bench_aiter_fused_add_rmsnorm(M: int, N: int, dtype: str): + torch_dtype = _torch_dtype(dtype) + + try: + from aiter.ops.triton.normalization.rmsnorm import ( + rmsnorm2d_fwd_with_add as aiter_fused_add_rmsnorm, + ) + except Exception as e: + print(f"[Perf] AIter fused_add rmsnorm skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual_in = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + out = torch.empty((M, N), device="cuda", dtype=torch_dtype) + residual_out = torch.empty((M, N), device="cuda", dtype=torch_dtype) + + def run_aiter(): + aiter_fused_add_rmsnorm(out, x, residual_in, residual_out, w, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter fused_add rmsnorm gpu: {aiter_us:.1f} us") + return aiter_us + + +def _bench_aiter_fused_add_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + torch_dtype = _torch_dtype(dtype) + + try: + if is_smooth: + from aiter.ops.triton.normalization.rmsnorm import ( + rmsnorm2d_fwd_with_add_smoothquant as aiter_fused_add_rmsnorm_quant, + ) + else: + from aiter.ops.triton.normalization.rmsnorm import ( + rmsnorm2d_fwd_with_add_dynamicquant as aiter_fused_add_rmsnorm_quant, + ) + except Exception as e: + print(f"[Perf] AIter fused_add rmsnorm {mode} skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual_in = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + y = torch.empty((M, N), dtype=torch.int8, device="cuda") + residual_out = torch.empty((M, N), device="cuda", dtype=torch_dtype) + yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") + + if is_smooth: + xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + + def run_aiter(): + aiter_fused_add_rmsnorm_quant(y, x, residual_in, residual_out, xscale, yscale, w, EPS) + + else: + + def run_aiter(): + aiter_fused_add_rmsnorm_quant(y, x, residual_in, residual_out, yscale, w, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter fused_add rmsnorm {mode} gpu: {aiter_us:.1f} us") + return aiter_us + + +def test_rmsnorm(): + print("=" * 80) + print("Running RMSNorm Tests") + print("=" * 80) + + configs = _get_rmsnorm_configs() + + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + + failures = 0 + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_test(M, N, dtype) + if not ok: + failures += 1 + + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + try: + from aiter.ops.triton.rmsnorm import rms_norm as aiter_rms_norm + + torch_dtype = _torch_dtype(dtype) + x = torch.randn((M, N), device="cuda", dtype=torch_dtype) + w = torch.rand((N,), device="cuda", dtype=torch_dtype) + + def run_aiter(): + aiter_rms_norm(x, w, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter rmsnorm gpu: {aiter_us:.1f} us") + except Exception as e: + print(f"[Perf] AIter rmsnorm skipped: {type(e).__name__}: {e!r}") + + perf_rows.append( + PerfRow(op="rmsnorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us) + ) + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: print(f"{failures} TESTS FAILED") print("=" * 80) if do_compare and perf_rows: print_perf_table(perf_rows) + # Ensure a non-zero exit code on failure for shell wrappers. if failures != 0: raise SystemExit(1) -def test_rmsnorm_fused_add_smoothquant(): +def test_rmsnorm_dynamicquant(): print("=" * 80) - print("Running FusedAdd RMSNorm SmoothQuant Tests") + print("Running RMSNorm DynamicQuant Tests") print("=" * 80) + configs = _get_rmsnorm_configs() + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] + failures = 0 + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False) + if not ok: + failures += 1 + + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_rmsnorm_quant(M, N, dtype, is_smooth=False) + + perf_rows.append( + PerfRow( + op="rmsnorm_dynamicquant", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + # Ensure a non-zero exit code on failure for shell wrappers. + if failures != 0: + raise SystemExit(1) + + +def test_rmsnorm_smoothquant(): + print("=" * 80) + print("Running RMSNorm SmoothQuant Tests") + print("=" * 80) + + configs = _get_rmsnorm_configs() + + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + + failures = 0 + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True) + if not ok: + failures += 1 + + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_rmsnorm_quant(M, N, dtype, is_smooth=True) + + perf_rows.append( + PerfRow( + op="rmsnorm_smoothquant", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + # Ensure a non-zero exit code on failure for shell wrappers. + if failures != 0: + raise SystemExit(1) + + +def test_fused_add_rmsnorm(): + print("=" * 80) + print("Running FusedAdd RMSNorm Tests") + print("=" * 80) + + configs = _get_rmsnorm_configs() + + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] - for M, N, dtype in _get_rmsnorm_configs(): + failures = 0 + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_fused_add_test(M, N, dtype) + if not ok: + failures += 1 + + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_fused_add_rmsnorm(M, N, dtype) + perf_rows.append( + PerfRow( + op="rmsnorm_add", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + # Ensure a non-zero exit code on failure for shell wrappers. + if failures != 0: + raise SystemExit(1) + + +def test_fused_add_rmsnorm_dynamicquant(): + print("=" * 80) + print("Running FusedAdd RMSNorm DynamicQuant Tests") + print("=" * 80) + + configs = _get_rmsnorm_configs() + + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + + failures = 0 + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=False) + if not ok: + failures += 1 + + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_fused_add_rmsnorm_quant(M, N, dtype, is_smooth=False) + perf_rows.append( + PerfRow( + op="rmsnorm_add_dynamicquant", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) + + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + if failures != 0: + raise SystemExit(1) + + +def test_fused_add_rmsnorm_smoothquant(): + print("=" * 80) + print("Running FusedAdd RMSNorm SmoothQuant Tests") + print("=" * 80) + + configs = _get_rmsnorm_configs() + + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + + failures = 0 + for M, N, dtype in configs: ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=True) if not ok: failures += 1 @@ -941,7 +914,7 @@ def test_rmsnorm_fused_add_smoothquant(): aiter_us = _bench_aiter_fused_add_rmsnorm_quant(M, N, dtype, is_smooth=True) perf_rows.append( PerfRow( - op="rmsnorm_add_sq", + op="rmsnorm_add_smoothquant", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, @@ -962,4 +935,9 @@ def test_rmsnorm_fused_add_smoothquant(): if __name__ == "__main__": - test_all() + test_rmsnorm() + test_rmsnorm_dynamicquant() + test_rmsnorm_smoothquant() + test_fused_add_rmsnorm() + test_fused_add_rmsnorm_dynamicquant() + test_fused_add_rmsnorm_smoothquant() From 2b28750a4dcbb096991106268f551bb4a80a6367 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 4 Jun 2026 16:10:06 +0800 Subject: [PATCH 02/17] add vector fast paths for layernorm variants --- kernels/layernorm_kernel.py | 617 +++++++++++++++++++++++++----------- 1 file changed, 425 insertions(+), 192 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index ffc3530ad..8334aebba 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -17,6 +17,7 @@ import flydsl.expr as fx from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath +from flydsl.expr.typing import Vector as Vec from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.kernels_common import dtype_to_elem_type, get_warp_size @@ -32,6 +33,26 @@ VEC_ALIGN = 16 +def _to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): + if const_expr(dtype_str == "bf16"): + if const_expr(use_hw_cvt_bf16): + return y.to(elem_dtype) + u = y.bitcast(fx.Uint32) + upper = u >> 16 + lsb = upper & 1 + bias = lsb + 0x7FFF + u_round = y.bitcast(fx.Uint32) + bias + bf16_bits = u_round >> 16 + even = bf16_bits.shuffle(bf16_bits, [0, 2, 4, 6]) + odd = bf16_bits.shuffle(bf16_bits, [1, 3, 5, 7]) + odd_sh = odd << 16 + packed = even | odd_sh + return packed.bitcast(elem_dtype) + if const_expr(dtype_str == "f32"): + return y + return y.to(elem_dtype) + + def build_layernorm_module(M: int, N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -329,6 +350,9 @@ def _quant_dtype_max(dtype_str: str) -> float: def build_fused_add_layernorm_module(M: int, N: int, dtype_str: str): + arch = get_hip_arch() + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") + RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 @@ -404,82 +428,154 @@ def compute_mean_rstd(sum_val, sumsq_val): var = (var < 0.0).select(0.0, var) return mean, fmath.rsqrt(var + eps_c, fastmath=fm_fast) - Input_buf = fx.rocdl.make_buffer_tensor(Input) - ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) - Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) - Beta_buf = fx.rocdl.make_buffer_tensor(Beta) - Output_buf = fx.rocdl.make_buffer_tensor(Output) - ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) - - row_in = fx.slice(Input_buf, (bid, None)) - row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) - row_out = fx.slice(Output_buf, (bid, None)) - row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) + # ================================================================== + # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # ================================================================== + if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): + num_tiles_py = 4 + c_zero_f = fx.Float32(0.0) + thread_sum = c_zero_f + thread_sumsq = c_zero_f + added_local = [] - in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - - def _load_scalar(divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy_atom_call(copy_atom_s, view, r) - return fx.memref_load_vec(r)[0] - - def _store_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, elem_dtype(val), elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_s, r, view) + Input_buf = fx.rocdl.make_buffer_tensor(Input) + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) - c_zero_f = fx.Float32(0.0) - thread_sum = c_zero_f - thread_sumsq = c_zero_f - - for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(in_div, idx_safe) - r_e = _load_scalar(residual_in_div, idx_safe) - x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) - added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) - added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) - added_safe = is_valid.select(added, c_zero_f) - thread_sum = thread_sum + added_safe - thread_sumsq = thread_sumsq + is_valid.select(added * added, c_zero_f) - if idx < N: - _store_scalar(residual_out_div, idx, added_e) - - sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) - mean, rstd = compute_mean_rstd(sum_val, sumsq_val) - - for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base_idx_int - if idx < N: - added_e = _load_scalar(residual_out_div, idx) - g_e = _load_scalar(gamma_div, idx) - b_e = _load_scalar(beta_div, idx) - added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) + row_in = fx.slice(Input_buf, (bid, None)) + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) + + in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(VEC_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) + + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + + def _load_vec(div_tensor, idx): + r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) + fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) + return fx.memref_load_vec(r) + + def _store_vec(val, div_tensor, idx): + r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) + fx.memref_store_vec(val, r) + fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) + + # Pass 1: add residual, cache/store it, and accumulate sum/sumsq. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + x = _load_vec(in_div, idx).to(fx.Float32) + residual = _load_vec(residual_in_div, idx).to(fx.Float32) + added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) + added_local.append(added_e) + added = added_e.to(fx.Float32) + added2 = added * added + thread_sum = thread_sum + added.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = thread_sumsq + added2.reduce(ReductionOp.ADD, fastmath=fm_fast) + _store_vec(added_e, residual_out_div, idx) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean, rstd = compute_mean_rstd(sum_val, sumsq_val) + + # Pass 2: normalize + affine + store, reusing cached added values. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + added = added_local[tile_i].to(fx.Float32) + g = _load_vec(gamma_div, idx).to(fx.Float32) + b = _load_vec(beta_div, idx).to(fx.Float32) y = (added - mean) * rstd y = y * g + b - if const_expr(dtype_str == "f32"): - y_e = y - else: - y_e = y.to(elem_dtype) - _store_scalar(out_div, idx, y_e) + y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(y_e, out_div, idx) + + else: + # ============================================================== + # Generic path: scalar 2-pass implementation for arbitrary N + # ============================================================== + Input_buf = fx.rocdl.make_buffer_tensor(Input) + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) + + row_in = fx.slice(Input_buf, (bid, None)) + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) + + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + + in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) + + def _load_scalar(divided_tensor, index): + view = fx.slice(divided_tensor, (None, index)) + r = fx.make_rmem_tensor(1, elem_dtype) + fx.copy_atom_call(copy_atom_s, view, r) + return fx.memref_load_vec(r)[0] + + def _store_scalar(divided_tensor, index, val): + r = fx.make_rmem_tensor(1, elem_dtype) + ts = full(1, elem_dtype(val), elem_dtype) + fx.memref_store_vec(ts, r) + view = fx.slice(divided_tensor, (None, index)) + fx.copy_atom_call(copy_atom_s, r, view) + + c_zero_f = fx.Float32(0.0) + thread_sum = c_zero_f + thread_sumsq = c_zero_f + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(in_div, idx_safe) + r_e = _load_scalar(residual_in_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) + added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) + added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) + added_safe = is_valid.select(added, c_zero_f) + thread_sum = thread_sum + added_safe + thread_sumsq = thread_sumsq + is_valid.select(added * added, c_zero_f) + if idx < N: + _store_scalar(residual_out_div, idx, added_e) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean, rstd = compute_mean_rstd(sum_val, sumsq_val) + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + if idx < N: + added_e = _load_scalar(residual_out_div, idx) + g_e = _load_scalar(gamma_div, idx) + b_e = _load_scalar(beta_div, idx) + added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) + y = (added - mean) * rstd + y = y * g + b + if const_expr(dtype_str == "f32"): + y_e = y + else: + y_e = y.to(elem_dtype) + _store_scalar(out_div, idx, y_e) @flyc.jit def launch_fused_add_layernorm( @@ -511,6 +607,9 @@ def _build_layernorm_quant_module( is_fused_add: bool, quant_dtype_str: str = "i8", ): + arch = get_hip_arch() + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") + RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 quant_dtype_max = _quant_dtype_max(quant_dtype_str) @@ -628,145 +727,279 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) - Input_buf = fx.rocdl.make_buffer_tensor(Input) - Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) - Beta_buf = fx.rocdl.make_buffer_tensor(Beta) - Output_buf = fx.rocdl.make_buffer_tensor(Output) - if const_expr(is_fused_add): - ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) - ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) - if const_expr(is_smooth): - XScale_buf = fx.rocdl.make_buffer_tensor(XScale) + # ================================================================== + # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # ================================================================== + if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): + num_tiles_py = 4 + quant_half_width = VEC_WIDTH // 2 + xscale_vec_width = 4 + abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) - row_in = fx.slice(Input_buf, (bid, None)) - row_out = fx.slice(Output_buf, (bid, None)) - if const_expr(is_fused_add): - row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) - row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) + Input_buf = fx.rocdl.make_buffer_tensor(Input) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + if const_expr(is_fused_add): + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) + if const_expr(is_smooth): + XScale_buf = fx.rocdl.make_buffer_tensor(XScale) - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) - copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) + row_in = fx.slice(Input_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) + if const_expr(is_fused_add): + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - if const_expr(is_fused_add): - residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) - - def _load_scalar(divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy_atom_call(copy_atom_s, view, r) - return fx.memref_load_vec(r)[0] - - def _store_elem_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, elem_dtype(val), elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_s, r, view) + in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div_q = fx.logical_divide(row_out, fx.make_layout(quant_half_width, 1)) + if const_expr(is_fused_add): + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(VEC_WIDTH, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) + if const_expr(is_smooth): + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(xscale_vec_width, 1)) - def _store_quant_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, quant_dtype) - ts = full(1, quant_dtype(val), quant_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_qs, r, view) + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + copy_atom_q = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 8) + if const_expr(is_smooth): + copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32) + + def _load_elem_vec(div_tensor, idx): + r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) + fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) + return fx.memref_load_vec(r) + + def _store_elem_vec(val, div_tensor, idx): + r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) + fx.memref_store_vec(val, r) + fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) + + def _load_vec_xscale(idx): + r = fx.make_rmem_tensor(xscale_vec_width, fx.Float32) + fx.copy_atom_call(copy_atom_xs, fx.slice(xscale_div, (None, idx)), r) + return fx.memref_load_vec(r) + + def _load_xscale_vec(idx): + s_lo = _load_vec_xscale(idx * 2) + s_hi = _load_vec_xscale(idx * 2 + 1) + return Vec(s_lo).shuffle(Vec(s_hi), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() - def _abs_scalar(val): - is_neg = val < c_zero_f - neg_val = c_zero_f - val - return is_neg.select(neg_val, val) + def _store_quant_vec(val, div_tensor, idx): + r = fx.make_rmem_tensor(quant_half_width, quant_dtype) + fx.memref_store_vec(val, r) + fx.copy_atom_call(copy_atom_q, r, fx.slice(div_tensor, (None, idx))) + + thread_sum = c_zero_f + thread_sumsq = c_zero_f + norm_input_local = [] - def _load_base_input_value(index): - x_e = _load_scalar(in_div, index) - return x_e if dtype_str == "f32" else x_e.to(fx.Float32) + # Pass 1: prepare normalization input and accumulate sum/sumsq. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + if const_expr(is_fused_add): + x = _load_elem_vec(in_div, idx).to(fx.Float32) + residual = _load_elem_vec(residual_in_div, idx).to(fx.Float32) + added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) + norm_input_local.append(added_e) + x_norm = added_e.to(fx.Float32) + _store_elem_vec(added_e, residual_out_div, idx) + else: + x_e = _load_elem_vec(in_div, idx) + norm_input_local.append(x_e) + x_norm = x_e.to(fx.Float32) + x2 = x_norm * x_norm + thread_sum = thread_sum + x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = thread_sumsq + x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean = sum_val / n_float + var = sumsq_val / n_float - mean * mean + var = (var < c_zero_f).select(c_zero_f, var) + rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + + thread_row_max = c_zero_f + y_local = [] + + # Pass 2: affine (+ optional smooth scale), cache y, accumulate row max. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + x = norm_input_local[tile_i].to(fx.Float32) + g = _load_elem_vec(gamma_div, idx).to(fx.Float32) + b = _load_elem_vec(beta_div, idx).to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + if const_expr(is_smooth): + y = y * _load_xscale_vec(idx) + y_local.append(y) + y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) + tile_max = y_abs.reduce(ReductionOp.MAX) + thread_row_max = thread_row_max.maximumf(tile_max) + + row_max = block_reduce_max(thread_row_max) + scale = row_max / c_dtype_max + final_scale = (scale == c_zero_f).select(c_one_f, scale) + + if tid == 0: + _store_yscale(bid, final_scale) + + inv_scale = c_one_f / final_scale + + # Pass 3: quantize + store using per-row scale. + for tile_i in range_constexpr(num_tiles_py): + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 + _store_quant_vec(q_lo, out_div_q, out_idx) + _store_quant_vec(q_hi, out_div_q, out_idx + 1) + + else: + # ============================================================== + # Generic path: scalar 3-pass implementation for arbitrary N + # ============================================================== + Input_buf = fx.rocdl.make_buffer_tensor(Input) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + if const_expr(is_fused_add): + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) + if const_expr(is_smooth): + XScale_buf = fx.rocdl.make_buffer_tensor(XScale) - def _load_norm_input_value(index): + row_in = fx.slice(Input_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) if const_expr(is_fused_add): - added_e = _load_scalar(residual_out_div, index) - return added_e if dtype_str == "f32" else added_e.to(fx.Float32) - return _load_base_input_value(index) + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) - thread_sum = c_zero_f - thread_sumsq = c_zero_f + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) - for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) + in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) if const_expr(is_fused_add): - x = _load_base_input_value(idx_safe) - r_e = _load_scalar(residual_in_div, idx_safe) - residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) - added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) - if idx < N: - _store_elem_scalar(residual_out_div, idx, added_e) - x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) - else: - x = _load_norm_input_value(idx_safe) - x2 = x * x - thread_sum = thread_sum + is_valid.select(x, c_zero_f) - thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) - - sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) - mean = sum_val / n_float - var = sumsq_val / n_float - mean * mean - var = (var < c_zero_f).select(c_zero_f, var) - rstd = (var + eps_c).rsqrt(fastmath=fm_fast) - - thread_row_max = c_zero_f - for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - x = _load_norm_input_value(idx_safe) - g_e = _load_scalar(gamma_div, idx_safe) - b_e = _load_scalar(beta_div, idx_safe) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) - b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) - y = (x - mean) * rstd - y = y * g + b + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) if const_expr(is_smooth): - s_e = _load_scalar(xscale_div, idx_safe) - s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) - y = y * s - y_abs = _abs_scalar(y) - thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f)) - - row_max = block_reduce_max(thread_row_max) - scale = row_max / c_dtype_max - final_scale = (scale == c_zero_f).select(c_one_f, scale) - - if tid == 0: - _store_yscale(bid, final_scale) - - inv_scale = c_one_f / final_scale - - for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): - idx = tid + base_idx_int - if idx < N: - x = _load_norm_input_value(idx) - g_e = _load_scalar(gamma_div, idx) - b_e = _load_scalar(beta_div, idx) + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) + + def _load_scalar(divided_tensor, index): + view = fx.slice(divided_tensor, (None, index)) + r = fx.make_rmem_tensor(1, elem_dtype) + fx.copy_atom_call(copy_atom_s, view, r) + return fx.memref_load_vec(r)[0] + + def _store_elem_scalar(divided_tensor, index, val): + r = fx.make_rmem_tensor(1, elem_dtype) + ts = full(1, elem_dtype(val), elem_dtype) + fx.memref_store_vec(ts, r) + view = fx.slice(divided_tensor, (None, index)) + fx.copy_atom_call(copy_atom_s, r, view) + + def _store_quant_scalar(divided_tensor, index, val): + r = fx.make_rmem_tensor(1, quant_dtype) + ts = full(1, quant_dtype(val), quant_dtype) + fx.memref_store_vec(ts, r) + view = fx.slice(divided_tensor, (None, index)) + fx.copy_atom_call(copy_atom_qs, r, view) + + def _abs_scalar(val): + is_neg = val < c_zero_f + neg_val = c_zero_f - val + return is_neg.select(neg_val, val) + + def _load_base_input_value(index): + x_e = _load_scalar(in_div, index) + return x_e if dtype_str == "f32" else x_e.to(fx.Float32) + + def _load_norm_input_value(index): + if const_expr(is_fused_add): + added_e = _load_scalar(residual_out_div, index) + return added_e if dtype_str == "f32" else added_e.to(fx.Float32) + return _load_base_input_value(index) + + thread_sum = c_zero_f + thread_sumsq = c_zero_f + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + if const_expr(is_fused_add): + x = _load_base_input_value(idx_safe) + r_e = _load_scalar(residual_in_div, idx_safe) + residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) + added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) + if idx < N: + _store_elem_scalar(residual_out_div, idx, added_e) + x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) + else: + x = _load_norm_input_value(idx_safe) + x2 = x * x + thread_sum = thread_sum + is_valid.select(x, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean = sum_val / n_float + var = sumsq_val / n_float - mean * mean + var = (var < c_zero_f).select(c_zero_f, var) + rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + + thread_row_max = c_zero_f + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x = _load_norm_input_value(idx_safe) + g_e = _load_scalar(gamma_div, idx_safe) + b_e = _load_scalar(beta_div, idx_safe) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(xscale_div, idx) + s_e = _load_scalar(xscale_div, idx_safe) s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s - q = y * inv_scale - q_i8 = q.to(quant_dtype) - _store_quant_scalar(out_div, idx, q_i8) + y_abs = _abs_scalar(y) + thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f)) + + row_max = block_reduce_max(thread_row_max) + scale = row_max / c_dtype_max + final_scale = (scale == c_zero_f).select(c_one_f, scale) + + if tid == 0: + _store_yscale(bid, final_scale) + + inv_scale = c_one_f / final_scale + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + if idx < N: + x = _load_norm_input_value(idx) + g_e = _load_scalar(gamma_div, idx) + b_e = _load_scalar(beta_div, idx) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + if const_expr(is_smooth): + s_e = _load_scalar(xscale_div, idx) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) + y = y * s + q = y * inv_scale + q_i8 = q.to(quant_dtype) + _store_quant_scalar(out_div, idx, q_i8) if is_fused_add: if is_smooth: From d670743ae3129cef8cc6a2b98518854539011d12 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 4 Jun 2026 17:32:26 +0800 Subject: [PATCH 03/17] refactor: extract scalar/vector load-store helpers to reduce duplicated codes --- kernels/layernorm_kernel.py | 294 +++++++++++++----------------------- 1 file changed, 105 insertions(+), 189 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index 8334aebba..a0590ef23 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -29,8 +29,49 @@ BLOCK_THREADS = 256 WARP_SIZE = get_warp_size() VEC_WIDTH = 8 -USE_NONTEMPORAL = True -VEC_ALIGN = 16 + + +# ── Shared-memory allocation for block reductions ───────────────────── +def _make_reduction_storage(red_slots: int): + @fx.struct + class SharedStorage: + s_red: fx.Array[fx.Float32, red_slots, 16] + s_red2: fx.Array[fx.Float32, red_slots, 16] + + return SharedStorage + + +def _load_scalar(copy_atom, elem_dtype, divided_tensor, index): + view = fx.slice(divided_tensor, (None, index)) + r = fx.make_rmem_tensor(1, elem_dtype) + fx.copy_atom_call(copy_atom, view, r) + return fx.memref_load_vec(r)[0] + + +def _store_scalar(copy_atom, elem_dtype, store_dtype, divided_tensor, index, val): + r = fx.make_rmem_tensor(1, elem_dtype) + ts = full(1, store_dtype(val), store_dtype) + fx.memref_store_vec(ts, r) + view = fx.slice(divided_tensor, (None, index)) + fx.copy_atom_call(copy_atom, r, view) + + +def _load_vec(copy_atom, vec_width, elem_dtype, div_tensor, idx): + r = fx.make_rmem_tensor(vec_width, elem_dtype) + fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) + return fx.memref_load_vec(r) + + +def _store_vec(copy_atom, vec_width, elem_dtype, val, div_tensor, idx): + r = fx.make_rmem_tensor(vec_width, elem_dtype) + fx.memref_store_vec(val, r) + fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) + + +def _to_elem_scalar(dtype_str: str, elem_dtype, y): + if const_expr(dtype_str == "f32"): + return y + return y.to(elem_dtype) def _to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): @@ -53,19 +94,21 @@ def _to_elem_vec(dtype_str: str, elem_dtype, use_hw_cvt_bf16: bool, y): return y.to(elem_dtype) +def _store_yscale(scale_copy_atom, yscale_div, index, val): + r = fx.make_rmem_tensor(1, fx.Float32) + ts = full(1, fx.Float32(val), fx.Float32) + fx.memref_store_vec(ts, r) + fx.copy_atom_call(scale_copy_atom, r, fx.slice(yscale_div, (None, index))) + + def build_layernorm_module(M: int, N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) - elem_bits = 32 if dtype_str == "f32" else 16 - # ── Shared-memory allocation for block reductions ───────────────────── - @fx.struct - class SharedStorage: - s_sum: fx.Array[fx.Float32, RED_SLOTS, 16] - s_sumsq: fx.Array[fx.Float32, RED_SLOTS, 16] + SharedStorage = _make_reduction_storage(RED_SLOTS) # ── GPU kernel ──────────────────────────────────────────────────────── @flyc.kernel @@ -167,20 +210,10 @@ def compute_mean_rstd(sum_val, sumsq_val): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) - def _load_vec(div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) - return fx.memref_load_vec(r) - - def _store_vec(val, div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.memref_store_vec(val, r) - fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) - # ── Pass 1: load input, accumulate sum / sumsq ─────────────── for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS - vec = _load_vec(in_div, idx) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) in_local.append(vec) x = vec.to(fx.Float32) @@ -193,8 +226,8 @@ def _store_vec(val, div_tensor, idx): sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) - g_cur = _load_vec(gamma_div, tid).to(fx.Float32) - b_cur = _load_vec(beta_div, tid).to(fx.Float32) + g_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, tid).to(fx.Float32) + b_cur = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, tid).to(fx.Float32) # ── Pass 2: normalize + affine + store ─────────────────────── for tile_i in range_constexpr(num_tiles_py): @@ -202,8 +235,8 @@ def _store_vec(val, div_tensor, idx): b_next = b_cur if const_expr(tile_i + 1 < num_tiles_py): next_idx = tid + (tile_i + 1) * BLOCK_THREADS - g_next = _load_vec(gamma_div, next_idx).to(fx.Float32) - b_next = _load_vec(beta_div, next_idx).to(fx.Float32) + g_next = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, next_idx).to(fx.Float32) + b_next = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, next_idx).to(fx.Float32) else: g_next = g_cur b_next = b_cur @@ -212,29 +245,9 @@ def _store_vec(val, div_tensor, idx): y = (x - mean) * rstd y = y * g_cur + b_cur - out_e = y.to(elem_dtype) - if const_expr(dtype_str == "bf16"): - if const_expr(USE_HW_CVT_PK_BF16_F32): - out_e = y.to(elem_dtype) - else: - u = y.bitcast(fx.Uint32) - upper = u >> 16 - lsb = upper & 1 - bias = lsb + 0x7FFF - u_round = y.bitcast(fx.Uint32) + bias - bf16_bits = u_round >> 16 - even = bf16_bits.shuffle(bf16_bits, [0, 2, 4, 6]) - odd = bf16_bits.shuffle(bf16_bits, [1, 3, 5, 7]) - odd_sh = odd << 16 - packed = even | odd_sh - out_e = packed.bitcast(elem_dtype) - elif const_expr(dtype_str == "f32"): - out_e = y - else: - out_e = y.to(elem_dtype) - + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) out_idx = tid + tile_i * BLOCK_THREADS - _store_vec(out_e, out_div, out_idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, out_idx) g_cur = g_next b_cur = b_next @@ -265,25 +278,12 @@ def _store_vec(val, div_tensor, idx): beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - def _load_scalar(divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy_atom_call(copy_atom_s, view, r) - return fx.memref_load_vec(r)[0] - - def _store_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, elem_dtype(val), elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_s, r, view) - # ── Pass 1: sum + sumsq ────────────────────────────────────── for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(row_div, idx_safe) + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x x_safe = is_valid.select(x, c_zero_f) @@ -298,9 +298,9 @@ def _store_scalar(divided_tensor, index, val): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - x_e = _load_scalar(row_div, idx) - g_e = _load_scalar(gamma_div, idx) - b_e = _load_scalar(beta_div, idx) + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) @@ -308,14 +308,8 @@ def _store_scalar(divided_tensor, index, val): norm = diff * rstd scaled = norm * g y = scaled + b - y_e = y - if const_expr(dtype_str == "bf16"): - y_e = y.to(elem_dtype) - elif const_expr(dtype_str == "f32"): - y_e = y - else: - y_e = y.to(elem_dtype) - _store_scalar(out_div, idx, y_e) + y_e = _to_elem_scalar(dtype_str, elem_dtype, y) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) # ── JIT host launcher ───────────────────────────────────────────────── @flyc.jit @@ -356,10 +350,7 @@ def build_fused_add_layernorm_module(M: int, N: int, dtype_str: str): RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 - @fx.struct - class SharedStorage: - s_sum: fx.Array[fx.Float32, RED_SLOTS, 16] - s_sumsq: fx.Array[fx.Float32, RED_SLOTS, 16] + SharedStorage = _make_reduction_storage(RED_SLOTS) @flyc.kernel def fused_add_layernorm_kernel( @@ -459,28 +450,18 @@ def compute_mean_rstd(sum_val, sumsq_val): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) - def _load_vec(div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) - return fx.memref_load_vec(r) - - def _store_vec(val, div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.memref_store_vec(val, r) - fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) - # Pass 1: add residual, cache/store it, and accumulate sum/sumsq. for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS - x = _load_vec(in_div, idx).to(fx.Float32) - residual = _load_vec(residual_in_div, idx).to(fx.Float32) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) added_local.append(added_e) added = added_e.to(fx.Float32) added2 = added * added thread_sum = thread_sum + added.reduce(ReductionOp.ADD, fastmath=fm_fast) thread_sumsq = thread_sumsq + added2.reduce(ReductionOp.ADD, fastmath=fm_fast) - _store_vec(added_e, residual_out_div, idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) @@ -489,12 +470,12 @@ def _store_vec(val, div_tensor, idx): for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS added = added_local[tile_i].to(fx.Float32) - g = _load_vec(gamma_div, idx).to(fx.Float32) - b = _load_vec(beta_div, idx).to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) y = (added - mean) * rstd y = y * g + b y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) - _store_vec(y_e, out_div, idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, idx) else: # ============================================================== @@ -524,19 +505,6 @@ def _store_vec(val, div_tensor, idx): out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) - def _load_scalar(divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy_atom_call(copy_atom_s, view, r) - return fx.memref_load_vec(r)[0] - - def _store_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, elem_dtype(val), elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_s, r, view) - c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -545,17 +513,17 @@ def _store_scalar(divided_tensor, index, val): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(in_div, idx_safe) - r_e = _load_scalar(residual_in_div, idx_safe) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) + r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) - added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) + added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) added_safe = is_valid.select(added, c_zero_f) thread_sum = thread_sum + added_safe thread_sumsq = thread_sumsq + is_valid.select(added * added, c_zero_f) if idx < N: - _store_scalar(residual_out_div, idx, added_e) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) mean, rstd = compute_mean_rstd(sum_val, sumsq_val) @@ -563,19 +531,16 @@ def _store_scalar(divided_tensor, index, val): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - added_e = _load_scalar(residual_out_div, idx) - g_e = _load_scalar(gamma_div, idx) - b_e = _load_scalar(beta_div, idx) + added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (added - mean) * rstd y = y * g + b - if const_expr(dtype_str == "f32"): - y_e = y - else: - y_e = y.to(elem_dtype) - _store_scalar(out_div, idx, y_e) + y_e = _to_elem_scalar(dtype_str, elem_dtype, y) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) @flyc.jit def launch_fused_add_layernorm( @@ -614,10 +579,7 @@ def _build_layernorm_quant_module( elem_bits = 32 if dtype_str == "f32" else 16 quant_dtype_max = _quant_dtype_max(quant_dtype_str) - @fx.struct - class SharedStorage: - s_sum: fx.Array[fx.Float32, RED_SLOTS, 16] - s_sumsq: fx.Array[fx.Float32, RED_SLOTS, 16] + SharedStorage = _make_reduction_storage(RED_SLOTS) @flyc.kernel def layernorm_quant_kernel( @@ -652,12 +614,6 @@ def layernorm_quant_kernel( yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) - def _store_yscale(index, val): - r = fx.make_rmem_tensor(1, fx.Float32) - ts = full(1, fx.Float32(val), fx.Float32) - fx.memref_store_vec(ts, r) - fx.copy_atom_call(scale_copy_atom, r, fx.slice(yscale_div, (None, index))) - def wave_reduce_add(x): w = x for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): @@ -767,31 +723,11 @@ def block_reduce_max(val): if const_expr(is_smooth): copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32) - def _load_elem_vec(div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, idx)), r) - return fx.memref_load_vec(r) - - def _store_elem_vec(val, div_tensor, idx): - r = fx.make_rmem_tensor(VEC_WIDTH, elem_dtype) - fx.memref_store_vec(val, r) - fx.copy_atom_call(copy_atom, r, fx.slice(div_tensor, (None, idx))) - - def _load_vec_xscale(idx): - r = fx.make_rmem_tensor(xscale_vec_width, fx.Float32) - fx.copy_atom_call(copy_atom_xs, fx.slice(xscale_div, (None, idx)), r) - return fx.memref_load_vec(r) - def _load_xscale_vec(idx): - s_lo = _load_vec_xscale(idx * 2) - s_hi = _load_vec_xscale(idx * 2 + 1) + s_lo = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2) + s_hi = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2 + 1) return Vec(s_lo).shuffle(Vec(s_hi), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() - def _store_quant_vec(val, div_tensor, idx): - r = fx.make_rmem_tensor(quant_half_width, quant_dtype) - fx.memref_store_vec(val, r) - fx.copy_atom_call(copy_atom_q, r, fx.slice(div_tensor, (None, idx))) - thread_sum = c_zero_f thread_sumsq = c_zero_f norm_input_local = [] @@ -800,14 +736,14 @@ def _store_quant_vec(val, div_tensor, idx): for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS if const_expr(is_fused_add): - x = _load_elem_vec(in_div, idx).to(fx.Float32) - residual = _load_elem_vec(residual_in_div, idx).to(fx.Float32) + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) norm_input_local.append(added_e) x_norm = added_e.to(fx.Float32) - _store_elem_vec(added_e, residual_out_div, idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) else: - x_e = _load_elem_vec(in_div, idx) + x_e = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) norm_input_local.append(x_e) x_norm = x_e.to(fx.Float32) x2 = x_norm * x_norm @@ -827,8 +763,8 @@ def _store_quant_vec(val, div_tensor, idx): for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS x = norm_input_local[tile_i].to(fx.Float32) - g = _load_elem_vec(gamma_div, idx).to(fx.Float32) - b = _load_elem_vec(beta_div, idx).to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): @@ -843,7 +779,7 @@ def _store_quant_vec(val, div_tensor, idx): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(bid, final_scale) + _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) inv_scale = c_one_f / final_scale @@ -854,8 +790,8 @@ def _store_quant_vec(val, div_tensor, idx): q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 - _store_quant_vec(q_lo, out_div_q, out_idx) - _store_quant_vec(q_hi, out_div_q, out_idx + 1) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) else: # ============================================================== @@ -893,38 +829,18 @@ def _store_quant_vec(val, div_tensor, idx): if const_expr(is_smooth): xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) - def _load_scalar(divided_tensor, index): - view = fx.slice(divided_tensor, (None, index)) - r = fx.make_rmem_tensor(1, elem_dtype) - fx.copy_atom_call(copy_atom_s, view, r) - return fx.memref_load_vec(r)[0] - - def _store_elem_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, elem_dtype) - ts = full(1, elem_dtype(val), elem_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_s, r, view) - - def _store_quant_scalar(divided_tensor, index, val): - r = fx.make_rmem_tensor(1, quant_dtype) - ts = full(1, quant_dtype(val), quant_dtype) - fx.memref_store_vec(ts, r) - view = fx.slice(divided_tensor, (None, index)) - fx.copy_atom_call(copy_atom_qs, r, view) - def _abs_scalar(val): is_neg = val < c_zero_f neg_val = c_zero_f - val return is_neg.select(neg_val, val) def _load_base_input_value(index): - x_e = _load_scalar(in_div, index) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, index) return x_e if dtype_str == "f32" else x_e.to(fx.Float32) def _load_norm_input_value(index): if const_expr(is_fused_add): - added_e = _load_scalar(residual_out_div, index) + added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, index) return added_e if dtype_str == "f32" else added_e.to(fx.Float32) return _load_base_input_value(index) @@ -937,11 +853,11 @@ def _load_norm_input_value(index): idx_safe = is_valid.select(idx, 0) if const_expr(is_fused_add): x = _load_base_input_value(idx_safe) - r_e = _load_scalar(residual_in_div, idx_safe) + r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) - added_e = (x + residual) if dtype_str == "f32" else (x + residual).to(elem_dtype) + added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) if idx < N: - _store_elem_scalar(residual_out_div, idx, added_e) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) else: x = _load_norm_input_value(idx_safe) @@ -961,14 +877,14 @@ def _load_norm_input_value(index): is_valid = idx < N idx_safe = is_valid.select(idx, 0) x = _load_norm_input_value(idx_safe) - g_e = _load_scalar(gamma_div, idx_safe) - b_e = _load_scalar(beta_div, idx_safe) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx_safe) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(xscale_div, idx_safe) + s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx_safe) s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) @@ -979,7 +895,7 @@ def _load_norm_input_value(index): final_scale = (scale == c_zero_f).select(c_one_f, scale) if tid == 0: - _store_yscale(bid, final_scale) + _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) inv_scale = c_one_f / final_scale @@ -987,19 +903,19 @@ def _load_norm_input_value(index): idx = tid + base_idx_int if idx < N: x = _load_norm_input_value(idx) - g_e = _load_scalar(gamma_div, idx) - b_e = _load_scalar(beta_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - s_e = _load_scalar(xscale_div, idx) + s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx) s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) - _store_quant_scalar(out_div, idx, q_i8) + _store_scalar(copy_atom_qs, quant_dtype, quant_dtype, out_div, idx, q_i8) if is_fused_add: if is_smooth: From f6c474cc3ec86a109f9934ccaeef6454b2dfdc15 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 12 Jun 2026 15:30:42 +0800 Subject: [PATCH 04/17] fix layernorm shared storage name error, use helpers to simplify tests --- kernels/layernorm_kernel.py | 4 +- tests/kernels/test_layernorm.py | 165 +++++++++----------------------- 2 files changed, 47 insertions(+), 122 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index a0590ef23..97aa81ed9 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -35,8 +35,8 @@ def _make_reduction_storage(red_slots: int): @fx.struct class SharedStorage: - s_red: fx.Array[fx.Float32, red_slots, 16] - s_red2: fx.Array[fx.Float32, red_slots, 16] + s_sum: fx.Array[fx.Float32, red_slots, 16] + s_sumsq: fx.Array[fx.Float32, red_slots, 16] return SharedStorage diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 35814ba6b..517f424bc 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -53,6 +53,39 @@ BENCH_ITERS = 100 +def _torch_dtype(dtype: str): + if dtype == "f32": + return DTYPE_FP32 + if dtype == "f16": + return DTYPE_FP16 + if dtype == "bf16": + return DTYPE_BF16 + raise ValueError(f"unsupported dtype: {dtype}") + + +def _get_layernorm_configs(): + shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() + if shapes_env: + configs = [] + for part in shapes_env.split(";"): + p = part.strip() + if not p: + continue + m_s, n_s, dt = [x.strip() for x in p.split(",")] + configs.append((int(m_s), int(n_s), dt)) + else: + configs = [ + (64, 256, "f32"), # Aligned + (128, 1024, "f32"), # Aligned + (32, 128, "f16"), # Aligned + (64, 2000, "f32"), # Unaligned (tail handling) + (16, 512, "bf16"), # BF16 + (1024, 8192, "bf16"), # BF16 + (32768, 8192, "bf16"), + ] + return configs + + def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting LayerNorm (M={M}, N={N}, dtype={dtype})") @@ -159,25 +192,7 @@ def test_layernorm(): print("Running LayerNorm Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -199,7 +214,7 @@ def test_layernorm(): x = torch.randn( (M, N), device="cuda", - dtype=DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32), + dtype=_torch_dtype(dtype), ) w = torch.rand((N,), device="cuda", dtype=x.dtype) b = torch.rand((N,), device="cuda", dtype=x.dtype) @@ -346,25 +361,7 @@ def test_fused_add_layernorm(): print("Running FusedAdd LayerNorm Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -383,7 +380,7 @@ def test_fused_add_layernorm(): try: from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add - torch_dtype = DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual_out = torch.empty_like(x) @@ -542,25 +539,7 @@ def test_layernorm_dynamicquant(): print("Running LayerNorm DynamicQuant Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -579,7 +558,7 @@ def test_layernorm_dynamicquant(): try: from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_dynamicquant - torch_dtype = DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() @@ -743,25 +722,7 @@ def test_layernorm_smoothquant(): print("Running LayerNorm SmoothQuant Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -780,7 +741,7 @@ def test_layernorm_smoothquant(): try: from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_smoothquant - torch_dtype = DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() @@ -957,25 +918,7 @@ def test_fused_add_layernorm_dynamicquant(): print("Running FusedAdd LayerNorm DynamicQuant Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -994,7 +937,7 @@ def test_fused_add_layernorm_dynamicquant(): try: from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add_dynamicquant - torch_dtype = DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual_out = torch.empty_like(x) @@ -1189,25 +1132,7 @@ def test_fused_add_layernorm_smoothquant(): print("Running FusedAdd LayerNorm SmoothQuant Tests") print("=" * 80) - shapes_env = os.environ.get("ROCDSL_LAYERNORM_SHAPES", "").strip() - if shapes_env: - configs = [] - for part in shapes_env.split(";"): - p = part.strip() - if not p: - continue - m_s, n_s, dt = [x.strip() for x in p.split(",")] - configs.append((int(m_s), int(n_s), dt)) - else: - configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), - ] + configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] @@ -1226,7 +1151,7 @@ def test_fused_add_layernorm_smoothquant(): try: from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add_smoothquant - torch_dtype = DTYPE_BF16 if dtype == "bf16" else (DTYPE_FP16 if dtype == "f16" else DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() residual_out = torch.empty_like(x) From d48c011ce3b2bc131421f385f1152128a54f1d9f Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 16 Jun 2026 15:39:16 +0800 Subject: [PATCH 05/17] fix norm xscale dtype issue and extract test helpers --- kernels/layernorm_kernel.py | 534 +++++++++++---- kernels/rmsnorm_kernel.py | 66 +- tests/kernels/test_layernorm.py | 1131 ++++++++++++------------------- tests/kernels/test_rmsnorm.py | 149 ++-- 4 files changed, 923 insertions(+), 957 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index 97aa81ed9..84fa49ed0 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -101,6 +101,18 @@ def _store_yscale(scale_copy_atom, yscale_div, index, val): fx.copy_atom_call(scale_copy_atom, r, fx.slice(yscale_div, (None, index))) +def _quant_dtype_to_elem_type(dtype_str: str): + if dtype_str in ("i8", "int8"): + return fx.Int8 + raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") + + +def _quant_dtype_max(dtype_str: str) -> float: + if dtype_str in ("i8", "int8"): + return 127.0 + raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") + + def build_layernorm_module(M: int, N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -331,18 +343,6 @@ def launch_layernorm( return launch_layernorm -def _quant_dtype_to_elem_type(dtype_str: str): - if dtype_str in ("i8", "int8"): - return fx.Int8 - raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") - - -def _quant_dtype_max(dtype_str: str) -> float: - if dtype_str in ("i8", "int8"): - return 127.0 - raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") - - def build_fused_add_layernorm_module(M: int, N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -569,12 +569,8 @@ def _build_layernorm_quant_module( dtype_str: str, *, is_smooth: bool, - is_fused_add: bool, quant_dtype_str: str = "i8", ): - arch = get_hip_arch() - USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") - RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 quant_dtype_max = _quant_dtype_max(quant_dtype_str) @@ -584,13 +580,11 @@ def _build_layernorm_quant_module( @flyc.kernel def layernorm_quant_kernel( Input: fx.Tensor, - ResidualIn: fx.Tensor, Gamma: fx.Tensor, Beta: fx.Tensor, XScale: fx.Tensor, YScale: fx.Tensor, Output: fx.Tensor, - ResidualOut: fx.Tensor, ): bid = fx.block_idx.x tid = fx.thread_idx.x @@ -689,44 +683,29 @@ def block_reduce_max(val): if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): num_tiles_py = 4 quant_half_width = VEC_WIDTH // 2 - xscale_vec_width = 4 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) Input_buf = fx.rocdl.make_buffer_tensor(Input) Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) Beta_buf = fx.rocdl.make_buffer_tensor(Beta) Output_buf = fx.rocdl.make_buffer_tensor(Output) - if const_expr(is_fused_add): - ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) - ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) if const_expr(is_smooth): XScale_buf = fx.rocdl.make_buffer_tensor(XScale) row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) - if const_expr(is_fused_add): - row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) - row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) beta_div = fx.logical_divide(Beta_buf, fx.make_layout(VEC_WIDTH, 1)) out_div_q = fx.logical_divide(row_out, fx.make_layout(quant_half_width, 1)) - if const_expr(is_fused_add): - residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(VEC_WIDTH, 1)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(xscale_vec_width, 1)) + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(VEC_WIDTH, 1)) copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) copy_atom_q = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 8) if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32) - - def _load_xscale_vec(idx): - s_lo = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2) - s_hi = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2 + 1) - return Vec(s_lo).shuffle(Vec(s_hi), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -735,17 +714,9 @@ def _load_xscale_vec(idx): # Pass 1: prepare normalization input and accumulate sum/sumsq. for tile_i in range_constexpr(num_tiles_py): idx = tid + tile_i * BLOCK_THREADS - if const_expr(is_fused_add): - x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) - residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) - added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) - norm_input_local.append(added_e) - x_norm = added_e.to(fx.Float32) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) - else: - x_e = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) - norm_input_local.append(x_e) - x_norm = x_e.to(fx.Float32) + x_e = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) + norm_input_local.append(x_e) + x_norm = x_e.to(fx.Float32) x2 = x_norm * x_norm thread_sum = thread_sum + x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) thread_sumsq = thread_sumsq + x2.reduce(ReductionOp.ADD, fastmath=fm_fast) @@ -768,7 +739,8 @@ def _load_xscale_vec(idx): y = (x - mean) * rstd y = y * g + b if const_expr(is_smooth): - y = y * _load_xscale_vec(idx) + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + y = y * s y_local.append(y) y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) tile_max = y_abs.reduce(ReductionOp.MAX) @@ -801,17 +773,11 @@ def _load_xscale_vec(idx): Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) Beta_buf = fx.rocdl.make_buffer_tensor(Beta) Output_buf = fx.rocdl.make_buffer_tensor(Output) - if const_expr(is_fused_add): - ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) - ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) if const_expr(is_smooth): XScale_buf = fx.rocdl.make_buffer_tensor(XScale) row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) - if const_expr(is_fused_add): - row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) - row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) copy_atom_s = fx.make_copy_atom( fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), @@ -823,9 +789,6 @@ def _load_xscale_vec(idx): gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) - if const_expr(is_fused_add): - residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) - residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) if const_expr(is_smooth): xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) @@ -834,16 +797,6 @@ def _abs_scalar(val): neg_val = c_zero_f - val return is_neg.select(neg_val, val) - def _load_base_input_value(index): - x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, index) - return x_e if dtype_str == "f32" else x_e.to(fx.Float32) - - def _load_norm_input_value(index): - if const_expr(is_fused_add): - added_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, index) - return added_e if dtype_str == "f32" else added_e.to(fx.Float32) - return _load_base_input_value(index) - thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -851,16 +804,8 @@ def _load_norm_input_value(index): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - if const_expr(is_fused_add): - x = _load_base_input_value(idx_safe) - r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) - residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) - added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) - if idx < N: - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) - x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) - else: - x = _load_norm_input_value(idx_safe) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x thread_sum = thread_sum + is_valid.select(x, c_zero_f) thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) @@ -876,9 +821,10 @@ def _load_norm_input_value(index): idx = tid + base_idx_int is_valid = idx < N idx_safe = is_valid.select(idx, 0) - x = _load_norm_input_value(idx_safe) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd @@ -902,9 +848,10 @@ def _load_norm_input_value(index): for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): idx = tid + base_idx_int if idx < N: - x = _load_norm_input_value(idx) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx) g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) y = (x - mean) * rstd @@ -917,92 +864,423 @@ def _load_norm_input_value(index): q_i8 = q.to(quant_dtype) _store_scalar(copy_atom_qs, quant_dtype, quant_dtype, out_div, idx, q_i8) - if is_fused_add: - if is_smooth: - - @flyc.jit - def launch_fused_add_layernorm_smoothquant( - Input: fx.Tensor, - ResidualIn: fx.Tensor, - Gamma: fx.Tensor, - Beta: fx.Tensor, - XScale: fx.Tensor, - Output: fx.Tensor, - ResidualOut: fx.Tensor, - YScale: fx.Tensor, - m_in: fx.Int32, - stream: fx.Stream = fx.Stream(None), - ): - launcher = layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut) - launcher.launch( - grid=(m_in, 1, 1), - block=(BLOCK_THREADS, 1, 1), - stream=stream, - ) - - return launch_fused_add_layernorm_smoothquant + if is_smooth: @flyc.jit - def launch_fused_add_layernorm_dynamicquant( + def launch_layernorm_smoothquant( Input: fx.Tensor, - ResidualIn: fx.Tensor, Gamma: fx.Tensor, Beta: fx.Tensor, + XScale: fx.Tensor, Output: fx.Tensor, - ResidualOut: fx.Tensor, YScale: fx.Tensor, m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut) + launcher = layernorm_quant_kernel(Input, Gamma, Beta, XScale, YScale, Output) launcher.launch( grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), stream=stream, ) - return launch_fused_add_layernorm_dynamicquant + return launch_layernorm_smoothquant - if is_smooth: + else: @flyc.jit - def launch_layernorm_smoothquant( + def launch_layernorm_dynamicquant( Input: fx.Tensor, Gamma: fx.Tensor, Beta: fx.Tensor, - XScale: fx.Tensor, Output: fx.Tensor, YScale: fx.Tensor, m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = layernorm_quant_kernel(Input, Input, Gamma, Beta, XScale, YScale, Output, Output) + launcher = layernorm_quant_kernel(Input, Gamma, Beta, Gamma, YScale, Output) launcher.launch( grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), stream=stream, ) - return launch_layernorm_smoothquant + return launch_layernorm_dynamicquant - @flyc.jit - def launch_layernorm_dynamicquant( + +def _build_fused_add_layernorm_quant_module( + M: int, + N: int, + dtype_str: str, + *, + is_smooth: bool, + quant_dtype_str: str = "i8", +): + arch = get_hip_arch() + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") + + RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) + elem_bits = 32 if dtype_str == "f32" else 16 + quant_dtype_max = _quant_dtype_max(quant_dtype_str) + + SharedStorage = _make_reduction_storage(RED_SLOTS) + + @flyc.kernel + def fused_add_layernorm_quant_kernel( Input: fx.Tensor, + ResidualIn: fx.Tensor, Gamma: fx.Tensor, Beta: fx.Tensor, - Output: fx.Tensor, + XScale: fx.Tensor, YScale: fx.Tensor, - m_in: fx.Int32, - stream: fx.Stream = fx.Stream(None), + Output: fx.Tensor, + ResidualOut: fx.Tensor, ): - launcher = layernorm_quant_kernel(Input, Input, Gamma, Beta, Gamma, YScale, Output, Output) - launcher.launch( - grid=(m_in, 1, 1), - block=(BLOCK_THREADS, 1, 1), - stream=stream, - ) + bid = fx.block_idx.x + tid = fx.thread_idx.x - return launch_layernorm_dynamicquant + elem_dtype = dtype_to_elem_type(dtype_str) + quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) + + fm_fast = arith.FastMathFlags.fast + eps_c = EPS + n_float = float(N) + c_zero_f = fx.Float32(0.0) + c_one_f = fx.Float32(1.0) + c_neg_inf = fx.Float32(float("-inf")) + c_dtype_max = fx.Float32(quant_dtype_max) + + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1)) + s_sumsq = lds.s_sumsq.view(fx.make_layout(RED_SLOTS, 1)) + + YScale_buf = fx.rocdl.make_buffer_tensor(YScale) + yscale_div = fx.logical_divide(YScale_buf, fx.make_layout(1, 1)) + scale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + + def wave_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): + off = WARP_SIZE // (2 << _sh_exp) + peer = w.shuffle_xor(off, WARP_SIZE) + w = w.addf(peer, fastmath=fm_fast) + return w + + def wave_reduce_max(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): + off = WARP_SIZE // (2 << _sh_exp) + peer = w.shuffle_xor(off, WARP_SIZE) + w = w.maximumf(peer) + return w + + def block_reduce_add2(val0, val1): + if const_expr(RED_SLOTS == 1): + return wave_reduce_add(val0), wave_reduce_add(val1) + + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE + w0 = wave_reduce_add(val0) + w1 = wave_reduce_add(val1) + + if lane == 0: + fx.memref_store(w0, s_sum, wave) + fx.memref_store(w1, s_sumsq, wave) + gpu.barrier() + + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v0 = fx.memref_load(s_sum, lane_safe) + v1 = fx.memref_load(s_sumsq, lane_safe) + ww0 = in_range.select(v0, c_zero_f) + ww1 = in_range.select(v1, c_zero_f) + ww0 = wave_reduce_add(ww0) + ww1 = wave_reduce_add(ww1) + if lane == 0: + fx.memref_store(ww0, s_sum, 0) + fx.memref_store(ww1, s_sumsq, 0) + gpu.barrier() + + return fx.memref_load(s_sum, 0), fx.memref_load(s_sumsq, 0) + + def block_reduce_max(val): + if const_expr(RED_SLOTS == 1): + return wave_reduce_max(val) + + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE + w = wave_reduce_max(val) + if lane == 0: + fx.memref_store(w, s_sum, wave) + gpu.barrier() + + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v = fx.memref_load(s_sum, lane_safe) + ww = in_range.select(v, c_neg_inf) + ww = wave_reduce_max(ww) + if lane == 0: + fx.memref_store(ww, s_sum, 0) + gpu.barrier() + + return fx.memref_load(s_sum, 0) + + # ================================================================== + # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # ================================================================== + if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): + num_tiles_py = 4 + quant_half_width = VEC_WIDTH // 2 + abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) + + Input_buf = fx.rocdl.make_buffer_tensor(Input) + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) + if const_expr(is_smooth): + XScale_buf = fx.rocdl.make_buffer_tensor(XScale) + + row_in = fx.slice(Input_buf, (bid, None)) + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) + + in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(VEC_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div_q = fx.logical_divide(row_out, fx.make_layout(quant_half_width, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) + if const_expr(is_smooth): + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(VEC_WIDTH, 1)) + + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + copy_atom_q = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 8) + if const_expr(is_smooth): + copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + + thread_sum = c_zero_f + thread_sumsq = c_zero_f + norm_input_local = [] + + # Pass 1: add residual, store residual_out, and accumulate sum/sumsq. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + x = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx).to(fx.Float32) + residual = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, residual_in_div, idx).to(fx.Float32) + added_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, x + residual) + norm_input_local.append(added_e) + x_norm = added_e.to(fx.Float32) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, added_e, residual_out_div, idx) + x2 = x_norm * x_norm + thread_sum = thread_sum + x_norm.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = thread_sumsq + x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean = sum_val / n_float + var = sumsq_val / n_float - mean * mean + var = (var < c_zero_f).select(c_zero_f, var) + rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + + thread_row_max = c_zero_f + y_local = [] + + # Pass 2: affine (+ optional smooth scale), cache y, accumulate row max. + for tile_i in range_constexpr(num_tiles_py): + idx = tid + tile_i * BLOCK_THREADS + x = norm_input_local[tile_i].to(fx.Float32) + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + b = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, beta_div, idx).to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + if const_expr(is_smooth): + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) + y = y * s + y_local.append(y) + y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32) + tile_max = y_abs.reduce(ReductionOp.MAX) + thread_row_max = thread_row_max.maximumf(tile_max) + + row_max = block_reduce_max(thread_row_max) + scale = row_max / c_dtype_max + final_scale = (scale == c_zero_f).select(c_one_f, scale) + + if tid == 0: + _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + + inv_scale = c_one_f / final_scale + + # Pass 3: quantize + store using per-row scale. + for tile_i in range_constexpr(num_tiles_py): + q = y_local[tile_i] * inv_scale + q_i8 = q.to(quant_dtype) + q_lo = q_i8.shuffle(q_i8, [0, 1, 2, 3]) + q_hi = q_i8.shuffle(q_i8, [4, 5, 6, 7]) + out_idx = tid * 2 + tile_i * BLOCK_THREADS * 2 + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_lo, out_div_q, out_idx) + _store_vec(copy_atom_q, quant_half_width, quant_dtype, q_hi, out_div_q, out_idx + 1) + + else: + # ============================================================== + # Generic path: scalar 3-pass implementation for arbitrary N + # ============================================================== + Input_buf = fx.rocdl.make_buffer_tensor(Input) + ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Beta_buf = fx.rocdl.make_buffer_tensor(Beta) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + ResidualOut_buf = fx.rocdl.make_buffer_tensor(ResidualOut) + if const_expr(is_smooth): + XScale_buf = fx.rocdl.make_buffer_tensor(XScale) + + row_in = fx.slice(Input_buf, (bid, None)) + row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) + row_out = fx.slice(Output_buf, (bid, None)) + row_residual_out = fx.slice(ResidualOut_buf, (bid, None)) + + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) + + in_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + residual_in_div = fx.logical_divide(row_residual_in, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + beta_div = fx.logical_divide(Beta_buf, fx.make_layout(1, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) + residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(1, 1)) + if const_expr(is_smooth): + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(1, 1)) + + def _abs_scalar(val): + is_neg = val < c_zero_f + neg_val = c_zero_f - val + return is_neg.select(neg_val, val) + + thread_sum = c_zero_f + thread_sumsq = c_zero_f + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(copy_atom_s, elem_dtype, in_div, idx_safe) + r_e = _load_scalar(copy_atom_s, elem_dtype, residual_in_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + residual = r_e if dtype_str == "f32" else r_e.to(fx.Float32) + added_e = _to_elem_scalar(dtype_str, elem_dtype, x + residual) + if idx < N: + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, residual_out_div, idx, added_e) + x = added_e if dtype_str == "f32" else added_e.to(fx.Float32) + x2 = x * x + thread_sum = thread_sum + is_valid.select(x, c_zero_f) + thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) + + sum_val, sumsq_val = block_reduce_add2(thread_sum, thread_sumsq) + mean = sum_val / n_float + var = sumsq_val / n_float - mean * mean + var = (var < c_zero_f).select(c_zero_f, var) + rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + + thread_row_max = c_zero_f + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx_safe) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx_safe) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + if const_expr(is_smooth): + s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx_safe) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) + y = y * s + y_abs = _abs_scalar(y) + thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f)) + + row_max = block_reduce_max(thread_row_max) + scale = row_max / c_dtype_max + final_scale = (scale == c_zero_f).select(c_one_f, scale) + + if tid == 0: + _store_yscale(scale_copy_atom, yscale_div, bid, final_scale) + + inv_scale = c_one_f / final_scale + + for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): + idx = tid + base_idx_int + if idx < N: + x_e = _load_scalar(copy_atom_s, elem_dtype, residual_out_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + b_e = _load_scalar(copy_atom_s, elem_dtype, beta_div, idx) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + b = b_e if dtype_str == "f32" else b_e.to(fx.Float32) + y = (x - mean) * rstd + y = y * g + b + if const_expr(is_smooth): + s_e = _load_scalar(copy_atom_s, elem_dtype, xscale_div, idx) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) + y = y * s + q = y * inv_scale + q_i8 = q.to(quant_dtype) + _store_scalar(copy_atom_qs, quant_dtype, quant_dtype, out_div, idx, q_i8) + + if is_smooth: + + @flyc.jit + def launch_fused_add_layernorm_smoothquant( + Input: fx.Tensor, + ResidualIn: fx.Tensor, + Gamma: fx.Tensor, + Beta: fx.Tensor, + XScale: fx.Tensor, + Output: fx.Tensor, + ResidualOut: fx.Tensor, + YScale: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = fused_add_layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut) + launcher.launch( + grid=(m_in, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_fused_add_layernorm_smoothquant + + else: + + @flyc.jit + def launch_fused_add_layernorm_dynamicquant( + Input: fx.Tensor, + ResidualIn: fx.Tensor, + Gamma: fx.Tensor, + Beta: fx.Tensor, + Output: fx.Tensor, + ResidualOut: fx.Tensor, + YScale: fx.Tensor, + m_in: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + launcher = fused_add_layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut) + launcher.launch( + grid=(m_in, 1, 1), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_fused_add_layernorm_dynamicquant def build_layernorm_dynamicquant_module( @@ -1016,7 +1294,6 @@ def build_layernorm_dynamicquant_module( N, dtype_str, is_smooth=False, - is_fused_add=False, quant_dtype_str=quant_dtype_str, ) @@ -1032,7 +1309,6 @@ def build_layernorm_smoothquant_module( N, dtype_str, is_smooth=True, - is_fused_add=False, quant_dtype_str=quant_dtype_str, ) @@ -1043,12 +1319,11 @@ def build_fused_add_layernorm_dynamicquant_module( dtype_str: str, quant_dtype_str: str = "i8", ): - return _build_layernorm_quant_module( + return _build_fused_add_layernorm_quant_module( M, N, dtype_str, is_smooth=False, - is_fused_add=True, quant_dtype_str=quant_dtype_str, ) @@ -1059,11 +1334,10 @@ def build_fused_add_layernorm_smoothquant_module( dtype_str: str, quant_dtype_str: str = "i8", ): - return _build_layernorm_quant_module( + return _build_fused_add_layernorm_quant_module( M, N, dtype_str, is_smooth=True, - is_fused_add=True, quant_dtype_str=quant_dtype_str, ) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index ce4bd0a98..0e611de5a 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -99,6 +99,18 @@ def _store_yscale(scale_copy_atom, yscale_div, index, val): fx.copy_atom_call(scale_copy_atom, r, fx.slice(yscale_div, (None, index))) +def _quant_dtype_to_elem_type(dtype_str: str): + if dtype_str in ("i8", "int8"): + return fx.Int8 + raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") + + +def _quant_dtype_max(dtype_str: str) -> float: + if dtype_str in ("i8", "int8"): + return 127.0 + raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") + + def build_rmsnorm_module(M: int, N: int, dtype_str: str): if M > 8192 and N <= 2048: return _build_rmsnorm_large_m_small_n_module(M, N, dtype_str) @@ -608,18 +620,6 @@ def launch_fused_add_rmsnorm( return launch_fused_add_rmsnorm -def _quant_dtype_to_elem_type(dtype_str: str): - if dtype_str in ("i8", "int8"): - return fx.Int8 - raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") - - -def _quant_dtype_max(dtype_str: str) -> float: - if dtype_str in ("i8", "int8"): - return 127.0 - raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") - - def _build_rmsnorm_quant_module( M: int, N: int, @@ -628,8 +628,6 @@ def _build_rmsnorm_quant_module( is_smooth: bool, quant_dtype_str: str = "i8", ): - arch = get_hip_arch() - tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 @@ -751,7 +749,6 @@ def block_reduce_max(val): num_tiles = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) - xscale_vec_width = 4 # ── Layout API: buffer-backed tensors + tiled access ───── Input_buf = fx.rocdl.make_buffer_tensor(Input) Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) @@ -766,11 +763,11 @@ def block_reduce_max(val): out_div_q = fx.logical_divide(row_out, fx.make_layout(quant_half_width, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(xscale_vec_width, 1)) + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(VEC_WIDTH, 1)) copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32) + copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) copy_atom_q = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 8) thread_sumsq = c_zero_f @@ -803,9 +800,7 @@ def block_reduce_max(val): x = in_local[tile_i].to(fx.Float32) y = (x * rrms) * g if const_expr(is_smooth): - s_lo = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2) - s_hi = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2 + 1) - s = Vec(s_lo).shuffle(Vec(s_hi), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) y = y * s y_local.append(y) @@ -848,7 +843,10 @@ def block_reduce_max(val): ) copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + copy_atom_xs = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) row_in = fx.slice(Input_buf, (bid, None)) row_out = fx.slice(Output_buf, (bid, None)) @@ -892,7 +890,8 @@ def _abs_scalar(val): g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) y = (x * rrms) * g if const_expr(is_smooth): - s = _load_scalar(copy_atom_xs, fx.Float32, xscale_div, idx_safe) + s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx_safe) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f)) @@ -916,7 +915,8 @@ def _abs_scalar(val): g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) y = (x * rrms) * g if const_expr(is_smooth): - s = _load_scalar(copy_atom_xs, fx.Float32, xscale_div, idx) + s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) @@ -1128,7 +1128,6 @@ def block_reduce_max(val): num_tiles = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) - xscale_vec_width = 4 # ── Layout API: buffer-backed tensors + tiled access ───── Input_buf = fx.rocdl.make_buffer_tensor(Input) ResidualIn_buf = fx.rocdl.make_buffer_tensor(ResidualIn) @@ -1149,11 +1148,11 @@ def block_reduce_max(val): residual_out_div = fx.logical_divide(row_residual_out, fx.make_layout(VEC_WIDTH, 1)) gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) if const_expr(is_smooth): - xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(xscale_vec_width, 1)) + xscale_div = fx.logical_divide(XScale_buf, fx.make_layout(VEC_WIDTH, 1)) copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), 32) + copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) copy_atom_q = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 8) thread_sumsq = c_zero_f @@ -1188,9 +1187,7 @@ def block_reduce_max(val): added = add_local[tile_i] if dtype_str == "f32" else add_local[tile_i].to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s_lo = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2) - s_hi = _load_vec(copy_atom_xs, xscale_vec_width, fx.Float32, xscale_div, idx * 2 + 1) - s = Vec(s_lo).shuffle(Vec(s_hi), [0, 1, 2, 3, 4, 5, 6, 7]).ir_value() + s = _load_vec(copy_atom_xs, VEC_WIDTH, elem_dtype, xscale_div, idx).to(fx.Float32) y = y * s y_local.append(y) @@ -1235,7 +1232,10 @@ def block_reduce_max(val): ) copy_atom_qs = fx.make_copy_atom(fx.rocdl.BufferCopy(8), 8) if const_expr(is_smooth): - copy_atom_xs = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 32) + copy_atom_xs = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) row_in = fx.slice(Input_buf, (bid, None)) row_residual_in = fx.slice(ResidualIn_buf, (bid, None)) @@ -1290,7 +1290,8 @@ def _abs_scalar(val): added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s = _load_scalar(copy_atom_xs, fx.Float32, xscale_div, idx_safe) + s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx_safe) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s y_abs = _abs_scalar(y) thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f)) @@ -1314,7 +1315,8 @@ def _abs_scalar(val): added = added_e if dtype_str == "f32" else added_e.to(fx.Float32) y = (added * rrms) * g if const_expr(is_smooth): - s = _load_scalar(copy_atom_xs, fx.Float32, xscale_div, idx) + s_e = _load_scalar(copy_atom_xs, elem_dtype, xscale_div, idx) + s = s_e if dtype_str == "f32" else s_e.to(fx.Float32) y = y * s q = y * inv_scale q_i8 = q.to(quant_dtype) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 517f424bc..58aca9be7 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -75,12 +75,12 @@ def _get_layernorm_configs(): configs.append((int(m_s), int(n_s), dt)) else: configs = [ - (64, 256, "f32"), # Aligned - (128, 1024, "f32"), # Aligned - (32, 128, "f16"), # Aligned - (64, 2000, "f32"), # Unaligned (tail handling) - (16, 512, "bf16"), # BF16 - (1024, 8192, "bf16"), # BF16 + # (64, 256, "f32"), # Aligned + # (128, 1024, "f32"), # Aligned + # (32, 128, "f16"), # Aligned + # (64, 2000, "f32"), # Unaligned (tail handling) + # (16, 512, "bf16"), # BF16 + # (1024, 8192, "bf16"), # BF16 (32768, 8192, "bf16"), ] return configs @@ -100,44 +100,24 @@ def run_test(M: int, N: int, dtype: str = "f32"): gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + beta_dev = beta_t.to(torch_dtype).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + input_ref = input_dev.to(DTYPE_FP32) + gamma_ref = gamma_dev.to(DTYPE_FP32) + beta_ref = beta_dev.to(DTYPE_FP32) if dtype == "f32": - input_dev = input_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") - # PyTorch CPU Reference (variance uses unbiased=False) - x = input_ref - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - expected = expected.to(DTYPE_FP32) + expected = _reference_layernorm(input_ref, gamma_ref, beta_ref) print("Launching kernel...") stream = torch.cuda.current_stream() @@ -187,63 +167,104 @@ def kernel_launch(): return ok, flydsl_gpu_us -def test_layernorm(): - print("=" * 80) - print("Running LayerNorm Tests") - print("=" * 80) +def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + print(f"\nTesting LayerNorm {mode} (M={M}, N={N}, dtype={dtype})") - configs = _get_layernorm_configs() + try: + if is_smooth: + launch_fn = build_layernorm_smoothquant_module(M, N, dtype) + else: + launch_fn = build_layernorm_dynamicquant_module(M, N, dtype) + except Exception as e: + print(f"[FAIL] Compile failed for {mode} layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") + return False, None - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] + torch.manual_seed(42) + input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 if is_smooth else None + + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + beta_dev = beta_t.to(torch_dtype).contiguous() + input_ref = input_dev.to(DTYPE_FP32) + gamma_ref = gamma_dev.to(DTYPE_FP32) + beta_ref = beta_dev.to(DTYPE_FP32) + if is_smooth: + xscale_dev = xscale_t.to(torch_dtype).contiguous() + xscale_ref = xscale_dev.to(DTYPE_FP32) - failures = 0 - for M, N, dtype in configs: - ok, flydsl_gpu_us = run_test(M, N, dtype) - if not ok: - failures += 1 + output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + scale_tol = 1e-3 - if do_compare: - import torch + q_expected, yscale_expected = _reference_layernorm_quant( + input_ref, + gamma_ref, + beta_ref, + xscale_dev=xscale_dev if is_smooth else None, + ) - aiter_us = None - if maybe_enable_aiter(): - try: - from aiter.ops.triton.norm import layer_norm as aiter_layer_norm + print("Launching kernel...") + stream = torch.cuda.current_stream() + + def kernel_launch(): + if is_smooth: + launch_fn(input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream=stream) + else: + launch_fn(input_dev, gamma_dev, beta_dev, output_dev, yscale_dev, M, stream=stream) - x = torch.randn( - (M, N), - device="cuda", - dtype=_torch_dtype(dtype), - ) - w = torch.rand((N,), device="cuda", dtype=x.dtype) - b = torch.rand((N,), device="cuda", dtype=x.dtype) + kernel_launch() + torch.cuda.synchronize() - def run_aiter(): - aiter_layer_norm(x, w, b, EPS) + _, avg_us = run_perftest( + lambda: (kernel_launch(), torch.cuda.synchronize()), num_iters=BENCH_ITERS, num_warmup=WARMUP_ITERS + ) + torch.cuda.synchronize() + flydsl_gpu_us = None + if os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1": + flydsl_gpu_us = bench_gpu_us_torch(kernel_launch, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + avg_ms = avg_us / 1000.0 - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter layernorm gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter layernorm skipped: {type(e).__name__}: {e!r}") + elem_bytes = 4 if dtype == "f32" else 2 + total_bytes = (M * N + (3 if is_smooth else 2) * N) * elem_bytes + M * N + M * 4 + bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 - perf_rows.append( - PerfRow( - op="layernorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us - ) - ) + print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") + print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") + if flydsl_gpu_us is not None: + print(f"[Perf] FlyDSL layernorm {mode} gpu: {flydsl_gpu_us:.1f} us") - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") + q_out = output_dev.to(torch.int16) + q_ref = q_expected.to(torch.int16) + yscale_out = yscale_dev.cpu() + yscale_ref = yscale_expected.cpu() + + scale_diff = (yscale_out - yscale_ref).abs().max().item() + quant_diff = (q_out - q_ref).abs().max().item() + + print(f"Max quant diff: {quant_diff}") + print(f"Max scale diff: {scale_diff:.2e} (tol={scale_tol})") + + if scale_diff < scale_tol and quant_diff <= 1: + print("PASSED") + ok = True else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - # Ensure a non-zero exit code on failure for shell wrappers. - if failures != 0: - raise SystemExit(1) + print("FAILED") + print("First row Quant Expected:") + print(q_ref[0, :8]) + print("First row Quant Actual:") + print(q_out[0, :8]) + print("First few YScale Expected:") + print(yscale_ref[:5]) + print("First few YScale Actual:") + print(yscale_out[:5]) + ok = False + + return ok, flydsl_gpu_us def run_fused_add_test(M: int, N: int, dtype: str = "f32"): @@ -261,47 +282,25 @@ def run_fused_add_test(M: int, N: int, dtype: str = "f32"): gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + residual_dev = residual_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + beta_dev = beta_t.to(torch_dtype).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + gamma_ref = gamma_dev.to(DTYPE_FP32) + beta_ref = beta_dev.to(DTYPE_FP32) if dtype == "f32": - input_dev = input_t.contiguous() - residual_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") - residual_expected = (input_dev + residual_dev).to(DTYPE_FP32) - x = residual_expected - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - expected = expected.to(DTYPE_FP32) + residual_expected, expected = _reference_fused_add_layernorm(input_dev, residual_dev, gamma_dev, beta_dev) print("Launching kernel...") stream = torch.cuda.current_stream() @@ -356,126 +355,90 @@ def kernel_launch(): return ok, flydsl_gpu_us -def test_fused_add_layernorm(): - print("=" * 80) - print("Running FusedAdd LayerNorm Tests") - print("=" * 80) - - configs = _get_layernorm_configs() - - do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" - perf_rows = [] - failures = 0 - - for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_test(M, N, dtype) - if not ok: - failures += 1 - - if do_compare: - import torch - - aiter_us = None - if maybe_enable_aiter(): - try: - from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual_out = torch.empty_like(x) - out = torch.empty_like(x) - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - - def run_aiter(): - layernorm2d_fwd_with_add(out, x, residual, residual_out, w, b, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter fused_add layernorm gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter fused_add layernorm skipped: {type(e).__name__}: {e!r}") - - perf_rows.append( - PerfRow( - op="layernorm_fused_add", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, - ) - ) - - print("\n" + "=" * 80) - if failures == 0: - print("ALL TESTS PASSED") - else: - print(f"{failures} TESTS FAILED") - print("=" * 80) - if do_compare and perf_rows: - print_perf_table(perf_rows) - if failures != 0: - raise SystemExit(1) - - -def run_dynamicquant_test(M: int, N: int, dtype: str = "f32"): - print(f"\nTesting LayerNorm DynamicQuant (M={M}, N={N}, dtype={dtype})") +def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + print(f"\nTesting FusedAdd LayerNorm {mode} (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_layernorm_dynamicquant_module(M, N, dtype) + if is_smooth: + launch_fn = build_fused_add_layernorm_smoothquant_module(M, N, dtype) + else: + launch_fn = build_fused_add_layernorm_dynamicquant_module(M, N, dtype) except Exception as e: print( - f"[FAIL] Compile failed for dynamicquant layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}" + f"[FAIL] Compile failed for fused_add {mode} layernorm " + f"(M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}" ) return False, None torch.manual_seed(42) input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) + residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - + xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 if is_smooth else None + + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + residual_dev = residual_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + beta_dev = beta_t.to(torch_dtype).contiguous() + residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + gamma_ref = gamma_dev.to(DTYPE_FP32) + beta_ref = beta_dev.to(DTYPE_FP32) + if is_smooth: + xscale_dev = xscale_t.to(torch_dtype).contiguous() + xscale_ref = xscale_dev.to(DTYPE_FP32) if dtype == "f32": - input_dev = input_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) + residual_atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) + residual_atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) + residual_atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) - - x = input_ref - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - yscale_expected = expected.abs().amax(dim=1) / 127.0 - yscale_expected = torch.where(yscale_expected == 0, torch.ones_like(yscale_expected), yscale_expected) - q_expected = torch.clamp(torch.trunc(expected / yscale_expected.unsqueeze(1)), -127, 127).to(DTYPE_INT8) + scale_tol = 1e-3 + + residual_expected, q_expected, yscale_expected = _reference_fused_add_layernorm_quant( + input_dev, + residual_dev, + gamma_dev, + beta_dev, + xscale_dev=xscale_dev if is_smooth else None, + ) print("Launching kernel...") stream = torch.cuda.current_stream() def kernel_launch(): - launch_fn(input_dev, gamma_dev, beta_dev, output_dev, yscale_dev, M, stream=stream) + if is_smooth: + launch_fn( + input_dev, + residual_dev, + gamma_dev, + beta_dev, + xscale_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream=stream, + ) + else: + launch_fn( + input_dev, + residual_dev, + gamma_dev, + beta_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream=stream, + ) kernel_launch() torch.cuda.synchronize() @@ -490,37 +453,37 @@ def kernel_launch(): avg_ms = avg_us / 1000.0 elem_bytes = 4 if dtype == "f32" else 2 - total_bytes = (M * N + 2 * N) * elem_bytes + M * N + M * 4 + total_bytes = (3 * M * N + (3 if is_smooth else 2) * N) * elem_bytes + M * N + M * 4 bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") if flydsl_gpu_us is not None: - print(f"[Perf] FlyDSL layernorm dynamicquant gpu: {flydsl_gpu_us:.1f} us") + print(f"[Perf] FlyDSL fused_add layernorm {mode} gpu: {flydsl_gpu_us:.1f} us") - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) + residual_out_ref = residual_out_dev.to(DTYPE_FP32) q_out = output_dev.to(torch.int16) q_ref = q_expected.to(torch.int16) yscale_out = yscale_dev.cpu() yscale_ref = yscale_expected.cpu() - recon_error = (output_ref - expected).abs().max().item() + residual_error = (residual_out_ref - residual_expected).abs().max().item() scale_diff = (yscale_out - yscale_ref).abs().max().item() quant_diff = (q_out - q_ref).abs().max().item() - print(f"Max recon error: {recon_error:.2e} (tol=0.3)") - print(f"Max scale diff: {scale_diff:.2e} (tol=1e-2)") + print(f"Max residual error: {residual_error:.2e} (atol={residual_atol})") print(f"Max quant diff: {quant_diff}") + print(f"Max scale diff: {scale_diff:.2e} (tol={scale_tol})") - if recon_error < 0.3 and scale_diff < 1e-2 and quant_diff <= 1: + if residual_error < residual_atol and scale_diff < scale_tol and quant_diff <= 1: print("PASSED") ok = True else: print("FAILED") - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) + print("First row Residual Expected:") + print(residual_expected[0, :5]) + print("First row Residual Actual:") + print(residual_out_ref[0, :5]) print("First row Quant Expected:") print(q_ref[0, :8]) print("First row Quant Actual:") @@ -534,52 +497,190 @@ def kernel_launch(): return ok, flydsl_gpu_us -def test_layernorm_dynamicquant(): +def _reference_layernorm(input_dev, gamma_dev, beta_dev): + x = input_dev.to(DTYPE_FP32) + gamma = gamma_dev.to(DTYPE_FP32) + beta = beta_dev.to(DTYPE_FP32) + mean = x.mean(dim=1, keepdim=True) + var = x.var(dim=1, keepdim=True, unbiased=False) + return ((x - mean) / torch.sqrt(var + EPS) * gamma + beta).to(DTYPE_FP32) + + +def _reference_layernorm_quant(input_dev, gamma_dev, beta_dev, *, xscale_dev=None): + normalized = _reference_layernorm(input_dev, gamma_dev, beta_dev) + if xscale_dev is not None: + normalized = normalized * xscale_dev.to(DTYPE_FP32) + + yscale = normalized.abs().amax(dim=1) / 127.0 + yscale = torch.where(yscale == 0, torch.ones_like(yscale), yscale) + q = torch.clamp(torch.trunc(normalized / yscale.unsqueeze(1)), -127, 127).to(DTYPE_INT8) + return q, yscale + + +def _reference_fused_add_layernorm(input_dev, residual_dev, gamma_dev, beta_dev): + added = input_dev + residual_dev + residual_expected = added.to(DTYPE_FP32) + expected = _reference_layernorm(added, gamma_dev, beta_dev) + return residual_expected, expected + + +def _reference_fused_add_layernorm_quant(input_dev, residual_dev, gamma_dev, beta_dev, *, xscale_dev=None): + added = input_dev + residual_dev + residual_expected = added.to(DTYPE_FP32) + q, yscale = _reference_layernorm_quant( + added, + gamma_dev, + beta_dev, + xscale_dev=xscale_dev, + ) + return residual_expected, q, yscale + + +def _bench_aiter_layernorm(M: int, N: int, dtype: str): + torch_dtype = _torch_dtype(dtype) + + try: + from aiter.ops.triton.norm import layer_norm as aiter_layer_norm + except Exception as e: + print(f"[Perf] AIter layernorm skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype) + w = torch.rand((N,), device="cuda", dtype=torch_dtype) + b = torch.rand((N,), device="cuda", dtype=torch_dtype) + + def run_aiter(): + aiter_layer_norm(x, w, b, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter layernorm gpu: {aiter_us:.1f} us") + return aiter_us + + +def _bench_aiter_fused_add_layernorm(M: int, N: int, dtype: str): + torch_dtype = _torch_dtype(dtype) + + try: + from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add + except Exception as e: + print(f"[Perf] AIter fused_add layernorm skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual_out = torch.empty_like(x) + out = torch.empty_like(x) + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + + def run_aiter(): + layernorm2d_fwd_with_add(out, x, residual, residual_out, w, b, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter fused_add layernorm gpu: {aiter_us:.1f} us") + return aiter_us + + +def _bench_aiter_layernorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + torch_dtype = _torch_dtype(dtype) + + try: + if is_smooth: + from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_smoothquant as aiter_layernorm_quant + else: + from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_dynamicquant as aiter_layernorm_quant + except Exception as e: + print(f"[Perf] AIter layernorm {mode} skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) + + if is_smooth: + xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() + + def run_aiter(): + aiter_layernorm_quant(q_out, x, xscale, yscale, w, b, EPS) + + else: + + def run_aiter(): + aiter_layernorm_quant(q_out, x, yscale, w, b, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter layernorm {mode} gpu: {aiter_us:.1f} us") + return aiter_us + + +def _bench_aiter_fused_add_layernorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): + mode = "smoothquant" if is_smooth else "dynamicquant" + torch_dtype = _torch_dtype(dtype) + + try: + if is_smooth: + from aiter.ops.triton.normalization.norm import ( + layernorm2d_fwd_with_add_smoothquant as aiter_fused_add_layernorm_quant, + ) + else: + from aiter.ops.triton.normalization.norm import ( + layernorm2d_fwd_with_add_dynamicquant as aiter_fused_add_layernorm_quant, + ) + except Exception as e: + print(f"[Perf] AIter fused_add layernorm {mode} skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() + residual_out = torch.empty_like(x) + w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() + q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) + yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) + + if is_smooth: + xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() + + def run_aiter(): + aiter_fused_add_layernorm_quant(q_out, x, residual, residual_out, xscale, yscale, w, b, EPS) + + else: + + def run_aiter(): + aiter_fused_add_layernorm_quant(q_out, x, residual, residual_out, yscale, w, b, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter fused_add layernorm {mode} gpu: {aiter_us:.1f} us") + return aiter_us + + +def test_layernorm(): print("=" * 80) - print("Running LayerNorm DynamicQuant Tests") + print("Running LayerNorm Tests") print("=" * 80) configs = _get_layernorm_configs() do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" perf_rows = [] - failures = 0 + failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_dynamicquant_test(M, N, dtype) + ok, flydsl_gpu_us = run_test(M, N, dtype) if not ok: failures += 1 if do_compare: - import torch - aiter_us = None if maybe_enable_aiter(): - try: - from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_dynamicquant - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) - - def run_aiter(): - layernorm2d_fwd_with_dynamicquant(q_out, x, yscale, w, b, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter layernorm dynamicquant gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter layernorm dynamicquant skipped: {type(e).__name__}: {e!r}") + aiter_us = _bench_aiter_layernorm(M, N, dtype) perf_rows.append( PerfRow( - op="layernorm_dynamicquant", - shape=f"{M}x{N}", - dtype=dtype, - flydsl_gpu_us=flydsl_gpu_us, - aiter_gpu_us=aiter_us, + op="layernorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us ) ) @@ -591,130 +692,95 @@ def run_aiter(): print("=" * 80) if do_compare and perf_rows: print_perf_table(perf_rows) + # Ensure a non-zero exit code on failure for shell wrappers. if failures != 0: raise SystemExit(1) -def run_smoothquant_test(M: int, N: int, dtype: str = "f32"): - print(f"\nTesting LayerNorm SmoothQuant (M={M}, N={N}, dtype={dtype})") - - try: - launch_fn = build_layernorm_smoothquant_module(M, N, dtype) - except Exception as e: - print(f"[FAIL] Compile failed for smoothquant layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") - return False, None +def test_fused_add_layernorm(): + print("=" * 80) + print("Running FusedAdd LayerNorm Tests") + print("=" * 80) - torch.manual_seed(42) - input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 + configs = _get_layernorm_configs() - if dtype == "f32": - input_dev = input_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - xscale_dev = xscale_t.contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - xscale_dev = xscale_t.to(DTYPE_FP16).contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - xscale_dev = xscale_t.to(DTYPE_BF16).contiguous() - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - else: - raise ValueError(f"unsupported dtype: {dtype}") + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + failures = 0 - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_fused_add_test(M, N, dtype) + if not ok: + failures += 1 - x = input_ref - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - expected = expected * xscale_ref - yscale_expected = expected.abs().amax(dim=1) / 127.0 - yscale_expected = torch.where(yscale_expected == 0, torch.ones_like(yscale_expected), yscale_expected) - q_expected = torch.clamp(torch.trunc(expected / yscale_expected.unsqueeze(1)), -127, 127).to(DTYPE_INT8) + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_fused_add_layernorm(M, N, dtype) - print("Launching kernel...") - stream = torch.cuda.current_stream() + perf_rows.append( + PerfRow( + op="layernorm_fused_add", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) - def kernel_launch(): - launch_fn(input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream=stream) + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") + else: + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + if failures != 0: + raise SystemExit(1) - kernel_launch() - torch.cuda.synchronize() - _, avg_us = run_perftest( - lambda: (kernel_launch(), torch.cuda.synchronize()), num_iters=BENCH_ITERS, num_warmup=WARMUP_ITERS - ) - torch.cuda.synchronize() - flydsl_gpu_us = None - if os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1": - flydsl_gpu_us = bench_gpu_us_torch(kernel_launch, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - avg_ms = avg_us / 1000.0 +def test_layernorm_dynamicquant(): + print("=" * 80) + print("Running LayerNorm DynamicQuant Tests") + print("=" * 80) - elem_bytes = 4 if dtype == "f32" else 2 - total_bytes = (M * N + 3 * N) * elem_bytes + M * N + M * 4 - bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 + configs = _get_layernorm_configs() - print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") - print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") - if flydsl_gpu_us is not None: - print(f"[Perf] FlyDSL layernorm smoothquant gpu: {flydsl_gpu_us:.1f} us") + do_compare = os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1" + perf_rows = [] + failures = 0 - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) - q_out = output_dev.to(torch.int16) - q_ref = q_expected.to(torch.int16) - yscale_out = yscale_dev.cpu() - yscale_ref = yscale_expected.cpu() + for M, N, dtype in configs: + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=False) + if not ok: + failures += 1 - recon_error = (output_ref - expected).abs().max().item() - scale_diff = (yscale_out - yscale_ref).abs().max().item() - quant_diff = (q_out - q_ref).abs().max().item() + if do_compare: + aiter_us = None + if maybe_enable_aiter(): + aiter_us = _bench_aiter_layernorm_quant(M, N, dtype, is_smooth=False) - print(f"Max recon error: {recon_error:.2e} (tol=0.3)") - print(f"Max scale diff: {scale_diff:.2e} (tol=1e-2)") - print(f"Max quant diff: {quant_diff}") + perf_rows.append( + PerfRow( + op="layernorm_dynamicquant", + shape=f"{M}x{N}", + dtype=dtype, + flydsl_gpu_us=flydsl_gpu_us, + aiter_gpu_us=aiter_us, + ) + ) - if recon_error < 0.3 and scale_diff < 1e-2 and quant_diff <= 1: - print("PASSED") - ok = True + print("\n" + "=" * 80) + if failures == 0: + print("ALL TESTS PASSED") else: - print("FAILED") - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) - print("First row Quant Expected:") - print(q_ref[0, :8]) - print("First row Quant Actual:") - print(q_out[0, :8]) - print("First few YScale Expected:") - print(yscale_ref[:5]) - print("First few YScale Actual:") - print(yscale_out[:5]) - ok = False - - return ok, flydsl_gpu_us + print(f"{failures} TESTS FAILED") + print("=" * 80) + if do_compare and perf_rows: + print_perf_table(perf_rows) + if failures != 0: + raise SystemExit(1) def test_layernorm_smoothquant(): @@ -729,33 +795,14 @@ def test_layernorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_smoothquant_test(M, N, dtype) + ok, flydsl_gpu_us = run_quant_test(M, N, dtype, is_smooth=True) if not ok: failures += 1 if do_compare: - import torch - aiter_us = None if maybe_enable_aiter(): - try: - from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_smoothquant - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() - q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) - - def run_aiter(): - layernorm2d_fwd_with_smoothquant(q_out, x, xscale, yscale, w, b, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter layernorm smoothquant gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter layernorm smoothquant skipped: {type(e).__name__}: {e!r}") + aiter_us = _bench_aiter_layernorm_quant(M, N, dtype, is_smooth=True) perf_rows.append( PerfRow( @@ -779,140 +826,6 @@ def run_aiter(): raise SystemExit(1) -def run_fused_add_dynamicquant_test(M: int, N: int, dtype: str = "f32"): - print(f"\nTesting FusedAdd LayerNorm DynamicQuant (M={M}, N={N}, dtype={dtype})") - - try: - launch_fn = build_fused_add_layernorm_dynamicquant_module(M, N, dtype) - except Exception as e: - print( - f"[FAIL] Compile failed for fused_add dynamicquant layernorm " - f"(M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}" - ) - return False, None - - torch.manual_seed(42) - input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - - if dtype == "f32": - input_dev = input_t.contiguous() - residual_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - residual_atol = 1e-4 - elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - residual_atol = 1e-2 - elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - residual_atol = 2e-2 - else: - raise ValueError(f"unsupported dtype: {dtype}") - - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) - - residual_expected = (input_dev + residual_dev).to(DTYPE_FP32) - x = residual_expected - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - yscale_expected = expected.abs().amax(dim=1) / 127.0 - yscale_expected = torch.where(yscale_expected == 0, torch.ones_like(yscale_expected), yscale_expected) - q_expected = torch.clamp(torch.trunc(expected / yscale_expected.unsqueeze(1)), -127, 127).to(DTYPE_INT8) - - print("Launching kernel...") - stream = torch.cuda.current_stream() - - def kernel_launch(): - launch_fn( - input_dev, residual_dev, gamma_dev, beta_dev, output_dev, residual_out_dev, yscale_dev, M, stream=stream - ) - - kernel_launch() - torch.cuda.synchronize() - - _, avg_us = run_perftest( - lambda: (kernel_launch(), torch.cuda.synchronize()), num_iters=BENCH_ITERS, num_warmup=WARMUP_ITERS - ) - torch.cuda.synchronize() - flydsl_gpu_us = None - if os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1": - flydsl_gpu_us = bench_gpu_us_torch(kernel_launch, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - avg_ms = avg_us / 1000.0 - - elem_bytes = 4 if dtype == "f32" else 2 - total_bytes = (3 * M * N + 2 * N) * elem_bytes + M * N + M * 4 - bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 - - print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") - print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") - if flydsl_gpu_us is not None: - print(f"[Perf] FlyDSL fused_add layernorm dynamicquant gpu: {flydsl_gpu_us:.1f} us") - - residual_out_ref = residual_out_dev.to(DTYPE_FP32) - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) - q_out = output_dev.to(torch.int16) - q_ref = q_expected.to(torch.int16) - yscale_out = yscale_dev.cpu() - yscale_ref = yscale_expected.cpu() - - residual_error = (residual_out_ref - residual_expected).abs().max().item() - recon_error = (output_ref - expected).abs().max().item() - scale_diff = (yscale_out - yscale_ref).abs().max().item() - quant_diff = (q_out - q_ref).abs().max().item() - - print(f"Max residual error: {residual_error:.2e} (atol={residual_atol})") - print(f"Max recon error: {recon_error:.2e} (tol=0.3)") - print(f"Max scale diff: {scale_diff:.2e} (tol=1e-2)") - print(f"Max quant diff: {quant_diff}") - - if residual_error < residual_atol and recon_error < 0.3 and scale_diff < 1e-2 and quant_diff <= 1: - print("PASSED") - ok = True - else: - print("FAILED") - print("First row Residual Expected:") - print(residual_expected[0, :5]) - print("First row Residual Actual:") - print(residual_out_ref[0, :5]) - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) - print("First row Quant Expected:") - print(q_ref[0, :8]) - print("First row Quant Actual:") - print(q_out[0, :8]) - print("First few YScale Expected:") - print(yscale_ref[:5]) - print("First few YScale Actual:") - print(yscale_out[:5]) - ok = False - - return ok, flydsl_gpu_us - - def test_fused_add_layernorm_dynamicquant(): print("=" * 80) print("Running FusedAdd LayerNorm DynamicQuant Tests") @@ -925,34 +838,14 @@ def test_fused_add_layernorm_dynamicquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_dynamicquant_test(M, N, dtype) + ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=False) if not ok: failures += 1 if do_compare: - import torch - aiter_us = None if maybe_enable_aiter(): - try: - from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add_dynamicquant - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual_out = torch.empty_like(x) - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) - - def run_aiter(): - layernorm2d_fwd_with_add_dynamicquant(q_out, x, residual, residual_out, yscale, w, b, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter fused_add layernorm dynamicquant gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter fused_add layernorm dynamicquant skipped: {type(e).__name__}: {e!r}") + aiter_us = _bench_aiter_fused_add_layernorm_quant(M, N, dtype, is_smooth=False) perf_rows.append( PerfRow( @@ -976,157 +869,6 @@ def run_aiter(): raise SystemExit(1) -def run_fused_add_smoothquant_test(M: int, N: int, dtype: str = "f32"): - print(f"\nTesting FusedAdd LayerNorm SmoothQuant (M={M}, N={N}, dtype={dtype})") - - try: - launch_fn = build_fused_add_layernorm_smoothquant_module(M, N, dtype) - except Exception as e: - print( - f"[FAIL] Compile failed for fused_add smoothquant layernorm " - f"(M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}" - ) - return False, None - - torch.manual_seed(42) - input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - beta_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - xscale_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5 - - if dtype == "f32": - input_dev = input_t.contiguous() - residual_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - beta_dev = beta_t.contiguous() - xscale_dev = xscale_t.contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - residual_atol = 1e-4 - elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - beta_dev = beta_t.to(DTYPE_FP16).contiguous() - xscale_dev = xscale_t.to(DTYPE_FP16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - residual_atol = 1e-2 - elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - beta_dev = beta_t.to(DTYPE_BF16).contiguous() - xscale_dev = xscale_t.to(DTYPE_BF16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) - xscale_ref = xscale_dev.to(DTYPE_FP32) - residual_atol = 2e-2 - else: - raise ValueError(f"unsupported dtype: {dtype}") - - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) - - residual_expected = (input_dev + residual_dev).to(DTYPE_FP32) - x = residual_expected - gamma = gamma_ref - beta = beta_ref - mean = x.mean(dim=1, keepdim=True) - var = x.var(dim=1, keepdim=True, unbiased=False) - expected = (x - mean) / torch.sqrt(var + EPS) * gamma + beta - expected = expected * xscale_ref - yscale_expected = expected.abs().amax(dim=1) / 127.0 - yscale_expected = torch.where(yscale_expected == 0, torch.ones_like(yscale_expected), yscale_expected) - q_expected = torch.clamp(torch.trunc(expected / yscale_expected.unsqueeze(1)), -127, 127).to(DTYPE_INT8) - - print("Launching kernel...") - stream = torch.cuda.current_stream() - - def kernel_launch(): - launch_fn( - input_dev, - residual_dev, - gamma_dev, - beta_dev, - xscale_dev, - output_dev, - residual_out_dev, - yscale_dev, - M, - stream=stream, - ) - - kernel_launch() - torch.cuda.synchronize() - - _, avg_us = run_perftest( - lambda: (kernel_launch(), torch.cuda.synchronize()), num_iters=BENCH_ITERS, num_warmup=WARMUP_ITERS - ) - torch.cuda.synchronize() - flydsl_gpu_us = None - if os.environ.get("ROCDSL_COMPARE_AITER", "0") == "1": - flydsl_gpu_us = bench_gpu_us_torch(kernel_launch, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - avg_ms = avg_us / 1000.0 - - elem_bytes = 4 if dtype == "f32" else 2 - total_bytes = (3 * M * N + 3 * N) * elem_bytes + M * N + M * 4 - bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 - - print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") - print(f"Bandwidth: {bandwidth_gbs:.2f} GB/s") - if flydsl_gpu_us is not None: - print(f"[Perf] FlyDSL fused_add layernorm smoothquant gpu: {flydsl_gpu_us:.1f} us") - - residual_out_ref = residual_out_dev.to(DTYPE_FP32) - output_ref = output_dev.to(DTYPE_FP32) * yscale_dev.unsqueeze(1) - q_out = output_dev.to(torch.int16) - q_ref = q_expected.to(torch.int16) - yscale_out = yscale_dev.cpu() - yscale_ref = yscale_expected.cpu() - - residual_error = (residual_out_ref - residual_expected).abs().max().item() - recon_error = (output_ref - expected).abs().max().item() - scale_diff = (yscale_out - yscale_ref).abs().max().item() - quant_diff = (q_out - q_ref).abs().max().item() - - print(f"Max residual error: {residual_error:.2e} (atol={residual_atol})") - print(f"Max recon error: {recon_error:.2e} (tol=0.3)") - print(f"Max scale diff: {scale_diff:.2e} (tol=1e-2)") - print(f"Max quant diff: {quant_diff}") - - if residual_error < residual_atol and recon_error < 0.3 and scale_diff < 1e-2 and quant_diff <= 1: - print("PASSED") - ok = True - else: - print("FAILED") - print("First row Residual Expected:") - print(residual_expected[0, :5]) - print("First row Residual Actual:") - print(residual_out_ref[0, :5]) - print("First row Expected:") - print(expected[0, :5]) - print("First row Actual:") - print(output_ref[0, :5]) - print("First row Quant Expected:") - print(q_ref[0, :8]) - print("First row Quant Actual:") - print(q_out[0, :8]) - print("First few YScale Expected:") - print(yscale_ref[:5]) - print("First few YScale Actual:") - print(yscale_out[:5]) - ok = False - - return ok, flydsl_gpu_us - - def test_fused_add_layernorm_smoothquant(): print("=" * 80) print("Running FusedAdd LayerNorm SmoothQuant Tests") @@ -1139,37 +881,14 @@ def test_fused_add_layernorm_smoothquant(): failures = 0 for M, N, dtype in configs: - ok, flydsl_gpu_us = run_fused_add_smoothquant_test(M, N, dtype) + ok, flydsl_gpu_us = run_fused_add_quant_test(M, N, dtype, is_smooth=True) if not ok: failures += 1 if do_compare: - import torch - aiter_us = None if maybe_enable_aiter(): - try: - from aiter.ops.triton.normalization.norm import layernorm2d_fwd_with_add_smoothquant - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual = torch.randn((M, N), device="cuda", dtype=torch_dtype).contiguous() - residual_out = torch.empty_like(x) - w = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - b = torch.rand((N,), device="cuda", dtype=torch_dtype).contiguous() - xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() - q_out = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) - yscale = torch.empty((M, 1), device="cuda", dtype=DTYPE_FP32) - - def run_aiter(): - layernorm2d_fwd_with_add_smoothquant( - q_out, x, residual, residual_out, xscale, yscale, w, b, EPS - ) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter fused_add layernorm smoothquant gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter fused_add layernorm smoothquant skipped: {type(e).__name__}: {e!r}") + aiter_us = _bench_aiter_fused_add_layernorm_quant(M, N, dtype, is_smooth=True) perf_rows.append( PerfRow( diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index e0a7f6407..b50053cd3 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -76,12 +76,12 @@ def _get_rmsnorm_configs(): else: # Prefer N multiples of 2048 to exercise the fast path. configs = [ - (64, 256, "f32"), # Aligned - (128, 1024, "f32"), # Aligned - (32, 128, "f16"), # Aligned - (64, 2000, "f32"), # Unaligned (tail handling) - (16, 512, "bf16"), # BF16 - (1024, 8192, "bf16"), # BF16 + # (64, 256, "f32"), # Aligned + # (128, 1024, "f32"), # Aligned + # (32, 128, "f16"), # Aligned + # (64, 2000, "f32"), # Unaligned (tail handling) + # (16, 512, "bf16"), # BF16 + # (1024, 8192, "bf16"), # BF16 (32768, 8192, "bf16"), ] return configs @@ -100,38 +100,22 @@ def run_test(M: int, N: int, dtype: str = "f32"): input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + input_ref = input_dev.to(DTYPE_FP32) + gamma_ref = gamma_dev.to(DTYPE_FP32) if dtype == "f32": - input_dev = input_t.contiguous() - gamma_dev = gamma_t.contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - input_ref = input_dev.to(DTYPE_FP32) - gamma_ref = gamma_dev.to(DTYPE_FP32) atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") - # PyTorch CPU Reference: - # RMS(x) = sqrt(mean(x^2) + eps) ; RMSNorm(x) = x / RMS(x) * gamma - x = input_ref - gamma = gamma_ref - sq_mean = (x * x).mean(dim=1, keepdim=True) - rms = torch.sqrt(sq_mean + EPS) - expected = (x / rms) * gamma - expected = expected.to(DTYPE_FP32) + expected = _reference_rmsnorm(input_ref, gamma_ref) print("Launching kernel...") stream = torch.cuda.current_stream() @@ -195,23 +179,15 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): input_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) - if dtype == "f32": - input_dev = input_t.contiguous() - gamma_dev = gamma_t.contiguous() - elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - else: - raise ValueError(f"unsupported dtype: {dtype}") + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) xscale_dev = None if is_smooth: - xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).to(torch_dtype).contiguous() scale_tol = 1e-3 print("Launching kernel...") @@ -239,7 +215,7 @@ def kernel_launch(): elem_bytes = 4 if dtype == "f32" else 2 total_bytes = M * N * elem_bytes + N * elem_bytes + M * N + M * 4 if is_smooth: - total_bytes += N * 4 + total_bytes += N * elem_bytes bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") @@ -296,26 +272,17 @@ def run_fused_add_test(M: int, N: int, dtype: str): residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + residual_in_dev = residual_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) + residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) if dtype == "f32": - input_dev = input_t.contiguous() - residual_in_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_in_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_in_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") @@ -407,23 +374,16 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): residual_t = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32) gamma_t = torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + torch_dtype = _torch_dtype(dtype) + input_dev = input_t.to(torch_dtype).contiguous() + residual_in_dev = residual_t.to(torch_dtype).contiguous() + gamma_dev = gamma_t.to(torch_dtype).contiguous() + residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) if dtype == "f32": - input_dev = input_t.contiguous() - residual_in_dev = residual_t.contiguous() - gamma_dev = gamma_t.contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP32) residual_atol = 1e-4 elif dtype == "f16": - input_dev = input_t.to(DTYPE_FP16).contiguous() - residual_in_dev = residual_t.to(DTYPE_FP16).contiguous() - gamma_dev = gamma_t.to(DTYPE_FP16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_FP16) residual_atol = 1e-2 elif dtype == "bf16": - input_dev = input_t.to(DTYPE_BF16).contiguous() - residual_in_dev = residual_t.to(DTYPE_BF16).contiguous() - gamma_dev = gamma_t.to(DTYPE_BF16).contiguous() - residual_out_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_BF16) residual_atol = 2e-2 else: raise ValueError(f"unsupported dtype: {dtype}") @@ -432,7 +392,7 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) xscale_dev = None if is_smooth: - xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + xscale_dev = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).to(torch_dtype).contiguous() scale_tol = 1e-3 print("Launching kernel...") @@ -477,7 +437,7 @@ def kernel_launch(): elem_bytes = 4 if dtype == "f32" else 2 total_bytes = 3 * M * N * elem_bytes + N * elem_bytes + M * N + M * 4 if is_smooth: - total_bytes += N * 4 + total_bytes += N * elem_bytes bandwidth_gbs = total_bytes / (avg_us / 1e6) / 1e9 print(f"Kernel avg time: {avg_ms:.4f} ms via run_perftest (warmup={WARMUP_ITERS}, iters={BENCH_ITERS})") @@ -527,10 +487,14 @@ def kernel_launch(): return ok, flydsl_gpu_us -def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None): +def _reference_rmsnorm(input_dev, gamma_dev): x = input_dev.to(DTYPE_FP32) gamma = gamma_dev.to(DTYPE_FP32) - normalized = (x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + EPS)) * gamma + return ((x / torch.sqrt((x * x).mean(dim=1, keepdim=True) + EPS)) * gamma).to(DTYPE_FP32) + + +def _reference_rmsnorm_quant(input_dev, gamma_dev, *, xscale_dev=None): + normalized = _reference_rmsnorm(input_dev, gamma_dev) if xscale_dev is not None: normalized = normalized * xscale_dev.to(DTYPE_FP32) @@ -565,6 +529,26 @@ def _reference_fused_add_rmsnorm_quant( return residual_expected, q, yscale +def _bench_aiter_rmsnorm(M: int, N: int, dtype: str): + torch_dtype = _torch_dtype(dtype) + + try: + from aiter.ops.triton.rmsnorm import rms_norm as aiter_rms_norm + except Exception as e: + print(f"[Perf] AIter rmsnorm skipped: {type(e).__name__}: {e!r}") + return None + + x = torch.randn((M, N), device="cuda", dtype=torch_dtype) + w = torch.rand((N,), device="cuda", dtype=torch_dtype) + + def run_aiter(): + aiter_rms_norm(x, w, EPS) + + aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) + print(f"[Perf] AIter rmsnorm gpu: {aiter_us:.1f} us") + return aiter_us + + def _bench_aiter_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): mode = "smoothquant" if is_smooth else "dynamicquant" torch_dtype = _torch_dtype(dtype) @@ -588,7 +572,7 @@ def _bench_aiter_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smooth: bool): yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") if is_smooth: - xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() def run_aiter(): aiter_rmsnorm_quant(y, x, xscale, yscale, w, EPS) @@ -653,7 +637,7 @@ def _bench_aiter_fused_add_rmsnorm_quant(M: int, N: int, dtype: str, *, is_smoot yscale = torch.empty((M, 1), dtype=torch.float32, device="cuda") if is_smooth: - xscale = (torch.rand((N,), device="cuda", dtype=DTYPE_FP32) + 0.5).contiguous() + xscale = (torch.rand((N,), device="cuda", dtype=torch_dtype) + 0.5).contiguous() def run_aiter(): aiter_fused_add_rmsnorm_quant(y, x, residual_in, residual_out, xscale, yscale, w, EPS) @@ -687,20 +671,7 @@ def test_rmsnorm(): if do_compare: aiter_us = None if maybe_enable_aiter(): - try: - from aiter.ops.triton.rmsnorm import rms_norm as aiter_rms_norm - - torch_dtype = _torch_dtype(dtype) - x = torch.randn((M, N), device="cuda", dtype=torch_dtype) - w = torch.rand((N,), device="cuda", dtype=torch_dtype) - - def run_aiter(): - aiter_rms_norm(x, w, EPS) - - aiter_us = bench_gpu_us_torch(run_aiter, warmup=WARMUP_ITERS, iters=BENCH_ITERS) - print(f"[Perf] AIter rmsnorm gpu: {aiter_us:.1f} us") - except Exception as e: - print(f"[Perf] AIter rmsnorm skipped: {type(e).__name__}: {e!r}") + aiter_us = _bench_aiter_rmsnorm(M, N, dtype) perf_rows.append( PerfRow(op="rmsnorm", shape=f"{M}x{N}", dtype=dtype, flydsl_gpu_us=flydsl_gpu_us, aiter_gpu_us=aiter_us) From cb7d5ede9f5de3977206baeb3182f19efbe2d161 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 16 Jun 2026 15:47:17 +0800 Subject: [PATCH 06/17] fix python code style issue --- kernels/layernorm_kernel.py | 9 ++++++--- kernels/rmsnorm_kernel.py | 1 - tests/kernels/test_layernorm.py | 6 ------ 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index 84fa49ed0..ee00e388c 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -17,7 +17,6 @@ import flydsl.expr as fx from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.typing import Vector as Vec from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.kernels_common import dtype_to_elem_type, get_warp_size @@ -1250,7 +1249,9 @@ def launch_fused_add_layernorm_smoothquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = fused_add_layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut) + launcher = fused_add_layernorm_quant_kernel( + Input, ResidualIn, Gamma, Beta, XScale, YScale, Output, ResidualOut + ) launcher.launch( grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), @@ -1273,7 +1274,9 @@ def launch_fused_add_layernorm_dynamicquant( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = fused_add_layernorm_quant_kernel(Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut) + launcher = fused_add_layernorm_quant_kernel( + Input, ResidualIn, Gamma, Beta, Gamma, YScale, Output, ResidualOut + ) launcher.launch( grid=(m_in, 1, 1), block=(BLOCK_THREADS, 1, 1), diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 0e611de5a..836093073 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -16,7 +16,6 @@ import flydsl.expr as fx from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.typing import Vector as Vec from flydsl.expr.vector import ReductionOp, full from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.kernels_common import dtype_to_elem_type, get_warp_size diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 58aca9be7..1a0eb1953 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -195,7 +195,6 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): beta_ref = beta_dev.to(DTYPE_FP32) if is_smooth: xscale_dev = xscale_t.to(torch_dtype).contiguous() - xscale_ref = xscale_dev.to(DTYPE_FP32) output_dev = torch.empty((M, N), device="cuda", dtype=DTYPE_INT8) yscale_dev = torch.empty((M,), device="cuda", dtype=DTYPE_FP32) @@ -289,8 +288,6 @@ def run_fused_add_test(M: int, N: int, dtype: str = "f32"): beta_dev = beta_t.to(torch_dtype).contiguous() output_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) if dtype == "f32": atol = 1e-4 elif dtype == "f16": @@ -384,11 +381,8 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): gamma_dev = gamma_t.to(torch_dtype).contiguous() beta_dev = beta_t.to(torch_dtype).contiguous() residual_out_dev = torch.empty((M, N), device="cuda", dtype=torch_dtype) - gamma_ref = gamma_dev.to(DTYPE_FP32) - beta_ref = beta_dev.to(DTYPE_FP32) if is_smooth: xscale_dev = xscale_t.to(torch_dtype).contiguous() - xscale_ref = xscale_dev.to(DTYPE_FP32) if dtype == "f32": residual_atol = 1e-4 elif dtype == "f16": From 947d00b52617eef3648e5ec79bb7d85b84230616 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Wed, 17 Jun 2026 13:22:34 +0800 Subject: [PATCH 07/17] replace rsqrt API with fmath.rsqrt --- kernels/layernorm_kernel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index ee00e388c..020bd5253 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -724,7 +724,7 @@ def block_reduce_max(val): mean = sum_val / n_float var = sumsq_val / n_float - mean * mean var = (var < c_zero_f).select(c_zero_f, var) - rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + rstd = fmath.rsqrt(var + eps_c, fastmath=fm_fast) thread_row_max = c_zero_f y_local = [] @@ -813,7 +813,7 @@ def _abs_scalar(val): mean = sum_val / n_float var = sumsq_val / n_float - mean * mean var = (var < c_zero_f).select(c_zero_f, var) - rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + rstd = fmath.rsqrt(var + eps_c, fastmath=fm_fast) thread_row_max = c_zero_f for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): @@ -1083,7 +1083,7 @@ def block_reduce_max(val): mean = sum_val / n_float var = sumsq_val / n_float - mean * mean var = (var < c_zero_f).select(c_zero_f, var) - rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + rstd = fmath.rsqrt(var + eps_c, fastmath=fm_fast) thread_row_max = c_zero_f y_local = [] @@ -1184,7 +1184,7 @@ def _abs_scalar(val): mean = sum_val / n_float var = sumsq_val / n_float - mean * mean var = (var < c_zero_f).select(c_zero_f, var) - rstd = (var + eps_c).rsqrt(fastmath=fm_fast) + rstd = fmath.rsqrt(var + eps_c, fastmath=fm_fast) thread_row_max = c_zero_f for base_idx_int in range_constexpr(0, N, BLOCK_THREADS): From c1b12117d81b724d5659fe75038ec4f0eb8e5795 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 22 Jun 2026 14:53:16 +0800 Subject: [PATCH 08/17] remove M from builder --- kernels/layernorm_kernel.py | 14 ++------------ kernels/rmsnorm_kernel.py | 27 +++++++++------------------ tests/kernels/benchmark_common.py | 4 ++-- tests/kernels/test_layernorm.py | 12 ++++++------ tests/kernels/test_rmsnorm.py | 12 ++++++------ 5 files changed, 25 insertions(+), 44 deletions(-) diff --git a/kernels/layernorm_kernel.py b/kernels/layernorm_kernel.py index 020bd5253..c525fb9d9 100644 --- a/kernels/layernorm_kernel.py +++ b/kernels/layernorm_kernel.py @@ -112,7 +112,7 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_layernorm_module(M: int, N: int, dtype_str: str): +def build_layernorm_module(N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -342,7 +342,7 @@ def launch_layernorm( return launch_layernorm -def build_fused_add_layernorm_module(M: int, N: int, dtype_str: str): +def build_fused_add_layernorm_module(N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -563,7 +563,6 @@ def launch_fused_add_layernorm( def _build_layernorm_quant_module( - M: int, N: int, dtype_str: str, *, @@ -908,7 +907,6 @@ def launch_layernorm_dynamicquant( def _build_fused_add_layernorm_quant_module( - M: int, N: int, dtype_str: str, *, @@ -1287,13 +1285,11 @@ def launch_fused_add_layernorm_dynamicquant( def build_layernorm_dynamicquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_layernorm_quant_module( - M, N, dtype_str, is_smooth=False, @@ -1302,13 +1298,11 @@ def build_layernorm_dynamicquant_module( def build_layernorm_smoothquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_layernorm_quant_module( - M, N, dtype_str, is_smooth=True, @@ -1317,13 +1311,11 @@ def build_layernorm_smoothquant_module( def build_fused_add_layernorm_dynamicquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_fused_add_layernorm_quant_module( - M, N, dtype_str, is_smooth=False, @@ -1332,13 +1324,11 @@ def build_fused_add_layernorm_dynamicquant_module( def build_fused_add_layernorm_smoothquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_fused_add_layernorm_quant_module( - M, N, dtype_str, is_smooth=True, diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 2d8221ec0..17af22fe6 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -110,9 +110,9 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_rmsnorm_module(M: int, N: int, dtype_str: str): - if M > 8192 and N <= 2048: - return _build_rmsnorm_large_m_small_n_module(M, N, dtype_str) +def build_rmsnorm_module(N: int, dtype_str: str): + if N <= 2048: + return _build_rmsnorm_large_m_small_n_module(N, dtype_str) arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -308,7 +308,7 @@ def launch_rmsnorm( return launch_rmsnorm -def _build_rmsnorm_large_m_small_n_module(M: int, N: int, dtype_str: str): +def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str): BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) @@ -321,6 +321,7 @@ def rmsnorm_large_m_small_n_kernel( Gamma: fx.Tensor, _Unused: fx.Tensor, Output: fx.Tensor, + MIn: fx.Int32, ): bid = fx.block_idx.x tid = fx.thread_idx.x @@ -329,7 +330,7 @@ def rmsnorm_large_m_small_n_kernel( row_local = tid // THREADS_PER_ROW row = bid * fx.Int32(BLOCK_M) + row_local - if row < M: + if row < MIn: elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast eps_c = EPS @@ -395,9 +396,9 @@ def launch_rmsnorm_large_m_small_n( m_in: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output) + launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output, m_in) launcher.launch( - grid=((M + BLOCK_M - 1) // BLOCK_M, 1, 1), + grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), block=(BLOCK_THREADS_SPECIAL, 1, 1), stream=stream, ) @@ -405,7 +406,7 @@ def launch_rmsnorm_large_m_small_n( return launch_rmsnorm_large_m_small_n -def build_fused_add_rmsnorm_module(M: int, N: int, dtype_str: str): +def build_fused_add_rmsnorm_module(N: int, dtype_str: str): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -620,7 +621,6 @@ def launch_fused_add_rmsnorm( def _build_rmsnorm_quant_module( - M: int, N: int, dtype_str: str, *, @@ -964,13 +964,11 @@ def launch_rmsnorm_dynamicquant( def build_rmsnorm_dynamicquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_rmsnorm_quant_module( - M, N, dtype_str, is_smooth=False, @@ -979,13 +977,11 @@ def build_rmsnorm_dynamicquant_module( def build_rmsnorm_smoothquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_rmsnorm_quant_module( - M, N, dtype_str, is_smooth=True, @@ -994,7 +990,6 @@ def build_rmsnorm_smoothquant_module( def _build_fused_add_rmsnorm_quant_module( - M: int, N: int, dtype_str: str, *, @@ -1368,13 +1363,11 @@ def launch_fused_add_rmsnorm_dynamicquant( def build_fused_add_rmsnorm_dynamicquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_fused_add_rmsnorm_quant_module( - M, N, dtype_str, is_smooth=False, @@ -1383,13 +1376,11 @@ def build_fused_add_rmsnorm_dynamicquant_module( def build_fused_add_rmsnorm_smoothquant_module( - M: int, N: int, dtype_str: str, quant_dtype_str: str = "i8", ): return _build_fused_add_rmsnorm_quant_module( - M, N, dtype_str, is_smooth=True, diff --git a/tests/kernels/benchmark_common.py b/tests/kernels/benchmark_common.py index 86d6ff79d..05f04026d 100644 --- a/tests/kernels/benchmark_common.py +++ b/tests/kernels/benchmark_common.py @@ -198,7 +198,7 @@ def _bench_flydsl_torch(*, op: str, M: int, N: int, dtype: str, warmup: int, ite if op == "layernorm": from kernels.layernorm_kernel import build_layernorm_module - m = build_layernorm_module(1, N, dtype) + m = build_layernorm_module(N, dtype) exe = flydsl.compile(m) x = torch.randn((M, N), device="cuda", dtype=torch_dtype) gamma = torch.randn((N,), device="cuda", dtype=torch_dtype) @@ -209,7 +209,7 @@ def _bench_flydsl_torch(*, op: str, M: int, N: int, dtype: str, warmup: int, ite if op == "rmsnorm": from kernels.rmsnorm_kernel import build_rmsnorm_module - m = build_rmsnorm_module(1, N, dtype) + m = build_rmsnorm_module(N, dtype) exe = flydsl.compile(m) x = torch.randn((M, N), device="cuda", dtype=torch_dtype) gamma = torch.randn((N,), device="cuda", dtype=torch_dtype) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 1a0eb1953..b64a71dd2 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -90,7 +90,7 @@ def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting LayerNorm (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_layernorm_module(M, N, dtype) + launch_fn = build_layernorm_module(N, dtype) except ValueError as e: print(f"[FAIL] Compile failed: {e}") return False, None @@ -173,9 +173,9 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): try: if is_smooth: - launch_fn = build_layernorm_smoothquant_module(M, N, dtype) + launch_fn = build_layernorm_smoothquant_module(N, dtype) else: - launch_fn = build_layernorm_dynamicquant_module(M, N, dtype) + launch_fn = build_layernorm_dynamicquant_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for {mode} layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None @@ -270,7 +270,7 @@ def run_fused_add_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting FusedAdd LayerNorm (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_fused_add_layernorm_module(M, N, dtype) + launch_fn = build_fused_add_layernorm_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for fused_add layernorm (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None @@ -358,9 +358,9 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): try: if is_smooth: - launch_fn = build_fused_add_layernorm_smoothquant_module(M, N, dtype) + launch_fn = build_fused_add_layernorm_smoothquant_module(N, dtype) else: - launch_fn = build_fused_add_layernorm_dynamicquant_module(M, N, dtype) + launch_fn = build_fused_add_layernorm_dynamicquant_module(N, dtype) except Exception as e: print( f"[FAIL] Compile failed for fused_add {mode} layernorm " diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index b50053cd3..2132f5b1c 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -91,7 +91,7 @@ def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting RMSNorm (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_rmsnorm_module(M, N, dtype) + launch_fn = build_rmsnorm_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for (M={M}, N={N}, dtype={dtype}): {type(e).__name__}: {e}") return False, None @@ -168,9 +168,9 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): try: if is_smooth: - launch_fn = build_rmsnorm_smoothquant_module(M, N, dtype) + launch_fn = build_rmsnorm_smoothquant_module(N, dtype) else: - launch_fn = build_rmsnorm_dynamicquant_module(M, N, dtype) + launch_fn = build_rmsnorm_dynamicquant_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for {mode} (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") return False, None @@ -262,7 +262,7 @@ def run_fused_add_test(M: int, N: int, dtype: str): print(f"\nTesting FusedAdd RMSNorm (M={M}, N={N}, dtype={dtype})") try: - launch_fn = build_fused_add_rmsnorm_module(M, N, dtype) + launch_fn = build_fused_add_rmsnorm_module(N, dtype) except Exception as e: print(f"[FAIL] Compile failed for fused_add rmsnorm (M={M}, N={N}, dtype={dtype}): " f"{type(e).__name__}: {e}") return False, None @@ -359,9 +359,9 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): try: if is_smooth: - launch_fn = build_fused_add_rmsnorm_smoothquant_module(M, N, dtype) + launch_fn = build_fused_add_rmsnorm_smoothquant_module(N, dtype) else: - launch_fn = build_fused_add_rmsnorm_dynamicquant_module(M, N, dtype) + launch_fn = build_fused_add_rmsnorm_dynamicquant_module(N, dtype) except Exception as e: print( f"[FAIL] Compile failed for fused_add rmsnorm {mode} " From 24f5dd35cd41a24b2d0aa37b3ecf697ef1e29fd0 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 22 Jun 2026 16:46:33 +0800 Subject: [PATCH 09/17] use fly.compile, and split large shape from default test configs --- tests/kernels/test_layernorm.py | 105 ++++++++++++++++++++++++------ tests/kernels/test_rmsnorm.py | 111 +++++++++++++++++++++++--------- 2 files changed, 165 insertions(+), 51 deletions(-) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index b64a71dd2..0c491be5e 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -15,6 +15,7 @@ import os +import flydsl.compiler as flyc import pytest from kernels.layernorm_kernel import ( @@ -75,17 +76,21 @@ def _get_layernorm_configs(): configs.append((int(m_s), int(n_s), dt)) else: configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), + (64, 256, "f32"), # f32 aligned + (32, 128, "f16"), # f16 aligned + (64, 2000, "f32"), # unaligned tail handling + (16, 512, "bf16"), # bf16 small shape + (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs +def _get_layernorm_large_configs(): + return [ + (32768, 8192, "bf16"), + ] + + def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting LayerNorm (M={M}, N={N}, dtype={dtype})") @@ -121,9 +126,10 @@ def run_test(M: int, N: int, dtype: str = "f32"): print("Launching kernel...") stream = torch.cuda.current_stream() + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, beta_dev, output_dev, M, stream) def kernel_launch(): - launch_fn(input_dev, gamma_dev, beta_dev, output_dev, M, stream=stream) + compiled_fn(input_dev, gamma_dev, beta_dev, output_dev, M, stream) # One run for correctness visibility, then benchmark via shared harness. kernel_launch() @@ -210,11 +216,17 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): print("Launching kernel...") stream = torch.cuda.current_stream() - def kernel_launch(): - if is_smooth: - launch_fn(input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream=stream) - else: - launch_fn(input_dev, gamma_dev, beta_dev, output_dev, yscale_dev, M, stream=stream) + if is_smooth: + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream) + + def kernel_launch(): + compiled_fn(input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream) + + else: + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, beta_dev, output_dev, yscale_dev, M, stream) + + def kernel_launch(): + compiled_fn(input_dev, gamma_dev, beta_dev, output_dev, yscale_dev, M, stream) kernel_launch() torch.cuda.synchronize() @@ -301,9 +313,20 @@ def run_fused_add_test(M: int, N: int, dtype: str = "f32"): print("Launching kernel...") stream = torch.cuda.current_stream() + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_dev, + gamma_dev, + beta_dev, + output_dev, + residual_out_dev, + M, + stream, + ) def kernel_launch(): - launch_fn(input_dev, residual_dev, gamma_dev, beta_dev, output_dev, residual_out_dev, M, stream=stream) + compiled_fn(input_dev, residual_dev, gamma_dev, beta_dev, output_dev, residual_out_dev, M, stream) kernel_launch() torch.cuda.synchronize() @@ -407,9 +430,23 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): print("Launching kernel...") stream = torch.cuda.current_stream() - def kernel_launch(): - if is_smooth: - launch_fn( + if is_smooth: + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_dev, + gamma_dev, + beta_dev, + xscale_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + + def kernel_launch(): + compiled_fn( input_dev, residual_dev, gamma_dev, @@ -419,10 +456,25 @@ def kernel_launch(): residual_out_dev, yscale_dev, M, - stream=stream, + stream, ) - else: - launch_fn( + + else: + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_dev, + gamma_dev, + beta_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + + def kernel_launch(): + compiled_fn( input_dev, residual_dev, gamma_dev, @@ -431,7 +483,7 @@ def kernel_launch(): residual_out_dev, yscale_dev, M, - stream=stream, + stream, ) kernel_launch() @@ -691,6 +743,17 @@ def test_layernorm(): raise SystemExit(1) +@pytest.mark.large_shape +def test_layernorm_large_shape(): + print("=" * 80) + print("Running LayerNorm Large Shape Tests") + print("=" * 80) + + for M, N, dtype in _get_layernorm_large_configs(): + ok, _ = run_test(M, N, dtype) + assert ok + + def test_fused_add_layernorm(): print("=" * 80) print("Running FusedAdd LayerNorm Tests") diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 2132f5b1c..eedab247d 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -15,6 +15,7 @@ import os +import flydsl.compiler as flyc import pytest from kernels.rmsnorm_kernel import ( @@ -74,19 +75,22 @@ def _get_rmsnorm_configs(): m_s, n_s, dt = [x.strip() for x in p.split(",")] configs.append((int(m_s), int(n_s), dt)) else: - # Prefer N multiples of 2048 to exercise the fast path. configs = [ - # (64, 256, "f32"), # Aligned - # (128, 1024, "f32"), # Aligned - # (32, 128, "f16"), # Aligned - # (64, 2000, "f32"), # Unaligned (tail handling) - # (16, 512, "bf16"), # BF16 - # (1024, 8192, "bf16"), # BF16 - (32768, 8192, "bf16"), + (64, 256, "f32"), # f32 aligned + (32, 128, "f16"), # f16 aligned + (64, 2000, "f32"), # unaligned tail handling + (16, 512, "bf16"), # bf16 small shape + (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs +def _get_rmsnorm_large_configs(): + return [ + (32768, 8192, "bf16"), + ] + + def run_test(M: int, N: int, dtype: str = "f32"): print(f"\nTesting RMSNorm (M={M}, N={N}, dtype={dtype})") @@ -119,9 +123,10 @@ def run_test(M: int, N: int, dtype: str = "f32"): print("Launching kernel...") stream = torch.cuda.current_stream() + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, output_dev, M, stream) def kernel_launch(): - launch_fn(input_dev, gamma_dev, output_dev, M, stream=stream) + compiled_fn(input_dev, gamma_dev, output_dev, M, stream) # run_perftest returns (data, avg_us) _, avg_us = run_perftest( @@ -193,11 +198,17 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): print("Launching kernel...") stream = torch.cuda.current_stream() - def kernel_launch(): - if is_smooth: - launch_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream=stream) - else: - launch_fn(input_dev, gamma_dev, output_dev, yscale_dev, M, stream=stream) + if is_smooth: + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) + + def kernel_launch(): + compiled_fn(input_dev, gamma_dev, xscale_dev, output_dev, yscale_dev, M, stream) + + else: + compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, output_dev, yscale_dev, M, stream) + + def kernel_launch(): + compiled_fn(input_dev, gamma_dev, output_dev, yscale_dev, M, stream) # run_perftest returns (data, avg_us) _, avg_us = run_perftest( @@ -289,17 +300,19 @@ def run_fused_add_test(M: int, N: int, dtype: str): print("Launching kernel...") stream = torch.cuda.current_stream() + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_in_dev, + gamma_dev, + output_dev, + residual_out_dev, + M, + stream, + ) def kernel_launch(): - launch_fn( - input_dev, - residual_in_dev, - gamma_dev, - output_dev, - residual_out_dev, - M, - stream=stream, - ) + compiled_fn(input_dev, residual_in_dev, gamma_dev, output_dev, residual_out_dev, M, stream) _, avg_us = run_perftest( lambda: (kernel_launch(), torch.cuda.synchronize()), @@ -398,9 +411,22 @@ def run_fused_add_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): print("Launching kernel...") stream = torch.cuda.current_stream() - def kernel_launch(): - if is_smooth: - launch_fn( + if is_smooth: + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_in_dev, + gamma_dev, + xscale_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + + def kernel_launch(): + compiled_fn( input_dev, residual_in_dev, gamma_dev, @@ -409,10 +435,24 @@ def kernel_launch(): residual_out_dev, yscale_dev, M, - stream=stream, + stream, ) - else: - launch_fn( + + else: + compiled_fn = flyc.compile( + launch_fn, + input_dev, + residual_in_dev, + gamma_dev, + output_dev, + residual_out_dev, + yscale_dev, + M, + stream, + ) + + def kernel_launch(): + compiled_fn( input_dev, residual_in_dev, gamma_dev, @@ -420,7 +460,7 @@ def kernel_launch(): residual_out_dev, yscale_dev, M, - stream=stream, + stream, ) _, avg_us = run_perftest( @@ -690,6 +730,17 @@ def test_rmsnorm(): raise SystemExit(1) +@pytest.mark.large_shape +def test_rmsnorm_large_shape(): + print("=" * 80) + print("Running RMSNorm Large Shape Tests") + print("=" * 80) + + for M, N, dtype in _get_rmsnorm_large_configs(): + ok, _ = run_test(M, N, dtype) + assert ok + + def test_rmsnorm_dynamicquant(): print("=" * 80) print("Running RMSNorm DynamicQuant Tests") From 084527f363ef430253802274a5ec12a2a4b43e99 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 22 Jun 2026 16:55:55 +0800 Subject: [PATCH 10/17] fix python code style issue --- tests/kernels/test_layernorm.py | 12 +++++++----- tests/kernels/test_rmsnorm.py | 8 ++++---- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 0c491be5e..ba15e803c 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -76,10 +76,10 @@ def _get_layernorm_configs(): configs.append((int(m_s), int(n_s), dt)) else: configs = [ - (64, 256, "f32"), # f32 aligned - (32, 128, "f16"), # f16 aligned - (64, 2000, "f32"), # unaligned tail handling - (16, 512, "bf16"), # bf16 small shape + (64, 256, "f32"), # f32 aligned + (32, 128, "f16"), # f16 aligned + (64, 2000, "f32"), # unaligned tail handling + (16, 512, "bf16"), # bf16 small shape (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs @@ -217,7 +217,9 @@ def run_quant_test(M: int, N: int, dtype: str, *, is_smooth: bool): stream = torch.cuda.current_stream() if is_smooth: - compiled_fn = flyc.compile(launch_fn, input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream) + compiled_fn = flyc.compile( + launch_fn, input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream + ) def kernel_launch(): compiled_fn(input_dev, gamma_dev, beta_dev, xscale_dev, output_dev, yscale_dev, M, stream) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index eedab247d..ecf36d494 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -76,10 +76,10 @@ def _get_rmsnorm_configs(): configs.append((int(m_s), int(n_s), dt)) else: configs = [ - (64, 256, "f32"), # f32 aligned - (32, 128, "f16"), # f16 aligned - (64, 2000, "f32"), # unaligned tail handling - (16, 512, "bf16"), # bf16 small shape + (64, 256, "f32"), # f32 aligned + (32, 128, "f16"), # f16 aligned + (64, 2000, "f32"), # unaligned tail handling + (16, 512, "bf16"), # bf16 small shape (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs From 0e740cf1340a63359ea6207a068390156c26f936 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 22 Jun 2026 17:26:38 +0800 Subject: [PATCH 11/17] fix python style issue --- tests/kernels/test_layernorm.py | 2 +- tests/kernels/test_rmsnorm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index ba15e803c..c4813636d 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -15,9 +15,9 @@ import os -import flydsl.compiler as flyc import pytest +import flydsl.compiler as flyc from kernels.layernorm_kernel import ( build_fused_add_layernorm_dynamicquant_module, build_fused_add_layernorm_module, diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index ecf36d494..023280d57 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -15,9 +15,9 @@ import os -import flydsl.compiler as flyc import pytest +import flydsl.compiler as flyc from kernels.rmsnorm_kernel import ( build_fused_add_rmsnorm_dynamicquant_module, build_fused_add_rmsnorm_module, From 0b2e5356d4646cd3fbd03aa2024cf2f1674af09c Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Wed, 8 Jul 2026 20:09:34 +0800 Subject: [PATCH 12/17] make eps as build function parameter --- kernels/rmsnorm_kernel.py | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 17af22fe6..511129b25 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -22,7 +22,7 @@ KERNEL_NAME = "rmsnorm" -EPS = 1e-5 +DEFAULT_EPS = 1e-5 BLOCK_THREADS = 256 WARP_SIZE = get_warp_size() @@ -110,9 +110,9 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_rmsnorm_module(N: int, dtype_str: str): +def build_rmsnorm_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): if N <= 2048: - return _build_rmsnorm_large_m_small_n_module(N, dtype_str) + return _build_rmsnorm_large_m_small_n_module(N, dtype_str, eps=eps) arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -135,7 +135,7 @@ def rmsnorm_kernel( elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) lds = fx.SharedAllocator().allocate(SharedStorage).peek() @@ -308,7 +308,7 @@ def launch_rmsnorm( return launch_rmsnorm -def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str): +def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) @@ -333,7 +333,7 @@ def rmsnorm_large_m_small_n_kernel( if row < MIn: elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) Input_buf = fx.rocdl.make_buffer_tensor(Input) @@ -406,7 +406,7 @@ def launch_rmsnorm_large_m_small_n( return launch_rmsnorm_large_m_small_n -def build_fused_add_rmsnorm_module(N: int, dtype_str: str): +def build_fused_add_rmsnorm_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -429,7 +429,7 @@ def fused_add_rmsnorm_kernel( elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) lds = fx.SharedAllocator().allocate(SharedStorage).peek() @@ -626,6 +626,7 @@ def _build_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) @@ -649,7 +650,7 @@ def rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -967,12 +968,14 @@ def build_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -980,12 +983,14 @@ def build_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -995,6 +1000,7 @@ def _build_fused_add_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -1023,7 +1029,7 @@ def fused_add_rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1366,12 +1372,14 @@ def build_fused_add_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1379,10 +1387,12 @@ def build_fused_add_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = DEFAULT_EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) From 536b0ba6b1068f9dfef4e94d78cb4ab3c09fc524 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Thu, 9 Jul 2026 23:11:40 +0800 Subject: [PATCH 13/17] WIP rmsnorm: use vec8 loads for aligned bf16/f16 cases --- kernels/rmsnorm_kernel.py | 252 +++++++++++++++++++++++++++----------- 1 file changed, 180 insertions(+), 72 deletions(-) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 511129b25..0eca332b4 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -117,9 +117,13 @@ def build_rmsnorm_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): arch = get_hip_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") - tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + USE_VEC_N = elem_bits <= 16 and N % VEC_WIDTH == 0 + VEC_TILES = N // VEC_WIDTH if USE_VEC_N else 0 + NUM_VEC_ITERS = ( + (VEC_TILES + BLOCK_THREADS - 1) // BLOCK_THREADS if USE_VEC_N else 0 + ) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -188,10 +192,9 @@ def block_reduce_add2(val0, val1): return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) # ================================================================== - # Fast path: N is a multiple of tile_cols + # Fast path: vectorized full-row body for f16/bf16 N aligned to vec8 # ================================================================== - if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): - num_tiles = N // tile_cols + if const_expr(USE_VEC_N): # ── Layout API: buffer-backed tensors + tiled access ───── Input_buf = fx.rocdl.make_buffer_tensor(Input) Output_buf = fx.rocdl.make_buffer_tensor(Output) @@ -212,15 +215,17 @@ def block_reduce_add2(val0, val1): in_local = [] # Pass 1: load + cache + sumsq - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx) + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) in_local.append(vec) x = vec.to(fx.Float32) x2 = x * x red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = thread_sumsq + red2 + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) mean_sq = sum_sq / n_float @@ -228,17 +233,17 @@ def block_reduce_add2(val0, val1): rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) # Pass 2: normalize + gamma + store (reuse cached input) - for tile_i in range_constexpr(num_tiles): + for tile_i in range_constexpr(NUM_VEC_ITERS): idx = tid + tile_i * BLOCK_THREADS - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) - x = in_local[tile_i].to(fx.Float32) + if idx < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + x = in_local[tile_i].to(fx.Float32) - y = (x * rrms) * g - out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + y = (x * rrms) * g + out_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) - out_idx = tid + tile_i * BLOCK_THREADS - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, out_idx) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, out_e, out_div, idx) else: # ============================================================== @@ -309,11 +314,24 @@ def launch_rmsnorm( def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): - BLOCK_N = 1 << (N - 1).bit_length() - BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) - THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) - BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW elem_bits = 32 if dtype_str == "f32" else 16 + USE_VEC_SMALL_N = elem_bits <= 16 and N % VEC_WIDTH == 0 + VEC_TILES = N // VEC_WIDTH if USE_VEC_SMALL_N else 0 + + if USE_VEC_SMALL_N: + BLOCK_M = 1 + THREADS_PER_ROW = max(64, 1 << (VEC_TILES - 1).bit_length()) + BLOCK_THREADS_SPECIAL = THREADS_PER_ROW + else: + BLOCK_N = 1 << (N - 1).bit_length() + BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) + THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) + BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW + + RED_SLOTS = max(1, (BLOCK_THREADS_SPECIAL + WARP_SIZE - 1) // WARP_SIZE) + arch = get_hip_arch() + USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") + SharedStorage = _make_reduction_storage(RED_SLOTS) @flyc.kernel(known_block_size=[BLOCK_THREADS_SPECIAL, 1, 1]) def rmsnorm_large_m_small_n_kernel( @@ -326,67 +344,150 @@ def rmsnorm_large_m_small_n_kernel( bid = fx.block_idx.x tid = fx.thread_idx.x - lane = tid % THREADS_PER_ROW - row_local = tid // THREADS_PER_ROW - row = bid * fx.Int32(BLOCK_M) + row_local + elem_dtype = dtype_to_elem_type(dtype_str) + fm_fast = arith.FastMathFlags.fast + eps_c = eps + n_float = float(N) - if row < MIn: - elem_dtype = dtype_to_elem_type(dtype_str) - fm_fast = arith.FastMathFlags.fast - eps_c = eps - n_float = float(N) + def wave_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))): + off = WARP_SIZE // (2 << _sh_exp) + peer = w.shuffle_xor(off, WARP_SIZE) + w = w.addf(peer, fastmath=fm_fast) + return w - Input_buf = fx.rocdl.make_buffer_tensor(Input) - Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) - Output_buf = fx.rocdl.make_buffer_tensor(Output) + def block_reduce_add2(val0, val1): + if const_expr(RED_SLOTS == 1): + return wave_reduce_add(val0), wave_reduce_add(val1) - row_in = fx.slice(Input_buf, (row, None)) - row_out = fx.slice(Output_buf, (row, None)) + lane = tid % WARP_SIZE + wave = tid // WARP_SIZE - copy_atom_s = fx.make_copy_atom( - fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), - elem_bits, - ) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + s_red = lds.s_red.view(fx.make_layout(RED_SLOTS, 1)) + s_red2 = lds.s_red2.view(fx.make_layout(RED_SLOTS, 1)) - row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) - gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) - out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) + w0 = wave_reduce_add(val0) + w1 = wave_reduce_add(val1) - def group_reduce_add(x): - w = x - for _sh_exp in range_constexpr(int(math.log2(THREADS_PER_ROW))): - off = THREADS_PER_ROW // (2 << _sh_exp) - peer = w.shuffle_xor(off, fx.Int32(THREADS_PER_ROW)) - w = w.addf(peer, fastmath=fm_fast) - return w + if lane == 0: + fx.memref_store(w0, s_red, wave) + fx.memref_store(w1, s_red2, wave) + gpu.barrier() - c_zero_f = fx.Float32(0.0) - thread_sumsq = c_zero_f + if wave == 0: + in_range = lane < RED_SLOTS + lane_safe = in_range.select(lane, 0) + v0 = fx.memref_load(s_red, lane_safe) + v1 = fx.memref_load(s_red2, lane_safe) + ww0 = in_range.select(v0, 0.0) + ww1 = in_range.select(v1, 0.0) + ww0 = wave_reduce_add(ww0) + ww1 = wave_reduce_add(ww1) + + if lane == 0: + fx.memref_store(ww0, s_red, 0) + fx.memref_store(ww1, s_red2, 0) + gpu.barrier() + + return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) + + if const_expr(USE_VEC_SMALL_N): + row = bid + + if row < MIn: + Input_buf = fx.rocdl.make_buffer_tensor(Input) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + + row_in = fx.slice(Input_buf, (row, None)) + row_out = fx.slice(Output_buf, (row, None)) + + in_div = fx.logical_divide(row_in, fx.make_layout(VEC_WIDTH, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(VEC_WIDTH, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) + + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + + c_zero_f = fx.Float32(0.0) + thread_dummy = c_zero_f + is_valid = tid < VEC_TILES + idx_safe = is_valid.select(tid, 0) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) + x = vec.to(fx.Float32) - for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): - idx = lane + base_idx_int - is_valid = idx < N - idx_safe = is_valid.select(idx, 0) - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) - x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) x2 = x * x - thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) + red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = is_valid.select(red2, c_zero_f) - sum_sq = group_reduce_add(thread_sumsq) - mean_sq = sum_sq / n_float - ms_eps = mean_sq + eps_c - rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) + mean_sq = sum_sq / n_float + ms_eps = mean_sq + eps_c + rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) - for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): - idx = lane + base_idx_int - if idx < N: - x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) - g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) - x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) - g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + if tid < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, tid).to(fx.Float32) y = (x * rrms) * g - y_e = _to_elem_scalar(dtype_str, elem_dtype, y) - _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) + y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, tid) + else: + lane = tid % THREADS_PER_ROW + row_local = tid // THREADS_PER_ROW + row = bid * fx.Int32(BLOCK_M) + row_local + + if row < MIn: + Input_buf = fx.rocdl.make_buffer_tensor(Input) + Gamma_buf = fx.rocdl.make_buffer_tensor(Gamma) + Output_buf = fx.rocdl.make_buffer_tensor(Output) + + row_in = fx.slice(Input_buf, (row, None)) + row_out = fx.slice(Output_buf, (row, None)) + + copy_atom_s = fx.make_copy_atom( + fx.rocdl.BufferCopy16b() if elem_bits <= 16 else fx.rocdl.BufferCopy32b(), + elem_bits, + ) + + row_div = fx.logical_divide(row_in, fx.make_layout(1, 1)) + gamma_div = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + out_div = fx.logical_divide(row_out, fx.make_layout(1, 1)) + + def group_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(THREADS_PER_ROW))): + off = THREADS_PER_ROW // (2 << _sh_exp) + peer = w.shuffle_xor(off, fx.Int32(THREADS_PER_ROW)) + w = w.addf(peer, fastmath=fm_fast) + return w + + c_zero_f = fx.Float32(0.0) + thread_sumsq = c_zero_f + + for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): + idx = lane + base_idx_int + is_valid = idx < N + idx_safe = is_valid.select(idx, 0) + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx_safe) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + x2 = x * x + thread_sumsq = thread_sumsq + is_valid.select(x2, c_zero_f) + + sum_sq = group_reduce_add(thread_sumsq) + mean_sq = sum_sq / n_float + ms_eps = mean_sq + eps_c + rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) + + for base_idx_int in range_constexpr(0, BLOCK_N, THREADS_PER_ROW): + idx = lane + base_idx_int + if idx < N: + x_e = _load_scalar(copy_atom_s, elem_dtype, row_div, idx) + g_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div, idx) + x = x_e if dtype_str == "f32" else x_e.to(fx.Float32) + g = g_e if dtype_str == "f32" else g_e.to(fx.Float32) + y = (x * rrms) * g + y_e = _to_elem_scalar(dtype_str, elem_dtype, y) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div, idx, y_e) @flyc.jit def launch_rmsnorm_large_m_small_n( @@ -397,11 +498,18 @@ def launch_rmsnorm_large_m_small_n( stream: fx.Stream = fx.Stream(None), ): launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output, m_in) - launcher.launch( - grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), - block=(BLOCK_THREADS_SPECIAL, 1, 1), - stream=stream, - ) + if USE_VEC_SMALL_N: + launcher.launch( + grid=(m_in, 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) + else: + launcher.launch( + grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) return launch_rmsnorm_large_m_small_n From 6a55fbaf5453484236a73bc9012825aeb65aae1f Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 13 Jul 2026 16:49:08 +0800 Subject: [PATCH 14/17] each block process multiple rows in aligned small-N vec8 path --- kernels/rmsnorm_kernel.py | 74 +++++++++++++++++++++++---------------- 1 file changed, 44 insertions(+), 30 deletions(-) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 0eca332b4..5b1fb2b92 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -319,14 +319,17 @@ def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, eps: float = D VEC_TILES = N // VEC_WIDTH if USE_VEC_SMALL_N else 0 if USE_VEC_SMALL_N: - BLOCK_M = 1 - THREADS_PER_ROW = max(64, 1 << (VEC_TILES - 1).bit_length()) - BLOCK_THREADS_SPECIAL = THREADS_PER_ROW + BLOCK_N = 1 << (N - 1).bit_length() + BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) + THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) + BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW + NUM_VEC_ROW_ITERS = (VEC_TILES + THREADS_PER_ROW - 1) // THREADS_PER_ROW else: BLOCK_N = 1 << (N - 1).bit_length() BLOCK_M = max(min(16384 // BLOCK_N, 32), 8) THREADS_PER_ROW = min(WARP_SIZE, 1024 // BLOCK_M) BLOCK_THREADS_SPECIAL = BLOCK_M * THREADS_PER_ROW + NUM_VEC_ROW_ITERS = 0 RED_SLOTS = max(1, (BLOCK_THREADS_SPECIAL + WARP_SIZE - 1) // WARP_SIZE) arch = get_hip_arch() @@ -394,7 +397,9 @@ def block_reduce_add2(val0, val1): return fx.memref_load(s_red, 0), fx.memref_load(s_red2, 0) if const_expr(USE_VEC_SMALL_N): - row = bid + lane = tid % THREADS_PER_ROW + row_local = tid // THREADS_PER_ROW + row = bid * fx.Int32(BLOCK_M) + row_local if row < MIn: Input_buf = fx.rocdl.make_buffer_tensor(Input) @@ -411,26 +416,42 @@ def block_reduce_add2(val0, val1): copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) c_zero_f = fx.Float32(0.0) - thread_dummy = c_zero_f - is_valid = tid < VEC_TILES - idx_safe = is_valid.select(tid, 0) - vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) - x = vec.to(fx.Float32) + thread_sumsq = c_zero_f + in_local = [] - x2 = x * x - red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) - thread_sumsq = is_valid.select(red2, c_zero_f) + def group_reduce_add(x): + w = x + for _sh_exp in range_constexpr(int(math.log2(THREADS_PER_ROW))): + off = THREADS_PER_ROW // (2 << _sh_exp) + peer = w.shuffle_xor(off, fx.Int32(THREADS_PER_ROW)) + w = w.addf(peer, fastmath=fm_fast) + return w - _, sum_sq = block_reduce_add2(thread_dummy, thread_sumsq) + for tile_i in range_constexpr(NUM_VEC_ROW_ITERS): + idx = lane + tile_i * THREADS_PER_ROW + is_valid = idx < VEC_TILES + idx_safe = is_valid.select(idx, 0) + vec = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, in_div, idx_safe) + in_local.append(vec) + x = vec.to(fx.Float32) + + x2 = x * x + red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) + thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) + + sum_sq = group_reduce_add(thread_sumsq) mean_sq = sum_sq / n_float ms_eps = mean_sq + eps_c rrms = fmath.rsqrt(ms_eps, fastmath=fm_fast) - if tid < VEC_TILES: - g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, tid).to(fx.Float32) - y = (x * rrms) * g - y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) - _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, tid) + for tile_i in range_constexpr(NUM_VEC_ROW_ITERS): + idx = lane + tile_i * THREADS_PER_ROW + if idx < VEC_TILES: + g = _load_vec(copy_atom, VEC_WIDTH, elem_dtype, gamma_div, idx).to(fx.Float32) + x = in_local[tile_i].to(fx.Float32) + y = (x * rrms) * g + y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) + _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, idx) else: lane = tid % THREADS_PER_ROW row_local = tid // THREADS_PER_ROW @@ -498,18 +519,11 @@ def launch_rmsnorm_large_m_small_n( stream: fx.Stream = fx.Stream(None), ): launcher = rmsnorm_large_m_small_n_kernel(Input, Gamma, Gamma, Output, m_in) - if USE_VEC_SMALL_N: - launcher.launch( - grid=(m_in, 1, 1), - block=(BLOCK_THREADS_SPECIAL, 1, 1), - stream=stream, - ) - else: - launcher.launch( - grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), - block=(BLOCK_THREADS_SPECIAL, 1, 1), - stream=stream, - ) + launcher.launch( + grid=((m_in + fx.Int32(BLOCK_M - 1)) // fx.Int32(BLOCK_M), 1, 1), + block=(BLOCK_THREADS_SPECIAL, 1, 1), + stream=stream, + ) return launch_rmsnorm_large_m_small_n From fa4ba4cabc2b8f48014731ba76a58efa095baea6 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Mon, 13 Jul 2026 18:05:51 +0800 Subject: [PATCH 15/17] test_rmsnorm: add unaligned cases for tuning --- tests/kernels/test_rmsnorm.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 023280d57..d34addfae 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -76,11 +76,21 @@ def _get_rmsnorm_configs(): configs.append((int(m_s), int(n_s), dt)) else: configs = [ - (64, 256, "f32"), # f32 aligned - (32, 128, "f16"), # f16 aligned - (64, 2000, "f32"), # unaligned tail handling - (16, 512, "bf16"), # bf16 small shape - (64, 8192, "bf16"), # bf16 fast-path N with small M + (16, 512, "bf16"), + (1024, 512, "bf16"), + (32768, 512, "bf16"), + (16, 2000, "bf16"), + (1024, 2000, "bf16"), + (32768, 2000, "bf16"), + (16, 2001, "bf16"), + (1024, 2001, "bf16"), + (32768, 2001, "bf16"), + (16, 3072, "bf16"), + (1024, 3072, "bf16"), + (32768, 3072, "bf16"), + (16, 8192, "bf16"), + (1024, 8192, "bf16"), + (32768, 8192, "bf16"), ] return configs From e6f20456ff93a9701f9d8bbb1e1fe8941d2bb68f Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Fri, 17 Jul 2026 08:11:27 +0000 Subject: [PATCH 16/17] rmsnorm: support unaligned small-N vector path --- kernels/rmsnorm_kernel.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/kernels/rmsnorm_kernel.py b/kernels/rmsnorm_kernel.py index 5b1fb2b92..0eadab3e3 100644 --- a/kernels/rmsnorm_kernel.py +++ b/kernels/rmsnorm_kernel.py @@ -315,8 +315,10 @@ def launch_rmsnorm( def _build_rmsnorm_large_m_small_n_module(N: int, dtype_str: str, eps: float = DEFAULT_EPS): elem_bits = 32 if dtype_str == "f32" else 16 - USE_VEC_SMALL_N = elem_bits <= 16 and N % VEC_WIDTH == 0 + USE_VEC_SMALL_N = elem_bits <= 16 and N >= VEC_WIDTH VEC_TILES = N // VEC_WIDTH if USE_VEC_SMALL_N else 0 + VEC_ELEMS = VEC_TILES * VEC_WIDTH + TAIL_ELEMS = N - VEC_ELEMS if USE_VEC_SMALL_N: BLOCK_N = 1 << (N - 1).bit_length() @@ -414,6 +416,11 @@ def block_reduce_add2(val0, val1): out_div = fx.logical_divide(row_out, fx.make_layout(VEC_WIDTH, 1)) copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + if const_expr(TAIL_ELEMS > 0): + row_div_s = fx.logical_divide(row_in, fx.make_layout(1, 1)) + gamma_div_s = fx.logical_divide(Gamma_buf, fx.make_layout(1, 1)) + out_div_s = fx.logical_divide(row_out, fx.make_layout(1, 1)) + copy_atom_s = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), elem_bits) c_zero_f = fx.Float32(0.0) thread_sumsq = c_zero_f @@ -439,6 +446,13 @@ def group_reduce_add(x): red2 = x2.reduce(ReductionOp.ADD, fastmath=fm_fast) thread_sumsq = thread_sumsq + is_valid.select(red2, c_zero_f) + if const_expr(TAIL_ELEMS > 0): + if lane < TAIL_ELEMS: + tail_idx = lane + VEC_ELEMS + x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + thread_sumsq = thread_sumsq + (x_tail * x_tail) + sum_sq = group_reduce_add(thread_sumsq) mean_sq = sum_sq / n_float ms_eps = mean_sq + eps_c @@ -452,6 +466,17 @@ def group_reduce_add(x): y = (x * rrms) * g y_e = _to_elem_vec(dtype_str, elem_dtype, USE_HW_CVT_PK_BF16_F32, y) _store_vec(copy_atom, VEC_WIDTH, elem_dtype, y_e, out_div, idx) + + if const_expr(TAIL_ELEMS > 0): + if lane < TAIL_ELEMS: + tail_idx = lane + VEC_ELEMS + x_tail_e = _load_scalar(copy_atom_s, elem_dtype, row_div_s, tail_idx) + g_tail_e = _load_scalar(copy_atom_s, elem_dtype, gamma_div_s, tail_idx) + x_tail = x_tail_e.to(fx.Float32) + g_tail = g_tail_e.to(fx.Float32) + y_tail = (x_tail * rrms) * g_tail + y_tail_e = _to_elem_scalar(dtype_str, elem_dtype, y_tail) + _store_scalar(copy_atom_s, elem_dtype, elem_dtype, out_div_s, tail_idx, y_tail_e) else: lane = tid % THREADS_PER_ROW row_local = tid // THREADS_PER_ROW From 4172c758fdc444b8bc323d7161f95a806f807393 Mon Sep 17 00:00:00 2001 From: Junlin Chen Date: Tue, 21 Jul 2026 13:56:40 +0800 Subject: [PATCH 17/17] norm: align layernorm and rmsnorm eps/vector coverage --- kernels/norm/layernorm_kernel.py | 54 ++++++++++++++++++++------------ kernels/norm/rmsnorm_kernel.py | 14 +++++++-- tests/kernels/test_layernorm.py | 46 +++++++++++++++++++++++++-- tests/kernels/test_rmsnorm.py | 2 ++ 4 files changed, 92 insertions(+), 24 deletions(-) diff --git a/kernels/norm/layernorm_kernel.py b/kernels/norm/layernorm_kernel.py index b7be89f8a..2798ff0ec 100644 --- a/kernels/norm/layernorm_kernel.py +++ b/kernels/norm/layernorm_kernel.py @@ -6,8 +6,8 @@ LayerNorm(x) = (x - mean) / sqrt(var + eps) * gamma + beta Two paths: - - Fast path (N == BLOCK_THREADS * VEC_WIDTH * 4): vectorised tiled copy, - register caching, pipelined gamma/beta loads. + - Fast path (N is a multiple of BLOCK_THREADS * VEC_WIDTH): vectorised + tiled copy, register caching, pipelined gamma/beta loads. - Generic path (arbitrary N): scalar 2-pass implementation. """ @@ -112,12 +112,13 @@ def _quant_dtype_max(dtype_str: str) -> float: raise ValueError(f"unsupported quant dtype: {dtype_str!r} (expected 'i8' or 'int8')") -def build_layernorm_module(N: int, dtype_str: str): +def build_layernorm_module(N: int, dtype_str: str, eps: float = EPS): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -134,7 +135,7 @@ def layernorm_kernel( elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps lds = fx.SharedAllocator().allocate(SharedStorage).peek() s_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1)) @@ -194,12 +195,12 @@ def compute_mean_rstd(sum_val, sumsq_val): return mean, rstd # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # Uses buffer_load / buffer_store for high-bandwidth vectorised # memory access (same approach as preshuffle_gemm). # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -342,12 +343,13 @@ def launch_layernorm( return launch_layernorm -def build_fused_add_layernorm_module(N: int, dtype_str: str): +def build_fused_add_layernorm_module(N: int, dtype_str: str, eps: float = EPS): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -365,7 +367,7 @@ def fused_add_layernorm_kernel( elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps lds = fx.SharedAllocator().allocate(SharedStorage).peek() s_sum = lds.s_sum.view(fx.make_layout(RED_SLOTS, 1)) @@ -419,10 +421,10 @@ def compute_mean_rstd(sum_val, sumsq_val): return mean, fmath.rsqrt(var + eps_c, fastmath=fm_fast) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols c_zero_f = fx.Float32(0.0) thread_sum = c_zero_f thread_sumsq = c_zero_f @@ -568,9 +570,11 @@ def _build_layernorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -591,7 +595,7 @@ def layernorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -676,10 +680,10 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -912,12 +916,14 @@ def _build_fused_add_layernorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) elem_bits = 32 if dtype_str == "f32" else 16 + tile_cols = BLOCK_THREADS * VEC_WIDTH quant_dtype_max = _quant_dtype_max(quant_dtype_str) SharedStorage = _make_reduction_storage(RED_SLOTS) @@ -940,7 +946,7 @@ def fused_add_layernorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1025,10 +1031,10 @@ def block_reduce_max(val): return fx.memref_load(s_sum, 0) # ================================================================== - # Fast path: N == BLOCK_THREADS * VEC_WIDTH * 4 + # Fast path: N is a multiple of BLOCK_THREADS * VEC_WIDTH # ================================================================== - if const_expr(N == (BLOCK_THREADS * VEC_WIDTH * 4) and elem_bits <= 16): - num_tiles_py = 4 + if const_expr(N >= tile_cols and N % tile_cols == 0 and elem_bits <= 16): + num_tiles_py = N // tile_cols quant_half_width = VEC_WIDTH // 2 abs_mask = full(VEC_WIDTH, fx.Uint32(0x7FFFFFFF), fx.Uint32) @@ -1288,12 +1294,14 @@ def build_layernorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_layernorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1301,12 +1309,14 @@ def build_layernorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_layernorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1314,12 +1324,14 @@ def build_fused_add_layernorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_layernorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1327,10 +1339,12 @@ def build_fused_add_layernorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_layernorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/kernels/norm/rmsnorm_kernel.py b/kernels/norm/rmsnorm_kernel.py index b9de9eb86..776d477d0 100644 --- a/kernels/norm/rmsnorm_kernel.py +++ b/kernels/norm/rmsnorm_kernel.py @@ -782,6 +782,7 @@ def _build_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): tile_cols = BLOCK_THREADS * VEC_WIDTH RED_SLOTS = max(1, (BLOCK_THREADS + WARP_SIZE - 1) // WARP_SIZE) @@ -805,7 +806,7 @@ def rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1123,12 +1124,14 @@ def build_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1136,12 +1139,14 @@ def build_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1151,6 +1156,7 @@ def _build_fused_add_rmsnorm_quant_module( *, is_smooth: bool, quant_dtype_str: str = "i8", + eps: float = EPS, ): arch = get_rocm_arch() USE_HW_CVT_PK_BF16_F32 = (arch == "gfx950") or str(arch).startswith("gfx95") @@ -1179,7 +1185,7 @@ def fused_add_rmsnorm_quant_kernel( quant_dtype = _quant_dtype_to_elem_type(quant_dtype_str) fm_fast = arith.FastMathFlags.fast - eps_c = EPS + eps_c = eps n_float = float(N) c_zero_f = fx.Float32(0.0) c_one_f = fx.Float32(1.0) @@ -1522,12 +1528,14 @@ def build_fused_add_rmsnorm_dynamicquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=False, quant_dtype_str=quant_dtype_str, + eps=eps, ) @@ -1535,12 +1543,14 @@ def build_fused_add_rmsnorm_smoothquant_module( N: int, dtype_str: str, quant_dtype_str: str = "i8", + eps: float = EPS, ): return _build_fused_add_rmsnorm_quant_module( N, dtype_str, is_smooth=True, quant_dtype_str=quant_dtype_str, + eps=eps, ) diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py index 2e00795f9..bc40ecf9d 100644 --- a/tests/kernels/test_layernorm.py +++ b/tests/kernels/test_layernorm.py @@ -80,6 +80,8 @@ def _get_layernorm_configs(): (32, 128, "f16"), # f16 aligned (64, 2000, "f32"), # unaligned tail handling (16, 512, "bf16"), # bf16 small shape + (64, 4096, "f16"), # f16 vectorized fast path below 8192 + (64, 8192, "f16"), # f16 fast-path N (64, 8192, "bf16"), # bf16 fast-path N with small M ] return configs @@ -545,13 +547,13 @@ def kernel_launch(): return ok, flydsl_gpu_us -def _reference_layernorm(input_dev, gamma_dev, beta_dev): +def _reference_layernorm(input_dev, gamma_dev, beta_dev, eps: float = EPS): x = input_dev.to(DTYPE_FP32) gamma = gamma_dev.to(DTYPE_FP32) beta = beta_dev.to(DTYPE_FP32) mean = x.mean(dim=1, keepdim=True) var = x.var(dim=1, keepdim=True, unbiased=False) - return ((x - mean) / torch.sqrt(var + EPS) * gamma + beta).to(DTYPE_FP32) + return ((x - mean) / torch.sqrt(var + eps) * gamma + beta).to(DTYPE_FP32) def _reference_layernorm_quant(input_dev, gamma_dev, beta_dev, *, xscale_dev=None): @@ -971,8 +973,48 @@ def test_fused_add_layernorm_smoothquant(): raise SystemExit(1) +def test_layernorm_eps_honored(): + """eps must be baked into LayerNorm kernels, matching RMSNorm behavior.""" + print("=" * 80) + print("Running LayerNorm eps-honored Test") + print("=" * 80) + torch.manual_seed(0) + + for M, N in ((32, 256), (32, 3000)): + x = torch.randn((M, N), device="cuda", dtype=DTYPE_FP32).contiguous() + gamma = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).contiguous() + beta = torch.rand((N,), device="cuda", dtype=DTYPE_FP32).contiguous() + out = torch.empty_like(x) + stream = torch.cuda.current_stream() + + for eps in (1e-5, 1e-6, 1e-2): + launch_fn = build_layernorm_module(N, "f32", eps=eps) + compiled_fn = flyc.compile(launch_fn, x, gamma, beta, out, M, stream) + compiled_fn(x, gamma, beta, out, M, stream) + torch.cuda.synchronize() + + ref = _reference_layernorm(x, gamma, beta, eps=eps) + err = (out - ref).abs().max().item() + print(f" N={N} eps={eps:g}: max err vs torch ref = {err:.3e}") + assert err < 1e-4, f"N={N} eps={eps} not honored (err={err})" + + y_large_eps = torch.empty_like(x) + y_small_eps = torch.empty_like(x) + compiled_large = flyc.compile(build_layernorm_module(N, "f32", eps=1e-2), x, gamma, beta, y_large_eps, M, stream) + compiled_small = flyc.compile(build_layernorm_module(N, "f32", eps=1e-6), x, gamma, beta, y_small_eps, M, stream) + compiled_large(x, gamma, beta, y_large_eps, M, stream) + compiled_small(x, gamma, beta, y_small_eps, M, stream) + torch.cuda.synchronize() + diff = (y_large_eps - y_small_eps).abs().max().item() + print(f" N={N} eps 1e-2 vs 1e-6 output diff = {diff:.3e} (must be > 0)") + assert diff > 0, f"N={N}: eps appears to be ignored" + + print(" -> PASSED") + + if __name__ == "__main__": test_layernorm() + test_layernorm_eps_honored() test_fused_add_layernorm() test_layernorm_dynamicquant() test_layernorm_smoothquant() diff --git a/tests/kernels/test_rmsnorm.py b/tests/kernels/test_rmsnorm.py index 4d9b424c0..8464e2f89 100644 --- a/tests/kernels/test_rmsnorm.py +++ b/tests/kernels/test_rmsnorm.py @@ -128,6 +128,8 @@ def _get_rmsnorm_configs(): (32, 128, "f16"), # f16 aligned (64, 2000, "f32"), # unaligned tail handling (16, 512, "bf16"), # bf16 small shape + (64, 4096, "f16"), # f16 vectorized fast path below 8192 + (64, 8192, "f16"), # f16 fast-path N (64, 8192, "bf16"), # bf16 fast-path N with small M ]