[ROCm][Kernel] W4A16 skinny GEMM: hand-asm kernel for gfx1151

[mgehre-amd]

Hand-editable AMDGCN implementation of the W4A16 (int4 weight / bf16
activation) skinny GEMM, tuned for gfx1151. It targets the [128,64,32]
8-warp tile, so all four Qwen3-VL-4B-AWQ decoder GEMMs are routed to it at
deep prefill -- each ~1.2x faster than the Triton path. The dispatcher is
also extended to route down_proj (K=9728, N=2560) to this tile instead of
its previous tall-K default (which the hand-asm did not cover).

Key techniques

Cumulative impact, measured by editing the raw Triton AMDGCN dump one
technique at a time and benchmarking each state (gate_up 19456x2560,
M=1322, gfx1151, clock-pinned, median of 3, TFLOP/s). Correctness held at
every step (baseline cos=1.000000, final cos=0.999995).

step	technique added	TFLOP/s	delta
0	baseline (raw Triton dump)	23.8	--
1	+ magic-constant dequant	25.3	+6.3%
2	+ software-pipelined loads	27.1	+7.0%
3	+ exact descriptors (drop OOB guards)	27.1	+0.0%
4	+ WMMA SRC0/SRC1 bank split	27.2	+0.4%
5	+ scalar-derived load offsets	27.1	-0.1%
	total		+14.1%

The two real wins are the dequant rewrite and software pipelining. Steps
3-5 are perf-neutral at this operating point (the loop is occupancy-bound,
so their savings are hidden) but are kept for other reasons: exact
descriptors remove page-fault risk, and the bank split + scalar offsets
free VGPR-bank cycles and in-loop VALU that would otherwise cap a future
deeper-pipelined version. (Steps are listed in applied order; this swaps
the numbering of "scalar offsets" and "WMMA banks" relative to a purely
logical grouping.)

1. Magic-constant bf16 dequant -- +6.3% (no int->float convert, no v_perm)

The stock int4->bf16 dequant masks each nibble, subtracts 8, and runs
v_cvt_f32_i32 per element. Instead, OR the 4-bit nibble into a fixed bf16
bit pattern so the raw value already reads as bf16(128 + n); a paired
constant supplies bf16(-136). A single v_dot2_bf16_bf16 with the
broadcast (scale, scale) then computes scale * ((128 + n) - 136) = scale * (n - 8) directly -- fusing the -8 offset, the scale multiply, and
the pair-accumulate into one op, and packing two K into one register via
op_sel (which also removes the v_perm_b32 repack).

; before: mask + (-8) + int->float, per element
v_and_b32_e32     v61, 15, v59
v_add_nc_u32_e32  v84, -8, v61
v_cvt_f32_i32_e32 v84, v84
...
v_dot2_bf16_bf16  v92, v86, v83, 0        ; then v_perm_b32 to repack halves
v_perm_b32        v86, v90, v86, 0x5040100

; after: OR nibble into a bf16 magic constant, dot2 does -8 + scale
v_and_or_b32      v84, v59, v92, 0xc3084300   ; -> (bf16(128+n), bf16(-136))
v_dot2_bf16_bf16  v94, v84, v124, 0            ; = scale*(n-8), low half
v_dot2_bf16_bf16  v94, v85, v124, 0 op_sel:[0,0,0,1]   ; second K -> high half

2. Software-pipelined global loads -- +7.0%

The loop-top s_waitcnt vmcnt(0) was stalling on global loads issued in
the same iteration -- ATT attributed ~40% of stall cycles to it, and the
activation load alone has ~250-cycle exposed latency. Prefetch iteration
k+1's weight/scale/activation tiles into a second register set while
iteration k's LDS stores and WMMAs run, so the wait is (mostly) already
satisfied by the time the loop reaches it.

; before: load and immediately consume within the same iteration
.LBB0_2:                                 ; loop header
    buffer_load_b128 v[51:54], v52, s[12:15], 0 offen
    ...
    ds_store_b128    v42, v[51:54]        ; waits on the load just issued

; after: prologue prefetches iter 0; loop body prefetches iter k+1
    ; --- prologue: prefetch iter 0 ---
    buffer_load_b128 v[115:118], v126, s[12:15], 0 offen
.LBB0_2:                                 ; loop header
    ds_store_b128    v42, v[115:118]      ; consumes the PREVIOUS prefetch
    ; --- prefetch iter k+1 (overlaps ds_store B + ds_load B + WMMA) ---
    buffer_load_b128 v[115:118], v126, s[12:15], 0 offen

3. Exact buffer-descriptor sizing (drop in-loop OOB guards) -- +0.0% (fault-safety)

The Triton output uses an "unbounded" descriptor (num_records = 0x7ffffffe) and guards every in-loop global load with a v_cmp +
v_cndmask that rewrites the address to 0x80000000 for OOB lanes. That
is 3 VALU ops + a compare per load, every iteration. Instead, set each
buffer's num_records to its exact byte size in the prologue; the hardware
then clamps OOB accesses to 0 with no fault, and the per-iteration guards
are deleted.

