Pull requests / #3102

#3102 perf(hip): enable wave32 QSA decode attention and MMVQ row reuse on gfx12

open · @bsorensen110 · 0 comentários · No GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quants

Descrição

## Summary

This PR combines two bitwise-exact RDNA4 (`gfx12`, wave32) decode/speculative-verification optimizations and enables their parity/benchmark harnesses (`qsa_merge_daily_parity`, `mmvq_multi_parity`, `b6_mmvq_bench`) under `STRATA_ENABLE_HIP`:

1. **D5 — Wave32 QSA decode attention (`src/kernels/cuda/qsa_decode_attn.cu`):**
   - Replaces cross-warp `__syncthreads()` calls around shared-memory reductions in `attn_chunk_kernel` and `attn_chunk_kernel_pre75` with `qda_sync_lds()` (`s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx12`; `s_waitcnt lgkmcnt(0); s_barrier` on `gfx1030`/`gfx11xx`; `__syncthreads()` elsewhere) so shared-memory barriers do not drain in-flight global KV VMEM loads.
   - Enables the existing wave32 `attn_chunk_kernel_pre75<M, 8, true>` kernel by default on `gfx12` (`STRATA_ATTN_PRE75=0` disables; `=1` forces on). Each 32-lane wave processes 8 KV cells per step, caches `q[t]` in registers across the chunk, and uses the exact-order wave32 DPP reduction (`0` register spills on `gfx1201`).
   - Enables the bounded `attn_merge_v2_kernel<16>` merge kernel by default on `gfx12` (`STRATA_ATTN_MERGE_V2=0|1` overrides), skipping empty chunks while preserving the exact sequential chunk merge order.

2. **D6 — Native MMVQ cross-row fragment reuse & LDS barriers (`src/kernels/cuda/native_mmvq.cu`, `include/strata/kernels/native_mmvq.hpp`, `src/kernels/b6_mmvq_bench.cpp`):**
   - Replaces cross-warp `partial` LDS reduction barriers with `mmvq_sync_lds()` on `gfx12` and `gfx1030`/`gfx11xx`.
   - Enables the `v_perm_b32` (`__builtin_amdgcn_perm`) IQ4 lookup table path across HIP targets (`STRATA_USE_HIP`) rather than restricting it to `STRATA_HIP_GFX906`.
   - Extends `b6_mmvq_rows_kernel<F, NCOLS, ROWS>` to `SmallTraits` (`Q8_0` and `IQ4_NL`) by adding `PreloadX` (`load_x` / `dot_preloaded`), preloading quantized activation fragments and reusable weight-code transforms once across `ROWS` output rows while preserving the exact per-row reduction order (`0` register spills across all `b6_mmvq_rows_kernel` instantiations on `gfx1201`).
   - Adds `b6_default_rows_gfx12<F, NCOLS>(n_out)` when `STRATA_B6_MMVQ_ROWS=-1` (the new default; `gfx12` uses the measured per-shape table, all other architectures remain `0`/off; `0` disables, `1` preserves explicit 2-row mode, `2`/`4`/`8` force that row count). Gated by `multi_exact()` and `ncols > 1`.

---

## Root Cause & Kernel Resource Evidence (`gfx1201`)

1. **QSA decode attention (`qsa_decode_attn.cu`):**
   - On `gfx12` prior to this PR, `pre75_attn()` only matched `gfx103x` on HIP and `merge_v2_on()` defaulted to `0`. Consequently, `gfx12` ran `attn_chunk_kernel<M, true, false, false>` (one cell per warp per step, reloading `q` from global/L1 each step) followed by `attn_merge_kernel` (a 1-thread-per-head serial loop over all 64 chunk slots, including empty chunks at short/medium context).
   - Furthermore, the three `__syncthreads()` barriers inside each chunk step emit `s_wait_bvhcnt 0x0; s_wait_samplecnt 0x0; s_wait_loadcnt 0x0; s_wait_storecnt 0x0; s_wait_dscnt 0x0` on `gfx12`, stalling on outstanding VMEM loads when only LDS visibility is required.
   - `llvm-readobj --notes` on `qsa_decode_attn.cu.o` (`hipv4-amdgcn-amd-amdhsa--gfx1201`):
     - `attn_chunk_kernel_pre75<0..3, 8, true>`: VGPRs `141 / 149 / 183 / 149`, SGPRs `34 / 40 / 38 / 38`, LDS `15,872 B`, spills `0/0`.
     - `attn_merge_v2_kernel<16>`: VGPRs `21`, SGPRs `78`, LDS `15,872 B`, spills `0/0`.

