diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp
index ba087d3ce..b0f5b487e 100644
--- a/src/core/ggml_extend.hpp
+++ b/src/core/ggml_extend.hpp
@@ -1362,6 +1362,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         }
 
         auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
         ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
         if (kv_scale != 1.0f) {
             out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
@@ -1379,9 +1382,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         if (can_use_flash_attn) {
             kqv = build_kqv(q, k, v, mask);
-            if (!ggml_backend_supports_op(backend, kqv)) {
-                kqv = nullptr;
-            } else {
+            if (kqv != nullptr) {
                 kqv = ggml_view_4d(ctx,
                                    kqv,
                                    d_head,