[ROCm][Perf] W4A16: packed fp16 prefill dequant

[mgehre-amd]

[ROCm][Perf] W4A16: packed fp16 prefill dequant and gfx1151 tile selection

Purpose

RDNAHybridW4A16LinearKernel routes prefill (M > 5) through a Triton fused-dequant GEMM.
On RDNA3 that inner loop is VALU-issue-bound: the dequant instructions, not the WMMAs
and not memory, set the rate. This PR attacks the instruction count and retunes the tiles
that the new dequant wants.

Performance

Same build, only rdna_hybrid_w4a16.py swapped between arms; arms interleaved
(before, after, before, ...) so thermal and clock drift hit both. 3 reps per arm,
--input-len 2048 --output-len 1 --num-prompts 10 (prefill-weighted, so M is well past the
skinny-decode threshold). Medians.

model	dtype	quant	metric	before	after	delta	rep ranges
RedHatAI/Qwen3-8B-quantized.w4a16	fp16	asym	prefill tok/s	1408.77	1519.02	+7.83%	disjoint
			TTFT ms	1454	1348	−7.29%	disjoint
RedHatAI/Qwen3-4B-quantized.w4a16	fp16	sym	prefill tok/s	2717.46	3071.02	+13.01%	disjoint
			TTFT ms	754	667	−11.54%	disjoint
RedHatAI/Qwen3-8B-quantized.w4a16	bf16	asym	prefill tok/s	1348.53	1348.74	+0.02%	overlap
			TTFT ms	1519	1518	−0.07%	overlap

The bf16 row is the control, it is not affected by this PR.

1. Packed fp16 dequant

For a 4-bit code n, 0x6400 | n reinterpreted as float16 is exactly 1024 + n
(fp16's 11-bit mantissa represents every integer below 2048, so nothing rounds). Applying
that to a whole 32-bit lane with a single v_and_or_b32 — (x & 0x000F000F) | 0x64006400
— dequantizes a nibble pair per instruction, against the scalar v_and_b16 + v_or_b16
pair Triton emits from the elementwise form.

The ExLlama shuffle layout this kernel already uses makes the pairing free: an int32 holds
val[2p] at bits [4p:4p+4] and val[2p+1] at bits [16+4p:20+4p], so one pre-shift by
4p (p = 0..3) puts both halves of a half2 in K order with no cross-lane shuffling. The
downstream affine then also packs, into v_pk_fma_f16.

The dequant arithmetic is unchanged. The nibble arrives as 1024 + n, so the subtrahend
absorbs the 1024: (b_raw - (1024 + zp)) == (nibble - zp), exactly, and the multiply that
follows rounds once as before.

2. dtype-aware gfx1151 tile selection

The packed fp16 path and the scalar bf16 path want materially different tiles — most visibly
BLOCK_N at deep M, 256 against 64 — so one M-ladder cannot serve both. BLOCK_K stays
capped to group_size so a K-block never straddles a quant group.

The packed dequant is enabled for the whole gfx11 family; only the tile table, which is
what actually gets tuned, stays gfx1151-gated. num_stages=1 is applied to the fp16 arm
only. Applying it to bf16 as well measured
−15.8% end-to-end prefill on an asymmetric bf16 model: the packed fp16 path issues one
per-group load and does not miss the pipelining, while the scalar asymmetric path issues two
(scale and zero point) and needs the software pipeline to hide the second gather.

The bf16 arm is consequently byte-for-byte the pre-existing scalar-tuned ladder, its
per-shape override table, and its pipeline depth — verified identical to the previous
selection across all 576 shape × group-size combinations swept. The measured bf16
control below is the null this predicts.

Duplicate-work check

Two open PRs touch rdna_hybrid_w4a16.py; neither overlaps this change:

• [ROCm][Perf] W4A16: keep skinny GEMM zero-points packed 4-bit #54965 (mine) keeps the skinny GEMM's zero-points in packed 4-bit form. It is C++ plus
  the weight-processing path; it does not touch the Triton prefill dequant or the tile table.
  Whichever of the two lands second rebases — the conflict is mechanical.
• [ROCm][Bugfix] Ignore redundant qzeros in symmetric RDNAHybridW4A16 #52637 ignores redundant qzeros on the symmetric path, in process_weights_after_loading.
  Orthogonal to the prefill kernel body.

No open PR changes the prefill dequant or the gfx1151 tile selection.

Scope

Deliberately gfx1151 only.

Testing

pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py \
       tests/kernels/quantization/test_w4a16_kernel_selection.py -q
# 103 passed

New tests:

• Prefill GEMM against a float32 oracle, over both dtypes (hence both unpacks) and both the
  symmetric and asymmetric dequants. A wrong magic constant or wrong ExLlama bit position
  moves the dequantized weight by whole integers, so this catches it comfortably.
• The tile table never returns BLOCK_K > group_size, swept in Python so the invariant holds
  for shapes there is no hardware test for.

Hardware: gfx1151 (Radeon 8060S, 40 CU), torch 2.11.0+rocm7.15, Triton 3.8.0.

Model evaluation

GSM8K, 5-shot, 500 questions, via the in-tree tests/evals/gsm8k harness, same before/after
file swap as above:

model	dtype	quant	before	after	delta
RedHatAI/Qwen3-8B-quantized.w4a16	fp16	asym	0.884 (442/500)	0.884 (442/500)	0.0 pp
RedHatAI/Qwen3-4B-quantized.w4a16	fp16	sym	0.848 (424/500)	0.850 (425/500)	+0.2 pp

Greedy decode (temperature 0, 6 fixed prompts, 64 tokens each), before vs after:

model	dtype	quant	token-ids identical	text identical
RedHatAI/Qwen3-8B-quantized.w4a16	fp16	asym	6/6	6/6
RedHatAI/Qwen3-4B-quantized.w4a16	fp16	sym	6/6	6/6
RedHatAI/Qwen3-8B-quantized.w4a16	bf16	asym	6/6	6/6

Byte-identical output is the observed result here, not a guarantee: the fp16 tiles move
BLOCK_K from 64 to 32 on most shapes, so the fp32 accumulation order does change and a
different prompt set could land differently. The probe was confirmed deterministic run to run
(a same-arm repeat is 6/6), so these are real comparisons rather than sampling noise.

AI assistance

AI assistance (Claude) was used for this change: to derive the split from a larger internal
branch, write the code and tests, and run the measurements. I have reviewed every changed
line, and ran the tests and benchmarks reported above myself.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgehre-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork