Pull requests / #3024
#3024 perf(hip): optimize EXPDEC routed-expert kernels and enable by default on gfx12
open · @bsorensen110 · 0 comentarios · En GitHub
BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quants
Descripción
## Summary Optimize the grouped native routed-expert decode kernels (`src/kernels/cuda/iq_kernels.cu`) for RDNA4 (`gfx1200` / `gfx1201`), enable `EXPDEC` (`expdec_gu_kernel` and `expdec_down_kernel` on `Q2_0`) by default on `gfx12`, and enable `expdec_bench` on HIP builds. ### Changes 1. **Fix HIP `__launch_bounds__(256, MINB)` wave-occupancy mapping on `expdec_gu_kernel` (`src/kernels/cuda/iq_kernels.cu`)**: - On CUDA, `__launch_bounds__(256, MINB)` specifies `MINB` blocks per SM (`4` blocks = $\le 64$ registers). On HIP (`amd_hip_runtime.h`), the second argument of `__launch_bounds__(max_threads, min_blocks)` maps directly to `amdgpu_waves_per_eu(min_blocks)` (waves per SIMD, not blocks per WGP). Since a 256-thread CTA has 8 waves = 2 waves/SIMD on RDNA3/RDNA4, `MINB = 4` requested only 4 waves/SIMD (allowing up to 256 VGPRs), so `expdec_gu_kernel<16..18, 1, 4>` compiled to **`222–223 VGPRs`** on `gfx1201` (capping occupancy at 6 waves/SIMD = 3 CTAs/WGP instead of 8 waves/SIMD = 4 CTAs/WGP). - Map the second `__launch_bounds__` argument to `MINB * 2` on HIP (`8` waves/SIMD for `MINB = 4`), dropping `gfx1201` register usage to **`171–187 VGPRs`** (`0` spills, $\le 192$ VGPRs = **4 CTAs/WGP** across `IQ2_XXS`, `IQ2_XS`, `IQ3_XXS`, and `IQ3_S`). 2. **Vectorize `ed_stage_entry` in `expdec_gu_kernel` and `expdec_down_kernel`**: - Previously, `ed_stage_entry<XB, HALVES>` looped over `w = tid; w < XB * 9; w += 256`, executing 3 integer divisions by 9 (`w / 9`), 3 intra-warp divergent `if (j == 0)` branches, and 3 scattered 4-byte `ds_store_b32` stores per thread per staged entry. - Because `XB <= 80 <= 128` (`ED_XB = 80`, `ED_HB = 20`), map the 256 threads directly as 128 thread-pairs `(k = tid >> 1, h = tid & 1)` where each active pair loads one 16-byte half-block (`int4`) and writes a single 16-byte `ds_store_b128` to `s_u`, while `tid < XB` writes `s_ds[c * XB + tid]` with zero loops, zero `/ 9` divisions, and zero intra-warp branch divergence. 3. **LDS-only barriers (`iq_sync_lds()`) and elided first-pass pre-staging barriers**: - Add `iq_sync_lds()` (`s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx1200`/`gfx1201`, `s_waitcnt lgkmcnt(0); s_barrier` on other HIP targets, `__syncthreads()` on CUDA) in `stage_iq_grid`, `native_down_multi_kernel`, `expdec_gu_kernel`, and `expdec_down_kernel`, avoiding `global_inv scope:SCOPE_SE` L0 vector-cache invalidations and global load-queue drains (`s_wait_loadcnt_dscnt 0x0`). - Guard the pre-staging barrier inside `for (int e = e0; e < e1; ...)` with `if (g > (int) blockIdx.y || e > e0) iq_sync_lds();` so the first pass (`g == blockIdx.y && e == e0`, which covers all `T <= 4` decode groups when `gridDim.y == cap_groups`) skips a redundant barrier. 4. **Enable `EXPDEC` (`STRATA_EXPDEC=1`, `STRATA_EXPDEC_DOWN=1`) by default on `gfx12` and enable `expdec_bench` on HIP**: - Default `expdec_on()` and `expdec_down_on()` to `1` on `gfx1200` / `gfx1201` (`expdec_default_gfx12()`) when `STRATA_EXPDEC` / `STRATA_EXPDEC_DOWN` are unset (`STRATA_EXPDEC=0` or `STRATA_EXPDEC_DOWN=0` turns them off; non-`gfx12` architectures remain unchanged). - Gate `expdec_note` stderr logging on explicit `STRATA_EXPDEC=1`. - Enable the `expdec_bench` target on HIP builds in `CMakeLists.txt` and map `cudaDevAttrMemoryClockRate` / `cudaDevAttrGlobalMemoryBusWidth` in `include/strata/hip_compat/cuda_runtime.h`. - Extend `iq_multi_parity` (`check_grouped`) to verify bitwise equality of both `multi` (`EXPDEC=0`) and `expdec` (`EXPDEC=1, EXPDEC_DOWN=1, EXPDEC_D20=1`) against the reference per-entry kernels across all format pairs. --- ## Validation & Benchmarks (`gfx1201`, AMD Radeon AI PRO R9700) ### 1. Bitwise Parity & Multi-Arch Compilation - `./build-copilot-hip/iq_multi_parity`: `0 failures` across all 18 `iq_mmvq` checks, 24 `native_expert_grouped` checks (testing both `multi` and `expdec` against `old`), and 3 `q8_1 finite` checks. - `./build-copilot-hip/expdec_bench 150 0 20 18/20,21/20,22/20,18/42,21/42,22/42 def,dec,dec:w1 real`: `identical` (`0 choice comparisons differ`) across all 6 format pairs and `T = 1..4`. - Multi-architecture compile check: verified clean compilation on AMD HIP (`gfx1100`, `gfx1151`, `gfx1200`, `gfx1201`) and NVIDIA CUDA (`sm_75`, `sm_89`, `sm_100`). ### 2. Cold Routed-Expert Microbenchmark (`./build-copilot-hip/expdec_bench 150 0 20 ...`, 150 reps $\times$ 5 rounds over 121–186 MB cold expert pools) | Format Pair (`gu/down`) | `T` | Baseline `def` (`all / gu / down`, us) | `gfx12` Default `EXPDEC` (`all / gu / down`, us) | Total Speedup | Phase Speedup | |---|---:|---|---|---:|---| | `18/20` (`IQ3_XXS / IQ4_NL`, 14 layers) | `1` | `58.7 / 34.6 / 18.9` | **`54.7 / 31.0 / 18.9`** | **`-6.8%` (`1.073x`)** | `gu`: **`-10.4%`** | | `18/20` (`IQ3_XXS / IQ4_NL`, 14 layers) | `2` | `103.0 / 62.0 / 35.5` | **`94.1 / 53.4 / 35.6`** | **`-8.6%` (`1.095x`)** | `gu`: **`-13.9%`** | | `18/20` (`IQ3_XXS / IQ4_NL`, 14 layers) | `3` | `107.6 / 65.0 / 37.2` | **`98.7 / 56.3 / 37.1`** | **`-8.3%` (`1.090x`)** | `gu`: **`-13.4%`** | | `18/20` (`IQ3_XXS / IQ4_NL`, 14 layers) | `4` | `111.6 / 68.0 / 38.6` | **`102.8 / 59.4 / 38.3`** | **`-7.9%` (`1.086x`)** | `gu`: **`-12.6%`** | | `21/20` (`IQ3_S / IQ4_NL`, 9 layers) | `1` | `62.6 / 37.8 / 18.9` | **`58.8 / 35.5 / 19.5`** | **`-6.1%` (`1.065x`)** | `gu`: **`-6.1%`** | | `21/20` (`IQ3_S / IQ4_NL`, 9 layers) | `2` | `108.8 / 68.5 / 35.7` | **`101.2 / 60.6 / 35.8`** | **`-7.0%` (`1.075x`)** | `gu`: **`-11.5%`** | | `21/20` (`IQ3_S / IQ4_NL`, 9 layers) | `3` | `112.7 / 71.0 / 37.2` | **`103.5 / 61.6 / 37.5`** | **`-8.2%` (`1.089x`)** | `gu`: **`-13.2%`** | | `21/20` (`IQ3_S / IQ4_NL`, 9 layers) | `4` | `116.0 / 72.9 / 38.8` | **`107.6 / 64.8 / 39.0`** | **`-7.2%` (`1.078x`)** | `gu`: **`-11.1%`** | | `22/20` (`IQ2_S / IQ4_NL`, 15 layers) | `2` | `101.5 / 59.4 / 35.7` | **`96.4 / 54.7 / 36.2`** | **`-5.0%` (`1.053x`)** | `gu`: **`-7.9%`** | | `22/20` (`IQ2_S / IQ4_NL`, 15 layers) | `4` | `111.5 / 67.4 / 38.8` | **`106.7 / 63.0 / 39.0`** | **`-4.3%` (`1.045x`)** | `gu`: **`-6.5%`** | | `18/42` (`IQ3_XXS / Q2_0`, 3 layers) | `2` | `97.7 / 62.9 / 31.3` | **`79.3 / 54.0 / 22.7`** | **`-18.8%` (`1.232x`)** | `gu`: **`-14.1%`**, `down`: **`-27.5%`** | | `18/42` (`IQ3_XXS / Q2_0`, 3 layers) | `4` | `107.9 / 68.7 / 35.3` | **`88.9 / 60.1 / 26.3`** | **`-17.6%` (`1.214x`)** | `gu`: **`-12.5%`**, `down`: **`-25.5%`** | | `21/42` (`IQ3_S / Q2_0`, 1 layer) | `2` | `104.0 / 68.2 / 32.1` | **`86.5 / 60.7 / 23.2`** | **`-16.8%` (`1.202x`)** | `gu`: **`-11.0%`**, `down`: **`-27.7%`** | | `21/42` (`IQ3_S / Q2_0`, 1 layer) | `4` | `112.2 / 72.5 / 36.3` | **`93.8 / 64.9 / 26.6`** | **`-16.4%` (`1.196x`)** | `gu`: **`-10.5%`**, `down`: **`-26.7%`** | | `22/42` (`IQ2_S / Q2_0`, 5 layers) | `2` | `96.9 / 60.8 / 31.9` | **`81.6 / 54.9 / 22.8`** | **`-15.8%` (`1.188x`)** | `gu`: **`-9.7%`**, `down`: **`-28.5%`** | | `22/42` (`IQ2_S / Q2_0`, 5 layers) | `4` | `108.6 / 68.6 / 35.7` | **`92.7 / 64.5 / 26.4`** | **`-14.6%` (`1.172x`)** | `gu`: **`-6.0%`**, `down`: **`-26.1%`** | ### 3. End-to-End 397B `--serve` Decode A/B (`Swift-Qwen3.8-Flash-Next-IQ3_S`, `--spec 4 --kv k8v4`, 256 tokens, 5 matched deterministic prompts) - **`STRATA_EXPDEC=0` (legacy `native_gu_multi_kernel` + `native_down_multi_kernel`):** - Per-prompt decode times (`256` tokens each, identical draft acceptance on all 5 prompts): `2766.0 ms` (`92.60 tok/s`), `2696.0 ms` (`95.00 tok/s`), `2731.0 ms` (`93.80 tok/s`), `1989.0 ms` (`128.70 tok/s`), `2118.0 ms` (`120.90 tok/s`) - Aggregate across 5 prompts (`1,280` generated tokens): **`12,300.0 ms` (`104.07 tok/s`)** - **`gfx12` Default (`STRATA_EXPDEC=1 STRATA_EXPDEC_DOWN=1`):** - Per-prompt decode times (`256` tokens each, winning `5/5` prompts by `-21.0` to `-25.0 ms`): **`2741.0 ms` (`93.40 tok/s`)**, **`2671.0 ms` (`95.90 tok/s`)**, **`2706.0 ms` (`94.60 tok/s`)**, **`1965.0 ms` (`130.20 tok/s`)**, **`2097.0 ms` (`122.10 tok/s`)** - Aggregate across 5 prompts (`1,280` generated tokens): **`12,180.0 ms` (`105.09 tok/s`, `-120.0 ms` / `+0.98%`)** - Production `--adapt-async 1` 5-trial median (`256` tokens): **`103.40 -> 106.20 tok/s` (`2475.0 -> 2412.0 ms`, `+2.71%` / `-63.0 ms`)**.
En el sitio
Enlaces a install, modelos, releases.