改动 / #3069

#3069 perf(hip): enable wave32 head ops, MTP q8_mixed/HEAD_MIX_MULTI, and S2 <1> specialization

open · @bsorensen110 · 0 评论 · 去 GitHub 看

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsWindows

说明

## Summary

This PR enables and optimizes several bitwise-identical paths across the MTP draft layer, the verify/draft head greedy sampler & top-probability kernels, the S2 grouped expert kernels, and the shared expert on RDNA HIP (`gfx1201`):

1. **Wave32 multi-block `argmax_rows` and `row_top_prob_split` on RDNA HIP (`src/kernels/cuda/verify_kernels.cu`, `CMakeLists.txt`):**
   - `argmax_rows_kernel` and `row_top_prob_split_kernel` were previously compiled out under `#if !defined(__HIPCC__)`, forcing wave32 RDNA GPUs (`warpSize == 32`) onto the single-block `sampler_greedy_kernel<<<n_rows, 1024>>>` and `row_top_prob_kernel<<<n_rows, 1024>>>` (1..5 blocks total across 48 WGPs to scan 425 KB..4.97 MB of logits).
   - Guards the wave32 multi-block kernels with `#if !defined(STRATA_HIP_GFX906)` (keeping wave64 `gfx906` on the one-block kernels) and enables `multi_block_head_ops()` / `argmax_rows_wanted()` per device on HIP when `p.warpSize == 32` (`STRATA_MULTI_BLOCK_ARGMAX=0` and `STRATA_ARGMAX_ROWS=0` force them off).
   - Sets `kTopBlocks = 32` (`32` threads = 1 warp/block across 32 WGPs, avoiding intra-block `__syncthreads()` barriers when `blockDim.x == 32`) and adds `#pragma unroll 8` to the stride-1024 reduction loop in `row_top_prob_kernel` and `row_top_prob_split_kernel` (preserving exact per-thread summation order while overlapping 8 global loads in flight).
   - Enables the `verify_parity` target on HIP (`STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP`) and extends `verify_parity` with `copy_from_mapped_multi`, `quantize_q8_0_scaled`, and `--bench` checks.

2. **Fused HC `q8_mixed`, `STRATA_HEAD_MIX_MULTI`, and shared-expert `x_q8_1` reuse in MTP and Verify (`src/core/mtp.cpp`, `src/core/verify.cpp`, `src/kernels/cuda/fused_gr.cu`, `src/kernels/cuda/shared_expert.cu`):**
   - Allows `fused_gr_read_multi` (`src/kernels/cuda/fused_gr.cu`) to write `q8_mixed` without a `q8_cnt` buffer when `qfuse_hc_local() && a[0].s_down == nullptr` (`gr_up_q32_kernel` and `gr_up_q8_q32_kernel` assign one warp per 32-column `q8_1` block and never touch `q8_cnt`).
   - Enables `STRATA_HEAD_MIX_MULTI` by default on `gfx12` in `verify.cpp` and `mtp.cpp` (`STRATA_HEAD_MIX_MULTI=0` disables), replacing the 4-launch `gr_read` fallback on the MTP final mixer with `fused_gr_read_multi`.
   - Passes `q8_mixed = xq_ + t * (N / 32) * 36` to all 3 `fused_gr_read_multi` calls in `MtpDrafter` (`attn_hyper_connection`, `mlp_hyper_connection`, and `hyper_connection_mixer`) and to the final `output_hc_*` `fused_gr_read_multi` in `Verifier`, skipping redundant standalone `native_quantize_q8_1` launches when `fused_gr_read_multi` writes `xq_`.
   - Updates `shared_expert()` (`src/kernels/cuda/shared_expert.cu`) to honor `native->x_q8_1` (set by `MtpDrafter` when `q8_mlp && T == 1`), matching `shared_expert_multi`.

