diff --git a/csrc/cpu/layernorm.cpp b/csrc/cpu/layernorm.cpp index a76ad08928..704fb14633 100644 --- a/csrc/cpu/layernorm.cpp +++ b/csrc/cpu/layernorm.cpp @@ -4,8 +4,9 @@ namespace { template void rms_norm_impl(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, - const scalar_t* __restrict__ weight, const float epsilon, - const int num_tokens, const int hidden_size) { + const scalar_t* __restrict__ weight, const bool has_weight, + const float epsilon, const int num_tokens, + const int hidden_size) { using scalar_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0); @@ -27,12 +28,15 @@ void rms_norm_impl(scalar_t* __restrict__ out, for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) { scalar_vec_t x(input_p + j); - scalar_vec_t w(weight + j); - vec_op::FP32Vec8 fp32_x(x); - vec_op::FP32Vec8 fp32_w(w); - - vec_op::FP32Vec8 fp32_out = fp32_x * fp32_s_variance * fp32_w; + vec_op::FP32Vec8 fp32_out; + if (has_weight) { + scalar_vec_t w(weight + j); + vec_op::FP32Vec8 fp32_w(w); + fp32_out = fp32_x * fp32_s_variance * fp32_w; + } else { + fp32_out = fp32_x * fp32_s_variance; + } scalar_vec_t out(fp32_out); out.save(output_p + j); @@ -44,8 +48,8 @@ template void fused_add_rms_norm_impl(scalar_t* __restrict__ input, scalar_t* __restrict__ residual, const scalar_t* __restrict__ weight, - const float epsilon, const int num_tokens, - const int hidden_size) { + const bool has_weight, const float epsilon, + const int num_tokens, const int hidden_size) { using scalar_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0); @@ -72,13 +76,18 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input, vec_op::FP32Vec8 fp32_s_variance(s_variance); for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) { - scalar_vec_t w(weight + j); - scalar_vec_t res(residual_p + j); - - vec_op::FP32Vec8 fp32_w(w); - vec_op::FP32Vec8 fp32_res(res); - - vec_op::FP32Vec8 fp32_out = fp32_res * fp32_s_variance * fp32_w; + vec_op::FP32Vec8 fp32_out; + if (has_weight) { + scalar_vec_t w(weight + j); + scalar_vec_t res(residual_p + j); + vec_op::FP32Vec8 fp32_w(w); + vec_op::FP32Vec8 fp32_res(res); + fp32_out = fp32_res * fp32_s_variance * fp32_w; + } else { + scalar_vec_t res(residual_p + j); + vec_op::FP32Vec8 fp32_res(res); + fp32_out = fp32_res * fp32_s_variance; + } scalar_vec_t out(fp32_out); out.save(input_p + j); @@ -87,31 +96,41 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input, } } // namespace -void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, - double epsilon) { +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon) { int hidden_size = input.size(-1); int num_tokens = input.numel() / hidden_size; + const bool has_weight = weight.has_value(); + if (has_weight) { + TORCH_CHECK(weight->is_contiguous()); + } VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] { CPU_KERNEL_GUARD_IN(rms_norm_impl) rms_norm_impl(out.data_ptr(), input.data_ptr(), - weight.data_ptr(), epsilon, num_tokens, - hidden_size); + has_weight ? weight->data_ptr() : nullptr, + has_weight, epsilon, num_tokens, hidden_size); CPU_KERNEL_GUARD_OUT(rms_norm_impl) }); } void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - torch::Tensor& weight, double epsilon) { + std::optional weight, double epsilon) { int hidden_size = input.size(-1); int num_tokens = input.numel() / hidden_size; + const bool has_weight = weight.has_value(); + if (has_weight) { + TORCH_CHECK(weight->scalar_type() == input.scalar_type()); + TORCH_CHECK(weight->is_contiguous()); + } VLLM_DISPATCH_FLOATING_TYPES( input.scalar_type(), "fused_add_rms_norm_impl", [&] { CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl) fused_add_rms_norm_impl( input.data_ptr(), residual.data_ptr(), - weight.data_ptr(), epsilon, num_tokens, hidden_size); + has_weight ? weight->data_ptr() : nullptr, has_weight, + epsilon, num_tokens, hidden_size); CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl) }); } diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index b1a9342dee..2aad5e2387 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -310,13 +310,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Layernorm // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( - "rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> " + "rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> " "()"); ops.impl("rms_norm", torch::kCPU, &rms_norm); // In-place fused Add and RMS Normalization. ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, " + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " "float epsilon) -> ()"); ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm); diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 37df6be329..eb121b0b88 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -11,7 +11,7 @@ namespace vllm { // TODO(woosuk): Further optimize this kernel. -template +template __global__ void rms_norm_kernel( scalar_t* __restrict__ out, // [..., hidden_size] const scalar_t* __restrict__ input, // [..., hidden_size] @@ -20,7 +20,7 @@ __global__ void rms_norm_kernel( const int64_t input_stride_d4, // input.stride(-4) const int64_t input_shape_d2, // input.size(-2) const int64_t input_shape_d3, // input.size(-3) - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; @@ -74,11 +74,19 @@ __global__ void rms_norm_kernel( for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { vec_n_t dst; vec_n_t src1 = v_in[i]; - vec_n_t src2 = v_w[i]; + vec_n_t src2; + if constexpr (HasWeight) { + src2 = v_w[i]; + } #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - dst.val[j] = static_cast(x * s_variance) * src2.val[j]; + scalar_t normalized = static_cast(x * s_variance); + if constexpr (HasWeight) { + dst.val[j] = normalized * src2.val[j]; + } else { + dst.val[j] = normalized; + } } v_out[i] = dst; } @@ -88,13 +96,13 @@ __global__ void rms_norm_kernel( Additional optimizations we can make in this case are packed and vectorized operations, which help with the memory latency bottleneck. */ -template +template __global__ std::enable_if_t<(width > 0) && _typeConvert::exists> fused_add_rms_norm_kernel( scalar_t* __restrict__ input, // [..., hidden_size] const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight const float epsilon, const int num_tokens, const int hidden_size) { // Sanity checks on our vector struct and type-punned pointer arithmetic static_assert(std::is_pod_v<_f16Vec>); @@ -136,13 +144,21 @@ fused_add_rms_norm_kernel( int id = blockIdx.x * vec_hidden_size + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec res = residual_v[id]; - _f16Vec w = weight_v[idx]; _f16Vec out; using Converter = _typeConvert; + if constexpr (HasWeight) { + _f16Vec w = weight_v[idx]; #pragma unroll - for (int j = 0; j < width; ++j) { - float x = Converter::convert(res.data[j]); - out.data[j] = Converter::convert(x * s_variance) * w.data[j]; + for (int j = 0; j < width; ++j) { + float x = Converter::convert(res.data[j]); + out.data[j] = Converter::convert(x * s_variance) * w.data[j]; + } + } else { +#pragma unroll + for (int j = 0; j < width; ++j) { + float x = Converter::convert(res.data[j]); + out.data[j] = Converter::convert(x * s_variance); + } } input_v[strided_id] = out; } @@ -151,13 +167,13 @@ fused_add_rms_norm_kernel( /* Generic fused_add_rms_norm_kernel The width field is not used here but necessary for other specializations. */ -template +template __global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> fused_add_rms_norm_kernel( scalar_t* __restrict__ input, // [..., hidden_size] const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; @@ -181,23 +197,29 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - input[blockIdx.x * input_stride + idx] = - (scalar_t)(x * s_variance) * weight[idx]; + if constexpr (HasWeight) { + input[blockIdx.x * input_stride + idx] = + (scalar_t)(x * s_variance) * weight[idx]; + } else { + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance); + } } } } // namespace vllm -void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] - torch::stable::Tensor& input, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] +void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] + torch::stable::Tensor& input, // [..., hidden_size] + std::optional weight, // [hidden_size] double epsilon) { STD_TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { input = torch::stable::contiguous(input); } STD_TORCH_CHECK(input.stride(-1) == 1); - STD_TORCH_CHECK(weight.is_contiguous()); + if (weight.has_value()) { + STD_TORCH_CHECK(weight->is_contiguous()); + } int hidden_size = input.size(-1); @@ -215,46 +237,69 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); + const bool has_weight = weight.has_value(); VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] { VLLM_STABLE_DISPATCH_FLOATING_TYPES( input.scalar_type(), "rms_norm_kernel", [&] { + const scalar_t* weight_ptr = + has_weight ? weight->const_data_ptr() : nullptr; const int calculated_vec_size = std::gcd(16 / sizeof(scalar_t), hidden_size); const int block_size = std::min(hidden_size / calculated_vec_size, max_block_size); dim3 block(block_size); VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] { - vllm::rms_norm_kernel - <<>>( - out.mutable_data_ptr(), - input.const_data_ptr(), input_stride_d2, - input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight.const_data_ptr(), epsilon, - num_tokens, hidden_size); + if (has_weight) { + vllm::rms_norm_kernel + <<>>( + out.mutable_data_ptr(), + input.const_data_ptr(), input_stride_d2, + input_stride_d3, input_stride_d4, input_shape_d2, + input_shape_d3, weight_ptr, epsilon, num_tokens, + hidden_size); + } else { + vllm::rms_norm_kernel + <<>>( + out.mutable_data_ptr(), + input.const_data_ptr(), input_stride_d2, + input_stride_d3, input_stride_d4, input_shape_d2, + input_shape_d3, weight_ptr, epsilon, num_tokens, + hidden_size); + } }); }); }); } -#define LAUNCH_FUSED_ADD_RMS_NORM(width) \ - VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ - input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ - vllm::fused_add_rms_norm_kernel \ - <<>>( \ - input.mutable_data_ptr(), input_stride, \ - residual.mutable_data_ptr(), \ - weight.const_data_ptr(), epsilon, num_tokens, \ - hidden_size); \ +#define LAUNCH_FUSED_ADD_RMS_NORM(width, has_weight) \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ + if (has_weight) { \ + vllm::fused_add_rms_norm_kernel \ + <<>>( \ + input.mutable_data_ptr(), input_stride, \ + residual.mutable_data_ptr(), \ + weight->const_data_ptr(), epsilon, num_tokens, \ + hidden_size); \ + } else { \ + vllm::fused_add_rms_norm_kernel \ + <<>>( \ + input.mutable_data_ptr(), input_stride, \ + residual.mutable_data_ptr(), nullptr, epsilon, \ + num_tokens, hidden_size); \ + } \ }); void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] torch::stable::Tensor& residual, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] + std::optional weight, double epsilon) { - STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type()); STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type()); STD_TORCH_CHECK(residual.is_contiguous()); - STD_TORCH_CHECK(weight.is_contiguous()); + if (weight.has_value()) { + STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type()); + STD_TORCH_CHECK(weight->is_contiguous()); + } int hidden_size = input.size(-1); int64_t input_stride = input.stride(-2); int num_tokens = input.numel() / hidden_size; @@ -269,30 +314,33 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); - /*If the tensor types are FP16/BF16, try to use the optimized kernel - with packed + vectorized ops. - Max optimization is achieved with a width-8 vector of FP16/BF16s - since we can load at most 128 bits at once in a global memory op. - However, this requires each tensor's data to be aligned to 16 - bytes. - */ + constexpr int vector_width = 8; + constexpr int req_alignment_bytes = vector_width * 2; auto inp_ptr = reinterpret_cast(input.data_ptr()); auto res_ptr = reinterpret_cast(residual.data_ptr()); - auto wt_ptr = reinterpret_cast(weight.data_ptr()); - constexpr int vector_width = 8; - constexpr int req_alignment_bytes = - vector_width * 2; // vector_width * sizeof(bfloat16 or float16) (float32 - // falls back to non-vectorized version anyway) - bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && - res_ptr % req_alignment_bytes == 0 && - wt_ptr % req_alignment_bytes == 0; bool offsets_are_multiple_of_vector_width = hidden_size % vector_width == 0 && input_stride % vector_width == 0; bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); - if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && - !batch_invariant_launch) { - LAUNCH_FUSED_ADD_RMS_NORM(8); + const bool has_weight = weight.has_value(); + if (has_weight) { + auto wt_ptr = reinterpret_cast(weight->data_ptr()); + bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && + res_ptr % req_alignment_bytes == 0 && + wt_ptr % req_alignment_bytes == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && + !batch_invariant_launch) { + LAUNCH_FUSED_ADD_RMS_NORM(8, true); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0, true); + } } else { - LAUNCH_FUSED_ADD_RMS_NORM(0); + bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && + res_ptr % req_alignment_bytes == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && + !batch_invariant_launch) { + LAUNCH_FUSED_ADD_RMS_NORM(8, false); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0, false); + } } } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 49ce174400..9efc12e9f4 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -188,11 +188,12 @@ torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, // Layernorm kernels (shared CUDA/ROCm) void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - torch::stable::Tensor& weight, double epsilon); + std::optional weight, double epsilon); void fused_add_rms_norm(torch::stable::Tensor& input, torch::stable::Tensor& residual, - torch::stable::Tensor& weight, double epsilon); + std::optional weight, + double epsilon); // Layernorm-quant kernels (shared CUDA/ROCm) void rms_norm_static_fp8_quant(torch::stable::Tensor& out, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index f360c83b8a..0aabcc757d 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -369,12 +369,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> " + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " "()"); // In-place fused Add and RMS Normalization. ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, " + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " "float epsilon) -> ()"); // Layernorm-quant diff --git a/csrc/ops.h b/csrc/ops.h index b909c5711d..ec3f5e187c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -34,11 +34,11 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. -void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, - double epsilon); +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon); void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - torch::Tensor& weight, double epsilon); + std::optional weight, double epsilon); // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable // ABI for CUDA). It remains here because the CPU build still uses these diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index c39d42c759..fde09710b5 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -6,6 +6,7 @@ import torch from tests.kernels.quant_utils import FP8_DTYPE from tests.kernels.utils import opcheck +from vllm import ir from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -27,6 +28,10 @@ CUDA_DEVICES = [ ] +def _rms_norm_tolerance(dtype: torch.dtype) -> dict[str, float]: + return ir.ops.rms_norm.get_tolerance(dtype) + + @pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) @pytest.mark.parametrize("add_residual", ADD_RESIDUAL) @@ -81,6 +86,49 @@ def test_rms_norm( ) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_rms_norm_weightless( + default_vllm_config, + num_tokens: int, + hidden_size: int, + add_residual: bool, + dtype: torch.dtype, + seed: int, + device: str, +) -> None: + set_random_seed(seed) + torch.set_default_device(device) + layer = RMSNorm(hidden_size, has_weight=False).to(dtype=dtype) + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + residual = torch.randn_like(x) if add_residual else None + + ref_out = layer.forward_native(x, residual) + out = layer(x, residual) + tol = _rms_norm_tolerance(dtype) + if add_residual: + torch.testing.assert_close(out[0], ref_out[0], **tol) + torch.testing.assert_close(out[1], ref_out[1], **tol) + else: + torch.testing.assert_close(out, ref_out, **tol) + + if residual is not None: + opcheck( + torch.ops._C.fused_add_rms_norm, + (x, residual, None, layer.variance_epsilon), + ) + else: + opcheck( + torch.ops._C.rms_norm, + (out, x, None, layer.variance_epsilon), + ) + + @pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) @pytest.mark.parametrize("add_residual", ADD_RESIDUAL) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 5f52bc50eb..1b49c9159d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -315,13 +315,19 @@ def rotary_embedding( # layer norm ops def rms_norm( - out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor | None, + epsilon: float, ) -> None: torch.ops._C.rms_norm(out, input, weight, epsilon) def fused_add_rms_norm( - input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, epsilon: float + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor | None, + epsilon: float, ) -> None: # Note: this func is batch invariant torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon) diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index 3b194b2ab9..6ae5d9939e 100644 --- a/vllm/kernels/vllm_c.py +++ b/vllm/kernels/vllm_c.py @@ -25,9 +25,6 @@ rms_no_var_size = lambda x, weight, epsilon, variance_size=None: ( def rms_norm( x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None ) -> Tensor: - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) assert variance_size is None # ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors. # Higher-rank callers still normalize over the last dimension, so flatten @@ -64,10 +61,6 @@ def fused_add_rms_norm( epsilon: float, variance_size: int | None = None, ) -> tuple[Tensor, Tensor]: - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) - assert variance_size is None if IS_ROCM and (not x.is_contiguous() or not x_residual.is_contiguous()): output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn( diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py index 5e7f90f708..8a86b1226b 100644 --- a/vllm/kernels/xpu_ops.py +++ b/vllm/kernels/xpu_ops.py @@ -29,10 +29,12 @@ rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is Non def rms_norm( x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None ) -> Tensor: - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) assert variance_size is None + if weight is None: + # Weightless _C ops are CUDA-only; native skips the multiply on XPU. + return ir.ops.rms_norm.impls["native"].impl_fn( + x, weight, epsilon, variance_size + ) output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output @@ -57,10 +59,14 @@ def fused_add_rms_norm( epsilon: float, variance_size: int | None = None, ) -> tuple[Tensor, Tensor]: - if weight is None: - # Kernel requires weight tensor, pass ones - weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) - assert variance_size is None + if weight is None: + # Weightless _C ops are CUDA-only; native skips the multiply on XPU. + output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn( + x, x_residual, weight, epsilon, variance_size + ) + x.copy_(output) + x_residual.copy_(residual) + return x, x_residual torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) return x, x_residual diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 13b0ae7813..8418245b82 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -9,7 +9,6 @@ import torch.nn.functional as F # Import kernels import vllm.kernels # noqa: F401 from vllm import envs, ir -from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant @@ -65,19 +64,12 @@ class RMSNorm(CustomOp): if self.has_weight: self.weight = nn.Parameter(self.weight) - # Do not pass identity weight to native implementation (causes issue on TPU). - # Other implementations require weight to be passed even if all ones. - # Cheat and predict if native will be dispatched to: - # 1) if native is first in priority list - # 2) if variance_size_override is given (only supported by native impl) - # TODO(luka): address weight passing inconsistency: - # https://github.com/vllm-project/vllm/issues/39370 - priority = get_current_vllm_config().kernel_config.ir_op_priority - var_override = self.variance_size_override is not None - native_rms_norm = priority.rms_norm[0] == "native" or var_override - native_add_rms_norm = priority.fused_add_rms_norm[0] == "native" or var_override - self.pass_weight = self.has_weight or not native_rms_norm - self.pass_weight_add = self.has_weight or not native_add_rms_norm + # When has_weight=False, pass weight=None so implementations that + # support a weightless path can skip the per-channel multiply. + # Implementations that require weight (e.g. oink) fall back via IR + # op priority when weight=None is unsupported. + self.pass_weight = self.has_weight + self.pass_weight_add = self.has_weight def forward_native( self,