[None][perf] Enable PDL for DeepGEMM and related kernels

cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu

@@ -25,6 +25,7 @@
 #include "fp8_blockscale_quant_packed.h"
 
 #include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/envUtils.h"
 
 #include <cstdint>
 #include <cuda_bf16.h>
@@ -61,8 +62,15 @@ __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict_
     int const lane_id = static_cast<int>(threadIdx.x) & 31;
     int const m_idx = static_cast<int>(blockIdx.y) * WarpsPerBlock + warp_id;
 
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
     if (m_idx >= m)
     {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        cudaTriggerProgrammaticLaunchCompletion();
+#endif
         return;
     }
 
@@ -175,6 +183,10 @@ __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict_
         // Layout: packed_scale[packed_sf_k_idx, m_idx]
         packed_scale_output[static_cast<int64_t>(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed;
     }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }
 
 } // namespace
@@ -187,8 +199,10 @@ void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32
     int const m_blocks = (m + kWarpsPerBlock - 1) / kWarpsPerBlock;
     dim3 const grid(num_packed_sf_k, m_blocks, 1);
     dim3 const block(kWarpsPerBlock * 32, 1, 1);
-    fp8_quantize_1x128_packed_kernel_impl<kWarpsPerBlock>
-        <<<grid, block, 0, stream>>>(fp8_output, packed_scale_output, input, m, k, scale_leading_dim_uint32);
+
+    tensorrt_llm::common::launchWithPdlWhenEnabled("fp8_quantize_1x128_packed_kernel_impl",
+        fp8_quantize_1x128_packed_kernel_impl<kWarpsPerBlock>, grid, block, 0, stream, fp8_output, packed_scale_output,
+        input, m, k, scale_leading_dim_uint32);
 }
 
 } // namespace kernels::fp8_blockscale_gemm

cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu

@@ -33,6 +33,7 @@
 
 #include "tensorrt_llm/common/assert.h"
 #include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
 
 #include <cuda.h>
 #include <cuda_bf16.h>
@@ -540,8 +541,9 @@ void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* resid
     FmaKsplitFn fn = pickFhcFma(tile_n, num_k_splits);
     dim3 const grid(static_cast<unsigned>(M), static_cast<unsigned>(N / tile_n), static_cast<unsigned>(num_k_splits));
     dim3 const block(256);
-    fn<<<grid, block, 0, stream>>>(residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, y_acc_workspace,
-        r_acc_workspace, hidden_size, N, K, residual_cur);
+    tensorrt_llm::common::launchWithPdlWhenEnabled("fused_pmap_gemm_fma_ksplit", fn, grid, block, 0, stream,
+        residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, y_acc_workspace, r_acc_workspace, hidden_size, N, K,
+        residual_cur);
 
     // ---- Step 2: big-fuse postlogue (reduces ks splits internally) ----
     mhcBigFuseLaunch(y_acc_workspace, r_acc_workspace, residual_cur, hc_scale, hc_base, post_mix_cur, comb_mix_cur,
@@ -817,10 +819,10 @@ void mhcFusedHcFmaAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons
     dim3 const grid(
         static_cast<unsigned>(m_batches), static_cast<unsigned>(N / tile_n), static_cast<unsigned>(num_k_splits));
     dim3 const block(256);
-    fn<<<grid, block, 0, stream>>>(residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale, hc_base,
-        residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace,
-        done_counter_workspace, M, K, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value,
-        sinkhorn_repeat, norm_weight, norm_eps);
+    tensorrt_llm::common::launchWithPdlWhenEnabled("fused_pmap_gemm_fma_allinone", fn, grid, block, 0, stream,
+        residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale, hc_base, residual_cur, post_mix_cur,
+        comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, done_counter_workspace, M, K, hidden_size,
+        rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps);
 }
 
 } // namespace kernels::mhc

cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu

@@ -17,6 +17,7 @@
 #include "mhcKernels.h"
 
 #include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/envUtils.h"
 
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
@@ -62,6 +63,10 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re
     int const warp_id = tid / WARP_SIZE;
     int const lane = tid % WARP_SIZE;
 
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
     __shared__ float s_pre_mix[HC_MULT];
 
     float cm[HC_MULT];
