Pull requests / #3001

#3001 perf(prefill): use single LDS-only barrier per stage in native_w11_kernel

open · @bsorensen110 · 0 comments · View on GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsLinux

Description

## Summary

Optimizes `native_w11_kernel` (`src/prefill/moe_fused_iq.cu`), the default `gfx12` (`RDNA4`) and opt-in `gfx11` (`STRATA_PF_FUSED=1`) int8-WMMA fused MoE prefill kernel:

1. **Removes the duplicate per-stage barrier (`src/prefill/moe_fused_iq.cu:913, 964`):**
   - Stage `0`'s decoded weights (`wt[0]`, `ws[0]`) are written by `put(raw, 0)` before the prologue barrier (`line 882`), and each subsequent stage `s + 1`'s decoded weights (`wt[(s + 1) & 1]`, `ws[(s + 1) & 1]`) are written by `put(raw, (s + 1) & 1)` before the end-of-stage barrier added in `#1180` (`line 964`).
   - Because `wt` and `ws` are double-buffered (`wt[2][NW_ROWS][WLD]`, `ws[2][NW_ROWS][4]`), buffer `(s + 1) & 1` is only read in stage `s - 1` (before the barrier at the end of stage `s - 1`). Keeping both the top-of-loop barrier (`line 913`) and the end-of-stage barrier (`line 964`) executed two workgroup barriers per stage (`80` barriers in `GU` across `NS = 40`, `20` in `down` across `NS = 10`) with only `load_unit(s + 2)` and register moves between them.
2. **Replaces full-memory `__syncthreads()` with an LDS-only barrier (`nw_sync_lds()`):**
   - Every barrier in `native_w11_kernel` only synchronizes LDS (`wt`, `ws`, `srow`, `hs`).
   - On HIP, `__syncthreads()` lowers to a workgroup acquire/release fence over all memory (`s_wait_loadcnt_dscnt 0x0; s_barrier; global_inv scope:SCOPE_SE` on `gfx1200`/`gfx1201`, and `s_waitcnt vmcnt(0) lgkmcnt(0); s_barrier` on `gfx11`).
   - Waiting on `loadcnt(0)` / `vmcnt(0)` at every barrier forced `load_unit(s + 2)` and `fetch_b(s + 1)` global loads to complete at the barrier instead of overlapping the next stage's 16 `v_wmma_i32_16x16x16_iu8` instructions, and `global_inv scope:SCOPE_SE` on `gfx12` flushed the L0 vector cache on every stage even though 4 consecutive `GU` stages (`s >> 2`) read sub-blocks from the same 256-value weight superblock.
   - `nw_sync_lds()` uses `s_waitcnt lgkmcnt(0); s_barrier` on `STRATA_NAT_W11` (matching `gw_lds_barrier()` in `src/prefill/kernels.cu:1572`) and `s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `STRATA_NAT_W12`.

---

## Root cause

In `src/prefill/moe_fused_iq.cu:880–976`:
- Commit `3b499f17` (`#1180`) added `__syncthreads()` after `put(raw, (s + 1) & 1)` at line 964 while leaving the original `__syncthreads()` at the top of `for (int s = 0; s < NS; ++s)` (`line 913`).
- In the `gfx1201` disassembly of `native_w11_kernel<T_IQ3_S, true, true, 8>`:
  - At the end of stage `s`, `put(raw, (s + 1) & 1)` stored to LDS and executed `s_wait_loadcnt_dscnt 0x0; s_barrier_signal -1; s_barrier_wait 0xffff; global_inv scope:SCOPE_SE` (`line 964`).
  - Immediately after, `load_unit(s + 2)` issued 5 `global_load` instructions for stage `s + 2`'s compressed weights, and the loop backedge immediately hit the second `__syncthreads()` (`line 913`), emitting `s_wait_loadcnt 0x0; s_barrier_signal -1; s_barrier_wait 0xffff; global_inv scope:SCOPE_SE` — stalling on `load_unit(s + 2)` and invalidating the L0 vector cache before stage `s + 1`'s WMMAs began.

---

## The change

- `src/prefill/moe_fused_iq.cu`:
  - Define `nw_sync_lds()` (`s_waitcnt lgkmcnt(0)\n\ts_barrier` on `STRATA_NAT_W11`, `s_wait_dscnt 0x0\n\ts_barrier_signal -1\n\ts_barrier_wait -1` on `STRATA_NAT_W12`, `__syncthreads()` fallback).
  - Remove the redundant top-of-loop barrier at the start of `for (int s = 0; s < NS; ++s)` and replace the remaining `__syncthreads()` calls in `native_w11_kernel` with `nw_sync_lds()`.

