[ROCm] Fused 4-bit SIMT GEMM

[sstamenk]

Summary

This PR adds a fused 4-bit SIMT dequant+GEMM path for ROCm/RDNA and wires it into gemm_4bit dispatch for small inference batch sizes. The kernel avoids the dequant+BLAS fallback overhead by decoding NF4/FP4 weights inside the GEMM, reusing decoded values across batch rows, and using RDNA-friendly primitives such as DPP reductions and LDS centroid lookup.

The implementation is validated across gfx1100, gfx1201, and gfx1151 using a 40-shape LLM sweep over bf16, fp16, and fp32, comparing SIMT against dequant+BLAS, the baseline gemv implementation, and the optimized gemv from #1920. The results show SIMT consistently wins in the small-M inference regime, with architecture-specific crossover behavior. This PR does not include dispatch heuristics that take into account the architecture, shape, data type and other data but instead opts for a single value per architecture that is determined from the 40 shape median crossover values. Future work will add more intelligent heuristics that will cover the WMMA/MFMA kernels as well.

Across the 40-shape RDNA sweep, the fused SIMT kernel delivers up to 44.8x speedup over dequant+BLAS (gfx1100, fp32, small GQA k/v, M=1), up to 9.9x over the baseline gemv (gfx1201, bf16, qwen3-235b-moe expert down, M=64), and up to 4.5x over the optimized #1920 gemv (gfx1100, bf16, small GQA k/v, M=64).

Technical details

• HIP fused SIMT enablement: compiles and exposes cgemm_4bit_* on ROCm using shared bnb_bfloat16, bnb_bfloat162, and bnb_stream_t compatibility aliases.

• RDNA bf16 dot2 path: bf16 uses native v_dot2_f32_bf16 on gfx11/gfx12, accumulating bf16 pairs into fp32 and applying scale once per chunk.

• RDNA fp16 dot2 path: fp16 uses native v_dot2_f32_f16 via __ockl_fdot2, reducing the previous packed-multiply/mix/add overhead.

• CDNA/gfx9 SIMT tuning: gfx9 keeps bf16 VDOT2 disabled, uses fp32 FMA accumulation for fp16/bf16 pair paths, unscaled centroids where beneficial, and packed fp16 accumulation.

• LDS centroid lookup: HIP stages NF4/FP4 centroids in LDS for decode-heavy paths, avoiding repeated shuffle/bpermute lookup overhead.

• DPP reduction: HIP uses DPP row-shift reductions plus a cross-row exchange. The kernel uses 32-lane logical warps, so this works on both RDNA wave32 and CDNA wave64.

• Register prefetch: HIP prefetches the next packed-weight chunk and scale into registers while accumulating the current K group.

• ROCm compatibility: supports older ROCm 6.x builds by using HIP_ENABLE_WARP_SYNC_BUILTINS, host-safe bf16 aliases, and avoiding host inclusion of ROCm 6.x hip_bf16.h.

• Raw-pointer safety: the Python wrapper makes A contiguous before passing A.data_ptr() into the fused kernel, which assumes row-major contiguous activation rows.

Testing plan

• 40-shape performance sweep across RDNA architectures
  Run the full cold-cache sweep over the 40 canonical LLM shapes for bf16, fp16, and fp32, comparing:
  • gemv (baseline) — current gemv implementation slightly modified into a tiled kernel that supports M > 1
  • gemv ([ROCm] Optimize kgemm_4bit_inference_naive for ROCm, use it for batch sizes other than 1 #1920) — optimized multi-row gemv
  • SIMT — fused 4-bit dequant + GEMM kernel
  • dequant+BLAS — current fallback baseline
    The sweep is run across the tested RDNA targets (gfx1100, gfx1201, gfx1151, gfx942) and summarized with per-shape grids plus median-over-shapes aggregate plots.
• NVIDIA regression testing
  Run the same gemm_4bit correctness and smoke-performance checks on NVIDIA hardware to ensure the shared C++/Python dispatch changes do not regress the existing CUDA path or alter the CUDA MMA/SIMT selection behavior.
• Unit test / correctness validation
  Run the relevant unit tests for gemm_4bit, gemv_4bit, quantization layouts, and fallback correctness to verify outputs match dequantized reference results across supported dtypes, quantization modes, nested/non-nested absmax, bias handling, and dispatch fallback cases.

Testing results

Median results over the 40 comparing different kernel performance

• gfx1100 - full per shape data [fp16, bf16, fp32]

• gfx1201 - full per shape data [fp16, bf16, fp32]

• gfx1151 - full per shape data [fp16, bf16, fp32]

• gfx942 - full per shape data [fp16, bf16, fp32]

Nvidia regression testing

• No performance regressions observed

Correctness validation

• All tests successfully passed

Known limitations

• RDNA memory-camping effects can dominate specific N values. The SIMT kernel is bandwidth-bound at small M, so some output dimensions alias poorly onto the memory-channel/partition layout and create large fixed stalls that are mostly independent of work size. This is visible as isolated outlier shapes in the per-shape charts: for example, gfx1201 shows the strongest camp behavior on power-of-two N values such as the small square / N=2048 case, while gfx1100 is more sensitive around multiples of 3072. Strix Halo (gfx1151) shows both families in the resonance sweep.

• Example RDNA memory-camping outlier: on gfx1201, the N=2048 small square shape aliases poorly and forces the SIMT kernel into a fixed ~27 µs latency floor extending up to M = 8.

• Latency shape sweep showing spikes at specific sizes

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[matthewdouglas]

This looks quite promising, thank you! I may have a few small nits to add, but also the builds seem to have an issue right now. There's also some very minor linting issues needing to be fixed.

[sstamenk]

I have some MI related changes pending, I also managed to squeeze out a bit more performance

[sstamenk]

Added the gfx942 data

[sstamenk]

@matthewdouglas can you retrigger the builds? I pushed some fixes for older ROCm versions

[matthewdouglas]

@sstamenk It looks like some of the 6.x builds are working now, but still some issues with 7.x.

[sstamenk]

@matthewdouglas should be fixed now

edit: misread the logs, looking into the failure right now

[sstamenk]

Fixed the header include issues.

I've also added some data for e2e decode performance showing +50% decode throughput improvement on gfx1201 for Qwen 3.6 27b in 4bit vs the current gemv kernel.