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 3, 20:15 UTC
[Perf] Use 2D-grid to eliminate divmod in W8W8 group quant (#42153)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Jiahan Chang (Cyrus) committed
May 12, 2026 at 22:01 UTC
dd6b3a5ef548e9349e6c103dfa7b1b484c9ca41a
1 file changed
+69
-40
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+69
-40
@@ -156,6 +156,17 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
156
return 1;
157
}
158
159
+// Largest divisor of padded_groups_per_row that is <= 16. ry = 16 / kx.
160
+inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
161
+ if (padded_groups_per_row % 16 == 0) {
162
+ return 16;
163
+ }
164
+ if (padded_groups_per_row % 8 == 0) {
165
+ return 8;
166
+ }
167
+ return 4;
168
+}
169
+
170
void per_token_group_quant_8bit(const torch::stable::Tensor& input,
171
torch::stable::Tensor& output_q,
172
torch::stable::Tensor& output_s,
@@ -247,11 +258,11 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
258
//
259
// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
260
// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
250
-template <typename T, typename DST_DTYPE, int GROUP_SIZE>
261
+template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
262
+ int kRowsPerBlock>
263
__global__ void per_token_group_quant_8bit_packed_register_kernel(
264
const T* __restrict__ input, void* __restrict__ output_q,
253
- unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
254
- const int groups_per_block, const int padded_groups_per_row,
265
+ unsigned int* __restrict__ output_s_packed, const int padded_groups_per_row,
266
const int groups_per_row, const int mn, const int output_q_mn_extent,
267
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
268
const float min_8bit, const float max_8bit) {
@@ -260,27 +271,25 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
271
constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
272
static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
273
"GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
263
- // Each group's 8 threads must live in a single warp octet so the
264
- // 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
265
- // that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
266
- // to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
267
- // via num_threads = groups_per_block * THREADS_PER_GROUP).
274
static_assert(32 % THREADS_PER_GROUP == 0,
275
"THREADS_PER_GROUP must divide warp size for the shuffle "
276
"mask to be valid");
277
+ static_assert(
278
+ kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
279
+ "kGroupsPerBlockX must be a positive power of 2");
280
+ static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
281
282
const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
283
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
284
275
- const int64_t block_group_id = blockIdx.x * groups_per_block;
276
- const int64_t global_group_id = block_group_id + local_group_id;
277
- if (global_group_id >= num_groups_padded) {
285
+ const int sf_k_local = local_group_id % kGroupsPerBlockX;
286
+ const int row_local = local_group_id / kGroupsPerBlockX;
287
+ const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
288
+ const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
289
+
290
+ if (mn_idx >= tma_aligned_mn) {
291
return;
292
}
280
-
281
- const int sf_k_idx =
282
- static_cast<int>(global_group_id % padded_groups_per_row);
283
- const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
293
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
294
295
// Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -443,34 +452,53 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
452
453
constexpr int THREADS_PER_GROUP = 8;
454
const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
446
- const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
455
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
448
- const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
456
+
457
+ STD_TORCH_CHECK(padded_groups_per_row % 4 == 0,
458
+ "padded_groups_per_row=", padded_groups_per_row,
459
+ " is not a multiple of 4.");
460
+ const int kx = GetGroupsPerBlockX(padded_groups_per_row);
461
+ const int ry = 16 / kx;
462
+ const int64_t blocks_x = padded_groups_per_row / kx;
463
+ const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
464
+ const int num_threads = (kx * ry) * THREADS_PER_GROUP;
465
+ // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
466
+ STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
467
+ blocks_y <= static_cast<int64_t>(INT32_MAX),
468
+ "per_token_group_quant_8bit_packed grid too large: (",
469
+ blocks_x, ", ", blocks_y, ").");
470
471
auto dst_type = output_q.scalar_type();
451
- const int64_t num_blocks = num_groups_padded / groups_per_block;
452
- const int num_threads = groups_per_block * THREADS_PER_GROUP;
453
- // CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
454
- // against pathological inputs.
455
- STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
456
- "per_token_group_quant_8bit_packed grid too large: ",
457
- num_blocks, " blocks (max ", INT32_MAX, ").");
458
-
459
-#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
460
- do { \
461
- dim3 grid(static_cast<unsigned int>(num_blocks)); \
462
- dim3 block(num_threads); \
463
- per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128> \
464
- <<<grid, block, 0, stream>>>( \
465
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
466
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
467
- num_groups_padded, groups_per_block, \
468
- static_cast<int>(padded_groups_per_row), \
469
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
470
- static_cast<int>(output_q_mn_extent), \
471
- static_cast<int>(tma_aligned_mn), num_scale_elems, \
472
- static_cast<float>(eps), static_cast<float>(min_8bit), \
473
- static_cast<float>(max_8bit)); \
472
+
473
+#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
474
+ do { \
475
+ dim3 grid(static_cast<unsigned int>(blocks_x), \
476
+ static_cast<unsigned int>(blocks_y)); \
477
+ dim3 block(num_threads); \
478
+ per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
479
+ RY> \
480
+ <<<grid, block, 0, stream>>>( \
481
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
482
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
483
+ static_cast<int>(padded_groups_per_row), \
484
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
485
+ static_cast<int>(output_q_mn_extent), \
486
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
487
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
488
+ static_cast<float>(max_8bit)); \
489
+ } while (0)
490
+
491
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
492
+ do { \
493
+ if (kx == 16) { \
494
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1); \
495
+ } else if (kx == 8) { \
496
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2); \
497
+ } else if (kx == 4) { \
498
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4); \
499
+ } else { \
500
+ STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
501
+ } \
502
} while (0)
503
504
VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -488,6 +516,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
516
}));
517
518
#undef LAUNCH_REG_KERNEL
519
+#undef LAUNCH_REG_KERNEL_INST
520
}
521
522
void per_token_group_quant_fp8(const torch::stable::Tensor& input,