Pull requests / #3108

#3108 perf(hip): optimize exact QSA block scoring and top-k selection on gfx12

open · @bsorensen110 · 0 comments · View on GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quants

Description

## Summary

This PR optimizes the exact QSA block-scoring and radix top-k selection kernels in `src/kernels/cuda/qsa_select.cu` on RDNA4 (`gfx1200`/`gfx1201`) and wave32 HIP targets with **100% bitwise parity** (`0` differing scores across all 12 `qsa_scores_parity` cases / `51,069,696` scores, `0` differing selected IDs across all 20 `qsa_topk_active_parity` cases, and `0` differing scores or IDs across all 8 `qsa_rows_bench --rows` cases at `4K`, `32K`, and `128K` context):

1. **Wave32 VALU DPP reductions & workgroup-scoped barriers (`sel_xor_f32`, `sel_xor_i32`, `sel_sync_lds`)**:
   - On wave32 HIP targets, `__shfl_xor_sync` lowers through `ds_bpermute_b32` on the LDS/DS pipe (`5` `ds_bpermute_b32` + `5` `s_wait_dscnt` per 32-lane reduction, `20` per 4-head block score). Replacing power-of-two XOR butterfly steps with `v_permlanex16_b32` (`MASK=16`, `0x76543210` / `0xfedcba98`) and `v_mov_b32_dpp` (`row_xmask:1|2|4|8`, `0x160 | MASK`) keeps reductions on the VALU pipe while summing the exact same lane pairs in the exact same tree order.
   - Replaces device-scope `__syncthreads()` inside LDS-only selection/scoring loops with workgroup-scoped `sel_sync_lds()` (`s_wait_loadcnt_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx1200`/`gfx1201`; `s_waitcnt vmcnt(0) lgkmcnt(0); s_barrier` on `gfx103x`/`gfx11xx`). Waiting on both `loadcnt(0)` and `dscnt(0)` avoids L2 cache-invalidation overhead while remaining safe when divergent `s_cbranch_execz` paths skip a VMEM basic block.

2. **Smaller-`PER` specializations for `block_topk_reg_kernel` (`PER = 2, 9, 17`) & parallelized scans in `block_topk_kernel` / `block_topk_wide_kernel`**:
   - Previously, `qsa_block_topk` only instantiated `block_topk_reg_kernel<33>` (`151 VGPRs, 107 SGPRs, 7 SGPR spills`) and `block_topk_reg_kernel<66>` (`192 VGPRs, 107 SGPRs, 88 VGPR / 106 SGPR spills`). Every decode call up to `33,792` blocks ran `<33>`.
   - Adds `block_topk_reg_kernel<2>` (`reach <= 2,048` blocks / `8K` cells: `29 VGPRs, 26 SGPRs, 0 spills`), `block_topk_reg_kernel<9>` (`reach <= 9,216` blocks / `36K` cells: `50 VGPRs, 43 SGPRs, 0 spills`), and `block_topk_reg_kernel<17>` (`reach <= 17,408` blocks / `69K` cells: `91 VGPRs, 67 SGPRs, 0 spills`), and lowers the HIP prefill crossover `kRegMinBlocks` from `7,168` (`~28K` cells) to `4,096` (`~16K` cells, matching `gfx1151`).
   - Replaces the single-thread (`if (t == 0)`) 256-iteration radix-bin search and 32-warp prefix loops in `block_topk_wide_kernel` (`77 → 28 VGPRs`) and `block_topk_kernel` (`35 → 25 VGPRs`, `3,080 → 1,064 B` LDS) with parallel 8-warp / 1-warp prefix scans.

3. **Prompt exact block-scorer (`block_scores_thread_kernel`) grid granularity & threshold on `gfx12`**:
   - Previously, `block_scores_thread_kernel` used a fixed `TS_BR = 512` blocks per CTA along `X` and only engaged at `active_blocks >= 16,385` (`ctx > 64K`). At `32K` context (`8,193` blocks) with `256` queries, `tgrid` launched only `17 × 4 = 68` CTAs (`__launch_bounds__(256, 1)`, `177 VGPRs`), leaving 28 of 48 WGPs idle once the first 48 CTAs finished.
   - Passes a dynamic `br` (`64` for `reach <= 2,048`, `128` for `reach <= 8,192`, `256` for `reach <= 32,768`, `512` above) with `__launch_bounds__(256, 2)` (`175 VGPRs`) and lowers the default `gfx12` activation threshold to `active_blocks >= 512` (`ctx >= 2,048`).

