fix MHA KV-cache allocation and singleton mask broadcast across backends

[ghost]

Summary

This PR fixes persistent KV-cache allocation and singleton 3D attention-mask broadcasting across the generic, x86, ARM, MIPS,
LoongArch, and Vulkan attention implementations.

It also fixes the equivalent singleton-mask issue in Vulkan SDPA.

Problems

Persistent MHA KV caches used the wrong allocator

The architecture-specific MHA implementations concatenated cached and current K/V tensors with:

  k_affine.create(dst_seqlen, embed_dim, elemsize);
  v_affine.create(dst_seqlen, embed_dim, elemsize);

These tensors are returned through top_blobs[1] and top_blobs[2] and persist across subsequent decoding steps. They must
therefore use the caller-provided blob allocator.

Affected optimized implementations:

• x86
• ARM
• MIPS
• LoongArch

The concatenated caches are now created with:

k_affine.create(dst_seqlen, embed_dim, elemsize, opt.blob_allocator);
v_affine.create(dst_seqlen, embed_dim, elemsize, opt.blob_allocator);

Singleton 3D masks were treated as per-head masks

An attention mask with shape:

(dst_seqlen, src_seqlen, 1)

is intended to be broadcast across every attention head.

Several implementations treated every 3D mask as a per-head mask and selected:

attn_mask_blob.channel(head)

For head > 0, this accesses a channel that does not exist.

Affected paths included:

• generic MHA FP32
• generic MHA INT8
• generic weight block-quantized MHA
• x86 MHA
• ARM MHA
• MIPS MHA
• LoongArch MHA
• Vulkan MHA
• Vulkan SDPA normal path
• Vulkan SDPA Flash Attention path
• Vulkan SDPA cooperative-matrix path

Singleton masks now always select channel zero, while masks with multiple channels retain their per-head behavior.

Vulkan mask indexing

The Vulkan SDPA shaders use the following addressing logic:

mask_head = attn_mask_dims == 3 ? head : 0;
mask_index = mask_head * mask_cstep + row * dst_seqlen + column;

Previously, a dims=3, c=1 mask was still passed to the shader as a 3D per-head mask. Every head greater than zero therefore
read beyond the only allocated mask channel.

This was reproduced on an NVIDIA RTX 4060:

past_seqlen=2560
num_heads=4
mask_channels=1

value not match at c:1 h:0 w:1
expect 0.016805 but got 0.015332

The old implementation failed with exit code 1. After treating a singleton 3D mask as a 2D broadcast mask, the same test passes
with exit code 0.

Vulkan MHA additionally keeps 3D masks unpacked so that head packing cannot be confused with query-position packing.

Regression coverage

The MHA and SDPA tests now cover:

• single-token decoding
• multi-token decoding with 3 and 4 tokens
• long KV caches up to 2560 tokens
• singleton masks
• per-head masks
• FP32 MHA
• INT8 MHA
• weight block-quantized MHA
• cross-attention cache reuse
• 257-step rolling KV-cache decoding
• blob allocator ownership for returned K/V caches

Verification

The following tests pass:

test_multiheadattention_kvcache CPU: PASS
test_sdpa_kvcache CPU: PASS
test_multiheadattention_block_quant CPU: PASS
test_multiheadattention_kvcache Vulkan: PASS
test_sdpa_kvcache Vulkan: PASS

Additional validation:

• ARM MHA translation unit compiled with an ARM64 target.
• MIPS and LoongArch MHA translation units compiled successfully.
• git diff --check passes.
• Downstream LiteOCR LaTeX OCR S and M models complete successfully with Vulkan enabled.

The optimized MHA implementations remain enabled for both single-token and multi-token KV-cache decoding. No reference fallback
is introduced.

[tencent-adm]

Thank you for your submission, we really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

fdph seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}