hip : skip fully masked KV tiles in WMMA Flash Attention

[moristeven477-ship-it]

Overview

Improve HIP prefill when server slots share a unified KV cache. The AMD WMMA attention path still loads K/V and performs matrix multiplication for interior tiles masked out for the current sequence, as reported in #28495.

The 19-line change skips a non-sparse AMD WMMA iteration only when every relevant mask element is exactly -INFINITY. A finite probe takes the existing path; otherwise a cooperative full-tile check and __syncthreads_or decide whether to skip. The outer kernel still handles sinks and partial-result fixups. Ten operator fixtures cover interior masks, an off-probe visible element, and masked queries with sinks.

Credit to DanoPTT and alankila for the reproductions and diagnosis in #28495 and #14924. Vulkan already has masked-block skipping (#17186). The WMMA dispatch changes in #26419 affect whether this path is selected on RDNA3/head256.

Additional information

RX 7900 XTX, Ubuntu 24.04, ROCm 7.2.1; Qwen3.8-27B-UD-Q4_K_M with full offload and Flash Attention. Context 98,304, two unified-cache slots, 40,960 input tokens and 64 output tokens, slots 0/1/0, prompt reuse disabled. Third-request means from three independent restarts per variant, in stock/patch/patch/stock/stock/patch order:

KV	Stock request time	Patched request time	Time reduction	Prefill throughput gain
F16	77.414 s	57.833 s	25.29%	35.17%
Q8_0	79.793 s	59.488 s	25.45%	35.85%

Paired benchmark base: 41abbfd599fbdd3470fcae0a1fb6530ad8403cd7. The benefit targets shared-cache prefill. Ordinary decode stays close to stock; split KV remains faster here when the prompts fit its fixed slot capacity.

Validation:

• Submitted base 4c9233c034fc450dcf34c7c0988aebe6da5cdf1a: all 11 local CI stages passed, including 3,988 executed Flash Attention CPU-reference checks. The 10 new fixtures pass on both stock and patch.
• The original paired runs have identical generated token IDs in all 36 request-position comparisons. Printed Wikitext-2 PPL matches stock: F16 6.4117 +/- 0.12601, Q8_0 6.4165 +/- 0.12618 (eight chunks, context 4096).
• The submitted base also passed llama-bench at depths 0/8192/32768, F16 perplexity, and three 40K-input requests with outputs matching the original patched run. These follow-up checks are separate from the paired benchmarks.

Runtime coverage is Linux gfx1100. Concurrent controls and output variability are documented in the evidence. The skip assumes finite attention arithmetic; overflowing inputs and future asynchronous HIP loading require separate consideration.

Reproduction scripts, hashes, raw measurements, profiles, controls, and CI logs: validation artifacts. The primary archive contains paired experiments; the supplement validates the submitted base.

Requirements

• I have read and agree with the contributing guidelines.
• AI usage disclosure: YES. Codex investigated and implemented the patch and tests, ran the local experiments, and assembled the evidence. The contributor reviewed the code, supplied the initial summary, and explicitly authorized Codex to expand the technical description and submit the patch. This description and the commit message were prepared with Codex.

[JohannesGaessler]

There already is an auxiliary kernel flash_attn_mask_to_KV_max that makes it possible to skip fully masked KV slices for multiple sequences. A solution for unified KV cache should look something like this. Avoid dropping a bunch of extra instructions into an inner loop as was done in this PR.

[JohannesGaessler]

According to the llama.cpp AI usage policy:

It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...).

[DanoPTT]

In my config (-ub 512, --kv-unified) flash_attn_mask_to_KV_max never runs:
the gate at fattn-common.cuh:1108 requires Q->ne[1] >= 1024 || Q->ne[3] > 1,
but -ub 512 caps Q->ne[1] at 512 and --kv-unified keeps Q->ne[3] == 1.
So KV_max stays null and the kernel never skips a single tile.
Verified against current master (60081bb2).