Pull requests / #3224
#3224 [Engine] perf(hip): tune RDNA4 QSA prompt WMMA path (STRATA_HIP_WMMA stays opt-in)
open · @bsorensen110 · 0 comentários · No GitHub
BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quantsSecurity
Descrição
## Summary
Optimizes the HIP `int8`-KV QSA prompt-attention matrix-core kernel (`prompt_attn_wmma_kernel<SPLIT>` in `src/kernels/cuda/qsa_prompt_attn.cu`).
**`STRATA_HIP_WMMA=1` remains opt-in on every architecture (`gfx12` and `gfx11` included).** While this PR's changes inside `prompt_attn_wmma_kernel<SPLIT>` are **100% bitwise identical** to `main`'s `prompt_attn_wmma_kernel<SPLIT>` (`NEWHASH` `2c5cab7eb6601f66` / `deb4f727fbac41af` / `2d10cd9d6d1fd5b5`), the matrix-core path itself differs in its last bits from the default FP32 fallback (`qsa_decode_attn_batch`) due to `FP16 hi + lo` splitting and reduction order. Keeping `STRATA_HIP_WMMA` out of `arch_defaults.cpp` preserves the exact-only policy of `arch_default_env("gfx1201")` and avoids process-wide enablement on `gfx11` stages in mixed-GPU runs.
1. **LDS-only workgroup barriers (`pa_sync_lds`):**
- Replaces the `__syncthreads()` barriers inside `prompt_attn_wmma_kernel` (which synchronize only LDS `S`) with `s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx12` and `s_waitcnt lgkmcnt(0); s_barrier` on `gfx11` (generic `__syncthreads()` elsewhere), avoiding the `s_wait_loadcnt_dscnt 0x0` global-load drain and `global_inv scope:SCOPE_SE` L0 vector-cache invalidation emitted by generic `__syncthreads()` on `gfx12`.
2. **Global-load / WMMA overlap (`fetch` K-before-V + deferred FP16 scale conversion + post-`q·K` V staging):**
- Keeps `k_scale` and `v_scale` in raw `uint16_t` registers (`Chunk::ksc_u16[2]`, `Chunk::vsc_u16`) during `fetch()` and converts them via `u16_scale()` (with an empty `asm volatile("" : "+v"(u))` barrier on AMDGPU) only at their point of use after the `q·K` WMMAs, preventing LLVM from hoisting `v_cvt_f32_f16` into the conditional global-load blocks and forcing an early `s_wait_loadcnt 0x0`.
- Issues the `K` global loads before the `V` global loads in `fetch()`, and moves LDS `V` staging (`S.v` / `S.vs`) inside `chunk()` to *after* the `q·K` WMMAs, so `V` global loads stay in flight while the matrix cores compute `q·K`.
3. **Vectorized 32-bit LDS `V` gather + in-register DPP byte transpose (`dpp_tr4`) + VALU wave butterfly reductions (`pa_xor_u32` / `pa_xor_f32`):**
- Replaces the 64 scalar 1-byte `ds_load_i8` reads per wave per chunk in the `p·V` loop (`S.v[warp][cb + i][d]`) with 16 aligned 32-bit `ds_load_b32` reads (`*reinterpret_cast<const uint32_t*>(&S.v[warp][cb + 4 * q4 + qs4][j * 16 + qm4])`) followed by an in-register 4×4 byte transpose across the 4 lanes of a hardware DPP quad (`dpp_tr4`: `__builtin_amdgcn_mov_dpp` quad_perm `0xb1` / `0x4e` + `__builtin_amdgcn_perm`).
- Replaces `__shfl_xor_sync` in the `q`-max, row-max, row-sum, `row_x16`, and `V`-scale reductions with VALU lane permutations (`pa_xor_u32` / `pa_xor_f32`: `mov_dpp` quad_perm `0xb1`/`0x4e` for masks `1`/`2`, `permlane16` for `4`/`8`, `permlanex16` for `16`), avoiding `ds_bpermute_b32` traffic through the LDS pipe and cutting `gfx1201` register pressure from `191` to `168` VGPRs (`SPLIT=true`) and `160` to `147` VGPRs (`SPLIT=false`) with `0` spills (`SmemW` remains `21,648 B`).
---
## The change
Three files changed vs `main` (`477ac13c`) at `f9b11142`:
- `src/kernels/cuda/qsa_prompt_attn.cu` (+110/−38):
- Adds `pa_sync_lds()`, `dpp_tr4()`, `pa_xor_u32()`, and `pa_xor_f32()` inside `#if defined(__HIPCC__)`.
- Updates `prompt_attn_wmma_kernel<SPLIT>`: raw `uint16_t` scales (`Chunk::vsc_u16`, `Chunk::ksc_u16[2]`) + `u16_scale()`, K-before-V global loads in `fetch()`, V slice staged into `S.v`/`S.vs` after the `q·K` WMMAs in `chunk()`, `pa_sync_lds()` barriers, `pa_xor_f32()` butterfly reductions, and aligned 32-bit `S.v` loads + `dpp_tr4()` in `p·V`.
- `SmemW` layout is unchanged (`uint8_t v[4][WCH][WVS]` with `WVS = 80`, `21,648 B` total).
- `src/core/arch_defaults.cpp` (+2/−0):
- Comment-only note documenting why `STRATA_HIP_WMMA` is intentionally excluded from the `gfx12` default table.
- `src/core/arch_defaults_test.cpp` (+14/−1):
- Asserts `STRATA_HIP_WMMA` is absent from `arch_default_env` on `gfx1200`/`gfx1201`, clears all seven `gfx12` env vars (`STRATA_GFX12_DEFAULTS`, `STRATA_HIP_WMMA`, `STRATA_GDN_HEAD`, `STRATA_GDN_PP`, `STRATA_GDN_CONVL2`, `STRATA_GDN_NOY`, `STRATA_CVEC_FUSE`) before asserting `apply_arch_defaults("gfx1201").size() == 5`, and verifies that `STRATA_HIP_WMMA` stays unset by default while an explicit `STRATA_HIP_WMMA=1` is preserved.
---
## Evidence
All runtime measurements below were taken on commit `f9b11142` vs `main` (`477ac13c`) on an uncontended AMD Radeon AI PRO R9700 (`gfx1201`, 32 GB, `strata.service` stopped).
### 1. ELF Resource Usage (`llvm-readobj --notes`)
| Target | Kernel | VGPRs (`main` → this PR) | SGPRs (`main` → this PR) | LDS (B) | Spills |
|---|---|---|---|---:|---:|
| `gfx1201` / `gfx1200` | `prompt_attn_wmma_kernel<true>` | `191` → **`168`** (`-23`) | `32` → `32` | `21,648` | `0` |
| `gfx1201` / `gfx1200` | `prompt_attn_wmma_kernel<false>` | `160` → **`147`** (`-13`) | `32` → `32` | `21,648` | `0` |
| `gfx1100` / `gfx1151` | `prompt_attn_wmma_kernel<true>` | `243` → **`235`** (`-8`) | `32` → `34` | `21,648` | `0` |
| `gfx1100` / `gfx1151` | `prompt_attn_wmma_kernel<false>` | `197` → **`177`** (`-20`) | `32` → `34` | `21,648` | `0` |
### 2. Kernel Microbenchmark & Parity: `hip_prompt_attn_wmma` (`main` @ `477ac13c` vs this PR @ `f9b11142`)
1 warmup pair + 5 interleaved pairs of separately built `hip_prompt_attn_wmma` binaries (`main` vs `f9b11142`). On HIP, `hip_prompt_attn_wmma` sets `STRATA_HIP_WMMA=1` in `main()` and times `prompt_attn_wmma_kernel<true>` against the FP32 fallback (`qsa_decode_attn_batch`) and an FP64 host reference.
Output FNV-1a checksums (`NEWHASH`) of `prompt_attn_wmma_kernel<true>` are **100% bitwise identical** between `main` and `f9b11142` across all three shapes (`FAILURES: 0`):
| Shape (`int8` KV) | `main` WMMA (`477ac13c`, 5 trials / median) | This PR WMMA (`f9b11142`, 5 trials / median) | vs `main` WMMA | vs FP32 fallback | `NEWHASH` (`main` == this PR) |
|---|---|---|---|---:|---|
| `ctx 32768, 2048 q` | `5.605, 5.586, 5.582, 5.608, 5.603 ms` (**`5.603 ms`**) | `4.653, 4.666, 4.674, 4.668, 4.659 ms` (**`4.666 ms`**) | **`-16.7%` (`1.20×`)** | `8.28×` (`38.61 ms`) | `2c5cab7eb6601f66` (exact) |
| `ctx 1500, 1500 q` | `1.755, 1.754, 1.729, 1.753, 1.652 ms` (**`1.753 ms`**) | `1.462, 1.470, 1.461, 1.456, 1.459 ms` (**`1.461 ms`**) | **`-16.7%` (`1.20×`)** | `8.66×` (`12.65 ms`) | `deb4f727fbac41af` (exact) |
| `ctx 2100, 256 q` | `0.742, 0.748, 0.742, 0.743, 0.738 ms` (**`0.742 ms`**) | `0.589, 0.590, 0.595, 0.606, 0.590 ms` (**`0.590 ms`**) | **`-20.5%` (`1.26×`)** | `8.89×` (`5.25 ms`) | `2d10cd9d6d1fd5b5` (exact) |
Parity against FP64 and the FP32 kernel (`FAILURES: 0` on all three shapes):
| Shape | vs FP64 (this PR WMMA) | vs FP64 (FP32 kernel) | WMMA vs FP32 kernel |
|---|---:|---:|---:|
| `ctx 32768, 2048 q` | `2.64e-06` | `2.13e-06` | `6.32e-06` (`1.7e-06` of scale `3.62`) |
| `ctx 1500, 1500 q` | `3.57e-06` | `1.78e-06` | `5.25e-06` (`1.4e-06` of scale `3.57`) |
| `ctx 2100, 256 q` | `3.13e-06` | `2.03e-06` | `4.05e-06` (`1.1e-06` of scale `3.61`) |
### 3. End-to-End 397B `--serve` Prefill A/B (`Swift-Qwen3.8-Flash-Next-IQ3_S`, 60 Layers / 12 QSA Layers, `T = 4096`)
Measured via `strata --serve --kv int8 --prefill 4096 --prompt-cache 0` with `STRATA_PREFILL_TIMING=1`, `STRATA_HIPBLASLT_TUNING`, `STRATA_DENSE_MMQ=1`, `STRATA_PF_FUSED=1`, `STRATA_HC_Q8=1` (`1` warmup + `3` rotated-token trials per arm on idle `gfx1201`):
| Arm | `qsa attn` GPU Phase (12 QSA layers, trials / median) | Total GPU Timeline (trials / median) | `DONE` Prefill Read Wall (trials / median) | Prefill Throughput (median) |
|---|---|---|---|---:|
| `main` (`477ac13c`), `STRATA_HIP_WMMA=0` (default FP32) | `706, 707, 710 ms` (**`707 ms`**) | `1922, 1925, 1930 ms` (**`1925 ms`**) | `1987.0, 1990.4, 1995.6 ms` (**`1990.4 ms`**) | `2,057.9 tok/s` |
| `main` (`477ac13c`), `STRATA_HIP_WMMA=1` (baseline WMMA) | `105, 105, 106 ms` (**`105 ms`**) | `1367, 1420, 1371 ms` (**`1371 ms`**) | `1410.7, 1463.0, 1414.2 ms` (**`1414.2 ms`**) | `2,896.4 tok/s` |
| **This PR (`f9b11142`), `STRATA_HIP_WMMA=1` (P4 WMMA)** | `87, 87, 87 ms` (**`87 ms`**, **`-18 ms` / `-17.1%`** vs `main` WMMA; **`8.13×`** vs FP32) | `1353, 1355, 1356 ms` (**`1355 ms`**, **`-16 ms`** vs `main` WMMA) | `1395.9, 1397.6, 1398.4 ms` (**`1397.6 ms`**, **`-16.6 ms`** vs `main` WMMA) | **`2,930.8 tok/s`** (`+1.2%` vs `main` WMMA; `+42.4%` vs FP32) |
### 4. Multi-Architecture Compilation & Policy Checks
- **`arch_defaults_test`:** `arch_defaults_test: OK` (`gfx12 (RDNA4): 5 exact speed switches on by default: GDN_HEAD GDN_PP GDN_CONVL2 GDN_NOY CVEC_FUSE`; `STRATA_HIP_WMMA` unset by default, preserved when explicitly set).
- **HIP (`hipcc` / Clang 22.0.0):** Compiled `src/kernels/cuda/qsa_prompt_attn.cu` cleanly across `--offload-arch=gfx1100,gfx1101,gfx1102,gfx1103,gfx1151,gfx1200,gfx1201` with `0` register spills.
- **CUDA (`nvcc` 13.4):** Compiled `src/kernels/cuda/qsa_prompt_attn.cu` cleanly across `-gencode arch=compute_75,code=sm_75` (Turing), `-gencode arch=compute_89,code=sm_89` (Ada), `-gencode arch=compute_100,code=sm_100` (Blackwell), and `-gencode arch=compute_120,code=sm_120` (Blackwell).
---
## What this does NOT fix
- **Last-bit rounding difference vs the FP32 fallback (`STRATA_HIP_WMMA=0`):** While this PR is **bit-for-bit identical** to `main`'s `STRATA_HIP_WMMA=1` implementation (`2c5cab7eb6601f66` / `deb4f727fbac41af` / `2d10cd9d6d1fd5b5`), the matrix-core path splits FP32 `q` and `p` into `FP16 hi + lo` pairs (`2.64e-06 .. 3.57e-06` vs FP64, compared to `1.78e-06 .. 2.13e-06` for the FP32 fallback). For that reason `STRATA_HIP_WMMA=1` remains opt-in on every architecture.
- **No RDNA3 runtime measurement on this host:** `gfx11` compiles cleanly (`235` VGPRs for `SPLIT=true`, `177` VGPRs for `SPLIT=false`, `0` spills), but no RDNA3 hardware runtime measurement was taken on this host.
- **Does not change `--kv k8v4` prompt attention:** `prompt_attn_wmma_kernel` applies to `--kv int8` (`pools.k_q4 == nullptr && pools.v_q4 == nullptr`).
---
## Test plan
- [x] `arch_defaults_test` passes (`OK`, `gfx12` table = 5 exact switches, `STRATA_HIP_WMMA` opt-in)
- [x] `hip_prompt_attn_wmma` passes all 3 shapes (`FAILURES: 0`) with identical `NEWHASH` checksums to `main` on `gfx1201`
- [x] Uncontended 5-pair interleaved `hip_prompt_attn_wmma` A/B against `main` (`-16.7%`, `-16.7%`, `-20.5%`)
- [x] End-to-end 397B `--serve --kv int8` 4K prefill A/B (`qsa attn` `105 → 87 ms`, `-17.1%`; total prefill `1,414.2 → 1,397.6 ms`)
- [x] Multi-arch HIP compile check (`gfx1100,gfx1101,gfx1102,gfx1103,gfx1151,gfx1200,gfx1201`)
- [x] Multi-arch CUDA compile check (`sm_75`, `sm_89`, `sm_100`, `sm_120`)
- [ ] CI
---
## Reproduction
```sh
cmake -B build-hip -DSTRATA_ENABLE_HIP=ON -DSTRATA_PREFILL_MMQ=ON \
-DCMAKE_HIP_ARCHITECTURES="gfx1200;gfx1201" -DSTRATA_BUILD_TESTS=ON
cmake --build build-hip --target arch_defaults_test hip_prompt_attn_wmma strata -j
./build-hip/arch_defaults_test
./build-hip/hip_prompt_attn_wmma 32768 2048 5
```
## Validation
**Passed:**
- Build strata target: success
- arch_defaults_test: PASS
**Failed:**
- none observed
**Not run:**
- CUDA builds (no CUDA toolchain available)
- SYCL builds (no SYCL toolchain available)
- Metal builds (no Metal support)
No site
Links install, modelos, releases.