vulkan : keep masked keys from leaking NaN through flash attention

[Nathanw1014]

This fixes the poison=1 QSA test failures on Vulkan. If a masked key's K or V row holds a NaN, flash attention still lets it through (the mask is added to the score, and V is multiplied by a probability of 0, which stays NaN).

• Masked scores are now set straight to -inf. That covers K and costs nothing.
• Zeroing masked V rows is a separate pipeline variant, since it costs occupancy. It's used for selected-key attention (the QSA ops, where selections can point at any cell), or everywhere with GGML_VK_FA_NAN_SAFE=1 for a unified multi-sequence cache. Normal attention stays as fast as before and keeps relying on freed KV cells being zeroed (Correctness: repeatable output on gfx1151 (freed-cell zeroing, deterministic top-k), plus two upstream picks #20).

On a Strix Halo (gfx1151, RADV): QSA_PREFILL 22/26 -> 26/26, QSA_DECODE 82/92 -> 92/92, FLASH_ATTN_EXT all passing. Prefill on 27B and 35B was unchanged at depth 0, 12k and 32k when I measured it before #111.

Written with Claude, tested by me.

Assisted-by: Claude Opus 5.5

[gaetan-puleo]

Independent verification on a second Strix Halo (strix-halo-002, gfx1151, RADV)

Built test-backend-ops with -DGGML_VULKAN=ON at the PR head 07afdcf0c. To compare against master I reverted the five changed files to 1531bf5c4 and relinked in the same tree, so both runs use identical build settings (GGML_NLOOP=1, GGML_N_THREADS=1, coopmat enabled).

test	master 1531bf5c4	PR 07afdcf0c
-o QSA_PREFILL	22/26	26/26
-o QSA_DECODE	82/92	92/92
-o FLASH_ATTN_EXT	—	5344/5344

The four prefill failures on master are exactly the poison=1 cases (513,512,17), (640,4096,257), (128,512,17), (129,4096,257) with ERR ≈ 1.0. This reproduces your 22→26 and 82→92 numbers.

Code review

The diagnosis is right, and so is the fix. The masked score is dot(Q, K) computed before the bias is added; a NaN in the poisoned K row makes it NaN, and adding an -inf mask bias leaves it NaN (then exp and 0 * NaN propagate it). Setting masked scores straight to -inf and dropping V contributions for keys hidden from every row is the correct way to break the chain.

• FA_MASK_DEAD = -65504 catches both f16 -inf and the saturated lowest finite value; no legitimate additive bias is that negative.
• The dead_stamp[c] == j + 1 / dead_blk == j + 1 stamping avoids clearing shared state every block. The writes land before the existing barrier() (line ~254) that precedes the dead_blk read, and dead_stamp is read in the V phase after later barriers, so the cross-thread access is ordered. atomicMax handles the multiple-keys-per-tile case.
• Forcing mv[k] = FA_MASK_DEAD on rows past N is a good catch: without it a partial tile's zero mask fill would count as visibility and the key would never be marked dead.
• The cm1 per-key vote assumes Br / 4 == 4, i.e. Br == 16. That holds today because the coopmat path fixes block_rows = 16, and the shuffle groups don't straddle a subgroup boundary (subgroup size is a multiple of 4). Worth a comment/assert if cm1 Br is ever allowed to change, since subgroupShuffleXor(1,2) would then be wrong.
• ggml_vk_flash_attn_coopmat_shmem_support now counts (Bc + 1) * sizeof(uint32_t) for every coopmat FA pipeline, not just NAN_SAFE_V ones. That mirrors the unconditional dead_stamp/dead_blk declarations, but it can in principle flip a borderline config off the coopmat path; I didn't hit one. Normal attention is untouched behaviourally (nan_safe_v is only set with a mask and src[5], or via GGML_VK_FA_NAN_SAFE=1), which the 5344/5344 FLASH_ATTN_EXT run supports.

Note: the windows (CI vulkan) check is red, but only on the pre-existing test-qsa-fast.cpp setenv build error — it is not from this PR and #138 fixes it.

LGTM. Real correctness fix, covered by tests that fail without it, reproduced on independent hardware.