改动 / #2988
#2988 perf(hip): vectorize HC_Q8 prefill dequant and enable gfx12 Q8_0 gr_upmix
open · @bsorensen110 · 0 评论 · 去 GitHub 看
BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quantsLinux
说明
## Summary Optimizes the `STRATA_HC_Q8=1` hyper-connection (`hc_down` / `hc_up`) prefill read path on RDNA4 (`gfx1200` / `gfx1201`) and CUDA/HIP in two complementary steps: 1. **Default path (`STRATA_HC_Q8=1`, bitwise-identical):** - Vectorizes `gr_dequant_q8_0_kernel` (`src/kernels/cuda/fused_gr.cu`) from 1 `int8` element per lane (`global_load_u8/i8` + `global_store_b16`) to 8 `int8` elements per lane (`dequant_q8_0_vec8_item` + `bf16x2_from_f32_finite` with `uint4` 16-byte stores, 64 `Q8_0` blocks per 256-thread CTA), and adds pair-fused `gr_dequant_q8_0_pair_kernel` (`fused_gr_dequant_q8_0_pair`) plus `NativeDense::hc_bf16_view_pair` / `hc_bf16_release_pair` (`src/core/native_dense.cpp`). - Tracks the last recording stream per scratch slot (`g_hc_ev_stream[dev][slot]`) so same-stream `cudaStreamWaitEvent` calls are elided during single-stream prefill chunks. - Moves `hc_bf16_view(*wd)` after `gr_norm_rs` (immediately before `bf16_proj(wd)`) and `hc_bf16_view(*wu)` after `gr_silu` (immediately before `bf16_proj(wu)`) in `src/prefill/prefill.cpp` so `gr_norm_rs` (`125.8 MB` at `T=2048`, `251.7 MB` at `T=4096`) and `wu` do not evict `wd` (`6.55 MB`) from the 8 MB L2 / 64 MB Infinity Cache before `bf16_proj(wd)` reads it. 2. **Opt-in `STRATA_HC_UPMIX=1` path on `gfx12` (`gfx1200` / `gfx1201`) + fused `Q8_0` `w_up` dequantization:** - Extends `gr_upmix_kernel` and `gr_upmix()` (`src/prefill/kernels.cu`) to `gfx1200` / `gfx1201` (`__builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12`, 8-element `um_b16` fragments per lane, `t = t0 + 8 * hi + i` wave32 output row mapping, and `s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1`). - Adds an optional `const uint8_t* w_up_q8 = nullptr` parameter to `gr_upmix` so that when `STRATA_HC_Q8=1` and `STRATA_HC_UPMIX=1` are active (`wu->data == nullptr && wu->hc_q8 != nullptr`), `gr_upmix_kernel` dequantizes each CTA's 64 `w_up` rows (`64 × 320` = 640 `Q8_0` blocks) directly into shared memory `sB` once per CTA. `prefill.cpp` skips `hc_bf16_view(*wu)` entirely, eliminating all 96 `wu` full-matrix dequantizations (`1.26 GB` of BF16 scratch write+read per chunk) in addition to eliminating the `m.gated` (`T × 10,240` FP32 = `16.1 GB` at `T=2048`, `32.2 GB` at `T=4096`) global memory round-trip across all 96 halves. - Enables `gr_upmix_parity` (`src/prefill/gr_upmix_parity.cu`, `CMakeLists.txt`) on HIP (`gfx11` / `gfx12`) and adds a bitwise parity check verifying `gr_upmix(..., d_wu_q8)` matches `fused_gr_dequant_q8_0 + gr_upmix(..., d_wu_deq)` bit-for-bit across all 21 token counts (`T = 1 .. 32768`). --- ## Root cause 1. **`STRATA_HC_Q8=1` prefill dequantization overhead and cache eviction (`src/core/native_dense.cpp`, `src/kernels/cuda/fused_gr.cu`, `src/prefill/prefill.cpp`):** - When `STRATA_HC_Q8=1` is set (`NativeDense::add_hc_skip`), the BF16 hyper-connection matrices (`blk.*.attn_hc_down/up.weight`, `blk.*.ffn_hc_down/up.weight`) are skipped at load (`w.data == nullptr`), leaving only `w.hc_q8` (`Q8_0`, `320 × 10,240` = `3,276,800` elements / `102,400` blocks per matrix). - On every prefill chunk (`prefill.cpp:2943`), each of the 96 halves called `NativeDense::hc_bf16_view` twice (`192` launches of `gr_dequant_q8_0_kernel`), where each thread loaded a single `int8` quant (`global_load_u8`, `global_load_u8`, `global_load_i8`) and issued a 2-byte `global_store_b16` (`601.5 GB/s` on `gfx1201`). - Worse, `prefill.cpp` dequantized both `wd` (`6.55 MB` BF16) and `wu` (`6.55 MB` BF16) before `gr_norm_rs`. Running `gr_norm_rs` (`251.7 MB` of `m.R` read + `m.xn16` write at `T=4096`) between dequantization and `bf16_proj(wd)` evicted the freshly dequantized BF16 scratch from the 8 MB L2 / 64 MB Infinity Cache before `bf16_proj` read it. 2. **`STRATA_HC_UPMIX=1` gated off on `gfx12` (`src/prefill/kernels.cu:2112, 2398`):** - `gr_upmix()` checked only `gfx_arch_is_gfx11_wmma(p.gcnArchName)`, returning `false` on `gfx1200` / `gfx1201`. Furthermore, `prefill.cpp:2972` required `wu->data != nullptr` before calling `hc_bf16_view`, or else required dequantizing `wu` (`6.55 MB` per half) into global BF16 scratch first. - As a result, `gfx1201` always fell back to `hc_bf16_view(wu)` + `hipBLASLt` (`wu` GEMM) + `gr_mix_r`, writing and re-reading `T × 10,240` FP32 `m.gated` (`80 MB` per half = `7.68 GB` read + `7.68 GB` write at `T=2048`; `30.7 GB` total at `T=4096`). --- ## The change - `include/strata/kernels/fused_gr.hpp`, `src/kernels/cuda/fused_gr.cu`: - Add `bf16x2_from_f32_finite` and `dequant_q8_0_vec8_item` (8 elements per thread, handling both 2-byte aligned and arbitrary unaligned pointers, using round-to-nearest-even finite packing when `(d_u16 & 0x7C00u) != 0x7C00u` and falling back to `bf16_from_f32` for `±Inf`/`NaN` scales). - Update `gr_dequant_q8_0_kernel` (`64` `Q8_0` blocks per 256-thread CTA, `uint4` 16B stores) and add `gr_dequant_q8_0_pair_kernel` (`fused_gr_dequant_q8_0_pair`). - `include/strata/core/native_dense.hpp`, `src/core/native_dense.cpp`, `src/core/layer.cpp`: - Add `NativeDense::hc_bf16_view_pair` and `NativeDense::hc_bf16_release_pair`, plus `g_hc_ev_stream[64][2]` tracking to elide redundant same-stream `cudaStreamWaitEvent` calls. - `include/strata/prefill/kernels.hpp`, `src/prefill/kernels.cu`, `src/prefill/prefill.cpp`: - Port `gr_upmix_kernel` and `gr_upmix` to `gfx1200` / `gfx1201` (`__builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12`, wave32 `gfx12` fragment/row layout, `s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1`). - Add `const uint8_t* w_up_q8 = nullptr` to `gr_upmix` so `gr_upmix_kernel` dequantizes `w_up_q8` (`Q8_0`) directly into shared memory `sB`, and wire `prefill.cpp` to skip `hc_bf16_view(*wu)` when `gr_upmix` runs. - Move `hc_bf16_view(*wd)` after `gr_norm_rs` and `hc_bf16_view(*wu)` after `gr_silu` in `prefill.cpp`. - `src/kernels/gr_parity.cpp`, `src/prefill/gr_upmix_parity.cu`, `CMakeLists.txt`: - Add bitwise parity self-test for `fused_gr_dequant_q8_0` and `fused_gr_dequant_q8_0_pair` across finite, subnormal, `±0`, `±Inf`, `qNaN`, `sNaN` scales and aligned/unaligned pointers in `gr_parity --selftest`. - Enable `gr_upmix_parity` on HIP (`gfx11` / `gfx12`) and verify that `gr_upmix(..., d_wu_q8)` is bitwise identical to `fused_gr_dequant_q8_0 + gr_upmix(..., d_wu_deq)` across all 21 token counts (`T = 1 .. 32768`). --- ## Evidence ### 1. Standalone Kernel Microbenchmarks & Parity (`gfx1201`) #### A. `gr_dequant_q8_0` (`192` matrices per chunk, `1.93 GB` total read+write traffic on Radeon AI PRO R9700 `gfx1201`) - Exhaustive bitwise check across all `65,536 × 256 = 16,777,216` `(fp16 d, int8 q)` pairs (`±0`, subnormals, normals, `±Inf`, `qNaN`, `sNaN`): **0 bitwise mismatches** vs scalar `gr_dequant_q8_0_kernel`. | Kernel variant | Launches / chunk | Median GPU time (192 matrices) | Effective bandwidth | Speedup | |---|---:|---:|---:|---:| | Baseline `gr_dequant_q8_0_kernel` (scalar 1B loads, 2B store) | 192 | `3.204 ms` | `601.5 GB/s` | `1.00×` | | Vectorized `gr_dequant_q8_0_kernel` (`vec8`, 16B `uint4` store) | 192 | `1.701 ms` | `1,132.7 GB/s` | `1.88×` | | Pair-fused `gr_dequant_q8_0_pair_kernel` (`pair_vec8`) | 96 | **`1.646 ms`** | **`1,170.7 GB/s`** | **`1.95×` (`-1.56 ms`)** | #### B. `gr_upmix_parity --bench` on `gfx1201` (`STRATA_HIPBLASLT_TUNING=tools/hip/gfx1201-hipblaslt-100500.txt`) - All 21 token counts (`T = 1, 7, 16, 31, 32, 33, 48, 63, 64, 65, 96, 126, 127, 128, 129, 600, 1000, 2048, 4099, 8192, 32768`) pass (`gr_upmix_parity: 0 failures`): - `mixed max |diff| / max |value|`: `7.10e-08 .. 1.86e-07` vs `hipBLASLt + gr_mix_r`; sampled FP64 relative error `5.60e-08 .. 1.92e-07` (matching `hipBLASLt + gr_mix_r`). - Direct `Q8_0` (`gr_upmix(..., d_wu_q8)`) vs `fused_gr_dequant_q8_0 + gr_upmix(..., d_wu_deq)`: **0 bitwise diffs across all 21 token counts** (`mixed` FP32, `mixed16` BF16, and `mixed_h` FP16 images). | Token count `T` | Default pair (`hipBLASLt` up GEMM + `gr_mix_r`) | Fused `gr_upmix_kernel` (`gfx1201`) | Speedup | Fused bandwidth / TFLOPS | |---:|---:|---:|---:|---:| | `T = 2048` | `0.508 ms` (`0.122 + 0.384 ms`) | **`0.325 ms`** | **`1.57×`** | `444 GB/s`, `41 TFLOPS` | | `T = 8192` | `1.942 ms` (`0.518 + 1.422 ms`) | **`1.208 ms`** | **`1.61×`** | `461 GB/s`, `44 TFLOPS` | | `T = 32768` | `7.820 ms` (`2.255 + 5.567 ms`) | **`4.826 ms`** | **`1.62×`** | `458 GB/s`, `45 TFLOPS` | --- ### 2. End-to-End `--serve` Prefill A/B Benchmark (`gfx1201`) - **Workload & harness:** `--serve` with `GEN 1 <tokens>` (`Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S`, `--prefill 4096 --spec 4 --kv k8v4 --prompt-cache 0 --expert-cache auto --expert-cache-per-layer`), `STRATA_HC_Q8=1`, `STRATA_HIPBLASLT_TUNING=tools/hip/gfx1201-hipblaslt-100500.txt`, and RDNA4 defaults. - **Warmup & trials:** 11 warmup requests per process (all 5 rotated 64-token-vocabulary sequences at `T=4096` plus 1 request at each length `T=128, 256, 512, 1024, 2048, 4096`) so expert cache and `hipBLASLt` plan cache are warm, followed by `N = 5` measured trials per length. - **Output parity:** With `STRATA_HC_UPMIX=1 STRATA_HC_UPMIX_CHECK=4` at `T=4095`, `mixed rel RMS = 4.480e-08 .. 8.184e-08` (`max |diff| = 1.088e-06 .. 1.967e-06`), and all 30 measured requests across `T=128..4096` produce identical output token IDs across all three arms. #### A. Phase Breakdown (`STRATA_PREFILL_TIMING=1`, warm runs at `T = 2047` and `T = 4095`) | Arm | `T` | `hc read` (ms) | `combine` (ms) | `gdn out proj` (ms) | HC read + adjacent sum (ms) | Total GPU timeline (ms) | Total wall (ms) | |---|---:|---:|---:|---:|---:|---:|---:| | **Arm 1:** Upstream `origin/main` (`477ac13c`, `STRATA_HC_Q8=1`) | `2047` | `80, 80` | `42, 42` | `45, 45` | `167.0` | `980, 988` | `984, 992` | | **Arm 2:** Candidate (`STRATA_HC_Q8=1, STRATA_HC_UPMIX=0` — P1a only) | `2047` | `83, 83` | `38, 38` | `43, 42` | `163.5` (`-3.5 ms`) | `972` (warm) | `976` (warm) | | **Arm 3:** Candidate (`STRATA_HC_Q8=1, STRATA_HC_UPMIX=1` — P1a + P1b) | `2047` | **`66, 66`** | **`38, 38`** | **`42, 42`** | **`146.0` (`-21.0 ms`, `-12.6%`)** | **`958, 958` (`-26 ms`)** | **`961, 962` (`-26.5 ms`)** | | **Arm 1:** Upstream `origin/main` (`477ac13c`, `STRATA_HC_Q8=1`) | `4095` | `149, 149` | `83, 83` | `77, 77` | `309.0` | `1648, 1663` | `1655, 1671` | | **Arm 2:** Candidate (`STRATA_HC_Q8=1, STRATA_HC_UPMIX=0` — P1a only) | `4095` | `153, 153` | `79, 79` | `74, 74` | `306.0` (`-3.0 ms`) | `1642, 1646` (`-11.5 ms`) | `1649, 1653` (`-12.0 ms`) | | **Arm 3:** Candidate (`STRATA_HC_Q8=1, STRATA_HC_UPMIX=1` — P1a + P1b) | `4095` | **`123, 123`** | **`79, 80`** | **`74, 74`** | **`276.5` (`-32.5 ms`, `-10.5%`)** | **`1616, 1619` (`-38 ms`)** | **`1624, 1627` (`-37.5 ms`)** | #### B. End-to-End `--serve` Wall Latency (`timing=0`, `N = 5` trials per length) | Length | Arm 1: `origin/main` (`477ac13c`) median ms (range) [trials] | Arm 2: Candidate `STRATA_HC_UPMIX=0` median ms (range) [trials] | Arm 3: Candidate `STRATA_HC_UPMIX=1` median ms (range) [trials] | Same-binary (`UPMIX=1` vs `UPMIX=0`) | Net vs `origin/main` (`Arm 3` vs `Arm 1`) | |---:|---|---|---|---|---| | `T=128` | `229.10 ms` (`223.4..235.9`) `[229.1, 235.9, 230.9, 223.4, 228.8]` | `228.80 ms` (`222.0..235.2`) `[228.8, 235.2, 230.7, 222.0, 227.4]` | **`228.40 ms`** (`220.4..231.3`) `[226.2, 231.3, 229.4, 220.4, 228.4]` | `-0.40 ms` (`+0.18%`) | `-0.70 ms` (`+0.31%`) | | `T=256` | `330.80 ms` (`327.2..335.8`) `[328.3, 333.2, 327.2, 335.8, 330.8]` | `329.40 ms` (`326.6..332.7`) `[327.4, 332.1, 326.6, 332.7, 329.4]` | **`326.50 ms`** (`322.8..334.6`) `[326.5, 334.5, 322.8, 334.6, 324.9]` | `-2.90 ms` (`+0.89%`) | `-4.30 ms` (`+1.32%`) | | `T=512` | `461.50 ms` (`453.1..480.8`) `[461.5, 480.8, 455.0, 470.8, 453.1]` | `460.90 ms` (`452.6..479.7`) `[460.9, 479.7, 455.2, 470.7, 452.6]` | `464.60 ms` (`452.0..476.6`) `[465.2, 476.6, 452.5, 464.6, 452.0]` | `+3.70 ms` (overlapping) | `+3.10 ms` (overlapping) | | `T=1024` | `668.80 ms` (`662.1..6
本站相关内容
相关页面的快捷入口。