A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 12:00 UTC
[Kernel] Optimize rms_norm kernel (#27931)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Nov 11, 2025 at 10:02 UTC
6c3c0f8235cacce28982687e362b80d953ea7617
3 files changed
+86
-25
csrc/dispatch_utils.h
+29
@@ -88,3 +88,32 @@
88
#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
89
AT_DISPATCH_SWITCH( \
90
TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
91
+
92
+#define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \
93
+ switch (VEC_SIZE) { \
94
+ case 16: { \
95
+ constexpr int vec_size = 16; \
96
+ __VA_ARGS__(); \
97
+ break; \
98
+ } \
99
+ case 8: { \
100
+ constexpr int vec_size = 8; \
101
+ __VA_ARGS__(); \
102
+ break; \
103
+ } \
104
+ case 4: { \
105
+ constexpr int vec_size = 4; \
106
+ __VA_ARGS__(); \
107
+ break; \
108
+ } \
109
+ case 2: { \
110
+ constexpr int vec_size = 2; \
111
+ __VA_ARGS__(); \
112
+ break; \
113
+ } \
114
+ default: { \
115
+ constexpr int vec_size = 1; \
116
+ __VA_ARGS__(); \
117
+ break; \
118
+ } \
119
+ }
csrc/layernorm_kernels.cu
+28
-11
@@ -10,7 +10,7 @@
10
namespace vllm {
11
12
// TODO(woosuk): Further optimize this kernel.
13
-template <typename scalar_t>
13
+template <typename scalar_t, int VEC_SIZE>
14
__global__ void rms_norm_kernel(
15
scalar_t* __restrict__ out, // [..., hidden_size]
16
const scalar_t* __restrict__ input, // [..., hidden_size]
@@ -21,7 +21,6 @@ __global__ void rms_norm_kernel(
21
float variance = 0.0f;
22
const scalar_t* input_row = input + blockIdx.x * input_stride;
23
24
- constexpr int VEC_SIZE = 8;
24
auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
25
#pragma unroll
26
for (int i = 0; i < VEC_SIZE; ++i) {
@@ -45,10 +44,20 @@ __global__ void rms_norm_kernel(
44
}
45
__syncthreads();
46
48
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
49
- float x = (float)input[blockIdx.x * input_stride + idx];
50
- out[blockIdx.x * hidden_size + idx] =
51
- ((scalar_t)(x * s_variance)) * weight[idx];
47
+ scalar_t* out_row = out + blockIdx.x * hidden_size;
48
+ auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
49
+ auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
50
+ auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row);
51
+ for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) {
52
+ vec_n_t<scalar_t, VEC_SIZE> dst;
53
+ vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[i];
54
+ vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[i];
55
+#pragma unroll
56
+ for (int j = 0; j < VEC_SIZE; j++) {
57
+ float x = static_cast<float>(src1.val[j]);
58
+ dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
59
+ }
60
+ v_out[i] = dst;
61
}
62
}
63
@@ -168,16 +177,24 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]
177
int num_tokens = input_view.numel() / hidden_size;
178
int64_t input_stride = input_view.stride(-2);
179
180
+ // For large num_tokens, use smaller blocks to increase SM concurrency.
181
+ const int max_block_size = (num_tokens < 256) ? 1024 : 256;
182
dim3 grid(num_tokens);
172
- dim3 block(std::min(hidden_size, 1024));
183
const at::cuda::OptionalCUDAGuard device_guard(device_of(input_view));
184
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
185
VLLM_DISPATCH_FLOATING_TYPES(
186
input_view.scalar_type(), "rms_norm_kernel", [&] {
177
- vllm::rms_norm_kernel<scalar_t><<<grid, block, 0, stream>>>(
178
- out.data_ptr<scalar_t>(), input_view.data_ptr<scalar_t>(),
179
- input_stride, weight.data_ptr<scalar_t>(), epsilon, num_tokens,
180
- hidden_size);
187
+ const int calculated_vec_size =
188
+ std::gcd(16 / sizeof(scalar_t), hidden_size);
189
+ const int block_size =
190
+ std::min(hidden_size / calculated_vec_size, max_block_size);
191
+ dim3 block(block_size);
192
+ VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
193
+ vllm::rms_norm_kernel<scalar_t, vec_size><<<grid, block, 0, stream>>>(
194
+ out.data_ptr<scalar_t>(), input_view.data_ptr<scalar_t>(),
195
+ input_stride, weight.data_ptr<scalar_t>(), epsilon, num_tokens,
196
+ hidden_size);
197
+ });
198
});
199
}
200
csrc/layernorm_quant_kernels.cu
+29
-14
@@ -18,7 +18,7 @@
18
namespace vllm {
19
20
// TODO(woosuk): Further optimize this kernel.
21
-template <typename scalar_t, typename fp8_type>
21
+template <typename scalar_t, typename fp8_type, int VEC_SIZE>
22
__global__ void rms_norm_static_fp8_quant_kernel(
23
fp8_type* __restrict__ out, // [..., hidden_size]
24
const scalar_t* __restrict__ input, // [..., hidden_size]
@@ -31,7 +31,6 @@ __global__ void rms_norm_static_fp8_quant_kernel(
31
32
const scalar_t* input_row = input + blockIdx.x * input_stride;
33
34
- constexpr int VEC_SIZE = 8;
34
auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
35
#pragma unroll
36
for (int i = 0; i < VEC_SIZE; ++i) {
@@ -58,11 +57,18 @@ __global__ void rms_norm_static_fp8_quant_kernel(
57
// invert scale to avoid division
58
float const scale_inv = 1.0f / *scale;
59
61
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
62
- float x = (float)input[blockIdx.x * input_stride + idx];
63
- float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
64
- out[blockIdx.x * hidden_size + idx] =
65
- scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
60
+ auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
61
+ auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
62
+ for (int idx = threadIdx.x; idx < hidden_size / VEC_SIZE; idx += blockDim.x) {
63
+ vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[idx];
64
+ vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[idx];
65
+#pragma unroll
66
+ for (int j = 0; j < VEC_SIZE; j++) {
67
+ float x = static_cast<float>(src1.val[j]);
68
+ float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
69
+ out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
70
+ scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
71
+ }
72
}
73
}
74
@@ -188,20 +194,29 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
194
int input_stride = input.stride(-2);
195
int num_tokens = input.numel() / hidden_size;
196
197
+ // For large num_tokens, use smaller blocks to increase SM concurrency.
198
+ const int max_block_size = (num_tokens < 256) ? 1024 : 256;
199
dim3 grid(num_tokens);
192
- dim3 block(std::min(hidden_size, 1024));
200
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
201
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
202
VLLM_DISPATCH_FLOATING_TYPES(
203
input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
204
VLLM_DISPATCH_FP8_TYPES(
205
out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
199
- vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t>
200
- <<<grid, block, 0, stream>>>(
201
- out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
202
- input_stride, weight.data_ptr<scalar_t>(),
203
- scale.data_ptr<float>(), epsilon, num_tokens,
204
- hidden_size);
206
+ const int calculated_vec_size =
207
+ std::gcd(16 / sizeof(scalar_t), hidden_size);
208
+ const int block_size =
209
+ std::min(hidden_size / calculated_vec_size, max_block_size);
210
+ dim3 block(block_size);
211
+ VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
212
+ vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t,
213
+ vec_size>
214
+ <<<grid, block, 0, stream>>>(
215
+ out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
216
+ input_stride, weight.data_ptr<scalar_t>(),
217
+ scale.data_ptr<float>(), epsilon, num_tokens,
218
+ hidden_size);
219
+ });
220
});
221
});
222
}