hip : skip fully masked KV tiles in WMMA Flash Attention

[moristeven477-ship-it]

What changed

When server slots share a unified KV cache, the AMD WMMA Flash Attention kernel can still compute interior KV tiles that are fully masked for the current request. Following the review suggestion, this revision moves mask scanning and reduction out of the attention loop and into an auxiliary preprocessing kernel.

The new flash_attn_mask_to_KV_blocks classifies each block as fully masked, all-zero, or general. The attention kernel skips runs of fully masked blocks, omits mask loading and mask addition for all-zero runs, and keeps normal mask evaluation for general blocks. Sixteen two-bit classes pack into a single 32-bit word and are reused across attention heads.

This path is limited to synchronous, dense AMD WMMA. It covers the measured -ub 512 --kv-unified configuration, in which the existing flash_attn_mask_to_KV_max helper is not launched. The patch also includes the broadcast-mask and query-tail fixes plus additional operator tests.

On RX 7900 XTX / Ubuntu 24.04 / ROCm 7.2.1, third-request time drops by approximately 33.3% versus upstream, and by a further 10% versus V1. The gain is primarily in prefill; this is not a general decode or concurrent-server throughput claim.

The model is Qwen3.8-27B-UD-Q4_K_M, fully offloaded, with Flash Attention enabled. The workload uses 98,304 context capacity, 40,960 input / 64 output tokens, and two unified-cache slots requested sequentially in 0/1/0 order. Ubatch is 512, with no prompt cache reuse and no speculative decoding.

KV cache	Upstream	V2	95% CI for relative time change, V2 vs upstream
F16	77.8 s	51.9 s	[-33.447%, -33.225%]
Q8_0	80.3 s	53.5 s	[-33.330%, -33.257%]

Times are geometric means over six restarts per variant, covering all six U/V1/V2 execution orders on the common base c6824a9. Confidence intervals use paired log-ratios. The current source is ported to ec91ab5; the three-arm performance comparison has not been repeated on that base.

Validation

• Operator tests on ec91: 4,007 passed, 0 failed, 1,323 unsupported. The extended suite adds boundary cases. Upstream and V1 each fail six of the new broadcast-mask cases that V2 fixes; these are not failures in pre-existing upstream CI tests.
• Model checks: fixed F16/Q8_0 perplexity matches the per-chunk reference values at printed precision, without claiming bitwise-identical internals. Sequential requests match the reference token IDs.
• Kernel trace: all 3,840 WMMA calls in the instrumented run had a preceding classifier launch on the same stream. Peak dispatch metadata was 5,184 bytes; this is not total extra VRAM.
• Local HIP CI: passed in aggregate, after one HTTP-test failure was re-run and passed and the model-download interruptions were resolved. This was not an uninterrupted all-green run. The full operator suite also relied on a separate TOPK comparator fix, which is excluded from this attention patch.

The V2 evidence release contains the patch, commands, hashes, raw data, and logs. The V1 artifacts do not substitute for this evidence.

• V2 patch on ec91; SHA256 checksums.
• Reproduction commands: c682 performance / ec91 validation.
• Raw performance data and primary-workload findings.
• CI and model logs, plus CI scope.
• Matched operator tests, calibration, and raw kernel traces.

Credit & inspiration

Thanks to JohannesGaessler for the auxiliary-preprocessing suggestion, to DanoPTT for the V1 testing and the helper-launch diagnosis, and to alankila together with the earlier reports in #28495 and #14924. The masked-block preprocessing idea follows other backends, including Vulkan (#17186); the ec91 port retains ynankani's swizzle work (#28536).

[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).

[moristeven477-ship-it]

https://github.com/moristeven477-ship-it/llama.cpp/releases/tag/rdna3-masked-kv-v2-20260921

[DanoPTT]

As a happy user of the V1 patch, I tested the V2 patch on gfx1201 (Radeon AI PRO R9700, ROCm 10.0.0, Windows), so this is
a different WMMA generation than the gfx1100 results above.
llama-server, Qwen3.8-27B, Q5_K_XL, -np 3 --kv-unified -ub 512 -c 161280 --no-cache-idle-slots,
sequential requests filling slots 0, 1, 2, 0, 1, 2 so the mask gets interior -INF blocks while each request still runs alone.
Three builds off one tree, differing only in the patch: no skip, V1, V2.
8 visits in A,B,B,A,A,B,B,A order, so each median below is over the 4 visits of that arm.

Prompt processing, t/s. Step 0 has the other slots empty, so there is nothing to skip; by step 2 all three slots hold content.

step	~5.4k tokens			~11k tokens
	no skip	V1	V2	no skip	V1	V2
0	894.3	892.0	901.0	783.6	841.6	870.7
1	846.5	874.2	887.5	746.6	826.9	864.6
2	803.2	857.0	884.0	712.2	813.1	858.5
3	785.3	849.1	880.9	690.4	804.9	855.1
4	780.6	847.8	877.7	688.5	804.3	853.7
5	795.4	848.7	876.2	708.6	810.1	855.5

Two runs against the same no-skip build (V2 in one, V1 in the other); its control column agreed between them within 0.5%.
Over steps 2-5 that is +10.1-12.4% for V2 against no skip at 5.4k and +20.5-24.0% at 11k,
of which +2.4-3.1% and +5.0-5.8% is V2 over V1.
More useful than the percentages is the shape: across the steps prompt processing falls 12.1% without the patch, 4.4% with V1 and 2.0% with V2 (11k column), so V2 largely removes the decay reported in #28495.

Worth noting for this config: the separate use_mask_blocks path is not behind the Q->ne[1] >= 1024 || Q->ne[3] > 1 condition, so it actually runs here.
As I mentioned earlier in this thread, that condition means flash_attn_mask_to_KV_max never fires with -ub 512 and --kv-unified, so on this setup the existing helper skips nothing at all, not even the mask tail.

Probe cost: step 0 is the case with the other slots empty, where there is least to gain, and V2 is not slower there in either column.

Output is unchanged: perplexity is bit-identical and token IDs match over 640 tokens in the deterministic sequential-slot shape. test-backend-ops FLASH_ATTN_EXT passes 4599/4599 including the new mask-class cases.

Important note: this tree has no sparse FA (#27970), so I dropped the use_sparse template parameter and the indices arguments when applying the patch. The numbers are for that adaptation, not for the patch verbatim.

[JohannesGaessler]

The way this is implemented still introduces too much technical debt. This needs a human to do the software design.

[moristeven477-ship-it]

The way this is implemented still introduces too much technical debt. This needs a human to do the software design.这种实现方式仍然引入了过多的技术负债。这需要有人来完成软件设计。

Can you give me some hint or target please?

[JohannesGaessler]

To make it clear: my impression is that the design in this PR is mostly driven by a language model and significantly more human effort is needed to come up with a better design. This should be done by someone else.