Pull requests / #260
#260 RDNA3 WMMA kernels for the gfx1100 prefill (dense GEMMs and prompt attention)
closed · @StevenChenSE · 0 comentarios · En GitHub
BenchmarksAMD / HIPNVIDIA / CUDAModels & quants
Descripción
gfx11's matrix instructions (`v_wmma_*`) are unused by this backend: dense projections reach BLAS, and the prompt
attention uses the ordered FP32 fallback, which exists because the QSA matrix instructions it was written for are
CUDA-only. This adds two native paths for gfx1100.
## What
* `src/prefill/wmma_gemm.{cu,h}` - dense GEMM over the RDNA3 WMMA fragment layout (16x16x16, wave32 "doubled"
inputs), templated for `_Float16` and `__bfloat16`, dispatched from `Gemm::f16` and `Gemm::bf16` ahead of the
BLAS fallback.
* `src/kernels/cuda/qsa_prompt_attn.cu` - the prompt attention itself on WMMA (scores 16x32 in two waves, output
four warps over 64 dimensions) instead of the emulated scalar path.
* `tests/hip/prefill_wmma_gemm_parity.cpp` - parity against a double-precision host reference over 440 shapes x 4
beta/ldy configurations x 2 dtypes, including partial tiles, padding columns, end guards and the declined
shapes; registered for ctest like the neighbouring `hip_*` tests.
## A/B switches
`STRATA_WMMA_GEMM=0` (both dense GEMMs), `STRATA_WMMA_BF16=0` (bf16 only), `STRATA_PA_WMMA=0` (attention).
## Measured
Against an unmodified build of this branch with the same flags, one RX 7900 XTX, 6-core host, PCIe 4.0 x16, using
the documented measured configuration (`--prefill 8192 --spec 4 --kv int8 --kv-resident 32768 --adapt-every 0
--pcie-frac 0`, `STRATA_PREFILL_MMQ=1`):
| prefill tier | before | after | decode tier | before | after |
| --- | ---: | ---: | --- | ---: | ---: |
| 1K | 444 tok/s | **601** (+35 %) | 1K | 42.6 | 41.7 |
| 4K | 778 | **1178** (+51 %) | 4K | 42.7 | 41.0 |
| 32K | 778 | **1527** (+96 %) | 32K | 55.2 | 54.9 |
Decode is unchanged inside the noise. The gain is prefill and it grows with context, which is where the dense
projections and the attention dominate. Under the kernels' own switches (1K prompt, 32 generated tokens): 616.5
tok/s with both on, 597.6 with the attention off, 387.1 with the dense GEMMs off, 369.0 with both off.
## Numerical note
The fp16 parity comparison allows 4 ULP: the matrix core accumulates in fp32 but rounds toward zero, while a host
reference rounds to nearest. Every delta observed on this hardware is <= 1 ULP and toward zero, and bf16 compares
exactly; `STRATA_WMMA_PARITY_EXACT=1` demands bit equality instead. The fp16 and bf16 paths both leave an 8-token
seed run bit-identical, and the decode side is unaffected.
## Three integration details
Without these the kernels compile but never run, so they are part of the change:
1. this tree's HIP macro is `STRATA_USE_HIP`. The kernels were written against a private `STRATA_BACKEND_HIP`, and
with it undefined the WMMA dispatch was silently compiled out.
2. `cmake/hip_backend.cmake` supplies `STRATA_WMMA_GFX11`, because the **host** pass of a HIP compile does not
define `__gfx1100__` even though it does define `__HIP_DEVICE_COMPILE__`. A compiler-macro guard therefore
built the returning-false stubs for the symbol the engine links and the WMMA path never ran.
3. `hip_compat` needs `CUDART_INF_F`, CUDA's `math_constants.h` spelling, mapped to `HIP_INF_F` - in the
force-included `cuda_runtime.h` and in a standalone header for sources that include it by name.
## Not claimed
* These were **not** compared against the calibrated hipBLASLt table, which needs a ROCm release matching it (the
table's version guard rejects this one); the BLAS fallback here is `hipblasGemmEx`.
* The numbers are one host and one model pack. The host is described above because it matters: expert offload over
PCIe and a 6-core CPU are part of the configuration.
En el sitio
Enlaces a install, modelos, releases.