[ROCm] Converge the two W4A16 linear kernels onto the upstream shell

[roberteg16]

Summary

Follow-up to #1175, which flagged the duplication and deliberately left the call to someone else:

The two kernels are duplicated work by the same author, and one of them is now dead weight on this branch. Converging them is a real decision, not a merge resolution [...] That call wants a benchmark and an owner.

This is that call, with the benchmark.

The two kernels were not interchangeable

Investigating turned up something #1175 did not: upstream's Python is ABI-incompatible with this fork's C++. Upstream hands the skinny GEMM unpacked activation-dtype zero points [N, groups]; the fork's C++ moved to packed int32 zp [N/8, groups] to cut 4x the DRAM traffic on a memory-bound kernel. Each side is self-consistent — the cross pairing #1175 created is not:

RuntimeError: Zero points must be int32 or uint32 (packed 8x uint4 along dim 0),
              got c10::Half

So two suites were already red on rogarcia.merge-upstream-96:

suite	before	after
test_rdna_hybrid_w4a16.py	71 passed, 12 failed	83 passed
test_w4a16_kernel_selection.py	0 passed, 2 failed	2 passed

The 12 are every asymmetric decode case. The 2 are upstream's test asserting a kernel selection that would have crashed AWQ decode — it only stayed latent because the fork's kernel sat first in the ROCm priority list.

This means "just delete the fork's kernel and keep upstream's" was never available: it would have shipped that crash.

What this does

Keeps upstream's module, class and op name (rdna_hybrid_w4a16.py, RDNAHybridW4A16LinearKernel, rdna_hybrid_w4a16_apply) as the single implementation, and ports the fork's contract and tuning into it:

• packed zero points (correctness — the blocker above);
• gfx1151 cache-cliff padding of the weight and metadata rows;
• the v_and_or_b32 packed fp16 dequant and its dtype-aware tile table;
• the scale/zp carrier that folds two per-group loads into one;
• the optional dequantized prefill cache behind --w4a16-prefill-dequant (otherwise that CLI flag becomes dead config).

Upstream's gfx12x tile branch is kept and is now actually reachable — the fork tested on_gfx1x() first, which covers gfx12, so the fork's own gfx12 branch was dead code.

Two upstream pieces are dropped on purpose:

• the K*M <= 32768 LDS gate, which guards a C++ constraint this fork's kernel no longer has (skinny_gemms_int4.cu streams the overflow from global). It would push M<=5, K>=8192 decode onto Triton for nothing.
• _GFX1X_PREFILL_OVERRIDES, tuned against upstream's scalar dequant body; it does not transfer to the packed one and would need re-sweeping to reapply.

Net: -1302 / +805 lines, one Triton kernel instead of two.

Performance: no regression, measured with a control

Ran test_hybrid_w4a16_perf.py (43 shapes x 14 batch sizes x 4 providers) against the unchanged golden/hybrid_w4a16_gfx1151.json, then re-ran the identical test on rogarcia.merge-upstream-96 as a control:

	base (fork kernel, untouched)	this PR
passed / failed	40 / 100	43 / 97
out-of-band improvements	304	290
out-of-band regressions	2	0

The control matters: the base branch already fails ~100 cases against these goldens, all improvements (median +13.7%), and 250 of them are on wvsplitk_int4 — the C++ decode kernel neither branch touches. The goldens are stale relative to the current ROCm stack, not broken by this change. Against that baseline the convergence is neutral-to-better, and it clears the 2 regressions the base had.

I did not regenerate the goldens: doing so in this PR would bake in the staleness and hide exactly what the test exists to catch. Rebaselining is worth a separate commit.

Test plan

• test_rdna_hybrid_w4a16.py — 83 passed (was 71 passed / 12 failed)
• test_w4a16_kernel_selection.py — 2 passed (was 0 passed / 2 failed)
• test_hybrid_w4a16_triton.py + test_rocm_compressed_tensors_w4a16.py — 43 passed
• test_hip_w4a16.py + test_moe_gemm_w4a16.py + test_hybrid_w4a16_moe.py — 196 passed
• test_hybrid_w4a16_perf.py vs unchanged goldens, with the base-branch control above
• ruff@0.14.0 check / format --check clean
• End-to-end vllm bench on AWQ models, against a base-branch control:

model	base merge-upstream-96	this PR
trymirai/SmolLM2-1.7B-Instruct-AWQ	166.2 tok/s	165.9 tok/s
hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4	47.0 tok/s	47.0 tok/s
Qwen3-VL-4B-Instruct-AWQ-4bit-lm_head_int8 (fp16, multimodal)	77.7 tok/s	77.8 tok/s
cyankiwi/Qwen3.6-27B-AWQ-INT4	fails: KV cache -1.01 GiB	fails identically

Decode tok/s (1000/median TPOT), --max-num-seqs 1. Coverage is deliberate on two axes: Llama-8B and Qwen3-VL run fp16 (the v_and_or_b32 packed dequant path) while SmolLM2 runs bf16 (the scalar path), and the models reach the kernel through both AutoAWQMarlinLinearMethod and CompressedTensorsWNA16. Every run logs the kernel actually selected, so this exercises the converged kernel end to end. Throughput is unchanged within noise. The Qwen 27B failure is pre-existing and byte-identical on both branches (Available KV cache memory: -1.01 GiB at --target-gpu-memory-gb 28) — a config headroom issue, not a regression from this PR.

All measurements on gfx1151 (AMD Radeon 8060S, ROCm 7.15, torch 2.12.0+rocm10.1.0).

[mgehre-amd]

@roberteg16 could you please rebase this and hand over to @eble-amd while you are away? We need to unify the w4a16 kernels into what we started to upstream to make further upstreaming possible.

[roberteg16]

@roberteg16 could you please rebase this and hand over to @eble-amd while you are away? We need to unify the w4a16 kernels into what we started to upstream to make further upstreaming possible.

Should be done

[mgehre-amd]

Thanks! I assumed you verified that there is no performance difference?

[roberteg16]

Thanks! I assumed you verified that there is no performance difference?

From the body of the PR:

model	base merge-upstream-96	this PR
trymirai/SmolLM2-1.7B-Instruct-AWQ	166.2 tok/s	165.9 tok/s
hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4	47.0 tok/s	47.0 tok/s
Qwen3-VL-4B-Instruct-AWQ-4bit-lm_head_int8 (fp16, multimodal)	77.7 tok/s	77.8 tok/s
cyankiwi/Qwen3.6-27B-AWQ-INT4	fails: KV cache -1.01 GiB	fails identically

I could expand the test suite for more coverage.

Edit: Let's wait for CI to build a wheel and I launch a bit wider e2e check

[roberteg16]

Running here: https://gitenterprise.xilinx.com/FaaSApps/rocm-scripts/actions/runs/3341745

[roberteg16]

Running here: https://gitenterprise.xilinx.com/FaaSApps/rocm-scripts/actions/runs/3341745

No regressions @mgehre-amd , the improvement of +48.5% on Gemma-4-31B-IT_VLM_OpenNav_AWQ-4bit_MTP_prefill-dequant and (other improvements) are just hiccups between runs.

Merging.