A high-throughput and memory-efficient inference and serving engine for LLMs
[Kernel] Add weightless RMSNorm CUDA kernels for has_weight=False (#41430) (#44109)
Signed-off-by: hello-args <args.sarkar@gmail.com>
arghyadeep sarkar committed
Jun 17, 2026 at 12:15 UTC
93bbe94d3a3a49bc5a845edfa8ac75ada54614e1
11 files changed
+231
-117
csrc/cpu/layernorm.cpp
+41
-22
@@ -4,8 +4,9 @@ namespace {
4
template <typename scalar_t>
5
void rms_norm_impl(scalar_t* __restrict__ out,
6
const scalar_t* __restrict__ input,
7
- const scalar_t* __restrict__ weight, const float epsilon,
8
- const int num_tokens, const int hidden_size) {
7
+ const scalar_t* __restrict__ weight, const bool has_weight,
8
+ const float epsilon, const int num_tokens,
9
+ const int hidden_size) {
10
using scalar_vec_t = vec_op::vec_t<scalar_t>;
11
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
12
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -27,12 +28,15 @@ void rms_norm_impl(scalar_t* __restrict__ out,
28
29
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
30
scalar_vec_t x(input_p + j);
30
- scalar_vec_t w(weight + j);
31
-
31
vec_op::FP32Vec8 fp32_x(x);
33
- vec_op::FP32Vec8 fp32_w(w);
34
-
35
- vec_op::FP32Vec8 fp32_out = fp32_x * fp32_s_variance * fp32_w;
32
+ vec_op::FP32Vec8 fp32_out;
33
+ if (has_weight) {
34
+ scalar_vec_t w(weight + j);
35
+ vec_op::FP32Vec8 fp32_w(w);
36
+ fp32_out = fp32_x * fp32_s_variance * fp32_w;
37
+ } else {
38
+ fp32_out = fp32_x * fp32_s_variance;
39
+ }
40
41
scalar_vec_t out(fp32_out);
42
out.save(output_p + j);
@@ -44,8 +48,8 @@ template <typename scalar_t>
48
void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
49
scalar_t* __restrict__ residual,
50
const scalar_t* __restrict__ weight,
47
- const float epsilon, const int num_tokens,
48
- const int hidden_size) {
51
+ const bool has_weight, const float epsilon,
52
+ const int num_tokens, const int hidden_size) {
53
using scalar_vec_t = vec_op::vec_t<scalar_t>;
54
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
55
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -72,13 +76,18 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
76
vec_op::FP32Vec8 fp32_s_variance(s_variance);
77
78
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
75
- scalar_vec_t w(weight + j);
76
- scalar_vec_t res(residual_p + j);
77
-
78
- vec_op::FP32Vec8 fp32_w(w);
79
- vec_op::FP32Vec8 fp32_res(res);
80
-
81
- vec_op::FP32Vec8 fp32_out = fp32_res * fp32_s_variance * fp32_w;
79
+ vec_op::FP32Vec8 fp32_out;
80
+ if (has_weight) {
81
+ scalar_vec_t w(weight + j);
82
+ scalar_vec_t res(residual_p + j);
83
+ vec_op::FP32Vec8 fp32_w(w);
84
+ vec_op::FP32Vec8 fp32_res(res);
85
+ fp32_out = fp32_res * fp32_s_variance * fp32_w;
86
+ } else {
87
+ scalar_vec_t res(residual_p + j);
88
+ vec_op::FP32Vec8 fp32_res(res);
89
+ fp32_out = fp32_res * fp32_s_variance;
90
+ }
91
92
scalar_vec_t out(fp32_out);
93
out.save(input_p + j);
@@ -87,31 +96,41 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
96
}
97
} // namespace
98
90
-void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
91
- double epsilon) {
99
+void rms_norm(torch::Tensor& out, torch::Tensor& input,
100
+ std::optional<torch::Tensor> weight, double epsilon) {
101
int hidden_size = input.size(-1);
102
int num_tokens = input.numel() / hidden_size;
103
+ const bool has_weight = weight.has_value();
104
+ if (has_weight) {
105
+ TORCH_CHECK(weight->is_contiguous());
106
+ }
107
108
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] {
109
CPU_KERNEL_GUARD_IN(rms_norm_impl)
110
rms_norm_impl(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
98
- weight.data_ptr<scalar_t>(), epsilon, num_tokens,
99
- hidden_size);
111
+ has_weight ? weight->data_ptr<scalar_t>() : nullptr,
112
+ has_weight, epsilon, num_tokens, hidden_size);
113
CPU_KERNEL_GUARD_OUT(rms_norm_impl)
114
});
115
}
116
117
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
105
- torch::Tensor& weight, double epsilon) {
118
+ std::optional<torch::Tensor> weight, double epsilon) {
119
int hidden_size = input.size(-1);
120
int num_tokens = input.numel() / hidden_size;
121
+ const bool has_weight = weight.has_value();
122
+ if (has_weight) {
123
+ TORCH_CHECK(weight->scalar_type() == input.scalar_type());
124
+ TORCH_CHECK(weight->is_contiguous());
125
+ }
126
127
VLLM_DISPATCH_FLOATING_TYPES(
128
input.scalar_type(), "fused_add_rms_norm_impl", [&] {
129
CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl)
130
fused_add_rms_norm_impl(
131
input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(),
114
- weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
132
+ has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
133
+ epsilon, num_tokens, hidden_size);
134
CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
135
});
136
}
csrc/cpu/torch_bindings.cpp
+2
-2
@@ -310,13 +310,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
310
// Layernorm
311
// Apply Root Mean Square (RMS) Normalization to the input tensor.
312
ops.def(
313
- "rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> "
313
+ "rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> "
314
"()");
315
ops.impl("rms_norm", torch::kCPU, &rms_norm);
316
317
// In-place fused Add and RMS Normalization.
318
ops.def(
319
- "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
319
+ "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
320
"float epsilon) -> ()");
321
ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm);
322
csrc/libtorch_stable/layernorm_kernels.cu
+104
-56
@@ -11,7 +11,7 @@
11
namespace vllm {
12
13
// TODO(woosuk): Further optimize this kernel.
14
-template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
14
+template <typename scalar_t, int VEC_SIZE, int NUM_DIMS, bool HasWeight>
15
__global__ void rms_norm_kernel(
16
scalar_t* __restrict__ out, // [..., hidden_size]
17
const scalar_t* __restrict__ input, // [..., hidden_size]
@@ -20,7 +20,7 @@ __global__ void rms_norm_kernel(
20
const int64_t input_stride_d4, // input.stride(-4)
21
const int64_t input_shape_d2, // input.size(-2)
22
const int64_t input_shape_d3, // input.size(-3)
23
- const scalar_t* __restrict__ weight, // [hidden_size]
23
+ const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight
24
const float epsilon, const int num_tokens, const int hidden_size) {
25
__shared__ float s_variance;
26
float variance = 0.0f;
@@ -74,11 +74,19 @@ __global__ void rms_norm_kernel(
74
for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) {
75
vec_n_t<scalar_t, VEC_SIZE> dst;
76
vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[i];
77
- vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[i];
77
+ vec_n_t<scalar_t, VEC_SIZE> src2;
78
+ if constexpr (HasWeight) {
79
+ src2 = v_w[i];
80
+ }
81
#pragma unroll
82
for (int j = 0; j < VEC_SIZE; j++) {
83
float x = static_cast<float>(src1.val[j]);
81
- dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
84
+ scalar_t normalized = static_cast<scalar_t>(x * s_variance);
85
+ if constexpr (HasWeight) {
86
+ dst.val[j] = normalized * src2.val[j];
87
+ } else {
88
+ dst.val[j] = normalized;
89
+ }
90
}
91
v_out[i] = dst;
92
}
@@ -88,13 +96,13 @@ __global__ void rms_norm_kernel(
96
Additional optimizations we can make in this case are
97
packed and vectorized operations, which help with the
98
memory latency bottleneck. */
91
-template <typename scalar_t, int width>
99
+template <typename scalar_t, int width, bool HasWeight>
100
__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
101
fused_add_rms_norm_kernel(
102
scalar_t* __restrict__ input, // [..., hidden_size]
103
const int64_t input_stride,
104
scalar_t* __restrict__ residual, // [..., hidden_size]
97
- const scalar_t* __restrict__ weight, // [hidden_size]
105
+ const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight
106
const float epsilon, const int num_tokens, const int hidden_size) {
107
// Sanity checks on our vector struct and type-punned pointer arithmetic
108
static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
@@ -136,13 +144,21 @@ fused_add_rms_norm_kernel(
144
int id = blockIdx.x * vec_hidden_size + idx;
145
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
146
_f16Vec<scalar_t, width> res = residual_v[id];
139
- _f16Vec<scalar_t, width> w = weight_v[idx];
147
_f16Vec<scalar_t, width> out;
148
using Converter = _typeConvert<scalar_t>;
149
+ if constexpr (HasWeight) {
150
+ _f16Vec<scalar_t, width> w = weight_v[idx];
151
+#pragma unroll
152
+ for (int j = 0; j < width; ++j) {
153
+ float x = Converter::convert(res.data[j]);
154
+ out.data[j] = Converter::convert(x * s_variance) * w.data[j];
155
+ }
156
+ } else {
157
#pragma unroll
143
- for (int j = 0; j < width; ++j) {
144
- float x = Converter::convert(res.data[j]);
145
- out.data[j] = Converter::convert(x * s_variance) * w.data[j];
158
+ for (int j = 0; j < width; ++j) {
159
+ float x = Converter::convert(res.data[j]);
160
+ out.data[j] = Converter::convert(x * s_variance);
161
+ }
162
}
163
input_v[strided_id] = out;
164
}
@@ -151,13 +167,13 @@ fused_add_rms_norm_kernel(
167
/* Generic fused_add_rms_norm_kernel
168
The width field is not used here but necessary for other specializations.
169
*/
154
-template <typename scalar_t, int width>
170
+template <typename scalar_t, int width, bool HasWeight>
171
__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
172
fused_add_rms_norm_kernel(
173
scalar_t* __restrict__ input, // [..., hidden_size]
174
const int64_t input_stride,
175
scalar_t* __restrict__ residual, // [..., hidden_size]
160
- const scalar_t* __restrict__ weight, // [hidden_size]
176
+ const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight
177
const float epsilon, const int num_tokens, const int hidden_size) {
178
__shared__ float s_variance;
179
float variance = 0.0f;
@@ -181,23 +197,29 @@ fused_add_rms_norm_kernel(
197
198
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
199
float x = (float)residual[blockIdx.x * hidden_size + idx];
184
- input[blockIdx.x * input_stride + idx] =
185
- (scalar_t)(x * s_variance) * weight[idx];
200
+ if constexpr (HasWeight) {
201
+ input[blockIdx.x * input_stride + idx] =
202
+ (scalar_t)(x * s_variance) * weight[idx];
203
+ } else {
204
+ input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
205
+ }
206
}
207
}
208
209
} // namespace vllm
210
191
-void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
192
- torch::stable::Tensor& input, // [..., hidden_size]
193
- torch::stable::Tensor& weight, // [hidden_size]
211
+void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
212
+ torch::stable::Tensor& input, // [..., hidden_size]
213
+ std::optional<torch::stable::Tensor> weight, // [hidden_size]
214
double epsilon) {
215
STD_TORCH_CHECK(out.is_contiguous());
216
if (input.stride(-1) != 1) {
217
input = torch::stable::contiguous(input);
218
}
219
STD_TORCH_CHECK(input.stride(-1) == 1);
200
- STD_TORCH_CHECK(weight.is_contiguous());
220
+ if (weight.has_value()) {
221
+ STD_TORCH_CHECK(weight->is_contiguous());
222
+ }
223
224
int hidden_size = input.size(-1);
225
@@ -215,46 +237,69 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
237
const torch::stable::accelerator::DeviceGuard device_guard(
238
input.get_device_index());
239
const cudaStream_t stream = get_current_cuda_stream();
240
+ const bool has_weight = weight.has_value();
241
VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] {
242
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
243
input.scalar_type(), "rms_norm_kernel", [&] {
244
+ const scalar_t* weight_ptr =
245
+ has_weight ? weight->const_data_ptr<scalar_t>() : nullptr;
246
const int calculated_vec_size =
247
std::gcd(16 / sizeof(scalar_t), hidden_size);
248
const int block_size =
249
std::min(hidden_size / calculated_vec_size, max_block_size);
250
dim3 block(block_size);
251
VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
227
- vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
228
- <<<grid, block, 0, stream>>>(
229
- out.mutable_data_ptr<scalar_t>(),
230
- input.const_data_ptr<scalar_t>(), input_stride_d2,
231
- input_stride_d3, input_stride_d4, input_shape_d2,
232
- input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
233
- num_tokens, hidden_size);
252
+ if (has_weight) {
253
+ vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, true>
254
+ <<<grid, block, 0, stream>>>(
255
+ out.mutable_data_ptr<scalar_t>(),
256
+ input.const_data_ptr<scalar_t>(), input_stride_d2,
257
+ input_stride_d3, input_stride_d4, input_shape_d2,
258
+ input_shape_d3, weight_ptr, epsilon, num_tokens,
259
+ hidden_size);
260
+ } else {
261
+ vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, false>
262
+ <<<grid, block, 0, stream>>>(
263
+ out.mutable_data_ptr<scalar_t>(),
264
+ input.const_data_ptr<scalar_t>(), input_stride_d2,
265
+ input_stride_d3, input_stride_d4, input_shape_d2,
266
+ input_shape_d3, weight_ptr, epsilon, num_tokens,
267
+ hidden_size);
268
+ }
269
});
270
});
271
});
272
}
273
239
-#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
240
- VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
241
- input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
242
- vllm::fused_add_rms_norm_kernel<scalar_t, width> \
243
- <<<grid, block, 0, stream>>>( \
244
- input.mutable_data_ptr<scalar_t>(), input_stride, \
245
- residual.mutable_data_ptr<scalar_t>(), \
246
- weight.const_data_ptr<scalar_t>(), epsilon, num_tokens, \
247
- hidden_size); \
274
+#define LAUNCH_FUSED_ADD_RMS_NORM(width, has_weight) \
275
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
276
+ input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
277
+ if (has_weight) { \
278
+ vllm::fused_add_rms_norm_kernel<scalar_t, width, true> \
279
+ <<<grid, block, 0, stream>>>( \
280
+ input.mutable_data_ptr<scalar_t>(), input_stride, \
281
+ residual.mutable_data_ptr<scalar_t>(), \
282
+ weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \
283
+ hidden_size); \
284
+ } else { \
285
+ vllm::fused_add_rms_norm_kernel<scalar_t, width, false> \
286
+ <<<grid, block, 0, stream>>>( \
287
+ input.mutable_data_ptr<scalar_t>(), input_stride, \
288
+ residual.mutable_data_ptr<scalar_t>(), nullptr, epsilon, \
289
+ num_tokens, hidden_size); \
290
+ } \
291
});
292
293
void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
294
torch::stable::Tensor& residual, // [..., hidden_size]
252
- torch::stable::Tensor& weight, // [hidden_size]
295
+ std::optional<torch::stable::Tensor> weight,
296
double epsilon) {
254
- STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
297
STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());
298
STD_TORCH_CHECK(residual.is_contiguous());
257
- STD_TORCH_CHECK(weight.is_contiguous());
299
+ if (weight.has_value()) {
300
+ STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type());
301
+ STD_TORCH_CHECK(weight->is_contiguous());
302
+ }
303
int hidden_size = input.size(-1);
304
int64_t input_stride = input.stride(-2);
305
int num_tokens = input.numel() / hidden_size;
@@ -269,30 +314,33 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
314
const torch::stable::accelerator::DeviceGuard device_guard(
315
input.get_device_index());
316
const cudaStream_t stream = get_current_cuda_stream();
272
- /*If the tensor types are FP16/BF16, try to use the optimized kernel
273
- with packed + vectorized ops.
274
- Max optimization is achieved with a width-8 vector of FP16/BF16s
275
- since we can load at most 128 bits at once in a global memory op.
276
- However, this requires each tensor's data to be aligned to 16
277
- bytes.
278
- */
317
+ constexpr int vector_width = 8;
318
+ constexpr int req_alignment_bytes = vector_width * 2;
319
auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
320
auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
281
- auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
282
- constexpr int vector_width = 8;
283
- constexpr int req_alignment_bytes =
284
- vector_width * 2; // vector_width * sizeof(bfloat16 or float16) (float32
285
- // falls back to non-vectorized version anyway)
286
- bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
287
- res_ptr % req_alignment_bytes == 0 &&
288
- wt_ptr % req_alignment_bytes == 0;
321
bool offsets_are_multiple_of_vector_width =
322
hidden_size % vector_width == 0 && input_stride % vector_width == 0;
323
bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
292
- if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
293
- !batch_invariant_launch) {
294
- LAUNCH_FUSED_ADD_RMS_NORM(8);
324
+ const bool has_weight = weight.has_value();
325
+ if (has_weight) {
326
+ auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());
327
+ bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
328
+ res_ptr % req_alignment_bytes == 0 &&
329
+ wt_ptr % req_alignment_bytes == 0;
330
+ if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
331
+ !batch_invariant_launch) {
332
+ LAUNCH_FUSED_ADD_RMS_NORM(8, true);
333
+ } else {
334
+ LAUNCH_FUSED_ADD_RMS_NORM(0, true);
335
+ }
336
} else {
296
- LAUNCH_FUSED_ADD_RMS_NORM(0);
337
+ bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
338
+ res_ptr % req_alignment_bytes == 0;
339
+ if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
340
+ !batch_invariant_launch) {
341
+ LAUNCH_FUSED_ADD_RMS_NORM(8, false);
342
+ } else {
343
+ LAUNCH_FUSED_ADD_RMS_NORM(0, false);
344
+ }
345
}
346
}
csrc/libtorch_stable/ops.h
+3
-2
@@ -188,11 +188,12 @@ torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
188
189
// Layernorm kernels (shared CUDA/ROCm)
190
void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
191
- torch::stable::Tensor& weight, double epsilon);
191
+ std::optional<torch::stable::Tensor> weight, double epsilon);
192
193
void fused_add_rms_norm(torch::stable::Tensor& input,
194
torch::stable::Tensor& residual,
195
- torch::stable::Tensor& weight, double epsilon);
195
+ std::optional<torch::stable::Tensor> weight,
196
+ double epsilon);
197
198
// Layernorm-quant kernels (shared CUDA/ROCm)
199
void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
csrc/libtorch_stable/torch_bindings.cpp
+3
-2
@@ -369,12 +369,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
369
370
// Apply Root Mean Square (RMS) Normalization to the input tensor.
371
ops.def(
372
- "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
372
+ "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) "
373
+ "-> "
374
"()");
375
376
// In-place fused Add and RMS Normalization.
377
ops.def(
377
- "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
378
+ "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
379
"float epsilon) -> ()");
380
381
// Layernorm-quant
csrc/ops.h
+3
-3
@@ -34,11 +34,11 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
34
// rms_norm and fused_add_rms_norm declarations also exist in
35
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
36
// because the CPU build still uses these torch::Tensor declarations.
37
-void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
38
- double epsilon);
37
+void rms_norm(torch::Tensor& out, torch::Tensor& input,
38
+ std::optional<torch::Tensor> weight, double epsilon);
39
40
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
41
- torch::Tensor& weight, double epsilon);
41
+ std::optional<torch::Tensor> weight, double epsilon);
42
43
// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
44
// ABI for CUDA). It remains here because the CPU build still uses these
tests/kernels/core/test_layernorm.py
+48
@@ -6,6 +6,7 @@ import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
from tests.kernels.utils import opcheck
9
+from vllm import ir
10
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
11
from vllm.platforms import current_platform
12
from vllm.utils.torch_utils import set_random_seed
@@ -27,6 +28,10 @@ CUDA_DEVICES = [
28
]
29
30
31
+def _rms_norm_tolerance(dtype: torch.dtype) -> dict[str, float]:
32
+ return ir.ops.rms_norm.get_tolerance(dtype)
33
+
34
+
35
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
36
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
37
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
@@ -81,6 +86,49 @@ def test_rms_norm(
86
)
87
88
89
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
90
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
91
+@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
92
+@pytest.mark.parametrize("dtype", DTYPES)
93
+@pytest.mark.parametrize("seed", SEEDS)
94
+@pytest.mark.parametrize("device", CUDA_DEVICES)
95
+@torch.inference_mode()
96
+def test_rms_norm_weightless(
97
+ default_vllm_config,
98
+ num_tokens: int,
99
+ hidden_size: int,
100
+ add_residual: bool,
101
+ dtype: torch.dtype,
102
+ seed: int,
103
+ device: str,
104
+) -> None:
105
+ set_random_seed(seed)
106
+ torch.set_default_device(device)
107
+ layer = RMSNorm(hidden_size, has_weight=False).to(dtype=dtype)
108
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype)
109
+ residual = torch.randn_like(x) if add_residual else None
110
+
111
+ ref_out = layer.forward_native(x, residual)
112
+ out = layer(x, residual)
113
+ tol = _rms_norm_tolerance(dtype)
114
+ if add_residual:
115
+ torch.testing.assert_close(out[0], ref_out[0], **tol)
116
+ torch.testing.assert_close(out[1], ref_out[1], **tol)
117
+ else:
118
+ torch.testing.assert_close(out, ref_out, **tol)
119
+
120
+ if residual is not None:
121
+ opcheck(
122
+ torch.ops._C.fused_add_rms_norm,
123
+ (x, residual, None, layer.variance_epsilon),
124
+ )
125
+ else:
126
+ opcheck(
127
+ torch.ops._C.rms_norm,
128
+ (out, x, None, layer.variance_epsilon),
129
+ )
130
+
131
+
132
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
133
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
134
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
vllm/_custom_ops.py
+8
-2
@@ -315,13 +315,19 @@ def rotary_embedding(
315
316
# layer norm ops
317
def rms_norm(
318
- out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float
318
+ out: torch.Tensor,
319
+ input: torch.Tensor,
320
+ weight: torch.Tensor | None,
321
+ epsilon: float,
322
) -> None:
323
torch.ops._C.rms_norm(out, input, weight, epsilon)
324
325
326
def fused_add_rms_norm(
324
- input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, epsilon: float
327
+ input: torch.Tensor,
328
+ residual: torch.Tensor,
329
+ weight: torch.Tensor | None,
330
+ epsilon: float,
331
) -> None:
332
# Note: this func is batch invariant
333
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
vllm/kernels/vllm_c.py
-7
@@ -25,9 +25,6 @@ rms_no_var_size = lambda x, weight, epsilon, variance_size=None: (
25
def rms_norm(
26
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
27
) -> Tensor:
28
- if weight is None:
29
- # Kernel requires weight tensor, pass ones
30
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
28
assert variance_size is None
29
# ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors.
30
# Higher-rank callers still normalize over the last dimension, so flatten
@@ -64,10 +61,6 @@ def fused_add_rms_norm(
61
epsilon: float,
62
variance_size: int | None = None,
63
) -> tuple[Tensor, Tensor]:
67
- if weight is None:
68
- # Kernel requires weight tensor, pass ones
69
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
70
-
64
assert variance_size is None
65
if IS_ROCM and (not x.is_contiguous() or not x_residual.is_contiguous()):
66
output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
vllm/kernels/xpu_ops.py
+13
-7
@@ -29,10 +29,12 @@ rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is Non
29
def rms_norm(
30
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
31
) -> Tensor:
32
- if weight is None:
33
- # Kernel requires weight tensor, pass ones
34
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
32
assert variance_size is None
33
+ if weight is None:
34
+ # Weightless _C ops are CUDA-only; native skips the multiply on XPU.
35
+ return ir.ops.rms_norm.impls["native"].impl_fn(
36
+ x, weight, epsilon, variance_size
37
+ )
38
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
39
torch.ops._C.rms_norm(output, x, weight, epsilon)
40
return output
@@ -57,10 +59,14 @@ def fused_add_rms_norm(
59
epsilon: float,
60
variance_size: int | None = None,
61
) -> tuple[Tensor, Tensor]:
60
- if weight is None:
61
- # Kernel requires weight tensor, pass ones
62
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
63
-
62
assert variance_size is None
63
+ if weight is None:
64
+ # Weightless _C ops are CUDA-only; native skips the multiply on XPU.
65
+ output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
66
+ x, x_residual, weight, epsilon, variance_size
67
+ )
68
+ x.copy_(output)
69
+ x_residual.copy_(residual)
70
+ return x, x_residual
71
torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
72
return x, x_residual
vllm/model_executor/layers/layernorm.py
+6
-14
@@ -9,7 +9,6 @@ import torch.nn.functional as F
9
# Import kernels
10
import vllm.kernels # noqa: F401
11
from vllm import envs, ir
12
-from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.custom_op import CustomOp
14
from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant
@@ -65,19 +64,12 @@ class RMSNorm(CustomOp):
64
if self.has_weight:
65
self.weight = nn.Parameter(self.weight)
66
68
- # Do not pass identity weight to native implementation (causes issue on TPU).
69
- # Other implementations require weight to be passed even if all ones.
70
- # Cheat and predict if native will be dispatched to:
71
- # 1) if native is first in priority list
72
- # 2) if variance_size_override is given (only supported by native impl)
73
- # TODO(luka): address weight passing inconsistency:
74
- # https://github.com/vllm-project/vllm/issues/39370
75
- priority = get_current_vllm_config().kernel_config.ir_op_priority
76
- var_override = self.variance_size_override is not None
77
- native_rms_norm = priority.rms_norm[0] == "native" or var_override
78
- native_add_rms_norm = priority.fused_add_rms_norm[0] == "native" or var_override
79
- self.pass_weight = self.has_weight or not native_rms_norm
80
- self.pass_weight_add = self.has_weight or not native_add_rms_norm
67
+ # When has_weight=False, pass weight=None so implementations that
68
+ # support a weightless path can skip the per-channel multiply.
69
+ # Implementations that require weight (e.g. oink) fall back via IR
70
+ # op priority when weight=None is unsupported.
71
+ self.pass_weight = self.has_weight
72
+ self.pass_weight_add = self.has_weight
73
74
def forward_native(
75
self,