HIP: fix corrupted flash-attention output at fp32 precision - #24984
HIP: fix corrupted flash-attention output at fp32 precision#24984RapidMark wants to merge 1 commit into
Conversation
The tile kernel picked fp16 vs fp32 storage from the arch capability alone and ignored the op's requested precision, so an fp32 request was silently down-cast to fp16. On long key lengths that rounding compounds into a visibly corrupted (ghosted) result for non-GQA attention — the path AMD uses without rocWMMA flash-attention. Honor GGML_PREC_F32: store Q/K/KQ in fp32 when it's requested and fits shared memory; keep the existing fp16 path for shapes where it doesn't (byte-for-byte unchanged). Gated to HIP.
github-actions
Bot
added
ggml
|
Hi @RapidMark, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
Sorry again about the multiple PRs... but I'm finding some important issues and trying to fix them... As for AI... I did use AI to run the multiple tests to verify everything is working and to create the image... but that was about it. |
1 participant
Why this exists
On AMD, image generation came out visibly wrong at 1024×1024 and larger: a ghosted, doubled subject that got worse as the resolution grew. Smaller resolutions were clean, and the same job ran clean on CPU and on NVIDIA, so it was specific to AMD and only showed above a certain resolution.
The flash-attention kernel was quietly ignoring the precision the model asked for and running in half precision instead. That is usually harmless, but at large resolutions the rounding piles up until it shows in the image. Worse, it fails silently: no error, just a wrong, ghosted picture.
What it does
When the model asks for fp32 precision, the kernel now uses it, instead of quietly running in half precision. Where fp32 will not fit, it falls back to the existing path, so those cases stay unchanged.
This only affects AMD/HIP. NVIDIA never uses this kernel, so CUDA builds are untouched.
Why it's the right fix
Before this, there was no way to get a correct image out of an affected AMD GPU. The kernel did not fall back to anything; it just silently produced a garbled result. The only way to a correct image was to give up on the GPU and run on the CPU, which is far slower. So this is not a slower-path tradeoff: it makes the AMD GPU usable for these images, where before it was quietly broken. Builds with rocWMMA enabled already had a correct GPU path; this fix is for the AMD setups that do not.
How I found it
I reproduced it with Chroma (a non-GQA model, head size 128) in stable-diffusion.cpp on an RDNA4 GPU (a Radeon AI PRO R9700), built without rocWMMA flash-attention. Comparing the same generation on the CPU and on the GPU step by step, everything matched until the attention step (the
FLASH_ATTN_EXTop), which was the one place the GPU result diverged. It also does not show up intest-backend-ops, because that test's synthetic inputs are too uniform to trigger the rounding.It surfaced alongside a separate bug on the same path, an incorrect mask shape in masked flash attention, which I fixed upstream in leejet/stable-diffusion.cpp#1625 (merged). Those two changes are independent; together they make Chroma's masked flash attention correct on HIP.
Validation
The clearest way to see the bug is to render the same prompt and seed on several backends and compare. The correct backends agree on one image; the unpatched AMD tile kernel is the lone outlier, and the fix brings it back into agreement.
Chroma1-HD, 1280², 20 steps, euler, flash attention on, fixed seed:
That is the point: the unpatched AMD result is not "just a different valid render." Three correct paths converge on one image: a different vendor (NVIDIA) and two different AMD kernels, the matrix-core rocWMMA path and the fp32 tile fix. Only the unpatched fp16 tile diverges. With the fix, that path rejoins them. At higher resolution the unpatched path degrades further, to a fully black image at 2048, while the others stay correct.
Environment:
test-backend-opsis unchanged; head sizes 512 and 576 fall back to fp16.Reproduction
Reproduced via stable-diffusion.cpp (commit
4e85e07), which vendors this ggml kernel.Build
sd-clifor AMD (HIP):The wrong result is this build as-is (the unpatched fp16 tile path); the correct result is the same build with this PR's ggml patch applied. For the NVIDIA reference, configure with
-DSD_CUDA=ONinstead.Render (identical command on every backend, only the build differs):
Performance note
This is the tile flash-attention kernel, the path used whenever the matrix-core path (rocWMMA flash-attention) is not in use. That path is opt-in, a build flag that is off by default, so a default HIP build runs the tile kernel on any AMD GPU, with or without matrix cores. The tile kernel has no hardware-accelerated fp32. On RDNA the fast fp16 path uses a packed dot instruction that fp32 cannot use, and fp32 storage roughly doubles the shared memory the kernel needs, which lowers occupancy. So running this kernel in fp32 is inherently slower than fp16. It is still far faster than the only correct option this path had before, which was to fall back to the CPU.
That slowdown is not a reason to leave it broken. The kernel is an active code path, and this is current hardware, not a legacy quirk: it was caught on an RDNA4 Radeon AI PRO R9700. Any build without rocWMMA flash-attention lands here, and the hardware affected worst is anything with no matrix cores at all, such as the RDNA2 generation, which cannot use rocWMMA and so has the tile kernel as its only attention path, with no faster option to switch to. All of them silently produce a ghosted image. GPUs that do have the matrix-core path keep using it and are unaffected. For everyone on the fallback, a small and correct slowdown is far better than a fast and wrong result.