vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-27 01:04:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[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 16:19:12 +05:30 • committed by GitHub
parent 8144f3192e
commit d3371929bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 113 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

29
ci/run.sh
View file

@ -300,6 +300,31 @@ function gg_sum_ctest_release {
gg_printf '```\n'
}
# test_llama_archs_tensor_split
function gg_run_test_llama_archs_tensor_split {
cd ${SRC}
set -e
GGML_CUDA_DEVICES=1 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1
set +e
}
function gg_sum_test_llama_archs_tensor_split {
gg_printf '### %s\n\n' "${ci}"
gg_printf 'Runs test-llama-archs with 1 to 4 CUDA devices\n'
gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
gg_printf '```\n'
gg_printf '%s\n' "$(cat $OUT/${ci}.log)"
gg_printf '```\n'
}
# test_scripts
function gg_run_test_scripts {
@ -751,6 +776,10 @@ ret=0
test $ret -eq 0 && gg_run ctest_debug
test $ret -eq 0 && gg_run ctest_release
if [ ! -z ${GG_BUILD_CUDA} ]; then
test $ret -eq 0 && gg_run test_llama_archs_tensor_split
fi
if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
test $ret -eq 0 && gg_run test_backend_ops_cpu
fi

68
ggml/src/ggml-backend-meta.cpp
View file

@ -602,27 +602,40 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
case GGML_BACKEND_SPLIT_AXIS_1:
case GGML_BACKEND_SPLIT_AXIS_2:
case GGML_BACKEND_SPLIT_AXIS_3: {
GGML_ASSERT(src_ss[0].n_segments == 1);
if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) {
return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1};
}
int64_t base_ne_in = tensor->src[0]->ne[0];
for (int dim = 1; dim <= src_ss[0].axis; dim++) {
int64_t base_ne_in = 1;
for (int dim = 0; dim <= src_ss[0].axis; dim++) {
base_ne_in *= tensor->src[0]->ne[dim];
}
base_ne_in /= src_ss[0].nr[0];
if (src_ss[0].n_segments == 1) {
base_ne_in /= src_ss[0].nr[0];
if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) {
return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1};
}
if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && tensor->ne[0] == tensor->src[0]->ne[0] &&
tensor->ne[1] == 1 && src_ss[0].nr[0] == 1) {
bool complete_rows = true;
for (size_t j = 0; j < n_bufs; j++) {
const int64_t ne = src_ss[0].ne[j];
complete_rows = complete_rows && (ne == 0 || ne == tensor->src[0]->ne[0]);
}
if (complete_rows) {
// Move a complete dim-0 split to the following singleton dimension.
return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1};
}
}
}
// Reshape outputs use one segment; split-state propagation merges source segments.
int64_t base_ne_out = 1;
for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim];
if (base_ne_out_next % base_ne_in == 0) {
return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1};
base_ne_out *= tensor->ne[dim];
if (base_ne_out % base_ne_in == 0) {
return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out/base_ne_in)}, 1};
}
if (base_ne_out_next > base_ne_in) {
if (base_ne_out > base_ne_in) {
GGML_ASSERT(src_ss[0].n_segments == 1);
GGML_ASSERT(src_ss[0].nr[0] == 1);
return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
}
base_ne_out = base_ne_out_next;
}
GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op));
}
@ -792,7 +805,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud);
if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) {
if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) {
const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1;
int64_t ne_sum = 0;
for (size_t s = 0; s < ret.n_segments; s++) {
@ -802,6 +815,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
}
}
GGML_ASSERT(ne_sum == tensor->ne[ret.axis]);
} else if (ret.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
GGML_ASSERT(ret.n_segments == 1);
GGML_ASSERT(ret.nr[0] == 1);
}
return ret;
}
@ -1352,15 +1368,29 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
} break;
case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
GGML_ASSERT(tensor->type == GGML_TYPE_F32);
const int64_t ne = ggml_nelements(tensor);
std::vector<float> tmp;
tmp.reserve(ne);
for (int64_t i = 0; i < ne; i++) {
tmp.push_back(((const float *) data)[i] / n_bufs);
GGML_ASSERT(offset % sizeof(float) == 0);
GGML_ASSERT(size % sizeof(float) == 0);
const size_t n_values = size / sizeof(float);
size_t n_contributors = 0;
for (size_t j = 0; j < n_bufs; j++) {
n_contributors += split_state.ne[j] != 0;
}
const bool has_contributor_mask = n_contributors != 0;
if (!has_contributor_mask) {
n_contributors = n_bufs;
}
std::vector<float> tmp(n_values);
for (size_t i = 0; i < n_values; i++) {
tmp[i] = ((const float *) data)[i] / n_contributors;
}
std::vector<float> zero;
if (has_contributor_mask) {
zero.resize(n_values, 0.0f);
}
for (size_t j = 0; j < n_bufs; j++) {
ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size);
const float * partial = has_contributor_mask && split_state.ne[j] == 0 ? zero.data() : tmp.data();
ggml_backend_tensor_set(simple_tensor, partial, offset, size);
}
} break;
default: {

28
src/llama-model.cpp
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);

12
tests/test-llama-archs.cpp
View file

@ -101,6 +101,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
n_head = 1;
n_ff = 96;
n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded
} else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) {
n_embd = 160; // exercise per-head tensor split granularity with head size 80
} else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
n_head = 4;
} else if (arch == LLM_ARCH_DEEPSEEK2
|| arch == LLM_ARCH_DEEPSEEK32
|| arch == LLM_ARCH_GLM_DSA
@ -120,6 +124,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
}
uint32_t n_head_kv = n_head;
if (arch == LLM_ARCH_QWEN3) {
n_head_kv = 1; // MQA coverage
} else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
n_head_kv = 2; // GQA coverage
}
const uint32_t n_embd_head = n_embd / n_head;
ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE, llm_arch_name(arch));
@ -160,7 +170,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
} else {
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head);
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_kv);
}
ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);