From f68920ab407fb46234c4a49b2b82f522420f9c96 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 24 Jun 2026 04:58:56 +0000 Subject: [PATCH] perf(rope): fix token/head-parallel dispatch threshold on CDNA3 The RoPE dispatchers chose between the token-parallel and head-parallel cos/sin kernels using an occupancy-derived threshold (num_blocks_per_sm * num_sms), which set the bar at ~8-16x the CU count. This effectively never selected the token-parallel kernel even at sequence lengths where it is substantially faster: it keeps each block's accesses contiguous across heads, whereas the head-parallel kernel scatters one head across tokens (strided DRAM). Retune to nblks_x >= 2 * num_sms, the measured crossover on gfx942 where the token-parallel grid has enough blocks to fill the device. Also drops the per-launch gpuOccupancyMaxActiveBlocksPerMultiprocessor query from the hot path. Measured 1.4-1.8x speedup at seq_len 8192-16384 on MI300X for both the cos/sin-cache and on-the-fly paths; small and very large sizes unchanged. Co-Authored-By: Claude Opus 4.7 --- .../flashinfer/attention/generic/pos_enc.cuh | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/include/flashinfer/attention/generic/pos_enc.cuh b/include/flashinfer/attention/generic/pos_enc.cuh index f22d2c90b6..4d7ed202de 100644 --- a/include/flashinfer/attention/generic/pos_enc.cuh +++ b/include/flashinfer/attention/generic/pos_enc.cuh @@ -606,12 +606,14 @@ gpuError_t BatchQKApplyRotaryPosIdsCosSinCache( auto kernel_0 = BatchQKApplyRotaryPosIdsCosSinCacheKernel; - int num_blocks_per_sm_0 = 0; - FI_GPU_CALL(gpuOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm_0, kernel_0, - num_threads, /*smem_size=*/0)); - uint32_t num_ctas_0 = num_blocks_per_sm_0 * num_sms; - - if ((nnz + bdy - 1) / bdy >= num_ctas_0) { + // Pick token-parallel (kernel_0) once it launches enough blocks to fill the + // device. kernel_0 gives each block bdy contiguous tokens across all heads + // (contiguous DRAM), while kernel_1 scatters one head across tokens (strided). + // The token-parallel layout is markedly faster (measured ~1.4x at seq_len 8192 + // on gfx942) whenever it has enough blocks; below that it under-fills the GPU + // and the head-parallel form (num_heads x more blocks) wins. Crossover measured + // at nblks_x ~= 2 * num_sms on CDNA3 (independent of head count). + if (nblks_x >= 2u * static_cast(num_sms)) { dim3 nblks(nblks_x); dim3 nthrs(bdx, bdy); FI_GPU_CALL(gpuLaunchKernel((void*)kernel_0, nblks, nthrs, args, 0, stream)); @@ -677,11 +679,10 @@ gpuError_t BatchQKApplyRotaryPosIds( auto kernel_0 = BatchQKApplyRotaryPosIdsKernel; - int num_blocks_per_sm_0 = 0; - FI_GPU_CALL(gpuOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm_0, kernel_0, - num_threads, /*smem_size=*/0)); - uint32_t num_ctas_0 = num_blocks_per_sm_0 * num_sms; - if (nblks_x >= num_ctas_0) { + // Same token- vs head-parallel tradeoff as the cos/sin-cache path: token-parallel + // (kernel_0) keeps each block's memory accesses contiguous across heads and wins + // once it has enough blocks to fill the device (~2 * num_sms on CDNA3). + if (nblks_x >= 2u * static_cast(num_sms)) { dim3 nblks(nblks_x); dim3 nthrs(bdx, bdy);