Use bool flag for PDL

Author: Trevor Morris

csrc/nv_internal/cpp/common/envUtils.cpp

@@ -356,16 +356,4 @@ int getEnvMoeA2ACombineBlockSize() {
 
 bool getEnvEplbForceGdrcopy() { return getBoolEnv("TRTLLM_EPLB_FORCE_GDRCOPY"); }
 
-bool getEnvEnablePDL() {
-  static bool const kEnablePDL = []() -> bool {
-    auto const val = getIntEnv("TRTLLM_ENABLE_PDL");
-    if (!val.has_value()) {
-      // Default: enabled on SM>=90 (Hopper, Blackwell, Grace-Blackwell)
-      return getSMVersion() >= 90;
-    }
-    return val.value() != 0;
-  }();
-  return kEnablePDL;
-}
-
 }  // namespace tensorrt_llm::common

csrc/nv_internal/tensorrt_llm/common/envUtils.h

@@ -105,13 +105,9 @@ int getEnvMoeA2ADispatchBlockSize();
 // Block size (threads per block) for MoE A2A Combine kernels (default 256 if unset or invalid)
 int getEnvMoeA2ACombineBlockSize();
 
-// Whether PDL (Programmatic Dependent Launch) is enabled (default true on SM>=90).
-// Controlled by TRTLLM_ENABLE_PDL env var: "0" disables, "1" force-enables.
-bool getEnvEnablePDL();
-
 template <typename KernelFn, typename... Args>
-inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block,
-    size_t dynamicShmSize, cudaStream_t stream, Args&&... args)
+inline void launchWithPdlWhenEnabled(char const* name, bool enable_pdl, KernelFn kernelFn,
+    dim3 grid, dim3 block, size_t dynamicShmSize, cudaStream_t stream, Args&&... args)
 {
     cudaLaunchConfig_t kernelConfig;
     kernelConfig.gridDim = grid;
@@ -120,7 +116,7 @@ inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 g
     kernelConfig.stream = stream;
     cudaLaunchAttribute attrs[1];
     attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+    attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
     kernelConfig.attrs = attrs;
     kernelConfig.numAttrs = 1;
     cudaError_t e = cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward<Args>(args)...);

csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu

