[Perf][CUDA] Programmatic dependent launch for the DSA decode kernels (#50230)
Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
xiaozhoupy committed
Aug 6, 2026 at 15:53 UTC
2dfb8ba59098eb489197e1b4c643addffd51592e
2 files changed
+54
-10
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+40
-10
@@ -49,6 +49,11 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
49
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
50
"Vec size is not matched.");
51
52
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
53
+ cudaGridDependencySynchronize();
54
+ cudaTriggerProgrammaticLaunchCompletion();
55
+#endif
56
+
57
// Precompute SF layout parameter (constant for entire kernel).
58
int32_t const numKTiles = (outputCols + 63) / 64;
59
@@ -123,6 +128,11 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
128
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
129
"Vec size is not matched.");
130
131
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
132
+ cudaGridDependencySynchronize();
133
+ cudaTriggerProgrammaticLaunchCompletion();
134
+#endif
135
+
136
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
137
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
138
@@ -224,10 +234,21 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
234
input.scalar_type(), "nvfp4_quant_kernel", [&] {
235
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
236
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
227
- vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
228
- m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr,
229
- reinterpret_cast<uint32_t*>(output_ptr),
230
- reinterpret_cast<uint32_t*>(sf_out));
237
+ cudaLaunchConfig_t config = {};
238
+ config.gridDim = grid;
239
+ config.blockDim = block;
240
+ config.dynamicSmemBytes = 0;
241
+ config.stream = stream;
242
+ cudaLaunchAttribute attrs[1];
243
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
244
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
245
+ config.numAttrs = 1;
246
+ config.attrs = attrs;
247
+ cudaLaunchKernelEx(&config, vllm::cvt_fp16_to_fp4<cuda_type, false>,
248
+ m, n, output_n, num_padded_cols, input_ptr,
249
+ input_sf_ptr,
250
+ reinterpret_cast<uint32_t*>(output_ptr),
251
+ reinterpret_cast<uint32_t*>(sf_out));
252
});
253
} else {
254
int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
@@ -240,12 +261,21 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
261
input.scalar_type(), "nvfp4_quant_kernel", [&] {
262
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
263
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
243
- vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
244
- <<<grid, block, 0, stream>>>(
245
- m, n, output_n, output_sf_n_unpadded, num_packed_cols,
246
- input_ptr, input_sf_ptr,
247
- reinterpret_cast<uint32_t*>(output_ptr),
248
- reinterpret_cast<uint32_t*>(sf_out));
264
+ cudaLaunchConfig_t config = {};
265
+ config.gridDim = grid;
266
+ config.blockDim = block;
267
+ config.dynamicSmemBytes = 0;
268
+ config.stream = stream;
269
+ cudaLaunchAttribute attrs[1];
270
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
271
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
272
+ config.numAttrs = 1;
273
+ config.attrs = attrs;
274
+ cudaLaunchKernelEx(
275
+ &config, vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>, m, n,
276
+ output_n, output_sf_n_unpadded, num_packed_cols, input_ptr,
277
+ input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
278
+ reinterpret_cast<uint32_t*>(sf_out));
279
});
280
}
281
}
vllm/models/deepseek_v32/common/kernels.py
+14
@@ -150,9 +150,13 @@ def _fused_norm_rope_kernel(
150
TOPK_BLOCK_SIZE: tl.constexpr,
151
HAS_INDEXER: tl.constexpr,
152
INDEX_ROPE_INTERLEAVE: tl.constexpr,
153
+ USE_PDL: tl.constexpr,
154
):
155
pid = tl.program_id(0)
156
tok_idx = tl.program_id(1)
157
+ if USE_PDL:
158
+ tl.extra.cuda.gdc_wait()
159
+ tl.extra.cuda.gdc_launch_dependents()
160
if pid == 3:
161
if not HAS_INDEXER:
162
# Shared layer: reuse the previous indexer layer's top-k; do not
@@ -474,6 +478,7 @@ def fused_norm_rope(
478
479
if q_c_out is None:
480
q_c_out = torch.empty_like(q_c)
481
+ use_pdl = current_platform.is_arch_support_pdl()
482
_fused_norm_rope_kernel[(4, num_tokens)](
483
positions,
484
# Q RMS norm
@@ -532,6 +537,8 @@ def fused_norm_rope(
537
TOPK_BLOCK_SIZE=1024,
538
HAS_INDEXER=has_indexer,
539
INDEX_ROPE_INTERLEAVE=index_rope_interleave,
540
+ USE_PDL=use_pdl,
541
+ launch_pdl=use_pdl,
542
)
543
return q_c_out
544
@@ -585,10 +592,14 @@ def _fused_q_kernel(
592
HAS_INDEXER: tl.constexpr,
593
INDEX_ROPE_INTERLEAVE: tl.constexpr,
594
QUANTIZE_MQA: tl.constexpr,
595
+ USE_PDL: tl.constexpr,
596
):
597
pid = tl.program_id(0)
598
tok_idx = tl.program_id(1)
599
head_idx = tl.program_id(2)
600
+ if USE_PDL:
601
+ tl.extra.cuda.gdc_wait()
602
+ tl.extra.cuda.gdc_launch_dependents()
603
604
if pid == 2:
605
# ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -857,6 +868,7 @@ def fused_q(
868
)
869
return index_q_fp8, index_weights_out, mqa_q
870
871
+ use_pdl = current_platform.is_arch_support_pdl()
872
_fused_q_kernel[(3, num_tokens, grid_heads)](
873
positions,
874
q_pe,
@@ -898,6 +910,8 @@ def fused_q(
910
HAS_INDEXER=has_indexer,
911
INDEX_ROPE_INTERLEAVE=index_rope_interleave,
912
QUANTIZE_MQA=quantize_mqa,
913
+ USE_PDL=use_pdl,
914
+ launch_pdl=use_pdl,
915
# num_warps=1 is optimal here: each program is a single 128-element
916
# rope+quant, so the kernel is program-count/occupancy bound, not
917
# per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).