; before: unbounded descriptor + per-load address guards in the loop
s_mov_b32        s14, 0x7ffffffe
...
v_cmp_gt_i32_e64 s3, s26, v125
v_cndmask_b32_e64 v125, 0x80000000, v39, s2
buffer_load_b32  v123, v125, s[16:19], 0 offen

; after: exact per-buffer num_records in prologue, bare loads in the loop
s_lshl_b32       s18, s37, ...   ; weight num_records = N*stride_bn*4 (bytes)
s_lshl_b32       s10, s38, ...   ; scale  num_records = N*num_groups*2 (bytes)
buffer_load_b32  v123, v39, s[16:19], 0 offen   ; no guard; HW clamps OOB

4. WMMA SRC0/SRC1 VGPR bank separation -- +0.4% (within noise)

All 8 WMMAs had SRC0 (weight frag) and SRC1 (activation frag) starting in
the same VGPR bank (reg mod 4 == 3), which costs extra cycles per WMMA on
gfx11. The register file is fully packed (128/128 VGPRs at the 12-wave
boundary), so rather than add registers, the activation fragments are
shifted down by one register (v51-82 -> v50-81) to land in bank 2, and the
freed collision (the weight LDS load address) is relocated to the now-free
v82. VGPR count and occupancy are unchanged.

; before: SRC0 and SRC1 both start in bank 3 (reg%4==3)
ds_load_b128            v[51:54], v43
v_wmma_f32_16x16x16_bf16 v[25:32], v[83:90], v[51:58], v[25:32]

; after: SRC1 shifted to bank 2 (v50..), distinct from SRC0/SRC2 banks
ds_load_b128            v[50:53], v43
v_wmma_f32_16x16x16_bf16 v[25:32], v[83:90], v[50:57], v[25:32]

5. Scalar-derived load offsets -- -0.1% (within noise)

Each iteration bumped the per-lane VGPR base pointers with v_add
(v36 += 4, v39 += 16, v40 += 64). Since the per-K stride is uniform
across lanes, compute the byte offsets once as scalars and pass them in the
instruction's soffset field instead, removing the in-loop address VALU.

; before: VGPR pointer bumps every iteration
v_add_nc_u32_e32 v39, 16, v39
v_add_nc_u32_e32 v40, 64, v40
buffer_load_b32  v123, v39, s[16:19], 0  offen
buffer_load_b128 v[115:118], v40, s[12:15], 0  offen

; after: uniform K stride precomputed in SGPRs, used as soffset
s_lshr_b32 s37, s5, 1     ; weight K byte offset = s5/2 (4-bit packed)
s_lshl_b32 s38, s5, 1     ; act    K byte offset = s5*2 (bf16)
buffer_load_b32  v123, v39, s[16:19], s37 offen
buffer_load_b128 v[115:118], v40, s[12:15], s38 offen

Benchmark results (per-GEMM)

Qwen3-VL-4B-Instruct-AWQ-4bit decoder GEMMs, M=1322, group_size=32,
symmetric, bf16, on gfx1151 (Strix Halo). Clock-pinned, TFLOP/s (higher is
better), median of 3 reps. "Triton (before)" is the pre-PR production path.

GEMM	shape (N x K)	Triton (before)	hand-asm (PR)	speedup
gate_up_proj	19456 x 2560	22.1	27.0	1.22x
qkv_proj	6144 x 2560	21.9	26.8	1.23x
o_proj	4096 x 2560	22.5	27.2	1.21x
down_proj	2560 x 9728	18.6	27.4	1.47x
per-layer	(aggregate)	21.1	27.1	1.28x

down_proj's 1.47x is two effects: routing it to the [128,64,32]x8 tile is
1.22x on Triton alone (18.6 -> 22.6 TFLOP/s), and the hand-asm kernel adds
a further 1.22x (22.6 -> 27.4 TFLOP/s).

Benchmark results (end-to-end)

Full-model prefill on Qwen3-VL-4B-Instruct-AWQ-4bit, gfx1151 (Strix Halo),
hand-asm kernel toggled via VLLM_W4A16_HANDASM. Identical args otherwise:
input-len 2048, output-len 1, num-prompts 10, max-num-batched-tokens 4096
(prefill M ~2858 > 1024, so the hand-asm tile is selected). Clock-pinned.

Metric	Triton (off)	hand-asm (on)	speedup
Median TTFT	1406.3 ms	1208.9 ms	1.16x (-14%)

The ~14% median TTFT gain is well above the ~1% Strix Halo e2e variance.
It is diluted from the ~1.28x per-layer GEMM speedup by attention, norms,
and (in this multimodal run) the vision encoder.

Testing

• benchmarks/kernels/benchmark_hybrid_w4a16.py --check --batches 1322
  -> cos=0.999994 vs Triton reference on all four shapes, PASS
• Benchmarked clock-pinned on gfx1151, numbers above.

AI assistance was used for this change.