Pull requests / #944

#944 hip: gfx11 WMMA prompt attention and selection scorer, gfx1151 hipBLASLt table

closed · @nsbradley88 · 0 Kommentare · Auf GitHub

BenchmarksSetup & installServer & APIAMD / HIPModels & quantsLinux

Beschreibung

## Summary

Three changes that take Strata's prompt path on **Strix Halo (gfx1151, Ryzen AI Max+ 395 / Radeon 8060S)** from ~230 to ~700 tok/s at 16K–250K context:

1. **gfx11 WMMA prompt attention.** `STRATA_HIP_WMMA=1` runs the RDNA4 matrix-core prompt attention on gfx11 (RDNA3 / RDNA3.5) too. gfx11's `v_wmma_f32_16x16x16_f16` keeps all 16 k of a lane's A row / B column in that lane (both wave halves the same), and puts accumulator element `i` in row `2*i + lane/16`. Every fragment in the kernel is loaded directly, so the two layouts differ only in `FK` / `PA_KOFF` (which k a lane loads) and `PA_ROW` (an accumulator element's row). The gfx12 code path is unchanged.
2. **gfx11 WMMA selection scorer.** `STRATA_SELECT_WMMA=1` gets the same treatment (`SEL_FK` / `SEL_KOFF` / `SEL_ROW`). The 32-byte gfx11 B fragment is read from LDS as two 16-byte halves, because the LDS rows are only 16-byte aligned.
3. **`tools/hip/gfx1151-hipblaslt-100500.txt`.** `tune_hipblaslt` was run on the 32 shapes of the gfx1201-100500 table, with ROCm 10.2.0a20260930 (whl-next) and hipBLASLt 1.5.0 `efafc1c1`.

Both kernels stay opt-in, as on RDNA4, and nothing changes for other architectures. The parity test's skip check now also accepts gfx11.

This PR does **not** add gfx1151 to setup.py or the CMake arch list, since you mentioned bringing up Strix Halo yourselves in #422. The changes are meant to drop into that work. To build for gfx1151, I used #422's `CMAKE_HIP_ARCHITECTURES` change locally.

## Validation (gfx1151, ROCm 10.2.0a20260930, Linux 7.0)

| Test | Result |
| --- | --- |
| `hip_prompt_attn_wmma 32768 2048 3` | PASS, vs FP64 3.3e-6 / 3.8e-6 / 4.5e-6 (old kernel 2.1e-6 / 1.8e-6 / 2.0e-6); **3.4–4.0x** per chunk (85.6 → 25.3 ms at 32K) |
| `qsa_select_bench 32768 256 3 262144` | PASS, 2.8e-7 of the score scale vs FP64; selections identical 256/256; scores **4.2x** |
| `qsa_select_bench 262144 256 3 262144` | PASS, 2.3e-7 of scale; selections identical 256/256; scores **4.1x** (19.8 → 4.9 ms) |
| hipBLASLt table | 3.9–15x vs plain hipBLAS per shape; full prefill with `STRATA_HIPBLASLT_VERBOSE=1`: `launches=3138 fallbacks=0` |
| End to end | 244K-token prompt, greedy, 128 tokens: identical to the default kernels for the first 40 tokens, then diverges (both WMMA arms; the kernels are documented as not bitwise equal to the default ones) |

gfx1100 and gfx1201: the branch compiles with `-DCMAKE_HIP_ARCHITECTURES="gfx1100;gfx1201"` (`strata`, `hip_prompt_attn_wmma`, `qsa_select_bench`). **Not run on RDNA3 or RDNA4**, since I only have the gfx1151 machine.

## Prefill on gfx1151 (UD-IQ4_XS, engine `--stats`, 256K context)

| Step | 16K prompt | 64K | 244K |
| --- | ---: | ---: | ---: |
| ROCm 7.10.0a20251120 (setup's pin) | 231 | | |
| ROCm 10.2.0a20260930 | 276 | | |
| + gfx1151 hipBLASLt table | 499 | | 396 |
| + `--prefill auto:16384` | 555 | 562 | |
| + WMMA prompt attention | | 683 | 556 |
| + WMMA selection scorer | | | 681 |

## Through the server (OpenAI API, 512-token answers, 256K context; prefill / decode tok/s)

Same UD-IQ4_XS file in both engines. llama.cpp is ROCm master `f7b384c1e`, `-fa on -ub 2048 -b 2048`, one slot.

| Prompt | llama.cpp | Strata (this PR + ROCm 10.2) |
| ---: | ---: | ---: |
| 2K | 361 / 21.1 | 507 / 30.0 |
| 16K | 325 / 18.1 | 693 / 30.9 |
| 32K | 308 / 15.4 | 750 / 31.7 |
| 64K | 235 / 11.0 | 732 / 31.1 |
| 136K | 167 / 7.1 | 700 / 26.9 |
| 250K | 115 / 4.4 | 648 / 25.9 |

## Notes for the Strix Halo bring-up

- **ROCm 7.10.0a20251120 (TheRock gfx1151 wheels, the Linux pin) segfaults on this machine** in `rocr::AMD::GpuAgent::QueueCreate` → `ReleaseQueueMainScratch`, on the first DMA queue (called from `probe_pcie_h2d_gbps`). ROCm 10.2.0a20260930 from `whl-next` (`rocm[libraries,devel,device-gfx1151]`) has no such crash, gives +20% prefill on its own, and is what the table above was tuned with.
- On the APU, `mem_info_vram_total` is the BIOS carve-out (512 MB here); the GPU's real pool is `mem_info_gtt_total`.
- Prefill time at 244K after these changes: expert gate/up 15%, gdn recurrence 11%, selection 11%, attention 10.5%, gdn 9.7%, hc read 9.5%, expert down 8.8%.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Mehr auf der Site

Links zu Install, Modellen, Releases.