4. **Bitwise-exact `block_scores_rows_kernel<1..8>` on HIP (`STRATA_QSA_ROWS=1`)**:
   - On `gfx1201`, LLVM's `GCNCreateVOPD` pairs heads `(0, 1)` and `(2, 3)` in `block_scores_multi_kernel` so even heads multiply `x*x` first (`fmaf(w, fmaf(z, fmaf(y, x*x)))`) and odd heads multiply `y*y` first (`fmaf(w, fmaf(z, fmaf(x, y*y)))`), whereas `block_scores_rows_kernel<NQ>` had a different register schedule and 1-ULP divergence on ~3.4% of blocks, plus unrolled `nq` branches bloating `<4>` to `60 KB` of code.
   - Aligns the 4-head FMA chain via `rows_dot4(k4, q4, h)`, stages `q_idx` cooperatively through shared memory `qs`, uses `NQ` as the compile-time bound, and extends `qsa_block_scores_rows` to `nq = 1..8`.
   - Registers `qsa_scores_parity`, `qsa_topk_active_parity`, and `qsa_rows_bench` under `STRATA_ENABLE_HIP` in `CMakeLists.txt`.

---

## Kernel Resource Audit (`gfx1201`, `llvm-readobj --notes`)

| Kernel | `origin/main` (`477ac13c`) | This PR |
|---|---|---|
| `block_scores_kernel` | `26 VGPR / 20 SGPR / 0 B LDS / 0 spill` | `24 VGPR / 20 SGPR / 0 B LDS / 0 spill` |
| `block_scores_tail_kernel` | `24 VGPR / 18 SGPR / 0 B LDS / 0 spill` | `21 VGPR / 18 SGPR / 0 B LDS / 0 spill` |
| `block_scores_multi_kernel` | `65 VGPR / 84 SGPR / 16,512 B LDS / 0 spill` | `65 VGPR / 84 SGPR / 16,512 B LDS / 0 spill` |
| `block_scores_slots_kernel` | `40 VGPR / 39 SGPR / 16,512 B LDS / 0 spill` | `38 VGPR / 39 SGPR / 16,512 B LDS / 0 spill` |
| `block_scores_thread_kernel` | `177 VGPR / 22 SGPR / 25,344 B LDS / 0 spill` | `175 VGPR / 24 SGPR / 25,344 B LDS / 0 spill` |
| `block_scores_rows_kernel<1..4>` | `65..98 VGPR / 84 SGPR / 16..64 B LDS` | `39..95 VGPR / 29 SGPR / 2,064..8,256 B LDS` |
| `block_scores_rows_kernel<5..8>` | N/A (`nq <= 4` only) | `113..167 VGPR / 30..42 SGPR / 10,328..16,512 B LDS` |
| `block_topk_kernel` (`ref`) | `35 VGPR / 44 SGPR / 3,080 B LDS / 0 spill` | `25 VGPR / 28 SGPR / 1,064 B LDS / 0 spill` |
| `block_topk_wide_kernel<false>` | `77 VGPR / 43 SGPR / 33,032 B LDS / 0 spill` | `28 VGPR / 31 SGPR / 33,032 B LDS / 0 spill` |
| `block_topk_reg_kernel<2>` | N/A (used `<33>`) | `29 VGPR / 26 SGPR / 32,908 B LDS / 0 spill` |
| `block_topk_reg_kernel<9>` | N/A (used `<33>`) | `50 VGPR / 43 SGPR / 32,908 B LDS / 0 spill` |
| `block_topk_reg_kernel<17>` | N/A (used `<33>`) | `91 VGPR / 67 SGPR / 32,908 B LDS / 0 spill` |
| `block_topk_reg_kernel<33>` | `151 VGPR / 107 SGPR / 32,908 B LDS / 0/7 spill` | `151 VGPR / 107 SGPR / 32,908 B LDS / 0/7 spill` |

---

## Microbenchmark & Parity Evidence (`AMD Radeon AI PRO R9700`, `gfx1201`)

### 1. Prefill Exact Block Scoring (`qsa_scores_parity` & `qsa_select_bench`, `capacity = 262,144`)

All `51,069,696` scores across the 12 `qsa_scores_parity` cases and all `256/256` top-k selections in `qsa_select_bench` match `block_scores_kernel` / `qsa_block_topk_ref` **bit-for-bit**:

