Pull requests / #3264

#3264 [Engine] perf(hip): optimize MTP draft, shared expert, and verify/sampler head kernels on RDNA4

open · @bsorensen110 · 0 comments · View on GitHub

BenchmarksServer & APIMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux

Description

## Problem and result

Supersedes #3069 (closed during the `0.2.0-beta` tracker reset), rebased onto `0.2.0-beta` (`e88ee39d`).

During speculative decode and verification on RDNA4 (`gfx1200`/`gfx1201`), six bottlenecks in the MTP drafter, shared expert, S2 routed experts, and verify/sampler head operations accounted for significant per-window overhead while having **100% bitwise-identical** replacements:
1. `multi_block_head_ops()` and `argmax_rows_wanted()` in `src/kernels/cuda/verify_kernels.cu` unconditionally returned `false` under `#if defined(__HIPCC__)`, forcing `argmax_rows` (`T x 248,320` logits) and `row_top_prob_split` onto single-WGP (`<<<n_rows, 1024>>>`) launches instead of 128-block (`argmax_rows_kernel`) and 32-block (`row_top_prob_split_kernel`) wave32 kernels (`402.0 -> 35.1 us` at `T=5`, **11.45× speedup**).
2. `quantize_q8_0_scaled` in `src/kernels/cuda/quantize_act.cu` used a 1-thread-per-block serial 32-element loop under `#if defined(__HIPCC__)` instead of the warp-per-block shuffle-reduction kernel (`quantize_q8_0_scaled_kernel`).
3. `moe_hit_grouped_s2`, `moe_hit_grouped_s2_dev`, and `moe_hit_grouped_s2_multi` in `src/kernels/cuda/s2_expert_grouped.cu` launched `swiglu_kernel` and `quantize_q8_0_scaled` separately instead of the fused `swiglu_quantize_q8_0_scaled_kernel` already used by `moe_grouped_s2`, and `gu_grouped_t_kernel`/`down_grouped_t_kernel` allocated 25 KB / 6.25 KB LDS (`GMAX=8`) even for `cap_entries <= 16` (`T=1` MTP draft), capping occupancy at 2 blocks/WGP.
4. `MtpDrafter::forward_body` (`src/core/mtp.cpp`) and `Verifier::forward_range` (`src/core/verify.cpp`) quantized `mixed` to `Q8_1` twice per layer (once in `fused_gr_read_multi` / `native_quantize_q8_1` for routed experts, and a second time inside `shared_expert`), and `fused_gr_read_multi` declined `q8_mixed` when `q8_cnt == nullptr` (`STRATA_QFUSE_HC_LOCAL=1`).
5. `Verifier::forward_range` issued up to 8 separate `cudaMemcpyAsync(..., cudaMemcpyHostToDevice)` calls per layer to copy pinned expert blobs for misses instead of a single batched `copy_from_mapped_multi` device kernel reading mapped host pointers.

Every optimization in this PR is **100% bitwise identical** (`0` differing bits across all outputs).

## Related discussion

Supersedes #3069 (includes the `copy_from_mapped_multi_fallback` alignment-safe scalar fallback for `bytes % 4 != 0` or `4 mod 8` pointers and the `verify_parity` unaligned fallback tests from review on #3069).

## Approach

- Enable wave32 `argmax_rows_kernel` (`128` blocks/row) and `row_top_prob_split_kernel` (`kTopBlocks = 32` single-warp blocks/row) on RDNA HIP (`p.warpSize == 32`, cached per device in `arch[64]`) in `src/kernels/cuda/verify_kernels.cu`.
- Enable warp-per-block `quantize_q8_0_scaled_kernel` on HIP in `src/kernels/cuda/quantize_act.cu` (`STRATA_OLD_Q80_SCALED=1` fallback).
- Route `moe_hit_grouped_s2{,_dev,_multi}` through `launch_hit_swiglu_q8`, specialize `gu_grouped_t_kernel<1>` / `down_grouped_t_kernel<1>` for `cap_entries <= 16` (reducing LDS from 25 KB to 3.125 KB), and use `s2_sync_lds()` in `src/kernels/cuda/s2_expert_grouped.cu`.
- Reuse pre-quantized `x_q8_1` across routed and shared experts in `src/kernels/cuda/shared_expert.cu`, `src/core/mtp.cpp`, and `src/core/verify.cpp`, and allow counter-free `q8_mixed` in `fused_gr_read_multi` (`src/kernels/cuda/fused_gr.cu`) when `q8_cnt == nullptr`.
- Add `copy_from_mapped_multi` in `src/kernels/cuda/elementwise.cu` (with alignment-safe fallback for non-8B-aligned pointers/sizes) to batch mapped-host miss expert copies in `src/core/verify.cpp`.

## Validation

**Passed:**
- `cmake --build build-copilot-hip --target strata verify_parity gr_parity shared_expert_parity s2_expert_grouped_parity arch_defaults_test -j16`: revision `d79fe8e8` on `0.2.0-beta` (`e88ee39d`), AMD Radeon AI PRO R9700 (`gfx1201`, 32 GB, Linux `7.2.9-2-cachyos-server`, ROCm `10.2.0a20260914`), clean build & link.
- `./build-copilot-hip/arch_defaults_test`: revision `d79fe8e8` (`0.2.0-beta`), `arch_defaults_test: OK`.
- `./build-copilot-hip/verify_parity`: revision `d79fe8e8` (`0.2.0-beta`), `gfx1201`, `argmax_rows, row_top_prob_split: bitwise equal`, `copy_from_mapped_multi: bitwise equal (multi=1 and fallback=0, aligned + 1025w + 4B-aligned)`, `quantize_q8_0_scaled: bitwise equal`, `verify_parity: PASS`.
- `./build-copilot-hip/gr_parity --selftest`: revision `d79fe8e8` (`0.2.0-beta`), `gfx1201`, `fused GR counter-free q8_1 (q8_cnt == nullptr), T 1..8, direct/graph + LOCAL=0 guard: pass`, `gr_read/gr_write: 0 failures`, `gr_parity OK`.
- `./build-copilot-hip/shared_expert_parity`: revision `d79fe8e8` (`0.2.0-beta`), `gfx1201`, `native shared x_q8_1 reuse: byte-identical`, `shared_expert: 0 failures`.
- `./build-copilot-hip/s2_expert_grouped_parity` (default, `STRATA_S2_SWIGLU_Q8=0`, and `STRATA_OLD_Q80_SCALED=1`): revision `d79fe8e8` (`0.2.0-beta`), `gfx1201`, all 14 cases bitwise identical (`s2_expert_grouped_parity: 0 failures`, `PASS`).

**Failed:**
- none observed

**Not run:**
- Target-device runtime execution on `gfx11` / `gfx906` and CUDA hardware (no `gfx11`/`gfx906` or NVIDIA GPU in this host).
- SYCL and Metal builds (changes are in CUDA/HIP kernels and engine verify/MTP paths).

## Limits

- Wave64 `gfx906` keeps the single-block `row_top_prob` and `sampler_greedy_kernel` paths (`#if !defined(STRATA_HIP_GFX906)` and `p.warpSize == 32` check).

---

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

Related on this site

Quick links to related pages.