Pull requests / #2299
#2299 prefill: STRATA_GDN_CHUNKED=1 on gfx11 WMMA (opt-in; gfx1151: recurrence 1.6x, a warm 4K-128K prompt 1.4-2.1%)
open · @libratechw · 0 コメント · GitHub で見る
BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDADocumentationWindowsLinux
本文
## Title
Issue: none (follows #1372, which added `STRATA_GDN_CHUNKED=1` for CUDA; related: #1296, the gfx1100 WY arm)
## Summary
`STRATA_GDN_CHUNKED=1` exists for CUDA (#1372) and is compiled out on HIP. On gfx1151 the prompt's recurrence is
`gdn_rec_quad2c_kernel` plus a norm kernel: token by token, FP32, on the vector ALUs. This PR makes the same switch take a gfx11
path: the WY form of the gated delta rule in chunks of 16 tokens, every product one 16x16x16 WMMA (FP16 operands, FP32
accumulation, the state in the waves' accumulators), with the output norm and the z gate fused into the scan.
It is off by default; the default path is unchanged. It rounds differently from the FP32 recurrence (FP16 operands), which is why
it is opt-in; the KL numbers are below.
## What changed
- `src/prefill/kernels.cu`, HIP only, behind the `STRATA_GDN_CHUNKED` switch in `gdn_recurrence_variant` (variants 0 and 1, the
ones `gdn_recurrence` uses; `STRATA_GDN_REC_HEADS` is read once at the top of the function now, one declaration moved up):
- `gdn_wmma_prep_kernel`: T = (I + A)^-1 and P of every chunk at once (they do not depend on the state); a block per
(key head, chunk), a wave per value head.
- `gdn_wmma_scan_kernel`: a block per value head, a wave per 16 value columns; walks the chunks and loads the next chunk's k/q/v
while this one computes. Its barriers wait for LDS only (`__syncthreads()` also waits for every global load in flight, which
would put the prefetch in front of every barrier), and the prefetched rows are `ext_vector_type` values (HIP's `float4` copies
go through a stack object, which also waits for the load); both are in the code comments.
- `gdn_rec_wmma`: the launch. It returns false, and the existing kernels run, for a part outside `gfx_arch_is_gfx11_wmma`, fewer
than 64 tokens and an output row stride not divisible by 8; the caller skips it under `STRATA_GDN_REC_HEADS`. The part is
judged per call from the current device of the calling thread. The gfx11-only instructions sit inside the gfx11 device passes.
- `tests/hip/gdn_chunked.cpp` and CMake: `hip_gdn_chunked`, `hip_gdn_chunked_off`, `hip_gdn_chunked_default`, and
`hip_gdn_chunked --bench` / `--bench-default` (one layer's time on this path / on gfx1151's default kernels). `tests/hip/gdn_rec_head.cpp` unsets `STRATA_GDN_CHUNKED` (its variants 0 and 1 are the references there).
- `include/strata/prefill/kernels.hpp` (the variant comment), `docs/AMD_HIP.md` (conditions, numbers, how to run the checks),
`docs/DETAILS.md` (the CUDA entry now says what HIP does).
Not touched: the CUDA kernels (the `#if !defined(__HIPCC__)` block), the SYCL port, any default. On HIP `STRATA_GDN_CHUNKED=2` is
the same as 1 (the docs say so).
## Checks
Ryzen AI Max+ 395 (gfx1151), ROCm 7.14.1, Linux 7.0, UD-IQ4_XS unless stated, all on this commit. "Section 5" is the nine variables of
`docs/STRIX_HALO.md` section 5 (with `STRATA_PF_SWITCH_MIN_T=2048` where the doc has 4096); "defaults" is the engine without them
(section 4's gfx1151 defaults stay).
**The default path is unchanged.** With the switch unset or 0, for a chunk of fewer than 64 tokens and for a row stride the path
does not take, the output and the carried state are bit-identical to the existing recurrence (`hip_gdn_chunked_off`,
`hip_gdn_chunked_default`; both set `STRATA_GDN_HEAD=1` and `STRATA_GDN_PP=2`, gfx1151's defaults).
**Numerics (`hip_gdn_chunked`).** 0-16,385 tokens, the 64-token edge, 16-token remainders, row strides 6144 / 6152 / 6145, weak
decay, beta = 0, and a request split at 129 + 128 tokens with the state carried. Bounds: relative RMS of the FP16 output and of the
final state <= 0.003; max |diff| <= 2% of the reference's max + 1e-5; all finite; guard values around the output unchanged. Two
short cases are also compared with a double-precision CPU recurrence of three value heads.
**First-token KL against the same run without the switch** (`STRATA_DUMP_FIRST_LOGITS`; 12 natural-text prompts of 1.5K-24K tokens,
nats, full vocabulary):
| | UD-IQ4_XS mean / max | UD-Q4_K_XL mean / max | top-1 |
|---|---:|---:|---:|
| defaults, with vs without the switch | 0.0067 / 0.044 | 0.0063 / 0.043 | 12 of 12 |
| section 5, with vs without the switch | 0.0039 / 0.027 | 0.0040 / 0.014 | 12 of 12 |
Section 5's own switches against the defaults are 0.0068 (IQ4_XS) and 0.0031 (Q4_K_XL); that is a different change, given for scale
and not as a limit for this one. #1372 reports 0.002-0.009 for the CUDA path. The mean is carried by a few prompts (one at 0.044;
the median is 0.002-0.003). Only the first generated token is compared, and across the 12 prompts it takes only two distinct values,
so top-1 agreement is a weak check: a long generation, coding / agent use and perplexity have **not** been measured, which is the
reason for opt-in.
**Speed, one layer** (`hip_gdn_chunked --bench` against `--bench-default`, which runs gfx1151's default kernels with
`STRATA_GDN_PP=2` and `STRATA_GDN_NOY=1`; the switch is read once per process, so the two are separate runs, alternated three
times; each run reports the GPU-event median of 7 runs after 2 warmups with the same nonzero start state; recurrence + norm;
medians of the three):
| tokens | default kernels | this PR | |
|---:|---:|---:|---:|
| 64 | 0.111 ms (0.082-0.119) | 0.050 ms | 2.2x |
| 256 | 0.279 ms | 0.192 ms | 1.45x |
| 2,048 | 2.603 ms | 1.810 ms | 1.44x |
| 16,384 | 20.924 ms | 13.111 ms | 1.60x |
| 32,768 | 42.142 ms | 26.298 ms | 1.60x |
**Speed, whole prompt.** The same binary on both sides, only `STRATA_GDN_CHUNKED` differing, a fresh engine per side, resident
budget 55 GiB, int8 KV, 11 pairs in alternating order after a discarded one, 256-token answers, the engine's `prompt_ms`. Every
prompt has its own nonce at the front and was read in full (0 reused tokens in every request). Each engine reads a discarded 4K
prompt first: on those first prompts the same pairs gave 95% intervals of [1.004, 1.032] (section 5) and [0.981, 1.015]
(defaults). Before each prompt the client waits until the page-migration and direct-reclaim counters of `/proc/vmstat` have
stood still for 3 s; a pair in which either prompt still met page migration (over 50,000 pages) or direct reclaim was dropped
and measured again (3 of 69 pairs, all 32K and all with the switch on; the dropped prompts were 0.12 s, 0.27 s and 0.00 s slower than the median
of their side, and with them kept 32K gives 1.020 (13 of 13 faster) and 1.015 (12 of 12)). Times are
medians; off / on is the geometric mean of the 11 per-pair ratios with a 95% bootstrap interval:
| tokens | arm | off -> on | off / on | faster | sign test p |
|---:|---|---:|---:|---:|---:|
| 4,096 | section 5 | 3.703 -> 3.641 s | 1.018 [1.016, 1.020] | 11 of 11 | 0.001 |
| 4,096 | defaults | 5.914 -> 5.829 s | 1.014 [1.010, 1.017] | 10 of 11 | 0.012 |
| 32,768 | section 5 | 26.157 -> 25.589 s | 1.021 [1.020, 1.023] | 11 of 11 | 0.001 |
| 32,768 | defaults | 44.082 -> 43.479 s | 1.015 [1.012, 1.017] | 11 of 11 | 0.001 |
| 128,000 | section 5 | 107.149 -> 104.866 s | 1.021 [1.020, 1.023] | 11 of 11 | 0.001 |
| 128,000 | defaults | 193.439 -> 191.055 s | 1.014 [1.012, 1.016] | 11 of 11 | 0.001 |
So about 2% on top of section 5 and about 1.5% on the defaults, for a prompt that is not the first one after a start. With 36 gated-DeltaNet layers the layer figures predict about
0.56 s saved on a 32K prompt; 0.57 s (section 5) and 0.60 s (defaults) were measured.
**Builds.** HIP: `kernels.cu` compiles for gfx1030, gfx1100, gfx1151, gfx1152 and gfx1201; the four HIP tests above (the three
`hip_gdn_chunked*` and `hip_gdn_rec_head`) pass on gfx1151. Run on gfx1151 only: nothing here has been run on another gfx11 card or
with more than one GPU. CUDA (RTX 3060, CUDA 13.0, sm_86, Windows, this commit): `strata` builds (140 of 140 steps, no error). `gdn_rec_parity` passes
too, but it is a self-contained copy of the CUDA recurrence that does not link the engine, so it says nothing about this change.
Not built: SYCL (its own sources, none touched).
## Extra Notes
- #1296 adds a WY arm (`STRATA_GDN_WY`) to the same dispatch for gfx1100, where it loses to the serial kernel. This PR starts from
gfx1151's `gdn_rec_quad2c_kernel`, uses 16x16x16 WMMA tiles for every product and reuses the existing switch instead of adding a
second one. They share no code that I found; whichever lands second needs a small rebase of the dispatch.
- Open question for review: the 64-token minimum and the 16-token chunk come from gfx1151 only. Whether gfx1100 / gfx1150 want the
same values is untested.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
関連リンク
インストール・モデル・リリースへの站内リンク。