[Kernel] PDL enablement for fusedQKNormRopeKernel (#55755)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Jee Jee Li committed
Sep 8, 2026 at 11:49 UTC
9c297e3b8bc8957d591ef7c244890c543cb9b19f
1 file changed
+112
-48
csrc/libtorch_stable/fused_qknorm_rope_kernel.cu
+112
-48
@@ -189,6 +189,12 @@ __global__ void fusedQKNormRopeKernel(
189
}
190
int offsetThread = offsetWarp + laneId * numElemsPerThread;
191
192
+ // PDL: wait for the predecessor kernel (the QKV projection that produces
193
+ // `qkv`) to complete before touching any global memory.
194
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
195
+ cudaGridDependencySynchronize();
196
+#endif
197
+
198
// Sum of squares for RMSNorm
199
float sumOfSquares = 0.0f;
200
@@ -300,6 +306,11 @@ __global__ void fusedQKNormRopeKernel(
306
*reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
307
}
308
309
+ // PDL: signal completion so a dependent successor may launch early.
310
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
311
+ cudaTriggerProgrammaticLaunchCompletion();
312
+#endif
313
+
314
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
315
}
316
#endif
@@ -385,6 +396,13 @@ __global__ void fusedQKNormRopeKernelNTokenHeads(
396
smem_storage + cos_sin_bytes +
397
warpId * (HEADS_PER_WARP * qkv_tile_bytes);
398
399
+ // PDL: wait for the predecessor kernel (the QKV projection that produces
400
+ // `qkv`) to complete before touching any global memory -- cp.async reads
401
+ // global memory, so this must precede the first commit group.
402
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
403
+ cudaGridDependencySynchronize();
404
+#endif
405
+
406
// === Group 0: async load all heads' QKV into smem (issued first). ===
407
for (int k = 0; k < num_heads_this_warp; ++k) {
408
int const localHeadIdx = first_head + k;
@@ -531,6 +549,11 @@ __global__ void fusedQKNormRopeKernelNTokenHeads(
549
}
550
}
551
552
+ // PDL: signal completion so a dependent successor may launch early.
553
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
554
+ cudaTriggerProgrammaticLaunchCompletion();
555
+#endif
556
+
557
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
558
}
559
#endif
@@ -547,6 +570,42 @@ __global__ void fusedQKNormRopeKernelNTokenHeads(
570
__VA_ARGS__ \
571
}
572
573
+namespace {
574
+// The empirical value for small batch: overlapping with the predecessor QKV
575
+// projection only pays off while this kernel is launch-bound.
576
+constexpr int kPDLEnableTokens = 16;
577
+
578
+// Launch helper for PDL (programmatic dependent launch), which lets this
579
+// kernel start while the predecessor QKV projection is still draining. The
580
+// launch attribute is the only switch: with it off, the kernels'
581
+// cudaGridDependencySynchronize() / cudaTriggerProgrammaticLaunchCompletion()
582
+// are no-ops.
583
+template <typename KernelFn, typename... Args>
584
+void launchWithPDL(KernelFn kernel, dim3 gridDim, dim3 blockDim,
585
+ size_t smem_bytes, cudaStream_t stream, bool enable_pdl,
586
+ Args... args) {
587
+#ifndef USE_ROCM
588
+ cudaLaunchConfig_t config;
589
+ config.gridDim = gridDim;
590
+ config.blockDim = blockDim;
591
+ config.dynamicSmemBytes = smem_bytes;
592
+ config.stream = stream;
593
+ cudaLaunchAttribute attrs[1];
594
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
595
+ attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
596
+ config.numAttrs = 1;
597
+ config.attrs = attrs;
598
+ cudaLaunchKernelEx(&config, kernel, args...);
599
+#else
600
+ // ROCm: standard kernel launch syntax (no PDL/stream serialization).
601
+ (void)enable_pdl;
602
+ // clang-format off
603
+ kernel<<<gridDim, blockDim, smem_bytes, stream>>>(args...);
604
+ // clang-format on
605
+#endif
606
+}
607
+} // namespace
608
+
609
template <typename scalar_t_in, typename scalar_t_cache>
610
void launchFusedQKNormRope(void* qkv, int const num_tokens,
611
int const num_heads_q, int const num_heads_k,
@@ -562,29 +621,33 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
621
int const gridSize = common::divUp(totalWarps, warpsPerBlock);
622
dim3 gridDim(gridSize);
623
dim3 blockDim(blockSize);
624
+ bool const enable_pdl = num_tokens <= kPDLEnableTokens;
625
switch (head_dim) {
626
case 64:
627
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
568
- fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
569
- <<<gridDim, blockDim, 0, stream>>>(
570
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
571
- k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
628
+ launchWithPDL(
629
+ fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>,
630
+ gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q,
631
+ num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache,
632
+ position_ids, num_tokens, rotary_dim);
633
});
634
break;
635
case 128:
636
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
576
- fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
577
- <<<gridDim, blockDim, 0, stream>>>(
578
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
579
- k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
637
+ launchWithPDL(
638
+ fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>,
639
+ gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q,
640
+ num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache,
641
+ position_ids, num_tokens, rotary_dim);
642
});
643
break;
644
case 256:
645
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
584
- fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
585
- <<<gridDim, blockDim, 0, stream>>>(
586
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
587
- k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
646
+ launchWithPDL(
647
+ fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>,
648
+ gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q,
649
+ num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache,
650
+ position_ids, num_tokens, rotary_dim);
651
});
652
break;
653
default:
@@ -644,6 +707,7 @@ void launchFusedQKNormRopeNTokenHeads(
707
int const gridSize = common::divUp(total_warps, warpsPerBlock);
708
dim3 gridDim(gridSize);
709
dim3 blockDim(blockSize);
710
+ bool const enable_pdl = num_tokens <= kPDLEnableTokens;
711
// Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same
712
// layout).
713
size_t const cache_elem_size =
@@ -656,42 +720,42 @@ void launchFusedQKNormRopeNTokenHeads(
720
warpsPerBlock * static_cast<size_t>(rotary_dim) * cache_elem_size +
721
warpsPerBlock * qkv_smem_per_warp;
722
659
-#define LAUNCH_N_TOKEN_HEADS(N) \
660
- do { \
661
- switch (head_dim) { \
662
- case 64: \
663
- DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
664
- fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 64, \
665
- INTERLEAVE, (N)> \
666
- <<<gridDim, blockDim, smem_bytes, stream>>>( \
667
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
668
- k_weight, cos_sin_cache, position_ids, num_tokens, \
669
- rotary_dim); \
670
- }); \
671
- break; \
672
- case 128: \
673
- DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
674
- fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 128, \
675
- INTERLEAVE, (N)> \
676
- <<<gridDim, blockDim, smem_bytes, stream>>>( \
677
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
678
- k_weight, cos_sin_cache, position_ids, num_tokens, \
679
- rotary_dim); \
680
- }); \
681
- break; \
682
- case 256: \
683
- DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
684
- fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 256, \
685
- INTERLEAVE, (N)> \
686
- <<<gridDim, blockDim, smem_bytes, stream>>>( \
687
- qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
688
- k_weight, cos_sin_cache, position_ids, num_tokens, \
689
- rotary_dim); \
690
- }); \
691
- break; \
692
- default: \
693
- STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
694
- } \
723
+#define LAUNCH_N_TOKEN_HEADS(N) \
724
+ do { \
725
+ switch (head_dim) { \
726
+ case 64: \
727
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
728
+ launchWithPDL( \
729
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \
730
+ 64, INTERLEAVE, (N)>, \
731
+ gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \
732
+ num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \
733
+ cos_sin_cache, position_ids, num_tokens, rotary_dim); \
734
+ }); \
735
+ break; \
736
+ case 128: \
737
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
738
+ launchWithPDL( \
739
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \
740
+ 128, INTERLEAVE, (N)>, \
741
+ gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \
742
+ num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \
743
+ cos_sin_cache, position_ids, num_tokens, rotary_dim); \
744
+ }); \
745
+ break; \
746
+ case 256: \
747
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
748
+ launchWithPDL( \
749
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \
750
+ 256, INTERLEAVE, (N)>, \
751
+ gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \
752
+ num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \
753
+ cos_sin_cache, position_ids, num_tokens, rotary_dim); \
754
+ }); \
755
+ break; \
756
+ default: \
757
+ STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
758
+ } \
759
} while (0)
760
761
if (token_heads_per_warp == 2) {