Pull requests / #2017

#2017 prompt attention (gfx11 / gfx12 WMMA): issue a chunk's reads before its compute

open · @libratechw · 0 comentários · No GitHub

AMD / HIPNVIDIA / CUDAModels & quants

Descrição

## Title
Issue: none


## Summary

`prompt_attn_wmma_kernel` interleaved a chunk's global reads with its compute: the V slice was staged first, then each 16-cell
tile loaded its K scale and fragments between the q.k WMMAs, so a wave waited on each read in turn. A thread trace on gfx1151
showed 56% of the kernel's time in `s_waitcnt vmcnt` and 0.4% in WMMA. The chunk's reads now go out together (`fetch`), then
the same compute runs on them (`chunk`). Values and the order of the arithmetic are unchanged.

## What changed

`src/kernels/cuda/qsa_prompt_attn.cu` only (the WMMA kernel, HIP).

Ryzen AI Max+ 395 (gfx1151), ROCm 7.14.1 (clang 23), `hip_prompt_attn_wmma` with `STRATA_PA_FAST=1`, three interleaved runs
against the unchanged kernel, ms per chunk: 32K context / 2,048 queries 12.50 -> 11.29 (-9.7%), 1,500 / 1,500 3.39 -> 3.20,
2,100 / 256 1.30 -> 1.23.

Bitwise: the output hash (`NEWHASH`) of the unchanged kernel and of this one is the same with `STRATA_PA_FAST=1` and `=0`
(re-checked on v0.1.43: 3dc8754c6a35763e / aa99bad211c243a0 in both builds). gfx12 compiles the same change; it was not run there.

## Extra Notes

The kernel is on a path that is opt-in (`STRATA_HIP_WMMA=1`, `STRATA_PA_FAST`), so the default path is untouched.
Not built: CUDA, SYCL.

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

No site

Links install, modelos, releases.