| Benchmark Case | `origin/main` (`477ac13c`) | This PR (`STRATA_SCORES_THREAD=0`, warp) | This PR (`STRATA_SCORES_THREAD=1`, default) | Speedup vs `origin/main` |
|---|---:|---:|---:|---:|
| `qsa_select_bench` `ctx=2048, nq=256` (`scores`) | `0.055 ms` | `0.051 ms` | **`0.034 ms`** | **`1.62×` (`-38.2%`)** |
| `qsa_select_bench` `ctx=4096, nq=256` (`scores`) | `0.108 ms` | `0.101 ms` | **`0.046 ms`** | **`2.35×` (`-57.4%`)** |
| `qsa_select_bench` `ctx=16384, nq=256` (`scores`) | `0.422 ms` | `0.395 ms` | **`0.108 ms`** | **`3.91×` (`-74.4%`)** |
| `qsa_select_bench` `ctx=16384, nq=256` (`top-k`) | `0.088 ms` | `0.042 ms` | **`0.042 ms`** | **`2.10×` (`-52.3%`)** |
| `qsa_select_bench` `ctx=32768, nq=256` (`scores`) | `0.832 ms` | `0.772 ms` | **`0.197 ms`** | **`4.22×` (`-76.3%`)** |
| `qsa_select_bench` `ctx=32768, nq=256` (`top-k`) | `0.103 ms` | `0.058 ms` | **`0.058 ms`** | **`1.78×` (`-43.7%`)** |
| `qsa_scores_parity` `ctx=40000, nq=2048` | `8.08 ms` (`thread: 8.11 ms`) | `7.31 ms` | **`1.93 ms`** | **`4.19×` (`-76.1%`)** |
| `qsa_scores_parity` `ctx=140000, nq=1024` | `14.36 ms` (`thread: 3.20 ms`) | `13.82 ms` | **`3.08 ms`** | **`4.66×` vs warp (`-3.8%` vs thread)** |
| `qsa_scores_parity` `ctx=200000, nq=1024` | `20.29 ms` (`thread: 4.27 ms`) | `19.57 ms` | **`4.19 ms`** | **`4.84×` vs warp (`-1.9%` vs thread)** |

### 2. Decode Block Scoring & Top-K (`qsa_rows_bench --iters 100 --rows --no-dense`)

In `qsa_rows_bench --rows`, `origin/main` previously **failed** bitwise parity (`scores DIFFER (max abs 1.53e-05), ids DIFFER`). With this PR, all **8/8 cases (`T = 1..4` at target context and short context) are 100% bit-identical**:

| Context / Capacity | Kernel Stage | `origin/main` (`T=1 / 2 / 3 / 4` us) | This PR (`T=1 / 2 / 3 / 4` us) | Delta at `T=4` |
|---|---|---|---|---|
| `ctx=4096, max_ctx=32768` | `scores (multi)` | `5.4 / 5.8 / 6.8 / 7.9` | **`5.3 / 5.4 / 6.2 / 7.0`** | **`-11.4%`** |
| `ctx=4096, max_ctx=32768` | `topk` (`<33>` → `<9>`) | `11.7 / 10.6 / 10.6 / 10.7` | **`8.7 / 7.8 / 7.9 / 7.9`** | **`-26.2%` (`1.35×`)** |
| `ctx=4096, max_ctx=32768` | `layer (append+scores+topk+attn)` | `65.2 / 75.5 / 90.0 / 112.1` | **`62.0 / 71.9 / 86.7 / 108.3`** | **`-3.8 us` (`-3.4%`)** |
| `ctx=4096, max_ctx=262144` | `scores (multi)` | `5.5 / 5.8 / 6.8 / 7.9` | **`5.1 / 5.4 / 6.2 / 7.0`** | **`-11.4%`** |
| `ctx=4096, max_ctx=262144` | `topk` (`wide`, active-bounded) | `47.0 / 42.5 / 43.9 / 47.2` | **`9.2 / 8.3 / 8.3 / 8.3`** | **`-82.4%` (`5.69×`)** |
| `ctx=4096, max_ctx=262144` | `layer (append+scores+topk+attn)` | `92.7 / 97.6 / 113.8 / 140.4` | **`62.8 / 72.4 / 86.5 / 108.7`** | **`-31.7 us` (`-22.6%`)** |
| `ctx=32768, max_ctx=262144` | `scores (multi)` | `12.7 / 16.3 / 21.6 / 26.4` | **`11.0 / 12.4 / 15.4 / 17.9`** (`rows: 13.2`) | **`-32.2%`** (`rows: -50.0%`) |
| `ctx=32768, max_ctx=262144` | `topk` | `57.4 / 55.2 / 56.4 / 60.1` | **`39.4 / 21.2 / 23.8 / 29.9`** | **`-50.2%` (`2.01×`)** |
| `ctx=32768, max_ctx=262144` | `layer (append+scores+topk+attn)` | `106.7 / 121.0 / 138.9 / 169.3` | **`76.3 / 88.2 / 103.5 / 128.4`** | **`-40.9 us` (`-24.2%`)** |
| `ctx=131072, max_ctx=262144` | `scores (multi)` | `35.3 / 51.5 / 69.9 / 89.0` | **`27.6 / 35.2 / 44.5 / 53.6`** (`rows: 32.8`) | **`-39.8%`** (`rows: -63.1%`) |
| `ctx=131072, max_ctx=262144` | `topk` (`wide`) | `79.1 / 81.8 / 84.1 / 83.8` | **`69.5 / 71.3 / 71.4 / 71.6`** | **`-14.6%`** |
| `ctx=131072, max_ctx=262144` | `layer (append+scores+topk+attn)` | `149.1 / 172.0 / 201.6 / 233.6` | **`145.3 / 154.6 / 176.8 / 210.3`** | **`-23.3 us` (`-10.0%`)** |

