Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 3 commits into from
Sep 16, 2026
Merged

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

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3da324e
model: calculate split states for attn_qkv from n_head * n_embd_head_k
dfriehs Sep 15, 2026
b1a9fc9
model: handle fused full attention layers for qwen35/qwen35moe
dfriehs Sep 15, 2026
c2651dc
model: add TODO: [TAG_SPLIT_QGATE_QWEN]
dfriehs Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -597,8 +597,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) {
Comment on lines 602 to 603

@ggerganov ggerganov Sep 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
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) {
// TODO: clarify why this is necessary specifically for these models
// TODO: deduplicate [TAG_SPLIT_FUSED_QKV_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;
Expand Down Expand Up @@ -648,9 +660,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}};
}
Expand Down Expand Up @@ -736,6 +748,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)};
Expand Down Expand Up @@ -763,6 +776,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) {
Comment on lines +779 to +782

@ggerganov ggerganov Sep 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should deduplicate this condition if it has to match the one earlier:

Suggested change
// fused full attention layers need Q gate tensors handled like above:
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 need Q gate tensors handled like above:
// TODO: deduplicate with the condition in `get_split_segments` [TAG_SPLIT_FUSED_QKV_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};
}
}
Expand Down