From a34f2ac53a5fbb50ae056d93e8ad3a8c8b562dbe Mon Sep 17 00:00:00 2001 From: Jun Yamog Date: Thu, 16 Jul 2026 23:16:31 +1200 Subject: [PATCH 1/4] feat(meta): support selected-device tensor parallelism --- ggml/include/ggml-alloc.h | 1 + ggml/include/ggml-backend.h | 20 +- ggml/src/ggml-backend-meta.cpp | 1181 ++++++++++++++++++++----------- ggml/src/ggml-backend.cpp | 7 + ggml/src/ggml-cuda/common.cuh | 5 +- ggml/src/ggml-cuda/ggml-cuda.cu | 150 +++- ggml/src/ggml-impl.h | 6 + ggml/src/ggml.c | 12 + tests/CMakeLists.txt | 1 + tests/test-backend-meta.cpp | 299 ++++++++ 10 files changed, 1228 insertions(+), 454 deletions(-) create mode 100644 tests/test-backend-meta.cpp diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h index 78aa059dde38..a7926a21a9a2 100644 --- a/ggml/include/ggml-alloc.h +++ b/ggml/include/ggml-alloc.h @@ -76,6 +76,7 @@ GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_i // Utils // Create a buffer and allocate all the tensors in a ggml_context // ggml_backend_alloc_ctx_tensors_from_buft_size returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft +// ggml_backend_alloc_ctx_tensors_from_buft returns NULL on failure or if all tensors in ctx are already allocated or zero-sized GGML_API size_t ggml_backend_alloc_ctx_tensors_from_buft_size(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend); diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 4a8f6d4287da..48ebcafd4823 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -202,8 +202,11 @@ extern "C" { // Common functions that may be obtained using ggml_backend_reg_get_proc_address - // AllReduce operation for tensor parallelism (meta backend) - typedef bool (*ggml_backend_allreduce_tensor_t)(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends); + // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend) + typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends); + typedef void (*ggml_backend_comm_free_t)(void * comm_ctx); + typedef bool (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors); + // Split buffer type for tensor parallelism (old) typedef ggml_backend_buffer_type_t (*ggml_backend_split_buffer_type_t)(int main_device, const float * tensor_split); // Set the number of threads for the backend @@ -378,11 +381,15 @@ extern "C" { // - most tensors have n_segments == 1 and a contiguous slice of the tensor data // - some tensors have an inhomogenenous data layout along the split axis, // those tensors are divided into segments which are each individually split across devices - // - ne has one entry per segment and device that add up to ggml_tensor::ne for that axis, - // the outer/inner loops are over segments/devices like [seg0_dev0, seg0_dev1, seg1_dev0, seg1_dev1], + // - ne has one entry per segment and device and that segment repeats nr times, + // in total when accounting for repetitions the segments add up to ggml_tensor::ne for that axis, + // the outer/inner loops are over segments/devices like [seg0_dev0_r0, seg0_dev1_r0, seg0_dev0_r1, seg0_dev1_r1, seg1_dev0_r0, seg1_dev1_r0], // - for example, a transformer may have a fused QKV matrix rather than 3 matrices, those would be 3 separate segments - // that each need to be split individually across devices so that each device gets a slice of Q, K, and V + // that each need to be split individually across devices so that each device gets a slice of Q, K, and V, + // the Q matrix can be larger than the K and V matrices so this can either be expressed as 3 segments or as 2 segments + // where the segment for K/V repeats twice int64_t ne[16*GGML_BACKEND_META_MAX_DEVICES]; + uint32_t nr[16]; uint32_t n_segments; }; @@ -395,6 +402,9 @@ extern "C" { GGML_API ggml_backend_dev_t ggml_backend_meta_device( ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud); + // True when a buffer type is owned by a tensor-parallel meta device. + GGML_API bool ggml_backend_buft_is_meta(ggml_backend_buffer_type_t buft); + // // Utils // diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 8e4273212003..01b5df7c710c 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -392,64 +393,103 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( // meta backend buffer // +// Container to hold the tensor slices per simple ggml backend buffer. +struct ggml_backend_meta_simple_tensor_container { + std::vector ctxs; + std::map> simple_tensors; + + ggml_backend_meta_simple_tensor_container(const ggml_init_params & params, const int n_simple) { + ctxs.reserve(n_simple); + for (int i = 0; i < n_simple; i++) { + ctxs.emplace_back(ggml_init(params)); + } + } + ggml_backend_meta_simple_tensor_container() {} +}; + struct ggml_backend_meta_buffer_context { + // FIXME + // Most tensors can simply be stored statically in their own buffer. + // Externally created views however also need a mapping to simple tensors but they use the buffer of the view source. + // If external views are simply using that buffer they will slowly deplete its memory. + // Current solution: rotating set of 2 "compute" containers to hold external views, works correctly for llama.cpp. + // Long-term: tie the lifetime of external views to the meta backend executing the graph instead, + // currently not possible due to graph-external operations in the backend scheduler. + ggml_backend_meta_simple_tensor_container stc_static; + ggml_backend_meta_simple_tensor_container stc_compute[2]; + int stc_compute_index = 0; + int stc_compute_index_next = 0; + std::vector bufs; + + // FIXME + // The size of the split state cache is unbounded and can theoretically grow infinitely large. + // However, it is also expensive to build and clearing it on every rebuild in ggml_backend_meta_graph_compute is too expensive. static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); - std::map, std::pair> split_state_cache; - std::map< const ggml_tensor *, std::vector> simple_tensors; - - struct buffer_config { - ggml_context * ctx; - ggml_backend_buffer_t buf; - - buffer_config(ggml_context * ctx, ggml_backend_buffer_t buf) : ctx(ctx), buf(buf) {} - }; - std::vector buf_configs; int debug; - ggml_backend_meta_buffer_context() { + ggml_backend_meta_buffer_context( + ggml_backend_meta_simple_tensor_container & stc_static, + ggml_backend_meta_simple_tensor_container & stc_compute_0, + ggml_backend_meta_simple_tensor_container & stc_compute_1, + const std::vector & bufs) + : stc_static(std::move(stc_static)), stc_compute{std::move(stc_compute_0), std::move(stc_compute_1)} { + this->bufs.reserve(bufs.size()); + for (ggml_backend_buffer_t buf : bufs) { + this->bufs.emplace_back(buf); + } const char * GGML_META_DEBUG = getenv("GGML_META_DEBUG"); debug = GGML_META_DEBUG ? atoi(GGML_META_DEBUG) : 0; } + + ggml_backend_meta_simple_tensor_container & get_simple_tensor_container(const ggml_tensor * tensor) { + if (stc_static.simple_tensors.find(tensor) != stc_static.simple_tensors.end()) { + return stc_static; + } + return stc_compute[stc_compute_index]; + } }; static void ggml_backend_meta_buffer_free_buffer(ggml_backend_buffer_t buffer) { GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; - for (auto & [ctx, buf] : buf_ctx->buf_configs) { - ggml_backend_buffer_free(buf); - ggml_free(ctx); - } delete buf_ctx; } static size_t ggml_backend_meta_buffer_n_bufs(ggml_backend_buffer_t meta_buf) { GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; - return buf_ctx->buf_configs.size(); + return buf_ctx->bufs.size(); } static ggml_backend_buffer_t ggml_backend_meta_buffer_simple_buffer(ggml_backend_buffer_t meta_buf, size_t index) { GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; - GGML_ASSERT(index < buf_ctx->buf_configs.size()); - return buf_ctx->buf_configs[index].buf; + GGML_ASSERT(index < buf_ctx->bufs.size()); + return buf_ctx->bufs[index].get(); } static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct ggml_tensor * tensor, size_t index) { GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; - GGML_ASSERT(index < buf_ctx->buf_configs.size()); + GGML_ASSERT(index < buf_ctx->bufs.size()); - auto it = buf_ctx->simple_tensors.find(tensor); - if (it == buf_ctx->simple_tensors.end()) { + ggml_backend_meta_simple_tensor_container & stc = buf_ctx->get_simple_tensor_container(tensor); + auto it = stc.simple_tensors.find(tensor); + if (it == stc.simple_tensors.end()) { return nullptr; } return it->second[index]; } -static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { +static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); + +static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( + ggml_backend_meta_simple_tensor_container & stc, const struct ggml_tensor * tensor, bool assume_sync) { + // FIXME Currently this function preserves/erases the information in n_segments and nr in an inconsistent way. + // Since the operations in question are developed specifically for llama.cpp this currently does not manifest as a bug there. + // However, in a broader ggml context with arbitrary ggml graphs this can lead to unexpected results. const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; @@ -460,11 +500,11 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co for (size_t j = 0; j < n_bufs; j++) { int64_t sum_a = 0; for (size_t s = 0; s < a.n_segments; s++) { - sum_a += a.ne[s*n_bufs + j]; + sum_a += a.ne[s*n_bufs + j] * a.nr[s]; } int64_t sum_b = 0; for (size_t s = 0; s < b.n_segments; s++) { - sum_b += b.ne[s*n_bufs + j]; + sum_b += b.ne[s*n_bufs + j] * b.nr[s]; } if (sum_a != sum_b) { return false; @@ -474,7 +514,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co }; auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { - ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}; + ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}; for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { continue; @@ -482,15 +522,15 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { ret = src_ss[i]; } else if (!split_states_equal(src_ss[i], ret)) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; break; } } if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); return ret; @@ -498,6 +538,20 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co // Some ops process data on a per-row bases: auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { + std::fprintf(stderr, + "meta per-row split violation: op=%s tensor=%s shape=[%lld,%lld,%lld,%lld] " + "src=%s src_shape=[%lld,%lld,%lld,%lld] axis=%d\n", + ggml_op_name(tensor->op), tensor->name, + (long long) tensor->ne[0], (long long) tensor->ne[1], + (long long) tensor->ne[2], (long long) tensor->ne[3], + tensor->src[0] ? tensor->src[0]->name : "(null)", + tensor->src[0] ? (long long) tensor->src[0]->ne[0] : 0, + tensor->src[0] ? (long long) tensor->src[0]->ne[1] : 0, + tensor->src[0] ? (long long) tensor->src[0]->ne[2] : 0, + tensor->src[0] ? (long long) tensor->src[0]->ne[3] : 0, + (int) src_ss[0].axis); + } GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); return src_ss[0]; }; @@ -534,42 +588,24 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co auto handle_mul_mat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { ggml_backend_meta_split_state ret = src_ss[0]; ret.axis = GGML_BACKEND_SPLIT_AXIS_0; + ret.nr[0] = 1; ret.n_segments = 1; return ret; } if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - ggml_backend_meta_split_state ret = src_ss[1]; - ret.n_segments = 1; - return ret; + return src_ss[1]; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); - return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, 1}; + return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; } GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - }; - - auto handle_cpy = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - int64_t ne_split_src = tensor->src[0]->ne[0]; - for (int dim = 1; dim <= src_ss[0].axis; dim++) { - ne_split_src *= tensor->src[0]->ne[dim]; - } - int64_t ne_split_dst = 1; - for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - ne_split_dst *= tensor->ne[dim]; - if (ne_split_dst == ne_split_src) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; - } - } - } - return handle_generic(src_ss, /*scalar_only =*/ false); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; auto handle_reshape = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { @@ -578,33 +614,25 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - GGML_ASSERT(!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0])); - if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1) { - return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, 1}; + 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}; } - std::vector base_ne_in; - base_ne_in.reserve(GGML_MAX_DIMS - src_ss[0].axis); - { - base_ne_in.push_back(1); - int dim = 0; - for (; dim <= src_ss[0].axis; dim++) { - base_ne_in[0] *= tensor->src[0]->ne[dim]; - } - for (; dim <= GGML_MAX_DIMS; dim++) { - base_ne_in.push_back(base_ne_in.back() * tensor->src[0]->ne[dim]); - } + int64_t base_ne_in = tensor->src[0]->ne[0]; + for (int dim = 1; dim <= src_ss[0].axis; dim++) { + base_ne_in *= tensor->src[0]->ne[dim]; } + base_ne_in /= src_ss[0].nr[0]; 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]; - for (const int64_t & bni : base_ne_in) { - if (bni == base_ne_out_next) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; - } + 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}; } - if (base_ne_out_next > base_ne_in[0]) { - GGML_ASSERT(dim + 1 < GGML_MAX_DIMS); - return {ggml_backend_meta_split_axis(dim + 1), {0}, 1}; + if (base_ne_out_next > 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; } @@ -616,11 +644,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } } }; + auto handle_cpy = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { + return handle_reshape(src_ss); + } + return handle_generic(src_ss, /*scalar_only =*/ false); + }; + auto handle_view = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) { return handle_reshape(src_ss); @@ -644,7 +679,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co if (!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { for (int dim = 0; dim < GGML_MAX_DIMS-1; dim++) { if (tensor->nb[dim+1] == tensor->src[0]->nb[axis+1]) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; + return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; } } GGML_ABORT("fatal error"); @@ -653,7 +688,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co return src_ss[0]; } GGML_ABORT("view of permuted tensor not implemented"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; auto handle_permute = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { @@ -662,7 +697,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: { - return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, 1}; + GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); + return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { @@ -670,7 +706,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } } }; @@ -679,7 +715,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co switch (src_ss[0].axis) { case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: { - return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, 1}; + GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); + return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, {src_ss[0].nr[0]}, 1}; } case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: @@ -689,7 +726,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } } }; @@ -727,16 +764,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; }; auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; } } return handle_generic(src_ss, /*scalar_only =*/ false); @@ -744,8 +781,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co auto handle_gated_delta_net = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { return src_ss[0]; } GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); @@ -753,44 +790,55 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2); - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; + // state shape is (S_v*S_v*H, K, n_seqs); the heads dim is nested inside axis 0, + // so a head-aligned split on the input cache reshapes to axis 0 here (not axis 2). + GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0); + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; }; auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { if (ggml_nelements(tensor) == 0) { - return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { 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); + for (size_t s = 0; s < ret.n_segments; ++s) { + // nr predates multi-segment repetition in callback users. Keep + // zero-initialized legacy callbacks equivalent to one repeat. + if (ret.nr[s] == 0) { + ret.nr[s] = 1; + } + } 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 sj = 0; sj < ret.n_segments*n_bufs; sj++) { - GGML_ASSERT(ret.ne[sj] % granularity == 0); - ne_sum += ret.ne[sj]; + for (size_t s = 0; s < ret.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + GGML_ASSERT(ret.ne[s*n_bufs + j] % granularity == 0); + ne_sum += ret.ne[s*n_bufs + j] * ret.nr[s]; + } } GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); } return ret; } - std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}); + std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { - src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; continue; } - src_ss[i] = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); + src_ss[i] = ggml_backend_meta_get_split_state(stc, tensor->src[i], /*assume_sync =*/ true); GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); } ggml_backend_meta_split_state split_state; switch (tensor->op) { case GGML_OP_NONE: { - split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; + split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; } break; case GGML_OP_DUP: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); @@ -815,6 +863,11 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_OP_COS: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; + case GGML_OP_TURBO_WHT: { + // The transform is independent for each row and preserves the + // tensor shape, so the source partition remains valid. + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; case GGML_OP_SUM: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); } break; @@ -982,7 +1035,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co } break; default: { GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); - split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } break; } if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { @@ -1000,23 +1053,25 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co split_state.ne[s*n_bufs + j] = 0; } for (size_t s = 0; s < src_ss[i].n_segments; s++) { - split_state.ne[j] += src_ss[i].ne[s*n_bufs + j]; + split_state.ne[j] += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; } split_state.ne[j] *= tensor->ne[split_state.axis]; if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { - GGML_ASSERT(split_state.ne[j] % tensor->src[i]->ne[src_ss[i].axis] == 0); - split_state.ne[j] /= tensor->src[i]->ne[src_ss[i].axis]; + const int64_t div = tensor->src[i]->ne[src_ss[i].axis] * split_state.nr[0]; + GGML_ASSERT(split_state.ne[j] % div == 0); + split_state.ne[j] /= div; } } } else { + GGML_ASSERT(split_state.n_segments == 1); for (size_t j = 0; j < n_bufs; j++) { + // Assert that ratio is consistent: int64_t sum = 0; for (size_t s = 0; s < src_ss[i].n_segments; s++) { - sum += src_ss[i].ne[s*n_bufs + j]; + sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; } - // Assert that ratio is consistent: - GGML_ASSERT(split_state.ne[j] * tensor->src[i]->ne[src_ss[i].axis] - == sum * tensor->ne[split_state.axis]); + GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis] + == sum * tensor->ne[split_state.axis]); } } first_src_split_by_axis = false; @@ -1046,13 +1101,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co srcs_info += ", "; } const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); + GGML_ASSERT(split_state.n_segments == 1); const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); std::string ne_info; for (size_t j = 0; j < n_bufs; j++) { if (!ne_info.empty()) { ne_info += ", "; } - ne_info += std::to_string(split_state.ne[j]); + ne_info += std::to_string(split_state.ne[j]) + "x" + std::to_string(split_state.nr[0]); } srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; } @@ -1061,7 +1117,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co if (!ne_info.empty()) { ne_info += ", "; } - ne_info += std::to_string(buf_ctx->split_state_cache[key].first.ne[j]); + const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; + ne_info += std::to_string(ss.ne[j]) + "x" + std::to_string(ss.nr[0]); } GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); @@ -1073,8 +1130,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co #ifndef NDEBUG if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { int64_t ne_ret = 0; - for (size_t sj = 0; sj < ret.n_segments*n_bufs; sj++) { - ne_ret += ret.ne[sj]; + for (size_t s = 0; s < ret.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + ne_ret += ret.ne[s*n_bufs + j] * ret.nr[s]; + } } assert(ne_ret == tensor->ne[int(ret.axis)]); } @@ -1082,17 +1141,23 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co return ret; } +static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { + GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; + return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync); +} + static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { GGML_UNUSED(buffer); return (void *) 0x1000000000000000; // FIXME } -static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; - const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(buffer); +static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_meta_simple_tensor_container & stc, ggml_tensor * tensor) { + GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; + const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ true); + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(stc, tensor, /*assume_sync =*/ true); GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); GGML_ASSERT(split_state.n_segments <= 16); @@ -1107,15 +1172,15 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer std::vector simple_tensors; simple_tensors.reserve(n_simple_bufs); for (size_t j = 0; j < n_simple_bufs; j++) { - ggml_context * simple_ctx = buf_ctx->buf_configs[j].ctx; - ggml_backend_buffer_t simple_buf = buf_ctx->buf_configs[j].buf; + ggml_context * simple_ctx = stc.ctxs[j].get(); + ggml_backend_buffer_t simple_buf = buf_ctx->bufs[j].get(); if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { // TODO: the following assert fails for llama-parallel even though the results are correct: // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); ne[split_dim] = 0; for (size_t s = 0; s < split_state.n_segments; s++) { - ne[split_dim] += split_state.ne[s*n_simple_bufs + j]; + ne[split_dim] += split_state.ne[s*n_simple_bufs + j] * split_state.nr[s]; } for (int i = 0; i < GGML_MAX_DIMS; i++) { if (tensor->nb[i] > tensor->nb[split_dim]) { @@ -1138,7 +1203,7 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { - GGML_ASSERT(ne[split_dim] != 0 && tensor->ne[split_dim] != 0); + GGML_ASSERT(tensor->ne[split_dim] != 0); const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); @@ -1161,7 +1226,7 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer t_ij->data = (char *) t_ij->view_src->data + t_ij->view_offs; } else if (simple_buf != nullptr) { t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) - + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(buffer)); + + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer)); } t_ij->extra = tensor->extra; for (int i = 0; i < GGML_MAX_SRC; i++) { @@ -1175,53 +1240,110 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer simple_tensors.push_back(t_ij); } - buf_ctx->simple_tensors[tensor] = simple_tensors; + + // If one of the sources has a zero-sized slice, disable the computation: + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || !ggml_backend_buffer_is_meta(tensor->src[i]->buffer)) { + continue; + } + + const ggml_backend_meta_split_state split_state_src = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); + if (split_state_src.axis < 0 || split_state_src.axis >= GGML_MAX_DIMS) { + continue; + } + for (size_t j = 0; j < n_simple_bufs; j++) { + int64_t ne_sum = 0; + for (size_t s = 0; s < split_state_src.n_segments; s++) { + ne_sum += split_state_src.ne[s*n_simple_bufs + j] * split_state_src.nr[s]; + } + if (ne_sum == 0) { + simple_tensors[j]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + } + } + } + + stc.simple_tensors[tensor] = simple_tensors; return GGML_STATUS_SUCCESS; } +static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; + buf_ctx->stc_compute_index = buf_ctx->stc_compute_index_next; + return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor); +} + static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + if (size == 0) return; const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); GGML_ASSERT(ggml_is_contiguous(tensor)); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - if (split_state.n_segments != 1) { + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - GGML_ASSERT(tensor->ne[3] == 1); + GGML_ASSERT(split_state.nr[0] != 0); size_t offset_data = 0; std::vector simple_offsets(n_bufs, 0); if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[1]); + const int64_t blck_size = ggml_blck_size(tensor->type); for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); - const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, simple_offsets[j], nbytes, - tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); - offset_data += nbytes; - simple_offsets[j] += nbytes; + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, + simple_offsets[j] + row_start * simple_tensor->nb[1], nbytes, + row_count, simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } } } - GGML_ASSERT(offset_data*tensor->ne[1] == size); + GGML_ASSERT(offset_data*row_count == size); return; } - GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1 || + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2); + const int axis = split_state.axis; + const size_t row_stride = tensor->nb[axis + 1]; + if (offset % row_stride != 0 || size % row_stride != 0) { + GGML_LOG_ERROR("meta set_tensor unaligned: tensor=%s axis=%d offset=%zu size=%zu stride=%zu ne=[%lld,%lld,%lld,%lld]\n", + tensor->name, axis, offset, size, row_stride, + (long long) tensor->ne[0], (long long) tensor->ne[1], + (long long) tensor->ne[2], (long long) tensor->ne[3]); + } + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[axis + 1]); + for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, simple_offsets[j], nbytes, - tensor->ne[2], simple_tensor->nb[2], tensor->nb[2]); - offset_data += nbytes; - simple_offsets[j] += nbytes; + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[axis]; + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, + simple_offsets[j] + row_start * simple_tensor->nb[axis + 1], nbytes, + row_count, simple_tensor->nb[axis + 1], tensor->nb[axis + 1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } } } - GGML_ASSERT(offset_data*tensor->ne[2] == size); + GGML_ASSERT(offset_data*row_count == size); return; } @@ -1229,21 +1351,32 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - const size_t simple_offset = i_start * chunk_size_j; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; + const size_t request_end = offset + size; + const size_t outer_first = offset / chunk_size_full; + const size_t outer_last = (request_end - 1) / chunk_size_full; + for (size_t outer = outer_first; outer <= outer_last; ++outer) { + const size_t logical_chunk = outer * chunk_size_full; + size_t rank_begin = logical_chunk; + for (size_t j = 0; j < n_bufs; ++j) { + ggml_tensor * simple_tensor = + ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = + simple_tensor->nb[split_state.axis + 1]; + const size_t rank_end = rank_begin + chunk_size_j; + const size_t copy_begin = std::max(offset, rank_begin); + const size_t copy_end = std::min(request_end, rank_end); + if (copy_begin < copy_end) { + ggml_backend_tensor_set( + simple_tensor, + (const char *) data + copy_begin - offset, + outer * chunk_size_j + copy_begin - rank_begin, + copy_end - copy_begin); + } + rank_begin = rank_end; + } + GGML_ASSERT(rank_begin == logical_chunk + chunk_size_full); } - GGML_ASSERT(offset_j == chunk_size_full); } break; case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { for (size_t j = 0; j < n_bufs; j++) { @@ -1271,31 +1404,108 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg } static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + if (size == 0) return; const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); GGML_ASSERT(ggml_is_contiguous(tensor)); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - GGML_ASSERT(split_state.n_segments == 1); + + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { + GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); + GGML_ASSERT(split_state.nr[0] != 0); + size_t offset_data = 0; + std::vector simple_offsets(n_bufs, 0); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[1]); + + const int64_t blck_size = ggml_blck_size(tensor->type); + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, + simple_offsets[j] + row_start * simple_tensor->nb[1], nbytes, + row_count, simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + } + GGML_ASSERT(offset_data*row_count == size); + return; + } + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1 || + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2); + const int axis = split_state.axis; + const size_t row_stride = tensor->nb[axis + 1]; + if (offset % row_stride != 0 || size % row_stride != 0) { + GGML_LOG_ERROR("meta get_tensor unaligned: tensor=%s axis=%d offset=%zu size=%zu stride=%zu ne=[%lld,%lld,%lld,%lld]\n", + tensor->name, axis, offset, size, row_stride, + (long long) tensor->ne[0], (long long) tensor->ne[1], + (long long) tensor->ne[2], (long long) tensor->ne[3]); + } + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[axis + 1]); + + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[axis]; + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, + simple_offsets[j] + row_start * simple_tensor->nb[axis + 1], nbytes, + row_count, simple_tensor->nb[axis + 1], tensor->nb[axis + 1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + } + GGML_ASSERT(offset_data*row_count == size); + return; + } switch (split_state.axis) { case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_bufs; j++){ - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - const size_t simple_offset = i_start * chunk_size_j; - ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; + const size_t request_end = offset + size; + const size_t outer_first = offset / chunk_size_full; + const size_t outer_last = (request_end - 1) / chunk_size_full; + for (size_t outer = outer_first; outer <= outer_last; ++outer) { + const size_t logical_chunk = outer * chunk_size_full; + size_t rank_begin = logical_chunk; + for (size_t j = 0; j < n_bufs; ++j) { + const ggml_tensor * simple_tensor = + ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = + simple_tensor->nb[split_state.axis + 1]; + const size_t rank_end = rank_begin + chunk_size_j; + const size_t copy_begin = std::max(offset, rank_begin); + const size_t copy_end = std::min(request_end, rank_end); + if (copy_begin < copy_end) { + ggml_backend_tensor_get( + simple_tensor, + (char *) data + copy_begin - offset, + outer * chunk_size_j + copy_begin - rank_begin, + copy_end - copy_begin); + } + rank_begin = rank_end; + } + GGML_ASSERT(rank_begin == logical_chunk + chunk_size_full); } - GGML_ASSERT(offset_j == chunk_size_full); } break; case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { // TODO other simple backend may be better @@ -1316,8 +1526,9 @@ static void ggml_backend_meta_buffer_clear(ggml_backend_buffer_t buffer, uint8_t } static void ggml_backend_meta_buffer_reset(ggml_backend_buffer_t buffer) { - const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); - for (size_t i = 0; i < n_buffers; i++) { + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; + for (size_t i = 0; i < buf_ctx->bufs.size(); i++) { ggml_backend_buffer_reset(ggml_backend_meta_buffer_simple_buffer(buffer, i)); } } @@ -1343,20 +1554,24 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) { static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - ggml_init_params params = { - /*.mem_size =*/ 1024*1024*1024, // FIXME + const ggml_init_params params = { + /*.mem_size =*/ 1024*1024*ggml_tensor_overhead(), // FIXME /*.mem_buffer =*/ nullptr, /*.no_alloc =*/ true, }; + ggml_backend_meta_simple_tensor_container stc_static; + ggml_backend_meta_simple_tensor_container stc_compute_0(params, n_simple_bufts); + ggml_backend_meta_simple_tensor_container stc_compute_1(params, n_simple_bufts); - ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(); size_t max_size = 0; - buf_ctx->buf_configs.reserve(n_simple_bufts); + std::vector bufs; + bufs.reserve(n_simple_bufts); for (size_t i = 0; i < n_simple_bufts; i++) { - ggml_backend_buffer_t simple_buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size); - max_size = std::max(max_size, ggml_backend_buffer_get_size(simple_buf)); - buf_ctx->buf_configs.emplace_back(ggml_init(params), simple_buf); + bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size)); + GGML_ASSERT(bufs.back() != nullptr); + max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back())); } + ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size); } @@ -1364,28 +1579,53 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - ggml_init_params params = { - /*.mem_size =*/ 1024*1024*1024, // FIXME + constexpr size_t compute_headroom = 16; // Maximum number of views per statically allocated tensor that can be created between evals. + const ggml_init_params params_static = { + /*.mem_size =*/ ggml_get_mem_size(ctx), /*.mem_buffer =*/ nullptr, /*.no_alloc =*/ true, }; + const ggml_init_params params_compute = { + /*.mem_size =*/ compute_headroom*ggml_get_mem_size(ctx), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_backend_meta_simple_tensor_container stc_static (params_static, n_simple_bufts); + ggml_backend_meta_simple_tensor_container stc_compute_0(params_compute, n_simple_bufts); + ggml_backend_meta_simple_tensor_container stc_compute_1(params_compute, n_simple_bufts); - ggml_backend_meta_buffer_context * meta_buf_ctx = new ggml_backend_meta_buffer_context(); - meta_buf_ctx->buf_configs.reserve(n_simple_bufts); - for (size_t i = 0; i < n_simple_bufts; i++) { - meta_buf_ctx->buf_configs.emplace_back(ggml_init(params), nullptr); - } + std::vector bufs(n_simple_bufts, nullptr); + ggml_backend_meta_buffer_context * meta_buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); ggml_backend_buffer_t meta_buf = ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, meta_buf_ctx, 0); for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { t->buffer = meta_buf; - ggml_backend_meta_buffer_init_tensor(meta_buf, t); + ggml_backend_meta_buffer_init_tensor_impl(meta_buf_ctx->stc_static, t); t->data = (void *) 0x2000000000000000; // FIXME } for (size_t i = 0; i < n_simple_bufts; i++) { - meta_buf_ctx->buf_configs[i].buf = ggml_backend_alloc_ctx_tensors_from_buft( - meta_buf_ctx->buf_configs[i].ctx, ggml_backend_meta_buft_simple_buft(buft, i)); - meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->buf_configs[i].buf)); + ggml_context * ctx = meta_buf_ctx->stc_static.ctxs[i].get(); + ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, i); + + // If a ggml_context only has zero-sized tensors, ggml_backend_alloc_ctx_tensors_from_buft returns NULL. + // For those edge cases, allocate a dummy buffer instead. + bool any_nonzero_slice = false; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (ggml_nelements(t) != 0) { + any_nonzero_slice = true; + break; + } + } + if (any_nonzero_slice) { + meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, simple_buft)); + } else { + meta_buf_ctx->bufs[i].reset(ggml_backend_buft_alloc_buffer(simple_buft, 0)); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = meta_buf_ctx->bufs[i].get(); + } + } + GGML_ASSERT(meta_buf_ctx->bufs[i]); + meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get())); } return meta_buf; } @@ -1409,45 +1649,73 @@ struct ggml_backend_meta_context { struct backend_config { ggml_backend_t backend; - std::vector cgraphs; - std::vector nodes; - ggml_backend_buffer_ptr buf; + std::vector cgraphs; + std::vector nodes; + std::vector bufs; - backend_config(ggml_backend_t backend) : backend(backend) {} + backend_config(ggml_backend_t backend, const size_t n_reduce_steps) : backend(backend) { + bufs.resize(n_reduce_steps); + } }; std::string name; std::vector backend_configs; ggml_context_ptr ctx; std::vector cgraphs_aux; std::vector nodes_aux; + size_t n_reduce_steps; int max_nnodes = 0; size_t max_tmp_size = 0; size_t max_subgraphs = 0; + size_t n_subgraphs = 0; + uint64_t uid = 0; + + void * comm_ctx = nullptr; + ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); + n_reduce_steps = std::ceil(std::log2(n_devs)); name = "Meta("; + std::vector simple_backends; backend_configs.reserve(n_devs); + simple_backends.reserve(n_devs); for (size_t i = 0; i < n_devs; i++) { ggml_backend_dev_t simple_dev = ggml_backend_meta_dev_simple_dev(meta_dev, i); if (i > 0) { name += ","; } name += ggml_backend_dev_name(simple_dev); - backend_configs.emplace_back(ggml_backend_dev_init(simple_dev, params)); + simple_backends.push_back(ggml_backend_dev_init(simple_dev, params)); + backend_configs.emplace_back(simple_backends.back(), n_reduce_steps); } name += ")"; + + if (n_devs > 1) { + ggml_backend_comm_init_t comm_init = (ggml_backend_comm_init_t) ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_init"); + if (comm_init != nullptr) { + comm_ctx = comm_init(simple_backends.data(), simple_backends.size()); + } + } + if (comm_ctx != nullptr) { + comm_allreduce = (ggml_backend_comm_allreduce_tensor_t) + ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg( + ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_allreduce_tensor"); + GGML_ASSERT(comm_allreduce != nullptr); + } } ~ggml_backend_meta_context() { + if (comm_ctx != nullptr) { + ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free"); + GGML_ASSERT(comm_free != nullptr); + comm_free(comm_ctx); + } for (auto & bc : backend_configs) { ggml_backend_free(bc.backend); } } - - size_t n_reduce_steps() const { - return std::ceil(std::log2(backend_configs.size())); - } }; static const char * ggml_backend_meta_get_name(ggml_backend_t backend) { @@ -1470,6 +1738,7 @@ static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tens const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); GGML_ASSERT(split_state.n_segments == 1); + GGML_ASSERT(split_state.nr[0] == 1); switch (split_state.axis) { case GGML_BACKEND_SPLIT_AXIS_0: @@ -1486,6 +1755,9 @@ static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tens ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + if (chunk_size_j == 0) { + continue; + } ggml_backend_tensor_set_2d_async(simple_backend, simple_tensor, (const char *) data + offset_j, offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); offset_j += chunk_size_j; @@ -1511,6 +1783,7 @@ static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggm const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); GGML_ASSERT(split_state.n_segments == 1); + GGML_ASSERT(split_state.nr[0] == 1); switch (split_state.axis) { case GGML_BACKEND_SPLIT_AXIS_0: @@ -1527,6 +1800,9 @@ static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggm ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + if (chunk_size_j == 0) { + continue; + } ggml_backend_tensor_get_2d_async(simple_backend, simple_tensor, (char *) data + offset_j, offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); offset_j += chunk_size_j; @@ -1557,6 +1833,9 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, const size_t n_backends = ggml_backend_meta_n_backends(backend); ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; + // If the previous cgraph had a defined UID it can be used to skip rebuilding the subgraphs per simple backend. + const bool needs_rebuild = (cgraph->uid == 0) || (cgraph->uid != backend_ctx->uid); + bool max_nnodes_raised = false; if (cgraph->n_nodes > backend_ctx->max_nnodes) { for (size_t j = 0; j < n_backends; j++) { @@ -1566,173 +1845,219 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, } backend_ctx->max_nnodes = cgraph->n_nodes; max_nnodes_raised = true; + assert(needs_rebuild); } - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { - // FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes. - // For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash. - bcj.nodes[i] = node; - continue; + if (needs_rebuild) { + std::set used_buffers; + for (int i = 0; i < cgraph->n_leafs; i++) { + if (ggml_backend_buffer_is_meta(cgraph->leafs[i]->buffer)) { + used_buffers.emplace(cgraph->leafs[i]->buffer); } - bcj.nodes[i] = ggml_backend_meta_buffer_simple_tensor(node, j); - GGML_ASSERT(bcj.nodes[i]); } - } - - size_t n_subgraphs = 0; - size_t max_tmp_size = 0; - { - // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: - auto get_i_delayed = [&](const int i) -> int { - int id = i; // i_delayed - int idr = i; // i_delayed return, last safe return value - - ggml_tensor * node = cgraph->nodes[id]; - int32_t n_used = ggml_node_get_use_count(cgraph, id); - if (id + 1 >= cgraph->n_nodes) { - return idr; - } - { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op == GGML_OP_ADD_ID && next->src[0] == node && - ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL && - ggml_backend_meta_get_split_state(next->src[2], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - node = next; - id++; - idr = id; - n_used = ggml_node_get_use_count(cgraph, id); - } + for (int i = 0; i < cgraph->n_nodes; i++) { + if (ggml_backend_buffer_is_meta(cgraph->nodes[i]->buffer)) { + used_buffers.emplace(cgraph->nodes[i]->buffer); } - if (id + 1 >= cgraph->n_nodes) { - return idr; + } + for (ggml_backend_buffer_t buf : used_buffers) { + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buf->context; + buf_ctx->stc_compute_index_next = buf_ctx->stc_compute_index ^ 1; + ggml_backend_meta_simple_tensor_container & stc = buf_ctx->stc_compute[buf_ctx->stc_compute_index_next]; + for (ggml_context_ptr & ctx : stc.ctxs) { + ggml_reset(ctx.get()); } - { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op == GGML_OP_MUL && next->src[0] == node && - ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - node = next; - id++; - idr = id; - n_used = ggml_node_get_use_count(cgraph, id); + stc.simple_tensors.clear(); + } + size_t n_subgraphs = 0; + size_t max_tmp_size = 0; + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { + // FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes. + // For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash. + bcj.nodes[i] = node; + continue; } + bcj.nodes[i] = ggml_backend_meta_buffer_simple_tensor(node, j); + GGML_ASSERT(bcj.nodes[i]); } + } - if (n_used != node->ne[1] || id + 2*n_used-1 >= cgraph->n_nodes) { - return idr; - } - for (int32_t k = 0; k < n_used; k++) { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_VIEW || next->view_src != node || next->view_offs != k*node->nb[1] || - next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || next->nb[1] != node->nb[2] || - ggml_node_get_use_count(cgraph, id+1) != 1) { + { + // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: + auto get_i_delayed = [&](const int i) -> int { + int id = i; // i_delayed + int idr = i; // i_delayed return, last safe return value + + ggml_tensor * node = cgraph->nodes[id]; + int32_t n_used = ggml_node_get_use_count(cgraph, id); + if (id + 1 >= cgraph->n_nodes) { return idr; } - id++; - } - { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id - (n_used-1)] || - next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op == GGML_OP_ADD_ID && next->src[0] == node && + ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL && + ggml_backend_meta_get_split_state(next->src[2], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + node = next; + id++; + idr = id; + n_used = ggml_node_get_use_count(cgraph, id); + } + } + if (id + 1 >= cgraph->n_nodes) { return idr; } - id++; - } - for (int32_t k = 0; k < n_used - 2; k++) { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id] || - next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op == GGML_OP_MUL && next->src[0] == node && + ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + node = next; + id++; + idr = id; + n_used = ggml_node_get_use_count(cgraph, id); + } + } + + if (n_used != node->ne[1] || id + 2*n_used-1 >= cgraph->n_nodes) { return idr; } - id++; - } - idr = id; - return idr; - }; + for (int32_t k = 0; k < n_used; k++) { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_VIEW || next->view_src != node || next->view_offs != k*node->nb[1] || + next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || next->nb[1] != node->nb[2] || + ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id - (n_used-1)] || + next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + for (int32_t k = 0; k < n_used - 2; k++) { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id] || + next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + idr = id; + return idr; + }; - int i_start = 0; - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { - continue; - } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); - if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { - max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); - } - const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; - if (!new_subgraph) { - continue; + int i_start = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { + continue; + } + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); + } + const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; + if (!new_subgraph) { + continue; + } + + const int i_delayed = get_i_delayed(i); + + // If we can delay the AllReduce we need to consider the interaction with zero-sized tensor slices. + // A backend with such a slice would normally have valid data after participating in the AllReduce with a node that has + // its compute flag disabled and thus gets its data zeroed out. + // If the AllReduce is delayed then the nodes until that point also need to have their compute flag disabled. + if (i_delayed > i) { + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + if ((bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + for (int ii = i + 1; ii <= i_delayed; ii++) { + bcj.nodes[ii]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + } + } + } + } + + i = i_delayed; + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + bcj.cgraphs[n_subgraphs].offset = i_start; + } + n_subgraphs++; + i_start = i + 1; } + GGML_ASSERT(i_start == cgraph->n_nodes); + } - i = get_i_delayed(i); + backend_ctx->uid = cgraph->uid; + backend_ctx->n_subgraphs = n_subgraphs; + if (max_tmp_size > backend_ctx->max_tmp_size) { for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; - bcj.cgraphs[n_subgraphs].offset = i_start; + for (size_t i = 0; i < backend_ctx->n_reduce_steps; i++) { + bcj.bufs[i].reset(ggml_backend_alloc_buffer(bcj.backend, max_tmp_size)); + } + } + backend_ctx->max_tmp_size = max_tmp_size; + } + + if (max_nnodes_raised || n_subgraphs > backend_ctx->max_subgraphs) { + backend_ctx->max_subgraphs = std::max(backend_ctx->max_subgraphs, n_subgraphs); + const size_t n_nodes_per_device = 3 * backend_ctx->n_reduce_steps; // tmp + ADD (+zeroing) graph per step and device + const size_t n_cgraphs_per_device = 2 * backend_ctx->n_reduce_steps; // ADD ( + zeroing) graph per step and device + const size_t mem_per_device_graphs_main = backend_ctx->max_subgraphs*ggml_graph_overhead_custom(backend_ctx->max_nnodes, cgraph->grads); + const size_t mem_per_device_graphs_aux = n_cgraphs_per_device*backend_ctx->max_subgraphs*ggml_graph_overhead_custom(1, cgraph->grads); + const size_t mem_per_device_nodes_aux = n_nodes_per_device*backend_ctx->max_subgraphs*ggml_tensor_overhead(); + const ggml_init_params params = { + /*.mem_size =*/ n_backends * (mem_per_device_graphs_main + mem_per_device_graphs_aux + mem_per_device_nodes_aux), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + backend_ctx->ctx.reset(ggml_init(params)); + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + for (size_t i = 0; i < n_subgraphs; i++) { + bcj.cgraphs[i].cgraph_main = ggml_new_graph_custom(backend_ctx->ctx.get(), cgraph->n_nodes, /*grads =*/ false); + } + } + backend_ctx->cgraphs_aux.resize(n_backends*n_cgraphs_per_device*backend_ctx->max_subgraphs); + for (size_t k = 0; k < backend_ctx->cgraphs_aux.size(); k++) { + backend_ctx->cgraphs_aux[k] = ggml_new_graph_custom(backend_ctx->ctx.get(), 1, cgraph->grads); + } + backend_ctx->nodes_aux.resize(n_backends*n_nodes_per_device*backend_ctx->max_subgraphs); + for (size_t k = 0; k < backend_ctx->nodes_aux.size(); k++) { + backend_ctx->nodes_aux[k] = ggml_new_tensor_1d(backend_ctx->ctx.get(), GGML_TYPE_F32, 1); } - n_subgraphs++; - i_start = i + 1; - } - GGML_ASSERT(i_start == cgraph->n_nodes); - } - - if (max_tmp_size > backend_ctx->max_tmp_size) { - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - bcj.buf.reset(ggml_backend_alloc_buffer(bcj.backend, max_tmp_size)); } - backend_ctx->max_tmp_size = max_tmp_size; - } - - if (max_nnodes_raised || n_subgraphs > backend_ctx->max_subgraphs) { - backend_ctx->max_subgraphs = std::max(backend_ctx->max_subgraphs, n_subgraphs); - const size_t n_reduce_steps = backend_ctx->n_reduce_steps(); - const size_t n_nodes_per_device = 2 * n_reduce_steps; // tmp + ADD per step - const size_t n_cgraphs_per_device = n_reduce_steps; // 1 ADD graph per step - const size_t mem_per_device_graphs_main = backend_ctx->max_subgraphs*ggml_graph_overhead_custom(backend_ctx->max_nnodes, cgraph->grads); - const size_t mem_per_device_graphs_aux = n_cgraphs_per_device*backend_ctx->max_subgraphs*ggml_graph_overhead_custom(1, cgraph->grads); - const size_t mem_per_device_nodes_aux = n_nodes_per_device*backend_ctx->max_subgraphs*ggml_tensor_overhead(); - ggml_init_params params = { - /*.mem_size =*/ n_backends * (mem_per_device_graphs_main + mem_per_device_graphs_aux + mem_per_device_nodes_aux), - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - backend_ctx->ctx.reset(ggml_init(params)); for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; - for (size_t i = 0; i < n_subgraphs; i++) { - bcj.cgraphs[i].cgraph_main = ggml_new_graph_custom(backend_ctx->ctx.get(), cgraph->n_nodes, /*grads =*/ false); - } - } - backend_ctx->cgraphs_aux.resize(n_backends*n_cgraphs_per_device*backend_ctx->max_subgraphs); - for (size_t k = 0; k < backend_ctx->cgraphs_aux.size(); k++) { - backend_ctx->cgraphs_aux[k] = ggml_new_graph_custom(backend_ctx->ctx.get(), 1, cgraph->grads); - } - backend_ctx->nodes_aux.resize(n_backends*n_nodes_per_device*backend_ctx->max_subgraphs); - for (size_t k = 0; k < backend_ctx->nodes_aux.size(); k++) { - backend_ctx->nodes_aux[k] = ggml_new_tensor_1d(backend_ctx->ctx.get(), GGML_TYPE_F32, 1); - } - } - - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { - ggml_cgraph * cgraph_ij = bcj.cgraphs[i_graph].cgraph_main; - const size_t i_node_start = bcj.cgraphs[i_graph].offset; - const size_t i_node_stop = i_graph + 1 < n_subgraphs ? bcj.cgraphs[i_graph + 1].offset : cgraph->n_nodes; - cgraph_ij->n_nodes = i_node_stop - i_node_start; - ggml_hash_set_reset(&cgraph_ij->visited_hash_set); - for (size_t i_node = i_node_start; i_node < i_node_stop; i_node++) { - ggml_tensor * node_ij = bcj.nodes[i_node]; - cgraph_ij->nodes[i_node - i_node_start] = node_ij; - const size_t hash_pos_orig = ggml_hash_find(&cgraph->visited_hash_set, cgraph->nodes[i_node]); - const size_t hash_pos_ij = ggml_hash_insert(&cgraph_ij->visited_hash_set, node_ij); - cgraph_ij->use_counts[hash_pos_ij] = cgraph->use_counts[hash_pos_orig]; + for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { + ggml_cgraph * cgraph_ij = bcj.cgraphs[i_graph].cgraph_main; + const size_t i_node_start = bcj.cgraphs[i_graph].offset; + const size_t i_node_stop = i_graph + 1 < n_subgraphs ? bcj.cgraphs[i_graph + 1].offset : cgraph->n_nodes; + cgraph_ij->n_nodes = i_node_stop - i_node_start; + ggml_hash_set_reset(&cgraph_ij->visited_hash_set); + for (size_t i_node = i_node_start; i_node < i_node_stop; i_node++) { + ggml_tensor * node_ij = bcj.nodes[i_node]; + cgraph_ij->nodes[i_node - i_node_start] = node_ij; + const size_t hash_pos_orig = ggml_hash_find(&cgraph->visited_hash_set, cgraph->nodes[i_node]); + const size_t hash_pos_ij = ggml_hash_insert(&cgraph_ij->visited_hash_set, node_ij); + cgraph_ij->use_counts[hash_pos_ij] = cgraph->use_counts[hash_pos_orig]; + } + cgraph_ij->uid = ggml_graph_next_uid(); } } } @@ -1740,11 +2065,6 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, size_t iga = 0; // i graph aux size_t ina = 0; // i node aux - // FIXME usage_counts - auto get_cgraph_aux = [&]() -> ggml_cgraph * { - ggml_cgraph * ret = backend_ctx->cgraphs_aux[iga++]; - return ret; - }; auto get_node_aux = [&](ggml_tensor * t) -> ggml_tensor * { ggml_tensor * ret = backend_ctx->nodes_aux[ina++]; memset(ret, 0, sizeof(ggml_tensor)); @@ -1756,75 +2076,111 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, } return ret; }; + auto set_tmp_data = [&](ggml_tensor * tensor, const size_t j, const size_t i_buf) { + auto & bcj = backend_ctx->backend_configs[j]; + ggml_backend_buffer_ptr & buf_ptr = bcj.bufs[i_buf]; + if (!buf_ptr || ggml_backend_buffer_get_size(buf_ptr.get()) < backend_ctx->max_tmp_size) { + buf_ptr.reset(ggml_backend_alloc_buffer(bcj.backend, backend_ctx->max_tmp_size)); + } + tensor->buffer = buf_ptr.get(); + tensor->data = ggml_backend_buffer_get_base(buf_ptr.get()); + }; + // FIXME usage_counts + auto get_cgraph_aux = [&]() -> ggml_cgraph * { + ggml_cgraph * ret = backend_ctx->cgraphs_aux[iga++]; + return ret; + }; // Preferentially use backend-specific allreduce_tensor_async (e.g. NCCL for CUDA), use a generic fallback if unavailable: auto allreduce_fallback = [&](size_t i) -> ggml_status { std::vector step_cgraphs(n_backends, nullptr); - for (size_t offset_j = 1; offset_j < n_backends; offset_j *= 2) { + // Zero out nodes that were disabled due to having a zero-sized slice: + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + ggml_tensor * node = bcj.cgraphs[i].cgraph_main->nodes[bcj.cgraphs[i].cgraph_main->n_nodes - 1]; + if (node->flags & GGML_TENSOR_FLAG_COMPUTE) { + continue; + } + ggml_tensor * node_zero = get_node_aux(node); + node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN + node_zero->src[0] = node; + ggml_set_op_params_f32(node_zero, 0, 0.0f); + node_zero->data = node->data; + node_zero->buffer = node->buffer; + node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE; + + step_cgraphs[j] = get_cgraph_aux(); + step_cgraphs[j]->nodes[0] = node_zero; + step_cgraphs[j]->n_nodes = 1; + const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + } + std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); + + auto push_data = [&](const size_t j_src, const size_t j_dst, const size_t i_buf) { + assert(step_cgraphs[j_dst] == nullptr); + auto & bcj_src = backend_ctx->backend_configs[j_src]; + auto & bcj_dst = backend_ctx->backend_configs[j_dst]; + + ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; + GGML_ASSERT(ggml_is_contiguous(node_src)); + GGML_ASSERT(ggml_is_contiguous(node_dst)); + + ggml_tensor * node_tmp = get_node_aux(node_dst); + set_tmp_data(node_tmp, j_dst, i_buf); + + ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_tmp); + + ggml_tensor * node_red = get_node_aux(node_dst); + node_red->view_src = node_dst->view_src == nullptr ? node_dst : node_dst->view_src; + node_red->view_offs = node_dst->view_offs; + node_red->op = GGML_OP_ADD; + node_red->src[0] = node_dst; + node_red->src[1] = node_tmp; + node_red->flags |= GGML_TENSOR_FLAG_COMPUTE; + ggml_backend_view_init(node_red); + + ggml_cgraph * cgraph_aux = get_cgraph_aux(); + cgraph_aux->nodes[0] = node_red; + cgraph_aux->n_nodes = 1; + step_cgraphs[j_dst] = cgraph_aux; + }; + + size_t offset_j = n_backends/2; + while ((offset_j & (offset_j - 1)) != 0) { + offset_j--; + } + const size_t offset_j_max = offset_j; + size_t i_buf = 0; + + // If n_backends is not a power of 2, fold in the excess prior to butterfly reduction: + for (size_t j_src = 2*offset_j_max; j_src < n_backends; j_src++) { + const size_t j_dst = j_src - 2*offset_j_max; + push_data(j_src, j_dst, i_buf); + const ggml_status status = ggml_backend_graph_compute_async(backend_ctx->backend_configs[j_dst].backend, step_cgraphs[j_dst]); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + i_buf = 1; + } + + // Butterfly reduction: + for (; offset_j >= 1; offset_j /= 2) { std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); - for (size_t j = 0; j < n_backends; j++) { + for (size_t j = 0; j < 2*offset_j_max; j++) { const size_t j_other = j ^ offset_j; - if (j_other > j) { + if (j_other >= n_backends) { continue; } + push_data(j, j_other, i_buf); + } - auto & bcj1 = backend_ctx->backend_configs[j]; - auto & bcj2 = backend_ctx->backend_configs[j_other]; - - ggml_tensor * node1 = bcj1.cgraphs[i].cgraph_main->nodes[bcj1.cgraphs[i].cgraph_main->n_nodes - 1]; - ggml_tensor * node2 = bcj2.cgraphs[i].cgraph_main->nodes[bcj2.cgraphs[i].cgraph_main->n_nodes - 1]; - GGML_ASSERT(ggml_is_contiguous(node1)); - GGML_ASSERT(ggml_is_contiguous(node2)); - - // Tmp tensors to receive P2P copies - ggml_tensor * node_tmp_1 = get_node_aux(node1); - node_tmp_1->buffer = bcj1.buf.get(); - node_tmp_1->data = ggml_backend_buffer_get_base(bcj1.buf.get()); - - ggml_tensor * node_tmp_2 = get_node_aux(node2); - node_tmp_2->buffer = bcj2.buf.get(); - node_tmp_2->data = ggml_backend_buffer_get_base(bcj2.buf.get()); - - // 2 P2P copies: exchange full buffers - ggml_backend_tensor_copy_async(bcj1.backend, bcj2.backend, node1, node_tmp_2); - ggml_backend_tensor_copy_async(bcj2.backend, bcj1.backend, node2, node_tmp_1); - - // Local ADD: node1 += tmp1 (in-place via view) - ggml_tensor * node_red_1 = get_node_aux(node1); - node_red_1->view_src = node1->view_src == nullptr ? node1 : node1->view_src; - node_red_1->view_offs = node1->view_offs; - node_red_1->op = GGML_OP_ADD; - node_red_1->src[0] = node1; - node_red_1->src[1] = node_tmp_1; - node_red_1->flags |= GGML_TENSOR_FLAG_COMPUTE; - ggml_backend_view_init(node_red_1); - - // Local ADD: node2 += tmp2 (in-place via view) - ggml_tensor * node_red_2 = get_node_aux(node2); - node_red_2->view_src = node2->view_src == nullptr ? node2 : node2->view_src; - node_red_2->view_offs = node2->view_offs; - node_red_2->op = GGML_OP_ADD; - node_red_2->src[0] = node2; - node_red_2->src[1] = node_tmp_2; - node_red_2->flags |= GGML_TENSOR_FLAG_COMPUTE; - ggml_backend_view_init(node_red_2); - - // Build 1-node cgraphs for the ADD ops - ggml_cgraph * cgraph_aux_1 = get_cgraph_aux(); - cgraph_aux_1->nodes[0] = node_red_1; - cgraph_aux_1->n_nodes = 1; - step_cgraphs[j] = cgraph_aux_1; - - ggml_cgraph * cgraph_aux_2 = get_cgraph_aux(); - cgraph_aux_2->nodes[0] = node_red_2; - cgraph_aux_2->n_nodes = 1; - step_cgraphs[j_other] = cgraph_aux_2; - } - - // Execute local ADDs for this step - for (size_t j = 0; j < n_backends; j++) { + for (size_t j = 0; j < 2*offset_j_max; j++) { if (step_cgraphs[j] == nullptr) { continue; } @@ -1834,12 +2190,25 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, return status; } } + i_buf++; } + assert(i_buf == backend_ctx->n_reduce_steps); + + // If n_backends is not a power of 2, copy back the reduced tensors to the excess: + for (size_t j = 2*offset_j_max; j < n_backends; j++) { + auto & bcj_src = backend_ctx->backend_configs[j - 2*offset_j_max]; + auto & bcj_dst = backend_ctx->backend_configs[j]; + + ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_dst); + } + return GGML_STATUS_SUCCESS; }; - for (size_t i = 0; i < n_subgraphs; i++) { + for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) { for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main); @@ -1848,22 +2217,17 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, } } - if (n_backends > 1 && i < n_subgraphs - 1) { + if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) { bool backend_allreduce_success = false; - ggml_backend_allreduce_tensor_t allreduce_tensor = (ggml_backend_allreduce_tensor_t) ggml_backend_reg_get_proc_address( - ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_ctx->backend_configs[0].backend)), "ggml_backend_allreduce_tensor"); - if (allreduce_tensor) { - std::vector backends; - backends.reserve(n_backends); + if (backend_ctx->comm_ctx) { std::vector nodes; nodes.reserve(n_backends); for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; - backends.push_back(bcj.backend); ggml_cgraph * cgraph_ij = bcj.cgraphs[i].cgraph_main; nodes.push_back(cgraph_ij->nodes[cgraph_ij->n_nodes-1]); } - backend_allreduce_success = allreduce_tensor(backends.data(), nodes.data(), n_backends); + backend_allreduce_success = backend_ctx->comm_allreduce(backend_ctx->comm_ctx, nodes.data()); } if (!backend_allreduce_success) { @@ -1922,4 +2286,3 @@ ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, siz const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; return backend_ctx->backend_configs[index].backend; } - diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 01073f63120e..7f1b57b1a263 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1066,6 +1066,8 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra GGML_ABORT("%s: failed to initialize context\n", __func__); } + graph->uid = ggml_graph_next_uid(); + // pass 1: assign backends to ops with pre-allocated inputs for (int i = 0; i < graph->n_leafs; i++) { struct ggml_tensor * leaf = graph->leafs[i]; @@ -1513,6 +1515,11 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra assert(graph_copy->size > graph_copy->n_leafs); graph_copy->leafs[graph_copy->n_leafs++] = leaf; } + + // set ids for all splits + for (int i = 0; i < sched->n_splits; ++i) { + sched->splits[i].graph.uid = ggml_graph_next_uid(); + } } static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9df1461a5749..dc959bbfeaa6 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1092,10 +1092,6 @@ struct ggml_cuda_device_info { cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {}; std::array default_tensor_split = {}; - -#ifdef GGML_USE_NCCL - ncclComm_t comms[GGML_CUDA_MAX_DEVICES]; -#endif // GGML_USE_NCCL }; const ggml_cuda_device_info & ggml_cuda_info(); @@ -1185,6 +1181,7 @@ struct ggml_cuda_graph { bool warmup_complete = false; // GGML_CUDA_GRAPH_STATS=1 counters uint64_t stat_total = 0, stat_replay = 0, stat_capture = 0, stat_eager = 0; + uint64_t uid = 0; struct node_properties { ggml_tensor node; void * node_src_data_ptrs[GGML_MAX_SRC]; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 29e77ae72117..6ea8683bce81 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -419,14 +419,6 @@ static ggml_cuda_device_info ggml_cuda_init() { } } -#ifdef GGML_USE_NCCL - int dev_ids[GGML_CUDA_MAX_DEVICES]; - for (int id = 0; id < info.device_count; ++id) { - dev_ids[id] = id; - } - NCCL_CHECK(ncclCommInitAll(info.comms, info.device_count, dev_ids)); -#endif // GGML_USE_NCCL - return info; } @@ -1276,7 +1268,69 @@ static const ggml_backend_buffer_type_i ggml_backend_cuda_split_buffer_type_inte /* .is_host = */ ggml_backend_cuda_split_buffer_type_is_host, }; -bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends) { +#ifdef GGML_USE_NCCL +struct ggml_backend_cuda_comm_context { + std::vector backends; + std::vector comms; + + ~ggml_backend_cuda_comm_context() { + for (ncclComm_t comm : comms) { + NCCL_CHECK(ncclCommDestroy(comm)); + } + } +}; +#endif // GGML_USE_NCCL + +static void ggml_backend_cuda_comm_free(void * comm_ctx_v) { +#ifdef GGML_USE_NCCL + if (comm_ctx_v == nullptr) { + return; + } + ggml_backend_cuda_comm_context * comm_ctx = (ggml_backend_cuda_comm_context *) comm_ctx_v; + delete comm_ctx; +#else + GGML_UNUSED(comm_ctx_v); +#endif // GGML_USE_NCCL +} + +static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_backends) { +#ifdef GGML_USE_NCCL + for (size_t i = 0; i < n_backends; i++) { + if (!ggml_backend_is_cuda(backends[i])) { + return nullptr; + } + } + ggml_backend_cuda_comm_context * ret = new ggml_backend_cuda_comm_context; + std::vector dev_ids; + ret->backends.reserve(n_backends); + dev_ids.reserve(n_backends); + for (size_t i = 0; i < n_backends; i++) { + ret->backends.push_back(backends[i]); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + dev_ids.push_back(cuda_ctx->device); + } + + ret->comms.resize(n_backends); + NCCL_CHECK(ncclCommInitAll(ret->comms.data(), n_backends, dev_ids.data())); + return ret; +#else + // If NCCL is installed it is used by default for optimal performance. + // However, NVIDIA does not distribute NCCL with CUDA so users may be unwittingly missing this package. + // RCCL is disabled by default, users are explicitly opting in. + // Therefore print no warning for RCCL. +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + static bool warning_printed = false; + if (!warning_printed) { + GGML_LOG_WARN("%s: NVIDIA Collective Communications Library (NCCL) is unavailable, multi GPU performance will be suboptimal\n", __func__); + warning_printed = true; + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(backends, n_backends); + return nullptr; +#endif // GGML_USE_NCCL +} + +static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) { #ifdef GGML_USE_NCCL const int64_t ne = ggml_nelements(tensors[0]); // FIXME the input of llm_graph_context::build_in_out_ids can produce a tensor with 0 elements if n_outputs == 0 @@ -1284,21 +1338,31 @@ bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_t if (ne == 0) { return true; } + + GGML_ASSERT(comm_ctx_v != nullptr); + ggml_backend_cuda_comm_context * comm_ctx = (ggml_backend_cuda_comm_context *) comm_ctx_v; + const size_t n_backends = comm_ctx->backends.size(); + for (size_t i = 0; i < n_backends; ++i) { GGML_ASSERT(tensors[i] != nullptr); GGML_ASSERT(ggml_nelements(tensors[i]) == ne); GGML_ASSERT(ggml_is_contiguously_allocated(tensors[i])); } - const ggml_cuda_device_info info = ggml_cuda_info(); - // For small tensors, simply reduce them as FP32. // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0. if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) { + for (size_t i = 0; i < n_backends; ++i) { + if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + ggml_cuda_set_device(cuda_ctx->device); + CUDA_CHECK(cudaMemsetAsync(tensors[i]->data, 0, ggml_nbytes(tensors[i]), cuda_ctx->stream())); + } + } NCCL_CHECK(ncclGroupStart()); for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); } NCCL_CHECK(ncclGroupEnd()); @@ -1311,44 +1375,37 @@ bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_t ggml_cuda_pool_alloc tmp[GGML_CUDA_MAX_DEVICES]; for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; tmp[i].pool = &cuda_ctx->pool(); tmp[i].alloc(ne); - ggml_cuda_set_device(i); - to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); + ggml_cuda_set_device(cuda_ctx->device); + if (tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) { + to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); + } else { + CUDA_CHECK(cudaMemsetAsync(tmp[i].get(), 0, ne * sizeof(nv_bfloat16), cuda_ctx->stream())); + } CUDA_CHECK(cudaGetLastError()); } NCCL_CHECK(ncclGroupStart()); for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); } NCCL_CHECK(ncclGroupEnd()); for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - ggml_cuda_set_device(i); + ggml_cuda_set_device(cuda_ctx->device); to_fp32(tmp[i].get(), (float *) tensors[i]->data, ne, cuda_ctx->stream()); CUDA_CHECK(cudaGetLastError()); } return true; #else - // If NCCL is installed it is used by default for optimal performance. - // However, NVIDIA does not distribute NCCL with CUDA so users may be unwittingly missing this package. - // RCCL is disabled by default, users are explicitly opting in. - // Therefore print no warning for RCCL. -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - static bool warning_printed = false; - if (!warning_printed) { - GGML_LOG_WARN("%s: NVIDIA Collective Communications Library (NCCL) is unavailable, multi GPU performance will be suboptimal\n", __func__); - warning_printed = true; - } -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - GGML_UNUSED_VARS(backends, tensors, n_backends); + GGML_UNUSED_VARS(comm_ctx_v, tensors); return false; #endif // GGML_USE_NCCL } @@ -3246,6 +3303,15 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx const void * graph_key = ggml_cuda_graph_get_key(cgraph); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (cgraph->uid != 0 && + cgraph->uid == graph->uid) { + GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); + GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); + return false; + } + + graph->uid = cgraph->uid; + // Check if the graph size has changed if ((int)graph->node_props.size() != cgraph->n_nodes) { res = true; @@ -4330,6 +4396,8 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (graph->is_enabled()) { + static const bool log_graph_warmup = + getenv("GGML_CUDA_GRAPH_WARMUP_LOG") != nullptr; const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); if (graph_compatible) { const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); @@ -4338,7 +4406,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, // Warmup: need at least 2 calls with no property change on the 2nd call if (!properties_changed) { graph->warmup_complete = true; - GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); + if (log_graph_warmup) { + GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); + } use_cuda_graph = true; cuda_graph_update_required = true; } @@ -4348,7 +4418,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, if (properties_changed) { // Properties changed - reset warmup, execute directly until stable again graph->warmup_complete = false; - GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + if (log_graph_warmup) { + GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + } } else { use_cuda_graph = true; cuda_graph_update_required = graph->instance == nullptr; @@ -5470,8 +5542,14 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); - if (strcmp(name, "ggml_backend_allreduce_tensor") == 0) { - return (void *)ggml_backend_cuda_allreduce_tensor; + if (strcmp(name, "ggml_backend_comm_init") == 0) { + return (void *)ggml_backend_cuda_comm_init; + } + if (strcmp(name, "ggml_backend_comm_free") == 0) { + return (void *)ggml_backend_cuda_comm_free; + } + if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) { + return (void *)ggml_backend_cuda_comm_allreduce_tensor; } if (strcmp(name, "ggml_backend_split_buffer_type") == 0) { return (void *)ggml_backend_cuda_split_buffer_type; diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 0639db362e74..62b76abbcec9 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -30,6 +30,8 @@ extern "C" { void ggml_print_backtrace(void); +uint64_t ggml_graph_next_uid(void); + #ifndef MIN # define MIN(a, b) ((a) < (b) ? (a) : (b)) #endif @@ -338,6 +340,10 @@ struct ggml_cgraph { struct ggml_hash_set visited_hash_set; enum ggml_cgraph_eval_order order; + + // an optional identifier that can be utilized to recognize same graphs if two non-zero values match + // a value of 0 means it is not set and should be ignored + uint64_t uid; }; // returns a slice of cgraph with nodes [i0, i1) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c1ffccafb7e3..0e3117e99e0c 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -53,6 +53,16 @@ #define UNUSED GGML_UNUSED +uint64_t ggml_graph_next_uid(void) { +#ifdef _MSC_VER + static volatile long long counter = 1; + return (uint64_t) _InterlockedIncrement64(&counter) - 1; +#else + static uint64_t counter = 1; + return __atomic_fetch_add(&counter, 1, __ATOMIC_RELAXED); +#endif +} + // Needed for ggml_fp32_to_bf16_row() #if defined(__AVX512BF16__) #if defined(_MSC_VER) @@ -7253,6 +7263,7 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz /*.use_counts =*/ use_counts_ptr, /*.hash_table =*/ { hash_size, hash_used, hash_keys_ptr }, /*.order =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT, + /*.uid =*/ 0, }; ggml_hash_set_reset(&cgraph->visited_hash_set); @@ -7280,6 +7291,7 @@ struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1) /*.use_counts =*/ cgraph0->use_counts, /*.visited_hash_set =*/ cgraph0->visited_hash_set, /*.order =*/ cgraph0->order, + /*.uid =*/ 0 }; return cgraph; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index cd4bc5ef1d31..c417cfefe181 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -239,6 +239,7 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) endif() llama_build_and_test(test-gguf.cpp) llama_build_and_test(test-backend-ops.cpp) +llama_build_and_test(test-backend-meta.cpp) llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") diff --git a/tests/test-backend-meta.cpp b/tests/test-backend-meta.cpp new file mode 100644 index 000000000000..dd4e0be65441 --- /dev/null +++ b/tests/test-backend-meta.cpp @@ -0,0 +1,299 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +#define CHECK(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "CHECK failed: %s (%s:%d)\n", \ + #condition, __FILE__, __LINE__); \ + std::abort(); \ + } \ +} while (0) + +static ggml_backend_meta_split_state split_state( + const ggml_tensor * tensor, + void * userdata) { + const size_t n_devices = *static_cast(userdata); + ggml_backend_meta_split_state state{}; + if (std::strcmp(tensor->name, "repeated_axis2") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_2; + state.ne[0] = 2; + state.ne[1] = 2; + state.nr[0] = 3; + state.n_segments = 1; + return state; + } + if (std::strcmp(tensor->name, "axis2") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_2; + state.ne[0] = tensor->ne[2] / 2; + state.ne[1] = tensor->ne[2] - state.ne[0]; + // Leave nr zero to cover callbacks compiled before repetition counts + // were added to the split-state contract. + state.n_segments = 1; + return state; + } + if (std::strcmp(tensor->name, "row_weight") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_0; + state.ne[0] = tensor->ne[0] / 2; + state.ne[1] = tensor->ne[0] - state.ne[0]; + state.nr[0] = 1; + state.n_segments = 1; + return state; + } + if (std::strcmp(tensor->name, "column_weight") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_1; + state.ne[0] = tensor->ne[1] / 2; + state.ne[1] = tensor->ne[1] - state.ne[0]; + state.nr[0] = 1; + state.n_segments = 1; + return state; + } + if (std::strcmp(tensor->name, "repeated") != 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_MIRRORED; + state.nr[0] = 1; + state.n_segments = 1; + return state; + } + + if (n_devices != 2) { + std::fprintf(stderr, "test fixture requires exactly two devices\n"); + std::abort(); + } + state.axis = GGML_BACKEND_SPLIT_AXIS_0; + state.ne[0] = 2; + state.ne[1] = 2; + state.ne[2] = 4; + state.ne[3] = 4; + state.nr[0] = 2; + state.nr[1] = 2; + state.n_segments = 2; + return state; +} + +int main() { + ggml_backend_load_all(); + + const char * first_name = std::getenv("GGML_META_TEST_DEVICE_0"); + const char * second_name = std::getenv("GGML_META_TEST_DEVICE_1"); + first_name = first_name ? first_name : "CUDA0"; + second_name = second_name ? second_name : "CUDA1"; + + ggml_backend_dev_t devices[] = { + ggml_backend_dev_by_name(first_name), + ggml_backend_dev_by_name(second_name), + }; + if (!devices[0] || !devices[1]) { + std::printf("SKIP: two requested devices are not available (%s,%s)\n", + first_name, second_name); + return 0; + } + + size_t n_devices = 2; + ggml_backend_dev_t meta_device = + ggml_backend_meta_device(devices, n_devices, split_state, &n_devices); + CHECK(meta_device); + ggml_backend_t backend = ggml_backend_dev_init(meta_device, nullptr); + CHECK(backend); + + ggml_init_params params{}; + constexpr size_t graph_nodes = 64; + params.mem_size = 8 * ggml_tensor_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + CHECK(ctx); + + ggml_tensor * repeated = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 24, 4); + ggml_set_name(repeated, "repeated"); + ggml_tensor * mirrored = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 24, 4); + ggml_set_name(mirrored, "mirrored"); + ggml_tensor * repeated_axis2 = + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 3, 12, 2); + ggml_set_name(repeated_axis2, "repeated_axis2"); + ggml_tensor * axis2 = + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 3, 12, 2); + ggml_set_name(axis2, "axis2"); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + CHECK(buffer); + + std::vector input((size_t) ggml_nelements(repeated)); + for (size_t i = 0; i < input.size(); ++i) input[i] = (float) i + 0.25f; + ggml_backend_tensor_set(repeated, input.data(), 0, input.size() * sizeof(float)); + ggml_backend_tensor_set(mirrored, input.data(), 0, input.size() * sizeof(float)); + + std::vector output(input.size(), 0.0f); + ggml_backend_tensor_get(repeated, output.data(), 0, output.size() * sizeof(float)); + CHECK(output == input); + std::fill(output.begin(), output.end(), 0.0f); + ggml_backend_tensor_get(mirrored, output.data(), 0, output.size() * sizeof(float)); + CHECK(output == input); + + std::vector partial(48, 0.0f); + ggml_backend_tensor_get(repeated, partial.data(), repeated->nb[1], + partial.size() * sizeof(float)); + for (size_t i = 0; i < partial.size(); ++i) { + CHECK(partial[i] == input[i + 24]); + } + + std::vector axis2_input((size_t) ggml_nelements(repeated_axis2)); + for (size_t i = 0; i < axis2_input.size(); ++i) axis2_input[i] = (float) i + 0.5f; + ggml_backend_tensor_set(repeated_axis2, axis2_input.data(), 0, + axis2_input.size() * sizeof(float)); + std::vector axis2_output(axis2_input.size(), 0.0f); + ggml_backend_tensor_get(repeated_axis2, axis2_output.data(), 0, + axis2_output.size() * sizeof(float)); + CHECK(axis2_output == axis2_input); + + std::vector axis2_partial((size_t) repeated_axis2->ne[0] * + repeated_axis2->ne[1] * repeated_axis2->ne[2]); + ggml_backend_tensor_get(repeated_axis2, axis2_partial.data(), + repeated_axis2->nb[3], repeated_axis2->nb[3]); + for (size_t i = 0; i < axis2_partial.size(); ++i) { + CHECK(axis2_partial[i] == axis2_input[i + axis2_partial.size()]); + } + + std::vector simple_axis2_input((size_t) ggml_nelements(axis2)); + for (size_t i = 0; i < simple_axis2_input.size(); ++i) { + simple_axis2_input[i] = (float) i + 0.75f; + } + ggml_backend_tensor_set(axis2, simple_axis2_input.data(), 0, + simple_axis2_input.size() * sizeof(float)); + std::vector simple_axis2_output(simple_axis2_input.size(), 0.0f); + ggml_backend_tensor_get(axis2, simple_axis2_output.data(), 0, + simple_axis2_output.size() * sizeof(float)); + CHECK(simple_axis2_output == simple_axis2_input); + + const int64_t test_head = 7; + const size_t partial_bytes = 2 * axis2->nb[1]; + std::vector simple_axis2_partial( + partial_bytes / sizeof(float), 0.0f); + ggml_backend_tensor_get(axis2, simple_axis2_partial.data(), + (size_t) test_head * axis2->nb[2], + partial_bytes); + const size_t logical_start = + (size_t) test_head * axis2->ne[0] * axis2->ne[1]; + for (size_t i = 0; i < simple_axis2_partial.size(); ++i) { + CHECK(simple_axis2_partial[i] == simple_axis2_input[logical_start + i]); + simple_axis2_partial[i] = -(float) i - 1.0f; + } + ggml_backend_tensor_set(axis2, simple_axis2_partial.data(), + (size_t) test_head * axis2->nb[2], + partial_bytes); + ggml_backend_tensor_get(axis2, simple_axis2_output.data(), 0, + simple_axis2_output.size() * sizeof(float)); + for (size_t i = 0; i < simple_axis2_partial.size(); ++i) { + CHECK(simple_axis2_output[logical_start + i] == + simple_axis2_partial[i]); + } + + ggml_init_params weight_params{}; + weight_params.mem_size = 8 * ggml_tensor_overhead(); + weight_params.no_alloc = true; + ggml_context * weight_ctx = ggml_init(weight_params); + CHECK(weight_ctx); + ggml_tensor * column_weight = + ggml_new_tensor_2d(weight_ctx, GGML_TYPE_F32, 4, 8); + ggml_set_name(column_weight, "column_weight"); + ggml_tensor * row_weight = + ggml_new_tensor_2d(weight_ctx, GGML_TYPE_F32, 8, 4); + ggml_set_name(row_weight, "row_weight"); + ggml_backend_buffer_t weight_buffer = + ggml_backend_alloc_ctx_tensors(weight_ctx, backend); + CHECK(weight_buffer); + ggml_backend_buffer_set_usage(weight_buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + ggml_init_params graph_params{}; + graph_params.mem_size = 16 * ggml_tensor_overhead() + + ggml_graph_overhead_custom(graph_nodes, false); + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx); + ggml_tensor * mat_input = + ggml_new_tensor_2d(graph_ctx, GGML_TYPE_F32, 4, 2); + ggml_set_name(mat_input, "mat_input"); + ggml_tensor * mat_hidden = + ggml_mul_mat(graph_ctx, column_weight, mat_input); + ggml_set_name(mat_hidden, "mat_hidden"); + ggml_tensor * mat_result = ggml_mul_mat(graph_ctx, row_weight, mat_hidden); + ggml_set_name(mat_result, "mat_result"); + ggml_tensor * mat_output = ggml_scale(graph_ctx, mat_result, 1.0f); + ggml_set_name(mat_output, "mat_output"); + ggml_set_output(mat_output); + ggml_cgraph * graph = + ggml_new_graph_custom(graph_ctx, graph_nodes, false); + ggml_build_forward_expand(graph, mat_output); + ggml_gallocr_t graph_alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + CHECK(graph_alloc); + CHECK(ggml_gallocr_alloc_graph(graph_alloc, graph)); + + std::vector column_data((size_t) ggml_nelements(column_weight)); + std::vector row_data((size_t) ggml_nelements(row_weight)); + std::vector mat_input_data((size_t) ggml_nelements(mat_input)); + for (size_t i = 0; i < column_data.size(); ++i) { + column_data[i] = (float) ((int) (i % 11) - 5) / 8.0f; + } + for (size_t i = 0; i < row_data.size(); ++i) { + row_data[i] = (float) ((int) (i % 13) - 6) / 9.0f; + } + for (size_t i = 0; i < mat_input_data.size(); ++i) { + mat_input_data[i] = (float) ((int) (i % 7) - 3) / 4.0f; + } + ggml_backend_tensor_set(column_weight, column_data.data(), 0, + column_data.size() * sizeof(float)); + ggml_backend_tensor_set(row_weight, row_data.data(), 0, + row_data.size() * sizeof(float)); + ggml_backend_tensor_set(mat_input, mat_input_data.data(), 0, + mat_input_data.size() * sizeof(float)); + CHECK(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + + std::vector result((size_t) ggml_nelements(mat_output)); + ggml_backend_tensor_get(mat_output, result.data(), 0, + result.size() * sizeof(float)); + std::vector hidden(16, 0.0f); + for (int column = 0; column < 2; ++column) { + for (int row = 0; row < 8; ++row) { + for (int k = 0; k < 4; ++k) { + hidden[(size_t) column * 8 + row] += + column_data[(size_t) row * 4 + k] * + mat_input_data[(size_t) column * 4 + k]; + } + } + } + for (int column = 0; column < 2; ++column) { + for (int row = 0; row < 4; ++row) { + float expected = 0.0f; + for (int k = 0; k < 8; ++k) { + expected += row_data[(size_t) row * 8 + k] * + hidden[(size_t) column * 8 + k]; + } + const float actual = result[(size_t) column * 4 + row]; + if (std::fabs(actual - expected) >= 1e-5f) { + std::fprintf(stderr, + "matmul mismatch column=%d row=%d actual=%f expected=%f\n", + column, row, actual, expected); + } + CHECK(std::fabs(actual - expected) < 1e-5f); + } + } + + ggml_gallocr_free(graph_alloc); + ggml_free(graph_ctx); + ggml_backend_buffer_free(weight_buffer); + ggml_free(weight_ctx); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + std::printf("meta backend selected-device round trip passed (%s,%s)\n", + first_name, second_name); + return 0; +} From 871055ef424599eaa52c0ed95f48be65d6840f0e Mon Sep 17 00:00:00 2001 From: Jun Yamog Date: Fri, 17 Jul 2026 03:46:17 +1200 Subject: [PATCH 2/4] feat(meta): expose rank-local tensor views --- ggml/src/ggml-backend-impl.h | 3 +++ ggml/src/ggml-backend-meta.cpp | 13 +++++++++++++ tests/CMakeLists.txt | 1 + tests/test-backend-meta.cpp | 12 ++++++++++++ 4 files changed, 29 insertions(+) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 9c56ec30c5f1..e22b7338bfe2 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -94,6 +94,9 @@ extern "C" { GGML_API size_t ggml_backend_meta_n_backends (ggml_backend_t meta_backend); GGML_API ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index); + // Non-owning rank-local view. Compute-tensor views remain valid only until + // the meta backend rebuilds its graph tensor container. + GGML_API struct ggml_tensor * ggml_backend_meta_simple_tensor(const struct ggml_tensor * tensor, size_t index); // temporary workaround to statically allocate tensors from a context in a deduplicated way: GGML_API struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 01b5df7c710c..09d38a4cede6 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -483,6 +483,19 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct return it->second[index]; } +struct ggml_tensor * ggml_backend_meta_simple_tensor(const struct ggml_tensor * tensor, size_t index) { + if (tensor == nullptr || tensor->buffer == nullptr || + !ggml_backend_buffer_is_meta(tensor->buffer)) { + return nullptr; + } + ggml_backend_meta_buffer_context * buf_ctx = + (ggml_backend_meta_buffer_context *) tensor->buffer->context; + if (index >= buf_ctx->bufs.size()) { + return nullptr; + } + return ggml_backend_meta_buffer_simple_tensor(tensor, index); +} + static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c417cfefe181..77690706b8b2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -240,6 +240,7 @@ endif() llama_build_and_test(test-gguf.cpp) llama_build_and_test(test-backend-ops.cpp) llama_build_and_test(test-backend-meta.cpp) +target_include_directories(test-backend-meta PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../ggml/src) llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") diff --git a/tests/test-backend-meta.cpp b/tests/test-backend-meta.cpp index dd4e0be65441..678ab9e754ee 100644 --- a/tests/test-backend-meta.cpp +++ b/tests/test-backend-meta.cpp @@ -1,6 +1,7 @@ #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-backend-impl.h" #include #include @@ -123,6 +124,15 @@ int main() { ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); CHECK(buffer); + ggml_tensor * repeated_0 = ggml_backend_meta_simple_tensor(repeated, 0); + ggml_tensor * repeated_1 = ggml_backend_meta_simple_tensor(repeated, 1); + CHECK(repeated_0); + CHECK(repeated_1); + CHECK(repeated_0 != repeated_1); + CHECK(ggml_backend_meta_simple_tensor(repeated, 2) == nullptr); + CHECK(ggml_backend_meta_simple_tensor(nullptr, 0) == nullptr); + CHECK(ggml_nelements(repeated_0) + ggml_nelements(repeated_1) == + ggml_nelements(repeated)); std::vector input((size_t) ggml_nelements(repeated)); for (size_t i = 0; i < input.size(); ++i) input[i] = (float) i + 0.25f; @@ -235,6 +245,8 @@ int main() { ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); CHECK(graph_alloc); CHECK(ggml_gallocr_alloc_graph(graph_alloc, graph)); + CHECK(ggml_backend_meta_simple_tensor(mat_output, 0)); + CHECK(ggml_backend_meta_simple_tensor(mat_output, 1)); std::vector column_data((size_t) ggml_nelements(column_weight)); std::vector row_data((size_t) ggml_nelements(row_weight)); From e0dc5509d8ae9cfbf234cb6c1d74ab70f60bcabe Mon Sep 17 00:00:00 2001 From: Jun Yamog Date: Fri, 17 Jul 2026 07:57:32 +1200 Subject: [PATCH 3/4] fix(meta): preserve repeated split layouts --- ggml/include/ggml-backend.h | 2 +- ggml/src/ggml-backend-meta.cpp | 153 ++++++++++++++++++--------------- tests/test-backend-meta.cpp | 124 ++++++++++++++++++++++++++ 3 files changed, 209 insertions(+), 70 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 48ebcafd4823..e32a5d4ca082 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -389,8 +389,8 @@ extern "C" { // the Q matrix can be larger than the K and V matrices so this can either be expressed as 3 segments or as 2 segments // where the segment for K/V repeats twice int64_t ne[16*GGML_BACKEND_META_MAX_DEVICES]; - uint32_t nr[16]; uint32_t n_segments; + uint32_t nr[16]; }; // function to assign split states for statically allocated tensors, compute tensor split states will be assigned to be compatible: diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 09d38a4cede6..0552a3fcb3d2 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -212,7 +212,15 @@ static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t me ggml_backend_dev_t ggml_backend_meta_device( ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) { - GGML_ASSERT(n_devs <= GGML_BACKEND_META_MAX_DEVICES); + if (devs == nullptr || get_split_state == nullptr || n_devs == 0 || + n_devs > GGML_BACKEND_META_MAX_DEVICES) { + return nullptr; + } + for (size_t i = 0; i < n_devs; ++i) { + if (devs[i] == nullptr) { + return nullptr; + } + } // TODO: this is not thread-safe - needs to be fixed static std::vector> ctxs; static std::map meta_devs; @@ -527,7 +535,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { - ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}; + ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1, {1}}; for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { continue; @@ -535,15 +543,15 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { ret = src_ss[i]; } else if (!split_states_equal(src_ss[i], ret)) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; break; } } if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); return ret; @@ -601,7 +609,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( auto handle_mul_mat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1, {1}}; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { ggml_backend_meta_split_state ret = src_ss[0]; @@ -615,10 +623,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); - return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; + return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, 1, {1}}; } GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; }; auto handle_reshape = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { @@ -627,9 +635,12 @@ 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: { + if (ggml_are_same_shape(tensor, tensor->src[0])) { + return src_ss[0]; + } 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}; + 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++) { @@ -640,12 +651,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( 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}; + return {ggml_backend_meta_split_axis(dim), {0}, 1, {uint32_t(base_ne_out_next/base_ne_in)}}; } if (base_ne_out_next > 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}; + return {ggml_backend_meta_split_axis(dim), {0}, 1, {1}}; } base_ne_out = base_ne_out_next; } @@ -657,7 +668,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } } }; @@ -692,7 +703,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { for (int dim = 0; dim < GGML_MAX_DIMS-1; dim++) { if (tensor->nb[dim+1] == tensor->src[0]->nb[axis+1]) { - return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; + return {ggml_backend_meta_split_axis(dim), {0}, 1, {1}}; } } GGML_ABORT("fatal error"); @@ -701,7 +712,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return src_ss[0]; } GGML_ABORT("view of permuted tensor not implemented"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; }; auto handle_permute = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { @@ -710,8 +721,10 @@ 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 || src_ss[0].nr[0] == 1); - return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, {src_ss[0].nr[0]}, 1}; + ggml_backend_meta_split_state result = src_ss[0]; + result.axis = ggml_backend_meta_split_axis( + tensor->op_params[src_ss[0].axis]); + return result; } case GGML_BACKEND_SPLIT_AXIS_MIRRORED: case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { @@ -719,7 +732,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } } }; @@ -728,8 +741,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( switch (src_ss[0].axis) { case GGML_BACKEND_SPLIT_AXIS_0: case GGML_BACKEND_SPLIT_AXIS_1: { - GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); - return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, {src_ss[0].nr[0]}, 1}; + ggml_backend_meta_split_state result = src_ss[0]; + result.axis = ggml_backend_meta_split_axis( + int(src_ss[0].axis) ^ 1); + return result; } case GGML_BACKEND_SPLIT_AXIS_2: case GGML_BACKEND_SPLIT_AXIS_3: @@ -739,7 +754,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } default: { GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } } }; @@ -777,16 +792,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1, {1}}; }; auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1, {1}}; } if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1, {1}}; } } return handle_generic(src_ss, /*scalar_only =*/ false); @@ -806,12 +821,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( // state shape is (S_v*S_v*H, K, n_seqs); the heads dim is nested inside axis 0, // so a head-aligned split on the input cache reshapes to axis 0 here (not axis 2). GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0); - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1, {1}}; }; auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { if (ggml_nelements(tensor) == 0) { - return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); @@ -838,10 +853,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return ret; } - std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); + std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1, {1}}); for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { - src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; continue; } src_ss[i] = ggml_backend_meta_get_split_state(stc, tensor->src[i], /*assume_sync =*/ true); @@ -851,7 +866,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( ggml_backend_meta_split_state split_state; switch (tensor->op) { case GGML_OP_NONE: { - split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1, {1}}; } break; case GGML_OP_DUP: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); @@ -1048,10 +1063,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } break; default: { GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); - split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; } break; } - if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { + // Multi-segment handlers propagate a complete physical layout. The + // ratio pass below is only for one-segment states whose per-rank sizes + // still need to be derived from their inputs. + if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS && + split_state.n_segments == 1) { bool first_src_split_by_axis = true; const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); @@ -1105,6 +1124,23 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( buf_ctx->split_state_cache[key].first = calculate_split_state(); memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); if (buf_ctx->debug > 0) { + auto format_split_state = [&](const ggml_backend_meta_split_state & split_state) { + std::string result; + for (size_t s = 0; s < split_state.n_segments; ++s) { + if (!result.empty()) { + result += "; "; + } + result += "{"; + for (size_t j = 0; j < n_bufs; ++j) { + if (j > 0) { + result += ", "; + } + result += std::to_string(split_state.ne[s*n_bufs + j]); + } + result += "}x" + std::to_string(split_state.nr[s]); + } + return result; + }; std::string srcs_info; for (size_t i = 0; i < GGML_MAX_SRC; i++) { if (tensor->src[i] == nullptr) { @@ -1113,26 +1149,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (!srcs_info.empty()) { srcs_info += ", "; } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); - GGML_ASSERT(split_state.n_segments == 1); + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(tensor->src[i], true); const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - ne_info += std::to_string(split_state.ne[j]) + "x" + std::to_string(split_state.nr[0]); - } + const std::string ne_info = format_split_state(split_state); srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; } - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; - ne_info += std::to_string(ss.ne[j]) + "x" + std::to_string(ss.nr[0]); - } + const std::string ne_info = + format_split_state(buf_ctx->split_state_cache[key].first); GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); } @@ -1746,12 +1770,14 @@ static void ggml_backend_meta_free(ggml_backend_t backend) { static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_backends = ggml_backend_meta_n_backends(backend); - GGML_ASSERT(offset == 0); GGML_ASSERT(ggml_is_contiguous(tensor)); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - GGML_ASSERT(split_state.n_segments == 1); - GGML_ASSERT(split_state.nr[0] == 1); + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { + ggml_backend_tensor_set(tensor, data, offset, size); + return; + } + GGML_ASSERT(offset == 0); switch (split_state.axis) { case GGML_BACKEND_SPLIT_AXIS_0: @@ -1791,12 +1817,14 @@ static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tens static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { const size_t n_backends = ggml_backend_meta_n_backends(backend); - GGML_ASSERT(offset == 0); GGML_ASSERT(ggml_is_contiguous(tensor)); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - GGML_ASSERT(split_state.n_segments == 1); - GGML_ASSERT(split_state.nr[0] == 1); + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { + ggml_backend_tensor_get(tensor, data, offset, size); + return; + } + GGML_ASSERT(offset == 0); switch (split_state.axis) { case GGML_BACKEND_SPLIT_AXIS_0: @@ -2108,28 +2136,15 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, auto allreduce_fallback = [&](size_t i) -> ggml_status { std::vector step_cgraphs(n_backends, nullptr); - // Zero out nodes that were disabled due to having a zero-sized slice: + // Zero out nodes that were disabled due to having a zero-sized slice. + // A scale-by-zero graph is not sufficient because 0 * NaN is NaN. for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; ggml_tensor * node = bcj.cgraphs[i].cgraph_main->nodes[bcj.cgraphs[i].cgraph_main->n_nodes - 1]; if (node->flags & GGML_TENSOR_FLAG_COMPUTE) { continue; } - ggml_tensor * node_zero = get_node_aux(node); - node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN - node_zero->src[0] = node; - ggml_set_op_params_f32(node_zero, 0, 0.0f); - node_zero->data = node->data; - node_zero->buffer = node->buffer; - node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE; - - step_cgraphs[j] = get_cgraph_aux(); - step_cgraphs[j]->nodes[0] = node_zero; - step_cgraphs[j]->n_nodes = 1; - const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); - if (status != GGML_STATUS_SUCCESS) { - return status; - } + ggml_backend_tensor_memset(node, 0, 0, ggml_nbytes(node)); } std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); diff --git a/tests/test-backend-meta.cpp b/tests/test-backend-meta.cpp index 678ab9e754ee..a64519052164 100644 --- a/tests/test-backend-meta.cpp +++ b/tests/test-backend-meta.cpp @@ -32,6 +32,17 @@ static ggml_backend_meta_split_state split_state( state.n_segments = 1; return state; } + if (std::strcmp(tensor->name, "multi_repeat") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_0; + state.ne[0] = 2; + state.ne[1] = 2; + state.ne[2] = 4; + state.ne[3] = 4; + state.nr[0] = 2; + state.nr[1] = 3; + state.n_segments = 2; + return state; + } if (std::strcmp(tensor->name, "axis2") == 0) { state.axis = GGML_BACKEND_SPLIT_AXIS_2; state.ne[0] = tensor->ne[2] / 2; @@ -57,6 +68,22 @@ static ggml_backend_meta_split_state split_state( state.n_segments = 1; return state; } + if (std::strcmp(tensor->name, "zero_column_weight") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_1; + state.ne[0] = tensor->ne[1]; + state.ne[1] = 0; + state.nr[0] = 1; + state.n_segments = 1; + return state; + } + if (std::strcmp(tensor->name, "zero_row_weight") == 0) { + state.axis = GGML_BACKEND_SPLIT_AXIS_0; + state.ne[0] = tensor->ne[0]; + state.ne[1] = 0; + state.nr[0] = 1; + state.n_segments = 1; + return state; + } if (std::strcmp(tensor->name, "repeated") != 0) { state.axis = GGML_BACKEND_SPLIT_AXIS_MIRRORED; state.nr[0] = 1; @@ -82,6 +109,19 @@ static ggml_backend_meta_split_state split_state( int main() { ggml_backend_load_all(); +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wmissing-field-initializers" +#endif + const ggml_backend_meta_split_state legacy_positional = { + GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1 + }; +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + CHECK(legacy_positional.n_segments == 1); + CHECK(legacy_positional.nr[0] == 0); + const char * first_name = std::getenv("GGML_META_TEST_DEVICE_0"); const char * second_name = std::getenv("GGML_META_TEST_DEVICE_1"); first_name = first_name ? first_name : "CUDA0"; @@ -98,6 +138,11 @@ int main() { } size_t n_devices = 2; + CHECK(ggml_backend_meta_device(nullptr, 0, split_state, &n_devices) == nullptr); + std::vector too_many( + GGML_BACKEND_META_MAX_DEVICES + 1, devices[0]); + CHECK(ggml_backend_meta_device(too_many.data(), too_many.size(), + split_state, &n_devices) == nullptr); ggml_backend_dev_t meta_device = ggml_backend_meta_device(devices, n_devices, split_state, &n_devices); CHECK(meta_device); @@ -121,6 +166,9 @@ int main() { ggml_tensor * axis2 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 3, 12, 2); ggml_set_name(axis2, "axis2"); + ggml_tensor * multi_repeat = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + ggml_set_name(multi_repeat, "multi_repeat"); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); CHECK(buffer); @@ -138,10 +186,32 @@ int main() { for (size_t i = 0; i < input.size(); ++i) input[i] = (float) i + 0.25f; ggml_backend_tensor_set(repeated, input.data(), 0, input.size() * sizeof(float)); ggml_backend_tensor_set(mirrored, input.data(), 0, input.size() * sizeof(float)); + std::vector multi_repeat_input( + (size_t) ggml_nelements(multi_repeat)); + for (size_t i = 0; i < multi_repeat_input.size(); ++i) { + multi_repeat_input[i] = (float) i + 1.25f; + } + ggml_backend_tensor_set(multi_repeat, multi_repeat_input.data(), 0, + multi_repeat_input.size() * sizeof(float)); std::vector output(input.size(), 0.0f); ggml_backend_tensor_get(repeated, output.data(), 0, output.size() * sizeof(float)); CHECK(output == input); + + std::vector async_input(input.size()); + for (size_t i = 0; i < async_input.size(); ++i) { + async_input[i] = -(float) i - 0.5f; + } + ggml_backend_tensor_set_async(backend, repeated, async_input.data(), 0, + async_input.size() * sizeof(float)); + ggml_backend_synchronize(backend); + std::fill(output.begin(), output.end(), 0.0f); + ggml_backend_tensor_get_async(backend, repeated, output.data(), 0, + output.size() * sizeof(float)); + ggml_backend_synchronize(backend); + CHECK(output == async_input); + ggml_backend_tensor_set(repeated, input.data(), 0, + input.size() * sizeof(float)); std::fill(output.begin(), output.end(), 0.0f); ggml_backend_tensor_get(mirrored, output.data(), 0, output.size() * sizeof(float)); CHECK(output == input); @@ -215,6 +285,12 @@ int main() { ggml_tensor * row_weight = ggml_new_tensor_2d(weight_ctx, GGML_TYPE_F32, 8, 4); ggml_set_name(row_weight, "row_weight"); + ggml_tensor * zero_column_weight = + ggml_new_tensor_2d(weight_ctx, GGML_TYPE_F32, 4, 8); + ggml_set_name(zero_column_weight, "zero_column_weight"); + ggml_tensor * zero_row_weight = + ggml_new_tensor_2d(weight_ctx, GGML_TYPE_F32, 8, 4); + ggml_set_name(zero_row_weight, "zero_row_weight"); ggml_backend_buffer_t weight_buffer = ggml_backend_alloc_ctx_tensors(weight_ctx, backend); CHECK(weight_buffer); @@ -230,6 +306,13 @@ int main() { ggml_tensor * mat_input = ggml_new_tensor_2d(graph_ctx, GGML_TYPE_F32, 4, 2); ggml_set_name(mat_input, "mat_input"); + ggml_tensor * multi_repeat_transposed = + ggml_transpose(graph_ctx, multi_repeat); + ggml_set_name(multi_repeat_transposed, "multi_repeat_transposed"); + ggml_tensor * multi_repeat_contiguous = + ggml_cont(graph_ctx, multi_repeat_transposed); + ggml_set_name(multi_repeat_contiguous, "multi_repeat_contiguous"); + ggml_set_output(multi_repeat_contiguous); ggml_tensor * mat_hidden = ggml_mul_mat(graph_ctx, column_weight, mat_input); ggml_set_name(mat_hidden, "mat_hidden"); @@ -238,15 +321,33 @@ int main() { ggml_tensor * mat_output = ggml_scale(graph_ctx, mat_result, 1.0f); ggml_set_name(mat_output, "mat_output"); ggml_set_output(mat_output); + ggml_tensor * zero_hidden = + ggml_mul_mat(graph_ctx, zero_column_weight, mat_input); + ggml_set_name(zero_hidden, "zero_hidden"); + ggml_tensor * zero_result = + ggml_mul_mat(graph_ctx, zero_row_weight, zero_hidden); + ggml_set_name(zero_result, "zero_result"); + ggml_tensor * zero_output = ggml_scale(graph_ctx, zero_result, 1.0f); + ggml_set_name(zero_output, "zero_output"); + ggml_set_output(zero_output); ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx, graph_nodes, false); + ggml_build_forward_expand(graph, multi_repeat_contiguous); ggml_build_forward_expand(graph, mat_output); + ggml_build_forward_expand(graph, zero_output); ggml_gallocr_t graph_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); CHECK(graph_alloc); CHECK(ggml_gallocr_alloc_graph(graph_alloc, graph)); CHECK(ggml_backend_meta_simple_tensor(mat_output, 0)); CHECK(ggml_backend_meta_simple_tensor(mat_output, 1)); + ggml_tensor * zero_result_rank_1 = + ggml_backend_meta_simple_tensor(zero_result, 1); + CHECK(zero_result_rank_1); + std::vector poisoned((size_t) ggml_nelements(zero_result_rank_1), + NAN); + ggml_backend_tensor_set(zero_result_rank_1, poisoned.data(), 0, + poisoned.size() * sizeof(float)); std::vector column_data((size_t) ggml_nelements(column_weight)); std::vector row_data((size_t) ggml_nelements(row_weight)); @@ -264,13 +365,32 @@ int main() { column_data.size() * sizeof(float)); ggml_backend_tensor_set(row_weight, row_data.data(), 0, row_data.size() * sizeof(float)); + ggml_backend_tensor_set(zero_column_weight, column_data.data(), 0, + column_data.size() * sizeof(float)); + ggml_backend_tensor_set(zero_row_weight, row_data.data(), 0, + row_data.size() * sizeof(float)); ggml_backend_tensor_set(mat_input, mat_input_data.data(), 0, mat_input_data.size() * sizeof(float)); CHECK(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + std::vector multi_repeat_output( + (size_t) ggml_nelements(multi_repeat_contiguous)); + ggml_backend_tensor_get(multi_repeat_contiguous, + multi_repeat_output.data(), 0, + multi_repeat_output.size() * sizeof(float)); + for (int64_t row = 0; row < multi_repeat->ne[1]; ++row) { + for (int64_t column = 0; column < multi_repeat->ne[0]; ++column) { + CHECK(multi_repeat_output[(size_t) column * multi_repeat->ne[1] + row] == + multi_repeat_input[(size_t) row * multi_repeat->ne[0] + column]); + } + } + std::vector result((size_t) ggml_nelements(mat_output)); ggml_backend_tensor_get(mat_output, result.data(), 0, result.size() * sizeof(float)); + std::vector zero_result_host((size_t) ggml_nelements(zero_output)); + ggml_backend_tensor_get(zero_output, zero_result_host.data(), 0, + zero_result_host.size() * sizeof(float)); std::vector hidden(16, 0.0f); for (int column = 0; column < 2; ++column) { for (int row = 0; row < 8; ++row) { @@ -295,6 +415,10 @@ int main() { column, row, actual, expected); } CHECK(std::fabs(actual - expected) < 1e-5f); + const float zero_actual = + zero_result_host[(size_t) column * 4 + row]; + CHECK(std::isfinite(zero_actual)); + CHECK(std::fabs(zero_actual - expected) < 1e-5f); } } From a6eb14b8d678c23f111b7acfcfe6b51b2ea95c46 Mon Sep 17 00:00:00 2001 From: Jun Yamog Date: Sat, 18 Jul 2026 21:55:08 +1200 Subject: [PATCH 4/4] fix(cuda): validate communicator devices and preserve API --- ggml/src/ggml-cuda/ggml-cuda.cu | 26 ++++++++++++++++++++++++++ tests/CMakeLists.txt | 6 ++++++ tests/test-cuda-comm-api.cpp | 31 +++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 tests/test-cuda-comm-api.cpp diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6ea8683bce81..592b03e19166 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1295,11 +1295,23 @@ static void ggml_backend_cuda_comm_free(void * comm_ctx_v) { static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_backends) { #ifdef GGML_USE_NCCL + if (n_backends == 0 || n_backends > GGML_CUDA_MAX_DEVICES) { + return nullptr; + } + + std::array seen_devices = {}; for (size_t i = 0; i < n_backends; i++) { if (!ggml_backend_is_cuda(backends[i])) { return nullptr; } + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + if (cuda_ctx->device < 0 || cuda_ctx->device >= GGML_CUDA_MAX_DEVICES || + seen_devices[cuda_ctx->device]) { + return nullptr; + } + seen_devices[cuda_ctx->device] = true; } + ggml_backend_cuda_comm_context * ret = new ggml_backend_cuda_comm_context; std::vector dev_ids; ret->backends.reserve(n_backends); @@ -1410,6 +1422,20 @@ static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct gg #endif // GGML_USE_NCCL } +bool ggml_backend_cuda_allreduce_tensor( + ggml_backend_t * backends, + struct ggml_tensor ** tensors, + size_t n_backends) { + void * comm_ctx = ggml_backend_cuda_comm_init(backends, n_backends); + if (comm_ctx == nullptr) { + return false; + } + + const bool result = ggml_backend_cuda_comm_allreduce_tensor(comm_ctx, tensors); + ggml_backend_cuda_comm_free(comm_ctx); + return result; +} + ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split) { static std::mutex mutex; std::lock_guard lock(mutex); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 77690706b8b2..e66d9927a4f3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -241,6 +241,12 @@ llama_build_and_test(test-gguf.cpp) llama_build_and_test(test-backend-ops.cpp) llama_build_and_test(test-backend-meta.cpp) target_include_directories(test-backend-meta PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../ggml/src) +if (GGML_CUDA) + add_executable(test-cuda-comm-api test-cuda-comm-api.cpp) + target_link_libraries(test-cuda-comm-api PRIVATE ggml-cuda ggml-base) + target_include_directories(test-cuda-comm-api PRIVATE ${PROJECT_SOURCE_DIR}/ggml/include) + llama_test(test-cuda-comm-api LABEL "backend") +endif() llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") diff --git a/tests/test-cuda-comm-api.cpp b/tests/test-cuda-comm-api.cpp new file mode 100644 index 000000000000..f783d190fc18 --- /dev/null +++ b/tests/test-cuda-comm-api.cpp @@ -0,0 +1,31 @@ +#include "ggml-backend.h" +#include "ggml-cuda.h" + +#include + +int main() { + if (ggml_backend_cuda_get_device_count() == 0) { + std::puts("CUDA device unavailable; skipping communicator API test"); + return 0; + } + + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (backend == nullptr) { + std::fputs("failed to initialize CUDA backend 0\n", stderr); + return 1; + } + + ggml_backend_t backends[] = {backend, backend}; + ggml_tensor * tensors[] = {nullptr, nullptr}; + const bool result = + ggml_backend_cuda_allreduce_tensor(backends, tensors, 2); + ggml_backend_free(backend); + + if (result) { + std::fputs("duplicate CUDA devices unexpectedly initialized a communicator\n", stderr); + return 1; + } + + std::puts("CUDA communicator API compatibility test passed"); + return 0; +}