vulkan: fix OOB check in flash_attn_mask_opt (llama/20296)

This commit is contained in:
Jeff Bolz
2026-03-16 13:10:15 +02:00
committed by Georgi Gerganov
parent 26ee4f7362
commit 6c5e3aac3e
2 changed files with 60 additions and 40 deletions
+1 -1
View File
@@ -8840,7 +8840,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
}
// Only use mask opt when the mask is fairly large. This hasn't been tuned extensively.
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768;
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16;
vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
mask != nullptr, use_mask_opt, logit_softcap != 0);