meta: fix tensor split metadata for GQA attention

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]);
                     }

src/llama-model.cpp

@@ -1003,9 +1003,17 @@ struct llama_model::impl {
     std::vector<layer_dev> dev_layer;
 
     bool has_tensor_overrides;
+
+    std::vector<float> tensor_split_owned;
 };
 
 llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique<impl>()) {
+    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;
 }