---

## Evidence

### 1. `gfx1201` Disassembly Verification (`llvm-objdump -d`)

- Eliminates all `global_inv scope:SCOPE_SE` instructions from `native_w11_kernel` (previously `82` L0 cache invalidations per 64-row tile in `GU` and `22` in `down`).
- Halves per-stage barrier count (`80 → 39` in `GU`, `20 → 9` in `down`).
- In the steady-state stage loop (`s = 1..NS-2`), `fetch_b(s + 1)`'s 10 `global_load_b64` instructions execute concurrently with all 16 `v_wmma_i32_16x16x16_iu8` instructions of stage `s` with zero `s_wait_loadcnt` instructions inside the WMMA loop body.

### 2. Bitwise Parity & Single-Layer Prefill Benchmark (`hip_prefill_fused_iq`, `S20_NOMMQ=1`, 3 alternating rounds × 8 back-to-back runs per round on `gfx1201`)

All 21 FNV-1a 64-bit output hashes (`7` format pairs × `3` chunk sizes) are **100% bit-for-bit identical** between `origin/main` (`477ac13c`) and this PR:

| Format Pair (`GU / down`) | Tokens (`T × top 10`, `512` experts) | Baseline (`477ac13c`) median (range) | With PR median (range) | Delta (median) | Output FNV-1a Hash (both arms) |
|---|---:|---|---|---|---|
| `IQ2_S / Q2_0` | `2048` | `5.924 ms` (`5.876..5.981`) | **`5.733 ms`** (`5.701..5.739`) | **`-0.191 ms` (`-3.2%`)** | `a00480861b802319` |
| `IQ2_S / Q2_0` | `4096` | `10.053 ms` (`9.960..10.102`) | **`9.763 ms`** (`9.743..9.784`) | **`-0.290 ms` (`-2.9%`)** | `aee2d4d83555ed14` |
| `IQ2_S / Q2_0` | `8192` | `17.087 ms` (`16.955..17.259`) | **`16.580 ms`** (`16.580..16.602`) | **`-0.507 ms` (`-3.0%`)** | `00e028dc81b0f8ad` |
| `IQ2_XXS / Q2_0` | `2048` | `5.193 ms` (`5.142..5.229`) | **`5.168 ms`** (`5.144..5.195`) | **`-0.025 ms` (`-0.5%`)** | `5d683c2b3687e30b` |
| `IQ2_XXS / Q2_0` | `4096` | `8.805 ms` (`8.693..8.809`) | **`8.662 ms`** (`8.646..8.701`) | **`-0.143 ms` (`-1.6%`)** | `989bce78f95d1314` |
| `IQ2_XXS / Q2_0` | `8192` | `15.053 ms` (`14.983..15.093`) | **`14.930 ms`** (`14.926..14.982`) | **`-0.123 ms` (`-0.8%`)** | `13e842c9319ee999` |
| `IQ2_XS / IQ4_NL` | `2048` | `5.920 ms` (`5.915..5.956`) | **`5.802 ms`** (`5.789..5.804`) | **`-0.118 ms` (`-2.0%`)** | `bdaafb7393e0eba7` |
| `IQ2_XS / IQ4_NL` | `4096` | `9.964 ms` (`9.959..9.975`) | **`9.775 ms`** (`9.748..9.779`) | **`-0.189 ms` (`-1.9%`)** | `d5fe1b375e1c0a51` |
| `IQ2_XS / IQ4_NL` | `8192` | `16.963 ms` (`16.947..16.979`) | **`16.738 ms`** (`16.686..16.746`) | **`-0.225 ms` (`-1.3%`)** | `30a69a735319617a` |
| `IQ3_S / IQ4_NL` | `2048` | `5.600 ms` (`5.573..5.604`) | **`5.577 ms`** (`5.563..5.586`) | **`-0.023 ms` (`-0.4%`)** | `2facbc39b2f3ccd6` |
| `IQ3_S / IQ4_NL` | `4096` | `9.388 ms` (`9.363..9.396`) | **`9.283 ms`** (`9.249..9.298`) | **`-0.105 ms` (`-1.1%`)** | `c3656c1659ed6fc9` |
| `IQ3_S / IQ4_NL` | `8192` | `16.095 ms` (`16.060..16.152`) | **`15.807 ms`** (`15.743..15.820`) | **`-0.288 ms` (`-1.8%`)** | `2b9894f699b0d225` |
| `IQ3_XXS / IQ4_NL` | `2048` | `5.458 ms` (`5.441..5.487`) | **`5.381 ms`** (`5.360..5.409`) | **`-0.077 ms` (`-1.4%`)** | `51f5c4d34bc14422` |
| `IQ3_XXS / IQ4_NL` | `4096` | `9.312 ms` (`9.257..9.339`) | **`9.095 ms`** (`9.094..9.098`) | **`-0.217 ms` (`-2.3%`)** | `523b47c2778f2f67` |
| `IQ3_XXS / IQ4_NL` | `8192` | `15.595 ms` (`15.531..15.609`) | **`15.397 ms`** (`15.311..15.404`) | **`-0.198 ms` (`-1.3%`)** | `ebcc0b6618e1fefa` |
| `IQ4_XS / Q2_0` | `2048` | `5.009 ms` (`4.987..5.070`) | **`4.969 ms`** (`4.942..4.976`) | **`-0.040 ms` (`-0.8%`)** | `b9b1de12520448fb` |
| `IQ4_XS / Q2_0` | `4096` | `8.478 ms` (`8.458..8.508`) | **`8.350 ms`** (`8.340..8.355`) | **`-0.128 ms` (`-1.5%`)** | `60bc8a0be37c16f2` |
| `IQ4_XS / Q2_0` | `8192` | `14.389 ms` (`14.380..14.408`) | **`14.191 ms`** (`14.180..14.200`) | **`-0.198 ms` (`-1.4%`)** | `954126a6868b77c0` |
| `Q4_K / Q5_1` | `2048` | `6.485 ms` (`6.472..6.516`) | **`6.183 ms`** (`6.181..6.215`) | **`-0.302 ms` (`-4.7%`)** | `1088e4f56d22033e` |
| `Q4_K / Q5_1` | `4096` | `10.903 ms` (`10.869..10.920`) | **`10.292 ms`** (`10.289..10.304`) | **`-0.611 ms` (`-5.6%`)** | `47cb583723e30c68` |
| `Q4_K / Q5_1` | `8192` | `19.256 ms` (`19.161..19.300`) | **`18.546 ms`** (`18.543..18.565`) | **`-0.710 ms` (`-3.7%`)** | `3b8f0ec84cbb4490` |

