From 6388a0ed238f67dc33335d32592ea49ced1dbe40 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Thu, 1 Oct 2026 16:05:08 -0700 Subject: [PATCH 1/5] [Kernel] Enable PDL for fused MoE top-k gating Allow SM90+ fused softmax and sigmoid gating to participate in programmatic dependent launch. Wait before producer-dependent reads and retain the ordinary launch for unsupported platforms. Co-authored-by: OpenAI Codex Signed-off-by: Hanlin Bi --- .../moe/topk_softmax_kernels.cu | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 098644a1b901..f5e1d5100b40 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -276,7 +276,7 @@ __launch_bounds__(TPB) __global__ void moeTopK( */ template + 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, + 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<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding); } From c53ac814e6284e7fbf3c52d5db13e9cb2d00a322 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Fri, 9 Oct 2026 17:42:12 -0700 Subject: [PATCH 2/5] Remove redundant CUDA version guards for top-k PDL Co-authored-by: OpenAI Codex Signed-off-by: Hanlin Bi --- csrc/libtorch_stable/moe/topk_softmax_kernels.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index f5e1d5100b40..1c0d1ed21ef5 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -329,7 +329,7 @@ __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 !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 if constexpr (ENABLE_PDL) { // PDL permits early launch, but logits/metadata still depend on the producer. cudaGridDependencySynchronize(); @@ -628,7 +628,7 @@ 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 !defined(USE_ROCM) if (get_device_prop()->major >= 9) { cudaLaunchConfig_t config{}; config.gridDim = dim3(num_blocks); From 713192babcf878fd0b4ec6c269963b591081a877 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Fri, 9 Oct 2026 17:43:38 -0700 Subject: [PATCH 3/5] Simplify device-side PDL guards in fused top-k gating Use the SM90 compile guard for dependency synchronization and launch completion, matching the adjacent softplus-sqrt gating kernel. Remove the redundant ENABLE_PDL template specialization. Co-authored-by: OpenAI Codex Signed-off-by: Hanlin Bi --- csrc/libtorch_stable/moe/topk_softmax_kernels.cu | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 1c0d1ed21ef5..84094cc7731d 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -276,7 +276,7 @@ __launch_bounds__(TPB) __global__ void moeTopK( */ template + typename InputType = float, ScoringFunc SF> __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,12 +329,10 @@ __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 - if constexpr (ENABLE_PDL) { - // PDL permits early launch, but logits/metadata still depend on the producer. - cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); - } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + // 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. @@ -641,7 +639,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa config.numAttrs = 1; auto err = cudaLaunchKernelEx( &config, - topkGating, + topkGating, 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)); From 09647c0364e536612dd378f83a25145a958e5bb2 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Mon, 12 Oct 2026 10:38:17 -0700 Subject: [PATCH 4/5] Remove host architecture check from top-k PDL launch Set the PDL launch attribute on every CUDA launch, matching topk_softplus_sqrt and per_token_group_quant. Keep the device-side SM90 guard and the ROCm launch path. Co-authored-by: OpenAI Codex Signed-off-by: Hanlin Bi --- .../moe/topk_softmax_kernels.cu | 36 +++++++++---------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 84094cc7731d..dd2df4e3ee6e 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -627,27 +627,25 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); #if !defined(USE_ROCM) - 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, - 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 + 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, + 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)); +#else topkGating<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding); +#endif } #ifndef USE_ROCM From 30ec272dd3d9ab80cb50b7225622d114ef8bb1fa Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Mon, 12 Oct 2026 10:51:53 -0700 Subject: [PATCH 5/5] Match neighboring kernel ROCm guard style Co-authored-by: OpenAI Codex Signed-off-by: Hanlin Bi --- csrc/libtorch_stable/moe/topk_softmax_kernels.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index dd2df4e3ee6e..cf3e4dd55d53 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -626,7 +626,7 @@ 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) +#ifndef USE_ROCM cudaLaunchConfig_t config{}; config.gridDim = dim3(num_blocks); config.blockDim = block_dim;