From 443994d5ada3403f0423d1f631cab1fe652e6374 Mon Sep 17 00:00:00 2001
From: Hongqiang Wang <wangh@qti.qualcomm.com>
Date: Thu, 30 Jul 2026 20:43:02 -0700
Subject: [PATCH] llama: re-create the KV cache when flash attention resolves
 to disabled

With -fa auto, the KV cache is created before flash attention support is
probed: cparams.flash_attn still holds the requested value, so V is laid
out for flash attention. The probe in resolve_fused_ops runs later,
during the first sched_reserve, and may turn flash attention off -
nothing updates the cache layout when it does.

Unfused attention needs a transposed V, so every graph then inserts a
cont(transpose(v)) per layer - a full copy of the V cache on each decode
step, in the branch llama-graph.cpp marks "note: avoid this branch".

Re-create the memory module once flash attention has resolved. This can
only trigger on the first reserve (resolve_fused_ops clears
cparams.auto_fa), when the cache is still empty, so no cached data is
moved or lost. The re-reserve rebuilds the worst-case graphs against the
corrected layout. params_mem is hoisted to the enclosing scope so the
re-creation can reuse it.

Qwen3-4B-Q4_0, tg32 at depth 4096, -fa auto resolved to disabled vs an
explicit -fa 0:

  Adreno 840     3.72 -> 10.20 t/s  (-69.0% -> parity)
  Adreno X2-90   8.01 -> 24.54 t/s  (-67.5% -> parity)

Output is byte-identical across -fa 0, -fa 1, auto-resolved-on and
auto-resolved-off, before and after the change.
---
 src/llama-context.cpp | 38 ++++++++++++++++++++++++++++++--------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 5ef7becf6f29..780d4ff2f9f4 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -378,15 +378,15 @@ llama_context::llama_context(
     }
 
     // init the memory module
-    if (!hparams.vocab_only) {
-        llama_memory_params params_mem = {
-            /*.type_k    =*/ params.type_k,
-            /*.type_v    =*/ params.type_v,
-            /*.swa_full  =*/ params.swa_full,
-            /*.ctx_type  =*/ cparams.ctx_type,
-            /*.mem_other =*/ llama_get_memory(cparams.ctx_other),
-        };
+    const llama_memory_params params_mem = {
+        /*.type_k    =*/ params.type_k,
+        /*.type_v    =*/ params.type_v,
+        /*.swa_full  =*/ params.swa_full,
+        /*.ctx_type  =*/ cparams.ctx_type,
+        /*.mem_other =*/ llama_get_memory(cparams.ctx_other),
+    };
 
+    if (!hparams.vocab_only) {
         memory.reset(model.create_memory(params_mem, cparams));
     }
 
@@ -453,8 +453,30 @@ llama_context::llama_context(
             LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__);
         }
 
+        const bool flash_attn_requested = cparams.flash_attn;
+
         sched_reserve();
 
+        // The memory module above was created with attn_v_trans = !flash_attn, but under
+        // LLAMA_FLASH_ATTN_TYPE_AUTO that read the *requested* value: the probe inside
+        // resolve_fused_ops (reached via sched_reserve) may only have turned flash attention
+        // off afterwards. When it does, V was laid out for flash attention while attention
+        // will actually run unfused, and the unfused path then needs a transposed V - so ggml
+        // inserts a cont(transpose(v)) per layer per token, which is expensive at depth.
+        //
+        // Re-create the memory module so the layout matches the resolved value. This can only
+        // trigger on the first reserve, since resolve_fused_ops clears cparams.auto_fa, and at
+        // this point the cache is still empty, so no cached data is moved or lost.
+        if (memory && flash_attn_requested && !cparams.flash_attn) {
+            LLAMA_LOG_INFO("%s: flash attention resolved to disabled - re-creating the memory "
+                    "module so the V cache is transposed for the unfused path\n", __func__);
+
+            memory.reset(model.create_memory(params_mem, cparams));
+
+            sched_need_reserve = true;
+            sched_reserve();
+        }
+
         if (!cparams.flash_attn) {
             if (ggml_is_quantized(params.type_v)) {
                 throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention");