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 12, 19:45 UTC
[kernel][perf] support uncontiguous input for rms_norm kernel (#28103)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Signed-off-by: izhuhaoran <izhuhaoran@qq.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
zhrrr committed
Nov 21, 2025 at 11:39 UTC
a982f5b5ea4a1932424927ea357b532d0e45caf1
4 files changed
+77
-33
csrc/dispatch_utils.h
+21
@@ -117,3 +117,24 @@
117
break; \
118
} \
119
}
120
+
121
+#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \
122
+ switch (NUM_DIMS) { \
123
+ case 2: { \
124
+ constexpr int tensor_rank = 2; \
125
+ __VA_ARGS__(); \
126
+ break; \
127
+ } \
128
+ case 3: { \
129
+ constexpr int tensor_rank = 3; \
130
+ __VA_ARGS__(); \
131
+ break; \
132
+ } \
133
+ case 4: { \
134
+ constexpr int tensor_rank = 4; \
135
+ __VA_ARGS__(); \
136
+ break; \
137
+ } \
138
+ default: \
139
+ TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \
140
+ }
csrc/layernorm_kernels.cu
+54
-26
@@ -10,16 +10,38 @@
10
namespace vllm {
11
12
// TODO(woosuk): Further optimize this kernel.
13
-template <typename scalar_t, int VEC_SIZE>
13
+template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
14
__global__ void rms_norm_kernel(
15
- scalar_t* __restrict__ out, // [..., hidden_size]
16
- const scalar_t* __restrict__ input, // [..., hidden_size]
17
- const int64_t input_stride,
15
+ scalar_t* __restrict__ out, // [..., hidden_size]
16
+ const scalar_t* __restrict__ input, // [..., hidden_size]
17
+ const int64_t input_stride_d2, // input.stride(-2)
18
+ const int64_t input_stride_d3, // input.stride(-3)
19
+ const int64_t input_stride_d4, // input.stride(-4)
20
+ const int64_t input_shape_d2, // input.size(-2)
21
+ const int64_t input_shape_d3, // input.size(-3)
22
const scalar_t* __restrict__ weight, // [hidden_size]
23
const float epsilon, const int num_tokens, const int hidden_size) {
24
__shared__ float s_variance;
25
float variance = 0.0f;
22
- const scalar_t* input_row = input + blockIdx.x * input_stride;
26
+ const scalar_t* input_row;
27
+ if constexpr (NUM_DIMS == 2) {
28
+ // 2D for layernorm normal case [batch_size, hidden]
29
+ input_row = input + blockIdx.x * input_stride_d2;
30
+ } else if constexpr (NUM_DIMS == 3) {
31
+ // 3D for q/k norm [batch_size, num_heads, head_size]
32
+ int batch_idx = blockIdx.x / input_shape_d2;
33
+ int head_idx = blockIdx.x % input_shape_d2;
34
+ input_row =
35
+ input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
36
+ } else if constexpr (NUM_DIMS == 4) {
37
+ // 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
38
+ int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2);
39
+ int remaining = blockIdx.x % (input_shape_d3 * input_shape_d2);
40
+ int seq_idx = remaining / input_shape_d2;
41
+ int head_idx = remaining % input_shape_d2;
42
+ input_row = input + batch_idx * input_stride_d4 +
43
+ seq_idx * input_stride_d3 + head_idx * input_stride_d2;
44
+ }
45
46
auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
47
#pragma unroll
@@ -164,38 +186,44 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]
186
torch::Tensor& weight, // [hidden_size]
187
double epsilon) {
188
TORCH_CHECK(out.is_contiguous());
189
+ if (input.stride(-1) != 1) {
190
+ input = input.contiguous();
191
+ }
192
TORCH_CHECK(input.stride(-1) == 1);
193
TORCH_CHECK(weight.is_contiguous());
194
195
int hidden_size = input.size(-1);
196
172
- // We cannot just use `input.stride(-2)` if the tensor is not row-major.
173
- // Instead, we use a 2d view to get the second-innermost stride.
174
- // That way the dimensions (except the last one) can be arbitrarily permuted.
175
- torch::Tensor input_view = input.view({-1, hidden_size});
176
-
177
- int num_tokens = input_view.numel() / hidden_size;
178
- int64_t input_stride = input_view.stride(-2);
197
+ int num_tokens = input.numel() / hidden_size;
198
+ int num_dims = input.dim();
199
+ int64_t input_stride_d2 = input.stride(-2);
200
+ int64_t input_stride_d3 = (num_dims >= 3) ? input.stride(-3) : 0;
201
+ int64_t input_stride_d4 = (num_dims >= 4) ? input.stride(-4) : 0;
202
+ int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
203
+ int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
204
205
// For large num_tokens, use smaller blocks to increase SM concurrency.
206
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
207
dim3 grid(num_tokens);
183
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input_view));
208
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
209
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
185
- VLLM_DISPATCH_FLOATING_TYPES(
186
- input_view.scalar_type(), "rms_norm_kernel", [&] {
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
- });
210
+ VLLM_DISPATCH_RANK234(num_dims, [&] {
211
+ VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
212
+ const int calculated_vec_size =
213
+ std::gcd(16 / sizeof(scalar_t), hidden_size);
214
+ const int block_size =
215
+ std::min(hidden_size / calculated_vec_size, max_block_size);
216
+ dim3 block(block_size);
217
+ VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
218
+ vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
219
+ <<<grid, block, 0, stream>>>(
220
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
221
+ input_stride_d2, input_stride_d3, input_stride_d4,
222
+ input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
223
+ epsilon, num_tokens, hidden_size);
224
});
225
+ });
226
+ });
227
}
228
229
#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
vllm/_custom_ops.py
+1
-4
@@ -328,10 +328,7 @@ def rotary_embedding(
328
def rms_norm(
329
out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float
330
) -> None:
331
- # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
332
- # If removed, also need to remove contiguous in MatcherRMSNorm
333
- input_contiguous = input.contiguous()
334
- torch.ops._C.rms_norm(out, input_contiguous, weight, epsilon)
331
+ torch.ops._C.rms_norm(out, input, weight, epsilon)
332
333
334
def fused_add_rms_norm(
vllm/compilation/matcher_utils.py
+1
-3
@@ -162,12 +162,10 @@ class MatcherRMSNorm(MatcherCustomOp):
162
weight: torch.Tensor,
163
) -> torch.Tensor:
164
result = torch.empty_like(input)
165
- # TODO: support non-contiguous input for RMSNorm and remove this
166
- input_contiguous = input.contiguous()
165
_, result = auto_functionalized(
166
RMS_OP,
167
result=result,
170
- input=input_contiguous,
168
+ input=input,
169
weight=weight,
170
epsilon=self.epsilon,
171
)