@@ -312,6 +317,9 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re
             }
         }
     }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }
 
 #define INST_BIGFUSE(NS, BS)                                                                                           \
@@ -777,9 +785,9 @@ static void mhcBigFuseDispatch(float const* y_acc, float const* r_acc, __nv_bflo
     dim3 grid(static_cast<unsigned int>(M));
 
 #define LAUNCH_BF(BS)                                                                                                  \
-    mhcBigFuseKernel<NUM_SPLITS, BS, kFuseNorm><<<grid, BS, 0, stream>>>(y_acc, r_acc, residual, hc_scale, hc_base,    \
-        post_mix, comb_mix, layer_input, M, K, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value,  \
-        sinkhorn_repeat, norm_weight, norm_eps)
+    tensorrt_llm::common::launchWithPdlWhenEnabled("mhcBigFuseKernel", mhcBigFuseKernel<NUM_SPLITS, BS, kFuseNorm>,    \
+        grid, dim3(BS), 0, stream, y_acc, r_acc, residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K,   \
+        hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps)
 
     if (block_size >= 512)
     {

cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh

@@ -78,6 +78,10 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_ksplit(__nv_bfloat16
     int const lane = tid % WARP_SIZE;
     bool const do_sqr = (n_start == 0);
 
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
     int const hidden_per_split = hidden_size / NUM_K_SPLITS;
     int const h_start = k_split * hidden_per_split;
     int const h_end = h_start + hidden_per_split;
@@ -363,6 +367,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_ksplit(__nv_bfloat16
         }
         __syncthreads();
     }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }
 
 // ===================================================================
@@ -419,8 +426,17 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1
     int const lane = tid % WARP_SIZE;
     bool const is_n0 = (n_start == 0);
 
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
     if (base_tok >= M)
+    {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        cudaTriggerProgrammaticLaunchCompletion();
+#endif
         return;
+    }
 
     int const hidden_per_split = hidden_size / KS;
     int const h_lo = k_split * hidden_per_split;
@@ -773,7 +789,12 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1
         __syncthreads();
     }
     if (!s_is_last)
+    {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        cudaTriggerProgrammaticLaunchCompletion();
+#endif
         return;
+    }
 
     // Phase 4: inline bigFuse for the TM tokens in this batch.
     // Layout: FULL_N = HC_MULT*(2+HC_MULT) = 24
@@ -963,6 +984,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1
         }
         __syncthreads();
     }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }
 
 } // namespace fused_fma_kernels

cpp/tensorrt_llm/kernels/mlaKernels.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -1208,8 +1208,16 @@ __global__ void mlaRoPEInplaceKernel(T* __restrict__ data, int32_t const* __rest
     int const vecs_per_rope
         = IS_NEOX ? (half_rope * BYTES_PER_ELT / BYTES_PER_LOAD) : (rope_dim * BYTES_PER_ELT / BYTES_PER_LOAD);
     int const head_idx = blockIdx.y * HEADS_PER_BLOCK + threadIdx.y;
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
     if (head_idx >= num_heads || tid >= vecs_per_rope)
+    {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        cudaTriggerProgrammaticLaunchCompletion();
+#endif
         return;
+    }
 
     int const head_size = nope_dim + rope_dim;
     T* head_ptr = data + (static_cast<int64_t>(blockIdx.x) * num_heads + head_idx) * head_size;
@@ -1273,6 +1281,9 @@ __global__ void mlaRoPEInplaceKernel(T* __restrict__ data, int32_t const* __rest
 
         *reinterpret_cast<VecT*>(&head_ptr[nope_dim + elem_offset]) = v;
     }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
 }
 
 template <typename T>
