Pull requests / #3225

#3225 [Engine] perf(hip): enable fused_gr tsum on gfx12, fix q8_inject guard, and use LDS-only barriers

open · @bsorensen110 · 0 comments · View on GitHub

needs-format

Description

## Summary

Optimize the multi-token fused hyperconnection (HC) decode read path (`src/kernels/cuda/fused_gr.cu`) across both `STRATA_HC_Q8=1` (`launch_q8_t`) and `STRATA_HC_Q8=0` (`gr_up_q32_kernel`), and enable the transposed butterfly reduction (`s26ts::tsum`) inside `fused_gr.cu` by default on RDNA4 (`gfx1200` / `gfx1201`) via `gr_tsum_on()` without affecting `STRATA_TSUM` in `native_mmvq` / `native_bf16`.

### Changes

1. **Fix `q8_inject` guard in Q8 up-projection kernels (`STRATA_HC_Q8=1`)**:
   - In `gr_up_q8_kernel`, `gr_up_q8_q32_kernel`, and `gr_up_q8_fast_kernel` (`src/kernels/cuda/fused_gr.cu`), check `(m.a[k].w_inject == nullptr && m.a[k].q8_inject == nullptr)` instead of `m.a[k].w_inject == nullptr` when reducing the `HC` inject rows in `blockIdx.x == 0`, so callers that pass `q8_inject` with `w_inject == nullptr` still write `inject_out`.
   - Remove the redundant `cudaMemsetAsync(ssg, 0, ...)` and intermediate `cudaStreamSynchronize(st)` in `fused_gr_quantize_q8_0`.

2. **Transposed butterfly reduction (`s26ts::tsum`) + parallel `owner`-lane preloads in `gr_up_q32_kernel` and `gr_up_q8_q32_kernel`**:
   - Add `template <int T, bool TS = false> gr_up_q8_q32_kernel` and `template <int MAX_T, bool EXACT_T, bool TS = false> gr_up_q32_kernel`:
     - Replace `T` sequential `warp_sum(acc)` reductions with a single transposed butterfly reduction `s26ts::tsum<P>(acc, lane)` across power-of-two `P = pow2_ceil(T)` tokens (`v_permlane16_swap_b32` + DPP row butterfly on RDNA3/RDNA4).
     - Preload `R[i]`, `w_norm[i]`, `bo_prev[d0 + dd]`, and `inj_prev[c]` in parallel across `owner` lanes (`lane == s26ts::tsum_lane<P>(tk) && tk < T`) before the `lo[k]` dot loop so global memory latency is hidden behind LDS dot products.
     - Add `#pragma unroll 2` to the `8`-row per-warp loop.
   - On `gfx1201`, this lowers register usage across `T = 1..6` and keeps `T = 5, 6` under the `96 VGPR` threshold needed for **2 CTAs/WGP** (16 waves/SIMD, 512 threads/CTA):
     - `gr_up_q8_q32_kernel<T, false>` (`53 / 69 / 78 / 89 / 99 / 109 VGPRs`) $\rightarrow$ `gr_up_q8_q32_kernel<T, true>` (**`45 / 59 / 68 / 77 / 86 / 95 VGPRs`**).
     - `gr_up_q32_kernel<T, true, false>` (`48 / 70 / 80 / 90 / 100 / 110 VGPRs`) $\rightarrow$ `gr_up_q32_kernel<T, true, true>` (**`40 / 61 / 69 / 78 / 87 / 96 VGPRs`**).
     - `gr_down_q8_kernel<T>` (`93 / 119 / 139 / 159 / 179 / 199 / 208 VGPRs`) $\rightarrow$ `gr_down_q8_fast_kernel<T, 4>` (**`88 / 106 / 123 / 142 / 163 / 178 / 189 VGPRs`**, keeping `T = 6, 7` $\le 192$ VGPRs for **2 CTAs/WGP**).

