diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 0a36f09900..99dbb0f14d 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1039,6 +1039,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( for (size_t s = 0; s < src_ss[i].n_segments; s++) { sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; } + const bool is_flash_attn_gqa_kv = + tensor->op == GGML_OP_FLASH_ATTN_EXT && (i == 1 || i == 2) && + tensor->src[i]->ne[src_ss[i].axis] != tensor->ne[split_state.axis]; + if (is_flash_attn_gqa_kv) { + // GQA/MQA splits Q/output over query heads, but K/V over fewer KV heads. + continue; + } GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis] == sum * tensor->ne[split_state.axis]); } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c528755339..63cb486ad9 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1003,9 +1003,17 @@ struct llama_model::impl { std::vector dev_layer; bool has_tensor_overrides; + + std::vector tensor_split_owned; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { + if (params.tensor_split != nullptr) { + // llama_model_params stores tensor_split as a borrowed pointer, but the model + // may need it later for tensor-parallel KV-cache split metadata. + pimpl->tensor_split_owned.assign(params.tensor_split, params.tensor_split + llama_max_devices()); + this->params.tensor_split = pimpl->tensor_split_owned.data(); + } pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern; }