@@ -1299,18 +1310,21 @@ void invokeMLARoPEInplace(T* data, int32_t const* position_ids, float const* cos
 
         if (hpb <= 4)
         {
-            mlaRoPEInplaceKernel<T, INV, NEOX, 4><<<grid, dim3(vecs_per_rope, 4), 0, stream>>>(
-                data, position_ids, cos_sin_cache, num_heads, nope_dim, rope_dim);
+            tensorrt_llm::common::launchWithPdlWhenEnabled("mlaRoPEInplaceKernel",
+                mlaRoPEInplaceKernel<T, INV, NEOX, 4>, grid, dim3(vecs_per_rope, 4), 0, stream, data, position_ids,
+                cos_sin_cache, num_heads, nope_dim, rope_dim);
         }
         else if (hpb <= 8)
         {
-            mlaRoPEInplaceKernel<T, INV, NEOX, 8><<<grid, dim3(vecs_per_rope, 8), 0, stream>>>(
-                data, position_ids, cos_sin_cache, num_heads, nope_dim, rope_dim);
+            tensorrt_llm::common::launchWithPdlWhenEnabled("mlaRoPEInplaceKernel",
+                mlaRoPEInplaceKernel<T, INV, NEOX, 8>, grid, dim3(vecs_per_rope, 8), 0, stream, data, position_ids,
+                cos_sin_cache, num_heads, nope_dim, rope_dim);
         }
         else
         {
-            mlaRoPEInplaceKernel<T, INV, NEOX, 16><<<grid, dim3(vecs_per_rope, 16), 0, stream>>>(
-                data, position_ids, cos_sin_cache, num_heads, nope_dim, rope_dim);
+            tensorrt_llm::common::launchWithPdlWhenEnabled("mlaRoPEInplaceKernel",
+                mlaRoPEInplaceKernel<T, INV, NEOX, 16>, grid, dim3(vecs_per_rope, 16), 0, stream, data, position_ids,
+                cos_sin_cache, num_heads, nope_dim, rope_dim);
         }
     };
 

tensorrt_llm/_torch/attention_backend/sparse/kernel.py

@@ -1,4 +1,5 @@
 import math
+import os
 
 import torch
 import triton
@@ -2114,6 +2115,7 @@ def _deepseek_v4_local_to_global_kernel(
     compressed_indices_stride1,
     out_stride0,
     out_stride1,
+    LAUNCH_WITH_PDL: tl.constexpr,
 ):
     """
     Triton kernel for converting local indices to global KV cache pool indices.
@@ -2127,6 +2129,9 @@ def _deepseek_v4_local_to_global_kernel(
     This enables the FMHA kernel to determine which TMA descriptor to use based
     solely on tile index (tile 0 = SWA via tmaKSecondary_, rest = compress via tmaK_).
     """
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_wait()
+
     token_id = tl.program_id(0)
 
     # Load request ID for this token
@@ -2188,6 +2193,9 @@ def _deepseek_v4_local_to_global_kernel(
         compressed_out_ptr = out_ptr + token_id * out_stride0 + compressed_write_pos * out_stride1
         tl.store(compressed_out_ptr, compressed_global_index)
 
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 def deepseek_v4_local_to_global_indices(
         req_id: torch.Tensor,  # int32 [num_tokens]
@@ -2310,6 +2318,7 @@ def deepseek_v4_local_to_global_indices(
     compressed_indices_stride0, compressed_indices_stride1 = compressed_local_indices_c.stride(
     )
     out_stride0, out_stride1 = out.stride()
+    launch_with_pdl = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1"
 
     # Launch kernel
     _deepseek_v4_local_to_global_kernel[grid](
@@ -2339,6 +2348,8 @@ def deepseek_v4_local_to_global_indices(
         compressed_indices_stride1,
         out_stride0,
         out_stride1,
+        LAUNCH_WITH_PDL=launch_with_pdl,
+        launch_pdl=launch_with_pdl,
     )
 
     return out

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -55,6 +55,8 @@
 # BufferKind is bound from C++; see cpp/tensorrt_llm/thop/outputTensor.h (torch_ext::BufferKind).
 from tensorrt_llm.bindings.internal.thop import BufferKind
 
+deep_gemm.set_pdl(get_env_enable_pdl())
+
 
 # Used to WAR an issue in torch.bmm that it would break the graph when the out is not contiguous.
 @torch.library.custom_op("trtllm::bmm_out", mutates_args=("out", ))