@@ -290,9 +290,9 @@ __device__ void vectorized_dispatch(uint8_t const* src_ptr, int bytes_per_token,
 }
 
 __global__ void moeA2APrepareDispatchKernel(int* send_counters, int* local_token_counter,
-                                            int ep_size, uint32_t* flag_val_ptr) {
+                                            int ep_size, uint32_t* flag_val_ptr, bool enable_pdl) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  if (enable_pdl) cudaGridDependencySynchronize();
   cudaTriggerProgrammaticLaunchCompletion();
 #endif
   int idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -322,7 +322,7 @@ __global__ void moeA2ADispatchKernel(
     const DispatchKernelPointers ptrs,      // Struct containing all kernel pointers
     int num_payloads,                       // Number of payloads
     int max_tokens_per_rank,                // Maximum tokens per rank
-    int local_num_tokens, int rank_id, int ep_size, int num_experts_per_rank) {
+    int local_num_tokens, int rank_id, int ep_size, int num_experts_per_rank, bool enable_pdl) {
   int thread_idx = ThreadingPolicy::offset();
   int local_token_idx = ThreadingPolicy::token_idx();
 
@@ -332,14 +332,14 @@ __global__ void moeA2ADispatchKernel(
     // synchronization. Other threads should return.
     if (local_token_idx > 0) return;
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaGridDependencySynchronize();
+    if (enable_pdl) cudaGridDependencySynchronize();
 #endif
   } else {
     // Threads that do not have a token to process should return.
     if (local_token_idx >= local_num_tokens) return;
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaGridDependencySynchronize();
+    if (enable_pdl) cudaGridDependencySynchronize();
 #endif
 
     // Prepare per-policy shared-memory tiles for this token
@@ -491,9 +491,9 @@ __global__ void moeA2ADispatchKernel(
 }
 
 void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params) {
-  launchWithPdlWhenEnabled("moeA2APrepareDispatchKernel", moeA2APrepareDispatchKernel, 1,
-      params.ep_size, 0, params.stream, params.send_counters, params.local_token_counter,
-      params.ep_size, params.flag_val);
+  launchWithPdlWhenEnabled("moeA2APrepareDispatchKernel", params.enable_pdl,
+      moeA2APrepareDispatchKernel, 1, params.ep_size, 0, params.stream, params.send_counters,
+      params.local_token_counter, params.ep_size, params.flag_val, params.enable_pdl);
 }
 
 // ============================================================================
@@ -552,10 +552,10 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) {
     int shared_bytes = 2 * params.top_k * (int)sizeof(int);
     SWITCH_TOP_K(params.top_k, TOP_K, {
       auto kernel_fn = moeA2ADispatchKernel<BlockPolicy, TOP_K>;
-      launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize,
-          shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs,
+      launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size,
+          kBlockSize, shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs,
           params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank,
-          params.ep_size, params.num_experts_per_rank);
+          params.ep_size, params.num_experts_per_rank, params.enable_pdl);
     })
   } else {
     int grid_size = ceilDiv(params.local_num_tokens, kWarpsPerBlock);
@@ -567,10 +567,10 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) {
     int shared_bytes = 2 * kWarpsPerBlock * params.top_k * (int)sizeof(int);
     SWITCH_TOP_K(params.top_k, TOP_K, {
       auto kernel_fn = moeA2ADispatchKernel<WarpPolicy, TOP_K>;
-      launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize,
-          shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs,
+      launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size,
+          kBlockSize, shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs,
           params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank,
-          params.ep_size, params.num_experts_per_rank);
+          params.ep_size, params.num_experts_per_rank, params.enable_pdl);
     })
   }
 }
@@ -919,9 +919,10 @@ template <typename ThreadingPolicy, bool LOW_PRECISION, typename SrcT>
 __global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void const* payload,
                                            int elements_per_token, int ep_size,
                                            int max_tokens_per_rank, uint32_t* flag_val_ptr,
-                                           int const* recv_counters, int stride_per_token) {
+                                           int const* recv_counters, int stride_per_token,
+                                           bool enable_pdl) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  if (enable_pdl) cudaGridDependencySynchronize();
   cudaTriggerProgrammaticLaunchCompletion();
 #endif
 
@@ -977,7 +978,7 @@ template <typename T, typename ThreadingPolicy, int TOP_K>
 __global__ void moeA2ACombineKernel(
     const CombineKernelPointers ptrs,  // Combine-specific struct, src_data_ptrs[0] is output
     int max_tokens_per_rank, int elements_per_token, int stride_per_token, int local_num_tokens,
-    int rank_id, int ep_size) {
+    int rank_id, int ep_size, bool enable_pdl) {
   int local_token_idx = ThreadingPolicy::token_idx();
   int const size_per_token = elements_per_token * static_cast<int>(sizeof(T));
 
@@ -992,7 +993,7 @@ __global__ void moeA2ACombineKernel(
   }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  if (enable_pdl) cudaGridDependencySynchronize();
 #endif
 
 #if !DISABLE_SYNC_FOR_PROFILING
@@ -1108,9 +1109,10 @@ void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params) {
           params.one_block_per_token
               ? moeA2APrepareCombineKernel<BlockPolicy, LOW_PRECISION, SrcT>
               : moeA2APrepareCombineKernel<WarpPolicy, LOW_PRECISION, SrcT>;
-      launchWithPdlWhenEnabled("moeA2APrepareCombineKernel", kernel_fn, grid, kBlockSize, 0,
-          params.stream, recv_buffer_bytes, payload, params.elements_per_token, params.ep_size,
-          params.max_tokens_per_rank, params.flag_val, params.recv_counters, stride_per_token);
+      launchWithPdlWhenEnabled("moeA2APrepareCombineKernel", params.enable_pdl, kernel_fn, grid,
+          kBlockSize, 0, params.stream, recv_buffer_bytes, payload, params.elements_per_token,
+          params.ep_size, params.max_tokens_per_rank, params.flag_val, params.recv_counters,
+          stride_per_token, params.enable_pdl);
     });
   });
 }
