vulkan: Handle FA with all -inf mask values (#16447)

2025-10-20 22:16:08 -05:00
parent 6de8ed7519
commit fb349848f3
4 changed files with 8 additions and 4 deletions
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
@@ -345,7 +345,7 @@ void main() {

    float Lfrcp[Br];
    [[unroll]] for (uint32_t r = 0; r < Br; ++r) {
-        Lfrcp[r] = 1.0 / Lf[r];
+        Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]);
    }

    [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {