diff --git a/CMakeLists.txt b/CMakeLists.txt index 29d914eb5..ac939e6e4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2363,6 +2363,35 @@ if (ENGINE_BUILD_TESTS) COMMAND gguf_tensor_source_test ) + add_engine_unittest(minimax_music3_lm_head_test tests/unittests/test_minimax_music3_lm_head.cpp) + target_include_directories(minimax_music3_lm_head_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + + add_test( + NAME minimax_music3_lm_head_test + COMMAND minimax_music3_lm_head_test + ) + + add_engine_unittest( + minimax_music3_pipeline_buffers_test + tests/unittests/test_minimax_music3_pipeline_buffers.cpp) + target_include_directories( + minimax_music3_pipeline_buffers_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + + add_test( + NAME minimax_music3_pipeline_buffers_test + COMMAND minimax_music3_pipeline_buffers_test + ) + + add_engine_unittest( + minimax_music3_graph_release_policy_test + tests/unittests/test_minimax_music3_graph_release_policy.cpp) + + add_test( + NAME minimax_music3_graph_release_policy_test + COMMAND minimax_music3_graph_release_policy_test + ) + add_engine_unittest(model_spec_system_test tests/unittests/test_model_spec_system.cpp) target_include_directories(model_spec_system_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) diff --git a/external/ggml/include/ggml-cuda.h b/external/ggml/include/ggml-cuda.h index 87fbaa151..d5afa8f8b 100644 --- a/external/ggml/include/ggml-cuda.h +++ b/external/ggml/include/ggml-cuda.h @@ -1,52 +1,55 @@ -#pragma once - -#include "ggml.h" -#include "ggml-backend.h" - -#ifdef __cplusplus -extern "C" { -#endif - -#ifdef GGML_USE_HIP -#define GGML_CUDA_NAME "ROCm" -#define GGML_CUBLAS_NAME "hipBLAS" -#elif defined(GGML_USE_MUSA) -#define GGML_CUDA_NAME "MUSA" -#define GGML_CUBLAS_NAME "muBLAS" -#else -#define GGML_CUDA_NAME "CUDA" -#define GGML_CUBLAS_NAME "cuBLAS" -#endif -#define GGML_CUDA_MAX_DEVICES 16 - -// backend API -GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device); - -GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#ifdef GGML_USE_HIP +#define GGML_CUDA_NAME "ROCm" +#define GGML_CUBLAS_NAME "hipBLAS" +#elif defined(GGML_USE_MUSA) +#define GGML_CUDA_NAME "MUSA" +#define GGML_CUBLAS_NAME "muBLAS" +#else +#define GGML_CUDA_NAME "CUDA" +#define GGML_CUBLAS_NAME "cuBLAS" +#endif +#define GGML_CUDA_MAX_DEVICES 16 + +// backend API +GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device); + +GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); GGML_BACKEND_API void ggml_backend_cuda_trim_pools(ggml_backend_t backend); +// Returns the backend's current compute CUDA stream (cudaStream_t) so host +// code can enqueue its own kernels/copies ordered with graph computes. +GGML_BACKEND_API void * ggml_backend_cuda_get_stream(ggml_backend_t backend); GGML_BACKEND_API void ggml_backend_cuda_clear_graph(ggml_backend_t backend, const struct ggml_cgraph * graph); - -// 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); - -// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU -GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void); - -GGML_BACKEND_API int ggml_backend_cuda_get_device_count(void); -GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size); -GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); - -GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); -GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); - -GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); - -#ifdef __cplusplus -} -#endif + +// 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); + +// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void); + +GGML_BACKEND_API int ggml_backend_cuda_get_device_count(void); +GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size); +GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); + +GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); +GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); + +GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); + +#ifdef __cplusplus +} +#endif diff --git a/external/ggml/src/ggml-cuda/common.cuh b/external/ggml/src/ggml-cuda/common.cuh index 79053a82b..ab65ab604 100644 --- a/external/ggml/src/ggml-cuda/common.cuh +++ b/external/ggml/src/ggml-cuda/common.cuh @@ -1451,7 +1451,16 @@ struct ggml_backend_cuda_context { cudaStream_t stream(int device, int stream) { if (streams[device][stream] == nullptr) { ggml_cuda_set_device(device); - CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking)); + // GGML_CUDA_STREAM_PRIORITY (read at stream creation, not cached): + // lets a host create backend instances whose streams differ in + // scheduling priority (CUDA: numerically lower = higher priority). + const char * priority_env = getenv("GGML_CUDA_STREAM_PRIORITY"); + if (priority_env != nullptr && atoi(priority_env) != 0) { + CUDA_CHECK(cudaStreamCreateWithPriority( + &streams[device][stream], cudaStreamNonBlocking, atoi(priority_env))); + } else { + CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking)); + } } return streams[device][stream]; } @@ -1521,12 +1530,12 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * gate = nullptr; const ggml_tensor * gate_bias = nullptr; ggml_glu_op glu_op; - bool residual_only = false; + bool residual_only = false; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; const void * gate = nullptr; const void * gate_bias = nullptr; ggml_glu_op glu_op; - bool residual_only = false; + bool residual_only = false; }; diff --git a/external/ggml/src/ggml-cuda/ggml-cuda.cu b/external/ggml/src/ggml-cuda/ggml-cuda.cu index ea67a93f4..ec46187ad 100644 --- a/external/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/external/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1,5043 +1,5050 @@ -#include "ggml-cuda.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" - -#include "ggml-cuda/allreduce.cuh" -#include "ggml-cuda/common.cuh" -#include "ggml-cuda/acc.cuh" -#include "ggml-cuda/add-id.cuh" -#include "ggml-cuda/arange.cuh" -#include "ggml-cuda/argmax.cuh" -#include "ggml-cuda/argsort.cuh" -#include "ggml-cuda/binbcast.cuh" +#include "ggml-cuda.h" +#include "ggml-impl.h" +#include "ggml-backend-impl.h" + +#include "ggml-cuda/allreduce.cuh" +#include "ggml-cuda/common.cuh" +#include "ggml-cuda/acc.cuh" +#include "ggml-cuda/add-id.cuh" +#include "ggml-cuda/arange.cuh" +#include "ggml-cuda/argmax.cuh" +#include "ggml-cuda/argsort.cuh" +#include "ggml-cuda/binbcast.cuh" #include "ggml-cuda/clamp.cuh" #include "ggml-cuda/col2im-1d.cuh" #include "ggml-cuda/concat.cuh" #include "ggml-cuda/convrot-linear.cuh" #include "ggml-cuda/conv-transpose-1d.cuh" -#include "ggml-cuda/conv2d.cuh" -#include "ggml-cuda/conv2d-dw.cuh" -#include "ggml-cuda/conv2d-transpose.cuh" -#include "ggml-cuda/convert.cuh" -#include "ggml-cuda/count-equal.cuh" -#include "ggml-cuda/cpy.cuh" -#include "ggml-cuda/cross-entropy-loss.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/diagmask.cuh" -#include "ggml-cuda/diag.cuh" -#include "ggml-cuda/fattn.cuh" -#include "ggml-cuda/getrows.cuh" -#include "ggml-cuda/im2col.cuh" -#include "ggml-cuda/mmf.cuh" -#include "ggml-cuda/mmq.cuh" -#include "ggml-cuda/mmvf.cuh" -#include "ggml-cuda/mmvq.cuh" -#include "ggml-cuda/norm.cuh" -#include "ggml-cuda/opt-step-adamw.cuh" -#include "ggml-cuda/opt-step-sgd.cuh" -#include "ggml-cuda/out-prod.cuh" -#include "ggml-cuda/pad.cuh" -#include "ggml-cuda/pool2d.cuh" -#include "ggml-cuda/quantize.cuh" -#include "ggml-cuda/rope.cuh" +#include "ggml-cuda/conv2d.cuh" +#include "ggml-cuda/conv2d-dw.cuh" +#include "ggml-cuda/conv2d-transpose.cuh" +#include "ggml-cuda/convert.cuh" +#include "ggml-cuda/count-equal.cuh" +#include "ggml-cuda/cpy.cuh" +#include "ggml-cuda/cross-entropy-loss.cuh" +#include "ggml-cuda/cumsum.cuh" +#include "ggml-cuda/diagmask.cuh" +#include "ggml-cuda/diag.cuh" +#include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/getrows.cuh" +#include "ggml-cuda/im2col.cuh" +#include "ggml-cuda/mmf.cuh" +#include "ggml-cuda/mmq.cuh" +#include "ggml-cuda/mmvf.cuh" +#include "ggml-cuda/mmvq.cuh" +#include "ggml-cuda/norm.cuh" +#include "ggml-cuda/opt-step-adamw.cuh" +#include "ggml-cuda/opt-step-sgd.cuh" +#include "ggml-cuda/out-prod.cuh" +#include "ggml-cuda/pad.cuh" +#include "ggml-cuda/pool2d.cuh" +#include "ggml-cuda/quantize.cuh" +#include "ggml-cuda/rope.cuh" #include "ggml-cuda/roll.cuh" #include "ggml-cuda/scale.cuh" #include "ggml-cuda/sage-attn2.cuh" #include "ggml-cuda/snake.cuh" -#include "ggml-cuda/softcap.cuh" -#include "ggml-cuda/softmax.cuh" -#include "ggml-cuda/ssm-conv.cuh" -#include "ggml-cuda/ssm-scan.cuh" -#include "ggml-cuda/sum.cuh" -#include "ggml-cuda/sumrows.cuh" -#include "ggml-cuda/top-k.cuh" -#include "ggml-cuda/mean.cuh" -#include "ggml-cuda/tsembd.cuh" -#include "ggml-cuda/topk-moe.cuh" -#include "ggml-cuda/unary.cuh" -#include "ggml-cuda/upscale.cuh" -#include "ggml-cuda/wkv.cuh" -#include "ggml-cuda/gla.cuh" -#include "ggml-cuda/gated_delta_net.cuh" -#include "ggml-cuda/set.cuh" -#include "ggml-cuda/set-rows.cuh" -#include "ggml-cuda/pad_reflect_1d.cuh" -#include "ggml-cuda/solve_tri.cuh" -#include "ggml-cuda/tri.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/fill.cuh" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); - -#define GGML_LOG_WARN_ONCE(str) \ - { static std::once_flag warn_flag; std::call_once(warn_flag, []() { GGML_LOG_WARN(str); }); } - -[[noreturn]] -void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { - int id = -1; // in case cudaGetDevice fails - (void)cudaGetDevice(&id); - - GGML_LOG_ERROR(GGML_CUDA_NAME " error: %s\n", msg); - GGML_LOG_ERROR(" current device: %d, in function %s at %s:%d\n", id, func, file, line); - GGML_LOG_ERROR(" %s\n", stmt); - // abort with GGML_ABORT to get a stack trace - GGML_ABORT(GGML_CUDA_NAME " error"); -} - -// this is faster on Windows -// probably because the Windows CUDA libraries forget to make this check before invoking the drivers -void ggml_cuda_set_device(int device) { - int current_device; - CUDA_CHECK(cudaGetDevice(¤t_device)); - - if (device == current_device) { - return; - } - - CUDA_CHECK(cudaSetDevice(device)); -} - -int ggml_cuda_get_device() { - int id; - CUDA_CHECK(cudaGetDevice(&id)); - return id; -} - -static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device) { - ggml_cuda_set_device(device); - cudaError_t err; - if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr) { - err = cudaMallocManaged(ptr, size); -#if defined(GGML_USE_HIP) - if (err == hipSuccess) { - // hipMemAdviseSetCoarseGrain is an optional performance hint; - // ignore errors (e.g. hipErrorInvalidValue on some APU/iGPU configs). - (void)cudaMemAdvise(*ptr, size, hipMemAdviseSetCoarseGrain, device); - (void)hipGetLastError(); // clear any error - } - - // fall back to cudaMalloc if not supported (e.g. on Windows) - if (err == hipErrorNotSupported) { - static bool warned_unsupported = false; - if (!warned_unsupported) { - GGML_LOG_WARN("hipMallocManaged unsupported, falling back to hipMalloc.\n"); - warned_unsupported = true; - } - - err = cudaMalloc(ptr, size); - } -#endif // defined(GGML_USE_HIP) - } else { - err = cudaMalloc(ptr, size); - } - return err; -} - -#if defined(GGML_USE_HIP) -static int ggml_cuda_parse_id(char devName[]) { - // A list of possible Target IDs can be found under the rocclr/clr repo in device.cpp - // these values are not stable so this is susceptible to breakage - // https://github.com/ROCm/clr/blob/amd-staging/rocclr/device/device.cpp - int archMajor = 0x0; - int archMinor = 0x0; - int archNum = GGML_CUDA_CC_OFFSET_AMD; - int archLen = strlen(devName); - char archName[archLen + 1]; - - // strip leading 'gfx' while copying into our buffer - if (archLen > 3) { - strcpy(archName, &devName[3]); - archLen -= 3; - } - - // trim trailing :xnack- or :sramecc- statuses - archLen = strcspn(archName, ":"); - archName[archLen] = '\0'; - - // tease out the version information - if (archLen > 8) { - // versions labeled generic use '-' as delimiter - // strip the trailing "-generic" then iterate through what remains - if ((strstr(archName, "-generic"))) { - archName[archLen - 8] = '\0'; - char * pch; - if ((pch = strtok(archName, "-"))) { - archMajor = (int)strtoul(pch, 0, 16); - if ((pch = strtok(NULL, "-"))) { - archMinor = 0x10 * (int)strtoul(pch, 0, 16); - } - } - } - } else if (archLen >= 3) { - // last two digits should be the minor * 0x10 + stepping - archMinor = (int)strtoul(&archName[archLen - 2], 0, 16); - archName[archLen - 2] = '\0'; - - // only the major version remains - archMajor = (int)strtoul(archName, 0, 16); - } - archNum += archMajor * 0x100; - archNum += archMinor; - return archNum; -} -#endif // defined(GGML_USE_HIP) - -static ggml_cuda_device_info ggml_cuda_init() { - ggml_cuda_device_info info = {}; - - cudaError_t err = cudaGetDeviceCount(&info.device_count); - if (err != cudaSuccess) { - GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err)); - return info; - } - - GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES); - - int64_t total_vram = 0; - for (int id = 0; id < info.device_count; ++id) { - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); - total_vram += prop.totalGlobalMem; - } - GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n", - __func__, info.device_count, (size_t)(total_vram / (1024 * 1024))); - total_vram = 0; - - std::vector> turing_devices_without_mma; - for (int id = 0; id < info.device_count; ++id) { - int device_vmm = 0; - -#if defined(GGML_USE_VMM) - CUdevice device; - CU_CHECK(cuDeviceGet(&device, id)); - CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device)); - - if (device_vmm) { - CUmemAllocationProp alloc_prop = {}; - alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - alloc_prop.location.id = id; - CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); - } -#endif // defined(GGML_USE_VMM) - info.devices[id].vmm = !!device_vmm; - - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); - - info.default_tensor_split[id] = total_vram; - total_vram += prop.totalGlobalMem; - info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034) - info.devices[id].nsm = prop.multiProcessorCount; - info.devices[id].smpb = prop.sharedMemPerBlock; - info.devices[id].warp_size = prop.warpSize; - -#ifndef GGML_USE_MUSA - int supports_coop_launch = 0; - CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id)); - info.devices[id].supports_cooperative_launch = !!supports_coop_launch; -#else - info.devices[id].supports_cooperative_launch = false; -#endif // !(GGML_USE_MUSA) - -#if defined(GGML_USE_HIP) - info.devices[id].smpbo = prop.sharedMemPerBlock; - - info.devices[id].cc = ggml_cuda_parse_id(prop.gcnArchName); - if ((info.devices[id].cc & 0xff00) == 0x0) { - GGML_LOG_WARN("invalid architecture ID received for device %d %s: %s cc %d.%d\n", - id, prop.name, prop.gcnArchName, prop.major, prop.minor); - - // Fallback to prop.major and prop.minor - if (prop.major > 0) { - info.devices[id].cc = GGML_CUDA_CC_OFFSET_AMD + prop.major * 0x100; - info.devices[id].cc += prop.minor * 0x10; - } - } - GGML_LOG_INFO(" Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n", - id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff, - device_vmm ? "yes" : "no", prop.warpSize, - (size_t)(prop.totalGlobalMem / (1024 * 1024))); -#elif defined(GGML_USE_MUSA) - // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs. - info.devices[id].warp_size = 32; - info.devices[id].smpbo = prop.sharedMemPerBlockOptin; - info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100; - info.devices[id].cc += prop.minor * 0x10; - GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", - id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", - (size_t)(prop.totalGlobalMem / (1024 * 1024))); -#else - info.devices[id].smpbo = prop.sharedMemPerBlockOptin; - info.devices[id].cc = 100*prop.major + 10*prop.minor; - GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", - id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", - (size_t)(prop.totalGlobalMem / (1024 * 1024))); - std::string device_name(prop.name); - if (device_name == "NVIDIA GeForce MX450") { - turing_devices_without_mma.push_back({ id, device_name }); - } else if (device_name == "NVIDIA GeForce MX550") { - turing_devices_without_mma.push_back({ id, device_name }); - } else if (device_name.substr(0, 21) == "NVIDIA GeForce GTX 16") { - turing_devices_without_mma.push_back({ id, device_name }); - } - - // Temporary performance fix: - // Setting device scheduling strategy for iGPUs with cc121 to "spinning" to avoid delays in cuda synchronize calls. - // TODO: Check for future drivers the default scheduling strategy and - // remove this call again when cudaDeviceScheduleSpin is default. - if (prop.major == 12 && prop.minor == 1) { - CUDA_CHECK(cudaSetDevice(id)); - CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin)); - } - -#endif // defined(GGML_USE_HIP) - } - - if (ggml_cuda_highest_compiled_arch(GGML_CUDA_CC_TURING) >= GGML_CUDA_CC_TURING && !turing_devices_without_mma.empty()) { - GGML_LOG_INFO("The following devices will have suboptimal performance due to a lack of tensor cores:\n"); - for (size_t device_pos = 0; device_pos < turing_devices_without_mma.size(); device_pos++) { - GGML_LOG_INFO( - " Device %d: %s\n", turing_devices_without_mma[device_pos].first, turing_devices_without_mma[device_pos].second.c_str()); - } - GGML_LOG_INFO( - "Consider compiling with CMAKE_CUDA_ARCHITECTURES=61-virtual;80-virtual and DGGML_CUDA_FORCE_MMQ to force the use of the Pascal code for Turing.\n"); - } - - for (int id = 0; id < info.device_count; ++id) { - info.default_tensor_split[id] /= total_vram; - } - - // configure logging to stdout - // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr)); - - if (getenv("GGML_CUDA_P2P") != 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)); - } - } - } - } - - return info; -} - -const ggml_cuda_device_info & ggml_cuda_info() { - static ggml_cuda_device_info info = ggml_cuda_init(); - return info; -} - -// #define DEBUG_CUDA_MALLOC - -// buffer pool for cuda (legacy) -struct ggml_cuda_pool_leg : public ggml_cuda_pool { - static const int MAX_BUFFERS = 256; - - int device; - struct ggml_cuda_buffer { - void * ptr = nullptr; - size_t size = 0; - }; - - ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {}; - size_t pool_size = 0; - - explicit ggml_cuda_pool_leg(int device) : - device(device) { - } - - ~ggml_cuda_pool_leg() { - clear_pool(); - GGML_ASSERT(pool_size == 0); - } - - void clear() override { - clear_pool(); - } - - void clear_pool() { - ggml_cuda_set_device(device); - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer & b = buffer_pool[i]; - if (b.ptr != nullptr) { - CUDA_CHECK(cudaFree(b.ptr)); - pool_size -= b.size; - b.ptr = nullptr; - b.size = 0; - } - } - } - - void * alloc(size_t size, size_t * actual_size) override { -#ifdef DEBUG_CUDA_MALLOC - int nnz = 0; - size_t max_size = 0; -#endif - size_t best_diff = 1ull << 36; - int ibest = -1; - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr != nullptr) { -#ifdef DEBUG_CUDA_MALLOC - ++nnz; - if (b.size > max_size) max_size = b.size; -#endif - if (b.size >= size) { - size_t diff = b.size - size; - if (diff < best_diff) { - best_diff = diff; - ibest = i; - if (!best_diff) { - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; - return ptr; - } - } - } - } - } - if (ibest >= 0) { - ggml_cuda_buffer& b = buffer_pool[ibest]; - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; - return ptr; - } - void * ptr; - size_t look_ahead_size = (size_t) (1.05 * size); - look_ahead_size = 256 * ((look_ahead_size + 255)/256); - ggml_cuda_set_device(device); - cudaError_t err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); - if (err == cudaErrorMemoryAllocation) { - (void)cudaGetLastError(); - const size_t cached_bytes = pool_size; - GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: alloc of %.2f MiB failed, flushing %.2f MiB of cached buffers and retrying\n", - device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0); - CUDA_CHECK(cudaDeviceSynchronize()); - clear_pool(); - err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); - if (err == cudaSuccess) { - GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device); - } - } - CUDA_CHECK(err); - *actual_size = look_ahead_size; - pool_size += look_ahead_size; -#ifdef DEBUG_CUDA_MALLOC - GGML_LOG_INFO("%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, requested %u MB\n", __func__, device, nnz, - (uint32_t)(max_size / 1024 / 1024), (uint32_t)(pool_size / 1024 / 1024), (uint32_t)(size / 1024 / 1024)); -#endif - return ptr; - } - - void free(void * ptr, size_t size) override { - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr == nullptr) { - b.ptr = ptr; - b.size = size; - return; - } - } - GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n"); - ggml_cuda_set_device(device); - CUDA_CHECK(cudaFree(ptr)); - pool_size -= size; - } -}; - -// pool with virtual memory -#if defined(GGML_USE_VMM) -struct ggml_cuda_pool_vmm : public ggml_cuda_pool { - static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB - - int device; - CUdeviceptr pool_addr = 0; - size_t pool_used = 0; - size_t pool_size = 0; - size_t granularity; -#if defined(GGML_USE_HIP) - std::vector> mappings; -#endif - - explicit ggml_cuda_pool_vmm(int device) : - device(device), - granularity(ggml_cuda_info().devices[device].vmm_granularity) { - } - - ~ggml_cuda_pool_vmm() { - if (pool_addr != 0) { -#if defined(GGML_USE_HIP) - // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285 - for (std::pair & mapping : mappings) { - CU_CHECK(cuMemUnmap(mapping.first, mapping.second)); - } -#else - CU_CHECK(cuMemUnmap(pool_addr, pool_size)); -#endif - CU_CHECK(cuMemAddressFree(pool_addr, CUDA_POOL_VMM_MAX_SIZE)); - } - } - - void * alloc(size_t size, size_t * actual_size) override { - // round up the allocation size to the alignment to ensure that all allocations are aligned for all data types - const size_t alignment = 128; - size = alignment * ((size + alignment - 1) / alignment); - - size_t avail = pool_size - pool_used; - - if (size > avail) { - // round up to the next multiple of the granularity - size_t reserve_size = size - avail; - reserve_size = granularity * ((reserve_size + granularity - 1) / granularity); - - GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE); - - // allocate more physical memory - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = device; - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0)); - - // reserve virtual address space (if not already reserved) - if (pool_addr == 0) { - CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0)); - } - - // map at the end of the pool - CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size); - CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0)); -#if defined(GGML_USE_HIP) - mappings.push_back({start_ptr, reserve_size}); -#endif - - // the memory allocation handle is no longer needed after mapping - CU_CHECK(cuMemRelease(handle)); - - // set access - CUmemAccessDesc access = {}; - access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access.location.id = device; - access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - CU_CHECK(cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1)); - - // add to the pool - pool_size += reserve_size; - - //printf("cuda pool[%d]: size increased to %llu MB (reserved %llu MB)\n", - // device, (unsigned long long) (pool_size/1024/1024), - // (unsigned long long) (reserve_size/1024/1024)); - } - - GGML_ASSERT(pool_addr != 0); - - void * ptr = (void *) ((CUdeviceptr)((char *)(pool_addr) + pool_used)); - *actual_size = size; - pool_used += size; - -#ifdef DEBUG_CUDA_MALLOC - printf("cuda pool[%d]: allocated %llu bytes at %llx\n", device, (unsigned long long) size, ptr); -#endif - - return ptr; - } - - void free(void * ptr, size_t size) override { -#ifdef DEBUG_CUDA_MALLOC - printf("cuda pool[%d]: freed %llu bytes at %llx\n", device, (unsigned long long) size, ptr); -#endif - - pool_used -= size; - - // all deallocations must be in reverse order of the allocations - GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used)); - } -}; -#endif // defined(GGML_USE_VMM) - -std::unique_ptr ggml_backend_cuda_context::new_pool_for_device(int device, - [[maybe_unused]] int stream_no) { -#if defined(GGML_USE_VMM) - if (ggml_cuda_info().devices[device].vmm) { - return std::unique_ptr(new ggml_cuda_pool_vmm(device)); - } -#endif // defined(GGML_USE_VMM) - return std::unique_ptr(new ggml_cuda_pool_leg(device)); -} - -// destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error -// this lock is used to ensure that no cuBLAS handle is destroyed while a graph is being captured - -static std::mutex ggml_cuda_lock; -static std::condition_variable ggml_cuda_lock_cv; -static std::atomic ggml_cuda_lock_counter; - -ggml_backend_cuda_context::~ggml_backend_cuda_context() { - std::unique_lock lock(ggml_cuda_lock); - ggml_cuda_lock_cv.wait(lock, []{ return ggml_cuda_lock_counter.load(std::memory_order_relaxed) == 0; }); - - if (copy_event != nullptr) { - CUDA_CHECK(cudaEventDestroy(copy_event)); - } - for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { - for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { - if (streams[i][j] != nullptr) { - CUDA_CHECK(cudaStreamDestroy(streams[i][j])); - } - } - if (cublas_handles[i] != nullptr) { - CUBLAS_CHECK(cublasDestroy(cublas_handles[i])); - } -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - if (hipblaslt_handles[i] != nullptr) { - HIPBLASLT_CHECK(hipblasLtDestroy(hipblaslt_handles[i])); - } - if (hipblaslt_workspaces[i] != nullptr) { - CUDA_CHECK(cudaFree(hipblaslt_workspaces[i])); - } -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } -} - - -// cuda buffer - -struct ggml_backend_cuda_buffer_context { - int device; - void * dev_ptr = nullptr; - std::string name; - - ggml_backend_cuda_buffer_context(int device, void * dev_ptr) : - device(device), dev_ptr(dev_ptr), - name(GGML_CUDA_NAME + std::to_string(device)) { - } - - ~ggml_backend_cuda_buffer_context() { - CUDA_CHECK(cudaFree(dev_ptr)); - } -}; - -static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - delete ctx; -} - -static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer) { - return buffer->iface.free_buffer == ggml_backend_cuda_buffer_free_buffer; -} - -static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - return ctx->dev_ptr; -} - -static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - - if (tensor->view_src != NULL) { - assert(tensor->view_src->buffer->buft == buffer->buft); - return GGML_STATUS_SUCCESS; - } - - if (ggml_is_quantized(tensor->type) && tensor->view_src == nullptr && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) { - // initialize padding to 0 to avoid possible NaN values - const size_t original_size = ggml_nbytes(tensor); - const size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor); - - if (padded_size > original_size) { - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); - } - } - return GGML_STATUS_SUCCESS; -} - -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_cuda_set_device(ctx->device); - 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_cuda_set_device(ctx->device); - 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(cudaMemcpy2DAsync( - data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { - if (ggml_backend_buffer_is_cuda(src->buffer)) { - ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context; - ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context; - if (src_ctx->device == dst_ctx->device) { - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread)); - } else { -#ifdef GGML_CUDA_NO_PEER_COPY - return false; -#else - CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread)); -#endif - } - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - return true; - } - return false; - - GGML_UNUSED(buffer); -} - -static void ggml_backend_cuda_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemsetAsync(ctx->dev_ptr, value, buffer->size, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = { - /* .free_buffer = */ ggml_backend_cuda_buffer_free_buffer, - /* .get_base = */ ggml_backend_cuda_buffer_get_base, - /* .init_tensor = */ ggml_backend_cuda_buffer_init_tensor, - /* .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, -}; - -// cuda buffer type -struct ggml_backend_cuda_buffer_type_context { - int device; - std::string name; -}; - -static const char * ggml_backend_cuda_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_buffer_type_context * ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; - - return ctx->name.c_str(); -} - -static bool ggml_backend_buft_is_cuda(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_buffer_type_get_name; -} - -static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; - - ggml_cuda_set_device(buft_ctx->device); - - void * dev_ptr; - cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - GGML_LOG_ERROR("%s: allocating %.2f MiB on device %d: cudaMalloc failed: %s\n", __func__, size / 1024.0 / 1024.0, buft_ctx->device, cudaGetErrorString(err)); - return nullptr; - } - - ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr); - - return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size); -} - -static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return 128; - - GGML_UNUSED(buft); -} - -static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - size_t size = ggml_nbytes(tensor); - int64_t ne0 = tensor->ne[0]; - - if (ggml_is_quantized(tensor->type)) { - if (ne0 % MATRIX_ROW_PADDING != 0) { - GGML_ASSERT(tensor->nb[0] == ggml_element_size(tensor)); - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - } - - return size; - - GGML_UNUSED(buft); -} - -static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface = { - /* .get_name = */ ggml_backend_cuda_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cuda_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cuda_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cuda_buffer_type_get_alloc_size, - /* .is_host = */ NULL, -}; - -ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - if (device >= ggml_backend_cuda_get_device_count()) { - return nullptr; - } - - static ggml_backend_buffer_type ggml_backend_cuda_buffer_types[GGML_CUDA_MAX_DEVICES]; - - static bool ggml_backend_cuda_buffer_type_initialized = false; - - if (!ggml_backend_cuda_buffer_type_initialized) { - for (int i = 0; i < ggml_backend_cuda_get_device_count(); i++) { - ggml_backend_cuda_buffer_types[i] = { - /* .iface = */ ggml_backend_cuda_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), i), - /* .context = */ new ggml_backend_cuda_buffer_type_context{i, GGML_CUDA_NAME + std::to_string(i)}, - }; - } - ggml_backend_cuda_buffer_type_initialized = true; - } - - return &ggml_backend_cuda_buffer_types[device]; -} - -// cuda split buffer - -static int64_t get_row_rounding(const std::array & tensor_split) { - int64_t row_rounding = 0; - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { - continue; - } - - const int cc = ggml_cuda_info().devices[id].cc; - row_rounding = std::max(row_rounding, (int64_t)get_mmq_y_host(cc)); - } - return row_rounding; -} - -static void get_row_split(int64_t * row_low, int64_t * row_high, const ggml_tensor * tensor, const std::array & tensor_split, int id) { - const int64_t nrows = ggml_nrows(tensor); - const int64_t rounding = get_row_rounding(tensor_split); - - *row_low = id == 0 ? 0 : nrows*tensor_split[id]; - *row_low -= *row_low % rounding; - - if (id == ggml_backend_cuda_get_device_count() - 1) { - *row_high = nrows; - } else { - *row_high = nrows*tensor_split[id + 1]; - *row_high -= *row_high % rounding; - } -} - -static size_t ggml_nbytes_split(const struct ggml_tensor * tensor, int nrows_split) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return nrows_split*ggml_row_size(tensor->type, tensor->ne[0]); -} - -struct ggml_backend_cuda_split_buffer_type_context { - int main_device; - std::array tensor_split; - std::string name; -}; - -struct ggml_backend_cuda_split_buffer_context { - ~ggml_backend_cuda_split_buffer_context() { - for (ggml_tensor_extra_gpu * extra : tensor_extras) { - for (int id = 0; id < GGML_CUDA_MAX_DEVICES; ++id) { - for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { - if (extra->events[id][is] != nullptr) { - CUDA_CHECK(cudaEventDestroy(extra->events[id][is])); - } - } - if (extra->data_device[id] != nullptr) { - CUDA_CHECK(cudaFree(extra->data_device[id])); - } - } - delete extra; - } - } - - std::vector tensor_extras; -}; - - -static void ggml_backend_cuda_split_buffer_free_buffer(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; - delete ctx; -} - -static void * ggml_backend_cuda_split_buffer_get_base(ggml_backend_buffer_t buffer) { - // the pointers are stored in the tensor extras, this is just a dummy address and never dereferenced - return (void *)0x1000; - - GGML_UNUSED(buffer); -} - -static enum ggml_status ggml_backend_cuda_split_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - GGML_ASSERT(tensor->view_src == nullptr); // views of split tensors are not supported - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - - ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; - ctx->tensor_extras.push_back(extra); - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - // FIXME: do not crash if cudaMalloc fails - // currently, init_tensor cannot fail, it needs to be fixed in ggml-backend first - ggml_cuda_set_device(id); - char * buf; - CUDA_CHECK(ggml_cuda_device_malloc((void**)&buf, size, id)); - - // set padding to 0 to avoid possible NaN values - if (size > original_size) { - CUDA_CHECK(cudaMemset(buf + original_size, 0, size - original_size)); - } - - extra->data_device[id] = buf; - - for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { - CUDA_CHECK(cudaEventCreateWithFlags(&extra->events[id][is], cudaEventDisableTiming)); - } - } - tensor->extra = extra; - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cuda_split_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - // split tensors must always be set in their entirety at once - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - const size_t nb1 = tensor->nb[1]; - ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - const size_t offset_split = row_low*nb1; - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - const char * buf_host = (const char *)data + offset_split; - CUDA_CHECK(cudaMemcpyAsync(extra->data_device[id], buf_host, original_size, cudaMemcpyHostToDevice, cudaStreamPerThread)); - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - } -} - -static void ggml_backend_cuda_split_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - // split tensors must always be set in their entirety at once - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - const size_t nb1 = tensor->nb[1]; - ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - const size_t offset_split = row_low*nb1; - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - char * buf_host = (char *)data + offset_split; - CUDA_CHECK(cudaMemcpyAsync(buf_host, extra->data_device[id], original_size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - } -} - -static void ggml_backend_cuda_split_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - GGML_UNUSED(buffer); - GGML_UNUSED(value); -} - -static const ggml_backend_buffer_i ggml_backend_cuda_split_buffer_interface = { - /* .free_buffer = */ ggml_backend_cuda_split_buffer_free_buffer, - /* .get_base = */ ggml_backend_cuda_split_buffer_get_base, - /* .init_tensor = */ ggml_backend_cuda_split_buffer_init_tensor, - /* .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, -}; - -// cuda split buffer type - -static const char * ggml_backend_cuda_split_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; - - return ctx->name.c_str(); -} - -static bool ggml_backend_buft_is_cuda_split(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_split_buffer_type_get_name; -} - -static ggml_backend_buffer_t ggml_backend_cuda_split_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - // since we don't know the exact split after rounding, we cannot allocate the device buffers at this point - // instead, we allocate them for each tensor separately in init_tensor - // however, the size still represents the maximum cumulative size of all the device buffers after the tensors are allocated, - // as returned by get_alloc_size. this limit is enforced during tensor allocation by ggml-alloc, so it must be correct. - ggml_backend_cuda_split_buffer_context * ctx = new ggml_backend_cuda_split_buffer_context(); - - return ggml_backend_buffer_init(buft, ggml_backend_cuda_split_buffer_interface, ctx, size); -} - -static size_t ggml_backend_cuda_split_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return 128; - - GGML_UNUSED(buft); -} - -static size_t ggml_backend_cuda_split_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - size_t total_size = 0; - - const int64_t ne0 = tensor->ne[0]; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - total_size += ggml_nbytes_split(tensor, nrows_split); - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - total_size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - } - - return total_size; -} - -static bool ggml_backend_cuda_split_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return false; - - GGML_UNUSED(buft); -} - -static const ggml_backend_buffer_type_i ggml_backend_cuda_split_buffer_type_interface = { - /* .get_name = */ ggml_backend_cuda_split_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cuda_split_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cuda_split_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cuda_split_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_cuda_split_buffer_type_is_host, -}; - -// Communication context for multi-GPU AllReduce during tensor parallelism. -// -// Created once per meta backend instance. Resources for the selected mode -// (NCCL communicators or the internal AllReduce pipeline) are initialised -// eagerly during comm_init so any init failure surfaces at startup rather -// than mid-run. -struct ggml_backend_cuda_comm_context { - using try_allreduce_fn = bool(*)(ggml_backend_cuda_comm_context *, struct ggml_tensor **); - - std::vector backends; - std::vector dev_ids; - - // Set by the init chain (comm_init_{nccl, internal, none}) to one of - // try_allreduce_{nccl, internal, butterfly}. nccl needs `comms`, - // internal needs `ar_pipeline`, butterfly needs nothing. Per-call - // failures return false; the meta backend's generic implementation then - // handles that call. - try_allreduce_fn try_allreduce = nullptr; - - ggml_cuda_ar_pipeline * ar_pipeline = nullptr; - -#ifdef GGML_USE_NCCL - std::vector comms; -#endif // GGML_USE_NCCL - - ~ggml_backend_cuda_comm_context() { -#ifdef GGML_USE_NCCL - for (ncclComm_t comm : comms) { - NCCL_CHECK(ncclCommDestroy(comm)); - } -#endif // GGML_USE_NCCL - ggml_cuda_ar_pipeline_free(ar_pipeline); - } -}; - -#ifdef GGML_USE_NCCL -// AllReduce via NCCL. Reduces as FP32 for small tensors and BF16 for large -// tensors (bandwidth-bound), then converts back to FP32. -static bool ggml_backend_cuda_comm_allreduce_nccl( - ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { - 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; - } - - const size_t n_backends = comm_ctx->backends.size(); - - for (size_t i = 0; i < n_backends; ++i) { - GGML_ASSERT(tensors[i] != nullptr); - GGML_ASSERT(ggml_nelements(tensors[i]) == ne); - GGML_ASSERT(ggml_is_contiguously_allocated(tensors[i])); - } - - // For small tensors, simply reduce them as FP32. - // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0. - if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) { - for (size_t i = 0; i < n_backends; ++i) { - if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - ggml_cuda_set_device(cuda_ctx->device); - CUDA_CHECK(cudaMemsetAsync(tensors[i]->data, 0, ggml_nbytes(tensors[i]), cuda_ctx->stream())); - } - } - NCCL_CHECK(ncclGroupStart()); - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); - } - NCCL_CHECK(ncclGroupEnd()); - 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 tmp[GGML_CUDA_MAX_DEVICES]; - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - tmp[i].pool = &cuda_ctx->pool(); - tmp[i].alloc(ne); - - ggml_cuda_set_device(cuda_ctx->device); - if (tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) { - to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); - } else { - CUDA_CHECK(cudaMemsetAsync(tmp[i].get(), 0, ne * sizeof(nv_bfloat16), cuda_ctx->stream())); - } - CUDA_CHECK(cudaGetLastError()); - } - - NCCL_CHECK(ncclGroupStart()); - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); - } - NCCL_CHECK(ncclGroupEnd()); - - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; - - ggml_cuda_set_device(cuda_ctx->device); - to_fp32(tmp[i].get(), (float *) tensors[i]->data, ne, cuda_ctx->stream()); - CUDA_CHECK(cudaGetLastError()); - } - - return true; -} -#endif // GGML_USE_NCCL - -// Run the internal AR pipeline. Returns false on unsupported / failed input -// -- the caller decides whether to abort (env-forced) or fall back silently. -static bool ggml_backend_cuda_comm_allreduce_internal( - ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { - GGML_ASSERT(comm_ctx->ar_pipeline != nullptr); - - const size_t n_backends = comm_ctx->backends.size(); - GGML_ASSERT(n_backends == 2); - GGML_ASSERT(tensors[0] != nullptr); - - const int64_t ne = ggml_nelements(tensors[0]); - const ggml_type type = tensors[0]->type; - - if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) { - GGML_LOG_DEBUG("%s: internal unsupported: type=%d\n", __func__, (int) type); - return false; - } - - if (ne == 0) { - return true; - } - - for (size_t i = 0; i < n_backends; ++i) { - if (tensors[i] == nullptr) { - GGML_LOG_ERROR("%s: internal failed: tensor[%zu] is null\n", __func__, i); - return false; - } - if (ggml_nelements(tensors[i]) != ne || tensors[i]->type != type) { - GGML_LOG_ERROR("%s: internal failed: tensor[%zu] ne=%" PRId64 " type=%d expected ne=%" PRId64 " type=%d\n", - __func__, i, ggml_nelements(tensors[i]), (int) tensors[i]->type, ne, (int) type); - return false; - } - if (!ggml_is_contiguously_allocated(tensors[i])) { - GGML_LOG_DEBUG("%s: internal unsupported: tensor[%zu] is not contiguously allocated: ne=%" PRId64 " nbytes=%zu packed=%zu type=%d\n", - __func__, i, ne, ggml_nbytes(tensors[i]), - (size_t) ne * ggml_type_size(type) / ggml_blck_size(type), (int) type); - return false; - } - if (((uintptr_t) tensors[i]->data & 0xF) != 0) { - GGML_LOG_DEBUG("%s: internal unsupported: tensor[%zu] data pointer is not 16-byte aligned: %p type=%d ne=%" PRId64 "\n", - __func__, i, tensors[i]->data, (int) type, ne); - return false; - } - GGML_ASSERT((ggml_nbytes(tensors[i]) & 0xF) == 0); - } - - return ggml_cuda_ar_allreduce(comm_ctx->ar_pipeline, comm_ctx->backends.data(), tensors); -} - -// --------------------------------------------------------------------------- -// Per-call dispatch -- three variants, one per backend. Each is set as -// comm_ctx->try_allreduce by the matching init step. Per-call failure -// returns false; the meta backend's generic implementation handles that call. -// --------------------------------------------------------------------------- - -#ifdef GGML_USE_NCCL -static bool ggml_backend_cuda_comm_try_allreduce_nccl( - ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { - return ggml_backend_cuda_comm_allreduce_nccl(comm_ctx, tensors); -} -#endif // GGML_USE_NCCL - -static bool ggml_backend_cuda_comm_try_allreduce_internal( - ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { - return ggml_backend_cuda_comm_allreduce_internal(comm_ctx, tensors); -} - -static bool ggml_backend_cuda_comm_try_allreduce_butterfly( - ggml_backend_cuda_comm_context *, struct ggml_tensor **) { - return false; -} - -static void ggml_backend_cuda_comm_free(void * comm_ctx_v) { - if (comm_ctx_v == nullptr) { - return; - } - delete static_cast(comm_ctx_v); -} - -// --------------------------------------------------------------------------- -// Init -- chained nccl -> internal -> none. Each step tries to bring up its -// resource; on failure it warns and recurses into the next step. -// --------------------------------------------------------------------------- -static void ggml_backend_cuda_comm_init_none(ggml_backend_cuda_comm_context * ret) { - ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_butterfly; -} - -static void ggml_backend_cuda_comm_init_internal(ggml_backend_cuda_comm_context * ret) { - ret->ar_pipeline = ggml_cuda_ar_pipeline_init(ret->dev_ids.data(), ret->dev_ids.size()); - if (ret->ar_pipeline) { - ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_internal; - return; - } - - // Clear sticky CUDA error from the failed init. - (void) cudaGetLastError(); - GGML_LOG_WARN("internal AllReduce init failed (n_devices != 2?); " - "falling back to meta-backend butterfly\n"); - ggml_backend_cuda_comm_init_none(ret); -} - -static void ggml_backend_cuda_comm_init_nccl(ggml_backend_cuda_comm_context * ret) { -#ifdef GGML_USE_NCCL - const size_t n = ret->dev_ids.size(); - ret->comms.resize(n); - ncclResult_t rc = ncclCommInitAll(ret->comms.data(), (int) n, ret->dev_ids.data()); - if (rc == ncclSuccess) { - ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl; - return; - } - - ret->comms.clear(); - GGML_LOG_WARN("NCCL init failed (%s); falling back to internal AllReduce\n", - ncclGetErrorString(rc)); -#else // GGML_USE_NCCL -#ifndef GGML_USE_HIP - GGML_LOG_WARN("NCCL not compiled in; falling back to internal AllReduce. " - "Recompile with -DGGML_CUDA_NCCL=ON for best multi-GPU performance.\n"); -#endif // !GGML_USE_HIP -#endif // GGML_USE_NCCL - - ggml_backend_cuda_comm_init_internal(ret); -} - -// Top-level init. Picks one of the three init paths based on -// GGML_CUDA_ALLREDUCE (or the platform default) and lets the chain handle -// any fallback. Unrecognised env values warn and fall through to the -// platform default. -static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_backends) { - for (size_t i = 0; i < n_backends; i++) { - if (!ggml_backend_is_cuda(backends[i])) { - return nullptr; - } - } - - auto * ret = new ggml_backend_cuda_comm_context; - ret->backends.assign(backends, backends + n_backends); - ret->dev_ids.reserve(n_backends); - for (size_t i = 0; i < n_backends; i++) { - ret->dev_ids.push_back(static_cast(backends[i]->context)->device); - } - - const char * env = getenv("GGML_CUDA_ALLREDUCE"); - if (!env) { - // Platform default: Linux uses NCCL, otherwise (generally Windows) internal -#if defined(__linux__) - ggml_backend_cuda_comm_init_nccl(ret); -#else - ggml_backend_cuda_comm_init_internal(ret); -#endif // defined(__linux__) - } else { - std::string env_str(env); - if (env_str == "nccl") { - ggml_backend_cuda_comm_init_nccl(ret); - } else if (env_str == "internal") { - ggml_backend_cuda_comm_init_internal(ret); - } else if (env_str == "none") { - ggml_backend_cuda_comm_init_none(ret); - } else { - GGML_LOG_WARN("unknown GGML_CUDA_ALLREDUCE value: %s\n", env); - ggml_backend_cuda_comm_init_none(ret); - } - } - - return ret; -} - -// Top-level dispatch -- calls the function pointer chosen by comm_init. -// Returns false to let the meta-backend's butterfly run. -static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) { - if (comm_ctx_v == nullptr) { - return false; - } - auto * comm_ctx = static_cast(comm_ctx_v); - return comm_ctx->try_allreduce(comm_ctx, tensors); -} - -ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - static std::map>, struct ggml_backend_buffer_type> buft_map; - - std::array tensor_split_arr = {}; - - bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + GGML_CUDA_MAX_DEVICES, [](float x) { return x == 0.0f; }); - if (all_zero) { - tensor_split_arr = ggml_cuda_info().default_tensor_split; - } else { - float split_sum = 0.0f; - for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { - tensor_split_arr[i] = split_sum; - split_sum += tensor_split[i]; - } - for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { - tensor_split_arr[i] /= split_sum; - } - } - - auto it = buft_map.find({main_device, tensor_split_arr}); - if (it != buft_map.end()) { - return &it->second; - } - auto * ctx = new ggml_backend_cuda_split_buffer_type_context{ - main_device, - tensor_split_arr, - GGML_CUDA_NAME + std::to_string(main_device) + "_Split", - }; - - struct ggml_backend_buffer_type buft { - /* .iface = */ ggml_backend_cuda_split_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), main_device), - /* .context = */ ctx, - }; - - auto result = buft_map.emplace(std::make_pair(main_device, tensor_split_arr), buft); - return &result.first->second; -} - -// host buffer type - -static const char * ggml_backend_cuda_host_buffer_type_name(ggml_backend_buffer_type_t buft) { - return GGML_CUDA_NAME "_Host"; - - GGML_UNUSED(buft); -} - -static bool ggml_backend_buft_is_cuda_host(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; -} - -static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { - CUDA_CHECK(cudaFreeHost(buffer->context)); -} - -static void * ggml_cuda_host_malloc(size_t size) { - if (getenv("GGML_CUDA_NO_PINNED") != nullptr) { - return nullptr; - } - - void * ptr = nullptr; - cudaError_t err = cudaMallocHost((void **) &ptr, size); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - GGML_LOG_DEBUG("%s: failed to allocate %.2f MiB of pinned memory: %s\n", __func__, - size / 1024.0 / 1024.0, cudaGetErrorString(err)); - return nullptr; - } - - return ptr; -} - -static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - void * ptr = ggml_cuda_host_malloc(size); - - if (ptr == nullptr) { - // fallback to cpu buffer - return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); - } - - ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); - buffer->buft = buft; - buffer->iface.free_buffer = ggml_backend_cuda_host_buffer_free_buffer; - - return buffer; -} - -ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { - static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cuda_host_buffer_type_name, - /* .alloc_buffer = */ ggml_backend_cuda_host_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, - /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, - }, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), 0), - /* .context = */ nullptr, - }; - - return &ggml_backend_cuda_buffer_type_host; -} - -//static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) { -// return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; -//} - -/// kernels - -typedef void (*ggml_cuda_op_mul_mat_t)( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, - const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, - const int64_t src1_padded_row_size, cudaStream_t stream); - -#ifndef GGML_CUDA_PEER_MAX_BATCH_SIZE -#define GGML_CUDA_PEER_MAX_BATCH_SIZE 128 -#endif // GGML_CUDA_PEER_MAX_BATCH_SIZE - -#define MUL_MAT_SRC1_COL_STRIDE 128 - -static cudaError_t ggml_cuda_cpy_tensor_2d( - void * dst, const struct ggml_tensor * src, int64_t i3, int64_t i2, int64_t i1_low, int64_t i1_high, cudaStream_t stream) { - - const char * src_ptr = (const char *) src->data; - char * dst_ptr = (char *) dst; - - const int64_t ne0 = src->ne[0]; - const int64_t nb0 = src->nb[0]; - const int64_t nb1 = src->nb[1]; - const int64_t nb2 = src->nb[2]; - const int64_t nb3 = src->nb[3]; - const enum ggml_type type = src->type; - const int64_t ts = ggml_type_size(type); - const int64_t bs = ggml_blck_size(type); - const int64_t i1_diff = i1_high - i1_low; - - const char * x = src_ptr + i1_low*nb1 + i2*nb2 + i3*nb3; - if (nb0 == ts && nb1 == ts*ne0/bs) { - return cudaMemcpyAsync(dst_ptr, x, i1_diff*nb1, cudaMemcpyDeviceToDevice, stream); - } else if (nb0 == ts) { - return cudaMemcpy2DAsync(dst_ptr, ts*ne0/bs, x, nb1, ts*ne0/bs, i1_diff, cudaMemcpyDeviceToDevice, stream); - } else { - for (int64_t i1 = 0; i1 < i1_diff; i1++) { - const void * rx = (const void *) ((const char *) x + i1*nb1); - void * rd = (void *) (dst_ptr + i1*ts*ne0/bs); - // pretend the row is a matrix with cols=1 - cudaError_t r = cudaMemcpy2DAsync(rd, ts/bs, rx, nb0, ts/bs, ne0, cudaMemcpyDeviceToDevice, stream); - if (r != cudaSuccess) { - return r; - } - } - return cudaSuccess; - } -} - -struct cublas_force_compute_type { - bool fp32 = false; - bool fp16 = false; -}; - -static const cublas_force_compute_type & ggml_cuda_cublas_get_force_compute_type() { - static const cublas_force_compute_type compute_type = [] { - cublas_force_compute_type result; - - const bool ggml_cuda_force_cublas_compute_32f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F") != nullptr; - const bool ggml_cuda_force_cublas_compute_16f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F") != nullptr; - - GGML_ASSERT(ggml_cuda_force_cublas_compute_16f_env == false || ggml_cuda_force_cublas_compute_32f_env == false); - - if (ggml_cuda_force_cublas_compute_32f_env) { - GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F\n"); - result.fp32 = true; - } else if (ggml_cuda_force_cublas_compute_16f_env) { - GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F\n"); - result.fp16 = true; - } - - return result; - }(); - - return compute_type; -} - -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) -// hipBLASLt equivalent of the cublasGemm* calls used below. -// rocBLAS does not ship Tensile kernels for every AMD GPU arch (e.g. gfx1103 on Windows), -// while hipBLASLt covers them, so HIP builds route GEMM through hipBLASLt when available. -// Computes C = op(A) * op(B) with op(A) = A^T, op(B) = B (column-major, same as the cublas calls). -// hipBLASLt only accepts hipDataType. ROCm < 6.5 routes cudaDataType_t to the legacy -// hipblasDatatype_t enum (150/151/168), while ROCm >= 6.5 uses hipDataType (0/2/14) directly. -// Accept the raw integer value and map both numbering schemes, so this compiles on all ROCm versions. -static hipDataType ggml_hipblaslt_convert_type(int type) { - switch (type) { - case 150: return HIP_R_16F; // legacy HIPBLAS_R_16F - case 151: return HIP_R_32F; // legacy HIPBLAS_R_32F - case 168: return HIP_R_16BF; // legacy HIPBLAS_R_16B - default: - GGML_ASSERT(type == HIP_R_16F || type == HIP_R_32F || type == HIP_R_16BF); - return (hipDataType) type; +#include "ggml-cuda/softcap.cuh" +#include "ggml-cuda/softmax.cuh" +#include "ggml-cuda/ssm-conv.cuh" +#include "ggml-cuda/ssm-scan.cuh" +#include "ggml-cuda/sum.cuh" +#include "ggml-cuda/sumrows.cuh" +#include "ggml-cuda/top-k.cuh" +#include "ggml-cuda/mean.cuh" +#include "ggml-cuda/tsembd.cuh" +#include "ggml-cuda/topk-moe.cuh" +#include "ggml-cuda/unary.cuh" +#include "ggml-cuda/upscale.cuh" +#include "ggml-cuda/wkv.cuh" +#include "ggml-cuda/gla.cuh" +#include "ggml-cuda/gated_delta_net.cuh" +#include "ggml-cuda/set.cuh" +#include "ggml-cuda/set-rows.cuh" +#include "ggml-cuda/pad_reflect_1d.cuh" +#include "ggml-cuda/solve_tri.cuh" +#include "ggml-cuda/tri.cuh" +#include "ggml-cuda/cumsum.cuh" +#include "ggml-cuda/fill.cuh" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); + +#define GGML_LOG_WARN_ONCE(str) \ + { static std::once_flag warn_flag; std::call_once(warn_flag, []() { GGML_LOG_WARN(str); }); } + +[[noreturn]] +void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { + int id = -1; // in case cudaGetDevice fails + (void)cudaGetDevice(&id); + + GGML_LOG_ERROR(GGML_CUDA_NAME " error: %s\n", msg); + GGML_LOG_ERROR(" current device: %d, in function %s at %s:%d\n", id, func, file, line); + GGML_LOG_ERROR(" %s\n", stmt); + // abort with GGML_ABORT to get a stack trace + GGML_ABORT(GGML_CUDA_NAME " error"); +} + +// this is faster on Windows +// probably because the Windows CUDA libraries forget to make this check before invoking the drivers +void ggml_cuda_set_device(int device) { + int current_device; + CUDA_CHECK(cudaGetDevice(¤t_device)); + + if (device == current_device) { + return; } + + CUDA_CHECK(cudaSetDevice(device)); } -static void ggml_hipblaslt_gemm( - ggml_backend_cuda_context & ctx, cudaStream_t stream, - int64_t m, int64_t n, int64_t k, - const void * A, int type_a, int64_t lda, int64_t stride_a, - const void * B, int type_b, int64_t ldb, int64_t stride_b, - void * C, int type_c, int64_t ldc, int64_t stride_c, - int64_t batch_count) { +int ggml_cuda_get_device() { + int id; + CUDA_CHECK(cudaGetDevice(&id)); + return id; +} - const hipblasOperation_t trans_a = HIPBLAS_OP_T; - const hipblasOperation_t trans_b = HIPBLAS_OP_N; +static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device) { + ggml_cuda_set_device(device); + cudaError_t err; + if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr) { + err = cudaMallocManaged(ptr, size); +#if defined(GGML_USE_HIP) + if (err == hipSuccess) { + // hipMemAdviseSetCoarseGrain is an optional performance hint; + // ignore errors (e.g. hipErrorInvalidValue on some APU/iGPU configs). + (void)cudaMemAdvise(*ptr, size, hipMemAdviseSetCoarseGrain, device); + (void)hipGetLastError(); // clear any error + } - const float alpha = 1.0f; - const float beta = 0.0f; + // fall back to cudaMalloc if not supported (e.g. on Windows) + if (err == hipErrorNotSupported) { + static bool warned_unsupported = false; + if (!warned_unsupported) { + GGML_LOG_WARN("hipMallocManaged unsupported, falling back to hipMalloc.\n"); + warned_unsupported = true; + } - hipblasLtHandle_t lt = ctx.hipblaslt_handle(); - void * workspace = ctx.hipblaslt_workspace(ctx.device); + err = cudaMalloc(ptr, size); + } +#endif // defined(GGML_USE_HIP) + } else { + err = cudaMalloc(ptr, size); + } + return err; +} - hipblasLtMatmulDesc_t matmul_desc; - hipblasLtMatrixLayout_t layout_a, layout_b, layout_c; - hipblasLtMatmulPreference_t pref; +#if defined(GGML_USE_HIP) +static int ggml_cuda_parse_id(char devName[]) { + // A list of possible Target IDs can be found under the rocclr/clr repo in device.cpp + // these values are not stable so this is susceptible to breakage + // https://github.com/ROCm/clr/blob/amd-staging/rocclr/device/device.cpp + int archMajor = 0x0; + int archMinor = 0x0; + int archNum = GGML_CUDA_CC_OFFSET_AMD; + int archLen = strlen(devName); + char archName[archLen + 1]; - HIPBLASLT_CHECK(hipblasLtMatmulDescCreate(&matmul_desc, HIPBLAS_COMPUTE_32F, HIP_R_32F)); - HIPBLASLT_CHECK(hipblasLtMatmulDescSetAttribute(matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a))); - HIPBLASLT_CHECK(hipblasLtMatmulDescSetAttribute(matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b))); + // strip leading 'gfx' while copying into our buffer + if (archLen > 3) { + strcpy(archName, &devName[3]); + archLen -= 3; + } - // layout dims describe the stored (pre-op) matrix: A is stored [k, m], B is stored [k, n], C is [m, n] - HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_a, ggml_hipblaslt_convert_type(type_a), k, m, lda)); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_b, ggml_hipblaslt_convert_type(type_b), k, n, ldb)); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_c, ggml_hipblaslt_convert_type(type_c), m, n, ldc)); + // trim trailing :xnack- or :sramecc- statuses + archLen = strcspn(archName, ":"); + archName[archLen] = '\0'; - if (batch_count > 1) { - int batch_count_i32 = (int) batch_count; - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_a, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_a, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_a, sizeof(stride_a))); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_b, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_b, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_b, sizeof(stride_b))); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_c, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_c, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_c, sizeof(stride_c))); + // tease out the version information + if (archLen > 8) { + // versions labeled generic use '-' as delimiter + // strip the trailing "-generic" then iterate through what remains + if ((strstr(archName, "-generic"))) { + archName[archLen - 8] = '\0'; + char * pch; + if ((pch = strtok(archName, "-"))) { + archMajor = (int)strtoul(pch, 0, 16); + if ((pch = strtok(NULL, "-"))) { + archMinor = 0x10 * (int)strtoul(pch, 0, 16); + } + } + } + } else if (archLen >= 3) { + // last two digits should be the minor * 0x10 + stepping + archMinor = (int)strtoul(&archName[archLen - 2], 0, 16); + archName[archLen - 2] = '\0'; + + // only the major version remains + archMajor = (int)strtoul(archName, 0, 16); } + archNum += archMajor * 0x100; + archNum += archMinor; + return archNum; +} +#endif // defined(GGML_USE_HIP) - HIPBLASLT_CHECK(hipblasLtMatmulPreferenceCreate(&pref)); - size_t max_workspace = HIPBLASLT_WORKSPACE_SIZE; - HIPBLASLT_CHECK(hipblasLtMatmulPreferenceSetAttribute(pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace, sizeof(max_workspace))); +static ggml_cuda_device_info ggml_cuda_init() { + ggml_cuda_device_info info = {}; - hipblasLtMatmulHeuristicResult_t heuristic; - int algo_count = 0; - HIPBLASLT_CHECK(hipblasLtMatmulAlgoGetHeuristic(lt, matmul_desc, layout_a, layout_b, layout_c, layout_c, - pref, 1, &heuristic, &algo_count)); - GGML_ASSERT(algo_count > 0); + cudaError_t err = cudaGetDeviceCount(&info.device_count); + if (err != cudaSuccess) { + GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err)); + return info; + } - HIPBLASLT_CHECK(hipblasLtMatmul(lt, matmul_desc, - &alpha, A, layout_a, B, layout_b, - &beta, C, layout_c, C, layout_c, - &heuristic.algo, workspace, max_workspace, stream)); + GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES); - HIPBLASLT_CHECK(hipblasLtMatmulPreferenceDestroy(pref)); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_a)); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_b)); - HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_c)); - HIPBLASLT_CHECK(hipblasLtMatmulDescDestroy(matmul_desc)); + int64_t total_vram = 0; + for (int id = 0; id < info.device_count; ++id) { + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); + total_vram += prop.totalGlobalMem; + } + GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n", + __func__, info.device_count, (size_t)(total_vram / (1024 * 1024))); + total_vram = 0; + + std::vector> turing_devices_without_mma; + for (int id = 0; id < info.device_count; ++id) { + int device_vmm = 0; + +#if defined(GGML_USE_VMM) + CUdevice device; + CU_CHECK(cuDeviceGet(&device, id)); + CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device)); + + if (device_vmm) { + CUmemAllocationProp alloc_prop = {}; + alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + alloc_prop.location.id = id; + CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + } +#endif // defined(GGML_USE_VMM) + info.devices[id].vmm = !!device_vmm; + + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); + + info.default_tensor_split[id] = total_vram; + total_vram += prop.totalGlobalMem; + info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034) + info.devices[id].nsm = prop.multiProcessorCount; + info.devices[id].smpb = prop.sharedMemPerBlock; + info.devices[id].warp_size = prop.warpSize; + +#ifndef GGML_USE_MUSA + int supports_coop_launch = 0; + CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id)); + info.devices[id].supports_cooperative_launch = !!supports_coop_launch; +#else + info.devices[id].supports_cooperative_launch = false; +#endif // !(GGML_USE_MUSA) + +#if defined(GGML_USE_HIP) + info.devices[id].smpbo = prop.sharedMemPerBlock; + + info.devices[id].cc = ggml_cuda_parse_id(prop.gcnArchName); + if ((info.devices[id].cc & 0xff00) == 0x0) { + GGML_LOG_WARN("invalid architecture ID received for device %d %s: %s cc %d.%d\n", + id, prop.name, prop.gcnArchName, prop.major, prop.minor); + + // Fallback to prop.major and prop.minor + if (prop.major > 0) { + info.devices[id].cc = GGML_CUDA_CC_OFFSET_AMD + prop.major * 0x100; + info.devices[id].cc += prop.minor * 0x10; + } + } + GGML_LOG_INFO(" Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n", + id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff, + device_vmm ? "yes" : "no", prop.warpSize, + (size_t)(prop.totalGlobalMem / (1024 * 1024))); +#elif defined(GGML_USE_MUSA) + // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs. + info.devices[id].warp_size = 32; + info.devices[id].smpbo = prop.sharedMemPerBlockOptin; + info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100; + info.devices[id].cc += prop.minor * 0x10; + GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", + id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", + (size_t)(prop.totalGlobalMem / (1024 * 1024))); +#else + info.devices[id].smpbo = prop.sharedMemPerBlockOptin; + info.devices[id].cc = 100*prop.major + 10*prop.minor; + GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", + id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", + (size_t)(prop.totalGlobalMem / (1024 * 1024))); + std::string device_name(prop.name); + if (device_name == "NVIDIA GeForce MX450") { + turing_devices_without_mma.push_back({ id, device_name }); + } else if (device_name == "NVIDIA GeForce MX550") { + turing_devices_without_mma.push_back({ id, device_name }); + } else if (device_name.substr(0, 21) == "NVIDIA GeForce GTX 16") { + turing_devices_without_mma.push_back({ id, device_name }); + } + + // Temporary performance fix: + // Setting device scheduling strategy for iGPUs with cc121 to "spinning" to avoid delays in cuda synchronize calls. + // TODO: Check for future drivers the default scheduling strategy and + // remove this call again when cudaDeviceScheduleSpin is default. + if (prop.major == 12 && prop.minor == 1) { + CUDA_CHECK(cudaSetDevice(id)); + CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin)); + } + +#endif // defined(GGML_USE_HIP) + } + + if (ggml_cuda_highest_compiled_arch(GGML_CUDA_CC_TURING) >= GGML_CUDA_CC_TURING && !turing_devices_without_mma.empty()) { + GGML_LOG_INFO("The following devices will have suboptimal performance due to a lack of tensor cores:\n"); + for (size_t device_pos = 0; device_pos < turing_devices_without_mma.size(); device_pos++) { + GGML_LOG_INFO( + " Device %d: %s\n", turing_devices_without_mma[device_pos].first, turing_devices_without_mma[device_pos].second.c_str()); + } + GGML_LOG_INFO( + "Consider compiling with CMAKE_CUDA_ARCHITECTURES=61-virtual;80-virtual and DGGML_CUDA_FORCE_MMQ to force the use of the Pascal code for Turing.\n"); + } + + for (int id = 0; id < info.device_count; ++id) { + info.default_tensor_split[id] /= total_vram; + } + + // configure logging to stdout + // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr)); + + if (getenv("GGML_CUDA_P2P") != 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)); + } + } + } + } + + return info; } -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) -static void ggml_cuda_op_mul_mat_cublas( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, - const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, - const int64_t src1_padded_row_size, cudaStream_t stream) { - - GGML_ASSERT(src0_dd_i != nullptr); - GGML_ASSERT(src1_ddf_i != nullptr); - GGML_ASSERT(dst_dd_i != nullptr); - - const int64_t ne00 = src0->ne[0]; - const int64_t ne10 = src1->ne[0]; - - const int64_t ne0 = dst->ne[0]; - - const int64_t row_diff = row_high - row_low; - - int id = ggml_cuda_get_device(); - - // the main device has a larger memory buffer to hold the results from all GPUs - // ldc == nrows of the matrix that cuBLAS writes into - int64_t ldc = id == ctx.device ? ne0 : row_diff; - - const int cc = ggml_cuda_info().devices[id].cc; - - const bool supports_bf16 = - (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc) || - (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); - - const bool use_fp16 = - src0->type != GGML_TYPE_NVFP4 && - (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && - ggml_is_contiguous(src0) && - row_diff == src0->ne[1] && - dst->op_params[0] == GGML_PREC_DEFAULT; - - if (supports_bf16 && src0->type == GGML_TYPE_BF16 && ggml_is_contiguous(src0) && row_diff == src0->ne[1]) { - ggml_cuda_pool_alloc src1_as_bf16(ctx.pool(id)); - if (src1->type != GGML_TYPE_BF16) { - const to_bf16_cuda_t to_bf16_cuda = ggml_get_to_bf16_cuda(src1->type); - GGML_ASSERT(to_bf16_cuda != nullptr); - size_t ne = src1_ncols*ne10; - src1_as_bf16.alloc(ne); - to_bf16_cuda(src1_ddf_i, src1_as_bf16.get(), ne, stream); - } - const nv_bfloat16 * src1_ptr = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1_ddf_i : src1_as_bf16.get(); - const nv_bfloat16 * src0_ptr = (const nv_bfloat16 *)src0_dd_i; - const float alpha_f32 = 1.0f; - const float beta_f32 = 0.0f; +const ggml_cuda_device_info & ggml_cuda_info() { + static ggml_cuda_device_info info = ggml_cuda_init(); + return info; +} + +// #define DEBUG_CUDA_MALLOC + +// buffer pool for cuda (legacy) +struct ggml_cuda_pool_leg : public ggml_cuda_pool { + static const int MAX_BUFFERS = 256; + + int device; + struct ggml_cuda_buffer { + void * ptr = nullptr; + size_t size = 0; + }; + ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {}; + size_t pool_size = 0; + + explicit ggml_cuda_pool_leg(int device) : + device(device) { + } + + ~ggml_cuda_pool_leg() { + clear_pool(); + GGML_ASSERT(pool_size == 0); + } + + void clear() override { + clear_pool(); + } + + void clear_pool() { + ggml_cuda_set_device(device); + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer & b = buffer_pool[i]; + if (b.ptr != nullptr) { + CUDA_CHECK(cudaFree(b.ptr)); + pool_size -= b.size; + b.ptr = nullptr; + b.size = 0; + } + } + } + + void * alloc(size_t size, size_t * actual_size) override { +#ifdef DEBUG_CUDA_MALLOC + int nnz = 0; + size_t max_size = 0; +#endif + size_t best_diff = 1ull << 36; + int ibest = -1; + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer& b = buffer_pool[i]; + if (b.ptr != nullptr) { +#ifdef DEBUG_CUDA_MALLOC + ++nnz; + if (b.size > max_size) max_size = b.size; +#endif + if (b.size >= size) { + size_t diff = b.size - size; + if (diff < best_diff) { + best_diff = diff; + ibest = i; + if (!best_diff) { + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + } + } + } + } + if (ibest >= 0) { + ggml_cuda_buffer& b = buffer_pool[ibest]; + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + void * ptr; + size_t look_ahead_size = (size_t) (1.05 * size); + look_ahead_size = 256 * ((look_ahead_size + 255)/256); + ggml_cuda_set_device(device); + cudaError_t err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); + if (err == cudaErrorMemoryAllocation) { + (void)cudaGetLastError(); + const size_t cached_bytes = pool_size; + GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: alloc of %.2f MiB failed, flushing %.2f MiB of cached buffers and retrying\n", + device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0); + CUDA_CHECK(cudaDeviceSynchronize()); + clear_pool(); + err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); + if (err == cudaSuccess) { + GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device); + } + } + CUDA_CHECK(err); + *actual_size = look_ahead_size; + pool_size += look_ahead_size; +#ifdef DEBUG_CUDA_MALLOC + GGML_LOG_INFO("%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, requested %u MB\n", __func__, device, nnz, + (uint32_t)(max_size / 1024 / 1024), (uint32_t)(pool_size / 1024 / 1024), (uint32_t)(size / 1024 / 1024)); +#endif + return ptr; + } + + void free(void * ptr, size_t size) override { + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer& b = buffer_pool[i]; + if (b.ptr == nullptr) { + b.ptr = ptr; + b.size = size; + return; + } + } + GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n"); + ggml_cuda_set_device(device); + CUDA_CHECK(cudaFree(ptr)); + pool_size -= size; + } +}; + +// pool with virtual memory +#if defined(GGML_USE_VMM) +struct ggml_cuda_pool_vmm : public ggml_cuda_pool { + static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB + + int device; + CUdeviceptr pool_addr = 0; + size_t pool_used = 0; + size_t pool_size = 0; + size_t granularity; +#if defined(GGML_USE_HIP) + std::vector> mappings; +#endif + + explicit ggml_cuda_pool_vmm(int device) : + device(device), + granularity(ggml_cuda_info().devices[device].vmm_granularity) { + } + + ~ggml_cuda_pool_vmm() { + if (pool_addr != 0) { +#if defined(GGML_USE_HIP) + // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285 + for (std::pair & mapping : mappings) { + CU_CHECK(cuMemUnmap(mapping.first, mapping.second)); + } +#else + CU_CHECK(cuMemUnmap(pool_addr, pool_size)); +#endif + CU_CHECK(cuMemAddressFree(pool_addr, CUDA_POOL_VMM_MAX_SIZE)); + } + } + + void * alloc(size_t size, size_t * actual_size) override { + // round up the allocation size to the alignment to ensure that all allocations are aligned for all data types + const size_t alignment = 128; + size = alignment * ((size + alignment - 1) / alignment); + + size_t avail = pool_size - pool_used; + + if (size > avail) { + // round up to the next multiple of the granularity + size_t reserve_size = size - avail; + reserve_size = granularity * ((reserve_size + granularity - 1) / granularity); + + GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE); + + // allocate more physical memory + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = device; + CUmemGenericAllocationHandle handle; + CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0)); + + // reserve virtual address space (if not already reserved) + if (pool_addr == 0) { + CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0)); + } + + // map at the end of the pool + CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size); + CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0)); +#if defined(GGML_USE_HIP) + mappings.push_back({start_ptr, reserve_size}); +#endif + + // the memory allocation handle is no longer needed after mapping + CU_CHECK(cuMemRelease(handle)); + + // set access + CUmemAccessDesc access = {}; + access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access.location.id = device; + access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + CU_CHECK(cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1)); + + // add to the pool + pool_size += reserve_size; + + //printf("cuda pool[%d]: size increased to %llu MB (reserved %llu MB)\n", + // device, (unsigned long long) (pool_size/1024/1024), + // (unsigned long long) (reserve_size/1024/1024)); + } + + GGML_ASSERT(pool_addr != 0); + + void * ptr = (void *) ((CUdeviceptr)((char *)(pool_addr) + pool_used)); + *actual_size = size; + pool_used += size; + +#ifdef DEBUG_CUDA_MALLOC + printf("cuda pool[%d]: allocated %llu bytes at %llx\n", device, (unsigned long long) size, ptr); +#endif + + return ptr; + } + + void free(void * ptr, size_t size) override { +#ifdef DEBUG_CUDA_MALLOC + printf("cuda pool[%d]: freed %llu bytes at %llx\n", device, (unsigned long long) size, ptr); +#endif + + pool_used -= size; + + // all deallocations must be in reverse order of the allocations + GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used)); + } +}; +#endif // defined(GGML_USE_VMM) + +std::unique_ptr ggml_backend_cuda_context::new_pool_for_device(int device, + [[maybe_unused]] int stream_no) { +#if defined(GGML_USE_VMM) + if (ggml_cuda_info().devices[device].vmm) { + return std::unique_ptr(new ggml_cuda_pool_vmm(device)); + } +#endif // defined(GGML_USE_VMM) + return std::unique_ptr(new ggml_cuda_pool_leg(device)); +} + +// destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error +// this lock is used to ensure that no cuBLAS handle is destroyed while a graph is being captured + +static std::mutex ggml_cuda_lock; +static std::condition_variable ggml_cuda_lock_cv; +static std::atomic ggml_cuda_lock_counter; + +ggml_backend_cuda_context::~ggml_backend_cuda_context() { + std::unique_lock lock(ggml_cuda_lock); + ggml_cuda_lock_cv.wait(lock, []{ return ggml_cuda_lock_counter.load(std::memory_order_relaxed) == 0; }); + + if (copy_event != nullptr) { + CUDA_CHECK(cudaEventDestroy(copy_event)); + } + for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { + for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { + if (streams[i][j] != nullptr) { + CUDA_CHECK(cudaStreamDestroy(streams[i][j])); + } + } + if (cublas_handles[i] != nullptr) { + CUBLAS_CHECK(cublasDestroy(cublas_handles[i])); + } #if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - ggml_cuda_pool_alloc dst_bf16(ctx.pool(id), row_diff*src1_ncols); - ggml_hipblaslt_gemm(ctx, stream, - row_diff, src1_ncols, ne10, - src0_ptr, CUDA_R_16BF, ne00, 0, - src1_ptr, CUDA_R_16BF, ne10, 0, - dst_bf16.get(), CUDA_R_16BF, ldc, 0, - 1); - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); - to_fp32_cuda(dst_bf16.get(), dst_dd_i, row_diff*src1_ncols, stream); - GGML_UNUSED_VARS(alpha_f32, beta_f32); -#else - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha_f32, src0_ptr, CUDA_R_16BF, ne00, - src1_ptr, CUDA_R_16BF, ne10, - &beta_f32, dst_dd_i, CUDA_R_32F, ldc, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + if (hipblaslt_handles[i] != nullptr) { + HIPBLASLT_CHECK(hipblasLtDestroy(hipblaslt_handles[i])); + } + if (hipblaslt_workspaces[i] != nullptr) { + CUDA_CHECK(cudaFree(hipblaslt_workspaces[i])); + } #endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } else if (fast_fp16_hardware_available(cc) && use_fp16) { - // convert src0 and src1 to fp16, multiply as fp16, convert dst to fp32 - ggml_cuda_pool_alloc src0_as_f16(ctx.pool(id)); - if (src0->type != GGML_TYPE_F16) { - const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src0->type); - GGML_ASSERT(to_fp16_cuda != nullptr); - size_t ne = row_diff*ne00; - src0_as_f16.alloc(ne); - to_fp16_cuda(src0_dd_i, src0_as_f16.get(), ne, stream); - } - const half * src0_ptr = src0->type == GGML_TYPE_F16 ? (const half *) src0_dd_i : src0_as_f16.get(); - - ggml_cuda_pool_alloc src1_as_f16(ctx.pool(id)); - if (src1->type != GGML_TYPE_F16) { - const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src1->type); - GGML_ASSERT(to_fp16_cuda != nullptr); - size_t ne = src1_ncols*ne10; - src1_as_f16.alloc(ne); - to_fp16_cuda(src1_ddf_i, src1_as_f16.get(), ne, stream); - } - const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16.get(); - - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - - const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); - - if (!force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) - || GGML_CUDA_CC_IS_RDNA4(cc) - || cc == GGML_CUDA_CC_VOLTA - || force_compute_type.fp32)) - { - const float alpha = 1.0f; - const float beta = 0.0f; -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - GGML_UNUSED_VARS(alpha, beta); - ggml_hipblaslt_gemm(ctx, stream, - row_diff, src1_ncols, ne10, - src0_ptr, CUDA_R_16F, ne00, 0, - src1_ptr, CUDA_R_16F, ne10, 0, - dst_dd_i, CUDA_R_32F, ldc, 0, - 1); -#else - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha, src0_ptr, CUDA_R_16F, ne00, - src1_ptr, CUDA_R_16F, ne10, - &beta, dst_dd_i, CUDA_R_32F, ldc, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } else { - ggml_cuda_pool_alloc dst_f16(ctx.pool(id), row_diff*src1_ncols); - - const half alpha_f16 = 1.0f; - const half beta_f16 = 0.0f; - -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - GGML_UNUSED_VARS(alpha_f16, beta_f16); - ggml_hipblaslt_gemm(ctx, stream, - row_diff, src1_ncols, ne10, - src0_ptr, CUDA_R_16F, ne00, 0, - src1_ptr, CUDA_R_16F, ne10, 0, - dst_f16.get(), CUDA_R_16F, ldc, 0, - 1); -#else - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha_f16, src0_ptr, CUDA_R_16F, ne00, - src1_ptr, CUDA_R_16F, ne10, - &beta_f16, dst_f16.get(), CUDA_R_16F, ldc, - CUBLAS_COMPUTE_16F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16); - to_fp32_cuda(dst_f16.get(), dst_dd_i, row_diff*src1_ncols, stream); - } - } else { - ggml_cuda_pool_alloc src0_ddq_as_f32(ctx.pool(id)); - ggml_cuda_pool_alloc src1_ddq_as_f32(ctx.pool(id)); - - if (src0->type != GGML_TYPE_F32) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src0->type); - GGML_ASSERT(to_fp32_cuda != nullptr); - src0_ddq_as_f32.alloc(row_diff*ne00); - to_fp32_cuda(src0_dd_i, src0_ddq_as_f32.get(), row_diff*ne00, stream); - } - if (src1->type != GGML_TYPE_F32) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src1->type); - GGML_ASSERT(to_fp32_cuda != nullptr); - src1_ddq_as_f32.alloc(src1_ncols*ne10); - to_fp32_cuda(src1_ddf_i, src1_ddq_as_f32.get(), src1_ncols*ne10, stream); - } - - const float * src0_ddf_i = src0->type == GGML_TYPE_F32 ? (const float *) src0_dd_i : src0_ddq_as_f32.get(); - const float * src1_ddf1_i = src1->type == GGML_TYPE_F32 ? (const float *) src1_ddf_i : src1_ddq_as_f32.get(); - - const float alpha = 1.0f; - const float beta = 0.0f; - -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - GGML_UNUSED_VARS(alpha, beta); - ggml_hipblaslt_gemm(ctx, stream, - row_diff, src1_ncols, ne10, - src0_ddf_i, CUDA_R_32F, ne00, 0, - src1_ddf1_i, CUDA_R_32F, ne10, 0, - dst_dd_i, CUDA_R_32F, ldc, 0, - 1); -#else - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - CUBLAS_CHECK( - cublasSgemm(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha, src0_ddf_i, ne00, - src1_ddf1_i, ne10, - &beta, dst_dd_i, ldc)); -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } - - GGML_UNUSED_VARS(dst, src1_ddq_i, src1_padded_row_size); -} - -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) { - -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - // cudaMemcpy2DAsync may fail with copies between vmm pools of different devices - cudaMemcpy3DPeerParms p = {}; - p.dstDevice = dstDevice; - p.dstPtr = make_cudaPitchedPtr(dst, dpitch, dpitch, height); - p.srcDevice = srcDevice; - p.srcPtr = make_cudaPitchedPtr(src, spitch, spitch, height); - p.extent = make_cudaExtent(width, height, 1); - return cudaMemcpy3DPeerAsync(&p, stream); -#else - // HIP does not support cudaMemcpy3DPeerAsync or vmm pools - GGML_UNUSED(dstDevice); - GGML_UNUSED(srcDevice); - return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, cudaMemcpyDeviceToDevice, stream); -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -} - -static void ggml_cuda_op_mul_mat( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ggml_cuda_op_mul_mat_t op, - quantize_cuda_t quantize_src1) { - - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - const int64_t ne02 = src0->ne[2]; - const int64_t ne03 = src0->ne[3]; - - const int64_t ne10 = src1->ne[0]; - const int64_t ne11 = src1->ne[1]; - const int64_t ne12 = src1->ne[2]; - const int64_t ne13 = src1->ne[3]; - const int64_t nrows1 = ggml_nrows(src1); - - const int64_t ne0 = dst->ne[0]; - const int64_t ne1 = dst->ne[1]; - - // const int64_t nb10 = src1->nb[0]; - const int64_t nb11 = src1->nb[1]; - const int64_t nb12 = src1->nb[2]; - const int64_t nb13 = src1->nb[3]; - - const int64_t nb2 = dst->nb[2]; - const int64_t nb3 = dst->nb[3]; - - ggml_backend_cuda_buffer_context * src1_ctx = (ggml_backend_cuda_buffer_context *) src1->buffer->context; - ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *) dst->buffer->context; - - GGML_ASSERT(src1->type == GGML_TYPE_F32 || (src1->ne[2] == 1 && src1->ne[3] == 1)); - - GGML_ASSERT(ne12 % ne02 == 0); - GGML_ASSERT(ne13 % ne03 == 0); - - const int64_t i02_divisor = ne12 / ne02; - const int64_t i03_divisor = ne13 / ne03; - - const size_t src0_ts = ggml_type_size(src0->type); - const size_t src0_bs = ggml_blck_size(src0->type); - const size_t q8_1_ts = sizeof(block_q8_1); - const size_t q8_1_bs = QK8_1; - - const bool src0_is_contiguous = ggml_is_contiguous(src0); - const bool src1_is_contiguous = ggml_is_contiguous(src1); - - const int64_t src1_padded_col_size = GGML_PAD(ne10, MATRIX_ROW_PADDING); - - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); - GGML_ASSERT(!(split && ne02 > 1)); - GGML_ASSERT(!(split && ne03 > 1)); - GGML_ASSERT(!(split && ne02 < ne12)); - GGML_ASSERT(!(split && ne03 < ne13)); - - ggml_tensor_extra_gpu * src0_extra = split ? (ggml_tensor_extra_gpu *) src0->extra : nullptr; - - - std::array tensor_split; - if (split) { - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; - tensor_split = buft_ctx->tensor_split; - } - - struct dev_data { - int cc; - - ggml_cuda_pool_alloc src0_dd_alloc; - ggml_cuda_pool_alloc src1_ddf_alloc; - ggml_cuda_pool_alloc src1_ddq_alloc; - ggml_cuda_pool_alloc dst_dd_alloc; - - char * src0_dd = nullptr; - float * src1_ddf = nullptr; // float - char * src1_ddq = nullptr; // q8_1 - float * dst_dd = nullptr; - - int64_t row_low; - int64_t row_high; - }; - - dev_data dev[GGML_CUDA_MAX_DEVICES]; - - int used_devices = 0; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - dev[id].cc = ggml_cuda_info().devices[id].cc; - - // by default, use all rows - dev[id].row_low = 0; - dev[id].row_high = ne01; - - // for multi GPU, get the row boundaries from tensor split - // and round to mul_mat_q tile sizes - if (split) { - const int64_t rounding = get_row_rounding(tensor_split); - - if (id != 0) { - dev[id].row_low = ne01*tensor_split[id]; - if (dev[id].row_low < ne01) { - dev[id].row_low -= dev[id].row_low % rounding; - } - } - - if (id != ggml_backend_cuda_get_device_count() - 1) { - dev[id].row_high = ne01*tensor_split[id + 1]; - if (dev[id].row_high < ne01) { - dev[id].row_high -= dev[id].row_high % rounding; - } - } - } - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { - continue; - } - - used_devices++; - - const bool src1_on_device = id == src1_ctx->device; - const bool dst_on_device = id == dst_ctx->device; - - ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, 0); - - if (src0_is_contiguous) { - dev[id].src0_dd = split ? (char *) src0_extra->data_device[id] : (char *) src0->data; - } else { - // If src0 is not contiguous it will be copied to a temporary buffer. - // This buffer needs to be cleared entirely because multiple regions will function as padding. - const size_t nbytes_data = ggml_nbytes(src0); - const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); - dev[id].src0_dd = dev[id].src0_dd_alloc.alloc(ctx.pool(id), nbytes_data + nbytes_padding); - CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd, 0, nbytes_data + nbytes_padding, stream)); - } - - // If src0 is on a temporary compute buffer (partial offloading) there may be some padding that needs to be cleared: - if (ne00 % MATRIX_ROW_PADDING != 0 && ggml_is_quantized(src0->type) && ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && src0->view_src == nullptr) { - GGML_ASSERT(ggml_is_contiguously_allocated(src0)); - GGML_ASSERT(!src0->view_src); - const size_t nbytes_data = ggml_row_size(src0->type, (dev[id].row_high - dev[id].row_low)*ne00); - const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); - CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd + nbytes_data, 0, nbytes_padding, stream)); - } - - if (src1_on_device && src1_is_contiguous) { - dev[id].src1_ddf = (float *) src1->data; - } else { - dev[id].src1_ddf = dev[id].src1_ddf_alloc.alloc(ctx.pool(id), ggml_nelements(src1)); - } - - if (quantize_src1) { - size_t src_1_ddq_size = nrows1*src1_padded_col_size*q8_1_ts/q8_1_bs; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - src_1_ddq_size += get_mmq_x_max_host(dev[id].cc)*sizeof(block_q8_1_mmq); - } - dev[id].src1_ddq = dev[id].src1_ddq_alloc.alloc(ctx.pool(id), src_1_ddq_size); - - if (src1_on_device && src1_is_contiguous) { - quantize_src1( - dev[id].src1_ddf, nullptr, dev[id].src1_ddq, src0->type, ne10, - nb11/sizeof(float), nb12/sizeof(float), nb13/sizeof(float), - src1_padded_col_size, ne11, ne12, ne13, stream); - CUDA_CHECK(cudaGetLastError()); - } - } - - if (dst_on_device) { - dev[id].dst_dd = (float *) dst->data; - } else { - const size_t size_dst_ddf = split ? (dev[id].row_high - dev[id].row_low)*ne1 : ggml_nelements(dst); - dev[id].dst_dd = dev[id].dst_dd_alloc.alloc(ctx.pool(id), size_dst_ddf); - } - } - - // if multiple devices are used they need to wait for the main device - // here an event is recorded that signals that the main device has finished calculating the input data - if (split && used_devices > 1) { - ggml_cuda_set_device(ctx.device); - CUDA_CHECK(cudaEventRecord(src0_extra->events[ctx.device][0], ctx.stream())); - } - - const int64_t src1_col_stride = split && used_devices > 1 ? MUL_MAT_SRC1_COL_STRIDE : ne11; - for (int64_t src1_col_0 = 0; src1_col_0 < ne11; src1_col_0 += src1_col_stride) { - const int64_t is = split ? (src1_col_0/src1_col_stride) % GGML_CUDA_MAX_STREAMS : 0; - const int64_t src1_ncols = src1_col_0 + src1_col_stride > ne11 ? ne11 - src1_col_0 : src1_col_stride; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { - continue; - } - - const bool src1_on_device = id == src1_ctx->device; - const bool dst_on_device = id == dst_ctx->device; - const int64_t row_diff = dev[id].row_high - dev[id].row_low; - - ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, is); - - // wait for main GPU data if necessary - if (split && (id != ctx.device || is != 0)) { - CUDA_CHECK(cudaStreamWaitEvent(stream, src0_extra->events[ctx.device][0], 0)); - } - - for (int64_t i0 = 0; i0 < ne13*ne12; ++i0) { - const int64_t i03 = i0 / ne12; - const int64_t i02 = i0 % ne12; - - size_t src1_ddq_i_offset = i0*ne11 * src1_padded_col_size*q8_1_ts/q8_1_bs; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - src1_ddq_i_offset += src1_col_0 * sizeof(block_q8_1_mmq); - } else { - src1_ddq_i_offset += src1_col_0 * src1_padded_col_size*q8_1_ts/q8_1_bs; - } - - // for split tensors the data begins at i0 == i0_offset_low - const size_t nbytes_src0_matrix = ne01*ne00*src0_ts / src0_bs; - char * src0_dd_i = dev[id].src0_dd + ((i03/i03_divisor)*ne02 + (i02/i02_divisor)) * nbytes_src0_matrix; - float * src1_ddf_i = dev[id].src1_ddf + (i0*ne11 + src1_col_0) * ne10; - char * src1_ddq_i = dev[id].src1_ddq + src1_ddq_i_offset; - float * dst_dd_i = dev[id].dst_dd + (i0*ne1 + src1_col_0) * (dst_on_device ? ne0 : row_diff); - - // the main device memory buffer can be on VRAM scratch, with space for all partial results - // in that case an offset on dst_ddf_i is needed - if (id == ctx.device) { - dst_dd_i += dev[id].row_low; // offset is 0 if no tensor split - } - - // copy src0, src1 to device if necessary - if (src1_is_contiguous) { - if (id != ctx.device) { - if (quantize_src1) { - char * src1_ddq_i_source = dev[ctx.device].src1_ddq + src1_ddq_i_offset; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - const size_t pitch = ne11*sizeof(block_q8_1_mmq); - const size_t width = src1_ncols*sizeof(block_q8_1_mmq); - const size_t height = src1_padded_col_size/(4*QK8_1); - CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync(src1_ddq_i, id, pitch, src1_ddq_i_source, ctx.device, pitch, width, height, stream)); - } else { - CUDA_CHECK(cudaMemcpyPeerAsync( - src1_ddq_i, id, src1_ddq_i_source, ctx.device, src1_ncols*src1_padded_col_size*q8_1_ts/q8_1_bs, stream)); - } - } else { - float * src1_ddf_i_source = (float *) src1->data; - src1_ddf_i_source += (i0*ne11 + src1_col_0) * ne10; - CUDA_CHECK(cudaMemcpyPeerAsync(src1_ddf_i, id, src1_ddf_i_source, ctx.device, - src1_ncols*ne10*sizeof(float), stream)); - } - } - } else if (src1_on_device && !src1_is_contiguous) { - CUDA_CHECK(ggml_cuda_cpy_tensor_2d( - src1_ddf_i, src1, i03, i02, src1_col_0, src1_col_0+src1_ncols, stream)); - } else { - GGML_ABORT("fatal error"); - } - - if (quantize_src1 && !src1_is_contiguous) { - quantize_src1( - src1_ddf_i, nullptr, src1_ddq_i, src0->type, ne10, ne10, ne11*ne10, ne12*ne11*ne10, - src1_padded_col_size, src1_ncols, 1, 1, stream); - CUDA_CHECK(cudaGetLastError()); - } - - if (src1_col_0 == 0 && !src0_is_contiguous && i03 % i03_divisor == 0 && i02 % i02_divisor == 0) { - CUDA_CHECK(ggml_cuda_cpy_tensor_2d( - src0_dd_i, src0, i03/i03_divisor, i02/i02_divisor, dev[id].row_low, dev[id].row_high, stream)); - } - - // do the computation - op(ctx, src0, src1, dst, src0_dd_i, src1_ddf_i, src1_ddq_i, dst_dd_i, - dev[id].row_low, dev[id].row_high, src1_ncols, src1_padded_col_size, stream); - CUDA_CHECK(cudaGetLastError()); - - // copy dst to host or other device if necessary - if (!dst_on_device) { - void * dst_off_device = dst->data; - if (split) { - // src0 = weight matrix is saved as a transposed matrix for better memory layout. - // dst is NOT transposed. - // The outputs of matrix matrix multiplications can therefore NOT simply be concatenated for >1 GPU. - // Instead they need to be copied to the correct slice in ne0 = dst row index. - // If dst is a vector with ne0 == 1 then you don't have to do this but it still produces correct results. - float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); - GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); - dhf_dst_i += src1_col_0*ne0 + dev[id].row_low; - CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync( - dhf_dst_i, ctx.device, ne0*sizeof(float), dst_dd_i, id, row_diff*sizeof(float), row_diff*sizeof(float), src1_ncols, stream)); - } else { - float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); - GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); - dhf_dst_i += src1_col_0*ne0; - CUDA_CHECK(cudaMemcpyAsync(dhf_dst_i, dst_dd_i, src1_ncols*ne0*sizeof(float), cudaMemcpyDeviceToDevice, stream)); - } - } - - // add event for the main device to wait on until other device is done - if (split && (id != ctx.device || is != 0)) { - CUDA_CHECK(cudaEventRecord(src0_extra->events[id][is], stream)); - } - } - } - } - - // main device waits for all other devices to be finished - if (split && ggml_backend_cuda_get_device_count() > 1) { - int64_t is_max = (ne11 + MUL_MAT_SRC1_COL_STRIDE - 1) / MUL_MAT_SRC1_COL_STRIDE; - is_max = is_max <= GGML_CUDA_MAX_STREAMS ? is_max : GGML_CUDA_MAX_STREAMS; - - ggml_cuda_set_device(ctx.device); - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if (dev[id].row_low == dev[id].row_high) { - continue; - } - for (int64_t is = 0; is < is_max; ++is) { - CUDA_CHECK(cudaStreamWaitEvent(ctx.stream(), src0_extra->events[id][is], 0)); - } - } - } -} - -static __global__ void k_compute_batched_ptrs( - const void * src0_as_f16, const void * src1_as_f16, char * dst, - const void ** ptrs_src, void ** ptrs_dst, - int64_t ne12, int64_t ne13, - int64_t ne23, - size_t nb02, size_t nb03, - size_t nb12, size_t nb13, - size_t nbd2, size_t nbd3, - int64_t r2, int64_t r3) { - const int64_t i13 = blockIdx.x * blockDim.x + threadIdx.x; - const int64_t i12 = blockIdx.y * blockDim.y + threadIdx.y; - - if (i13 >= ne13 || i12 >= ne12) { - return; - } - - const int64_t i03 = i13 / r3; - const int64_t i02 = i12 / r2; - - ptrs_src[0*ne23 + i12 + i13*ne12] = (const char *) src0_as_f16 + i02*nb02 + i03*nb03; - ptrs_src[1*ne23 + i12 + i13*ne12] = (const char *) src1_as_f16 + i12*nb12 + i13*nb13; - ptrs_dst[0*ne23 + i12 + i13*ne12] = ( char *) dst + i12*nbd2 + i13*nbd3; -} - -// Type traits for mapping ggml types to CUDA/cuBLAS types -template -struct batched_mul_mat_traits; - -template<> -struct batched_mul_mat_traits { - using cuda_type = float; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; - static inline const cudaDataType_t data_type = CUDA_R_32F; - static inline const ggml_type ggml_type_val = GGML_TYPE_F32; - static inline const float alpha = 1.0f; - static inline const float beta = 0.0f; - static inline const void* get_alpha() { static const float val = alpha; return &val; } - static inline const void* get_beta() { static const float val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp32_nc_cuda(src_type); } -}; - -template<> -struct batched_mul_mat_traits { - using cuda_type = nv_bfloat16; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; - static inline const cudaDataType_t data_type = CUDA_R_16BF; - static inline const ggml_type ggml_type_val = GGML_TYPE_BF16; - static inline const float alpha = 1.0f; - static inline const float beta = 0.0f; - static inline const void* get_alpha() { static const float val = alpha; return &val; } - static inline const void* get_beta() { static const float val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_bf16_nc_cuda(src_type); } -}; - -template<> -struct batched_mul_mat_traits { - using cuda_type = half; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_16F; - static inline const cudaDataType_t data_type = CUDA_R_16F; - static inline const ggml_type ggml_type_val = GGML_TYPE_F16; - static inline const half alpha = 1.0; - static inline const half beta = 0.0; - static inline const void* get_alpha() { static const half val = alpha; return &val; } - static inline const void* get_beta() { static const half val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } -}; - -template -static void ggml_cuda_mul_mat_batched_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - using traits = batched_mul_mat_traits; - using cuda_t = typename traits::cuda_type; - - GGML_ASSERT(!ggml_is_transposed(src0)); - GGML_ASSERT(!ggml_is_transposed(src1)); - GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft)); - GGML_ASSERT(src0->type == src0_type); - GGML_ASSERT(ggml_is_contiguous(dst)); - - // Byte offsets and tensor dimensions are currently used in an inconsistent way for dst. - // As long as dst is contiguous this does not matter though. - - GGML_TENSOR_BINARY_OP_LOCALS - - const int64_t ne_dst = ggml_nelements(dst); - cudaStream_t main_stream = ctx.stream(); - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), main_stream)); - - float * dst_ddf = (float *) dst->data; - const size_t ts_src1 = ggml_type_size(src1->type); - GGML_ASSERT(nb10 == ts_src1); - int64_t s11 = nb11 / ts_src1; - int64_t s12 = nb12 / ts_src1; - int64_t s13 = nb13 / ts_src1; - - const cuda_t * src0_ptr = nullptr; - const cuda_t * src1_ptr = nullptr; - - ggml_cuda_pool_alloc src0_alloc(ctx.pool()); - ggml_cuda_pool_alloc src1_alloc(ctx.pool()); - - bool is_src0_cont_2 = ggml_is_contiguous_2(src0); - bool is_src1_cont_2 = ggml_is_contiguous_2(src1); - - // Handle src0 - src0_ptr = (const cuda_t *) src0->data; - - // Handle src1 - convert if necessary - if (src1->type == src0_type) { - src1_ptr = (const cuda_t *) src1->data; - } else { - // Convert src1 to target type using traits conversion functions - const int64_t ne_src1 = ggml_nelements(src1); - src1_alloc.alloc(ne_src1); - - const auto convert_func = traits::get_nc_converter(src1->type); - GGML_ASSERT(convert_func != nullptr); - convert_func(src1->data, src1_alloc.get(), ne10, ne11, ne12, ne13, s11, s12, s13, main_stream); - src1_ptr = src1_alloc.get(); - s11 = ne10; - s12 = ne11*s11; - s13 = ne12*s12; - - is_src1_cont_2 = true; - } - - // Setup destination buffer - ggml_cuda_pool_alloc dst_temp(ctx.pool()); - char * dst_t; - size_t nbd2 = dst->nb[2]; - size_t nbd3 = dst->nb[3]; - - cublasComputeType_t cu_compute_type = traits::compute_type; -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - GGML_UNUSED(cu_compute_type); // only referenced by the cublas fallback paths -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - cudaDataType_t cu_data_type = traits::data_type; - cudaDataType_t cu_data_type_a = traits::data_type; - cudaDataType_t cu_data_type_b = traits::data_type; - const void * alpha = traits::get_alpha(); - const void * beta = traits::get_beta(); - - const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); - - int id = ggml_cuda_get_device(); - const int cc = ggml_cuda_info().devices[id].cc; - static constexpr bool is_src0_type_f16 = src0_type == GGML_TYPE_F16; - - // bf16 and fp32 are already being computed in fp32 (ensure it using static_assert), - // so checking necessity of forced fp32 only for fp16 src0_type - static_assert(is_src0_type_f16 || traits::compute_type == CUBLAS_COMPUTE_32F); - - const bool need_compute_32f = is_src0_type_f16 && !force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) - || GGML_CUDA_CC_IS_RDNA4(cc) - || cc == GGML_CUDA_CC_VOLTA - || force_compute_type.fp32); - - if (dst->op_params[0] == GGML_PREC_DEFAULT && !need_compute_32f) { - if constexpr (src0_type == GGML_TYPE_F32) { - dst_t = (char *) dst_ddf; // Direct F32 output - } else { - dst_t = (char *) dst_temp.alloc(ne_dst); - nbd2 /= sizeof(float) / sizeof(cuda_t); - nbd3 /= sizeof(float) / sizeof(cuda_t); - } - } else { - dst_t = (char *) dst_ddf; - cu_compute_type = batched_mul_mat_traits::compute_type; - cu_data_type = batched_mul_mat_traits::data_type; - alpha = batched_mul_mat_traits::get_alpha(); - beta = batched_mul_mat_traits::get_beta(); - } - - GGML_ASSERT(ne12 % ne02 == 0); - GGML_ASSERT(ne13 % ne03 == 0); - - // broadcast factors - const int64_t r2 = ne12/ne02; - const int64_t r3 = ne13/ne03; - - if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { - // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: - const int64_t sma = ne02 == 1 ? nb03/nb00 : nb02/nb00; - const int64_t smb = ne12 == 1 ? s13 : s12; - - // there is no broadcast and src0, src1 are contiguous across dims 2, 3 - // use cublasGemmStridedBatchedEx -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - GGML_UNUSED_VARS(alpha, beta); - ggml_hipblaslt_gemm(ctx, main_stream, - ne01, ne11, ne10, - src0_ptr, cu_data_type_a, nb01/nb00, sma, - src1_ptr, cu_data_type_b, s11, smb, - dst_t, cu_data_type, ne0, ne1*ne0, - ne12*ne13); -#else - CUBLAS_CHECK( - cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, - ne01, ne11, ne10, - alpha, src0_ptr, cu_data_type_a, nb01/nb00, sma, // strideA - src1_ptr, cu_data_type_b, s11, smb, // strideB - beta, dst_t, cu_data_type, ne0, ne1*ne0, // strideC - ne12*ne13, - cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } else { -#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - // hipBLASLt has no pointer-array batched GEMM; issue one GEMM per batch element instead. - GGML_UNUSED_VARS(alpha, beta); - const size_t src1_nb2 = (src1->type == src0_type) ? nb12 : s12*sizeof(cuda_t); - const size_t src1_nb3 = (src1->type == src0_type) ? nb13 : s13*sizeof(cuda_t); - for (int64_t i13 = 0; i13 < ne13; i13++) { - for (int64_t i12 = 0; i12 < ne12; i12++) { - const char * ptr_a = (const char *) src0_ptr + (i12/r2)*nb02 + (i13/r3)*nb03; - const char * ptr_b = (const char *) src1_ptr + i12*src1_nb2 + i13*src1_nb3; - char * ptr_c = ( char *) dst_t + i12*nbd2 + i13*nbd3; - ggml_hipblaslt_gemm(ctx, main_stream, - ne01, ne11, ne10, - ptr_a, cu_data_type_a, nb01/nb00, 0, - ptr_b, cu_data_type_b, s11, 0, - ptr_c, cu_data_type, ne0, 0, - 1); - } - } -#else - // use cublasGemmBatchedEx - const int64_t ne23 = ne12*ne13; - - ggml_cuda_pool_alloc ptrs_src(ctx.pool(), 2*ne23); - ggml_cuda_pool_alloc< void *> ptrs_dst(ctx.pool(), 1*ne23); - - size_t src1_stride_size = sizeof(cuda_t); - - const int threads_x = 16; - const int threads_y = 16; - dim3 block_dims(threads_x, threads_y); - - dim3 grid_dims( - (ne13 + threads_x - 1) / threads_x, - (ne12 + threads_y - 1) / threads_y - ); - k_compute_batched_ptrs<<>>( - src0_ptr, src1_ptr, dst_t, - ptrs_src.get(), ptrs_dst.get(), - ne12, ne13, - ne23, - nb02, nb03, - (src1->type == src0_type) ? nb12 : s12*src1_stride_size, - (src1->type == src0_type) ? nb13 : s13*src1_stride_size, - nbd2, nbd3, - r2, r3); - - CUDA_CHECK(cudaGetLastError()); - - CUBLAS_CHECK( - cublasGemmBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, - ne01, ne11, ne10, - alpha, (const void **) (ptrs_src.get() + 0*ne23), cu_data_type_a, nb01/nb00, - (const void **) (ptrs_src.get() + 1*ne23), cu_data_type_b, s11, - beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, - ne23, - cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); -#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) - } - - // Convert output back to F32 if needed - if (dst->op_params[0] == GGML_PREC_DEFAULT && cu_data_type != CUDA_R_32F) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); - to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); - } -} - -static void ggml_cuda_mul_mat_batched_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || src0->type == GGML_TYPE_F32); - - switch (src0->type) { - case GGML_TYPE_F32: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - case GGML_TYPE_BF16: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - case GGML_TYPE_F16: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - default: - GGML_ABORT("Unsupported type"); - } -} - -static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, - const ggml_tensor * ffn_gate, - const ggml_tensor * glu, - const ggml_tensor * ffn_up_bias = nullptr, - const ggml_tensor * ffn_gate_bias = nullptr) { - const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; - - if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { - return false; - } - - const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU; - const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU; - - GGML_ASSERT(ffn_up && ffn_gate && glu); - - if (!is_mul_mat && !is_mul_mat_id) { - return false; - } - - const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; - - if (has_bias) { - if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) { - return false; - } - - if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) { - return false; - } - - if (expected_bias_op == GGML_OP_ADD) { - const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up; - const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate; - if (!up_has_mul || !gate_has_mul) { - return false; - } - } else { // GGML_OP_ADD_ID - if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) { - return false; - } - if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) { - return false; - } - } - } else { - if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { - return false; - } - } - - if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) || - !ggml_are_same_stride(ffn_up->src[0], ffn_gate->src[0])) { - return false; - } - - if (ffn_up->src[1] != ffn_gate->src[1]) { - return false; - } - - if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) { - return false; - } - - static constexpr std::array valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI }; - - if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { - return false; - } - - if (const bool swapped = ggml_get_op_params_i32(glu, 1); swapped) { - return false; - } - - const bool split = ggml_backend_buft_is_cuda_split(ffn_up->src[0]->buffer->buft) || - ggml_backend_buft_is_cuda_split(ffn_gate->src[0]->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - return true; -} - -static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { - ggml_tensor * src0 = tensor->src[0]; - ggml_tensor * src1 = tensor->src[1]; - const ggml_tensor * dst = tensor; - - const bool is_mul_mat = tensor->op == GGML_OP_MUL_MAT || - tensor->op == GGML_OP_MUL_MAT_PACK4; - const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; - - bool use_mul_mat_vec_f = - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && - src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]); - - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_backend_buft_is_cuda_split(src1->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - //we only support fusion for ncols_dst = 1 - if (is_mul_mat && dst->ne[1] != 1) { - return false; - } - - if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { - return false; - } - - - return use_mul_mat_vec_f; -} - -static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { - ggml_tensor * src0 = tensor->src[0]; - ggml_tensor * src1 = tensor->src[1]; - const ggml_tensor * dst = tensor; - - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && - ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && - src0->view_src; - - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && - dst->type == GGML_TYPE_F32 && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; - - // fusion is not universally faster on Pascal - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - if (cc <= GGML_CUDA_CC_PASCAL) { - return false; - } - //we only support fusion for ncols_dst = 1 - if ((tensor->op == GGML_OP_MUL_MAT || - tensor->op == GGML_OP_MUL_MAT_PACK4) && dst->ne[1] != 1) { - return false; - } - - if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { - return false; - } - - - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_backend_buft_is_cuda_split(src1->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - return use_mul_mat_vec_q; -} - -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); - - // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. - // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. - // Therefore, in such cases use cuBLAS. - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE - && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; - - bool use_mul_mat_vec_f = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - bool use_mul_mat_f = !ggml_is_quantized(src0->type) - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 - && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; - bool use_mul_mat_q = ggml_is_quantized(src0->type) && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - - bool any_gpus_with_slow_fp16 = false; - - if (split) { - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; - auto & tensor_split = buft_ctx->tensor_split; - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - // skip devices that are not going to do any work: - if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { - continue; - } - - const int cc = ggml_cuda_info().devices[id].cc; - const int warp_size = ggml_cuda_info().devices[id].warp_size; - use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); - use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); - any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); - } - } else { - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); - use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); - any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); - } - - // debug helpers - //printf("src0: %8d %8d %8d %8d\n", src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]); - //printf(" %8d %8d %8d %8d\n", src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3]); - //printf("src1: %8d %8d %8d %8d\n", src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]); - //printf(" %8d %8d %8d %8d\n", src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3]); - //printf("src0 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src0), ggml_is_transposed(src0), ggml_type_name(src0->type), src0->name); - //printf("src1 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src1), ggml_is_transposed(src1), ggml_type_name(src1->type), src1->name); - - //TODO update for generic tensor parallelism - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - bool use_batched_cublas_f16 = src0->type == GGML_TYPE_F16 && (src1->type == GGML_TYPE_F16 || !any_gpus_with_slow_fp16); - bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); - bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; - - if (!split && use_mul_mat_vec_f) { - // the custom F16 vector kernel can be used over batched cuBLAS GEMM - // but this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_f) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_vec_q) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_q) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); - } else if (!split && (use_batched_cublas_f16 || use_batched_cublas_bf16 || use_batched_cublas_f32) - && !ggml_is_transposed(src0) && !ggml_is_transposed(src1) && src1->ne[2]*src1->ne[3] > 1) { - // general KQ + KQV multi-batch without FlashAttention - ggml_cuda_mul_mat_batched_cublas(ctx, src0, src1, dst); - } else if (use_mul_mat_vec_f) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_f, nullptr); - } else if (use_mul_mat_vec_q) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_q, quantize_row_q8_1_cuda); - } else if (use_mul_mat_q) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_q, quantize_mmq_q8_1_cuda); - } else { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_cublas, nullptr); - } -} - -static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * ids = dst->src[2]; - - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft) && "mul_mat_id does not support split buffers"); - - GGML_TENSOR_BINARY_OP_LOCALS - - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - - // [TAG_MUL_MAT_ID_CUDA_GRAPHS] - if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { - static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); - if (ne2 <= MMVQ_MAX_BATCH_SIZE) { - if (ggml_is_quantized(src0->type)) { - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); - if (ne2 <= mmvq_mmid_max) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); - return; - } - } else { - if (GGML_CUDA_CC_IS_AMD(cc)) { - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, dst); - return; - } - } - } - - if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); - return; - } - - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); - return; - } - } - - // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization - // TODO: add asserts to verify this. should work with CUDA, HIP, etc. - cudaStream_t stream = ctx.stream(); - - GGML_ASSERT(nb12 % nb11 == 0); - GGML_ASSERT(nb2 % nb1 == 0); - - const ggml_type type_src1_sorted = (src0->type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) - || ggml_is_quantized(src0->type) ? GGML_TYPE_F32 : src0->type; - const ggml_type type_dst_sorted = GGML_TYPE_F32; - const size_t ts_src1_sorted = ggml_type_size(type_src1_sorted); - const size_t ts_dst_sorted = ggml_type_size(type_dst_sorted); - - const int64_t n_expert_used = ids->ne[0]; - const int64_t ne_get_rows = ne12 * n_expert_used; - - std::vector ids_to_sorted_host; - ids_to_sorted_host.reserve(2*ne_get_rows); - std::vector ids_from_sorted_host(ne_get_rows); - - ggml_cuda_pool_alloc ids_buf_dev(ctx.pool(), 2*ne_get_rows); - - std::vector tokens_per_expert(ne02); - - ggml_cuda_pool_alloc src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted); - ggml_cuda_pool_alloc dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted); - - std::vector ids_host(ggml_nbytes(ids)); - CUDA_CHECK(cudaMemcpyAsync(ids_host.data(), ids->data, ggml_nbytes(ids), cudaMemcpyDeviceToHost, stream)); - CUDA_CHECK(cudaStreamSynchronize(stream)); - - for (int64_t i02 = 0; i02 < ne02; ++i02) { // expert matrices - for (int64_t i12 = 0; i12 < ne12; ++i12) { // tokens - for (int64_t iex = 0; iex < n_expert_used; ++iex) { - const int32_t expert_to_use = *(const int32_t *)(ids_host.data() + i12*ids->nb[1] + iex*ids->nb[0]); - assert(expert_to_use >= 0 && expert_to_use < ne02); - if (expert_to_use == i02) { - ids_from_sorted_host[i12*n_expert_used + iex] = ids_to_sorted_host.size(); - ids_to_sorted_host.push_back(i12*ne11 + iex % ne11); - tokens_per_expert[i02]++; - break; - } - } - } - } - GGML_ASSERT(ids_to_sorted_host.size() == size_t(ne_get_rows)); - - ids_to_sorted_host.insert(ids_to_sorted_host.end(), ids_from_sorted_host.begin(), ids_from_sorted_host.end()); - - CUDA_CHECK(cudaMemcpyAsync(ids_buf_dev.ptr, ids_to_sorted_host.data(), 2*ne_get_rows*sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - CUDA_CHECK(cudaStreamSynchronize(stream)); - - const int32_t * ids_to_sorted = ids_buf_dev.ptr + 0*ne_get_rows; - const int32_t * ids_from_sorted = ids_buf_dev.ptr + 1*ne_get_rows; - - get_rows_cuda(src1->data, src1->type, ids_to_sorted, src1_sorted.ptr, type_src1_sorted, - ne10, nb11, nb12, nb13, - ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), - ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, stream); - CUDA_CHECK(cudaGetLastError()); - - char * src1_data_cur = (char *) src1_sorted.ptr; - char * dst_data_cur = (char *) dst_sorted.ptr; - for (int64_t i02 = 0; i02 < ne02; ++i02) { - if (tokens_per_expert[i02] == 0) { - continue; - } - - ggml_tensor src0_slice = *src0; - src0_slice.ne[2] = 1; - src0_slice.nb[3] = src0_slice.nb[2]; - src0_slice.op = GGML_OP_VIEW; - src0_slice.view_src = dst->src[0]; // non-const pointer to src0 - src0_slice.data = (char *) src0->data + i02*nb02; - - ggml_tensor src1_slice; - memset(&src1_slice, 0, sizeof(src1_slice)); - src1_slice.buffer = src1->buffer; - src1_slice.type = type_src1_sorted; - src1_slice.ne[0] = ne10; - src1_slice.ne[1] = tokens_per_expert[i02]; - src1_slice.ne[2] = 1; - src1_slice.ne[3] = 1; - src1_slice.nb[0] = ts_src1_sorted; - src1_slice.nb[1] = src1_slice.ne[0] * src1_slice.nb[0]; - src1_slice.nb[2] = src1_slice.ne[1] * src1_slice.nb[1]; - src1_slice.nb[3] = src1_slice.ne[2] * src1_slice.nb[2]; - src1_slice.data = src1_data_cur; - - ggml_tensor dst_slice; - memset(&dst_slice, 0, sizeof(dst_slice)); - dst_slice.buffer = dst->buffer; - dst_slice.type = type_dst_sorted; - dst_slice.ne[0] = ne0; - dst_slice.ne[1] = tokens_per_expert[i02]; - dst_slice.ne[2] = 1; - dst_slice.ne[3] = 1; - dst_slice.nb[0] = ts_dst_sorted; - dst_slice.nb[1] = dst_slice.ne[0] * dst_slice.nb[0]; - dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; - dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; - dst_slice.data = dst_data_cur; - - ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); - CUDA_CHECK(cudaGetLastError()); - - src1_data_cur += src1_slice.nb[2]; - dst_data_cur += dst_slice.nb[2]; - } - - get_rows_cuda(dst_sorted.ptr, type_dst_sorted, ids_from_sorted, dst->data, dst->type, - ne0, ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, - ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), - nb1, nb2, nb3, stream); -} - -static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { - switch (dst->op) { - case GGML_OP_ARGMAX: - ggml_cuda_argmax(ctx, dst); - break; - case GGML_OP_COUNT_EQUAL: - ggml_cuda_count_equal(ctx, dst); - break; - case GGML_OP_REPEAT: - ggml_cuda_op_repeat(ctx, dst); - break; - case GGML_OP_REPEAT_BACK: - ggml_cuda_op_repeat_back(ctx, dst); - break; - case GGML_OP_GET_ROWS: - ggml_cuda_op_get_rows(ctx, dst); - break; - case GGML_OP_GET_ROWS_BACK: - ggml_cuda_op_get_rows_back(ctx, dst); - break; - case GGML_OP_SET_ROWS: - ggml_cuda_op_set_rows(ctx, dst); - break; - case GGML_OP_SET: - ggml_cuda_op_set(ctx, dst); - break; - case GGML_OP_DUP: - ggml_cuda_dup(ctx, dst); - break; - case GGML_OP_CPY: - ggml_cuda_cpy(ctx, dst->src[0], dst->src[1]); - break; - case GGML_OP_CONT: - ggml_cuda_dup(ctx, dst); - break; - case GGML_OP_ADD: - case GGML_OP_ADD1: // TODO: more efficient implementation - ggml_cuda_op_add(ctx, dst); - break; - case GGML_OP_ADD_ID: - ggml_cuda_op_add_id(ctx, dst); - break; - case GGML_OP_SUB: - ggml_cuda_op_sub(ctx, dst); - break; - case GGML_OP_ACC: - ggml_cuda_op_acc(ctx, dst); - break; - case GGML_OP_MUL: - ggml_cuda_op_mul(ctx, dst); - break; - case GGML_OP_DIV: - ggml_cuda_op_div(ctx, dst); - break; - case GGML_OP_UNARY: - switch (ggml_get_unary_op(dst)) { - case GGML_UNARY_OP_ABS: - ggml_cuda_op_abs(ctx, dst); - break; - case GGML_UNARY_OP_SGN: - ggml_cuda_op_sgn(ctx, dst); - break; - case GGML_UNARY_OP_NEG: - ggml_cuda_op_neg(ctx, dst); - break; - case GGML_UNARY_OP_STEP: - ggml_cuda_op_step(ctx, dst); - break; - case GGML_UNARY_OP_GELU: - ggml_cuda_op_gelu(ctx, dst); - break; - case GGML_UNARY_OP_SILU: - ggml_cuda_op_silu(ctx, dst); - break; - case GGML_UNARY_OP_GELU_ERF: - ggml_cuda_op_gelu_erf(ctx, dst); - break; - case GGML_UNARY_OP_GELU_QUICK: - ggml_cuda_op_gelu_quick(ctx, dst); - break; - case GGML_UNARY_OP_TANH: - ggml_cuda_op_tanh(ctx, dst); - break; - case GGML_UNARY_OP_RELU: - ggml_cuda_op_relu(ctx, dst); - break; - case GGML_UNARY_OP_SIGMOID: - ggml_cuda_op_sigmoid(ctx, dst); - break; - case GGML_UNARY_OP_HARDSIGMOID: - ggml_cuda_op_hardsigmoid(ctx, dst); - break; - case GGML_UNARY_OP_HARDSWISH: - ggml_cuda_op_hardswish(ctx, dst); - break; - case GGML_UNARY_OP_EXP: - ggml_cuda_op_exp(ctx, dst); - break; - case GGML_UNARY_OP_ELU: - ggml_cuda_op_elu(ctx, dst); - break; - case GGML_UNARY_OP_XIELU: - ggml_cuda_op_xielu(ctx, dst); - break; - case GGML_UNARY_OP_FLOOR: - ggml_cuda_op_floor(ctx, dst); - break; - case GGML_UNARY_OP_CEIL: - ggml_cuda_op_ceil(ctx, dst); - break; - case GGML_UNARY_OP_ROUND: - ggml_cuda_op_round(ctx, dst); - break; - case GGML_UNARY_OP_TRUNC: - ggml_cuda_op_trunc(ctx, dst); - break; - case GGML_UNARY_OP_EXPM1: - ggml_cuda_op_expm1(ctx, dst); - break; - case GGML_UNARY_OP_SOFTPLUS: - ggml_cuda_op_softplus(ctx, dst); - break; - default: - return false; - } - break; - case GGML_OP_GLU: - switch (ggml_get_glu_op(dst)) { - case GGML_GLU_OP_REGLU: - ggml_cuda_op_reglu(ctx, dst); - break; - case GGML_GLU_OP_GEGLU: - ggml_cuda_op_geglu(ctx, dst); - break; - case GGML_GLU_OP_SWIGLU: - ggml_cuda_op_swiglu(ctx, dst); - break; - case GGML_GLU_OP_SWIGLU_OAI: - ggml_cuda_op_swiglu_oai(ctx, dst); - break; - case GGML_GLU_OP_GEGLU_ERF: - ggml_cuda_op_geglu_erf(ctx, dst); - break; - case GGML_GLU_OP_GEGLU_QUICK: - ggml_cuda_op_geglu_quick(ctx, dst); - break; - default: - return false; - } - break; - case GGML_OP_NORM: - ggml_cuda_op_norm(ctx, dst); - break; - case GGML_OP_GROUP_NORM: - ggml_cuda_op_group_norm(ctx, dst); - break; - case GGML_OP_L2_NORM: - ggml_cuda_op_l2_norm(ctx, dst); - break; - case GGML_OP_CONCAT: - ggml_cuda_op_concat(ctx, dst); - break; - case GGML_OP_UPSCALE: - ggml_cuda_op_upscale(ctx, dst); - break; - case GGML_OP_PAD: - ggml_cuda_op_pad(ctx, dst); - break; - case GGML_OP_PAD_REFLECT_1D: - ggml_cuda_op_pad_reflect_1d(ctx, dst); - break; - case GGML_OP_ARANGE: - ggml_cuda_op_arange(ctx, dst); - break; - case GGML_OP_TIMESTEP_EMBEDDING: - ggml_cuda_op_timestep_embedding(ctx, dst); - break; - case GGML_OP_LEAKY_RELU: - ggml_cuda_op_leaky_relu(ctx, dst); - break; - case GGML_OP_SILU_BACK: - ggml_cuda_op_silu_back(ctx, dst); - break; - case GGML_OP_RMS_NORM: - ggml_cuda_op_rms_norm(ctx, dst); - break; - case GGML_OP_RMS_NORM_BACK: - ggml_cuda_op_rms_norm_back(ctx, dst); - break; - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_PACK4: - ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); - break; - case GGML_OP_MUL_MAT_ID: - ggml_cuda_mul_mat_id(ctx, dst); - break; - case GGML_OP_OUT_PROD: - ggml_cuda_out_prod(ctx, dst); - break; - case GGML_OP_SCALE: - ggml_cuda_op_scale(ctx, dst); - break; - case GGML_OP_SQR: - ggml_cuda_op_sqr(ctx, dst); - break; - case GGML_OP_SQRT: - ggml_cuda_op_sqrt(ctx, dst); - break; - case GGML_OP_SIN: - ggml_cuda_op_sin(ctx, dst); - break; - case GGML_OP_COS: - ggml_cuda_op_cos(ctx, dst); - break; - case GGML_OP_CLAMP: - ggml_cuda_op_clamp(ctx, dst); - break; - case GGML_OP_LOG: - ggml_cuda_op_log(ctx, dst); - break; - case GGML_OP_NONE: - case GGML_OP_RESHAPE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - case GGML_OP_TRANSPOSE: - break; - case GGML_OP_DIAG: - ggml_cuda_op_diag(ctx, dst); - break; - case GGML_OP_DIAG_MASK_INF: - ggml_cuda_op_diag_mask_inf(ctx, dst); - break; - case GGML_OP_SOFT_MAX: - ggml_cuda_op_soft_max(ctx, dst); - break; - case GGML_OP_SOFT_MAX_BACK: - ggml_cuda_op_soft_max_back(ctx, dst); - break; - case GGML_OP_ROPE: - ggml_cuda_op_rope(ctx, dst); - break; - case GGML_OP_ROPE_BACK: - ggml_cuda_op_rope_back(ctx, dst); - break; - case GGML_OP_ROLL: - ggml_cuda_op_roll(ctx, dst); - break; - case GGML_OP_IM2COL: - case GGML_OP_IM2COL_FAST_1D: - ggml_cuda_op_im2col(ctx, dst); - break; - case GGML_OP_IM2COL_3D: - ggml_cuda_op_im2col_3d(ctx, dst); - break; - case GGML_OP_COL2IM_1D: - ggml_cuda_op_col2im_1d(ctx, dst); - break; - case GGML_OP_CONV_2D: - ggml_cuda_op_conv2d(ctx, dst); - break; - case GGML_OP_CONV_2D_DW: - ggml_cuda_op_conv2d_dw(ctx, dst); - break; - case GGML_OP_CONV_TRANSPOSE_2D: - ggml_cuda_conv_2d_transpose_p0(ctx, dst); - break; - case GGML_OP_CONV_TRANSPOSE_1D: - ggml_cuda_op_conv_transpose_1d(ctx,dst); - break; - case GGML_OP_POOL_2D: - ggml_cuda_op_pool2d(ctx, dst); - break; - case GGML_OP_SUM: - ggml_cuda_op_sum(ctx, dst); - break; - case GGML_OP_CUMSUM: - ggml_cuda_op_cumsum(ctx, dst); - break; - case GGML_OP_SUM_ROWS: - ggml_cuda_op_sum_rows(ctx, dst); - break; - case GGML_OP_MEAN: - ggml_cuda_op_mean(ctx, dst); - break; - case GGML_OP_SSM_CONV: - ggml_cuda_op_ssm_conv(ctx, dst); - break; - case GGML_OP_SSM_SCAN: - ggml_cuda_op_ssm_scan(ctx, dst); - break; - case GGML_OP_TOP_K: - ggml_cuda_op_top_k(ctx, dst); - break; - case GGML_OP_ARGSORT: - ggml_cuda_op_argsort(ctx, dst); - break; - case GGML_OP_FLASH_ATTN_EXT: - ggml_cuda_flash_attn_ext(ctx, dst); - break; - case GGML_OP_SAGE_ATTN2: - ggml_cuda_sage_attn2(ctx, dst); - break; - case GGML_OP_SAGE_ATTN2_I8: - ggml_cuda_sage_attn2_i8(ctx, dst); - break; - case GGML_OP_CONVROT_LINEAR: - ggml_cuda_convrot_linear(ctx, dst); - break; - case GGML_OP_CROSS_ENTROPY_LOSS: - ggml_cuda_cross_entropy_loss(ctx, dst); - break; - case GGML_OP_TRI: - ggml_cuda_op_tri(ctx, dst); - break; - case GGML_OP_RWKV_WKV6: - ggml_cuda_op_rwkv_wkv6(ctx, dst); - break; - case GGML_OP_GATED_LINEAR_ATTN: - ggml_cuda_op_gated_linear_attn(ctx, dst); - break; - case GGML_OP_GATED_DELTA_NET: - ggml_cuda_op_gated_delta_net(ctx, dst); - break; - case GGML_OP_RWKV_WKV7: - ggml_cuda_op_rwkv_wkv7(ctx, dst); - break; - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - ggml_cuda_cross_entropy_loss_back(ctx, dst); - break; - case GGML_OP_OPT_STEP_ADAMW: - ggml_cuda_opt_step_adamw(ctx, dst); - break; - case GGML_OP_OPT_STEP_SGD: - ggml_cuda_opt_step_sgd(ctx, dst); - break; - case GGML_OP_SOLVE_TRI: - ggml_cuda_op_solve_tri(ctx, dst); - break; - case GGML_OP_FILL: - ggml_cuda_op_fill(ctx, dst); - break; - default: - return false; - } - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - GGML_LOG_ERROR("%s: %s failed\n", __func__, ggml_op_desc(dst)); - CUDA_CHECK(err); - } - - return true; -} - -//////////////////////////////////////////////////////////////////////////////// - -// backend - -static const char * ggml_backend_cuda_get_name(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - return cuda_ctx->name.c_str(); -} - -static void ggml_backend_cuda_free(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - delete cuda_ctx; - delete backend; -} - -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_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())); -} - -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_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())); -} - -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) { - ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; - ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; - - if (!ggml_backend_is_cuda(backend_src) || !ggml_backend_is_cuda(backend_dst)) { - return false; - } - - 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_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 // NDEBUG - return false; - } - - if (backend_src != backend_dst) { - // copy on src stream - if (cuda_ctx_src->device == cuda_ctx_dst->device) { - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); - } else { -#ifdef GGML_CUDA_NO_PEER_COPY - 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 // GGML_CUDA_NO_PEER_COPY - } - - // record event on src stream after the copy - if (!cuda_ctx_src->copy_event) { - ggml_cuda_set_device(cuda_ctx_src->device); - CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); - } - - CUDA_CHECK(cudaEventRecord(cuda_ctx_src->copy_event, cuda_ctx_src->stream())); - - // wait on dst stream for the copy to complete - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); - } else { - // src and dst are on the same backend - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); - } - return true; -} - -static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); - - GGML_UNUSED(backend); -} - -#ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { - - bool use_cuda_graph = true; - // Loop over nodes in GGML graph to obtain info needed for CUDA graph - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - - if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { - continue; - } - - if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) { - use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to split buffer\n", __func__); -#endif - } - - // [TAG_MUL_MAT_ID_CUDA_GRAPHS] - if (node->op == GGML_OP_MUL_MAT_ID) { - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); - if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { - // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs - // TODO: figure out a way to enable for larger batch sizes, without hurting performance - // ref: https://github.com/ggml-org/llama.cpp/pull/18958 - use_cuda_graph = false; -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__); -#endif - } - } - - if (!use_cuda_graph) { - break; - } - } - - return use_cuda_graph; -} - -static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { - return cgraph->nodes[0]; -} - -static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { - bool res = false; - - const void * graph_key = ggml_cuda_graph_get_key(cgraph); - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - - if (cgraph->uid != 0 && - cgraph->uid == graph->uid) { - GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); - GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); - return false; - } - - graph->uid = cgraph->uid; - - // Check if the graph size has changed - if ((int)graph->node_props.size() != cgraph->n_nodes) { - res = true; - graph->node_props.resize(cgraph->n_nodes); - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_cuda_graph::node_properties prop = {}; - memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); - - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (cgraph->nodes[i]->src[j]) { - prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; - memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); - memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); - } - } - - if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { - graph->node_props[i] = prop; - res = true; - } - } - - return res; -} - -static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - -#if CUDART_VERSION >= 12000 - cudaGraphExecUpdateResultInfo result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info); -#else - cudaGraphNode_t errorNode; - cudaGraphExecUpdateResult result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info); -#endif // CUDART_VERSION >= 12000 - - if (stat == cudaErrorGraphExecUpdateFailure) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__); -#endif - - // The pre-existing graph exec cannot be updated due to violated constraints - // so instead clear error and re-instantiate - (void)cudaGetLastError(); - CUDA_CHECK(cudaGraphExecDestroy(graph->instance)); - graph->instance = nullptr; - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); - } else { - GGML_ASSERT(stat == cudaSuccess); - } -} -#endif // USE_CUDA_GRAPH - -static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope, - const ggml_tensor * view, - const ggml_tensor * set_rows) { - - if (rope->op != GGML_OP_ROPE || view->op != GGML_OP_VIEW || set_rows->op != GGML_OP_SET_ROWS) { - return false; - } - // ne3 not tested - if (rope->src[0]->ne[3] != 1) { - return false; - } - - if (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) { - return false; - } - - if (set_rows->src[1]->type != GGML_TYPE_I64) { - return false; - } - - // The view should flatten two dims of rope into one dim - if (!ggml_is_contiguous(view) || view->ne[0] != rope->ne[0] * rope->ne[1]) { - return false; - } - - // Only norm/neox shaders have the fusion code - const int mode = ((const int32_t *) rope->op_params)[2]; - if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) { - return false; - } - - return true; -} - -static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { - args.sigmoid = false; - args.softmax = false; - args.delayed_softmax = false; - args.prob_bias = false; - args.norm = false; - - const int n_nodes = cgraph->n_nodes; - ggml_tensor ** nodes = cgraph->nodes; - - if (nodes[node_idx]->op == GGML_OP_SOFT_MAX) { - args.softmax = true; - } - - if (nodes[node_idx]->op == GGML_OP_UNARY) { - if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) { - return false; - } - args.sigmoid = true; - } - - if (nodes[node_idx]->op == GGML_OP_ARGSORT) { - args.delayed_softmax = true; - } - - node_idx++; - - if (args.sigmoid || args.softmax) { - // SOFTMAX -> RESHAPE - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - ggml_tensor * probs_reshaped = nodes[node_idx]; - node_idx++; - - if (node_idx >= n_nodes) { - return false; - } - - // src of bias add is the unreshaped probs (-2 instead of -1) - if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) { - args.prob_bias = true; - node_idx++; - } - // RESHAPE/ADD -> ARGSORT - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) { - return false; - } - - if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) { - return false; - } - - node_idx++; - - // ARGSORT-> VIEW - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) { - return false; - } - - // GET_ROWS - if (nodes[node_idx]->src[0] != probs_reshaped || nodes[node_idx]->src[1] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - } else if (args.delayed_softmax) { - if (node_idx - 2 < 0) { - return false; - } - ggml_tensor * probs_reshaped = nodes[node_idx - 2]; - - // VIEW->ARGSORT - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - - // GET_ROWS - if (node_idx >= n_nodes || nodes[node_idx]->src[1] != nodes[node_idx - 1] || - nodes[node_idx]->src[0] != probs_reshaped) { - return false; - } - node_idx++; - - static const std::vector remaining_ops = { GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }; - - for (const ggml_op op : remaining_ops) { - if (node_idx >= n_nodes || nodes[node_idx]->op != op || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - } - } - - // At this point we can check for norm + scale. Everything is now at least valid till the norm - if (node_idx >= n_nodes) { - return true; - } - - if (nodes[node_idx]->op == GGML_OP_RESHAPE) { - //check RESHAPE->SUM_ROWS->CLAMP->DIV->RESHAPE - static const std::vector norm_ops = { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP }; - - args.norm = true; - for (const ggml_op op : norm_ops) { - if (nodes[node_idx]->op == op && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { - node_idx++; - } else { - args.norm = false; - return true; - } - } - - // DIV <- CLAMP, RESHAPE - if (nodes[node_idx]->op != GGML_OP_DIV || nodes[node_idx]->src[1] != nodes[node_idx - 1] || - nodes[node_idx]->src[0] != nodes[node_idx - 3]) { - args.norm = false; - return true; - } - node_idx++; - - if (nodes[node_idx]->op != GGML_OP_RESHAPE || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - args.norm = false; - return true; - } - - node_idx++; - } - - if (nodes[node_idx]->op == GGML_OP_SCALE && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { - args.scale = true; - } - - return true; -} - -// returns whether the write (out) nodes overwrite the read nodes in operation -static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, - const int node_idx, - const int node_count, - const int * out_nodes, - const int out_count, - const bool is_topk_moe = false) { - auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { - const int64_t a_start = (int64_t) a->data; - const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); - - const int64_t b_start = (int64_t) b->data; - const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); - - if ((b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end)) { - return true; - } - - return false; - }; - - bool is_ok = true; - // exception for topk-moe, as each row is read entirely before writing - if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) { - return true; - } - - for (int i = 0; i < out_count; ++i) { - const ggml_tensor * dst = cgraph->nodes[out_nodes[i]]; - - for (int j = node_idx; j < node_idx + node_count; ++j) { - // Loop over all srcs of all nodes in the fusion. If the src overlaps - // the destination and the src is not an intermediate node that's being - // elided, then disable fusion. - - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - const ggml_tensor * src = cgraph->nodes[j]->src[src_idx]; - - if (!src || src->op == GGML_OP_NONE) { - continue; - } - - if (nodes_overlap(dst, src)) { - bool found = false; - - for (int k = node_idx; k < j; ++k) { - if (cgraph->nodes[k] == src) { - found = true; - break; - } - } - - if (!found) { - is_ok = false; - break; - } - } - } - } - } - - return is_ok; -} - -// Some model graphs reshape a matvec result before adding the residual. RESHAPE -// is metadata-only and therefore cannot pass the generic compute-node fusion -// validator. Validate this exact chain explicitly so the residual-only Q8_0 -// specialization can write the final result directly. -static bool ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add( - const struct ggml_cgraph * cgraph, int node_idx) { - if (node_idx + 2 >= cgraph->n_nodes) { - return false; } +} + + +// cuda buffer + +struct ggml_backend_cuda_buffer_context { + int device; + void * dev_ptr = nullptr; + std::string name; + + ggml_backend_cuda_buffer_context(int device, void * dev_ptr) : + device(device), dev_ptr(dev_ptr), + name(GGML_CUDA_NAME + std::to_string(device)) { + } + + ~ggml_backend_cuda_buffer_context() { + CUDA_CHECK(cudaFree(dev_ptr)); + } +}; + +static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + delete ctx; +} + +static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer) { + return buffer->iface.free_buffer == ggml_backend_cuda_buffer_free_buffer; +} + +static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + return ctx->dev_ptr; +} + +static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + + if (tensor->view_src != NULL) { + assert(tensor->view_src->buffer->buft == buffer->buft); + return GGML_STATUS_SUCCESS; + } + + if (ggml_is_quantized(tensor->type) && tensor->view_src == nullptr && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + // initialize padding to 0 to avoid possible NaN values + const size_t original_size = ggml_nbytes(tensor); + const size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor); + + if (padded_size > original_size) { + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); + } + } + return GGML_STATUS_SUCCESS; +} + +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_cuda_set_device(ctx->device); + 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_cuda_set_device(ctx->device); + 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(cudaMemcpy2DAsync( + data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { + if (ggml_backend_buffer_is_cuda(src->buffer)) { + ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context; + ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context; + if (src_ctx->device == dst_ctx->device) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread)); + } else { +#ifdef GGML_CUDA_NO_PEER_COPY + return false; +#else + CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread)); +#endif + } + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + return true; + } + return false; + + GGML_UNUSED(buffer); +} + +static void ggml_backend_cuda_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemsetAsync(ctx->dev_ptr, value, buffer->size, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = { + /* .free_buffer = */ ggml_backend_cuda_buffer_free_buffer, + /* .get_base = */ ggml_backend_cuda_buffer_get_base, + /* .init_tensor = */ ggml_backend_cuda_buffer_init_tensor, + /* .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, +}; + +// cuda buffer type +struct ggml_backend_cuda_buffer_type_context { + int device; + std::string name; +}; + +static const char * ggml_backend_cuda_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_buffer_type_context * ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static bool ggml_backend_buft_is_cuda(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_buffer_type_get_name; +} + +static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; + + ggml_cuda_set_device(buft_ctx->device); + + void * dev_ptr; + cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + GGML_LOG_ERROR("%s: allocating %.2f MiB on device %d: cudaMalloc failed: %s\n", __func__, size / 1024.0 / 1024.0, buft_ctx->device, cudaGetErrorString(err)); + return nullptr; + } + + ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr); + + return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size); +} + +static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return 128; + + GGML_UNUSED(buft); +} + +static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + size_t size = ggml_nbytes(tensor); + int64_t ne0 = tensor->ne[0]; + + if (ggml_is_quantized(tensor->type)) { + if (ne0 % MATRIX_ROW_PADDING != 0) { + GGML_ASSERT(tensor->nb[0] == ggml_element_size(tensor)); + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + } + + return size; + + GGML_UNUSED(buft); +} + +static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface = { + /* .get_name = */ ggml_backend_cuda_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cuda_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cuda_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cuda_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + +ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + if (device >= ggml_backend_cuda_get_device_count()) { + return nullptr; + } + + static ggml_backend_buffer_type ggml_backend_cuda_buffer_types[GGML_CUDA_MAX_DEVICES]; + + static bool ggml_backend_cuda_buffer_type_initialized = false; + + if (!ggml_backend_cuda_buffer_type_initialized) { + for (int i = 0; i < ggml_backend_cuda_get_device_count(); i++) { + ggml_backend_cuda_buffer_types[i] = { + /* .iface = */ ggml_backend_cuda_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), i), + /* .context = */ new ggml_backend_cuda_buffer_type_context{i, GGML_CUDA_NAME + std::to_string(i)}, + }; + } + ggml_backend_cuda_buffer_type_initialized = true; + } + + return &ggml_backend_cuda_buffer_types[device]; +} + +// cuda split buffer + +static int64_t get_row_rounding(const std::array & tensor_split) { + int64_t row_rounding = 0; + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { + continue; + } + + const int cc = ggml_cuda_info().devices[id].cc; + row_rounding = std::max(row_rounding, (int64_t)get_mmq_y_host(cc)); + } + return row_rounding; +} + +static void get_row_split(int64_t * row_low, int64_t * row_high, const ggml_tensor * tensor, const std::array & tensor_split, int id) { + const int64_t nrows = ggml_nrows(tensor); + const int64_t rounding = get_row_rounding(tensor_split); + + *row_low = id == 0 ? 0 : nrows*tensor_split[id]; + *row_low -= *row_low % rounding; + + if (id == ggml_backend_cuda_get_device_count() - 1) { + *row_high = nrows; + } else { + *row_high = nrows*tensor_split[id + 1]; + *row_high -= *row_high % rounding; + } +} + +static size_t ggml_nbytes_split(const struct ggml_tensor * tensor, int nrows_split) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return nrows_split*ggml_row_size(tensor->type, tensor->ne[0]); +} + +struct ggml_backend_cuda_split_buffer_type_context { + int main_device; + std::array tensor_split; + std::string name; +}; + +struct ggml_backend_cuda_split_buffer_context { + ~ggml_backend_cuda_split_buffer_context() { + for (ggml_tensor_extra_gpu * extra : tensor_extras) { + for (int id = 0; id < GGML_CUDA_MAX_DEVICES; ++id) { + for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { + if (extra->events[id][is] != nullptr) { + CUDA_CHECK(cudaEventDestroy(extra->events[id][is])); + } + } + if (extra->data_device[id] != nullptr) { + CUDA_CHECK(cudaFree(extra->data_device[id])); + } + } + delete extra; + } + } + + std::vector tensor_extras; +}; + + +static void ggml_backend_cuda_split_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; + delete ctx; +} + +static void * ggml_backend_cuda_split_buffer_get_base(ggml_backend_buffer_t buffer) { + // the pointers are stored in the tensor extras, this is just a dummy address and never dereferenced + return (void *)0x1000; + + GGML_UNUSED(buffer); +} + +static enum ggml_status ggml_backend_cuda_split_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + GGML_ASSERT(tensor->view_src == nullptr); // views of split tensors are not supported + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + + ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; + ctx->tensor_extras.push_back(extra); + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + // FIXME: do not crash if cudaMalloc fails + // currently, init_tensor cannot fail, it needs to be fixed in ggml-backend first + ggml_cuda_set_device(id); + char * buf; + CUDA_CHECK(ggml_cuda_device_malloc((void**)&buf, size, id)); + + // set padding to 0 to avoid possible NaN values + if (size > original_size) { + CUDA_CHECK(cudaMemset(buf + original_size, 0, size - original_size)); + } + + extra->data_device[id] = buf; + + for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { + CUDA_CHECK(cudaEventCreateWithFlags(&extra->events[id][is], cudaEventDisableTiming)); + } + } + tensor->extra = extra; + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cuda_split_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + // split tensors must always be set in their entirety at once + GGML_ASSERT(offset == 0); + GGML_ASSERT(size == ggml_nbytes(tensor)); + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + const size_t nb1 = tensor->nb[1]; + ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + const size_t offset_split = row_low*nb1; + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + const char * buf_host = (const char *)data + offset_split; + CUDA_CHECK(cudaMemcpyAsync(extra->data_device[id], buf_host, original_size, cudaMemcpyHostToDevice, cudaStreamPerThread)); + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + } +} + +static void ggml_backend_cuda_split_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + // split tensors must always be set in their entirety at once + GGML_ASSERT(offset == 0); + GGML_ASSERT(size == ggml_nbytes(tensor)); + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + const size_t nb1 = tensor->nb[1]; + ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + const size_t offset_split = row_low*nb1; + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + char * buf_host = (char *)data + offset_split; + CUDA_CHECK(cudaMemcpyAsync(buf_host, extra->data_device[id], original_size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + } +} + +static void ggml_backend_cuda_split_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + GGML_UNUSED(buffer); + GGML_UNUSED(value); +} + +static const ggml_backend_buffer_i ggml_backend_cuda_split_buffer_interface = { + /* .free_buffer = */ ggml_backend_cuda_split_buffer_free_buffer, + /* .get_base = */ ggml_backend_cuda_split_buffer_get_base, + /* .init_tensor = */ ggml_backend_cuda_split_buffer_init_tensor, + /* .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, +}; + +// cuda split buffer type + +static const char * ggml_backend_cuda_split_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static bool ggml_backend_buft_is_cuda_split(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_split_buffer_type_get_name; +} + +static ggml_backend_buffer_t ggml_backend_cuda_split_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + // since we don't know the exact split after rounding, we cannot allocate the device buffers at this point + // instead, we allocate them for each tensor separately in init_tensor + // however, the size still represents the maximum cumulative size of all the device buffers after the tensors are allocated, + // as returned by get_alloc_size. this limit is enforced during tensor allocation by ggml-alloc, so it must be correct. + ggml_backend_cuda_split_buffer_context * ctx = new ggml_backend_cuda_split_buffer_context(); + + return ggml_backend_buffer_init(buft, ggml_backend_cuda_split_buffer_interface, ctx, size); +} + +static size_t ggml_backend_cuda_split_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return 128; + + GGML_UNUSED(buft); +} + +static size_t ggml_backend_cuda_split_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + size_t total_size = 0; + + const int64_t ne0 = tensor->ne[0]; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + total_size += ggml_nbytes_split(tensor, nrows_split); + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + total_size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + } + + return total_size; +} + +static bool ggml_backend_cuda_split_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return false; + + GGML_UNUSED(buft); +} + +static const ggml_backend_buffer_type_i ggml_backend_cuda_split_buffer_type_interface = { + /* .get_name = */ ggml_backend_cuda_split_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cuda_split_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cuda_split_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cuda_split_buffer_type_get_alloc_size, + /* .is_host = */ ggml_backend_cuda_split_buffer_type_is_host, +}; + +// Communication context for multi-GPU AllReduce during tensor parallelism. +// +// Created once per meta backend instance. Resources for the selected mode +// (NCCL communicators or the internal AllReduce pipeline) are initialised +// eagerly during comm_init so any init failure surfaces at startup rather +// than mid-run. +struct ggml_backend_cuda_comm_context { + using try_allreduce_fn = bool(*)(ggml_backend_cuda_comm_context *, struct ggml_tensor **); + + std::vector backends; + std::vector dev_ids; + + // Set by the init chain (comm_init_{nccl, internal, none}) to one of + // try_allreduce_{nccl, internal, butterfly}. nccl needs `comms`, + // internal needs `ar_pipeline`, butterfly needs nothing. Per-call + // failures return false; the meta backend's generic implementation then + // handles that call. + try_allreduce_fn try_allreduce = nullptr; + + ggml_cuda_ar_pipeline * ar_pipeline = nullptr; + +#ifdef GGML_USE_NCCL + std::vector comms; +#endif // GGML_USE_NCCL + + ~ggml_backend_cuda_comm_context() { +#ifdef GGML_USE_NCCL + for (ncclComm_t comm : comms) { + NCCL_CHECK(ncclCommDestroy(comm)); + } +#endif // GGML_USE_NCCL + ggml_cuda_ar_pipeline_free(ar_pipeline); + } +}; + +#ifdef GGML_USE_NCCL +// AllReduce via NCCL. Reduces as FP32 for small tensors and BF16 for large +// tensors (bandwidth-bound), then converts back to FP32. +static bool ggml_backend_cuda_comm_allreduce_nccl( + ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { + 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; + } + + const size_t n_backends = comm_ctx->backends.size(); + + for (size_t i = 0; i < n_backends; ++i) { + GGML_ASSERT(tensors[i] != nullptr); + GGML_ASSERT(ggml_nelements(tensors[i]) == ne); + GGML_ASSERT(ggml_is_contiguously_allocated(tensors[i])); + } + + // For small tensors, simply reduce them as FP32. + // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0. + if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) { + for (size_t i = 0; i < n_backends; ++i) { + if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + ggml_cuda_set_device(cuda_ctx->device); + CUDA_CHECK(cudaMemsetAsync(tensors[i]->data, 0, ggml_nbytes(tensors[i]), cuda_ctx->stream())); + } + } + NCCL_CHECK(ncclGroupStart()); + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); + } + NCCL_CHECK(ncclGroupEnd()); + 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 tmp[GGML_CUDA_MAX_DEVICES]; + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + tmp[i].pool = &cuda_ctx->pool(); + tmp[i].alloc(ne); + + ggml_cuda_set_device(cuda_ctx->device); + if (tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) { + to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); + } else { + CUDA_CHECK(cudaMemsetAsync(tmp[i].get(), 0, ne * sizeof(nv_bfloat16), cuda_ctx->stream())); + } + CUDA_CHECK(cudaGetLastError()); + } + + NCCL_CHECK(ncclGroupStart()); + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, comm_ctx->comms[i], cuda_ctx->stream())); + } + NCCL_CHECK(ncclGroupEnd()); + + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context; + + ggml_cuda_set_device(cuda_ctx->device); + to_fp32(tmp[i].get(), (float *) tensors[i]->data, ne, cuda_ctx->stream()); + CUDA_CHECK(cudaGetLastError()); + } + + return true; +} +#endif // GGML_USE_NCCL + +// Run the internal AR pipeline. Returns false on unsupported / failed input +// -- the caller decides whether to abort (env-forced) or fall back silently. +static bool ggml_backend_cuda_comm_allreduce_internal( + ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { + GGML_ASSERT(comm_ctx->ar_pipeline != nullptr); + + const size_t n_backends = comm_ctx->backends.size(); + GGML_ASSERT(n_backends == 2); + GGML_ASSERT(tensors[0] != nullptr); + + const int64_t ne = ggml_nelements(tensors[0]); + const ggml_type type = tensors[0]->type; + + if (type != GGML_TYPE_F32 && type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) { + GGML_LOG_DEBUG("%s: internal unsupported: type=%d\n", __func__, (int) type); + return false; + } + + if (ne == 0) { + return true; + } + + for (size_t i = 0; i < n_backends; ++i) { + if (tensors[i] == nullptr) { + GGML_LOG_ERROR("%s: internal failed: tensor[%zu] is null\n", __func__, i); + return false; + } + if (ggml_nelements(tensors[i]) != ne || tensors[i]->type != type) { + GGML_LOG_ERROR("%s: internal failed: tensor[%zu] ne=%" PRId64 " type=%d expected ne=%" PRId64 " type=%d\n", + __func__, i, ggml_nelements(tensors[i]), (int) tensors[i]->type, ne, (int) type); + return false; + } + if (!ggml_is_contiguously_allocated(tensors[i])) { + GGML_LOG_DEBUG("%s: internal unsupported: tensor[%zu] is not contiguously allocated: ne=%" PRId64 " nbytes=%zu packed=%zu type=%d\n", + __func__, i, ne, ggml_nbytes(tensors[i]), + (size_t) ne * ggml_type_size(type) / ggml_blck_size(type), (int) type); + return false; + } + if (((uintptr_t) tensors[i]->data & 0xF) != 0) { + GGML_LOG_DEBUG("%s: internal unsupported: tensor[%zu] data pointer is not 16-byte aligned: %p type=%d ne=%" PRId64 "\n", + __func__, i, tensors[i]->data, (int) type, ne); + return false; + } + GGML_ASSERT((ggml_nbytes(tensors[i]) & 0xF) == 0); + } + + return ggml_cuda_ar_allreduce(comm_ctx->ar_pipeline, comm_ctx->backends.data(), tensors); +} + +// --------------------------------------------------------------------------- +// Per-call dispatch -- three variants, one per backend. Each is set as +// comm_ctx->try_allreduce by the matching init step. Per-call failure +// returns false; the meta backend's generic implementation handles that call. +// --------------------------------------------------------------------------- + +#ifdef GGML_USE_NCCL +static bool ggml_backend_cuda_comm_try_allreduce_nccl( + ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { + return ggml_backend_cuda_comm_allreduce_nccl(comm_ctx, tensors); +} +#endif // GGML_USE_NCCL + +static bool ggml_backend_cuda_comm_try_allreduce_internal( + ggml_backend_cuda_comm_context * comm_ctx, struct ggml_tensor ** tensors) { + return ggml_backend_cuda_comm_allreduce_internal(comm_ctx, tensors); +} + +static bool ggml_backend_cuda_comm_try_allreduce_butterfly( + ggml_backend_cuda_comm_context *, struct ggml_tensor **) { + return false; +} + +static void ggml_backend_cuda_comm_free(void * comm_ctx_v) { + if (comm_ctx_v == nullptr) { + return; + } + delete static_cast(comm_ctx_v); +} + +// --------------------------------------------------------------------------- +// Init -- chained nccl -> internal -> none. Each step tries to bring up its +// resource; on failure it warns and recurses into the next step. +// --------------------------------------------------------------------------- +static void ggml_backend_cuda_comm_init_none(ggml_backend_cuda_comm_context * ret) { + ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_butterfly; +} + +static void ggml_backend_cuda_comm_init_internal(ggml_backend_cuda_comm_context * ret) { + ret->ar_pipeline = ggml_cuda_ar_pipeline_init(ret->dev_ids.data(), ret->dev_ids.size()); + if (ret->ar_pipeline) { + ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_internal; + return; + } + + // Clear sticky CUDA error from the failed init. + (void) cudaGetLastError(); + GGML_LOG_WARN("internal AllReduce init failed (n_devices != 2?); " + "falling back to meta-backend butterfly\n"); + ggml_backend_cuda_comm_init_none(ret); +} + +static void ggml_backend_cuda_comm_init_nccl(ggml_backend_cuda_comm_context * ret) { +#ifdef GGML_USE_NCCL + const size_t n = ret->dev_ids.size(); + ret->comms.resize(n); + ncclResult_t rc = ncclCommInitAll(ret->comms.data(), (int) n, ret->dev_ids.data()); + if (rc == ncclSuccess) { + ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl; + return; + } + + ret->comms.clear(); + GGML_LOG_WARN("NCCL init failed (%s); falling back to internal AllReduce\n", + ncclGetErrorString(rc)); +#else // GGML_USE_NCCL +#ifndef GGML_USE_HIP + GGML_LOG_WARN("NCCL not compiled in; falling back to internal AllReduce. " + "Recompile with -DGGML_CUDA_NCCL=ON for best multi-GPU performance.\n"); +#endif // !GGML_USE_HIP +#endif // GGML_USE_NCCL + + ggml_backend_cuda_comm_init_internal(ret); +} + +// Top-level init. Picks one of the three init paths based on +// GGML_CUDA_ALLREDUCE (or the platform default) and lets the chain handle +// any fallback. Unrecognised env values warn and fall through to the +// platform default. +static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_backends) { + for (size_t i = 0; i < n_backends; i++) { + if (!ggml_backend_is_cuda(backends[i])) { + return nullptr; + } + } + + auto * ret = new ggml_backend_cuda_comm_context; + ret->backends.assign(backends, backends + n_backends); + ret->dev_ids.reserve(n_backends); + for (size_t i = 0; i < n_backends; i++) { + ret->dev_ids.push_back(static_cast(backends[i]->context)->device); + } + + const char * env = getenv("GGML_CUDA_ALLREDUCE"); + if (!env) { + // Platform default: Linux uses NCCL, otherwise (generally Windows) internal +#if defined(__linux__) + ggml_backend_cuda_comm_init_nccl(ret); +#else + ggml_backend_cuda_comm_init_internal(ret); +#endif // defined(__linux__) + } else { + std::string env_str(env); + if (env_str == "nccl") { + ggml_backend_cuda_comm_init_nccl(ret); + } else if (env_str == "internal") { + ggml_backend_cuda_comm_init_internal(ret); + } else if (env_str == "none") { + ggml_backend_cuda_comm_init_none(ret); + } else { + GGML_LOG_WARN("unknown GGML_CUDA_ALLREDUCE value: %s\n", env); + ggml_backend_cuda_comm_init_none(ret); + } + } + + return ret; +} + +// Top-level dispatch -- calls the function pointer chosen by comm_init. +// Returns false to let the meta-backend's butterfly run. +static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) { + if (comm_ctx_v == nullptr) { + return false; + } + auto * comm_ctx = static_cast(comm_ctx_v); + return comm_ctx->try_allreduce(comm_ctx, tensors); +} + +ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + static std::map>, struct ggml_backend_buffer_type> buft_map; + + std::array tensor_split_arr = {}; + + bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + GGML_CUDA_MAX_DEVICES, [](float x) { return x == 0.0f; }); + if (all_zero) { + tensor_split_arr = ggml_cuda_info().default_tensor_split; + } else { + float split_sum = 0.0f; + for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { + tensor_split_arr[i] = split_sum; + split_sum += tensor_split[i]; + } + for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { + tensor_split_arr[i] /= split_sum; + } + } + + auto it = buft_map.find({main_device, tensor_split_arr}); + if (it != buft_map.end()) { + return &it->second; + } + auto * ctx = new ggml_backend_cuda_split_buffer_type_context{ + main_device, + tensor_split_arr, + GGML_CUDA_NAME + std::to_string(main_device) + "_Split", + }; + + struct ggml_backend_buffer_type buft { + /* .iface = */ ggml_backend_cuda_split_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), main_device), + /* .context = */ ctx, + }; + + auto result = buft_map.emplace(std::make_pair(main_device, tensor_split_arr), buft); + return &result.first->second; +} + +// host buffer type + +static const char * ggml_backend_cuda_host_buffer_type_name(ggml_backend_buffer_type_t buft) { + return GGML_CUDA_NAME "_Host"; + + GGML_UNUSED(buft); +} + +static bool ggml_backend_buft_is_cuda_host(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; +} + +static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { + CUDA_CHECK(cudaFreeHost(buffer->context)); +} + +static void * ggml_cuda_host_malloc(size_t size) { + if (getenv("GGML_CUDA_NO_PINNED") != nullptr) { + return nullptr; + } + + void * ptr = nullptr; + cudaError_t err = cudaMallocHost((void **) &ptr, size); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + GGML_LOG_DEBUG("%s: failed to allocate %.2f MiB of pinned memory: %s\n", __func__, + size / 1024.0 / 1024.0, cudaGetErrorString(err)); + return nullptr; + } + + return ptr; +} + +static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + void * ptr = ggml_cuda_host_malloc(size); + + if (ptr == nullptr) { + // fallback to cpu buffer + return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); + } + + ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); + buffer->buft = buft; + buffer->iface.free_buffer = ggml_backend_cuda_host_buffer_free_buffer; + + return buffer; +} + +ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { + static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cuda_host_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_cuda_host_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, + /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), 0), + /* .context = */ nullptr, + }; + + return &ggml_backend_cuda_buffer_type_host; +} + +//static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) { +// return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; +//} + +/// kernels + +typedef void (*ggml_cuda_op_mul_mat_t)( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream); + +#ifndef GGML_CUDA_PEER_MAX_BATCH_SIZE +#define GGML_CUDA_PEER_MAX_BATCH_SIZE 128 +#endif // GGML_CUDA_PEER_MAX_BATCH_SIZE + +#define MUL_MAT_SRC1_COL_STRIDE 128 + +static cudaError_t ggml_cuda_cpy_tensor_2d( + void * dst, const struct ggml_tensor * src, int64_t i3, int64_t i2, int64_t i1_low, int64_t i1_high, cudaStream_t stream) { + + const char * src_ptr = (const char *) src->data; + char * dst_ptr = (char *) dst; + + const int64_t ne0 = src->ne[0]; + const int64_t nb0 = src->nb[0]; + const int64_t nb1 = src->nb[1]; + const int64_t nb2 = src->nb[2]; + const int64_t nb3 = src->nb[3]; + const enum ggml_type type = src->type; + const int64_t ts = ggml_type_size(type); + const int64_t bs = ggml_blck_size(type); + const int64_t i1_diff = i1_high - i1_low; + + const char * x = src_ptr + i1_low*nb1 + i2*nb2 + i3*nb3; + if (nb0 == ts && nb1 == ts*ne0/bs) { + return cudaMemcpyAsync(dst_ptr, x, i1_diff*nb1, cudaMemcpyDeviceToDevice, stream); + } else if (nb0 == ts) { + return cudaMemcpy2DAsync(dst_ptr, ts*ne0/bs, x, nb1, ts*ne0/bs, i1_diff, cudaMemcpyDeviceToDevice, stream); + } else { + for (int64_t i1 = 0; i1 < i1_diff; i1++) { + const void * rx = (const void *) ((const char *) x + i1*nb1); + void * rd = (void *) (dst_ptr + i1*ts*ne0/bs); + // pretend the row is a matrix with cols=1 + cudaError_t r = cudaMemcpy2DAsync(rd, ts/bs, rx, nb0, ts/bs, ne0, cudaMemcpyDeviceToDevice, stream); + if (r != cudaSuccess) { + return r; + } + } + return cudaSuccess; + } +} + +struct cublas_force_compute_type { + bool fp32 = false; + bool fp16 = false; +}; + +static const cublas_force_compute_type & ggml_cuda_cublas_get_force_compute_type() { + static const cublas_force_compute_type compute_type = [] { + cublas_force_compute_type result; + + const bool ggml_cuda_force_cublas_compute_32f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F") != nullptr; + const bool ggml_cuda_force_cublas_compute_16f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F") != nullptr; + + GGML_ASSERT(ggml_cuda_force_cublas_compute_16f_env == false || ggml_cuda_force_cublas_compute_32f_env == false); + + if (ggml_cuda_force_cublas_compute_32f_env) { + GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F\n"); + result.fp32 = true; + } else if (ggml_cuda_force_cublas_compute_16f_env) { + GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F\n"); + result.fp16 = true; + } + + return result; + }(); + + return compute_type; +} + +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) +// hipBLASLt equivalent of the cublasGemm* calls used below. +// rocBLAS does not ship Tensile kernels for every AMD GPU arch (e.g. gfx1103 on Windows), +// while hipBLASLt covers them, so HIP builds route GEMM through hipBLASLt when available. +// Computes C = op(A) * op(B) with op(A) = A^T, op(B) = B (column-major, same as the cublas calls). +// hipBLASLt only accepts hipDataType. ROCm < 6.5 routes cudaDataType_t to the legacy +// hipblasDatatype_t enum (150/151/168), while ROCm >= 6.5 uses hipDataType (0/2/14) directly. +// Accept the raw integer value and map both numbering schemes, so this compiles on all ROCm versions. +static hipDataType ggml_hipblaslt_convert_type(int type) { + switch (type) { + case 150: return HIP_R_16F; // legacy HIPBLAS_R_16F + case 151: return HIP_R_32F; // legacy HIPBLAS_R_32F + case 168: return HIP_R_16BF; // legacy HIPBLAS_R_16B + default: + GGML_ASSERT(type == HIP_R_16F || type == HIP_R_32F || type == HIP_R_16BF); + return (hipDataType) type; + } +} + +static void ggml_hipblaslt_gemm( + ggml_backend_cuda_context & ctx, cudaStream_t stream, + int64_t m, int64_t n, int64_t k, + const void * A, int type_a, int64_t lda, int64_t stride_a, + const void * B, int type_b, int64_t ldb, int64_t stride_b, + void * C, int type_c, int64_t ldc, int64_t stride_c, + int64_t batch_count) { + + const hipblasOperation_t trans_a = HIPBLAS_OP_T; + const hipblasOperation_t trans_b = HIPBLAS_OP_N; + + const float alpha = 1.0f; + const float beta = 0.0f; + + hipblasLtHandle_t lt = ctx.hipblaslt_handle(); + void * workspace = ctx.hipblaslt_workspace(ctx.device); + + hipblasLtMatmulDesc_t matmul_desc; + hipblasLtMatrixLayout_t layout_a, layout_b, layout_c; + hipblasLtMatmulPreference_t pref; + + HIPBLASLT_CHECK(hipblasLtMatmulDescCreate(&matmul_desc, HIPBLAS_COMPUTE_32F, HIP_R_32F)); + HIPBLASLT_CHECK(hipblasLtMatmulDescSetAttribute(matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a))); + HIPBLASLT_CHECK(hipblasLtMatmulDescSetAttribute(matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b))); + + // layout dims describe the stored (pre-op) matrix: A is stored [k, m], B is stored [k, n], C is [m, n] + HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_a, ggml_hipblaslt_convert_type(type_a), k, m, lda)); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_b, ggml_hipblaslt_convert_type(type_b), k, n, ldb)); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutCreate(&layout_c, ggml_hipblaslt_convert_type(type_c), m, n, ldc)); + + if (batch_count > 1) { + int batch_count_i32 = (int) batch_count; + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_a, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_a, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_a, sizeof(stride_a))); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_b, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_b, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_b, sizeof(stride_b))); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_c, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count_i32, sizeof(batch_count_i32))); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutSetAttribute(layout_c, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_c, sizeof(stride_c))); + } + + HIPBLASLT_CHECK(hipblasLtMatmulPreferenceCreate(&pref)); + size_t max_workspace = HIPBLASLT_WORKSPACE_SIZE; + HIPBLASLT_CHECK(hipblasLtMatmulPreferenceSetAttribute(pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace, sizeof(max_workspace))); + + hipblasLtMatmulHeuristicResult_t heuristic; + int algo_count = 0; + HIPBLASLT_CHECK(hipblasLtMatmulAlgoGetHeuristic(lt, matmul_desc, layout_a, layout_b, layout_c, layout_c, + pref, 1, &heuristic, &algo_count)); + GGML_ASSERT(algo_count > 0); + + HIPBLASLT_CHECK(hipblasLtMatmul(lt, matmul_desc, + &alpha, A, layout_a, B, layout_b, + &beta, C, layout_c, C, layout_c, + &heuristic.algo, workspace, max_workspace, stream)); + + HIPBLASLT_CHECK(hipblasLtMatmulPreferenceDestroy(pref)); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_a)); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_b)); + HIPBLASLT_CHECK(hipblasLtMatrixLayoutDestroy(layout_c)); + HIPBLASLT_CHECK(hipblasLtMatmulDescDestroy(matmul_desc)); +} +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + +static void ggml_cuda_op_mul_mat_cublas( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream) { + + GGML_ASSERT(src0_dd_i != nullptr); + GGML_ASSERT(src1_ddf_i != nullptr); + GGML_ASSERT(dst_dd_i != nullptr); + + const int64_t ne00 = src0->ne[0]; + const int64_t ne10 = src1->ne[0]; + + const int64_t ne0 = dst->ne[0]; + + const int64_t row_diff = row_high - row_low; + + int id = ggml_cuda_get_device(); + + // the main device has a larger memory buffer to hold the results from all GPUs + // ldc == nrows of the matrix that cuBLAS writes into + int64_t ldc = id == ctx.device ? ne0 : row_diff; + + const int cc = ggml_cuda_info().devices[id].cc; + + const bool supports_bf16 = + (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc) || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); + + const bool use_fp16 = + src0->type != GGML_TYPE_NVFP4 && + (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && + ggml_is_contiguous(src0) && + row_diff == src0->ne[1] && + dst->op_params[0] == GGML_PREC_DEFAULT; + + if (supports_bf16 && src0->type == GGML_TYPE_BF16 && ggml_is_contiguous(src0) && row_diff == src0->ne[1]) { + ggml_cuda_pool_alloc src1_as_bf16(ctx.pool(id)); + if (src1->type != GGML_TYPE_BF16) { + const to_bf16_cuda_t to_bf16_cuda = ggml_get_to_bf16_cuda(src1->type); + GGML_ASSERT(to_bf16_cuda != nullptr); + size_t ne = src1_ncols*ne10; + src1_as_bf16.alloc(ne); + to_bf16_cuda(src1_ddf_i, src1_as_bf16.get(), ne, stream); + } + const nv_bfloat16 * src1_ptr = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1_ddf_i : src1_as_bf16.get(); + const nv_bfloat16 * src0_ptr = (const nv_bfloat16 *)src0_dd_i; + const float alpha_f32 = 1.0f; + const float beta_f32 = 0.0f; + +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + ggml_cuda_pool_alloc dst_bf16(ctx.pool(id), row_diff*src1_ncols); + ggml_hipblaslt_gemm(ctx, stream, + row_diff, src1_ncols, ne10, + src0_ptr, CUDA_R_16BF, ne00, 0, + src1_ptr, CUDA_R_16BF, ne10, 0, + dst_bf16.get(), CUDA_R_16BF, ldc, 0, + 1); + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); + to_fp32_cuda(dst_bf16.get(), dst_dd_i, row_diff*src1_ncols, stream); + GGML_UNUSED_VARS(alpha_f32, beta_f32); +#else + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha_f32, src0_ptr, CUDA_R_16BF, ne00, + src1_ptr, CUDA_R_16BF, ne10, + &beta_f32, dst_dd_i, CUDA_R_32F, ldc, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + } else if (fast_fp16_hardware_available(cc) && use_fp16) { + // convert src0 and src1 to fp16, multiply as fp16, convert dst to fp32 + ggml_cuda_pool_alloc src0_as_f16(ctx.pool(id)); + if (src0->type != GGML_TYPE_F16) { + const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src0->type); + GGML_ASSERT(to_fp16_cuda != nullptr); + size_t ne = row_diff*ne00; + src0_as_f16.alloc(ne); + to_fp16_cuda(src0_dd_i, src0_as_f16.get(), ne, stream); + } + const half * src0_ptr = src0->type == GGML_TYPE_F16 ? (const half *) src0_dd_i : src0_as_f16.get(); + + ggml_cuda_pool_alloc src1_as_f16(ctx.pool(id)); + if (src1->type != GGML_TYPE_F16) { + const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src1->type); + GGML_ASSERT(to_fp16_cuda != nullptr); + size_t ne = src1_ncols*ne10; + src1_as_f16.alloc(ne); + to_fp16_cuda(src1_ddf_i, src1_as_f16.get(), ne, stream); + } + const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16.get(); + + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + + const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); + + if (!force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) + || GGML_CUDA_CC_IS_RDNA4(cc) + || cc == GGML_CUDA_CC_VOLTA + || force_compute_type.fp32)) + { + const float alpha = 1.0f; + const float beta = 0.0f; +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + GGML_UNUSED_VARS(alpha, beta); + ggml_hipblaslt_gemm(ctx, stream, + row_diff, src1_ncols, ne10, + src0_ptr, CUDA_R_16F, ne00, 0, + src1_ptr, CUDA_R_16F, ne10, 0, + dst_dd_i, CUDA_R_32F, ldc, 0, + 1); +#else + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha, src0_ptr, CUDA_R_16F, ne00, + src1_ptr, CUDA_R_16F, ne10, + &beta, dst_dd_i, CUDA_R_32F, ldc, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + } else { + ggml_cuda_pool_alloc dst_f16(ctx.pool(id), row_diff*src1_ncols); + + const half alpha_f16 = 1.0f; + const half beta_f16 = 0.0f; + +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + GGML_UNUSED_VARS(alpha_f16, beta_f16); + ggml_hipblaslt_gemm(ctx, stream, + row_diff, src1_ncols, ne10, + src0_ptr, CUDA_R_16F, ne00, 0, + src1_ptr, CUDA_R_16F, ne10, 0, + dst_f16.get(), CUDA_R_16F, ldc, 0, + 1); +#else + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha_f16, src0_ptr, CUDA_R_16F, ne00, + src1_ptr, CUDA_R_16F, ne10, + &beta_f16, dst_f16.get(), CUDA_R_16F, ldc, + CUBLAS_COMPUTE_16F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16); + to_fp32_cuda(dst_f16.get(), dst_dd_i, row_diff*src1_ncols, stream); + } + } else { + ggml_cuda_pool_alloc src0_ddq_as_f32(ctx.pool(id)); + ggml_cuda_pool_alloc src1_ddq_as_f32(ctx.pool(id)); + + if (src0->type != GGML_TYPE_F32) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src0->type); + GGML_ASSERT(to_fp32_cuda != nullptr); + src0_ddq_as_f32.alloc(row_diff*ne00); + to_fp32_cuda(src0_dd_i, src0_ddq_as_f32.get(), row_diff*ne00, stream); + } + if (src1->type != GGML_TYPE_F32) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src1->type); + GGML_ASSERT(to_fp32_cuda != nullptr); + src1_ddq_as_f32.alloc(src1_ncols*ne10); + to_fp32_cuda(src1_ddf_i, src1_ddq_as_f32.get(), src1_ncols*ne10, stream); + } + + const float * src0_ddf_i = src0->type == GGML_TYPE_F32 ? (const float *) src0_dd_i : src0_ddq_as_f32.get(); + const float * src1_ddf1_i = src1->type == GGML_TYPE_F32 ? (const float *) src1_ddf_i : src1_ddq_as_f32.get(); + + const float alpha = 1.0f; + const float beta = 0.0f; + +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + GGML_UNUSED_VARS(alpha, beta); + ggml_hipblaslt_gemm(ctx, stream, + row_diff, src1_ncols, ne10, + src0_ddf_i, CUDA_R_32F, ne00, 0, + src1_ddf1_i, CUDA_R_32F, ne10, 0, + dst_dd_i, CUDA_R_32F, ldc, 0, + 1); +#else + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + CUBLAS_CHECK( + cublasSgemm(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha, src0_ddf_i, ne00, + src1_ddf1_i, ne10, + &beta, dst_dd_i, ldc)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + } + + GGML_UNUSED_VARS(dst, src1_ddq_i, src1_padded_row_size); +} + +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) { + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // cudaMemcpy2DAsync may fail with copies between vmm pools of different devices + cudaMemcpy3DPeerParms p = {}; + p.dstDevice = dstDevice; + p.dstPtr = make_cudaPitchedPtr(dst, dpitch, dpitch, height); + p.srcDevice = srcDevice; + p.srcPtr = make_cudaPitchedPtr(src, spitch, spitch, height); + p.extent = make_cudaExtent(width, height, 1); + return cudaMemcpy3DPeerAsync(&p, stream); +#else + // HIP does not support cudaMemcpy3DPeerAsync or vmm pools + GGML_UNUSED(dstDevice); + GGML_UNUSED(srcDevice); + return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, cudaMemcpyDeviceToDevice, stream); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +} + +static void ggml_cuda_op_mul_mat( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ggml_cuda_op_mul_mat_t op, + quantize_cuda_t quantize_src1) { + + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; + const int64_t ne02 = src0->ne[2]; + const int64_t ne03 = src0->ne[3]; + + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + const int64_t ne12 = src1->ne[2]; + const int64_t ne13 = src1->ne[3]; + const int64_t nrows1 = ggml_nrows(src1); + + const int64_t ne0 = dst->ne[0]; + const int64_t ne1 = dst->ne[1]; + + // const int64_t nb10 = src1->nb[0]; + const int64_t nb11 = src1->nb[1]; + const int64_t nb12 = src1->nb[2]; + const int64_t nb13 = src1->nb[3]; + + const int64_t nb2 = dst->nb[2]; + const int64_t nb3 = dst->nb[3]; + + ggml_backend_cuda_buffer_context * src1_ctx = (ggml_backend_cuda_buffer_context *) src1->buffer->context; + ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *) dst->buffer->context; + + GGML_ASSERT(src1->type == GGML_TYPE_F32 || (src1->ne[2] == 1 && src1->ne[3] == 1)); + + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + + const int64_t i02_divisor = ne12 / ne02; + const int64_t i03_divisor = ne13 / ne03; + + const size_t src0_ts = ggml_type_size(src0->type); + const size_t src0_bs = ggml_blck_size(src0->type); + const size_t q8_1_ts = sizeof(block_q8_1); + const size_t q8_1_bs = QK8_1; + + const bool src0_is_contiguous = ggml_is_contiguous(src0); + const bool src1_is_contiguous = ggml_is_contiguous(src1); + + const int64_t src1_padded_col_size = GGML_PAD(ne10, MATRIX_ROW_PADDING); + + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); + GGML_ASSERT(!(split && ne02 > 1)); + GGML_ASSERT(!(split && ne03 > 1)); + GGML_ASSERT(!(split && ne02 < ne12)); + GGML_ASSERT(!(split && ne03 < ne13)); + + ggml_tensor_extra_gpu * src0_extra = split ? (ggml_tensor_extra_gpu *) src0->extra : nullptr; + + + std::array tensor_split; + if (split) { + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; + tensor_split = buft_ctx->tensor_split; + } + + struct dev_data { + int cc; + + ggml_cuda_pool_alloc src0_dd_alloc; + ggml_cuda_pool_alloc src1_ddf_alloc; + ggml_cuda_pool_alloc src1_ddq_alloc; + ggml_cuda_pool_alloc dst_dd_alloc; + + char * src0_dd = nullptr; + float * src1_ddf = nullptr; // float + char * src1_ddq = nullptr; // q8_1 + float * dst_dd = nullptr; + + int64_t row_low; + int64_t row_high; + }; + + dev_data dev[GGML_CUDA_MAX_DEVICES]; + + int used_devices = 0; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + dev[id].cc = ggml_cuda_info().devices[id].cc; + + // by default, use all rows + dev[id].row_low = 0; + dev[id].row_high = ne01; + + // for multi GPU, get the row boundaries from tensor split + // and round to mul_mat_q tile sizes + if (split) { + const int64_t rounding = get_row_rounding(tensor_split); + + if (id != 0) { + dev[id].row_low = ne01*tensor_split[id]; + if (dev[id].row_low < ne01) { + dev[id].row_low -= dev[id].row_low % rounding; + } + } + + if (id != ggml_backend_cuda_get_device_count() - 1) { + dev[id].row_high = ne01*tensor_split[id + 1]; + if (dev[id].row_high < ne01) { + dev[id].row_high -= dev[id].row_high % rounding; + } + } + } + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { + continue; + } + + used_devices++; + + const bool src1_on_device = id == src1_ctx->device; + const bool dst_on_device = id == dst_ctx->device; + + ggml_cuda_set_device(id); + cudaStream_t stream = ctx.stream(id, 0); + + if (src0_is_contiguous) { + dev[id].src0_dd = split ? (char *) src0_extra->data_device[id] : (char *) src0->data; + } else { + // If src0 is not contiguous it will be copied to a temporary buffer. + // This buffer needs to be cleared entirely because multiple regions will function as padding. + const size_t nbytes_data = ggml_nbytes(src0); + const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); + dev[id].src0_dd = dev[id].src0_dd_alloc.alloc(ctx.pool(id), nbytes_data + nbytes_padding); + CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd, 0, nbytes_data + nbytes_padding, stream)); + } + + // If src0 is on a temporary compute buffer (partial offloading) there may be some padding that needs to be cleared: + if (ne00 % MATRIX_ROW_PADDING != 0 && ggml_is_quantized(src0->type) && ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && src0->view_src == nullptr) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + const size_t nbytes_data = ggml_row_size(src0->type, (dev[id].row_high - dev[id].row_low)*ne00); + const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); + CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd + nbytes_data, 0, nbytes_padding, stream)); + } + + if (src1_on_device && src1_is_contiguous) { + dev[id].src1_ddf = (float *) src1->data; + } else { + dev[id].src1_ddf = dev[id].src1_ddf_alloc.alloc(ctx.pool(id), ggml_nelements(src1)); + } + + if (quantize_src1) { + size_t src_1_ddq_size = nrows1*src1_padded_col_size*q8_1_ts/q8_1_bs; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + src_1_ddq_size += get_mmq_x_max_host(dev[id].cc)*sizeof(block_q8_1_mmq); + } + dev[id].src1_ddq = dev[id].src1_ddq_alloc.alloc(ctx.pool(id), src_1_ddq_size); + + if (src1_on_device && src1_is_contiguous) { + quantize_src1( + dev[id].src1_ddf, nullptr, dev[id].src1_ddq, src0->type, ne10, + nb11/sizeof(float), nb12/sizeof(float), nb13/sizeof(float), + src1_padded_col_size, ne11, ne12, ne13, stream); + CUDA_CHECK(cudaGetLastError()); + } + } + + if (dst_on_device) { + dev[id].dst_dd = (float *) dst->data; + } else { + const size_t size_dst_ddf = split ? (dev[id].row_high - dev[id].row_low)*ne1 : ggml_nelements(dst); + dev[id].dst_dd = dev[id].dst_dd_alloc.alloc(ctx.pool(id), size_dst_ddf); + } + } + + // if multiple devices are used they need to wait for the main device + // here an event is recorded that signals that the main device has finished calculating the input data + if (split && used_devices > 1) { + ggml_cuda_set_device(ctx.device); + CUDA_CHECK(cudaEventRecord(src0_extra->events[ctx.device][0], ctx.stream())); + } + + const int64_t src1_col_stride = split && used_devices > 1 ? MUL_MAT_SRC1_COL_STRIDE : ne11; + for (int64_t src1_col_0 = 0; src1_col_0 < ne11; src1_col_0 += src1_col_stride) { + const int64_t is = split ? (src1_col_0/src1_col_stride) % GGML_CUDA_MAX_STREAMS : 0; + const int64_t src1_ncols = src1_col_0 + src1_col_stride > ne11 ? ne11 - src1_col_0 : src1_col_stride; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { + continue; + } + + const bool src1_on_device = id == src1_ctx->device; + const bool dst_on_device = id == dst_ctx->device; + const int64_t row_diff = dev[id].row_high - dev[id].row_low; + + ggml_cuda_set_device(id); + cudaStream_t stream = ctx.stream(id, is); + + // wait for main GPU data if necessary + if (split && (id != ctx.device || is != 0)) { + CUDA_CHECK(cudaStreamWaitEvent(stream, src0_extra->events[ctx.device][0], 0)); + } + + for (int64_t i0 = 0; i0 < ne13*ne12; ++i0) { + const int64_t i03 = i0 / ne12; + const int64_t i02 = i0 % ne12; + + size_t src1_ddq_i_offset = i0*ne11 * src1_padded_col_size*q8_1_ts/q8_1_bs; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + src1_ddq_i_offset += src1_col_0 * sizeof(block_q8_1_mmq); + } else { + src1_ddq_i_offset += src1_col_0 * src1_padded_col_size*q8_1_ts/q8_1_bs; + } + + // for split tensors the data begins at i0 == i0_offset_low + const size_t nbytes_src0_matrix = ne01*ne00*src0_ts / src0_bs; + char * src0_dd_i = dev[id].src0_dd + ((i03/i03_divisor)*ne02 + (i02/i02_divisor)) * nbytes_src0_matrix; + float * src1_ddf_i = dev[id].src1_ddf + (i0*ne11 + src1_col_0) * ne10; + char * src1_ddq_i = dev[id].src1_ddq + src1_ddq_i_offset; + float * dst_dd_i = dev[id].dst_dd + (i0*ne1 + src1_col_0) * (dst_on_device ? ne0 : row_diff); + + // the main device memory buffer can be on VRAM scratch, with space for all partial results + // in that case an offset on dst_ddf_i is needed + if (id == ctx.device) { + dst_dd_i += dev[id].row_low; // offset is 0 if no tensor split + } + + // copy src0, src1 to device if necessary + if (src1_is_contiguous) { + if (id != ctx.device) { + if (quantize_src1) { + char * src1_ddq_i_source = dev[ctx.device].src1_ddq + src1_ddq_i_offset; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + const size_t pitch = ne11*sizeof(block_q8_1_mmq); + const size_t width = src1_ncols*sizeof(block_q8_1_mmq); + const size_t height = src1_padded_col_size/(4*QK8_1); + CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync(src1_ddq_i, id, pitch, src1_ddq_i_source, ctx.device, pitch, width, height, stream)); + } else { + CUDA_CHECK(cudaMemcpyPeerAsync( + src1_ddq_i, id, src1_ddq_i_source, ctx.device, src1_ncols*src1_padded_col_size*q8_1_ts/q8_1_bs, stream)); + } + } else { + float * src1_ddf_i_source = (float *) src1->data; + src1_ddf_i_source += (i0*ne11 + src1_col_0) * ne10; + CUDA_CHECK(cudaMemcpyPeerAsync(src1_ddf_i, id, src1_ddf_i_source, ctx.device, + src1_ncols*ne10*sizeof(float), stream)); + } + } + } else if (src1_on_device && !src1_is_contiguous) { + CUDA_CHECK(ggml_cuda_cpy_tensor_2d( + src1_ddf_i, src1, i03, i02, src1_col_0, src1_col_0+src1_ncols, stream)); + } else { + GGML_ABORT("fatal error"); + } + + if (quantize_src1 && !src1_is_contiguous) { + quantize_src1( + src1_ddf_i, nullptr, src1_ddq_i, src0->type, ne10, ne10, ne11*ne10, ne12*ne11*ne10, + src1_padded_col_size, src1_ncols, 1, 1, stream); + CUDA_CHECK(cudaGetLastError()); + } + + if (src1_col_0 == 0 && !src0_is_contiguous && i03 % i03_divisor == 0 && i02 % i02_divisor == 0) { + CUDA_CHECK(ggml_cuda_cpy_tensor_2d( + src0_dd_i, src0, i03/i03_divisor, i02/i02_divisor, dev[id].row_low, dev[id].row_high, stream)); + } + + // do the computation + op(ctx, src0, src1, dst, src0_dd_i, src1_ddf_i, src1_ddq_i, dst_dd_i, + dev[id].row_low, dev[id].row_high, src1_ncols, src1_padded_col_size, stream); + CUDA_CHECK(cudaGetLastError()); + + // copy dst to host or other device if necessary + if (!dst_on_device) { + void * dst_off_device = dst->data; + if (split) { + // src0 = weight matrix is saved as a transposed matrix for better memory layout. + // dst is NOT transposed. + // The outputs of matrix matrix multiplications can therefore NOT simply be concatenated for >1 GPU. + // Instead they need to be copied to the correct slice in ne0 = dst row index. + // If dst is a vector with ne0 == 1 then you don't have to do this but it still produces correct results. + float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); + GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); + dhf_dst_i += src1_col_0*ne0 + dev[id].row_low; + CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync( + dhf_dst_i, ctx.device, ne0*sizeof(float), dst_dd_i, id, row_diff*sizeof(float), row_diff*sizeof(float), src1_ncols, stream)); + } else { + float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); + GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); + dhf_dst_i += src1_col_0*ne0; + CUDA_CHECK(cudaMemcpyAsync(dhf_dst_i, dst_dd_i, src1_ncols*ne0*sizeof(float), cudaMemcpyDeviceToDevice, stream)); + } + } + + // add event for the main device to wait on until other device is done + if (split && (id != ctx.device || is != 0)) { + CUDA_CHECK(cudaEventRecord(src0_extra->events[id][is], stream)); + } + } + } + } + + // main device waits for all other devices to be finished + if (split && ggml_backend_cuda_get_device_count() > 1) { + int64_t is_max = (ne11 + MUL_MAT_SRC1_COL_STRIDE - 1) / MUL_MAT_SRC1_COL_STRIDE; + is_max = is_max <= GGML_CUDA_MAX_STREAMS ? is_max : GGML_CUDA_MAX_STREAMS; + + ggml_cuda_set_device(ctx.device); + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if (dev[id].row_low == dev[id].row_high) { + continue; + } + for (int64_t is = 0; is < is_max; ++is) { + CUDA_CHECK(cudaStreamWaitEvent(ctx.stream(), src0_extra->events[id][is], 0)); + } + } + } +} + +static __global__ void k_compute_batched_ptrs( + const void * src0_as_f16, const void * src1_as_f16, char * dst, + const void ** ptrs_src, void ** ptrs_dst, + int64_t ne12, int64_t ne13, + int64_t ne23, + size_t nb02, size_t nb03, + size_t nb12, size_t nb13, + size_t nbd2, size_t nbd3, + int64_t r2, int64_t r3) { + const int64_t i13 = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t i12 = blockIdx.y * blockDim.y + threadIdx.y; + + if (i13 >= ne13 || i12 >= ne12) { + return; + } + + const int64_t i03 = i13 / r3; + const int64_t i02 = i12 / r2; + + ptrs_src[0*ne23 + i12 + i13*ne12] = (const char *) src0_as_f16 + i02*nb02 + i03*nb03; + ptrs_src[1*ne23 + i12 + i13*ne12] = (const char *) src1_as_f16 + i12*nb12 + i13*nb13; + ptrs_dst[0*ne23 + i12 + i13*ne12] = ( char *) dst + i12*nbd2 + i13*nbd3; +} + +// Type traits for mapping ggml types to CUDA/cuBLAS types +template +struct batched_mul_mat_traits; + +template<> +struct batched_mul_mat_traits { + using cuda_type = float; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + static inline const cudaDataType_t data_type = CUDA_R_32F; + static inline const ggml_type ggml_type_val = GGML_TYPE_F32; + static inline const float alpha = 1.0f; + static inline const float beta = 0.0f; + static inline const void* get_alpha() { static const float val = alpha; return &val; } + static inline const void* get_beta() { static const float val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp32_nc_cuda(src_type); } +}; + +template<> +struct batched_mul_mat_traits { + using cuda_type = nv_bfloat16; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + static inline const cudaDataType_t data_type = CUDA_R_16BF; + static inline const ggml_type ggml_type_val = GGML_TYPE_BF16; + static inline const float alpha = 1.0f; + static inline const float beta = 0.0f; + static inline const void* get_alpha() { static const float val = alpha; return &val; } + static inline const void* get_beta() { static const float val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_bf16_nc_cuda(src_type); } +}; + +template<> +struct batched_mul_mat_traits { + using cuda_type = half; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_16F; + static inline const cudaDataType_t data_type = CUDA_R_16F; + static inline const ggml_type ggml_type_val = GGML_TYPE_F16; + static inline const half alpha = 1.0; + static inline const half beta = 0.0; + static inline const void* get_alpha() { static const half val = alpha; return &val; } + static inline const void* get_beta() { static const half val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } +}; + +template +static void ggml_cuda_mul_mat_batched_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + using traits = batched_mul_mat_traits; + using cuda_t = typename traits::cuda_type; + + GGML_ASSERT(!ggml_is_transposed(src0)); + GGML_ASSERT(!ggml_is_transposed(src1)); + GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft)); + GGML_ASSERT(src0->type == src0_type); + GGML_ASSERT(ggml_is_contiguous(dst)); + + // Byte offsets and tensor dimensions are currently used in an inconsistent way for dst. + // As long as dst is contiguous this does not matter though. + + GGML_TENSOR_BINARY_OP_LOCALS + + const int64_t ne_dst = ggml_nelements(dst); + cudaStream_t main_stream = ctx.stream(); + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), main_stream)); + + float * dst_ddf = (float *) dst->data; + const size_t ts_src1 = ggml_type_size(src1->type); + GGML_ASSERT(nb10 == ts_src1); + int64_t s11 = nb11 / ts_src1; + int64_t s12 = nb12 / ts_src1; + int64_t s13 = nb13 / ts_src1; + + const cuda_t * src0_ptr = nullptr; + const cuda_t * src1_ptr = nullptr; + + ggml_cuda_pool_alloc src0_alloc(ctx.pool()); + ggml_cuda_pool_alloc src1_alloc(ctx.pool()); + + bool is_src0_cont_2 = ggml_is_contiguous_2(src0); + bool is_src1_cont_2 = ggml_is_contiguous_2(src1); + + // Handle src0 + src0_ptr = (const cuda_t *) src0->data; + + // Handle src1 - convert if necessary + if (src1->type == src0_type) { + src1_ptr = (const cuda_t *) src1->data; + } else { + // Convert src1 to target type using traits conversion functions + const int64_t ne_src1 = ggml_nelements(src1); + src1_alloc.alloc(ne_src1); + + const auto convert_func = traits::get_nc_converter(src1->type); + GGML_ASSERT(convert_func != nullptr); + convert_func(src1->data, src1_alloc.get(), ne10, ne11, ne12, ne13, s11, s12, s13, main_stream); + src1_ptr = src1_alloc.get(); + s11 = ne10; + s12 = ne11*s11; + s13 = ne12*s12; + + is_src1_cont_2 = true; + } + + // Setup destination buffer + ggml_cuda_pool_alloc dst_temp(ctx.pool()); + char * dst_t; + size_t nbd2 = dst->nb[2]; + size_t nbd3 = dst->nb[3]; + + cublasComputeType_t cu_compute_type = traits::compute_type; +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + GGML_UNUSED(cu_compute_type); // only referenced by the cublas fallback paths +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + cudaDataType_t cu_data_type = traits::data_type; + cudaDataType_t cu_data_type_a = traits::data_type; + cudaDataType_t cu_data_type_b = traits::data_type; + const void * alpha = traits::get_alpha(); + const void * beta = traits::get_beta(); + + const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); + + int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + static constexpr bool is_src0_type_f16 = src0_type == GGML_TYPE_F16; + + // bf16 and fp32 are already being computed in fp32 (ensure it using static_assert), + // so checking necessity of forced fp32 only for fp16 src0_type + static_assert(is_src0_type_f16 || traits::compute_type == CUBLAS_COMPUTE_32F); + + const bool need_compute_32f = is_src0_type_f16 && !force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) + || GGML_CUDA_CC_IS_RDNA4(cc) + || cc == GGML_CUDA_CC_VOLTA + || force_compute_type.fp32); + + if (dst->op_params[0] == GGML_PREC_DEFAULT && !need_compute_32f) { + if constexpr (src0_type == GGML_TYPE_F32) { + dst_t = (char *) dst_ddf; // Direct F32 output + } else { + dst_t = (char *) dst_temp.alloc(ne_dst); + nbd2 /= sizeof(float) / sizeof(cuda_t); + nbd3 /= sizeof(float) / sizeof(cuda_t); + } + } else { + dst_t = (char *) dst_ddf; + cu_compute_type = batched_mul_mat_traits::compute_type; + cu_data_type = batched_mul_mat_traits::data_type; + alpha = batched_mul_mat_traits::get_alpha(); + beta = batched_mul_mat_traits::get_beta(); + } + + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + + // broadcast factors + const int64_t r2 = ne12/ne02; + const int64_t r3 = ne13/ne03; + + if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { + // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: + const int64_t sma = ne02 == 1 ? nb03/nb00 : nb02/nb00; + const int64_t smb = ne12 == 1 ? s13 : s12; + + // there is no broadcast and src0, src1 are contiguous across dims 2, 3 + // use cublasGemmStridedBatchedEx +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + GGML_UNUSED_VARS(alpha, beta); + ggml_hipblaslt_gemm(ctx, main_stream, + ne01, ne11, ne10, + src0_ptr, cu_data_type_a, nb01/nb00, sma, + src1_ptr, cu_data_type_b, s11, smb, + dst_t, cu_data_type, ne0, ne1*ne0, + ne12*ne13); +#else + CUBLAS_CHECK( + cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + alpha, src0_ptr, cu_data_type_a, nb01/nb00, sma, // strideA + src1_ptr, cu_data_type_b, s11, smb, // strideB + beta, dst_t, cu_data_type, ne0, ne1*ne0, // strideC + ne12*ne13, + cu_compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + } else { +#if defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + // hipBLASLt has no pointer-array batched GEMM; issue one GEMM per batch element instead. + GGML_UNUSED_VARS(alpha, beta); + const size_t src1_nb2 = (src1->type == src0_type) ? nb12 : s12*sizeof(cuda_t); + const size_t src1_nb3 = (src1->type == src0_type) ? nb13 : s13*sizeof(cuda_t); + for (int64_t i13 = 0; i13 < ne13; i13++) { + for (int64_t i12 = 0; i12 < ne12; i12++) { + const char * ptr_a = (const char *) src0_ptr + (i12/r2)*nb02 + (i13/r3)*nb03; + const char * ptr_b = (const char *) src1_ptr + i12*src1_nb2 + i13*src1_nb3; + char * ptr_c = ( char *) dst_t + i12*nbd2 + i13*nbd3; + ggml_hipblaslt_gemm(ctx, main_stream, + ne01, ne11, ne10, + ptr_a, cu_data_type_a, nb01/nb00, 0, + ptr_b, cu_data_type_b, s11, 0, + ptr_c, cu_data_type, ne0, 0, + 1); + } + } +#else + // use cublasGemmBatchedEx + const int64_t ne23 = ne12*ne13; + + ggml_cuda_pool_alloc ptrs_src(ctx.pool(), 2*ne23); + ggml_cuda_pool_alloc< void *> ptrs_dst(ctx.pool(), 1*ne23); + + size_t src1_stride_size = sizeof(cuda_t); + + const int threads_x = 16; + const int threads_y = 16; + dim3 block_dims(threads_x, threads_y); + + dim3 grid_dims( + (ne13 + threads_x - 1) / threads_x, + (ne12 + threads_y - 1) / threads_y + ); + k_compute_batched_ptrs<<>>( + src0_ptr, src1_ptr, dst_t, + ptrs_src.get(), ptrs_dst.get(), + ne12, ne13, + ne23, + nb02, nb03, + (src1->type == src0_type) ? nb12 : s12*src1_stride_size, + (src1->type == src0_type) ? nb13 : s13*src1_stride_size, + nbd2, nbd3, + r2, r3); + + CUDA_CHECK(cudaGetLastError()); + + CUBLAS_CHECK( + cublasGemmBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + alpha, (const void **) (ptrs_src.get() + 0*ne23), cu_data_type_a, nb01/nb00, + (const void **) (ptrs_src.get() + 1*ne23), cu_data_type_b, s11, + beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, + ne23, + cu_compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif // defined(GGML_USE_HIP) && defined(GGML_HIP_USE_HIPBLASLT) + } + + // Convert output back to F32 if needed + if (dst->op_params[0] == GGML_PREC_DEFAULT && cu_data_type != CUDA_R_32F) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); + to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); + } +} + +static void ggml_cuda_mul_mat_batched_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || src0->type == GGML_TYPE_F32); + + switch (src0->type) { + case GGML_TYPE_F32: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + case GGML_TYPE_BF16: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + case GGML_TYPE_F16: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + default: + GGML_ABORT("Unsupported type"); + } +} + +static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, + const ggml_tensor * ffn_gate, + const ggml_tensor * glu, + const ggml_tensor * ffn_up_bias = nullptr, + const ggml_tensor * ffn_gate_bias = nullptr) { + const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; + + if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { + return false; + } + + const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU; + const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU; + + GGML_ASSERT(ffn_up && ffn_gate && glu); + + if (!is_mul_mat && !is_mul_mat_id) { + return false; + } + + const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; + + if (has_bias) { + if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) { + return false; + } + + if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) { + return false; + } + + if (expected_bias_op == GGML_OP_ADD) { + const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up; + const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate; + if (!up_has_mul || !gate_has_mul) { + return false; + } + } else { // GGML_OP_ADD_ID + if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) { + return false; + } + if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) { + return false; + } + } + } else { + if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { + return false; + } + } + + if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) || + !ggml_are_same_stride(ffn_up->src[0], ffn_gate->src[0])) { + return false; + } + + if (ffn_up->src[1] != ffn_gate->src[1]) { + return false; + } + + if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) { + return false; + } + + static constexpr std::array valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI }; + + if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { + return false; + } + + if (const bool swapped = ggml_get_op_params_i32(glu, 1); swapped) { + return false; + } + + const bool split = ggml_backend_buft_is_cuda_split(ffn_up->src[0]->buffer->buft) || + ggml_backend_buft_is_cuda_split(ffn_gate->src[0]->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + return true; +} + +static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + const ggml_tensor * dst = tensor; + + const bool is_mul_mat = tensor->op == GGML_OP_MUL_MAT || + tensor->op == GGML_OP_MUL_MAT_PACK4; + const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; + + bool use_mul_mat_vec_f = + (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && + src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]); + + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_backend_buft_is_cuda_split(src1->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + //we only support fusion for ncols_dst = 1 + if (is_mul_mat && dst->ne[1] != 1) { + return false; + } + + if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { + return false; + } + + + return use_mul_mat_vec_f; +} + +static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + const ggml_tensor * dst = tensor; + + const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && + src0->view_src; + + bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_F32 && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; + + // fusion is not universally faster on Pascal + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (cc <= GGML_CUDA_CC_PASCAL) { + return false; + } + //we only support fusion for ncols_dst = 1 + if ((tensor->op == GGML_OP_MUL_MAT || + tensor->op == GGML_OP_MUL_MAT_PACK4) && dst->ne[1] != 1) { + return false; + } + + if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { + return false; + } + + + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_backend_buft_is_cuda_split(src1->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + return use_mul_mat_vec_q; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); + + // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. + // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. + // Therefore, in such cases use cuBLAS. + const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE + && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; + + bool use_mul_mat_vec_f = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + bool use_mul_mat_f = !ggml_is_quantized(src0->type) + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 + && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; + bool use_mul_mat_q = ggml_is_quantized(src0->type) && !bad_padding_clear + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + + bool any_gpus_with_slow_fp16 = false; + + if (split) { + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; + auto & tensor_split = buft_ctx->tensor_split; + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + // skip devices that are not going to do any work: + if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { + continue; + } + + const int cc = ggml_cuda_info().devices[id].cc; + const int warp_size = ggml_cuda_info().devices[id].warp_size; + use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); + use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); + any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); + } + } else { + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); + use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); + any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); + } + + // debug helpers + //printf("src0: %8d %8d %8d %8d\n", src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]); + //printf(" %8d %8d %8d %8d\n", src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3]); + //printf("src1: %8d %8d %8d %8d\n", src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]); + //printf(" %8d %8d %8d %8d\n", src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3]); + //printf("src0 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src0), ggml_is_transposed(src0), ggml_type_name(src0->type), src0->name); + //printf("src1 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src1), ggml_is_transposed(src1), ggml_type_name(src1->type), src1->name); + + //TODO update for generic tensor parallelism + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + bool use_batched_cublas_f16 = src0->type == GGML_TYPE_F16 && (src1->type == GGML_TYPE_F16 || !any_gpus_with_slow_fp16); + bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); + bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; + + if (!split && use_mul_mat_vec_f) { + // the custom F16 vector kernel can be used over batched cuBLAS GEMM + // but this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_f) { + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_vec_q) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_q) { + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + } else if (!split && (use_batched_cublas_f16 || use_batched_cublas_bf16 || use_batched_cublas_f32) + && !ggml_is_transposed(src0) && !ggml_is_transposed(src1) && src1->ne[2]*src1->ne[3] > 1) { + // general KQ + KQV multi-batch without FlashAttention + ggml_cuda_mul_mat_batched_cublas(ctx, src0, src1, dst); + } else if (use_mul_mat_vec_f) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_f, nullptr); + } else if (use_mul_mat_vec_q) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_q, quantize_row_q8_1_cuda); + } else if (use_mul_mat_q) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_q, quantize_mmq_q8_1_cuda); + } else { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_cublas, nullptr); + } +} + +static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft) && "mul_mat_id does not support split buffers"); + + GGML_TENSOR_BINARY_OP_LOCALS + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); + if (ne2 <= MMVQ_MAX_BATCH_SIZE) { + if (ggml_is_quantized(src0->type)) { + const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); + if (ne2 <= mmvq_mmid_max) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } + } else { + if (GGML_CUDA_CC_IS_AMD(cc)) { + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, dst); + return; + } + } + } + + if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { + ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); + return; + } + + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); + return; + } + } + + // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization + // TODO: add asserts to verify this. should work with CUDA, HIP, etc. + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(nb12 % nb11 == 0); + GGML_ASSERT(nb2 % nb1 == 0); + + const ggml_type type_src1_sorted = (src0->type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) + || ggml_is_quantized(src0->type) ? GGML_TYPE_F32 : src0->type; + const ggml_type type_dst_sorted = GGML_TYPE_F32; + const size_t ts_src1_sorted = ggml_type_size(type_src1_sorted); + const size_t ts_dst_sorted = ggml_type_size(type_dst_sorted); + + const int64_t n_expert_used = ids->ne[0]; + const int64_t ne_get_rows = ne12 * n_expert_used; + + std::vector ids_to_sorted_host; + ids_to_sorted_host.reserve(2*ne_get_rows); + std::vector ids_from_sorted_host(ne_get_rows); + + ggml_cuda_pool_alloc ids_buf_dev(ctx.pool(), 2*ne_get_rows); + + std::vector tokens_per_expert(ne02); + + ggml_cuda_pool_alloc src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted); + ggml_cuda_pool_alloc dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted); + + std::vector ids_host(ggml_nbytes(ids)); + CUDA_CHECK(cudaMemcpyAsync(ids_host.data(), ids->data, ggml_nbytes(ids), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int64_t i02 = 0; i02 < ne02; ++i02) { // expert matrices + for (int64_t i12 = 0; i12 < ne12; ++i12) { // tokens + for (int64_t iex = 0; iex < n_expert_used; ++iex) { + const int32_t expert_to_use = *(const int32_t *)(ids_host.data() + i12*ids->nb[1] + iex*ids->nb[0]); + assert(expert_to_use >= 0 && expert_to_use < ne02); + if (expert_to_use == i02) { + ids_from_sorted_host[i12*n_expert_used + iex] = ids_to_sorted_host.size(); + ids_to_sorted_host.push_back(i12*ne11 + iex % ne11); + tokens_per_expert[i02]++; + break; + } + } + } + } + GGML_ASSERT(ids_to_sorted_host.size() == size_t(ne_get_rows)); + + ids_to_sorted_host.insert(ids_to_sorted_host.end(), ids_from_sorted_host.begin(), ids_from_sorted_host.end()); + + CUDA_CHECK(cudaMemcpyAsync(ids_buf_dev.ptr, ids_to_sorted_host.data(), 2*ne_get_rows*sizeof(int32_t), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + const int32_t * ids_to_sorted = ids_buf_dev.ptr + 0*ne_get_rows; + const int32_t * ids_from_sorted = ids_buf_dev.ptr + 1*ne_get_rows; + + get_rows_cuda(src1->data, src1->type, ids_to_sorted, src1_sorted.ptr, type_src1_sorted, + ne10, nb11, nb12, nb13, + ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), + ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, stream); + CUDA_CHECK(cudaGetLastError()); + + char * src1_data_cur = (char *) src1_sorted.ptr; + char * dst_data_cur = (char *) dst_sorted.ptr; + for (int64_t i02 = 0; i02 < ne02; ++i02) { + if (tokens_per_expert[i02] == 0) { + continue; + } + + ggml_tensor src0_slice = *src0; + src0_slice.ne[2] = 1; + src0_slice.nb[3] = src0_slice.nb[2]; + src0_slice.op = GGML_OP_VIEW; + src0_slice.view_src = dst->src[0]; // non-const pointer to src0 + src0_slice.data = (char *) src0->data + i02*nb02; + + ggml_tensor src1_slice; + memset(&src1_slice, 0, sizeof(src1_slice)); + src1_slice.buffer = src1->buffer; + src1_slice.type = type_src1_sorted; + src1_slice.ne[0] = ne10; + src1_slice.ne[1] = tokens_per_expert[i02]; + src1_slice.ne[2] = 1; + src1_slice.ne[3] = 1; + src1_slice.nb[0] = ts_src1_sorted; + src1_slice.nb[1] = src1_slice.ne[0] * src1_slice.nb[0]; + src1_slice.nb[2] = src1_slice.ne[1] * src1_slice.nb[1]; + src1_slice.nb[3] = src1_slice.ne[2] * src1_slice.nb[2]; + src1_slice.data = src1_data_cur; + + ggml_tensor dst_slice; + memset(&dst_slice, 0, sizeof(dst_slice)); + dst_slice.buffer = dst->buffer; + dst_slice.type = type_dst_sorted; + dst_slice.ne[0] = ne0; + dst_slice.ne[1] = tokens_per_expert[i02]; + dst_slice.ne[2] = 1; + dst_slice.ne[3] = 1; + dst_slice.nb[0] = ts_dst_sorted; + dst_slice.nb[1] = dst_slice.ne[0] * dst_slice.nb[0]; + dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; + dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; + dst_slice.data = dst_data_cur; + + ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); + CUDA_CHECK(cudaGetLastError()); + + src1_data_cur += src1_slice.nb[2]; + dst_data_cur += dst_slice.nb[2]; + } + + get_rows_cuda(dst_sorted.ptr, type_dst_sorted, ids_from_sorted, dst->data, dst->type, + ne0, ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, + ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), + nb1, nb2, nb3, stream); +} + +static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { + switch (dst->op) { + case GGML_OP_ARGMAX: + ggml_cuda_argmax(ctx, dst); + break; + case GGML_OP_COUNT_EQUAL: + ggml_cuda_count_equal(ctx, dst); + break; + case GGML_OP_REPEAT: + ggml_cuda_op_repeat(ctx, dst); + break; + case GGML_OP_REPEAT_BACK: + ggml_cuda_op_repeat_back(ctx, dst); + break; + case GGML_OP_GET_ROWS: + ggml_cuda_op_get_rows(ctx, dst); + break; + case GGML_OP_GET_ROWS_BACK: + ggml_cuda_op_get_rows_back(ctx, dst); + break; + case GGML_OP_SET_ROWS: + ggml_cuda_op_set_rows(ctx, dst); + break; + case GGML_OP_SET: + ggml_cuda_op_set(ctx, dst); + break; + case GGML_OP_DUP: + ggml_cuda_dup(ctx, dst); + break; + case GGML_OP_CPY: + ggml_cuda_cpy(ctx, dst->src[0], dst->src[1]); + break; + case GGML_OP_CONT: + ggml_cuda_dup(ctx, dst); + break; + case GGML_OP_ADD: + case GGML_OP_ADD1: // TODO: more efficient implementation + ggml_cuda_op_add(ctx, dst); + break; + case GGML_OP_ADD_ID: + ggml_cuda_op_add_id(ctx, dst); + break; + case GGML_OP_SUB: + ggml_cuda_op_sub(ctx, dst); + break; + case GGML_OP_ACC: + ggml_cuda_op_acc(ctx, dst); + break; + case GGML_OP_MUL: + ggml_cuda_op_mul(ctx, dst); + break; + case GGML_OP_DIV: + ggml_cuda_op_div(ctx, dst); + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(dst)) { + case GGML_UNARY_OP_ABS: + ggml_cuda_op_abs(ctx, dst); + break; + case GGML_UNARY_OP_SGN: + ggml_cuda_op_sgn(ctx, dst); + break; + case GGML_UNARY_OP_NEG: + ggml_cuda_op_neg(ctx, dst); + break; + case GGML_UNARY_OP_STEP: + ggml_cuda_op_step(ctx, dst); + break; + case GGML_UNARY_OP_GELU: + ggml_cuda_op_gelu(ctx, dst); + break; + case GGML_UNARY_OP_SILU: + ggml_cuda_op_silu(ctx, dst); + break; + case GGML_UNARY_OP_GELU_ERF: + ggml_cuda_op_gelu_erf(ctx, dst); + break; + case GGML_UNARY_OP_GELU_QUICK: + ggml_cuda_op_gelu_quick(ctx, dst); + break; + case GGML_UNARY_OP_TANH: + ggml_cuda_op_tanh(ctx, dst); + break; + case GGML_UNARY_OP_RELU: + ggml_cuda_op_relu(ctx, dst); + break; + case GGML_UNARY_OP_SIGMOID: + ggml_cuda_op_sigmoid(ctx, dst); + break; + case GGML_UNARY_OP_HARDSIGMOID: + ggml_cuda_op_hardsigmoid(ctx, dst); + break; + case GGML_UNARY_OP_HARDSWISH: + ggml_cuda_op_hardswish(ctx, dst); + break; + case GGML_UNARY_OP_EXP: + ggml_cuda_op_exp(ctx, dst); + break; + case GGML_UNARY_OP_ELU: + ggml_cuda_op_elu(ctx, dst); + break; + case GGML_UNARY_OP_XIELU: + ggml_cuda_op_xielu(ctx, dst); + break; + case GGML_UNARY_OP_FLOOR: + ggml_cuda_op_floor(ctx, dst); + break; + case GGML_UNARY_OP_CEIL: + ggml_cuda_op_ceil(ctx, dst); + break; + case GGML_UNARY_OP_ROUND: + ggml_cuda_op_round(ctx, dst); + break; + case GGML_UNARY_OP_TRUNC: + ggml_cuda_op_trunc(ctx, dst); + break; + case GGML_UNARY_OP_EXPM1: + ggml_cuda_op_expm1(ctx, dst); + break; + case GGML_UNARY_OP_SOFTPLUS: + ggml_cuda_op_softplus(ctx, dst); + break; + default: + return false; + } + break; + case GGML_OP_GLU: + switch (ggml_get_glu_op(dst)) { + case GGML_GLU_OP_REGLU: + ggml_cuda_op_reglu(ctx, dst); + break; + case GGML_GLU_OP_GEGLU: + ggml_cuda_op_geglu(ctx, dst); + break; + case GGML_GLU_OP_SWIGLU: + ggml_cuda_op_swiglu(ctx, dst); + break; + case GGML_GLU_OP_SWIGLU_OAI: + ggml_cuda_op_swiglu_oai(ctx, dst); + break; + case GGML_GLU_OP_GEGLU_ERF: + ggml_cuda_op_geglu_erf(ctx, dst); + break; + case GGML_GLU_OP_GEGLU_QUICK: + ggml_cuda_op_geglu_quick(ctx, dst); + break; + default: + return false; + } + break; + case GGML_OP_NORM: + ggml_cuda_op_norm(ctx, dst); + break; + case GGML_OP_GROUP_NORM: + ggml_cuda_op_group_norm(ctx, dst); + break; + case GGML_OP_L2_NORM: + ggml_cuda_op_l2_norm(ctx, dst); + break; + case GGML_OP_CONCAT: + ggml_cuda_op_concat(ctx, dst); + break; + case GGML_OP_UPSCALE: + ggml_cuda_op_upscale(ctx, dst); + break; + case GGML_OP_PAD: + ggml_cuda_op_pad(ctx, dst); + break; + case GGML_OP_PAD_REFLECT_1D: + ggml_cuda_op_pad_reflect_1d(ctx, dst); + break; + case GGML_OP_ARANGE: + ggml_cuda_op_arange(ctx, dst); + break; + case GGML_OP_TIMESTEP_EMBEDDING: + ggml_cuda_op_timestep_embedding(ctx, dst); + break; + case GGML_OP_LEAKY_RELU: + ggml_cuda_op_leaky_relu(ctx, dst); + break; + case GGML_OP_SILU_BACK: + ggml_cuda_op_silu_back(ctx, dst); + break; + case GGML_OP_RMS_NORM: + ggml_cuda_op_rms_norm(ctx, dst); + break; + case GGML_OP_RMS_NORM_BACK: + ggml_cuda_op_rms_norm_back(ctx, dst); + break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_PACK4: + ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); + break; + case GGML_OP_MUL_MAT_ID: + ggml_cuda_mul_mat_id(ctx, dst); + break; + case GGML_OP_OUT_PROD: + ggml_cuda_out_prod(ctx, dst); + break; + case GGML_OP_SCALE: + ggml_cuda_op_scale(ctx, dst); + break; + case GGML_OP_SQR: + ggml_cuda_op_sqr(ctx, dst); + break; + case GGML_OP_SQRT: + ggml_cuda_op_sqrt(ctx, dst); + break; + case GGML_OP_SIN: + ggml_cuda_op_sin(ctx, dst); + break; + case GGML_OP_COS: + ggml_cuda_op_cos(ctx, dst); + break; + case GGML_OP_CLAMP: + ggml_cuda_op_clamp(ctx, dst); + break; + case GGML_OP_LOG: + ggml_cuda_op_log(ctx, dst); + break; + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + break; + case GGML_OP_DIAG: + ggml_cuda_op_diag(ctx, dst); + break; + case GGML_OP_DIAG_MASK_INF: + ggml_cuda_op_diag_mask_inf(ctx, dst); + break; + case GGML_OP_SOFT_MAX: + ggml_cuda_op_soft_max(ctx, dst); + break; + case GGML_OP_SOFT_MAX_BACK: + ggml_cuda_op_soft_max_back(ctx, dst); + break; + case GGML_OP_ROPE: + ggml_cuda_op_rope(ctx, dst); + break; + case GGML_OP_ROPE_BACK: + ggml_cuda_op_rope_back(ctx, dst); + break; + case GGML_OP_ROLL: + ggml_cuda_op_roll(ctx, dst); + break; + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_FAST_1D: + ggml_cuda_op_im2col(ctx, dst); + break; + case GGML_OP_IM2COL_3D: + ggml_cuda_op_im2col_3d(ctx, dst); + break; + case GGML_OP_COL2IM_1D: + ggml_cuda_op_col2im_1d(ctx, dst); + break; + case GGML_OP_CONV_2D: + ggml_cuda_op_conv2d(ctx, dst); + break; + case GGML_OP_CONV_2D_DW: + ggml_cuda_op_conv2d_dw(ctx, dst); + break; + case GGML_OP_CONV_TRANSPOSE_2D: + ggml_cuda_conv_2d_transpose_p0(ctx, dst); + break; + case GGML_OP_CONV_TRANSPOSE_1D: + ggml_cuda_op_conv_transpose_1d(ctx,dst); + break; + case GGML_OP_POOL_2D: + ggml_cuda_op_pool2d(ctx, dst); + break; + case GGML_OP_SUM: + ggml_cuda_op_sum(ctx, dst); + break; + case GGML_OP_CUMSUM: + ggml_cuda_op_cumsum(ctx, dst); + break; + case GGML_OP_SUM_ROWS: + ggml_cuda_op_sum_rows(ctx, dst); + break; + case GGML_OP_MEAN: + ggml_cuda_op_mean(ctx, dst); + break; + case GGML_OP_SSM_CONV: + ggml_cuda_op_ssm_conv(ctx, dst); + break; + case GGML_OP_SSM_SCAN: + ggml_cuda_op_ssm_scan(ctx, dst); + break; + case GGML_OP_TOP_K: + ggml_cuda_op_top_k(ctx, dst); + break; + case GGML_OP_ARGSORT: + ggml_cuda_op_argsort(ctx, dst); + break; + case GGML_OP_FLASH_ATTN_EXT: + ggml_cuda_flash_attn_ext(ctx, dst); + break; + case GGML_OP_SAGE_ATTN2: + ggml_cuda_sage_attn2(ctx, dst); + break; + case GGML_OP_SAGE_ATTN2_I8: + ggml_cuda_sage_attn2_i8(ctx, dst); + break; + case GGML_OP_CONVROT_LINEAR: + ggml_cuda_convrot_linear(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); + break; + case GGML_OP_TRI: + ggml_cuda_op_tri(ctx, dst); + break; + case GGML_OP_RWKV_WKV6: + ggml_cuda_op_rwkv_wkv6(ctx, dst); + break; + case GGML_OP_GATED_LINEAR_ATTN: + ggml_cuda_op_gated_linear_attn(ctx, dst); + break; + case GGML_OP_GATED_DELTA_NET: + ggml_cuda_op_gated_delta_net(ctx, dst); + break; + case GGML_OP_RWKV_WKV7: + ggml_cuda_op_rwkv_wkv7(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + ggml_cuda_cross_entropy_loss_back(ctx, dst); + break; + case GGML_OP_OPT_STEP_ADAMW: + ggml_cuda_opt_step_adamw(ctx, dst); + break; + case GGML_OP_OPT_STEP_SGD: + ggml_cuda_opt_step_sgd(ctx, dst); + break; + case GGML_OP_SOLVE_TRI: + ggml_cuda_op_solve_tri(ctx, dst); + break; + case GGML_OP_FILL: + ggml_cuda_op_fill(ctx, dst); + break; + default: + return false; + } + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + GGML_LOG_ERROR("%s: %s failed\n", __func__, ggml_op_desc(dst)); + CUDA_CHECK(err); + } + + return true; +} + +//////////////////////////////////////////////////////////////////////////////// + +// backend + +static const char * ggml_backend_cuda_get_name(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + return cuda_ctx->name.c_str(); +} + +static void ggml_backend_cuda_free(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + delete cuda_ctx; + delete backend; +} + +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_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())); +} + +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_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())); +} + +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) { + ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; + ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; + + if (!ggml_backend_is_cuda(backend_src) || !ggml_backend_is_cuda(backend_dst)) { + return false; + } + + 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_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 // NDEBUG + return false; + } + + if (backend_src != backend_dst) { + // copy on src stream + if (cuda_ctx_src->device == cuda_ctx_dst->device) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); + } else { +#ifdef GGML_CUDA_NO_PEER_COPY + 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 // GGML_CUDA_NO_PEER_COPY + } + + // record event on src stream after the copy + if (!cuda_ctx_src->copy_event) { + ggml_cuda_set_device(cuda_ctx_src->device); + CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); + } + + CUDA_CHECK(cudaEventRecord(cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + + // wait on dst stream for the copy to complete + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); + } else { + // src and dst are on the same backend + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); + } + return true; +} + +static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); + + GGML_UNUSED(backend); +} + +#ifdef USE_CUDA_GRAPH +static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { + + bool use_cuda_graph = true; + // Loop over nodes in GGML graph to obtain info needed for CUDA graph + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + + if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { + continue; + } + + if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) { + use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to split buffer\n", __func__); +#endif + } + + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] + if (node->op == GGML_OP_MUL_MAT_ID) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); + if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { + // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs + // TODO: figure out a way to enable for larger batch sizes, without hurting performance + // ref: https://github.com/ggml-org/llama.cpp/pull/18958 + use_cuda_graph = false; +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__); +#endif + } + } + + if (!use_cuda_graph) { + break; + } + } + + return use_cuda_graph; +} + +static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + return cgraph->nodes[0]; +} + +static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { + bool res = false; + + const void * graph_key = ggml_cuda_graph_get_key(cgraph); + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + if (cgraph->uid != 0 && + cgraph->uid == graph->uid) { + GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); + GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); + return false; + } + + graph->uid = cgraph->uid; + + // Check if the graph size has changed + if ((int)graph->node_props.size() != cgraph->n_nodes) { + res = true; + graph->node_props.resize(cgraph->n_nodes); + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_cuda_graph::node_properties prop = {}; + memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); + + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (cgraph->nodes[i]->src[j]) { + prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; + memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); + memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); + } + } + + if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { + graph->node_props[i] = prop; + res = true; + } + } + + return res; +} + +static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + +#if CUDART_VERSION >= 12000 + cudaGraphExecUpdateResultInfo result_info; + cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info); +#else + cudaGraphNode_t errorNode; + cudaGraphExecUpdateResult result_info; + cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info); +#endif // CUDART_VERSION >= 12000 + + if (stat == cudaErrorGraphExecUpdateFailure) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__); +#endif + + // The pre-existing graph exec cannot be updated due to violated constraints + // so instead clear error and re-instantiate + (void)cudaGetLastError(); + CUDA_CHECK(cudaGraphExecDestroy(graph->instance)); + graph->instance = nullptr; + CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + } else { + GGML_ASSERT(stat == cudaSuccess); + } +} +#endif // USE_CUDA_GRAPH + +static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope, + const ggml_tensor * view, + const ggml_tensor * set_rows) { + + if (rope->op != GGML_OP_ROPE || view->op != GGML_OP_VIEW || set_rows->op != GGML_OP_SET_ROWS) { + return false; + } + // ne3 not tested + if (rope->src[0]->ne[3] != 1) { + return false; + } + + if (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) { + return false; + } + + if (set_rows->src[1]->type != GGML_TYPE_I64) { + return false; + } + + // The view should flatten two dims of rope into one dim + if (!ggml_is_contiguous(view) || view->ne[0] != rope->ne[0] * rope->ne[1]) { + return false; + } + + // Only norm/neox shaders have the fusion code + const int mode = ((const int32_t *) rope->op_params)[2]; + if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) { + return false; + } + + return true; +} + +static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { + args.sigmoid = false; + args.softmax = false; + args.delayed_softmax = false; + args.prob_bias = false; + args.norm = false; + + const int n_nodes = cgraph->n_nodes; + ggml_tensor ** nodes = cgraph->nodes; + + if (nodes[node_idx]->op == GGML_OP_SOFT_MAX) { + args.softmax = true; + } + + if (nodes[node_idx]->op == GGML_OP_UNARY) { + if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) { + return false; + } + args.sigmoid = true; + } + + if (nodes[node_idx]->op == GGML_OP_ARGSORT) { + args.delayed_softmax = true; + } + + node_idx++; + + if (args.sigmoid || args.softmax) { + // SOFTMAX -> RESHAPE + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + ggml_tensor * probs_reshaped = nodes[node_idx]; + node_idx++; + + if (node_idx >= n_nodes) { + return false; + } + + // src of bias add is the unreshaped probs (-2 instead of -1) + if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) { + args.prob_bias = true; + node_idx++; + } + // RESHAPE/ADD -> ARGSORT + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) { + return false; + } + + if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) { + return false; + } + + node_idx++; + + // ARGSORT-> VIEW + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) { + return false; + } + + // GET_ROWS + if (nodes[node_idx]->src[0] != probs_reshaped || nodes[node_idx]->src[1] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + } else if (args.delayed_softmax) { + if (node_idx - 2 < 0) { + return false; + } + ggml_tensor * probs_reshaped = nodes[node_idx - 2]; + + // VIEW->ARGSORT + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + + // GET_ROWS + if (node_idx >= n_nodes || nodes[node_idx]->src[1] != nodes[node_idx - 1] || + nodes[node_idx]->src[0] != probs_reshaped) { + return false; + } + node_idx++; + + static const std::vector remaining_ops = { GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }; + + for (const ggml_op op : remaining_ops) { + if (node_idx >= n_nodes || nodes[node_idx]->op != op || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + } + } + + // At this point we can check for norm + scale. Everything is now at least valid till the norm + if (node_idx >= n_nodes) { + return true; + } + + if (nodes[node_idx]->op == GGML_OP_RESHAPE) { + //check RESHAPE->SUM_ROWS->CLAMP->DIV->RESHAPE + static const std::vector norm_ops = { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP }; + + args.norm = true; + for (const ggml_op op : norm_ops) { + if (nodes[node_idx]->op == op && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { + node_idx++; + } else { + args.norm = false; + return true; + } + } + + // DIV <- CLAMP, RESHAPE + if (nodes[node_idx]->op != GGML_OP_DIV || nodes[node_idx]->src[1] != nodes[node_idx - 1] || + nodes[node_idx]->src[0] != nodes[node_idx - 3]) { + args.norm = false; + return true; + } + node_idx++; + + if (nodes[node_idx]->op != GGML_OP_RESHAPE || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + args.norm = false; + return true; + } + + node_idx++; + } + + if (nodes[node_idx]->op == GGML_OP_SCALE && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { + args.scale = true; + } + + return true; +} + +// returns whether the write (out) nodes overwrite the read nodes in operation +static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, + const int node_idx, + const int node_count, + const int * out_nodes, + const int out_count, + const bool is_topk_moe = false) { + auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { + const int64_t a_start = (int64_t) a->data; + const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + + const int64_t b_start = (int64_t) b->data; + const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + + if ((b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end)) { + return true; + } + + return false; + }; + + bool is_ok = true; + // exception for topk-moe, as each row is read entirely before writing + if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) { + return true; + } + + for (int i = 0; i < out_count; ++i) { + const ggml_tensor * dst = cgraph->nodes[out_nodes[i]]; + + for (int j = node_idx; j < node_idx + node_count; ++j) { + // Loop over all srcs of all nodes in the fusion. If the src overlaps + // the destination and the src is not an intermediate node that's being + // elided, then disable fusion. + + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + const ggml_tensor * src = cgraph->nodes[j]->src[src_idx]; + + if (!src || src->op == GGML_OP_NONE) { + continue; + } + + if (nodes_overlap(dst, src)) { + bool found = false; + + for (int k = node_idx; k < j; ++k) { + if (cgraph->nodes[k] == src) { + found = true; + break; + } + } + + if (!found) { + is_ok = false; + break; + } + } + } + } + } + + return is_ok; +} + +// Some model graphs reshape a matvec result before adding the residual. RESHAPE +// is metadata-only and therefore cannot pass the generic compute-node fusion +// validator. Validate this exact chain explicitly so the residual-only Q8_0 +// specialization can write the final result directly. +static bool ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add( + const struct ggml_cgraph * cgraph, int node_idx) { + if (node_idx + 2 >= cgraph->n_nodes) { + return false; + } + + const ggml_tensor * mul_mat = cgraph->nodes[node_idx + 0]; + const ggml_tensor * reshape = cgraph->nodes[node_idx + 1]; + const ggml_tensor * add = cgraph->nodes[node_idx + 2]; + + if (mul_mat->op != GGML_OP_MUL_MAT || + !mul_mat->src[0] || + mul_mat->src[0]->type != GGML_TYPE_Q8_0 || + reshape->op != GGML_OP_RESHAPE || + reshape->src[0] != mul_mat || + add->op != GGML_OP_ADD || + (add->src[0] != reshape && add->src[1] != reshape)) { + return false; + } + + if (ggml_nelements(mul_mat) != ggml_nelements(reshape) || + ggml_nelements(reshape) != ggml_nelements(add) || + ggml_node_get_use_count(cgraph, node_idx + 0) != 1 || + ggml_node_get_use_count(cgraph, node_idx + 1) != 1 || + (mul_mat->flags & GGML_TENSOR_FLAG_OUTPUT) || + (reshape->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + + const int out_nodes[] = { node_idx + 2 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, 3, out_nodes, 1); +} + + +static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, + int node_idx, + std::initializer_list ops, + std::initializer_list unary_ops) { +#ifndef NDEBUG + const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY); + GGML_ASSERT(unary_ops.size() == num_unary); +#endif + + const auto is_equal = [](const std::initializer_list & list1, + const std::initializer_list & list2) { + return std::equal(list1.begin(), list1.end(), list2.begin(), list2.end()); + }; + + std::initializer_list mul_mat_bias_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_GLU }; + std::initializer_list mul_mat_id_bias_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_GLU }; + + std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; + std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + + if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_gate_bias = cgraph->nodes[node_idx + 1]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; + const ggml_tensor * ffn_up_bias = cgraph->nodes[node_idx + 3]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; + + if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu, ffn_up_bias, ffn_gate_bias)) { + int out_nodes[] = { node_idx + 4 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); + } + } + + if ((is_equal(mul_mat_id_glu_ops, ops) || is_equal(mul_mat_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 1]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 2]; + + if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu)) { + int out_nodes[] = { node_idx + 2 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); + } + } + + std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; + + if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { + const ggml_tensor * rope = cgraph->nodes[node_idx]; + const ggml_tensor * view = cgraph->nodes[node_idx + 1]; + const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + + if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) { + return true; + } + } + + if (!ggml_can_fuse(cgraph, node_idx, ops)) { + return false; + } + + if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor *mul = cgraph->nodes[node_idx+1]; + const ggml_tensor *add = nullptr; + + if (ops.size() == 3 && ops.begin()[2] == GGML_OP_ADD) { + add = cgraph->nodes[node_idx+2]; + } + + GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); + + //rms norm only supports F32 + if (mul->src[0]->type != GGML_TYPE_F32 || + mul->src[1]->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32) { + return false; + } + + if (add && (add->src[0]->type != GGML_TYPE_F32 || + add->src[1]->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32) ) { + return false; + } + + //if rms norm is the B operand, then we don't handle broadcast + if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { + return false; + } + + //rms_norm kernel assumes contiguous rows + if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { + return false; + } + + if (add && (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous_rows(add->src[1]))) { + return false; + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY + && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { + const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; + const ggml_tensor * silu = cgraph->nodes[node_idx+1]; + if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) { + return false; + } + + if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { + return false; + } + + return true; + } + + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_ADD + && ops.begin()[2] == GGML_OP_UNARY && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { + const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; + const ggml_tensor * add = cgraph->nodes[node_idx+1]; + const ggml_tensor * silu = cgraph->nodes[node_idx+2]; + if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) { + return false; + } + + if (ssm_conv->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { + return false; + } + + // ADD must consume ssm_conv's output and broadcast a 1-D channel-wise bias. + const ggml_tensor * bias = (add->src[0] == ssm_conv) ? add->src[1] : add->src[0]; + if (bias->type != GGML_TYPE_F32 || !ggml_is_contiguous(bias)) { + return false; + } + if (ggml_nelements(bias) != ssm_conv->ne[0] || bias->ne[0] != ssm_conv->ne[0]) { + return false; + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL + && unary_ops.size() == 1 && (unary_ops.begin()[0] == GGML_UNARY_OP_SILU || unary_ops.begin()[0] == GGML_UNARY_OP_SIGMOID || unary_ops.begin()[0] == GGML_UNARY_OP_SOFTPLUS)) { + const ggml_tensor * unary = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx+1]; + + if (ggml_get_unary_op(unary) != unary_ops.begin()[0]) { + return false; + } + + if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { + return false; + } + + if (unary->type != mul->type) { + return false; + } + + const ggml_tensor * other = (mul->src[0] == unary) ? mul->src[1] : mul->src[0]; + if (other->type != unary->type) { + return false; + } + if (!ggml_is_contiguous_1(other) || !ggml_is_contiguous_1(unary->src[0]) || !ggml_are_same_shape(other, unary)) { + return false; + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_SQR + && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_RELU) { + const ggml_tensor * unary = cgraph->nodes[node_idx]; + const ggml_tensor * sqr = cgraph->nodes[node_idx+1]; + + if (ggml_get_unary_op(unary) != GGML_UNARY_OP_RELU) { + return false; + } + + if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { + return false; + } + + if (unary->type != sqr->type) { + return false; + } + + if (!ggml_is_contiguous(unary->src[0])) { + return false; + } + + return true; + } + + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SCALE && ops.begin()[1] == GGML_OP_UNARY && ops.begin()[2] == GGML_OP_SCALE + && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_TANH) { + const ggml_tensor *scale = cgraph->nodes[node_idx]; + const ggml_tensor *tanh = cgraph->nodes[node_idx+1]; + const ggml_tensor *scale2 = cgraph->nodes[node_idx+2]; + + GGML_ASSERT(scale->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + + if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) { + return false; + } + + // Check for bias + if (ggml_get_op_params_f32(scale, 1) != 0.0f || ggml_get_op_params_f32(scale2, 1) != 0.0f) { + return false; + } + + return true; + } + + return false; +} + +// try and fuse nodes and return the number of nodes to skip +static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { + + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); + if (disable_fusion) { + return 0; + } + + ggml_tensor * node = cgraph->nodes[i]; + + //topk-moe + if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || + cgraph->nodes[i]->op == GGML_OP_ARGSORT) { + ggml_cuda_topk_moe_args args; + const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); + std::vector ops; + + if (can_fuse) { + const ggml_tensor * logits = node->src[0]; + ggml_tensor * weights = nullptr; + ggml_tensor * ids = nullptr; + const ggml_tensor * bias = nullptr; + const ggml_tensor * clamp = nullptr; + const ggml_tensor * scale = nullptr; + + if (!args.delayed_softmax) { + ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX; + int out_nodes[2]; // nodes which can't be elided + + if (args.prob_bias) { + bias = cgraph->nodes[i + 2]->src[1]; + ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, + GGML_OP_GET_ROWS }); + out_nodes[0] = i + 4; + ids = cgraph->nodes[i + 4]; + } else { + ops.insert(ops.end(), + { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); + out_nodes[0] = i + 3; + ids = cgraph->nodes[i + 3]; + } + + if (args.norm) { + ops.insert(ops.end(), + { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }); + clamp = cgraph->nodes[i + ops.size() - 3]; + } + if (args.scale) { + ops.insert(ops.end(), { GGML_OP_SCALE }); + scale = cgraph->nodes[i + ops.size() - 1]; + } + + weights = cgraph->nodes[i + ops.size() - 1]; + out_nodes[1] = i + ops.size() - 1; + + if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && + ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { + ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + return ops.size() - 1; + } + } else if (!args.norm && !args.prob_bias) { + //special case gpt-oss, no norm, no bias. + ops.insert(ops.end(), { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, + GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }); + weights = cgraph->nodes[i + 5]; + ids = cgraph->nodes[i + 1]; + const ggml_tensor * softmax = cgraph->nodes[i + 4]; + + int out_nodes[2] = { i + 1, i + 5 }; + if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && + ggml_cuda_should_use_topk_moe(softmax, logits, weights, ids) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { + ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + return ops.size() - 1; + } + } + } + } + + //RoPE + view + set-rows + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { + ggml_tensor * rope = cgraph->nodes[i]; + ggml_tensor * set_rows = cgraph->nodes[i + 2]; + + ggml_cuda_op_rope_fused(*cuda_ctx, rope, set_rows); + return 2; + } + + // Snake activation: y = x + sin(a*x)^2 * inv_b + // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add + if (ggml_can_fuse_subgraph(cgraph, i, + { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }, + { i + 4 })) { + const ggml_tensor * mul0 = cgraph->nodes[i]; + const ggml_tensor * sqr = cgraph->nodes[i + 2]; + const ggml_tensor * mul1 = cgraph->nodes[i + 3]; + ggml_tensor * add = cgraph->nodes[i + 4]; + + // x carries the full activation shape, a is the broadcast operand + const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; + const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; + + // mul1 reads sqr and inv_b in either operand order + const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; + + // closure check: the trailing add must read the same x as the leading mul + const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; + + // Kernel iterates over total = T * C, so x and add must be 2D and + // a / inv_b must collapse to [1, C, 1, 1]. Higher dims are not handled. + const bool dim_ok = (x->ne[2] == 1 && x->ne[3] == 1) && + (add->ne[2] == 1 && add->ne[3] == 1) && + (a->ne[2] == 1 && a->ne[3] == 1); + const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; + + // x must be in the supported whitelist and every operand / intermediate + // result must share x's type, since launch_snake casts a / inv_b as + // float and templates the kernel on a single T. Mixed precision chains + // fall back to the naive path. + const ggml_tensor * sin1 = cgraph->nodes[i + 1]; + const bool types_ok = (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && + (a->type == x->type) && (inv_b->type == x->type) && + (mul0->type == x->type) && (sin1->type == x->type) && + (sqr->type == x->type) && (mul1->type == x->type) && + (add->type == x->type); + + if (types_ok && shape_ok && dim_ok && x_in_add == x) { + ggml_cuda_op_snake_fused(*cuda_ctx, x, a, inv_b, add); + return 4; + } + } + + // multi-(add or mul) + if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { + int n_fuse = 0; + ggml_op ops[8]; + std::fill(ops, ops + 8, node->op); + + for (; n_fuse <= 6; ++n_fuse) { + if (!ggml_can_fuse(cgraph, i + n_fuse, ops + n_fuse, 2)) { + break; + } + if (cgraph->nodes[i + n_fuse] != cgraph->nodes[i + n_fuse + 1]->src[0]) { + break; + } + if (!ggml_are_same_layout(cgraph->nodes[i + n_fuse]->src[1], cgraph->nodes[i + n_fuse + 1]->src[1])) { + break; + } + } + + n_fuse++; + + if (n_fuse > 1) { + ggml_tensor fused_node; + memcpy(&fused_node, node, sizeof(ggml_tensor)); + for (int j = 0; j < n_fuse - 1; ++j) { + fused_node.src[j + 2] = cgraph->nodes[i + j + 1]->src[1]; + } + fused_node.data = cgraph->nodes[i + n_fuse - 1]->data; + if (node->op == GGML_OP_ADD) { + ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); + } else { + ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); + } + return n_fuse - 1; + } + } + + bool fused_mul_mat_vec = false; + int fused_node_count = 0; + + // gate + glu + up + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { + const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + + if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 4]; + ggml_tensor * gate_bias_n = glu->src[0]; + ggml_tensor * up_bias_n = glu->src[1]; + + //we don't assume the order for {gate, up}. Instead infer it from the bias tensor + ggml_tensor * gate_n = nullptr; + ggml_tensor * up_n = nullptr; + + if (gate_bias_n->src[0] == cgraph->nodes[i] || gate_bias_n->src[1] == cgraph->nodes[i]) { + gate_n = cgraph->nodes[i]; + up_n = cgraph->nodes[i + 2]; + } else if (gate_bias_n->src[0] == cgraph->nodes[i + 2] || gate_bias_n->src[1] == cgraph->nodes[i + 2]) { + gate_n = cgraph->nodes[i + 2]; + up_n = cgraph->nodes[i]; + } else { + continue; + } + + auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) { + if (op_bias == GGML_OP_ADD) { + if (bias_node->src[0] == mul_node) { + return bias_node->src[1]; + } + if (bias_node->src[1] == mul_node) { + return bias_node->src[0]; + } + return (ggml_tensor *) nullptr; + } + GGML_ASSERT(op_bias == GGML_OP_ADD_ID); + GGML_ASSERT(bias_node->src[0] == mul_node); + return bias_node->src[1]; + }; + + ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); + ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); + + if (!up_bias_tensor || !gate_bias_tensor) { + continue; + } + + // we don't support repeating adds + if (bias_op == GGML_OP_ADD && (!ggml_are_same_shape(gate_bias_n->src[0], gate_bias_n->src[1]) || + !ggml_are_same_shape(up_bias_n->src[0], up_bias_n->src[1]))) { + continue; + } + + const ggml_tensor * src0 = up_n->src[0]; + const ggml_tensor * src1 = up_n->src[1]; + const ggml_tensor * ids = up_n->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up_n)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias_tensor; + fusion_data.gate_bias = gate_bias_tensor; + fusion_data.glu_op = ggml_get_glu_op(glu); + + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias_tensor; + fusion_data.gate_bias = gate_bias_tensor; + fusion_data.glu_op = ggml_get_glu_op(glu); + + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } + } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 2]; + ggml_tensor * gate = glu->src[0]; + ggml_tensor * up = glu->src[1]; + + bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) || + (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]); + + if (!ok) { + continue; + } + + const ggml_tensor * src0 = up->src[0]; + const ggml_tensor * src1 = up->src[1]; + const ggml_tensor * ids = up->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + fusion_data.glu_op = ggml_get_glu_op(glu); + + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + fusion_data.glu_op = ggml_get_glu_op(glu); + + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } + } + } + + if (fused_mul_mat_vec) { + return fused_node_count - 1; + } + + fused_mul_mat_vec = false; + fused_node_count = 0; + + // mul_mat + optional metadata-only reshape + add + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { + const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + + const bool reshape_bridge = + op == GGML_OP_MUL_MAT && + ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add(cgraph, i); + if (!reshape_bridge && !ggml_can_fuse(cgraph, i, { op, bias_op })) { + continue; + } + + ggml_tensor * mm_node = cgraph->nodes[i]; + ggml_tensor * mm_output = reshape_bridge ? cgraph->nodes[i + 1] : mm_node; + ggml_tensor * bias_node = cgraph->nodes[i + (reshape_bridge ? 2 : 1)]; + if (reshape_bridge && mm_output->src[0] != mm_node) { + continue; + } + + ggml_tensor * bias_tensor = nullptr; + if (bias_op == GGML_OP_ADD) { + if (bias_node->src[0] == mm_output) { + bias_tensor = bias_node->src[1]; + } else if (bias_node->src[1] == mm_output) { + bias_tensor = bias_node->src[0]; + } else { + continue; + } + } else { + if (bias_node->src[0] != mm_node) { + continue; + } + bias_tensor = bias_node->src[1]; + } + + const ggml_tensor * src0 = mm_node->src[0]; + const ggml_tensor * src1 = mm_node->src[1]; + const ggml_tensor * ids = mm_node->src[2]; + + if (bias_op == GGML_OP_ADD_ID && bias_node->src[2] != ids) { + continue; + } + + if (bias_op == GGML_OP_ADD && !ggml_are_same_shape(bias_node->src[0], bias_node->src[1])) { + continue; + } + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.x_bias = bias_tensor; + fusion_data.residual_only = reshape_bridge; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = reshape_bridge ? 3 : 2; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = reshape_bridge ? 3 : 2; + break; + } + } + + if (fused_mul_mat_vec) { + return fused_node_count - 1; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { + ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + return 2; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { + ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); + return 1; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { + ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + return 2; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { + ggml_cuda_op_ssm_conv(*cuda_ctx, node, /*bias_add_node=*/ nullptr, cgraph->nodes[i + 1]); + return 1; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || + ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || + ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { + ggml_cuda_op_unary_mul(*cuda_ctx, node, cgraph->nodes[i + 1]); + return 1; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_SQR }, { GGML_UNARY_OP_RELU })) { + ggml_cuda_op_relu_sqr(*cuda_ctx, node, cgraph->nodes[i + 1]); + return 1; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { GGML_UNARY_OP_TANH })) { + ggml_cuda_op_softcap(*cuda_ctx, cgraph->nodes[i + 2], node); + return 2; + } + + return 0; +} + +static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { + bool graph_evaluated_or_captured = false; + + // flag used to determine whether it is an integrated_gpu + const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated; + + ggml_cuda_stream_context & stream_ctx = cuda_ctx->stream_context(); + bool is_concurrent_event_active = false; + ggml_cuda_concurrent_event * concurrent_event = nullptr; + bool should_launch_concurrent_events = false; + + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { + if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { + concurrent_event = &stream_ctx.concurrent_events[node]; + + is_concurrent_event_active = true; + + GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); + + cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 + GGML_ASSERT(cuda_ctx->curr_stream_no == 0); + CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); + + for (int i = 1; i <= concurrent_event->n_streams; ++i) { + cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); + CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); + } + } + }; + + while (!graph_evaluated_or_captured) { + // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. + // With the use of CUDA graphs, the execution will be performed by the graph launch. + if (!use_cuda_graph || cuda_graph_update_required) { + [[maybe_unused]] int prev_i = 0; + + if (stream_ctx.concurrent_events.size() > 0) { + should_launch_concurrent_events = true; + for (const auto & [tensor, event] : stream_ctx.concurrent_events) { + should_launch_concurrent_events = should_launch_concurrent_events && event.is_valid(); + } + } + + if (should_launch_concurrent_events) { + // Restore original node order within each concurrent region to enable fusion within streams + + std::unordered_map node_to_idx; + node_to_idx.reserve(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; ++i) { + node_to_idx[cgraph->nodes[i]] = i; + } + + for (auto & [fork_node, event] : stream_ctx.concurrent_events) { + // Find positions of all nodes from this event in the current graph + std::vector positions; + positions.reserve(event.original_order.size()); + + bool all_found = true; + for (const ggml_tensor * orig_node : event.original_order) { + auto it = node_to_idx.find(orig_node); + if (it != node_to_idx.end()) { + positions.push_back(it->second); + } else { + all_found = false; + break; + } + } + + if (!all_found || positions.size() != event.original_order.size()) { + continue; + } + + // Sort positions to get contiguous range + std::vector sorted_positions = positions; + std::sort(sorted_positions.begin(), sorted_positions.end()); + + bool is_contiguous = true; + for (size_t i = 1; i < sorted_positions.size(); ++i) { + if (sorted_positions[i] != sorted_positions[i-1] + 1) { + is_contiguous = false; + break; + } + } + + if (!is_contiguous) { + continue; + } + + // Restore original order at the sorted positions + int start_pos = sorted_positions[0]; + for (size_t i = 0; i < event.original_order.size(); ++i) { + cgraph->nodes[start_pos + i] = const_cast(event.original_order[i]); + } + } + } else { + stream_ctx.concurrent_events.clear(); + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (is_concurrent_event_active) { + GGML_ASSERT(concurrent_event); + + if (node == concurrent_event->join_node) { + cuda_ctx->curr_stream_no = 0; + for (int i = 1; i <= concurrent_event->n_streams; ++i) { + // Wait on join events of forked streams in the main stream + CUDA_CHECK(cudaEventRecord(concurrent_event->join_events[i - 1], + cuda_ctx->stream(cuda_ctx->device, i))); + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), concurrent_event->join_events[i - 1])); + } + + is_concurrent_event_active = false; + concurrent_event = nullptr; + } else { + GGML_ASSERT (concurrent_event->stream_mapping.find(node) != concurrent_event->stream_mapping.end()); + cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); + } + } else if (i - prev_i > 1) { + //the previous node was fused + const ggml_tensor * prev_node = cgraph->nodes[i - 1]; + try_launch_concurrent_event(prev_node); + + if (is_concurrent_event_active) { + cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); + } + } + +#ifdef GGML_CUDA_DEBUG + const int nodes_fused = i - prev_i - 1; + if (nodes_fused > 0) { + GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused); + } +#endif + prev_i = i; + + if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { + continue; + } + + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + continue; + } + + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); + + if (nodes_to_skip != 0) { + i += nodes_to_skip; + continue; + } +#ifndef NDEBUG + assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device)); + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j] != nullptr) { + assert(node->src[j]->buffer); + assert(node->src[j]->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || + ggml_backend_buft_is_cuda_split(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); + } + } +#else + GGML_UNUSED(integrated); +#endif // NDEBUG + + bool ok = ggml_cuda_compute_forward(*cuda_ctx, node); + if (!ok) { + GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); + } + GGML_ASSERT(ok); + + if (!is_concurrent_event_active) { + try_launch_concurrent_event(node); + } + } + } + +#ifdef USE_CUDA_GRAPH + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture + if (graph->graph != nullptr) { + CUDA_CHECK(cudaGraphDestroy(graph->graph)); + graph->graph = nullptr; + } + + CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph)); + graph_evaluated_or_captured = true; // CUDA graph has been captured + + std::lock_guard lock(ggml_cuda_lock); + if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) { + ggml_cuda_lock_cv.notify_all(); + } + } else { + graph_evaluated_or_captured = true; // ggml graph has been directly evaluated + } + } + + if (use_cuda_graph) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->instance == nullptr) { // Create executable graph from captured graph. + CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + } + if (cuda_graph_update_required) { // Update graph executable + ggml_cuda_graph_update_executable(cuda_ctx, graph_key); + } + // Launch graph + CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); +#else + GGML_UNUSED(graph_key); + graph_evaluated_or_captured = true; +#endif // USE_CUDA_GRAPH + } +} + +#ifdef USE_CUDA_GRAPH +static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + if (graph->graph == nullptr) { + if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { + if (!graph->disable_due_to_gpu_arch) { + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); + } + graph->disable_due_to_gpu_arch = true; + } + } + + return graph->is_enabled(); +} +#endif // USE_CUDA_GRAPH + +static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + + ggml_cuda_set_device(cuda_ctx->device); + + bool use_cuda_graph = false; + bool cuda_graph_update_required = false; + const void * graph_key = nullptr; + +#ifdef USE_CUDA_GRAPH + graph_key = ggml_cuda_graph_get_key(cgraph); + + ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); + + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->is_enabled()) { + const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); + if (graph_compatible) { + const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); + + if (!graph->warmup_complete) { + // Warmup: need at least 2 calls with no property change on the 2nd call + if (!properties_changed) { + graph->warmup_complete = true; + GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); + use_cuda_graph = true; + cuda_graph_update_required = true; + } + // else: properties changed or first call - execute directly (use_cuda_graph stays false) + } else { + // Post-warmup: normal CUDA graph operation + if (properties_changed) { + // Properties changed - reset warmup, execute directly until stable again + graph->warmup_complete = false; + GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + } else { + use_cuda_graph = true; + cuda_graph_update_required = graph->instance == nullptr; + } + } + } + } +#endif // USE_CUDA_GRAPH + + if (use_cuda_graph && cuda_graph_update_required) { + // Start CUDA graph capture + { + std::lock_guard lock(ggml_cuda_lock); + ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed); + } + + CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); + } + + ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); + + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + CUDA_CHECK(cudaEventRecord((cudaEvent_t)event->context, cuda_ctx->stream())); +} + +static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + if (ggml_backend_is_cuda(backend)) { + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), (cudaEvent_t)event->context, 0)); + } else { +#if 0 + // untested + auto wait_fn = [](void * user_data) { + ggml_backend_event_t event = (ggml_backend_event_t)user_data; + ggml_backend_event_synchronize(event); + }; + + CUDA_CHECK(cudaLaunchHostFunc(cuda_ctx->stream(), wait_fn, event)); +#endif + GGML_ABORT("fatal error"); + } +} + +static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + +#ifdef USE_CUDA_GRAPH + const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); +#else + const bool use_cuda_graph = false; + GGML_UNUSED(cuda_ctx); + GGML_UNUSED(cgraph); +#endif + + static bool enable_graph_optimization = [] { + const char * env = getenv("GGML_CUDA_GRAPH_OPT"); + return env != nullptr && atoi(env) == 1; + }(); + + if (!enable_graph_optimization) { + return; + } + + ggml_cuda_stream_context & stream_context = cuda_ctx->stream_context(); + stream_context.reset(); + + if (!use_cuda_graph || ggml_backend_cuda_get_device_count() != 1) { + return; + } + + // number of out-degrees for a particular node + std::unordered_map fan_out; + // reverse mapping of node to index in the cgraph + std::unordered_map node_indices; + + const auto & is_noop = [](const ggml_tensor * node) -> bool { + return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; + }; + + const auto & depends_on = [](const ggml_tensor * dst, const ggml_tensor * src) -> bool { + for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { + if (dst->src[s] == src) { + return true; + } + } + // implicit dependency if they view the same tensor + const ggml_tensor * dst2 = dst->view_src ? dst->view_src : dst; + const ggml_tensor * src2 = src->view_src ? src->view_src : src; + if (dst2 == src2) { + return true; + } + return false; + }; + + for (int node_idx = 0; node_idx < cgraph->n_nodes; node_idx++) { + const ggml_tensor * node = cgraph->nodes[node_idx]; + node_indices[node] = node_idx; + + if (is_noop(node)) { + continue; + } + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + const ggml_tensor * src = cgraph->nodes[node_idx]->src[src_idx]; + //TODO: check why nrows > 1 fails + if (node && !is_noop(node) && ggml_nrows(node) <= 1) { + fan_out[src] += 1; + } + } + } + + // Target Q, K, V for concurrency + // this is a more general way to find nodes which can be candidates for concurrency (although it has not been tested for anything else): + // 1. find fan-out (fork) nodes where the same input is used at least N times (in QKV, it would be "attn-norm") + // 2. find the join node, where 2 or more of the outputs are required (in QKV, this would "KQ" or "flash-attn") + // 3. account for all branches from the fork to the join + // 4. To extend lifetimes of the tensors, we interleave the branches (see below for more details) + // 5. save the original cgraph and restore it in graph_compute, to enable fusion within streams + // See discussion: https://github.com/ggml-org/llama.cpp/pull/16991#issuecomment-3522620030 + + const int min_fan_out = 3; + const int max_fan_out = 3; + + // store {fork_idx, join_idx} + std::vector> concurrent_node_ranges; + + for (const auto & [root_node, count] : fan_out) { + if (count >= min_fan_out && count <= max_fan_out) { + const int root_node_idx = node_indices[root_node]; + + // only optimize for attn_norm + // TODO: make this more generic + if (!strstr(root_node->name, "attn_norm")) { + continue; + } + + bool is_part_of_event = false; + for (const auto & [start, end] : concurrent_node_ranges) { + if (root_node_idx >= start && root_node_idx <= end) { + is_part_of_event = true; + } + } + + if (is_part_of_event) { + continue; + } + + std::vector> nodes_per_branch; + for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (!is_noop(node) && depends_on(node, root_node)) { + nodes_per_branch.push_back({ node }); + } + } + + GGML_ASSERT(nodes_per_branch.size() == (size_t) count); + + //find the join point + const ggml_tensor * join_node = nullptr; + + const auto & belongs_to_branch = [&](const ggml_tensor * node, + const std::vector & branch) -> bool { + for (const ggml_tensor * n : branch) { + if (depends_on(node, n)) { + return true; + } + } + return false; + }; + + for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * curr_node = cgraph->nodes[i]; + + int num_joins = 0; + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + if (belongs_to_branch(curr_node, nodes_per_branch[branch_idx])) { + num_joins++; + } + } + + if (num_joins >= 2) { + join_node = curr_node; + break; + } - const ggml_tensor * mul_mat = cgraph->nodes[node_idx + 0]; - const ggml_tensor * reshape = cgraph->nodes[node_idx + 1]; - const ggml_tensor * add = cgraph->nodes[node_idx + 2]; + bool found_branch = false; + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + std::vector & branch_vec = nodes_per_branch[branch_idx]; + if (belongs_to_branch(curr_node, branch_vec)) { + //continue accumulating + if (std::find(branch_vec.begin(), branch_vec.end(), curr_node) == branch_vec.end()) { + branch_vec.push_back(curr_node); + } + found_branch = true; + } + } - if (mul_mat->op != GGML_OP_MUL_MAT || - !mul_mat->src[0] || - mul_mat->src[0]->type != GGML_TYPE_Q8_0 || - reshape->op != GGML_OP_RESHAPE || - reshape->src[0] != mul_mat || - add->op != GGML_OP_ADD || - (add->src[0] != reshape && add->src[1] != reshape)) { - return false; - } + if (!found_branch && is_noop(curr_node)) { + // we can put it in any branch because it will be ignored + nodes_per_branch[0].push_back({ curr_node }); + } + } - if (ggml_nelements(mul_mat) != ggml_nelements(reshape) || - ggml_nelements(reshape) != ggml_nelements(add) || - ggml_node_get_use_count(cgraph, node_idx + 0) != 1 || - ggml_node_get_use_count(cgraph, node_idx + 1) != 1 || - (mul_mat->flags & GGML_TENSOR_FLAG_OUTPUT) || - (reshape->flags & GGML_TENSOR_FLAG_OUTPUT)) { - return false; + if (join_node) { + //Create ggml_cuda_concurrent_event + ggml_cuda_concurrent_event concurrent_event(nodes_per_branch.size()); + concurrent_event.join_node = join_node; + + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + for (const ggml_tensor * n : nodes_per_branch[branch_idx]) { + concurrent_event.stream_mapping[n] = branch_idx + 1; + } + } + + int fork_node_idx = node_indices[root_node]; + int join_node_idx = node_indices[join_node]; + + int current_branch_idx = 0; + int current_node_idx = fork_node_idx + 1; + const int n_branches = nodes_per_branch.size(); + + int total_branch_nodes = 0; + for (std::vector branch_nodes : nodes_per_branch) { + total_branch_nodes += branch_nodes.size(); + } + + // there are other nodes in the middle which are unaccounted for + // usually (cpy) nodes, then ignore this fork + if (join_node_idx - fork_node_idx - 1 != total_branch_nodes) { + GGML_LOG_DEBUG( + "Skipping %s because the number of nodes in the middle is not equal to the total number of " + "branch nodes %d != %d\n", + root_node->name, join_node_idx - fork_node_idx - 1, total_branch_nodes); + continue; + } + + // Save the original order of nodes in this region before interleaving + // This is used later to restore grouping for fusion within streams + concurrent_event.original_order.reserve(total_branch_nodes); + for (int i = fork_node_idx + 1; i < join_node_idx; ++i) { + concurrent_event.original_order.push_back(cgraph->nodes[i]); + } + + std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; + GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); + concurrent_events.emplace(root_node, std::move(concurrent_event)); + GGML_LOG_DEBUG("Adding stream at node %s %p\n", root_node->name, root_node); + concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx); + + // interleave tensors to extend lifetimes so that ggml graph doesn't recycle them + // example transformation: + // [attn-norm, QMul, QNorm, QRope, KMul, KNorm, KRope, VMul, attn] -> + // [attn-norm, QMul, KMul, VMul, QNorm, VNorm, QRope, KRope, attn] + while (current_node_idx < join_node_idx) { + std::vector & branch_nodes = nodes_per_branch[current_branch_idx]; + + bool has_node = false; + for (std::vector branch_node : nodes_per_branch) { + has_node |= branch_node.size() > 0; + } + + GGML_ASSERT(has_node); + + if (branch_nodes.empty()) { + current_branch_idx = (current_branch_idx + 1) % n_branches; + continue; + } + + cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); + current_node_idx++; + branch_nodes.erase(branch_nodes.begin()); + + // append all empty nodes + while (!branch_nodes.empty() && is_noop(branch_nodes.front())) { + cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); + current_node_idx++; + branch_nodes.erase(branch_nodes.begin()); + } + + current_branch_idx = (current_branch_idx + 1) % n_branches; + } + } + } } +} - const int out_nodes[] = { node_idx + 2 }; - return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, 3, out_nodes, 1); +static const ggml_backend_i ggml_backend_cuda_interface = { + /* .get_name = */ ggml_backend_cuda_get_name, + /* .free = */ ggml_backend_cuda_free, + /* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async, + /* .set_tensor_2d_async = */ ggml_backend_cuda_set_tensor_2d_async, + /* .get_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, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_cuda_graph_compute, + /* .event_record = */ ggml_backend_cuda_event_record, + /* .event_wait = */ ggml_backend_cuda_event_wait, + /* .graph_optimize = */ ggml_backend_cuda_graph_optimize, +}; + +static ggml_guid_t ggml_backend_cuda_guid() { + static ggml_guid guid = { 0x2c, 0xdd, 0xe8, 0x1c, 0x65, 0xb3, 0x65, 0x73, 0x6a, 0x12, 0x88, 0x61, 0x1c, 0xc9, 0xdc, 0x25 }; + return &guid; +} + +void * ggml_backend_cuda_get_stream(ggml_backend_t backend) { + GGML_ASSERT(ggml_backend_is_cuda(backend)); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + ggml_cuda_set_device(cuda_ctx->device); + return (void *) cuda_ctx->stream(); +} + +bool ggml_backend_is_cuda(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); } - -static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, - int node_idx, - std::initializer_list ops, - std::initializer_list unary_ops) { -#ifndef NDEBUG - const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY); - GGML_ASSERT(unary_ops.size() == num_unary); -#endif - - const auto is_equal = [](const std::initializer_list & list1, - const std::initializer_list & list2) { - return std::equal(list1.begin(), list1.end(), list2.begin(), list2.end()); - }; - - std::initializer_list mul_mat_bias_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_GLU }; - std::initializer_list mul_mat_id_bias_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_GLU }; - - std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; - std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; - - if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && - ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { - const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; - const ggml_tensor * ffn_gate_bias = cgraph->nodes[node_idx + 1]; - const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; - const ggml_tensor * ffn_up_bias = cgraph->nodes[node_idx + 3]; - const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; - - if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu, ffn_up_bias, ffn_gate_bias)) { - int out_nodes[] = { node_idx + 4 }; - return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); - } - } - - if ((is_equal(mul_mat_id_glu_ops, ops) || is_equal(mul_mat_glu_ops, ops)) && - ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { - const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; - const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 1]; - const ggml_tensor * glu = cgraph->nodes[node_idx + 2]; - - if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu)) { - int out_nodes[] = { node_idx + 2 }; - return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); - } - } - - std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; - - if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { - const ggml_tensor * rope = cgraph->nodes[node_idx]; - const ggml_tensor * view = cgraph->nodes[node_idx + 1]; - const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; - - if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) { - return true; - } - } - - if (!ggml_can_fuse(cgraph, node_idx, ops)) { - return false; - } - - if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { - const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; - const ggml_tensor *mul = cgraph->nodes[node_idx+1]; - const ggml_tensor *add = nullptr; - - if (ops.size() == 3 && ops.begin()[2] == GGML_OP_ADD) { - add = cgraph->nodes[node_idx+2]; - } - - GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); - - //rms norm only supports F32 - if (mul->src[0]->type != GGML_TYPE_F32 || - mul->src[1]->type != GGML_TYPE_F32 || - mul->type != GGML_TYPE_F32) { - return false; - } - - if (add && (add->src[0]->type != GGML_TYPE_F32 || - add->src[1]->type != GGML_TYPE_F32 || - add->type != GGML_TYPE_F32) ) { - return false; - } - - //if rms norm is the B operand, then we don't handle broadcast - if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { - return false; - } - - //rms_norm kernel assumes contiguous rows - if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { - return false; - } - - if (add && (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous_rows(add->src[1]))) { - return false; - } - - return true; - } - - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY - && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { - const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; - const ggml_tensor * silu = cgraph->nodes[node_idx+1]; - if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) { - return false; - } - - if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { - return false; - } - - return true; - } - - if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_ADD - && ops.begin()[2] == GGML_OP_UNARY && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { - const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; - const ggml_tensor * add = cgraph->nodes[node_idx+1]; - const ggml_tensor * silu = cgraph->nodes[node_idx+2]; - if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) { - return false; - } - - if (ssm_conv->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { - return false; - } - - // ADD must consume ssm_conv's output and broadcast a 1-D channel-wise bias. - const ggml_tensor * bias = (add->src[0] == ssm_conv) ? add->src[1] : add->src[0]; - if (bias->type != GGML_TYPE_F32 || !ggml_is_contiguous(bias)) { - return false; - } - if (ggml_nelements(bias) != ssm_conv->ne[0] || bias->ne[0] != ssm_conv->ne[0]) { - return false; - } - - return true; - } - - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL - && unary_ops.size() == 1 && (unary_ops.begin()[0] == GGML_UNARY_OP_SILU || unary_ops.begin()[0] == GGML_UNARY_OP_SIGMOID || unary_ops.begin()[0] == GGML_UNARY_OP_SOFTPLUS)) { - const ggml_tensor * unary = cgraph->nodes[node_idx]; - const ggml_tensor * mul = cgraph->nodes[node_idx+1]; - - if (ggml_get_unary_op(unary) != unary_ops.begin()[0]) { - return false; - } - - if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { - return false; - } - - if (unary->type != mul->type) { - return false; - } - - const ggml_tensor * other = (mul->src[0] == unary) ? mul->src[1] : mul->src[0]; - if (other->type != unary->type) { - return false; - } - if (!ggml_is_contiguous_1(other) || !ggml_is_contiguous_1(unary->src[0]) || !ggml_are_same_shape(other, unary)) { - return false; - } - - return true; - } - - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_SQR - && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_RELU) { - const ggml_tensor * unary = cgraph->nodes[node_idx]; - const ggml_tensor * sqr = cgraph->nodes[node_idx+1]; - - if (ggml_get_unary_op(unary) != GGML_UNARY_OP_RELU) { - return false; - } - - if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { - return false; - } - - if (unary->type != sqr->type) { - return false; - } - - if (!ggml_is_contiguous(unary->src[0])) { - return false; - } - - return true; - } - - if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SCALE && ops.begin()[1] == GGML_OP_UNARY && ops.begin()[2] == GGML_OP_SCALE - && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_TANH) { - const ggml_tensor *scale = cgraph->nodes[node_idx]; - const ggml_tensor *tanh = cgraph->nodes[node_idx+1]; - const ggml_tensor *scale2 = cgraph->nodes[node_idx+2]; - - GGML_ASSERT(scale->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(scale->type == GGML_TYPE_F32); - - if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) { - return false; - } - - // Check for bias - if (ggml_get_op_params_f32(scale, 1) != 0.0f || ggml_get_op_params_f32(scale2, 1) != 0.0f) { - return false; - } - - return true; - } - - return false; -} - -// try and fuse nodes and return the number of nodes to skip -static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { - - static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); - if (disable_fusion) { - return 0; - } - - ggml_tensor * node = cgraph->nodes[i]; - - //topk-moe - if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || - cgraph->nodes[i]->op == GGML_OP_ARGSORT) { - ggml_cuda_topk_moe_args args; - const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); - std::vector ops; - - if (can_fuse) { - const ggml_tensor * logits = node->src[0]; - ggml_tensor * weights = nullptr; - ggml_tensor * ids = nullptr; - const ggml_tensor * bias = nullptr; - const ggml_tensor * clamp = nullptr; - const ggml_tensor * scale = nullptr; - - if (!args.delayed_softmax) { - ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX; - int out_nodes[2]; // nodes which can't be elided - - if (args.prob_bias) { - bias = cgraph->nodes[i + 2]->src[1]; - ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, - GGML_OP_GET_ROWS }); - out_nodes[0] = i + 4; - ids = cgraph->nodes[i + 4]; - } else { - ops.insert(ops.end(), - { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); - out_nodes[0] = i + 3; - ids = cgraph->nodes[i + 3]; - } - - if (args.norm) { - ops.insert(ops.end(), - { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }); - clamp = cgraph->nodes[i + ops.size() - 3]; - } - if (args.scale) { - ops.insert(ops.end(), { GGML_OP_SCALE }); - scale = cgraph->nodes[i + ops.size() - 1]; - } - - weights = cgraph->nodes[i + ops.size() - 1]; - out_nodes[1] = i + ops.size() - 1; - - if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && - ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); - return ops.size() - 1; - } - } else if (!args.norm && !args.prob_bias) { - //special case gpt-oss, no norm, no bias. - ops.insert(ops.end(), { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, - GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }); - weights = cgraph->nodes[i + 5]; - ids = cgraph->nodes[i + 1]; - const ggml_tensor * softmax = cgraph->nodes[i + 4]; - - int out_nodes[2] = { i + 1, i + 5 }; - if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && - ggml_cuda_should_use_topk_moe(softmax, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); - return ops.size() - 1; - } - } - } - } - - //RoPE + view + set-rows - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { - ggml_tensor * rope = cgraph->nodes[i]; - ggml_tensor * set_rows = cgraph->nodes[i + 2]; - - ggml_cuda_op_rope_fused(*cuda_ctx, rope, set_rows); - return 2; - } - - // Snake activation: y = x + sin(a*x)^2 * inv_b - // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add - if (ggml_can_fuse_subgraph(cgraph, i, - { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }, - { i + 4 })) { - const ggml_tensor * mul0 = cgraph->nodes[i]; - const ggml_tensor * sqr = cgraph->nodes[i + 2]; - const ggml_tensor * mul1 = cgraph->nodes[i + 3]; - ggml_tensor * add = cgraph->nodes[i + 4]; - - // x carries the full activation shape, a is the broadcast operand - const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; - const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; - - // mul1 reads sqr and inv_b in either operand order - const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; - - // closure check: the trailing add must read the same x as the leading mul - const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; - - // Kernel iterates over total = T * C, so x and add must be 2D and - // a / inv_b must collapse to [1, C, 1, 1]. Higher dims are not handled. - const bool dim_ok = (x->ne[2] == 1 && x->ne[3] == 1) && - (add->ne[2] == 1 && add->ne[3] == 1) && - (a->ne[2] == 1 && a->ne[3] == 1); - const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; - - // x must be in the supported whitelist and every operand / intermediate - // result must share x's type, since launch_snake casts a / inv_b as - // float and templates the kernel on a single T. Mixed precision chains - // fall back to the naive path. - const ggml_tensor * sin1 = cgraph->nodes[i + 1]; - const bool types_ok = (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && - (a->type == x->type) && (inv_b->type == x->type) && - (mul0->type == x->type) && (sin1->type == x->type) && - (sqr->type == x->type) && (mul1->type == x->type) && - (add->type == x->type); - - if (types_ok && shape_ok && dim_ok && x_in_add == x) { - ggml_cuda_op_snake_fused(*cuda_ctx, x, a, inv_b, add); - return 4; - } - } - - // multi-(add or mul) - if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { - int n_fuse = 0; - ggml_op ops[8]; - std::fill(ops, ops + 8, node->op); - - for (; n_fuse <= 6; ++n_fuse) { - if (!ggml_can_fuse(cgraph, i + n_fuse, ops + n_fuse, 2)) { - break; - } - if (cgraph->nodes[i + n_fuse] != cgraph->nodes[i + n_fuse + 1]->src[0]) { - break; - } - if (!ggml_are_same_layout(cgraph->nodes[i + n_fuse]->src[1], cgraph->nodes[i + n_fuse + 1]->src[1])) { - break; - } - } - - n_fuse++; - - if (n_fuse > 1) { - ggml_tensor fused_node; - memcpy(&fused_node, node, sizeof(ggml_tensor)); - for (int j = 0; j < n_fuse - 1; ++j) { - fused_node.src[j + 2] = cgraph->nodes[i + j + 1]->src[1]; - } - fused_node.data = cgraph->nodes[i + n_fuse - 1]->data; - if (node->op == GGML_OP_ADD) { - ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); - } else { - ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); - } - return n_fuse - 1; - } - } - - bool fused_mul_mat_vec = false; - int fused_node_count = 0; - - // gate + glu + up - for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { - const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; - - if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) { - ggml_tensor * glu = cgraph->nodes[i + 4]; - ggml_tensor * gate_bias_n = glu->src[0]; - ggml_tensor * up_bias_n = glu->src[1]; - - //we don't assume the order for {gate, up}. Instead infer it from the bias tensor - ggml_tensor * gate_n = nullptr; - ggml_tensor * up_n = nullptr; - - if (gate_bias_n->src[0] == cgraph->nodes[i] || gate_bias_n->src[1] == cgraph->nodes[i]) { - gate_n = cgraph->nodes[i]; - up_n = cgraph->nodes[i + 2]; - } else if (gate_bias_n->src[0] == cgraph->nodes[i + 2] || gate_bias_n->src[1] == cgraph->nodes[i + 2]) { - gate_n = cgraph->nodes[i + 2]; - up_n = cgraph->nodes[i]; - } else { - continue; - } - - auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) { - if (op_bias == GGML_OP_ADD) { - if (bias_node->src[0] == mul_node) { - return bias_node->src[1]; - } - if (bias_node->src[1] == mul_node) { - return bias_node->src[0]; - } - return (ggml_tensor *) nullptr; - } - GGML_ASSERT(op_bias == GGML_OP_ADD_ID); - GGML_ASSERT(bias_node->src[0] == mul_node); - return bias_node->src[1]; - }; - - ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); - ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); - - if (!up_bias_tensor || !gate_bias_tensor) { - continue; - } - - // we don't support repeating adds - if (bias_op == GGML_OP_ADD && (!ggml_are_same_shape(gate_bias_n->src[0], gate_bias_n->src[1]) || - !ggml_are_same_shape(up_bias_n->src[0], up_bias_n->src[1]))) { - continue; - } - - const ggml_tensor * src0 = up_n->src[0]; - const ggml_tensor * src1 = up_n->src[1]; - const ggml_tensor * ids = up_n->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up_n)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate_n->src[0]; - fusion_data.x_bias = up_bias_tensor; - fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 5; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate_n->src[0]; - fusion_data.x_bias = up_bias_tensor; - fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 5; - break; - } - } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { - ggml_tensor * glu = cgraph->nodes[i + 2]; - ggml_tensor * gate = glu->src[0]; - ggml_tensor * up = glu->src[1]; - - bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) || - (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]); - - if (!ok) { - continue; - } - - const ggml_tensor * src0 = up->src[0]; - const ggml_tensor * src1 = up->src[1]; - const ggml_tensor * ids = up->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - } - } - - if (fused_mul_mat_vec) { - return fused_node_count - 1; - } - - fused_mul_mat_vec = false; - fused_node_count = 0; - - // mul_mat + optional metadata-only reshape + add - for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { - const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; - - const bool reshape_bridge = - op == GGML_OP_MUL_MAT && - ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add(cgraph, i); - if (!reshape_bridge && !ggml_can_fuse(cgraph, i, { op, bias_op })) { - continue; - } - - ggml_tensor * mm_node = cgraph->nodes[i]; - ggml_tensor * mm_output = reshape_bridge ? cgraph->nodes[i + 1] : mm_node; - ggml_tensor * bias_node = cgraph->nodes[i + (reshape_bridge ? 2 : 1)]; - if (reshape_bridge && mm_output->src[0] != mm_node) { - continue; - } - - ggml_tensor * bias_tensor = nullptr; - if (bias_op == GGML_OP_ADD) { - if (bias_node->src[0] == mm_output) { - bias_tensor = bias_node->src[1]; - } else if (bias_node->src[1] == mm_output) { - bias_tensor = bias_node->src[0]; - } else { - continue; - } - } else { - if (bias_node->src[0] != mm_node) { - continue; - } - bias_tensor = bias_node->src[1]; - } - - const ggml_tensor * src0 = mm_node->src[0]; - const ggml_tensor * src1 = mm_node->src[1]; - const ggml_tensor * ids = mm_node->src[2]; - - if (bias_op == GGML_OP_ADD_ID && bias_node->src[2] != ids) { - continue; - } - - if (bias_op == GGML_OP_ADD && !ggml_are_same_shape(bias_node->src[0], bias_node->src[1])) { - continue; - } - - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.x_bias = bias_tensor; - fusion_data.residual_only = reshape_bridge; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = reshape_bridge ? 3 : 2; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = reshape_bridge ? 3 : 2; - break; - } - } - - if (fused_mul_mat_vec) { - return fused_node_count - 1; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { - ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); - return 2; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { - ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); - return 1; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { - ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); - return 2; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { - ggml_cuda_op_ssm_conv(*cuda_ctx, node, /*bias_add_node=*/ nullptr, cgraph->nodes[i + 1]); - return 1; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || - ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || - ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { - ggml_cuda_op_unary_mul(*cuda_ctx, node, cgraph->nodes[i + 1]); - return 1; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_SQR }, { GGML_UNARY_OP_RELU })) { - ggml_cuda_op_relu_sqr(*cuda_ctx, node, cgraph->nodes[i + 1]); - return 1; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { GGML_UNARY_OP_TANH })) { - ggml_cuda_op_softcap(*cuda_ctx, cgraph->nodes[i + 2], node); - return 2; - } - - return 0; -} - -static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { - bool graph_evaluated_or_captured = false; - - // flag used to determine whether it is an integrated_gpu - const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated; - - ggml_cuda_stream_context & stream_ctx = cuda_ctx->stream_context(); - bool is_concurrent_event_active = false; - ggml_cuda_concurrent_event * concurrent_event = nullptr; - bool should_launch_concurrent_events = false; - - const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { - if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { - concurrent_event = &stream_ctx.concurrent_events[node]; - - is_concurrent_event_active = true; - - GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); - - cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 - GGML_ASSERT(cuda_ctx->curr_stream_no == 0); - CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); - - for (int i = 1; i <= concurrent_event->n_streams; ++i) { - cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); - CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); - } - } - }; - - while (!graph_evaluated_or_captured) { - // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. - // With the use of CUDA graphs, the execution will be performed by the graph launch. - if (!use_cuda_graph || cuda_graph_update_required) { - [[maybe_unused]] int prev_i = 0; - - if (stream_ctx.concurrent_events.size() > 0) { - should_launch_concurrent_events = true; - for (const auto & [tensor, event] : stream_ctx.concurrent_events) { - should_launch_concurrent_events = should_launch_concurrent_events && event.is_valid(); - } - } - - if (should_launch_concurrent_events) { - // Restore original node order within each concurrent region to enable fusion within streams - - std::unordered_map node_to_idx; - node_to_idx.reserve(cgraph->n_nodes); - for (int i = 0; i < cgraph->n_nodes; ++i) { - node_to_idx[cgraph->nodes[i]] = i; - } - - for (auto & [fork_node, event] : stream_ctx.concurrent_events) { - // Find positions of all nodes from this event in the current graph - std::vector positions; - positions.reserve(event.original_order.size()); - - bool all_found = true; - for (const ggml_tensor * orig_node : event.original_order) { - auto it = node_to_idx.find(orig_node); - if (it != node_to_idx.end()) { - positions.push_back(it->second); - } else { - all_found = false; - break; - } - } - - if (!all_found || positions.size() != event.original_order.size()) { - continue; - } - - // Sort positions to get contiguous range - std::vector sorted_positions = positions; - std::sort(sorted_positions.begin(), sorted_positions.end()); - - bool is_contiguous = true; - for (size_t i = 1; i < sorted_positions.size(); ++i) { - if (sorted_positions[i] != sorted_positions[i-1] + 1) { - is_contiguous = false; - break; - } - } - - if (!is_contiguous) { - continue; - } - - // Restore original order at the sorted positions - int start_pos = sorted_positions[0]; - for (size_t i = 0; i < event.original_order.size(); ++i) { - cgraph->nodes[start_pos + i] = const_cast(event.original_order[i]); - } - } - } else { - stream_ctx.concurrent_events.clear(); - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (is_concurrent_event_active) { - GGML_ASSERT(concurrent_event); - - if (node == concurrent_event->join_node) { - cuda_ctx->curr_stream_no = 0; - for (int i = 1; i <= concurrent_event->n_streams; ++i) { - // Wait on join events of forked streams in the main stream - CUDA_CHECK(cudaEventRecord(concurrent_event->join_events[i - 1], - cuda_ctx->stream(cuda_ctx->device, i))); - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), concurrent_event->join_events[i - 1])); - } - - is_concurrent_event_active = false; - concurrent_event = nullptr; - } else { - GGML_ASSERT (concurrent_event->stream_mapping.find(node) != concurrent_event->stream_mapping.end()); - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; - GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); - } - } else if (i - prev_i > 1) { - //the previous node was fused - const ggml_tensor * prev_node = cgraph->nodes[i - 1]; - try_launch_concurrent_event(prev_node); - - if (is_concurrent_event_active) { - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; - GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); - } - } - -#ifdef GGML_CUDA_DEBUG - const int nodes_fused = i - prev_i - 1; - if (nodes_fused > 0) { - GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused); - } -#endif - prev_i = i; - - if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { - continue; - } - - if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - continue; - } - - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); - - if (nodes_to_skip != 0) { - i += nodes_to_skip; - continue; - } -#ifndef NDEBUG - assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device)); - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j] != nullptr) { - assert(node->src[j]->buffer); - assert(node->src[j]->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || - ggml_backend_buft_is_cuda_split(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); - } - } -#else - GGML_UNUSED(integrated); -#endif // NDEBUG - - bool ok = ggml_cuda_compute_forward(*cuda_ctx, node); - if (!ok) { - GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); - } - GGML_ASSERT(ok); - - if (!is_concurrent_event_active) { - try_launch_concurrent_event(node); - } - } - } - -#ifdef USE_CUDA_GRAPH - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture - if (graph->graph != nullptr) { - CUDA_CHECK(cudaGraphDestroy(graph->graph)); - graph->graph = nullptr; - } - - CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph)); - graph_evaluated_or_captured = true; // CUDA graph has been captured - - std::lock_guard lock(ggml_cuda_lock); - if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) { - ggml_cuda_lock_cv.notify_all(); - } - } else { - graph_evaluated_or_captured = true; // ggml graph has been directly evaluated - } - } - - if (use_cuda_graph) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->instance == nullptr) { // Create executable graph from captured graph. - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); - } - if (cuda_graph_update_required) { // Update graph executable - ggml_cuda_graph_update_executable(cuda_ctx, graph_key); - } - // Launch graph - CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); -#else - GGML_UNUSED(graph_key); - graph_evaluated_or_captured = true; -#endif // USE_CUDA_GRAPH - } -} - -#ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - - if (graph->graph == nullptr) { - if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { - if (!graph->disable_due_to_gpu_arch) { - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); - } - graph->disable_due_to_gpu_arch = true; - } - } - - return graph->is_enabled(); -} -#endif // USE_CUDA_GRAPH - -static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - - ggml_cuda_set_device(cuda_ctx->device); - - bool use_cuda_graph = false; - bool cuda_graph_update_required = false; - const void * graph_key = nullptr; - -#ifdef USE_CUDA_GRAPH - graph_key = ggml_cuda_graph_get_key(cgraph); - - ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); - - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->is_enabled()) { - const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); - if (graph_compatible) { - const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); - - if (!graph->warmup_complete) { - // Warmup: need at least 2 calls with no property change on the 2nd call - if (!properties_changed) { - graph->warmup_complete = true; - GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); - use_cuda_graph = true; - cuda_graph_update_required = true; - } - // else: properties changed or first call - execute directly (use_cuda_graph stays false) - } else { - // Post-warmup: normal CUDA graph operation - if (properties_changed) { - // Properties changed - reset warmup, execute directly until stable again - graph->warmup_complete = false; - GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); - } else { - use_cuda_graph = true; - cuda_graph_update_required = graph->instance == nullptr; - } - } - } - } -#endif // USE_CUDA_GRAPH - - if (use_cuda_graph && cuda_graph_update_required) { - // Start CUDA graph capture - { - std::lock_guard lock(ggml_cuda_lock); - ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed); - } - - CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); - } - - ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); - - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - CUDA_CHECK(cudaEventRecord((cudaEvent_t)event->context, cuda_ctx->stream())); -} - -static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - if (ggml_backend_is_cuda(backend)) { - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), (cudaEvent_t)event->context, 0)); - } else { -#if 0 - // untested - auto wait_fn = [](void * user_data) { - ggml_backend_event_t event = (ggml_backend_event_t)user_data; - ggml_backend_event_synchronize(event); - }; - - CUDA_CHECK(cudaLaunchHostFunc(cuda_ctx->stream(), wait_fn, event)); -#endif - GGML_ABORT("fatal error"); - } -} - -static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - -#ifdef USE_CUDA_GRAPH - const void * graph_key = ggml_cuda_graph_get_key(cgraph); - const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); -#else - const bool use_cuda_graph = false; - GGML_UNUSED(cuda_ctx); - GGML_UNUSED(cgraph); -#endif - - static bool enable_graph_optimization = [] { - const char * env = getenv("GGML_CUDA_GRAPH_OPT"); - return env != nullptr && atoi(env) == 1; - }(); - - if (!enable_graph_optimization) { - return; - } - - ggml_cuda_stream_context & stream_context = cuda_ctx->stream_context(); - stream_context.reset(); - - if (!use_cuda_graph || ggml_backend_cuda_get_device_count() != 1) { - return; - } - - // number of out-degrees for a particular node - std::unordered_map fan_out; - // reverse mapping of node to index in the cgraph - std::unordered_map node_indices; - - const auto & is_noop = [](const ggml_tensor * node) -> bool { - return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || - node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; - }; - - const auto & depends_on = [](const ggml_tensor * dst, const ggml_tensor * src) -> bool { - for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { - if (dst->src[s] == src) { - return true; - } - } - // implicit dependency if they view the same tensor - const ggml_tensor * dst2 = dst->view_src ? dst->view_src : dst; - const ggml_tensor * src2 = src->view_src ? src->view_src : src; - if (dst2 == src2) { - return true; - } - return false; - }; - - for (int node_idx = 0; node_idx < cgraph->n_nodes; node_idx++) { - const ggml_tensor * node = cgraph->nodes[node_idx]; - node_indices[node] = node_idx; - - if (is_noop(node)) { - continue; - } - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - const ggml_tensor * src = cgraph->nodes[node_idx]->src[src_idx]; - //TODO: check why nrows > 1 fails - if (node && !is_noop(node) && ggml_nrows(node) <= 1) { - fan_out[src] += 1; - } - } - } - - // Target Q, K, V for concurrency - // this is a more general way to find nodes which can be candidates for concurrency (although it has not been tested for anything else): - // 1. find fan-out (fork) nodes where the same input is used at least N times (in QKV, it would be "attn-norm") - // 2. find the join node, where 2 or more of the outputs are required (in QKV, this would "KQ" or "flash-attn") - // 3. account for all branches from the fork to the join - // 4. To extend lifetimes of the tensors, we interleave the branches (see below for more details) - // 5. save the original cgraph and restore it in graph_compute, to enable fusion within streams - // See discussion: https://github.com/ggml-org/llama.cpp/pull/16991#issuecomment-3522620030 - - const int min_fan_out = 3; - const int max_fan_out = 3; - - // store {fork_idx, join_idx} - std::vector> concurrent_node_ranges; - - for (const auto & [root_node, count] : fan_out) { - if (count >= min_fan_out && count <= max_fan_out) { - const int root_node_idx = node_indices[root_node]; - - // only optimize for attn_norm - // TODO: make this more generic - if (!strstr(root_node->name, "attn_norm")) { - continue; - } - - bool is_part_of_event = false; - for (const auto & [start, end] : concurrent_node_ranges) { - if (root_node_idx >= start && root_node_idx <= end) { - is_part_of_event = true; - } - } - - if (is_part_of_event) { - continue; - } - - std::vector> nodes_per_branch; - for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * node = cgraph->nodes[i]; - if (!is_noop(node) && depends_on(node, root_node)) { - nodes_per_branch.push_back({ node }); - } - } - - GGML_ASSERT(nodes_per_branch.size() == (size_t) count); - - //find the join point - const ggml_tensor * join_node = nullptr; - - const auto & belongs_to_branch = [&](const ggml_tensor * node, - const std::vector & branch) -> bool { - for (const ggml_tensor * n : branch) { - if (depends_on(node, n)) { - return true; - } - } - return false; - }; - - for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * curr_node = cgraph->nodes[i]; - - int num_joins = 0; - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - if (belongs_to_branch(curr_node, nodes_per_branch[branch_idx])) { - num_joins++; - } - } - - if (num_joins >= 2) { - join_node = curr_node; - break; - } - - bool found_branch = false; - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - std::vector & branch_vec = nodes_per_branch[branch_idx]; - if (belongs_to_branch(curr_node, branch_vec)) { - //continue accumulating - if (std::find(branch_vec.begin(), branch_vec.end(), curr_node) == branch_vec.end()) { - branch_vec.push_back(curr_node); - } - found_branch = true; - } - } - - if (!found_branch && is_noop(curr_node)) { - // we can put it in any branch because it will be ignored - nodes_per_branch[0].push_back({ curr_node }); - } - } - - if (join_node) { - //Create ggml_cuda_concurrent_event - ggml_cuda_concurrent_event concurrent_event(nodes_per_branch.size()); - concurrent_event.join_node = join_node; - - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - for (const ggml_tensor * n : nodes_per_branch[branch_idx]) { - concurrent_event.stream_mapping[n] = branch_idx + 1; - } - } - - int fork_node_idx = node_indices[root_node]; - int join_node_idx = node_indices[join_node]; - - int current_branch_idx = 0; - int current_node_idx = fork_node_idx + 1; - const int n_branches = nodes_per_branch.size(); - - int total_branch_nodes = 0; - for (std::vector branch_nodes : nodes_per_branch) { - total_branch_nodes += branch_nodes.size(); - } - - // there are other nodes in the middle which are unaccounted for - // usually (cpy) nodes, then ignore this fork - if (join_node_idx - fork_node_idx - 1 != total_branch_nodes) { - GGML_LOG_DEBUG( - "Skipping %s because the number of nodes in the middle is not equal to the total number of " - "branch nodes %d != %d\n", - root_node->name, join_node_idx - fork_node_idx - 1, total_branch_nodes); - continue; - } - - // Save the original order of nodes in this region before interleaving - // This is used later to restore grouping for fusion within streams - concurrent_event.original_order.reserve(total_branch_nodes); - for (int i = fork_node_idx + 1; i < join_node_idx; ++i) { - concurrent_event.original_order.push_back(cgraph->nodes[i]); - } - - std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; - GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); - concurrent_events.emplace(root_node, std::move(concurrent_event)); - GGML_LOG_DEBUG("Adding stream at node %s %p\n", root_node->name, root_node); - concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx); - - // interleave tensors to extend lifetimes so that ggml graph doesn't recycle them - // example transformation: - // [attn-norm, QMul, QNorm, QRope, KMul, KNorm, KRope, VMul, attn] -> - // [attn-norm, QMul, KMul, VMul, QNorm, VNorm, QRope, KRope, attn] - while (current_node_idx < join_node_idx) { - std::vector & branch_nodes = nodes_per_branch[current_branch_idx]; - - bool has_node = false; - for (std::vector branch_node : nodes_per_branch) { - has_node |= branch_node.size() > 0; - } - - GGML_ASSERT(has_node); - - if (branch_nodes.empty()) { - current_branch_idx = (current_branch_idx + 1) % n_branches; - continue; - } - - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); - - // append all empty nodes - while (!branch_nodes.empty() && is_noop(branch_nodes.front())) { - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); - } - - current_branch_idx = (current_branch_idx + 1) % n_branches; - } - } - } - } -} - -static const ggml_backend_i ggml_backend_cuda_interface = { - /* .get_name = */ ggml_backend_cuda_get_name, - /* .free = */ ggml_backend_cuda_free, - /* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async, - /* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async, - /* .set_tensor_2d_async = */ ggml_backend_cuda_set_tensor_2d_async, - /* .get_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, - /* .graph_plan_free = */ NULL, - /* .graph_plan_update = */ NULL, - /* .graph_plan_compute = */ NULL, - /* .graph_compute = */ ggml_backend_cuda_graph_compute, - /* .event_record = */ ggml_backend_cuda_event_record, - /* .event_wait = */ ggml_backend_cuda_event_wait, - /* .graph_optimize = */ ggml_backend_cuda_graph_optimize, -}; - -static ggml_guid_t ggml_backend_cuda_guid() { - static ggml_guid guid = { 0x2c, 0xdd, 0xe8, 0x1c, 0x65, 0xb3, 0x65, 0x73, 0x6a, 0x12, 0x88, 0x61, 0x1c, 0xc9, 0xdc, 0x25 }; - return &guid; -} - -bool ggml_backend_is_cuda(ggml_backend_t backend) { - return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); -} - void ggml_backend_cuda_trim_pools(ggml_backend_t backend) { if (!ggml_backend_is_cuda(backend)) { return; @@ -5067,558 +5074,558 @@ void ggml_backend_cuda_clear_graph(ggml_backend_t backend, const ggml_cgraph * g #endif } -int ggml_backend_cuda_get_device_count() { - return ggml_cuda_info().device_count; -} - -void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) { - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); - snprintf(description, description_size, "%s", prop.name); -} - -void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) { - ggml_cuda_set_device(device); - - CUDA_CHECK(cudaMemGetInfo(free, total)); -} - -bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) { - if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { - return false; - } - -#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) || defined(GGML_USE_HIP) - cudaError_t err = cudaHostRegister(buffer, size, cudaHostRegisterPortable | cudaHostRegisterReadOnly); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - - GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__, - size / 1024.0 / 1024.0, cudaGetErrorString(err)); - return false; - } - return true; -#else - GGML_UNUSED(buffer); - GGML_UNUSED(size); - return false; -#endif // CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) -} - -void ggml_backend_cuda_unregister_host_buffer(void * buffer) { - if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { - return; - } - - cudaError_t err = cudaHostUnregister(buffer); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - } -} - - -// backend device - -struct ggml_backend_cuda_device_context { - int device; - std::string name; - std::string description; - std::string pci_bus_id; - int op_offload_min_batch_size; -}; - -static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ctx->name.c_str(); -} - -static const char * ggml_backend_cuda_device_get_description(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ctx->description.c_str(); -} - -#if defined(__linux__) -// Helper function to get available memory from /proc/meminfo for UMA systems -static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_kb, long * free_swap_kb) { - FILE * meminfo_file = nullptr; - // 2KB buffer for reading /proc/meminfo since it does not report size info, should be enough - const size_t BUFFER_SIZE = 2048; - auto file_buffer = std::make_unique(BUFFER_SIZE); - size_t bytes_read = 0; - long huge_tlb_total_pages = -1; - long huge_tlb_free_pages = -1; - long huge_tlb_page_size = -1; - - if (available_memory_kb == nullptr || free_swap_kb == nullptr) { - return false; - } - - meminfo_file = fopen("/proc/meminfo", "r"); - if (meminfo_file == nullptr) { - GGML_LOG_ERROR("%s: failed to open /proc/meminfo\n", __func__); - return false; - } - - // Read file into buffer - bytes_read = fread(file_buffer.get(), 1, BUFFER_SIZE - 1, meminfo_file); - fclose(meminfo_file); - - if (bytes_read == 0) { - GGML_LOG_ERROR("%s: failed to read from /proc/meminfo\n", __func__); - return false; - } - file_buffer[bytes_read] = '\0'; - - *available_memory_kb = -1; - *free_swap_kb = -1; - - // Parse the file buffer line by line - char * line = file_buffer.get(); - char * line_next; - while (line < file_buffer.get() + bytes_read) { - // Find the end of the current line - line_next = strchr(line, '\n'); - if (line_next != nullptr) { - *line_next = '\0'; - line_next++; - } else { - line_next = file_buffer.get() + bytes_read; - } - - long value; - if (sscanf(line, "MemAvailable: %ld kB", &value) == 1) { - *available_memory_kb = value; - } else if (sscanf(line, "SwapFree: %ld kB", &value) == 1) { - *free_swap_kb = value; - } else if (sscanf(line, "HugePages_Total: %ld", &value) == 1) { - huge_tlb_total_pages = value; - } else if (sscanf(line, "HugePages_Free: %ld", &value) == 1) { - huge_tlb_free_pages = value; - } else if (sscanf(line, "Hugepagesize: %ld kB", &value) == 1) { - huge_tlb_page_size = value; - } - - line = line_next; - } - - if (huge_tlb_total_pages != 0 && huge_tlb_total_pages != -1) { - *available_memory_kb = huge_tlb_free_pages * huge_tlb_page_size; - - // Hugetlbfs pages are not swappable. - *free_swap_kb = 0; - } - - GGML_LOG_DEBUG("%s: final available_memory_kb: %ld\n", __func__, *available_memory_kb); - return true; -} -#endif // defined(__linux__) - -static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemGetInfo(free, total)); - -// ref: https://github.com/ggml-org/llama.cpp/pull/17368 -#if defined(__linux__) - // Check if this is a UMA (Unified Memory Architecture) system - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device)); - - // Check if UMA is explicitly enabled via environment variable - bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr; - bool is_uma = prop.integrated > 0 || uma_env; - - if (is_uma) { - // For UMA systems (like DGX Spark), use system memory info - long available_memory_kb = 0; - long free_swap_kb = 0; - - if (ggml_backend_cuda_get_available_uma_memory(&available_memory_kb, &free_swap_kb) && available_memory_kb > 0) { - *free = (size_t)available_memory_kb * 1024; - } else { - GGML_LOG_ERROR("%s: /proc/meminfo reading failed, using cudaMemGetInfo\n", __func__); - } - } -#endif // defined(__linux__) - -} - -static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return GGML_BACKEND_DEVICE_TYPE_GPU; -} - -static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - - props->name = ggml_backend_cuda_device_get_name(dev); - props->description = ggml_backend_cuda_device_get_description(dev); - props->type = ggml_backend_cuda_device_get_type(dev); - props->device_id = ctx->pci_bus_id.empty() ? nullptr : ctx->pci_bus_id.c_str(); - ggml_backend_cuda_device_get_memory(dev, &props->memory_free, &props->memory_total); - - bool host_buffer = getenv("GGML_CUDA_NO_PINNED") == nullptr; -#ifdef GGML_CUDA_NO_PEER_COPY - bool events = false; -#else - bool events = true; -#endif - - props->caps = { - /* .async = */ true, - /* .host_buffer = */ host_buffer, - /* .buffer_from_host_ptr = */ false, - /* .events = */ events, - }; -} - -static ggml_backend_t ggml_backend_cuda_device_init_backend(ggml_backend_dev_t dev, const char * params) { - GGML_UNUSED(params); - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ggml_backend_cuda_init(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_buffer_type(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ggml_backend_cuda_buffer_type(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_host_buffer_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return ggml_backend_cuda_host_buffer_type(); -} - -// TODO: move these functions here -static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - - // split buffers can only be used with GGML_OP_MUL_MAT +int ggml_backend_cuda_get_device_count() { + return ggml_cuda_info().device_count; +} + +void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) { + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); + snprintf(description, description_size, "%s", prop.name); +} + +void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) { + ggml_cuda_set_device(device); + + CUDA_CHECK(cudaMemGetInfo(free, total)); +} + +bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) { + if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { + return false; + } + +#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) || defined(GGML_USE_HIP) + cudaError_t err = cudaHostRegister(buffer, size, cudaHostRegisterPortable | cudaHostRegisterReadOnly); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + + GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__, + size / 1024.0 / 1024.0, cudaGetErrorString(err)); + return false; + } + return true; +#else + GGML_UNUSED(buffer); + GGML_UNUSED(size); + return false; +#endif // CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) +} + +void ggml_backend_cuda_unregister_host_buffer(void * buffer) { + if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { + return; + } + + cudaError_t err = cudaHostUnregister(buffer); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + } +} + + +// backend device + +struct ggml_backend_cuda_device_context { + int device; + std::string name; + std::string description; + std::string pci_bus_id; + int op_offload_min_batch_size; +}; + +static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ctx->name.c_str(); +} + +static const char * ggml_backend_cuda_device_get_description(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ctx->description.c_str(); +} + +#if defined(__linux__) +// Helper function to get available memory from /proc/meminfo for UMA systems +static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_kb, long * free_swap_kb) { + FILE * meminfo_file = nullptr; + // 2KB buffer for reading /proc/meminfo since it does not report size info, should be enough + const size_t BUFFER_SIZE = 2048; + auto file_buffer = std::make_unique(BUFFER_SIZE); + size_t bytes_read = 0; + long huge_tlb_total_pages = -1; + long huge_tlb_free_pages = -1; + long huge_tlb_page_size = -1; + + if (available_memory_kb == nullptr || free_swap_kb == nullptr) { + return false; + } + + meminfo_file = fopen("/proc/meminfo", "r"); + if (meminfo_file == nullptr) { + GGML_LOG_ERROR("%s: failed to open /proc/meminfo\n", __func__); + return false; + } + + // Read file into buffer + bytes_read = fread(file_buffer.get(), 1, BUFFER_SIZE - 1, meminfo_file); + fclose(meminfo_file); + + if (bytes_read == 0) { + GGML_LOG_ERROR("%s: failed to read from /proc/meminfo\n", __func__); + return false; + } + file_buffer[bytes_read] = '\0'; + + *available_memory_kb = -1; + *free_swap_kb = -1; + + // Parse the file buffer line by line + char * line = file_buffer.get(); + char * line_next; + while (line < file_buffer.get() + bytes_read) { + // Find the end of the current line + line_next = strchr(line, '\n'); + if (line_next != nullptr) { + *line_next = '\0'; + line_next++; + } else { + line_next = file_buffer.get() + bytes_read; + } + + long value; + if (sscanf(line, "MemAvailable: %ld kB", &value) == 1) { + *available_memory_kb = value; + } else if (sscanf(line, "SwapFree: %ld kB", &value) == 1) { + *free_swap_kb = value; + } else if (sscanf(line, "HugePages_Total: %ld", &value) == 1) { + huge_tlb_total_pages = value; + } else if (sscanf(line, "HugePages_Free: %ld", &value) == 1) { + huge_tlb_free_pages = value; + } else if (sscanf(line, "Hugepagesize: %ld kB", &value) == 1) { + huge_tlb_page_size = value; + } + + line = line_next; + } + + if (huge_tlb_total_pages != 0 && huge_tlb_total_pages != -1) { + *available_memory_kb = huge_tlb_free_pages * huge_tlb_page_size; + + // Hugetlbfs pages are not swappable. + *free_swap_kb = 0; + } + + GGML_LOG_DEBUG("%s: final available_memory_kb: %ld\n", __func__, *available_memory_kb); + return true; +} +#endif // defined(__linux__) + +static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemGetInfo(free, total)); + +// ref: https://github.com/ggml-org/llama.cpp/pull/17368 +#if defined(__linux__) + // Check if this is a UMA (Unified Memory Architecture) system + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device)); + + // Check if UMA is explicitly enabled via environment variable + bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr; + bool is_uma = prop.integrated > 0 || uma_env; + + if (is_uma) { + // For UMA systems (like DGX Spark), use system memory info + long available_memory_kb = 0; + long free_swap_kb = 0; + + if (ggml_backend_cuda_get_available_uma_memory(&available_memory_kb, &free_swap_kb) && available_memory_kb > 0) { + *free = (size_t)available_memory_kb * 1024; + } else { + GGML_LOG_ERROR("%s: /proc/meminfo reading failed, using cudaMemGetInfo\n", __func__); + } + } +#endif // defined(__linux__) + +} + +static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { + GGML_UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + + props->name = ggml_backend_cuda_device_get_name(dev); + props->description = ggml_backend_cuda_device_get_description(dev); + props->type = ggml_backend_cuda_device_get_type(dev); + props->device_id = ctx->pci_bus_id.empty() ? nullptr : ctx->pci_bus_id.c_str(); + ggml_backend_cuda_device_get_memory(dev, &props->memory_free, &props->memory_total); + + bool host_buffer = getenv("GGML_CUDA_NO_PINNED") == nullptr; +#ifdef GGML_CUDA_NO_PEER_COPY + bool events = false; +#else + bool events = true; +#endif + + props->caps = { + /* .async = */ true, + /* .host_buffer = */ host_buffer, + /* .buffer_from_host_ptr = */ false, + /* .events = */ events, + }; +} + +static ggml_backend_t ggml_backend_cuda_device_init_backend(ggml_backend_dev_t dev, const char * params) { + GGML_UNUSED(params); + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ggml_backend_cuda_init(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ggml_backend_cuda_buffer_type(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_host_buffer_type(ggml_backend_dev_t dev) { + GGML_UNUSED(dev); + return ggml_backend_cuda_host_buffer_type(); +} + +// TODO: move these functions here +static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + + // split buffers can only be used with GGML_OP_MUL_MAT if (op->op != GGML_OP_MUL_MAT && op->op != GGML_OP_MUL_MAT_PACK4) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda_split(op->src[i]->buffer->buft)) { - return false; - } - } - } - - // check if all the sources are allocated on this device - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda(op->src[i]->buffer->buft)) { - ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)op->src[i]->buffer->buft->context; - if (buft_ctx->device != dev_ctx->device) { - return false; - } - } - } - - switch (op->op) { - case GGML_OP_UNARY: - switch (ggml_get_unary_op(op)) { - case GGML_UNARY_OP_ABS: - case GGML_UNARY_OP_SGN: - case GGML_UNARY_OP_NEG: - case GGML_UNARY_OP_STEP: - case GGML_UNARY_OP_GELU: - case GGML_UNARY_OP_SILU: - case GGML_UNARY_OP_RELU: - case GGML_UNARY_OP_SIGMOID: - case GGML_UNARY_OP_HARDSIGMOID: - case GGML_UNARY_OP_HARDSWISH: - case GGML_UNARY_OP_GELU_ERF: - case GGML_UNARY_OP_GELU_QUICK: - case GGML_UNARY_OP_TANH: - case GGML_UNARY_OP_EXP: - case GGML_UNARY_OP_EXPM1: - case GGML_UNARY_OP_SOFTPLUS: - case GGML_UNARY_OP_ELU: - case GGML_UNARY_OP_XIELU: - case GGML_UNARY_OP_FLOOR: - case GGML_UNARY_OP_CEIL: - case GGML_UNARY_OP_ROUND: - case GGML_UNARY_OP_TRUNC: - // TODO: should become: - //return ggml_is_contiguous_rows(op->src[0]); - return ggml_is_contiguous(op->src[0]); - default: - return false; - } - break; - case GGML_OP_GLU: - switch (ggml_get_glu_op(op)) { - case GGML_GLU_OP_REGLU: - case GGML_GLU_OP_GEGLU: - case GGML_GLU_OP_SWIGLU: - case GGML_GLU_OP_SWIGLU_OAI: - case GGML_GLU_OP_GEGLU_ERF: - case GGML_GLU_OP_GEGLU_QUICK: - return ggml_is_contiguous_1(op->src[0]); - default: - return false; - } - break; - case GGML_OP_MUL_MAT: + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda_split(op->src[i]->buffer->buft)) { + return false; + } + } + } + + // check if all the sources are allocated on this device + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda(op->src[i]->buffer->buft)) { + ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)op->src[i]->buffer->buft->context; + if (buft_ctx->device != dev_ctx->device) { + return false; + } + } + } + + switch (op->op) { + case GGML_OP_UNARY: + switch (ggml_get_unary_op(op)) { + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: + case GGML_UNARY_OP_NEG: + case GGML_UNARY_OP_STEP: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_SIGMOID: + case GGML_UNARY_OP_HARDSIGMOID: + case GGML_UNARY_OP_HARDSWISH: + case GGML_UNARY_OP_GELU_ERF: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_EXP: + case GGML_UNARY_OP_EXPM1: + case GGML_UNARY_OP_SOFTPLUS: + case GGML_UNARY_OP_ELU: + case GGML_UNARY_OP_XIELU: + case GGML_UNARY_OP_FLOOR: + case GGML_UNARY_OP_CEIL: + case GGML_UNARY_OP_ROUND: + case GGML_UNARY_OP_TRUNC: + // TODO: should become: + //return ggml_is_contiguous_rows(op->src[0]); + return ggml_is_contiguous(op->src[0]); + default: + return false; + } + break; + case GGML_OP_GLU: + switch (ggml_get_glu_op(op)) { + case GGML_GLU_OP_REGLU: + case GGML_GLU_OP_GEGLU: + case GGML_GLU_OP_SWIGLU: + case GGML_GLU_OP_SWIGLU_OAI: + case GGML_GLU_OP_GEGLU_ERF: + case GGML_GLU_OP_GEGLU_QUICK: + return ggml_is_contiguous_1(op->src[0]); + default: + return false; + } + break; + case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_PACK4: - case GGML_OP_MUL_MAT_ID: - { - struct ggml_tensor * a = op->src[0]; - struct ggml_tensor * b = op->src[1]; - if (a->buffer && ggml_backend_buft_is_cuda_split(a->buffer->buft)) { - if (a->ne[2] > 1 || a->ne[3] > 1) { - return false; - } - // for small weight matrices the active device can end up without any rows, don't use row split in those cases - // this avoids some edge cases (and the performance would not be good anyways) - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) a->buffer->buft->context; - int64_t row_low; - int64_t row_high; - get_row_split(&row_low, &row_high, a, buft_ctx->tensor_split, dev_ctx->device); - if (row_low == row_high) { - return false; - } - } - if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { - return false; - } -#ifdef GGML_USE_MUSA - const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; - if (b->ne[2]*b->ne[3] > 1 && !ggml_is_transposed(a) && !ggml_is_transposed(b)) { - if (GGML_CUDA_CC_IS_QY1(cc) && op->op == GGML_OP_MUL_MAT && - a->type == GGML_TYPE_F16 && b->type == GGML_TYPE_F16) { - return false; - } - if (GGML_CUDA_CC_IS_QY2(cc) && op->op == GGML_OP_MUL_MAT_ID && - a->type == GGML_TYPE_Q2_K && b->type == GGML_TYPE_F32) { - return false; - } - } -#endif // GGML_USE_MUSA - switch (a->type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_Q8_K: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_BF16: - return true; - default: - return false; - } - } break; - case GGML_OP_OUT_PROD: - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; - case GGML_OP_GET_ROWS: - { - switch (op->src[0]->type) { - case GGML_TYPE_F16: - case GGML_TYPE_F32: - case GGML_TYPE_BF16: - case GGML_TYPE_I32: - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - return true; - default: - return false; - } - } break; - case GGML_OP_GET_ROWS_BACK: - { - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; - } break; - case GGML_OP_SET_ROWS: - { - return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || - op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || - op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && - op->src[0]->type == GGML_TYPE_F32 && - (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); - } break; - case GGML_OP_SET: - { - const ggml_type t = op->type; - return (t == GGML_TYPE_F32 || t == GGML_TYPE_I32) && - t == op->src[0]->type && - t == op->src[1]->type; - } break; - case GGML_OP_CPY: - { - ggml_type src0_type = op->src[0]->type; - ggml_type src1_type = op->src[1]->type; - if ((src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16 || src0_type == GGML_TYPE_F16) && - (src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_BF16 || src1_type == GGML_TYPE_F16) - ) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_0) { - return true; - } - if (src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_0) { - return true; - } - if (src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_1) { - return true; - } - if (src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_0) { - return true; - } - if (src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_1) { - return true; - } - if (src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_IQ4_NL) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_I32) { - return true; - } - if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) { - return true; - } - if (src0_type == src1_type && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) { - return true; - } - return false; - } break; - case GGML_OP_DUP: - { - ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; - } break; - case GGML_OP_ARGMAX: - case GGML_OP_COUNT_EQUAL: - { - return true; - } break; - case GGML_OP_REPEAT: - { - ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; - } break; - case GGML_OP_REPEAT_BACK: - return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); - case GGML_OP_CONCAT: - { - ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; - } break; - case GGML_OP_CONV_TRANSPOSE_1D: - { - ggml_type src0_type = op->src[0]->type; - ggml_type src1_type = op->src[1]->type; - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { - return true; - } - return false; - } break; - case GGML_OP_SILU_BACK: - return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; - break; - case GGML_OP_NORM: - case GGML_OP_RMS_NORM: - case GGML_OP_L2_NORM: - return true; - case GGML_OP_RMS_NORM_BACK: - return ggml_is_contiguous(op->src[0]); - break; - case GGML_OP_NONE: - case GGML_OP_RESHAPE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - case GGML_OP_TRANSPOSE: - case GGML_OP_ADD_ID: - case GGML_OP_ADD1: - case GGML_OP_SCALE: - case GGML_OP_SQR: - case GGML_OP_SQRT: - case GGML_OP_SIN: - case GGML_OP_COS: - case GGML_OP_CLAMP: - case GGML_OP_LOG: - return true; - case GGML_OP_ADD: - case GGML_OP_SUB: - case GGML_OP_MUL: - case GGML_OP_DIV: - return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && - (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && - (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); - case GGML_OP_SSM_SCAN: { - if (op->src[3]->ne[0] == 1) { - // Mamba2 - // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) - return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; - } else { - // Mamba - // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) - return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; - } - } - case GGML_OP_SSM_CONV: { - // assumes d_inner % threads == 0 - return op->src[0]->ne[1] % 128 == 0; - } - case GGML_OP_CONT: - return true; - case GGML_OP_DIAG_MASK_INF: - return true; - case GGML_OP_SOFT_MAX: - return true; - case GGML_OP_SOFT_MAX_BACK: { - float max_bias = 0.0f; - memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); - return max_bias == 0.0f; - } - case GGML_OP_ROLL: - if(op->src[0]->type == GGML_TYPE_F32) { - return true; - } - return false; - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: { - return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]); - } + case GGML_OP_MUL_MAT_ID: + { + struct ggml_tensor * a = op->src[0]; + struct ggml_tensor * b = op->src[1]; + if (a->buffer && ggml_backend_buft_is_cuda_split(a->buffer->buft)) { + if (a->ne[2] > 1 || a->ne[3] > 1) { + return false; + } + // for small weight matrices the active device can end up without any rows, don't use row split in those cases + // this avoids some edge cases (and the performance would not be good anyways) + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) a->buffer->buft->context; + int64_t row_low; + int64_t row_high; + get_row_split(&row_low, &row_high, a, buft_ctx->tensor_split, dev_ctx->device); + if (row_low == row_high) { + return false; + } + } + if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { + return false; + } +#ifdef GGML_USE_MUSA + const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; + if (b->ne[2]*b->ne[3] > 1 && !ggml_is_transposed(a) && !ggml_is_transposed(b)) { + if (GGML_CUDA_CC_IS_QY1(cc) && op->op == GGML_OP_MUL_MAT && + a->type == GGML_TYPE_F16 && b->type == GGML_TYPE_F16) { + return false; + } + if (GGML_CUDA_CC_IS_QY2(cc) && op->op == GGML_OP_MUL_MAT_ID && + a->type == GGML_TYPE_Q2_K && b->type == GGML_TYPE_F32) { + return false; + } + } +#endif // GGML_USE_MUSA + switch (a->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q8_K: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_BF16: + return true; + default: + return false; + } + } break; + case GGML_OP_OUT_PROD: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; + case GGML_OP_GET_ROWS: + { + switch (op->src[0]->type) { + case GGML_TYPE_F16: + case GGML_TYPE_F32: + case GGML_TYPE_BF16: + case GGML_TYPE_I32: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + return true; + default: + return false; + } + } break; + case GGML_OP_GET_ROWS_BACK: + { + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; + } break; + case GGML_OP_SET_ROWS: + { + return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || + op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || + op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && + op->src[0]->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); + } break; + case GGML_OP_SET: + { + const ggml_type t = op->type; + return (t == GGML_TYPE_F32 || t == GGML_TYPE_I32) && + t == op->src[0]->type && + t == op->src[1]->type; + } break; + case GGML_OP_CPY: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1]->type; + if ((src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16 || src0_type == GGML_TYPE_F16) && + (src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_BF16 || src1_type == GGML_TYPE_F16) + ) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_0) { + return true; + } + if (src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_0) { + return true; + } + if (src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_1) { + return true; + } + if (src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_0) { + return true; + } + if (src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_1) { + return true; + } + if (src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_IQ4_NL) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_I32) { + return true; + } + if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) { + return true; + } + if (src0_type == src1_type && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) { + return true; + } + return false; + } break; + case GGML_OP_DUP: + { + ggml_type src0_type = op->src[0]->type; + return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + } break; + case GGML_OP_ARGMAX: + case GGML_OP_COUNT_EQUAL: + { + return true; + } break; + case GGML_OP_REPEAT: + { + ggml_type src0_type = op->src[0]->type; + return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + } break; + case GGML_OP_REPEAT_BACK: + return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); + case GGML_OP_CONCAT: + { + ggml_type src0_type = op->src[0]->type; + return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1]->type; + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return true; + } + return false; + } break; + case GGML_OP_SILU_BACK: + return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; + break; + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_L2_NORM: + return true; + case GGML_OP_RMS_NORM_BACK: + return ggml_is_contiguous(op->src[0]); + break; + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_LOG: + return true; + case GGML_OP_ADD: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); + case GGML_OP_SSM_SCAN: { + if (op->src[3]->ne[0] == 1) { + // Mamba2 + // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) + return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; + } else { + // Mamba + // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) + return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; + } + } + case GGML_OP_SSM_CONV: { + // assumes d_inner % threads == 0 + return op->src[0]->ne[1] % 128 == 0; + } + case GGML_OP_CONT: + return true; + case GGML_OP_DIAG_MASK_INF: + return true; + case GGML_OP_SOFT_MAX: + return true; + case GGML_OP_SOFT_MAX_BACK: { + float max_bias = 0.0f; + memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); + return max_bias == 0.0f; + } + case GGML_OP_ROLL: + if(op->src[0]->type == GGML_TYPE_F32) { + return true; + } + return false; + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: { + return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]); + } case GGML_OP_IM2COL: case GGML_OP_IM2COL_FAST_1D: case GGML_OP_IM2COL_3D: @@ -5633,40 +5640,40 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16); case GGML_OP_ACC: - // TODO: extend support like so: - //return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]); - return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); - case GGML_OP_SUM: - return ggml_is_contiguous_rows(op->src[0]); - case GGML_OP_TOP_K: - case GGML_OP_ARGSORT: -#ifndef GGML_CUDA_USE_CUB - return op->src[0]->ne[0] <= 1024; -#else - return true; -#endif - case GGML_OP_SUM_ROWS: - case GGML_OP_MEAN: - case GGML_OP_GROUP_NORM: - return ggml_is_contiguous(op->src[0]); - case GGML_OP_PAD: - return true; - case GGML_OP_UPSCALE: - case GGML_OP_PAD_REFLECT_1D: - case GGML_OP_ARANGE: - case GGML_OP_TIMESTEP_EMBEDDING: - case GGML_OP_LEAKY_RELU: - case GGML_OP_RWKV_WKV6: - case GGML_OP_GATED_LINEAR_ATTN: - case GGML_OP_RWKV_WKV7: - return true; - case GGML_OP_GATED_DELTA_NET: - //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 -#ifdef GGML_USE_MUSA - return false; -#else - return true; -#endif // GGML_USE_MUSA + // TODO: extend support like so: + //return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]); + return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + case GGML_OP_SUM: + return ggml_is_contiguous_rows(op->src[0]); + case GGML_OP_TOP_K: + case GGML_OP_ARGSORT: +#ifndef GGML_CUDA_USE_CUB + return op->src[0]->ne[0] <= 1024; +#else + return true; +#endif + case GGML_OP_SUM_ROWS: + case GGML_OP_MEAN: + case GGML_OP_GROUP_NORM: + return ggml_is_contiguous(op->src[0]); + case GGML_OP_PAD: + return true; + case GGML_OP_UPSCALE: + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_ARANGE: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_LEAKY_RELU: + case GGML_OP_RWKV_WKV6: + case GGML_OP_GATED_LINEAR_ATTN: + case GGML_OP_RWKV_WKV7: + return true; + case GGML_OP_GATED_DELTA_NET: + //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 +#ifdef GGML_USE_MUSA + return false; +#else + return true; +#endif // GGML_USE_MUSA case GGML_OP_FLASH_ATTN_EXT: return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); case GGML_OP_SAGE_ATTN2: @@ -5676,289 +5683,289 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_CONVROT_LINEAR: return ggml_cuda_convrot_linear_supported(dev_ctx->device, op); case GGML_OP_CROSS_ENTROPY_LOSS: - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - case GGML_OP_OPT_STEP_ADAMW: - case GGML_OP_OPT_STEP_SGD: - case GGML_OP_FILL: - case GGML_OP_CUMSUM: - case GGML_OP_TRI: - case GGML_OP_DIAG: - case GGML_OP_SOLVE_TRI: - return true; - - default: - return false; - } -} - -static bool ggml_backend_cuda_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - const bool integrated = ggml_cuda_info().devices[dev_ctx->device].integrated; - return (((ggml_backend_buft_is_cuda(buft) || ggml_backend_buft_is_cuda_split(buft)) && buft->device == dev) || (integrated && ggml_backend_buft_is_cuda_host(buft))); -} - -static int64_t get_op_batch_size(const ggml_tensor * op) { - switch (op->op) { - case GGML_OP_GET_ROWS: - return 0; - case GGML_OP_MUL_MAT: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: + case GGML_OP_FILL: + case GGML_OP_CUMSUM: + case GGML_OP_TRI: + case GGML_OP_DIAG: + case GGML_OP_SOLVE_TRI: + return true; + + default: + return false; + } +} + +static bool ggml_backend_cuda_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + const bool integrated = ggml_cuda_info().devices[dev_ctx->device].integrated; + return (((ggml_backend_buft_is_cuda(buft) || ggml_backend_buft_is_cuda_split(buft)) && buft->device == dev) || (integrated && ggml_backend_buft_is_cuda_host(buft))); +} + +static int64_t get_op_batch_size(const ggml_tensor * op) { + switch (op->op) { + case GGML_OP_GET_ROWS: + return 0; + case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_PACK4: - return op->ne[1]; - case GGML_OP_MUL_MAT_ID: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - return op->ne[2]; - default: - return ggml_nrows(op); - } -} - -static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - - return get_op_batch_size(op) >= dev_ctx->op_offload_min_batch_size; -} - -static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) { -#ifdef GGML_CUDA_NO_PEER_COPY - return nullptr; -#else - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context; - - ggml_cuda_set_device(dev_ctx->device); - - cudaEvent_t event; - CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); - - return new ggml_backend_event { - /* .device = */ dev, - /* .context = */ event, - }; -#endif -} - -static void ggml_backend_cuda_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { - GGML_UNUSED(dev); - - CUDA_CHECK(cudaEventDestroy((cudaEvent_t)event->context)); - delete event; -} - -static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { - GGML_UNUSED(dev); - CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); -} - -static const ggml_backend_device_i ggml_backend_cuda_device_interface = { - /* .get_name = */ ggml_backend_cuda_device_get_name, - /* .get_description = */ ggml_backend_cuda_device_get_description, - /* .get_memory = */ ggml_backend_cuda_device_get_memory, - /* .get_type = */ ggml_backend_cuda_device_get_type, - /* .get_props = */ ggml_backend_cuda_device_get_props, - /* .init_backend = */ ggml_backend_cuda_device_init_backend, - /* .get_buffer_type = */ ggml_backend_cuda_device_get_buffer_type, - /* .get_host_buffer_type = */ ggml_backend_cuda_device_get_host_buffer_type, - /* .buffer_from_host_ptr = */ NULL, - /* .supports_op = */ ggml_backend_cuda_device_supports_op, - /* .supports_buft = */ ggml_backend_cuda_device_supports_buft, - /* .offload_op = */ ggml_backend_cuda_device_offload_op, - /* .event_new = */ ggml_backend_cuda_device_event_new, - /* .event_free = */ ggml_backend_cuda_device_event_free, - /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, -}; - -// backend reg - -struct ggml_backend_cuda_reg_context { - std::vector devices; -}; - -static const char * ggml_backend_cuda_reg_get_name(ggml_backend_reg_t reg) { - GGML_UNUSED(reg); - return GGML_CUDA_NAME; -} - -static size_t ggml_backend_cuda_reg_get_device_count(ggml_backend_reg_t reg) { - ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; - return ctx->devices.size(); -} - -static ggml_backend_dev_t ggml_backend_cuda_reg_get_device(ggml_backend_reg_t reg, size_t index) { - ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; - GGML_ASSERT(index < ctx->devices.size()); - return ctx->devices[index]; -} - -static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t reg) { - static std::vector features = []() { - std::vector features; - #define _STRINGIFY(...) #__VA_ARGS__ - #define STRINGIFY(...) _STRINGIFY(__VA_ARGS__) - - #ifdef __CUDA_ARCH_LIST__ - features.push_back({ "ARCHS", STRINGIFY(__CUDA_ARCH_LIST__) }); - #endif - - #ifdef GGML_CUDA_FORCE_MMQ - features.push_back({ "FORCE_MMQ", "1" }); - #endif - - #ifdef GGML_CUDA_FORCE_CUBLAS - features.push_back({ "FORCE_CUBLAS", "1" }); - #endif - - #ifndef GGML_USE_VMM - features.push_back({ "NO_VMM", "1" }); - #endif - - #ifdef GGML_CUDA_NO_PEER_COPY - features.push_back({ "NO_PEER_COPY", "1" }); - #endif - - #ifdef GGML_CUDA_USE_GRAPHS - features.push_back({ "USE_GRAPHS", "1" }); - #endif - - #ifdef GGML_CUDA_PEER_MAX_BATCH_SIZE - features.push_back({ "PEER_MAX_BATCH_SIZE", STRINGIFY(GGML_CUDA_PEER_MAX_BATCH_SIZE) }); - #endif - - #ifdef GGML_CUDA_FA_ALL_QUANTS - features.push_back({ "FA_ALL_QUANTS", "1" }); - #endif - - { - const auto & info = ggml_cuda_info(); - for (int id = 0; id < info.device_count; ++id) { - if (blackwell_mma_available(info.devices[id].cc)) { - features.push_back({ "BLACKWELL_NATIVE_FP4", "1"}); - break; - } - } - } - - #undef _STRINGIFY - #undef STRINGIFY - - features.push_back({ nullptr, nullptr }); - - return features; - }(); - - return features.data(); - - GGML_UNUSED(reg); -} - -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_comm_init") == 0) { - return (void *)ggml_backend_cuda_comm_init; - } - if (strcmp(name, "ggml_backend_comm_free") == 0) { - return (void *)ggml_backend_cuda_comm_free; - } - if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) { - return (void *)ggml_backend_cuda_comm_allreduce_tensor; - } - if (strcmp(name, "ggml_backend_split_buffer_type") == 0) { - return (void *)ggml_backend_cuda_split_buffer_type; - } - if (strcmp(name, "ggml_backend_register_host_buffer") == 0) { - return (void *)ggml_backend_cuda_register_host_buffer; - } - if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) { - return (void *)ggml_backend_cuda_unregister_host_buffer; - } - if (strcmp(name, "ggml_backend_get_features") == 0) { - return (void *)ggml_backend_cuda_get_features; - } - if (strcmp(name, "ggml_backend_cuda_clear_graph") == 0) { - return (void *)ggml_backend_cuda_clear_graph; - } - if (strcmp(name, "ggml_backend_cuda_trim_pools") == 0) { - return (void *)ggml_backend_cuda_trim_pools; - } - return nullptr; -} - -static const ggml_backend_reg_i ggml_backend_cuda_reg_interface = { - /* .get_name = */ ggml_backend_cuda_reg_get_name, - /* .get_device_count = */ ggml_backend_cuda_reg_get_device_count, - /* .get_device = */ ggml_backend_cuda_reg_get_device, - /* .get_proc_address = */ ggml_backend_cuda_reg_get_proc_address, -}; - -// backend registry -ggml_backend_reg_t ggml_backend_cuda_reg() { - static ggml_backend_reg reg; - static bool initialized = false; - - { - static std::mutex mutex; - std::lock_guard lock(mutex); - if (!initialized) { - ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context; - const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; - - for (int i = 0; i < ggml_cuda_info().device_count; i++) { - ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context; - dev_ctx->device = i; - dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); - - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, i)); - dev_ctx->description = prop.name; - - char pci_bus_id[32] = {}; - CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), i)); - dev_ctx->pci_bus_id = pci_bus_id; - for (char & c : dev_ctx->pci_bus_id) { - c = std::tolower(c); - } - dev_ctx->op_offload_min_batch_size = min_batch_size; - - ggml_backend_dev_t dev = new ggml_backend_device { - /* .iface = */ ggml_backend_cuda_device_interface, - /* .reg = */ ®, - /* .context = */ dev_ctx - }; - ctx->devices.push_back(dev); - } - - reg = ggml_backend_reg { - /* .api_version = */ GGML_BACKEND_API_VERSION, - /* .iface = */ ggml_backend_cuda_reg_interface, - /* .context = */ ctx - }; - } - - initialized = true; - } - - return ® -} - -ggml_backend_t ggml_backend_cuda_init(int device) { - if (device < 0 || device >= ggml_backend_cuda_get_device_count()) { - GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); - return nullptr; - } - - ggml_backend_cuda_context * ctx = new ggml_backend_cuda_context(device); - if (ctx == nullptr) { - GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); - return nullptr; - } - - ggml_backend_t cuda_backend = new ggml_backend { - /* .guid = */ ggml_backend_cuda_guid(), - /* .iface = */ ggml_backend_cuda_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device), - /* .context = */ ctx, - }; - - return cuda_backend; -} - -GGML_BACKEND_DL_IMPL(ggml_backend_cuda_reg) + return op->ne[1]; + case GGML_OP_MUL_MAT_ID: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + return op->ne[2]; + default: + return ggml_nrows(op); + } +} + +static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + + return get_op_batch_size(op) >= dev_ctx->op_offload_min_batch_size; +} + +static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) { +#ifdef GGML_CUDA_NO_PEER_COPY + return nullptr; +#else + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context; + + ggml_cuda_set_device(dev_ctx->device); + + cudaEvent_t event; + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ event, + }; +#endif +} + +static void ggml_backend_cuda_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + CUDA_CHECK(cudaEventDestroy((cudaEvent_t)event->context)); + delete event; +} + +static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); +} + +static const ggml_backend_device_i ggml_backend_cuda_device_interface = { + /* .get_name = */ ggml_backend_cuda_device_get_name, + /* .get_description = */ ggml_backend_cuda_device_get_description, + /* .get_memory = */ ggml_backend_cuda_device_get_memory, + /* .get_type = */ ggml_backend_cuda_device_get_type, + /* .get_props = */ ggml_backend_cuda_device_get_props, + /* .init_backend = */ ggml_backend_cuda_device_init_backend, + /* .get_buffer_type = */ ggml_backend_cuda_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_cuda_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_cuda_device_supports_op, + /* .supports_buft = */ ggml_backend_cuda_device_supports_buft, + /* .offload_op = */ ggml_backend_cuda_device_offload_op, + /* .event_new = */ ggml_backend_cuda_device_event_new, + /* .event_free = */ ggml_backend_cuda_device_event_free, + /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, +}; + +// backend reg + +struct ggml_backend_cuda_reg_context { + std::vector devices; +}; + +static const char * ggml_backend_cuda_reg_get_name(ggml_backend_reg_t reg) { + GGML_UNUSED(reg); + return GGML_CUDA_NAME; +} + +static size_t ggml_backend_cuda_reg_get_device_count(ggml_backend_reg_t reg) { + ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; + return ctx->devices.size(); +} + +static ggml_backend_dev_t ggml_backend_cuda_reg_get_device(ggml_backend_reg_t reg, size_t index) { + ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; + GGML_ASSERT(index < ctx->devices.size()); + return ctx->devices[index]; +} + +static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t reg) { + static std::vector features = []() { + std::vector features; + #define _STRINGIFY(...) #__VA_ARGS__ + #define STRINGIFY(...) _STRINGIFY(__VA_ARGS__) + + #ifdef __CUDA_ARCH_LIST__ + features.push_back({ "ARCHS", STRINGIFY(__CUDA_ARCH_LIST__) }); + #endif + + #ifdef GGML_CUDA_FORCE_MMQ + features.push_back({ "FORCE_MMQ", "1" }); + #endif + + #ifdef GGML_CUDA_FORCE_CUBLAS + features.push_back({ "FORCE_CUBLAS", "1" }); + #endif + + #ifndef GGML_USE_VMM + features.push_back({ "NO_VMM", "1" }); + #endif + + #ifdef GGML_CUDA_NO_PEER_COPY + features.push_back({ "NO_PEER_COPY", "1" }); + #endif + + #ifdef GGML_CUDA_USE_GRAPHS + features.push_back({ "USE_GRAPHS", "1" }); + #endif + + #ifdef GGML_CUDA_PEER_MAX_BATCH_SIZE + features.push_back({ "PEER_MAX_BATCH_SIZE", STRINGIFY(GGML_CUDA_PEER_MAX_BATCH_SIZE) }); + #endif + + #ifdef GGML_CUDA_FA_ALL_QUANTS + features.push_back({ "FA_ALL_QUANTS", "1" }); + #endif + + { + const auto & info = ggml_cuda_info(); + for (int id = 0; id < info.device_count; ++id) { + if (blackwell_mma_available(info.devices[id].cc)) { + features.push_back({ "BLACKWELL_NATIVE_FP4", "1"}); + break; + } + } + } + + #undef _STRINGIFY + #undef STRINGIFY + + features.push_back({ nullptr, nullptr }); + + return features; + }(); + + return features.data(); + + GGML_UNUSED(reg); +} + +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_comm_init") == 0) { + return (void *)ggml_backend_cuda_comm_init; + } + if (strcmp(name, "ggml_backend_comm_free") == 0) { + return (void *)ggml_backend_cuda_comm_free; + } + if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) { + return (void *)ggml_backend_cuda_comm_allreduce_tensor; + } + if (strcmp(name, "ggml_backend_split_buffer_type") == 0) { + return (void *)ggml_backend_cuda_split_buffer_type; + } + if (strcmp(name, "ggml_backend_register_host_buffer") == 0) { + return (void *)ggml_backend_cuda_register_host_buffer; + } + if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) { + return (void *)ggml_backend_cuda_unregister_host_buffer; + } + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cuda_get_features; + } + if (strcmp(name, "ggml_backend_cuda_clear_graph") == 0) { + return (void *)ggml_backend_cuda_clear_graph; + } + if (strcmp(name, "ggml_backend_cuda_trim_pools") == 0) { + return (void *)ggml_backend_cuda_trim_pools; + } + return nullptr; +} + +static const ggml_backend_reg_i ggml_backend_cuda_reg_interface = { + /* .get_name = */ ggml_backend_cuda_reg_get_name, + /* .get_device_count = */ ggml_backend_cuda_reg_get_device_count, + /* .get_device = */ ggml_backend_cuda_reg_get_device, + /* .get_proc_address = */ ggml_backend_cuda_reg_get_proc_address, +}; + +// backend registry +ggml_backend_reg_t ggml_backend_cuda_reg() { + static ggml_backend_reg reg; + static bool initialized = false; + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context; + const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; + + for (int i = 0; i < ggml_cuda_info().device_count; i++) { + ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context; + dev_ctx->device = i; + dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); + + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, i)); + dev_ctx->description = prop.name; + + char pci_bus_id[32] = {}; + CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), i)); + dev_ctx->pci_bus_id = pci_bus_id; + for (char & c : dev_ctx->pci_bus_id) { + c = std::tolower(c); + } + dev_ctx->op_offload_min_batch_size = min_batch_size; + + ggml_backend_dev_t dev = new ggml_backend_device { + /* .iface = */ ggml_backend_cuda_device_interface, + /* .reg = */ ®, + /* .context = */ dev_ctx + }; + ctx->devices.push_back(dev); + } + + reg = ggml_backend_reg { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_cuda_reg_interface, + /* .context = */ ctx + }; + } + + initialized = true; + } + + return ® +} + +ggml_backend_t ggml_backend_cuda_init(int device) { + if (device < 0 || device >= ggml_backend_cuda_get_device_count()) { + GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); + return nullptr; + } + + ggml_backend_cuda_context * ctx = new ggml_backend_cuda_context(device); + if (ctx == nullptr) { + GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); + return nullptr; + } + + ggml_backend_t cuda_backend = new ggml_backend { + /* .guid = */ ggml_backend_cuda_guid(), + /* .iface = */ ggml_backend_cuda_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device), + /* .context = */ ctx, + }; + + return cuda_backend; +} + +GGML_BACKEND_DL_IMPL(ggml_backend_cuda_reg) diff --git a/include/engine/community_models/minimax_music3/ar_runtime.h b/include/engine/community_models/minimax_music3/ar_runtime.h index f6b1fc7ac..fea14fbca 100644 --- a/include/engine/community_models/minimax_music3/ar_runtime.h +++ b/include/engine/community_models/minimax_music3/ar_runtime.h @@ -8,6 +8,7 @@ #include "engine/framework/sampling/torch_random.h" #include +#include #include #include @@ -20,7 +21,8 @@ class MiniMaxMusic3ArRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3ArRuntime(); std::vector generate_frame_hiddens( @@ -28,6 +30,29 @@ class MiniMaxMusic3ArRuntime { int64_t target_frames, uint64_t & rng_offset_blocks); + // Identical generation, but appends into a caller-owned buffer (which must + // survive without reallocation: callers reserve target capacity up front) + // and reports completed hidden rows so a consumer thread can start work on + // finished frames while later frames are still decoding. + void generate_frame_hiddens_into( + const MiniMaxMusic3Request & request, + int64_t target_frames, + uint64_t & rng_offset_blocks, + std::vector & frame_hiddens, + const std::function * progress); + + // Ensemble decode: K independent takes of the same prompt advance in one + // batched pass (global LM and depth run at batch 2K), each take sampling + // with its own seed. Weight reads are amortized across takes, which is + // where the bandwidth-bound AR stage spends its time. Returns per-take + // frame hiddens; rng_offset_blocks[i] carries take i's counter onward. + std::vector> generate_frame_hiddens_ensemble( + const MiniMaxMusic3Request & request, + int64_t target_frames, + const std::vector & take_seeds, + std::vector & rng_offset_blocks, + int64_t prefix_frames = 0); + void release_runtime_graphs(); private: diff --git a/include/engine/community_models/minimax_music3/condition_encoder.h b/include/engine/community_models/minimax_music3/condition_encoder.h index a6f2a9b1b..7a540e279 100644 --- a/include/engine/community_models/minimax_music3/condition_encoder.h +++ b/include/engine/community_models/minimax_music3/condition_encoder.h @@ -5,11 +5,25 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/modules/conv_modules.h" +#include +#include #include #include namespace engine::models::minimax_music3 { +namespace detail { + +std::vector project_frame_hiddens( + const float * frame_hiddens, + size_t value_count, + int64_t frames, + int64_t layers, + int64_t hidden_size, + const std::vector & layer_weights); + +} // namespace detail + struct MiniMaxMusic3ConditionWeights { std::shared_ptr store; std::vector layer_weights; @@ -23,9 +37,15 @@ class MiniMaxMusic3ConditionEncoderRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3ConditionEncoderRuntime(); + std::vector encode( + const float * frame_hiddens, + size_t value_count, + int64_t frames, + int64_t & condition_frames); std::vector encode(const std::vector & frame_hiddens, int64_t frames, int64_t & condition_frames); void release_runtime_graphs(); diff --git a/include/engine/community_models/minimax_music3/depth_decoder.h b/include/engine/community_models/minimax_music3/depth_decoder.h index 5e25c4f15..8e33b7a81 100644 --- a/include/engine/community_models/minimax_music3/depth_decoder.h +++ b/include/engine/community_models/minimax_music3/depth_decoder.h @@ -34,7 +34,8 @@ class MiniMaxMusic3DepthDecoderRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3DepthDecoderRuntime(); MiniMaxMusic3DepthCodes generate( @@ -47,6 +48,20 @@ class MiniMaxMusic3DepthDecoderRuntime { uint64_t & sample_call_index, uint64_t & rng_offset_blocks); + // Ensemble variant: decodes the depth chain for `songs` independent takes + // in one batched pass (rows are [cond_0, uncond_0, cond_1, ...]). Each + // take samples with its own seed/counters, so results are identical to + // running the single-take path per song. + std::vector generate_batch( + const std::vector & interleaved_hiddens, + int64_t songs, + const std::vector & semantic_codes, + float guidance_scale, + int64_t top_k, + const std::vector & seeds, + std::vector & sample_call_indices, + std::vector & rng_offset_blocks); + std::vector feedback_embedding(const std::vector & codes) const; void release_runtime_graphs(); diff --git a/include/engine/community_models/minimax_music3/flow_sampler.h b/include/engine/community_models/minimax_music3/flow_sampler.h index 4be88d150..99b63f29b 100644 --- a/include/engine/community_models/minimax_music3/flow_sampler.h +++ b/include/engine/community_models/minimax_music3/flow_sampler.h @@ -15,7 +15,8 @@ class MiniMaxMusic3FlowSamplerRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3FlowSamplerRuntime(); std::vector denoise_chunk( diff --git a/include/engine/community_models/minimax_music3/flow_transformer.h b/include/engine/community_models/minimax_music3/flow_transformer.h index d02f34eaa..8ca3fdcf2 100644 --- a/include/engine/community_models/minimax_music3/flow_transformer.h +++ b/include/engine/community_models/minimax_music3/flow_transformer.h @@ -42,7 +42,8 @@ class MiniMaxMusic3FlowTransformerRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3FlowTransformerRuntime(); std::vector predict_velocity_branches( @@ -50,6 +51,13 @@ class MiniMaxMusic3FlowTransformerRuntime { const std::vector & condition, int64_t latent_frames, float timestep); + // Evaluates only the conditional branch (batch 1) for guidance-delta + // reuse steps; returns one branch worth of velocity values. + std::vector predict_velocity_cond( + const std::vector & latents, + const std::vector & condition, + int64_t latent_frames, + float timestep); void prepare_chunk_condition( const std::vector & condition, int64_t latent_frames); diff --git a/include/engine/community_models/minimax_music3/global_lm.h b/include/engine/community_models/minimax_music3/global_lm.h index 86cfc1b78..6b4b2e94d 100644 --- a/include/engine/community_models/minimax_music3/global_lm.h +++ b/include/engine/community_models/minimax_music3/global_lm.h @@ -12,10 +12,25 @@ namespace engine::models::minimax_music3 { +enum class MiniMaxMusic3LmHeadLayout { + FullVocab, + SemanticCompactV1, +}; + +MiniMaxMusic3LmHeadLayout classify_minimax_music3_lm_head_shape( + const std::vector & shape, + int64_t vocab_size, + int64_t hidden_size); + +int64_t minimax_music3_lm_head_output_size( + MiniMaxMusic3LmHeadLayout layout, + int64_t vocab_size) noexcept; + struct MiniMaxMusic3GlobalLMWeights { std::shared_ptr store; core::TensorValue token_embedding; modules::QwenCausalDecodeRuntimeWeights qwen; + MiniMaxMusic3LmHeadLayout lm_head_layout = MiniMaxMusic3LmHeadLayout::FullVocab; }; MiniMaxMusic3GlobalLMWeights load_minimax_music3_global_lm_weights( @@ -26,6 +41,7 @@ MiniMaxMusic3GlobalLMWeights load_minimax_music3_global_lm_weights( modules::QwenCausalDecodeRuntimeConfig make_minimax_music3_global_lm_runtime_config( const MiniMaxMusic3Config & config, + MiniMaxMusic3LmHeadLayout lm_head_layout, core::BackendType backend_type, size_t prefill_graph_arena_bytes, size_t decode_graph_arena_bytes); diff --git a/include/engine/community_models/minimax_music3/pipeline.h b/include/engine/community_models/minimax_music3/pipeline.h index cd6504dce..755d5684f 100644 --- a/include/engine/community_models/minimax_music3/pipeline.h +++ b/include/engine/community_models/minimax_music3/pipeline.h @@ -8,9 +8,20 @@ #include "engine/framework/runtime/model.h" #include +#include namespace engine::models::minimax_music3 { +namespace detail { + +void append_cropped_interleaved_audio( + runtime::AudioBuffer & destination, + const runtime::AudioBuffer & chunk, + int64_t left_frames, + int64_t right_frames); + +} // namespace detail + class MiniMaxMusic3PipelineRuntime { public: MiniMaxMusic3PipelineRuntime( @@ -19,10 +30,16 @@ class MiniMaxMusic3PipelineRuntime { size_t graph_arena_bytes, size_t weight_context_bytes, assets::TensorStorageType storage_type, - bool memory_saver); + bool memory_saver, + bool pipeline_overlap = false); ~MiniMaxMusic3PipelineRuntime(); runtime::AudioBuffer generate(const MiniMaxMusic3Request & request); + // K independent takes of the same prompt sharing one batched AR pass; + // take i uses seed take_seeds[i] end to end (AR, depth, flow noise). + std::vector generate_ensemble( + const MiniMaxMusic3Request & request, + const std::vector & take_seeds); void release_runtime_graphs(); private: diff --git a/include/engine/community_models/minimax_music3/types.h b/include/engine/community_models/minimax_music3/types.h index a90f7410e..5aab46f69 100644 --- a/include/engine/community_models/minimax_music3/types.h +++ b/include/engine/community_models/minimax_music3/types.h @@ -85,6 +85,33 @@ struct MiniMaxMusic3Request { float ar_guidance_scale = 1.5F; int64_t top_k = 50; uint64_t seed = 0; + // Flow CFG guidance-delta reuse: when interval > 1 the unconditional + // branch is evaluated only on warmup steps, every interval-th step and + // the final step; other steps reuse the cached (cond - uncond) delta. + // The default of 2 is the listening-accepted recipe (mel-L1 ~0.4 dB to + // the exact reference at -15..-21% wall); set 1 for the exact-reference + // trajectory. + int64_t flow_uncond_interval = 1; + int64_t flow_uncond_warmup = 2; + // Number of independent takes decoded together in one batched AR pass + // (per-take seeds seed, seed+1, ...). The global LM and depth decoder are + // bandwidth-bound, so K takes cost far less than K runs; flow/vocoder run + // per take. 1 keeps the plain single-song path. + int64_t ensemble_takes = 1; + // Intro-lock fork: the first N frames are decoded once at batch 2 (one + // master trajectory shared by every take), then the batched decode KV is + // replicated to 2K rows and the takes diverge with their own seeds. Take + // 0 continues the master trajectory exactly. 0 disables the fork. + int64_t ensemble_prefix_frames = 0; + // Flow chunk hop in AR frames (0 = the model config's 100). A larger hop + // means fewer chunks and less double-denoising: hop 150 cuts flow ~-28%. + // Crops and the carry window are rederived from the hop so the seams stay + // consistent: kept = latents(hop), overlap = (chunk - kept)/2, left crop + // = overlap/2 (hop 100 reproduces the historical 86/258/172 exactly). + int64_t flow_chunk_hop_frames = 0; + // Derived internally by the pipeline from the hop; 0 keeps the config's + // overlap_latent_length. Not a user knob. + int64_t flow_overlap_latent_length = 0; }; } // namespace engine::models::minimax_music3 diff --git a/include/engine/community_models/minimax_music3/vocoder.h b/include/engine/community_models/minimax_music3/vocoder.h index 562a96d0f..e2f396442 100644 --- a/include/engine/community_models/minimax_music3/vocoder.h +++ b/include/engine/community_models/minimax_music3/vocoder.h @@ -16,7 +16,8 @@ class MiniMaxMusic3VocoderRuntime { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type); + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release = false); ~MiniMaxMusic3VocoderRuntime(); runtime::AudioBuffer decode(const std::vector & latents, int64_t latent_frames); @@ -28,4 +29,3 @@ class MiniMaxMusic3VocoderRuntime { }; } // namespace engine::models::minimax_music3 - diff --git a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h index a02848fa9..15063d25a 100644 --- a/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h +++ b/include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h @@ -28,6 +28,7 @@ struct QwenCausalDecodeRuntimeConfig { std::optional readback_round_type; std::vector logits_readback_token_ids; int64_t sliding_window = 0; + bool evict_cuda_graph_cache_on_release = false; }; struct QwenCausalDecodeRuntimeWeights { @@ -93,6 +94,11 @@ class QwenCausalDecodeRuntime { const std::vector & embeddings, int64_t batch_size); + // Snapshot of the batched decode KV cache (host vectors), suitable for + // replication and re-import via start_decode_*_batched with a different + // batch size — the runtime rebuilds its decode graphs for the new batch. + runtime::TransformerBatchedKVState export_batched_decode_state() const; + int64_t decode_cache_steps() const noexcept; int64_t decode_current_end() const noexcept; int64_t decode_valid_steps() const noexcept; diff --git a/include/engine/framework/sampling/torch_random.h b/include/engine/framework/sampling/torch_random.h index f4fc808cf..10365cbbe 100644 --- a/include/engine/framework/sampling/torch_random.h +++ b/include/engine/framework/sampling/torch_random.h @@ -73,6 +73,65 @@ void fill_torch_cuda_uniform( std::vector generate_torch_cuda_uniform(size_t count, uint64_t seed, uint64_t start_index = 0); +// True when the CUDA build carries the in-graph-side top-k exponential +// sampler (device logits in, sampled codes out; no host logits readback). +bool torch_cuda_sample_topk_exponential_pairs_available(); + +// Samples one code per song from device-resident logits laid out as +// [cond_0; uncond_0; cond_1; ...] rows: bf16-rounds both branches, mixes with +// guidance_scale, applies top-k and the torch exponential ranking, matching +// the CPU path sample-for-sample (up to logf ULP differences). +// seeds/offset_blocks may be null after the first call of a frame: the +// device keeps the frame constants and offset_step_blocks advances the +// per-call RNG offset (call_step * blocks-per-call). +void torch_cuda_sample_topk_exponential_pairs( + const void * device_logits_f32, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const uint64_t * seeds, + const uint64_t * offset_blocks, + uint64_t offset_step_blocks, + const TorchCudaSamplingPolicy & policy, + int32_t * out_codes); + +// GPU-resident depth frame: per-codebook sampled codes and cond-row hiddens +// stay on device (residual ids for the next codebook are filled device-side); +// one host sync per frame in torch_cuda_depth_frame_end. All calls enqueue on +// the given backend stream. +void * torch_cuda_backend_stream(void * ggml_backend); +void torch_cuda_depth_frame_ensure(int64_t songs, int64_t levels, int64_t hidden_size, const TorchCudaSamplingPolicy & policy); +void torch_cuda_depth_frame_begin(const uint64_t * seeds, const uint64_t * offset_blocks, int64_t songs, void * stream); +void torch_cuda_depth_frame_sample( + const void * device_logits_f32, + int64_t level_index, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const TorchCudaSamplingPolicy & policy, + void * stream); +void torch_cuda_depth_frame_residual_fill( + void * residual_ids_i32, + int64_t previous_levels, + int64_t songs, + int64_t audio_vocab, + void * stream); +void torch_cuda_depth_frame_accumulate_hidden( + const void * hidden_f32, + int64_t level_index, + int64_t songs, + int64_t hidden_size, + void * stream); +void torch_cuda_depth_frame_end( + int32_t * host_codes, + float * host_hidden, + int64_t levels, + int64_t songs, + int64_t hidden_size, + void * stream); + float torch_cuda_tensor_iterator_exponential_element( uint64_t seed, uint64_t total_elements, diff --git a/scripts/minimax_music3/repack_lm_head_gguf.py b/scripts/minimax_music3/repack_lm_head_gguf.py new file mode 100755 index 000000000..623297f0f --- /dev/null +++ b/scripts/minimax_music3/repack_lm_head_gguf.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Repack a MiniMax Music 3 GGUF with an exact compact semantic LM head. + +The operation never dequantizes the head. It copies the already-quantized rows +used by the native MiniMax semantic sampler and leaves every other tensor byte +sequence unchanged. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np + + +LM_HEAD_NAME = "lm_head.weight" +SOURCE_VOCAB_SIZE = 200_000 +HIDDEN_SIZE = 4_096 +SEMANTIC_OFFSET = 151_675 +SEMANTIC_VOCAB_SIZE = 16_384 +EOS_TOKEN_ID = 151_670 +LAYOUT = "semantic_compact_v1" + + +def compact_token_ids() -> list[int]: + return list(range(SEMANTIC_OFFSET, SEMANTIC_OFFSET + SEMANTIC_VOCAB_SIZE)) + [EOS_TOKEN_ID] + + +def compact_lm_head_rows(source: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]: + if source.ndim != 2 or source.shape[0] != SOURCE_VOCAB_SIZE: + raise ValueError( + f"MiniMax Music 3 full lm_head must have {SOURCE_VOCAB_SIZE} rows; got {source.shape}" + ) + return np.ascontiguousarray(source[compact_token_ids()]) + + +def compact_logical_shapes(shapes: dict[str, tuple[int, ...]]) -> dict[str, tuple[int, ...]]: + if shapes.get(LM_HEAD_NAME) != (SOURCE_VOCAB_SIZE, HIDDEN_SIZE): + raise ValueError( + f"expected logical {LM_HEAD_NAME} shape {(SOURCE_VOCAB_SIZE, HIDDEN_SIZE)}; " + f"got {shapes.get(LM_HEAD_NAME)}" + ) + out = dict(shapes) + out[LM_HEAD_NAME] = (SEMANTIC_VOCAB_SIZE + 1, HIDDEN_SIZE) + return out + + +def _array_strings(reader: Any, key: str) -> list[str]: + parts = reader.fields[key].parts + count = int(parts[4][0]) + out: list[str] = [] + index = 5 + for _ in range(count): + size = int(parts[index][0]) + out.append(bytes(parts[index + 1][:size]).decode("utf-8")) + index += 2 + return out + + +def _array_i32(reader: Any, key: str) -> list[int]: + parts = reader.fields[key].parts + count = int(parts[4][0]) + return [int(parts[5 + index][0]) for index in range(count)] + + +def _array_i64(reader: Any, key: str) -> list[int]: + parts = reader.fields[key].parts + count = int(parts[4][0]) + return [int(parts[5 + index][0]) for index in range(count)] + + +def _field_string(reader: Any, key: str) -> str: + field = reader.fields[key] + size = int(field.parts[-2][0]) + return bytes(field.parts[-1][:size]).decode("utf-8") + + +def _logical_shapes(reader: Any) -> dict[str, tuple[int, ...]]: + names = _array_strings(reader, "audiocpp.tensor_names") + ranks = _array_i32(reader, "audiocpp.tensor_ranks") + flat = _array_i64(reader, "audiocpp.tensor_shapes") + if len(names) != len(ranks): + raise ValueError("GGUF tensor name/rank metadata length mismatch") + out: dict[str, tuple[int, ...]] = {} + offset = 0 + for name, rank in zip(names, ranks): + shape = tuple(flat[offset : offset + rank]) + if len(shape) != rank: + raise ValueError("GGUF logical tensor shapes are truncated") + out[name] = shape + offset += rank + if offset != len(flat): + raise ValueError("GGUF logical tensor shapes have trailing dimensions") + return out + + +def _sha256_array(data: np.ndarray[Any, Any]) -> str: + digest = hashlib.sha256() + view = memoryview(np.ascontiguousarray(data)).cast("B") + stride = 64 * 1024 * 1024 + for offset in range(0, len(view), stride): + digest.update(view[offset : offset + stride]) + return digest.hexdigest() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--manifest", type=Path) + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def main() -> int: + import gguf + + args = _parse_args() + source_path = args.input.expanduser().resolve() + output_path = args.output.expanduser().resolve() + manifest_path = ( + args.manifest.expanduser().resolve() + if args.manifest is not None + else output_path.with_suffix(output_path.suffix + ".manifest.json") + ) + if not source_path.is_file(): + raise FileNotFoundError(source_path) + if output_path.exists() and not args.overwrite: + raise FileExistsError(f"output already exists; pass --overwrite: {output_path}") + output_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + + reader = gguf.GGUFReader(source_path, "r") + shapes = _logical_shapes(reader) + output_shapes = compact_logical_shapes(shapes) + tensors = {tensor.name: tensor for tensor in reader.tensors} + if LM_HEAD_NAME not in tensors: + raise ValueError(f"source GGUF is missing {LM_HEAD_NAME}") + source_head = tensors[LM_HEAD_NAME] + if source_head.data.dtype != np.uint8: + raise ValueError(f"expected raw quantized lm_head bytes; got {source_head.data.dtype}") + compact_head = compact_lm_head_rows(source_head.data) + + tmp = output_path.with_name(output_path.name + ".tmp") + if tmp.exists(): + tmp.unlink() + writer = gguf.GGUFWriter(tmp, "audiocpp", use_temp_file=True) + try: + writer.add_name(output_path.stem) + writer.add_string("audiocpp.tensor_name_format", _field_string(reader, "audiocpp.tensor_name_format")) + writer.add_string("audiocpp.source_format", "gguf_compact_lm_head") + writer.add_string("audiocpp.weight_type", _field_string(reader, "audiocpp.weight_type")) + writer.add_array("audiocpp.tensor_sources.names", _array_strings(reader, "audiocpp.tensor_sources.names")) + writer.add_array("audiocpp.tensor_sources.paths", _array_strings(reader, "audiocpp.tensor_sources.paths")) + writer.add_string("standalone_gguf_converter.version", "1") + writer.add_string("minimax_music3.lm_head.layout", LAYOUT) + writer.add_uint64("minimax_music3.lm_head.source_vocab_size", SOURCE_VOCAB_SIZE) + writer.add_array("minimax_music3.lm_head.token_ids", compact_token_ids()) + + logical_names: list[str] = [] + logical_shape_list: list[tuple[int, ...]] = [] + for tensor in sorted(reader.tensors, key=lambda item: item.name): + data = compact_head if tensor.name == LM_HEAD_NAME else tensor.data + raw_dtype = tensor.tensor_type if data.dtype == np.uint8 else None + writer.add_tensor(tensor.name, data, raw_dtype=raw_dtype) + logical_names.append(tensor.name) + logical_shape_list.append(output_shapes[tensor.name]) + + writer.add_array("audiocpp.tensor_names", logical_names) + writer.add_key_value( + "audiocpp.tensor_ranks", + [len(shape) for shape in logical_shape_list], + gguf.GGUFValueType.ARRAY, + sub_type=gguf.GGUFValueType.INT32, + ) + writer.add_key_value( + "audiocpp.tensor_shapes", + [dim for shape in logical_shape_list for dim in shape], + gguf.GGUFValueType.ARRAY, + sub_type=gguf.GGUFValueType.INT64, + ) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + os.replace(tmp, output_path) + except Exception: + writer.close() + if tmp.exists(): + tmp.unlink() + raise + + verify = gguf.GGUFReader(output_path, "r") + output_head = next(tensor for tensor in verify.tensors if tensor.name == LM_HEAD_NAME) + if output_head.data.shape != compact_head.shape or not np.array_equal(output_head.data, compact_head): + raise RuntimeError("output compact lm_head does not match retained source rows byte-for-byte") + + manifest = { + "schema_version": 1, + "layout": LAYOUT, + "source": str(source_path), + "output": str(output_path), + "source_lm_head_shape": list(shapes[LM_HEAD_NAME]), + "output_lm_head_shape": list(output_shapes[LM_HEAD_NAME]), + "source_lm_head_type": str(source_head.tensor_type), + "source_lm_head_bytes": int(source_head.data.nbytes), + "output_lm_head_bytes": int(output_head.data.nbytes), + "source_lm_head_sha256": _sha256_array(source_head.data), + "output_lm_head_sha256": _sha256_array(output_head.data), + "token_ids": compact_token_ids(), + "row_bytes_equal": True, + "output_size_bytes": output_path.stat().st_size, + } + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") + print(f"output={output_path}") + print(f"manifest={manifest_path}") + print(f"lm_head_bytes={source_head.data.nbytes}->{output_head.data.nbytes}") + print("row_bytes_equal=true") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/community_models/minimax_music3/ar_runtime.cpp b/src/community_models/minimax_music3/ar_runtime.cpp index 7df9d9ab9..1b09aaab2 100644 --- a/src/community_models/minimax_music3/ar_runtime.cpp +++ b/src/community_models/minimax_music3/ar_runtime.cpp @@ -259,6 +259,15 @@ int32_t sample_semantic_token( return token; } +void assign_batch_row(const std::vector & values, int64_t row, int64_t width, std::vector & out) { + if (row < 0 || width <= 0 || static_cast(values.size()) < (row + 1) * width) { + throw std::runtime_error("MiniMax Music 3 batch output shape mismatch"); + } + out.assign( + values.begin() + static_cast(row * width), + values.begin() + static_cast((row + 1) * width)); +} + void assign_batch2_row(const std::vector & values, int64_t row, int64_t width, std::vector & out) { if (row < 0 || row >= 2 || width <= 0 || static_cast(values.size()) != 2 * width) { throw std::runtime_error("MiniMax Music 3 batch-2 output shape mismatch"); @@ -294,9 +303,11 @@ struct MiniMaxMusic3ArRuntime::Impl { core::ExecutionContext & input_execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), execution(input_execution), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), prompt_builder(assets), global_weights(load_minimax_music3_global_lm_weights(*assets, execution, weight_context_bytes, storage_type)), depth(std::make_unique( @@ -305,7 +316,8 @@ struct MiniMaxMusic3ArRuntime::Impl { execution, graph_arena_bytes, weight_context_bytes, - storage_type)), + storage_type, + evict_cuda_graph_cache_on_release)), sampling_policy(sampling::resolve_torch_cuda_sampling_policy( execution.backend_type(), execution.config().device, @@ -317,11 +329,15 @@ struct MiniMaxMusic3ArRuntime::Impl { } auto qwen_config = make_minimax_music3_global_lm_runtime_config( assets->config, + global_weights.lm_head_layout, execution.backend_type(), graph_arena_bytes, graph_arena_bytes); + qwen_config.evict_cuda_graph_cache_on_release = evict_cuda_graph_cache_on_release; qwen_config.return_hidden = true; - qwen_config.logits_readback_token_ids = semantic_logits_readback_token_ids(MiniMaxMusic3Prompt{}); + if (global_weights.lm_head_layout == MiniMaxMusic3LmHeadLayout::FullVocab) { + qwen_config.logits_readback_token_ids = semantic_logits_readback_token_ids(MiniMaxMusic3Prompt{}); + } global_runtime = std::make_unique( execution, qwen_config, @@ -348,10 +364,12 @@ struct MiniMaxMusic3ArRuntime::Impl { assign_batch2_row(hidden, 1, assets->config.qwen.hidden_size, state.uncond.hidden); } - std::vector generate_frame_hiddens( + void generate_frame_hiddens_into( const MiniMaxMusic3Request & request, int64_t target_frames, - uint64_t & rng_offset_blocks) { + uint64_t & rng_offset_blocks, + std::vector & frame_hiddens, + const std::function * progress) { const auto ar_start = Clock::now(); const auto prompt = prompt_from_request(request); const int64_t prompt_steps = static_cast(prompt.conditional_ids.size()); @@ -362,12 +380,17 @@ struct MiniMaxMusic3ArRuntime::Impl { ArStepState state; assign_step_outputs(std::move(prefill.logits), prefill.hidden, state); - std::vector frame_hiddens; + frame_hiddens.clear(); frame_hiddens.reserve(static_cast( target_frames * assets->config.condition.condition_layers * assets->config.qwen.hidden_size)); uint64_t sample_call_index = 0; std::mt19937 fallback_rng(static_cast(request.seed)); + double semantic_ms = 0.0; + double depth_ms = 0.0; + double feedback_ms = 0.0; + double lm_ms = 0.0; for (int64_t frame = 0; frame <= target_frames; ++frame) { + const auto t_sem = Clock::now(); const int32_t token = sample_semantic_token( prompt, state.logits, @@ -389,7 +412,9 @@ struct MiniMaxMusic3ArRuntime::Impl { token >= prompt.audio_code_offset + prompt.semantic_vocab_size) { throw std::runtime_error("MiniMax Music 3 sampled token outside semantic audio range"); } + semantic_ms += engine::debug::elapsed_ms(t_sem, Clock::now()); const int32_t semantic_code = token - prompt.audio_code_offset; + const auto t_depth = Clock::now(); auto depth_codes = depth->generate( state.cond.hidden, state.uncond.hidden, @@ -399,21 +424,343 @@ struct MiniMaxMusic3ArRuntime::Impl { request.seed, sample_call_index, rng_offset_blocks); + depth_ms += engine::debug::elapsed_ms(t_depth, Clock::now()); if (frame > 0) { frame_hiddens.insert(frame_hiddens.end(), state.cond.hidden.begin(), state.cond.hidden.end()); frame_hiddens.insert(frame_hiddens.end(), depth_codes.hidden.begin(), depth_codes.hidden.end()); + if (progress != nullptr && *progress) { + (*progress)(frame); + } } if (frame == target_frames) { break; } + const auto t_fb = Clock::now(); const auto feedback = depth->feedback_embedding(depth_codes.codes); duplicate_feedback_batch2(feedback, feedback_batch); + feedback_ms += engine::debug::elapsed_ms(t_fb, Clock::now()); + const auto t_lm = Clock::now(); auto step = global_runtime->decode_embeddings_batched(feedback_batch, 2); assign_step_outputs(std::move(step.logits), step.hidden, state); + lm_ms += engine::debug::elapsed_ms(t_lm, Clock::now()); } + engine::debug::timing_log_scalar("minimax_music3.ar.semantic_ms", semantic_ms); + engine::debug::timing_log_scalar("minimax_music3.ar.depth_ms", depth_ms); + engine::debug::timing_log_scalar("minimax_music3.ar.feedback_ms", feedback_ms); + engine::debug::timing_log_scalar("minimax_music3.ar.lm_decode_ms", lm_ms); engine::debug::timing_log_scalar( "minimax_music3.ar.total_ms", engine::debug::elapsed_ms(ar_start, Clock::now())); + } + + std::vector> generate_frame_hiddens_ensemble( + const MiniMaxMusic3Request & request, + int64_t target_frames, + const std::vector & take_seeds, + std::vector & rng_offset_blocks, + int64_t prefix_frames = 0) { + const auto ar_start = Clock::now(); + const int64_t takes = static_cast(take_seeds.size()); + if (takes <= 0 || static_cast(rng_offset_blocks.size()) != takes) { + throw std::runtime_error("MiniMax Music 3 AR ensemble take shape mismatch"); + } + const auto prompt = prompt_from_request(request); + const int64_t prompt_steps = static_cast(prompt.conditional_ids.size()); + const int64_t required_cache_steps = std::max(1, prompt_steps + target_frames); + const int64_t rows = 2 * takes; + const int64_t hidden_size = assets->config.qwen.hidden_size; + const int64_t prefix = std::min(std::max(0, prefix_frames), std::max(0, target_frames - 1)); + + // Intro-lock: run the shared prefix once at batch 2, then replicate + // the decode KV to all pairs so the takes fork mid-song. Take 0 keeps + // the master seed and counters, so it continues the master exactly. + std::vector master_logits; + std::vector master_hidden; + std::vector master_frame_hiddens; + uint64_t master_call = 0; + uint64_t master_offset = 0; + int64_t forked_at = 0; + bool master_finished = false; + if (prefix > 0 && takes > 1) { + MiniMaxMusic3Request master_request = request; + master_request.seed = take_seeds[0]; + auto prefill2 = global_runtime->prefill_tokens_batched(batch2_prompt_ids(prompt), 2, prompt_steps); + global_runtime->start_decode_embeddings_batched(prefill2.state, required_cache_steps); + ArStepState state; + assign_step_outputs(std::move(prefill2.logits), prefill2.hidden, state); + std::mt19937 master_rng(static_cast(master_request.seed)); + for (int64_t frame = 0; frame < prefix; ++frame) { + const int32_t token = sample_semantic_token( + prompt, state.logits, assets->config.qwen.vocab_size, + semantic_logits, topk_window, compact_candidates, + master_request, master_call, master_offset, + sampling_policy, semantic_scratch, semantic_weights, master_rng); + if (token == prompt.audio_end_token_id) { + master_finished = true; + break; + } + if (token < prompt.audio_code_offset || + token >= prompt.audio_code_offset + prompt.semantic_vocab_size) { + throw std::runtime_error("MiniMax Music 3 sampled token outside semantic audio range"); + } + auto depth_codes = depth->generate( + state.cond.hidden, state.uncond.hidden, + token - prompt.audio_code_offset, + master_request.ar_guidance_scale, master_request.top_k, + master_request.seed, master_call, master_offset); + if (frame > 0) { + master_frame_hiddens.insert( + master_frame_hiddens.end(), state.cond.hidden.begin(), state.cond.hidden.end()); + master_frame_hiddens.insert( + master_frame_hiddens.end(), depth_codes.hidden.begin(), depth_codes.hidden.end()); + } + const auto feedback = depth->feedback_embedding(depth_codes.codes); + duplicate_feedback_batch2(feedback, feedback_batch); + auto step = global_runtime->decode_embeddings_batched(feedback_batch, 2); + assign_step_outputs(std::move(step.logits), step.hidden, state); + forked_at = frame + 1; + } + if (master_finished) { + engine::debug::timing_log_scalar("minimax_music3.ar.prefix_eos", 1.0); + } else { + master_logits = state.logits; + master_hidden.clear(); + master_hidden.reserve(static_cast(2 * hidden_size)); + master_hidden.insert(master_hidden.end(), state.cond.hidden.begin(), state.cond.hidden.end()); + master_hidden.insert(master_hidden.end(), state.uncond.hidden.begin(), state.uncond.hidden.end()); + auto master_state = global_runtime->export_batched_decode_state(); + runtime::TransformerBatchedKVState forked; + forked.batch_size = rows; + forked.current_end = master_state.current_end; + forked.layers.resize(master_state.layers.size()); + for (size_t layer = 0; layer < master_state.layers.size(); ++layer) { + const auto & src = master_state.layers[layer]; + auto & dst = forked.layers[layer]; + dst.valid_steps = src.valid_steps; + const size_t pair_elems = src.key.size() / 2; + dst.key.resize(pair_elems * static_cast(rows)); + dst.value.resize(pair_elems * static_cast(rows)); + for (int64_t row = 0; row < rows; ++row) { + const size_t src_off = static_cast(row % 2) * pair_elems; + const size_t dst_off = static_cast(row) * pair_elems; + std::copy(src.key.begin() + src_off, src.key.begin() + src_off + pair_elems, + dst.key.begin() + dst_off); + std::copy(src.value.begin() + src_off, src.value.begin() + src_off + pair_elems, + dst.value.begin() + dst_off); + } + } + global_runtime->start_decode_embeddings_batched(forked, required_cache_steps); + engine::debug::timing_log_scalar("minimax_music3.ar.prefix_forked_frames", static_cast(forked_at)); + } + } + const bool forked = prefix > 0 && takes > 1 && !master_finished && forked_at > 0; + if (!forked) { + const auto pair_ids = batch2_prompt_ids(prompt); + std::vector prompt_ids; + prompt_ids.reserve(pair_ids.size() * static_cast(takes)); + for (int64_t take = 0; take < takes; ++take) { + prompt_ids.insert(prompt_ids.end(), pair_ids.begin(), pair_ids.end()); + } + auto prefill = global_runtime->prefill_tokens_batched(prompt_ids, rows, prompt_steps); + global_runtime->start_decode_embeddings_batched(prefill.state, required_cache_steps); + master_logits = std::move(prefill.logits); + master_hidden = std::move(prefill.hidden); + forked_at = 0; + master_call = 0; + master_offset = 0; + } + + // Per-take state. The shared graph advances every take each frame; + // takes that already sampled EOS keep occupying their rows (frozen + // feedback, discarded samples with restored counters) so graph shapes + // stay stable for CUDA-graph reuse. + struct TakeState { + MiniMaxMusic3Request request; + uint64_t sample_call_index = 0; + std::mt19937 fallback_rng; + std::vector cond_hidden; + std::vector uncond_hidden; + std::vector frame_hiddens; + std::vector last_feedback; + int32_t semantic_code = 0; + bool finished = false; + }; + std::vector take_states(static_cast(takes)); + for (int64_t take = 0; take < takes; ++take) { + auto & ts = take_states[static_cast(take)]; + ts.request = request; + ts.request.seed = take_seeds[static_cast(take)]; + ts.fallback_rng.seed(static_cast(ts.request.seed)); + ts.frame_hiddens.reserve(static_cast( + target_frames * assets->config.condition.condition_layers * hidden_size)); + } + + std::vector batch_logits; + std::vector batch_hidden; + if (forked) { + // Tile the master pair's outputs to every take: all takes see the + // same logits/hidden at the fork frame, then diverge by sampling. + const size_t pair_logits_elems = master_logits.size(); + const size_t pair_hidden_elems = master_hidden.size(); + batch_logits.resize(pair_logits_elems * static_cast(takes)); + batch_hidden.resize(pair_hidden_elems * static_cast(takes)); + for (int64_t take = 0; take < takes; ++take) { + std::copy(master_logits.begin(), master_logits.end(), + batch_logits.begin() + static_cast(take * pair_logits_elems)); + std::copy(master_hidden.begin(), master_hidden.end(), + batch_hidden.begin() + static_cast(take * pair_hidden_elems)); + } + for (int64_t take = 0; take < takes; ++take) { + rng_offset_blocks[static_cast(take)] = master_offset; + take_states[static_cast(take)].sample_call_index = master_call; + take_states[static_cast(take)].frame_hiddens = master_frame_hiddens; + } + } else { + batch_logits = std::move(master_logits); + batch_hidden = std::move(master_hidden); + } + const int64_t logits_width = static_cast(batch_logits.size()) / rows; + if (static_cast(batch_logits.size()) != rows * logits_width || + static_cast(batch_hidden.size()) != rows * hidden_size) { + throw std::runtime_error("MiniMax Music 3 AR ensemble prefill shape mismatch"); + } + + std::vector pair_logits; + std::vector depth_hiddens(static_cast(rows * hidden_size)); + std::vector depth_semantic(static_cast(takes)); + std::vector depth_calls(static_cast(takes)); + std::vector depth_offsets(static_cast(takes)); + + for (int64_t frame = forked_at; frame <= target_frames; ++frame) { + for (int64_t take = 0; take < takes; ++take) { + auto & ts = take_states[static_cast(take)]; + assign_batch_row(batch_hidden, 2 * take, hidden_size, ts.cond_hidden); + assign_batch_row(batch_hidden, 2 * take + 1, hidden_size, ts.uncond_hidden); + if (ts.finished) { + continue; + } + pair_logits.resize(static_cast(2 * logits_width)); + std::copy_n( + batch_logits.begin() + static_cast(2 * take * logits_width), + static_cast(2 * logits_width), + pair_logits.begin()); + const int32_t token = sample_semantic_token( + prompt, + pair_logits, + assets->config.qwen.vocab_size, + semantic_logits, + topk_window, + compact_candidates, + ts.request, + ts.sample_call_index, + rng_offset_blocks[static_cast(take)], + sampling_policy, + semantic_scratch, + semantic_weights, + ts.fallback_rng); + if (token == prompt.audio_end_token_id) { + ts.finished = true; + continue; + } + if (token < prompt.audio_code_offset || + token >= prompt.audio_code_offset + prompt.semantic_vocab_size) { + throw std::runtime_error("MiniMax Music 3 sampled token outside semantic audio range"); + } + ts.semantic_code = token - prompt.audio_code_offset; + } + bool any_active = false; + for (const auto & ts : take_states) { + any_active = any_active || !ts.finished; + } + if (!any_active) { + break; + } + for (int64_t take = 0; take < takes; ++take) { + const auto & ts = take_states[static_cast(take)]; + std::copy( + ts.cond_hidden.begin(), ts.cond_hidden.end(), + depth_hiddens.begin() + static_cast((2 * take) * hidden_size)); + std::copy( + ts.uncond_hidden.begin(), ts.uncond_hidden.end(), + depth_hiddens.begin() + static_cast((2 * take + 1) * hidden_size)); + depth_semantic[static_cast(take)] = ts.semantic_code; + depth_calls[static_cast(take)] = ts.sample_call_index; + depth_offsets[static_cast(take)] = rng_offset_blocks[static_cast(take)]; + } + auto depth_codes = depth->generate_batch( + depth_hiddens, + takes, + depth_semantic, + request.ar_guidance_scale, + request.top_k, + take_seeds, + depth_calls, + depth_offsets); + for (int64_t take = 0; take < takes; ++take) { + auto & ts = take_states[static_cast(take)]; + if (ts.finished) { + continue; // discard the row's output, counters stay frozen + } + ts.sample_call_index = depth_calls[static_cast(take)]; + rng_offset_blocks[static_cast(take)] = depth_offsets[static_cast(take)]; + if (frame > 0) { + auto & take_codes = depth_codes[static_cast(take)]; + ts.frame_hiddens.insert( + ts.frame_hiddens.end(), ts.cond_hidden.begin(), ts.cond_hidden.end()); + ts.frame_hiddens.insert( + ts.frame_hiddens.end(), take_codes.hidden.begin(), take_codes.hidden.end()); + } + } + if (frame == target_frames) { + break; + } + feedback_batch.resize(static_cast(rows * hidden_size)); + for (int64_t take = 0; take < takes; ++take) { + auto & ts = take_states[static_cast(take)]; + if (!ts.finished) { + ts.last_feedback = depth->feedback_embedding(depth_codes[static_cast(take)].codes); + } + if (ts.last_feedback.empty()) { + ts.last_feedback.assign(static_cast(hidden_size), 0.0F); + } + std::copy( + ts.last_feedback.begin(), ts.last_feedback.end(), + feedback_batch.begin() + static_cast((2 * take) * hidden_size)); + std::copy( + ts.last_feedback.begin(), ts.last_feedback.end(), + feedback_batch.begin() + static_cast((2 * take + 1) * hidden_size)); + } + auto step = global_runtime->decode_embeddings_batched(feedback_batch, rows); + batch_logits = std::move(step.logits); + batch_hidden = std::move(step.hidden); + if (static_cast(batch_logits.size()) != rows * logits_width || + static_cast(batch_hidden.size()) != rows * hidden_size) { + throw std::runtime_error("MiniMax Music 3 AR ensemble step shape mismatch"); + } + } + if (prefix > 0 && takes > 1 && master_finished) { + for (int64_t take = 0; take < takes; ++take) { + take_states[static_cast(take)].frame_hiddens = master_frame_hiddens; + rng_offset_blocks[static_cast(take)] = master_offset; + } + } + engine::debug::timing_log_scalar( + "minimax_music3.ar.ensemble_total_ms", + engine::debug::elapsed_ms(ar_start, Clock::now())); + std::vector> out; + out.reserve(static_cast(takes)); + for (auto & ts : take_states) { + out.push_back(std::move(ts.frame_hiddens)); + } + return out; + } + + std::vector generate_frame_hiddens( + const MiniMaxMusic3Request & request, + int64_t target_frames, + uint64_t & rng_offset_blocks) { + std::vector frame_hiddens; + generate_frame_hiddens_into(request, target_frames, rng_offset_blocks, frame_hiddens, nullptr); return frame_hiddens; } @@ -428,6 +775,7 @@ struct MiniMaxMusic3ArRuntime::Impl { std::shared_ptr assets; core::ExecutionContext & execution; + bool evict_cuda_graph_cache_on_release = false; MiniMaxMusic3PromptBuilder prompt_builder; MiniMaxMusic3GlobalLMWeights global_weights; std::unique_ptr global_runtime; @@ -446,13 +794,15 @@ MiniMaxMusic3ArRuntime::MiniMaxMusic3ArRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3ArRuntime::~MiniMaxMusic3ArRuntime() = default; @@ -463,6 +813,25 @@ std::vector MiniMaxMusic3ArRuntime::generate_frame_hiddens( return impl_->generate_frame_hiddens(request, target_frames, rng_offset_blocks); } +void MiniMaxMusic3ArRuntime::generate_frame_hiddens_into( + const MiniMaxMusic3Request & request, + int64_t target_frames, + uint64_t & rng_offset_blocks, + std::vector & frame_hiddens, + const std::function * progress) { + impl_->generate_frame_hiddens_into(request, target_frames, rng_offset_blocks, frame_hiddens, progress); +} + +std::vector> MiniMaxMusic3ArRuntime::generate_frame_hiddens_ensemble( + const MiniMaxMusic3Request & request, + int64_t target_frames, + const std::vector & take_seeds, + std::vector & rng_offset_blocks, + int64_t prefix_frames) { + return impl_->generate_frame_hiddens_ensemble( + request, target_frames, take_seeds, rng_offset_blocks, prefix_frames); +} + void MiniMaxMusic3ArRuntime::release_runtime_graphs() { if (impl_ != nullptr) { impl_->release_runtime_graphs(); diff --git a/src/community_models/minimax_music3/condition_encoder.cpp b/src/community_models/minimax_music3/condition_encoder.cpp index 33f21b285..5e683058e 100644 --- a/src/community_models/minimax_music3/condition_encoder.cpp +++ b/src/community_models/minimax_music3/condition_encoder.cpp @@ -87,16 +87,47 @@ MiniMaxMusic3ConditionWeights load_condition_weights( } // namespace +std::vector detail::project_frame_hiddens( + const float * frame_hiddens, + size_t value_count, + int64_t frames, + int64_t layers, + int64_t hidden_size, + const std::vector & layer_weights) { + if (frames <= 0 || layers <= 0 || hidden_size <= 0) { + throw std::runtime_error("MiniMax Music 3 condition projection requires positive dimensions"); + } + const size_t expected = static_cast(frames * layers * hidden_size); + if (frame_hiddens == nullptr || value_count != expected || + layer_weights.size() != static_cast(layers)) { + throw std::runtime_error("MiniMax Music 3 condition frame hidden shape mismatch"); + } + std::vector projected(static_cast(hidden_size * frames), 0.0F); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t layer = 0; layer < layers; ++layer) { + const float weight = layer_weights[static_cast(layer)]; + const size_t src_base = static_cast((frame * layers + layer) * hidden_size); + for (int64_t channel = 0; channel < hidden_size; ++channel) { + projected[static_cast(channel * frames + frame)] += + frame_hiddens[src_base + static_cast(channel)] * weight; + } + } + } + return projected; +} + struct MiniMaxMusic3ConditionEncoderRuntime::Impl { Impl( std::shared_ptr input_assets, core::ExecutionContext & input_execution, size_t input_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), execution(input_execution), graph_arena_bytes(input_graph_arena_bytes), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), weights(load_condition_weights(*assets, execution, weight_context_bytes, storage_type)) { if (assets == nullptr) { throw std::runtime_error("MiniMax Music 3 condition encoder requires assets"); @@ -109,7 +140,8 @@ struct MiniMaxMusic3ConditionEncoderRuntime::Impl { void release_runtime_graphs() { if (graph != nullptr) { - core::release_backend_graph_resources(execution.backend(), graph); + core::release_backend_graph_resources( + execution.backend(), graph, evict_cuda_graph_cache_on_release); } graph = nullptr; input = {}; @@ -164,12 +196,12 @@ struct MiniMaxMusic3ConditionEncoderRuntime::Impl { condition_frames = output_frames; } - std::vector encode(const std::vector & frame_hiddens, int64_t input_frames, int64_t & out_frames) { + std::vector encode( + const float * frame_hiddens, + size_t value_count, + int64_t input_frames, + int64_t & out_frames) { const auto & config = assets->config.condition; - const int64_t expected = input_frames * config.condition_layers * config.condition_hidden_dim; - if (static_cast(frame_hiddens.size()) != expected) { - throw std::runtime_error("MiniMax Music 3 condition frame hidden shape mismatch"); - } out_frames = static_cast( static_cast(input_frames) * static_cast(config.output_sample_rate) / static_cast(config.input_sample_rate) * static_cast(config.input_hop_length) / @@ -177,17 +209,13 @@ struct MiniMaxMusic3ConditionEncoderRuntime::Impl { if (out_frames <= 0) { throw std::runtime_error("MiniMax Music 3 condition encoder produced non-positive frame count"); } - std::vector projected_input(static_cast(config.condition_hidden_dim * input_frames), 0.0F); - for (int64_t frame = 0; frame < input_frames; ++frame) { - for (int64_t layer = 0; layer < config.condition_layers; ++layer) { - const float weight = weights.layer_weights[static_cast(layer)]; - const size_t src_base = static_cast((frame * config.condition_layers + layer) * config.condition_hidden_dim); - for (int64_t channel = 0; channel < config.condition_hidden_dim; ++channel) { - projected_input[static_cast(channel * input_frames + frame)] += - frame_hiddens[src_base + static_cast(channel)] * weight; - } - } - } + auto projected_input = detail::project_frame_hiddens( + frame_hiddens, + value_count, + input_frames, + config.condition_layers, + config.condition_hidden_dim, + weights.layer_weights); ensure_graph(input_frames, out_frames); core::write_tensor_f32(input, projected_input); if (core::compute_graph(execution, graph, plan, "minimax_music3.condition") != GGML_STATUS_SUCCESS) { @@ -207,6 +235,7 @@ struct MiniMaxMusic3ConditionEncoderRuntime::Impl { std::shared_ptr assets; core::ExecutionContext & execution; size_t graph_arena_bytes = 0; + bool evict_cuda_graph_cache_on_release = false; MiniMaxMusic3ConditionWeights weights; int64_t frames = 0; int64_t condition_frames = 0; @@ -223,21 +252,31 @@ MiniMaxMusic3ConditionEncoderRuntime::MiniMaxMusic3ConditionEncoderRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3ConditionEncoderRuntime::~MiniMaxMusic3ConditionEncoderRuntime() = default; +std::vector MiniMaxMusic3ConditionEncoderRuntime::encode( + const float * frame_hiddens, + size_t value_count, + int64_t frames, + int64_t & condition_frames) { + return impl_->encode(frame_hiddens, value_count, frames, condition_frames); +} + std::vector MiniMaxMusic3ConditionEncoderRuntime::encode( const std::vector & frame_hiddens, int64_t frames, int64_t & condition_frames) { - return impl_->encode(frame_hiddens, frames, condition_frames); + return impl_->encode(frame_hiddens.data(), frame_hiddens.size(), frames, condition_frames); } void MiniMaxMusic3ConditionEncoderRuntime::release_runtime_graphs() { diff --git a/src/community_models/minimax_music3/depth_decoder.cpp b/src/community_models/minimax_music3/depth_decoder.cpp index c95c95b07..c5844973f 100644 --- a/src/community_models/minimax_music3/depth_decoder.cpp +++ b/src/community_models/minimax_music3/depth_decoder.cpp @@ -178,7 +178,7 @@ MiniMaxMusic3DepthWeights load_depth_weights( } int32_t sample_top_k( - std::vector logits, + std::vector & logits, int64_t top_k, uint64_t seed, uint64_t & sample_call_index, @@ -243,11 +243,13 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { core::ExecutionContext & input_execution, size_t input_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), global_token_embedding(input_global_token_embedding), execution(input_execution), graph_arena_bytes(input_graph_arena_bytes), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), weights(load_depth_weights(*assets, execution, weight_context_bytes, storage_type)), sampling_policy(sampling::resolve_torch_cuda_sampling_policy( execution.backend_type(), @@ -268,17 +270,21 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { release_runtime_graphs(); } - DecodeGraph & decode_graph(int64_t codebook) { + DecodeGraph & decode_graph(int64_t codebook, int64_t rows) { const int64_t sequence_steps = codebook + 1; auto & slot = decode_graphs[static_cast(codebook - 1)]; - if (slot.graph != nullptr) { + if (slot.graph != nullptr && graph_rows == rows) { return slot; } - build_decode_graph(slot, codebook, sequence_steps); + if (graph_rows != rows) { + release_runtime_graphs(); + graph_rows = rows; + } + build_decode_graph(slot, codebook, sequence_steps, rows); return slot; } - void build_decode_graph(DecodeGraph & out, int64_t codebook, int64_t sequence_steps) { + void build_decode_graph(DecodeGraph & out, int64_t codebook, int64_t sequence_steps, int64_t rows) { const auto & config = assets->config.depth; ggml_init_params params{graph_arena_bytes, nullptr, true}; out.ggml.reset(ggml_init(params)); @@ -288,8 +294,8 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { core::ModuleBuildContext ctx{out.ggml.get(), "minimax_music3.depth", execution.backend_type()}; out.codebook = codebook; out.sequence_steps = sequence_steps; - out.last_hidden = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, config.hidden_size})); - out.semantic_ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({2})); + out.last_hidden = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({rows, config.hidden_size})); + out.semantic_ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({rows})); out.positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({sequence_steps})); ggml_set_input(out.last_hidden.tensor); ggml_set_input(out.semantic_ids.tensor); @@ -299,7 +305,7 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { ctx, out.last_hidden, weights.projection); - hidden0 = core::reshape_tensor(ctx, hidden0, core::TensorShape::from_dims({2, 1, config.hidden_size})); + hidden0 = core::reshape_tensor(ctx, hidden0, core::TensorShape::from_dims({rows, 1, config.hidden_size})); auto semantic = modules::EmbeddingModule({assets->config.qwen.vocab_size, config.hidden_size}) .build(ctx, out.semantic_ids, global_token_embedding); @@ -307,13 +313,13 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { ctx, semantic, weights.projection); - semantic = core::reshape_tensor(ctx, semantic, core::TensorShape::from_dims({2, 1, config.hidden_size})); + semantic = core::reshape_tensor(ctx, semantic, core::TensorShape::from_dims({rows, 1, config.hidden_size})); auto sequence = modules::ConcatModule({1}).build(ctx, hidden0, semantic); if (codebook > 1) { out.residual_ids = core::make_tensor( ctx, GGML_TYPE_I32, - core::TensorShape::from_dims({2, codebook - 1})); + core::TensorShape::from_dims({rows, codebook - 1})); ggml_set_input(out.residual_ids.tensor); auto residual = modules::EmbeddingModule({ config.audio_vocab_size * (config.codebooks - 1), @@ -328,7 +334,7 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { auto pos = modules::EmbeddingModule({config.max_position_embeddings, config.hidden_size}) .build(ctx, out.positions, weights.position_embedding); pos = core::reshape_tensor(ctx, pos, core::TensorShape::from_dims({1, sequence_steps, config.hidden_size})); - pos = modules::RepeatModule({core::TensorShape::from_dims({2, sequence_steps, config.hidden_size})}).build(ctx, pos); + pos = modules::RepeatModule({core::TensorShape::from_dims({rows, sequence_steps, config.hidden_size})}).build(ctx, pos); sequence = core::wrap_tensor( ggml_add(ctx.ggml, sequence.tensor, pos.tensor), sequence.shape, @@ -342,7 +348,7 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { weights.norm); auto last = modules::SliceModule({1, sequence_steps - 1, 1}).build(ctx, normalized); last = core::ensure_backend_addressable_layout(ctx, last); - last = core::reshape_tensor(ctx, last, core::TensorShape::from_dims({2, config.hidden_size})); + last = core::reshape_tensor(ctx, last, core::TensorShape::from_dims({rows, config.hidden_size})); auto logits = modules::LinearModule({config.hidden_size, config.audio_vocab_size, false}).build( ctx, last, @@ -413,33 +419,204 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { static_cast(last_hidden_uncond.size()) != config.hidden_size) { throw std::runtime_error("MiniMax Music 3 depth hidden input shape mismatch"); } - last_hidden_scratch.resize(static_cast(2 * config.hidden_size)); - std::copy(last_hidden_cond.begin(), last_hidden_cond.end(), last_hidden_scratch.begin()); - std::copy(last_hidden_uncond.begin(), last_hidden_uncond.end(), last_hidden_scratch.begin() + config.hidden_size); - std::vector out_codes{semantic_code}; - std::vector out_hidden; - out_hidden.reserve(static_cast((config.codebooks - 1) * config.hidden_size)); - std::mt19937 fallback_rng(static_cast(seed)); + std::vector interleaved(static_cast(2 * config.hidden_size)); + std::copy(last_hidden_cond.begin(), last_hidden_cond.end(), interleaved.begin()); + std::copy(last_hidden_uncond.begin(), last_hidden_uncond.end(), interleaved.begin() + config.hidden_size); + std::vector seeds{seed}; + std::vector calls{sample_call_index}; + std::vector offsets{rng_offset_blocks}; + auto batch = generate_batch(interleaved, 1, {semantic_code}, guidance_scale, top_k, seeds, calls, offsets); + sample_call_index = calls[0]; + rng_offset_blocks = offsets[0]; + return std::move(batch.front()); + } + + void generate_batch_gpu_frame( + const std::vector & interleaved_hiddens, + int64_t songs, + const std::vector & semantic_codes, + float guidance_scale, + int64_t top_k, + const std::vector & seeds, + std::vector & sample_call_indices, + std::vector & rng_offset_blocks, + std::vector & out) { + const auto & config = assets->config.depth; + const int64_t rows = 2 * songs; + const int64_t levels = config.codebooks - 1; + const int64_t hidden_size = config.hidden_size; + sampling::torch_cuda_depth_frame_ensure(songs, levels, hidden_size, sampling_policy); + void * stream = sampling::torch_cuda_backend_stream(execution.backend()); + ggml_backend_t backend = execution.backend(); + + semantic_ids_scratch.resize(static_cast(rows)); + for (int64_t song = 0; song < songs; ++song) { + const int32_t id = semantic_codes[static_cast(song)] + 151675; + semantic_ids_scratch[static_cast(2 * song)] = id; + semantic_ids_scratch[static_cast(2 * song + 1)] = id; + } + positions_scratch.resize(static_cast(config.codebooks)); + for (int64_t i = 0; i < config.codebooks; ++i) { + positions_scratch[static_cast(i)] = static_cast(i); + } + sampling::torch_cuda_depth_frame_begin( + seeds.data(), rng_offset_blocks.data(), songs, stream); for (int64_t codebook = 1; codebook < config.codebooks; ++codebook) { - auto & graph = decode_graph(codebook); - const int32_t semantic_ids[2] = { - semantic_code + 151675, - semantic_code + 151675, - }; - active_residual_ids_scratch.assign(static_cast(2 * std::max(1, codebook - 1)), 0); - for (int64_t previous = 1; previous < codebook; ++previous) { - const int32_t id = out_codes[static_cast(previous)] + - static_cast((previous - 1) * config.audio_vocab_size); - active_residual_ids_scratch[static_cast(previous - 1)] = id; - active_residual_ids_scratch[static_cast((codebook - 1) + previous - 1)] = id; + auto & graph = decode_graph(codebook, rows); + ggml_backend_tensor_set_async( + backend, graph.last_hidden.tensor, + interleaved_hiddens.data(), 0, + interleaved_hiddens.size() * sizeof(float)); + ggml_backend_tensor_set_async( + backend, graph.semantic_ids.tensor, + semantic_ids_scratch.data(), 0, + semantic_ids_scratch.size() * sizeof(int32_t)); + if (codebook > 1) { + sampling::torch_cuda_depth_frame_residual_fill( + graph.residual_ids.tensor->data, + codebook - 1, + songs, + config.audio_vocab_size, + stream); + } + ggml_backend_tensor_set_async( + backend, graph.positions.tensor, + positions_scratch.data(), 0, + static_cast(codebook + 1) * sizeof(int32_t)); + if (ggml_backend_graph_compute_async(backend, graph.graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("MiniMax Music 3 depth graph compute failed"); + } + sampling::torch_cuda_depth_frame_sample( + graph.logits->data, + codebook - 1, + songs, + config.audio_vocab_size, + guidance_scale, + top_k, + sampling_policy, + stream); + sampling::torch_cuda_depth_frame_accumulate_hidden( + graph.hidden->data, + codebook - 1, + songs, + hidden_size, + stream); + } + frame_codes_scratch.resize(static_cast(levels * songs)); + frame_hidden_scratch.resize(static_cast(levels * songs * hidden_size)); + sampling::torch_cuda_depth_frame_end( + frame_codes_scratch.data(), + frame_hidden_scratch.data(), + levels, + songs, + hidden_size, + stream); + core::round_f32_to_bf16_in_place(frame_hidden_scratch); + const uint64_t step_blocks = sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.audio_vocab_size), + sampling_policy); + for (int64_t level = 0; level < levels; ++level) { + for (int64_t song = 0; song < songs; ++song) { + auto & take = out[static_cast(song)]; + const int32_t code = frame_codes_scratch[static_cast(level * songs + song)]; + if (code < 0) { + throw std::runtime_error("MiniMax Music 3 depth GPU frame sampler selected no token"); + } + take.codes.push_back(code); + const float * hidden_row = frame_hidden_scratch.data() + + static_cast((level * songs + song) * hidden_size); + take.hidden.insert(take.hidden.end(), hidden_row, hidden_row + hidden_size); + ++sample_call_indices[static_cast(song)]; + rng_offset_blocks[static_cast(song)] += step_blocks; + } + } + } + + std::vector generate_batch( + const std::vector & interleaved_hiddens, + int64_t songs, + const std::vector & semantic_codes, + float guidance_scale, + int64_t top_k, + const std::vector & seeds, + std::vector & sample_call_indices, + std::vector & rng_offset_blocks) { + const auto & config = assets->config.depth; + const int64_t rows = 2 * songs; + if (songs <= 0 || + static_cast(interleaved_hiddens.size()) != rows * config.hidden_size || + static_cast(semantic_codes.size()) != songs || + static_cast(seeds.size()) != songs || + static_cast(sample_call_indices.size()) != songs || + static_cast(rng_offset_blocks.size()) != songs) { + throw std::runtime_error("MiniMax Music 3 depth batch input shape mismatch"); + } + std::vector out(static_cast(songs)); + std::vector fallback_rngs; + fallback_rngs.reserve(static_cast(songs)); + for (int64_t song = 0; song < songs; ++song) { + out[static_cast(song)].codes = {semantic_codes[static_cast(song)]}; + out[static_cast(song)].hidden.reserve( + static_cast((config.codebooks - 1) * config.hidden_size)); + fallback_rngs.emplace_back(static_cast(seeds[static_cast(song)])); + } + + // The CUDA sampler consumes the logits where they already live and + // returns one code per song, replacing the logits readback plus the + // per-song CPU top-k/exponential scan (the depth stage's dominant + // cost). Sample-for-sample it matches the CPU path up to logf ULPs. + static const bool gpu_sampler_requested = [] { + const char * env = std::getenv("MM3_DEPTH_GPU_SAMPLE"); + return env != nullptr && env[0] == '1'; + }(); + const bool use_gpu_sampler = gpu_sampler_requested && + sampling_policy.cuda_fast_path && + sampling::torch_cuda_sample_topk_exponential_pairs_available(); + // v2: the whole codebook chain stays on the backend stream — inputs go + // in with async sets, sampled codes feed the next codebook's residual + // ids device-side, cond hiddens accumulate on device, and the frame + // ends with a single host sync. Latency was the depth stage's cost + // (write->launch->sync->readback per codebook), not compute. + static const bool gpu_frame_requested = [] { + const char * env = std::getenv("MM3_DEPTH_GPU_FRAME"); + return env != nullptr && env[0] == '1'; + }(); + const bool use_gpu_frame = gpu_frame_requested && + sampling_policy.cuda_fast_path && + sampling::torch_cuda_sample_topk_exponential_pairs_available(); + if (use_gpu_frame) { + generate_batch_gpu_frame( + interleaved_hiddens, songs, semantic_codes, guidance_scale, top_k, + seeds, sample_call_indices, rng_offset_blocks, out); + return out; + } + + for (int64_t codebook = 1; codebook < config.codebooks; ++codebook) { + auto & graph = decode_graph(codebook, rows); + semantic_ids_scratch.resize(static_cast(rows)); + for (int64_t song = 0; song < songs; ++song) { + const int32_t id = semantic_codes[static_cast(song)] + 151675; + semantic_ids_scratch[static_cast(2 * song)] = id; + semantic_ids_scratch[static_cast(2 * song + 1)] = id; + } + const int64_t residual_width = std::max(1, codebook - 1); + active_residual_ids_scratch.assign(static_cast(rows * residual_width), 0); + for (int64_t song = 0; song < songs; ++song) { + const auto & song_codes = out[static_cast(song)].codes; + for (int64_t previous = 1; previous < codebook; ++previous) { + const int32_t id = song_codes[static_cast(previous)] + + static_cast((previous - 1) * config.audio_vocab_size); + active_residual_ids_scratch[static_cast((2 * song) * (codebook - 1) + previous - 1)] = id; + active_residual_ids_scratch[static_cast((2 * song + 1) * (codebook - 1) + previous - 1)] = id; + } } positions_scratch.resize(static_cast(codebook + 1)); for (int64_t i = 0; i <= codebook; ++i) { positions_scratch[static_cast(i)] = static_cast(i); } - core::write_tensor_f32(graph.last_hidden, last_hidden_scratch); - core::write_tensor_i32(graph.semantic_ids, semantic_ids, 2); + core::write_tensor_f32(graph.last_hidden, interleaved_hiddens); + core::write_tensor_i32(graph.semantic_ids, semantic_ids_scratch); if (codebook > 1) { core::write_tensor_i32(graph.residual_ids, active_residual_ids_scratch); } @@ -447,30 +624,70 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { if (core::compute_graph(execution, graph.graph, graph.plan, "minimax_music3.depth") != GGML_STATUS_SUCCESS) { throw std::runtime_error("MiniMax Music 3 depth graph compute failed"); } + // Reading hidden synchronizes the backend, so the graph's logits + // are complete before the sampler kernel touches them in place. auto hidden = core::read_tensor_f32(graph.hidden); core::round_f32_to_bf16_in_place(hidden); - out_hidden.insert(out_hidden.end(), hidden.begin(), hidden.begin() + config.hidden_size); + if (use_gpu_sampler) { + gpu_codes_scratch.resize(static_cast(songs)); + const uint64_t step_blocks = sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.audio_vocab_size), + sampling_policy); + sampling::torch_cuda_sample_topk_exponential_pairs( + graph.logits->data, + songs, + config.audio_vocab_size, + guidance_scale, + top_k, + codebook == 1 ? seeds.data() : nullptr, + codebook == 1 ? rng_offset_blocks.data() : nullptr, + static_cast(codebook - 1) * step_blocks, + sampling_policy, + gpu_codes_scratch.data()); + for (int64_t song = 0; song < songs; ++song) { + auto & take = out[static_cast(song)]; + const float * cond_hidden = hidden.data() + static_cast(2 * song) * config.hidden_size; + take.hidden.insert(take.hidden.end(), cond_hidden, cond_hidden + config.hidden_size); + const int32_t code = gpu_codes_scratch[static_cast(song)]; + if (code < 0) { + throw std::runtime_error("MiniMax Music 3 depth GPU sampler selected no token"); + } + take.codes.push_back(code); + ++sample_call_indices[static_cast(song)]; + rng_offset_blocks[static_cast(song)] += + sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.audio_vocab_size), + sampling_policy); + } + continue; + } auto logits = core::read_tensor_f32(graph.logits); core::round_f32_to_bf16_in_place(logits); - for (int64_t i = 0; i < config.audio_vocab_size; ++i) { - const float cond = logits[static_cast(i)]; - const float uncond = logits[static_cast(config.audio_vocab_size + i)]; - logits[static_cast(i)] = uncond + (cond - uncond) * guidance_scale; + for (int64_t song = 0; song < songs; ++song) { + auto & take = out[static_cast(song)]; + const float * cond_hidden = hidden.data() + static_cast(2 * song) * config.hidden_size; + take.hidden.insert(take.hidden.end(), cond_hidden, cond_hidden + config.hidden_size); + song_logits_scratch.resize(static_cast(config.audio_vocab_size)); + const float * cond_logits = logits.data() + static_cast(2 * song) * config.audio_vocab_size; + const float * uncond_logits = logits.data() + static_cast(2 * song + 1) * config.audio_vocab_size; + for (int64_t i = 0; i < config.audio_vocab_size; ++i) { + song_logits_scratch[static_cast(i)] = + uncond_logits[i] + (cond_logits[i] - uncond_logits[i]) * guidance_scale; + } + const int32_t code = sample_top_k( + song_logits_scratch, + top_k, + seeds[static_cast(song)], + sample_call_indices[static_cast(song)], + rng_offset_blocks[static_cast(song)], + sampling_policy, + scratch, + fallback_rngs[static_cast(song)], + "MiniMax Music 3 depth"); + take.codes.push_back(code); } - logits.resize(static_cast(config.audio_vocab_size)); - const int32_t code = sample_top_k( - std::move(logits), - top_k, - seed, - sample_call_index, - rng_offset_blocks, - sampling_policy, - scratch, - fallback_rng, - "MiniMax Music 3 depth"); - out_codes.push_back(code); } - return {std::move(out_codes), std::move(out_hidden)}; + return out; } std::vector feedback_embedding(const std::vector & codes) { @@ -498,12 +715,14 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { void release_runtime_graphs() { for (auto & graph : decode_graphs) { if (graph.graph != nullptr) { - core::release_backend_graph_resources(execution.backend(), graph.graph); + core::release_backend_graph_resources( + execution.backend(), graph.graph, evict_cuda_graph_cache_on_release); } graph = {}; } if (feedback.graph != nullptr) { - core::release_backend_graph_resources(execution.backend(), feedback.graph); + core::release_backend_graph_resources( + execution.backend(), feedback.graph, evict_cuda_graph_cache_on_release); } feedback = {}; } @@ -512,12 +731,19 @@ struct MiniMaxMusic3DepthDecoderRuntime::Impl { core::TensorValue global_token_embedding; core::ExecutionContext & execution; size_t graph_arena_bytes = 0; + bool evict_cuda_graph_cache_on_release = false; MiniMaxMusic3DepthWeights weights; std::array decode_graphs; + int64_t graph_rows = 2; FeedbackGraph feedback; sampling::TorchCudaSamplingPolicy sampling_policy; sampling::HfSamplerScratch scratch; std::vector last_hidden_scratch; + std::vector gpu_codes_scratch; + std::vector frame_codes_scratch; + std::vector frame_hidden_scratch; + std::vector semantic_ids_scratch; + std::vector song_logits_scratch; std::vector active_residual_ids_scratch; std::vector positions_scratch; }; @@ -528,17 +754,33 @@ MiniMaxMusic3DepthDecoderRuntime::MiniMaxMusic3DepthDecoderRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), global_token_embedding, execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3DepthDecoderRuntime::~MiniMaxMusic3DepthDecoderRuntime() = default; +std::vector MiniMaxMusic3DepthDecoderRuntime::generate_batch( + const std::vector & interleaved_hiddens, + int64_t songs, + const std::vector & semantic_codes, + float guidance_scale, + int64_t top_k, + const std::vector & seeds, + std::vector & sample_call_indices, + std::vector & rng_offset_blocks) { + return impl_->generate_batch( + interleaved_hiddens, songs, semantic_codes, guidance_scale, top_k, + seeds, sample_call_indices, rng_offset_blocks); +} + MiniMaxMusic3DepthCodes MiniMaxMusic3DepthDecoderRuntime::generate( const std::vector & last_hidden_cond, const std::vector & last_hidden_uncond, diff --git a/src/community_models/minimax_music3/flow_sampler.cpp b/src/community_models/minimax_music3/flow_sampler.cpp index 629527e9a..930c768b1 100644 --- a/src/community_models/minimax_music3/flow_sampler.cpp +++ b/src/community_models/minimax_music3/flow_sampler.cpp @@ -5,6 +5,8 @@ #include #include +#include +#include #include #include #include @@ -145,7 +147,21 @@ class MiniMaxMusic3FlowDenoiserRuntime final : public modules::FlowSamplerDenois frames_ = frames; channels_ = channels; overlap_ = overlap; + delta_cache_.clear(); flow_.prepare_chunk_condition(condition, frames); + static const bool warm_cond = getenv("MM3_DC_B1_WARMUP") != nullptr; + if (uncond_interval_ > 1 && warm_cond) { + const std::vector zero_latent(static_cast(channels * frames), 0.0F); + (void) flow_.predict_velocity_cond(zero_latent, condition, frames, 0.0F); + (void) flow_.predict_velocity_cond(zero_latent, condition, frames, 0.0F); + (void) flow_.predict_velocity_cond(zero_latent, condition, frames, 0.0F); + } + } + + void set_guidance_reuse(int64_t uncond_interval, int64_t uncond_warmup) { + uncond_interval_ = uncond_interval; + uncond_warmup_ = uncond_warmup; + delta_cache_.clear(); } void reset_sampler_caches(const std::vector & caches) override { @@ -188,16 +204,76 @@ class MiniMaxMusic3FlowDenoiserRuntime final : public modules::FlowSamplerDenois if (overlap_ > 0) { apply_overlap_prompt(denoiser_latent, input.state.schedule.t); } + const size_t branch_size = static_cast(channels_ * frames_); + const int64_t step = input.state.schedule.index; + const bool reuse_delta = + uncond_interval_ > 1 && + delta_cache_.size() == branch_size && + step >= uncond_warmup_ && + step + 1 < active_schedule_steps_ && + (step % uncond_interval_) != 0; + modules::FlowSamplerDenoiserOutput output; + if (reuse_delta) { + auto cond = flow_.predict_velocity_cond( + denoiser_latent, + *condition_, + frames_, + input.state.schedule.t); + if (cond.size() != branch_size) { + throw std::runtime_error("MiniMax Music 3 flow cond velocity shape mismatch"); + } + static const bool verify = getenv("MM3_DC_VERIFY") != nullptr; + if (verify) { + const auto cond_repeat = flow_.predict_velocity_cond( + denoiser_latent, *condition_, frames_, input.state.schedule.t); + const auto full = flow_.predict_velocity_branches( + denoiser_latent, *condition_, frames_, input.state.schedule.t); + float cond_diff = 0.0F; + float repeat_diff = 0.0F; + float synth_diff = 0.0F; + float cond_mag = 0.0F; + for (size_t i = 0; i < branch_size; ++i) { + cond_diff = std::max(cond_diff, std::fabs(cond[i] - full[i])); + repeat_diff = std::max(repeat_diff, std::fabs(cond[i] - cond_repeat[i])); + const float synth = cond[i] - delta_cache_[i]; + synth_diff = std::max(synth_diff, std::fabs(synth - full[branch_size + i])); + cond_mag = std::max(cond_mag, std::fabs(full[i])); + } + fprintf(stderr, + "MM3_DC_VERIFY step=%lld cond_maxdiff=%.6f b1_repeat_maxdiff=%.6f synth_uncond_maxdiff=%.6f cond_maxabs=%.6f\n", + static_cast(step), cond_diff, repeat_diff, synth_diff, cond_mag); + output.predictions.push_back({ + "cond", + std::vector(full.begin(), full.begin() + static_cast(branch_size)), + }); + output.predictions.push_back({ + "uncond", + std::vector(full.begin() + static_cast(branch_size), full.end()), + }); + return output; + } + std::vector uncond(branch_size); + for (size_t i = 0; i < branch_size; ++i) { + uncond[i] = cond[i] - delta_cache_[i]; + } + output.predictions.push_back({"cond", std::move(cond)}); + output.predictions.push_back({"uncond", std::move(uncond)}); + return output; + } const auto branches = flow_.predict_velocity_branches( denoiser_latent, *condition_, frames_, input.state.schedule.t); - const size_t branch_size = static_cast(channels_ * frames_); if (branches.size() != 2 * branch_size) { throw std::runtime_error("MiniMax Music 3 flow velocity shape mismatch"); } - modules::FlowSamplerDenoiserOutput output; + if (uncond_interval_ > 1) { + delta_cache_.resize(branch_size); + for (size_t i = 0; i < branch_size; ++i) { + delta_cache_[i] = branches[i] - branches[branch_size + i]; + } + } output.predictions.push_back({ "cond", std::vector(branches.begin(), branches.begin() + static_cast(branch_size)), @@ -233,10 +309,13 @@ class MiniMaxMusic3FlowDenoiserRuntime final : public modules::FlowSamplerDenois const std::vector * condition_ = nullptr; const std::vector * previous_latent_ = nullptr; std::vector noise_prompt_; + std::vector delta_cache_; int64_t frames_ = 0; int64_t channels_ = 0; int64_t overlap_ = 0; int64_t active_schedule_steps_ = 0; + int64_t uncond_interval_ = 1; + int64_t uncond_warmup_ = 2; }; class MiniMaxMusic3FlowUpdateRuntime final : public modules::FlowSamplerUpdateRuntime { @@ -299,26 +378,35 @@ struct MiniMaxMusic3FlowSamplerRuntime::Impl { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), flow(std::make_unique( assets, execution, graph_arena_bytes, weight_context_bytes, - storage_type)) { + storage_type, + evict_cuda_graph_cache_on_release)) { if (assets == nullptr) { throw std::runtime_error("MiniMax Music 3 flow sampler requires assets"); } } - void ensure_sampler(int64_t latent_values, int64_t steps, float guidance_scale) { + void ensure_sampler( + int64_t latent_values, + int64_t steps, + float guidance_scale, + int64_t uncond_interval, + int64_t uncond_warmup) { if (sampler != nullptr && denoiser != nullptr && updater != nullptr && sampler_latent_values == latent_values && sampler_steps == steps && sampler_guidance_scale == guidance_scale) { + denoiser->set_guidance_reuse(uncond_interval, uncond_warmup); return; } auto new_denoiser = std::make_unique(*flow); @@ -344,6 +432,7 @@ struct MiniMaxMusic3FlowSamplerRuntime::Impl { sampler_latent_values = latent_values; sampler_steps = steps; sampler_guidance_scale = guidance_scale; + denoiser->set_guidance_reuse(uncond_interval, uncond_warmup); } std::vector denoise_chunk( @@ -375,7 +464,9 @@ struct MiniMaxMusic3FlowSamplerRuntime::Impl { ensure_sampler( config.flow.in_channels * frames, request.num_inference_steps, - request.guidance_scale); + request.guidance_scale, + request.flow_uncond_interval, + request.flow_uncond_warmup); denoiser->set_chunk_inputs( chunk_condition, frames, @@ -395,8 +486,11 @@ struct MiniMaxMusic3FlowSamplerRuntime::Impl { if (overlap > 0) { copy_latent_prefix(latents, previous_latent, overlap, frames, config.flow.in_channels); } - const int64_t overlap_start = std::max(0, frames - 2 * config.overlap_latent_length); - const int64_t overlap_end = std::max(overlap_start, frames - config.overlap_latent_length); + const int64_t overlap_latent = request.flow_overlap_latent_length > 0 + ? request.flow_overlap_latent_length + : config.overlap_latent_length; + const int64_t overlap_start = std::max(0, frames - 2 * overlap_latent); + const int64_t overlap_end = std::max(overlap_start, frames - overlap_latent); carry_latent = latent_tail_window(latents, frames, config.flow.in_channels, overlap_start, overlap_end); carry_condition = condition_tail_window(chunk_condition, frames, config.flow.condition_dim, overlap_start, overlap_end); return latents; @@ -411,6 +505,7 @@ struct MiniMaxMusic3FlowSamplerRuntime::Impl { } std::shared_ptr assets; + bool evict_cuda_graph_cache_on_release = false; std::unique_ptr flow; MiniMaxMusic3FlowDenoiserRuntime * denoiser = nullptr; MiniMaxMusic3FlowUpdateRuntime * updater = nullptr; @@ -425,13 +520,15 @@ MiniMaxMusic3FlowSamplerRuntime::MiniMaxMusic3FlowSamplerRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3FlowSamplerRuntime::~MiniMaxMusic3FlowSamplerRuntime() = default; diff --git a/src/community_models/minimax_music3/flow_transformer.cpp b/src/community_models/minimax_music3/flow_transformer.cpp index d37d67ff2..6dafd58fc 100644 --- a/src/community_models/minimax_music3/flow_transformer.cpp +++ b/src/community_models/minimax_music3/flow_transformer.cpp @@ -12,6 +12,8 @@ #include #include +#include +#include #include #include #include @@ -126,82 +128,223 @@ core::TensorValue apply_partial_rope( const core::TensorValue & cos, const core::TensorValue & sin, int64_t rotary_dim) { - const auto rotary = modules::SliceModule({3, 0, rotary_dim}).build(ctx, input); - const auto rest = modules::SliceModule({3, rotary_dim, input.shape.dims[3] - rotary_dim}).build(ctx, input); - const auto rotated = modules::SplitRoPEModule({rotary_dim}).build(ctx, rotary, cos, sin); + auto rotary = modules::SliceModule({3, 0, rotary_dim}).build(ctx, input); + auto rest = modules::SliceModule({3, rotary_dim, input.shape.dims[3] - rotary_dim}).build(ctx, input); + auto cos_used = cos; + auto sin_used = sin; + if (input.shape.dims[0] == 1) { + // Break every view chain feeding the rope elementwise ops: at batch 1 + // the graph allocator produced buffer aliasing along these views + // (bitwise-reproducible call-to-call divergence, absent at batch 2 + // where ggml_repeat materializes the chain). + rotary = core::wrap_tensor(ggml_cont(ctx.ggml, rotary.tensor), rotary.shape, GGML_TYPE_F32); + rest = core::wrap_tensor(ggml_cont(ctx.ggml, rest.tensor), rest.shape, GGML_TYPE_F32); + // At batch 1 the rope tables match the sliced shape exactly, so the + // generic rope helper would feed the graph-input tensors straight + // into the elementwise chain (its repeat short-circuit). That path + // produced nondeterministic outputs on CUDA; materializing a copy + // restores the batch>1 behavior where ggml_repeat isolates the + // inputs from downstream scheduling. + cos_used = core::wrap_tensor(ggml_cont(ctx.ggml, cos.tensor), cos.shape, GGML_TYPE_F32); + sin_used = core::wrap_tensor(ggml_cont(ctx.ggml, sin.tensor), sin.shape, GGML_TYPE_F32); + } + const auto rotated = modules::SplitRoPEModule({rotary_dim}).build(ctx, rotary, cos_used, sin_used); return modules::ConcatModule({3}).build(ctx, rotated, rest); } } // namespace struct MiniMaxMusic3FlowTransformerRuntime::Impl { + // Two independently cached graphs share the weights: slot 0 evaluates the + // usual cond+uncond batch, slot 1 evaluates the cond branch alone for + // guidance-delta reuse steps. Keeping both alive avoids rebuilding when a + // sampling schedule alternates between them. + struct GraphSlot { + int64_t batch = 0; + int64_t latent_frames = 0; + std::unique_ptr, GgmlContextDeleter> ggml; + ggml_cgraph * graph = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr; + core::HostGraphPlan plan; + core::TensorValue latents; + core::TensorValue condition; + core::TensorValue time_features; + core::TensorValue rope_cos; + core::TensorValue rope_sin; + core::TensorValue rope_positions; + ggml_tensor * output = nullptr; + }; + Impl( std::shared_ptr input_assets, core::ExecutionContext & input_execution, size_t input_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), execution(input_execution), graph_arena_bytes(input_graph_arena_bytes), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), weights(load_flow_weights(*assets, execution, weight_context_bytes, storage_type)) { time_proj = assets->transformer_weights->require_f32( "time_proj.weight", {assets->config.flow.fourier_embedding_dim / 2, 1}); + + if (getenv("MM3_DC_B1_SELFTEST") != nullptr) { + run_b1_selftest(); + } + } + + void run_b1_selftest() { + const auto & config = assets->config.flow; + const int64_t frames = 344; + std::vector latents(static_cast(config.in_channels * frames)); + std::vector condition(static_cast(config.condition_dim * frames)); + for (size_t i = 0; i < latents.size(); ++i) { + latents[i] = std::sin(0.001F * static_cast(i)); + } + for (size_t i = 0; i < condition.size(); ++i) { + condition[i] = std::cos(0.0007F * static_cast(i)); + } + prepare_chunk_condition(condition, frames); + std::vector batch2_cond; + { + // batch-2 determinism control at full depth + std::vector runs2[2]; + for (int r = 0; r < 2; ++r) { + runs2[r] = predict_velocity(latents, condition, frames, 0.85F, 2); + } + float d2 = 0.0F; + for (size_t i = 0; i < runs2[0].size(); ++i) { + d2 = std::max(d2, std::fabs(runs2[0][i] - runs2[1][i])); + } + fprintf(stderr, "MM3_B1_SELFTEST batch2_control rep01=%.6f\n", d2); + batch2_cond.assign(runs2[0].begin(), runs2[0].begin() + static_cast(runs2[0].size() / 2)); + } + for (const char * skip : {"", "rope", "attn", "ffn", "rope,attn", "attn,ffn"}) { +#ifndef _WIN32 + if (skip[0] != '\0') setenv("MM3_DC_B1_SKIP", skip, 1); else unsetenv("MM3_DC_B1_SKIP"); +#endif + release_slot(slots[1]); + std::vector runsk[2]; + for (int r = 0; r < 2; ++r) { + runsk[r] = predict_velocity(latents, condition, frames, 0.85F, 1); + } + float dk = 0.0F; + for (size_t i = 0; i < runsk[0].size(); ++i) { + dk = std::max(dk, std::fabs(runsk[0][i] - runsk[1][i])); + } + float db2 = 0.0F; + if (skip[0] == '\0' && batch2_cond.size() == runsk[0].size()) { + for (size_t i = 0; i < runsk[0].size(); ++i) { + db2 = std::max(db2, std::fabs(runsk[0][i] - batch2_cond[i])); + } + } + fprintf(stderr, "MM3_B1_SELFTEST skip='%s' rep01=%.6f b1_vs_b2cond=%.6f\n", skip, dk, db2); + } +#ifndef _WIN32 + unsetenv("MM3_DC_B1_SKIP"); +#endif + const char * caps_env = getenv("MM3_DC_B1_SELFTEST"); + std::string caps = caps_env != nullptr && std::string(caps_env) != "1" ? caps_env : "1,2,4,8,16,36"; + size_t pos = 0; + while (pos <= caps.size()) { + size_t comma = caps.find(',', pos); + if (comma == std::string::npos) comma = caps.size(); + const std::string tok = caps.substr(pos, comma - pos); + pos = comma + 1; + if (tok.empty()) continue; +#ifndef _WIN32 + setenv("MM3_DC_B1_LAYERS", tok.c_str(), 1); +#endif + release_slot(slots[1]); + std::vector runs[3]; + for (int r = 0; r < 3; ++r) { + runs[r] = predict_velocity(latents, condition, frames, 0.85F, 1); + } + float d01 = 0.0F, d12 = 0.0F, mag = 0.0F; + for (size_t i = 0; i < runs[0].size(); ++i) { + d01 = std::max(d01, std::fabs(runs[0][i] - runs[1][i])); + d12 = std::max(d12, std::fabs(runs[1][i] - runs[2][i])); + mag = std::max(mag, std::fabs(runs[2][i])); + } + fprintf(stderr, "MM3_B1_SELFTEST layers=%s rep01=%.6f rep12=%.6f maxabs=%.6f\n", + tok.c_str(), d01, d12, mag); + } +#ifndef _WIN32 + unsetenv("MM3_DC_B1_LAYERS"); +#endif + release_slot(slots[1]); } ~Impl() { release_runtime_graphs(); } - void release_runtime_graphs() { - if (graph != nullptr) { - core::release_backend_graph_resources(execution.backend(), graph); + void release_slot(GraphSlot & slot) { + if (slot.graph != nullptr) { + core::release_backend_graph_resources( + execution.backend(), slot.graph, evict_cuda_graph_cache_on_release); } - graph = nullptr; - latents = {}; - condition = {}; - time_features = {}; - rope_cos = {}; - rope_sin = {}; - output = nullptr; - gallocr.reset(); - ggml.reset(); - plan.reset(); - latent_frames = 0; + slot = {}; + } + + void release_runtime_graphs() { + release_slot(slots[0]); + release_slot(slots[1]); + condition_frames = 0; rope_cos_table.clear(); rope_sin_table.clear(); } - void ensure_graph(int64_t frames) { - if (graph != nullptr && latent_frames == frames) { - return; + GraphSlot & ensure_graph(int64_t frames, int64_t batch) { + GraphSlot & slot = slots[batch == 2 ? 0 : 1]; + if (slot.graph != nullptr && slot.latent_frames == frames && slot.batch == batch) { + return slot; } - release_runtime_graphs(); + release_slot(slot); const auto & config = assets->config.flow; const int64_t inner = config.attention_heads * config.head_dim; const int64_t steps = frames + 1; const int64_t concat_channels = 2 * config.in_channels + config.condition_dim; ggml_init_params params{graph_arena_bytes, nullptr, true}; - ggml.reset(ggml_init(params)); - if (ggml == nullptr) { + slot.ggml.reset(ggml_init(params)); + if (slot.ggml == nullptr) { throw std::runtime_error("failed to initialize MiniMax Music 3 flow graph context"); } - core::ModuleBuildContext ctx{ggml.get(), "minimax_music3.flow", execution.backend_type()}; - latents = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, config.in_channels, frames})); - condition = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, config.condition_dim, frames})); - time_features = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, config.fourier_embedding_dim})); - rope_cos = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.attention_heads, config.rotary_dim / 2})); - rope_sin = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.attention_heads, config.rotary_dim / 2})); - ggml_set_input(latents.tensor); - ggml_set_input(condition.tensor); - ggml_set_input(time_features.tensor); - ggml_set_input(rope_cos.tensor); - ggml_set_input(rope_sin.tensor); - - auto zeros = core::wrap_tensor(ggml_scale(ctx.ggml, latents.tensor, 0.0F), latents.shape, GGML_TYPE_F32); - auto x = modules::ConcatModule({1}).build(ctx, latents, zeros); - x = modules::ConcatModule({1}).build(ctx, x, condition); + int64_t layer_cap = static_cast(weights.blocks.size()); + bool skip_rope = false; + bool skip_attn = false; + bool skip_ffn = false; + if (batch == 1) { + if (const char * cap_env = getenv("MM3_DC_B1_LAYERS")) { + layer_cap = std::min(layer_cap, std::max(0, atoll(cap_env))); + } + if (const char * skip_env = getenv("MM3_DC_B1_SKIP")) { + const std::string skips = skip_env; + skip_rope = skips.find("rope") != std::string::npos; + skip_attn = skips.find("attn") != std::string::npos; + skip_ffn = skips.find("ffn") != std::string::npos; + } + } + core::ModuleBuildContext ctx{slot.ggml.get(), "minimax_music3.flow", execution.backend_type()}; + slot.latents = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, config.in_channels, frames})); + slot.condition = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, config.condition_dim, frames})); + slot.time_features = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch, config.fourier_embedding_dim})); + slot.rope_cos = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.attention_heads, config.rotary_dim / 2})); + slot.rope_sin = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.attention_heads, config.rotary_dim / 2})); + slot.rope_positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(slot.latents.tensor); + ggml_set_input(slot.condition.tensor); + ggml_set_input(slot.time_features.tensor); + ggml_set_input(slot.rope_cos.tensor); + ggml_set_input(slot.rope_sin.tensor); + ggml_set_input(slot.rope_positions.tensor); + + auto zeros = core::wrap_tensor(ggml_scale(ctx.ggml, slot.latents.tensor, 0.0F), slot.latents.shape, GGML_TYPE_F32); + auto x = modules::ConcatModule({1}).build(ctx, slot.latents, zeros); + x = modules::ConcatModule({1}).build(ctx, x, slot.condition); auto conv = modules::Conv1dModule({concat_channels, concat_channels, 1, 1, 0, 1, false}).build( ctx, x, @@ -212,29 +355,56 @@ struct MiniMaxMusic3FlowTransformerRuntime::Impl { auto temb = modules::LinearModule({config.fourier_embedding_dim, inner, true}).build( ctx, - time_features, + slot.time_features, weights.time_linear_1); temb = modules::SiluModule().build(ctx, temb); temb = modules::LinearModule({inner, inner, true}).build(ctx, temb, weights.time_linear_2); - temb = core::reshape_tensor(ctx, temb, core::TensorShape::from_dims({2, 1, inner})); + temb = core::reshape_tensor(ctx, temb, core::TensorShape::from_dims({batch, 1, inner})); x = modules::ConcatModule({1}).build(ctx, temb, x); + int64_t built_layers = 0; for (const auto & block : weights.blocks) { + if (built_layers++ >= layer_cap) { + break; + } auto normed = modules::LayerNormModule({inner, 1.0e-5F, true, true, false}).build(ctx, x, block.norm1); auto q = modules::LinearModule({inner, inner, false}).build(ctx, normed, block.q); auto k = modules::LinearModule({inner, inner, false}).build(ctx, normed, block.k); auto v = modules::LinearModule({inner, inner, false}).build(ctx, normed, block.v); - q = core::reshape_tensor(ctx, q, core::TensorShape::from_dims({2, steps, config.attention_heads, config.head_dim})); - k = core::reshape_tensor(ctx, k, core::TensorShape::from_dims({2, steps, config.attention_heads, config.head_dim})); - v = core::reshape_tensor(ctx, v, core::TensorShape::from_dims({2, steps, config.attention_heads, config.head_dim})); - q = apply_partial_rope(ctx, q, rope_cos, rope_sin, config.rotary_dim); - k = apply_partial_rope(ctx, k, rope_cos, rope_sin, config.rotary_dim); + q = core::reshape_tensor(ctx, q, core::TensorShape::from_dims({batch, steps, config.attention_heads, config.head_dim})); + k = core::reshape_tensor(ctx, k, core::TensorShape::from_dims({batch, steps, config.attention_heads, config.head_dim})); + v = core::reshape_tensor(ctx, v, core::TensorShape::from_dims({batch, steps, config.attention_heads, config.head_dim})); + if (!skip_rope) { + static const bool legacy_rope = getenv("MM3_LEGACY_ROPE") != nullptr; + if (!legacy_rope) { + // Native partial NEOX rope from integer positions: one + // fused kernel per projection instead of the split-table + // slice/mul/concat chain (which also aliased buffers at + // batch 1), for every batch size. + q = core::wrap_tensor( + ggml_rope_ext(ctx.ggml, q.tensor, slot.rope_positions.tensor, nullptr, + static_cast(config.rotary_dim), GGML_ROPE_TYPE_NEOX, 0, + 10000.0F, 1.0F, 0.0F, 1.0F, 0.0F, 0.0F), + q.shape, GGML_TYPE_F32); + k = core::wrap_tensor( + ggml_rope_ext(ctx.ggml, k.tensor, slot.rope_positions.tensor, nullptr, + static_cast(config.rotary_dim), GGML_ROPE_TYPE_NEOX, 0, + 10000.0F, 1.0F, 0.0F, 1.0F, 0.0F, 0.0F), + k.shape, GGML_TYPE_F32); + } else { + q = apply_partial_rope(ctx, q, slot.rope_cos, slot.rope_sin, config.rotary_dim); + k = apply_partial_rope(ctx, k, slot.rope_cos, slot.rope_sin, config.rotary_dim); + } + } q = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); k = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); v = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + const bool flash_ok = + core::uses_ggml_cuda_or_hip_backend(execution.backend_type()) && + !(batch == 1 && getenv("MM3_DC_B1_EXPLICIT_ATTN") != nullptr); auto attn = modules::ScaledDotProductAttentionModule({ config.head_dim, - core::uses_ggml_cuda_or_hip_backend(execution.backend_type()) + flash_ok ? modules::ScaledDotProductAttentionLowering::FlashPreserveViews : modules::ScaledDotProductAttentionLowering::Explicit, GGML_PREC_F32, @@ -243,21 +413,37 @@ struct MiniMaxMusic3FlowTransformerRuntime::Impl { attn = core::reshape_tensor( ctx, core::ensure_backend_addressable_layout(ctx, attn), - core::TensorShape::from_dims({2, steps, inner})); + core::TensorShape::from_dims({batch, steps, inner})); attn = modules::LinearModule({inner, inner, false}).build(ctx, attn, block.out); - x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, attn.tensor), x.shape, GGML_TYPE_F32); + if (!skip_attn) { + x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, attn.tensor), x.shape, GGML_TYPE_F32); + } auto ffn = modules::LayerNormModule({inner, 1.0e-5F, true, true, false}).build(ctx, x, block.norm2); ffn = modules::LinearModule({inner, 2 * config.ff_inner_dim, true}).build(ctx, ffn, block.ff_in); - const auto gate_states = modules::SliceModule({2, 0, config.ff_inner_dim}).build(ctx, ffn); - auto gate = modules::SliceModule({2, config.ff_inner_dim, config.ff_inner_dim}).build(ctx, ffn); - gate = modules::SiluModule().build(ctx, gate); - auto gated = core::wrap_tensor( - ggml_mul(ctx.ggml, gate_states.tensor, gate.tensor), - gate_states.shape, - GGML_TYPE_F32); + static const bool legacy_glu = getenv("MM3_LEGACY_GLU") != nullptr; + core::TensorValue gated; + if (!legacy_glu) { + // Fused SwiGLU over the packed [states | gate] projection (our gate + // half feeds silu, i.e. the swapped ggml convention): one + // kernel instead of slice+silu+mul. + gated = core::wrap_tensor( + ggml_swiglu_swapped(ctx.ggml, ffn.tensor), + ffn.shape.with_last_dim(config.ff_inner_dim), + GGML_TYPE_F32); + } else { + const auto gate_states = modules::SliceModule({2, 0, config.ff_inner_dim}).build(ctx, ffn); + auto gate = modules::SliceModule({2, config.ff_inner_dim, config.ff_inner_dim}).build(ctx, ffn); + gate = modules::SiluModule().build(ctx, gate); + gated = core::wrap_tensor( + ggml_mul(ctx.ggml, gate_states.tensor, gate.tensor), + gate_states.shape, + GGML_TYPE_F32); + } gated = modules::LinearModule({config.ff_inner_dim, inner, true}).build(ctx, gated, block.ff_out); - x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, gated.tensor), x.shape, GGML_TYPE_F32); + if (!skip_ffn) { + x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, gated.tensor), x.shape, GGML_TYPE_F32); + } } x = modules::SliceModule({1, 1, frames}).build(ctx, x); @@ -268,22 +454,53 @@ struct MiniMaxMusic3FlowTransformerRuntime::Impl { x, weights.postprocess_conv); x = core::wrap_tensor(ggml_add(ctx.ggml, post.tensor, x.tensor), post.shape, GGML_TYPE_F32); - output = x.tensor; - graph = ggml_new_graph_custom(ggml.get(), 524288, false); - ggml_build_forward_expand(graph, output); - gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution.backend()))); - if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr.get(), graph) || - !ggml_gallocr_alloc_graph(gallocr.get(), graph)) { + slot.output = x.tensor; + slot.graph = ggml_new_graph_custom(slot.ggml.get(), 524288, false); + ggml_build_forward_expand(slot.graph, slot.output); + slot.gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution.backend()))); + if (slot.gallocr == nullptr || !ggml_gallocr_reserve(slot.gallocr.get(), slot.graph) || + !ggml_gallocr_alloc_graph(slot.gallocr.get(), slot.graph)) { throw std::runtime_error("failed to allocate MiniMax Music 3 flow graph"); } - core::prepare_host_graph_plan(execution, graph, plan); - latent_frames = frames; - rope_cos_table = split_rope_table(steps, config.attention_heads, config.rotary_dim, true); - rope_sin_table = split_rope_table(steps, config.attention_heads, config.rotary_dim, false); - core::write_tensor_f32(rope_cos, rope_cos_table); - core::write_tensor_f32(rope_sin, rope_sin_table); + core::prepare_host_graph_plan(execution, slot.graph, slot.plan); + slot.latent_frames = frames; + slot.batch = batch; + if (rope_cos_table.size() != static_cast(steps * config.attention_heads * config.rotary_dim / 2)) { + rope_cos_table = split_rope_table(steps, config.attention_heads, config.rotary_dim, true); + rope_sin_table = split_rope_table(steps, config.attention_heads, config.rotary_dim, false); + } + static const bool trace = getenv("MM3_DC_B1_TRACE") != nullptr; + if (trace) { + int rope_as_src = 0; + int rope_as_leaf = 0; + for (int n = 0; n < ggml_graph_n_nodes(slot.graph); ++n) { + ggml_tensor * node = ggml_graph_node(slot.graph, n); + for (int si = 0; si < GGML_MAX_SRC; ++si) { + if (node->src[si] == slot.rope_cos.tensor || node->src[si] == slot.rope_sin.tensor) { + ++rope_as_src; + } + } + } + fprintf(stderr, "MM3_B1_TRACE build batch=%lld frames=%lld nodes=%d rope_src=%d rope_leaf=%d rope_buf=%p rope_flags=%d latents_buf=%p\n", + (long long) batch, (long long) frames, + ggml_graph_n_nodes(slot.graph), rope_as_src, rope_as_leaf, + (void *) slot.rope_cos.tensor->buffer, slot.rope_cos.tensor->flags, + (void *) slot.latents.tensor->buffer); + } + if (slot.rope_cos.tensor->buffer != nullptr) { + core::write_tensor_f32(slot.rope_cos, rope_cos_table); + core::write_tensor_f32(slot.rope_sin, rope_sin_table); + } + if (slot.rope_positions.tensor->buffer != nullptr) { + std::vector positions(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + positions[static_cast(i)] = static_cast(i); + } + core::write_tensor_i32(slot.rope_positions, positions); + } latent_batch.resize(static_cast(2 * config.in_channels * frames)); time_batch.resize(static_cast(2 * config.fourier_embedding_dim)); + return slot; } void prepare_chunk_condition( @@ -303,57 +520,59 @@ struct MiniMaxMusic3FlowTransformerRuntime::Impl { condition_frames = frames; } - std::vector predict_velocity_branches( + std::vector predict_velocity( const std::vector & input_latents, const std::vector & input_condition, int64_t frames, - float timestep) { + float timestep, + int64_t batch) { + static const bool b1_as_batch2 = getenv("MM3_DC_B1_BATCH2") != nullptr; + if (batch == 1 && b1_as_batch2) { + auto both = predict_velocity(input_latents, input_condition, frames, timestep, 2); + both.resize(both.size() / 2); + return both; + } const auto & config = assets->config.flow; if (static_cast(input_latents.size()) != config.in_channels * frames || static_cast(input_condition.size()) != config.condition_dim * frames) { throw std::runtime_error("MiniMax Music 3 flow input shape mismatch"); } - ensure_graph(frames); + auto & slot = ensure_graph(frames, batch); if (condition_frames != frames || condition_batch.size() != static_cast(2 * config.condition_dim * frames)) { throw std::runtime_error("MiniMax Music 3 flow condition was not prepared for this chunk"); } std::copy(input_latents.begin(), input_latents.end(), latent_batch.begin()); - std::copy(input_latents.begin(), input_latents.end(), latent_batch.begin() + input_latents.size()); + if (batch == 2) { + std::copy(input_latents.begin(), input_latents.end(), latent_batch.begin() + input_latents.size()); + } const auto time_feature = fourier_embedding(time_proj, config.fourier_embedding_dim, timestep); std::copy(time_feature.begin(), time_feature.end(), time_batch.begin()); - std::copy(time_feature.begin(), time_feature.end(), time_batch.begin() + time_feature.size()); - core::write_tensor_f32(latents, latent_batch); - core::write_tensor_f32(condition, condition_batch); - core::write_tensor_f32(time_features, time_batch); - if (core::compute_graph(execution, graph, plan, "minimax_music3.flow") != GGML_STATUS_SUCCESS) { + if (batch == 2) { + std::copy(time_feature.begin(), time_feature.end(), time_batch.begin() + time_feature.size()); + } + core::write_tensor_f32(slot.latents, latent_batch.data(), static_cast(batch * config.in_channels * frames)); + core::write_tensor_f32(slot.condition, condition_batch.data(), static_cast(batch * config.condition_dim * frames)); + core::write_tensor_f32(slot.time_features, time_batch.data(), static_cast(batch * config.fourier_embedding_dim)); + if (core::compute_graph(execution, slot.graph, slot.plan, "minimax_music3.flow") != GGML_STATUS_SUCCESS) { throw std::runtime_error("MiniMax Music 3 flow graph compute failed"); } - return core::read_tensor_f32(output); + return core::read_tensor_f32(slot.output); } std::shared_ptr assets; core::ExecutionContext & execution; size_t graph_arena_bytes = 0; + bool evict_cuda_graph_cache_on_release = false; MiniMaxMusic3FlowWeights weights; std::vector time_proj; - int64_t latent_frames = 0; int64_t condition_frames = 0; std::vector latent_batch; std::vector condition_batch; std::vector time_batch; std::vector rope_cos_table; std::vector rope_sin_table; - std::unique_ptr, GgmlContextDeleter> ggml; - ggml_cgraph * graph = nullptr; - std::unique_ptr, GgmlGallocrDeleter> gallocr; - core::HostGraphPlan plan; - core::TensorValue latents; - core::TensorValue condition; - core::TensorValue time_features; - core::TensorValue rope_cos; - core::TensorValue rope_sin; - ggml_tensor * output = nullptr; + GraphSlot slots[2]; }; MiniMaxMusic3FlowTransformerRuntime::MiniMaxMusic3FlowTransformerRuntime( @@ -361,13 +580,15 @@ MiniMaxMusic3FlowTransformerRuntime::MiniMaxMusic3FlowTransformerRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3FlowTransformerRuntime::~MiniMaxMusic3FlowTransformerRuntime() = default; @@ -376,7 +597,15 @@ std::vector MiniMaxMusic3FlowTransformerRuntime::predict_velocity_branche const std::vector & condition, int64_t latent_frames, float timestep) { - return impl_->predict_velocity_branches(latents, condition, latent_frames, timestep); + return impl_->predict_velocity(latents, condition, latent_frames, timestep, 2); +} + +std::vector MiniMaxMusic3FlowTransformerRuntime::predict_velocity_cond( + const std::vector & latents, + const std::vector & condition, + int64_t latent_frames, + float timestep) { + return impl_->predict_velocity(latents, condition, latent_frames, timestep, 1); } void MiniMaxMusic3FlowTransformerRuntime::prepare_chunk_condition( diff --git a/src/community_models/minimax_music3/global_lm.cpp b/src/community_models/minimax_music3/global_lm.cpp index f47a2e09c..abdc823a1 100644 --- a/src/community_models/minimax_music3/global_lm.cpp +++ b/src/community_models/minimax_music3/global_lm.cpp @@ -96,8 +96,35 @@ modules::QwenDecoderLayerWeights load_qwen_layer( } // namespace +MiniMaxMusic3LmHeadLayout classify_minimax_music3_lm_head_shape( + const std::vector & shape, + int64_t vocab_size, + int64_t hidden_size) { + if (shape == std::vector{vocab_size, hidden_size}) { + return MiniMaxMusic3LmHeadLayout::FullVocab; + } + const int64_t compact_size = MiniMaxMusic3Prompt{}.semantic_vocab_size + 1; + if (shape == std::vector{compact_size, hidden_size}) { + return MiniMaxMusic3LmHeadLayout::SemanticCompactV1; + } + throw std::runtime_error( + "MiniMax Music 3 lm_head must be full-vocab [" + std::to_string(vocab_size) + "," + + std::to_string(hidden_size) + "] or semantic_compact_v1 [" + + std::to_string(compact_size) + "," + std::to_string(hidden_size) + "]"); +} + +int64_t minimax_music3_lm_head_output_size( + MiniMaxMusic3LmHeadLayout layout, + int64_t vocab_size) noexcept { + if (layout == MiniMaxMusic3LmHeadLayout::SemanticCompactV1) { + return MiniMaxMusic3Prompt{}.semantic_vocab_size + 1; + } + return vocab_size; +} + modules::QwenCausalDecodeRuntimeConfig make_minimax_music3_global_lm_runtime_config( const MiniMaxMusic3Config & config, + MiniMaxMusic3LmHeadLayout lm_head_layout, core::BackendType backend_type, size_t prefill_graph_arena_bytes, size_t decode_graph_arena_bytes) { @@ -127,7 +154,7 @@ modules::QwenCausalDecodeRuntimeConfig make_minimax_music3_global_lm_runtime_con backend_type == core::BackendType::Vulkan) { out.decoder.static_cache_type = GGML_TYPE_F16; } - out.decoder.logits_size = config.qwen.vocab_size; + out.decoder.logits_size = minimax_music3_lm_head_output_size(lm_head_layout, config.qwen.vocab_size); out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; out.decoder.lm_head_precision = GGML_PREC_DEFAULT; out.readback_round_type = GGML_TYPE_BF16; @@ -150,6 +177,11 @@ MiniMaxMusic3GlobalLMWeights load_minimax_music3_global_lm_weights( const auto & config = assets.config.qwen; const auto & source = *assets.language_model_weights; MiniMaxMusic3GlobalLMWeights out; + out.lm_head_layout = classify_minimax_music3_lm_head_shape( + source.require_metadata("lm_head.weight").shape, + config.vocab_size, + config.hidden_size); + const int64_t lm_head_rows = minimax_music3_lm_head_output_size(out.lm_head_layout, config.vocab_size); out.store = std::make_shared( execution.backend(), execution.backend_type(), @@ -171,7 +203,7 @@ MiniMaxMusic3GlobalLMWeights load_minimax_music3_global_lm_weights( source, "lm_head", storage_type, - config.vocab_size, + lm_head_rows, config.hidden_size, false); out.store->upload(); diff --git a/src/community_models/minimax_music3/pipeline.cpp b/src/community_models/minimax_music3/pipeline.cpp index 632f42c8e..e590f33d6 100644 --- a/src/community_models/minimax_music3/pipeline.cpp +++ b/src/community_models/minimax_music3/pipeline.cpp @@ -1,12 +1,21 @@ #include "engine/community_models/minimax_music3/pipeline.h" +#include "engine/framework/core/backend.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/sampling/hf_sampler.h" +#include + #include +#include #include +#include +#include +#include #include +#include #include +#include #include #include @@ -18,7 +27,14 @@ using Clock = std::chrono::steady_clock; constexpr int64_t kCropLeftLatent = 86; constexpr int64_t kCropRightLatent = 344 - kCropLeftLatent; -std::vector chunk_starts(int64_t frames, const MiniMaxMusic3Config & config) { +int64_t effective_chunk_hop(const MiniMaxMusic3Config & config, const MiniMaxMusic3Request & request) { + if (request.flow_chunk_hop_frames <= 0) { + return config.chunk_hop_frames; + } + return std::min(request.flow_chunk_hop_frames, config.chunk_frames); +} + +std::vector chunk_starts(int64_t frames, const MiniMaxMusic3Config & config, int64_t hop_frames) { if (frames <= 0) { throw std::runtime_error("MiniMax Music 3 requires positive AR frame count"); } @@ -26,43 +42,21 @@ std::vector chunk_starts(int64_t frames, const MiniMaxMusic3Config & co return {0}; } std::vector out; - for (int64_t start = 0; start < frames - config.chunk_hop_frames; start += config.chunk_hop_frames) { + for (int64_t start = 0; start < frames - hop_frames; start += hop_frames) { out.push_back(start); } return out; } -std::vector condition_slice( - const std::vector & condition, - int64_t start, - int64_t frames, - int64_t dim) { - if (start < 0 || frames <= 0 || dim <= 0 || - static_cast(condition.size()) < (start + frames) * dim) { - throw std::runtime_error("MiniMax Music 3 condition slice is out of range"); - } - return std::vector( - condition.begin() + static_cast(start * dim), - condition.begin() + static_cast((start + frames) * dim)); -} - -std::vector crop_interleaved_audio( - const runtime::AudioBuffer & audio, - int64_t left_samples, - int64_t right_samples) { - if (audio.channels != 2 || audio.sample_rate <= 0 || - audio.samples.size() % static_cast(audio.channels) != 0) { - throw std::runtime_error("MiniMax Music 3 vocoder returned invalid stereo audio"); - } - const int64_t frames = static_cast(audio.samples.size()) / audio.channels; - const int64_t start = std::min(std::max(0, left_samples), frames); - const int64_t end = std::max(start, frames - std::max(0, right_samples)); - std::vector out(static_cast((end - start) * audio.channels)); - std::copy( - audio.samples.begin() + static_cast(start * audio.channels), - audio.samples.begin() + static_cast(end * audio.channels), - out.begin()); - return out; +// Mirrors the condition encoder's output-frame computation exactly so chunk +// latent sizes (and therefore flow noise RNG offsets) can be planned before +// the encoder runs. +int64_t predicted_condition_frames(const MiniMaxMusic3Config & config, int64_t input_frames) { + return static_cast( + static_cast(input_frames) * static_cast(config.condition.output_sample_rate) / + static_cast(config.condition.input_sample_rate) * + static_cast(config.condition.input_hop_length) / + static_cast(config.condition.output_hop_length)); } struct DenoisedChunk { @@ -72,6 +66,25 @@ struct DenoisedChunk { } // namespace +void detail::append_cropped_interleaved_audio( + runtime::AudioBuffer & destination, + const runtime::AudioBuffer & chunk, + int64_t left_frames, + int64_t right_frames) { + if (destination.channels != 2 || destination.sample_rate <= 0 || + chunk.channels != destination.channels || chunk.sample_rate != destination.sample_rate || + chunk.samples.size() % static_cast(chunk.channels) != 0) { + throw std::runtime_error("MiniMax Music 3 chunk audio format mismatch"); + } + const int64_t frames = static_cast(chunk.samples.size()) / chunk.channels; + const int64_t start = std::min(std::max(0, left_frames), frames); + const int64_t end = std::max(start, frames - std::max(0, right_frames)); + destination.samples.insert( + destination.samples.end(), + chunk.samples.begin() + static_cast(start * chunk.channels), + chunk.samples.begin() + static_cast(end * chunk.channels)); +} + struct MiniMaxMusic3PipelineRuntime::Impl { Impl( core::ExecutionContext & input_execution, @@ -79,13 +92,15 @@ struct MiniMaxMusic3PipelineRuntime::Impl { size_t graph_arena_bytes, size_t weight_context_bytes, assets::TensorStorageType storage_type, - bool memory_saver) + bool memory_saver, + bool pipeline_overlap) : execution(input_execution), assets(std::move(input_assets)), graph_arena_bytes(graph_arena_bytes), weight_context_bytes(weight_context_bytes), storage_type(storage_type), memory_saver(memory_saver), + pipeline_overlap(pipeline_overlap), sampling_policy(sampling::resolve_torch_cuda_sampling_policy( execution.backend_type(), execution.config().device, @@ -95,7 +110,26 @@ struct MiniMaxMusic3PipelineRuntime::Impl { if (assets == nullptr) { throw std::runtime_error("MiniMax Music 3 pipeline requires assets"); } - if (!memory_saver) { + if (pipeline_overlap && memory_saver) { + throw std::runtime_error( + "MiniMax Music 3 pipeline_overlap requires mem_saver=false (all stages stay resident)"); + } + if (pipeline_overlap) { + // The AR stage runs a stream of short kernels that must not queue + // behind the denoise stream's long GEMM waves, so its lazily + // created streams get the highest CUDA priority. The env is read + // by the CUDA backend at stream creation: AR streams materialize + // during the AR weight upload below, denoise-context streams + // afterwards at default priority. + set_stream_priority_env("-5"); + ggml_backend_synchronize(execution.backend()); + set_stream_priority_env(nullptr); + ar = make_ar(); + overlap_execution = std::make_unique(execution.config()); + condition = make_condition(); + flow = make_flow(); + vocoder = make_vocoder(); + } else if (!memory_saver) { ar = make_ar(); condition = make_condition(); flow = make_flow(); @@ -103,10 +137,26 @@ struct MiniMaxMusic3PipelineRuntime::Impl { } } + static void set_stream_priority_env(const char * value) { +#ifndef _WIN32 + if (value != nullptr) { + setenv("GGML_CUDA_STREAM_PRIORITY", value, 1); + } else { + unsetenv("GGML_CUDA_STREAM_PRIORITY"); + } +#else + (void) value; +#endif + } + ~Impl() { release_runtime_graphs(); } + core::ExecutionContext & denoise_execution() { + return overlap_execution != nullptr ? *overlap_execution : execution; + } + runtime::AudioBuffer generate(const MiniMaxMusic3Request & request) { if (request.duration_sec <= 0.0) { throw std::runtime_error("MiniMax Music 3 duration_sec must be positive"); @@ -123,6 +173,81 @@ struct MiniMaxMusic3PipelineRuntime::Impl { const int64_t target_frames = std::min( assets->config.max_audio_frames, static_cast(request.duration_sec * static_cast(assets->config.frame_rate))); + // Overlap pays off only while the flow tail can hide under AR slack; + // on long requests the constant SM contention degrades AR far more + // than the tail saves (measured +70% AR at 60 s), so it is applied + // to short requests only. + constexpr int64_t kOverlapMaxFrames = 600; + if (pipeline_overlap && target_frames <= kOverlapMaxFrames && + request.flow_chunk_hop_frames <= 0) { + return generate_overlapped(request, target_frames); + } + if (pipeline_overlap) { + engine::debug::timing_log_scalar("minimax_music3.pipeline.overlap_skipped_long", 1.0); + } + return generate_sequential(request, target_frames); + } + + std::vector generate_ensemble( + const MiniMaxMusic3Request & request, + const std::vector & take_seeds) { + if (request.duration_sec <= 0.0) { + throw std::runtime_error("MiniMax Music 3 duration_sec must be positive"); + } + if (request.num_inference_steps <= 0) { + throw std::runtime_error("MiniMax Music 3 num_inference_steps must be positive"); + } + if (request.guidance_scale <= 0.0F || request.ar_guidance_scale <= 0.0F) { + throw std::runtime_error("MiniMax Music 3 guidance scales must be positive"); + } + if (request.top_k <= 0) { + throw std::runtime_error("MiniMax Music 3 top_k must be positive"); + } + const int64_t takes = static_cast(take_seeds.size()); + if (takes <= 0) { + throw std::runtime_error("MiniMax Music 3 ensemble requires at least one take seed"); + } + const int64_t target_frames = std::min( + assets->config.max_audio_frames, + static_cast(request.duration_sec * static_cast(assets->config.frame_rate))); + std::vector rng_offsets(static_cast(takes), 0); + std::vector> take_hiddens; + { + auto & ar_runtime = ensure_ar(); + take_hiddens = ar_runtime.generate_frame_hiddens_ensemble( + request, target_frames, take_seeds, rng_offsets, + request.ensemble_prefix_frames); + release_ar_after_phase(); + } + const int64_t hidden_frame_width = + assets->config.condition.condition_layers * assets->config.qwen.hidden_size; + std::vector out; + out.reserve(static_cast(takes)); + for (int64_t take = 0; take < takes; ++take) { + const auto & frame_hiddens = take_hiddens[static_cast(take)]; + const int64_t generated_frames = + static_cast(frame_hiddens.size()) / hidden_frame_width; + if (static_cast(frame_hiddens.size()) != generated_frames * hidden_frame_width) { + throw std::runtime_error("MiniMax Music 3 frame hidden shape mismatch"); + } + if (generated_frames <= 0) { + throw std::runtime_error("MiniMax Music 3 AR produced no frames"); + } + MiniMaxMusic3Request take_request = request; + // With an intro-lock prefix the flow noise is shared (master + // seed): the locked intro renders identically for every take, + // and post-fork divergence comes from the conditions themselves. + take_request.seed = request.ensemble_prefix_frames > 0 + ? take_seeds[0] + : take_seeds[static_cast(take)]; + out.push_back(denoise_and_vocode( + frame_hiddens, generated_frames, take_request, rng_offsets[static_cast(take)])); + } + return out; + } + + // The exact pre-existing sequential pipeline, byte-for-byte. + runtime::AudioBuffer generate_sequential(const MiniMaxMusic3Request & request, int64_t target_frames) { uint64_t rng_offset_blocks = 0; std::vector frame_hiddens; { @@ -130,24 +255,44 @@ struct MiniMaxMusic3PipelineRuntime::Impl { frame_hiddens = ar_runtime.generate_frame_hiddens(request, target_frames, rng_offset_blocks); release_ar_after_phase(); } + const int64_t hidden_frame_width = + assets->config.condition.condition_layers * assets->config.qwen.hidden_size; const int64_t generated_frames = - static_cast(frame_hiddens.size()) / - (assets->config.condition.condition_layers * assets->config.qwen.hidden_size); - if (static_cast(frame_hiddens.size()) != - generated_frames * assets->config.condition.condition_layers * assets->config.qwen.hidden_size) { + static_cast(frame_hiddens.size()) / hidden_frame_width; + if (static_cast(frame_hiddens.size()) != generated_frames * hidden_frame_width) { throw std::runtime_error("MiniMax Music 3 frame hidden shape mismatch"); } if (generated_frames <= 0) { throw std::runtime_error("MiniMax Music 3 AR produced no frames"); } + return denoise_and_vocode(frame_hiddens, generated_frames, request, rng_offset_blocks); + } - const auto denoise_start = Clock::now(); - const auto starts = chunk_starts(generated_frames, assets->config); + runtime::AudioBuffer denoise_and_vocode( + const std::vector & frame_hiddens, + int64_t generated_frames, + const MiniMaxMusic3Request & caller_request, + uint64_t rng_offset_blocks) { + const int64_t hidden_frame_width = + assets->config.condition.condition_layers * assets->config.qwen.hidden_size; + // Chunk geometry from the effective hop: kept = latents(hop), + // overlap = (chunk - kept)/2, left crop = overlap/2. Hop 100 yields + // the historical 86/258/172 unchanged. + const int64_t hop_frames = effective_chunk_hop(assets->config, caller_request); + const int64_t chunk_latents = predicted_condition_frames(assets->config, assets->config.chunk_frames); + const int64_t kept_latents = predicted_condition_frames(assets->config, hop_frames); + const int64_t overlap_latents = std::max(0, (chunk_latents - kept_latents) / 2); + const int64_t crop_left = overlap_latents / 2; + const int64_t crop_right = std::max(0, chunk_latents - crop_left - kept_latents); + MiniMaxMusic3Request request = caller_request; + request.flow_overlap_latent_length = overlap_latents; std::vector denoised; - denoised.reserve(starts.size()); - std::vector previous_latent; - std::vector previous_condition; + const auto denoise_start = Clock::now(); { + const auto starts = chunk_starts(generated_frames, assets->config, hop_frames); + denoised.reserve(starts.size()); + std::vector previous_latent; + std::vector previous_condition; auto & condition_runtime = ensure_condition(); auto & flow_runtime = ensure_flow(); for (size_t chunk_index = 0; chunk_index < starts.size(); ++chunk_index) { @@ -157,11 +302,8 @@ struct MiniMaxMusic3PipelineRuntime::Impl { int64_t condition_frames = 0; const auto condition_start = Clock::now(); auto condition_values = condition_runtime.encode( - condition_slice( - frame_hiddens, - start, - frame_count, - assets->config.condition.condition_layers * assets->config.qwen.hidden_size), + frame_hiddens.data() + static_cast(start * hidden_frame_width), + static_cast(frame_count * hidden_frame_width), frame_count, condition_frames); core::round_f32_to_bf16_in_place(condition_values); @@ -187,29 +329,57 @@ struct MiniMaxMusic3PipelineRuntime::Impl { rng_offset_blocks += sampling::torch_cuda_tensor_iterator_offset_blocks( static_cast(latents.size()), sampling_policy); - previous_latent = std::move(carry_latent); - previous_condition = std::move(carry_condition); + // Carry-free probe: with MM3_FLOW_NO_CARRY=1 every chunk + // denoises independently (the crop overlap still smooths the + // seams). If listening/panel accept this, chunks can be + // denoised in parallel across GPUs. + static const bool no_carry = [] { + const char * env = std::getenv("MM3_FLOW_NO_CARRY"); + return env != nullptr && env[0] == '1'; + }(); + if (no_carry) { + previous_latent.clear(); + previous_condition.clear(); + } else { + previous_latent = std::move(carry_latent); + previous_condition = std::move(carry_condition); + } denoised.push_back({std::move(latents), condition_frames}); } release_flow_after_phase(); } - std::vector chunks; - chunks.reserve(denoised.size()); + + runtime::AudioBuffer out; + out.sample_rate = assets->config.vocoder.sample_rate; + out.channels = 2; + size_t output_samples = 0; + for (size_t chunk_index = 0; chunk_index < denoised.size(); ++chunk_index) { + const int64_t estimated_decoded_frames = + denoised[chunk_index].latent_frames * assets->config.vocoder.hop_length; + const int64_t left = + chunk_index == 0 ? 0 : crop_left * assets->config.vocoder.hop_length; + const int64_t right = + chunk_index + 1 == denoised.size() + ? 0 + : crop_right * assets->config.vocoder.hop_length; + const int64_t estimated_kept_frames = + std::max(0, estimated_decoded_frames - left - right); + output_samples += static_cast(estimated_kept_frames * out.channels); + } + out.samples.reserve(output_samples); { auto & vocoder_runtime = ensure_vocoder(); for (size_t chunk_index = 0; chunk_index < denoised.size(); ++chunk_index) { + auto chunk = std::move(denoised[chunk_index]); const auto vocoder_start = Clock::now(); - auto audio = vocoder_runtime.decode( - denoised[chunk_index].latents, - denoised[chunk_index].latent_frames); + const auto audio = vocoder_runtime.decode(chunk.latents, chunk.latent_frames); engine::debug::timing_log_scalar( "minimax_music3.vocoder.total_ms", engine::debug::elapsed_ms(vocoder_start, Clock::now())); - const int64_t left = chunk_index == 0 ? 0 : kCropLeftLatent * assets->config.vocoder.hop_length; + const int64_t left = chunk_index == 0 ? 0 : crop_left * assets->config.vocoder.hop_length; const int64_t right = - chunk_index + 1 == denoised.size() ? 0 : kCropRightLatent * assets->config.vocoder.hop_length; - audio.samples = crop_interleaved_audio(audio, left, right); - chunks.push_back(std::move(audio)); + chunk_index + 1 == denoised.size() ? 0 : crop_right * assets->config.vocoder.hop_length; + detail::append_cropped_interleaved_audio(out, audio, left, right); } release_vocoder_after_phase(); } @@ -217,21 +387,196 @@ struct MiniMaxMusic3PipelineRuntime::Impl { "minimax_music3.flow_vocoder.total_ms", engine::debug::elapsed_ms(denoise_start, Clock::now())); - runtime::AudioBuffer out; - out.sample_rate = assets->config.vocoder.sample_rate; - out.channels = 2; - for (const auto & chunk : chunks) { - if (chunk.sample_rate != out.sample_rate || chunk.channels != out.channels) { - throw std::runtime_error("MiniMax Music 3 chunk audio format mismatch"); - } - out.samples.insert(out.samples.end(), chunk.samples.begin(), chunk.samples.end()); - } for (float & sample : out.samples) { sample = std::clamp(sample, -1.0F, 1.0F); } return out; } + // Overlapped pipeline: AR keeps decoding on the primary context while a + // worker thread runs condition+flow+vocoder for every chunk whose frames + // are already available, on a second backend context (its own CUDA + // stream). Flow noise offsets are planned from the deterministic per-frame + // RNG consumption; if the AR stage ends early (EOS) or consumption differs + // from the plan, the worker result is discarded and the exact sequential + // stage runs instead, so output bytes never depend on the overlap. + runtime::AudioBuffer generate_overlapped(const MiniMaxMusic3Request & request, int64_t target_frames) { + const auto & config = assets->config; + const int64_t hidden_frame_width = config.condition.condition_layers * config.qwen.hidden_size; + + struct ChunkPlan { + int64_t start = 0; + int64_t end = 0; + uint64_t noise_offset_blocks = 0; + }; + const uint64_t semantic_blocks = sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.qwen.vocab_size), sampling_policy); + const uint64_t depth_blocks = + sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.depth.audio_vocab_size), sampling_policy) * + static_cast(config.depth.codebooks - 1); + const uint64_t predicted_after_ar = + static_cast(target_frames + 1) * (semantic_blocks + depth_blocks); + + const auto starts = chunk_starts(target_frames, config, config.chunk_hop_frames); + std::vector plans(starts.size()); + { + uint64_t running = predicted_after_ar; + for (size_t chunk_index = 0; chunk_index < starts.size(); ++chunk_index) { + auto & plan = plans[chunk_index]; + plan.start = starts[chunk_index]; + plan.end = std::min(plan.start + config.chunk_frames, target_frames); + plan.noise_offset_blocks = running; + const int64_t latent_frames = predicted_condition_frames(config, plan.end - plan.start); + running += sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(config.flow.in_channels * latent_frames), + sampling_policy); + } + } + + std::vector frame_hiddens; + frame_hiddens.reserve(static_cast(target_frames * hidden_frame_width)); + // The AR stage appends into this reserved buffer without reallocating, + // so the base pointer is stable for the whole overlapped run; the + // worker reads only rows already published through rows_done. + const float * const hidden_rows = frame_hiddens.data(); + + std::mutex progress_mutex; + std::condition_variable progress_cv; + int64_t rows_done = 0; + bool ar_finished = false; + + std::atomic worker_valid{true}; + std::exception_ptr worker_error; + runtime::AudioBuffer worker_audio; + worker_audio.sample_rate = config.vocoder.sample_rate; + worker_audio.channels = 2; + + const auto denoise_start = Clock::now(); + std::thread worker([&]() { + try { + std::vector previous_latent; + std::vector previous_condition; + auto & condition_runtime = ensure_condition(); + auto & flow_runtime = ensure_flow(); + auto & vocoder_runtime = ensure_vocoder(); + for (size_t chunk_index = 0; chunk_index < plans.size(); ++chunk_index) { + const auto & plan = plans[chunk_index]; + { + std::unique_lock lock(progress_mutex); + progress_cv.wait(lock, [&] { return rows_done >= plan.end || ar_finished; }); + if (rows_done < plan.end) { + worker_valid.store(false, std::memory_order_release); + return; + } + } + const int64_t frame_count = plan.end - plan.start; + int64_t condition_frames = 0; + const auto condition_start = Clock::now(); + auto condition_values = condition_runtime.encode( + hidden_rows + static_cast(plan.start * hidden_frame_width), + static_cast(frame_count * hidden_frame_width), + frame_count, + condition_frames); + core::round_f32_to_bf16_in_place(condition_values); + engine::debug::timing_log_scalar( + "minimax_music3.condition.total_ms", + engine::debug::elapsed_ms(condition_start, Clock::now())); + std::vector carry_condition; + std::vector carry_latent; + const auto flow_start = Clock::now(); + auto latents = flow_runtime.denoise_chunk( + condition_values, + condition_frames, + previous_latent, + previous_condition, + request, + plan.noise_offset_blocks, + sampling_policy, + carry_condition, + carry_latent); + engine::debug::timing_log_scalar( + "minimax_music3.flow.total_ms", + engine::debug::elapsed_ms(flow_start, Clock::now())); + previous_latent = std::move(carry_latent); + previous_condition = std::move(carry_condition); + const auto vocoder_start = Clock::now(); + const auto audio = vocoder_runtime.decode(latents, condition_frames); + engine::debug::timing_log_scalar( + "minimax_music3.vocoder.total_ms", + engine::debug::elapsed_ms(vocoder_start, Clock::now())); + const int64_t left = + chunk_index == 0 ? 0 : kCropLeftLatent * config.vocoder.hop_length; + const int64_t right = + chunk_index + 1 == plans.size() ? 0 : kCropRightLatent * config.vocoder.hop_length; + detail::append_cropped_interleaved_audio(worker_audio, audio, left, right); + } + } catch (...) { + worker_error = std::current_exception(); + worker_valid.store(false, std::memory_order_release); + } + }); + + uint64_t rng_offset_blocks = 0; + std::exception_ptr ar_error; + try { + const std::function progress = [&](int64_t frame) { + { + std::lock_guard lock(progress_mutex); + rows_done = frame; + } + progress_cv.notify_one(); + }; + auto & ar_runtime = ensure_ar(); + ar_runtime.generate_frame_hiddens_into( + request, target_frames, rng_offset_blocks, frame_hiddens, &progress); + } catch (...) { + ar_error = std::current_exception(); + } + { + std::lock_guard lock(progress_mutex); + ar_finished = true; + } + progress_cv.notify_one(); + worker.join(); + if (ar_error != nullptr) { + std::rethrow_exception(ar_error); + } + release_ar_after_phase(); + + const int64_t generated_frames = + static_cast(frame_hiddens.size()) / hidden_frame_width; + if (static_cast(frame_hiddens.size()) != generated_frames * hidden_frame_width) { + throw std::runtime_error("MiniMax Music 3 frame hidden shape mismatch"); + } + if (generated_frames <= 0) { + throw std::runtime_error("MiniMax Music 3 AR produced no frames"); + } + + const bool plan_held = + worker_valid.load(std::memory_order_acquire) && + worker_error == nullptr && + generated_frames == target_frames && + rng_offset_blocks == predicted_after_ar; + if (!plan_held) { + if (worker_error != nullptr) { + std::rethrow_exception(worker_error); + } + engine::debug::timing_log_scalar("minimax_music3.pipeline.overlap_fallback", 1.0); + return denoise_and_vocode(frame_hiddens, generated_frames, request, rng_offset_blocks); + } + + release_flow_after_phase(); + release_vocoder_after_phase(); + engine::debug::timing_log_scalar( + "minimax_music3.flow_vocoder.total_ms", + engine::debug::elapsed_ms(denoise_start, Clock::now())); + for (float & sample : worker_audio.samples) { + sample = std::clamp(sample, -1.0F, 1.0F); + } + return worker_audio; + } + void release_runtime_graphs() { if (ar != nullptr) { ar->release_runtime_graphs(); @@ -282,6 +627,7 @@ struct MiniMaxMusic3PipelineRuntime::Impl { ar->release_runtime_graphs(); if (memory_saver) { ar.reset(); + core::trim_backend_pools(execution.backend()); } } @@ -295,6 +641,7 @@ struct MiniMaxMusic3PipelineRuntime::Impl { if (memory_saver) { flow.reset(); condition.reset(); + core::trim_backend_pools(denoise_execution().backend()); } } @@ -305,6 +652,7 @@ struct MiniMaxMusic3PipelineRuntime::Impl { vocoder->release_runtime_graphs(); if (memory_saver) { vocoder.reset(); + core::trim_backend_pools(denoise_execution().backend()); } } @@ -314,34 +662,38 @@ struct MiniMaxMusic3PipelineRuntime::Impl { execution, graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + memory_saver); } std::unique_ptr make_condition() { return std::make_unique( assets, - execution, + denoise_execution(), graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + memory_saver); } std::unique_ptr make_flow() { return std::make_unique( assets, - execution, + denoise_execution(), graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + memory_saver); } std::unique_ptr make_vocoder() { return std::make_unique( assets, - execution, + denoise_execution(), graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + memory_saver); } core::ExecutionContext & execution; @@ -350,6 +702,8 @@ struct MiniMaxMusic3PipelineRuntime::Impl { size_t weight_context_bytes = 0; assets::TensorStorageType storage_type = assets::TensorStorageType::Native; bool memory_saver = true; + bool pipeline_overlap = false; + std::unique_ptr overlap_execution; std::unique_ptr ar; std::unique_ptr condition; std::unique_ptr flow; @@ -363,14 +717,16 @@ MiniMaxMusic3PipelineRuntime::MiniMaxMusic3PipelineRuntime( size_t graph_arena_bytes, size_t weight_context_bytes, assets::TensorStorageType storage_type, - bool memory_saver) + bool memory_saver, + bool pipeline_overlap) : impl_(std::make_unique( execution, std::move(assets), graph_arena_bytes, weight_context_bytes, storage_type, - memory_saver)) {} + memory_saver, + pipeline_overlap)) {} MiniMaxMusic3PipelineRuntime::~MiniMaxMusic3PipelineRuntime() = default; @@ -378,6 +734,12 @@ runtime::AudioBuffer MiniMaxMusic3PipelineRuntime::generate(const MiniMaxMusic3R return impl_->generate(request); } +std::vector MiniMaxMusic3PipelineRuntime::generate_ensemble( + const MiniMaxMusic3Request & request, + const std::vector & take_seeds) { + return impl_->generate_ensemble(request, take_seeds); +} + void MiniMaxMusic3PipelineRuntime::release_runtime_graphs() { if (impl_ != nullptr) { impl_->release_runtime_graphs(); diff --git a/src/community_models/minimax_music3/session.cpp b/src/community_models/minimax_music3/session.cpp index 758c6a529..52c8cd393 100644 --- a/src/community_models/minimax_music3/session.cpp +++ b/src/community_models/minimax_music3/session.cpp @@ -140,15 +140,21 @@ runtime::ModelCliInterface minimax_music3_cli_interface() { {"ar_guidance_scale", "float", "Autoregressive semantic and depth CFG scale.", false, "1.5", "0.0"}, {"top_k", "int", "Top-k sampling for semantic and residual code sampling.", false, "50", "1"}, {"seed", "int", "Generation seed.", false, "0", "0"}, + {"flow_uncond_interval", "int", "Evaluate the flow unconditional CFG branch only every N-th step and reuse the cached guidance delta in between (measured flow -25..-40% at mel-L1 ~0.3 dB for 2-3). Default 1 keeps the exact reference trajectory.", false, "1", "1"}, + {"flow_uncond_warmup", "int", "Number of initial flow steps that always evaluate both CFG branches when delta reuse is enabled.", false, "2", "0"}, + {"ensemble_takes", "int", "Decode N independent takes of the same prompt in one batched AR pass (seeds seed..seed+N-1); flow and vocoder run per take. Outputs are returned as named audio (take_01..take_NN) for --out-dir.", false, "1", "1"}, + {"ensemble_prefix_frames", "int", "Intro-lock for ensembles: decode the first N frames once as a shared master trajectory (~25 frames per second), then fork the takes (take_01 continues the master exactly). 0 disables.", false, "0", "0"}, + {"flow_chunk_hop_frames", "int", "Flow chunk hop in AR frames (~25/sec). Default 0 keeps the model config (100, 50% chunk overlap). 150 measures flow ~-35% with consistently rederived crops/carry.", false, "0", "0"}, }; out.session_options = { {"minimax_music3.weight_type", "native|bf16|f16|q8_0|q4_0|q4_k", "Shared weight storage type.", false, "native"}, {"minimax_music3.language_model_gguf", "string", "Language model component GGUF file relative to the model root.", false, "language_model_q4_0.gguf"}, - {"minimax_music3.rvq_depth_decoder_gguf", "string", "RVQ depth decoder component GGUF file relative to the model root.", false, "rvq_depth_decoder_q8_0.gguf"}, + {"minimax_music3.rvq_depth_decoder_gguf", "string", "RVQ depth decoder component GGUF file relative to the model root. q4_k measures ~-24% depth-stage time at instrument-panel-clean quality.", false, "rvq_depth_decoder_q8_0.gguf"}, {"minimax_music3.flow_transformer_gguf", "string", "Flow transformer component GGUF file relative to the model root.", false, "transformer_q4_0.gguf"}, {"minimax_music3.graph_context_mb", "int", "Runtime graph arena size in MiB.", false, "32", "1"}, {"minimax_music3.weight_context_mb", "int", "Weight context size in MiB.", false, "32", "1"}, {"minimax_music3.mem_saver", "bool", "Load large generation stages only while they are needed to reduce peak VRAM.", false, "true"}, + {"minimax_music3.pipeline_overlap", "bool", "Overlap AR decoding with per-chunk condition/flow/vocoder work on a second backend stream. Requires mem_saver=false; output is validated to stay identical to the sequential pipeline (automatic sequential fallback otherwise).", false, "false"}, }; return out; } @@ -200,13 +206,18 @@ MiniMaxMusic3Session::MiniMaxMusic3Session( if (const auto value = runtime::find_option(this->options().options, {"minimax_music3.mem_saver"})) { memory_saver = runtime::parse_bool_option(*value, "minimax_music3.mem_saver"); } + bool pipeline_overlap = false; + if (const auto value = runtime::find_option(this->options().options, {"minimax_music3.pipeline_overlap"})) { + pipeline_overlap = runtime::parse_bool_option(*value, "minimax_music3.pipeline_overlap"); + } pipeline_ = std::make_unique( execution_context(), assets_, graph_arena_bytes, weight_context_bytes, weight_type, - memory_saver); + memory_saver, + pipeline_overlap); } MiniMaxMusic3Session::~MiniMaxMusic3Session() = default; @@ -258,6 +269,30 @@ MiniMaxMusic3Request MiniMaxMusic3Session::parse_request(const runtime::TaskRequ if (const auto value = runtime::parse_u64_option(request.options, {"seed"})) { out.seed = *value; } + if (const auto value = runtime::parse_i64_option(request.options, {"flow_uncond_interval"})) { + out.flow_uncond_interval = *value; + } + if (const auto value = runtime::parse_i64_option(request.options, {"flow_uncond_warmup"})) { + out.flow_uncond_warmup = *value; + } + if (const auto value = runtime::parse_i64_option(request.options, {"ensemble_takes"})) { + out.ensemble_takes = *value; + } + if (out.ensemble_takes < 1 || out.ensemble_takes > 16) { + throw std::runtime_error("MiniMax Music 3 ensemble_takes must be in [1, 16]"); + } + if (const auto value = runtime::parse_i64_option(request.options, {"ensemble_prefix_frames"})) { + out.ensemble_prefix_frames = *value; + } + if (out.ensemble_prefix_frames < 0) { + throw std::runtime_error("MiniMax Music 3 ensemble_prefix_frames must be non-negative"); + } + if (const auto value = runtime::parse_i64_option(request.options, {"flow_chunk_hop_frames"})) { + out.flow_chunk_hop_frames = *value; + } + if (out.flow_chunk_hop_frames < 0) { + throw std::runtime_error("MiniMax Music 3 flow_chunk_hop_frames must be non-negative"); + } return out; } @@ -266,7 +301,29 @@ runtime::TaskResult MiniMaxMusic3Session::run(const runtime::TaskRequest & reque const auto parsed = parse_request(request); const auto start = Clock::now(); runtime::TaskResult result; - result.audio_output = pipeline_->generate(parsed); + // Debug aid: route K=1 through the ensemble machinery to verify the + // batched path reproduces the plain path (it must at batch 2). + const char * force_env = std::getenv("MM3_ENSEMBLE_FORCE"); + const bool force_ensemble = force_env != nullptr && force_env[0] == '1'; + if (parsed.ensemble_takes > 1 || force_ensemble) { + std::vector take_seeds(static_cast(parsed.ensemble_takes)); + for (size_t take = 0; take < take_seeds.size(); ++take) { + take_seeds[take] = parsed.seed + static_cast(take); + } + auto takes = pipeline_->generate_ensemble(parsed, take_seeds); + result.audio_output = takes.front(); + for (size_t take = 0; take < takes.size(); ++take) { + runtime::NamedAudioBuffer named; + char id[16]; + std::snprintf(id, sizeof(id), "take_%02zu", take + 1); + named.id = id; + named.audio = std::move(takes[take]); + named.meta.emplace("seed", std::to_string(take_seeds[take])); + result.named_audio_outputs.push_back(std::move(named)); + } + } else { + result.audio_output = pipeline_->generate(parsed); + } engine::debug::timing_log_scalar( "session.wall_ms", engine::debug::elapsed_ms(start, Clock::now())); diff --git a/src/community_models/minimax_music3/vocoder.cpp b/src/community_models/minimax_music3/vocoder.cpp index 484bb8bc8..0e71ff7ef 100644 --- a/src/community_models/minimax_music3/vocoder.cpp +++ b/src/community_models/minimax_music3/vocoder.cpp @@ -186,10 +186,12 @@ struct MiniMaxMusic3VocoderRuntime::Impl { core::ExecutionContext & input_execution, size_t input_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool input_evict_cuda_graph_cache_on_release) : assets(std::move(input_assets)), execution(input_execution), graph_arena_bytes(input_graph_arena_bytes), + evict_cuda_graph_cache_on_release(input_evict_cuda_graph_cache_on_release), weights(load_vocoder_weights(*assets, execution, weight_context_bytes, storage_type)) {} ~Impl() { @@ -198,7 +200,8 @@ struct MiniMaxMusic3VocoderRuntime::Impl { void release_runtime_graphs() { if (graph != nullptr) { - core::release_backend_graph_resources(execution.backend(), graph); + core::release_backend_graph_resources( + execution.backend(), graph, evict_cuda_graph_cache_on_release); } graph = nullptr; input = {}; @@ -299,6 +302,7 @@ struct MiniMaxMusic3VocoderRuntime::Impl { std::shared_ptr assets; core::ExecutionContext & execution; size_t graph_arena_bytes = 0; + bool evict_cuda_graph_cache_on_release = false; Music3VocoderWeights weights; int64_t latent_frames = 0; std::unique_ptr, GgmlContextDeleter> ggml; @@ -314,13 +318,15 @@ MiniMaxMusic3VocoderRuntime::MiniMaxMusic3VocoderRuntime( core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + bool evict_cuda_graph_cache_on_release) : impl_(std::make_unique( std::move(assets), execution, graph_arena_bytes, weight_context_bytes, - storage_type)) {} + storage_type, + evict_cuda_graph_cache_on_release)) {} MiniMaxMusic3VocoderRuntime::~MiniMaxMusic3VocoderRuntime() = default; diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index c18eaf903..a897eec07 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -568,6 +568,10 @@ class QwenCausalDecodeRuntime::Impl { batched_decode_cache_.import_state(state); } + runtime::TransformerBatchedKVState export_batched_decode_state() const { + return batched_decode_cache_.export_state(); + } + void start_decode_embeddings_batched( const runtime::TransformerBatchedKVState & state, int64_t required_cache_steps) { @@ -1326,7 +1330,8 @@ class QwenCausalDecodeRuntime::Impl { void release_prefill_graph() { if (prefill_graph_ != nullptr) { - core::release_backend_graph_resources(backend_, prefill_graph_); + core::release_backend_graph_resources( + backend_, prefill_graph_, config_.evict_cuda_graph_cache_on_release); } if (prefill_gallocr_ != nullptr) { ggml_gallocr_free(prefill_gallocr_); @@ -1348,7 +1353,8 @@ class QwenCausalDecodeRuntime::Impl { void release_batched_prefill_graph() { if (batched_prefill_graph_ != nullptr) { - core::release_backend_graph_resources(backend_, batched_prefill_graph_); + core::release_backend_graph_resources( + backend_, batched_prefill_graph_, config_.evict_cuda_graph_cache_on_release); } if (batched_prefill_gallocr_ != nullptr) { ggml_gallocr_free(batched_prefill_gallocr_); @@ -1371,7 +1377,8 @@ class QwenCausalDecodeRuntime::Impl { void release_decode_graph() { if (decode_graph_ != nullptr) { - core::release_backend_graph_resources(backend_, decode_graph_); + core::release_backend_graph_resources( + backend_, decode_graph_, config_.evict_cuda_graph_cache_on_release); } if (decode_buffer_ != nullptr) { ggml_backend_buffer_free(decode_buffer_); @@ -1394,7 +1401,8 @@ class QwenCausalDecodeRuntime::Impl { void release_batched_decode_graph() { if (batched_decode_graph_ != nullptr) { - core::release_backend_graph_resources(backend_, batched_decode_graph_); + core::release_backend_graph_resources( + backend_, batched_decode_graph_, config_.evict_cuda_graph_cache_on_release); } if (batched_decode_buffer_ != nullptr) { ggml_backend_buffer_free(batched_decode_buffer_); @@ -1559,6 +1567,10 @@ QwenCausalDecodeStepResult QwenCausalDecodeRuntime::decode_embeddings_batched( return impl_->decode_embeddings_batched(embeddings, batch_size); } +runtime::TransformerBatchedKVState QwenCausalDecodeRuntime::export_batched_decode_state() const { + return impl_->export_batched_decode_state(); +} + int64_t QwenCausalDecodeRuntime::decode_cache_steps() const noexcept { return impl_->decode_cache_steps(); } diff --git a/src/framework/sampling/torch_random.cpp b/src/framework/sampling/torch_random.cpp index 9bcf0f47b..09456b4e3 100644 --- a/src/framework/sampling/torch_random.cpp +++ b/src/framework/sampling/torch_random.cpp @@ -4,6 +4,11 @@ #include "engine/framework/io/dynamic_library.h" #ifdef ENGINE_HAS_CUDA_TORCH_RANDOM #include "torch_random_cuda_runtime.h" + +#ifdef ENGINE_HAS_CUDA_TORCH_RANDOM +#include "ggml-backend.h" +#include "ggml-cuda.h" +#endif #endif #include @@ -437,6 +442,118 @@ uint64_t torch_cuda_tensor_iterator_offset_blocks( return ((total_elements - 1) / (stride * unroll_factor) + 1); } +bool torch_cuda_sample_topk_exponential_pairs_available() { +#ifdef ENGINE_HAS_CUDA_TORCH_RANDOM + return true; +#else + return false; +#endif +} + +void torch_cuda_sample_topk_exponential_pairs( + const void * device_logits_f32, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const uint64_t * seeds, + const uint64_t * offset_blocks, + uint64_t offset_step_blocks, + const TorchCudaSamplingPolicy & policy, + int32_t * out_codes) { +#ifdef ENGINE_HAS_CUDA_TORCH_RANDOM + detail::sample_topk_exponential_pairs_cuda( + device_logits_f32, songs, vocab, guidance_scale, top_k, + seeds, offset_blocks, offset_step_blocks, policy, out_codes); +#else + (void) device_logits_f32; + (void) songs; + (void) vocab; + (void) guidance_scale; + (void) top_k; + (void) seeds; + (void) offset_blocks; + (void) policy; + (void) out_codes; + throw std::runtime_error("torch CUDA top-k exponential sampler is unavailable in this build"); +#endif +} + +#ifdef ENGINE_HAS_CUDA_TORCH_RANDOM +#define ENGINE_TORCH_RANDOM_CUDA_ONLY(...) __VA_ARGS__ +#else +#define ENGINE_TORCH_RANDOM_CUDA_ONLY(...) \ + throw std::runtime_error("torch CUDA depth frame runtime is unavailable in this build") +#endif + +void * torch_cuda_backend_stream(void * ggml_backend) { +#ifdef ENGINE_HAS_CUDA_TORCH_RANDOM + return ggml_backend_cuda_get_stream(static_cast(ggml_backend)); +#else + (void) ggml_backend; + return nullptr; +#endif +} + +void torch_cuda_depth_frame_ensure(int64_t songs, int64_t levels, int64_t hidden_size, const TorchCudaSamplingPolicy & policy) { + (void) songs; (void) levels; (void) hidden_size; (void) policy; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_ensure_cuda(songs, levels, hidden_size, policy)); +} + +void torch_cuda_depth_frame_begin(const uint64_t * seeds, const uint64_t * offset_blocks, int64_t songs, void * stream) { + (void) seeds; (void) offset_blocks; (void) songs; (void) stream; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_begin_cuda(seeds, offset_blocks, songs, stream)); +} + +void torch_cuda_depth_frame_sample( + const void * device_logits_f32, + int64_t level_index, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const TorchCudaSamplingPolicy & policy, + void * stream) { + (void) device_logits_f32; (void) level_index; (void) songs; (void) vocab; + (void) guidance_scale; (void) top_k; (void) policy; (void) stream; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_sample_cuda( + device_logits_f32, level_index, songs, vocab, guidance_scale, top_k, policy, stream)); +} + +void torch_cuda_depth_frame_residual_fill( + void * residual_ids_i32, + int64_t previous_levels, + int64_t songs, + int64_t audio_vocab, + void * stream) { + (void) residual_ids_i32; (void) previous_levels; (void) songs; (void) audio_vocab; (void) stream; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_residual_fill_cuda( + residual_ids_i32, previous_levels, songs, audio_vocab, stream)); +} + +void torch_cuda_depth_frame_accumulate_hidden( + const void * hidden_f32, + int64_t level_index, + int64_t songs, + int64_t hidden_size, + void * stream) { + (void) hidden_f32; (void) level_index; (void) songs; (void) hidden_size; (void) stream; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_accumulate_hidden_cuda( + hidden_f32, level_index, songs, hidden_size, stream)); +} + +void torch_cuda_depth_frame_end( + int32_t * host_codes, + float * host_hidden, + int64_t levels, + int64_t songs, + int64_t hidden_size, + void * stream) { + (void) host_codes; (void) host_hidden; (void) levels; (void) songs; (void) hidden_size; (void) stream; + ENGINE_TORCH_RANDOM_CUDA_ONLY(detail::depth_frame_end_cuda( + host_codes, host_hidden, levels, songs, hidden_size, stream)); +} + float torch_cuda_tensor_iterator_exponential_element( uint64_t seed, uint64_t total_elements, diff --git a/src/framework/sampling/torch_random_cuda_runtime.cu b/src/framework/sampling/torch_random_cuda_runtime.cu index f321c230b..6156bccf2 100644 --- a/src/framework/sampling/torch_random_cuda_runtime.cu +++ b/src/framework/sampling/torch_random_cuda_runtime.cu @@ -145,8 +145,400 @@ __global__ void fill_tensor_iterator_randn_kernel( } } +__device__ float tensor_iterator_exponential_element_device( + uint64_t seed, + uint64_t index, + uint64_t offset_blocks, + uint64_t stride) { + constexpr uint64_t unroll_factor = 4; + const uint64_t chunk = index / stride; + const int component = static_cast(chunk % unroll_factor); + const uint64_t loop_index = chunk / unroll_factor; + const uint64_t sequence = index % stride; + const Philox4 random = philox_4x32_10( + Philox4{ + static_cast(offset_blocks + loop_index), + static_cast((offset_blocks + loop_index) >> 32U), + static_cast(sequence), + static_cast(sequence >> 32U), + }, + seed); + uint32_t value = random.x; + if (component == 1) { + value = random.y; + } else if (component == 2) { + value = random.z; + } else if (component == 3) { + value = random.w; + } + const float uniform = static_cast(value) * kInvTwoPow32 + (kInvTwoPow32 * 0.5F); + return -logf(uniform); +} + +// One block per song. Rows come in [cond_i; uncond_i] pairs. Reproduces the +// CPU chain bit-for-bit up to logf rounding: bf16-round both branches, mix in +// f32, top-k threshold (value of the k-th largest), then Gumbel-style ranking +// rank = exp(double(score - max)) / double(exponential(token)) with the +// smallest token winning ties, exactly like the sequential CPU scan. +__global__ void sample_topk_exponential_pairs_kernel( + const float * logits, + int64_t vocab, + float guidance, + int64_t top_k, + const uint64_t * seeds, + const uint64_t * offsets, + uint64_t offset_step_blocks, + uint64_t stride, + int32_t * out_codes) { + extern __shared__ float shared_scores[]; // [vocab original | vocab workspace] + float * original = shared_scores; + float * workspace = shared_scores + vocab; + __shared__ float reduce_val[256]; + __shared__ int reduce_idx[256]; + __shared__ double reduce_rank[256]; + __shared__ float shared_threshold; + __shared__ float shared_max; + + const int song = blockIdx.x; + const int tid = static_cast(threadIdx.x); + const int threads = static_cast(blockDim.x); + const float * cond = logits + static_cast(2 * song) * vocab; + const float * uncond = logits + static_cast(2 * song + 1) * vocab; + for (int64_t i = tid; i < vocab; i += threads) { + const float c = round_to_bfloat16(cond[i]); + const float u = round_to_bfloat16(uncond[i]); + const float mixed = u + (c - u) * guidance; + original[i] = mixed; + workspace[i] = mixed; + } + __syncthreads(); + + const int64_t keep = top_k < vocab ? (top_k > 0 ? top_k : vocab) : vocab; + for (int64_t extract = 0; extract < keep; ++extract) { + float local_best = -INFINITY; + int local_idx = -1; + for (int64_t i = tid; i < vocab; i += threads) { + const float value = workspace[i]; + if (value > local_best) { + local_best = value; + local_idx = static_cast(i); + } + } + reduce_val[tid] = local_best; + reduce_idx[tid] = local_idx; + __syncthreads(); + for (int step = threads / 2; step > 0; step >>= 1) { + if (tid < step) { + const bool take = reduce_val[tid + step] > reduce_val[tid] || + (reduce_val[tid + step] == reduce_val[tid] && reduce_idx[tid + step] >= 0 && + (reduce_idx[tid] < 0 || reduce_idx[tid + step] < reduce_idx[tid])); + if (take) { + reduce_val[tid] = reduce_val[tid + step]; + reduce_idx[tid] = reduce_idx[tid + step]; + } + } + __syncthreads(); + } + if (tid == 0) { + if (extract == 0) { + shared_max = reduce_val[0]; + } + shared_threshold = reduce_val[0]; + if (reduce_idx[0] >= 0) { + workspace[reduce_idx[0]] = -INFINITY; + } + } + __syncthreads(); + } + + const uint64_t seed = seeds[song]; + const uint64_t offset = offsets[song] + offset_step_blocks; + double local_rank = -1.0; + int local_token = -1; + for (int64_t i = tid; i < vocab; i += threads) { + const float value = original[i]; + if (isfinite(value) && value >= shared_threshold) { + const float exponential = tensor_iterator_exponential_element_device( + seed, static_cast(i), offset, stride); + const double weight = exp(static_cast(value - shared_max)); + const double rank = weight / static_cast(exponential); + if (local_token < 0 || rank > local_rank) { + local_rank = rank; + local_token = static_cast(i); + } + } + } + reduce_rank[tid] = local_rank; + reduce_idx[tid] = local_token; + __syncthreads(); + for (int step = threads / 2; step > 0; step >>= 1) { + if (tid < step) { + const double other = reduce_rank[tid + step]; + const int other_idx = reduce_idx[tid + step]; + const bool take = other_idx >= 0 && + (reduce_idx[tid] < 0 || other > reduce_rank[tid] || + (other == reduce_rank[tid] && other_idx < reduce_idx[tid])); + if (take) { + reduce_rank[tid] = other; + reduce_idx[tid] = other_idx; + } + } + __syncthreads(); + } + if (tid == 0) { + out_codes[song] = reduce_idx[0]; + } +} + +} // namespace + +void sample_topk_exponential_pairs_cuda( + const void * device_logits_f32, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const uint64_t * seeds, + const uint64_t * offset_blocks, + uint64_t offset_step_blocks, + const TorchCudaSamplingPolicy & policy, + int32_t * out_codes) { + if (device_logits_f32 == nullptr || songs <= 0 || vocab <= 0 || out_codes == nullptr) { + throw std::invalid_argument("torch CUDA top-k exponential sampler input is invalid"); + } + constexpr int64_t kMaxVocab = 6144; // 2 shared copies must fit in 48 KiB + if (vocab > kMaxVocab) { + throw std::invalid_argument("torch CUDA top-k exponential sampler vocab exceeds shared memory"); + } + if (policy.multiprocessor_count <= 0 || policy.max_threads_per_multiprocessor <= 0) { + throw std::invalid_argument("torch CUDA top-k exponential sampler requires CUDA device properties"); + } + const uint64_t stride = tensor_iterator_stride(static_cast(vocab), policy); + check_cuda(cudaSetDevice(policy.cuda_device_index), "cudaSetDevice"); + + constexpr int64_t kMaxSongs = 64; + if (songs > kMaxSongs) { + throw std::invalid_argument("torch CUDA top-k exponential sampler song count is too large"); + } + static thread_local uint64_t * device_seeds = nullptr; + static thread_local uint64_t * device_offsets = nullptr; + static thread_local int32_t * device_codes = nullptr; + if (device_seeds == nullptr) { + check_cuda(cudaMalloc(&device_seeds, kMaxSongs * sizeof(uint64_t)), "cudaMalloc sampler seeds"); + check_cuda(cudaMalloc(&device_offsets, kMaxSongs * sizeof(uint64_t)), "cudaMalloc sampler offsets"); + check_cuda(cudaMalloc(&device_codes, kMaxSongs * sizeof(int32_t)), "cudaMalloc sampler codes"); + } + // seeds/offsets are frame constants: callers pass them only on the first + // call of a frame (offset_step_blocks advances later calls on device). + if (seeds != nullptr && offset_blocks != nullptr) { + check_cuda( + cudaMemcpy(device_seeds, seeds, static_cast(songs) * sizeof(uint64_t), cudaMemcpyHostToDevice), + "cudaMemcpy sampler seeds"); + check_cuda( + cudaMemcpy(device_offsets, offset_blocks, static_cast(songs) * sizeof(uint64_t), cudaMemcpyHostToDevice), + "cudaMemcpy sampler offsets"); + } + const size_t shared_bytes = static_cast(vocab) * 2 * sizeof(float); + sample_topk_exponential_pairs_kernel<<(songs), 256, shared_bytes>>>( + static_cast(device_logits_f32), + vocab, + guidance_scale, + top_k, + device_seeds, + device_offsets, + offset_step_blocks, + stride, + device_codes); + check_cuda(cudaGetLastError(), "sample_topk_exponential_pairs_kernel"); + check_cuda( + cudaMemcpy(out_codes, device_codes, static_cast(songs) * sizeof(int32_t), cudaMemcpyDeviceToHost), + "cudaMemcpy sampler codes"); +} + +namespace { + +struct DepthFrameBuffers { + int32_t * codes = nullptr; // [levels][songs] + float * hidden = nullptr; // [levels][songs][hidden_size] + uint64_t * seeds = nullptr; // [songs] + uint64_t * offsets = nullptr; // [songs] + int64_t levels = 0; + int64_t songs = 0; + int64_t hidden_size = 0; +}; + +DepthFrameBuffers & depth_frame_buffers() { + static thread_local DepthFrameBuffers buffers; + return buffers; +} + +__global__ void depth_frame_residual_kernel( + int32_t * residual_ids, + const int32_t * frame_codes, + int64_t songs, + int64_t previous_levels, + int64_t audio_vocab) { + const int64_t total = 2 * songs * previous_levels; + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) { + return; + } + const int64_t row = index / previous_levels; + const int64_t previous = index % previous_levels; // 0-based: codebook previous+1 + const int64_t song = row / 2; + const int32_t code = frame_codes[previous * songs + song]; + residual_ids[index] = code + static_cast(previous) * static_cast(audio_vocab); +} + } // namespace +void depth_frame_ensure_cuda(int64_t songs, int64_t levels, int64_t hidden_size, const TorchCudaSamplingPolicy & policy) { + auto & buffers = depth_frame_buffers(); + if (buffers.songs >= songs && buffers.levels >= levels && buffers.hidden_size == hidden_size) { + return; + } + check_cuda(cudaSetDevice(policy.cuda_device_index), "cudaSetDevice"); + if (buffers.codes != nullptr) { + cudaFree(buffers.codes); + cudaFree(buffers.hidden); + cudaFree(buffers.seeds); + cudaFree(buffers.offsets); + } + buffers.levels = levels; + buffers.songs = songs; + buffers.hidden_size = hidden_size; + check_cuda(cudaMalloc(&buffers.codes, static_cast(levels * songs) * sizeof(int32_t)), "cudaMalloc frame codes"); + check_cuda( + cudaMalloc(&buffers.hidden, static_cast(levels * songs * hidden_size) * sizeof(float)), + "cudaMalloc frame hidden"); + check_cuda(cudaMalloc(&buffers.seeds, static_cast(songs) * sizeof(uint64_t)), "cudaMalloc frame seeds"); + check_cuda(cudaMalloc(&buffers.offsets, static_cast(songs) * sizeof(uint64_t)), "cudaMalloc frame offsets"); +} + +void depth_frame_begin_cuda( + const uint64_t * seeds, + const uint64_t * offset_blocks, + int64_t songs, + void * stream) { + auto & buffers = depth_frame_buffers(); + cudaStream_t cuda_stream = static_cast(stream); + check_cuda( + cudaMemcpyAsync(buffers.seeds, seeds, static_cast(songs) * sizeof(uint64_t), cudaMemcpyHostToDevice, cuda_stream), + "cudaMemcpyAsync frame seeds"); + check_cuda( + cudaMemcpyAsync(buffers.offsets, offset_blocks, static_cast(songs) * sizeof(uint64_t), cudaMemcpyHostToDevice, cuda_stream), + "cudaMemcpyAsync frame offsets"); +} + +void depth_frame_sample_cuda( + const void * device_logits_f32, + int64_t level_index, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const TorchCudaSamplingPolicy & policy, + void * stream) { + auto & buffers = depth_frame_buffers(); + constexpr int64_t kMaxVocab = 6144; + if (vocab > kMaxVocab || level_index >= buffers.levels || songs > buffers.songs) { + throw std::invalid_argument("torch CUDA depth frame sampler shape is invalid"); + } + const uint64_t stride = tensor_iterator_stride(static_cast(vocab), policy); + const uint64_t step_blocks = + static_cast(level_index) * + (((static_cast(vocab) - 1) / (stride * 4) + 1)); + const size_t shared_bytes = static_cast(vocab) * 2 * sizeof(float); + cudaStream_t cuda_stream = static_cast(stream); + sample_topk_exponential_pairs_kernel<<(songs), 256, shared_bytes, cuda_stream>>>( + static_cast(device_logits_f32), + vocab, + guidance_scale, + top_k, + buffers.seeds, + buffers.offsets, + step_blocks, + stride, + buffers.codes + level_index * buffers.songs); + check_cuda(cudaGetLastError(), "depth frame sample kernel"); +} + +void depth_frame_residual_fill_cuda( + void * residual_ids_i32, + int64_t previous_levels, + int64_t songs, + int64_t audio_vocab, + void * stream) { + auto & buffers = depth_frame_buffers(); + const int64_t total = 2 * songs * previous_levels; + const int threads = 128; + const int blocks = static_cast((total + threads - 1) / threads); + depth_frame_residual_kernel<<(stream)>>>( + static_cast(residual_ids_i32), + buffers.codes, + buffers.songs, + previous_levels, + audio_vocab); + check_cuda(cudaGetLastError(), "depth frame residual kernel"); +} + +void depth_frame_accumulate_hidden_cuda( + const void * hidden_f32, + int64_t level_index, + int64_t songs, + int64_t hidden_size, + void * stream) { + auto & buffers = depth_frame_buffers(); + // cond rows only: source pitch is two rows, destination is packed. + check_cuda( + cudaMemcpy2DAsync( + buffers.hidden + (level_index * buffers.songs) * hidden_size, + static_cast(hidden_size) * sizeof(float), + hidden_f32, + static_cast(2 * hidden_size) * sizeof(float), + static_cast(hidden_size) * sizeof(float), + static_cast(songs), + cudaMemcpyDeviceToDevice, + static_cast(stream)), + "cudaMemcpy2DAsync frame hidden"); +} + +void depth_frame_end_cuda( + int32_t * host_codes, + float * host_hidden, + int64_t levels, + int64_t songs, + int64_t hidden_size, + void * stream) { + auto & buffers = depth_frame_buffers(); + cudaStream_t cuda_stream = static_cast(stream); + // The frame code matrix is [levels][buffer_songs]; rows are contiguous per + // level, so a 2D copy trims to the active song count. + check_cuda( + cudaMemcpy2DAsync( + host_codes, + static_cast(songs) * sizeof(int32_t), + buffers.codes, + static_cast(buffers.songs) * sizeof(int32_t), + static_cast(songs) * sizeof(int32_t), + static_cast(levels), + cudaMemcpyDeviceToHost, + cuda_stream), + "cudaMemcpy2DAsync frame codes out"); + check_cuda( + cudaMemcpy2DAsync( + host_hidden, + static_cast(songs * hidden_size) * sizeof(float), + buffers.hidden, + static_cast(buffers.songs * hidden_size) * sizeof(float), + static_cast(songs * hidden_size) * sizeof(float), + static_cast(levels), + cudaMemcpyDeviceToHost, + cuda_stream), + "cudaMemcpy2DAsync frame hidden out"); + check_cuda(cudaStreamSynchronize(cuda_stream), "cudaStreamSynchronize depth frame"); +} + void fill_torch_cuda_tensor_iterator_randn_cuda( float * output, size_t count, diff --git a/src/framework/sampling/torch_random_cuda_runtime.h b/src/framework/sampling/torch_random_cuda_runtime.h index 4c930da80..2792bc0e7 100644 --- a/src/framework/sampling/torch_random_cuda_runtime.h +++ b/src/framework/sampling/torch_random_cuda_runtime.h @@ -7,6 +7,49 @@ namespace engine::sampling::detail { +void sample_topk_exponential_pairs_cuda( + const void * device_logits_f32, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const uint64_t * seeds, + const uint64_t * offset_blocks, + uint64_t offset_step_blocks, + const TorchCudaSamplingPolicy & policy, + int32_t * out_codes); + +void depth_frame_ensure_cuda(int64_t songs, int64_t levels, int64_t hidden_size, const TorchCudaSamplingPolicy & policy); +void depth_frame_begin_cuda(const uint64_t * seeds, const uint64_t * offset_blocks, int64_t songs, void * stream); +void depth_frame_sample_cuda( + const void * device_logits_f32, + int64_t level_index, + int64_t songs, + int64_t vocab, + float guidance_scale, + int64_t top_k, + const TorchCudaSamplingPolicy & policy, + void * stream); +void depth_frame_residual_fill_cuda( + void * residual_ids_i32, + int64_t previous_levels, + int64_t songs, + int64_t audio_vocab, + void * stream); +void depth_frame_accumulate_hidden_cuda( + const void * hidden_f32, + int64_t level_index, + int64_t songs, + int64_t hidden_size, + void * stream); +void depth_frame_end_cuda( + int32_t * host_codes, + float * host_hidden, + int64_t levels, + int64_t songs, + int64_t hidden_size, + void * stream); + void fill_torch_cuda_tensor_iterator_randn_cuda( float * output, size_t count, diff --git a/tests/minimax_music3/test_repack_lm_head_gguf.py b/tests/minimax_music3/test_repack_lm_head_gguf.py new file mode 100644 index 000000000..0ecb1e599 --- /dev/null +++ b/tests/minimax_music3/test_repack_lm_head_gguf.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import importlib.util +import unittest +from pathlib import Path + +import numpy as np + + +SCRIPT = Path(__file__).resolve().parents[2] / "scripts" / "minimax_music3" / "repack_lm_head_gguf.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("repack_lm_head_gguf", SCRIPT) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {SCRIPT}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class CompactLmHeadTest(unittest.TestCase): + def test_token_ids_match_native_compact_sampler_order(self) -> None: + module = load_module() + token_ids = module.compact_token_ids() + self.assertEqual(len(token_ids), 16_385) + self.assertEqual(token_ids[0], 151_675) + self.assertEqual(token_ids[-2], 168_058) + self.assertEqual(token_ids[-1], 151_670) + self.assertEqual(len(set(token_ids)), len(token_ids)) + + def test_repack_copies_quantized_rows_verbatim(self) -> None: + module = load_module() + rows = 200_000 + row_bytes = 3 + source = np.arange(rows * row_bytes, dtype=np.uint8).reshape(rows, row_bytes) + compact = module.compact_lm_head_rows(source) + expected_ids = module.compact_token_ids() + self.assertEqual(compact.shape, (16_385, row_bytes)) + np.testing.assert_array_equal(compact, source[expected_ids]) + self.assertTrue(compact.flags.c_contiguous) + + def test_repack_rejects_wrong_source_vocab(self) -> None: + module = load_module() + source = np.zeros((199_999, 3), dtype=np.uint8) + with self.assertRaisesRegex(ValueError, "200000"): + module.compact_lm_head_rows(source) + + def test_replace_logical_shape_preserves_hidden_width(self) -> None: + module = load_module() + shapes = { + "lm_head.weight": (200_000, 4_096), + "model.embed_tokens.weight": (200_000, 4_096), + } + replaced = module.compact_logical_shapes(shapes) + self.assertEqual(replaced["lm_head.weight"], (16_385, 4_096)) + self.assertEqual(replaced["model.embed_tokens.weight"], (200_000, 4_096)) + self.assertEqual(shapes["lm_head.weight"], (200_000, 4_096)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittests/test_minimax_music3_graph_release_policy.cpp b/tests/unittests/test_minimax_music3_graph_release_policy.cpp new file mode 100644 index 000000000..37214ebb6 --- /dev/null +++ b/tests/unittests/test_minimax_music3_graph_release_policy.cpp @@ -0,0 +1,85 @@ +#include "engine/community_models/minimax_music3/ar_runtime.h" +#include "engine/community_models/minimax_music3/condition_encoder.h" +#include "engine/community_models/minimax_music3/depth_decoder.h" +#include "engine/community_models/minimax_music3/flow_sampler.h" +#include "engine/community_models/minimax_music3/flow_transformer.h" +#include "engine/community_models/minimax_music3/vocoder.h" +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" + +#include +#include +#include + +namespace { + +using engine::assets::TensorStorageType; +using engine::core::ExecutionContext; +using engine::core::TensorValue; +using engine::models::minimax_music3::MiniMaxMusic3ArRuntime; +using engine::models::minimax_music3::MiniMaxMusic3Assets; +using engine::models::minimax_music3::MiniMaxMusic3ConditionEncoderRuntime; +using engine::models::minimax_music3::MiniMaxMusic3DepthDecoderRuntime; +using engine::models::minimax_music3::MiniMaxMusic3FlowSamplerRuntime; +using engine::models::minimax_music3::MiniMaxMusic3FlowTransformerRuntime; +using engine::models::minimax_music3::MiniMaxMusic3VocoderRuntime; +using engine::modules::QwenCausalDecodeRuntimeConfig; + +using Assets = std::shared_ptr; + +static_assert(std::is_same_v); +static_assert(std::is_constructible_v< + MiniMaxMusic3ArRuntime, + Assets, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); +static_assert(std::is_constructible_v< + MiniMaxMusic3DepthDecoderRuntime, + Assets, + TensorValue, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); +static_assert(std::is_constructible_v< + MiniMaxMusic3ConditionEncoderRuntime, + Assets, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); +static_assert(std::is_constructible_v< + MiniMaxMusic3FlowSamplerRuntime, + Assets, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); +static_assert(std::is_constructible_v< + MiniMaxMusic3FlowTransformerRuntime, + Assets, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); +static_assert(std::is_constructible_v< + MiniMaxMusic3VocoderRuntime, + Assets, + ExecutionContext &, + size_t, + size_t, + TensorStorageType, + bool>); + +} // namespace + +int main() { + std::cout << "minimax_music3_graph_release_policy_test: ok\\n"; + return 0; +} diff --git a/tests/unittests/test_minimax_music3_lm_head.cpp b/tests/unittests/test_minimax_music3_lm_head.cpp new file mode 100644 index 000000000..a8128bc86 --- /dev/null +++ b/tests/unittests/test_minimax_music3_lm_head.cpp @@ -0,0 +1,59 @@ +#include "engine/community_models/minimax_music3/global_lm.h" +#include "test_assert.h" + +#include +#include +#include +#include + +namespace { + +using engine::models::minimax_music3::MiniMaxMusic3LmHeadLayout; +using engine::models::minimax_music3::classify_minimax_music3_lm_head_shape; +using engine::models::minimax_music3::minimax_music3_lm_head_output_size; + +void test_full_vocab_layout() { + const auto layout = classify_minimax_music3_lm_head_shape({200000, 4096}, 200000, 4096); + engine::test::require(layout == MiniMaxMusic3LmHeadLayout::FullVocab, "full-vocab layout"); + engine::test::require_eq(minimax_music3_lm_head_output_size(layout, 200000), int64_t{200000}, "full rows"); +} + +void test_semantic_compact_layout() { + const auto layout = classify_minimax_music3_lm_head_shape({16385, 4096}, 200000, 4096); + engine::test::require( + layout == MiniMaxMusic3LmHeadLayout::SemanticCompactV1, + "semantic compact layout"); + engine::test::require_eq(minimax_music3_lm_head_output_size(layout, 200000), int64_t{16385}, "compact rows"); +} + +void test_rejects_ambiguous_or_wrong_shapes() { + for (const auto & shape : std::vector>{ + {16389, 4096}, + {16385, 2048}, + {200000, 2048}, + {16385}, + }) { + bool rejected = false; + try { + (void)classify_minimax_music3_lm_head_shape(shape, 200000, 4096); + } catch (const std::runtime_error &) { + rejected = true; + } + engine::test::require(rejected, "invalid MiniMax lm_head shape was accepted"); + } +} + +} // namespace + +int main() { + try { + test_full_vocab_layout(); + test_semantic_compact_layout(); + test_rejects_ambiguous_or_wrong_shapes(); + std::cout << "minimax_music3_lm_head_test: ok\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << "minimax_music3_lm_head_test: " << error.what() << '\n'; + return 1; + } +} diff --git a/tests/unittests/test_minimax_music3_pipeline_buffers.cpp b/tests/unittests/test_minimax_music3_pipeline_buffers.cpp new file mode 100644 index 000000000..4378b72dc --- /dev/null +++ b/tests/unittests/test_minimax_music3_pipeline_buffers.cpp @@ -0,0 +1,105 @@ +#include "engine/community_models/minimax_music3/condition_encoder.h" +#include "engine/community_models/minimax_music3/pipeline.h" +#include "test_assert.h" + +#include +#include +#include +#include +#include + +namespace { + +using engine::models::minimax_music3::detail::append_cropped_interleaved_audio; +using engine::models::minimax_music3::detail::project_frame_hiddens; + +void require_vector_eq( + const std::vector & actual, + const std::vector & expected, + const char * label) { + engine::test::require_eq(actual.size(), expected.size(), std::string(label) + " size"); + for (size_t index = 0; index < actual.size(); ++index) { + engine::test::require_eq( + actual[index], + expected[index], + std::string(label) + "[" + std::to_string(index) + "]"); + } +} + +void test_projects_an_offset_frame_window_without_repacking() { + const std::vector all_frames{ + 0.0F, 0.0F, 0.0F, + 4.0F, 4.0F, 4.0F, + 10.0F, 20.0F, 30.0F, + 14.0F, 24.0F, 34.0F, + 100.0F, 200.0F, 300.0F, + 104.0F, 204.0F, 304.0F, + }; + const std::vector layer_weights{0.25F, 0.75F}; + const float * window = all_frames.data() + 6; + + const auto projected = project_frame_hiddens( + window, + 12, + 2, + 2, + 3, + layer_weights); + + require_vector_eq(projected, {13.0F, 103.0F, 23.0F, 203.0F, 33.0F, 303.0F}, "projected"); +} + +void test_projection_rejects_a_truncated_window() { + const std::vector values(11, 0.0F); + bool rejected = false; + try { + (void)project_frame_hiddens(values.data(), values.size(), 2, 2, 3, {0.25F, 0.75F}); + } catch (const std::runtime_error &) { + rejected = true; + } + engine::test::require(rejected, "truncated frame-hidden window was accepted"); +} + +void test_appends_only_the_uncropped_stereo_frames() { + engine::runtime::AudioBuffer destination{44100, 2, {-1.0F, -10.0F}}; + const engine::runtime::AudioBuffer chunk{ + 44100, + 2, + {1.0F, 10.0F, 2.0F, 20.0F, 3.0F, 30.0F, 4.0F, 40.0F}, + }; + + append_cropped_interleaved_audio(destination, chunk, 1, 1); + + require_vector_eq( + destination.samples, + {-1.0F, -10.0F, 2.0F, 20.0F, 3.0F, 30.0F}, + "cropped append"); +} + +void test_append_rejects_mismatched_output_format() { + engine::runtime::AudioBuffer destination{48000, 2, {}}; + const engine::runtime::AudioBuffer chunk{44100, 2, {1.0F, 2.0F}}; + bool rejected = false; + try { + append_cropped_interleaved_audio(destination, chunk, 0, 0); + } catch (const std::runtime_error &) { + rejected = true; + } + engine::test::require(rejected, "mismatched chunk format was accepted"); +} + +} // namespace + +int main() { + try { + test_projects_an_offset_frame_window_without_repacking(); + test_projection_rejects_a_truncated_window(); + test_appends_only_the_uncropped_stereo_frames(); + test_append_rejects_mismatched_output_format(); + std::cout << "minimax_music3_pipeline_buffers_test: ok\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << "minimax_music3_pipeline_buffers_test: " << error.what() << '\n'; + return 1; + } +}