@@ -1184,9 +1186,10 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) {
     SWITCH_POLICY(params.one_block_per_token, Policy, {
       SWITCH_TOP_K(params.top_k, TOP_K, {
         auto kernel_fn = moeA2ACombineKernel<TKernelType, Policy, TOP_K>;
-        launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0,
-            params.stream, kernel_ptrs, params.max_tokens_per_rank, params.elements_per_token,
-            stride_per_token, params.local_num_tokens, params.ep_rank, params.ep_size);
+        launchWithPdlWhenEnabled("moeA2ACombineKernel", params.enable_pdl, kernel_fn, grid,
+            kBlockSize, 0, params.stream, kernel_ptrs, params.max_tokens_per_rank,
+            params.elements_per_token, stride_per_token, params.local_num_tokens, params.ep_rank,
+            params.ep_size, params.enable_pdl);
       });
     });
   });
@@ -1196,9 +1199,9 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) {
 __global__ void moeA2ASanitizeExpertIdsKernel(int32_t* expert_ids_ptr,
                                               int32_t const* recv_counters_ptr, int ep_size,
                                               int max_tokens_per_rank, int top_k,
-                                              int32_t invalid_id) {
+                                              int32_t invalid_id, bool enable_pdl) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  cudaGridDependencySynchronize();
+  if (enable_pdl) cudaGridDependencySynchronize();
   cudaTriggerProgrammaticLaunchCompletion();
 #endif
   int tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -1218,13 +1221,13 @@ __global__ void moeA2ASanitizeExpertIdsKernel(int32_t* expert_ids_ptr,
 
 void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters,
                                         int32_t invalid_id, int ep_size, int max_tokens_per_rank,
-                                        int top_k, cudaStream_t stream) {
+                                        int top_k, cudaStream_t stream, bool enable_pdl) {
   constexpr int kBlockSize = 256;
   int total_tokens = ep_size * max_tokens_per_rank;
   int grid = ceilDiv(total_tokens, kBlockSize);
-  launchWithPdlWhenEnabled("moeA2ASanitizeExpertIdsKernel", moeA2ASanitizeExpertIdsKernel, grid,
-      kBlockSize, 0, stream, expert_ids, recv_counters, ep_size, max_tokens_per_rank, top_k,
-      invalid_id);
+  launchWithPdlWhenEnabled("moeA2ASanitizeExpertIdsKernel", enable_pdl,
+      moeA2ASanitizeExpertIdsKernel, grid, kBlockSize, 0, stream, expert_ids, recv_counters,
+      ep_size, max_tokens_per_rank, top_k, invalid_id, enable_pdl);
 }
 
 }  // namespace tensorrt_llm::kernels::moe_alltoall

csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h

@@ -77,6 +77,7 @@ struct CombineKernelPointers {
 // Dispatch phase parameters
 struct MoeA2ADispatchParams {
   bool one_block_per_token;  // True: one block per token, False: one warp per token
+  bool enable_pdl;           // True: launch with programmatic dependent launch
 
   // Threading policy
   // EP configuration
@@ -126,6 +127,7 @@ void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params);
 // Combine phase parameters
 struct MoeA2ACombineParams {
   bool one_block_per_token;  // True: one block per token, False: one warp per token
+  bool enable_pdl;           // True: launch with programmatic dependent launch
 
   // EP configuration
   int ep_size;  // Number of EP ranks
@@ -179,6 +181,6 @@ void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params);
 // invalid_id: value to fill for invalid tokens' expert ids
 void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters,
                                         int32_t invalid_id, int ep_size, int max_tokens_per_rank,
-                                        int top_k, cudaStream_t stream);
+                                        int top_k, cudaStream_t stream, bool enable_pdl);
 
 }  // namespace tensorrt_llm::kernels::moe_alltoall

