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 11, 13:45 UTC
[Refactor] Reduce duplicate code in `per_token_group_quant` cuda kernels (#30496)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 12, 2025 at 16:45 UTC
1e6b115300b8b3629d10e69db0933246fe3253af
1 file changed
+83
-98
csrc/quantization/w8a8/fp8/per_token_group_quant.cu
+83
-98
@@ -22,6 +22,62 @@ __device__ __forceinline__ float GroupReduceMax(float val) {
22
return val;
23
}
24
25
+template <typename T, bool SCALE_UE8M0>
26
+__device__ __forceinline__ float ComputeGroupScale(
27
+ const T* __restrict__ group_input, T* __restrict__ smem_group,
28
+ const int group_size, const int lane_id, const int threads_per_group,
29
+ const float eps, const float max_8bit) {
30
+ float local_absmax = eps;
31
+
32
+ constexpr int vec_size = 16 / sizeof(T);
33
+
34
+ // copy global -> shared & compute absmax
35
+ auto scalar_op_cache = [&] __device__(T & dst, const T& src) {
36
+ float abs_v = fabsf(static_cast<float>(src));
37
+ local_absmax = fmaxf(local_absmax, abs_v);
38
+ dst = src;
39
+ };
40
+
41
+ vllm::vectorize_with_alignment<vec_size>(
42
+ group_input, // in
43
+ smem_group, // out (shared)
44
+ group_size, // elements per group
45
+ lane_id, // thread id
46
+ threads_per_group, // stride in group
47
+ scalar_op_cache); // scalar handler
48
+
49
+ local_absmax = GroupReduceMax(local_absmax);
50
+
51
+ float y_s = local_absmax / max_8bit;
52
+ if constexpr (SCALE_UE8M0) {
53
+ y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f))));
54
+ }
55
+
56
+ return y_s;
57
+}
58
+
59
+template <typename T, typename DST_DTYPE>
60
+__device__ __forceinline__ void QuantizeGroup(
61
+ const T* __restrict__ smem_group, DST_DTYPE* __restrict__ group_output,
62
+ const int group_size, const int lane_id, const int threads_per_group,
63
+ const float y_s, const float min_8bit, const float max_8bit) {
64
+ constexpr int vec_size = 16 / sizeof(T);
65
+
66
+ // quantize shared -> global 8-bit
67
+ auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) {
68
+ float q = fminf(fmaxf(static_cast<float>(src) / y_s, min_8bit), max_8bit);
69
+ dst = DST_DTYPE(q);
70
+ };
71
+
72
+ vllm::vectorize_with_alignment<vec_size>(
73
+ smem_group, // in (shared)
74
+ group_output, // out (global quant tensor)
75
+ group_size, // elements
76
+ lane_id, // tid
77
+ threads_per_group, // stride
78
+ scalar_op_quant); // scalar handler
79
+}
80
+
81
template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
82
bool SCALE_UE8M0 = false, typename scale_packed_t = float>
83
__global__ void per_token_group_quant_8bit_kernel(
@@ -38,8 +94,6 @@ __global__ void per_token_group_quant_8bit_kernel(
94
const int64_t global_group_id = block_group_id + local_group_id;
95
const int64_t block_group_offset = global_group_id * group_size;
96
41
- float local_absmax = eps;
42
-
97
using scale_element_t = float;
98
static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
99
@@ -68,30 +122,9 @@ __global__ void per_token_group_quant_8bit_kernel(
122
T* smem = reinterpret_cast<T*>(smem_raw);
123
T* smem_group = smem + local_group_id * group_size;
124
71
- constexpr int vec_size = 16 / sizeof(T);
72
- using vec_t = vllm::vec_n_t<T, vec_size>;
73
-
74
- // copy global -> shared & compute absmax
75
- auto scalar_op_cache = [&] __device__(T & dst, const T& src) {
76
- float abs_v = fabsf(static_cast<float>(src));
77
- local_absmax = fmaxf(local_absmax, abs_v);
78
- dst = src;
79
- };
80
-
81
- vllm::vectorize_with_alignment<vec_size>(
82
- group_input, // in
83
- smem_group, // out (shared)
84
- group_size, // elements per group
85
- lane_id, // thread id
86
- threads_per_group, // stride in group
87
- scalar_op_cache); // scalar handler
88
-
89
- local_absmax = GroupReduceMax(local_absmax);
90
-
91
- float y_s = local_absmax / max_8bit;
92
- if constexpr (SCALE_UE8M0) {
93
- y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f))));
94
- }
125
+ const float y_s = ComputeGroupScale<T, SCALE_UE8M0>(
126
+ group_input, smem_group, group_size, lane_id, threads_per_group, eps,
127
+ max_8bit);
128
129
scale_element_t y_s_quant = y_s;
130
@@ -101,19 +134,24 @@ __global__ void per_token_group_quant_8bit_kernel(
134
135
__syncthreads();
136
104
- // quantize shared -> global 8-bit
105
- auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) {
106
- float q = fminf(fmaxf(static_cast<float>(src) / y_s, min_8bit), max_8bit);
107
- dst = DST_DTYPE(q);
108
- };
137
+ QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
138
+ threads_per_group, y_s, min_8bit, max_8bit);
139
+}
140
110
- vllm::vectorize_with_alignment<vec_size>(
111
- smem_group, // in (shared)
112
- group_output, // out (global quant tensor)
113
- group_size, // elements
114
- lane_id, // tid
115
- threads_per_group, // stride
116
- scalar_op_quant); // scalar handler
141
+inline int GetGroupsPerBlock(int64_t num_groups) {
142
+ if (num_groups % 16 == 0) {
143
+ return 16;
144
+ }
145
+ if (num_groups % 8 == 0) {
146
+ return 8;
147
+ }
148
+ if (num_groups % 4 == 0) {
149
+ return 4;
150
+ }
151
+ if (num_groups % 2 == 0) {
152
+ return 2;
153
+ }
154
+ return 1;
155
}
156
157
void per_token_group_quant_8bit(const torch::Tensor& input,
@@ -133,17 +171,7 @@ void per_token_group_quant_8bit(const torch::Tensor& input,
171
172
constexpr int THREADS_PER_GROUP = 16;
173
136
- int groups_per_block = 1;
137
-
138
- if (num_groups % 16 == 0) {
139
- groups_per_block = 16;
140
- } else if (num_groups % 8 == 0) {
141
- groups_per_block = 8;
142
- } else if (num_groups % 4 == 0) {
143
- groups_per_block = 4;
144
- } else if (num_groups % 2 == 0) {
145
- groups_per_block = 2;
146
- }
174
+ const int groups_per_block = GetGroupsPerBlock(num_groups);
175
176
auto dst_type = output_q.scalar_type();
177
const int num_blocks = num_groups / groups_per_block;
@@ -225,8 +253,6 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
253
254
const int64_t block_group_offset = global_group_id * group_size;
255
228
- float local_absmax = eps;
229
-
256
const T* group_input = input + block_group_offset;
257
DST_DTYPE* group_output =
258
static_cast<DST_DTYPE*>(output_q) + block_group_offset;
@@ -235,29 +261,9 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
261
extern __shared__ __align__(16) char smem_raw[];
262
T* smem = reinterpret_cast<T*>(smem_raw);
263
T* smem_group = smem + local_group_id * group_size;
238
-
239
- constexpr int vec_size = 16 / sizeof(T);
240
- using vec_t = vllm::vec_n_t<T, vec_size>;
241
-
242
- // copy global -> shared & compute absmax
243
- auto scalar_op_cache = [&] __device__(T & dst, const T& src) {
244
- float abs_v = fabsf(static_cast<float>(src));
245
- local_absmax = fmaxf(local_absmax, abs_v);
246
- dst = src;
247
- };
248
-
249
- vllm::vectorize_with_alignment<vec_size>(
250
- group_input, // in
251
- smem_group, // out (shared)
252
- group_size, // elements per group
253
- lane_id, // thread id
254
- threads_per_group, // stride in group
255
- scalar_op_cache); // scalar handler
256
-
257
- local_absmax = GroupReduceMax(local_absmax);
258
-
259
- float y_s = local_absmax / max_8bit;
260
- y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f))));
264
+ const float y_s =
265
+ ComputeGroupScale<T, true>(group_input, smem_group, group_size, lane_id,
266
+ threads_per_group, eps, max_8bit);
267
268
// pack 4 scales into a uint32
269
if (lane_id == 0) {
@@ -284,19 +290,8 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
290
291
__syncthreads();
292
287
- // quantize shared -> global 8-bit
288
- auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) {
289
- float q = fminf(fmaxf(static_cast<float>(src) / y_s, min_8bit), max_8bit);
290
- dst = DST_DTYPE(q);
291
- };
292
-
293
- vllm::vectorize_with_alignment<vec_size>(
294
- smem_group, // in (shared)
295
- group_output, // out (global quant tensor)
296
- group_size, // elements
297
- lane_id, // tid
298
- threads_per_group, // stride
299
- scalar_op_quant); // scalar handler
293
+ QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
294
+ threads_per_group, y_s, min_8bit, max_8bit);
295
}
296
297
void per_token_group_quant_8bit_packed(const torch::Tensor& input,
@@ -337,17 +332,7 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
332
333
constexpr int THREADS_PER_GROUP = 16;
334
340
- int groups_per_block = 1;
341
-
342
- if (num_groups % 16 == 0) {
343
- groups_per_block = 16;
344
- } else if (num_groups % 8 == 0) {
345
- groups_per_block = 8;
346
- } else if (num_groups % 4 == 0) {
347
- groups_per_block = 4;
348
- } else if (num_groups % 2 == 0) {
349
- groups_per_block = 2;
350
- }
335
+ const int groups_per_block = GetGroupsPerBlock(num_groups);
336
337
auto dst_type = output_q.scalar_type();
338
const int num_blocks = num_groups / groups_per_block;