Pull requests / #3010
#3010 perf(prefill): port wmma_gemm and gdn_chunked WMMA kernels to RDNA4 gfx12
open · @bsorensen110 · 0 comentarios · En GitHub
BenchmarksServer & APIAMD / HIPModels & quantsLinux
Descripción
## Summary
Ports Strata's prefill FP16/BF16 WMMA kernels in `src/prefill/wmma_gemm.cu` and the WY-chunked Gated DeltaNet (GDN) prefill recurrence in `src/prefill/kernels.cu` (`gdn_chunked`) from RDNA3-only (`gfx11`) to RDNA4 (`gfx1200` / `gfx1201`):
1. **Enables `STRATA_WMMA_GFX11=1` for `gfx12` builds (`cmake/hip_backend.cmake:105`):**
- Previously `cmake/hip_backend.cmake` matched only `gfx11`, so on `gfx1200;gfx1201` builds `src/prefill/wmma_gemm.cu` compiled to `return false;` stubs for `strata_wmma_gemm_f16`, `strata_wmma_gemm_bf16`, `strata_pf_gemm_f16` / `strata_pf_gemm_f16_ld`, `strata_pf_hcdown_bf16`, and `strata_pf_hcdown_exact_bf16`.
2. **Ports `src/prefill/wmma_gemm.cu` to RDNA4 (`gfx1200` / `gfx1201`) Wave32 16×16×16 WMMA (P3a):**
- Extends `strata_wmma_gfx11_device()` to accept `gfx_arch_is_gfx12_wmma()`.
- Updates `WmmaTraits<_Float16 / __bf16>`, `gemm_wmma_16x16_1w`, `gemm_wmma_64x64_4w`, `pfg::kernel`, `pfg::kernel_bk64` (`STRATA_PF_GEMM=1`), `pfg::kernel_hcdown` (`STRATA_PF_HCDOWN=1`), and `pfg::kernel_hcd_exact` (`STRATA_HCD_EXACT=1`) to use RDNA4 8-element half-wave fragments (`k = 8 * hi + 0..7`), `__builtin_amdgcn_wmma_f32_16x16x16_{f16,bf16}_w32_gfx12`, contiguous half-wave accumulator row mapping (`row = 8 * hi + i`), and LDS-only barriers (`s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx12`, `s_waitcnt lgkmcnt(0); s_barrier` on `gfx11`).
3. **Ports WY-chunked GDN prefill recurrence (`gdn_wmma_prep_kernel`, `gdn_wmma_scan_kernel`, `gdn_rec_wmma` in `src/prefill/kernels.cu:1513–1875`, `STRATA_GDN_CHUNKED=1`) to RDNA4 (`gfx1200` / `gfx1201`):**
- On `gfx12`, the Wave32 WMMA accumulator layout `(row = 8 * hi + v, col = lane & 15)` is already identical to the B-operand fragment layout `(k = 8 * hi + v, n = lane & 15)`. Converting an accumulator tile `gw_f8` into a B-operand fragment (`gw_acc_to_b`) requires **zero cross-lane permutes** (`4× permlanex16` + `8× v_perm_b32` on `gfx11` become a direct 8-element FP32→FP16 cast on `gfx12`), keeping `gdn_wmma_scan_kernel` at `239 VGPRs` with `0` spills on `gfx1201`.
4. **Fixes host-only BF16 fixture encoding and non-finite checks in `tests/hip/prefill_wmma_gemm_parity.cpp`:**
- Uses `strata::kernels::bf16_from_f32(value)` in `encode(value, true)` instead of `hip_bfloat16(value).data` (which in a host-only `.cpp` translation unit without `__HIPCC__` performs C++20 parenthesized aggregate initialization `(uint16_t)value` and truncated fractional BF16 inputs to zero), and rejects non-finite `delta` values (`!std::isfinite(delta) || delta > 0.0`).
---
## Root cause
1. **`cmake/hip_backend.cmake:105` & `src/prefill/wmma_gemm.cu`:**
- `if(CMAKE_HIP_ARCHITECTURES MATCHES "gfx11")` excluded `gfx12`, leaving `STRATA_WMMA_GFX11` undefined on RDNA4 builds so `src/prefill/wmma_gemm.cu` compiled only the `#else` stub section (`return false;`).
- Within `src/prefill/wmma_gemm.cu`, device passes checked `#if defined(__gfx1100__) ...` and used RDNA3 16-element replicated fragments (`ext_vector_type(16)`) and interleaved accumulator row indexing (`2 * i + hi`), whereas RDNA4 (`gfx1200` / `gfx1201`) Wave32 16×16×16 WMMA instructions (`__builtin_amdgcn_wmma_f32_16x16x16_{f16,bf16}_w32_gfx12`) split the 16-element K-dimension across the two half-waves (`8` elements per lane: `hi = 0` holds `k = 0..7`, `hi = 1` holds `k = 8..15`) and map accumulator elements `c[0..7]` to contiguous half-wave rows `8 * hi + i`.
2. **`src/prefill/kernels.cu:1513–1875` (`gdn_chunked`):**
- `STRATA_GDN_WMMA` and `gdn_rec_wmma()` checked only `gfx11` targets (`gfx_arch_is_gfx11_wmma`), so `STRATA_GDN_CHUNKED=1` always returned `false` on `gfx1200` / `gfx1201` and fell back to the token-serial `gdn_rec_fused_kernel`.
---
## The change
- `cmake/hip_backend.cmake`:
- Match `gfx11|gfx12` when defining `STRATA_WMMA_GFX11=1`.
- `src/prefill/wmma_gemm.cu`:
- Define `STRATA_WMMA_DEV_GFX11`, `STRATA_WMMA_DEV_GFX12`, `WMMA_ROW(i, hi)`, `WMMA_K_OFF(hi)`, `wmma_sync_lds()`, and `pfg::pfg_sync_lds()`.
- Support `gfx1200` / `gfx1201` in `strata_wmma_gfx11_device()`, `WmmaTraits<_Float16>`, `WmmaTraits<__bf16>`, `gemm_wmma_16x16_1w`, `gemm_wmma_64x64_4w`, `pfg::kernel`, `pfg::kernel_bk64`, `pfg::kernel_hcdown`, and `pfg::kernel_hcd_exact`.
- `src/prefill/kernels.cu`:
- Define `STRATA_GDN_WMMA_GFX11`, `STRATA_GDN_WMMA_GFX12`, and `GW_ROW(v, hi)` for `gdn_wmma_prep_kernel` and `gdn_wmma_scan_kernel`, implement zero-permute `gw_acc_to_b` and `s_wait_dscnt 0x0` barriers on `gfx12`, and allow `gfx_arch_is_gfx12_wmma()` in `gdn_rec_wmma()`.
- `tests/hip/gdn_chunked.cpp` & `tests/hip/prefill_wmma_gemm_parity.cpp`:
- Enable `gfx_arch_is_gfx12_wmma()` in both test harnesses, fix `encode(value, true)` in `prefill_wmma_gemm_parity.cpp` to use `strata::kernels::bf16_from_f32(value)`, and add exact binary-fraction parity tests (with `!std::isfinite(delta)` guards) for `strata_pf_gemm_f16_ld`, `strata_pf_hcdown_bf16`, and `strata_pf_hcdown_exact_bf16`.
---
## Evidence
### 1. ELF Resource Usage on `gfx1201` (`llvm-readobj --notes`)
All 11 device kernels across `wmma_gemm.cu.o` and `kernels.cu.o` compile on `gfx1201` with **0 register spills (`vgpr_spill_count = 0`, `private_segment_fixed_size = 0`)**:
| Object | Kernel | Threads (`Waves`) | VGPRs | SGPRs | LDS (B) | Spills |
|---|---|---:|---:|---:|---:|---:|
| `wmma_gemm.cu.o` | `gemm_wmma_16x16_1w<_Float16 / bfloat16_t>` | 32 (`1w`) | 23 | 24 | 0 | 0 |
| `wmma_gemm.cu.o` | `gemm_wmma_64x64_4w<_Float16 / bfloat16_t>` | 128 (`4w`) | 64 | 26 | 4,096 | 0 |
| `wmma_gemm.cu.o` | `pfg::kernel<128, 32>` | 256 (`8w`) | 143 | 22 | 40,960 | 0 |
| `wmma_gemm.cu.o` | `pfg::kernel<256, 64>` | 256 (`8w`) | 236 | 20 | 61,440 | 0 |
| `wmma_gemm.cu.o` | `pfg::kernel_bk64` | 256 (`8w`) | 229 | 23 | 55,296 | 0 |
| `wmma_gemm.cu.o` | `pfg::kernel_hcdown` | 256 (`8w`) | 191 | 25 | 36,864 | 0 |
| `wmma_gemm.cu.o` | `pfg::kernel_hcd_exact` | 192 (`6w`) | 165 | 19 | 36,864 | 0 |
| `kernels.cu.o` | `gdn_wmma_prep_kernel` | 128 (`4w`) | 133 | 16 | 3,456 | 0 |
| `kernels.cu.o` | `gdn_wmma_scan_kernel` | 256 (`8w`) | 239 | 52 | 31,360 | 0 |
### 2. Parity & Standalone Kernel Benchmarks on `gfx1201`
- **`hip_prefill_wmma_gemm_parity`:**
- Before (`origin/main` `477ac13c`): skipped with exit code `77` on `gfx1201` (`WMMA GEMM parity SKIPPED: gfx1201 has no WMMA matrix cores`).
- With PR: **3,529 test cases passed, 26 declined shapes skipped, 0 failures, `max_abs = 0`** across `strata_wmma_gemm_f16` (`1,760` cases), `strata_wmma_gemm_bf16` (`1,760` cases with non-zero `bf16_from_f32` inputs), `strata_pf_gemm_f16_ld` (`5` cases covering `kernel<128, 32>`, `kernel<256, 64>`, and `kernel_bk64`), `strata_pf_hcdown_bf16` (`2` cases), and `strata_pf_hcdown_exact_bf16` (`2` cases).
- **`hip_gdn_chunked` (1 GDN layer: `HK = 16, HV = 48, S = 128`, 7-repeat median):**
- Before (`origin/main` `477ac13c`): reported `wmma=0` on `gfx1201` (`gdn_rec_wmma` declined every `T`).
- With PR: **`wmma=1`, `FAILURES: 0` across all 45 test cases** (`T = 0..16385`, CPU reference, chunk continuation, weak-decay, beta-zero, and guard checks).
| Token Count (`T`) | Default `gdn_rec_fused` (`--bench-default`) | WMMA `gdn_rec_wmma` (`--bench`) | Speedup (1 Layer) | 48 GDN Layers Default | 48 GDN Layers WMMA | Saved per Prompt |
|---:|---:|---:|---:|---:|---:|---:|
| `64` | `0.0814 ms` | **`0.0342 ms`** | **`2.38x` (`-57.9%`)** | `3.91 ms` | **`1.64 ms`** | **`-2.26 ms`** |
| `256` | `0.2306 ms` | **`0.0790 ms`** | **`2.92x` (`-65.8%`)** | `11.07 ms` | **`3.79 ms`** | **`-7.28 ms`** |
| `2048` | `1.8829 ms` | **`0.5159 ms`** | **`3.65x` (`-72.6%`)** | `90.38 ms` | **`24.76 ms`** | **`-65.62 ms`** |
| `16384` | `12.4872 ms` | **`4.0308 ms`** | **`3.10x` (`-67.7%`)** | `599.38 ms` | **`193.48 ms`** | **`-405.91 ms`** |
| `32768` | `24.6364 ms` | **`7.6724 ms`** | **`3.21x` (`-68.9%`)** | `1,182.55 ms` | **`368.28 ms`** | **`-814.27 ms`** |
### 3. End-to-End 397B Prefill A/B (`Swift-Qwen3.8-Flash-Next-IQ3_S`, 60 Layers, Same Binary Selector Comparison)
Measured via `strata --serve --prompt-cache 0 --short-read 0` with `STRATA_PREFILL_TIMING=1` (`1` warmup + `3` timed trials per prompt length on uncontended `gfx1201`):
| Selector Configuration | Prompt (`T`) | `gdn recurrence` GPU | `hc read` GPU | Total GPU Timeline (trials / median) | `DONE` Prefill Wall (trials / median) | Prefill Rate (median) | Delta vs Default |
|---|---:|---:|---:|---|---|---:|---:|
| Default (opt-ins off) | `2048` | `58 ms` (`6.0%`) | `80 ms` (`8.4%`) | `956, 956, 957` (**`956 ms`**) | `983.8, 986.3, 985.8` (**`985.8 ms`**) | `2,077.5 tok/s` | baseline |
| `STRATA_GDN_CHUNKED=1` | `2048` | **`18 ms` (`2.0%`)** | `80 ms` (`8.7%`) | `921, 935, 923` (**`923 ms`**) | `948.8, 965.4, 950.0` (**`950.0 ms`**) | **`2,155.8 tok/s`** | **`-35.8 ms` (`+3.8%`)** |
| `STRATA_GDN_CHUNKED=1 STRATA_PF_HCDOWN=1` | `2048` | **`19 ms` (`2.0%`)** | **`73 ms` (`7.9%`)** | `955, 917, 917` (**`917 ms`**) | `983.0, 947.0, 944.3` (**`947.0 ms`**) | **`2,162.6 tok/s`** | **`-38.8 ms` (`+4.1%`)** |
| `+ STRATA_PF_GEMM=1` | `2048` | **`19 ms` (`2.0%`)** | **`72 ms` (`7.9%`)** | `917, 917, 917` (**`917 ms`**) | `944.0, 947.8, 944.8` (**`944.8 ms`**) | **`2,167.7 tok/s`** | **`-41.0 ms` (`+4.3%`)** |
| Default (opt-ins off) | `4096` | `121 ms` (`7.5%`) | `149 ms` (`9.2%`) | `1605, 1610, 1611` (**`1610 ms`**) | `1647.5, 1652.1, 1653.4` (**`1652.1 ms`**) | `2,479.3 tok/s` | baseline |
| `STRATA_PF_HCDOWN=1` (P3a isolated, `--vram-reserve-mib 1024`) | `4096` | `121 ms` (`7.5%`) | **`129 ms` (`8.0%`)** | `1602, 1606, 1607` (**`1606 ms`**) | `1645, 1649, 1649` (**`1649.0 ms`**) | **`2,484.3 tok/s`** | **`hc read` `-19 ms` (`-12.8%`)** |
| `STRATA_GDN_CHUNKED=1` | `4096` | **`39 ms` (`2.6%`)** | `149 ms` (`9.7%`) | `1536, 1539, 1540` (**`1539 ms`**) | `1578.7, 1582.7, 1582.5` (**`1582.5 ms`**) | **`2,588.3 tok/s`** | **`-69.6 ms` (`+4.4%`)** |
| `STRATA_GDN_CHUNKED=1 STRATA_PF_HCDOWN=1 STRATA_PF_GEMM=1` | `4096` | **`39 ms` (`2.6%`)** | **`129 ms` (`8.4%`)** | `1531, 1534, 1534` (**`1534 ms`**) | `1573.5, 1577.3, 1576.7` (**`1576.7 ms`**) | **`2,597.8 tok/s`** | **`-75.4 ms` (`+4.8%`)** |
*(Candidate evaluated and excluded — P3b in-kernel `Wd_q8` dequant inside `pfg::kernel_hcdown`: Under `STRATA_HC_Q8=1`, dequantizing `Wd` (`320 × 10240` BF16 = `6.55 MB`) once per half-layer via `hc_bf16_view` + `kernel_hcdown` takes `129 ms` across 120 half-layers at `T = 4096` (`-19 ms` vs `148 ms` hipBLASLt), because the `6.55 MB` BF16 buffer stays resident in the R9700's 64 MB Infinity Cache across all 32 M-tiles. Dequantizing `Wd_q8` on the fly inside `kernel_hcdown` repeats the 34-byte-strided `Q8_0` address calculation and 8× FP32→BF16 RNE rounding across every M-tile and measured `140 ms` at `T = 4096` (`+11 ms` slower than P3a's `129 ms`), so P3b was excluded.)*
---
## What this does NOT fix
- **Does not change default prefill dispatch when opt-in switches are unset:** `STRATA_GDN_CHUNKED=1`, `STRATA_PF_HCDOWN=1`, `STRATA_PF_GEMM=1`, and `STRATA_HCD_EXACT=1` remain behind their existing opt-in environment variables because their accumulation order / FP16 state rounding differs from default hipBLASLt and FP32 serial GDN recurrence.
- **`STRATA_HCD_EXACT=1` hipBLASLt solution IDs (`1176` / `1177` in `src/prefill/gemm_hip.cpp:623`):** `kernel_hcd_exact` is ported to `gfx12` and verified via `strata_pf_hcdown_exact_bf16`, but `gemm_hip.cpp` only dispatches `STRATA_HCD_EXACT=1` when the loaded tuning table selects `gfx1151` solution IDs `1176` or `1177`.
- **Target-device runtime on `gfx1100` / `gfx1151`:** Multi-arch device compilation was verified across `gfx1100,gfx1151,gfx1200,gfx1201`; runtime execution and benchmarks were performed on `gfx1201`, while `gfx11` hardware runtime remains unrun on this host.
---
## Test plan
- [x] `hip_prefill_wmma_gemm_parity` passes `3,529` cases (`max_abs = 0`) and `26` declined-shape checks on `gfx1201`
- [x] `hip_gdn_chunked` passes all `45` pEn el sitio
Enlaces a install, modelos, releases.