set coarser granularity for head splits

This commit is contained in:
Aman Gupta
2026-08-03 18:28:59 +02:00
parent 7b07e05c1e
commit 7c26c91500
+4 -2
View File
@@ -663,7 +663,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
if (std::regex_match(tensor_name, pattern_attn_sinks)) {
GGML_ASSERT(segments.size() == 1);
if (ud->model->arch == LLM_ARCH_DEEPSEEK4) {
return {1};
return {hparams.n_head(il) / hparams.dsv4_o_group_count};
}
return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa};
}
@@ -671,7 +671,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
if (ud->model->arch == LLM_ARCH_DEEPSEEK4) {
if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
GGML_ASSERT(segments.size() == 1);
return {hparams.n_embd_head_k(il)};
// the grouped output projection requires each device to hold whole groups of heads
const int64_t n_head_group = hparams.n_head(il) / hparams.dsv4_o_group_count;
return {n_head_group * hparams.n_embd_head_k(il)};
}
if (std::regex_match(tensor_name, pattern_attn_out_a_weight) ||
std::regex_match(tensor_name, pattern_attn_out_b_weight)) {