From 580e88d8b7dece7099d9b62323521d0254ff3615 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Mon, 31 Aug 2026 12:17:51 +0200 Subject: [PATCH 01/10] ci : add check for unzip (#28082) --- ci/run.sh | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ci/run.sh b/ci/run.sh index 1701bc7ed..3909283d1 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -732,6 +732,11 @@ function gg_check_build_requirements { gg_printf 'ctest not found, please install\n' exit 1 fi + + if ! command -v unzip &> /dev/null; then + gg_printf 'unzip not found, please install\n' + exit 1 + fi } function gg_run_test_backend_ops_cpu { From a32af33de2b5950e701578dc23a229e8e2c727b9 Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Mon, 31 Aug 2026 18:33:02 +0800 Subject: [PATCH 02/10] sycl : Enhance to get the free memory of Intel GPU (#27968) * enhance get mem info by l0 an SYCL API * remove debug code, format the code * update SYCL.md for GGML_SYCL_GET_MEM_API --- docs/backend/SYCL.md | 1 + ggml/src/ggml-sycl/base.hpp | 36 +++++++ ggml/src/ggml-sycl/common.hpp | 16 +-- ggml/src/ggml-sycl/ggml-sycl.cpp | 60 +++++++++--- ggml/src/ggml-sycl/mem.cpp | 162 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/mem.hpp | 16 +++ 6 files changed, 261 insertions(+), 30 deletions(-) create mode 100644 ggml/src/ggml-sycl/base.hpp create mode 100644 ggml/src/ggml-sycl/mem.cpp create mode 100644 ggml/src/ggml-sycl/mem.hpp diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index 8b68851ff..5554c1d18 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -796,6 +796,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) | | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. | | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU.| +| GGML_SYCL_GET_MEM_API | 0 (default) or 1 | Set to get memory info (free, total) by Level Zero or SYCL API:
0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.
1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return total size for free size.| | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).| | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. | | GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. | diff --git a/ggml/src/ggml-sycl/base.hpp b/ggml/src/ggml-sycl/base.hpp new file mode 100644 index 000000000..3afd57ccb --- /dev/null +++ b/ggml/src/ggml-sycl/base.hpp @@ -0,0 +1,36 @@ +#ifndef GGML_SYCL_BASE_HPP +#define GGML_SYCL_BASE_HPP + +/** + * Module: base + * + * Description: + * Provides zero-dependency, foundational primitives, core abstractions, + * and low-level system interfaces. This module acts as the lowest layer + * of the architecture and is consumed globally across all subsystems. + * + * Constraints: + * - STRICTLY zero upstream dependencies (leaf module). + * - High stability and backward compatibility required. + */ + +#include + +extern int g_ggml_sycl_debug; + +#if defined(__clang__) && __has_builtin(__builtin_expect) +// Hint the optimizer to pipeline the more likely following instruction in branches +# define LIKELY(expr) __builtin_expect(expr, true) +# define UNLIKELY(expr) __builtin_expect(expr, false) +#else +# define LIKELY(expr) (expr) +# define UNLIKELY(expr) (expr) +#endif + +#define GGML_SYCL_DEBUG(...) \ + do { \ + if (UNLIKELY(g_ggml_sycl_debug)) \ + fprintf(stderr, __VA_ARGS__); \ + } while (0) + +#endif // GGML_SYCL_BASE_HPP diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index 34de284d8..b5f75ca54 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -18,6 +18,7 @@ #include #include +#include "base.hpp" #include "dpct/helper.hpp" #include "ggml.h" #include "ggml-impl.h" @@ -69,21 +70,6 @@ extern int g_ggml_sycl_fa_onednn; extern int g_ggml_sycl_fa_onednn_max_kv; -#if defined(__clang__) && __has_builtin(__builtin_expect) -// Hint the optimizer to pipeline the more likely following instruction in branches -# define LIKELY(expr) __builtin_expect(expr, true) -# define UNLIKELY(expr) __builtin_expect(expr, false) -#else -# define LIKELY(expr) (expr) -# define UNLIKELY(expr) (expr) -#endif - -#define GGML_SYCL_DEBUG(...) \ - do { \ - if (UNLIKELY(g_ggml_sycl_debug)) \ - fprintf(stderr, __VA_ARGS__); \ - } while (0) - #define CHECK_TRY_ERROR(expr) \ [&]() { \ try { \ diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 2b2d26cf2..290fb4676 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -35,6 +35,7 @@ #include #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API #include +#include #endif #if defined(GGML_SYCL_GRAPH) && SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC # include @@ -61,6 +62,7 @@ #include "ggml-sycl/fwht.hpp" #include "ggml-sycl/gemm.hpp" #include "ggml-sycl/getrows.hpp" +#include "ggml-sycl/mem.hpp" #include "ggml-sycl/norm.hpp" #include "ggml-sycl/presets.hpp" #include "ggml-sycl/quantize.hpp" @@ -105,6 +107,8 @@ int g_ggml_sycl_enable_flash_attention = 1; int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL; int g_ggml_sycl_usm_system = 0; int g_ggml_sycl_enable_host_pinned_mem = 1; +int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO; + static ggml_sycl_device_info ggml_sycl_init() { ggml_sycl_device_info info = {}; @@ -301,10 +305,27 @@ static const char* dev2dev_int2str(int dev2dev) { } } +/* +* There are several entry APIs to be called as first function in SYCL backend in different cases. +* It's the first internal function to be called by them in SYCL backend. +* This function is used to do initialize work for the SYCL backend and set the global variables. +*/ +void initialize_sycl_begining() { +#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API + ze_result_t zes_init = zesInit(0); + if (zes_init != ZE_RESULT_SUCCESS) { + std::cerr << "Warning: zesInit failed [ggml_check_sycl] with code " << static_cast(zes_init) + << ". Sysman free-memory query may be unavailable.\n"; + } +#endif +} + static void ggml_check_sycl() try { static bool initialized = false; if (!initialized) { + initialize_sycl_begining(); + g_ggml_sycl_debug = ggml_sycl_get_env("GGML_SYCL_DEBUG", 0); g_ggml_sycl_enable_optimize = ggml_sycl_get_env("GGML_SYCL_ENABLE_OPT", 1); g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0); @@ -317,8 +338,11 @@ static void ggml_check_sycl() try { g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0); g_ggml_sycl_dev2dev_memcpy = ggml_sycl_get_env("GGML_SYCL_DEV2DEV_MEMCPY", DEV2DEV_MEMCPY_SYCL); + g_ggml_sycl_get_mem_api = ggml_sycl_get_env("GGML_SYCL_GET_MEM_API", MEMORY_API_TYPE_LEVEL_ZERO); + if (g_ggml_sycl_use_level_zero_api == 0) { g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL; + g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_SYCL; } #ifdef SYCL_FLASH_ATTN @@ -331,6 +355,7 @@ static void ggml_check_sycl() try { g_ggml_sycl_enable_host_pinned_mem = ggml_sycl_get_env("GGML_SYCL_ENABLE_HOST_PINNED_MEM", 1); + GGML_SYCL_DEBUG("[SYCL] call ggml_check_sycl\n"); GGML_LOG_INFO("Build with Macros:\n"); @@ -374,9 +399,12 @@ static void ggml_check_sycl() try { #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API GGML_LOG_INFO(" GGML_SYCL_DEV2DEV_MEMCPY: %d (%s)\n", g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy)); + GGML_LOG_INFO(" GGML_SYCL_GET_MEM_API: %d (%s)\n", g_ggml_sycl_get_mem_api, mem_api_int2str(g_ggml_sycl_get_mem_api)); #else GGML_LOG_INFO(" GGML_SYCL_DEV2DEV_MEMCPY: %d (%s), enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n", g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy)); + GGML_LOG_INFO(" GGML_SYCL_GET_MEM_API: %d (%s), enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n", + g_ggml_sycl_get_mem_api, mem_api_int2str(g_ggml_sycl_get_mem_api)); #endif #if defined(GGML_SYCL_DNNL) @@ -5208,6 +5236,7 @@ catch (sycl::exception const &exc) { static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct ggml_tensor * dst) try { if (!g_sycl_loaded) return false; + initialize_sycl_begining(); if (dst->src[0] != nullptr && ggml_backend_buffer_is_sycl_split(dst->src[0]->buffer)) { ggml_sycl_set_peer_access(dst->src[1]->ne[1], ctx.device); @@ -5590,18 +5619,16 @@ catch (sycl::exception const &exc) { std::exit(1); } -void ggml_backend_sycl_get_device_memory(int device, size_t *free, - size_t *total) try { +void ggml_backend_sycl_get_device_memory(int device, size_t * free, size_t * total) try { GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_get_device_memory\n"); - ggml_sycl_set_device(device); - - SYCL_CHECK(CHECK_TRY_ERROR( - dpct::dev_mgr::instance().get_device(device).get_memory_info(*free, *total))); -} -catch (sycl::exception const &exc) { - std::cerr << exc.what() << "Exception caught at file:" << __FILE__ - << ", line:" << __LINE__ << std::endl; - std::exit(1); + bool res = get_memory_size(dpct::dev_mgr::instance().get_device(device), *free, *total, + (MemoryAPIType) g_ggml_sycl_get_mem_api); + if (!res) { + GGML_ABORT("[%s] failed to get device memory size", __func__); + } +} catch (const sycl::exception & exc) { + std::cerr << exc.what() << "Exception caught at file:" << __FILE__ << ", line:" << __LINE__ << std::endl; + std::exit(1); } //////////////////////////////////////////////////////////////////////////////// @@ -6020,10 +6047,12 @@ static const char * ggml_backend_sycl_device_get_description(ggml_backend_dev_t } static void ggml_backend_sycl_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { - ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *)dev->context; - ggml_sycl_set_device(ctx->device); - SYCL_CHECK(CHECK_TRY_ERROR( - dpct::dev_mgr::instance().get_device(ctx->device).get_memory_info(*free, *total))); + ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *) dev->context; + bool res = get_memory_size(dpct::dev_mgr::instance().get_device(ctx->device), *free, *total, + (MemoryAPIType) g_ggml_sycl_get_mem_api); + if (!res) { + GGML_ABORT("[%s] failed to get device memory size", __func__); + } } static enum ggml_backend_dev_type ggml_backend_sycl_device_get_type(ggml_backend_dev_t dev) { @@ -6906,6 +6935,7 @@ ggml_backend_reg_t ggml_backend_sycl_reg() { static std::mutex mutex; std::lock_guard lock(mutex); if (!initialized) { + initialize_sycl_begining(); ggml_backend_sycl_reg_context * ctx = new ggml_backend_sycl_reg_context; const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; diff --git a/ggml/src/ggml-sycl/mem.cpp b/ggml/src/ggml-sycl/mem.cpp new file mode 100644 index 000000000..5ec466420 --- /dev/null +++ b/ggml/src/ggml-sycl/mem.cpp @@ -0,0 +1,162 @@ +#include +#include + +#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API +#include +#include +#endif + +#include "base.hpp" +#include "mem.hpp" + +#include +#include +#include + +const char * mem_api_int2str(int mem_api) { + if (mem_api == MEMORY_API_TYPE_SYCL) { + return "SYCL API"; + } else if (mem_api == MEMORY_API_TYPE_LEVEL_ZERO) { + return "Level Zero API"; + } else { + return "Unknown"; + } +} + +#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API +bool query_free_memory_by_ze(sycl::device dev, size_t & free_bytes, size_t & total_bytes) { + free_bytes = 0; + total_bytes = 0; + + uint32_t module_count = 0; + +#if defined(SYCL_EXT_ONEAPI_BACKEND_LEVEL_ZERO) + constexpr sycl::backend kL0Backend = sycl::backend::ext_oneapi_level_zero; +#else + constexpr sycl::backend kL0Backend = sycl::backend::level_zero; +#endif + + try { + ze_result_t zes_init = zesInit(0); + if (zes_init != ZE_RESULT_SUCCESS) { + std::cerr << "Warning: zesInit failed with code " << static_cast(zes_init) + << ". Sysman free-memory query may be unavailable.\n"; + } + + if (dev.get_platform().get_backend() != kL0Backend) { + GGML_SYCL_DEBUG("Device backend is not Level Zero; falling back to SYCL memory query.\n"); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } + + ze_device_handle_t ze_dev = sycl::get_native(dev); + if (ze_dev == nullptr) { + GGML_SYCL_DEBUG("Level Zero device handle is null; falling back to SYCL memory query.\n"); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } + + ze_result_t r = zesDeviceEnumMemoryModules(ze_dev, &module_count, nullptr); + if (r != ZE_RESULT_SUCCESS || module_count == 0) { + GGML_SYCL_DEBUG("Failed to enumerate Level Zero memory modules. Falling back to SYCL memory query.\n"); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } + + std::vector modules(module_count); + r = zesDeviceEnumMemoryModules(ze_dev, &module_count, modules.data()); + if (r != ZE_RESULT_SUCCESS || module_count == 0) { + GGML_SYCL_DEBUG("Failed to enumerate Level Zero memory modules. Falling back to SYCL memory query.\n"); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } + + for (uint32_t i = 0; i < module_count; ++i) { + zes_mem_state_t state = {}; + state.stype = ZES_STRUCTURE_TYPE_MEM_STATE; + state.pNext = nullptr; + + r = zesMemoryGetState(modules[i], &state); + if (r != ZE_RESULT_SUCCESS) { + continue; + } + + free_bytes += state.free; + total_bytes += state.size; + } + + if (total_bytes == 0) { + GGML_SYCL_DEBUG("Level Zero memory query returned zero total bytes. Falling back to SYCL memory query.\n"); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } + return true; + } catch (const sycl::exception & e) { + GGML_SYCL_DEBUG("Level Zero memory query failed: %s\n", e.what()); + total_bytes = dev.get_info(); + free_bytes = total_bytes; + return false; + } +} +#endif + +bool get_memory_size_by_sycl_api(sycl::device dev, size_t & free_bytes, size_t & total_bytes) { + GGML_SYCL_DEBUG("[%s]Querying free memory using SYCL API.\n", __func__); + total_bytes = dev.get_info(); + +#if (defined(__SYCL_COMPILER_VERSION) && __SYCL_COMPILER_VERSION >= 20221105) + if (dev.has(sycl::aspect::ext_intel_free_memory)) { + try { + GGML_SYCL_DEBUG("Querying free memory using SYCL aspect::ext_intel_free_memory."); + free_bytes = dev.get_info(); + return true; + } catch (const sycl::exception &) { + GGML_SYCL_DEBUG( + "Failed to query free memory using SYCL aspect::ext_intel_free_memory. Using total memory as free " + "memory."); + free_bytes = total_bytes; + return false; + } + } else { + GGML_SYCL_DEBUG( + "Device does not support SYCL aspect::ext_intel_free_memory. Using total memory as free memory."); + free_bytes = total_bytes; + } +#else + GGML_SYCL_DEBUG("SYCL Compiler version is older than 20221105. Using total memory as free memory."); + free_bytes = total_bytes; +#endif + return true; +} + +bool get_memory_size(sycl::device dev, size_t & free_bytes, size_t & total_bytes, MemoryAPIType api_type) { + const auto name = dev.get_info(); + const auto vendor = dev.get_info(); + const auto global_mem = dev.get_info(); + + GGML_SYCL_DEBUG("[%s]GPU Name: %s\n", __func__, name.c_str()); + GGML_SYCL_DEBUG("[%s]GPU Vendor: %s\n", __func__, vendor.c_str()); + GGML_SYCL_DEBUG("[%s]GPU Global Memory: %zu bytes\n", __func__, static_cast(global_mem)); + + if (api_type == MEMORY_API_TYPE_LEVEL_ZERO) { +#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API + GGML_SYCL_DEBUG("[%s]Querying free memory using Level Zero API.\n", __func__); + if (!query_free_memory_by_ze(dev, free_bytes, total_bytes)) { + //fallback to SYCL API if Level Zero API fails + GGML_SYCL_DEBUG("[%s]Falling back to SYCL API for memory query.\n", __func__); + return get_memory_size_by_sycl_api(dev, free_bytes, total_bytes); + } + return true; +#else + GGML_SYCL_DEBUG("[%s]Level Zero API support is not enabled. Please enable it to use this feature.\n", __func__); + return false; +#endif + } else { //MEMORY_API_TYPE_SYCL + return get_memory_size_by_sycl_api(dev, free_bytes, total_bytes); + } +} diff --git a/ggml/src/ggml-sycl/mem.hpp b/ggml/src/ggml-sycl/mem.hpp new file mode 100644 index 000000000..b3e45cfea --- /dev/null +++ b/ggml/src/ggml-sycl/mem.hpp @@ -0,0 +1,16 @@ +#ifndef GGML_SYCL_MEM_HPP +#define GGML_SYCL_MEM_HPP + +#include + +enum MemoryAPIType { + MEMORY_API_TYPE_LEVEL_ZERO = 0, + MEMORY_API_TYPE_SYCL = 1, +}; + +const char* mem_api_int2str(int mem_api); + +bool get_memory_size(sycl::device dev, size_t & free_bytes, size_t & total_bytes, + MemoryAPIType api_type); + +#endif // GGML_SYCL_MEM_HPP From 41ef91f7c8046087cdfbb276b79bff311ecf1c6d Mon Sep 17 00:00:00 2001 From: ynankani Date: Mon, 31 Aug 2026 11:22:28 +0000 Subject: [PATCH 03/10] CUDA: extend MOE fusion to specdec, earlier MOE glu fusion and topk-router fusion were restricted to 1 token (#27621) * CUDA: extend MOE fusion to specdec, earlier MOE glu fusion and topk-router fusion were resticted to 1 token Signed-off-by: ynankani * Address review comments Signed-off-by: ynankani * Add SWIGLU_CLAMP case to multi-token moe fusion Signed-off-by: ynankani --------- Signed-off-by: ynankani --- ggml/src/ggml-cuda/ggml-cuda.cu | 11 +-- ggml/src/ggml-cuda/mmvq.cu | 115 +++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/topk-moe.cu | 24 ++++--- ggml/src/ggml-cuda/topk-moe.cuh | 3 + tests/test-backend-ops.cpp | 14 ++-- 5 files changed, 135 insertions(+), 32 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a6fc655c4..53eccdd65 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1807,7 +1807,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return false; } - if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { + if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] > get_mmvq_mmid_max_batch(src0->type, cc)) { return false; } @@ -2983,9 +2983,10 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, }; 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; + // one block reads all logits before it writes, so logits may alias the out nodes + const ggml_tensor * logits_may_alias = nullptr; + if (is_topk_moe && ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK) { + logits_may_alias = cgraph->nodes[node_idx]->src[0]; } for (int i = 0; i < out_count; ++i) { @@ -2999,7 +3000,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, 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) { + if (!src || src->op == GGML_OP_NONE || src == logits_may_alias) { continue; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 79f7a3f6f..2be2f2491 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -773,10 +773,10 @@ static __global__ void mul_mat_vec_q( // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) // Block: (warp_size, ncols_dst) - each warp handles one token independently. // No shared memory reduction needed since each warp works alone. -template +template __launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q_moe( - const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, + const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, @@ -794,6 +794,29 @@ static __global__ void mul_mat_vec_q_moe( constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + // fuse gate, bias, scales, and glu_op into the up projection + bool use_gate = false; + const void * vgate = nullptr; + const float * x_bias = nullptr; + const float * gate_bias = nullptr; + const float * x_scale = nullptr; + const float * gate_scale = nullptr; + ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU; + float glu_limit = 0.0f; + + if constexpr (has_fusion) { + use_gate = fusion.gate != nullptr; + vgate = fusion.gate; + x_bias = (const float *) fusion.x_bias; + gate_bias = (const float *) fusion.gate_bias; + active_glu = fusion.glu_op; + glu_limit = fusion.glu_limit; + if constexpr (type == GGML_TYPE_NVFP4) { + x_scale = (const float *) fusion.x_scale; + gate_scale = (const float *) fusion.gate_scale; + } + } + const uint32_t token_idx = threadIdx.y; const int row0 = c_rows_per_block*blockIdx.x; const int blocks_per_row_x = ncols_x / qk; @@ -814,6 +837,7 @@ static __global__ void mul_mat_vec_q_moe( // partial sum for each thread float tmp[c_rows_per_block] = {0.0f}; + float tmp_gate[c_rows_per_block] = {0.0f}; for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { const int kby = kbx * (qk/QK8_1); @@ -822,6 +846,11 @@ static __global__ void mul_mat_vec_q_moe( #pragma unroll for (int i = 0; i < c_rows_per_block; ++i) { tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } } } @@ -831,11 +860,63 @@ static __global__ void mul_mat_vec_q_moe( #pragma unroll for (int i = 0; i < c_rows_per_block; ++i) { tmp[i] = warp_reduce_sum(tmp[i]); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[i] = warp_reduce_sum(tmp_gate[i]); + } + } } // Write results if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) { - dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x]; + float result = tmp[threadIdx.x]; + if constexpr (has_fusion) { + const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x; + + if constexpr (type == GGML_TYPE_NVFP4) { + if (x_scale) { + result *= x_scale[channel_x]; + } + } + if (x_bias) { + result += x_bias[bias_idx]; + } + if (use_gate) { + float gate_value = tmp_gate[threadIdx.x]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (gate_scale) { + gate_value *= gate_scale[channel_x]; + } + } + if (gate_bias) { + gate_value += gate_bias[bias_idx]; + } + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + case GGML_GLU_OP_SWIGLU_CLAMP: + result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit); + break; + default: + result = result * gate_value; + break; + } + } + } + dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result; + } + + if constexpr (!has_fusion) { + GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale); + } else if constexpr (type != GGML_TYPE_NVFP4) { + GGML_UNUSED_VARS(x_scale, gate_scale); } } @@ -885,7 +966,7 @@ static void mul_mat_vec_q_switch_fusion( template static void mul_mat_vec_q_moe_launch( - const void * vx, const void * vy, const int32_t * ids, float * dst, + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, @@ -898,11 +979,22 @@ static void mul_mat_vec_q_moe_launch( const dim3 block_dims(warp_size, ncols_dst); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, - vx, vy, ids, dst, ncols_x, nchannels_y, nrows_x, - stride_row_x, stride_col_y, stride_col_dst, - stride_channel_x, stride_channel_y, stride_channel_dst, - ncols_dst, ids_stride); + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + + if (has_fusion) { + ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride); + } else { + ggml_cuda_kernel_launch(mul_mat_vec_q_moe, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, ids_stride); + } } template @@ -998,7 +1090,7 @@ static void mul_mat_vec_q_switch_ncols_dst( if (has_ids && ncols_dst > 1) { // Multi-token MUL_MAT_ID path - dedicated MoE kernel mul_mat_vec_q_moe_launch( - vx, vy, ids, dst, ncols_x, nchannels_y_fd, nrows_x, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, ncols_dst, ids_stride, warp_size, nchannels_dst, stream); @@ -1280,7 +1372,8 @@ void ggml_cuda_mul_mat_vec_q( ggml_cuda_mm_fusion_args_device fusion_local{}; if (fusion) { - GGML_ASSERT( !ids || dst->ne[2] == 1); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)); GGML_ASSERT( ids || dst->ne[1] == 1); // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is // non-negligible for some models such as gpt-oss-20b diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index c8cec70bb..dadcd601c 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -88,15 +88,16 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models */ template -__launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits, - float * weights, - int32_t * ids, - float * bias, - const int n_rows, - const int n_expert_used, - const float clamp_val, - const float scale_val, - const topk_moe_config config) { +__launch_bounds__(TOPK_MOE_ROWS_PER_BLOCK * WARP_SIZE, 1) +__global__ void topk_moe_cuda(const float * logits, + float * weights, + int32_t * ids, + float * bias, + const int n_rows, + const int n_expert_used, + const float clamp_val, + const float scale_val, + const topk_moe_config config) { const int row = blockIdx.x * blockDim.y + threadIdx.y; if (row >= n_rows) { return; @@ -123,6 +124,9 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY; } + // Weights and IDs can alias logits, so wait until every row in the block reads its logits. + __syncthreads(); + if (!config.delayed_softmax) { if (config.use_sigmoid) { sigmoid_warp_inplace(wt, n_experts, threadIdx.x); @@ -282,7 +286,7 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, const topk_moe_config config) { GGML_ASSERT(!(config.with_norm && config.delayed_softmax) && "delayed softmax is not supported with weight normalization"); - const int rows_per_block = 4; + const int rows_per_block = TOPK_MOE_ROWS_PER_BLOCK; dim3 grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1); dim3 block_dims(WARP_SIZE, rows_per_block, 1); cudaStream_t stream = ctx.stream(); diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh index 091ef02a4..061b37e29 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh @@ -3,6 +3,9 @@ #include +// Rows that one CUDA block handles. +#define TOPK_MOE_ROWS_PER_BLOCK 8 + struct ggml_cuda_topk_moe_args { bool sigmoid{}; bool sqrt_softplus{}; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index becf6a555..a1875c8ed 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10205,12 +10205,10 @@ static std::vector> make_test_cases_eval() { use_id, 16, 8, b, with_bias, with_gate, with_lane_scale)); test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256, use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1})); - if (!use_id && with_gate && !with_bias && glu_op != GGML_GLU_OP_SWIGLU_CLAMP) { - // small multi-token batches (speculative decoding / MTP verify) - for (int64_t m_batch : { 2, 4, 8 }) { - test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256, - use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1})); - } + // multi-token batches (spec decoding) + for (int64_t m_batch : { 2, 4, 8 }) { + test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256, + use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1})); } } } @@ -10239,6 +10237,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4 test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7 + // rows at and just past the limit where one block still covers all rows + test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 4, with_norm, bias_probs, gate, scale_w)); + test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); + test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); } } } From 5d4a3be26de60093263b2d41ca24c3a49de95cd5 Mon Sep 17 00:00:00 2001 From: Niklas Wenzel Date: Mon, 31 Aug 2026 13:58:55 +0200 Subject: [PATCH 04/10] metal : add fa-vec tunings for M1 (#28078) --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 190 ++++++++++++++++++++++ 1 file changed, 190 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 4f26fd9c3..6a742bd0f 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -68,6 +68,196 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) { // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics. // ref: https://github.com/ggml-org/llama.cpp/pull/27824 constexpr fa_vec_entry_t fa_vec_tuned_table[] = { + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, 3, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 2, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 3, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, From f8dbcd61893702976f9ab03be89c2b9f436d532c Mon Sep 17 00:00:00 2001 From: Jaden_Mach <88880593+jadenmach2@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:00:04 -0400 Subject: [PATCH 05/10] ROCm: add radix TOP_K for long rows (#27466) * ROCm: add radix TOP_K for long rows --- ggml/src/ggml-cuda/ggml-cuda.cu | 5 + ggml/src/ggml-cuda/top-k.cu | 180 +++++++++++++++++++++++++++++++- 2 files changed, 180 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 53eccdd65..31f5aeeac 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5273,6 +5273,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_SUM: return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: +#if defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB) + return true; +#else + return op->src[0]->ne[0] <= 1024; +#endif // defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB) case GGML_OP_ARGSORT: #ifndef GGML_CUDA_USE_CUB return op->src[0]->ne[0] <= 1024; diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd293..c7a0c8317 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -48,6 +48,168 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE +#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) + +static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t mask = (uint32_t) (-(int32_t) (bits >> 31)) | 0x80000000U; + return bits ^ mask; +} + +struct top_k_radix_state { + uint32_t prefix; + uint32_t prefix_mask; + int rank; + int greater_count; + int equal_count; +}; + +static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < nrows) { + states[row] = {0, 0, k, 0, 0}; + } +} + +template +static __global__ void top_k_radix_histogram( + const float * __restrict__ src, + const top_k_radix_state * __restrict__ states, + int * __restrict__ block_histograms, + int ncols, + int blocks_per_row, + int shift) { + constexpr int NBINS = 1 << RADIX_BITS; + + const int row = blockIdx.x / blocks_per_row; + const int row_block = blockIdx.x % blocks_per_row; + const int tid = threadIdx.x; + const float * row_src = src + (size_t) row * ncols; + __shared__ int histogram[NBINS]; + + histogram[tid] = 0; + __syncthreads(); + + const top_k_radix_state state = states[row]; + for (int col = row_block * BLOCK_SIZE + tid; + col < ncols; + col += blocks_per_row * BLOCK_SIZE) { + const uint32_t key = top_k_float_to_ordered(row_src[col]); + if ((key & state.prefix_mask) == state.prefix) { + atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1); + } + } + __syncthreads(); + + const size_t histogram_offset = + ((size_t) row * blocks_per_row + row_block) * NBINS; + block_histograms[histogram_offset + tid] = histogram[tid]; +} + +template +static __global__ void top_k_radix_select( + const int * __restrict__ block_histograms, + top_k_radix_state * __restrict__ states, + int blocks_per_row, + int shift) { + constexpr int NBINS = 1 << RADIX_BITS; + + const int row = blockIdx.x; + const int tid = threadIdx.x; + __shared__ int histogram[NBINS]; + + int count = 0; + for (int row_block = 0; row_block < blocks_per_row; ++row_block) { + const size_t offset = ((size_t) row * blocks_per_row + row_block) * NBINS; + count += block_histograms[offset + tid]; + } + histogram[tid] = count; + __syncthreads(); + + if (tid == 0) { + top_k_radix_state state = states[row]; + int bin = NBINS - 1; + while (bin > 0 && histogram[bin] < state.rank) { + state.rank -= histogram[bin--]; + } + state.prefix |= (uint32_t) bin << shift; + state.prefix_mask |= (uint32_t) (NBINS - 1) << shift; + states[row] = state; + } +} + +static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < nrows) { + states[row].greater_count = 0; + states[row].equal_count = 0; + } +} + +template +static __global__ void top_k_radix_gather( + const float * __restrict__ src, + int * __restrict__ dst, + top_k_radix_state * __restrict__ states, + int ncols, + int k, + int blocks_per_row) { + const int row = blockIdx.x / blocks_per_row; + const int row_block = blockIdx.x % blocks_per_row; + const int tid = threadIdx.x; + const float * row_src = src + (size_t) row * ncols; + int * row_dst = dst + (size_t) row * k; + top_k_radix_state * state = &states[row]; + + for (int col = row_block * BLOCK_SIZE + tid; + col < ncols; + col += blocks_per_row * BLOCK_SIZE) { + const uint32_t key = top_k_float_to_ordered(row_src[col]); + if (key > state->prefix) { + const int pos = atomicAdd(&state->greater_count, 1); + row_dst[pos] = col; + } else if (key == state->prefix) { + const int pos = atomicAdd(&state->equal_count, 1); + if (pos < state->rank) { + row_dst[k - state->rank + pos] = col; + } + } + } +} + +static void top_k_radix_cuda( + ggml_cuda_pool & pool, + const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) { + constexpr int BLOCK_SIZE = 256; + constexpr int RADIX_BITS = 8; + constexpr int NBINS = 1 << RADIX_BITS; + const int blocks_per_row = std::min((ncols + 1023) / 1024, 64); + + ggml_cuda_pool_alloc states_alloc(pool, nrows); + ggml_cuda_pool_alloc histograms_alloc(pool, (size_t) nrows * blocks_per_row * NBINS); + top_k_radix_state * states = states_alloc.get(); + int * histograms = histograms_alloc.get(); + + top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k); + + const dim3 row_grid(blocks_per_row * nrows); + for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) { + top_k_radix_histogram + <<>>( + src, states, histograms, ncols, blocks_per_row, shift); + top_k_radix_select + <<>>(histograms, states, blocks_per_row, shift); + } + + top_k_radix_reset_counters + <<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows); + top_k_radix_gather + <<>>( + src, dst, states, ncols, k, blocks_per_row); +} + +#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -96,10 +258,18 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { dst_d += k * iter_nrows; } #else // GGML_CUDA_USE_CUB - ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); - int * tmp_dst = temp_dst_alloc.get(); - argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); - CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, - cudaMemcpyDeviceToDevice, stream)); +#if defined(GGML_USE_HIP) + if (ncols > 1024) { + top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); + } else { +#endif // defined(GGML_USE_HIP) + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + int * tmp_dst = temp_dst_alloc.get(); + argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); + CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, + cudaMemcpyDeviceToDevice, stream)); +#if defined(GGML_USE_HIP) + } +#endif // defined(GGML_USE_HIP) #endif } From 8e53fcefd2c01ff70434ab41866bfc2eca31fe90 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:04:38 +0200 Subject: [PATCH 06/10] webgpu : avoid crash when offset is not multiple of 4 in WebGPU ggml_backend_tensor_get() implementation (#28045) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * webgpu : avoid crash when offset is not multiple of 4 in WebGPU ggml_backend_tensor_get() implementation * chore : improve code readability Co-authored-by: Sigbjørn Skjæret --------- Co-authored-by: Stanisław Szymczyk Co-authored-by: Sigbjørn Skjæret --- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index b953118a7..1a43c7273 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -3713,11 +3713,18 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer, size_t total_offset = ggml_webgpu_tensor_offset(tensor) + offset; - size_t final_size = size; - if (size % 4 != 0) { + size_t local_offset = total_offset % 4; + if (local_offset != 0) { + // If offset is not a multiple of 4, we need to round it down to the previous + // multiple of 4 + total_offset -= local_offset; + } + + size_t final_size = size + local_offset; + if (final_size % 4 != 0) { // If size is not a multiple of 4, we need to round it up to the next // multiple of 4 - final_size = size + (4 - (size % 4)); + final_size += 4 - (final_size % 4); } std::lock_guard lock(buf_ctx->global_ctx->mutex); @@ -3748,7 +3755,7 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer, const void * mapped_range = buf_ctx->global_ctx->get_tensor_staging_buf.GetConstMappedRange(0, final_size); // Copy the data from the mapped range to the output buffer - std::memcpy(data, mapped_range, size); + std::memcpy(data, (const void *) ((const char *) mapped_range + local_offset), size); buf_ctx->global_ctx->get_tensor_staging_buf.Unmap(); WEBGPU_CPU_PROFILE_TOTAL_END(get_tensor, buf_ctx->global_ctx); } From 774ee0e20097e8febbef0d4f735e30dd922cc239 Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 31 Aug 2026 17:48:43 +0200 Subject: [PATCH 07/10] ui: copy the displayed text of grouped agentic responses (#27832) * ui: copy the displayed text of grouped agentic responses Agentic sessions render as a single entry anchored on the first assistant turn, whose content is typically just the first tool call, so the copy button wrote an empty string to the clipboard. Derive the text sections of the whole session and copy them joined, matching the visible response. Plain messages keep the previous behavior. * const --- .../ChatMessage/ChatMessage.svelte | 25 ++++++++++++++++++- .../app/chat/ChatMessages/ChatMessages.svelte | 4 +-- .../ui/src/lib/constants/agentic.constants.ts | 4 +++ tools/ui/src/lib/types/chat.d.ts | 2 +- 4 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte index 0c9ead61e..fa2a50bc5 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte @@ -7,7 +7,12 @@ ChatMessageSystem, ChatMessageUser } from '$lib/components/app/chat'; - import { REASONING_TAGS, ROUTES, SYSTEM_MESSAGE_PLACEHOLDER } from '$lib/constants'; + import { + AGENTIC_TEXT_COPY_SEPARATOR, + REASONING_TAGS, + ROUTES, + SYSTEM_MESSAGE_PLACEHOLDER + } from '$lib/constants'; import { setChatMessageActionsContext, setChatMessageEditContext } from '$lib/contexts'; import { AgenticSectionType, AttachmentType, MessageRole } from '$lib/enums'; import { DatabaseService } from '$lib/services/database.service'; @@ -237,6 +242,24 @@ } function handleCopy() { + // Agentic sessions render as a single entry anchored on the first assistant + // turn, whose own content is typically just the first tool call. Copy the + // text sections of the whole session so the clipboard matches the visible + // response instead of the anchor turn. + if (message.role === MessageRole.ASSISTANT) { + const sections = deriveAgenticSections(message, toolMessages, [], false); + const text = sections + .filter((section) => section.type === AgenticSectionType.TEXT) + .map((section) => section.content) + .join(AGENTIC_TEXT_COPY_SEPARATOR); + + if (text) { + chatActions.copy(message, text); + + return; + } + } + chatActions.copy(message); } diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte index 45b863d66..4750a9f7c 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte @@ -29,10 +29,10 @@ refreshAllMessages(); }, - copy: async (message: DatabaseMessage) => { + copy: async (message: DatabaseMessage, contentOverride?: string) => { const asPlainText = Boolean(currentConfig.copyTextAttachmentsAsPlainText); const clipboardContent = formatMessageForClipboard( - message.content, + contentOverride ?? message.content, message.extra, asPlainText ); diff --git a/tools/ui/src/lib/constants/agentic.constants.ts b/tools/ui/src/lib/constants/agentic.constants.ts index e57104e8a..33fb30f1b 100644 --- a/tools/ui/src/lib/constants/agentic.constants.ts +++ b/tools/ui/src/lib/constants/agentic.constants.ts @@ -20,6 +20,10 @@ export const SEARCH_SUMMARY = { // wraps mid-paragraph. export const RESULT_STAT_SEPARATOR = ' - '; +// Separator between the assistant text sections of a grouped agentic +// session when they are joined for the clipboard. +export const AGENTIC_TEXT_COPY_SEPARATOR = '\n\n'; + export const DEFAULT_AGENTIC_CONFIG: AgenticConfig = { enabled: true, maxTurns: 100 diff --git a/tools/ui/src/lib/types/chat.d.ts b/tools/ui/src/lib/types/chat.d.ts index 274131dd4..86a868c33 100644 --- a/tools/ui/src/lib/types/chat.d.ts +++ b/tools/ui/src/lib/types/chat.d.ts @@ -249,7 +249,7 @@ export interface ChatMessageDeletionInfo { * refresh + user-action notification), passed to each ChatMessage as a prop. */ export interface ChatMessageActions { - copy: (message: DatabaseMessage) => void; + copy: (message: DatabaseMessage, contentOverride?: string) => void; delete: (message: DatabaseMessage) => void; navigateToSibling: (siblingId: string) => void; editWithBranching: ( From 010be9683afabe14ce299197b38c329f94bae568 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Mon, 31 Aug 2026 08:56:22 -0700 Subject: [PATCH 08/10] opencl: tune the quant paths for Intel Xe-LP GPUs to improve its TG and PP performance (#26438) * opencl: Q4_K/Q5_K mul_mv N_DST 4->8 on Intel for 2x activation reuse * opencl: Q4_K mul_mm 8x8 tile fot Intel * opencl: Q5_K mul_mm 8x8 tile for Intel * opencl: Q4_K mul_mv N_DST 8->16 for Intel --- ggml/src/ggml-opencl/ggml-opencl.cpp | 9 +++++---- ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl | 10 ++++++++++ ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl | 10 ++++++++++ ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl | 2 +- ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl | 2 +- 5 files changed, 27 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 90635cc85..34d58f4ee 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -20002,7 +20002,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co } kernel = backend_ctx->kernel_mul_mm_q4_k_f32_l4_lm; - nth0 = 128; // calculated as (BM*BN)/(TM*TN) + // (BM*BN)/(TM*TN): Intel uses an 8x8 microtile (WG=64), others 4x8 (WG=128) + nth0 = (backend_ctx->gpu_family == INTEL) ? 64 : 128; int batch_stride_a = ne00*ne01; int batch_stride_b = ne10*ne11; @@ -20046,7 +20047,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co } kernel = backend_ctx->kernel_mul_mm_q5_k_f32_l4_lm; - nth0 = 128; // calculated as (BM*BN)/(TM*TN) + nth0 = (backend_ctx->gpu_family == INTEL) ? 64 : 128; // Intel 8x8 microtile int batch_stride_a = ne00*ne01; int batch_stride_b = ne10*ne11; @@ -20860,7 +20861,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co if (backend_ctx->gpu_family == INTEL) { nth0 = 16; nth1 = 1; - ndst = 4; + ndst = 16; // 8->16 rows per subgroup — matches N_DST in mul_mv_q4_k_f32_flat.cl (32 spills) } else if (backend_ctx->gpu_family == ADRENO) { nth0 = 64; nth1 = 2; @@ -20934,7 +20935,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co if (backend_ctx->gpu_family == INTEL) { nth0 = 16; nth1 = 1; - ndst = 4; + ndst = 8; // 4->8 rows per subgroup (2x activation reuse) } else if (backend_ctx->gpu_family == ADRENO) { nth0 = 64; nth1 = 2; diff --git a/ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl b/ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl index 2235b1ae8..a9c649a52 100644 --- a/ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl +++ b/ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl @@ -1,13 +1,23 @@ #pragma OPENCL EXTENSION cl_khr_fp16 : enable +#ifdef cl_intel_required_subgroup_size +#define INTEL_GPU 1 +#endif + #define LOAD_VEC_A 4 #define LOAD_VEC_B 4 #define BM 64 #define BN 64 #define BK 32 +#ifdef INTEL_GPU +// Intel Xe iGPU: 8x8 microtile (WG = BM*BN/(TM*TN) = 64) — ~+12% pp512 vs 4x8 +#define TM 8 +#define TN 8 +#else #define TM 4 #define TN 8 +#endif kernel void kernel_mul_mm_q4_k_f32_l4_lm( global uchar4 * src0_q, diff --git a/ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl b/ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl index 8e191f57e..a343b5c4c 100644 --- a/ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl +++ b/ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl @@ -1,13 +1,23 @@ #pragma OPENCL EXTENSION cl_khr_fp16 : enable +#ifdef cl_intel_required_subgroup_size +#define INTEL_GPU 1 +#endif + #define LOAD_VEC_A 4 #define LOAD_VEC_B 4 #define BM 64 #define BN 64 #define BK 32 +#ifdef INTEL_GPU +// Intel Xe iGPU: 8x8 microtile (WG=64) +#define TM 8 +#define TN 8 +#else #define TM 4 #define TN 8 +#endif kernel void kernel_mul_mm_q5_k_f32_l4_lm( global uchar4 * src0_q, diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl index 70391866c..5316bd363 100644 --- a/ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl +++ b/ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl @@ -40,7 +40,7 @@ typedef struct { #undef N_SIMDWIDTH #ifdef INTEL_GPU -#define N_DST 4 // number of rows each SIMD group works on +#define N_DST 16 // number of rows each SIMD group works on (Intel: 8->16, 2x further activation reuse; 32 spills registers) #define N_SIMDGROUP 1 // number of SIMD groups in a thread group #define N_SIMDWIDTH 16 // SIMD group size #elif defined (ADRENO_GPU) diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl index 6020364b5..ab2e1fab8 100644 --- a/ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl +++ b/ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl @@ -38,7 +38,7 @@ typedef struct { #undef N_SIMDWIDTH #ifdef INTEL_GPU -#define N_DST 4 +#define N_DST 8 // Intel: 4->8 for 2x activation reuse (see mul_mv_q4_k_f32_flat.cl) #define N_SIMDGROUP 1 #define N_SIMDWIDTH 16 #elif defined(ADRENO_GPU) From 2d8d612e4c68d3801e556a1b4a028f55ec33ecbb Mon Sep 17 00:00:00 2001 From: itsnotoger <19309683+itsnotoger@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:49:58 +0200 Subject: [PATCH 09/10] kv-cache : optimize restoring non-contiguous cells (#27991) * kv cache : batch state restore scatter reads per contiguous run When restoring state into non-contiguous destination cells (e.g. a prompt-cache snapshot into a fragmented ring), state_read_data issued one small copy per KV cell - ~1.4M copies of a few KiB each for a 40k+ token restore, taking 25-63 s on the CUDA backend. The snapshot stores cell rows in cell order, so a maximal run of consecutive destination indices maps to one contiguous block and can be restored with a single copy. Precompute the runs once and use them in all three scatter loops (K, V, transposed V). Byte-identical. The on-device reader copies with a byte cursor when the read and write chunking differs, so the batched reads are safe for it as well. Batching makes equal tensor counts with a different split reachable (save ranges [2,1] vs restore runs [1,2]); the next commit teaches the reader's 1:1 path to fall back to the byte cursor in that case. Verified in a production setup: 1,363,616 copies / 25-63 s -> 224 copies / 221-424 ms for the same restores (42,603 cells, 4 runs). Assisted-by: Claude Code (unsloth/qwen3.8-27b) * context : fall back to the byte cursor when read and write chunking differ the on-device reader copies saved state back with a 1:1 copy by tensor index whenever the write and read sides recorded the same number of tensors, guarded by a per-tensor size assert. equal tensor counts do not imply equal chunking: a state restore may batch its reads per contiguous run of destination cells while the save used per-range reads, so both sides can record two tensors that split the same data differently, and the assert aborts in all builds. compare the per-tensor sizes and only take the 1:1 path when the chunking actually matches, otherwise fall through to the existing byte-cursor copy. both sides enumerate the same logical data in the same order, so the cursor copy is well-defined across tensor boundaries. Assisted-by: Claude Code (unsloth/qwen3.8-27b) * tests : cover state restore scatter reads on host and on-device paths decode the same prefix on two sequences, interleaving the seq 0 cells between the seq 1 cells, so the seq 1 cells are isolated from each other in the kv cache (three cells, two saved ranges). save the seq 1 state, free the interleaved seq 0 cells, and restore: the destination is then non-contiguous (two runs), and the restore-side chunking has the same tensor count as the save-side with a different split, so the scatter path is batched per contiguous run and the on-device reader's byte-cursor fallback is exercised. the restored state is saved again on the host and compared byte for byte with the first save: the blob is serialized in sequence cell order, so the two saves are identical if and only if the scatter restore wrote exactly the same KV content. this documents the byte-identical guarantee of the run-batched scatter reads. one test per io backend: the host (CPU) path and the on-device path. Assisted-by: Claude Code (unsloth/qwen3.8-27b) --- src/llama-context.cpp | 21 +++++-- src/llama-kv-cache.cpp | 64 ++++++++------------ tests/test-save-load-state.cpp | 106 ++++++++++++++++++++++++++++++++- 3 files changed, 147 insertions(+), 44 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a6d2f264f..a920c4231 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2907,17 +2907,28 @@ public: } if (mbuf_cur.n_tensors == mbuf.n_tensors) { - // same chunking: copy 1:1 by index + // an equal tensor count does not imply the same chunking, e.g. save ranges [2,1] vs restore runs [1,2] + bool same_chunking = true; for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i])); - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + if (ggml_nbytes(mbuf_cur.cpy[i]) != ggml_nbytes(mbuf.org[i])) { + same_chunking = false; + break; + } + } + + if (same_chunking) { + // same chunking: copy 1:1 by index + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + } + continue; } - continue; } // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk - // it differently, so copy across tensor boundaries rather than 1:1 by index. + // it differently (even with an equal number of tensors), so copy across tensor boundaries rather than + // 1:1 by index. const size_t total = mbuf_cur.total_size; ggml_init_params params_scratch = { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 65afbd8c3..0e095f8b6 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2533,6 +2533,24 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) { auto & cells = v_cells[strm]; + // batch the scatter reads per contiguous run of destination indices + // from inclusive, to exclusive - same convention as cell_ranges_t + // contiguous cells yield a single run covering the whole block + struct cell_run { uint32_t from; uint32_t to; }; + std::vector runs; + if (cell_count > 0) { + const auto & idxs = sinfo.idxs[0]; + uint32_t i0 = 0; + while (i0 < cell_count) { + uint32_t i1 = i0 + 1; + while (i1 < cell_count && idxs[i1] == idxs[i1 - 1] + 1) { + ++i1; + } + runs.push_back({idxs[i0], idxs[i1 - 1] + 1}); + i0 = i1; + } + } + uint32_t v_trans; uint32_t n_layer; @@ -2580,17 +2598,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells, single memcpy - io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row); - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = sinfo.idxs[0][i] * k_size_row; - io.read_tensor(k, dst_offset, k_size_row); - } - } + for (const auto & r : runs) { + io.read_tensor(k, (size_t) r.from * k_size_row, (size_t) (r.to - r.from) * k_size_row); } } @@ -2623,17 +2632,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells, single memcpy - io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row); - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = sinfo.idxs[0][i] * v_size_row; - io.read_tensor(v, dst_offset, v_size_row); - } - } + for (const auto & r : runs) { + io.read_tensor(v, (size_t) r.from * v_size_row, (size_t) (r.to - r.from) * v_size_row); } } } else { @@ -2674,22 +2674,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (cell_count) { - if (sinfo.is_contiguous()) { - // Fast path: contiguous cells - const uint32_t h = sinfo.head(); - for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - const size_t dst_offset = (h + j * cells.size()) * v_size_el; - io.read_tensor(v, dst_offset, cell_count * v_size_el); - } - } else { - // Slow path: scatter to non-contiguous positions - for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - for (uint32_t i = 0; i < cell_count; ++i) { - const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el; - io.read_tensor(v, dst_offset, v_size_el); - } - } + for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { + for (const auto & r : runs) { + const size_t dst_offset = ((size_t) r.from + j * cells.size()) * v_size_el; + io.read_tensor(v, dst_offset, (size_t) (r.to - r.from) * v_size_el); } } } diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 0ceab7c54..35c769058 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -355,7 +355,101 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } -// Run the full save/load test suite (tests 1-5) for a single model. +// Test 6/7: seq copy (scatter) +// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells +// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path +// - the restore destination is non-contiguous: scatter reads are batched per contiguous run +// - save again on the host and compare the two blobs byte for byte +static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) { + auto params_ctx = common_context_params_to_llama(params); + params_ctx.n_ctx = 256; + params_ctx.n_seq_max = 2; + params_ctx.kv_unified = true; + auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; + + LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host"); + + const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE; + + auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) { + llama_batch_ptr batch(1, 0, 1); + common_batch_add(batch.get(), tok, pos, { seq }, false); + return llama_decode(ctx.get(), batch.get()) == 0; + }; + + // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5 + if (!decode_one(tokens[0], 0, 0) || + !decode_one(tokens[1], 1, 0) || + !decode_one(tokens[0], 0, 1) || + !decode_one(tokens[1], 1, 1) || + !decode_one(tokens[2], 2, 0) || + !decode_one(tokens[2], 2, 1)) { + LOG_ERR("%s: failed to build interleaved state\n", __func__); + return false; + } + + const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector & state) { + const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl); + if (state_size == 0) { + LOG_ERR("%s: sequence state is empty\n", __func__); + return false; + } + + state.resize(state_size); + const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl); + if (ncopy != state.size()) { + LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n", + __func__, ncopy, state.size()); + return false; + } + + return true; + }; + + // host blob: contains the KV data, used for the byte-for-byte comparison + std::vector state_before; + if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) { + return false; + } + + // save via the io path under test + std::vector state_save; + if (!get_seq_state(1, flags, state_save)) { + return false; + } + LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size()); + + // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous + if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) { + LOG_ERR("%s: failed to remove sequence 0\n", __func__); + return false; + } + + // restore via the io path under test + const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags); + if (nset != state_save.size()) { + LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size()); + return false; + } + LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset); + + std::vector state_after; + if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) { + return false; + } + + // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV + if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) { + LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__); + return false; + } + + LOG("\nPASS\n"); + return true; +} + + +// Run the full save/load test suite (tests 1-7) for a single model. // Returns true if all tests pass, false otherwise. static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { struct common_params params = base_params; @@ -422,6 +516,16 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const return false; } + // Test 6: seq copy (host, scatter) + if (!test_seq_cp_scatter(model, params, tokens, 6, false)) { + return false; + } + + // Test 7: seq copy (device, scatter) + if (!test_seq_cp_scatter(model, params, tokens, 7, true)) { + return false; + } + LOG("\nAll tests passed.\n"); return true; From 2a74817f93dea568bcbc58b3bd7ed89dc92306e3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 21:31:53 +0300 Subject: [PATCH 10/10] metal : add top-k radix implementation (#28073) Assisted-by: DeepSeek-v4-Flash-0731 --- ggml/src/ggml-metal/ggml-metal-device.cpp | 19 +++- ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-impl.h | 11 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 72 ++++++++++++++- ggml/src/ggml-metal/kernels/argsort.metal | 105 ++++++++++++++++++++++ 5 files changed, 206 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 5a2f01f1d..b8d2ef9ce 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1336,7 +1336,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_ return res; } -// note: reuse the argsort kernel for top_k +// note: reuse the argsort kernel for the bitonic top_k fallback ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); @@ -1364,6 +1364,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) { + assert(op->op == GGML_OP_TOP_K); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 003b688db..7f6520103 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -145,6 +145,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse ); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op); diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 49102afe9..bdcd9c9e3 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1189,6 +1189,17 @@ typedef struct { int32_t len; } ggml_metal_kargs_argsort_merge; +typedef struct { + int32_t ne00; // number of columns (elements per row) + int32_t ne01; // rows + int32_t ne02; + int32_t ne03; + uint64_t nb01; // row stride in src0 + uint64_t nb02; + uint64_t nb03; + int32_t top_k; // k +} ggml_metal_kargs_top_k; + typedef struct { int32_t nrows; } ggml_metal_kargs_fwht; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 7671d1d01..30ea4ec27 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -5091,7 +5091,9 @@ int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) { return 1; } -int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { +// bitonic-sort + merge fallback: efficient when k is small and there are few rows, +// where the single-workgroup-per-row radix-select cannot reach enough parallelism +static void ggml_metal_op_top_k_bitonic(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); ggml_metal_library_t lib = ctx->lib; @@ -5199,6 +5201,74 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { len <<= 1; } +} + +// radix-select: one workgroup per row. Maps each float to an order-preserving unsigned +// key, finds the k-th largest via 4 radix-8 histogram passes, then compacts the top-k +// indices. Fast for large k and/or many rows. +static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + + auto pipeline = ggml_metal_library_get_pipeline_top_k_radix(lib, op); + + // one workgroup per row; radix-select the k-th largest value + const int nth = std::min(1024, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + + ggml_metal_kargs_top_k args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, + /*.ne03 =*/ ne03, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, + /*.nb03 =*/ nb03, + /*.top_k =*/ (int32_t) op->ne[0], + }; + + // shared memory: 256-entry histogram + bucket/above scalars + output counter + const size_t smem_histo = GGML_PAD(256*sizeof(uint32_t), 16); + const size_t smem_bucket = GGML_PAD( sizeof(uint32_t), 16); + const size_t smem_above = GGML_PAD( sizeof(uint32_t), 16); + const size_t smem_out = GGML_PAD( sizeof(uint32_t), 16); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_histo, 0); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_bucket, 1); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_above, 2); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_out, 3); + + ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); +} + +int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + // radix-select has a fixed single-workgroup-per-row cost (~50-60us) that is only + // amortized for long rows, many rows, or a large k; otherwise the bitonic path wins + const int ncols = op->src[0]->ne[0]; + const int k = op->ne[0]; + const int nrows = ggml_nrows(op->src[0]); + + const bool use_radix = + ncols > 2048 && (k > 64 || (nrows > 4 && ncols >= 8192)); + + if (use_radix) { + ggml_metal_op_top_k_radix(ctx, idx); + } else { + ggml_metal_op_top_k_bitonic(ctx, idx); + } return 1; } diff --git a/ggml/src/ggml-metal/kernels/argsort.metal b/ggml/src/ggml-metal/kernels/argsort.metal index 7d144fbd7..e81d194c3 100644 --- a/ggml/src/ggml-metal/kernels/argsort.metal +++ b/ggml/src/ggml-metal/kernels/argsort.metal @@ -230,3 +230,108 @@ kernel void kernel_argsort_merge_f32_i32( template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; + +static inline uint ggml_top_k_f2ui(float x) { + uint y = as_type(x); + if ((y & 0x80000000u) != 0u) { + y ^= 0xFFFFFFFFu; // negative floats: flip all bits + } else { + y |= 0x80000000u; // positive floats: set the sign bit + } + return y; +} + +kernel void kernel_top_k_f32_i32( + constant ggml_metal_kargs_top_k & args, + device const char * src0, + device int32_t * dst, + threadgroup atomic_uint * histo [[threadgroup(0)]], + threadgroup uint * sh_bucket [[threadgroup(1)]], + threadgroup uint * sh_above [[threadgroup(2)]], + threadgroup atomic_uint * out_count [[threadgroup(3)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + const uint ncols = args.ne00; + const uint top_k = args.top_k; + const uint i01 = tgpig[0]; + const uint i02 = tgpig[1]; + const uint i03 = tgpig[2]; + + device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03); + + device int32_t * dst_row = dst + top_k*(i01 + args.ne01*i02 + args.ne01*args.ne02*i03); + + const uint tid = tpitg.x; + const uint ntg_x = ntg.x; + + uint prefix = 0; // fixed high bits of the threshold key + uint desired = top_k; // count still needed from the candidate range + + for (int shift = 24; shift >= 0; shift -= 8) { + for (uint i = tid; i < 256; i += ntg_x) { + atomic_store_explicit(&histo[i], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint hi_mask = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + 8)); + const uint prefix_hi = prefix & hi_mask; + + for (uint i = tid; i < ncols; i += ntg_x) { + const uint key = ggml_top_k_f2ui(src0_row[i]); + if ((key & hi_mask) == prefix_hi) { + atomic_fetch_add_explicit(&histo[(key >> uint(shift)) & 0xFFu], 1u, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // top-down scan for the bucket holding the k-th value + if (tid == 0) { + uint acc = 0; + uint b = 0; + for (int bb = 255; bb >= 0; --bb) { + const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed); + if (acc + c >= desired) { + b = uint(bb); + break; + } + acc += c; + } + *sh_bucket = b; + *sh_above = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + prefix |= *sh_bucket << uint(shift); + desired -= *sh_above; + + // ensure every thread has consumed sh_bucket/sh_above before the next pass + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + atomic_store_explicit(out_count, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // emit everything above the threshold, then fill the rest from ties + const uint threshold = prefix; + + for (uint i = tid; i < ncols; i += ntg_x) { + if (ggml_top_k_f2ui(src0_row[i]) > threshold) { + const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed); + dst_row[pos] = (int32_t) i; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint i = tid; i < ncols; i += ntg_x) { + if (ggml_top_k_f2ui(src0_row[i]) == threshold) { + const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed); + if (pos < top_k) { + dst_row[pos] = (int32_t) i; + } + } + } +}