[nisbenz]

Overview

Fixes #30048.

The FA dequant path (use_dequant_kv, added in #25494) dequantizes K and V to f16 into one scratch buffer (prealloc_x). The existing checks only compare K and V separately against maxStorageBufferRange. Nothing checks the combined K+V size against the device's max buffer / allocation size, which is what we actually allocate.

With a large context and a quantized KV cache K+V in f16 can go past that limit. In the issue, at 1M context each of K and V is ~1 GiB in f16. One more batch takes K+V to 2049 MiB, over the 2 GiB limit, and allocation fails

This PR adds the combined size to the use_dequant_kv conditions, checked against both max_buffer_size and max_memory_allocation_size. If it does not fit, FA falls back to the existing quantized path. Normal context sizes are unchanged.

Additional information

Why not maxStorageBufferRange for the sum: that limit is per binding, and K and V are bound separately, so each one alone is still the right check there. The failure here is the allocation of the shared buffer, so the sum has to be checked against the buffer and allocation limits. Both are checked because the scratch is one buffer backed by one allocation.

Testing: reproducing the original failure needs ~1M context, so I lowered the limits instead:

GGML_VK_FORCE_MAX_BUFFER_SIZE=327680 GGML_VK_FORCE_MAX_ALLOCATION_SIZE=327680 ./build/bin/test-backend-ops -o FLASH_ATTN_EXT

• Before: crash with the same ErrorOutOfDeviceMemory as the issue.
• After: passes (falls back to the quantized path).
• Tested on AMD Radeon RX 6750 XT (RADV, Mesa 26.2.4), limits forced down with the env vars above. The bug is not vendor-specific: it triggers on any device once f16 K+V exceeds max_buffer_size / max_memory_allocation_size.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - opus 5.5

[0cc4m]

This is valid, but should not be needed anymore after #30665

[0cc4m]

Can you rebase this? #30665 will take longer to merge, so this can be merged in the meantime.

[nisbenz]

Rebased onto current master.