3. **S2 grouped expert `<1>` specialization, LDS-only barriers, and fused SwiGLU+Q8_0 on hit paths (`src/kernels/cuda/s2_expert_grouped.cu`, `src/kernels/s2_expert_grouped_parity.cpp`):**
   - Adds `s2_sync_lds()` (`s_wait_dscnt 0x0; s_barrier_signal -1; s_barrier_wait -1` on `gfx12`; `s_waitcnt lgkmcnt(0); s_barrier` on `gfx1030`/`gfx11xx`) for LDS-only staging barriers.
   - Adds `MAX_E = 1` template specializations `gu_grouped_t_kernel<1>` (`LDS = 3,200 B`, `VGPR = 85` vs `<8>`'s `LDS = 25,600 B`, `VGPR = 89`) and `down_grouped_t_kernel<1>` (`LDS = 800 B`, `VGPR = 47` vs `<8>`'s `LDS = 6,400 B`, `VGPR = 40`) dispatched when `cap_entries <= 16` (covering single-token MTP `T = 1, K = 10`).
   - Uses `swiglu_quantize_q8_0_scaled_kernel` in `moe_hit_grouped_s2`, `moe_hit_grouped_s2_dev`, and `moe_hit_grouped_s2_multi` when `fast && x_scales != nullptr && swiglu_q8_on()` (`STRATA_S2_SWIGLU_Q8=0` disables), matching `moe_grouped_s2`.

4. **Warp-per-block `quantize_q8_0_scaled` and `copy_from_mapped_multi` on HIP (`src/kernels/cuda/quantize_act.cu`, `src/kernels/cuda/elementwise.cu`, `src/core/mtp.cpp`):**
   - Enables the warp-per-block `quantize_q8_0_scaled_kernel` on HIP (`STRATA_OLD_Q80_SCALED=1` restores the thread-per-block kernel).
   - Enables `copy_from_mapped_multi_kernel` on HIP (`STRATA_MTP_MULTI_COPY=0` restores individual copies) and uses `copy_from_mapped_multi` in `MtpDrafter::capture_round` and `MtpDrafter::capture_step`.

---

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

All 7 parity and self-test suites pass with **0 failures** and 100% bitwise equality:

```text
$ ./build-copilot-hip/arch_defaults_test
arch_defaults_test: OK

$ ./build-copilot-hip/verify_parity --selftest --bench
argmax_rows, row_top_prob_split: bitwise equal to the one-block kernels (ties, NaN, -inf, 1-8 rows on one scratch)
copy_from_mapped_multi: bitwise equal to individual mapped copies (multi=1 and fallback=0, aligned + 1025w + 4B-aligned)
quantize_q8_0_scaled: bitwise equal to CPU act_quant_q8_1 reference (zero, .5 ties, normal)
  V=106299 rows=1: greedy 1-blk  24.39 us -> argmax_rows   7.87 us (x3.10) | top_prob 1-blk  22.43 us -> split  11.24 us (x2.00)
  V=248320 rows=1: greedy 1-blk  46.94 us -> argmax_rows   7.17 us (x6.55) | top_prob 1-blk  44.95 us -> split  13.57 us (x3.31)
  V=248320 rows=4: greedy 1-blk  47.01 us -> argmax_rows   7.62 us (x6.17) | top_prob 1-blk  44.78 us -> split  13.69 us (x3.27)
  V=248320 rows=5: greedy 1-blk  47.12 us -> argmax_rows  11.06 us (x4.26) | top_prob 1-blk  44.78 us -> split  13.68 us (x3.27)
verify_parity: PASS

$ ./build-copilot-hip/s2_expert_grouped_parity --bench
  moe_hit_grouped_s2 (fp32 scales)             bitwise identical
  moe_hit_grouped_s2 (fp16 d)                  bitwise identical
  moe_hit_grouped_s2_dev (count 7 of 10)       bitwise identical
  moe_hit_grouped_s2_multi (4 tokens, 29 hits) bitwise identical
  moe_grouped_s2 (20 groups, fp32 scales)      bitwise identical
  moe_grouped_s2 (20 groups, fp16 d)           bitwise identical
  moe_group_resident + moe_grouped_s2 T=1 (fp32) bitwise identical
  moe_group_resident + moe_grouped_s2 T=1 (fp16) bitwise identical
  moe_group_resident + moe_grouped_s2 T=5 (fp32) bitwise identical
  moe_group_resident + moe_grouped_s2 T=5 (fp16) bitwise identical
bench: AMD Radeon AI PRO R9700, L2 8 MB, 48 blobs (66 MB) cycled
  moe_hit_grouped_s2, 9 hits               previous     61.4 us (   203 GB/s)   new     25.6 us (   486 GB/s)   x2.40
  moe_hit_grouped_s2, 5 hits               previous     45.2 us (   153 GB/s)   new     18.5 us (   375 GB/s)   x2.45
  moe_hit_grouped_s2, 3 hits               previous     39.2 us (   106 GB/s)   new     14.9 us (   278 GB/s)   x2.63
  moe_hit_grouped_s2, 2 hits               previous     35.9 us (    77 GB/s)   new     13.6 us (   203 GB/s)   x2.64
  moe_hit_grouped_s2, 1 hits               previous     32.8 us (    42 GB/s)   new     11.9 us (   116 GB/s)   x2.75
  moe_grouped_s2, 1 token, MTP shape (10 groups) previous     55.0 us (   251 GB/s)   new     21.1 us (   655 GB/s)   x2.60
  moe_grouped_s2, 4 tokens, no sharing (40 groups) previous    122.0 us (   453 GB/s)   new     55.0 us (  1006 GB/s)   x2.22
  moe_grouped_s2, 4 tokens, 30% shared (31 groups) previous    105.6 us (   402 GB/s)   new     48.0 us (   885 GB/s)   x2.20
  moe_grouped_s2, 8 tokens, 30% shared (59 groups) previous    173.0 us (   472 GB/s)   new     83.3 us (   981 GB/s)   x2.08
  moe_grouped_s2, 8 tokens, all shared (ne 8) (10 groups) previous    131.4 us (   105 GB/s)   new     65.7 us (   210 GB/s)   x2.00
s2_expert_grouped_parity: 0 failures
PASS

$ ./build-copilot-hip/shared_expert_parity --selftest
shared_expert: 0 failures
shared_expert_parity OK

$ ./build-copilot-hip/sampler_parity
sampler: 0 failures

$ ./build-copilot-hip/gr_parity
gr_read/gr_write: 0 failures

$ ./build-copilot-hip/quantize_act_parity
quantize_act: 0 failures over 5 Q8_0 + 4 Q8_K distributions
```

