[Kernel] Add PDL support for DeepGEMM kernel (#46006)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jun 18, 2026 at 20:49 UTC
22cc891108b1721959a4e346665b4c9cdddd3fb0
3 files changed
+87
-23
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+60
-17
@@ -304,9 +304,17 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
304
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
305
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
306
307
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
308
+ asm volatile("griddepcontrol.wait;");
309
+#endif
310
+
311
if (mn_idx >= tma_aligned_mn) {
312
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
313
+ asm volatile("griddepcontrol.launch_dependents;");
314
+#endif
315
return;
316
}
317
+
318
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
319
320
// Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -417,6 +425,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
425
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
426
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
427
*reinterpret_cast<uint4*>(group_output) = packed_out;
428
+
429
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
430
+ asm volatile("griddepcontrol.launch_dependents;");
431
+#endif
432
}
433
434
// Public entry point: register-resident packed quant kernel.
@@ -495,23 +507,54 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
507
508
auto dst_type = output_q.scalar_type();
509
498
-#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
499
- do { \
500
- dim3 grid(static_cast<unsigned int>(blocks_x), \
501
- static_cast<unsigned int>(blocks_y)); \
502
- dim3 block(num_threads); \
503
- per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
504
- RY> \
505
- <<<grid, block, 0, stream>>>( \
506
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
507
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
508
- static_cast<int>(padded_groups_per_row), \
509
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
510
- static_cast<int>(output_q_mn_extent), \
511
- static_cast<int>(tma_aligned_mn), num_scale_elems, \
512
- static_cast<float>(eps), static_cast<float>(min_8bit), \
513
- static_cast<float>(max_8bit)); \
514
- } while (0)
510
+// PDL (Programmatic Dependent Launch) is NVIDIA-only; ROCm/HIP has no
511
+// equivalent launch attribute, so fall back to a classic launch there.
512
+#ifndef USE_ROCM
513
+ #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
514
+ do { \
515
+ cudaLaunchConfig_t config = {}; \
516
+ config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
517
+ static_cast<unsigned int>(blocks_y)); \
518
+ config.blockDim = dim3(num_threads); \
519
+ config.dynamicSmemBytes = 0; \
520
+ config.stream = stream; \
521
+ cudaLaunchAttribute attrs[1]; \
522
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
523
+ attrs[0].val.programmaticStreamSerializationAllowed = 1; \
524
+ config.numAttrs = 1; \
525
+ config.attrs = attrs; \
526
+ cudaLaunchKernelEx( \
527
+ &config, \
528
+ per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, \
529
+ KX, RY>, \
530
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
531
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
532
+ static_cast<int>(padded_groups_per_row), \
533
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
534
+ static_cast<int>(output_q_mn_extent), \
535
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
536
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
537
+ static_cast<float>(max_8bit)); \
538
+ } while (0)
539
+#else
540
+ #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
541
+ do { \
542
+ dim3 grid(static_cast<unsigned int>(blocks_x), \
543
+ static_cast<unsigned int>(blocks_y)); \
544
+ dim3 block(num_threads); \
545
+ per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
546
+ RY> \
547
+ <<<grid, block, 0, stream>>>( \
548
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
549
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
550
+ static_cast<int>(padded_groups_per_row), \
551
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
552
+ static_cast<int>(output_q_mn_extent), \
553
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
554
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
555
+ static_cast<float>(max_8bit)); \
556
+ } while (0)
557
+#endif
558
559
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
560
do { \
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+8
-6
@@ -37,6 +37,8 @@ def _fused_inv_rope_fp8_quant_per_head(
37
ROPE_START: tl.constexpr,
38
HALF_ROPE: tl.constexpr,
39
TMA_ALIGNED_SCALES: tl.constexpr,
40
+ USE_GDC: tl.constexpr,
41
+ launch_pdl: tl.constexpr, # triton metadata
42
):
43
# int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
44
pid_token = tl.program_id(0).to(tl.int64)
@@ -46,7 +48,9 @@ def _fused_inv_rope_fp8_quant_per_head(
48
head_in_group = pid_gh % heads_per_group
49
global_head = pid_gh
50
qb_start = head_in_group * CHUNKS_PER_HEAD
49
-
51
+ if USE_GDC:
52
+ tl.extra.cuda.gdc_launch_dependents()
53
+ tl.extra.cuda.gdc_wait()
54
# Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
55
if pid_token >= num_tokens:
56
if TMA_ALIGNED_SCALES:
@@ -243,11 +247,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
247
(scale_inner * tma_aligned_T, 1, tma_aligned_T),
248
)
249
grid = (tma_aligned_T, n_groups * heads_per_group)
246
- pdl_kwargs = (
247
- {}
248
- if current_platform.is_rocm() or current_platform.is_xpu()
249
- else {"launch_pdl": False}
250
- )
250
+ use_gdc = current_platform.is_arch_support_pdl()
251
+ pdl_kwargs = {"launch_pdl": True} if use_gdc else {}
252
_fused_inv_rope_fp8_quant_per_head[grid](
253
o,
254
positions,
@@ -270,6 +271,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
271
ROPE_START=rope_start,
272
HALF_ROPE=half_rope,
273
TMA_ALIGNED_SCALES=tma_aligned_scales,
274
+ USE_GDC=use_gdc,
275
num_stages=1,
276
**pdl_kwargs,
277
num_warps=1,
vllm/utils/deep_gemm.py
+19
@@ -177,6 +177,22 @@ def _import_deep_gemm():
177
return None
178
179
180
+def _apply_pdl(mod, enable: bool = True) -> None:
181
+ mod_name = getattr(mod, "__name__", str(mod))
182
+ try:
183
+ set_pdl_fn = getattr(mod, "set_pdl", None)
184
+ if set_pdl_fn is None:
185
+ return
186
+ set_pdl_fn(enable)
187
+ logger.info_once(
188
+ "DeepGEMM PDL %s on %s.",
189
+ "enabled" if enable else "disabled",
190
+ mod_name,
191
+ )
192
+ except Exception as e: # noqa: BLE001
193
+ logger.warning_once("Failed to set DeepGEMM PDL on %s: %s", mod_name, e)
194
+
195
+
196
def _lazy_init() -> None:
197
"""Import deep_gemm and resolve symbols on first use."""
198
global _cublaslt_gemm_nt_impl
@@ -219,6 +235,9 @@ def _lazy_init() -> None:
235
if _dg is None:
236
return
237
238
+ # Enable PDL for DeepGEMM on architectures that support it (SM90+).
239
+ if current_platform.is_arch_support_pdl():
240
+ _apply_pdl(_dg, True)
241
_cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
242
_fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
243
_fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)