ClaudeBot/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-17 20:31:47 +02:00
Code Issues Packages Projects Releases Wiki Activity

TP: fix split state and granularity for fused QKV gemma4, qwen35 (#28965)

Browse Source 
* 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]
This commit is contained in:
David Friehs
2026-09-16 22:02:12 +03:00
committed by GitHub
parent c6824a9e42
commit fb27a525d2
1 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-model.cpp
+21 -2
View File
@@ -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<std::pair<int64_t, uint32_t>> {
// 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};
}
}