Multi-architecture compilation verified cleanly on:
- **HIP (`hipcc`):** `gfx1100`, `gfx1101`, `gfx1102`, `gfx1103`, `gfx1151`, `gfx1200`, `gfx1201`
- **CUDA (`nvcc`):** `sm_75`, `sm_89`, `sm_100`, `sm_120`

---

## End-to-End 397B Speculative Decode A/B (Radeon AI PRO R9700 `gfx1201`)

Workload: Qwen3.5-397B-A17B (`swift-iq3_s` pack + native GGUF + MTP drafter `--spec 4 --spec-min-p 0.5 --pcie-frac 0.30 --kv int8`), `512` prompt tokens, `256` generated tokens per trial, `1` warmup + `5` timed trials per arm (`111` verify windows/trial, `147 / 224` drafts accepted, identical token sequence SHA-256 `ced7b38252fd809d` across all trials and arms).

| Arm | Binary & Selectors | Per-trial Decode (ms) | Median Decode (ms) | Median (tok/s) | Range (tok/s) | Window (ms/win) | MTP Step Graph GPU (ms/step) |
|---|---|---|---:|---:|---|---:|---:|
| **Primary: `d4_on` (default)** | `strata.d4` (default switches) | `2747.3, 2765.6, 2769.8, 2778.3, 2778.5` | **2,769.8** | **92.43** | `92.14..93.18` | **24.95** | **0.947** |
| **Primary: `d4_all_off` (same binary)** | `strata.d4` (`STRATA_MULTI_BLOCK_ARGMAX=0 STRATA_HEAD_MIX_MULTI=0 STRATA_MTP_MULTI_COPY=0 STRATA_OLD_Q80_SCALED=1 STRATA_QFUSE_HC=0 STRATA_LFUSE=0 STRATA_S2_SWIGLU_Q8=0`) | `2826.2, 2817.2, 2818.8, 2832.3, 2826.9` | **2,826.2** | **90.58** | `90.39..90.87` | **25.46** | **1.016** |
| **Baseline: `origin/main`** | `strata.main` (`477ac13c`) | `2794.6, 2796.1, 2801.2, 2805.7, 2808.0` | **2,801.2** | **91.39** | `91.17..91.61` | **25.24** | **1.007** |

