[Kernel] Add PDL support for DeepGEMM kernel

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -287,9 +287,17 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
   const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
   const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
 
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
   if (mn_idx >= tma_aligned_mn) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.launch_dependents;");
+#endif
     return;
   }
+
   const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
 
   // Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -392,6 +400,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
       static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
       sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
   *reinterpret_cast<uint4*>(group_output) = packed_out;
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
 }
 
 // Public entry point: register-resident packed quant kernel.
@@ -472,20 +484,29 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 
 #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
   do {                                                                       \
-    dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-              static_cast<unsigned int>(blocks_y));                          \
-    dim3 block(num_threads);                                                 \
-    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
-                                                      RY>                    \
-        <<<grid, block, 0, stream>>>(                                        \
-            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),    \
-            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),     \
-            static_cast<int>(padded_groups_per_row),                         \
-            static_cast<int>(groups_per_row), static_cast<int>(mn),          \
-            static_cast<int>(output_q_mn_extent),                            \
-            static_cast<int>(tma_aligned_mn), num_scale_elems,               \
-            static_cast<float>(eps), static_cast<float>(min_8bit),           \
-            static_cast<float>(max_8bit));                                   \
+    cudaLaunchConfig_t config = {};                                          \
+    config.gridDim = dim3(static_cast<unsigned int>(blocks_x),               \
+                          static_cast<unsigned int>(blocks_y));              \
+    config.blockDim = dim3(num_threads);                                     \
+    config.dynamicSmemBytes = 0;                                             \
+    config.stream = stream;                                                  \
+    cudaLaunchAttribute attrs[1];                                            \
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;        \
+    attrs[0].val.programmaticStreamSerializationAllowed = 1;                 \
+    config.numAttrs = 1;                                                     \
+    config.attrs = attrs;                                                    \
+    cudaLaunchKernelEx(                                                      \
+        &config,                                                             \
+        per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, \
+                                                          KX, RY>,           \
+        static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),        \
+        reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),         \
+        static_cast<int>(padded_groups_per_row),                             \
+        static_cast<int>(groups_per_row), static_cast<int>(mn),              \
+        static_cast<int>(output_q_mn_extent),                                \
+        static_cast<int>(tma_aligned_mn), num_scale_elems,                   \
+        static_cast<float>(eps), static_cast<float>(min_8bit),               \
+        static_cast<float>(max_8bit));                                       \
   } while (0)
 
 #define LAUNCH_REG_KERNEL(T, DST_DTYPE)                    \

vllm/envs.py

@@ -174,6 +174,7 @@
     VLLM_MOE_USE_DEEP_GEMM: bool = True
     VLLM_USE_DEEP_GEMM_E8M0: bool = True
     VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True
+    VLLM_USE_DEEP_GEMM_PDL: bool = False
     VLLM_DEEP_GEMM_WARMUP: Literal[
         "skip",
         "full",
@@ -1402,6 +1403,11 @@ def _resolve_rust_frontend_path() -> str | None:
     "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool(
         int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1"))
     ),
+    # Enable PDL for DeepGEMM kernels o Hopper/Blackwell.Tthe flag is consumed by
+    # deep_gemm.set_pdl() once per process before any DeepGEMM kernel launches
+    "VLLM_USE_DEEP_GEMM_PDL": lambda: bool(
+        int(os.getenv("VLLM_USE_DEEP_GEMM_PDL", "0"))
+    ),
     # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm
     # JIT all the required kernels before model execution so there is no
     # JIT'ing in the hot-path. However, this warmup increases the engine

vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py

@@ -37,6 +37,8 @@ def _fused_inv_rope_fp8_quant_per_head(
     ROPE_START: tl.constexpr,
     HALF_ROPE: tl.constexpr,
     TMA_ALIGNED_SCALES: tl.constexpr,
+    USE_GDC: tl.constexpr,
+    launch_pdl: tl.constexpr,  # triton metadata
 ):
     # int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
     pid_token = tl.program_id(0).to(tl.int64)
