@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 309 lines 13.3 KB
Raw
1 /*
2 * This file contains the CUDA kernels for the fused quantized layernorm.
3 * The kernels correspond to the kernels in layernorm_kernels.cu, except they
4 * also produce quantized output directly.
5 * Currently, only static fp8 quantization is supported.
6 */
7
8 #include <numeric>
9
10 #include "torch_utils.h"
11
12 #include "cub_helpers.h"
13 #include "../core/batch_invariant.hpp"
14 #include "../quantization/w8a8/fp8/common.cuh"
15 #include "type_convert.cuh"
16 #include "dispatch_utils.h"
17 #include "quantization/vectorization_utils.cuh"
18
19 namespace vllm {
20
21 // TODO(woosuk): Further optimize this kernel.
22 template <typename scalar_t, typename fp8_type, int VEC_SIZE>
23 __global__ void rms_norm_static_fp8_quant_kernel(
24 fp8_type* __restrict__ out, // [..., hidden_size]
25 const scalar_t* __restrict__ input, // [..., hidden_size]
26 const int input_stride,
27 const scalar_t* __restrict__ weight, // [hidden_size]
28 const float* __restrict__ scale, // [1]
29 const float epsilon, const int num_tokens, const int hidden_size) {
30 __shared__ float s_variance;
31 float variance = 0.0f;
32
33 const scalar_t* input_row = input + blockIdx.x * input_stride;
34
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;
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;
51 variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
52
53 if (threadIdx.x == 0) {
54 s_variance = rsqrtf(variance / hidden_size + epsilon);
55 }
56 __syncthreads();
57
58 // invert scale to avoid division
59 float const scale_inv = 1.0f / *scale;
60
61 auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
62 auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
63 for (int idx = threadIdx.x; idx < hidden_size / VEC_SIZE; idx += blockDim.x) {
64 vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[idx];
65 vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[idx];
66 #pragma unroll
67 for (int j = 0; j < VEC_SIZE; j++) {
68 float x = static_cast<float>(src1.val[j]);
69 float w = static_cast<float>(src2.val[j]);
70 // Round normalized result through scalar_t to match the precision of the
71 // unfused composite (rms_norm writes scalar_t, then
72 // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
73 // Without this round, the fused path is strictly more accurate and
74 // disagrees with the composite at exact E4M3 quantization tie boundaries.
75 scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
76 out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
77 scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
78 scale_inv);
79 }
80 }
81 }
82
83 /* Function specialization in the case of FP16/BF16 tensors.
84 Additional optimizations we can make in this case are
85 packed and vectorized operations, which help with the
86 memory latency bottleneck. */
87 template <typename scalar_t, int width, typename fp8_type>
88 __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
89 fused_add_rms_norm_static_fp8_quant_kernel(
90 fp8_type* __restrict__ out, // [..., hidden_size]
91 scalar_t* __restrict__ input, // [..., hidden_size]
92 const int input_stride,
93 scalar_t* __restrict__ residual, // [..., hidden_size]
94 const scalar_t* __restrict__ weight, // [hidden_size]
95 const float* __restrict__ scale, // [1]
96 const float epsilon, const int num_tokens, const int hidden_size) {
97 // Sanity checks on our vector struct and type-punned pointer arithmetic
98 static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
99 static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
100
101 const int vec_hidden_size = hidden_size / width;
102 const int vec_input_stride = input_stride / width;
103 __shared__ float s_variance;
104 float variance = 0.0f;
105 /* These and the argument pointers are all declared `restrict` as they are
106 not aliased in practice. Argument pointers should not be dereferenced
107 in this kernel as that would be undefined behavior */
108 auto* __restrict__ input_v =
109 reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
110 auto* __restrict__ residual_v =
111 reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
112 auto* __restrict__ weight_v =
113 reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
114
115 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
116 int stride_id = blockIdx.x * vec_input_stride + idx;
117 int id = blockIdx.x * vec_hidden_size + idx;
118 _f16Vec<scalar_t, width> temp = input_v[stride_id];
119 temp += residual_v[id];
120 variance += temp.sum_squares();
121 residual_v[id] = temp;
122 }
123
124 using BlockReduce = cub::BlockReduce<float, 1024>;
125 __shared__ typename BlockReduce::TempStorage reduceStore;
126 variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
127
128 if (threadIdx.x == 0) {
129 s_variance = rsqrtf(variance / hidden_size + epsilon);
130 }
131 __syncthreads();
132
133 // invert scale to avoid division
134 float const scale_inv = 1.0f / *scale;
135
136 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
137 int id = blockIdx.x * vec_hidden_size + idx;
138 _f16Vec<scalar_t, width> res = residual_v[id];
139 _f16Vec<scalar_t, width> w = weight_v[idx];
140 using Converter = _typeConvert<scalar_t>;
141 using HipT = typename Converter::hip_type;
142 #pragma unroll
143 for (int i = 0; i < width; ++i) {
144 float x = Converter::convert(res.data[i]);
145 float wf = Converter::convert(w.data[i]);
146 // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
147 // to match the unfused composite path at FP8 boundaries. We use the
148 // backend's hip_type for the intermediate since c10::Half/BFloat16 has
149 // ambiguous conversions on CUDA and no implicit conversion on ROCm.
150 HipT out_norm_h = Converter::convert(x * s_variance * wf);
151 out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
152 Converter::convert(out_norm_h), scale_inv);
153 }
154 }
155 }
156
157 /* Generic fused_add_rms_norm_kernel
158 The width field is not used here but necessary for other specializations.
159 */
160 template <typename scalar_t, int width, typename fp8_type>
161 __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
162 fused_add_rms_norm_static_fp8_quant_kernel(
163 fp8_type* __restrict__ out, // [..., hidden_size]
164 scalar_t* __restrict__ input, // [..., hidden_size]
165 const int input_stride,
166 scalar_t* __restrict__ residual, // [..., hidden_size]
167 const scalar_t* __restrict__ weight, // [hidden_size]
168 const float* __restrict__ scale, // [1]
169 const float epsilon, const int num_tokens, const int hidden_size) {
170 __shared__ float s_variance;
171 float variance = 0.0f;
172
173 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
174 scalar_t z = input[blockIdx.x * input_stride + idx];
175 z += residual[blockIdx.x * hidden_size + idx];
176 float x = (float)z;
177 variance += x * x;
178 residual[blockIdx.x * hidden_size + idx] = z;
179 }
180
181 using BlockReduce = cub::BlockReduce<float, 1024>;
182 __shared__ typename BlockReduce::TempStorage reduceStore;
183 variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
184
185 if (threadIdx.x == 0) {
186 s_variance = rsqrtf(variance / hidden_size + epsilon);
187 }
188 __syncthreads();
189
190 // invert scale to avoid division
191 float const scale_inv = 1.0f / *scale;
192
193 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
194 float x = (float)residual[blockIdx.x * hidden_size + idx];
195 float w = (float)weight[idx];
196 // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
197 // to match the unfused composite path at FP8 boundaries.
198 scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
199 out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
200 static_cast<float>(out_norm), scale_inv);
201 }
202 }
203
204 } // namespace vllm
205
206 void rms_norm_static_fp8_quant(
207 torch::stable::Tensor& out, // [..., hidden_size]
208 torch::stable::Tensor& input, // [..., hidden_size]
209 torch::stable::Tensor& weight, // [hidden_size]
210 torch::stable::Tensor& scale, // [1]
211 double epsilon) {
212 STD_TORCH_CHECK(out.is_contiguous());
213 int hidden_size = input.size(-1);
214 int input_stride = input.stride(-2);
215 int num_tokens = input.numel() / hidden_size;
216
217 // For large num_tokens, use smaller blocks to increase SM concurrency.
218 const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
219 const int max_block_size =
220 batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
221 dim3 grid(num_tokens);
222 const torch::stable::accelerator::DeviceGuard device_guard(
223 input.get_device_index());
224 const cudaStream_t stream = get_current_cuda_stream();
225 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
226 input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
227 VLLM_STABLE_DISPATCH_FP8_TYPES(
228 out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
229 const int calculated_vec_size =
230 std::gcd(16 / sizeof(scalar_t), hidden_size);
231 const int block_size =
232 std::min(hidden_size / calculated_vec_size, max_block_size);
233 dim3 block(block_size);
234 VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
235 vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t,
236 vec_size>
237 <<<grid, block, 0, stream>>>(
238 out.mutable_data_ptr<fp8_t>(),
239 input.const_data_ptr<scalar_t>(), input_stride,
240 weight.const_data_ptr<scalar_t>(),
241 scale.const_data_ptr<float>(), epsilon, num_tokens,
242 hidden_size);
243 });
244 });
245 });
246 }
247
248 #define LAUNCH_FUSED_ADD_RMS_NORM(width) \
249 VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
250 input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] { \
251 VLLM_STABLE_DISPATCH_FP8_TYPES( \
252 out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
253 vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t, \
254 width, fp8_t> \
255 <<<grid, block, 0, stream>>>( \
256 out.mutable_data_ptr<fp8_t>(), \
257 input.mutable_data_ptr<scalar_t>(), input_stride, \
258 residual.mutable_data_ptr<scalar_t>(), \
259 weight.const_data_ptr<scalar_t>(), \
260 scale.const_data_ptr<float>(), epsilon, num_tokens, \
261 hidden_size); \
262 }); \
263 });
264 void fused_add_rms_norm_static_fp8_quant(
265 torch::stable::Tensor& out, // [..., hidden_size],
266 torch::stable::Tensor& input, // [..., hidden_size]
267 torch::stable::Tensor& residual, // [..., hidden_size]
268 torch::stable::Tensor& weight, // [hidden_size]
269 torch::stable::Tensor& scale, // [1]
270 double epsilon) {
271 STD_TORCH_CHECK(out.is_contiguous());
272 STD_TORCH_CHECK(residual.is_contiguous());
273 STD_TORCH_CHECK(residual.scalar_type() == input.scalar_type());
274 STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
275 int hidden_size = input.size(-1);
276 int input_stride = input.stride(-2);
277 int num_tokens = input.numel() / hidden_size;
278
279 dim3 grid(num_tokens);
280 /* This kernel is memory-latency bound in many scenarios.
281 When num_tokens is large, a smaller block size allows
282 for increased block occupancy on CUs and better latency
283 hiding on global mem ops. */
284 const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
285 const int max_block_size =
286 batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
287 dim3 block(std::min(hidden_size, max_block_size));
288 const torch::stable::accelerator::DeviceGuard device_guard(
289 input.get_device_index());
290 const cudaStream_t stream = get_current_cuda_stream();
291 /*If the tensor types are FP16/BF16, try to use the optimized kernel
292 with packed + vectorized ops.
293 Max optimization is achieved with a width-8 vector of FP16/BF16s
294 since we can load at most 128 bits at once in a global memory op.
295 However, this requires each tensor's data to be aligned to 16
296 bytes.
297 */
298 auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
299 auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
300 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
301 bool ptrs_are_aligned =
302 inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
303 if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
304 !batch_invariant_launch) {
305 LAUNCH_FUSED_ADD_RMS_NORM(8);
306 } else {
307 LAUNCH_FUSED_ADD_RMS_NORM(0);
308 }
309 }