From 6aa49b91b16256018f1369bdf628dcbc23829b30 Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Wed, 8 Apr 2026 22:22:07 +0800
Subject: [PATCH] fixed acestep bad on vulkan

---
 otherarch/acestep/dit.h | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/otherarch/acestep/dit.h b/otherarch/acestep/dit.h
index 7cd4a63b1..acbb4aeca 100644
--- a/otherarch/acestep/dit.h
+++ b/otherarch/acestep/dit.h
@@ -633,6 +633,10 @@ static struct ggml_tensor * dit_ggml_build_self_attn(
         ? ggml_flash_attn_ext(ctx, q, k, v, mask, scale, 0.0f, 0.0f)
         : dit_attn_f32(ctx, q, k, v, mask, scale);
 
+    if (m->use_flash_attn) {
+        ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
+    }
+
     // Both return [D, Nh, S, N]
     // Reshape: [D, Nh, S, N] -> [D*Nh, S, N] = [H, S, N]
     attn = ggml_reshape_3d(ctx, attn, Nh * D, S, N);
@@ -742,6 +746,10 @@ static struct ggml_tensor * dit_ggml_build_cross_attn(
         ? ggml_flash_attn_ext(ctx, q, k, v, NULL, scale, 0.0f, 0.0f)
         : dit_attn_f32(ctx, q, k, v, NULL, scale);
 
+    if (m->use_flash_attn) {
+        ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
+    }
+
     // Attention output: [D, Nh, S, N], reshape to [H, S, N]
     attn = ggml_reshape_3d(ctx, attn, Nh * D, S, N);