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 5, 07:15 UTC
[Perf] Fuse Zero Initializer for FP8 DeepGemm Block Quant Kernel (#39547)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Apr 11, 2026 at 10:16 UTC
59b2f7b64007c6eae7d8ec82cbf1f33500d1ad2d
2 files changed
+180
-49
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+70
-49
@@ -240,8 +240,9 @@ template <typename T, typename DST_DTYPE>
240
__global__ void per_token_group_quant_8bit_packed_kernel(
241
const T* __restrict__ input, void* __restrict__ output_q,
242
unsigned int* __restrict__ output_s_packed, const int group_size,
243
- const int num_groups, const int groups_per_block, const int groups_per_row,
244
- const int mn, const int tma_aligned_mn, const float eps,
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,
246
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,51 +250,62 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
250
251
const int64_t block_group_id = blockIdx.x * groups_per_block;
252
const int64_t global_group_id = block_group_id + local_group_id;
252
- if (global_group_id >= num_groups) {
253
+ if (global_group_id >= num_groups_padded) {
254
return;
255
}
256
256
- const int64_t block_group_offset = global_group_id * group_size;
257
+ // map flat group id to 2D indices (mn_idx, sf_k_idx)
258
+ const int sf_k_idx =
259
+ static_cast<int>(global_group_id % padded_groups_per_row);
260
+ const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
261
258
- const T* group_input = input + block_group_offset;
259
- DST_DTYPE* group_output =
260
- static_cast<DST_DTYPE*>(output_q) + block_group_offset;
262
+ // whether it is a valid group (not padding)
263
+ const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
264
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;
266
- const float y_s =
267
- ComputeGroupScale<T, true>(group_input, smem_group, group_size, lane_id,
268
- threads_per_group, eps, max_8bit);
269
270
- // pack 4 scales into a uint32
271
- if (lane_id == 0) {
272
- // map flat group id to 2D indices (mn_idx, sf_k_idx)
273
- const int sf_k_idx = static_cast<int>(global_group_id % groups_per_row);
274
- const int mn_idx = static_cast<int>(global_group_id / groups_per_row);
270
+ // compute scale for valid groups
271
+ float y_s = 0.f;
272
+ if (is_valid_group) {
273
+ 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);
278
+ }
279
276
- if (mn_idx < mn) {
277
- // each uint32 in output_s_packed stores 4 packed scales
278
- const int sf_k_pack_idx = sf_k_idx / 4;
279
- const int pos = sf_k_idx % 4;
280
+ // pack 4 scales into a uint32 exponent
281
+ if (lane_id == 0) {
282
+ // each uint32 in output_s_packed stores 4 packed scales
283
+ const int sf_k_pack_idx = sf_k_idx / 4;
284
+ const int pos = sf_k_idx % 4;
285
+ const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
286
287
+ 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);
284
- const unsigned int exponent = (bits >> 23u) & 0xffu;
285
- const unsigned int contrib = exponent << (pos * 8u);
286
-
287
- const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
288
- // atomically OR 8-bit exponent into the packed scales buffer
289
- atomicOr(output_s_packed + out_idx, contrib);
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;
293
+ } else if (out_idx < num_scale_elems) {
294
+ // write zero for padding groups if within bounds of output_s_packed
295
+ reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = 0;
296
}
297
}
298
299
__syncthreads();
300
295
- QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
296
- threads_per_group, y_s, min_8bit, max_8bit);
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);
308
+ }
309
}
310
311
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
@@ -310,7 +322,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
322
323
const int64_t mn = input.numel() / k;
324
const int64_t groups_per_row = k / group_size;
313
- const int64_t num_groups = mn * groups_per_row;
325
326
STD_TORCH_CHECK(output_s_packed.dim() == 2,
327
"output_s_packed must be 2D, got dim=", output_s_packed.dim(),
@@ -330,36 +341,46 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
341
"output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
342
"], but got [", output_s_packed.size(0), ", ",
343
output_s_packed.size(1), "].");
344
+ // Verify column-major TMA-aligned layout
345
+ STD_TORCH_CHECK(output_s_packed.stride(0) == 1 &&
346
+ 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), ", ",
349
+ output_s_packed.stride(1), "].");
350
351
cudaStream_t stream = get_current_cuda_stream();
352
353
constexpr int THREADS_PER_GROUP = 16;
354
338
- const int groups_per_block = GetGroupsPerBlock(num_groups);
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).
357
+ const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
358
+ const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
359
+ // Number of elements in output_s_packed.
360
+ const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
361
+
362
+ const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
363
364
auto dst_type = output_q.scalar_type();
341
- const int num_blocks = num_groups / groups_per_block;
365
+ const int num_blocks = num_groups_padded / groups_per_block;
366
const int num_threads = groups_per_block * THREADS_PER_GROUP;
367
344
- // zero-initialize packed scales, since we use atomicOr to accumulate
345
- // exponents from different groups.
346
- torch::stable::zero_(output_s_packed);
347
-
348
-#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \
349
- do { \
350
- dim3 grid(num_blocks); \
351
- dim3 block(num_threads); \
352
- size_t smem_bytes = \
353
- static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
354
- per_token_group_quant_8bit_packed_kernel<T, DST_DTYPE> \
355
- <<<grid, block, smem_bytes, stream>>>( \
356
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
357
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
358
- static_cast<int>(group_size), static_cast<int>(num_groups), \
359
- groups_per_block, static_cast<int>(groups_per_row), \
360
- static_cast<int>(mn), static_cast<int>(tma_aligned_mn), \
361
- static_cast<float>(eps), static_cast<float>(min_8bit), \
362
- static_cast<float>(max_8bit)); \
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)); \
384
} while (0)
385
386
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
tests/kernels/quantization/test_per_token_group_quant.py
+110
@@ -48,6 +48,116 @@ def test_per_token_group_quant_fp8(
48
assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)
49
50
51
+@pytest.mark.parametrize(
52
+ "num_tokens,hidden_dim,group_size",
53
+ [
54
+ # No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4)
55
+ (4, 7168, 128),
56
+ # MN padding only: mn=1, tma_aligned_mn=4
57
+ (1, 7168, 128),
58
+ # MN padding only: mn=3, tma_aligned_mn=4
59
+ (3, 7168, 128),
60
+ # K padding only: groups_per_row=5 (5%4=1)
61
+ (4, 640, 128),
62
+ # K padding only: groups_per_row=6 (6%4=2)
63
+ (4, 768, 128),
64
+ # Single packed column, no padding: k_num_packed=1, mn%4=0
65
+ (4, 384, 128),
66
+ # Both MN and K padding
67
+ (1, 384, 128),
68
+ (3, 640, 128),
69
+ # Larger shapes with no padding
70
+ (64, 7168, 128),
71
+ (128, 14336, 128),
72
+ # Larger shapes with padding
73
+ (127, 7168, 128),
74
+ (253, 640, 128),
75
+ # Non-power-of-2 group size
76
+ (4, 768, 96), # 768/96=8 groups, no padding
77
+ (3, 768, 96), # 768/96=8 groups, MN padding
78
+ (4, 480, 96), # 480/96=5 groups, K padding
79
+ (1, 480, 96), # both MN and K padding
80
+ ],
81
+)
82
+@pytest.mark.parametrize("poisoned_scales", [False, True])
83
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
84
+def test_per_token_group_quant_fp8_packed(
85
+ num_tokens, hidden_dim, group_size, poisoned_scales
86
+):
87
+ """Test the packed DeepGEMM quantization kernel against the Triton
88
+ reference (row-major, UE8M0 scales)."""
89
+
90
+ device = "cuda"
91
+ torch.manual_seed(42)
92
+
93
+ x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
94
+
95
+ mn = num_tokens
96
+ groups_per_row = hidden_dim // group_size
97
+ k_num_packed = (groups_per_row + 3) // 4
98
+ tma_aligned_mn = ((mn + 3) // 4) * 4
99
+ num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
100
+
101
+ if poisoned_scales:
102
+ # Call the kernel with poisoned scale buffer to
103
+ # ensure padded indices are correctly zeroed.
104
+ fp8_dtype = torch.float8_e4m3fn
105
+ finfo = torch.finfo(fp8_dtype)
106
+ out_q = torch.empty_like(x, dtype=fp8_dtype)
107
+ out_s_packed = torch.empty_strided(
108
+ (mn, k_num_packed),
109
+ (1, tma_aligned_mn),
110
+ device=device,
111
+ dtype=torch.int32,
112
+ )
113
+ torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F)
114
+ torch.ops._C.per_token_group_fp8_quant_packed(
115
+ x,
116
+ out_q,
117
+ out_s_packed,
118
+ group_size,
119
+ 1e-10,
120
+ finfo.min,
121
+ finfo.max,
122
+ )
123
+ else:
124
+ out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
125
+ x,
126
+ group_size=group_size,
127
+ use_ue8m0=True,
128
+ )
129
+
130
+ # Triton reference (row-major float32 scales, UE8M0)
131
+ with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
132
+ ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
133
+ x,
134
+ group_size,
135
+ use_ue8m0=True,
136
+ )
137
+
138
+ # Quantized values must match.
139
+ assert torch.equal(out_q, ref_q), "Quantized output mismatch"
140
+
141
+ # Verify packed scales (valid exponents + padding zeros).
142
+ ref_s_flat = ref_s.reshape(mn, groups_per_row)
143
+ ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF
144
+
145
+ expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
146
+ for row in range(mn):
147
+ for g in range(groups_per_row):
148
+ pack_col = g // 4
149
+ pos = g % 4
150
+ idx = pack_col * tma_aligned_mn + row
151
+ expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)
152
+
153
+ actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
154
+ assert torch.equal(actual, expected), (
155
+ f"Packed scale storage mismatch.\n"
156
+ f"First diff at index "
157
+ f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}"
158
+ )
159
+
160
+
161
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
162
@pytest.mark.parametrize("group_size", [64, 128])
163
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")