---

## End-to-End 397B Benchmark vs Clean `origin/main` (`477ac13c`)

- **Model**: `Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S` (`397B` MoE, 12 QSA layers + 36 GDN layers) + MTP (`--spec 4 --spec-min-p 0.5`), `--max-context 32768 --kv int8`, `1` warmup + `5` timed trials per arm (`20s` GPU cooldown between arms).
- **Exact token hash parity**: All 5 trials of `main_4096`, `d7_default_4096`, and `d7_thread_off_4096` produce identical generated token hashes (`b52b193b1859dbf7`), and all 5 trials at `512` context produce `ced7b38252fd809d`.

| Arm | Binary / Selector | Prompt / Gen | Prefill Median (ms, range) | Prefill tok/s | Decode Median (ms, range) | Decode tok/s | Verify / GPU-Wait per Window (ms) | Token Hash |
|---|---|---:|---:|---:|---:|---:|---:|---|
| `main_4096` | `origin/main` (`477ac13c`), `WMMA=1` | `4096 / 256` | `1,427.9` (`1421.0–1447.0`) | `2,868.5` | `2,511.8` (`2489.6–2617.5`) | `101.92` | `24.18` (`gpu 19.56`) | `b52b193b1859dbf7` (5/5) |
| **`d7_default_4096`** | **This PR (`default`), `WMMA=1`** | **`4096 / 256`** | **`1,420.9` (`1410.7–1883.7`)** | **`2,882.7` (`+0.50%`)** | **`2,505.2` (`2491.8–2528.0`)** | **`102.19` (`+0.26%`)** | **`23.93` (`gpu 19.38`, `-0.18 ms/win`)** | **`b52b193b1859dbf7` (5/5)** |
| `d7_thdoff_wmma0_4096` | This PR (`STRATA_SCORES_THREAD=0`), `WMMA=0` | `4096 / 256` | `2,037.5` (`2020.0–2047.3`) | `2,010.3` | `2,650.6` (`2530.0–2823.8`) | `96.58` | `23.97` (`gpu 19.31`) | `d699b83b..` / `e74180b0..` |
| **`d7_default_wmma0_4096`** | **This PR (`default`), `WMMA=0`** | **`4096 / 256`** | **`2,023.2` (`2011.2–2033.1`)** | **`2,024.5` (`+0.71%`)** | **`2,626.2` (`2600.5–2672.2`)** | **`97.48` (`+0.93%`)** | **`23.92` (`gpu 19.32`)** | **`d699b83b..` / `e74180b0..`** |
| `main_512` | `origin/main` (`477ac13c`), `WMMA=1` | `512 / 256` | `495.7` (`494.6–496.5`) | `1,032.9` | `2,818.5` (`2810.5–2835.5`) | `90.83` | `22.40` (`gpu 19.18`) | `ced7b38252fd809d` (5/5) |
| **`d7_default_512`** | **This PR (`default`), `WMMA=1`** | **`512 / 256`** | **`495.5` (`494.2–496.6`)** | **`1,033.3`** | **`2,818.0` (`2811.7–2841.9`)** | **`90.84`** | **`22.40` (`gpu 19.16`)** | **`ced7b38252fd809d` (5/5)** |

Related on this site

Quick links to related pages.