csrc/trtllm_moe_alltoall.cu

@@ -119,7 +119,7 @@ Tensor moeA2AInitializeOp(TensorView workspace, int64_t epRank, int64_t epSize,
 Tuple<Array<int64_t>, Array<int64_t>, int64_t> moeA2ADispatchOp(
     TensorView tokenSelectedExperts, Array<Tensor> inputPayloads, TensorView workspace,
     TensorView metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank, int64_t epSize,
-    int64_t topK, int64_t numExperts) {
+    int64_t topK, int64_t numExperts, bool enablePdl) {
   using tl_throughput::PayloadDescriptor;
 
   CHECK_INPUT(tokenSelectedExperts);
@@ -197,6 +197,7 @@ Tuple<Array<int64_t>, Array<int64_t>, int64_t> moeA2ADispatchOp(
 
   tl_throughput::MoeA2ADispatchParams params{};
   params.one_block_per_token = tensorrt_llm::common::getEnvMoeA2AOneBlockPerToken();
+  params.enable_pdl = enablePdl;
   params.ep_size = static_cast<int>(epSize);
   params.ep_rank = static_cast<int>(epRank);
   params.num_experts_per_rank = static_cast<int>(numExperts / epSize);
@@ -275,7 +276,7 @@ nvinfer1::DataType toNvDataType(DLDataType dtype) {
 Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView workspace,
                        TensorView metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank,
                        int64_t epSize, int64_t topK, int64_t combinePayloadOffset,
-                       bool payloadInWorkspace, bool useLowPrecision) {
+                       bool payloadInWorkspace, bool useLowPrecision, bool enablePdl) {
   using tl_throughput::MoeA2ACombineParams;
   CHECK_INPUT(payload);
   TVM_FFI_ICHECK_EQ(payload.ndim(), 3)
@@ -322,6 +323,7 @@ Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView wo
 
   MoeA2ACombineParams params{};
   params.one_block_per_token = tensorrt_llm::common::getEnvMoeA2AOneBlockPerToken();
+  params.enable_pdl = enablePdl;
   params.ep_size = static_cast<int>(epSize);
   params.ep_rank = static_cast<int>(epRank);
   params.local_num_tokens = static_cast<int>(localNumTokens);
@@ -359,7 +361,7 @@ Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView wo
 }
 
 void moeA2ASanitizeExpertIdsOp(TensorView expertIds, TensorView workspace, TensorView metainfo,
-                               int64_t epRank, int64_t invalidExpertId) {
+                               int64_t epRank, int64_t invalidExpertId, bool enablePdl) {
   CHECK_INPUT(expertIds);
   CHECK_INPUT_TYPE(expertIds, dl_int32);
   TVM_FFI_ICHECK_EQ(expertIds.ndim(), 3);
@@ -385,7 +387,8 @@ void moeA2ASanitizeExpertIdsOp(TensorView expertIds, TensorView workspace, Tenso
   tl_throughput::moe_a2a_sanitize_expert_ids_launch(
       static_cast<int32_t*>(expertIds.data_ptr()), recvCounters,
       static_cast<int32_t>(invalidExpertId), static_cast<int>(epSize),
-      static_cast<int>(runtimeMaxTokensPerRank), static_cast<int>(topK), get_current_stream());
+      static_cast<int>(runtimeMaxTokensPerRank), static_cast<int>(topK), get_current_stream(),
+      enablePdl);
 
   auto err = cudaGetLastError();
   TVM_FFI_ICHECK(err == cudaSuccess)