Revert "do not use softmax fast mode in FusedSDPA (#26)"

This reverts commit 4e911b4.
HabanaAI · Nov 14, 2024 · 3b928de · 3b928de
1 parent 4e911b4
commit 3b928de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm_hpu_extension/ops.py b/vllm_hpu_extension/ops.py
@@ -223,7 +223,7 @@ def prompt_attention(
         if query_heads != kv_heads:
             key = repeat_kv(key, int(query_heads // kv_heads))
             value = repeat_kv(value, int(query_heads // kv_heads))
-        softmax_mode = 'None'
+        softmax_mode = 'fast'
         recompute_mode = True
         attn_weights = FusedSDPA.apply(query, key, value, None, 0.0, True,
                                        scale, softmax_mode, recompute_mode,