- **Same-binary selector delta (`d4_on` vs `d4_all_off`):** **`-56.4 ms` (`+1.85 tok/s`, `+2.04%`)** end-to-end decode (`24.95` vs `25.46 ms/window`), **`-69 us/step` (`-6.8%`)** on the MTP step graph (`0.947` vs `1.016 ms/step`).
- **Delta vs `origin/main` (`d4_on` vs `main_base`):** **`-31.4 ms` (`+1.04 tok/s`, `+1.14%`)** end-to-end decode (`24.95` vs `25.24 ms/window`, non-overlapping 5-trial distributions `[2747.3, 2778.5] ms` vs `[2794.6, 2808.0] ms`), **`-60 us/step` (`-6.0%`)** on the MTP step graph (`0.947` vs `1.007 ms/step`).

### Per-Stage MTP Step Graph Breakdown (`STRATA_MTP_PROFILE=1`, 492 step graphs)

- `k/v+append` (skips `native_quantize_q8_1` via `q8_attn`): `0.046 ms` -> **`0.043 ms`** (`-3 us`)
- `experts` (`quantize_q8_0_scaled` warp-per-block + `gu/down_grouped_t_kernel<1>` + `s2_sync_lds()`): `0.046 ms` -> **`0.044 ms`** (`-2 us`)
- `shared` (skips `native_quantize_q8_1` via `nsw.x_q8_1` reuse): `0.036 ms` -> **`0.034 ms`** (`-2 us`)
- `hc-read mixer` (`STRATA_HEAD_MIX_MULTI=1` + `q8_head` fused into `fused_gr_read_multi`): `0.049 ms` -> **`0.035 ms`** (`-14 us`)
- `head gemv` (skips `native_quantize_q8_1(sample_, xq_)`): `0.344 ms` -> **`0.341 ms`** (`-3 us`)
- `argmax+prob` (wave32 `argmax_rows` + `row_top_prob_split`): `0.054 ms` -> **`0.025 ms`** (`-29 us`, **`2.16x`**)
- **Total per MTP step graph:** `1.031 ms` -> **`0.980 ms`** (`-51 us/step` in `mtp profile`; `1.007 ms` -> **`0.947 ms`** = `-60 us/step` without profile stamps).

---

## What This Does NOT Fix / Scope Notes

- Wave64 CDNA/Vega (`STRATA_HIP_GFX906` / `warpSize == 64`) remains on the single-block `sampler_greedy_kernel` and `row_top_prob_kernel`.
- Target-device hardware execution was validated on `gfx1201` (AMD Radeon AI PRO R9700); `gfx1100..gfx1151` and CUDA `sm_75..sm_120` were verified by compilation.
- Escape hatches are preserved for every path: `STRATA_MULTI_BLOCK_ARGMAX=0`, `STRATA_ARGMAX_ROWS=0`, `STRATA_HEAD_MIX_MULTI=0`, `STRATA_QFUSE_HC=0`, `STRATA_MTP_MULTI_COPY=0`, `STRATA_OLD_Q80_SCALED=1`, `STRATA_LFUSE=0`, `STRATA_S2_SWIGLU_Q8=0`.

## Reproduction Commands

```bash
cmake -B build-hip -DSTRATA_ENABLE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1201
cmake --build build-hip --target strata verify_parity s2_expert_grouped_parity shared_expert_parity sampler_parity gr_parity quantize_act_parity arch_defaults_test -j16
./build-hip/arch_defaults_test
./build-hip/verify_parity --selftest --bench
./build-hip/s2_expert_grouped_parity --bench
./build-hip/shared_expert_parity --selftest
./build-hip/sampler_parity
./build-hip/gr_parity
./build-hip/quantize_act_parity
```

---

## Full-Stack Combined Integration

本站相关内容

相关页面的快捷入口。