[ROCm][Kernel] Add GFX11 (RDNA3/4) support for wvSplitK skinny GEMM kernels

[mgehre-amd]

Summary

Enable the wvSplitK skinny GEMM kernels on all GFX11 (RDNA3) GPUs. These kernels accelerate decode-phase GEMMs (small batch size, large K) by using wave-level split-K with butterfly reduction.

Key adaptations for RDNA's wave32 execution model:

• Use v_dot2acc_f32_f16 (renamed from v_dot2c_f32_f16 on GFX11)
• Wave32 butterfly reduction via __shfl_xor instead of GFX9's DPP row_shr
• Use THRDS=32 on GFX11 (one wavefront per row); keep THRDS=64 on GFX9
• Add is_gfx11() runtime helper for host-side kernel dispatch
• Add on_gfx11() helper and use on_gfx9() or on_gfx11() to gate the Python dispatch path
• Relax wvSplitK test tolerance to account for fp16 accumulation error scaling with sqrt(K)

Performance

Benchmarked on AMD Strix Halo (gfx1151, LPDDR5X-8000 128 GB):

Qwen/Qwen3-4B (bf16, unquantized) — decode TPOT ~20% faster:

Workload	Metric	Baseline	Optimized	Change
input=128, output=128	Median TPOT	50.03 ms	40.18 ms	19.7% faster
input=1920, output=128	Median TPOT	51.39 ms	41.37 ms	19.5% faster

hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4 — minimal impact (AWQ layers bypass skinny GEMM path):

Workload	Metric	Baseline	Optimized	Change
input=128, output=128	Median TPOT	172.22 ms	169.49 ms	1.6% faster
input=1920, output=128	Median TPOT	173.33 ms	170.54 ms	1.6% faster

TTFT (prefill) is unaffected as the skinny kernels only target batch_size ≤ 4 GEMMs.

Test plan

• test_rocm_wvsplitk_kernel — all 22 parametrized cases pass on gfx1151
• test_rocm_wvsplitk_bias1D_kernel — all 22 parametrized cases pass
• test_rocm_wvsplitk_bias2D_kernel — all 22 parametrized cases pass
• Sanity check: model produces correct answers to math questions in both benchmarks

[gemini-code-assist]

Code Review

This pull request successfully adds support for GFX11 (RDNA3/4) GPUs to the wvSplitK skinny GEMM kernels, which should provide a significant performance boost for decode-phase GEMMs on these architectures. The implementation correctly handles the instruction renaming and introduces a new wave32 reduction path. The Python dispatch logic and test tolerances are also updated accordingly. My main feedback is regarding significant code duplication in the CUDA/HIP kernels. I've added a high-severity comment suggesting a refactoring to improve long-term maintainability.

I am having trouble creating individual review comments. Click here to see my feedback.

csrc/rocm/skinny_gemms.cu (511-545)

This block of code for GFX11 reduction is substantially duplicated across three kernels: wvSplitK_hf_sml_, wvSplitK_hf_, and wvSplitK_hf_big_. This duplication negatively impacts maintainability, as any future changes to this logic will need to be manually applied in three different places, increasing the risk of introducing inconsistencies or bugs.

To address this, I recommend refactoring the GFX11-specific reduction and output logic into a shared __device__ function. This function could be parameterized to handle the minor differences between kernels, such as the conditional commit based on commitColumn.

Additionally, the GFX11-specific shared memory allocation logic is also duplicated and could be similarly refactored into a helper struct or macro to improve code reuse and clarity.

[mergify]

⚠️ The sha of the head commit of this PR conflicts with #34177. Mergify cannot evaluate rules on this PR. ⚠️