perf(rope): fix token/head-parallel dispatch threshold on CDNA3

[demandal25]

Summary

The RoPE kernels pick between a token-parallel and a head-parallel implementation at launch time, but the selection threshold was miscalibrated so the faster token-parallel kernel almost never ran. Retuning the threshold yields a measured 1.4–1.8x speedup at mid-to-large sequence lengths on MI300X (gfx942), with no change to small or very large sizes and no kernel logic changes.

What changed

• include/flashinfer/attention/generic/pos_enc.cuh — corrected the token- vs head-parallel dispatch threshold in BatchQKApplyRotaryPosIdsCosSinCache (cos/sin-cache path) and BatchQKApplyRotaryPosIds (on-the-fly path). The old threshold was nblks_x >= num_blocks_per_sm * num_sms, where num_blocks_per_sm came from gpuOccupancyMaxActiveBlocksPerMultiprocessor (8–16 on gfx942) — putting the bar at ~8–16x the CU count. The new threshold is nblks_x >= 2 * num_sms. This also removes the per-launch occupancy query from the hot path.

Architecture / design notes

The two kernels differ only in their grid mapping and resulting memory access pattern:

Kernel	Grid	Per-block access pattern
token-parallel (kernel_0)	(nblks_x,)	bdy contiguous tokens across all heads → contiguous DRAM
head-parallel (kernel_1)	(nblks_x, num_heads)	one head scattered across tokens → strided DRAM

The token-parallel layout streams memory contiguously and is markedly faster, but it launches num_headsx fewer blocks, so below a certain size it under-fills the device and the head-parallel form wins. The crossover is where the token-parallel grid has roughly enough blocks to fill all CUs, measured at nblks_x ~= 2 * num_sms on CDNA3. The threshold is device-count-relative (num_sms), so it adapts to other architectures; the 2x constant was tuned on gfx942 and is worth reconfirming on gfx950.

The old occupancy-derived threshold was both too high (suppressing the fast path) and pathological at specific sizes — e.g. the baseline apply_rope_with_cos_sin_cache was slower at seq_len 16384 (395 us) than at 32768 (388 us) because 16384 fell on the wrong side of it.

Benchmark results

Shape: head_size=128, rotary_dim=128, num_q_heads=32, num_kv_heads=8, bf16, batch=2, MI300X / gfx942. Latency in microseconds, median of bench_gpu_time.

apply_rope_with_cos_sin_cache_inplace (neox):

seq_len	before	after	speedup
4096	76.4	70.8	1.08x
8192	168.5	117.3	1.44x
16384	395.4	216.3	1.83x
32768	387.8	387.5	1.00x
65536	731.3	731.3	1.00x

apply_rope_pos_ids (on-the-fly, neox):

seq_len	before	after	speedup
4096	101.2	69.8	1.45x
8192	198.5	123.4	1.61x
16384	394.6	215.2	1.83x

Sizes at or below 2048 are unchanged (they correctly remain head-parallel). The largest sizes were already crossing to token-parallel under the old threshold, so they are unchanged.

Test plan

• pytest tests/rocm_tests/test_rope_hip.py — 13080 passed
• Before/after benchmarks above are real A/B measurements (stash baseline vs. tuned), not estimates
• pre-commit run --files include/flashinfer/attention/generic/pos_enc.cuh — passed (incl. clang-format)

[Copilot]

Pull request overview

This PR retunes the ROCm RoPE kernel launch-time dispatch heuristic so the faster token-parallel implementation is selected at realistic mid/large problem sizes on CDNA3, removing an overly strict occupancy-derived threshold that almost always forced the slower head-parallel path.

Changes:

• Replaced the occupancy-based token/head-parallel crossover condition with a simpler nblks_x >= 2 * num_sms threshold in both cos/sin-cache and on-the-fly RoPE paths.
• Removed the per-launch gpuOccupancyMaxActiveBlocksPerMultiprocessor query from these dispatch sites and documented the rationale/tradeoff inline.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.