[Kernel][Perf] Register-resident path for per-token-group 8-bit quant (#55330)
Signed-off-by: chao.huan <chao.huan@nio.com> Co-authored-by: Misha Goin <mgoin64@gmail.com>
chuan932 committed
Sep 25, 2026 at 23:11 UTC
f550ad8bf7a3983a16aafd25ad18d8b48f4d5218
3 files changed
+295
-6
benchmarks/kernels/benchmark_per_token_group_quant.py
+12
-1
@@ -117,7 +117,18 @@ if __name__ == "__main__":
117
args = parse_args()
118
warmup_iters, bench_iters = args.warmup_iters, args.bench_iters
119
120
- shapes = [(32, 128), (64, 256), (16, 512)]
120
+ shapes = [
121
+ (32, 128),
122
+ (64, 256),
123
+ (16, 512),
124
+ # Serving-sized activations, where this op is memory-bound and its
125
+ # schedule dominates: a prefill chunk and a decode batch at two
126
+ # common hidden sizes.
127
+ (1024, 6144),
128
+ (8192, 6144),
129
+ (1024, 7168),
130
+ (8192, 7168),
131
+ ]
132
group_sizes = [64, 128]
133
134
dtypes = ["fp8", "int8"] if args.dtype == "both" else [args.dtype]
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+215
-4
@@ -169,6 +169,156 @@ __global__ void per_token_group_quant_8bit_kernel(
169
#endif
170
}
171
172
+// Register-resident variant for the unpacked float32 scale layout, mirroring
173
+// per_token_group_quant_8bit_packed_register_kernel below, which float32-scale
174
+// callers cannot reach because it is hardwired to UE8M0 exponents packed
175
+// 4-per-int32 with TMA-aligned strides.
176
+//
177
+// A group_size-128 group of 16-bit input is 256 B, so it fits in registers and
178
+// the smem staging plus the __syncthreads() that the kernel above needs between
179
+// its absmax and quantize passes are both avoidable.
180
+//
181
+// Requires GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE, i.e. group_size 128 with
182
+// 16-bit input; CanUseRegisterPath enforces that.
183
+template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
184
+ bool SCALE_UE8M0 = false, int GROUP_SIZE = 128>
185
+__global__ void per_token_group_quant_8bit_register_kernel(
186
+ const T* __restrict__ input, void* __restrict__ output_q,
187
+ float* __restrict__ output_s, const int64_t num_groups, const float eps,
188
+ const float min_8bit, const float max_8bit, const int scale_num_rows,
189
+ const int scale_stride) {
190
+ constexpr int THREADS_PER_GROUP = 8;
191
+ constexpr int VEC_SIZE = 32 / sizeof(T);
192
+ static_assert(32 % THREADS_PER_GROUP == 0,
193
+ "THREADS_PER_GROUP must divide the warp size so the shuffle "
194
+ "mask stays inside one group");
195
+ static_assert(sizeof(DST_DTYPE) == 1,
196
+ "the uint4 store packs one byte per element");
197
+
198
+ // CanUseRegisterPath only routes 16-bit inputs here, but
199
+ // VLLM_STABLE_DISPATCH_FLOATING_TYPES instantiates every floating type, so
200
+ // gate the body on the shape this schedule can actually serve instead of a
201
+ // static_assert.
202
+ if constexpr (GROUP_SIZE != THREADS_PER_GROUP * VEC_SIZE) {
203
+ return;
204
+ } else {
205
+ const int64_t group_id =
206
+ (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) /
207
+ THREADS_PER_GROUP;
208
+
209
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
210
+ cudaGridDependencySynchronize();
211
+#endif
212
+
213
+ if (group_id >= num_groups) {
214
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
215
+ cudaTriggerProgrammaticLaunchCompletion();
216
+#endif
217
+ return;
218
+ }
219
+
220
+ const int lane_id = threadIdx.x % THREADS_PER_GROUP;
221
+
222
+ // alignas(16) is required for the uint4 reinterpret_cast: the default
223
+ // alignof for 16-bit T is 2.
224
+ alignas(16) T regs[VEC_SIZE];
225
+ {
226
+ const T* group_input = input + group_id * GROUP_SIZE + lane_id * VEC_SIZE;
227
+ uint4* dst = reinterpret_cast<uint4*>(®s[0]);
228
+ const uint4* src = reinterpret_cast<const uint4*>(group_input);
229
+ dst[0] = src[0];
230
+ dst[1] = src[1];
231
+ }
232
+
233
+ float local_absmax = eps;
234
+#pragma unroll
235
+ for (int i = 0; i < VEC_SIZE; ++i) {
236
+ local_absmax = fmaxf(local_absmax, fabsf(static_cast<float>(regs[i])));
237
+ }
238
+
239
+ // Reduce within the 8-lane subgroup that shares this group.
240
+#ifdef USE_ROCM
241
+ const int lane_in_wave = threadIdx.x % warpSize;
242
+ const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
243
+ local_absmax =
244
+ fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
245
+ local_absmax =
246
+ fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
247
+ local_absmax =
248
+ fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
249
+#else
250
+ const unsigned mask = 0xffu << (threadIdx.x & 24u);
251
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
252
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
253
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
254
+#endif
255
+
256
+ float y_s = local_absmax / max_8bit;
257
+ if constexpr (SCALE_UE8M0) {
258
+ // Bit-exact with exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f)))): bump the
259
+ // exponent whenever the mantissa is non-zero, then clear the mantissa.
260
+ y_s = fmaxf(fabsf(y_s), 1e-10f);
261
+ const uint32_t bits = __float_as_uint(y_s);
262
+ const uint32_t exp_bits = (bits >> 23) & 0xffu;
263
+ const uint32_t mant_bits = bits & 0x7fffffu;
264
+ y_s = __uint_as_float((exp_bits + (mant_bits != 0u ? 1u : 0u)) << 23);
265
+ }
266
+
267
+ if (lane_id == 0) {
268
+ if constexpr (IS_COLUMN_MAJOR) {
269
+ const int64_t row_idx = group_id / scale_num_rows;
270
+ const int64_t col_idx = group_id % scale_num_rows;
271
+ output_s[col_idx * scale_stride + row_idx] = y_s;
272
+ } else {
273
+ output_s[group_id] = y_s;
274
+ }
275
+ }
276
+
277
+ // A UE8M0 scale is an exact power of two, so multiplying by the reciprocal
278
+ // matches dividing. Otherwise the reciprocal costs an ulp and can move
279
+ // values that land exactly on an 8-bit boundary, so divide instead to stay
280
+ // bit-exact with the smem kernel.
281
+ const float inv_y = 1.0f / y_s;
282
+ uint32_t packed[VEC_SIZE / 4] = {};
283
+#pragma unroll
284
+ for (int i = 0; i < VEC_SIZE; ++i) {
285
+ const float v = static_cast<float>(regs[i]);
286
+ const float scaled = SCALE_UE8M0 ? (v * inv_y) : (v / y_s);
287
+ const float q = fminf(fmaxf(scaled, min_8bit), max_8bit);
288
+ DST_DTYPE qb = DST_DTYPE(q);
289
+ const uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
290
+ packed[i >> 2] |= static_cast<uint32_t>(byte) << ((i & 3) * 8);
291
+ }
292
+
293
+ DST_DTYPE* group_output = static_cast<DST_DTYPE*>(output_q) +
294
+ group_id * GROUP_SIZE + lane_id * VEC_SIZE;
295
+ *reinterpret_cast<uint4*>(group_output) =
296
+ make_uint4(packed[0], packed[1], packed[2], packed[3]);
297
+
298
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
299
+ cudaTriggerProgrammaticLaunchCompletion();
300
+#endif
301
+ }
302
+}
303
+
304
+// The register path needs group_size 128 with 16-bit input so that one group is
305
+// exactly THREADS_PER_GROUP(8) x VEC_SIZE(16) elements, and it loads/stores
306
+// through uint4, so both buffers must be 16 B aligned.
307
+inline bool CanUseRegisterPath(const torch::stable::Tensor& input,
308
+ const torch::stable::Tensor& output_q,
309
+ int64_t group_size) {
310
+ const auto dtype = input.scalar_type();
311
+ if (group_size != 128) return false;
312
+ if (dtype != torch::headeronly::ScalarType::Half &&
313
+ dtype != torch::headeronly::ScalarType::BFloat16) {
314
+ return false;
315
+ }
316
+ if (!input.is_contiguous() || !output_q.is_contiguous()) return false;
317
+ const auto in_addr = reinterpret_cast<uintptr_t>(input.data_ptr());
318
+ const auto out_addr = reinterpret_cast<uintptr_t>(output_q.data_ptr());
319
+ return (in_addr % 16 == 0) && (out_addr % 16 == 0);
320
+}
321
+
322
inline int GetGroupsPerBlock(int64_t num_groups) {
323
if (num_groups % 16 == 0) {
324
return 16;
@@ -235,6 +385,15 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
385
const int scale_num_rows = output_s.size(1);
386
const int scale_stride = output_s.stride(1);
387
388
+ const bool use_register_path =
389
+ CanUseRegisterPath(input, output_q, group_size);
390
+ constexpr int REG_THREADS_PER_GROUP = 8;
391
+ constexpr int REG_BLOCK_THREADS = 256;
392
+ const int64_t reg_groups_per_block =
393
+ REG_BLOCK_THREADS / REG_THREADS_PER_GROUP;
394
+ const int reg_num_blocks = static_cast<int>(
395
+ (num_groups + reg_groups_per_block - 1) / reg_groups_per_block);
396
+
397
#ifndef USE_ROCM
398
#define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \
399
SMEM_BYTES) \
@@ -281,21 +440,72 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
440
} \
441
} while (0)
442
443
+#ifndef USE_ROCM
444
+ #define LAUNCH_REGISTER_INST(T, DST_DTYPE, COL_MAJOR, UE8M0) \
445
+ do { \
446
+ cudaLaunchConfig_t config = {}; \
447
+ config.gridDim = dim3(reg_num_blocks); \
448
+ config.blockDim = dim3(REG_BLOCK_THREADS); \
449
+ config.dynamicSmemBytes = 0; \
450
+ config.stream = stream; \
451
+ cudaLaunchAttribute attrs[1]; \
452
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
453
+ attrs[0].val.programmaticStreamSerializationAllowed = 1; \
454
+ config.numAttrs = 1; \
455
+ config.attrs = attrs; \
456
+ cudaLaunchKernelEx( \
457
+ &config, \
458
+ per_token_group_quant_8bit_register_kernel<T, DST_DTYPE, COL_MAJOR, \
459
+ UE8M0>, \
460
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
461
+ static_cast<float*>(output_s.data_ptr()), \
462
+ static_cast<int64_t>(num_groups), (float)eps, (float)min_8bit, \
463
+ (float)max_8bit, scale_num_rows, scale_stride); \
464
+ } while (0)
465
+#else
466
+ #define LAUNCH_REGISTER_INST(T, DST_DTYPE, COL_MAJOR, UE8M0) \
467
+ do { \
468
+ per_token_group_quant_8bit_register_kernel<T, DST_DTYPE, COL_MAJOR, \
469
+ UE8M0> \
470
+ <<<dim3(reg_num_blocks), dim3(REG_BLOCK_THREADS), 0, stream>>>( \
471
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
472
+ static_cast<float*>(output_s.data_ptr()), \
473
+ static_cast<int64_t>(num_groups), (float)eps, (float)min_8bit, \
474
+ (float)max_8bit, scale_num_rows, scale_stride); \
475
+ } while (0)
476
+#endif
477
+
478
#define LAUNCH_KERNEL(T, DST_DTYPE) \
479
do { \
480
size_t smem_bytes = \
481
static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
482
if (is_column_major) { \
483
if (scale_ue8m0) { \
290
- LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \
484
+ if (use_register_path) { \
485
+ LAUNCH_REGISTER_INST(T, DST_DTYPE, true, true); \
486
+ } else { \
487
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \
488
+ } \
489
} else { \
292
- LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \
490
+ if (use_register_path) { \
491
+ LAUNCH_REGISTER_INST(T, DST_DTYPE, true, false); \
492
+ } else { \
493
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \
494
+ } \
495
} \
496
} else { \
497
if (scale_ue8m0) { \
296
- LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \
498
+ if (use_register_path) { \
499
+ LAUNCH_REGISTER_INST(T, DST_DTYPE, false, true); \
500
+ } else { \
501
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \
502
+ } \
503
} else { \
298
- LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \
504
+ if (use_register_path) { \
505
+ LAUNCH_REGISTER_INST(T, DST_DTYPE, false, false); \
506
+ } else { \
507
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \
508
+ } \
509
} \
510
} \
511
} while (0)
@@ -314,6 +524,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
524
#undef LAUNCH_KERNEL
525
#undef LAUNCH_KERNEL_INST
526
#undef LAUNCH_KERNEL_IMPL
527
+#undef LAUNCH_REGISTER_INST
528
}
529
530
// Register-resident fast path for group_size==128.
tests/kernels/quantization/test_per_token_group_quant.py
+68
-1
@@ -26,6 +26,8 @@ from vllm.platforms import current_platform
26
(16, 512),
27
(3, 2688),
28
(3, 5376),
29
+ (333, 128),
30
+ (512, 6144),
31
],
32
)
33
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -39,7 +41,7 @@ from vllm.platforms import current_platform
41
)
42
def test_per_token_group_quant_fp8(
43
shape,
42
- dtype,
44
+ dtype: torch.dtype,
45
column_major: bool,
46
tma_aligned: bool,
47
scale_ue8m0: bool,
@@ -78,6 +80,50 @@ def test_per_token_group_quant_fp8(
80
assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)
81
82
83
+def _misaligned_clone(x: torch.Tensor) -> torch.Tensor:
84
+ """Contiguous copy of ``x`` whose ``data_ptr`` is not 16B aligned."""
85
+ numel = x.numel()
86
+ base = torch.empty(numel + 16 // x.element_size(), device=x.device, dtype=x.dtype)
87
+ view = base[1 : 1 + numel]
88
+ if view.data_ptr() % 16 == 0:
89
+ pytest.skip("allocator did not yield a misaligned offset")
90
+ return view.view(x.shape).copy_(x)
91
+
92
+
93
+@pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)])
94
+@pytest.mark.parametrize("column_major", [False, True])
95
+@pytest.mark.parametrize("scale_ue8m0", [False, True])
96
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
97
+@pytest.mark.skipif(
98
+ not current_platform.is_cuda_alike(),
99
+ reason="register-resident schedule is CUDA/ROCm only",
100
+)
101
+def test_per_token_group_quant_fp8_register_matches_smem(
102
+ shape, column_major: bool, scale_ue8m0: bool, dtype: torch.dtype
103
+):
104
+ """group_size=128 with 16-bit input takes a register-resident schedule that
105
+ needs 16B-aligned buffers; a misaligned input falls back to the
106
+ shared-memory schedule. The two must agree bit-for-bit, which the
107
+ tolerance-based comparison against Triton above cannot establish.
108
+ """
109
+ device = current_platform.device_type
110
+ torch.manual_seed(42)
111
+ x = torch.randn(shape, device=device, dtype=dtype) * 8
112
+
113
+ reg_q, reg_s = fp8_utils.per_token_group_quant_fp8(
114
+ x, 128, column_major_scales=column_major, use_ue8m0=scale_ue8m0
115
+ )
116
+ smem_q, smem_s = fp8_utils.per_token_group_quant_fp8(
117
+ _misaligned_clone(x),
118
+ 128,
119
+ column_major_scales=column_major,
120
+ use_ue8m0=scale_ue8m0,
121
+ )
122
+
123
+ assert torch.equal(reg_q.view(torch.uint8), smem_q.view(torch.uint8))
124
+ assert torch.equal(reg_s, smem_s)
125
+
126
+
127
@pytest.mark.parametrize(
128
"num_tokens,hidden_dim,group_size",
129
[
@@ -472,3 +518,24 @@ def test_per_token_group_quant_int8(shape, dtype, group_size: int):
518
519
assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15)
520
assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)
521
+
522
+
523
+@pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)])
524
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
525
+@pytest.mark.skipif(
526
+ not current_platform.is_cuda_alike(),
527
+ reason="register-resident schedule is CUDA/ROCm only",
528
+)
529
+def test_per_token_group_quant_int8_register_matches_smem(shape, dtype: torch.dtype):
530
+ """int8 output shares the register-resident schedule with fp8, so it needs
531
+ the same bit-exactness guard against the shared-memory schedule.
532
+ """
533
+ device = current_platform.device_type
534
+ torch.manual_seed(42)
535
+ x = torch.randn(shape, device=device, dtype=dtype) * 8
536
+
537
+ reg_q, reg_s = int8_utils.per_token_group_quant_int8(x, 128)
538
+ smem_q, smem_s = int8_utils.per_token_group_quant_int8(_misaligned_clone(x), 128)
539
+
540
+ assert torch.equal(reg_q, smem_q)
541
+ assert torch.equal(reg_s, smem_s)