[Misc] Use functions instead of PTX for the PDL instruction (#46984)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jee Jee Li committed
Jul 2, 2026 at 10:38 UTC
d0a2584773b2c182cef0056812b6ba87e6a0a5b7
7 files changed
+27
-25
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+2
-2
@@ -328,7 +328,7 @@ struct GmemLoaderB {
328
329
__device__ void issue_mainloop() {
330
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
331
- asm volatile("griddepcontrol.wait;");
331
+ cudaGridDependencySynchronize();
332
#pragma unroll 1
333
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
334
if (need_wait) {
@@ -643,7 +643,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
643
mma_computer.issue_mainloop();
644
mma_computer.epi();
645
}
646
- asm volatile("griddepcontrol.launch_dependents;");
646
+ cudaTriggerProgrammaticLaunchCompletion();
647
#endif
648
}
649
csrc/libtorch_stable/fp32_router_gemm.cu
+2
-2
@@ -100,7 +100,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
100
}
101
102
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
103
- asm volatile("griddepcontrol.wait;");
103
+ cudaGridDependencySynchronize();
104
#endif
105
106
for (int ki = 0; ki < k_iterations; ki++) {
@@ -146,7 +146,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
146
}
147
148
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
149
- asm volatile("griddepcontrol.launch_dependents;");
149
+ cudaTriggerProgrammaticLaunchCompletion();
150
#endif
151
}
152
csrc/libtorch_stable/minimax_reduce_rms_kernel.cu
+4
-4
@@ -249,7 +249,7 @@ __global__ void __launch_bounds__(1024)
249
LamportComm<NRanks> comm(params.workspace, params.rank);
250
int clear_access = comm.clear_size / kElemsPerAccess<DType>;
251
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
252
- asm volatile("griddepcontrol.wait;");
252
+ cudaGridDependencySynchronize();
253
#endif
254
for (int idx = access_id; idx < tot_access;
255
idx += access_stride, token_id += token_stride) {
@@ -313,7 +313,7 @@ __global__ void __launch_bounds__(1024)
313
}
314
comm.update(params.size_q * NRanks);
315
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
316
- asm volatile("griddepcontrol.launch_dependents;");
316
+ cudaTriggerProgrammaticLaunchCompletion();
317
#endif
318
}
319
@@ -384,7 +384,7 @@ __global__ void __launch_bounds__(1024)
384
385
DType norm_weight[kElemsPerAccess<DType>]{};
386
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
387
- asm volatile("griddepcontrol.wait;");
387
+ cudaGridDependencySynchronize();
388
#endif
389
if (is_q) {
390
if (is_valid_q) {
@@ -596,7 +596,7 @@ __global__ void __launch_bounds__(1024)
596
}
597
} // end group loop
598
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
599
- asm volatile("griddepcontrol.launch_dependents;");
599
+ cudaTriggerProgrammaticLaunchCompletion();
600
#endif
601
602
int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu
+2
-2
@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
78
}
79
80
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
81
- asm volatile("griddepcontrol.wait;");
81
+ cudaGridDependencySynchronize();
82
#endif
83
84
// Process the GEMM in chunks
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
163
}
164
}
165
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
166
- asm volatile("griddepcontrol.launch_dependents;");
166
+ cudaTriggerProgrammaticLaunchCompletion();
167
#endif
168
}
169
csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu
+2
-2
@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
78
}
79
80
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
81
- asm volatile("griddepcontrol.wait;");
81
+ cudaGridDependencySynchronize();
82
#endif
83
84
// Process the GEMM in chunks
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
163
}
164
}
165
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
166
- asm volatile("griddepcontrol.launch_dependents;");
166
+ cudaTriggerProgrammaticLaunchCompletion();
167
#endif
168
}
169
csrc/libtorch_stable/moe/grouped_topk_kernels.cu
+12
-10
@@ -48,7 +48,8 @@ static constexpr int NumTopGroupScores = 2;
48
static constexpr int DefaultMaxNumTopExperts = 8;
49
static constexpr int MaxSupportedTopExperts = 22;
50
static constexpr int MaxNumTopGroups = 4;
51
-
51
+// The empirical value for small batch
52
+static constexpr int PDLEnableTokens = 16;
53
namespace warp_topk {
54
55
template <int size, typename T>
@@ -564,8 +565,8 @@ __global__ void grouped_topk_fused_kernel(
565
T* s_group_scores = reinterpret_cast<T*>(ptr_u);
566
567
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
567
- asm volatile("griddepcontrol.wait;"); // I think all prolog can be put before
568
- // acqbulk because it's ptr arithmetic
568
+ cudaGridDependencySynchronize(); // I think all prolog can be put before
569
+ // acqbulk because it's ptr arithmetic
570
#endif
571
572
// phase 1: per-group scan
@@ -609,7 +610,7 @@ __global__ void grouped_topk_fused_kernel(
610
topk_values[i] = 1.0f / static_cast<float>(topk_i32);
611
}
612
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
612
- asm volatile("griddepcontrol.launch_dependents;");
613
+ cudaTriggerProgrammaticLaunchCompletion();
614
#endif
615
return;
616
}
@@ -670,7 +671,7 @@ __global__ void grouped_topk_fused_kernel(
671
}
672
673
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
673
- asm volatile("griddepcontrol.launch_dependents;");
674
+ cudaTriggerProgrammaticLaunchCompletion();
675
#endif
676
}
677
@@ -895,7 +896,8 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
896
int64_t const num_experts, int64_t const n_group,
897
int64_t const topk_group, int64_t const topk,
898
bool const renormalize, double const routed_scaling_factor,
898
- bool enable_pdl = false, cudaStream_t const stream = 0) {
899
+ const bool enable_pdl = false,
900
+ cudaStream_t const stream = 0) {
901
cudaLaunchConfig_t config;
902
config.stream = stream;
903
cudaLaunchAttribute attrs[1];
@@ -983,7 +985,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
985
int64_t const num_tokens, int64_t const num_experts, \
986
int64_t const n_group, int64_t const topk_group, int64_t const topk, \
987
bool const renormalize, double const routed_scaling_factor, \
986
- bool enable_pdl, cudaStream_t const stream);
988
+ const bool enable_pdl, cudaStream_t const stream);
989
990
INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
991
INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
@@ -1037,7 +1039,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
1039
scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
1040
auto topk_indices = torch::stable::new_empty(
1041
scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
1040
-
1042
+ const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens;
1043
const cudaStream_t stream =
1044
get_current_cuda_stream(scores.get_device_index());
1045
auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
@@ -1052,7 +1054,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
1054
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
1055
reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
1056
num_experts, n_group, topk_group, topk, renormalize, \
1055
- routed_scaling_factor, false, stream); \
1057
+ routed_scaling_factor, pdl_flag, stream); \
1058
break; \
1059
case vllm::moe::SCORING_SIGMOID: \
1060
vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
@@ -1061,7 +1063,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
1063
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
1064
reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
1065
num_experts, n_group, topk_group, topk, renormalize, \
1064
- routed_scaling_factor, false, stream); \
1066
+ routed_scaling_factor, pdl_flag, stream); \
1067
break; \
1068
default: \
1069
STD_TORCH_CHECK(false, "Unsupported scoring_func"); \
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+3
-3
@@ -173,7 +173,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
173
float row_chunk[VPT];
174
175
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
176
- asm volatile("griddepcontrol.wait;");
176
+ cudaGridDependencySynchronize();
177
#endif
178
179
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
@@ -300,7 +300,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
300
}
301
}
302
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
303
- asm volatile("griddepcontrol.launch_dependents;");
303
+ cudaTriggerProgrammaticLaunchCompletion();
304
#endif
305
return;
306
} else {
@@ -425,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
425
}
426
}
427
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
428
- asm volatile("griddepcontrol.launch_dependents;");
428
+ cudaTriggerProgrammaticLaunchCompletion();
429
#endif
430
}
431
}