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 15, 07:15 UTC
[Perf] Tune `scaled_fp8_quant` by increasing vectorization (#18844)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 3, 2025 at 16:48 UTC
e31446b6c8d887cdca031abf8527555adee46058
4 files changed
+115
-110
csrc/quantization/fp8/common.cu
+19
-16
@@ -39,8 +39,8 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
39
fp8_type* __restrict__ token_output = &out[offset];
40
41
// For vectorization, token_input and token_output pointers need to be
42
- // aligned at 8-byte and 4-byte addresses respectively.
43
- bool const can_vectorize = hidden_size % 4 == 0;
42
+ // aligned at 32-byte and 16-byte addresses respectively.
43
+ bool const can_vectorize = hidden_size % 16 == 0;
44
45
float absmax_val = 0.0f;
46
if (can_vectorize) {
@@ -48,24 +48,24 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
48
} else {
49
for (int i = tid; i < hidden_size; i += blockDim.x) {
50
float const x = static_cast<float>(token_input[i]);
51
- absmax_val = max(absmax_val, fabs(x));
51
+ absmax_val = fmaxf(absmax_val, fabsf(x));
52
}
53
}
54
55
- using BlockReduce = cub::BlockReduce<float, 1024>;
55
+ using BlockReduce = cub::BlockReduce<float, 256>;
56
__shared__ typename BlockReduce::TempStorage reduceStorage;
57
float const block_absmax_val_maybe =
58
BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
59
__shared__ float token_scale;
60
if (tid == 0) {
61
if (scale_ub) {
62
- token_scale = min(block_absmax_val_maybe, *scale_ub);
62
+ token_scale = fminf(block_absmax_val_maybe, *scale_ub);
63
} else {
64
token_scale = block_absmax_val_maybe;
65
}
66
// token scale computation
67
- token_scale = max(token_scale / quant_type_max_v<fp8_type>,
68
- min_scaling_factor<fp8_type>::val());
67
+ token_scale = fmaxf(token_scale / quant_type_max_v<fp8_type>,
68
+ min_scaling_factor<fp8_type>::val());
69
scale[token_idx] = token_scale;
70
}
71
__syncthreads();
@@ -88,10 +88,11 @@ void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
88
torch::Tensor const& input, // [..., d]
89
torch::Tensor const& scale) // [1]
90
{
91
- int64_t num_tokens = input.numel() / input.size(-1);
92
- int64_t num_elems = input.numel();
93
- dim3 grid(num_tokens);
94
- dim3 block(1024);
91
+ int const block_size = 256;
92
+ int const num_tokens = input.numel() / input.size(-1);
93
+ int const num_elems = input.numel();
94
+ dim3 const grid(num_tokens);
95
+ dim3 const block(block_size);
96
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
97
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
98
VLLM_DISPATCH_FLOATING_TYPES(
@@ -110,10 +111,11 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out, // [..., d]
111
torch::Tensor const& input, // [..., d]
112
torch::Tensor& scale) // [1]
113
{
113
- int64_t num_tokens = input.numel() / input.size(-1);
114
- int64_t num_elems = input.numel();
115
- dim3 grid(num_tokens);
116
- dim3 block(1024);
114
+ int const block_size = 256;
115
+ int const num_tokens = input.numel() / input.size(-1);
116
+ int const num_elems = input.numel();
117
+ dim3 const grid(num_tokens);
118
+ dim3 const block(block_size);
119
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
120
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
121
VLLM_DISPATCH_FLOATING_TYPES(
@@ -141,8 +143,9 @@ void dynamic_per_token_scaled_fp8_quant(
143
144
int const hidden_size = input.size(-1);
145
int const num_tokens = input.numel() / hidden_size;
146
+ int const block_size = 256;
147
dim3 const grid(num_tokens);
145
- dim3 const block(std::min(hidden_size, 1024));
148
+ dim3 const block(std::min(hidden_size, block_size));
149
150
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
151
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
csrc/quantization/fp8/common.cuh
+35
-33
@@ -46,7 +46,7 @@ __device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
46
}
47
48
float r =
49
- fmax(-quant_type_max_v<fp8_type>, fmin(x, quant_type_max_v<fp8_type>));
49
+ fmaxf(-quant_type_max_v<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
50
#ifndef USE_ROCM
51
return static_cast<fp8_type>(r);
52
#else
@@ -65,7 +65,7 @@ template <typename scalar_t, typename fp8_type>
65
__global__ void segmented_max_reduction(float* __restrict__ scale,
66
const scalar_t* __restrict__ input,
67
int64_t num_elems) {
68
- __shared__ float cache[1024];
68
+ __shared__ float cache[256];
69
int64_t i = blockDim.x * blockIdx.x + threadIdx.x;
70
71
// First store maximum for all values processes by
@@ -73,7 +73,7 @@ __global__ void segmented_max_reduction(float* __restrict__ scale,
73
scalar_t tmp = 0.0;
74
while (i < num_elems) {
75
float x = static_cast<float>(input[i]);
76
- tmp = max(tmp, fabs(x));
76
+ tmp = fmaxf(tmp, fabsf(x));
77
i += blockDim.x * gridDim.x;
78
}
79
cache[threadIdx.x] = tmp;
@@ -100,25 +100,27 @@ template <typename scalar_t>
100
__device__ float thread_max_vec(scalar_t const* __restrict__ input,
101
int64_t const num_elems, int const tid,
102
int const step) {
103
+ constexpr size_t VEC_SIZE = 16;
104
+ using scalarxN_t = vec_n_t<scalar_t, VEC_SIZE>;
105
// Vectorized input/output to better utilize memory bandwidth.
104
- vec4_t<scalar_t> const* vectorized_in =
105
- reinterpret_cast<vec4_t<scalar_t> const*>(input);
106
+ auto const* vectorized_in = reinterpret_cast<scalarxN_t const*>(input);
107
107
- int64_t const num_vec_elems = num_elems >> 2;
108
+ // num_elems / VEC_SIZE (which is 16)
109
+ int64_t const num_vec_elems = num_elems >> 4;
110
float absmax_val = 0.0f;
111
110
-#pragma unroll 4
112
+#pragma unroll
113
for (int64_t i = tid; i < num_vec_elems; i += step) {
112
- vec4_t<scalar_t> in_vec = vectorized_in[i];
113
- absmax_val = max(absmax_val, fabs(in_vec.x));
114
- absmax_val = max(absmax_val, fabs(in_vec.y));
115
- absmax_val = max(absmax_val, fabs(in_vec.z));
116
- absmax_val = max(absmax_val, fabs(in_vec.w));
114
+ scalarxN_t in_vec = vectorized_in[i];
115
+#pragma unroll
116
+ for (int j = 0; j < VEC_SIZE; ++j) {
117
+ absmax_val = fmaxf(absmax_val, fabsf(in_vec.val[j]));
118
+ }
119
}
120
119
- // Handle the remaining elements if num_elems is not divisible by 4
120
- for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
121
- absmax_val = max(absmax_val, fabs(input[i]));
121
+ // Handle the remaining elements if num_elems is not divisible by VEC_SIZE
122
+ for (int64_t i = num_vec_elems * VEC_SIZE + tid; i < num_elems; i += step) {
123
+ absmax_val = fmaxf(absmax_val, fabsf(input[i]));
124
}
125
126
return absmax_val;
@@ -130,31 +132,31 @@ __device__ void scaled_fp8_conversion_vec(fp8_type* __restrict__ out,
132
float const scale,
133
int64_t const num_elems,
134
int const tid, int const step) {
133
- using float8x4_t = q8x4_t<fp8_type>;
135
+ constexpr size_t VEC_SIZE = 16;
136
+ using scalarxN_t = vec_n_t<scalar_t, VEC_SIZE>;
137
+ using float8xN_t = q8_n_t<fp8_type, VEC_SIZE>;
138
// Vectorized input/output to better utilize memory bandwidth.
135
- auto const* vectorized_in = reinterpret_cast<vec4_t<scalar_t> const*>(input);
136
- auto* vectorized_out = reinterpret_cast<float8x4_t*>(out);
139
+ auto const* vectorized_in = reinterpret_cast<scalarxN_t const*>(input);
140
+ auto* vectorized_out = reinterpret_cast<float8xN_t*>(out);
141
138
- int64_t const num_vec_elems = num_elems >> 2;
142
+ // num_elems / VEC_SIZE (which is 16)
143
+ int64_t const num_vec_elems = num_elems >> 4;
144
140
-#pragma unroll 4
145
+#pragma unroll
146
for (int64_t i = tid; i < num_vec_elems; i += step) {
142
- vec4_t<scalar_t> in_vec = vectorized_in[i];
143
- float8x4_t out_vec;
144
-
145
- out_vec.x = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
146
- static_cast<float>(in_vec.x), scale);
147
- out_vec.y = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
148
- static_cast<float>(in_vec.y), scale);
149
- out_vec.z = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
150
- static_cast<float>(in_vec.z), scale);
151
- out_vec.w = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
152
- static_cast<float>(in_vec.w), scale);
147
+ scalarxN_t in_vec = vectorized_in[i];
148
+ float8xN_t out_vec;
149
+
150
+#pragma unroll
151
+ for (int j = 0; j < VEC_SIZE; ++j) {
152
+ out_vec.val[j] = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
153
+ static_cast<float>(in_vec.val[j]), scale);
154
+ }
155
vectorized_out[i] = out_vec;
156
}
157
156
- // Handle the remaining elements if num_elems is not divisible by 4
157
- for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
158
+ // Handle the remaining elements if num_elems is not divisible by VEC_SIZE
159
+ for (int64_t i = num_vec_elems * VEC_SIZE + tid; i < num_elems; i += step) {
160
out[i] = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
161
static_cast<float>(input[i]), scale);
162
}
csrc/quantization/fused_kernels/layernorm_utils.cuh
+50
-49
@@ -140,6 +140,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
140
// sum of squares
141
float ss = 0.0f;
142
143
+ const int VEC_SIZE = 4;
144
int32_t const num_vec_elems = hidden_size >> 2;
145
146
#pragma unroll 4
@@ -147,22 +148,23 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
148
vec4_t<scalar_t> in = vec_input[i];
149
150
vec4_t<float> x;
150
- x.x = static_cast<float>(in.x);
151
- x.y = static_cast<float>(in.y);
152
- x.z = static_cast<float>(in.z);
153
- x.w = static_cast<float>(in.w);
151
+#pragma unroll
152
+ for (int j = 0; j < VEC_SIZE; ++j) {
153
+ x.val[j] = static_cast<float>(in.val[j]);
154
+ }
155
+
156
if constexpr (has_residual) {
157
vec4_t<scalar_t> r = vec_residual[i];
156
- x.x += static_cast<float>(r.x);
157
- x.y += static_cast<float>(r.y);
158
- x.z += static_cast<float>(r.z);
159
- x.w += static_cast<float>(r.w);
158
+#pragma unroll
159
+ for (int j = 0; j < VEC_SIZE; ++j) {
160
+ x.val[j] += static_cast<float>(r.val[j]);
161
+ }
162
}
163
162
- ss += x.x * x.x;
163
- ss += x.y * x.y;
164
- ss += x.z * x.z;
165
- ss += x.w * x.w;
164
+#pragma unroll
165
+ for (int j = 0; j < VEC_SIZE; ++j) {
166
+ ss += x.val[j] * x.val[j];
167
+ }
168
}
169
170
using BlockReduce = cub::BlockReduce<float, 1024>;
@@ -203,6 +205,7 @@ __device__ void compute_dynamic_per_token_scales(
205
206
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
207
208
+ const int VEC_SIZE = 4;
209
int32_t const num_vec_elems = hidden_size >> 2;
210
float block_absmax_val_maybe = 0.0f;
211
@@ -212,26 +215,25 @@ __device__ void compute_dynamic_per_token_scales(
215
vec4_t<scalar_t> const w = vec_weight[i];
216
217
vec4_t<float> x;
215
- x.x = static_cast<float>(in.x);
216
- x.y = static_cast<float>(in.y);
217
- x.z = static_cast<float>(in.z);
218
- x.w = static_cast<float>(in.w);
218
+#pragma unroll
219
+ for (int j = 0; j < VEC_SIZE; ++j) {
220
+ x.val[j] = static_cast<float>(in.val[j]);
221
+ }
222
+
223
if constexpr (has_residual) {
224
vec4_t<scalar_t> r = vec_residual[i];
221
- x.x += static_cast<float>(r.x);
222
- x.y += static_cast<float>(r.y);
223
- x.z += static_cast<float>(r.z);
224
- x.w += static_cast<float>(r.w);
225
+#pragma unroll
226
+ for (int j = 0; j < VEC_SIZE; ++j) {
227
+ x.val[j] += static_cast<float>(r.val[j]);
228
+ }
229
}
230
227
- block_absmax_val_maybe = fmaxf(
228
- block_absmax_val_maybe, fabs(static_cast<scalar_t>(x.x * rms) * w.x));
229
- block_absmax_val_maybe = fmaxf(
230
- block_absmax_val_maybe, fabs(static_cast<scalar_t>(x.y * rms) * w.y));
231
- block_absmax_val_maybe = fmaxf(
232
- block_absmax_val_maybe, fabs(static_cast<scalar_t>(x.z * rms) * w.z));
233
- block_absmax_val_maybe = fmaxf(
234
- block_absmax_val_maybe, fabs(static_cast<scalar_t>(x.w * rms) * w.w));
231
+#pragma unroll
232
+ for (int j = 0; j < VEC_SIZE; ++j) {
233
+ block_absmax_val_maybe =
234
+ fmaxf(block_absmax_val_maybe,
235
+ fabs(static_cast<scalar_t>(x.val[j] * rms) * w.val[j]));
236
+ }
237
}
238
239
using BlockReduce = cub::BlockReduce<float, 1024>;
@@ -282,6 +284,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
284
vec_residual = reinterpret_cast<vec4_t<scalar_t>*>(&residual[token_offset]);
285
}
286
287
+ const int VEC_SIZE = 4;
288
int32_t const num_vec_elems = hidden_size >> 2;
289
290
// TODO(luka/varun) extract into type-agnostic vectorized quant function to
@@ -292,33 +295,31 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
295
vec4_t<scalar_t> const w = vec_weight[i];
296
297
vec4_t<float> x;
295
- x.x = static_cast<float>(in.x);
296
- x.y = static_cast<float>(in.y);
297
- x.z = static_cast<float>(in.z);
298
- x.w = static_cast<float>(in.w);
298
+#pragma unroll
299
+ for (int j = 0; j < VEC_SIZE; ++j) {
300
+ x.val[j] = static_cast<float>(in.val[j]);
301
+ }
302
+
303
if constexpr (has_residual) {
304
vec4_t<scalar_t> r = vec_residual[i];
301
- x.x += static_cast<float>(r.x);
302
- x.y += static_cast<float>(r.y);
303
- x.z += static_cast<float>(r.z);
304
- x.w += static_cast<float>(r.w);
305
- // Update residual
306
- r.x = static_cast<scalar_t>(x.x);
307
- r.y = static_cast<scalar_t>(x.y);
308
- r.z = static_cast<scalar_t>(x.z);
309
- r.w = static_cast<scalar_t>(x.w);
305
+#pragma unroll
306
+ for (int j = 0; j < VEC_SIZE; ++j) {
307
+ x.val[j] += static_cast<float>(r.val[j]);
308
+ }
309
+// Update residual
310
+#pragma unroll
311
+ for (int j = 0; j < VEC_SIZE; ++j) {
312
+ r.val[j] = static_cast<scalar_t>(x.val[j]);
313
+ }
314
vec_residual[i] = r;
315
}
316
317
q8x4_t<scalar_out_t> out;
314
- out.x = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
315
- static_cast<scalar_t>(x.x * rms) * w.x, scale);
316
- out.y = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
317
- static_cast<scalar_t>(x.y * rms) * w.y, scale);
318
- out.z = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
319
- static_cast<scalar_t>(x.z * rms) * w.z, scale);
320
- out.w = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
321
- static_cast<scalar_t>(x.w * rms) * w.w, scale);
318
+#pragma unroll
319
+ for (int j = 0; j < VEC_SIZE; ++j) {
320
+ out.val[j] = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
321
+ static_cast<scalar_t>(x.val[j] * rms) * w.val[j], scale);
322
+ }
323
vec_output[i] = out;
324
}
325
}
csrc/quantization/vectorization.cuh
+11
-12
@@ -10,23 +10,22 @@
10
namespace vllm {
11
12
// Vectorization containers
13
-template <typename scalar_t>
14
-struct __align__(8) vec4_t {
15
- scalar_t x;
16
- scalar_t y;
17
- scalar_t z;
18
- scalar_t w;
13
+template <typename scalar_t, size_t vec_size>
14
+struct __align__(vec_size * sizeof(scalar_t)) vec_n_t {
15
+ scalar_t val[vec_size];
16
};
17
21
-template <typename quant_type_t>
22
-struct __align__(4) q8x4_t {
18
+template <typename quant_type_t, size_t vec_size>
19
+struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t {
20
static_assert(std::is_same_v<quant_type_t, int8_t> ||
21
std::is_same_v<quant_type_t, c10::Float8_e4m3fn> ||
22
std::is_same_v<quant_type_t, c10::Float8_e4m3fnuz>);
26
- quant_type_t x;
27
- quant_type_t y;
28
- quant_type_t z;
29
- quant_type_t w;
23
+ quant_type_t val[vec_size];
24
};
25
26
+template <typename scalar_t>
27
+using vec4_t = vec_n_t<scalar_t, 4>;
28
+template <typename quant_type_t>
29
+using q8x4_t = q8_n_t<quant_type_t, 4>;
30
+
31
} // namespace vllm