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 14, 05:30 UTC
Vectorize RMS norm variance using vectorize_read_with_alignment (#26234)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Benji Beck committed
Oct 15, 2025 at 11:54 UTC
1f491aa0c80c2bf07e3ad37c4b6af8a869d48b5d
2 files changed
+31
-6
csrc/layernorm_kernels.cu
+15
-3
@@ -2,6 +2,7 @@
2
#include "dispatch_utils.h"
3
#include "cub_helpers.h"
4
#include "core/batch_invariant.hpp"
5
+#include "quantization/vectorization_utils.cuh"
6
7
#include <torch/cuda.h>
8
#include <c10/cuda/CUDAGuard.h>
@@ -18,11 +19,22 @@ __global__ void rms_norm_kernel(
19
const float epsilon, const int num_tokens, const int hidden_size) {
20
__shared__ float s_variance;
21
float variance = 0.0f;
22
+ const scalar_t* input_row = input + blockIdx.x * input_stride;
23
22
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
23
- const float x = (float)input[blockIdx.x * input_stride + idx];
24
+ constexpr int VEC_SIZE = 8;
25
+ auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
26
+#pragma unroll
27
+ for (int i = 0; i < VEC_SIZE; ++i) {
28
+ float x = static_cast<float>(vec.val[i]);
29
+ variance += x * x;
30
+ }
31
+ };
32
+ auto scalar_op = [&variance](const scalar_t& val) {
33
+ float x = static_cast<float>(val);
34
variance += x * x;
25
- }
35
+ };
36
+ vllm::vectorize_read_with_alignment<VEC_SIZE>(
37
+ input_row, hidden_size, threadIdx.x, blockDim.x, vec_op, scalar_op);
38
39
using BlockReduce = cub::BlockReduce<float, 1024>;
40
__shared__ typename BlockReduce::TempStorage reduceStore;
csrc/layernorm_quant_kernels.cu
+16
-3
@@ -10,6 +10,7 @@
10
#include "dispatch_utils.h"
11
#include "cub_helpers.h"
12
#include "core/batch_invariant.hpp"
13
+#include "quantization/vectorization_utils.cuh"
14
15
#include <torch/cuda.h>
16
#include <c10/cuda/CUDAGuard.h>
@@ -28,10 +29,22 @@ __global__ void rms_norm_static_fp8_quant_kernel(
29
__shared__ float s_variance;
30
float variance = 0.0f;
31
31
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
32
- const float x = (float)input[blockIdx.x * input_stride + idx];
32
+ const scalar_t* input_row = input + blockIdx.x * input_stride;
33
+
34
+ constexpr int VEC_SIZE = 8;
35
+ auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
36
+#pragma unroll
37
+ for (int i = 0; i < VEC_SIZE; ++i) {
38
+ float x = static_cast<float>(vec.val[i]);
39
+ variance += x * x;
40
+ }
41
+ };
42
+ auto scalar_op = [&variance](const scalar_t& val) {
43
+ float x = static_cast<float>(val);
44
variance += x * x;
34
- }
45
+ };
46
+ vllm::vectorize_read_with_alignment<VEC_SIZE>(
47
+ input_row, hidden_size, threadIdx.x, blockDim.x, vec_op, scalar_op);
48
49
using BlockReduce = cub::BlockReduce<float, 1024>;
50
__shared__ typename BlockReduce::TempStorage reduceStore;