---

## What this does NOT fix

- **Does not alter CUDA `native_kernel` (`sm_75`/`sm_80+`) or HIP `native_f16_kernel` (`STRATA_PF_F16=1`).**
- **Target-device runtime on `gfx1100`/`gfx1151`:** Device code was compiled and verified for `gfx1100`, `gfx1151`, `gfx1200`, and `gfx1201`; runtime parity and timing were executed on `gfx1201`, while `gfx11` hardware runtime remains unrun on this host.

---

## Test plan

- [x] `hip_prefill_fused_iq --chunks=2048,4096,8192` passes reference + MMQ parity on `gfx1201` (`prefill fused MoE (native formats) parity passed`)
- [x] Bit-for-bit identical 64-bit FNV-1a output hashes across all 21 `(format_pair, chunk_size)` combinations on `gfx1201`
- [x] Multi-arch HIP compile check (`--offload-arch=gfx1100,gfx1151,gfx1200,gfx1201`)
- [ ] CI

---

## Reproduction

```sh
cmake -B build-hip -DSTRATA_ENABLE_HIP=ON -DCMAKE_HIP_ARCHITECTURES="gfx1200;gfx1201" -DSTRATA_BUILD_TESTS=ON
cmake --build build-hip --target hip_prefill_fused_iq -j

./build-hip/hip_prefill_fused_iq --chunks=2048,4096,8192
S20_NOMMQ=1 ./build-hip/hip_prefill_fused_iq --no-ref --chunks=2048,4096,8192
```

Environment: AMD Radeon AI PRO R9700 (`gfx1201`, 32 GB), Linux `7.2.9-2-cachyos-server`, ROCm `10.2.0a20260914`, base `v0.1.45` (`477ac13c`).

Related on this site

Quick links to related pages.