Add Triton BF16xINT4 rowwise GEMM for ROCm gfx942/gfx950

[apicciau]

Summary

Adds Triton implementations of bf16i4bf16_shuffled, bf16i4bf16_shuffled_grouped,
and bf16i4bf16_shuffled_batched for AMD MI300X/MI350X. These ops exist on CUDA via
CUTLASS with a preshuffle weight layout; on ROCm the preshuffle layout is not
available, so this PR implements them using a row-major INT4 layout. The activation
matrix is split into even and odd K columns in Python (via optimised HIP strided
copies) before kernel launch — doing this split inside the kernel would require
stride-2 global memory loads, which halve effective memory bandwidth as adjacent
threads in a wavefront load non-adjacent addresses.

Context

• Hardware: AMD Instinct MI300X (gfx942), MI350X (gfx950, CDNA4)
• Depends on bf16i4bf16_rowwise (already merged)

Implementation

Fused A-operand dot. The straightforward port issues two tl.dot calls per
K-iteration — one against the lo-nibble weights with even-K activations, one against
the hi-nibble weights with odd-K activations. On gfx950 each dot requires a
ds_write_b128 → s_barrier → ds_read_b128 LDS round-trip to convert the A operand
from #blocked to #dot_op register layout. Concatenating both activation tiles
and both weight tiles along the K dimension via tl.cat before the dot halves the
number of LDS staging passes. This is mathematically equivalent:
[x_even | x_odd] @ [w_lo | w_hi]ᵀ = x_even @ w_lo.T + x_odd @ w_hi.T.

Split-K without atomics. When SPLIT_K > 1, partial sums are written into a
float32 workspace tensor [SPLIT_K, M, N] using direct stores — no
tl.atomic_add. A separate _bf16i4_splitk_reduce kernel sums the slices and
casts to bf16. This avoids atomic contention and is CUDA-graph safe.

The prune function restricts SPLIT_K > 1 to M < 512, so it has zero cost for
large-batch (prefill) shapes. For decode shapes the autotuner selects freely:
benchmarks on MI350X show the autotuner consistently picks SK=8 at M ≤ 128,
SK=4 at M=128 for large NK products, and SK=1 at M ≥ 2048 — matching the
occupancy argument (at M=1, N=8192, BLOCK_M=32, BLOCK_N=32 a non-split grid
has only 256 tiles on a 304-CU GPU; SK=8 raises this to 2048 tiles).

XCD remapping. Output tiles are redistributed across the 8 XCDs of MI300X/MI350X
before the L2-friendly M-grouping step, avoiding hot-spot cross-chiplet traffic when
the tile count is not a multiple of 8.

Pruned autotuner config sweep. The config space is structured to avoid redundant
combinations: num_warps is derived from tile area (≤2048 elems → 4 warps,
≥16384 → 8 warps, otherwise both), GROUP_SIZE_M is fixed at 8 (grid scheduling
hint, no inner-loop effect), and num_stages is fixed at 2 (the kernel does
explicit prefetching). This reduces the total config count from 1536 to 276 with no
loss of autotuning quality, cutting per-shape autotuning time by ~5×.

Native fused grouped kernel. The grouped op dispatches all G groups in a single
kernel launch (_bf16i4_grouped_kernel) using a 3D grid
(G, ceil(M_total/BLOCK_M), ceil(N/BLOCK_N)). The grid is over-provisioned using
M_total (a host int from tensor metadata) — tiles that exceed their group's row
count early-exit in the kernel. M_starts is computed on-device via
M_sizes.cumsum(0) - M_sizes, so there are no device-to-host synchronisations and
the wrapper captures cleanly inside a CUDA graph.

Schema deduplication. The bf16i4bf16_* op schemas and preshuffle_i4 were
duplicated inside both #ifdef USE_ROCM and #else in gemm_ops.cpp. They are
now registered once, above the #ifdef, since the schemas are shared between CUDA
and ROCm builds. TritonBF16Int4GroupedShuffled in gemm_ops.py inherits from
GemmOpBase directly — the previous CutlassFP8Int4Rowwise base was a copy-paste
artefact; the class overrides every method so no code was shared.

Testing

All 9 BF16×INT4 accuracy tests pass:

9 passed, 204 deselected in 90s

Tests cover test_rowwise_accuracy, test_rowwise_batched_accuracy, and
test_torch_op_dispatch, three shapes each. The grouped op is tested via the Triton
kernel directly (the bf16i4bf16_shuffled_grouped torch op schema requires the C++
build; the kernel itself is validated by the Python-level tests in
BF16Int4TritonROCmGroupedTests).

