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");