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 4, 11:15 UTC
Faster per-token fp8 group quant packed kernel for blackwell (#41326)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Yongye Zhu committed
Apr 30, 2026 at 21:09 UTC
a3c83ff2fd050bc6392260d6e84ee1150f238f26
3 files changed
+366
-89
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+174
-84
@@ -236,17 +236,41 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
236
#undef LAUNCH_KERNEL
237
}
238
239
-template <typename T, typename DST_DTYPE>
240
-__global__ void per_token_group_quant_8bit_packed_kernel(
239
+// Register-resident fast path for group_size==128.
240
+//
241
+// Each thread holds 16 source elements (32 B = uint4 x 2) in registers across
242
+// the absmax reduce -> scale compute -> quantize pipeline. No shared memory.
243
+// UE8M0 scale extracted via bit math (bit-exact with exp2f(ceilf(log2f))).
244
+//
245
+// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell
246
+// nvcc fuses these into a single 256-bit LDG.E.256.
247
+//
248
+// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
249
+// 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>
251
+__global__ void per_token_group_quant_8bit_packed_register_kernel(
252
const T* __restrict__ input, void* __restrict__ output_q,
242
- unsigned int* __restrict__ output_s_packed, const int group_size,
243
- const int num_groups_padded, const int groups_per_block,
244
- const int padded_groups_per_row, const int groups_per_row, const int mn,
245
- const int tma_aligned_mn, const int num_scale_elems, const float eps,
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,
255
+ const int groups_per_row, const int mn, const int output_q_mn_extent,
256
+ const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
257
const float min_8bit, const float max_8bit) {
247
- const int threads_per_group = 16;
248
- const int64_t local_group_id = threadIdx.x / threads_per_group;
249
- const int lane_id = threadIdx.x % threads_per_group;
258
+ static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128");
259
+ constexpr int THREADS_PER_GROUP = 8;
260
+ constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
261
+ static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
262
+ "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).
268
+ static_assert(32 % THREADS_PER_GROUP == 0,
269
+ "THREADS_PER_GROUP must divide warp size for the shuffle "
270
+ "mask to be valid");
271
+
272
+ const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
273
+ const int lane_id = threadIdx.x % THREADS_PER_GROUP;
274
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;
@@ -254,141 +278,207 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
278
return;
279
}
280
257
- // map flat group id to 2D indices (mn_idx, sf_k_idx)
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);
261
-
262
- // whether it is a valid group (not padding)
284
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
285
265
- // shared memory to cache each group's data to avoid double DRAM reads.
266
- extern __shared__ __align__(16) char smem_raw[];
267
- T* smem = reinterpret_cast<T*>(smem_raw);
268
- T* smem_group = smem + local_group_id * group_size;
269
-
270
- // compute scale for valid groups
271
- float y_s = 0.f;
286
+ // Load 16 input elements (32 B) into registers as two adjacent uint4
287
+ // loads. nvcc keeps these as 2x LDG.E.128 on sm_100; the per-thread cost
288
+ // is dominated by HBM bandwidth at large MN, so a fused 256-bit load via
289
+ // inline PTX gave no measurable speedup.
290
+ // alignas(16) is required so the uint4* reinterpret_cast below is
291
+ // well-defined for T == bf16/fp16 (default alignof is 2).
292
+ alignas(16) T regs[VEC_SIZE];
293
+ float local_absmax = eps;
294
if (is_valid_group) {
295
const T* group_input =
274
- input + static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
275
- sf_k_idx * group_size;
276
- y_s = ComputeGroupScale<T, true>(group_input, smem_group, group_size,
277
- lane_id, threads_per_group, eps, max_8bit);
296
+ input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
297
+ sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
298
+ uint4* dst = reinterpret_cast<uint4*>(®s[0]);
299
+ const uint4* src = reinterpret_cast<const uint4*>(group_input);
300
+ dst[0] = src[0];
301
+ dst[1] = src[1];
302
+#pragma unroll
303
+ for (int i = 0; i < VEC_SIZE; ++i) {
304
+ float v = fabsf(static_cast<float>(regs[i]));
305
+ local_absmax = fmaxf(local_absmax, v);
306
+ }
307
}
308
280
- // pack 4 scales into a uint32 exponent
309
+ // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
310
+ // 8 lanes within the warp that share a group.
311
+ unsigned mask = 0xffu << (threadIdx.x & 24u);
312
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
313
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
314
+ local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
315
+
316
+ float y_s = local_absmax / max_8bit;
317
+ y_s = fmaxf(y_s, 1e-10f);
318
+ uint32_t bits = __float_as_uint(y_s);
319
+ uint32_t exp_bits = (bits >> 23) & 0xffu;
320
+ uint32_t mant_bits = bits & 0x7fffffu;
321
+ uint8_t exp_byte =
322
+ static_cast<uint8_t>(exp_bits + (mant_bits != 0u ? 1u : 0u));
323
+
324
+ // Lane 0 writes the packed scale byte.
325
if (lane_id == 0) {
282
- // each uint32 in output_s_packed stores 4 packed scales
326
const int sf_k_pack_idx = sf_k_idx / 4;
327
const int pos = sf_k_idx % 4;
328
const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
286
-
329
if (is_valid_group) {
288
- // reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit
289
- // exponent, and place it into the correct byte of the 32-bit word.
290
- const unsigned int bits = __float_as_uint(y_s);
291
- const uint8_t exponent = static_cast<uint8_t>((bits >> 23u) & 0xffu);
292
- reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exponent;
330
+ reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exp_byte;
331
} else if (out_idx < num_scale_elems) {
294
- // write zero for padding groups if within bounds of output_s_packed
332
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = 0;
333
}
334
}
335
299
- __syncthreads();
336
+ // For padded mn rows that fall within output_q's allocated extent, write
337
+ // a uint4 of zeros to keep the buffer clean for downstream TMA loads.
338
+ // Skip writes for sf_k padding (those positions don't exist in output_q).
339
+ if (!is_valid_group) {
340
+ if (sf_k_idx < groups_per_row && mn_idx >= mn &&
341
+ mn_idx < output_q_mn_extent) {
342
+ DST_DTYPE* group_output =
343
+ static_cast<DST_DTYPE*>(output_q) +
344
+ static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
345
+ sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
346
+ *reinterpret_cast<uint4*>(group_output) = make_uint4(0, 0, 0, 0);
347
+ }
348
+ return;
349
+ }
350
301
- if (is_valid_group) {
302
- DST_DTYPE* group_output =
303
- static_cast<DST_DTYPE*>(output_q) +
304
- static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
305
- sf_k_idx * group_size;
306
- QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
307
- threads_per_group, y_s, min_8bit, max_8bit);
351
+ // Reconstruct y_s as a power-of-2 float and use its reciprocal.
352
+ float y_s_q = __uint_as_float(static_cast<uint32_t>(exp_byte) << 23);
353
+ float inv_y = 1.0f / y_s_q;
354
+
355
+ // Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we
356
+ // fill four 32-bit words, four bytes each.
357
+ uint32_t packed_lo = 0;
358
+ uint32_t packed_lo_hi = 0;
359
+ uint32_t packed_hi_lo = 0;
360
+ uint32_t packed_hi = 0;
361
+#pragma unroll
362
+ for (int i = 0; i < VEC_SIZE; ++i) {
363
+ float q =
364
+ fminf(fmaxf(static_cast<float>(regs[i]) * inv_y, min_8bit), max_8bit);
365
+ DST_DTYPE qb = DST_DTYPE(q);
366
+ uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
367
+ const int shift = (i & 3) * 8;
368
+ if (i < 4) {
369
+ packed_lo |= static_cast<uint32_t>(byte) << shift;
370
+ } else if (i < 8) {
371
+ packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
372
+ } else if (i < 12) {
373
+ packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
374
+ } else {
375
+ packed_hi |= static_cast<uint32_t>(byte) << shift;
376
+ }
377
}
378
+
379
+ uint4 packed_out =
380
+ make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
381
+ DST_DTYPE* group_output =
382
+ static_cast<DST_DTYPE*>(output_q) +
383
+ static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
384
+ sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
385
+ *reinterpret_cast<uint4*>(group_output) = packed_out;
386
}
387
388
+// Public entry point: register-resident packed quant kernel.
389
+// Constraints: group_size == 128 and bf16/fp16 input.
390
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
391
torch::stable::Tensor& output_q,
392
torch::stable::Tensor& output_s_packed,
393
int64_t group_size, double eps,
394
double min_8bit, double max_8bit) {
395
+ STD_TORCH_CHECK(group_size == 128,
396
+ "per_token_group_quant_8bit_packed only supports "
397
+ "group_size==128, got ",
398
+ group_size, ".");
399
+ const auto in_dtype = input.scalar_type();
400
+ STD_TORCH_CHECK(
401
+ in_dtype == torch::headeronly::ScalarType::Half ||
402
+ in_dtype == torch::headeronly::ScalarType::BFloat16,
403
+ "per_token_group_quant_8bit_packed only supports bf16/fp16 input.");
404
+
405
STD_TORCH_CHECK(input.is_contiguous());
406
STD_TORCH_CHECK(output_q.is_contiguous());
407
408
const int64_t k = input.size(-1);
320
- STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
321
- ") must be divisible by group_size (", group_size, ").");
409
+ STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k,
410
+ " is not divisible by group_size=", group_size, ".");
411
412
const int64_t mn = input.numel() / k;
413
const int64_t groups_per_row = k / group_size;
325
-
326
- STD_TORCH_CHECK(output_s_packed.dim() == 2,
327
- "output_s_packed must be 2D, got dim=", output_s_packed.dim(),
328
- ".");
329
-
414
const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4;
415
const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4;
416
417
+ // output_q may be allocated with extra padded mn rows (e.g.,
418
+ // (tma_aligned_mn, k)) so the kernel can zero-fill them in-line and the
419
+ // caller can use torch.empty instead of torch.zeros. The grid only covers
420
+ // up to tma_aligned_mn, so we cap the extent there.
421
+ const int64_t output_q_mn_actual = output_q.numel() / k;
422
+ STD_TORCH_CHECK(output_q_mn_actual >= mn,
423
+ "output_q must have at least mn rows; got ",
424
+ output_q_mn_actual, " rows for mn=", mn, ".");
425
+ const int64_t output_q_mn_extent =
426
+ output_q_mn_actual < tma_aligned_mn ? output_q_mn_actual : tma_aligned_mn;
427
+
428
STD_TORCH_CHECK(
429
output_s_packed.scalar_type() == torch::headeronly::ScalarType::Int,
335
- "output_s_packed must have dtype int32 for UE8M0-packed scales.");
336
- // DeepGEMM expects SFA scales in MN-major form with shape
337
- // [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last
338
- // dimension.
430
+ "output_s_packed must be int32 for UE8M0-packed scales.");
431
STD_TORCH_CHECK(output_s_packed.size(0) == mn &&
432
output_s_packed.size(1) == k_num_packed_sfk,
433
"output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
342
- "], but got [", output_s_packed.size(0), ", ",
434
+ "]; got [", output_s_packed.size(0), ", ",
435
output_s_packed.size(1), "].");
344
- // Verify column-major TMA-aligned layout
436
STD_TORCH_CHECK(output_s_packed.stride(0) == 1 &&
437
output_s_packed.stride(1) == tma_aligned_mn,
347
- "output_s_packed must have strides [1, ", tma_aligned_mn,
348
- "], but got [", output_s_packed.stride(0), ", ",
438
+ "output_s_packed strides must be [1, ", tma_aligned_mn,
439
+ "]; got [", output_s_packed.stride(0), ", ",
440
output_s_packed.stride(1), "].");
441
442
cudaStream_t stream = get_current_cuda_stream();
443
353
- constexpr int THREADS_PER_GROUP = 16;
354
-
355
- // Expand the grid to cover MN and K padding so every byte in
356
- // output_s_packed is written (padding bytes get zeroed by the kernel).
444
+ constexpr int THREADS_PER_GROUP = 8;
445
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;
359
- // Number of elements in output_s_packed.
447
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
361
-
448
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
449
450
auto dst_type = output_q.scalar_type();
365
- const int num_blocks = num_groups_padded / groups_per_block;
451
+ const int64_t num_blocks = num_groups_padded / groups_per_block;
452
const int num_threads = groups_per_block * THREADS_PER_GROUP;
367
-
368
-#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \
369
- do { \
370
- dim3 grid(num_blocks); \
371
- dim3 block(num_threads); \
372
- size_t smem_bytes = \
373
- static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
374
- per_token_group_quant_8bit_packed_kernel<T, DST_DTYPE> \
375
- <<<grid, block, smem_bytes, stream>>>( \
376
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
377
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
378
- static_cast<int>(group_size), static_cast<int>(num_groups_padded), \
379
- groups_per_block, static_cast<int>(padded_groups_per_row), \
380
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
381
- static_cast<int>(tma_aligned_mn), \
382
- static_cast<int>(num_scale_elems), static_cast<float>(eps), \
383
- static_cast<float>(min_8bit), static_cast<float>(max_8bit)); \
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)); \
474
} while (0)
475
386
- VLLM_STABLE_DISPATCH_FLOATING_TYPES(
387
- input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] {
476
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
477
+ input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
478
if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
389
- LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3);
479
+ LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
480
} else if (dst_type == torch::headeronly::ScalarType::Char) {
391
- LAUNCH_PACKED_KERNEL(scalar_t, int8_t);
481
+ LAUNCH_REG_KERNEL(scalar_t, int8_t);
482
} else {
483
STD_TORCH_CHECK(
484
false,
@@ -397,7 +487,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
487
}
488
}));
489
400
-#undef LAUNCH_PACKED_KERNEL
490
+#undef LAUNCH_REG_KERNEL
491
}
492
493
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h
+10
@@ -8,3 +8,13 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
8
torch::stable::Tensor& output_s,
9
int64_t group_size, double eps, double min_8bit,
10
double max_8bit, bool scale_ue8m0 = false);
11
+
12
+// Public op: register-resident packed quant for the DeepGEMM Blackwell path.
13
+// Restricted to group_size == 128 and bf16/fp16 input; other configurations
14
+// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because
15
+// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes.
16
+void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
17
+ torch::stable::Tensor& output_q,
18
+ torch::stable::Tensor& output_s_packed,
19
+ int64_t group_size, double eps,
20
+ double min_8bit, double max_8bit);
tests/kernels/quantization/test_per_token_group_quant.py
+182
-5
@@ -73,11 +73,6 @@ def test_per_token_group_quant_fp8(
73
# Larger shapes with padding
74
(127, 7168, 128),
75
(253, 640, 128),
76
- # Non-power-of-2 group size
77
- (4, 768, 96), # 768/96=8 groups, no padding
78
- (3, 768, 96), # 768/96=8 groups, MN padding
79
- (4, 480, 96), # 480/96=5 groups, K padding
80
- (1, 480, 96), # both MN and K padding
76
],
77
)
78
@pytest.mark.parametrize("poisoned_scales", [False, True])
@@ -161,6 +156,188 @@ def test_per_token_group_quant_fp8_packed(
156
)
157
158
159
+@pytest.mark.skipif(
160
+ not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
161
+)
162
+def test_per_token_group_quant_fp8_packed_all_zero():
163
+ """All-zero input must produce well-defined UE8M0 scale bytes via the eps
164
+ floor in the kernel's UE8M0 path. Locks down the all-zero behavior before
165
+ optimization.
166
+
167
+ The CUDA kernel computes:
168
+ y_s = eps / fp8_max
169
+ y_s = exp2(ceil(log2(fmax(y_s, 1e-10))))
170
+ For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to
171
+ 1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94).
172
+ """
173
+
174
+ device = "cuda"
175
+ num_tokens, hidden_dim, group_size = 4, 7168, 128
176
+ x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16)
177
+
178
+ out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
179
+ x,
180
+ group_size=group_size,
181
+ use_ue8m0=True,
182
+ )
183
+
184
+ # Quantized values must be all zero.
185
+ assert torch.equal(
186
+ out_q.view(torch.uint8),
187
+ torch.zeros_like(out_q, dtype=torch.uint8),
188
+ ), "All-zero input should produce all-zero FP8 output"
189
+
190
+ # UE8M0 byte produced by the kernel for all-zero input.
191
+ # The kernel's inner fmax(y_s, 1e-10) clamps eps/fp8_max back to 1e-10.
192
+ # 1e-10 as float32 has biased exponent 0x5D and a non-zero mantissa, so
193
+ # the kernel's bit-twiddle (exp_bits + (mant_bits != 0)) rounds up to
194
+ # 0x5E. This matches exp2(ceil(log2(1e-10))) = exp2(-33).
195
+ expected_exp_byte = 0x5E
196
+
197
+ mn = num_tokens
198
+ groups_per_row = hidden_dim // group_size
199
+ k_num_packed = (groups_per_row + 3) // 4
200
+ tma_aligned_mn = ((mn + 3) // 4) * 4
201
+ num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
202
+
203
+ # All valid scale slots must contain the expected packed value.
204
+ # Padding slots must be zero.
205
+ actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
206
+
207
+ expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
208
+ for row in range(mn):
209
+ for g in range(groups_per_row):
210
+ pack_col = g // 4
211
+ pos = g % 4
212
+ idx = pack_col * tma_aligned_mn + row
213
+ expected[idx] |= expected_exp_byte << (pos * 8)
214
+
215
+ assert torch.equal(actual, expected), "All-zero scale bytes mismatch"
216
+
217
+
218
+@pytest.mark.skipif(
219
+ not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
220
+)
221
+def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
222
+ """Inputs whose absmax/max_8bit produces a non-power-of-2 force the
223
+ mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior
224
+ before optimization."""
225
+
226
+ device = "cuda"
227
+ num_tokens, hidden_dim, group_size = 4, 7168, 128
228
+
229
+ # Build a tensor whose per-group absmax = 1.5 * fp8_max * 2^k for various k.
230
+ # fp8_max = torch.finfo(torch.float8_e4m3fn).max = 448.0.
231
+ # Then absmax/fp8_max = 1.5 * 2^k -> non-zero mantissa, triggers ceil
232
+ # rounding to 2^(k+1). Use k=0 for simplicity; the bf16 representation of
233
+ # 1.5*448=672.0 is exact.
234
+ x = torch.full(
235
+ (num_tokens, hidden_dim),
236
+ 672.0,
237
+ device=device,
238
+ dtype=torch.bfloat16,
239
+ )
240
+
241
+ out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
242
+ x,
243
+ group_size=group_size,
244
+ use_ue8m0=True,
245
+ )
246
+
247
+ with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
248
+ ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
249
+ x,
250
+ group_size,
251
+ use_ue8m0=True,
252
+ )
253
+
254
+ assert torch.equal(out_q, ref_q), "Quantized output mismatch"
255
+
256
+ mn = num_tokens
257
+ groups_per_row = hidden_dim // group_size
258
+ k_num_packed = (groups_per_row + 3) // 4
259
+ tma_aligned_mn = ((mn + 3) // 4) * 4
260
+ num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
261
+
262
+ ref_s_flat = ref_s.reshape(mn, groups_per_row)
263
+ ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF
264
+ expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
265
+ for row in range(mn):
266
+ for g in range(groups_per_row):
267
+ pack_col = g // 4
268
+ pos = g % 4
269
+ idx = pack_col * tma_aligned_mn + row
270
+ expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)
271
+
272
+ actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
273
+ assert torch.equal(actual, expected), "Scale bytes mismatch"
274
+
275
+
276
+@pytest.mark.parametrize(
277
+ "num_tokens,hidden_dim",
278
+ [
279
+ (1, 7168), # mn padded 1 -> 4
280
+ (2, 7168), # mn padded 2 -> 4
281
+ (3, 7168), # mn padded 3 -> 4
282
+ (5, 7168), # mn padded 5 -> 8
283
+ (127, 7168), # mn padded 127 -> 128
284
+ (253, 640), # both mn and groups padded
285
+ (1, 384), # extreme: 1 group, 1 mn row -> both axes padded
286
+ ],
287
+)
288
+@pytest.mark.skipif(
289
+ not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"
290
+)
291
+def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
292
+ num_tokens, hidden_dim
293
+):
294
+ """When output_q is allocated with shape (tma_aligned_mn, k) instead of
295
+ (mn, k), the kernel must overwrite the padded mn rows with zeros so
296
+ callers can use ``torch.empty`` instead of ``torch.zeros``."""
297
+
298
+ device = "cuda"
299
+ group_size = 128
300
+ torch.manual_seed(42)
301
+ x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
302
+
303
+ mn = num_tokens
304
+ groups_per_row = hidden_dim // group_size
305
+ k_num_packed = (groups_per_row + 3) // 4
306
+ tma_aligned_mn = ((mn + 3) // 4) * 4
307
+
308
+ fp8_dtype = torch.float8_e4m3fn
309
+ finfo = torch.finfo(fp8_dtype)
310
+ # Allocate output_q with the padded mn extent and pre-fill with 0xFF
311
+ # so the kernel cannot rely on a clean buffer.
312
+ out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype)
313
+ out_q.view(torch.uint8).fill_(0xFF)
314
+
315
+ out_s_packed = torch.empty_strided(
316
+ (mn, k_num_packed),
317
+ (1, tma_aligned_mn),
318
+ device=device,
319
+ dtype=torch.int32,
320
+ )
321
+
322
+ torch.ops._C.per_token_group_fp8_quant_packed(
323
+ x, out_q, out_s_packed, group_size, 1e-10, finfo.min, finfo.max
324
+ )
325
+
326
+ # Live rows must match the Triton reference.
327
+ with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
328
+ ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True)
329
+ assert torch.equal(out_q[:mn], ref_q), "Live region mismatch"
330
+
331
+ # Padded rows must be all-zero; without this, downstream TMA loads would
332
+ # see uninitialised data.
333
+ if tma_aligned_mn > mn:
334
+ padded_bytes = out_q[mn:tma_aligned_mn].view(torch.uint8)
335
+ assert padded_bytes.eq(0).all(), (
336
+ f"Padded rows [{mn}, {tma_aligned_mn}) not zeroed; "
337
+ f"{padded_bytes.ne(0).sum().item()} non-zero bytes"
338
+ )
339
+
340
+
341
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
342
@pytest.mark.parametrize("group_size", [64, 128])
343
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")