3. **LDS-only barriers (`gr_sync_lds()`)**:
   - Replace `__syncthreads()` inside `gr_up_q32_kernel`, `gr_down_q8_kernel`, `gr_down_q8_fast_kernel`, and `gr_up_q8_q32_kernel` with `gr_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).
   - Every intra-block barrier in these kernels guards only shared memory (`lo`, `rsS`, `red`, `g`); using `gr_sync_lds()` avoids emitting `global_inv scope:SCOPE_SE` (`960` Shader Engine L0 cache invalidations per `gr_up_q8_q32_kernel` launch and `324` per `gr_down_q8_fast_kernel` launch) and avoids stalling in-flight global loads (`s_wait_loadcnt_dscnt 0x0`).

4. **Default enablement on `gfx12` (`gr_tsum_on()`)**:
   - Add `gr_tsum_on()` in `src/kernels/cuda/fused_gr.cu`, which checks `STRATA_GR_TSUM` (`0|1`), then `STRATA_TSUM` (`0|1`), and otherwise defaults to `norm_fast_on()` (**ON by default on `gfx1200` / `gfx1201`**, OFF on other architectures).
   - Scoping the `gfx12` default to `fused_gr.cu` avoids turning on `STRATA_TSUM=1` in `native_mmvq.cu` and `native_bf16.cu`.

---

## Validation & Benchmarks (`gfx1201`, AMD Radeon AI PRO R9700)

### 1. Correctness (`gr_parity --selftest` & `fused_gr_bench`)
- Extended `src/kernels/gr_parity.cpp` (`fused_multi_lds_parity`) to test:
  - `STRATA_QFUSE` (`gr_up_q32_kernel`) bitwise equality across `STRATA_TSUM=0` and `STRATA_TSUM=1` for `T = 1..8`;
  - `STRATA_HC_Q8=1` (`q8_down` + `q8_up` + `q8_inject` with `w_inject == nullptr` and `q8_inject == nullptr` with `w_inject != nullptr`) across `T = 1..8`, asserting bitwise-identical `R_out`, `rs`, `inject_out`, `mixed`, and `q8_mixed` outputs between `STRATA_TSUM=0` and `STRATA_TSUM=1`.
- `./build-copilot-hip/gr_parity --selftest`:
  ```text
  fused GR read + q8_1 (STRATA_QFUSE), T 1..8, direct and graph replays: pass
  fused GR Q8_0 multi (STRATA_HC_Q8 + q8_inject + STRATA_TSUM), T 1..8: pass
  gr_read/gr_write: 0 failures
  gr_parity OK
  ```
- `STRATA_BENCH_QMIXED=1 ./build-copilot-hip/fused_gr_bench 200 1 6 16`: `0 differing` and identical output hashes across `T = 1..6` (`apply = 0, 1`) between `STRATA_GR_TSUM=0` and `STRATA_GR_TSUM=1`.
- Multi-architecture compile check: verified clean compilation on AMD HIP (`gfx1100`, `gfx1151`, `gfx1200`, `gfx1201`) and NVIDIA CUDA (`sm_75`, `sm_89`, `sm_100`, `sm_120`).

### 2. Microbenchmark (`STRATA_BENCH_QMIXED=1 ./build-copilot-hip/fused_gr_bench 200 1 6 16`, cold 16-set rotation)

| Path | `T` | `STRATA_GR_TSUM=0` read (us) | `STRATA_GR_TSUM=1` (`gfx12` default) read (us) | Delta | Stage stamps (`norm / down / up`, us) |
|---|---:|---:|---:|---:|---|
| `fast 2` (`STRATA_HC_Q8=1`) | `1` | `23.4` | `23.6` | `+0.9%` | `2.8 / 13.8 / 15.5` $\rightarrow$ `2.8 / 13.6 / 15.8` |
| `fast 2` (`STRATA_HC_Q8=1`) | `2` | `27.7` | **`26.6`** | **`-4.0%` (`1.041x`)** | `2.8 / 15.1 / 18.5` $\rightarrow$ `2.8 / 14.2 / 18.2` |
| `fast 2` (`STRATA_HC_Q8=1`) | `3` | `32.6` | **`30.5`** | **`-6.4%` (`1.069x`)** | `2.8 / 18.0 / 20.3` $\rightarrow$ `2.8 / 15.8 / 20.0` |
| `fast 2` (`STRATA_HC_Q8=1`) | `4` | `35.7` | **`31.6`** | **`-11.5%` (`1.130x`)** | `2.8 / 20.0 / 21.4` $\rightarrow$ `2.8 / 16.6 / 20.8` |
| `fast 2` (`STRATA_HC_Q8=1`) | `5` | `39.2` | **`33.8`** | **`-13.8%` (`1.160x`)** | `2.8 / 22.3 / 22.7` $\rightarrow$ `2.8 / 17.4 / 22.0` |
| `fast 2` (`STRATA_HC_Q8=1`) | `6` | `42.6` | **`35.7`** | **`-16.2%` (`1.193x`)** | `2.8 / 23.9 / 24.3` $\rightarrow$ `3.0 / 18.6 / 23.2` |
| `fast 1` (`STRATA_HC_Q8=0`, `gr_up_q32_kernel`) | `6` | `66.2` | **`64.1`** | **`-3.2%` (`1.033x`)** | `12.5 / 38.5 / 23.6` $\rightarrow$ `12.5 / 38.6 / 21.6` (`up`: **`-8.5%`**) |

### 3. End-to-End 397B `--serve` Decode A/B (`Swift-Qwen3.8-Flash-Next-IQ3_S`, `--spec 4 --kv k8v4`, 256 tokens, 2 warmups + 5 measured trials)
- **`STRATA_GR_TSUM=0` (legacy `warp_sum`):** median **`102.10 tok/s`** (`2508.0 ms`; trials: `106.60`, `100.60`, `101.10`, `118.70`, `102.10 tok/s`)
- **`STRATA_GR_TSUM=1` (`gfx12` default `tsum` + `gr_sync_lds`):** median **`103.20 tok/s`** (`2482.0 ms`, **`+1.08%` / `-26.0 ms`**; trials: `100.40`, `103.00`, `103.20`, `109.10`, `108.40 tok/s`)

## 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)

Related on this site

Quick links to related pages.