[Bugfix][Kernel] Fix fused RoPE for head dimension 96

[tangzzycc]

Motivation

The fused RoPE + KV cache kernel currently fails to compile for head_dim=96 on wave32.

The previous vector-width calculation selects VEC_WIDTH=3. The ds_bpermute_b32 path handles the vector as 32-bit chunks and reconstructs only two BF16/FP16 elements, which causes:

ValueError: cannot broadcast shapes (2,) and (3,)

This is a real model configuration used by the Phi-3 Mini family. The existing comments also describe D=96 as a supported configuration.

Technical Details

• Select the smallest copy width from 1/2/4/8 that divides both the complete head and each NeoX half while fitting the head into one wave.
• On wave32, D=96 now uses VEC_WIDTH=4 and 24 active lanes.
• Keep the existing XOR lane mapping when vecs_per_half is a power of two.
• Use explicit addition/subtraction to pair the two NeoX halves for other lane spans, such as 12 lanes per half for D=96 on wave32.
• Add D=96 correctness tests for BF16/FP16 and flash/non-flash KV cache layouts.

The D=64, D=128, and D=256 paths retain their existing vector widths and XOR mapping.

Test Plan

Run the fused RoPE + KV cache test suite:

PYTHONPATH=. python3 -m pytest tests/kernels/test_fused_rope_cache.py -q

The change was tested on an AMD Radeon 8060S using native gfx1151 wave32 execution.

In addition to correctness testing, compare the generated ISA and CUDA Graph performance of the existing D=64/128/256 paths against origin/main.

Test Result

74 passed, 248 skipped

All four newly added D=96 combinations passed, and Q, K, key cache, and value cache matched the PyTorch reference.

The final gfx1151 ISA for D=64, D=128, and D=256 is byte-for-byte identical to origin/main. Interleaved CUDA Graph measurements showed no performance regression in those existing paths.

Representative D=96 CUDA Graph replay results for BF16, QH=KH=32:

Tokens	Latency
1	4.025 us
128	10.637 us
512	35.252 us
2048	456.657 us
4096	909.765 us

The generated D=96 kernel uses 35 VGPRs, 26 SGPRs, no LDS, no scratch, and no register spills.

A direct D=96 comparison is unavailable because AITER's public Triton fused kernel currently rejects non-power-of-two head dimensions.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

🟢 Approval recommended

The change is narrowly scoped to the failing head-dimension case, preserves existing power-of-two behavior, and adds targeted tests for the new configuration.

Pull request overview

Fixes a wave32 compilation failure in the fused RoPE + KV-cache kernel for the real-world configuration head_dim=96 by selecting a supported vector width and correctly pairing NeoX halves when the half-wave span is not a power of two.

Changes:

• Update VEC_WIDTH selection to choose the smallest supported copy width in {1,2,4,8} that divides both the full head and each NeoX half while fitting in a single wave.
• Generalize rotary pair-lane mapping: keep XOR pairing for power-of-two half spans, otherwise use explicit +/- half-span pairing (needed for D=96 on wave32).
• Add correctness coverage for head_dim=96 across bf16/f16 and flash/non-flash layouts.

File summaries

File	Description
tests/kernels/test_fused_rope_cache.py	Adds new parametrized tests covering head_dim=96 for bf16/f16 and flash/non-flash KV cache layouts.
kernels/attention/fused_rope_cache_kernel.py	Fixes wave32 head_dim=96 by enforcing native vector widths and correcting the NeoX pair-lane mapping for non-power-of-two half spans.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[tangzzycc]

Hi @zhiding512 @coderfeli, could you review this when you get a chance? Thanks!