@@ -46,7 +48,9 @@ def _fused_inv_rope_fp8_quant_per_head(
     head_in_group = pid_gh % heads_per_group
     global_head = pid_gh
     qb_start = head_in_group * CHUNKS_PER_HEAD
-
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+        tl.extra.cuda.gdc_wait()
     # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
     if pid_token >= num_tokens:
         if TMA_ALIGNED_SCALES:
@@ -133,6 +137,8 @@ def _fused_inv_rope_fp8_quant_per_head(
             scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k
         )
         tl.store(scale_addrs, scales)
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
 
 
 def fused_inv_rope_fp8_quant(
@@ -243,11 +249,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
         (scale_inner * tma_aligned_T, 1, tma_aligned_T),
     )
     grid = (tma_aligned_T, n_groups * heads_per_group)
-    pdl_kwargs = (
-        {}
-        if current_platform.is_rocm() or current_platform.is_xpu()
-        else {"launch_pdl": False}
-    )
+    use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90)
+    pdl_kwargs = {"launch_pdl": True} if use_gdc else {}
     _fused_inv_rope_fp8_quant_per_head[grid](
         o,
         positions,
@@ -270,6 +273,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
         ROPE_START=rope_start,
         HALF_ROPE=half_rope,
         TMA_ALIGNED_SCALES=tma_aligned_scales,
+        USE_GDC=use_gdc,
         num_stages=1,
         **pdl_kwargs,
         num_warps=1,

vllm/utils/deep_gemm.py

@@ -175,6 +175,36 @@ def _import_deep_gemm():
     return None
 
 
+def _apply_pdl(enable: bool) -> None:
+    """Set PDL on every DeepGEMM module currently importable.
+
+    The external (pip-installed) ``deep_gemm`` and the vendored
+    ``vllm.third_party.deep_gemm`` are independent C extensions with
+    independent global PDL state. Apply the flag to both so model code
+    that imports either module sees the same setting.
+    """
+    applied_to: list[str] = []
+    for mod_name in ("deep_gemm", "vllm.third_party.deep_gemm"):
+        try:
+            mod = importlib.import_module(mod_name)
+        except Exception:  # noqa: BLE001
+            continue
+        set_pdl_fn = getattr(mod, "set_pdl", None)
+        if set_pdl_fn is None:
+            continue
+        try:
+            set_pdl_fn(enable)
+            applied_to.append(mod_name)
+        except Exception as e:  # noqa: BLE001
+            logger.warning_once("Failed to set DeepGEMM PDL on %s: %s", mod_name, e)
+    if applied_to:
+        logger.info_once(
+            "DeepGEMM PDL %s on %s.",
+            "enabled" if enable else "disabled",
+            ", ".join(applied_to),
+        )
+
+
 def _lazy_init() -> None:
     """Import deep_gemm and resolve symbols on first use."""
     global _cublaslt_gemm_nt_impl
@@ -217,6 +247,18 @@ def _lazy_init() -> None:
     if _dg is None:
         return
 
+    # Apply DeepGEMM PDL setting once per process, before any kernel launches
+    # or CUDA-graph capture. PDL state is global in DeepGEMM and is read at
+    # each kernel launch; flipping it later would diverge from launches
+    # captured into CUDA graphs.
+    #
+    # NOTE: vLLM may load two independent DeepGEMM C extensions in the same
+    # process — the pip-installed ``deep_gemm`` and the vendored
+    # ``vllm.third_party.deep_gemm`` (e.g. deepseek_v4.py imports the
+    # vendored module directly). Each has its own global PDL state, so we
+    # apply set_pdl to whichever modules are already importable.
+    _apply_pdl(envs.VLLM_USE_DEEP_GEMM_PDL)
+
     _cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
     _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
     _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)
@@ -243,6 +285,15 @@ def _lazy_init() -> None:
     DeepGemmQuantScaleFMT.init_oracle_cache()
 
 
+def configure_deep_gemm() -> None:
+    """Eagerly initialize DeepGEMM so process-global settings (PDL, JIT
+    cache dir) are applied before profile_run / warmup / CUDA graph
+    capture. Safe to call when DeepGEMM is unsupported — it becomes a
+    no-op.
+    """
+    _lazy_init()
+
+
 def get_num_sms() -> int:
     _lazy_init()
     dg = _import_deep_gemm()
@@ -563,6 +614,7 @@ def should_use_deepgemm_for_fp8_linear(
 
 __all__ = [
     "calc_diff",
+    "configure_deep_gemm",
     "DeepGemmQuantScaleFMT",
     "fp8_gemm_nt",
     "fp8_einsum",

vllm/v1/worker/gpu_worker.py

@@ -51,6 +51,7 @@
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.tracing import instrument
+from vllm.utils.deep_gemm import configure_deep_gemm
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
 from vllm.utils.torch_utils import set_random_seed
@@ -294,6 +295,10 @@ def init_device(self):
             # Set random seed.
             set_random_seed(self.model_config.seed)
 
+            # Apply DeepGEMM process-global settings (PDL, JIT cache dir)
+            # before profile_run / warmup / CUDA graph capture.
+            configure_deep_gemm()
+
             # Now take memory snapshot after NCCL is initialized
             gc.collect()
             torch.accelerator.empty_cache()