iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel] Add weightless RMSNorm CUDA kernels for has_weight=False (#41430) (#44109)

Browse Source 
Signed-off-by: hello-args <args.sarkar@gmail.com>
This commit is contained in:
arghyadeep sarkar
2026-06-16 23:45:55 -07:00
committed by GitHub
parent 17bc144556
commit 93bbe94d3a
11 changed files with 231 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/layernorm.cpp
+41 -22
View File
@@ -4,8 +4,9 @@ namespace {
template <typename scalar_t>
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<scalar_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 <typename scalar_t>
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<scalar_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<torch::Tensor> 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<scalar_t>(), input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(), epsilon, num_tokens,
hidden_size);
has_weight ? weight->data_ptr<scalar_t>() : 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<torch::Tensor> 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<scalar_t>(), residual.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
epsilon, num_tokens, hidden_size);
CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
});
}
csrc/cpu/torch_bindings.cpp
+2 -2
View File
@@ -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);
csrc/libtorch_stable/layernorm_kernels.cu
+104 -56
View File
@@ -11,7 +11,7 @@
namespace vllm {
// TODO(woosuk): Further optimize this kernel.
template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
template <typename scalar_t, int VEC_SIZE, int NUM_DIMS, bool HasWeight>
__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<scalar_t, VEC_SIZE> dst;
vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[i];
vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[i];
vec_n_t<scalar_t, VEC_SIZE> src2;
if constexpr (HasWeight) {
src2 = v_w[i];
}
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
scalar_t normalized = static_cast<scalar_t>(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 <typename scalar_t, int width>
template <typename scalar_t, int width, bool HasWeight>
__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::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<scalar_t, width>>);
@@ -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<scalar_t, width> res = residual_v[id];
_f16Vec<scalar_t, width> w = weight_v[idx];
_f16Vec<scalar_t, width> out;
using Converter = _typeConvert<scalar_t>;
if constexpr (HasWeight) {
_f16Vec<scalar_t, width> 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 <typename scalar_t, int width>
template <typename scalar_t, int width, bool HasWeight>
__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::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<torch::stable::Tensor> 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<scalar_t>() : 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<scalar_t, vec_size, tensor_rank>
<<<grid, block, 0, stream>>>(
out.mutable_data_ptr<scalar_t>(),
input.const_data_ptr<scalar_t>(), input_stride_d2,
input_stride_d3, input_stride_d4, input_shape_d2,
input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
num_tokens, hidden_size);
if (has_weight) {
vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, true>
<<<grid, block, 0, stream>>>(
out.mutable_data_ptr<scalar_t>(),
input.const_data_ptr<scalar_t>(), 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<scalar_t, vec_size, tensor_rank, false>
<<<grid, block, 0, stream>>>(
out.mutable_data_ptr<scalar_t>(),
input.const_data_ptr<scalar_t>(), 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<scalar_t, width> \
<<<grid, block, 0, stream>>>( \
input.mutable_data_ptr<scalar_t>(), input_stride, \
residual.mutable_data_ptr<scalar_t>(), \
weight.const_data_ptr<scalar_t>(), 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<scalar_t, width, true> \
<<<grid, block, 0, stream>>>( \
input.mutable_data_ptr<scalar_t>(), input_stride, \
residual.mutable_data_ptr<scalar_t>(), \
weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \
hidden_size); \
} else { \
vllm::fused_add_rms_norm_kernel<scalar_t, width, false> \
<<<grid, block, 0, stream>>>( \
input.mutable_data_ptr<scalar_t>(), input_stride, \
residual.mutable_data_ptr<scalar_t>(), 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<torch::stable::Tensor> 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<std::uintptr_t>(input.data_ptr());
auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
auto wt_ptr = reinterpret_cast<std::uintptr_t>(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<std::uintptr_t>(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);
}
}
}
csrc/libtorch_stable/ops.h
+3 -2
View File
@@ -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<torch::stable::Tensor> weight, double epsilon);
void fused_add_rms_norm(torch::stable::Tensor& input,
torch::stable::Tensor& residual,
torch::stable::Tensor& weight, double epsilon);
std::optional<torch::stable::Tensor> weight,
double epsilon);
// Layernorm-quant kernels (shared CUDA/ROCm)
void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
csrc/libtorch_stable/torch_bindings.cpp
+3 -2
View File
@@ -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
csrc/ops.h
+3 -3
View File
@@ -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<torch::Tensor> weight, double epsilon);
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
torch::Tensor& weight, double epsilon);
std::optional<torch::Tensor> 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
tests/kernels/core/test_layernorm.py
+48
View File
@@ -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)
vllm/_custom_ops.py
+8 -2
View File
@@ -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)
vllm/kernels/vllm_c.py
-7
View File
@@ -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(
vllm/kernels/xpu_ops.py
+13 -7
View File
@@ -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
vllm/model_executor/layers/layernorm.py
+6 -14
View File
@@ -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,