Performance

Hardware: AMD Instinct MI350X (gfx950), ROCm 7.1.1, Triton 3.5.1+rocm7.1.1.
triton.testing.do_bench_cudagraph, rep=200, return_mode=mean.
Peak: 2300 TFLOPS BF16 matrix, 8000 GB/s HBM.
Ridge point: 287.5 FLOP/byte — shapes below are memory-bandwidth bound (report
GB/s, % HBM), shapes above are compute bound (report TFLOPS, % BF16 peak).

bf16i4bf16_shuffled — decode shapes (memory-bandwidth bound)

Model-attributed shapes from the MSLK registered shape sets. "Before" is the
bf16i4bf16_rowwise kernel on main (no tl.cat fusion, no XCD remapping,
no split-K).

Model	M	N	K	SK	Before (μs)	After (μs)	Speedup	GB/s	% HBM
llama4	1	896	5120	8	42.4	15.6	2.72x	170	2.1%
llama4	16	896	5120	8	42.7	18.7	2.28x	197	2.5%
llama4	64	896	5120	8	43.8	19.2	2.28x	366	4.6%
llama4	128	2048	5120	8	45.1	30.0	1.50x	817	10.2%
llama3_70b	1	8192	3584	8	32.4	27.3	1.19x	625	7.8%
llama3_70b	16	8192	3584	8	33.3	32.0	1.04x	790	9.9%
llama3_70b	64	8192	3584	8	39.2	33.9	1.16x	1521	19.0%
llama3_70b	128	7168	8192	4	82.5	75.0	1.10x	884	11.1%

The autotuner selects SK>1 only where it helps occupancy (M≤128) and SK=1
for prefill shapes where the overhead would exceed the benefit.

bf16i4bf16_shuffled — prefill shapes (compute bound)

Model	M	N	K	Before (TFLOPS)	After (TFLOPS)	Speedup	% BF16 peak
llama3_70b	2048	8192	3584	511.7	591.2	1.16x	25.7%
llama3_70b	2048	7168	8192	496.8	619.9	1.25x	27.0%
llama3_70b	4096	7168	8192	534.2	609.9	1.14x	26.5%

bf16i4bf16_shuffled_grouped — new op on ROCm

Single fused launch for all G groups; CUDA-graph safe. M/g = rows per group.

Model	G	M/g	N	K	μs	TFLOPS	% BF16 peak
llama4	2	64	2048	5120	112.7	23.8	1.0%
llama4	4	64	2048	5120	140.3	38.3	1.7%
llama4	2	128	2048	5120	117.0	45.9	2.0%
llama4	4	128	2048	5120	141.5	75.9	3.3%
llama3_70b	2	128	8192	3584	159.2	94.4	4.1%
llama3_70b	4	128	8192	3584	321.0	93.7	4.1%
llama3_70b	2	2048	7168	8192	2327.4	206.7	9.0%
llama3_70b	4	2048	7168	8192	4522.8	212.7	9.2%

bf16i4bf16_shuffled_batched — new op on ROCm

Model	B	M	N	K	SK	μs	TFLOPS	% BF16 peak
llama4	2	128	2048	5120	8	65.3	82.2	3.6%
llama4	4	128	2048	5120	8	127.2	84.4	3.7%
llama3_70b	2	128	8192	3584	1	88.3	170.2	7.4%
llama3_70b	4	128	8192	3584	1	173.0	173.7	7.6%
llama3_70b	2	2048	7168	8192	1	818.5	587.7	25.6%
llama3_70b	4	2048	7168	8192	1	1689.2	569.5	24.8%

[cthi]

@apicciau can you rebase this branch? thanks.

[apicciau]

@cthi, I don't think the failures are related to this PR...the job fails before reaching any of the code I've edited. Can we rerun the CI?

[cthi]

I think there is some existing test failure, let me try to resolve it.

[meta-codesync]

@cthi has imported this pull request. If you are a Meta employee, you can view this in D112191149.

[apicciau]

Hi @cthi! Do you have any news on this? I have a follow-up PR that sits on top of this change. Let me know how to help

[cthi]

We only want to support cuda graphable APIs largely in MSLK. E.g. we should not have any D2H/H2D sync in the API.

[apicciau]

Hi @cthi, I think the new failures are ROCm 7.0-specific, because that image ship an older Triton. I'm adding a workaround for that case.

[cthi]

Thanks, left few more comments.

[apicciau]

Hi @cthi, let me know if there is any issue with this code

[q10]

@apicciau could you rebase onto latest main to resolve the merge conflict? Thanks