[Perf] Enable PDL for fused MoE top-k gating

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

@@ -276,7 +276,7 @@ __launch_bounds__(TPB) __global__ void moeTopK(
 */
 
 template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType,
-          typename InputType = float, ScoringFunc SF>
+          typename InputType = float, ScoringFunc SF, bool ENABLE_PDL = false>
 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
     void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
         int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
@@ -329,6 +329,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
     const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
     const int thread_row = warp_base_row + thread_row_in_warp;
 
+#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12000
+    if constexpr (ENABLE_PDL) {
+        // PDL permits early launch, but logits/metadata still depend on the producer.
+        cudaGridDependencySynchronize();
+        cudaTriggerProgrammaticLaunchCompletion();
+    }
+#endif
+
     // Threads with indices out of bounds should early exit here.
     if (thread_row >= num_rows)
     {
@@ -620,6 +628,26 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
     const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
 
     dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
+#if !defined(USE_ROCM) && CUDART_VERSION >= 12000
+    if (get_device_prop()->major >= 9) {
+        cudaLaunchConfig_t config{};
+        config.gridDim = dim3(num_blocks);
+        config.blockDim = block_dim;
+        config.stream = stream;
+        cudaLaunchAttribute attr{};
+        attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attr.val.programmaticStreamSerializationAllowed = 1;
+        config.attrs = &attr;
+        config.numAttrs = 1;
+        auto err = cudaLaunchKernelEx(
+            &config,
+            topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF, true>,
+            input, finished, output, num_rows, indices, source_row, k, start_expert,
+            end_expert, renormalize, bias, routed_scaling_factor, is_padding);
+        STD_TORCH_CHECK(err == cudaSuccess, "PDL topkGating launch failed: ", cudaGetErrorString(err));
+        return;
+    }
+#endif
     topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
         input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
 }