diff --git a/benchmark/bench_rmsnorm.py b/benchmark/bench_rmsnorm.py new file mode 100644 index 000000000..184db4668 --- /dev/null +++ b/benchmark/bench_rmsnorm.py @@ -0,0 +1,579 @@ +import itertools +import os + +import pandas as pd +import sgl_kernel +import torch +import triton + +# Supported dtypes for benchmarking +DTYPE_MAP = { + "fp16": torch.float16, + "bf16": torch.bfloat16, +} +DTYPE_BYTES = { + "fp16": 2, + "bf16": 2, +} + + +def make_3d_input(batch_size, seq_len, hidden_size, dtype): + """Create a 3D input tensor for row-wise RMSNorm benchmarks.""" + return torch.randn( + batch_size, + seq_len, + hidden_size, + device=torch.device("xpu"), + dtype=DTYPE_MAP[dtype], + ) + + +def make_non_flattenable_3d(num_tokens, num_heads, head_dim, dtype): + """Create a non-flattenable 3D tensor mimicking a QKV slice pattern.""" + assert num_tokens > 1 + total_heads = num_heads + 4 + full = torch.randn( + num_tokens, + total_heads * head_dim, + device=torch.device("xpu"), + dtype=DTYPE_MAP[dtype], + ) + q_flat = full[:, : num_heads * head_dim] + x = q_flat.unflatten(-1, (num_heads, head_dim)) + assert x.stride(0) != x.size(1) * x.stride(1) + return x + + +def rms_norm(x, w, eps=1e-6): + """PyTorch reference implementation of RMSNorm.""" + orig_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + x = x * w.to(torch.float32) + x = x.to(orig_dtype) + return x + + +def fused_add_rms_norm(x, residual, w, eps=1e-6): + """PyTorch reference implementation of RMSNorm fused with residual add.""" + orig_dtype = x.dtype + x = x.to(torch.float32) + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + x = (x * w.to(torch.float32)).to(orig_dtype) + return x, residual + + +def gemma_rms_norm(x, w, eps=1e-6): + """PyTorch reference implementation of Gemma-style RMSNorm.""" + orig_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + x = x * (1.0 + w.to(torch.float32)) + x = x.to(orig_dtype) + return x + + +def gemma_fused_add_rms_norm(x, residual, w, eps=1e-6): + """PyTorch reference implementation of Gemma-style RMSNorm fused with residual add.""" + orig_dtype = x.dtype + x = x.to(torch.float32) + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + x = x.to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + x = x * (1.0 + w.to(torch.float32)) + x = x.to(orig_dtype) + return x, residual + + +# Benchmark configurations +batch_size_range = [1, 19, 99, 989, 1989] +hidden_size_range = [16, 32, 111, 500, 1024, 4096, 8192] +dtype_range = ["fp16", "bf16"] + +norm_configs = list(itertools.product(batch_size_range, hidden_size_range, dtype_range)) +three_d_norm_configs = list( + itertools.product([1, 4, 19], [1, 7, 32], [111, 1024, 4096], dtype_range) +) +non_flattenable_3d_norm_configs = list( + itertools.product([7, 32], [4, 8], [64, 128], dtype_range) +) + +rmsnorm_results = [] +fused_add_rmsnorm_results = [] +gemma_rmsnorm_results = [] +gemma_fused_add_rmsnorm_results = [] +rmsnorm_3d_results = [] +gemma_rmsnorm_non_flattenable_3d_results = [] + + +def calculate_norm_flops(M, N, fused_add=False): + """FLOPs per RMSNorm call: 4*M*N (square, mean, rsqrt, weight-mul), plus + M*N for the residual add in fused-add variants.""" + flops = 4 * M * N + if fused_add: + flops += M * N + return flops + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hidden_size", "dtype"], + x_vals=norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="rmsnorm-performance", + args={}, + ) +) +def benchmark_rmsnorm(batch_size, hidden_size, dtype, provider): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = torch.randn(batch_size, hidden_size, device=device, dtype=torch_dtype) + w = torch.randn(hidden_size, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: rms_norm(x, w, eps) + elif provider == "sglang": + fn = lambda: sgl_kernel.rmsnorm(x, w, eps) + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + + # GB/s = logical_tensor_size / time + total_bytes = (2 * batch_size * hidden_size + hidden_size) * DTYPE_BYTES[dtype] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops(batch_size, hidden_size, fused_add=False) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + rmsnorm_results.append( + { + "batch_size": batch_size, + "hidden_size": hidden_size, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hidden_size", "dtype"], + x_vals=norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="fused-add-rmsnorm-performance", + args={}, + ) +) +def benchmark_fused_add_rmsnorm(batch_size, hidden_size, dtype, provider): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = torch.randn(batch_size, hidden_size, device=device, dtype=torch_dtype) + residual = torch.randn_like(x) + w = torch.randn(hidden_size, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: fused_add_rms_norm(x.clone(), residual.clone(), w, eps) + elif provider == "sglang": + + def fn(): + x_fused = x.clone() + residual_fused = residual.clone() + sgl_kernel.fused_add_rmsnorm(x_fused, residual_fused, w, eps) + + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + + # GB/s = logical_tensor_size / time + total_bytes = (4 * batch_size * hidden_size + hidden_size) * DTYPE_BYTES[dtype] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops(batch_size, hidden_size, fused_add=True) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + fused_add_rmsnorm_results.append( + { + "batch_size": batch_size, + "hidden_size": hidden_size, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hidden_size", "dtype"], + x_vals=norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="gemma-rmsnorm-performance", + args={}, + ) +) +def benchmark_gemma_rmsnorm(batch_size, hidden_size, dtype, provider): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = torch.randn(batch_size, hidden_size, device=device, dtype=torch_dtype) + w = torch.randn(hidden_size, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: gemma_rms_norm(x, w, eps) + elif provider == "sglang": + fn = lambda: sgl_kernel.gemma_rmsnorm(x, w, eps) + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + + # GB/s = logical_tensor_size / time + total_bytes = (2 * batch_size * hidden_size + hidden_size) * DTYPE_BYTES[dtype] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops(batch_size, hidden_size, fused_add=False) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + gemma_rmsnorm_results.append( + { + "batch_size": batch_size, + "hidden_size": hidden_size, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hidden_size", "dtype"], + x_vals=norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="gemma-fused-add-rmsnorm-performance", + args={}, + ) +) +def benchmark_gemma_fused_add_rmsnorm(batch_size, hidden_size, dtype, provider): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = torch.randn(batch_size, hidden_size, device=device, dtype=torch_dtype) + residual = torch.randn_like(x) + w = torch.randn(hidden_size, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: gemma_fused_add_rms_norm(x.clone(), residual.clone(), w, eps) + elif provider == "sglang": + + def fn(): + x_fused = x.clone() + residual_fused = residual.clone() + sgl_kernel.gemma_fused_add_rmsnorm(x_fused, residual_fused, w, eps) + + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + + # GB/s = logical_tensor_size / time + total_bytes = (4 * batch_size * hidden_size + hidden_size) * DTYPE_BYTES[dtype] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops(batch_size, hidden_size, fused_add=True) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + gemma_fused_add_rmsnorm_results.append( + { + "batch_size": batch_size, + "hidden_size": hidden_size, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len", "hidden_size", "dtype"], + x_vals=three_d_norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="rmsnorm-3d-performance", + args={}, + ) +) +def benchmark_rmsnorm_3d(batch_size, seq_len, hidden_size, dtype, provider): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = make_3d_input(batch_size, seq_len, hidden_size, dtype) + w = torch.randn(hidden_size, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: rms_norm(x, w, eps) + elif provider == "sglang": + fn = lambda: sgl_kernel.rmsnorm(x, w, eps) + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + rows = batch_size * seq_len + total_bytes = (2 * rows * hidden_size + hidden_size) * DTYPE_BYTES[dtype] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops(rows, hidden_size, fused_add=False) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + rmsnorm_3d_results.append( + { + "batch_size": batch_size, + "seq_len": seq_len, + "hidden_size": hidden_size, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "num_heads", "head_dim", "dtype"], + x_vals=non_flattenable_3d_norm_configs, + line_arg="provider", + line_vals=["torch", "sglang"], + line_names=["PyTorch", "SGL Kernel"], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="gemma-rmsnorm-non-flattenable-3d-performance", + args={}, + ) +) +def benchmark_gemma_rmsnorm_non_flattenable_3d( + num_tokens, num_heads, head_dim, dtype, provider +): + device = torch.device("xpu") + torch_dtype = DTYPE_MAP[dtype] + + x = make_non_flattenable_3d(num_tokens, num_heads, head_dim, dtype) + w = torch.randn(head_dim, device=device, dtype=torch_dtype) + eps = 1e-6 + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + fn = lambda: gemma_rms_norm(x, w, eps) + elif provider == "sglang": + fn = lambda: sgl_kernel.gemma_rmsnorm(x, w, eps) + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench( + fn, warmup=50, rep=200, quantiles=quantiles + ) + total_bytes = (2 * num_tokens * num_heads * head_dim + head_dim) * DTYPE_BYTES[ + dtype + ] + bandwidth_gbs = (total_bytes / 1e9) / (ms / 1000.0) + total_flops = calculate_norm_flops( + num_tokens * num_heads, head_dim, fused_add=False + ) + gflops = (total_flops / 1e9) / (ms / 1000.0) + + gemma_rmsnorm_non_flattenable_3d_results.append( + { + "num_tokens": num_tokens, + "num_heads": num_heads, + "head_dim": head_dim, + "dtype": dtype, + "provider": provider, + "time_us": 1000 * ms, + "bandwidth_gbs": bandwidth_gbs, + "total_bytes": total_bytes, + "total_flops": total_flops, + "gflops": gflops, + } + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +def _save_and_report(results, name, out_dir): + df = pd.DataFrame(results) + if df.empty: + print(f"{name}: no benchmark results collected") + return + + out_csv = os.path.join(out_dir, f"{name}.csv") + df.to_csv(out_csv, index=False) + print(f"Wrote results CSV: {out_csv}") + + df["time_us"] = df["time_us"].round(2) + df["bandwidth_gbs"] = df["bandwidth_gbs"].round(2) + df["total_bytes"] = df["total_bytes"].round(2) + df["total_flops"] = df["total_flops"].round(2) + df["gflops"] = df["gflops"].round(2) + print(df.to_markdown(index=False)) + + index_cols = [ + col for col in ["batch_size", "hidden_size", "dtype"] if col in df.columns + ] + if not index_cols: + index_cols = [ + col + for col in ["batch_size", "seq_len", "hidden_size", "dtype"] + if col in df.columns + ] + if not index_cols: + index_cols = [ + col + for col in ["num_tokens", "num_heads", "head_dim", "dtype"] + if col in df.columns + ] + if not index_cols: + return + + speed_pivot = df.pivot_table(index=index_cols, columns="provider", values="time_us") + if "torch" in speed_pivot.columns and "sglang" in speed_pivot.columns: + speed_pivot["speedup"] = speed_pivot["torch"] / speed_pivot["sglang"] + + avg_speedup = speed_pivot["speedup"].mean() + + print(f"\n{name} avg speedup: {avg_speedup:.2f}x") + print(f"{name} median speedup: {speed_pivot['speedup'].median():.2f}x") + print(f"{name} max speedup: {speed_pivot['speedup'].max():.2f}x") + print(f"{name} min speedup: {speed_pivot['speedup'].min():.2f}x") + + above_avg_count = (speed_pivot["speedup"] > avg_speedup).sum() + print(f"{name} speedups above avg: {above_avg_count}/{len(speed_pivot)}") + + +if __name__ == "__main__": + print("Running RMSNorm benchmarks...") + benchmark_rmsnorm.run(print_data=True) + benchmark_fused_add_rmsnorm.run(print_data=True) + benchmark_gemma_rmsnorm.run(print_data=True) + benchmark_gemma_fused_add_rmsnorm.run(print_data=True) + benchmark_rmsnorm_3d.run(print_data=True) + benchmark_gemma_rmsnorm_non_flattenable_3d.run(print_data=True) + + out_dir = "benchmark/results" + os.makedirs(out_dir, exist_ok=True) + + print("\n" + "=" * 80) + print("RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report(rmsnorm_results, "rmsnorm", out_dir) + + print("\n" + "=" * 80) + print("Fused-Add RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report(fused_add_rmsnorm_results, "fused_add_rmsnorm", out_dir) + print("\n" + "=" * 80) + print("3D RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report(rmsnorm_3d_results, "rmsnorm_3d", out_dir) + + print("\n" + "=" * 80) + print("Non-Flattenable 3D Gemma RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report( + gemma_rmsnorm_non_flattenable_3d_results, + "gemma_rmsnorm_non_flattenable_3d", + out_dir, + ) + print("\n" + "=" * 80) + print("Gemma RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report(gemma_rmsnorm_results, "gemma_rmsnorm", out_dir) + + print("\n" + "=" * 80) + print("Gemma Fused-Add RMSNorm Benchmark Results") + print("=" * 80) + _save_and_report( + gemma_fused_add_rmsnorm_results, "gemma_fused_add_rmsnorm", out_dir + ) diff --git a/src/sycl/Norm.h b/src/sycl/Norm.h index 6330c2d71..7aeea9bc4 100644 --- a/src/sycl/Norm.h +++ b/src/sycl/Norm.h @@ -391,9 +391,11 @@ class NormForward { return vec_size; } - int get_update_vec_size(int Plane, int vec_size) { - vec_size = get_min_vec_size(vec_size, X_data, Y_data, gamma_data, beta_data); - + // Aligns vec_size against an arbitrary set of pointers (e.g. X/Y/gamma/add_data), + // unlike get_update_vec_size above which is fixed to X/Y/gamma/beta. + template + int get_aligned_update_vec_size(int Plane, int vec_size, Args*... args) const { + vec_size = get_min_vec_size(vec_size, args...); while (Plane % vec_size != 0) { vec_size = vec_size >> 1; } diff --git a/src/sycl/RMSNorm.cpp b/src/sycl/RMSNorm.cpp index f414ad635..a0ba7daa8 100644 --- a/src/sycl/RMSNorm.cpp +++ b/src/sycl/RMSNorm.cpp @@ -71,271 +71,31 @@ class RMSNormForward : public NormForward { RMSNormForward() = delete; RMSNormForward( scalar_t* X_data, scalar_t* Y_data, mean_t* var_data, weight_t* gamma_data, accscalar_t eps, int64_t M, int64_t N) - : NormForward(X_data, Y_data, nullptr, var_data, gamma_data, nullptr, eps), M(M), N(N) { - numel = M * N; - }; - - template - void reduce_combine(nd_item_id item_id, const NormConfig& cfg, accscalar_t& sum_value, accscalar_t& sum_tmp) const { - auto group_id = item_id.get_group(0); - auto group_id_foreach = item_id.get_group(1); - auto local_id = item_id.get_local_id(2); - index_t group_offset = (group_id / cfg.input_inner_size) * cfg.input_batch_stride + - (group_id % cfg.input_inner_size) * cfg.input_inner_stride; - - for (index_t j = local_id * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { - index_t plane_offset = group_id_foreach * cfg.WGPlane + j; - if (plane_offset < cfg.Plane) { - vec_t value = *(reinterpret_cast(NF::X_data + group_offset + plane_offset)); - for (int v = 0; v < vec_size; ++v) { - sum_value += Numerics::pow(value[v], 2); - } - } - } - } - - template - void reduce_project(nd_item_id item_id, accscalar_t sum_value, accscalar_t sum_tmp, const NormConfig& cfg) const { - auto group_id = item_id.get_group(0); - accscalar_t scale = static_cast(cfg.Plane); - NF::var_data[group_id] = static_cast( - Numerics::rsqrt(sum_value < 0 ? 0 : sum_value / scale + static_cast(NF::eps))); - } - - template - void update(nd_item_id item_id, const NormConfig& cfg, accscalar_t sum_value = 0, accscalar_t sum_tmp = 0) const { - auto group_id = item_id.get_group(0); - auto group_id_foreach = item_id.get_group(1); - auto local_id = item_id.get_local_id(2); - - index_t x_group_offset = (group_id / cfg.input_inner_size) * cfg.input_batch_stride + - (group_id % cfg.input_inner_size) * cfg.input_inner_stride; - index_t y_group_offset = (group_id / cfg.output_inner_size) * cfg.output_batch_stride + - (group_id % cfg.output_inner_size) * cfg.output_inner_stride; - if (cfg.workgroup_num_foreach == 1) { - if (local_id == 0) { - reduce_project(item_id, sum_value, sum_tmp, cfg); - } - item_id.barrier(DECLARE_SYCL_GLOBAL_FENCE); - } - - auto var_val = NF::var_data[group_id]; - for (index_t j = local_id * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { - index_t plane_offset = group_id_foreach * cfg.WGPlane + j; - if (plane_offset < cfg.Plane) { - vec_t X_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); - vec_t Y_val; - weight_vec_t gamma_val = *(reinterpret_cast(NF::gamma_data + plane_offset)); - - for (int v = 0; v < vec_size; ++v) { - Y_val[v] = static_cast(gamma_val[v] * var_val * X_val[v]); - } - *(reinterpret_cast(NF::Y_data + y_group_offset + plane_offset)) = Y_val; - } - } - } - - int64_t M; - int64_t N; - int64_t numel; -}; - -template -class AddRMSNormForward : public RMSNormForward { - public: - using accscalar_t = acc_type; - typedef NormForward NF; - AddRMSNormForward() = delete; - AddRMSNormForward( - scalar_t* X_data, - scalar_t* Y_data, - mean_t* var_data, - weight_t* gamma_data, - accscalar_t eps, - scalar_t* add_data, - int64_t M, - int64_t N) - : RMSNormForward(X_data, Y_data, var_data, gamma_data, eps, M, N), add_data(add_data) {}; - template - void reduce_combine(nd_item_id item_id, const NormConfig& cfg, accscalar_t& sum_value, accscalar_t& sum_tmp) const { - auto group_id = item_id.get_group(0); - auto group_id_foreach = item_id.get_group(1); - auto local_id = item_id.get_local_id(2); - index_t group_offset = (group_id / cfg.input_inner_size) * cfg.input_batch_stride + - (group_id % cfg.input_inner_size) * cfg.input_inner_stride; - - for (index_t j = local_id * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { - index_t plane_offset = group_id_foreach * cfg.WGPlane + j; - if (plane_offset < cfg.Plane) { - vec_t X_value = *(reinterpret_cast(NF::X_data + group_offset + plane_offset)); - vec_t add_value = *(reinterpret_cast(add_data + group_offset + plane_offset)); - for (int v = 0; v < vec_size; ++v) { - X_value[v] += add_value[v]; - sum_value += Numerics::pow(X_value[v], 2); - } - *(reinterpret_cast(add_data + group_offset + plane_offset)) = X_value; - *(reinterpret_cast(NF::X_data + group_offset + plane_offset)) = X_value; - } - } - } - scalar_t* add_data; -}; - -template < - typename scalar_t, - typename weight_t, - int vec_size, - template - class Norm, - bool one_moment = false, - typename mean_t = float, - typename index_t = uint32_t> -struct FusedNormKernelFunctor { - using accscalar_t = acc_type; - using vec_t = aligned_vector_loop; - using weight_vec_t = aligned_vector_loop; - [[sycl::reqd_sub_group_size(NUM_REDUCE_STAGES)]] void operator()(sycl::nd_item<3> item_id) const { - accscalar_t sum1 = 0; - accscalar_t sum2 = 0; - norm.template reduce_combine(item_id, cfg, sum1, sum2); - - if constexpr (one_moment) { - sum1 = sycl::reduce_over_group(item_id.get_group(), sum1, sycl::plus()); - } else { - norm_group_reduce( - item_id, cfg.sub_group_num, sum1, sum2, local_sum1, local_sum2, [](accscalar_t a, accscalar_t b) { - return a + b; - }); - } - norm.template update(item_id, cfg, sum1, sum2); - } - FusedNormKernelFunctor( - sycl_local_acc_t local_sum1_, - sycl_local_acc_t local_sum2_, - Norm norm_, - NormConfig cfg_) - : local_sum1(local_sum1_), local_sum2(local_sum2_), norm(norm_), cfg(cfg_) {} - - private: - sycl_local_acc_t local_sum1; - sycl_local_acc_t local_sum2; - Norm norm; - const NormConfig cfg; -}; - -template < - typename scalar_t, - typename weight_t, - int vec_size, - template - class Norm, - bool one_moment = false, - typename mean_t = float, - typename index_t = uint32_t> -void fused_norm_kernel(Norm& norm, const NormConfig& cfg) { - using accscalar_t = acc_type; - sycl::range<3> local_range{ - 1, static_cast(cfg.workgroup_num_foreach), static_cast(cfg.workgroup_size)}; - sycl::range<3> global_range{ - static_cast(cfg.workgroup_num), - static_cast(cfg.workgroup_num_foreach), - static_cast(cfg.workgroup_size)}; - - auto stream = at::xpu::getCurrentXPUStream(); - auto queue = stream.queue(); - - auto cgf = [&](sycl::handler& cgh) { - sycl_local_acc_t local_sum1(cfg.sub_group_num, cgh); - sycl_local_acc_t local_sum2(cfg.sub_group_num, cgh); - FusedNormKernelFunctor kfn(local_sum1, local_sum2, norm, cfg); - cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(global_range), sycl::range<3>(local_range)), kfn); - }; - queue.submit(cgf); -} - -template < - typename scalar_t, - typename weight_t, - template - class Norm, - bool one_moment = false, - typename mean_t = float> -void launch_vectorized_fused_norm_kernel(Norm& norm, const NormConfig& config) { - int vec_size = config.get_stride_aligned_vec_size(norm.get_update_vec_size(config.WGPlane, config.max_vec_size)); -#define VECTORIZED_FUSED_NORM_KERNEL(vec_size) \ - { \ - fused_norm_kernel(norm, config); \ - break; \ - } - switch (vec_size) { - case 8: { - VECTORIZED_FUSED_NORM_KERNEL(8); - } - case 4: { - VECTORIZED_FUSED_NORM_KERNEL(4); - } - case 2: { - VECTORIZED_FUSED_NORM_KERNEL(2); - } - default: { - VECTORIZED_FUSED_NORM_KERNEL(1); - } - } -} - -template -class NormNoRstdForward { - public: - using accscalar_t = acc_type; - NormNoRstdForward() = delete; - NormNoRstdForward(scalar_t* X_data, scalar_t* Y_data, accscalar_t eps) : X_data(X_data), Y_data(Y_data), eps(eps) {} + : NormForward(X_data, Y_data, nullptr, var_data, gamma_data, nullptr, eps) {}; int get_update_vec_size(int Plane, int vec_size) const { - return get_aligned_update_vec_size(Plane, vec_size, X_data, Y_data); - } - - protected: - template - int get_aligned_update_vec_size(int Plane, int vec_size, Args*... args) const { - vec_size = get_min_vec_size(vec_size, args...); - while (Plane % vec_size != 0) { - vec_size = vec_size >> 1; - } - return vec_size; - } - - public: - scalar_t* X_data; - scalar_t* Y_data; - accscalar_t eps; -}; - -template -class RMSNormNoRstdForward : public NormNoRstdForward { - public: - using accscalar_t = acc_type; - typedef NormNoRstdForward NF; - RMSNormNoRstdForward() = delete; - RMSNormNoRstdForward(scalar_t* X_data, scalar_t* Y_data, weight_t* gamma_data, accscalar_t eps) - : NormNoRstdForward(X_data, Y_data, eps), gamma_data(gamma_data) {} - - int get_update_vec_size(int Plane, int vec_size) const { - return NF::get_aligned_update_vec_size(Plane, vec_size, NF::X_data, NF::Y_data, gamma_data); + return NF::get_aligned_update_vec_size(Plane, vec_size, NF::X_data, NF::Y_data, NF::gamma_data); } + // Default path (cache_inputs=true): single workgroup per row, inputs cached in + // registers, rstd stays local to the workgroup (no global write). Fallback path + // (cache_inputs=false, register cache doesn't fit): reload inputs from global + // memory each pass, same as the original implementation that published rstd + // through var_data + a workgroup barrier. template void reduce_combine( - sycl::nd_item<1> item_id, + sycl::nd_item<3> item_id, const NormConfig& cfg, index_t x_group_offset, accscalar_t& sum_value, vec_t (®)[cache_inputs ? ITERS : 1]) const { - const index_t lid = item_id.get_local_id(0); + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; if constexpr (cache_inputs) { #pragma unroll for (int it = 0; it < ITERS; ++it) { - const index_t plane_offset = (static_cast(it) * cfg.workgroup_size + lid) * vec_size; + const index_t plane_offset = foreach_offset + (static_cast(it) * cfg.workgroup_size + lid) * vec_size; if (plane_offset < cfg.Plane) { vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); reg[it] = x_val; @@ -347,42 +107,57 @@ class RMSNormNoRstdForward : public NormNoRstdForward { } } } else { - for (index_t plane_offset = lid * vec_size; plane_offset < cfg.Plane; - plane_offset += cfg.workgroup_size * vec_size) { - vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); #pragma unroll - for (int v = 0; v < vec_size; ++v) { - const accscalar_t x = static_cast(x_val[v]); - sum_value += x * x; + for (int v = 0; v < vec_size; ++v) { + const accscalar_t x = static_cast(x_val[v]); + sum_value += x * x; + } } } } } - accscalar_t reduce_project(sycl::nd_item<1> item_id, const NormConfig& cfg, accscalar_t sum_value) const { + template + accscalar_t reduce_project(sycl::nd_item<3> item_id, const NormConfig& cfg, accscalar_t sum_value) const { sum_value = sycl::reduce_over_group(item_id.get_group(), sum_value, sycl::plus()); sum_value = sum_value < static_cast(0) ? static_cast(0) : sum_value; - return Numerics::rsqrt( + accscalar_t rstd = Numerics::rsqrt( sum_value / static_cast(cfg.Plane) + static_cast(NF::eps)); + if constexpr (!cache_inputs) { + // Fallback path: publish rstd through global memory + a workgroup barrier, + // matching the original (pre-norstd) implementation. + const auto group_id = item_id.get_group(0); + if (item_id.get_local_id(1) == 0 && item_id.get_local_id(2) == 0) { + NF::var_data[group_id] = static_cast(rstd); + } + item_id.barrier(DECLARE_SYCL_GLOBAL_FENCE); + rstd = static_cast(NF::var_data[group_id]); + } + return rstd; } template void update( - sycl::nd_item<1> item_id, + sycl::nd_item<3> item_id, const NormConfig& cfg, index_t x_group_offset, index_t y_group_offset, accscalar_t rstd, const vec_t (®)[cache_inputs ? ITERS : 1]) const { - const index_t lid = item_id.get_local_id(0); + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; if constexpr (cache_inputs) { #pragma unroll for (int it = 0; it < ITERS; ++it) { - const index_t plane_offset = (static_cast(it) * cfg.workgroup_size + lid) * vec_size; + const index_t plane_offset = foreach_offset + (static_cast(it) * cfg.workgroup_size + lid) * vec_size; if (plane_offset < cfg.Plane) { vec_t x_val = reg[it]; - weight_vec_t gamma_val = *(reinterpret_cast(gamma_data + plane_offset)); + weight_vec_t gamma_val = *(reinterpret_cast(NF::gamma_data + plane_offset)); vec_t y_val; #pragma unroll for (int v = 0; v < vec_size; ++v) { @@ -393,47 +168,168 @@ class RMSNormNoRstdForward : public NormNoRstdForward { } } } else { - for (index_t plane_offset = lid * vec_size; plane_offset < cfg.Plane; - plane_offset += cfg.workgroup_size * vec_size) { - vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); - weight_vec_t gamma_val = *(reinterpret_cast(gamma_data + plane_offset)); - vec_t y_val; + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); + weight_vec_t gamma_val = *(reinterpret_cast(NF::gamma_data + plane_offset)); + vec_t y_val; #pragma unroll - for (int v = 0; v < vec_size; ++v) { - y_val[v] = - static_cast(static_cast(x_val[v]) * rstd * static_cast(gamma_val[v])); + for (int v = 0; v < vec_size; ++v) { + y_val[v] = static_cast( + static_cast(x_val[v]) * rstd * static_cast(gamma_val[v])); + } + *(reinterpret_cast(NF::Y_data + y_group_offset + plane_offset)) = y_val; } - *(reinterpret_cast(NF::Y_data + y_group_offset + plane_offset)) = y_val; } } } +}; +template +class AddRMSNormForward : public RMSNormForward { public: - weight_t* gamma_data; + using accscalar_t = acc_type; + typedef NormForward NF; + AddRMSNormForward() = delete; + AddRMSNormForward( + scalar_t* X_data, + scalar_t* Y_data, + mean_t* var_data, + weight_t* gamma_data, + accscalar_t eps, + scalar_t* add_data, + int64_t M, + int64_t N) + : RMSNormForward(X_data, Y_data, var_data, gamma_data, eps, M, N), + add_data(add_data) {}; + + int get_update_vec_size(int Plane, int vec_size) const { + return NF::get_aligned_update_vec_size(Plane, vec_size, NF::X_data, NF::Y_data, NF::gamma_data, add_data); + } + + // Folds the residual add into the same cached/reload pass used for the + // squared-sum reduction; update() is inherited unchanged from RMSNormForward. + template + void reduce_combine( + sycl::nd_item<3> item_id, + const NormConfig& cfg, + index_t x_group_offset, + accscalar_t& sum_value, + vec_t (®)[cache_inputs ? ITERS : 1]) const { + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; + + if constexpr (cache_inputs) { +#pragma unroll + for (int it = 0; it < ITERS; ++it) { + const index_t plane_offset = foreach_offset + (static_cast(it) * cfg.workgroup_size + lid) * vec_size; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); + vec_t add_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); +#pragma unroll + for (int v = 0; v < vec_size; ++v) { + x_val[v] = static_cast(static_cast(x_val[v]) + static_cast(add_val[v])); + } + *(reinterpret_cast(add_data + x_group_offset + plane_offset)) = x_val; + reg[it] = x_val; +#pragma unroll + for (int v = 0; v < vec_size; ++v) { + const accscalar_t x = static_cast(x_val[v]); + sum_value += x * x; + } + } + } + } else { + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); + vec_t add_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); +#pragma unroll + for (int v = 0; v < vec_size; ++v) { + x_val[v] = static_cast(static_cast(x_val[v]) + static_cast(add_val[v])); + } + *(reinterpret_cast(add_data + x_group_offset + plane_offset)) = x_val; + // No register cache here, so update()'s reload path must see x+residual. + *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)) = x_val; +#pragma unroll + for (int v = 0; v < vec_size; ++v) { + const accscalar_t x = static_cast(x_val[v]); + sum_value += x * x; + } + } + } + } + } + scalar_t* add_data; +}; + +// Single workgroup per row (group(1) kept symbolic, not hardcoded, for a +// possible future workgroup_num_foreach > 1 split). Drives reduce_combine -> +// reduce_project -> update; cache_inputs picks the norstd +// (registers, no global write) path vs. the norm fallback (reload, var_data + +// barrier) path. +template < + typename scalar_t, + typename weight_t, + int vec_size, + int ITERS, + typename Norm, + bool cache_inputs, + typename mean_t = float, + typename index_t = uint32_t> +struct NormKernelFunctor { + using accscalar_t = acc_type; + using vec_t = aligned_vector_loop; + using weight_vec_t = aligned_vector_loop; + + [[sycl::reqd_sub_group_size(NUM_REDUCE_STAGES)]] void operator()(sycl::nd_item<3> item_id) const { + const index_t group_id = item_id.get_group(0); + const index_t x_group_offset = (group_id / cfg.input_inner_size) * cfg.input_batch_stride + + (group_id % cfg.input_inner_size) * cfg.input_inner_stride; + const index_t y_group_offset = (group_id / cfg.output_inner_size) * cfg.output_batch_stride + + (group_id % cfg.output_inner_size) * cfg.output_inner_stride; + + accscalar_t sum_value = 0; + vec_t reg[cache_inputs ? ITERS : 1]; + norm.template reduce_combine( + item_id, cfg, x_group_offset, sum_value, reg); + const accscalar_t rstd = norm.template reduce_project(item_id, cfg, sum_value); + norm.template update( + item_id, cfg, x_group_offset, y_group_offset, rstd, reg); + } + + NormKernelFunctor(Norm norm_, NormConfig cfg_) : norm(norm_), cfg(cfg_) {} + + private: + Norm norm; + const NormConfig cfg; }; template -class GemmaRMSNormNoRstdForward : public RMSNormNoRstdForward { +class GemmaRMSNormForward : public RMSNormForward { public: using accscalar_t = acc_type; - typedef RMSNormNoRstdForward RNF; - GemmaRMSNormNoRstdForward() = delete; + typedef RMSNormForward RNF; + GemmaRMSNormForward() = delete; using RNF::RNF; template void update( - sycl::nd_item<1> item_id, + sycl::nd_item<3> item_id, const NormConfig& cfg, index_t x_group_offset, index_t y_group_offset, accscalar_t rstd, const vec_t (®)[cache_inputs ? ITERS : 1]) const { - const index_t lid = item_id.get_local_id(0); + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; if constexpr (cache_inputs) { #pragma unroll for (int it = 0; it < ITERS; ++it) { - const index_t plane_offset = (static_cast(it) * cfg.workgroup_size + lid) * vec_size; + const index_t plane_offset = foreach_offset + (static_cast(it) * cfg.workgroup_size + lid) * vec_size; if (plane_offset < cfg.Plane) { vec_t x_val = reg[it]; weight_vec_t gamma_val = *(reinterpret_cast(RNF::gamma_data + plane_offset)); @@ -448,33 +344,34 @@ class GemmaRMSNormNoRstdForward : public RMSNormNoRstdForward(RNF::X_data + x_group_offset + plane_offset)); - weight_vec_t gamma_val = *(reinterpret_cast(RNF::gamma_data + plane_offset)); - vec_t y_val; + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(RNF::X_data + x_group_offset + plane_offset)); + weight_vec_t gamma_val = *(reinterpret_cast(RNF::gamma_data + plane_offset)); + vec_t y_val; #pragma unroll - for (int v = 0; v < vec_size; ++v) { - y_val[v] = static_cast( - static_cast(x_val[v]) * rstd * - (static_cast(1.0) + static_cast(gamma_val[v]))); + for (int v = 0; v < vec_size; ++v) { + y_val[v] = static_cast( + static_cast(x_val[v]) * rstd * + (static_cast(1.0) + static_cast(gamma_val[v]))); + } + *(reinterpret_cast(RNF::Y_data + y_group_offset + plane_offset)) = y_val; } - *(reinterpret_cast(RNF::Y_data + y_group_offset + plane_offset)) = y_val; } } } }; template -class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward { +class GemmaAddRMSNormForward : public GemmaRMSNormForward { public: using accscalar_t = acc_type; - typedef GemmaRMSNormNoRstdForward Base; - typedef NormNoRstdForward NF; - GemmaAddRMSNormNoRstdForward() = delete; - GemmaAddRMSNormNoRstdForward( - scalar_t* X_data, scalar_t* Y_data, weight_t* gamma_data, accscalar_t eps, scalar_t* add_data) - : GemmaRMSNormNoRstdForward(X_data, Y_data, gamma_data, eps), add_data(add_data) {} + typedef GemmaRMSNormForward Base; + typedef NormForward NF; + GemmaAddRMSNormForward() = delete; + GemmaAddRMSNormForward(scalar_t* X_data, scalar_t* Y_data, weight_t* gamma_data, accscalar_t eps, scalar_t* add_data) + : GemmaRMSNormForward(X_data, Y_data, nullptr, gamma_data, eps, 0, 0), add_data(add_data) {} int get_update_vec_size(int Plane, int vec_size) const { return NF::get_aligned_update_vec_size(Plane, vec_size, NF::X_data, NF::Y_data, Base::gamma_data, add_data); @@ -482,17 +379,18 @@ class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward void reduce_combine( - sycl::nd_item<1> item_id, + sycl::nd_item<3> item_id, const NormConfig& cfg, index_t x_group_offset, accscalar_t& sum_value, vec_t (®)[cache_inputs ? ITERS : 1]) const { - const index_t lid = item_id.get_local_id(0); + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; if constexpr (cache_inputs) { #pragma unroll for (int it = 0; it < ITERS; ++it) { - const index_t plane_offset = (static_cast(it) * cfg.workgroup_size + lid) * vec_size; + const index_t plane_offset = foreach_offset + (static_cast(it) * cfg.workgroup_size + lid) * vec_size; if (plane_offset < cfg.Plane) { vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); vec_t add_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); @@ -510,19 +408,21 @@ class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward(NF::X_data + x_group_offset + plane_offset)); - vec_t add_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(NF::X_data + x_group_offset + plane_offset)); + vec_t add_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); #pragma unroll - for (int v = 0; v < vec_size; ++v) { - x_val[v] = static_cast(static_cast(x_val[v]) + static_cast(add_val[v])); - } - *(reinterpret_cast(add_data + x_group_offset + plane_offset)) = x_val; + for (int v = 0; v < vec_size; ++v) { + x_val[v] = static_cast(static_cast(x_val[v]) + static_cast(add_val[v])); + } + *(reinterpret_cast(add_data + x_group_offset + plane_offset)) = x_val; #pragma unroll - for (int v = 0; v < vec_size; ++v) { - const accscalar_t x = static_cast(x_val[v]); - sum_value += x * x; + for (int v = 0; v < vec_size; ++v) { + const accscalar_t x = static_cast(x_val[v]); + sum_value += x * x; + } } } } @@ -530,7 +430,7 @@ class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward void update( - sycl::nd_item<1> item_id, + sycl::nd_item<3> item_id, const NormConfig& cfg, index_t x_group_offset, index_t y_group_offset, @@ -540,19 +440,22 @@ class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward( item_id, cfg, x_group_offset, y_group_offset, rstd, reg); } else { - const index_t lid = item_id.get_local_id(0); - for (index_t plane_offset = lid * vec_size; plane_offset < cfg.Plane; - plane_offset += cfg.workgroup_size * vec_size) { - vec_t x_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); - weight_vec_t gamma_val = *(reinterpret_cast(Base::gamma_data + plane_offset)); - vec_t y_val; + const index_t lid = item_id.get_local_id(2); + const index_t foreach_offset = static_cast(item_id.get_group(1)) * cfg.WGPlane; + for (index_t j = lid * vec_size; j < cfg.WGPlane; j += cfg.workgroup_size * vec_size) { + const index_t plane_offset = foreach_offset + j; + if (plane_offset < cfg.Plane) { + vec_t x_val = *(reinterpret_cast(add_data + x_group_offset + plane_offset)); + weight_vec_t gamma_val = *(reinterpret_cast(Base::gamma_data + plane_offset)); + vec_t y_val; #pragma unroll - for (int v = 0; v < vec_size; ++v) { - y_val[v] = static_cast( - static_cast(x_val[v]) * rstd * - (static_cast(1.0) + static_cast(gamma_val[v]))); + for (int v = 0; v < vec_size; ++v) { + y_val[v] = static_cast( + static_cast(x_val[v]) * rstd * + (static_cast(1.0) + static_cast(gamma_val[v]))); + } + *(reinterpret_cast(NF::Y_data + y_group_offset + plane_offset)) = y_val; } - *(reinterpret_cast(NF::Y_data + y_group_offset + plane_offset)) = y_val; } } } @@ -561,52 +464,20 @@ class GemmaAddRMSNormNoRstdForward : public GemmaRMSNormNoRstdForward -struct RMSNormNoRstdKernelFunctor { - using accscalar_t = acc_type; - using vec_t = aligned_vector_loop; - using weight_vec_t = aligned_vector_loop; - - [[sycl::reqd_sub_group_size(NUM_REDUCE_STAGES)]] void operator()(sycl::nd_item<1> item_id) const { - const index_t row = item_id.get_group(0); - const index_t x_group_offset = - (row / cfg.input_inner_size) * cfg.input_batch_stride + (row % cfg.input_inner_size) * cfg.input_inner_stride; - const index_t y_group_offset = (row / cfg.output_inner_size) * cfg.output_batch_stride + - (row % cfg.output_inner_size) * cfg.output_inner_stride; - - accscalar_t sum_value = 0; - vec_t reg[cache_inputs ? ITERS : 1]; - norm.template reduce_combine( - item_id, cfg, x_group_offset, sum_value, reg); - const accscalar_t rstd = norm.reduce_project(item_id, cfg, sum_value); - norm.template update( - item_id, cfg, x_group_offset, y_group_offset, rstd, reg); - } - - RMSNormNoRstdKernelFunctor(Norm norm_, NormConfig cfg_) : norm(norm_), cfg(cfg_) {} - - private: - Norm norm; - const NormConfig cfg; -}; - template void rmsnorm_no_rstd_kernel(Norm& norm, const NormConfig& config) { auto stream = at::xpu::getCurrentXPUStream(); auto queue = stream.queue(); - using KernelFunctor = RMSNormNoRstdKernelFunctor; + using KernelFunctor = NormKernelFunctor; KernelFunctor kfn(norm, config); - sycl::range<1> local_range{static_cast(config.workgroup_size)}; - sycl::range<1> global_range{static_cast(config.workgroup_num * config.workgroup_size)}; + sycl::range<3> local_range{ + static_cast(1), static_cast(1), static_cast(config.workgroup_size)}; + sycl::range<3> global_range{ + static_cast(config.workgroup_num), + static_cast(config.workgroup_num_foreach), + static_cast(config.workgroup_size)}; sycl_kernel_submit(global_range, local_range, queue, kfn); } @@ -686,6 +557,7 @@ void RMSNormKernelImplInternal( mean_t* var_data = rstd.data_ptr(); weight_t* gemma_data = gemma.defined() ? gemma.data_ptr() : nullptr; + RMSNormForward rms_norm_forward(X_data, Y_data, var_data, gemma_data, eps, M, N); auto config = NormConfig( M, N, @@ -693,15 +565,13 @@ void RMSNormKernelImplInternal( sizeof(scalar_t), input_batch_stride, output_batch_stride, + [&](int plane, int max_vec_size) { return rms_norm_forward.get_update_vec_size(plane, max_vec_size); }, input_inner_size, input_inner_stride, output_inner_size, output_inner_stride); - RMSNormForward rms_norm_forward(X_data, Y_data, var_data, gemma_data, eps, M, N); - config.workgroup_num_foreach = 1; - config.WGPlane = config.Plane; - launch_vectorized_fused_norm_kernel(rms_norm_forward, config); + launch_vectorized_rmsnorm_no_rstd_kernel(rms_norm_forward, config); } template @@ -718,13 +588,13 @@ void FusedAddRMSNormKernelImplInternal( weight_t* gemma_data = gemma.defined() ? gemma.data_ptr() : nullptr; scalar_t* residual_data = residual.data_ptr(); - auto config = NormConfig(M, N, 1, sizeof(scalar_t), N, N); AddRMSNormForward add_rms_norm_forward( X_data, X_data, var_data, gemma_data, eps, residual_data, M, N); - config.workgroup_num_foreach = 1; - config.WGPlane = config.Plane; + auto config = NormConfig(M, N, 1, sizeof(scalar_t), N, N, [&](int plane, int max_vec_size) { + return add_rms_norm_forward.get_update_vec_size(plane, max_vec_size); + }); - launch_vectorized_fused_norm_kernel(add_rms_norm_forward, config); + launch_vectorized_rmsnorm_no_rstd_kernel(add_rms_norm_forward, config); } template @@ -744,7 +614,7 @@ void GemmaRMSNormKernelImplInternal( scalar_t* X_data = X.data_ptr(); scalar_t* Y_data = Y.data_ptr(); weight_t* gemma_data = gemma.data_ptr(); - GemmaRMSNormNoRstdForward gemma_rms_norm_no_rstd_forward(X_data, Y_data, gemma_data, eps); + GemmaRMSNormForward gemma_rms_norm_forward(X_data, Y_data, nullptr, gemma_data, eps, M, N); auto config = NormConfig( M, @@ -753,14 +623,12 @@ void GemmaRMSNormKernelImplInternal( sizeof(scalar_t), input_batch_stride, output_batch_stride, - [&](int plane, int max_vec_size) { - return gemma_rms_norm_no_rstd_forward.get_update_vec_size(plane, max_vec_size); - }, + [&](int plane, int max_vec_size) { return gemma_rms_norm_forward.get_update_vec_size(plane, max_vec_size); }, input_inner_size, input_inner_stride, output_inner_size, output_inner_stride); - launch_vectorized_rmsnorm_no_rstd_kernel(gemma_rms_norm_no_rstd_forward, config); + launch_vectorized_rmsnorm_no_rstd_kernel(gemma_rms_norm_forward, config); } template @@ -769,8 +637,7 @@ void GemmaFusedAddRMSNormKernelImplInternal( scalar_t* X_data = X.data_ptr(); weight_t* gemma_data = gemma.data_ptr(); scalar_t* residual_data = residual.data_ptr(); - GemmaAddRMSNormNoRstdForward gemma_add_rms_norm_no_rstd_forward( - X_data, X_data, gemma_data, eps, residual_data); + GemmaAddRMSNormForward gemma_add_rms_norm_forward(X_data, X_data, gemma_data, eps, residual_data); auto config = NormConfig( M, @@ -779,15 +646,13 @@ void GemmaFusedAddRMSNormKernelImplInternal( sizeof(scalar_t), N, N, - [&](int plane, int max_vec_size) { - return gemma_add_rms_norm_no_rstd_forward.get_update_vec_size(plane, max_vec_size); - }, + [&](int plane, int max_vec_size) { return gemma_add_rms_norm_forward.get_update_vec_size(plane, max_vec_size); }, 1, 0, 1, 0); - launch_vectorized_rmsnorm_no_rstd_kernel(gemma_add_rms_norm_no_rstd_forward, config); + launch_vectorized_rmsnorm_no_rstd_kernel(gemma_add_rms_norm_forward, config); } SGL_KERNEL_EXPORT void rmsnorm(torch::Tensor& output, torch::Tensor& input, torch::Tensor& weight, double eps) {