From fb27a525d28381a16a4bb038858a10e4927381ca Mon Sep 17 00:00:00 2001 From: David Friehs Date: Wed, 16 Sep 2026 21:02:12 +0200 Subject: [PATCH] TP: fix split state and granularity for fused QKV gemma4, qwen35 (#28965) * model: calculate split states for attn_qkv from n_head * n_embd_head_k required for gemma4 with --fuse-qkv, where n_embd is 5376 but Q is 8192. * model: handle fused full attention layers for qwen35/qwen35moe * model: add TODO: [TAG_SPLIT_QGATE_QWEN] --- src/llama-model.cpp | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3148f27816..3607bacd6e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -603,8 +603,20 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { + // TODO: clarify why this is necessary specifically for these models + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { + + // fused full attention layers with Q gate tensors that need n_embd doubled: + if (!hparams.is_recr(il) && (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias))) { + const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il) * 2; + const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); + GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); + GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); + return {{n_embd, 1}, {n_embd_gqa, 2}}; + } + const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; const int64_t n_k_heads = hparams.ssm_n_group; @@ -654,9 +666,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - const int64_t n_embd = hparams.n_embd; + const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); - GGML_ASSERT(hparams.n_embd_k_gqa() == n_embd_gqa); + GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); return {{n_embd, 1}, {n_embd_gqa, 2}}; } @@ -742,6 +754,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf)}; @@ -769,6 +782,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { GGML_ASSERT(segments.size() == 2); + // fused full attention layers need Q gate tensors handled like above: + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { + return {std::lcm(2*n_embd_q, blck_size_perf), granularity_kv}; + } return {granularity_q, granularity_kv}; } }