[Kernel][Perf][Quantization] Fix odd-row performance cliff in per-token-group quantization (#56478)
Signed-off-by: Ayushman Singh <ayush1399@users.noreply.github.com> Co-authored-by: Ayushman Singh <ayush1399@users.noreply.github.com>
Ayushman Singh committed
Sep 12, 2026 at 23:27 UTC
c2ea9f1a5f10caec36eac91c461fcec6564f5394
2 files changed
+131
-64
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+88
-58
@@ -96,7 +96,8 @@ __device__ __forceinline__ void QuantizeGroup(
96
}
97
98
template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
99
- bool SCALE_UE8M0 = false, typename scale_packed_t = float>
99
+ bool SCALE_UE8M0 = false, typename scale_packed_t = float,
100
+ bool CHECK_BOUNDS = false>
101
__global__ void per_token_group_quant_8bit_kernel(
102
const T* __restrict__ input, void* __restrict__ output_q,
103
scale_packed_t* __restrict__ output_s, const int group_size,
@@ -114,49 +115,54 @@ __global__ void per_token_group_quant_8bit_kernel(
115
using scale_element_t = float;
116
static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
117
117
- const T* group_input = input + block_group_offset;
118
- DST_DTYPE* group_output =
119
- static_cast<DST_DTYPE*>(output_q) + block_group_offset;
120
- scale_element_t* scale_output;
121
-
118
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
119
cudaGridDependencySynchronize();
120
#endif
121
126
- if constexpr (IS_COLUMN_MAJOR) {
127
- const int num_elems_per_pack =
128
- static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
129
- const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
130
- const int row_idx = global_group_id / scale_num_rows_element;
131
- const int col_idx_raw = global_group_id % scale_num_rows_element;
132
- const int col_idx = col_idx_raw / num_elems_per_pack;
133
- const int pack_idx = col_idx_raw % num_elems_per_pack;
134
- scale_output = reinterpret_cast<scale_element_t*>(output_s) +
135
- (col_idx * scale_stride * num_elems_per_pack +
136
- row_idx * num_elems_per_pack + pack_idx);
137
- } else {
138
- scale_output = output_s + global_group_id;
139
- }
140
-
122
// shared memory to cache each group's data to avoid double DRAM reads.
123
extern __shared__ __align__(16) char smem_raw[];
124
T* smem = reinterpret_cast<T*>(smem_raw);
125
T* smem_group = smem + local_group_id * group_size;
126
+ const bool valid_group = !CHECK_BOUNDS || global_group_id < num_groups;
127
+ float y_s = 1.0f;
128
+ if (valid_group) {
129
+ const T* group_input = input + block_group_offset;
130
+ scale_element_t* scale_output;
131
+
132
+ if constexpr (IS_COLUMN_MAJOR) {
133
+ const int num_elems_per_pack =
134
+ static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
135
+ const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
136
+ const int row_idx = global_group_id / scale_num_rows_element;
137
+ const int col_idx_raw = global_group_id % scale_num_rows_element;
138
+ const int col_idx = col_idx_raw / num_elems_per_pack;
139
+ const int pack_idx = col_idx_raw % num_elems_per_pack;
140
+ scale_output = reinterpret_cast<scale_element_t*>(output_s) +
141
+ (col_idx * scale_stride * num_elems_per_pack +
142
+ row_idx * num_elems_per_pack + pack_idx);
143
+ } else {
144
+ scale_output = output_s + global_group_id;
145
+ }
146
146
- const float y_s = ComputeGroupScale<T, SCALE_UE8M0>(
147
- group_input, smem_group, group_size, lane_id, threads_per_group, eps,
148
- max_8bit);
147
+ y_s = ComputeGroupScale<T, SCALE_UE8M0>(group_input, smem_group, group_size,
148
+ lane_id, threads_per_group, eps,
149
+ max_8bit);
150
150
- scale_element_t y_s_quant = y_s;
151
+ scale_element_t y_s_quant = y_s;
152
152
- if (lane_id == 0) {
153
- *scale_output = y_s_quant;
153
+ if (lane_id == 0) {
154
+ *scale_output = y_s_quant;
155
+ }
156
}
157
158
__syncthreads();
159
158
- QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
159
- threads_per_group, y_s, min_8bit, max_8bit);
160
+ if (valid_group) {
161
+ DST_DTYPE* group_output =
162
+ static_cast<DST_DTYPE*>(output_q) + block_group_offset;
163
+ QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
164
+ threads_per_group, y_s, min_8bit, max_8bit);
165
+ }
166
167
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
168
cudaTriggerProgrammaticLaunchCompletion();
@@ -209,10 +215,20 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
215
216
constexpr int THREADS_PER_GROUP = 16;
217
212
- const int groups_per_block = GetGroupsPerBlock(num_groups);
218
+ if (num_groups == 0) {
219
+ return;
220
+ }
221
+ STD_TORCH_CHECK(input.size(-1) % group_size == 0);
222
+ const int64_t groups_per_row = input.size(-1) / group_size;
223
+ const int64_t num_rows = num_groups / groups_per_row;
224
+ // Odd row counts inherit the next even row count's launch geometry.
225
+ const int64_t policy_num_groups =
226
+ (num_rows + (num_rows & 1)) * groups_per_row;
227
+ const int groups_per_block = GetGroupsPerBlock(policy_num_groups);
228
+ const bool needs_bounds_check = num_groups % groups_per_block != 0;
229
230
auto dst_type = output_q.scalar_type();
215
- const int num_blocks = num_groups / groups_per_block;
231
+ const int num_blocks = (num_groups + groups_per_block - 1) / groups_per_block;
232
const int num_threads = groups_per_block * THREADS_PER_GROUP;
233
234
const bool is_column_major = output_s.stride(0) < output_s.stride(1);
@@ -220,38 +236,51 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
236
const int scale_stride = output_s.stride(1);
237
238
#ifndef USE_ROCM
223
- #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
224
- do { \
225
- cudaLaunchConfig_t config = {}; \
226
- config.gridDim = dim3(num_blocks); \
227
- config.blockDim = dim3(num_threads); \
228
- config.dynamicSmemBytes = (SMEM_BYTES); \
229
- config.stream = stream; \
230
- cudaLaunchAttribute attrs[1]; \
231
- attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
232
- attrs[0].val.programmaticStreamSerializationAllowed = 1; \
233
- config.numAttrs = 1; \
234
- config.attrs = attrs; \
235
- cudaLaunchKernelEx( \
236
- &config, \
237
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
238
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
239
- static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
240
- groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
241
- scale_num_rows, scale_stride); \
239
+ #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \
240
+ SMEM_BYTES) \
241
+ do { \
242
+ cudaLaunchConfig_t config = {}; \
243
+ config.gridDim = dim3(num_blocks); \
244
+ config.blockDim = dim3(num_threads); \
245
+ config.dynamicSmemBytes = (SMEM_BYTES); \
246
+ config.stream = stream; \
247
+ cudaLaunchAttribute attrs[1]; \
248
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
249
+ attrs[0].val.programmaticStreamSerializationAllowed = 1; \
250
+ config.numAttrs = 1; \
251
+ config.attrs = attrs; \
252
+ cudaLaunchKernelEx( \
253
+ &config, \
254
+ per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, \
255
+ float, CHECK_BOUNDS>, \
256
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
257
+ static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
258
+ groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
259
+ scale_num_rows, scale_stride); \
260
} while (0)
261
#else
244
- #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
245
- do { \
246
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0> \
247
- <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
248
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
249
- static_cast<float*>(output_s.data_ptr()), group_size, \
250
- num_groups, groups_per_block, (float)eps, (float)min_8bit, \
251
- (float)max_8bit, scale_num_rows, scale_stride); \
262
+ #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \
263
+ SMEM_BYTES) \
264
+ do { \
265
+ per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, float, \
266
+ CHECK_BOUNDS> \
267
+ <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
268
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
269
+ static_cast<float*>(output_s.data_ptr()), group_size, \
270
+ num_groups, groups_per_block, (float)eps, (float)min_8bit, \
271
+ (float)max_8bit, scale_num_rows, scale_stride); \
272
} while (0)
273
#endif
274
275
+#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
276
+ do { \
277
+ if (needs_bounds_check) { \
278
+ LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, true, SMEM_BYTES); \
279
+ } else { \
280
+ LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, false, SMEM_BYTES); \
281
+ } \
282
+ } while (0)
283
+
284
#define LAUNCH_KERNEL(T, DST_DTYPE) \
285
do { \
286
size_t smem_bytes = \
@@ -284,6 +313,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
313
314
#undef LAUNCH_KERNEL
315
#undef LAUNCH_KERNEL_INST
316
+#undef LAUNCH_KERNEL_IMPL
317
}
318
319
// Register-resident fast path for group_size==128.
tests/kernels/quantization/test_per_token_group_quant.py
+43
-6
@@ -11,8 +11,24 @@ from vllm.platforms import current_platform
11
12
13
@pytest.mark.parametrize(
14
- "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
14
+ "shape",
15
+ [
16
+ (0, 128),
17
+ (1, 128),
18
+ (3, 384),
19
+ (5, 384),
20
+ (7, 384),
21
+ (15, 384),
22
+ (31, 128),
23
+ (32, 128),
24
+ (63, 256),
25
+ (64, 256),
26
+ (16, 512),
27
+ (3, 2688),
28
+ (3, 5376),
29
+ ],
30
)
31
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
32
@pytest.mark.parametrize("column_major", [False, True])
33
@pytest.mark.parametrize("tma_aligned", [False, True])
34
@pytest.mark.parametrize("scale_ue8m0", [False, True])
@@ -22,14 +38,20 @@ from vllm.platforms import current_platform
38
reason="Only test on CUDA/ROCm/XPU.",
39
)
40
def test_per_token_group_quant_fp8(
25
- shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
41
+ shape,
42
+ dtype,
43
+ column_major: bool,
44
+ tma_aligned: bool,
45
+ scale_ue8m0: bool,
46
+ group_size: int,
47
):
48
+ """Tail groups preserve FP8 values and scales in every supported layout."""
49
device = current_platform.device_type
50
51
torch.manual_seed(42)
52
num_tokens, hidden_dim = shape
53
32
- x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
54
+ x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
55
56
# native kernel path
57
out_q, scale = fp8_utils.per_token_group_quant_fp8(
@@ -411,19 +433,34 @@ def test_per_token_group_quant_fp8_packed_large_mn():
433
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
434
435
414
-@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
436
+@pytest.mark.parametrize(
437
+ "shape",
438
+ [
439
+ (0, 128),
440
+ (1, 128),
441
+ (3, 384),
442
+ (5, 384),
443
+ (7, 384),
444
+ (15, 384),
445
+ (32, 128),
446
+ (64, 256),
447
+ (16, 512),
448
+ ],
449
+)
450
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
451
@pytest.mark.parametrize("group_size", [64, 128])
452
@pytest.mark.skipif(
453
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
454
reason="Only test on CUDA/ROCm/XPU.",
455
)
420
-def test_per_token_group_quant_int8(shape, group_size: int):
456
+def test_per_token_group_quant_int8(shape, dtype, group_size: int):
457
+ """The shared launch policy also handles INT8 tail groups."""
458
device = current_platform.device_type
459
460
torch.manual_seed(42)
461
num_tokens, hidden_dim = shape
462
426
- x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
463
+ x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
464
465
# cuda path
466
out_q, scale = int8_utils.per_token_group_quant_int8(