Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Tensor parallel] Fix meta tensor split state propagation (#27574)

Browse Source 
* ggml : fix meta tensor split state propagation

* Add test-llama-archs to CI
This commit is contained in:
Gaurav Garg
2026-08-23 18:49:12 +08:00
committed by GitHub
parent 8144f3192e
commit d3371929bb
4 changed files with 113 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-model.cpp
+24 -4
View File
@@ -520,7 +520,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
}
if (std::regex_match(tensor_name, pattern_ffn_down_exps_bias)) {
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL);
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL, "ffn_down_exps.weight");
}
// output
@@ -554,6 +554,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
return {{key_dim, 2}, {value_dim, 1}};
}
if (std::regex_match(tensor_name, pattern_r_cache)) {
return {{key_dim * (hparams.ssm_d_conv - 1), 2}, {value_dim * (hparams.ssm_d_conv - 1), 1}};
}
} else {
const int64_t head_ratio = n_v_heads / n_k_heads;
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
@@ -642,12 +645,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
blck_size_perf *= 2;
}
const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf);
const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head
if (std::regex_match(tensor_name, pattern_attn_sinks)) {
GGML_ASSERT(segments.size() == 1);
return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa};
return {granularity_head};
}
const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf);
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:
@@ -660,6 +663,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
GGML_ASSERT(segments.size() == 1);
return {granularity_q};
}
if (std::regex_match(tensor_name, pattern_attn_gate_weight)) {
GGML_ASSERT(segments.size() == 1);
if (tensor->ne[1] == hparams.n_head(il)) {
return {granularity_head};
}
return {granularity_q};
}
const int64_t granularity_kv = granularity_q / n_gqa;
if (std::regex_match(tensor_name, pattern_kv_weight) ||
@@ -728,6 +738,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
memset(split_state.ne, 0, sizeof(split_state.ne));
split_state.nr[0] = 1;
split_state.n_segments = 1;
if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
GGML_ASSERT(tc.tensor_axis_0 != tensor);
const ggml_backend_meta_split_state source_split_state = llama_meta_device_get_split_state(tc.tensor_axis_0, userdata);
GGML_ASSERT(source_split_state.axis >= 0 && source_split_state.axis < GGML_MAX_DIMS);
for (size_t j = 0; j < ud->n_devices; j++) {
for (size_t is = 0; is < source_split_state.n_segments; is++) {
split_state.ne[j] += source_split_state.ne[is*ud->n_devices + j] * source_split_state.nr[is];
}
}
}
}
return split_state;
GGML_UNUSED(userdata);