}
).set_env("LLAMA_ARG_N_GPU_LAYERS"));
add_opt(common_arg(
- {"-sm", "--split-mode"}, "{none,layer,row}",
+ {"-sm", "--split-mode"}, "{none,layer,row,tensor}",
"how to split the model across multiple GPUs, one of:\n"
"- none: use one GPU only\n"
- "- layer (default): split layers and KV across GPUs\n"
- "- row: split rows across GPUs",
+ "- layer (default): split layers and KV across GPUs (pipelined)\n"
+ "- row: split weight across GPUs by rows (parallelized)\n"
+ "- tensor: split weights and KV across GPUs (parallelized)",
[](common_params & params, const std::string & value) {
- std::string arg_next = value;
- if (arg_next == "none") {
+ if (value == "none") {
params.split_mode = LLAMA_SPLIT_MODE_NONE;
- } else if (arg_next == "layer") {
+ } else if (value == "layer") {
params.split_mode = LLAMA_SPLIT_MODE_LAYER;
- } else if (arg_next == "row") {
+ } else if (value == "row") {
params.split_mode = LLAMA_SPLIT_MODE_ROW;
+ } else if (value == "tensor") {
+ params.split_mode = LLAMA_SPLIT_MODE_TENSOR;
} else {
throw std::invalid_argument("invalid value");
}
set(GGML_VERSION_PATCH 11)
set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
+list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")
+
find_program(GIT_EXE NAMES git git.exe NO_CMAKE_FIND_ROOT_PATH)
if(GIT_EXE)
# Get current git commit hash
option(GGML_CUDA_FA "ggml: compile ggml FlashAttention CUDA kernels" ON)
option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF)
option(GGML_CUDA_GRAPHS "ggml: use CUDA graphs (llama.cpp only)" ${GGML_CUDA_GRAPHS_DEFAULT})
+option(GGML_CUDA_NCCL "ggml: use NVIDIA Collective Comm. Library" ON)
set (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING
"ggml: cuda link binary compression mode; requires cuda 12.8+")
set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size")
option(GGML_HIP "ggml: use HIP" OFF)
option(GGML_HIP_GRAPHS "ggml: use HIP graph, experimental, slow" OFF)
+option(GGML_HIP_RCCL "ggml: use ROCm Collective Comm. Library" OFF)
option(GGML_HIP_NO_VMM "ggml: do not try to use HIP VMM" ON)
option(GGML_HIP_ROCWMMA_FATTN "ggml: enable rocWMMA for FlashAttention" OFF)
option(GGML_HIP_MMQ_MFMA "ggml: enable MFMA MMA for CDNA in MMQ" ON)
--- /dev/null
+# cmake/FindNCCL.cmake
+
+# NVIDIA does not distribute CMake files with NCCl, therefore use this file to find it instead.
+
+find_path(NCCL_INCLUDE_DIR
+ NAMES nccl.h
+ HINTS ${NCCL_ROOT} $ENV{NCCL_ROOT} $ENV{CUDA_HOME} /usr/local/cuda
+ PATH_SUFFIXES include
+)
+
+find_library(NCCL_LIBRARY
+ NAMES nccl
+ HINTS ${NCCL_ROOT} $ENV{NCCL_ROOT} $ENV{CUDA_HOME} /usr/local/cuda
+ PATH_SUFFIXES lib lib64
+)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL
+ DEFAULT_MSG
+ NCCL_LIBRARY NCCL_INCLUDE_DIR
+)
+
+if(NCCL_FOUND)
+ set(NCCL_LIBRARIES ${NCCL_LIBRARY})
+ set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
+
+ if(NOT TARGET NCCL::NCCL)
+ add_library(NCCL::NCCL UNKNOWN IMPORTED)
+ set_target_properties(NCCL::NCCL PROPERTIES
+ IMPORTED_LOCATION "${NCCL_LIBRARY}"
+ INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}"
+ )
+ endif()
+endif()
+
+mark_as_advanced(NCCL_INCLUDE_DIR NCCL_LIBRARY)
GGML_API void ggml_backend_buffer_reset (ggml_backend_buffer_t buffer);
// tensor copy between different backends
- GGML_API void ggml_backend_tensor_copy(struct ggml_tensor * src, struct ggml_tensor * dst);
+ GGML_API void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst);
//
// Backend (stream)
GGML_API size_t ggml_backend_get_alignment(ggml_backend_t backend);
GGML_API size_t ggml_backend_get_max_size(ggml_backend_t backend);
- GGML_API void ggml_backend_tensor_set_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
- GGML_API void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_set_async (ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_get_async (ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+ GGML_API void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
// "offset" refers to the offset in tensor->data for setting/getting data
- GGML_API void ggml_backend_tensor_set( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
- GGML_API void ggml_backend_tensor_get(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
- GGML_API void ggml_backend_tensor_memset( struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_set ( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_get (const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ GGML_API void ggml_backend_tensor_set_2d( struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+ GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+ GGML_API void ggml_backend_tensor_memset( struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size);
GGML_API void ggml_backend_synchronize(ggml_backend_t backend);
// the copy is performed after all the currently queued operations in backend_src
// backend_dst will wait for the copy to complete before performing other operations
// automatic fallback to sync copy if async is not supported
- GGML_API void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, struct ggml_tensor * src, struct ggml_tensor * dst);
+ GGML_API void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);
GGML_API ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend);
// integrated GPU device using host memory
GGML_BACKEND_DEVICE_TYPE_IGPU,
// accelerator devices intended to be used together with the CPU backend (e.g. BLAS or AMX)
- GGML_BACKEND_DEVICE_TYPE_ACCEL
+ GGML_BACKEND_DEVICE_TYPE_ACCEL,
+ // "meta" device wrapping multiple other devices for tensor parallelism
+ GGML_BACKEND_DEVICE_TYPE_META,
};
// functionality supported by the device
// Common functions that may be obtained using ggml_backend_reg_get_proc_address
- // Split buffer type for tensor parallelism
+ // 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);
+ // 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
typedef void (*ggml_backend_set_n_threads_t)(ggml_backend_t backend, int n_threads);
// device buffer
GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device);
+// conduct allreduce operation between devices
+GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends);
+
// split tensor buffer that splits matrices by rows across multiple devices
GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split);
ggml.cpp
ggml-alloc.c
ggml-backend.cpp
+ ggml-backend-meta.cpp
ggml-opt.cpp
ggml-threading.cpp
ggml-threading.h
ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) {
size_t nbytes_total = 0;
+ if (ggml_backend_buft_is_meta(buft)) {
+ return ggml_backend_meta_alloc_ctx_tensors_from_buft(ctx, buft);
+ }
return ggml_backend_alloc_ctx_tensors_from_buft_impl(ctx, buft, &nbytes_total, /*no_alloc =*/ false);
}
void (*memset_tensor)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size);
void (*set_tensor) (ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
void (*get_tensor) (ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ // (optional) 2d data copies
+ void (*set_tensor_2d)(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+ void (*get_tensor_2d)(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+
// (optional) tensor copy: dst is in the buffer, src may be in any buffer, including buffers from a different backend (return false if not supported)
bool (*cpy_tensor) (ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst);
// clear the entire buffer
GGML_API bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer);
GGML_API void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+ //
+ // Backend (meta)
+ //
+
+ GGML_API bool ggml_backend_is_meta (ggml_backend_t backend);
+ GGML_API bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf);
+ GGML_API bool ggml_backend_buft_is_meta (ggml_backend_buffer_type_t buft);
+
+ 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);
+
+ // 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);
+
//
// Backend (stream)
//
void (*free)(ggml_backend_t backend);
// (optional) asynchronous tensor data access
- void (*set_tensor_async)(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
- void (*get_tensor_async)(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ void (*set_tensor_async) (ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+ void (*get_tensor_async) (ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+ void (*set_tensor_2d_async)(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+ void (*get_tensor_2d_async)(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
bool (*cpy_tensor_async)(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);
// (optional) complete all pending operations (required if the backend supports async operations)
--- /dev/null
+#include "ggml.h"
+#include "ggml-impl.h"
+#include "ggml-backend.h"
+#include "ggml-backend-impl.h"
+#include "ggml-alloc.h"
+#include "ggml-cpp.h"
+
+// TODO: tmp
+#include "ggml-ext.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <map>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+struct ggml_backend_meta_device;
+struct ggml_backend_meta_buffer_type;
+struct ggml_backend_meta_buffer;
+struct ggml_backend_meta;
+
+const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis) {
+ switch (split_axis) {
+ case GGML_BACKEND_SPLIT_AXIS_0:
+ return "0";
+ case GGML_BACKEND_SPLIT_AXIS_1:
+ return "1";
+ case GGML_BACKEND_SPLIT_AXIS_2:
+ return "2";
+ case GGML_BACKEND_SPLIT_AXIS_3:
+ return "3";
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED:
+ return "MIRRORED";
+ case GGML_BACKEND_SPLIT_AXIS_PARTIAL:
+ return "PARTIAL";
+ case GGML_BACKEND_SPLIT_AXIS_NONE:
+ return "NONE";
+ case GGML_BACKEND_SPLIT_AXIS_UNKNOWN:
+ return "UNKNOWN";
+ default:
+ GGML_ABORT("fatal error");
+ }
+}
+
+//
+// meta backend device
+//
+
+struct ggml_backend_meta_device_context {
+ std::vector<ggml_backend_dev_t> simple_devs;
+ ggml_backend_meta_get_split_state_t get_split_state;
+ void * get_split_state_ud;
+
+ std::string name;
+ std::string description;
+
+ ggml_backend_meta_device_context(
+ std::vector<ggml_backend_dev_t> simple_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) :
+ simple_devs(std::move(simple_devs)), get_split_state(get_split_state), get_split_state_ud(get_split_state_ud) {
+ name = std::string("Meta(");
+ description = std::string("Meta(");
+ for (size_t i = 0; i < simple_devs.size(); i++) {
+ if (i > 0) {
+ name += ",";
+ description += ",";
+ }
+ name += ggml_backend_dev_name (simple_devs[i]);
+ description += ggml_backend_dev_description(simple_devs[i]);
+ }
+ name += ")";
+ description += ")";
+ }
+
+ bool operator<(const ggml_backend_meta_device_context & other) const {
+ return std::tie(simple_devs, get_split_state, get_split_state_ud)
+ < std::tie(other.simple_devs, other.get_split_state, other.get_split_state_ud);
+ }
+};
+
+static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev);
+
+static const char * ggml_backend_meta_device_get_name(ggml_backend_dev_t dev) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+ return meta_dev_ctx->name.c_str();
+}
+
+static const char * ggml_backend_meta_device_get_description(ggml_backend_dev_t dev) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+ return meta_dev_ctx->description.c_str();
+}
+
+static void ggml_backend_meta_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+ *free = 0;
+ *total = 0;
+ for (ggml_backend_dev_t dev : meta_dev_ctx->simple_devs) {
+ size_t tmp_free, tmp_total;
+ ggml_backend_dev_memory(dev, &tmp_free, &tmp_total);
+ *free += tmp_free;
+ *total += tmp_total;
+ }
+}
+
+static enum ggml_backend_dev_type ggml_backend_meta_device_get_type(ggml_backend_dev_t dev) {
+ return GGML_BACKEND_DEVICE_TYPE_META;
+
+ GGML_UNUSED(dev);
+}
+
+static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+
+ // TODO replace placeholders
+ props->name = ggml_backend_meta_device_get_name(dev);
+ props->description = ggml_backend_meta_device_get_description(dev);
+ props->type = ggml_backend_meta_device_get_type(dev);
+ props->device_id = 0;
+
+ ggml_backend_meta_device_get_memory(dev, &props->memory_free, &props->memory_total);
+
+ props->caps = {
+ /* .async = */ true,
+ /* .host_buffer = */ false, // Not implemented.
+ /* .buffer_from_host_ptr = */ false, // Not implemented.
+ /* .events = */ false, // Not implemented.
+ };
+ for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) {
+ ggml_backend_dev_props tmp_props;
+ ggml_backend_dev_get_props(simple_dev, &tmp_props);
+ props->caps.async = props->caps.async && tmp_props.caps.async;
+ props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer;
+ props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr;
+ props->caps.events = props->caps.events && tmp_props.caps.events;
+ }
+}
+
+static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params);
+
+static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev);
+
+static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev);
+
+static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+ return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(),
+ [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); });
+}
+
+static bool ggml_backend_meta_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ ggml_backend_dev_t dev_buft = ggml_backend_buft_get_device(buft);
+ if (!ggml_backend_dev_is_meta(dev_buft)) {
+ return false;
+ }
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+ const ggml_backend_meta_device_context * meta_buft_dev_ctx = (const ggml_backend_meta_device_context *) dev_buft->context;
+ if (meta_dev_ctx->simple_devs.size() != meta_buft_dev_ctx->simple_devs.size()) {
+ return false;
+ }
+ for (size_t i = 0; i < meta_dev_ctx->simple_devs.size(); i++) {
+ if (meta_dev_ctx->simple_devs[i] != meta_buft_dev_ctx->simple_devs[i]) {
+ return false;
+ }
+ }
+ return true;
+}
+
+static const ggml_backend_device_i ggml_backend_meta_device_iface = {
+ /* .get_name = */ ggml_backend_meta_device_get_name,
+ /* .get_description = */ ggml_backend_meta_device_get_description,
+ /* .get_memory = */ ggml_backend_meta_device_get_memory,
+ /* .get_type = */ ggml_backend_meta_device_get_type,
+ /* .get_props = */ ggml_backend_meta_device_get_props,
+ /* .init_backend = */ ggml_backend_meta_device_init_backend,
+ /* .get_buffer_type = */ ggml_backend_meta_device_get_buffer_type,
+ /* .get_host_buffer_type = */ ggml_backend_meta_device_get_host_buffer_type,
+ /* .buffer_from_host_ptr = */ nullptr,
+ /* .supports_op = */ ggml_backend_meta_device_supports_op,
+ /* .supports_buft = */ ggml_backend_meta_device_supports_buft,
+ /* .offload_op = */ nullptr,
+ /* .event_new = */ nullptr,
+ /* .event_free = */ nullptr,
+ /* .event_synchronize = */ nullptr,
+};
+
+static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) {
+ return dev != nullptr && dev->iface.get_name == ggml_backend_meta_device_iface.get_name;
+}
+
+static size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context;
+ return meta_dev_ctx->simple_devs.size();
+}
+
+static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context;
+ GGML_ASSERT(index < meta_dev_ctx->simple_devs.size());
+ return meta_dev_ctx->simple_devs[index];
+}
+
+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);
+ // TODO: this is not thread-safe - needs to be fixed
+ static std::vector<std::unique_ptr<ggml_backend_meta_device_context>> ctxs;
+ static std::map<ggml_backend_meta_device_context, struct ggml_backend_device> meta_devs;
+
+ std::vector<ggml_backend_dev_t> simple_devs;
+ simple_devs.reserve(n_devs);
+ for (size_t i = 0; i < n_devs; i++) {
+ simple_devs.push_back(devs[i]);
+ }
+ ggml_backend_meta_device_context ctx(simple_devs, get_split_state, get_split_state_ud);
+
+ {
+ auto it = meta_devs.find(ctx);
+ if (it != meta_devs.end()) {
+ return &it->second;
+ }
+ }
+ ctxs.push_back(std::make_unique<ggml_backend_meta_device_context>(ctx));
+
+ struct ggml_backend_device meta_dev = {
+ /*iface =*/ ggml_backend_meta_device_iface,
+ /*reg =*/ nullptr,
+ /*ctx =*/ ctxs.back().get(),
+ };
+
+ auto result = meta_devs.emplace(*ctxs.back(), meta_dev);
+ return &result.first->second;
+}
+
+//
+// meta backend buffer type
+//
+
+struct ggml_backend_meta_buffer_type_context {
+ std::vector<ggml_backend_buffer_type_t> simple_bufts;
+
+ std::string name;
+
+ ggml_backend_meta_buffer_type_context(std::vector<ggml_backend_buffer_type_t> simple_bufts) : simple_bufts(std::move(simple_bufts)) {
+ name = "Meta(";
+ for (size_t i = 0; i < simple_bufts.size(); i++) {
+ if (i > 0) {
+ name += ",";
+ }
+ name += ggml_backend_buft_name(simple_bufts[i]);
+ }
+ name += ")";
+ }
+
+ bool operator<(const ggml_backend_meta_buffer_type_context & other) const {
+ return simple_bufts < other.simple_bufts;
+ }
+};
+
+static size_t ggml_backend_meta_buft_n_bufts(ggml_backend_buffer_type_t meta_buft) {
+ GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft));
+ const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context;
+ return meta_buft_ctx->simple_bufts.size();
+}
+
+static const char * ggml_backend_meta_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
+ GGML_ASSERT(ggml_backend_buft_is_meta(buft));
+ const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) buft->context;
+ return meta_buft_ctx->name.c_str();
+}
+
+static ggml_backend_buffer_type_t ggml_backend_meta_buft_simple_buft(ggml_backend_buffer_type_t meta_buft, size_t index) {
+ GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft));
+ const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context;
+ GGML_ASSERT(index < meta_buft_ctx->simple_bufts.size());
+ return meta_buft_ctx->simple_bufts[index];
+}
+
+static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size);
+
+static size_t ggml_backend_meta_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
+ const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);
+ size_t max_alignment = 1;
+ for (size_t i = 0; i < n_simple_bufts; i++) {
+ const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_meta_buft_simple_buft(buft, i));
+ max_alignment = std::max(max_alignment, alignment);
+ GGML_ASSERT(max_alignment % alignment == 0);
+ }
+ return max_alignment;
+}
+
+static size_t ggml_backend_meta_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
+ const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);
+ size_t max_size = SIZE_MAX;
+ for (size_t i = 0; i < n_simple_bufts; i++) {
+ max_size = std::min(max_size, ggml_backend_buft_get_max_size(ggml_backend_meta_buft_simple_buft(buft, i)));
+ }
+ return max_size;
+}
+
+static size_t ggml_backend_meta_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
+ const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);
+ size_t max_alloc_size = 0;
+ for (size_t i = 0; i < n_simple_bufts; i++) {
+ const size_t alloc_size = ggml_backend_buft_get_alloc_size(ggml_backend_meta_buft_simple_buft(buft, i), tensor);
+ max_alloc_size = std::max(max_alloc_size, alloc_size);
+ }
+ return max_alloc_size;
+}
+
+static bool ggml_backend_meta_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
+ const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);
+ for (size_t i = 0; i < n_simple_bufts; i++) {
+ if (!ggml_backend_buft_is_host(ggml_backend_meta_buft_simple_buft(buft, i))) {
+ return false;
+ }
+ }
+ return true;
+}
+
+static const struct ggml_backend_buffer_type_i ggml_backend_meta_buffer_type_iface = {
+ /* .get_name = */ ggml_backend_meta_buffer_type_get_name,
+ /* .alloc_buffer = */ ggml_backend_meta_buffer_type_alloc_buffer,
+ /* .get_alignment = */ ggml_backend_meta_buffer_type_get_alignment,
+ /* .get_max_size = */ ggml_backend_meta_buffer_type_get_max_size,
+ /* .get_alloc_size = */ ggml_backend_meta_buffer_type_get_alloc_size,
+ /* .is_host = */ ggml_backend_meta_buffer_type_is_host,
+};
+
+bool ggml_backend_buft_is_meta(ggml_backend_buffer_type_t buft) {
+ return buft != nullptr && buft->iface.get_name == ggml_backend_meta_buffer_type_iface.get_name;
+}
+
+static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev) {
+ static std::map<ggml_backend_dev_t, struct ggml_backend_buffer_type> meta_bufts;
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ {
+ auto it = meta_bufts.find(dev);
+ if (it != meta_bufts.end()) {
+ return &it->second;
+ }
+ }
+
+ const size_t n_devs = ggml_backend_meta_dev_n_devs(dev);
+ std::vector<ggml_backend_buffer_type_t> simple_bufts;
+ simple_bufts.reserve(n_devs);
+ for (size_t i = 0; i < n_devs; i++) {
+ simple_bufts.push_back(ggml_backend_dev_buffer_type(ggml_backend_meta_dev_simple_dev(dev, i)));
+ }
+ ggml_backend_meta_buffer_type_context * buft_ctx = new ggml_backend_meta_buffer_type_context(simple_bufts);
+
+ struct ggml_backend_buffer_type meta_buft = {
+ /*iface =*/ ggml_backend_meta_buffer_type_iface,
+ /*device =*/ dev,
+ /*ctx =*/ buft_ctx,
+ };
+ auto result = meta_bufts.emplace(dev, meta_buft);
+ return &result.first->second;
+}
+
+static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev) {
+ GGML_ASSERT(ggml_backend_dev_is_meta(dev));
+ const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context;
+
+ ggml_backend_buffer_type_t host_buft = nullptr;
+ for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) {
+ ggml_backend_buffer_type_t simple_host_buft = ggml_backend_dev_host_buffer_type(simple_dev);
+ if (simple_host_buft == nullptr) {
+ return nullptr;
+ }
+ if (host_buft == nullptr) {
+ host_buft = simple_host_buft;
+ } else if (host_buft != simple_host_buft) {
+ // if different simple devices have different host buffer types,
+ // we cannot provide a single host buffer type for the meta device
+ return nullptr;
+ }
+ }
+ return host_buft;
+}
+
+//
+// meta backend buffer
+//
+
+struct ggml_backend_meta_buffer_context {
+ static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding);
+
+ std::map<std::pair<const ggml_tensor *, bool>, std::pair<ggml_backend_meta_split_state, char[nbtc]>> split_state_cache;
+ std::map< const ggml_tensor *, std::vector<ggml_tensor *>> 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<buffer_config> buf_configs;
+
+ int debug;
+
+ ggml_backend_meta_buffer_context() {
+ const char * GGML_META_DEBUG = getenv("GGML_META_DEBUG");
+ debug = GGML_META_DEBUG ? atoi(GGML_META_DEBUG) : 0;
+ }
+};
+
+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();
+}
+
+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;
+}
+
+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());
+
+ auto it = buf_ctx->simple_tensors.find(tensor);
+ if (it == buf_ctx->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) {
+ 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;
+
+ auto split_states_equal = [&](const ggml_backend_meta_split_state & a, const ggml_backend_meta_split_state & b) -> bool {
+ if (a.axis != b.axis) {
+ return false;
+ }
+ 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];
+ }
+ int64_t sum_b = 0;
+ for (size_t s = 0; s < b.n_segments; s++) {
+ sum_b += b.ne[s*n_bufs + j];
+ }
+ if (sum_a != sum_b) {
+ return false;
+ }
+ }
+ return true;
+ };
+
+ auto handle_generic = [&](const std::vector<ggml_backend_meta_split_state> & src_ss, bool scalar_only) -> ggml_backend_meta_split_state {
+ ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1};
+ for (size_t i = 0; i < GGML_MAX_SRC; i++) {
+ if (tensor->src[i] == nullptr || tensor->src[i] == tensor) {
+ continue;
+ }
+ 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};
+ break;
+ }
+ }
+ if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) {
+ ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ }
+ if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) {
+ ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ }
+ GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN);
+ return ret;
+ };
+
+ // Some ops process data on a per-row bases:
+ auto handle_per_row = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0);
+ return src_ss[0];
+ };
+
+ // Some ops broadcast the src1 data across src0:
+ auto handle_bin_bcast = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS &&
+ tensor->src[1]->ne[src_ss[0].axis] == 1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) {
+ return src_ss[0];
+ }
+ if (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[0].axis == src_ss[1].axis ||
+ (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) {
+ return src_ss[0]; // GGML_OP_ADD_ID
+ }
+ GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ return handle_generic(src_ss, /*scalar_only =*/ false);
+ };
+
+ auto handle_concat = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0));
+ if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) {
+ GGML_ASSERT(concat_axis != src_ss[1].axis);
+ return src_ss[1];
+ }
+ if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) {
+ GGML_ASSERT(concat_axis != src_ss[0].axis);
+ return src_ss[0];
+ }
+ if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) {
+ return src_ss[0];
+ }
+ return handle_generic(src_ss, /*scalar_only =*/ true);
+ };
+
+ auto handle_mul_mat = [&](const std::vector<ggml_backend_meta_split_state> & 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};
+ }
+ 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.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;
+ }
+ 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};
+ }
+ GGML_ABORT("fatal error");
+ //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ };
+
+ auto handle_cpy = [&](const std::vector<ggml_backend_meta_split_state> & 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);
+ };
+
+ auto handle_reshape = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ switch (src_ss[0].axis) {
+ case GGML_BACKEND_SPLIT_AXIS_0:
+ 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};
+ }
+ std::vector<int64_t> 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_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]) {
+ GGML_ASSERT(dim + 1 < GGML_MAX_DIMS);
+ return {ggml_backend_meta_split_axis(dim + 1), {0}, 1};
+ }
+ base_ne_out = base_ne_out_next;
+ }
+ GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op));
+ }
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED:
+ case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
+ return src_ss[0];
+ }
+ default: {
+ GGML_ABORT("fatal error");
+ //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ }
+ }
+ };
+
+ auto handle_view = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) {
+ return handle_reshape(src_ss);
+ }
+ const int axis = src_ss[0].axis;
+ {
+ bool all_strides_the_same = true;
+ for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+ if (tensor->ne[dim] == 1 && tensor->src[0]->ne[dim] == 1) {
+ continue;
+ }
+ if (tensor->nb[dim] != tensor->src[0]->nb[dim]) {
+ all_strides_the_same = false;
+ break;
+ }
+ }
+ if (all_strides_the_same) {
+ return src_ss[0];
+ }
+ }
+ 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};
+ }
+ }
+ GGML_ABORT("fatal error");
+ }
+ if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
+ return src_ss[0];
+ }
+ GGML_ABORT("view of permuted tensor not implemented");
+ //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ };
+
+ auto handle_permute = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ switch (src_ss[0].axis) {
+ case GGML_BACKEND_SPLIT_AXIS_0:
+ 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};
+ }
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED:
+ case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
+ return src_ss[0];
+ }
+ default: {
+ GGML_ABORT("fatal error");
+ //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ }
+ }
+ };
+
+ auto handle_transpose = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ 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};
+ }
+ case GGML_BACKEND_SPLIT_AXIS_2:
+ case GGML_BACKEND_SPLIT_AXIS_3:
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED:
+ case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
+ return src_ss[0];
+ }
+ default: {
+ GGML_ABORT("fatal error");
+ //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ }
+ }
+ };
+
+ auto handle_get_rows = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) {
+ return src_ss[0];
+ }
+ return handle_generic(src_ss, /*scalar_only =*/ true);
+ };
+
+ auto handle_set_rows = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1);
+ GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2]));
+ return src_ss[0];
+ };
+
+ auto handle_rope = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ return src_ss[0];
+ };
+
+ auto handle_pad = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) {
+ GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0);
+ GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0);
+ }
+ return src_ss[0];
+ };
+
+ auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2);
+ GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2);
+ 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};
+ };
+
+ auto handle_ssm_conv = [&](const std::vector<ggml_backend_meta_split_state> & 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};
+ }
+ if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) {
+ return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1};
+ }
+ }
+ return handle_generic(src_ss, /*scalar_only =*/ false);
+ };
+
+ auto handle_gated_delta_net = [&](const std::vector<ggml_backend_meta_split_state> & 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) {
+ return src_ss[0];
+ }
+ GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1);
+ GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1);
+ 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};
+ };
+
+ auto calculate_split_state = [&]() -> ggml_backend_meta_split_state {
+ if (ggml_nelements(tensor) == 0) {
+ return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 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);
+ 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];
+ }
+ GGML_ASSERT(ne_sum == tensor->ne[ret.axis]);
+ }
+ return ret;
+ }
+
+ std::vector<ggml_backend_meta_split_state> src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 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};
+ continue;
+ }
+ src_ss[i] = ggml_backend_meta_get_split_state(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};
+ } break;
+ case GGML_OP_DUP: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_ADD:
+ case GGML_OP_ADD_ID: {
+ split_state = handle_bin_bcast(src_ss);
+ } break;
+ case GGML_OP_ADD1:
+ case GGML_OP_ACC: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SUB:
+ case GGML_OP_MUL:
+ case GGML_OP_DIV: {
+ split_state = handle_bin_bcast(src_ss);
+ } break;
+ case GGML_OP_SQR:
+ case GGML_OP_SQRT:
+ case GGML_OP_LOG:
+ case GGML_OP_SIN:
+ case GGML_OP_COS: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_SUM: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SUM_ROWS:
+ case GGML_OP_CUMSUM:
+ case GGML_OP_MEAN:
+ case GGML_OP_ARGMAX:
+ case GGML_OP_COUNT_EQUAL: {
+ split_state = handle_per_row(src_ss);
+ } break;
+ case GGML_OP_REPEAT:
+ case GGML_OP_REPEAT_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_CONCAT: {
+ split_state = handle_concat(src_ss);
+ } break;
+ case GGML_OP_SILU_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_NORM:
+ case GGML_OP_RMS_NORM:
+ case GGML_OP_RMS_NORM_BACK:
+ case GGML_OP_GROUP_NORM:
+ case GGML_OP_L2_NORM: {
+ split_state = handle_per_row(src_ss);
+ } break;
+ case GGML_OP_MUL_MAT:
+ case GGML_OP_MUL_MAT_ID: {
+ split_state = handle_mul_mat(src_ss);
+ } break;
+ case GGML_OP_OUT_PROD: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SCALE: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_SET: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_CPY: {
+ split_state = handle_cpy(src_ss);
+ } break;
+ case GGML_OP_CONT:
+ case GGML_OP_RESHAPE: {
+ split_state = handle_reshape(src_ss);
+ } break;
+ case GGML_OP_VIEW: {
+ split_state = handle_view(src_ss);
+ } break;
+ case GGML_OP_PERMUTE: {
+ split_state = handle_permute(src_ss);
+ } break;
+ case GGML_OP_TRANSPOSE: {
+ split_state = handle_transpose(src_ss);
+ } break;
+ case GGML_OP_GET_ROWS: {
+ split_state = handle_get_rows(src_ss);
+ } break;
+ case GGML_OP_GET_ROWS_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SET_ROWS: {
+ split_state = handle_set_rows(src_ss);
+ } break;
+ case GGML_OP_DIAG:
+ case GGML_OP_DIAG_MASK_INF:
+ case GGML_OP_DIAG_MASK_ZERO: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SOFT_MAX:
+ case GGML_OP_SOFT_MAX_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_ROPE: {
+ split_state = handle_rope(src_ss);
+ } break;
+ case GGML_OP_ROPE_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_CLAMP: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_CONV_TRANSPOSE_1D:
+ case GGML_OP_IM2COL:
+ case GGML_OP_IM2COL_BACK:
+ case GGML_OP_IM2COL_3D:
+ case GGML_OP_CONV_2D:
+ case GGML_OP_CONV_3D:
+ case GGML_OP_CONV_2D_DW:
+ case GGML_OP_CONV_TRANSPOSE_2D:
+ case GGML_OP_POOL_1D:
+ case GGML_OP_POOL_2D:
+ case GGML_OP_POOL_2D_BACK:
+ case GGML_OP_UPSCALE: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_PAD: {
+ split_state = handle_pad(src_ss);
+ } break;
+ case GGML_OP_PAD_REFLECT_1D:
+ case GGML_OP_ROLL:
+ case GGML_OP_ARANGE:
+ case GGML_OP_TIMESTEP_EMBEDDING: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_ARGSORT:
+ case GGML_OP_TOP_K: {
+ split_state = handle_per_row(src_ss);
+ } break;
+ case GGML_OP_LEAKY_RELU: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_TRI: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_FILL: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_FLASH_ATTN_EXT: {
+ split_state = handle_flash_attn_ext(src_ss);
+ } break;
+ case GGML_OP_FLASH_ATTN_BACK: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_SSM_CONV: {
+ split_state = handle_ssm_conv(src_ss);
+ } break;
+ case GGML_OP_SSM_SCAN:
+ case GGML_OP_WIN_PART:
+ case GGML_OP_WIN_UNPART:
+ case GGML_OP_GET_REL_POS:
+ case GGML_OP_ADD_REL_POS:
+ case GGML_OP_RWKV_WKV6:
+ case GGML_OP_GATED_LINEAR_ATTN:
+ case GGML_OP_RWKV_WKV7:
+ case GGML_OP_SOLVE_TRI: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_GATED_DELTA_NET: {
+ split_state = handle_gated_delta_net(src_ss);
+ } break;
+ case GGML_OP_UNARY: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ case GGML_OP_MAP_CUSTOM1:
+ case GGML_OP_MAP_CUSTOM2:
+ case GGML_OP_MAP_CUSTOM3:
+ case GGML_OP_CUSTOM: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+ } break;
+ case GGML_OP_CROSS_ENTROPY_LOSS:
+ case GGML_OP_CROSS_ENTROPY_LOSS_BACK: {
+ split_state = handle_per_row(src_ss);
+ } break;
+ case GGML_OP_OPT_STEP_ADAMW:
+ case GGML_OP_OPT_STEP_SGD:
+ case GGML_OP_GLU: {
+ split_state = handle_generic(src_ss, /*scalar_only =*/ false);
+ } break;
+ default: {
+ GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op));
+ split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1};
+ } break;
+ }
+ if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) {
+ bool first_src_split_by_axis = true;
+ const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer);
+
+ for (size_t i = 0; i < GGML_MAX_SRC; i++) {
+ if (tensor->src[i] == nullptr || src_ss[i].axis < 0 || src_ss[i].axis >= GGML_MAX_DIMS) {
+ continue;
+ }
+ if (first_src_split_by_axis) {
+ for (size_t j = 0; j < n_bufs; j++) {
+ // Take over ratio from src:
+ for (size_t s = 0; s < src_ss[i].n_segments; s++) {
+ 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] *= 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];
+ }
+ }
+ } else {
+ for (size_t j = 0; j < n_bufs; j++) {
+ 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];
+ }
+ // 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]);
+ }
+ }
+ first_src_split_by_axis = false;
+ }
+ GGML_ASSERT(!first_src_split_by_axis);
+ }
+ return split_state;
+ };
+
+ const std::pair key = std::make_pair(tensor, assume_sync);
+ auto it = buf_ctx->split_state_cache.find(key);
+ if (it != buf_ctx->split_state_cache.end() && memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) {
+ buf_ctx->split_state_cache.clear();
+ it = buf_ctx->split_state_cache.end();
+ }
+
+ if (it == buf_ctx->split_state_cache.end()) {
+ 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) {
+ std::string srcs_info;
+ for (size_t i = 0; i < GGML_MAX_SRC; i++) {
+ if (tensor->src[i] == nullptr) {
+ continue;
+ }
+ if (!srcs_info.empty()) {
+ srcs_info += ", ";
+ }
+ const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], 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]);
+ }
+ 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 += ", ";
+ }
+ ne_info += std::to_string(buf_ctx->split_state_cache[key].first.ne[j]);
+ }
+ 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());
+ }
+ }
+
+ ggml_backend_meta_split_state ret = buf_ctx->split_state_cache[key].first;
+ GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_NONE);
+#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];
+ }
+ assert(ne_ret == tensor->ne[int(ret.axis)]);
+ }
+#endif // NDEBUG
+ return ret;
+}
+
+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);
+
+ const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(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);
+
+ int split_dim = split_state.axis;
+ int64_t ne[GGML_MAX_DIMS];
+ size_t nb[GGML_MAX_DIMS];
+ for (size_t k = 0; k < GGML_MAX_DIMS; k++) {
+ ne[k] = tensor->ne[k];
+ nb[k] = tensor->nb[k];
+ }
+
+ std::vector<ggml_tensor *> 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;
+
+ 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];
+ }
+ for (int i = 0; i < GGML_MAX_DIMS; i++) {
+ if (tensor->nb[i] > tensor->nb[split_dim]) {
+ nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim];
+ }
+ }
+ }
+
+ ggml_tensor * t_ij = ggml_new_tensor(simple_ctx, tensor->type, GGML_MAX_DIMS, ne);
+ t_ij->op = tensor->op;
+ for (int i = 0; i < GGML_MAX_DIMS; i++) {
+ t_ij->nb[i] = nb[i];
+ }
+ t_ij->flags = tensor->flags;
+ memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params));
+ ggml_set_name(t_ij, tensor->name);
+ t_ij->buffer = simple_buf;
+ t_ij->view_src = tensor->view_src;
+ t_ij->view_offs = tensor->view_offs;
+ 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);
+ 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);
+
+ // The offset can be internal to the data split, in those cases the view offset should not be scaled.
+ // If however, the offset is larger than the data split then it needs to be scaled proportionally.
+ bool split_internal_offset = t_ij->view_offs <= tensor->view_src->nb[split_dim_view_src];
+ for (int i = 0; i < GGML_MAX_DIMS; i++) {
+ const size_t dim_size = tensor->ne[i] * tensor->nb[i];
+ if (tensor->view_offs <= dim_size && dim_size < tensor->nb[split_dim]) {
+ split_internal_offset = true;
+ break;
+ }
+ }
+ if (!split_internal_offset) {
+ t_ij->view_offs = t_ij->view_offs * ne[split_dim]/tensor->ne[split_dim];
+ }
+ }
+ }
+ if (t_ij->view_src != nullptr) {
+ 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));
+ }
+ t_ij->extra = tensor->extra;
+ for (int i = 0; i < GGML_MAX_SRC; i++) {
+ t_ij->src[i] = tensor->src[i];
+ if (tensor->src[i] == tensor) {
+ t_ij->src[i] = t_ij;
+ } else if (t_ij->src[i] != nullptr && ggml_backend_buffer_is_meta(t_ij->src[i]->buffer)) {
+ t_ij->src[i] = ggml_backend_meta_buffer_simple_tensor(tensor->src[i], j);
+ }
+ }
+
+ simple_tensors.push_back(t_ij);
+ }
+ buf_ctx->simple_tensors[tensor] = simple_tensors;
+
+ return GGML_STATUS_SUCCESS;
+}
+
+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) {
+ 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) {
+ 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);
+ size_t offset_data = 0;
+ std::vector<size_t> simple_offsets(n_bufs, 0);
+ if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) {
+ GGML_ASSERT(tensor->ne[2] == 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;
+ }
+ }
+ GGML_ASSERT(offset_data*tensor->ne[1] == size);
+ return;
+ }
+ GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_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;
+ }
+ }
+ GGML_ASSERT(offset_data*tensor->ne[2] == 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++) {
+ 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;
+ }
+ GGML_ASSERT(offset_j == chunk_size_full);
+ } break;
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED: {
+ for (size_t j = 0; j < n_bufs; j++) {
+ ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+ ggml_backend_tensor_set(simple_tensor, data, offset, size);
+ }
+ } break;
+ case GGML_BACKEND_SPLIT_AXIS_PARTIAL: {
+ GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+ const int64_t ne = ggml_nelements(tensor);
+ std::vector<float> tmp;
+ tmp.reserve(ne);
+ for (int64_t i = 0; i < ne; i++) {
+ tmp.push_back(((const float *) data)[i] / n_bufs);
+ }
+ for (size_t j = 0; j < n_bufs; j++) {
+ ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
+ ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size);
+ }
+ } break;
+ default: {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+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) {
+ 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);
+
+ 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;
+ }
+ GGML_ASSERT(offset_j == chunk_size_full);
+ } break;
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED: {
+ // TODO other simple backend may be better
+ const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0);
+ ggml_backend_tensor_get(simple_tensor, data, offset, size);
+ } break;
+ default: {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+static void ggml_backend_meta_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
+ const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer);
+ for (size_t i = 0; i < n_buffers; i++) {
+ ggml_backend_buffer_clear(ggml_backend_meta_buffer_simple_buffer(buffer, i), value);
+ }
+}
+
+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_backend_buffer_reset(ggml_backend_meta_buffer_simple_buffer(buffer, i));
+ }
+}
+
+static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = {
+ /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer,
+ /* .get_base = */ ggml_backend_meta_buffer_get_base,
+ /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor,
+ /* .memset_tensor = */ nullptr, // TODO implement
+ /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor,
+ /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor,
+ /* .set_tensor_2d = */ nullptr,
+ /* .get_tensor_2d = */ nullptr,
+ /* .cpy_tensor = */ nullptr,
+ /* .clear = */ ggml_backend_meta_buffer_clear,
+ /* .reset = */ ggml_backend_meta_buffer_reset,
+};
+
+bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {
+ return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer;
+}
+
+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
+ /*.mem_buffer =*/ nullptr,
+ /*.no_alloc =*/ true,
+ };
+
+ 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);
+ 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);
+ }
+
+ return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size);
+}
+
+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
+ /*.mem_buffer =*/ nullptr,
+ /*.no_alloc =*/ true,
+ };
+
+ 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);
+ }
+
+ 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);
+ 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));
+ }
+ return meta_buf;
+}
+
+//
+// meta backend
+//
+
+static ggml_guid_t ggml_backend_meta_guid() {
+ static ggml_guid guid = {0xf1, 0x0e, 0x34, 0xcf, 0x9c, 0x6f, 0x43, 0xcb, 0x96, 0x92, 0xbe, 0x8e, 0xbb, 0x71, 0x3f, 0xda};
+ return &guid;
+}
+
+struct ggml_backend_meta_context {
+ struct cgraph_config {
+ ggml_cgraph * cgraph_main = nullptr;
+ int offset = 0; // Node offset vs. original graph
+
+ std::vector<ggml_cgraph *> cgraphs_aux;
+ };
+ struct backend_config {
+ ggml_backend_t backend;
+
+ std::vector<cgraph_config> cgraphs;
+ std::vector<ggml_tensor *> nodes;
+ ggml_backend_buffer_ptr buf;
+
+ backend_config(ggml_backend_t backend) : backend(backend) {}
+ };
+ std::string name;
+ std::vector<backend_config> backend_configs;
+ ggml_context_ptr ctx;
+ std::vector<ggml_cgraph *> cgraphs_aux;
+ std::vector<ggml_tensor *> nodes_aux;
+ int max_nnodes = 0;
+ size_t max_tmp_size = 0;
+ size_t max_subgraphs = 0;
+
+ 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);
+ name = "Meta(";
+ backend_configs.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));
+ }
+ name += ")";
+ }
+
+ ~ggml_backend_meta_context() {
+ 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) {
+ GGML_ASSERT(ggml_backend_is_meta(backend));
+ const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) backend->context;
+ return backend_ctx->name.c_str();
+}
+
+static void ggml_backend_meta_free(ggml_backend_t backend) {
+ GGML_ASSERT(ggml_backend_is_meta(backend));
+ ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context;
+ delete backend_ctx;
+ delete 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);
+
+ 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_backends; j++){
+ 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];
+ 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;
+ }
+ GGML_ASSERT(offset_j == chunk_size_full);
+ } break;
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED: {
+ for (size_t j = 0; j < n_backends; j++) {
+ ggml_backend_tensor_set_async(
+ ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), data, offset, size);
+ }
+ } break;
+ default: {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+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);
+
+ 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_backends; j++){
+ 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];
+ 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;
+ }
+ GGML_ASSERT(offset_j == chunk_size_full);
+ } break;
+ case GGML_BACKEND_SPLIT_AXIS_MIRRORED: {
+ // TODO other simple backend may be better
+ ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, 0);
+ const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0);
+ ggml_backend_tensor_get_async(simple_backend, simple_tensor, data, offset, size);
+ } break;
+ default: {
+ GGML_ABORT("fatal error");
+ }
+ }
+}
+
+static void ggml_backend_meta_synchronize(ggml_backend_t backend) {
+ const size_t n_backends = ggml_backend_meta_n_backends(backend);
+ for (size_t i = 0; i < n_backends; i++) {
+ ggml_backend_synchronize(ggml_backend_meta_simple_backend(backend, i));
+ }
+}
+
+static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
+ GGML_ASSERT(cgraph->grads == nullptr);
+ const size_t n_backends = ggml_backend_meta_n_backends(backend);
+ ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context;
+
+ bool max_nnodes_raised = false;
+ if (cgraph->n_nodes > backend_ctx->max_nnodes) {
+ for (size_t j = 0; j < n_backends; j++) {
+ auto & bcj = backend_ctx->backend_configs[j];
+ bcj.nodes.resize(cgraph->n_nodes);
+ bcj.cgraphs.resize(cgraph->n_nodes);
+ }
+ backend_ctx->max_nnodes = cgraph->n_nodes;
+ max_nnodes_raised = true;
+ }
+ 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]);
+ }
+ }
+
+ 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);
+ }
+ }
+ if (id + 1 >= cgraph->n_nodes) {
+ return idr;
+ }
+ {
+ 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;
+ }
+ 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;
+ }
+
+ i = get_i_delayed(i);
+
+ 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);
+ }
+
+ 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];
+ }
+ }
+ }
+
+ 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));
+ ret->op = GGML_OP_NONE;
+ ret->type = t->type;
+ for (size_t k = 0; k < GGML_MAX_DIMS; k++) {
+ ret->ne[k] = t->ne[k];
+ ret->nb[k] = t->nb[k];
+ }
+ 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<ggml_cgraph *> step_cgraphs(n_backends, nullptr);
+
+ for (size_t offset_j = 1; offset_j < n_backends; offset_j *= 2) {
+ std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr);
+
+ for (size_t j = 0; j < n_backends; j++) {
+ const size_t j_other = j ^ offset_j;
+ if (j_other > j) {
+ continue;
+ }
+
+ 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++) {
+ if (step_cgraphs[j] == nullptr) {
+ continue;
+ }
+ auto & bcj = backend_ctx->backend_configs[j];
+ const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]);
+ if (status != GGML_STATUS_SUCCESS) {
+ return status;
+ }
+ }
+ }
+ return GGML_STATUS_SUCCESS;
+ };
+
+
+ for (size_t i = 0; i < 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);
+ if (status != GGML_STATUS_SUCCESS) {
+ return status;
+ }
+ }
+
+ if (n_backends > 1 && i < 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<ggml_backend_t> backends;
+ backends.reserve(n_backends);
+ std::vector<ggml_tensor *> 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);
+ }
+
+ if (!backend_allreduce_success) {
+ const ggml_status status = allreduce_fallback(i);
+ if (status != GGML_STATUS_SUCCESS) {
+ return status;
+ }
+ }
+ }
+ }
+ return GGML_STATUS_SUCCESS;
+}
+
+static const ggml_backend_i ggml_backend_meta_i = {
+ /* .get_name = */ ggml_backend_meta_get_name,
+ /* .free = */ ggml_backend_meta_free,
+ /* .set_tensor_async = */ ggml_backend_meta_set_tensor_async,
+ /* .get_tensor_async = */ ggml_backend_meta_get_tensor_async,
+ /* .get_tensor_2d_async = */ nullptr,
+ /* .set_tensor_2d_async = */ nullptr,
+ /* .cpy_tensor_async = */ nullptr,
+ /* .synchronize = */ ggml_backend_meta_synchronize,
+ /* .graph_plan_create = */ nullptr,
+ /* .graph_plan_free = */ nullptr,
+ /* .graph_plan_update = */ nullptr,
+ /* .graph_plan_compute = */ nullptr,
+ /* .graph_compute = */ ggml_backend_meta_graph_compute,
+ /* .event_record = */ nullptr,
+ /* .event_wait = */ nullptr,
+ /* .graph_optimize = */ nullptr,
+};
+
+bool ggml_backend_is_meta(ggml_backend_t backend) {
+ return backend != nullptr && backend->iface.get_name == ggml_backend_meta_i.get_name;
+}
+
+static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) {
+ ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params);
+
+ ggml_backend_t backend = new struct ggml_backend;
+ backend->guid = ggml_backend_meta_guid();
+ backend->iface = ggml_backend_meta_i;
+ backend->device = dev;
+ backend->context = backend_ctx;
+ return backend;
+}
+
+size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) {
+ GGML_ASSERT(ggml_backend_is_meta(meta_backend));
+ const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context;
+ return backend_ctx->backend_configs.size();
+}
+
+ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) {
+ GGML_ASSERT(ggml_backend_is_meta(meta_backend));
+ const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context;
+ return backend_ctx->backend_configs[index].backend;
+}
+
void * ggml_backend_buffer_get_base(ggml_backend_buffer_t buffer) {
GGML_ASSERT(buffer);
// get_base is optional if the buffer is zero-sized
- if (buffer->size == 0) {
+ if (!ggml_backend_buffer_is_meta(buffer) && buffer->size == 0) {
return NULL;
}
}
}
+void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size,
+ size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ GGML_ASSERT(backend);
+ GGML_ASSERT(tensor);
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+
+ if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) {
+ for (size_t i = 0; i < n_copies; i++) {
+ ggml_backend_tensor_set_async(backend, tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size);
+ }
+ return;
+ }
+ if (size == 0) {
+ return;
+ }
+
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+ GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds");
+ backend->iface.set_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data);
+}
+
+void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size,
+ size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ GGML_ASSERT(backend);
+ GGML_ASSERT(tensor);
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+
+ if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) {
+ for (size_t i = 0; i < n_copies; i++) {
+ ggml_backend_tensor_get_async(backend, tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size);
+ }
+ return;
+ }
+ if (size == 0) {
+ return;
+ }
+
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+ GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds");
+ backend->iface.get_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data);
+}
+
void ggml_backend_tensor_set(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
GGML_ASSERT(tensor);
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+ GGML_ASSERT(buf != NULL && "tensor buffer not set");
if (size == 0) {
return;
}
- GGML_ASSERT(buf != NULL && "tensor buffer not set");
GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds");
void ggml_backend_tensor_get(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) {
GGML_ASSERT(tensor);
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+ GGML_ASSERT(buf != NULL && "tensor buffer not set");
if (size == 0) {
return;
}
- GGML_ASSERT(buf != NULL && "tensor buffer not set");
GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds");
buf->iface.get_tensor(buf, tensor, data, offset, size);
}
+void ggml_backend_tensor_set_2d(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size,
+ size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ GGML_ASSERT(tensor);
+ ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+ GGML_ASSERT(buf != NULL && "tensor buffer not set");
+
+ if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) {
+ for (size_t i = 0; i < n_copies; i++) {
+ ggml_backend_tensor_set(tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size);
+ }
+ return;
+ }
+ if (size == 0) {
+ return;
+ }
+
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+ GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds");
+
+ buf->iface.set_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data);
+}
+
+void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size,
+ size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ GGML_ASSERT(tensor);
+ ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+ GGML_ASSERT(buf != NULL && "tensor buffer not set");
+
+ if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) {
+ for (size_t i = 0; i < n_copies; i++) {
+ ggml_backend_tensor_get(tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size);
+ }
+ return;
+ }
+ if (size == 0) {
+ return;
+ }
+
+ GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
+ GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor read out of bounds");
+
+ buf->iface.get_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data);
+}
+
void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
GGML_ASSERT(tensor);
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
// backend copy
-void ggml_backend_tensor_copy(struct ggml_tensor * src, struct ggml_tensor * dst) {
+void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst) {
GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
if (src == dst) {
} else if (!ggml_backend_buffer_copy_tensor(src, dst)) {
#ifndef NDEBUG
GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer));
-#endif
+#endif // NDEBUG
size_t nbytes = ggml_nbytes(src);
void * data = malloc(nbytes);
ggml_backend_tensor_get(src, data, 0, nbytes);
}
}
-void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, struct ggml_tensor * src, struct ggml_tensor * dst) {
+void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) {
GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
if (src == dst) {
}
void ggml_backend_dev_get_props(ggml_backend_dev_t device, struct ggml_backend_dev_props * props) {
+ GGML_ASSERT(device);
memset(props, 0, sizeof(*props));
device->iface.get_props(device, props);
}
/* .memset_tensor = */ NULL,
/* .set_tensor = */ NULL,
/* .get_tensor = */ NULL,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL,
/* .clear = */ ggml_backend_multi_buffer_clear,
/* .reset = */ NULL,
GGML_ASSERT(tensor->data == NULL);
GGML_ASSERT(tensor->view_src == NULL);
GGML_ASSERT(addr >= ggml_backend_buffer_get_base(buffer));
- GGML_ASSERT((char *)addr + ggml_backend_buffer_get_alloc_size(buffer, tensor) <=
- (char *)ggml_backend_buffer_get_base(buffer) + ggml_backend_buffer_get_size(buffer));
+ GGML_ASSERT(ggml_backend_buffer_is_meta(buffer) ||
+ (char *) addr + ggml_backend_buffer_get_alloc_size(buffer, tensor) <=
+ (char *) ggml_backend_buffer_get_base(buffer) + ggml_backend_buffer_get_size(buffer));
tensor->buffer = buffer;
tensor->data = addr;
/* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor,
/* .clear = */ ggml_backend_cpu_buffer_clear,
/* .reset = */ NULL,
/* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor,
/* .clear = */ ggml_backend_cpu_buffer_clear,
/* .reset = */ NULL,
/* .get_name = */ ggml_backend_blas_get_name,
/* .free = */ ggml_backend_blas_free,
/* .set_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .get_tensor_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ NULL,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_cann_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_cann_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_cann_buffer_cpy_tensor,
/* .clear = */ ggml_backend_cann_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_cann_free,
/* .set_tensor_async = */ ggml_backend_cann_set_tensor_async,
/* .get_tensor_async = */ ggml_backend_cann_get_tensor_async,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ ggml_backend_cann_cpy_tensor_async,
/* .synchronize = */ ggml_backend_cann_synchronize,
/* .graph_plan_create = */ NULL,
/* .memset_tensor = */ ggml_backend_amx_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_amx_buffer_set_tensor,
/* .get_tensor = */ nullptr,
+ /* .set_tensor_2d = */ nullptr,
+ /* .get_tensor_2d = */ nullptr,
/* .cpy_tensor = */ nullptr,
/* .clear = */ ggml_backend_amx_buffer_clear,
/* .reset = */ nullptr,
/* .free = */ ggml_backend_cpu_free,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ NULL,
/* .graph_plan_create = */ ggml_backend_cpu_graph_plan_create,
target_link_libraries(ggml-cuda PRIVATE CUDA::cuda_driver)
endif()
+ if (GGML_CUDA_NCCL)
+ find_package(NCCL)
+ if (NCCL_FOUND)
+ add_compile_definitions(GGML_USE_NCCL)
+ target_link_libraries(ggml-cuda PRIVATE NCCL::NCCL)
+ else()
+ message(STATUS "Warning: NCCL not found, performance for multiple CUDA GPUs will be suboptimal")
+ endif()
+ endif()
+
set(CUDA_CXX_FLAGS "")
set(CUDA_FLAGS -use_fast_math -extended-lambda)
#define CUBLAS_CHECK(err) CUDA_CHECK_GEN(err, CUBLAS_STATUS_SUCCESS, cublas_get_error_str)
+#ifdef GGML_USE_NCCL
+#define NCCL_CHECK(err) CUDA_CHECK_GEN(err, ncclSuccess, ncclGetErrorString)
+#endif // GGML_USE_NCCL
+
#if !defined(GGML_USE_HIP) && !defined(GGML_CUDA_NO_VMM)
static const char * cu_get_error_str(CUresult err) {
const char * err_str;
cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};
std::array<float, GGML_CUDA_MAX_DEVICES> 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();
// configure logging to stdout
// CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr));
+ for (int id = 0; id < info.device_count; ++id) {
+ ggml_cuda_set_device(id);
+ for (int id_other = 0; id_other < info.device_count; ++id_other) {
+ if (id == id_other) {
+ continue;
+ }
+ int can_access_peer;
+ CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, id_other));
+ if (can_access_peer) {
+ CUDA_CHECK(cudaDeviceEnablePeerAccess(id_other, 0));
+ }
+ }
+ }
+
+#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;
}
}
static void ggml_backend_cuda_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
- ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+ ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
ggml_cuda_set_device(ctx->device);
- CUDA_CHECK(cudaMemsetAsync((char *)tensor->data + offset, value, size, cudaStreamPerThread));
+ CUDA_CHECK(cudaMemsetAsync((char *) tensor->data + offset, value, size, cudaStreamPerThread));
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
}
static void ggml_backend_cuda_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
- ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+ ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
ggml_cuda_set_device(ctx->device);
- CUDA_CHECK(cudaMemcpyAsync((char *)tensor->data + offset, data, size, cudaMemcpyHostToDevice, cudaStreamPerThread));
+ CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cudaStreamPerThread));
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
}
static void ggml_backend_cuda_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
+ ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
+
+ ggml_cuda_set_device(ctx->device);
+ CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cudaStreamPerThread));
+ CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+}
+
+static void ggml_backend_cuda_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data,
+ size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
+
+ ggml_cuda_set_device(ctx->device);
+ CUDA_CHECK(cudaMemcpy2DAsync(
+ (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread));
+ CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+}
+
+static void ggml_backend_cuda_buffer_get_tensor_2d(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data,
+ size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) {
ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
ggml_cuda_set_device(ctx->device);
- CUDA_CHECK(cudaMemcpyAsync(data, (const char *)tensor->data + offset, size, cudaMemcpyDeviceToHost, cudaStreamPerThread));
+ CUDA_CHECK(cudaMemcpy2DAsync(
+ data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cudaStreamPerThread));
CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
}
/* .memset_tensor = */ ggml_backend_cuda_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_cuda_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_cuda_buffer_get_tensor,
+ /* .set_tensor_2d = */ ggml_backend_cuda_buffer_set_tensor_2d,
+ /* .get_tensor_2d = */ ggml_backend_cuda_buffer_get_tensor_2d,
/* .cpy_tensor = */ ggml_backend_cuda_buffer_cpy_tensor,
/* .clear = */ ggml_backend_cuda_buffer_clear,
/* .reset = */ NULL,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_cuda_split_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_cuda_split_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL,
/* .clear = */ ggml_backend_cuda_split_buffer_clear,
/* .reset = */ NULL,
/* .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
+ 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
+ // This then causes a crash in this function
+ if (ne == 0) {
+ return true;
+ }
+ 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)) {
+ 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()));
+ }
+ NCCL_CHECK(ncclGroupEnd());
+
+ return true;
+ }
+
+ // For large tensors it's faster to compress them to BF16 for the reduction:
+ to_bf16_cuda_t to_bf16 = ggml_get_to_bf16_cuda(GGML_TYPE_F32);
+ to_fp32_cuda_t to_fp32 = ggml_get_to_fp32_cuda(GGML_TYPE_BF16);
+
+ ggml_cuda_pool_alloc<nv_bfloat16> 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;
+ 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());
+ 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()));
+ }
+ 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_cuda_set_device(i);
+ 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);
+ return false;
+#endif // GGML_USE_NCCL
+}
+
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<std::mutex> lock(mutex);
GGML_UNUSED_VARS(dst, src1_ddq_i, src1_padded_row_size);
}
-static void ggml_cuda_set_peer_access(const int n_tokens, int main_device) {
- static bool peer_access_enabled = false;
-
- const bool enable_peer_access = n_tokens <= GGML_CUDA_PEER_MAX_BATCH_SIZE;
-
- if (peer_access_enabled == enable_peer_access) {
- return;
- }
-
-#ifdef NDEBUG
- for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) {
- ggml_cuda_set_device(id);
- CUDA_CHECK(cudaDeviceSynchronize());
- }
-
- for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) {
- ggml_cuda_set_device(id);
-
- for (int id_other = 0; id_other < ggml_backend_cuda_get_device_count(); ++id_other) {
- if (id == id_other) {
- continue;
- }
- if (id != main_device && id_other != main_device) {
- continue;
- }
-
- int can_access_peer;
- CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, id_other));
- if (can_access_peer) {
- if (enable_peer_access) {
- cudaError_t err = cudaDeviceEnablePeerAccess(id_other, 0);
- if (err != cudaErrorPeerAccessAlreadyEnabled) {
- CUDA_CHECK(err);
- } else {
- // reset the error
- (void)cudaGetLastError();
- }
- } else {
- cudaError_t err = cudaDeviceDisablePeerAccess(id_other);
- if (err != cudaErrorPeerAccessNotEnabled) {
- CUDA_CHECK(err);
- } else {
- // reset the error
- (void)cudaGetLastError();
- }
- }
- }
- }
- }
-
- ggml_cuda_set_device(main_device);
-#endif // NDEBUG
-
- peer_access_enabled = enable_peer_access;
-
- GGML_UNUSED(main_device);
-}
-
static cudaError_t ggml_cuda_Memcpy2DPeerAsync(
void * dst, int dstDevice, size_t dpitch, void * src, int srcDevice, size_t spitch, size_t width, size_t height, cudaStream_t stream) {
}
static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) {
- // why is this here instead of mul_mat?
- if (dst->src[0] != nullptr && ggml_backend_buft_is_cuda_split(dst->src[0]->buffer->buft)) {
- ggml_cuda_set_peer_access(dst->src[1]->ne[1], ctx.device);
- }
-
switch (dst->op) {
case GGML_OP_ARGMAX:
ggml_cuda_argmax(ctx, dst);
}
static void ggml_backend_cuda_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
- ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
+ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type");
- CUDA_CHECK(cudaMemcpyAsync((char *)tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream()));
+ CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream()));
}
static void ggml_backend_cuda_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
- ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
+ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type");
- CUDA_CHECK(cudaMemcpyAsync(data, (const char *)tensor->data + offset, size, cudaMemcpyDeviceToHost, cuda_ctx->stream()));
+ CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cuda_ctx->stream()));
+}
+
+static void ggml_backend_cuda_set_tensor_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data,
+ size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+ ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+
+ GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type");
+
+ CUDA_CHECK(cudaMemcpy2DAsync(
+ (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cuda_ctx->stream()));
+}
+
+static void ggml_backend_cuda_get_tensor_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data,
+ size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) {
+ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+ ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+
+ GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type");
+
+ CUDA_CHECK(cudaMemcpy2DAsync(
+ data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cuda_ctx->stream()));
}
static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) {
return false;
}
- if (!ggml_backend_buffer_is_cuda(src->buffer) || !ggml_backend_buffer_is_cuda(dst->buffer)) {
+ if (!ggml_backend_buffer_is_cuda(buf_src) || !ggml_backend_buffer_is_cuda(buf_dst)) {
return false;
}
// device -> device copy
- ggml_backend_cuda_context * cuda_ctx_src = (ggml_backend_cuda_context *)backend_src->context;
- ggml_backend_cuda_context * cuda_ctx_dst = (ggml_backend_cuda_context *)backend_dst->context;
+ ggml_backend_cuda_context * cuda_ctx_src = (ggml_backend_cuda_context *) backend_src->context;
+ ggml_backend_cuda_context * cuda_ctx_dst = (ggml_backend_cuda_context *) backend_dst->context;
- ggml_backend_cuda_buffer_context * buf_ctx_src = (ggml_backend_cuda_buffer_context *)buf_src->context;
- ggml_backend_cuda_buffer_context * buf_ctx_dst = (ggml_backend_cuda_buffer_context *)buf_dst->context;
+ ggml_backend_cuda_buffer_context * buf_ctx_src = (ggml_backend_cuda_buffer_context *) buf_src->context;
+ ggml_backend_cuda_buffer_context * buf_ctx_dst = (ggml_backend_cuda_buffer_context *) buf_dst->context;
if (cuda_ctx_src->device != buf_ctx_src->device || cuda_ctx_dst->device != buf_ctx_dst->device) {
#ifndef NDEBUG
GGML_LOG_DEBUG("%s: backend and buffer devices do not match\n", __func__);
-#endif
+#endif // NDEBUG
return false;
}
return false;
#else
CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream()));
-#endif
+#endif // GGML_CUDA_NO_PEER_COPY
}
// record event on src stream after the copy
/* .free = */ ggml_backend_cuda_free,
/* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async,
/* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async,
+ /* .get_tensor_2d_async = */ ggml_backend_cuda_set_tensor_2d_async,
+ /* .set_tensor_2d_async = */ ggml_backend_cuda_get_tensor_2d_async,
/* .cpy_tensor_async = */ ggml_backend_cuda_cpy_tensor_async,
/* .synchronize = */ ggml_backend_cuda_synchronize,
/* .graph_plan_create = */ NULL,
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_split_buffer_type") == 0) {
return (void *)ggml_backend_cuda_split_buffer_type;
}
#include <cuda_bf16.h>
#include <cuda_fp16.h>
+#ifdef GGML_USE_NCCL
+#include <nccl.h>
+#endif // GGML_USE_NCCL
+
#if CUDART_VERSION >= 11080
#include <cuda_fp8.h>
#define FP8_AVAILABLE
#include <rocwmma/rocwmma-version.hpp>
#endif // defined(GGML_HIP_ROCWMMA_FATTN)
+#ifdef GGML_USE_NCCL
+#include <rccl/rccl.h>
+#endif // GGML_USE_NCCL
+
+
#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT
#define CUBLAS_OP_N HIPBLAS_OP_N
#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice
#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite
#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }}
+#define NCCL_CHECK(fn) {ncclResult_t err = fn; if(err != ncclSuccess) { GGML_ABORT("RCCL Failure RCCL returned: %i\n", err); }}
#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width)
#define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width)
#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)
--- /dev/null
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+// This is a "staging" header for new ggml API
+// It is not publicly available and it should not be used by 3rd party projects
+//
+// When the API matures enough, it will be moved to the official public API
+
+//
+// Meta backend
+//
+
+#define GGML_BACKEND_META_MAX_DEVICES 16
+
+enum ggml_backend_meta_split_axis {
+ // tensor split by tensor dimensions:
+ GGML_BACKEND_SPLIT_AXIS_0 = 0,
+ GGML_BACKEND_SPLIT_AXIS_1 = 1,
+ GGML_BACKEND_SPLIT_AXIS_2 = 2,
+ GGML_BACKEND_SPLIT_AXIS_3 = 3,
+
+ GGML_BACKEND_SPLIT_AXIS_MIRRORED = 10, // all values on all backends
+ GGML_BACKEND_SPLIT_AXIS_PARTIAL = 11, // each backend has a partial sum
+
+ // for internal bookkeeping only:
+ GGML_BACKEND_SPLIT_AXIS_NONE = 98,
+ GGML_BACKEND_SPLIT_AXIS_UNKNOWN = 99,
+};
+GGML_API const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis);
+
+struct ggml_backend_meta_split_state {
+ enum ggml_backend_meta_split_axis axis;
+
+ // for tensors with axis >= 0 && axis < GGML_MAX_DIMS:
+ // - each device has a slice of the tensor along the split axis
+ // - 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],
+ // - 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
+ int64_t ne[16*GGML_BACKEND_META_MAX_DEVICES];
+ uint32_t n_segments;
+};
+
+// function to assign split states for statically allocated tensors, compute tensor split states will be assigned to be compatible:
+typedef struct ggml_backend_meta_split_state(*ggml_backend_meta_get_split_state_t)(const struct ggml_tensor * tensor, void * userdata);
+
+// create a new meta device from "simple" devices, meta buffer type/buffer/backend is then derived from this:
+// TODO: this looks a bit strange - a backend API creates a device. I think we should try
+// express this as a backend registry functionality instead
+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);
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_hexagon_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_hexagon_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_hexagon_buffer_cpy_tensor,
/* .clear = */ ggml_backend_hexagon_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_hexagon_free,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ ggml_backend_hexagon_synchronize,
/* .graph_plan_create = */ NULL,
find_package(hipblas REQUIRED)
find_package(rocblas REQUIRED)
+if (GGML_HIP_RCCL)
+ find_package(rccl REQUIRED)
+endif()
+
if (${hip_VERSION} VERSION_LESS 6.1)
message(FATAL_ERROR "At least ROCM/HIP V6.1 is required")
endif()
add_compile_definitions(GGML_HIP_NO_MMQ_MFMA)
endif()
+if (GGML_HIP_RCCL)
+ add_compile_definitions(GGML_USE_NCCL) # RCCL has the same interface as NCCL.
+endif()
+
if (GGML_HIP_EXPORT_METRICS)
set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Rpass-analysis=kernel-resource-usage --save-temps")
endif()
message(FATAL_ERROR "Static linking not supported for HIP/ROCm")
endif()
+if (GGML_HIP_RCCL)
+ target_link_libraries(ggml-hip PRIVATE ggml-base roc::rccl)
+endif()
+
target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas)
/* .memset_tensor = */ ggml_backend_metal_buffer_shared_memset_tensor,
/* .set_tensor = */ ggml_backend_metal_buffer_shared_set_tensor,
/* .get_tensor = */ ggml_backend_metal_buffer_shared_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_metal_buffer_shared_cpy_tensor,
/* .clear = */ ggml_backend_metal_buffer_shared_clear,
/* .reset = */ NULL,
}
static ggml_backend_buffer_i ggml_backend_metal_buffer_private_i = {
- /* .free_buffer = */ ggml_backend_metal_buffer_private_free_buffer,
- /* .get_base = */ ggml_backend_metal_buffer_private_get_base,
- /* .init_tensor = */ NULL,
- /* .memset_tensor = */ ggml_backend_metal_buffer_private_memset_tensor,
- /* .set_tensor = */ ggml_backend_metal_buffer_private_set_tensor,
- /* .get_tensor = */ ggml_backend_metal_buffer_private_get_tensor,
- /* .cpy_tensor = */ ggml_backend_metal_buffer_private_cpy_tensor,
- /* .clear = */ ggml_backend_metal_buffer_private_clear,
- /* .reset = */ NULL,
+ /* .free_buffer = */ ggml_backend_metal_buffer_private_free_buffer,
+ /* .get_base = */ ggml_backend_metal_buffer_private_get_base,
+ /* .init_tensor = */ NULL,
+ /* .memset_tensor = */ ggml_backend_metal_buffer_private_memset_tensor,
+ /* .set_tensor = */ ggml_backend_metal_buffer_private_set_tensor,
+ /* .get_tensor = */ ggml_backend_metal_buffer_private_get_tensor,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
+ /* .cpy_tensor = */ ggml_backend_metal_buffer_private_cpy_tensor,
+ /* .clear = */ ggml_backend_metal_buffer_private_clear,
+ /* .reset = */ NULL,
};
static bool ggml_backend_buffer_is_metal(ggml_backend_buffer_t buffer) {
/* .free = */ ggml_backend_metal_free,
/* .set_tensor_async = */ ggml_backend_metal_set_tensor_async,
/* .get_tensor_async = */ ggml_backend_metal_get_tensor_async,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ ggml_backend_metal_cpy_tensor_async, // only needed for multi-GPU setups
/* .synchronize = */ ggml_backend_metal_synchronize,
/* .graph_plan_create = */ NULL,
/* .set_tensor_async = */ NULL, /* ggml_backend_opencl_set_tensor_async */
/* .get_tensor_async = */ NULL, /* ggml_backend_opencl_get_tensor_async */
/* .cpy_tensor_async = */ NULL, /* ggml_backend_opencl_cpy_tensor_async */
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .synchronize = */ ggml_backend_opencl_synchronize,
/* .graph_plan_create = */ NULL,
/* .graph_plan_free = */ NULL,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_opencl_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_opencl_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL,
/* .clear = */ ggml_backend_opencl_buffer_clear,
/* .reset = */ ggml_backend_opencl_buffer_reset,
/* .memset_tensor = */ ggml_backend_openvino_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_openvino_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_openvino_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_openvino_buffer_cpy_tensor,
/* .clear = */ ggml_backend_openvino_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_openvino_free,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ NULL,
/* .graph_plan_create = */ NULL,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_rpc_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_rpc_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_rpc_buffer_cpy_tensor,
/* .clear = */ ggml_backend_rpc_buffer_clear,
/* .reset = */ NULL,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .synchronize = */ ggml_backend_rpc_synchronize,
/* .graph_plan_create = */ NULL,
/* .graph_plan_free = */ NULL,
/* .memset_tensor = */ ggml_backend_sycl_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_sycl_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_sycl_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_sycl_buffer_cpy_tensor,
/* .clear = */ ggml_backend_sycl_buffer_clear,
/* .reset = */ ggml_backend_sycl_buffer_reset,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_sycl_split_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_sycl_split_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL,
/* .clear = */ ggml_backend_sycl_split_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_sycl_free,
/* .set_tensor_async = */ ggml_backend_sycl_set_tensor_async,
/* .get_tensor_async = */ ggml_backend_sycl_get_tensor_async,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL, // ggml_backend_sycl_cpy_tensor_async,
// // TODO: update for the new
// interface
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor,
/* .clear = */ ggml_backend_remoting_buffer_clear,
/* .reset = */ NULL,
/* .memset_tensor = */ NULL,
/* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor_from_ptr,
/* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor_from_ptr,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor,
/* .clear = */ ggml_backend_remoting_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_remoting_free,
/* .set_tensor_async = */ NULL, // ggml_backend_remoting_set_tensor_async,
/* .get_tensor_async = */ NULL, // ggml_backend_remoting_get_tensor_async,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL, // ggml_backend_remoting_cpy_tensor_async,
/* .synchronize = */ NULL, // ggml_backend_remoting_synchronize,
/* .graph_plan_create = */ NULL,
/* .memset_tensor = */ ggml_backend_vk_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_vk_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_vk_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ ggml_backend_vk_buffer_cpy_tensor,
/* .clear = */ ggml_backend_vk_buffer_clear,
/* .reset = */ NULL,
/* .free = */ ggml_backend_vk_free,
/* .set_tensor_async = */ ggml_backend_vk_set_tensor_async,
/* .get_tensor_async = */ ggml_backend_vk_get_tensor_async,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ ggml_backend_vk_cpy_tensor_async,
/* .synchronize = */ ggml_backend_vk_synchronize,
/* .graph_plan_create = */ NULL,
/* .free = */ ggml_backend_webgpu_free,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ NULL,
/* .graph_plan_create = */ NULL,
/* .memset_tensor = */ ggml_backend_webgpu_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_webgpu_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_webgpu_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL, // TODO: optional, implement this
/* .clear = */ ggml_backend_webgpu_buffer_clear,
/* .reset = */ NULL, // TODO: optional, think it coordinates with
/* .memset_tensor = */ ggml_backend_zdnn_buffer_memset_tensor,
/* .set_tensor = */ ggml_backend_zdnn_buffer_set_tensor,
/* .get_tensor = */ ggml_backend_zdnn_buffer_get_tensor,
+ /* .set_tensor_2d = */ NULL,
+ /* .get_tensor_2d = */ NULL,
/* .cpy_tensor = */ NULL,
/* .clear = */ ggml_backend_zdnn_buffer_clear,
/* .reset = */ NULL,
}
static ggml_backend_i ggml_backend_zdnn_i = {
- /* .get_name = */ ggml_backend_zdnn_name,
- /* .free = */ ggml_backend_zdnn_free,
- /* .set_tensor_async = */ NULL,
- /* .get_tensor_async = */ NULL,
- /* .cpy_tensor_async = */ NULL,
- /* .synchronize = */ NULL,
- /* .graph_plan_create = */ NULL,
- /* .graph_plan_free = */ NULL,
- /* .graph_plan_update = */ NULL,
- /* .graph_plan_compute = */ NULL,
- /* .graph_compute = */ ggml_backend_zdnn_graph_compute,
- /* .event_record = */ NULL,
- /* .event_wait = */ NULL,
- /* .graph_optimize = */ NULL,
+ /* .get_name = */ ggml_backend_zdnn_name,
+ /* .free = */ ggml_backend_zdnn_free,
+ /* .set_tensor_async = */ NULL,
+ /* .get_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
+ /* .cpy_tensor_async = */ NULL,
+ /* .synchronize = */ NULL,
+ /* .graph_plan_create = */ NULL,
+ /* .graph_plan_free = */ NULL,
+ /* .graph_plan_update = */ NULL,
+ /* .graph_plan_compute = */ NULL,
+ /* .graph_compute = */ ggml_backend_zdnn_graph_compute,
+ /* .event_record = */ NULL,
+ /* .event_wait = */ NULL,
+ /* .graph_optimize = */ NULL,
};
static ggml_guid_t ggml_backend_zdnn_guid(void) {
/* .free = */ ggml_backend_zendnn_free,
/* .set_tensor_async = */ NULL,
/* .get_tensor_async = */ NULL,
+ /* .get_tensor_2d_async = */ NULL,
+ /* .set_tensor_2d_async = */ NULL,
/* .cpy_tensor_async = */ NULL,
/* .synchronize = */ NULL,
/* .graph_plan_create = */ NULL,
LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type);
enum llama_split_mode {
- LLAMA_SPLIT_MODE_NONE = 0, // single GPU
- LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs
- LLAMA_SPLIT_MODE_ROW = 2, // split layers and KV across GPUs, use tensor parallelism if supported
+ LLAMA_SPLIT_MODE_NONE = 0, // single GPU
+ LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs
+ LLAMA_SPLIT_MODE_ROW = 2, // split layers and KV across GPUs, use tensor parallelism if supported
+ LLAMA_SPLIT_MODE_TENSOR = 3,
};
// TODO: simplify (https://github.com/ggml-org/llama.cpp/pull/9294#pullrequestreview-2286561979)
return false;
}
}
+
+bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
+ switch (arch) {
+ case LLM_ARCH_GROK:
+ case LLM_ARCH_MPT:
+ case LLM_ARCH_PLAMO2:
+ case LLM_ARCH_MINICPM3:
+ case LLM_ARCH_GEMMA3N:
+ case LLM_ARCH_MAMBA:
+ case LLM_ARCH_MAMBA2:
+ case LLM_ARCH_JAMBA:
+ case LLM_ARCH_FALCON_H1:
+ case LLM_ARCH_OLMO2:
+ case LLM_ARCH_OLMOE:
+ case LLM_ARCH_DEEPSEEK2:
+ case LLM_ARCH_GLM_DSA:
+ case LLM_ARCH_BITNET:
+ case LLM_ARCH_T5:
+ case LLM_ARCH_NEMOTRON_H:
+ case LLM_ARCH_NEMOTRON_H_MOE:
+ case LLM_ARCH_GRANITE_HYBRID:
+ case LLM_ARCH_LFM2:
+ case LLM_ARCH_LFM2MOE:
+ case LLM_ARCH_MINIMAX_M2:
+ case LLM_ARCH_MISTRAL4:
+ case LLM_ARCH_KIMI_LINEAR:
+ return false;
+ default:
+ return true;
+ }
+}
const llm_tensor_info & llm_tensor_info_for(llm_tensor tensor);
-bool llm_arch_is_recurrent(const llm_arch & arch);
-bool llm_arch_is_hybrid (const llm_arch & arch);
-bool llm_arch_is_diffusion(const llm_arch & arch);
+bool llm_arch_is_recurrent (const llm_arch & arch);
+bool llm_arch_is_hybrid (const llm_arch & arch);
+bool llm_arch_is_diffusion (const llm_arch & arch);
+bool llm_arch_supports_sm_tensor(const llm_arch & arch);
#include "llama-context.h"
+#include "ggml.h"
#include "llama-arch.h"
#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-mmap.h"
#include "llama-model.h"
#include "llama-ext.h"
+#include "llama.h"
#include <cinttypes>
#include <cmath>
if (!hparams.vocab_only) {
// GPU backends
- for (auto * dev : model.devices) {
- ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
+ for (const auto & dev : model.devices) {
+ ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr);
if (backend == nullptr) {
- throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev)));
+ throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev)));
}
backends.emplace_back(backend);
}
if (backend_type == GGML_BACKEND_DEVICE_TYPE_CPU && !model.devices.empty()) {
// use the host buffer of the first device CPU for faster transfer of the intermediate state
- auto * dev = model.devices[0];
- auto * host_buft = ggml_backend_dev_host_buffer_type(dev);
+ const auto & dev = model.devices[0];
+ auto * host_buft = ggml_backend_dev_host_buffer_type(dev.dev);
if (host_buft) {
buft = host_buft;
}
for (auto & backend : backends) {
auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get()));
- auto * set_abort_callback_fn = (ggml_backend_set_abort_callback_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_abort_callback");
- if (set_abort_callback_fn) {
- set_abort_callback_fn(backend.get(), this->abort_callback, this->abort_callback_data);
+ if (reg) {
+ auto * set_abort_callback_fn = (ggml_backend_set_abort_callback_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_abort_callback");
+ if (set_abort_callback_fn) {
+ set_abort_callback_fn(backend.get(), this->abort_callback, this->abort_callback_data);
+ }
}
}
}
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
}
+ if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+ if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
+ LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for SPLIT_MODE_TENSOR\n", __func__);
+ params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+ }
+ if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
+ LLAMA_LOG_ERROR("%s: SPLIT_MODE_TENSOR requires flash_attn to be enabled\n", __func__);
+ return nullptr;
+ }
+ if (ggml_is_quantized(params.type_k) || ggml_is_quantized(params.type_v)) {
+ LLAMA_LOG_ERROR("%s: simultaneous use of SPLIT_MODE_TENSOR and KV cache quantization not implemented\n", __func__);
+ return nullptr;
+ }
+ }
+
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) {
const uint32_t blck_size = ggml_blck_size(params.type_k);
for (uint32_t il = 0; il < model->hparams.n_layer; ++il) {
}
void llama_memory_breakdown_print(const struct llama_context * ctx) {
- const std::vector<ggml_backend_dev_t> & devices = ctx->get_model().devices;
+ const auto & devices = ctx->get_model().devices;
std::map<ggml_backend_buffer_type_t, llama_memory_breakdown_data> memory_breakdown = ctx->memory_breakdown();
if (dev) {
int i_dev = -1;
for (size_t i = 0; i < devices.size(); i++) {
- if (devices[i] == dev) {
+ if (devices[i].dev == dev) {
i_dev = i;
break;
}
// print memory breakdown for each device:
for (size_t i = 0; i < devices.size(); i++) {
- ggml_backend_dev_t dev = devices[i];
+ ggml_backend_dev_t dev = devices[i].dev;
llama_memory_breakdown_data mb = mb_dev[i];
const std::string name = ggml_backend_dev_name(dev);
cb(experts, "ffn_moe_weighted", il);
}
+ ggml_build_forward_expand(gf, experts);
+
ggml_tensor * cur_experts[LLAMA_MAX_EXPERTS] = { nullptr };
assert(n_expert_used > 0);
for (uint32_t i = 1; i < hparams.n_expert_used; ++i) {
moe_out = ggml_add(ctx0, moe_out, cur_experts[i]);
+
+ ggml_build_forward_expand(gf, moe_out);
}
if (hparams.n_expert_used == 1) {
ggml_build_forward_expand(gf,
ggml_cpy(ctx0,
states_extra,
- ggml_view_1d(ctx0, s, state_size*(n_rs - n_seqs), (rs_head + n_seqs)*state_size*ggml_element_size(s))));
+ ggml_view_2d(ctx0, s, state_size, (n_rs - n_seqs), s->nb[1], (rs_head + n_seqs)*s->nb[1])));
return output_states;
}
#include "llama-memory-recurrent.h"
+#include "ggml-backend.h"
#include "llama-impl.h"
#include "llama-io.h"
#include "llama-batch.h"
throw std::runtime_error("failed to create ggml context for rs cache");
}
- ggml_tensor * r = ggml_new_tensor_1d(ctx, type_r, hparams.n_embd_r()*mem_size);
- ggml_tensor * s = ggml_new_tensor_1d(ctx, type_s, hparams.n_embd_s()*mem_size);
+ ggml_tensor * r = ggml_new_tensor_2d(ctx, type_r, hparams.n_embd_r(), mem_size);
+ ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), mem_size);
ggml_format_name(r, "cache_r_l%d", i);
ggml_format_name(s, "cache_s_l%d", i);
r_l[i] = r;
#include "llama-model.h"
-#include "ggml.h"
+#include "llama-arch.h"
+#include "llama-hparams.h"
#include "llama-impl.h"
#include "llama-mmap.h"
#include "llama-cparams.h"
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-recurrent.h"
+#include "models/models.h"
+
+#include "ggml.h"
#include "ggml-cpp.h"
-#include "models/models.h"
+// TODO: tmp until the ggml meta backend matures and becomes public
+#include "../src/ggml-ext.h"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <functional>
#include <map>
+#include <numeric>
#include <regex>
#include <sstream>
#include <stdexcept>
+#include <string>
+#include <vector>
+
+struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const struct ggml_tensor * tensor, void * userdata) {
+ const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata;
+ const llama_hparams & hparams = ud->model->hparams;
+ const std::string tensor_name = tensor->name;
+
+ const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
+ const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
+ const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight");
+ const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias");
+ const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias");
+ const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
+ const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
+ const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
+ const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
+ const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
+ const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
+ const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
+
+ const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
+ const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a");
+ const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight");
+ const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight");
+ const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight");
+ const std::regex pattern_r_cache ("cache_r_l\\d*");
+ const std::regex pattern_s_cache ("cache_s_l\\d*");
+ const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight");
+ const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight");
+
+ const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
+ const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
+ const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
+ const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
+ const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
+ const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
+
+ const std::regex pattern_output_weight("output\\.weight");
+ const std::regex pattern_output_bias ("output\\.bias");
+
+ struct tensor_config {
+ ggml_backend_meta_split_axis axis;
+
+ const ggml_tensor * tensor_axis_0;
+
+ uint32_t il;
+ size_t rotation;
+ };
+
+ auto get_tensor_config_impl = [&](
+ const ggml_backend_meta_split_axis axis, const std::string & suffix = "", const std::string & suffix_fallback = "") -> tensor_config {
+ uint32_t il;
+ std::string prefix;
+ size_t rotation;
+ if (tensor_name.substr(0, 4) == "blk.") {
+ const size_t length_prefix = tensor_name.find('.', 4);
+ GGML_ASSERT(length_prefix != std::string::npos);
+ prefix = tensor_name.substr(0, length_prefix + 1);
+ il = std::stoull(tensor_name.substr(4, length_prefix));
+ rotation = il % ud->n_devices;
+ } else if (tensor_name.substr(0, 6) == "cache_") {
+ const size_t layer_index_start = tensor_name.find("_l", 6);
+ GGML_ASSERT(layer_index_start != std::string::npos);
+ il = std::stoull(tensor_name.substr(layer_index_start + 2));
+ prefix = "blk." + std::to_string(il) + ".";
+ rotation = il % ud->n_devices;
+ } else {
+ il = 0;
+ rotation = hparams.n_layer % ud->n_devices;
+ }
+ const ggml_tensor * tensor_axis_0 = suffix.empty() ? tensor : ud->model->get_tensor((prefix + suffix).c_str());
+ if (tensor_axis_0 == nullptr) {
+ GGML_ASSERT(!suffix_fallback.empty());
+ tensor_axis_0 = ud->model->get_tensor((prefix + suffix_fallback).c_str());
+ }
+ GGML_ASSERT(tensor_axis_0 != nullptr);
+ return {axis, tensor_axis_0, il, rotation};
+ };
+
+ auto get_tensor_config = [&]() -> tensor_config {
+ // standard attention
+ if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_q_bias) || std::regex_match(tensor_name, pattern_kv_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_qkv_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
+ }
+ if ( std::regex_match(tensor_name, pattern_qkv_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+ }
+ if (std::regex_match(tensor_name, pattern_qk_norm)) {
+ return get_tensor_config_impl(tensor->ne[1] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_attn_sinks)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_attn_out_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+ }
+ if (std::regex_match(tensor_name, pattern_attn_out_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ }
+
+ if (std::regex_match(tensor_name, pattern_attn_gate_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ssm_out.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta) ||
+ std::regex_match(tensor_name, pattern_ssm_beta_alpha)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ssm_out.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_r_cache) || std::regex_match(tensor_name, pattern_s_cache)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ssm_out.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_conv1d)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ssm_out.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_out_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+ }
+
+ // FFN
+ if (std::regex_match(tensor_name, pattern_ffn_up_gate_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down.weight", "ffn_down_exps.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_up_gate_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down.weight", "ffn_down_exps.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down.weight", "ffn_down_exps.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_down_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down.weight", "ffn_down_exps.weight");
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_down_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_down_exps_bias)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_PARTIAL);
+ }
+
+ // output
+ if (std::regex_match(tensor_name, pattern_output_weight)) {
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
+ }
+ if (std::regex_match(tensor_name, pattern_output_bias)) {
+ const ggml_tensor * output_weight = ud->model->get_tensor("output.weight");
+ GGML_ASSERT(output_weight != nullptr);
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+ }
+
+ // everything else
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+ };
+
+ auto get_split_segments = [&](int axis, uint32_t il) -> std::vector<int64_t> {
+ if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) {
+ const int64_t head_k_dim = hparams.ssm_d_state;
+ const int64_t head_v_dim = hparams.ssm_d_state;
+ const int64_t n_k_heads = hparams.ssm_n_group;
+ const int64_t n_v_heads = hparams.ssm_dt_rank;
+ const int64_t key_dim = head_k_dim * n_k_heads;
+ const int64_t value_dim = head_v_dim * n_v_heads;
+ const int64_t head_ratio = n_v_heads / n_k_heads;
+ if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
+ GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
+ return std::vector<int64_t>(2 + head_ratio, key_dim);
+ }
+ if (std::regex_match(tensor_name, pattern_attn_gate_weight) || std::regex_match(tensor_name, pattern_ssm_out_weight)) {
+ return std::vector<int64_t>(head_ratio, key_dim);
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) ||
+ std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta)) {
+ return std::vector<int64_t>(head_ratio, n_k_heads);
+ }
+ if (std::regex_match(tensor_name, pattern_r_cache)) {
+ return std::vector<int64_t>(2 + head_ratio, key_dim * (hparams.ssm_d_conv - 1));
+ }
+ if (std::regex_match(tensor_name, pattern_s_cache)) {
+ return std::vector<int64_t>(head_ratio, n_k_heads * head_v_dim * head_v_dim);
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
+ return {n_ff_exp, n_ff_exp};
+ }
+ return {tensor->ne[axis]};
+ }
+
+ if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
+ const int64_t n_embd = hparams.n_embd;
+ const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il);
+ GGML_ASSERT(hparams.n_embd_k_gqa() == n_embd_gqa);
+ GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);
+ return {n_embd, n_embd_gqa, n_embd_gqa};
+ }
+ if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
+ return {n_ff_exp, n_ff_exp};
+ }
+ return {tensor->ne[axis]};
+ };
+
+ auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector<int64_t> & segments) -> std::vector<int64_t> {
+ if (hparams.is_recurrent(il)) {
+ // linear attention
+ const int64_t head_dim = hparams.ssm_d_state;
+ const int64_t granularity_qkv = std::lcm(blck_size, head_dim);
+ if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_attn_gate_weight) ||
+ std::regex_match(tensor_name, pattern_ssm_conv1d) || std::regex_match(tensor_name, pattern_ssm_out_weight)) {
+ return std::vector<int64_t>(segments.size(), granularity_qkv);
+ }
+ if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) ||
+ std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta)) {
+ return std::vector<int64_t>(segments.size(), granularity_qkv / head_dim);
+ }
+ if (std::regex_match(tensor_name, pattern_r_cache)) {
+ return std::vector<int64_t>(segments.size(), granularity_qkv * (hparams.ssm_d_conv - 1));
+ }
+ if (std::regex_match(tensor_name, pattern_s_cache)) {
+ return std::vector<int64_t>(segments.size(), granularity_qkv * head_dim);
+ }
+ } else {
+ // regular attention
+ const uint32_t n_gqa = hparams.n_gqa(il);
+ const uint32_t n_embd_q = n_gqa * hparams.n_embd_head_k(il);
+ if (std::regex_match(tensor_name, pattern_attn_sinks)) {
+ GGML_ASSERT(segments.size() == 1);
+ return {std::lcm(n_embd_q, blck_size)/n_embd_q * n_gqa};
+ }
+
+ const int64_t granularity_q = std::lcm(n_embd_q, blck_size);
+ if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) {
+ GGML_ASSERT(segments.size() == 1);
+ // some models have Q gate tensors, for those cases the granularity needs to be doubled:
+ if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) {
+ return {std::lcm(2*n_embd_q, blck_size)};
+ }
+ return {granularity_q};
+ }
+ if (std::regex_match(tensor_name, pattern_attn_out_weight)) {
+ GGML_ASSERT(segments.size() == 1);
+ return {granularity_q};
+ }
+
+ const int64_t granularity_kv = granularity_q / n_gqa;
+ if (std::regex_match(tensor_name, pattern_kv_weight) ||
+ std::regex_match(tensor_name, pattern_kv_bias) ||
+ std::regex_match(tensor_name, pattern_kv_cache)) {
+ GGML_ASSERT(segments.size() == 1);
+ return {granularity_kv};
+ }
+ if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
+ GGML_ASSERT(segments.size() == 3);
+ return {granularity_q, granularity_kv, granularity_kv};
+ }
+ }
+
+ // FFN
+ if (std::regex_match(tensor_name, pattern_ffn_up_gate_weight) || std::regex_match(tensor_name, pattern_ffn_up_gate_bias) ||
+ std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) {
+ GGML_ASSERT(segments.size() <= 2);
+ return std::vector<int64_t>(segments.size(), blck_size);
+ }
+
+ // everything else
+ GGML_ASSERT(segments.size() == 1);
+ return {1};
+ };
+
+ ggml_backend_meta_split_state split_state;
+ memset(&split_state, 0, sizeof(split_state));
+ tensor_config tc = get_tensor_config();
+ split_state.axis = tc.axis;
+ if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) {
+ const int64_t ne_full = tensor->ne[split_state.axis];
+ const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type);
+ const float * tensor_split = ud->model->tensor_split();
+ std::vector<float> tensor_split_scan;
+ tensor_split_scan.reserve(ud->n_devices);
+ for (size_t j = 0; j < ud->n_devices; j++) {
+ tensor_split_scan.push_back(tensor_split == nullptr ? 0.0f : tensor_split[(j + tc.rotation) % ud->n_devices]);
+ if (j > 0) {
+ tensor_split_scan[j] += tensor_split_scan[j - 1];
+ }
+ }
+ const std::vector<int64_t> segments = get_split_segments(split_state.axis, tc.il);
+ const std::vector<int64_t> granularity = get_split_granularity(blck_size, tc.il, segments);
+ for (size_t is = 0; is < segments.size(); is++) {
+ const int64_t ne_s = segments[is];
+ const int64_t g_s = granularity[is];
+ GGML_ASSERT(ne_full % g_s == 0);
+ int64_t low = 0;
+ size_t j = 0;
+ for (; j < ud->n_devices - 1; j++) {
+ int64_t high = tensor_split_scan.back() == 0.0f ?
+ ne_s * (j+1)/ud->n_devices : ne_s * tensor_split_scan[j]/tensor_split_scan.back();
+ if (high % g_s != 0) {
+ high -= high % g_s;
+ }
+ split_state.ne[is*ud->n_devices + (j + tc.rotation) % ud->n_devices] = high - low;
+ low = high;
+ }
+ split_state.ne[is*ud->n_devices + (j + tc.rotation) % ud->n_devices] = ne_s - low;
+ }
+ split_state.n_segments = segments.size();
+ } else {
+ memset(split_state.ne, 0, sizeof(split_state.ne));
+ split_state.n_segments = 1;
+ }
+ return split_state;
+ GGML_UNUSED(userdata);
+}
const char * llm_type_name(llm_type type) {
switch (type) {
}
// CPU: ACCEL -> GPU host -> CPU extra -> CPU
-static buft_list_t make_cpu_buft_list(const std::vector<ggml_backend_dev_t> & devices, bool use_extra_bufts, bool no_host) {
+static buft_list_t make_cpu_buft_list(const std::vector<llama_device> & devices, bool use_extra_bufts, bool no_host) {
buft_list_t buft_list;
// add ACCEL buffer types
// a better approach would be to handle this on a weight-by-weight basis using the offload_op
// function of the device to determine if it would benefit from being stored in a host buffer
if (!no_host) {
- for (auto * dev : devices) {
- ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(dev);
+ for (const auto & dev : devices) {
+ ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(dev.dev);
if (buft) {
- buft_list.emplace_back(dev, buft);
+ buft_list.emplace_back(dev.dev, buft);
break;
}
}
// add the device extra buffer type (if any)
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
- auto ggml_backend_dev_get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)
- ggml_backend_reg_get_proc_address(reg, "ggml_backend_dev_get_extra_bufts");
-
- if (ggml_backend_dev_get_extra_bufts_fn) {
- ggml_backend_buffer_type_t * extra_bufts = ggml_backend_dev_get_extra_bufts_fn(dev);
- while (extra_bufts && *extra_bufts) {
- buft_list.emplace_back(dev, *extra_bufts);
- ++extra_bufts;
+ if (reg) {
+ auto ggml_backend_dev_get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)
+ ggml_backend_reg_get_proc_address(reg, "ggml_backend_dev_get_extra_bufts");
+
+ if (ggml_backend_dev_get_extra_bufts_fn) {
+ ggml_backend_buffer_type_t * extra_bufts = ggml_backend_dev_get_extra_bufts_fn(dev);
+ while (extra_bufts && *extra_bufts) {
+ buft_list.emplace_back(dev, *extra_bufts);
+ ++extra_bufts;
+ }
}
}
if (arch == LLM_ARCH_UNKNOWN) {
throw std::runtime_error("unknown model architecture: '" + ml.get_arch_name() + "'");
}
+ if (!devices.empty() && devices[0].is_meta && !llm_arch_supports_sm_tensor(arch)) {
+ throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'");
+ }
}
void llama_model::load_hparams(llama_model_loader & ml) {
// build a list of buffer types for the CPU and GPU devices
pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host);
- for (auto * dev : devices) {
- buft_list_t buft_list = make_gpu_buft_list(dev, split_mode, tensor_split);
+ for (const auto & dev : devices) {
+ buft_list_t buft_list = make_gpu_buft_list(dev.dev, split_mode, tensor_split);
// add CPU buffer types as a fallback
buft_list.insert(buft_list.end(), pimpl->cpu_buft_list.begin(), pimpl->cpu_buft_list.end());
- pimpl->gpu_buft_list.emplace(dev, std::move(buft_list));
+ pimpl->gpu_buft_list.emplace(dev.dev, std::move(buft_list));
}
ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
if (all_zero) {
// default split, by free memory
for (size_t i = 0; i < n_devices(); ++i) {
- ggml_backend_dev_t dev = devices[i];
+ ggml_backend_dev_t dev = devices[i].dev;
size_t total;
size_t free;
ggml_backend_dev_memory(dev, &free, &total);
return {cpu_dev, &pimpl->cpu_buft_list};
}
const int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + n_devices(), float(il - i_gpu_start)/act_gpu_layers) - splits.begin();
- auto * dev = devices.at(layer_gpu);
+ auto * dev = devices.at(layer_gpu).dev;
LLAMA_LOG_DEBUG("load_tensors: layer %3d assigned to device %s, is_swa = %d\n", il, ggml_backend_dev_name(dev), is_swa);
return {dev, &pimpl->gpu_buft_list.at(dev)};
};
ml.done_getting_tensors();
+ // populate tensors_by_name
+ for (auto & [_, ctx_ptr] : ml.ctx_map) {
+ for (auto * cur = ggml_get_first_tensor(ctx_ptr.get()); cur != NULL; cur = ggml_get_next_tensor(ctx_ptr.get(), cur)) {
+ tensors_by_name.emplace_back(ggml_get_name(cur), cur);
+ }
+ }
+
ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr);
pimpl->mappings.reserve(ml.mappings.size());
}
}
- // populate tensors_by_name
- for (auto & [ctx, _] : pimpl->ctxs_bufs) {
- for (auto * cur = ggml_get_first_tensor(ctx.get()); cur != NULL; cur = ggml_get_next_tensor(ctx.get(), cur)) {
- tensors_by_name.emplace_back(ggml_get_name(cur), cur);
- }
- }
-
if (ml.no_alloc) {
return true;
}
return devices.size();
}
+const float * llama_model::tensor_split() const {
+ return params.tensor_split;
+}
+
uint32_t llama_model::n_gpu_layers() const {
return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer + 1;
}
struct llama_layer_nextn nextn;
};
+struct llama_device {
+ bool is_meta;
+
+ ggml_backend_dev_t dev;
+};
+
+struct llama_meta_device_get_split_state_userdata {
+ size_t n_devices;
+ const struct llama_model * model;
+};
+
+struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const struct ggml_tensor * tensor, void * userdata);
+
struct llama_model {
llm_type type = LLM_TYPE_UNKNOWN;
llm_arch arch = LLM_ARCH_UNKNOWN;
std::unordered_map<std::string, std::string> gguf_kv;
// list of devices used in this model
- std::vector<ggml_backend_dev_t> devices;
+ std::vector<llama_device> devices;
// for quantize-stats only
std::vector<std::pair<std::string, struct ggml_tensor *>> tensors_by_name;
// for keeping track of associated LoRA adapters
std::unordered_set<llama_adapter_lora *> loras;
+ // statically allocated context for assigning
+ struct llama_meta_device_get_split_state_userdata get_split_state_ud;
+
int64_t t_load_us = 0;
int64_t t_start_us = 0;
size_t size() const; // file size
size_t n_tensors() const;
size_t n_devices() const;
+ const float * tensor_split() const;
uint32_t n_gpu_layers() const;
llama_split_mode split_mode() const;
#include "llama.h"
-#include "ggml-cpp.h"
#include "llama-impl.h"
#include "llama-chat.h"
#include "llama-model.h"
#include "ggml.h"
+#include "ggml-cpp.h"
#include "ggml-backend.h"
#include "gguf.h"
+// TODO: tmp until the ggml meta backend matures and becomes public
+#include "../src/ggml-ext.h"
+
#include <algorithm>
#include <cassert>
#include <cinttypes>
#include <cstring>
#include <ctime>
#include <stdexcept>
+#include <vector>
#if defined(_MSC_VER)
#pragma warning(disable: 4244 4267) // possible loss of data
static std::vector<llama_device_memory_data> llama_get_device_memory_data(
const char * path_model, const llama_model_params * mparams, const llama_context_params * cparams,
- std::vector<ggml_backend_dev_t> & devs, uint32_t & hp_ngl, uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert,
+ std::vector<llama_device> & devs, uint32_t & hp_ngl, uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert,
const ggml_log_level log_level) {
struct user_data_t {
struct {
continue;
}
for (size_t i = 0; i < ret.size(); i++) {
- if (model->devices[i] == dev) {
+ if (model->devices[i].dev == dev) {
ret[i].mb.model += mb.model;
ret[i].mb.context += mb.context;
ret[i].mb.compute += mb.compute;
for (size_t i = 0; i < ret.size(); i++) {
size_t free;
size_t total;
- ggml_backend_dev_memory(model->devices[i], &free, &total);
+ ggml_backend_dev_memory(model->devices[i].dev, &free, &total);
// devices can return 0 bytes for free and total memory if they do not
// have any to report. in this case, we will use the host memory as a fallback
const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) {
+ if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) {
+ throw llama_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort");
+ }
constexpr int64_t MiB = 1024*1024;
typedef std::vector<llama_device_memory_data> dmds_t;
const llama_model_params default_mparams = llama_model_default_params();
- std::vector<ggml_backend_dev_t> devs;
+ std::vector<llama_device> devs;
uint32_t hp_ngl = 0; // hparams.n_gpu_layers
uint32_t hp_nct = 0; // hparams.n_ctx_train
uint32_t hp_nex = 0; // hparams.n_expert
{
dev_names.reserve(nd);
size_t max_length = 0;
- for (ggml_backend_dev_t dev : devs) {
- std::string name = ggml_backend_dev_name(dev);
+ for (const llama_device & dev : devs) {
+ std::string name = ggml_backend_dev_name(dev.dev);
name += " (";
- name += ggml_backend_dev_description(dev);
+ name += ggml_backend_dev_description(dev.dev);
name += ")";
dev_names.push_back(name);
max_length = std::max(max_length, name.length());
ngl_per_device_test[id].overflow_type = LAYER_FRACTION_UP;
std::vector<ggml_backend_buffer_type_t> overflow_bufts_test = overflow_bufts;
if (id < nd - 1) {
- overflow_bufts_test[id] = ggml_backend_dev_buffer_type(devs[id + 1]);
+ overflow_bufts_test[id] = ggml_backend_dev_buffer_type(devs[id + 1].dev);
}
LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_UP\n", __func__);
std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test);
// create list of devices to use with this model
if (params.devices) {
- for (ggml_backend_dev_t * dev = params.devices; *dev; ++dev) {
- model->devices.push_back(*dev);
+ if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR) {
+ size_t n_devs = 0;
+ while (params.devices[n_devs]) {
+ n_devs++;
+ }
+ if (n_devs == 0) {
+ LLAMA_LOG_ERROR("%s: LLAMA_SPLIT_MODE_TENSOR needs >= 1 devices\n", __func__);
+ return nullptr;
+ }
+ LLAMA_LOG_INFO("%s: creating a Meta device with %zu devices\n", __func__, n_devs);
+ for (size_t i = 0; i < n_devs; ++i) {
+ LLAMA_LOG_INFO("%s: - device %zu: %s\n", __func__, i, ggml_backend_dev_name(params.devices[i]));
+ }
+ model->get_split_state_ud.n_devices = n_devs;
+ model->get_split_state_ud.model = model;
+ model->devices.push_back({
+ true, ggml_backend_meta_device(
+ params.devices, n_devs, llama_meta_device_get_split_state, &model->get_split_state_ud)
+ });
+ } else {
+ for (ggml_backend_dev_t * dev = params.devices; *dev; ++dev) {
+ model->devices.push_back({false, *dev});
+ }
}
} else {
// default device selection
// build list of available devices
- std::vector<ggml_backend_dev_t> gpus;
- std::vector<ggml_backend_dev_t> igpus;
- std::vector<ggml_backend_dev_t> rpc_servers;
-
- for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
- ggml_backend_dev_t dev = ggml_backend_dev_get(i);
- switch (ggml_backend_dev_type(dev)) {
- case GGML_BACKEND_DEVICE_TYPE_CPU:
- case GGML_BACKEND_DEVICE_TYPE_ACCEL:
- // skip CPU backends since they are handled separately
- break;
-
- case GGML_BACKEND_DEVICE_TYPE_GPU: {
- ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
- if (ggml_backend_reg_name(reg) == std::string("RPC")) {
- rpc_servers.push_back(dev);
- } else {
- // check if there is already a GPU with the same device id
- ggml_backend_dev_props props;
- ggml_backend_dev_get_props(dev, &props);
- auto it = std::find_if(gpus.begin(), gpus.end(), [&props](ggml_backend_dev_t d) {
- ggml_backend_dev_props d_props;
- ggml_backend_dev_get_props(d, &d_props);
- if (props.device_id && d_props.device_id) {
- return strcmp(props.device_id, d_props.device_id) == 0;
- }
- return false;
- });
-
- if (it != gpus.end()) {
- LLAMA_LOG_INFO("%s: skipping device %s (%s) with id %s - already using device %s (%s) with the same id\n",
- __func__,
- ggml_backend_dev_name(dev), ggml_backend_dev_description(dev),
- props.device_id ? props.device_id : "unknown id",
- ggml_backend_dev_name(*it), ggml_backend_dev_description(*it));
+ std::vector<llama_device> gpus;
+ std::vector<llama_device> igpus;
+ std::vector<llama_device> rpc_servers;
+
+ if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR) {
+ std::vector<ggml_backend_dev_t> devs;
+ devs.reserve(ggml_backend_dev_count());
+ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+ auto * dev = ggml_backend_dev_get(i);
+ if (ggml_backend_dev_buffer_type(dev) == ggml_backend_cpu_buffer_type()) {
+ LLAMA_LOG_INFO("%s: skipping %s (%s) for tensor parallelism\n", __func__, ggml_backend_dev_name(dev), ggml_backend_dev_description(dev));
+ continue;
+ }
+ devs.push_back(dev);
+ }
+ if (devs.empty()) {
+ LLAMA_LOG_ERROR("%s: LLAMA_SPLIT_MODE_TENSOR needs >= 1 devices\n", __func__);
+ return nullptr;
+ }
+
+ LLAMA_LOG_INFO("%s: creating a Meta device for tensor parallelism from %zu devices:\n", __func__, devs.size());
+ for (size_t i = 0; i < devs.size(); ++i) {
+ LLAMA_LOG_INFO("%s: - device %zu: %s (%s)\n", __func__, i, ggml_backend_dev_name(devs[i]), ggml_backend_dev_description(devs[i]));
+ }
+
+ GGML_ASSERT(!devs.empty());
+ model->get_split_state_ud.n_devices = devs.size();
+ model->get_split_state_ud.model = model;
+ gpus.push_back({
+ true, ggml_backend_meta_device(
+ devs.data(), devs.size(), llama_meta_device_get_split_state, &model->get_split_state_ud)
+ });
+ } else {
+ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+ ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+ switch (ggml_backend_dev_type(dev)) {
+ case GGML_BACKEND_DEVICE_TYPE_CPU:
+ case GGML_BACKEND_DEVICE_TYPE_ACCEL:
+ // skip CPU backends since they are handled separately
+ break;
+
+ case GGML_BACKEND_DEVICE_TYPE_GPU: {
+ ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+ if (ggml_backend_reg_name(reg) == std::string("RPC")) {
+ rpc_servers.push_back({false, dev});
} else {
- gpus.push_back(dev);
+ // check if there is already a GPU with the same device id
+ ggml_backend_dev_props props;
+ ggml_backend_dev_get_props(dev, &props);
+ auto it = std::find_if(gpus.begin(), gpus.end(), [&props](const llama_device & d) {
+ ggml_backend_dev_props d_props;
+ ggml_backend_dev_get_props(d.dev, &d_props);
+ if (props.device_id && d_props.device_id) {
+ return strcmp(props.device_id, d_props.device_id) == 0;
+ }
+ return false;
+ });
+
+ if (it != gpus.end()) {
+ LLAMA_LOG_INFO("%s: skipping device %s (%s) with id %s - already using device %s (%s) with the same id\n",
+ __func__,
+ ggml_backend_dev_name(dev), ggml_backend_dev_description(dev),
+ props.device_id ? props.device_id : "unknown id",
+ ggml_backend_dev_name(it->dev), ggml_backend_dev_description(it->dev));
+ } else {
+ gpus.push_back({false, dev});
+ }
}
+ break;
}
- break;
- }
- case GGML_BACKEND_DEVICE_TYPE_IGPU:
- igpus.push_back(dev);
- break;
+ case GGML_BACKEND_DEVICE_TYPE_IGPU:
+ igpus.push_back({false, dev});
+ break;
+ case GGML_BACKEND_DEVICE_TYPE_META:
+ GGML_ABORT("fatal error");
+ }
}
}
llama_model_free(model);
return nullptr;
}
- ggml_backend_dev_t main_gpu = model->devices[params.main_gpu];
+ llama_device main_gpu = model->devices[params.main_gpu];
model->devices.clear();
model->devices.push_back(main_gpu);
}
}
- for (auto * dev : model->devices) {
+ for (const auto & dev : model->devices) {
ggml_backend_dev_props props;
- ggml_backend_dev_get_props(dev, &props);
+ ggml_backend_dev_get_props(dev.dev, &props);
LLAMA_LOG_INFO("%s: using device %s (%s) (%s) - %zu MiB free\n", __func__,
- ggml_backend_dev_name(dev), ggml_backend_dev_description(dev),
+ ggml_backend_dev_name(dev.dev), ggml_backend_dev_description(dev.dev),
props.device_id ? props.device_id : "unknown id",
props.memory_free/1024/1024);
}
cb(beta, "beta", il);
beta = ggml_sigmoid(ctx0, beta);
+ cb(beta, "beta_sigmoid", il);
ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
cb(last_conv_states, "last_conv_states", il);
ggml_tensor * state_update_target =
- ggml_view_1d(ctx0, conv_states_all, (conv_kernel_size - 1) * conv_channels * n_seqs,
+ ggml_view_2d(ctx0, conv_states_all, (conv_kernel_size - 1) * conv_channels, n_seqs, conv_states_all->nb[1],
kv_head * (conv_kernel_size - 1) * conv_channels * ggml_element_size(conv_states_all));
cb(state_update_target, "state_update_target", il);
// Update the recurrent states
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, new_state,
- ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
+ ggml_view_2d(ctx0, ssm_states_all, hparams.n_embd_s(), n_seqs, ssm_states_all->nb[1],
kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
// z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim]
cb(beta, "beta", il);
beta = ggml_sigmoid(ctx0, beta);
+ cb(beta, "beta_sigmoid", il);
ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);
alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);
cb(last_conv_states, "last_conv_states", il);
ggml_tensor * state_update_target =
- ggml_view_1d(ctx0, conv_states_all, (conv_kernel_size - 1) * conv_channels * n_seqs,
+ ggml_view_2d(ctx0, conv_states_all, (conv_kernel_size - 1) * conv_channels, n_seqs, conv_states_all->nb[1],
kv_head * (conv_kernel_size - 1) * conv_channels * ggml_element_size(conv_states_all));
cb(state_update_target, "state_update_target", il);
// Update the recurrent states
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, new_state,
- ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
+ ggml_view_2d(ctx0, ssm_states_all, hparams.n_embd_s(), n_seqs, ssm_states_all->nb[1],
kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
// z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim]
GGML_ASSERT(num_v_heads % num_k_heads == 0);
int64_t repeat_factor = num_v_heads / num_k_heads;
- // repeat interleave: reshape to (repeat part, 1, remaining part), do repeat, then reshape back
- ggml_tensor * q_reshaped = ggml_reshape_3d(ctx0, q_conv, head_k_dim, 1, num_k_heads * n_seq_tokens * n_seqs);
- ggml_tensor * k_reshaped = ggml_reshape_3d(ctx0, k_conv, head_k_dim, 1, num_k_heads * n_seq_tokens * n_seqs);
+ // repeat interleave: reshape to (repeat part, 1, remaining part...), do repeat, then reshape back
+ ggml_tensor * q_reshaped = ggml_reshape_4d(ctx0, q_conv, head_k_dim, 1, num_k_heads, n_seq_tokens * n_seqs);
+ ggml_tensor * k_reshaped = ggml_reshape_4d(ctx0, k_conv, head_k_dim, 1, num_k_heads, n_seq_tokens * n_seqs);
// Repeat along the third dimension (the new dimension with size 1)
ggml_tensor * q_repeated =
- ggml_repeat_4d(ctx0, q_reshaped, head_k_dim, repeat_factor, num_k_heads * n_seq_tokens * n_seqs, 1);
+ ggml_repeat_4d(ctx0, q_reshaped, head_k_dim, repeat_factor, num_k_heads, n_seq_tokens * n_seqs);
ggml_tensor * k_repeated =
- ggml_repeat_4d(ctx0, k_reshaped, head_k_dim, repeat_factor, num_k_heads * n_seq_tokens * n_seqs, 1);
+ ggml_repeat_4d(ctx0, k_reshaped, head_k_dim, repeat_factor, num_k_heads, n_seq_tokens * n_seqs);
// Reshape back to merge the head and repeat dimensions
- // From [head_dim, num_k_heads, repeat_factor, n_seq_tokens * n_seqs]
- // Back to [head_dim, num_k_heads * repeat_factor, n_seq_tokens, n_seqs]
+ // From [head_dim, repeat_factor, num_k_heads, n_seq_tokens * n_seqs]
+ // Back to [head_dim, repeat_factor * num_k_heads, n_seq_tokens, n_seqs]
q_conv = ggml_reshape_4d(ctx0, q_repeated, head_k_dim, num_k_heads * repeat_factor, n_seq_tokens, n_seqs);
k_conv = ggml_reshape_4d(ctx0, k_repeated, head_k_dim, num_k_heads * repeat_factor, n_seq_tokens, n_seqs);
}
#include "ggml-cpp.h"
#include "llama.h"
#include "llama-cpp.h"
+
+// TODO: replace with #include "llama-ext.h" in the future
#include "../src/llama-arch.h"
#include "../src/llama-model-saver.h"
ms.add_kv(LLM_KV_XIELU_ALPHA_P, 1.0f);
ms.add_kv(LLM_KV_XIELU_BETA, 1.0f);
ms.add_kv(LLM_KV_XIELU_EPS, 1.0e-7f);
- ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 64 : 2*n_embd);
+ ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 256 : 2*n_embd);
ms.add_kv(LLM_KV_SSM_CONV_KERNEL, uint32_t(4));
- ms.add_kv(LLM_KV_SSM_STATE_SIZE, uint32_t(32));
+ ms.add_kv(LLM_KV_SSM_STATE_SIZE, uint32_t(128));
ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head);
ms.add_kv(LLM_KV_SSM_GROUP_COUNT, arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2));
ms.add_kv(LLM_KV_KDA_HEAD_DIM, uint32_t(128));
}
static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
- struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs) {
+ struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs,
+ const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) {
GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
llama_model_params model_params = llama_model_default_params();
model_params.progress_callback = silent_model_load_progress;
std::vector<ggml_backend_dev_t> devs_copy = devs;
devs_copy.push_back(nullptr);
model_params.devices = devs_copy.data();
+ model_params.split_mode = split_mode;
llama_context_params ctx_params = llama_context_default_params();
ctx_params.n_ctx = 0;
ctx_params.n_threads = 4;
ctx_params.n_threads_batch = 4;
+ if (!encode) {
+ ctx_params.n_ubatch = 64;
+ }
size_t tmp = seed;
llama_model_ptr model(gguf_ctx != nullptr ?
}
}
+static bool arch_supported(const llm_arch arch) {
+ if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {
+ return false; // These models don't have usable implementations.
+ }
+ if (arch == LLM_ARCH_CHAMELEON) {
+ return false; // Only half-implemented and to be removed in the future.
+ }
+ if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {
+ return false; // FIXME CUDA backend crashes.
+ }
+ if (arch == LLM_ARCH_GEMMA4) {
+ return false; // FIXME @ngxson
+ }
+ if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
+ return false; // FIXME Embedding (?) models produce inconsistent results.
+ }
+ if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {
+ return false; // FIXME RWKV models hang indefinitely.
+ }
+ if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||
+ arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {
+ return false; // TODO vocab
+ }
+ if (arch == LLM_ARCH_PLM) {
+ return false; // TODO tensor shapes
+ }
+ if (arch == LLM_ARCH_DEEPSEEK2OCR) {
+ return false;
+ }
+
+ // FIXME some models are segfaulting with WebGPU:
+#ifdef GGML_USE_WEBGPU
+ if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_KIMI_LINEAR) {
+ return false;
+ }
+#endif // GGML_USE_WEBGPU
+
+ return true;
+}
+
static int save_models(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level, const std::string & dir) {
struct user_data_t {
struct {
}, &ud);
for (const llm_arch & arch : llm_arch_all()) {
- if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
+ if (arch == LLM_ARCH_UNKNOWN) {
continue;
}
- if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {
- continue; // These models don't have usable implementations.
- }
- if (arch == LLM_ARCH_CHAMELEON) {
- continue; // Only half-implemented and to be removed in the future.
- }
- if (arch == LLM_ARCH_GEMMA4) {
- continue; // FIXME @ngxson
- }
- if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {
- continue; // FIXME
- }
- if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||
- arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {
- continue; // TODO vocab
- }
- if (arch == LLM_ARCH_PLM) {
- continue; // TODO tensor shapes
+ if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
+ continue;
}
for (bool moe : {false, true}) {
if (moe && !moe_implemented(arch)) {
const std::vector<llama_token> tokens = get_tokens(128, 128, seed);
+ struct device_config {
+ std::vector<ggml_backend_dev_t> devs;
+ std::string label;
+ llama_split_mode split_mode;
+
+ device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode)
+ : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {}
+ };
+
+ std::vector<device_config> dev_configs;
+ {
+ std::vector<ggml_backend_dev_t> devices_meta;
+ {
+ const size_t device_count = ggml_backend_dev_count();
+ for (size_t i = 0; i < device_count; i++) {
+ ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+ dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER);
+
+ // cpu-based devices cannot be used in tensor split mode
+ if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {
+ devices_meta.push_back(dev);
+ }
+ }
+ }
+
+ dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR);
+ }
+
bool all_ok = true;
common_log_flush(common_log_main());
- printf("|%15s|%30s|%6s|%15s|%9s|\n", "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip");
- printf("|---------------|------------------------------|------|---------------|---------|\n");
+ printf("|%16s|%30s|%6s|%15s|%9s|\n", "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip");
+ printf("|----------------|------------------------------|------|---------------|---------|\n");
for (const llm_arch & arch : llm_arch_all()) {
- if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
+ if (arch == LLM_ARCH_UNKNOWN) {
continue;
}
- if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {
- continue; // These models don't have usable implementations.
- }
- if (arch == LLM_ARCH_CHAMELEON) {
- continue; // Only half-implemented and to be removed in the future.
- }
- if (arch == LLM_ARCH_GEMMA4) {
- continue; // FIXME @ngxson
- }
- if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {
- continue; // FIXME CUDA backend crashes.
- }
- if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
- continue; // FIXME Embedding (?) models produce inconsistent results.
- }
- if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {
- continue; // FIXME RWKV models hang indefinitely.
- }
- if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||
- arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {
- continue; // TODO vocab
- }
- if (arch == LLM_ARCH_PLM) {
- continue; // TODO tensor shapes
- }
- if (arch == LLM_ARCH_DEEPSEEK2OCR) {
- continue; // TODO tensor shapes
- }
-
- // FIXME some models are segfaulting with WebGPU:
-#ifdef GGML_USE_WEBGPU
- if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_KIMI_LINEAR) {
+ if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
continue;
}
-#endif // GGML_USE_WEBGPU
- const bool encode = arch == LLM_ARCH_T5;
+ const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1;
for (bool moe : {false, true}) {
if (moe && !moe_implemented(arch)) {
continue;
if (!moe && moe_mandatory(arch)) {
continue;
}
+ const std::string config_name = moe ? "MoE" : "Dense";
gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
- auto model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {});
- const std::vector<float> logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);
- for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
- ggml_backend_dev_t dev = ggml_backend_dev_get(i);
- if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
- continue;
- }
- auto model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev});
- std::string config_name = moe ? "MoE" : "Dense";
- const std::vector<float> logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
- const double nmse_val = nmse(logits_cpu, logits_dev);
- char nmse_str[10];
- snprintf(nmse_str, sizeof(nmse_str), "%.2e", nmse_val);
- std::string status_nmse = "\033[1;32mOK\033[0m";
- if (nmse_val > 1e-4) {
- all_ok = false;
- status_nmse = "\033[1;31mFAIL\033[0m";
- }
-
+ std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_cpu;
+ std::vector<float> logits_cpu;
+ for (device_config & dc : dev_configs) {
+ std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_dev;
+ std::vector<float> logits_dev;
+ std::string status_nmse = "\033[1;33mSKIP\033[0m";
std::string status_roundtrip = "\033[1;33mSKIP\033[0m";
- FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.
- if (file != nullptr && llama_model_saver_supports_arch(arch)) {
- llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());
- ms.add_kv_from_model();
- ms.add_tensors_from_model();
- ms.save(file);
- rewind(file);
-
- auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, {dev});
- const std::vector<float> logits_roundtrip = get_logits(
- model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
- status_roundtrip = "\033[1;32mOK\033[0m";
- GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
- for (size_t i = 0; i < logits_roundtrip.size(); i++) {
- if (logits_roundtrip[i] != logits_dev[i]) {
+ char nmse_str[12] = {0};
+ bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());
+#if defined(GGML_USE_WEBGPU)
+ skip = true; // FIXME
+#endif // GGML_USE_WEBGPU
+ if (!skip) {
+ if (logits_cpu.empty()) {
+ model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode);
+ logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);
+ }
+ if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {
+ model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode);
+ logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
+ const double nmse_val = nmse(logits_cpu, logits_dev);
+ snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);
+ status_nmse = "\033[1;32mOK\033[0m";
+ if (nmse_val > 1e-4) {
all_ok = false;
- status_roundtrip = "\033[1;31mFAIL\033[0m";
- break;
+ status_nmse = "\033[1;31mFAIL\033[0m";
+ }
+ }
+
+ FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.
+ // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend
+ // in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible
+ if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {
+ GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second);
+ llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());
+ ms.add_kv_from_model();
+ ms.add_tensors_from_model();
+ ms.save(file);
+ rewind(file);
+
+ auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode);
+ const std::vector<float> logits_roundtrip = get_logits(
+ model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
+ status_roundtrip = "\033[1;32mOK\033[0m";
+ GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
+ for (size_t i = 0; i < logits_roundtrip.size(); i++) {
+ if (logits_roundtrip[i] != logits_dev[i]) {
+ all_ok = false;
+ status_roundtrip = "\033[1;31mFAIL\033[0m";
+ break;
+ }
}
}
}
- printf("|%15s|%30s|%6s|%15s (%8s)|%20s|\n", llm_arch_name(arch), ggml_backend_dev_description(dev),
+ printf("|%16s|%30s|%6s|%15s %10s|%20s|\n", llm_arch_name(arch), dc.label.c_str(),
config_name.c_str(), status_nmse.c_str(), nmse_str, status_roundtrip.c_str());
}
}
return "layer";
case LLAMA_SPLIT_MODE_ROW:
return "row";
+ case LLAMA_SPLIT_MODE_TENSOR:
+ return "tensor";
default:
GGML_ABORT("invalid split mode");
}
printf(" --poll <0...100> (default: %s)\n", join(cmd_params_defaults.poll, ",").c_str());
printf(" -ngl, --n-gpu-layers <n> (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str());
printf(" -ncmoe, --n-cpu-moe <n> (default: %s)\n", join(cmd_params_defaults.n_cpu_moe, ",").c_str());
- printf(" -sm, --split-mode <none|layer|row> (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str());
+ printf(" -sm, --split-mode <none|layer|row|tensor> (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str());
printf(" -mg, --main-gpu <i> (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str());
printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str());
printf(" -fa, --flash-attn <0|1> (default: %s)\n", join(cmd_params_defaults.flash_attn, ",").c_str());
mode = LLAMA_SPLIT_MODE_LAYER;
} else if (m == "row") {
mode = LLAMA_SPLIT_MODE_ROW;
+ } else if (m == "tensor") {
+ mode = LLAMA_SPLIT_MODE_TENSOR;
} else {
invalid_param = true;
break;
return 6;
}
if (field == "split_mode") {
- return 5;
+ return 6;
}
if (field == "flash_attn") {
return 2;
auto * model = llama_init->model();
auto * ctx = llama_init->context();
- if (model == NULL) {
+ if (model == nullptr) {
LOG_ERR("%s: unable to load model\n", __func__);
return 1;
}
+ if (ctx == nullptr) {
+ LOG_ERR("%s: failed to create context\n", __func__);
+ return 1;
+ }
+
const int n_ctx_train = llama_model_n_ctx_train(model);
if (params.n_ctx > n_ctx_train) {
overview / pkg / ggml / sources / llama.cpp / commitdiff