diff --git a/benchmarks/kernels/rms_norm.py b/benchmarks/kernels/rms_norm.py new file mode 100644 index 00000000..9b366bc7 --- /dev/null +++ b/benchmarks/kernels/rms_norm.py @@ -0,0 +1,193 @@ +"""Matched RMSNorm / RMSNorm-then-residual accuracy and latency comparison. + +The optional vLLM library must expose the source-built ``vllm_rms_review`` +namespace. Its RMSNorm multiplication rounds differently from Ouro, so its +error against the staged reference is reported separately from its latency. +""" + +import argparse +import hashlib +import json +import os +import subprocess +from pathlib import Path + +import torch + +from benchmarks.kernels.silu_and_mul import errors, measure, profile +from vllm_rlt.kernels._extension import load_extension +from vllm_rlt.kernels.normalization import _torch_order_validated, rms_norm, rms_norm_reference + + +def oracle(x, weight, epsilon, residual): + value = x.double() + normalized = (value * torch.rsqrt(value.square().mean(-1, keepdim=True) + epsilon)).to(x.dtype) + result = (normalized.double() * weight.double()).to(x.dtype) + return result if residual is None else (residual.double() + result.double()).to(x.dtype) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--vllm-library", type=Path) + parser.add_argument("--vllm-namespace", default="vllm_rms_review") + parser.add_argument( + "--vllm-source-commit", help="Source revision used to build the vLLM library" + ) + parser.add_argument("--tokens", default="1,2,4,8,16,32,64,128,256,1024") + parser.add_argument("--widths", default="128,768,2048,2049,4096") + parser.add_argument("--dtypes", default="bfloat16,float16,float32") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--iterations", type=int, default=30) + parser.add_argument("--profile-dir", type=Path) + args = parser.parse_args() + load_extension() + vllm_op = None + if args.vllm_library is not None: + torch.ops.load_library(str(args.vllm_library.resolve())) + vllm_op = getattr(getattr(torch.ops, args.vllm_namespace), "rms_norm") + torch.manual_seed(42) + torch.set_num_threads(1) + patch = subprocess.check_output(["git", "diff", "HEAD", "--binary"]) + # New kernels often have not been added to the Git index yet. Include their + # complete contents rather than recording a misleading, incomplete hash. + untracked = subprocess.check_output( + ["git", "ls-files", "--others", "--exclude-standard", "-z"] + ).split(b"\0") + for path in filter(None, untracked): + diff = subprocess.run( + ["git", "diff", "--no-index", "--binary", "--", "/dev/null", os.fsdecode(path)], + capture_output=True, + check=False, + ) + if diff.returncode not in (0, 1): + raise RuntimeError(diff.stderr.decode(errors="replace")) + patch += diff.stdout + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.with_suffix(".patch").write_bytes(patch) + report = { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "capability": torch.cuda.get_device_capability(), + "torch_order_native_validated": _torch_order_validated(torch.cuda.current_device()), + "git_head": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + "diff_sha256": hashlib.sha256(patch).hexdigest(), + "vllm_library": str(args.vllm_library.resolve()) if args.vllm_library else None, + "vllm_library_sha256": ( + hashlib.sha256(args.vllm_library.read_bytes()).hexdigest() + if args.vllm_library + else None + ), + "vllm_source_commit": args.vllm_source_commit, + "seed": 42, + "epsilon": 1e-6, + "repeats": args.repeats, + "iterations": args.iterations, + "graph_unroll": 32, + "device_status_before": subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,memory.used,utilization.gpu,clocks.sm", + "--format=csv", + ], + text=True, + ), + "rows": [], + } + args.output.write_text(json.dumps(report, indent=2)) + with torch.inference_mode(): + for dtype_name in args.dtypes.split(","): + dtype = getattr(torch, dtype_name) + for width in map(int, args.widths.split(",")): + for tokens in map(int, args.tokens.split(",")): + x = torch.randn(tokens, width, dtype=dtype, device="cuda") + residual = torch.randn_like(x) + weight = (torch.randn(width, device="cuda") * 0.1 + 1).to(dtype) + + def vllm(): + output = torch.empty(x.shape, device=x.device, dtype=x.dtype) + vllm_op(output, x, weight, 1e-6) + return output + + for add in (False, True): + res = residual if add else None + funcs = { + "torch": lambda: rms_norm_reference(x, weight, 1e-6, res), + "rlt_fused": lambda: rms_norm(x, weight, 1e-6, res), + } + if vllm_op is not None: + funcs["vllm"] = (lambda: vllm() + residual) if add else vllm + if ( + report["torch_order_native_validated"] + and width == 2048 + and dtype in (torch.float16, torch.bfloat16) + ): + funcs["rlt_torch_order"] = lambda: rms_norm( + x, weight, 1e-6, res, match_torch=True + ) + if add: + funcs["rlt_separate"] = lambda: rms_norm(x, weight, 1e-6) + residual + if "rlt_torch_order" in funcs: + funcs["rlt_torch_order_separate"] = ( + lambda: rms_norm(x, weight, 1e-6, match_torch=True) + residual + ) + ref = funcs["torch"]() + exact = oracle(x, weight, 1e-6, res) + accuracy = {} + for name, fn in funcs.items(): + result = fn() + accuracy[name] = { + "vs_ouro": errors(result, ref), + "vs_fp64_staged": errors(result, exact), + } + row = { + "dtype": dtype_name, + "tokens": tokens, + "width": width, + "residual": add, + "accuracy": accuracy, + "timing": measure(funcs, args.repeats, args.iterations), + } + report["rows"].append(row) + args.output.write_text(json.dumps(report, indent=2)) + print( + dtype_name, + width, + tokens, + "add" if add else "norm", + {k: round(v["graph_median_us"], 3) for k, v in row["timing"].items()}, + flush=True, + ) + del x, residual, weight, result, ref, exact, funcs + if args.profile_dir: + x = torch.randn(128, 2048, device="cuda", dtype=torch.bfloat16) + weight = torch.ones(2048, device="cuda", dtype=x.dtype) + residual = torch.randn_like(x) + + def vllm_profile(): + output = torch.empty_like(x) + vllm_op(output, x, weight, 1e-6) + return output + residual + + functions = { + "torch_norm_add": lambda: rms_norm_reference(x, weight, 1e-6, residual), + "rlt_norm_add": lambda: rms_norm(x, weight, 1e-6, residual), + } + if vllm_op is not None: + functions["vllm_norm_add"] = vllm_profile + if report["torch_order_native_validated"]: + functions["rlt_torch_order_norm_add"] = lambda: rms_norm( + x, weight, 1e-6, residual, match_torch=True + ) + profile(functions, args.profile_dir) + report["device_status_after"] = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,memory.used,utilization.gpu,clocks.sm", "--format=csv"], + text=True, + ) + report["complete"] = True + args.output.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/csrc/normalization/rms_norm.cu b/csrc/normalization/rms_norm.cu new file mode 100644 index 00000000..612db480 --- /dev/null +++ b/csrc/normalization/rms_norm.cu @@ -0,0 +1,575 @@ +// SPDX-License-Identifier: Apache-2.0 +#include "../ops.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace vllm_rlt { +namespace { + +template +struct alignas(16) Vector { + T val[16 / sizeof(T)]; +}; + +// Return one logical 16-byte vector without reading outside the input view. +// Aligned words are combined in registers when the logical start is shifted. +template +__device__ __forceinline__ Vector load_vector(const T* input, int index, int count) { + constexpr int elements = 16 / sizeof(T); + const int offset_bytes = reinterpret_cast(input + index) & 15; + Vector result; + if (offset_bytes == 0) { + *reinterpret_cast(&result) = *reinterpret_cast(input + index); + return result; + } + const int offset = offset_bytes / sizeof(T); + const int lower = index - offset; + Vector head{}, tail{}; + if (lower >= 0) { + *reinterpret_cast(&head) = *reinterpret_cast(input + lower); + } else { +#pragma unroll + for (int j = 0; j < elements; ++j) { + if (j >= offset) head.val[j] = input[lower + j]; + } + } + if (lower + 2 * elements <= count) { + *reinterpret_cast(&tail) = *reinterpret_cast(input + lower + elements); + } else { +#pragma unroll + for (int j = 0; j < elements; ++j) { + if (j < offset) tail.val[j] = input[lower + elements + j]; + } + } + uint4 lo = *reinterpret_cast(&head); + uint4 hi = *reinterpret_cast(&tail); + if (offset_bytes & 8) { + lo = make_uint4(lo.z, lo.w, hi.x, hi.y); + hi = make_uint4(hi.z, hi.w, 0, 0); + } + if (offset_bytes & 4) { + lo = make_uint4(lo.y, lo.z, lo.w, hi.x); + hi.x = hi.y; + } + if constexpr (sizeof(T) == 2) { + if (offset_bytes & 2) { + lo = make_uint4(__funnelshift_r(lo.x, lo.y, 16), __funnelshift_r(lo.y, lo.z, 16), + __funnelshift_r(lo.z, lo.w, 16), __funnelshift_r(lo.w, hi.x, 16)); + } + } + *reinterpret_cast(&result) = lo; + return result; +} + +__device__ __forceinline__ float2 mul2(float2 a, float2 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fmul2_rn(a, b); +#else + return make_float2(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)); +#endif +} + +__device__ __forceinline__ float2 add2(float2 a, float2 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fadd2_rn(a, b); +#else + return make_float2(__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y)); +#endif +} + +template +struct Arithmetic; + +template <> +struct Arithmetic<__half> { + using Pair = __half2; + __device__ static float scalar(__half x) { return __half2float(x); } + __device__ static __half cast(float x) { return __float2half_rn(x); } + __device__ static float2 pair(Pair x) { return __half22float2(x); } + __device__ static Pair cast2(float2 x) { return __floats2half2_rn(x.x, x.y); } + __device__ static __half mul(__half a, __half b) { return __hmul_rn(a, b); } + __device__ static __half add(__half a, __half b) { return __hadd_rn(a, b); } + __device__ static Pair mul_pair(Pair a, Pair b) { return __hmul2_rn(a, b); } + __device__ static Pair add_pair(Pair a, Pair b) { return __hadd2_rn(a, b); } +}; + +template <> +struct Arithmetic<__nv_bfloat16> { + using Pair = __nv_bfloat162; + __device__ static float scalar(__nv_bfloat16 x) { return __bfloat162float(x); } + __device__ static __nv_bfloat16 cast(float x) { return __float2bfloat16_rn(x); } + __device__ static float2 pair(Pair x) { return __bfloat1622float2(x); } + __device__ static Pair cast2(float2 x) { return __floats2bfloat162_rn(x.x, x.y); } + __device__ static __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) { + return __hmul_rn(a, b); + } + __device__ static __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) { + return __hadd_rn(a, b); + } + __device__ static Pair mul_pair(Pair a, Pair b) { return __hmul2_rn(a, b); } + __device__ static Pair add_pair(Pair a, Pair b) { return __hadd2_rn(a, b); } +}; + +template <> +struct Arithmetic { + using Pair = float2; + __device__ static float scalar(float x) { return x; } + __device__ static float cast(float x) { return x; } + __device__ static float2 pair(Pair x) { return x; } + __device__ static Pair cast2(float2 x) { return x; } + __device__ static float mul(float a, float b) { return __fmul_rn(a, b); } + __device__ static float add(float a, float b) { return __fadd_rn(a, b); } + __device__ static Pair mul_pair(Pair a, Pair b) { return mul2(a, b); } + __device__ static Pair add_pair(Pair a, Pair b) { return add2(a, b); } +}; + +struct RowLayout { + int64_t sizes[3]; + int64_t input_strides[3]; + int64_t residual_strides[3]; + int64_t input_linear_stride; + int64_t residual_linear_stride; + int hidden_size; +}; + +__device__ __forceinline__ int64_t row_offset(int row, const int64_t* sizes, + const int64_t* strides) { + const int64_t index2 = row % sizes[2]; + row /= sizes[2]; + const int64_t index1 = row % sizes[1]; + const int64_t index0 = row / sizes[1]; + return index0 * strides[0] + index1 * strides[1] + index2 * strides[2]; +} + +// Ouro's 2048-wide BF16/FP16 reference uses torch.mean on an FP32 squared +// tensor. Reproduce that reduction tree, including its row-count-dependent +// width, because small differences can change early-exit decisions. +template +__global__ void rms_norm_torch_order_kernel(T* out, const T* input, const T* weight, + const T* residual, int64_t input_stride, + int64_t residual_stride, float epsilon) { + using A = Arithmetic; + using Pair = typename A::Pair; + constexpr int H = 2048; + constexpr int Packs = H / (Threads * 4); + const T* x = input + static_cast(blockIdx.x) * input_stride; + const T* r = nullptr; + if constexpr (AddResidual) r = residual + static_cast(blockIdx.x) * residual_stride; + T* y = out + static_cast(blockIdx.x) * H; + uint2 cached[Packs]; + float sums[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < Packs; ++k) { + int index = 4 * (threadIdx.x + k * Threads); + cached[k] = *reinterpret_cast(x + index); + const Pair* pairs = reinterpret_cast(&cached[k]); + float2 a = A::pair(pairs[0]); + float2 b = A::pair(pairs[1]); + sums[0] = __fadd_rn(sums[0], __fmul_rn(a.x, a.x)); + sums[1] = __fadd_rn(sums[1], __fmul_rn(a.y, a.y)); + sums[2] = __fadd_rn(sums[2], __fmul_rn(b.x, b.x)); + sums[3] = __fadd_rn(sums[3], __fmul_rn(b.y, b.y)); + } + float sum = __fadd_rn(__fadd_rn(__fadd_rn(sums[0], sums[1]), sums[2]), sums[3]); + __shared__ float partial[Threads]; + if constexpr (Threads > 32) { + partial[threadIdx.x] = sum; + for (int offset = Threads / 2; offset >= 32; offset /= 2) { + __syncthreads(); + if (threadIdx.x < offset) { + sum = __fadd_rn(sum, partial[threadIdx.x + offset]); + partial[threadIdx.x] = sum; + } + } + __syncthreads(); + } +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + sum = __fadd_rn(sum, __shfl_down_sync(0xffffffff, sum, offset)); + } + float scale; + if constexpr (Threads == 32) { + scale = threadIdx.x == 0 ? rsqrtf(__fadd_rn(__fmul_rn(sum, 1.f / H), epsilon)) : 0.f; + scale = __shfl_sync(0xffffffff, scale, 0); + } else { + __shared__ float shared_scale; + if (threadIdx.x == 0) shared_scale = rsqrtf(__fadd_rn(__fmul_rn(sum, 1.f / H), epsilon)); + __syncthreads(); + scale = shared_scale; + } +#pragma unroll + for (int k = 0; k < Packs; ++k) { + int index = 4 * (threadIdx.x + k * Threads); + uint2 w = *reinterpret_cast(weight + index); + uint2 rv{}; + if constexpr (AddResidual) rv = *reinterpret_cast(r + index); + uint2 result; + const Pair* xp = reinterpret_cast(&cached[k]); + const Pair* wp = reinterpret_cast(&w); + const Pair* rp = reinterpret_cast(&rv); + Pair* yp = reinterpret_cast(&result); +#pragma unroll + for (int j = 0; j < 2; ++j) { + float2 f = A::pair(xp[j]); + Pair normalized = A::cast2(make_float2(__fmul_rn(f.x, scale), __fmul_rn(f.y, scale))); + Pair value = A::mul_pair(normalized, wp[j]); + if constexpr (AddResidual) value = A::add_pair(rp[j], value); + yp[j] = value; + } + *reinterpret_cast(y + index) = result; + } +} + +template +void launch_torch_order(at::Tensor& out, const at::Tensor& input, const at::Tensor& weight, + const std::optional& residual, float epsilon) { + const int rows = static_cast(input.size(0)); + const T* x = reinterpret_cast(input.const_data_ptr()); + const T* w = reinterpret_cast(weight.const_data_ptr()); + const T* r = AddResidual ? reinterpret_cast(residual->const_data_ptr()) : nullptr; + T* y = reinterpret_cast(out.mutable_data_ptr()); + const int64_t residual_stride = AddResidual ? residual->stride(0) : 0; + auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); +#define RUN_TORCH_ORDER(N) \ + rms_norm_torch_order_kernel<<>>( \ + y, x, w, r, input.stride(0), residual_stride, epsilon) + if (rows < 2) { RUN_TORCH_ORDER(512); } + else if (rows < 4) { RUN_TORCH_ORDER(256); } + else if (rows < 8) { RUN_TORCH_ORDER(128); } + else if (rows < 16) { RUN_TORCH_ORDER(64); } + else { RUN_TORCH_ORDER(32); } +#undef RUN_TORCH_ORDER + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +__device__ __forceinline__ float2 accumulate(Vector x, float2 accum) { + using A = Arithmetic; + using Pair = typename A::Pair; + constexpr int pairs = 16 / sizeof(Pair); +#pragma unroll + for (int j = 0; j < pairs; ++j) { + Pair p = reinterpret_cast(&x)[j]; + float2 f = A::pair(p); + accum = add2(accum, mul2(f, f)); + } + return accum; +} + +template +__device__ __forceinline__ T finish_scalar(T x, T w, T residual, float scale) { + using A = Arithmetic; + T norm = A::cast(__fmul_rn(A::scalar(x), scale)); + T result = A::mul(norm, w); + if constexpr (AddResidual) result = A::add(residual, result); + return result; +} + +template +__device__ __forceinline__ Vector finish_vector(Vector x, Vector w, + Vector residual, float scale) { + using A = Arithmetic; + using Pair = typename A::Pair; + constexpr int pairs = 16 / sizeof(Pair); + Vector result; +#pragma unroll + for (int j = 0; j < pairs; ++j) { + Pair x_pair = reinterpret_cast(&x)[j]; + Pair w_pair = reinterpret_cast(&w)[j]; + Pair normalized = A::cast2(mul2(A::pair(x_pair), make_float2(scale, scale))); + Pair value = A::mul_pair(normalized, w_pair); + if constexpr (AddResidual) { + Pair r_pair = reinterpret_cast(&residual)[j]; + value = A::add_pair(r_pair, value); + } + reinterpret_cast(&result)[j] = value; + } + return result; +} + +template +__device__ __forceinline__ void write_vector(T* out, const T* weight, const T* residual, + Vector x, int index, int H, float scale) { + Vector w, r{}; + if constexpr (Aligned) { + *reinterpret_cast(&w) = *reinterpret_cast(weight + index); + if constexpr (AddResidual) { + *reinterpret_cast(&r) = *reinterpret_cast(residual + index); + } + } else { + w = load_vector(weight, index, H); + if constexpr (AddResidual) r = load_vector(residual, index, H); + } + Vector result = finish_vector(x, w, r, scale); + *reinterpret_cast(out + index) = *reinterpret_cast(&result); +} + +template +__global__ void rms_norm_kernel(T* out, const T* input, const T* weight, + const T* residual, RowLayout layout, float epsilon) { + constexpr int elements = 16 / sizeof(T); + const int H = StaticH ? StaticH : layout.hidden_size; + const int row = blockIdx.x; + const T* x = input + (layout.input_linear_stride >= 0 + ? static_cast(row) * layout.input_linear_stride + : row_offset(row, layout.sizes, layout.input_strides)); + const T* r = nullptr; + if constexpr (AddResidual) { + r = residual + (layout.residual_linear_stride >= 0 + ? static_cast(row) * layout.residual_linear_stride + : row_offset(row, layout.sizes, layout.residual_strides)); + } + T* y = out + static_cast(row) * H; + int prefix = 0; + if constexpr (!Aligned) { + prefix = min(H, static_cast(((16 - (reinterpret_cast(y) & 15)) & 15) / + sizeof(T))); + } + const int vectors = (H - prefix) / elements; + const int tail_start = prefix + vectors * elements; + Vector cached[CachePacks > 0 ? CachePacks : 1]; + float2 accum = make_float2(0.0f, 0.0f); + T head_value{}, tail_value{}; + float edge_sum = 0.0f; + if (threadIdx.x < prefix) { + head_value = x[threadIdx.x]; + float f = Arithmetic::scalar(head_value); + edge_sum = __fmul_rn(f, f); + } + if (tail_start + threadIdx.x < H) { + tail_value = x[tail_start + threadIdx.x]; + float f = Arithmetic::scalar(tail_value); + edge_sum = __fadd_rn(edge_sum, __fmul_rn(f, f)); + } + if constexpr (CachePacks > 0) { +#pragma unroll + for (int k = 0; k < CachePacks; ++k) { + int index = threadIdx.x + k * Threads; + if (index < vectors) { + int j = prefix + index * elements; + if constexpr (Aligned) { + *reinterpret_cast(&cached[k]) = *reinterpret_cast(x + j); + } else { + cached[k] = load_vector(x, j, H); + } + accum = accumulate(cached[k], accum); + } + } + } else { + for (int index = threadIdx.x; index < vectors; index += Threads) { + int j = prefix + index * elements; + Vector value; + if constexpr (Aligned) { + *reinterpret_cast(&value) = *reinterpret_cast(x + j); + } else { + value = load_vector(x, j, H); + } + accum = accumulate(value, accum); + } + } + float sum = __fadd_rn(__fadd_rn(accum.x, accum.y), edge_sum); + float scale; + if constexpr (Threads == 32) { +#pragma unroll + for (int delta = 16; delta > 0; delta /= 2) { + sum = __fadd_rn(sum, __shfl_down_sync(0xffffffff, sum, delta)); + } + scale = threadIdx.x == 0 ? rsqrtf(sum / H + epsilon) : 0.0f; + scale = __shfl_sync(0xffffffff, scale, 0); + } else { + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage storage; + __shared__ float shared_scale; + sum = Reduce(storage).Sum(sum); + if (threadIdx.x == 0) shared_scale = rsqrtf(sum / H + epsilon); + __syncthreads(); + scale = shared_scale; + } + if (threadIdx.x < prefix) { + int j = threadIdx.x; + T rv{}; + if constexpr (AddResidual) rv = r[j]; + y[j] = finish_scalar(head_value, weight[j], rv, scale); + } + if (tail_start + threadIdx.x < H) { + int j = tail_start + threadIdx.x; + T rv{}; + if constexpr (AddResidual) rv = r[j]; + y[j] = finish_scalar(tail_value, weight[j], rv, scale); + } + if constexpr (CachePacks > 0) { + // Constant indices let nvcc keep each saved vector in registers. +#pragma unroll + for (int k = 0; k < CachePacks; ++k) { + int index = threadIdx.x + k * Threads; + if (index < vectors) { + int j = prefix + index * elements; + write_vector(y, weight, r, cached[k], j, H, scale); + } + } + } else { + for (int index = threadIdx.x; index < vectors; index += Threads) { + int j = prefix + index * elements; + Vector xv; + if constexpr (Aligned) { + *reinterpret_cast(&xv) = *reinterpret_cast(x + j); + } else { + xv = load_vector(x, j, H); + } + write_vector(y, weight, r, xv, j, H, scale); + } + } +} + +template +void launch_layout(at::Tensor& out, const at::Tensor& input, const at::Tensor& weight, + const std::optional& residual, RowLayout layout, int rows, + float epsilon, cudaStream_t stream) { + const T* x = reinterpret_cast(input.const_data_ptr()); + const T* w = reinterpret_cast(weight.const_data_ptr()); + const T* r = residual ? reinterpret_cast(residual->const_data_ptr()) : nullptr; + T* y = reinterpret_cast(out.mutable_data_ptr()); + constexpr int elements = 16 / sizeof(T); +#define RUN_RMS(H, K) \ + rms_norm_kernel<<>>( \ + y, x, w, r, layout, epsilon) + if constexpr (Aligned && Threads == 256) { + if (layout.hidden_size == 2048) { + RUN_RMS(2048, (2048 + Threads * elements - 1) / (Threads * elements)); + return; + } + } + // Cache small rows; stream very wide rows to avoid unbounded register pressure. + const int packs = (layout.hidden_size + Threads * elements - 1) / (Threads * elements); + if (packs <= 1) { RUN_RMS(0, 1); } + else if (packs <= 2) { RUN_RMS(0, 2); } + else if (packs <= 4) { RUN_RMS(0, 4); } + else { RUN_RMS(0, 0); } +#undef RUN_RMS +} + +template +void launch(at::Tensor& out, const at::Tensor& input, const at::Tensor& weight, + const std::optional& residual, float epsilon) { + constexpr int elements = 16 / sizeof(T); + RowLayout layout{}; + for (int i = 0; i < 3; ++i) layout.sizes[i] = 1; + const int outer = input.dim() - 1; + // A singleton penultimate dimension can have an arbitrary stride even when + // the tensor is contiguous, so a contiguous row advances by hidden_size. + layout.input_linear_stride = input.is_contiguous() + ? (outer ? input.size(-1) : 0) + : (input.dim() <= 2 ? input.stride(-2) : -1); + if constexpr (AddResidual) { + layout.residual_linear_stride = residual->is_contiguous() + ? (outer ? residual->size(-1) : 0) + : (residual->dim() <= 2 ? residual->stride(-2) : -1); + } + bool aligned = ((reinterpret_cast(input.const_data_ptr()) | + reinterpret_cast(weight.const_data_ptr()) | + reinterpret_cast(out.mutable_data_ptr())) & 15) == 0; + if constexpr (AddResidual) { + aligned = aligned && (reinterpret_cast(residual->const_data_ptr()) & 15) == 0; + } + for (int d = 0; d < outer; ++d) { + int target = 3 - outer + d; + layout.sizes[target] = input.size(d); + layout.input_strides[target] = input.stride(d); + aligned = aligned && (input.size(d) == 1 || input.stride(d) % elements == 0); + if constexpr (AddResidual) { + layout.residual_strides[target] = residual->stride(d); + aligned = aligned && (input.size(d) == 1 || residual->stride(d) % elements == 0); + } + } + layout.hidden_size = static_cast(input.size(-1)); + aligned = aligned && layout.hidden_size % elements == 0; + const int rows = static_cast(input.numel() / layout.hidden_size); + auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); + if (layout.hidden_size <= 256) { + if (aligned) launch_layout(out, input, weight, residual, layout, rows, epsilon, stream); + else launch_layout(out, input, weight, residual, layout, rows, epsilon, stream); + } else { + if (aligned) launch_layout(out, input, weight, residual, layout, rows, epsilon, stream); + else launch_layout(out, input, weight, residual, layout, rows, epsilon, stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +at::Tensor rms_norm_cuda(const at::Tensor& input, const at::Tensor& weight, double epsilon, + const std::optional& residual, bool match_torch) { + TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "rms_norm expects CUDA inputs"); + TORCH_CHECK(input.device() == weight.device(), "rms_norm expects inputs on the same device"); + TORCH_CHECK(input.dim() >= 1 && input.dim() <= 4, "rms_norm expects rank 1 through 4"); + TORCH_CHECK(input.stride(-1) == 1, "rms_norm expects a contiguous last dimension"); + TORCH_CHECK(weight.dim() == 1 && weight.is_contiguous() && weight.size(0) == input.size(-1), + "rms_norm expects contiguous weight with shape [hidden_size]"); + TORCH_CHECK(input.scalar_type() == weight.scalar_type(), "rms_norm expects matching dtypes"); + TORCH_CHECK(input.scalar_type() == at::kFloat || input.scalar_type() == at::kHalf || + input.scalar_type() == at::kBFloat16, + "rms_norm supports float32, float16 and bfloat16"); + if (residual) { + TORCH_CHECK(residual->is_cuda() && residual->device() == input.device(), + "rms_norm expects residual on the input device"); + TORCH_CHECK(residual->sizes() == input.sizes() && residual->stride(-1) == 1 && + residual->scalar_type() == input.scalar_type(), + "rms_norm expects matching residual shape, dtype and contiguous last dimension"); + } + TORCH_CHECK(!at::GradMode::is_enabled() || + (!input.requires_grad() && !weight.requires_grad() && + (!residual || !residual->requires_grad())), + "rms_norm is inference-only; use torch.no_grad() or the Torch reference"); + TORCH_CHECK(input.size(-1) <= std::numeric_limits::max(), "hidden size is too large"); + const c10::cuda::CUDAGuard guard(input.device()); + auto out = at::empty(input.sizes(), input.options()); + if (input.numel() == 0) return out; + TORCH_CHECK(input.numel() / input.size(-1) <= std::numeric_limits::max(), + "too many rms_norm rows"); + if (match_torch) { + TORCH_CHECK(input.dim() == 2 && input.size(-1) == 2048 && + (input.scalar_type() == at::kHalf || input.scalar_type() == at::kBFloat16), + "torch-order rms_norm requires 2D FP16/BF16 input of width 2048"); + TORCH_CHECK((reinterpret_cast(input.const_data_ptr()) & 7) == 0 && + (reinterpret_cast(weight.const_data_ptr()) & 7) == 0 && + (input.size(0) == 1 || input.stride(0) % 4 == 0), + "torch-order rms_norm requires 8-byte-aligned rows and weight"); + if (residual) { + TORCH_CHECK((reinterpret_cast(residual->const_data_ptr()) & 7) == 0 && + (residual->size(0) == 1 || residual->stride(0) % 4 == 0), + "torch-order rms_norm requires 8-byte-aligned residual rows"); + } +#define LAUNCH_TORCH_ORDER(T) \ + if (residual) launch_torch_order(out, input, weight, residual, epsilon); \ + else launch_torch_order(out, input, weight, residual, epsilon) + if (input.scalar_type() == at::kHalf) { LAUNCH_TORCH_ORDER(__half); } + else { LAUNCH_TORCH_ORDER(__nv_bfloat16); } +#undef LAUNCH_TORCH_ORDER + return out; + } +#define LAUNCH_RMS(T) \ + if (residual) launch(out, input, weight, residual, epsilon); \ + else launch(out, input, weight, residual, epsilon) + switch (input.scalar_type()) { + case at::kHalf: LAUNCH_RMS(__half); break; + case at::kBFloat16: LAUNCH_RMS(__nv_bfloat16); break; + case at::kFloat: LAUNCH_RMS(float); break; + default: TORCH_CHECK(false, "unsupported dtype"); + } +#undef LAUNCH_RMS + return out; +} + +} // namespace vllm_rlt diff --git a/csrc/ops.h b/csrc/ops.h index e00c054e..68fe28a3 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -6,5 +6,7 @@ namespace vllm_rlt { at::Tensor silu_and_mul_cuda(const at::Tensor& gate, const at::Tensor& up); +at::Tensor rms_norm_cuda(const at::Tensor& input, const at::Tensor& weight, double epsilon, + const std::optional& residual, bool match_torch); } // namespace vllm_rlt diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index afd5cf41..9376c633 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -6,10 +6,12 @@ TORCH_LIBRARY(vllm_rlt, ops) { ops.def("silu_and_mul(Tensor gate, Tensor up) -> Tensor"); + ops.def("rms_norm(Tensor input, Tensor weight, float epsilon, Tensor? residual=None, bool match_torch=False) -> Tensor"); } TORCH_LIBRARY_IMPL(vllm_rlt, CUDA, ops) { ops.impl("silu_and_mul", &vllm_rlt::silu_and_mul_cuda); + ops.impl("rms_norm", &vllm_rlt::rms_norm_cuda); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {} diff --git a/docs/assets/performance/ouro-rmsnorm-graph-b300.png b/docs/assets/performance/ouro-rmsnorm-graph-b300.png new file mode 100644 index 00000000..2e69d2c7 Binary files /dev/null and b/docs/assets/performance/ouro-rmsnorm-graph-b300.png differ diff --git a/docs/kernels.md b/docs/kernels.md index a1fda4e1..93483333 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -13,9 +13,13 @@ fallbacks. It neither compiles code nor creates a CUDA context. | `csrc/ops.h` | Native operator declarations | | `csrc/torch_bindings.cpp` | Operator schemas and CUDA dispatch registration | | `csrc/activation/silu_and_mul.cu` | SiLU-and-Mul CUDA validation, launch and computation | +| `csrc/normalization/rms_norm.cu` | RMSNorm and RMSNorm-then-residual CUDA kernels | +| `vllm_rlt/kernels/normalization.py` | RMSNorm Python entry and Torch reference | | `setup.py` | Optional build of `vllm_rlt._C` | | `tests/test_activation.py` | Reference, CUDA numerics, layouts, tails and graph replay | +| `tests/test_normalization.py` | RMSNorm numerics, strides, alignment and graph replay | | `benchmarks/kernels/silu_and_mul.py` | Matched Torch/native/vLLM accuracy and latency comparison | +| `benchmarks/kernels/rms_norm.py` | RMSNorm and residual accuracy and latency comparison | Add future kernels to the same extension and group sources by operation family. No backend registry or base class is required. Put future kernel performance @@ -73,6 +77,54 @@ are disabled, consistent with Ouro's inference-only model. `torch.compile` has no fake kernel registration. Tests cover the operator and Ouro MLP, including CUDA Graph capture/replay. +## RMSNorm and residual contract + +`rms_norm(hidden, weight, epsilon, residual=None)` returns a new tensor. Its +reference computes the mean square and inverse square root in FP32, multiplies +the input by that scale, rounds the normalized values to the input dtype, then +multiplies by `weight`. When supplied, the residual is added after the weight +multiplication. In Ouro, this fuses the existing `residual + RMSNorm(hidden)` +order; it does not implement `RMSNorm(residual + hidden)`. The weight is a +one-dimensional tensor of shape `[hidden_size]` and is shared across rows. + +With the extension installed, the native inference path supports CUDA FP32, +FP16 and BF16 tensors of rank 1 through 4, matching one-dimensional contiguous +weights, and matching residuals. The last input and residual dimensions must +have stride 1. Outer dimensions may be strided, including views into a larger +row; input, weight and residual addresses may have independent offsets. The +output is contiguous and none of the inputs is modified. Empty tensors return +without a kernel launch. CPU tensors, autograd, mixed dtypes and other layouts +use the Torch reference. Nanbeige and Huginn continue to use their existing +Torch RMSNorm path. + +The general CUDA kernel assigns one block per row. It sums squares in FP32, +uses a warp reduction for widths up to 256 and a block reduction otherwise, +then writes the scaled and weighted values, optionally with a residual. The +aligned body loads and stores 16 bytes at a time and uses paired FP16/BF16 +arithmetic. Unaligned output heads and tails use scalar accesses; independently +shifted input, weight and residual vectors are reconstructed from safe aligned +reads, with scalar reads at view boundaries. Small rows retain their input +vectors in registers across the reduction; wider rows stream them again to bound +register usage. Aligned width 2048 has a specialized path. FP32 pairs use paired +arithmetic on SM100 and newer, with scalar FP32 operations on earlier +architectures. These launch choices are implementation details and are not +claimed optimal for every GPU. + +The general native reduction can use a different floating-point summation order +from PyTorch, so its result need not be bitwise equal. Ouro uses a separate +2048-wide FP16/BF16 compatibility path that reproduces the tested PyTorch mean +reduction order, including its row-count-dependent block width. This path +fuses the residual, uses 8-byte vector loads and stores and paired arithmetic, +and keeps each thread's input values in registers across the reduction. It +supports two-dimensional inputs and residuals with 8-byte-aligned rows, including +strided views. Since the reduction tree is a PyTorch implementation detail, +this path is currently enabled only for PyTorch 2.13.0 on SM103; Ouro falls +back to the Torch reference for other versions, GPUs, shapes or alignments. +The general kernel remains available through `rms_norm(..., match_torch=False)`. +The tested compatibility path and the general path both retain the explicit +intermediate dtype roundings and residual order. CUDA Graph replay is supported; +a fake kernel for `torch.compile` is not registered. + ## Build and reference checks The extension builds automatically when the build environment has CUDA-enabled @@ -87,7 +139,7 @@ python -m pip install "torch>=2.5" setuptools wheel ninja python -m pip install --no-build-isolation -e . python -m pytest tests/test_activation.py -q CUDA_VISIBLE_DEVICES= python -m pytest \ - tests/test_activation.py --run-gpu -q + tests/test_activation.py tests/test_normalization.py --run-gpu -q ``` `pip` invokes setuptools, whose PyTorch `CUDAExtension`/`BuildExtension` compiles @@ -145,3 +197,24 @@ commit. The report retains library hashes, the local diff and raw measurements. Optional ops-only/full profiles are captured separately after timing; full capture enables shapes, memory and stacks. Large traces stay under ignored `artifacts/`. + +For RMSNorm, compare the Torch reference and both native paths without vLLM: + +```bash +CUDA_VISIBLE_DEVICES= python -m benchmarks.kernels.rms_norm \ + --output artifacts/rmsnorm/results.json \ + --profile-dir artifacts/rmsnorm/profiles +``` + +An optional source-built vLLM comparison can be supplied with +`--vllm-library --vllm-namespace --vllm-source-commit `. +Its different intermediate rounding is reported separately from the reference +semantics. The general native path can use a different reduction order than +PyTorch; use the `rlt_torch_order` measurements when assessing Ouro's actual +compatibility path. + +For one example of the model-level execution change, see the +[matched B300 Graph timeline](assets/performance/ouro-rmsnorm-graph-b300.png) +for an Ouro-1.4B BF16 request with 128 input tokens and 2 output tokens. The +profiler trace illustrates work and launch counts; use separate unprofiled +measurements for latency and throughput. diff --git a/setup.py b/setup.py index b83b309d..c72d9c75 100644 --- a/setup.py +++ b/setup.py @@ -37,6 +37,7 @@ def native_build_options(): sources=[ "csrc/torch_bindings.cpp", "csrc/activation/silu_and_mul.cu", + "csrc/normalization/rms_norm.cu", ], extra_compile_args={"cxx": ["-O3"], "nvcc": ["-O3"]}, ) diff --git a/tests/test_normalization.py b/tests/test_normalization.py index f7111244..3f67cd8e 100644 --- a/tests/test_normalization.py +++ b/tests/test_normalization.py @@ -4,6 +4,7 @@ import pytest import torch +from vllm_rlt.kernels.normalization import rms_norm, rms_norm_reference from vllm_rlt.layers import RMSNorm from vllm_rlt.models.nanbeige import NanbeigeRMSNorm from vllm_rlt.models.ouro import OuroRMSNorm @@ -69,3 +70,262 @@ def test_rmsnorm_explicit_epsilon(): # Variance of ones is 1.0; 1.0 / sqrt(1.0 + 1.0) = 1.0 / sqrt(2) ≈ 0.7071 expected_large = 1.0 / (2.0**0.5) assert torch.allclose(out_large, torch.full_like(x, expected_large), atol=1e-4) + + +def test_opt_in_residual_preserves_cpu_autograd_and_mixed_dtype(): + norm = RMSNorm(5, use_native=True) + x = torch.randn(2, 5, dtype=torch.bfloat16, requires_grad=True) + residual = torch.randn(2, 5, dtype=torch.float32, requires_grad=True) + actual = norm(x, residual=residual) + expected = rms_norm_reference(x, norm.weight, norm.variance_epsilon, residual) + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual.sum().backward() + assert x.grad is not None and residual.grad is not None and norm.weight.grad is not None + + +@pytest.fixture +def native_rms_norm(): + from vllm_rlt.kernels._extension import load_extension + + load_extension() + assert torch._C._dispatch_has_kernel_for_dispatch_key("vllm_rlt::rms_norm", "CUDA") + return rms_norm + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [8, 17, 128, 256, 257, 768, 2048, 2049, 4096, 8193]) +@pytest.mark.parametrize("add_residual", [False, True]) +@torch.inference_mode() +def test_rmsnorm_cuda_values_tail_and_staged_rounding(native_rms_norm, dtype, width, add_residual): + torch.manual_seed(42) + x = torch.randn(3, width, device="cuda", dtype=dtype) + weight = (torch.randn(width, device="cuda") * 0.1 + 1).to(dtype) + residual = torch.randn_like(x) if add_residual else None + x_before, w_before = x.clone(), weight.clone() + r_before = residual.clone() if residual is not None else None + result = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + # CUDA's reduction order differs from torch.mean; keep the dtype rounding + # boundaries and independently check the residual's input is untouched. + tolerance = { + torch.float32: (2e-6, 2e-6), + torch.float16: (2e-3, 2e-3), + torch.bfloat16: (1e-2, 1e-2), + }[dtype] + torch.testing.assert_close(result, expected, rtol=tolerance[0], atol=tolerance[1]) + assert result.shape == x.shape and result.dtype == x.dtype and result.is_contiguous() + assert result.data_ptr() not in (x.data_ptr(), weight.data_ptr()) + assert torch.equal(x, x_before) and torch.equal(weight, w_before) + if residual is not None: + assert result.data_ptr() != residual.data_ptr() + assert torch.equal(residual, r_before) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rows", [1, 2, 3, 4, 5, 8, 15, 16, 17, 128, 1024]) +@pytest.mark.parametrize("add_residual", [False, True]) +@torch.inference_mode() +def test_rmsnorm_cuda_torch_order_is_exact(native_rms_norm, dtype, rows, add_residual): + # The row count changes torch.mean's reduction tree. Cover the boundaries + # where the compatibility kernel changes its block width. + from vllm_rlt.kernels.normalization import _torch_order_validated + + if not _torch_order_validated(torch.cuda.current_device()): + pytest.skip("the Torch-order CUDA kernel is not enabled on this Torch/GPU combination") + x = torch.randn(rows, 2048, dtype=dtype, device="cuda") + weight = torch.randn(2048, dtype=dtype, device="cuda") + residual = torch.randn_like(x) if add_residual else None + actual = native_rms_norm(x, weight, 1e-6, residual, match_torch=True) + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_rmsnorm_cuda_torch_order_strides_and_fallback(native_rms_norm, dtype): + # An aligned SP-style slice can use the exact kernel, while an offset that + # cannot safely use its vector loads falls back to the Torch reference. + rows, width = 4, 2048 + base = torch.randn(rows, width * 2 + 8, dtype=dtype, device="cuda") + residual_base = torch.randn_like(base) + weight = torch.randn(width + 1, dtype=dtype, device="cuda") + for offset in (0, 1): + x = base[:, offset : offset + width] + residual = residual_base[:, offset : offset + width] + w = weight[offset : offset + width] + actual = native_rms_norm(x, w, 1e-6, residual, match_torch=True) + expected = rms_norm_reference(x, w, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_rmsnorm_cuda_torch_order_nonfinite(native_rms_norm, dtype): + x = torch.zeros(4, 2048, dtype=dtype, device="cuda") + x[1, 0] = float("inf") + x[2, 0] = float("-inf") + x[3, 0] = float("nan") + weight = torch.randn(2048, dtype=dtype, device="cuda") + residual = torch.randn_like(x) + actual = native_rms_norm(x, weight, 1e-6, residual, match_torch=True) + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.gpu +@torch.inference_mode() +def test_rmsnorm_cuda_torch_order_graph_replay(native_rms_norm): + dtype = torch.bfloat16 + x = torch.randn(17, 2048, dtype=dtype, device="cuda") + weight = torch.randn(2048, dtype=dtype, device="cuda") + residual = torch.randn_like(x) + native_rms_norm(x, weight, 1e-6, residual, match_torch=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = native_rms_norm(x, weight, 1e-6, residual, match_torch=True) + x.fill_(0.25) + residual.fill_(0.5) + graph.replay() + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rank", [1, 2, 3, 4]) +@pytest.mark.parametrize("add_residual", [False, True]) +@torch.inference_mode() +def test_rmsnorm_cuda_strides_and_independent_offsets(native_rms_norm, dtype, rank, add_residual): + torch.manual_seed(123) + width = 33 + shape = (width,) if rank == 1 else (*((2,) * (rank - 1)), width) + base_shape = (*shape[:-1], width + 3) + x = torch.randn(base_shape, device="cuda", dtype=dtype)[..., 1 : width + 1] + residual = torch.randn(base_shape, device="cuda", dtype=dtype)[..., 2 : width + 2] + weight = torch.randn(width + 1, device="cuda", dtype=dtype)[1:] + if not add_residual: + residual = None + actual = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + tol = 2e-6 if dtype == torch.float32 else 2e-3 if dtype == torch.float16 else 1e-2 + torch.testing.assert_close(actual, expected, rtol=tol, atol=tol) + assert actual.is_contiguous() and actual.shape == x.shape + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_rmsnorm_cuda_contiguous_with_arbitrary_singleton_stride(native_rms_norm, dtype): + # PyTorch ignores singleton strides in is_contiguous(); flattened rows must + # still advance by width, not by the arbitrary penultimate stride. + width = 33 + x = torch.randn(4, width, device="cuda", dtype=dtype).as_strided( + (4, 1, width), (width, 7 * width, 1) + ) + residual = torch.randn(4, width, device="cuda", dtype=dtype).as_strided( + (4, 1, width), (width, 5 * width, 1) + ) + assert x.is_contiguous() and residual.is_contiguous() + weight = torch.randn(width, device="cuda", dtype=dtype) + actual = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + tol = 2e-6 if dtype == torch.float32 else 2e-3 if dtype == torch.float16 else 1e-2 + torch.testing.assert_close(actual, expected, rtol=tol, atol=tol) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_rmsnorm_cuda_all_alignment_combinations(native_rms_norm, dtype): + # Each row begins at a different offset; input, residual and weight offsets + # vary independently across every possible 16-byte alignment. + torch.manual_seed(7) + elements = 16 // torch.empty((), dtype=dtype).element_size() + width = 2 * elements + 1 + x_base = torch.randn(2, width + elements, device="cuda", dtype=dtype) + r_base = torch.randn_like(x_base) + w_base = torch.randn(width + elements, device="cuda", dtype=dtype) + tol = 2e-6 if dtype == torch.float32 else 2e-3 if dtype == torch.float16 else 1e-2 + for x_offset in range(elements): + x = x_base[:, x_offset : x_offset + width] + for w_offset in range(elements): + weight = w_base[w_offset : w_offset + width] + for r_offset in range(elements): + residual = r_base[:, r_offset : r_offset + width] + actual = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=tol, atol=tol) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_rmsnorm_cuda_nonfinite_and_zero_rows(native_rms_norm, dtype): + x = torch.zeros(4, 33, dtype=dtype, device="cuda") + x[1, 0] = float("inf") + x[2, 0] = float("-inf") + x[3, 0] = float("nan") + weight = torch.randn(33, dtype=dtype, device="cuda") + residual = torch.randn_like(x) + actual = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.gpu +@torch.inference_mode() +def test_rmsnorm_cuda_falls_back_for_other_layouts_and_mixed_dtypes(native_rms_norm): + width = 33 + inputs = ( + torch.randn(2, width * 2, device="cuda")[..., ::2], + torch.randn(2, 1, 2, 1, width, device="cuda"), + torch.randn(2, width, device="cuda", dtype=torch.bfloat16), + ) + for x in inputs: + # The last case intentionally has FP32 weight and BF16 activations. + weight = torch.randn(width, device="cuda", dtype=torch.float32) + residual = torch.randn_like(x) + actual = native_rms_norm(x, weight, 1e-6, residual) + expected = rms_norm_reference(x, weight, 1e-6, residual) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(0, 2048), (2, 0)]) +@torch.inference_mode() +def test_rmsnorm_cuda_empty(native_rms_norm, dtype, shape): + x = torch.empty(shape, device="cuda", dtype=dtype) + weight = torch.empty(shape[-1], device="cuda", dtype=dtype) + actual = native_rms_norm(x, weight, 1e-6) + assert actual.shape == x.shape and actual.dtype == dtype and actual.numel() == 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("add_residual", [False, True]) +@torch.inference_mode() +def test_rmsnorm_cuda_current_stream_and_graph_replay(native_rms_norm, dtype, add_residual): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), torch.inference_mode(): + x = torch.randn(7, 2049, device="cuda", dtype=dtype) + weight = torch.randn(2049, device="cuda", dtype=dtype) + residual = torch.randn_like(x) if add_residual else None + for _ in range(3): + native_rms_norm(x, weight, 1e-6, residual) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, stream=stream): + actual = native_rms_norm(x, weight, 1e-6, residual) + x.fill_(0.5) + if residual is not None: + residual.fill_(0.25) + graph.replay() + expected = rms_norm_reference(x, weight, 1e-6, residual) + tol = 2e-6 if dtype == torch.float32 else 2e-3 if dtype == torch.float16 else 1e-2 + torch.testing.assert_close(actual, expected, rtol=tol, atol=tol) diff --git a/vllm_rlt/kernels/__init__.py b/vllm_rlt/kernels/__init__.py index 7cf7031f..a51348ab 100644 --- a/vllm_rlt/kernels/__init__.py +++ b/vllm_rlt/kernels/__init__.py @@ -1,5 +1,6 @@ """Register native and Torch kernels once without initializing a GPU.""" from .activation import silu_and_mul +from .normalization import rms_norm -__all__ = ["silu_and_mul"] +__all__ = ["rms_norm", "silu_and_mul"] diff --git a/vllm_rlt/kernels/normalization.py b/vllm_rlt/kernels/normalization.py new file mode 100644 index 00000000..a95b65fe --- /dev/null +++ b/vllm_rlt/kernels/normalization.py @@ -0,0 +1,109 @@ +"""Inference RMSNorm with an optional residual and a PyTorch reference.""" + +from functools import cache + +import torch + +from vllm_rlt.kernels._extension import try_load_extension + + +def rms_norm_reference( + hidden: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + residual: torch.Tensor | None = None, +) -> torch.Tensor: + value = hidden.float() + value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + epsilon) + result = weight * value.to(hidden.dtype) + return result if residual is None else residual + result + + +_library = torch.library.Library("vllm_rlt", "FRAGMENT") +if not try_load_extension(): + _library.define( + "rms_norm(Tensor input, Tensor weight, float epsilon, Tensor? residual=None, " + "bool match_torch=False) -> Tensor" + ) + + +def _reference_dispatch(hidden, weight, epsilon, residual=None, match_torch=False): + return rms_norm_reference(hidden, weight, epsilon, residual) + + +_library.impl("rms_norm", _reference_dispatch, "CompositeExplicitAutograd") + + +@cache +def _torch_order_validated(device_index: int) -> bool: + # torch.mean's reduction tree is an implementation detail. Only use the + # compatibility kernel for the framework and GPU on which it was checked. + return torch.__version__.split("+", 1)[0] == "2.13.0" and torch.cuda.get_device_capability( + device_index + ) == (10, 3) + + +def rms_norm( + hidden: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + residual: torch.Tensor | None = None, + *, + match_torch: bool = False, +) -> torch.Tensor: + """Use the native CUDA operator where its inference layout contract applies.""" + # The module historically accepts autograd, mixed dtypes and arbitrary views. + # Keep those semantics in Torch rather than copying or narrowing inputs. + supported = ( + hidden.is_cuda + and 1 <= hidden.ndim <= 4 + and hidden.stride(-1) == 1 + and hidden.dtype in (torch.float32, torch.float16, torch.bfloat16) + and weight.ndim == 1 + and weight.is_contiguous() + and weight.shape[0] == hidden.shape[-1] + and weight.dtype == hidden.dtype + and weight.device == hidden.device + and ( + residual is None + or ( + residual.shape == hidden.shape + and residual.device == hidden.device + and residual.dtype == hidden.dtype + and residual.stride(-1) == 1 + ) + ) + and ( + not torch.is_grad_enabled() + or not ( + hidden.requires_grad + or weight.requires_grad + or (residual is not None and residual.requires_grad) + ) + ) + ) + if not supported: + return rms_norm_reference(hidden, weight, epsilon, residual) + if match_torch: + # The model-compatible kernel reproduces torch.mean's CUDA reduction + # order only for these shapes and alignments. Preserve the Torch path + # elsewhere instead of silently changing early-exit behavior. + compatible = ( + hidden.ndim == 2 + and hidden.shape[-1] == 2048 + and hidden.dtype in (torch.float16, torch.bfloat16) + and _torch_order_validated(hidden.device.index) + and hidden.data_ptr() % 8 == 0 + and weight.data_ptr() % 8 == 0 + and (hidden.shape[0] <= 1 or hidden.stride(0) % 4 == 0) + and ( + residual is None + or ( + residual.data_ptr() % 8 == 0 + and (residual.shape[0] <= 1 or residual.stride(0) % 4 == 0) + ) + ) + ) + if not compatible: + return rms_norm_reference(hidden, weight, epsilon, residual) + return torch.ops.vllm_rlt.rms_norm(hidden, weight, epsilon, residual, match_torch) diff --git a/vllm_rlt/layers/normalization.py b/vllm_rlt/layers/normalization.py index 9d32247d..ac46b52b 100644 --- a/vllm_rlt/layers/normalization.py +++ b/vllm_rlt/layers/normalization.py @@ -4,6 +4,8 @@ import torch from torch import nn +from vllm_rlt.kernels import rms_norm + class RMSNorm(nn.Module): """Root Mean Square Layer Normalization. @@ -12,12 +14,16 @@ class RMSNorm(nn.Module): across all input activation dtypes (BF16/FP16/FP32). """ - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + def __init__(self, hidden_size: int, eps: float = 1e-6, *, use_native: bool = False) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps + self.use_native = use_native - def forward(self, hidden: torch.Tensor) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor: + if self.use_native: + return rms_norm(hidden, self.weight, self.variance_epsilon, residual, match_torch=True) value = hidden.float() value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + self.variance_epsilon) - return self.weight * value.to(hidden.dtype) + result = self.weight * value.to(hidden.dtype) + return result if residual is None else residual + result diff --git a/vllm_rlt/models/ouro.py b/vllm_rlt/models/ouro.py index 31b37f9c..90e00427 100644 --- a/vllm_rlt/models/ouro.py +++ b/vllm_rlt/models/ouro.py @@ -202,10 +202,16 @@ def __init__(self, config: OuroConfig, layer_idx: int) -> None: super().__init__() self.self_attn = OuroAttention(config, layer_idx) self.mlp = OuroMLP(config) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, use_native=True) + self.input_layernorm_2 = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, use_native=True + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, use_native=True + ) + self.post_attention_layernorm_2 = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, use_native=True + ) def forward( self, @@ -215,9 +221,9 @@ def forward( cache: "KVCacheManager", ) -> torch.Tensor: attention = self.self_attn(self.input_layernorm(hidden), position_embeddings, batch, cache) - hidden = hidden + self.input_layernorm_2(attention) - return hidden + self.post_attention_layernorm_2( - self.mlp(self.post_attention_layernorm(hidden)) + hidden = self.input_layernorm_2(attention, residual=hidden) + return self.post_attention_layernorm_2( + self.mlp(self.post_attention_layernorm(hidden)), residual=hidden ) @@ -228,7 +234,7 @@ def __init__(self, config: OuroConfig) -> None: self.layers = nn.ModuleList( OuroDecoderLayer(config, layer) for layer in range(config.num_hidden_layers) ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, use_native=True) self.rotary_emb = RotaryEmbedding(head_dim=config.head_dim, rope_theta=config.rope_theta) self.early_exit_gate = nn.Linear(config.hidden_size, 1, bias=True)