Pull requests / #2588
#2588 prompt attention (HIP): gate_attn in the WMMA kernel's epilogue (opt-in STRATA_ATTN_GATE=1, the same bits)
open · @libratechw · 0 commentaires · Sur GitHub
BenchmarksServer & APIAMD / HIPNVIDIA / CUDADocumentationWindowsLinux
Description
## Title Issue: none ## Summary On the full-attention layers' prompt path, the attention writes `attn` and `gate_attn` then reads it back, multiplies each value by the sigmoid of its gate and writes the o-projection's FP16 input. With `STRATA_ATTN_GATE=1` the HIP `STRATA_PA_FAST` WMMA kernel does that in its epilogue, with `gate_attn`'s expression, so `attn` is not written and read back. The bits are the same. The gain is small (+0.6% on a 32K prompt), so it is opt-in. Say if you would rather not carry it. (It was first part of #2019, which now carries only the q-preparation fusion.) ## What changed - `src/kernels/cuda/qsa_prompt_attn.cu`, `include/strata/kernels/qsa_prompt_attn.hpp`: `prompt_attn_wmma_kernel<SPLIT, GATED>` and `qsa_prompt_attn_batch_gated`, which returns false (nothing launched) wherever `qsa_prompt_attn_batch` would not run the `STRATA_PA_FAST` kernel. - `src/prefill/prefill.cpp`: with `STRATA_ATTN_GATE=1` (and unrotated K/V) it takes the gated launch; otherwise the attention and `gate_attn` run as before. - `src/kernels/qsa_prompt_attn_parity.cpp`, `CMakeLists.txt`: a gated check (the FP16 bits against the attention followed by `gate_attn`), and an argument `gated` that makes the exit status that check's alone; ctest `hip_prompt_attn_wmma_gated` runs it with `STRATA_PA_FAST=1`: 0 values differ on gfx1151. Without the switch, the WMMA attention's output hash (`hip_prompt_attn_wmma`, three cases, `STRATA_PA_FAST=0` and `=1`) is 296b59d2's in every run (two rounds; measured on the build that also carried the q-preparation fusion, whose attention files are this branch's). Ryzen AI Max+ 395 (gfx1151, ROCm 7.14.1, Linux), on 296b59d2, the branch's binary without and with `STRATA_ATTN_GATE=1`, UD-Q4_K_XL, `--serve --prefill 16384 --spec 4`, the section-5 switches of `docs/STRIX_HALO.md` with `STRATA_PF_SWITCH_MIN_T=2048` and `STRATA_PF_FUSED_KQ=1`, a 32K chat request, fresh servers alternating, 7 pairs: 1,243 -> 1,253 tok/s (median), the median of the pair ratios +0.6%, 7 of 7 pairs faster. The records and scripts are in [libratechw/Strata@686de2e9, `2019g/`](https://github.com/libratechw/Strata/tree/686de2e9eedd6e34c18e7d5fab7a2fd2fe3e4241/2019g). A check outside ctest (a local copy of the parity test, in the records): with an output row stride of 6,152 instead of 6,144, the gated epilogue matches attention + `gate_attn` and leaves the padding untouched in the tool's three cases. Builds and tests on this branch: - HIP gfx1151: ctest, 113 tests: 105 pass (`hip_prompt_attn_wmma_gated` among them), 5 skip, and the same three fail as on 296b59d2 (`ple_parity`, `expert_parity`, `pool_test`). - HIP: `strata_kernels`, `strata_prefill` and `strata_mmq` compile for gfx1012, gfx1030, gfx1100, gfx1103, gfx1150, gfx1151, gfx1200 and gfx1201. - CUDA 13.0, RTX 3060 (Windows): builds; ctest, 117 tests: 114 pass, 2 skip, and `ple_parity` fails as on 296b59d2. - SYCL: not affected (the SYCL port has its own prompt attention and `prefill.cpp`). Not run: other gfx11 / gfx12 cards, UD-IQ4_XS. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Sur le site
Liens install, modèles, releases.