llama: re-create the KV cache when flash attention resolves to disabled (performance and layout bug)

[wanghqc]

Overview

This PR is to fix a performance bug with FA set to auto.

Context

• The V cache has two possible layouts, chosen at allocation by attn_v_trans = !flash_attn:
  • flash attention consumes V row-major;
  • while the unfused path computes KQV as a mul_mat that needs V stored transposed.

Problem

Under LLAMA_FLASH_ATTN_TYPE_AUTO — the default in both llama_context_default_params and the CLI — these two steps happen in the wrong order.

• The memory module is created while cparams.flash_attn still holds the requested value (true), so V gets the flash-attention layout.
• The probe in resolve_fused_ops runs later, during the first sched_reserve, and may then turn flash attention off.
• But nothing re-creates the memory in the case.
• From then on every graph runs unfused against a V cache laid out for flash attention.
• The graph builder detects the mismatch from the strides and falls into the branch llama-graph.cpp itself marks "note: avoid this branch": a cont(transpose(v)) per layer, per graph — a materialized copy of the full V cache on every decode step, growing with context depth.

Fix

• After the first reserve, if flash attention was requested (auto) but resolved to disabled, re-create the memory module with the resolved cparams and reserve again. The cache is then laid out for the path that will actually run.
• Why the re-creation is safe
  • sched_reserve() also runs during decode, where discarding the cache would lose live KV data. That cannot happen here: the condition is gated on cparams.auto_fa, which is set once in the constructor and cleared inside resolve_fused_ops, so it can only be true on the first reserve — inside the constructor, before any token has been processed.
  • The cache is empty at that point; nothing is moved or lost. The follow-up reserve rebuilds the worst-case graphs against the corrected layout.

Performance

Qwen3-4B-Q4_0, tg32 at depth 4096. Each row compares -fa auto (resolved to disabled) against an explicit -fa 0, before and after this change:

device	backend	before	after
Adreno 840	OpenCL	-69.0%	+0.06%
Adreno X2-90	OpenCL	-67.5%	-0.9%
RTX 4060 Ti	CUDA	-30.4%	-0.15%

The fix restores parity with -fa 0 in every case.

How the CUDA row was measured

CUDA supports flash attention for every configuration tested (12 of 12 resolve to enabled), so the auto-resolved-off path cannot be reached there naturally. To measure it, a small test harness (not part of this PR) was applied on top of the fix, adding two environment toggles:

• GGML_CUDA_FA_FORCE_DECLINE=1 — the CUDA backend declines FLASH_ATTN_EXT at the decode probe shape (ne[1] == 1), so -fa auto resolves to disabled through the normal resolve_fused_ops path, the same way a real backend decline occurs.
• LLAMA_FA_NO_RECREATE=1 — skips the re-creation added by this PR, so the before/after arms come from a single binary.

The CUDA row therefore isn't a claim that CUDA users hit this today — they don't. It shows the cost is incurred entirely in the fallback branch in llama-graph.cpp, independent of which backend declined: any backend that ever resolves flash attention off pays it.

Correctness

Model output is byte-identical across -fa 0, -fa 1, auto-resolved-on and auto-resolved-off, before and after the change. This is a layout and performance fix.

Additional information

Requirements

• I have read and agree with the contributing guidelines Yes
• AI usage disclosure: Yes, prototyping and testing.