2. **Native MMVQ multi-column projections (`native_mmvq.cu`):**
   - During speculative verification (`ncols = 2..8`, especially `T = 4..5` with `--spec 4`) and 3-column GDN QKV / dense projections, `native_mmvq_multi_kernel` reloads the same `NCOLS` `Q8_1` activation blocks for every output row. `b6_mmvq_rows_kernel<F, NCOLS, ROWS>` loads the `NCOLS` `Q8_1` blocks into registers once and accumulates `ROWS` output rows in lockstep with the exact same per-lane reduction order as `ncols` single-column calls.
   - `llvm-readobj --notes` on `native_mmvq.cu.o` (`hipv4-amdgcn-amd-amdhsa--gfx1201`):
     - `b6_mmvq_rows_kernel<Q4KTraits, 3, 4>`: VGPRs `65`, SGPRs `28`, LDS `2,304 B`, spills `0/0`.
     - `b6_mmvq_rows_kernel<Q5KTraits, 4, 4>`: VGPRs `82`, SGPRs `30`, LDS `3,072 B`, spills `0/0`.
     - `b6_mmvq_rows_kernel<Q6KTraits, 4, 4>`: VGPRs `57`, SGPRs `30`, LDS `3,072 B`, spills `0/0`.
     - `b6_mmvq_rows_kernel<Q6KTraits, 8, 4>`: VGPRs `97`, SGPRs `30`, LDS `2,048 B`, spills `0/0`.
     - `b6_mmvq_rows_kernel<IQ4XSTraits, 4, 4>`: VGPRs `91`, SGPRs `33`, LDS `3,072 B`, spills `0/0`.

---

## Bitwise Parity & Microbenchmarks (`gfx1201`, AMD Radeon AI PRO R9700)

### 1. Parity Suites

- **`qsa_merge_daily_parity`:** **1,172 cases, 0 bitwise differences** across `fp16`, `int8`, `q4_0`, and `k8v4` pools under all four selector combinations:
  - default (`STRATA_ATTN_PRE75` unset, `STRATA_ATTN_MERGE_V2` unset)
  - `STRATA_ATTN_PRE75=1`
  - `STRATA_ATTN_MERGE_V2=0`
  - `STRATA_ATTN_PRE75=1 STRATA_ATTN_MERGE_V2=0`
- **`mmvq_multi_parity`:** **148,480 outputs, 0 bitwise differences** (`multi_exact on: 148480 outputs, 0 bit-different`; negative control `multi_exact off: 63830/148480 bit-different`, confirming sensitivity).
- **`b6_mmvq_bench 80 0 -1`:** **0 failed comparisons** across all model shapes (`Q4_K`, `Q5_K`, `Q6_K`, `IQ4_XS`, `IQ4_NL`, `BF16`) and column counts (`1, 2, 3, 4, 5, 8`), verifying bitwise identity against both default multi-column `native_mmvq` and `T` single-column calls.
- **`arch_defaults_test`:** passed.

### 2. QSA Decode Attention Microbenchmark (`qsa_merge_daily_parity --bench`)

200 event-timed calls per arm, warmed, interleaved B-A-B order on `gfx1201`:

| Pool | Capacity | `n_q` | Width | Standard Chunk (us) | `pre75` Chunk B1 / B2 (us) | Chunk Delta | Upstream Merge (us) | `merge_v2` B1 / B2 (us) | Merge Delta |
|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| `int8` | 32,768 | 1 | 2,051 | 39.12 | 25.72 / 25.43 | **−34.6%** | 108.77 | 6.17 / 6.19 | **−94.3% (17.6×)** |
| `int8` | 32,768 | 4 | 2,051 | 84.33 | 49.49 / 49.55 | **−41.3%** | 112.96 | 7.36 / 7.39 | **−93.5% (15.3×)** |
| `int8` | 32,768 | 1 | 32,768 | 517.85 | 314.07 / 313.86 | **−39.4%** | 116.58 | 31.83 / 31.80 | **−72.7% (3.7×)** |
| `int8` | 32,768 | 4 | 32,768 | 1099.07 | 599.25 / 598.89 | **−45.5%** | 148.69 | 39.30 / 39.00 | **−73.7% (3.8×)** |
| `fp16` | 32,768 | 1 | 2,051 | 52.84 | 36.12 / 36.16 | **−31.6%** | 108.88 | 6.23 / 6.18 | **−94.3% (17.5×)** |
| `q4_0` | 32,768 | 1 | 2,051 | 71.15 | 28.43 / 28.45 | **−60.0%** | 108.76 | 6.19 / 6.21 | **−94.3% (17.5×)** |
| `k8v4` | 32,768 | 1 | 2,051 | 62.66 | 27.60 / 27.60 | **−56.0%** | 108.93 | 6.17 / 6.16 | **−94.3% (17.7×)** |

### 3. Native MMVQ Microbenchmark (`b6_mmvq_bench 80 0 -1`)

Medians (ms/call; `A` = default exact `STRATA_B6_MMVQ_ROWS=0`, `B` = auto table `STRATA_B6_MMVQ_ROWS=-1`, `0` bitwise differences):

| Shape | Format | `n_in → n_out` | `ncols=3` (A → B) | `ncols=4` (A → B) | `ncols=5` (A → B) | `ncols=8` (A → B) |
|---|---|---|---|---|---|---|
| GDN QKV | `Q4_K` | `2560 → 10240` | `0.0534 → 0.0489 ms` (**−8.4%**) | `0.0628 → 0.0558 ms` (**−11.1%**) | `0.0680 → 0.0638 ms` (**−6.2%**) | `0.0984 → 0.0908 ms` (**−7.7%**) |
| GDN QKV | `Q5_K` | `2560 → 10240` | `0.0564 → 0.0506 ms` (**−10.3%**) | `0.0659 → 0.0604 ms` (**−8.3%**) | `0.0725 → 0.0663 ms` (**−8.6%**) | `0.1059 → 0.0968 ms` (**−8.6%**) |
| GDN QKV | `Q6_K` | `2560 → 10240` | `0.0543 → 0.0531 ms` (**−2.2%**) | `0.0695 → 0.0630 ms` (**−9.4%**) | `0.0757 → 0.0715 ms` (**−5.5%**) | `0.1112 → 0.1128 ms` (+1.4%) |
| Attn Q | `IQ4_XS` | `2560 → 6144` | `0.0398 → 0.0378 ms` (**−5.0%**) | `0.0445 → 0.0422 ms` (**−5.2%**) | `0.0472 → 0.0473 ms` (0.0%) | `0.0658 → 0.0663 ms` (+0.8%) |
| MTP head | `Q6_K` | `2560 → 106299` | `0.4648 → 0.4150 ms` (**−10.7%**) | `0.5430 → 0.4525 ms` (**−16.7%**) | `0.6305 → 0.5484 ms` (**−13.0%**) | `0.9203 → 0.8042 ms` (**−12.6%**) |
| Output head | `Q5_K` | `2560 → 248320` | `1.3917 → 1.1997 ms` (**−13.8%**) | `1.6350 → 1.4186 ms` (**−13.2%**) | `1.9708 → 1.7878 ms` (**−9.3%**) | `2.6483 → 2.3260 ms` (**−12.2%**) |
| Output head | `Q6_K` | `2560 → 248320` | `1.2981 → 0.9032 ms` (**−30.4%**) | `1.4592 → 1.0230 ms` (**−29.9%**) | `1.6368 → 1.3204 ms` (**−19.3%**) | `2.3104 → 2.0594 ms` (**−10.9%**) |

---

## End-to-End 397B Benchmarks (`gfx1201`, Swift Qwen3.8-Flash-Next IQ3_S)

**Workload & configuration:** `swift-iq3_s` pack, `--native --kv int8 --spec 4 --spec-min-p 0.5 --pcie-frac 0.30 --vram-reserve-mib 1024 --no-prompt-cache`, `STRATA_HIP_WMMA=1`, `STRATA_HIPBLASLT_TUNING=tools/hip/gfx1201-hipblaslt-100500.txt`, 1 warmup + 5 timed trials per arm, 256 generated tokens.

### 1. Isolated Branch (`perf/rdna4-d5-qsa-decode-mmvq`) vs Upstream `origin/main` (`477ac13c`, no other PRs)

#### Prompt 512 + Gen 256 (all 25 timed trials produced exact token hash `ced7b38252fd809d`)

