[Kernel] Enable PDL for per_token_group_quant_8bit_kernel (#46508)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jun 25, 2026 at 08:42 UTC
23aed9b0eedcad57266f7f9d7776b2c50ce40fb9
1 file changed
+62
-42
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+62
-42
@@ -119,6 +119,10 @@ __global__ void per_token_group_quant_8bit_kernel(
119
static_cast<DST_DTYPE*>(output_q) + block_group_offset;
120
scale_element_t* scale_output;
121
122
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
123
+ cudaGridDependencySynchronize();
124
+#endif
125
+
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));
@@ -153,6 +157,10 @@ __global__ void per_token_group_quant_8bit_kernel(
157
158
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
159
threads_per_group, y_s, min_8bit, max_8bit);
160
+
161
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
162
+ cudaTriggerProgrammaticLaunchCompletion();
163
+#endif
164
}
165
166
inline int GetGroupsPerBlock(int64_t num_groups) {
@@ -209,45 +217,56 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
217
const int scale_num_rows = output_s.size(1);
218
const int scale_stride = output_s.stride(1);
219
212
-#define LAUNCH_KERNEL(T, DST_DTYPE) \
213
- do { \
214
- dim3 grid(num_blocks); \
215
- dim3 block(num_threads); \
216
- size_t smem_bytes = \
217
- static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
218
- if (is_column_major) { \
219
- if (scale_ue8m0) { \
220
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, true> \
221
- <<<grid, block, smem_bytes, stream>>>( \
222
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
223
- static_cast<float*>(output_s.data_ptr()), group_size, \
224
- num_groups, groups_per_block, (float)eps, (float)min_8bit, \
225
- (float)max_8bit, scale_num_rows, scale_stride); \
226
- } else { \
227
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, false> \
228
- <<<grid, block, smem_bytes, stream>>>( \
229
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
230
- static_cast<float*>(output_s.data_ptr()), group_size, \
231
- num_groups, groups_per_block, (float)eps, (float)min_8bit, \
232
- (float)max_8bit, scale_num_rows, scale_stride); \
233
- } \
234
- } else { \
235
- if (scale_ue8m0) { \
236
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, true> \
237
- <<<grid, block, smem_bytes, stream>>>( \
238
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
239
- static_cast<float*>(output_s.data_ptr()), group_size, \
240
- num_groups, groups_per_block, (float)eps, (float)min_8bit, \
241
- (float)max_8bit); \
242
- } else { \
243
- per_token_group_quant_8bit_kernel<T, DST_DTYPE, false, false> \
244
- <<<grid, block, smem_bytes, stream>>>( \
245
- static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
246
- static_cast<float*>(output_s.data_ptr()), group_size, \
247
- num_groups, groups_per_block, (float)eps, (float)min_8bit, \
248
- (float)max_8bit); \
249
- } \
250
- } \
220
+#ifndef USE_ROCM
221
+ #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
222
+ do { \
223
+ cudaLaunchConfig_t config = {}; \
224
+ config.gridDim = dim3(num_blocks); \
225
+ config.blockDim = dim3(num_threads); \
226
+ config.dynamicSmemBytes = (SMEM_BYTES); \
227
+ config.stream = stream; \
228
+ cudaLaunchAttribute attrs[1]; \
229
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
230
+ attrs[0].val.programmaticStreamSerializationAllowed = 1; \
231
+ config.numAttrs = 1; \
232
+ config.attrs = attrs; \
233
+ cudaLaunchKernelEx( \
234
+ &config, \
235
+ per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
236
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
237
+ static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
238
+ groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
239
+ scale_num_rows, scale_stride); \
240
+ } while (0)
241
+#else
242
+ #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
243
+ do { \
244
+ per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0> \
245
+ <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
246
+ static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
247
+ static_cast<float*>(output_s.data_ptr()), group_size, \
248
+ num_groups, groups_per_block, (float)eps, (float)min_8bit, \
249
+ (float)max_8bit, scale_num_rows, scale_stride); \
250
+ } while (0)
251
+#endif
252
+
253
+#define LAUNCH_KERNEL(T, DST_DTYPE) \
254
+ do { \
255
+ size_t smem_bytes = \
256
+ static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
257
+ if (is_column_major) { \
258
+ if (scale_ue8m0) { \
259
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \
260
+ } else { \
261
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \
262
+ } \
263
+ } else { \
264
+ if (scale_ue8m0) { \
265
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \
266
+ } else { \
267
+ LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \
268
+ } \
269
+ } \
270
} while (0)
271
272
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
@@ -262,6 +281,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
281
}));
282
283
#undef LAUNCH_KERNEL
284
+#undef LAUNCH_KERNEL_INST
285
}
286
287
// Register-resident fast path for group_size==128.
@@ -306,12 +326,12 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
326
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
327
328
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
309
- asm volatile("griddepcontrol.wait;");
329
+ cudaGridDependencySynchronize();
330
#endif
331
332
if (mn_idx >= tma_aligned_mn) {
333
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
314
- asm volatile("griddepcontrol.launch_dependents;");
334
+ cudaTriggerProgrammaticLaunchCompletion();
335
#endif
336
return;
337
}
@@ -428,7 +448,7 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
448
*reinterpret_cast<uint4*>(group_output) = packed_out;
449
450
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
431
- asm volatile("griddepcontrol.launch_dependents;");
451
+ cudaTriggerProgrammaticLaunchCompletion();
452
#endif
453
}
454