Pull requests / #3482

#3482 [Engine] HIP: opt-in Q4_0-KV prompt attention on RDNA4 matrix cores (STRATA_HIP_WMMA_Q4)

open · draft · @shenmintao · 0 commentaires · Sur GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsWindows

Description

## Problem and result
With `--kv q4_0` (the long-context setting on 16 GB cards) the gfx12 matrix-core prompt attention (`STRATA_HIP_WMMA`,
int8 KV only) never runs, and the FP32 kernel is the largest phase of a long prompt: `STRATA_PREFILL_TIMING` on an
RX 9060 XT, 100K-token document, 8192-token chunk: qsa attn 2,231 ms of 8,088.

This adds an opt-in Q4_0-KV variant, `STRATA_HIP_WMMA_Q4=1` (gfx12 only). With it the same chunk spends 543 ms in the
attention, and on that card (IQ2_XS, `--kv q4_0 --max-context 524288 --batch 2`, interleaved fresh engines) a 4K prompt
reads ~830 -> ~945 tok/s and the 100K document 948 -> 1,150 tok/s (108 s -> 89 s). Decode is not touched. Unset, nothing
changes.

## Related discussion
#3480 (gfx1200 prompt-path measurements); standard benchmark of a configuration that uses it: #3481.

## Approach
`prompt_attn_wmma_q4_kernel` is the int8 WMMA kernel (S6) with q4_0 rows: a wave's 64 dims are two q4_0 blocks; q.k keeps
one WMMA partial per block times its K scale, p.v folds the two V scales (tiles 0-1 block 2w, tiles 2-3 block 2w+1); the
nibbles enter as exact FP16, q and p are split into FP16 hi + lo as in the int8 kernel, so the accuracy is FP32-level
(another summation order than the FP32 kernel). q4_0 scales are signed (ggml's d = max / -8), so the V range uses the
magnitude. The caller rotates q and un-rotates the output as for the FP32 kernel (kv_q4.hpp). A launch error switches it
off for the process like the int8 kernel; `STRATA_PA_FAST=1` selects the single-FP16 variant as for int8. All additions
are inside the existing `#if defined(__HIPCC__)` blocks; `hip_prompt_attn_wmma` now also runs the q4_0 cases on AMD
(`STRATA_PARITY_SKIP_Q4` skips them).

Second commit, easy to drop: `tools/hip/gfx1200-hipblaslt-100202-t8k.txt`, the gfx1200 table retuned with T
4096/6144/8192 buckets and the IQ2_XS expert geometries (the shipped table's buckets end at 2560, so 8192-token chunks
use the 2560 rows); +1-2% end to end here. Not wired into anything; a user points `STRATA_HIPBLASLT_TUNING` at it.

## Validation
**Passed:**
- `hip_prompt_attn_wmma 32768 2048 3` on an RX 9060 XT (gfx1200), ROCm 7.2.3, kernel and test sources identical to this
  branch (blobs 7f464222 / 1204c952; built in the image of 690e4f6f = 9ae93f7b + local commits that do not touch
  `src/kernels`): FAILURES: 0. q4_0: ctx 32768 / 2048 q vs FP64 6.3e-6 (FP32 kernel 4.06e-6, output scale 7.8),
  new vs FP32 kernel 2.4e-6 of the scale, 78.4 -> 10.8 ms per chunk (7.29x); ctx 1500 / 1500 q 7.74x; ctx 2100 / 256 q
  7.13x; the int8 cases unchanged.
- End to end with the engine built from the same kernel source (0.2.0 base and 0.2.1 base): `bench4k` checks pass, a
  6-question retrieval set over a 100K-token document 6/6, 4 picture questions 4/4, decode and two-slot batching equal
  with and without the switch; answer quality on the final configuration arith 60/60, GSM8K 97/100, HumanEval (every
  second problem) 78/82 (without it 60 / 96 / 77). Standard benchmark: #3481.

**Failed:** `STRATA_PA_FAST=1` with q4_0: parity FAIL as for int8 (error 1.4e-3 of the scale), as expected for the
single-FP16 variant; no end-to-end gain - left as it is.

**Not run:** CUDA and SYCL builds (the change is HIP-only by preprocessor, but not compiled there - hence draft);
gfx1201 / gfx11 hardware; Windows.

## Limits
Opt-in only; measured on one gfx1200 card. Not proposed as a default. Prepared with an AI assistant; the numbers come from
the logged runs above.

Sur le site

Liens install, modèles, releases.