| Arm | Binary & Selectors | 5-Trial Decode Wall (ms) | Median Decode (ms) | Median tok/s | Median Window / Verify / Draft (ms) | MTP GPU `ps` (ms) | Delta vs Selectors-Off | Delta vs `origin/main` |
|---|---|---|---:|---:|---|---:|---:|---:|
| `origin/main` | `477ac13c` | `2805.1, 2821.1, 2814.0, 2824.4, 2849.5` | **2,821.1** | **90.74** | `25.42 / 22.39 / 2.98` | `1.009` | +0.94% | baseline |
| Selectors OFF | `HEAD` (`STRATA_ATTN_PRE75=0 STRATA_ATTN_MERGE_V2=0 STRATA_B6_MMVQ_ROWS=0`) | `2783.8, 2787.7, 2805.0, 2794.7, 2816.2` | **2,794.7** | **91.60** | `25.18 / 22.18 / 2.98` | `1.008` | baseline | **−0.94% wall / +0.95% tok/s** |
| D5 only | `HEAD` (`STRATA_B6_MMVQ_ROWS=0`) | `2682.5, 2687.0, 2689.3, 2693.0, 2696.8` | **2,689.3** | **95.19** | `24.23 / 21.56 / 2.66` | `0.883` | **−3.77% wall / +3.92% tok/s** | **−4.67% wall / +4.90% tok/s** |
| D6 only | `HEAD` (`STRATA_ATTN_PRE75=0 STRATA_ATTN_MERGE_V2=0`) | `2769.7, 2760.4, 2762.6, 2778.9, 2766.0` | **2,766.0** | **92.55** | `24.92 / 21.77 / 3.14` | `1.005` | **−1.03% wall / +1.04% tok/s** | **−1.95% wall / +1.99% tok/s** |
| **Default (D5+D6 ON)** | `HEAD` (default env) | `2639.0, 2655.6, 2645.0, 2648.9, 2664.6` | **2,648.9** | **96.64** | `23.86 / 21.20 / 2.65` | `0.881` | **−5.22% wall / +5.50% tok/s** | **−6.10% wall / +6.50% tok/s** |

#### Prompt 4,096 + Gen 256 (all 15 timed trials produced exact token hash `b52b193b1859dbf7`)

| Arm | Binary & Selectors | 5-Trial Decode Wall (ms) | Median Decode (ms) | Median tok/s | Median Window / Verify / Draft (ms) | MTP GPU `ps` (ms) | Delta vs Selectors-Off | Delta vs `origin/main` |
|---|---|---|---:|---:|---|---:|---:|---:|
| `origin/main` | `477ac13c` | `2468.9, 2474.6, 2789.2, 2495.8, 2518.0` | **2,495.8** | **102.57** | `27.13 / 23.84 / 3.28` | `1.032` | +0.56% | baseline |
| Selectors OFF | `HEAD` (`STRATA_ATTN_PRE75=0 STRATA_ATTN_MERGE_V2=0 STRATA_B6_MMVQ_ROWS=0`) | `2464.6, 2482.0, 2473.0, 2485.6, 2506.0` | **2,482.0** | **103.14** | `26.98 / 23.73 / 3.14` | `1.032` | baseline | **−0.55% wall / +0.56% tok/s** |
| **Default (D5+D6 ON)** | `HEAD` (default env) | `2364.7, 2376.6, 2366.9, 2372.4, 2373.1` | **2,372.4** | **107.91** | `25.79 / 22.84 / 2.94` | `0.897` | **−4.42% wall / +4.62% tok/s** | **−4.94% wall / +5.21% tok/s** |

### 2. Stacked with Open RDNA4 PRs (#2988, #2997, #3001, #3010, #3013, #3024, #3064, #3069)

When merged on top of the other open RDNA4 PRs (`perf/rdna4-all-combined`):
- **512 prompt + 256 generated (`hash = ced7b38252fd809d`):**
  - Combined stack with D5/D6 selectors OFF: `2667.0, 2667.5, 2670.6, 2673.7, 2677.2 ms` → median **`2,670.6 ms` (`95.86 tok/s`)**
  - Combined stack with D5+D6 ON (default): `2525.7, 2526.7, 2530.4, 2532.6, 2534.4 ms` → median **`2,530.4 ms` (`101.17 tok/s`)** (**`−5.25%` wall / `+5.54%` tok/s** vs D5/D6 off; **`−10.30%` wall / `+11.49%` tok/s** vs `origin/main`).
- **4,096 prompt + 256 generated (`hash = b52b193b1859dbf7`):**
  - Combined stack with D5/D6 selectors OFF: `2374.6, 2410.5, 2391.2, 2393.6, 2415.7 ms` → median **`2,393.6 ms` (`106.95 tok/s`)**
  -

No site

Links install, modelos, releases.