From 0b36705fc9851ce3fc2ae87360bb1933145bba56 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:54 -0500 Subject: [PATCH 1/6] llama/cuda: tiered KV cache (--kv-vram-cells): VRAM head, pinned-host tail in one VMM range --kv-vram-cells N (llama_context_params.n_kv_vram_cells) keeps cells [0, N) of each layer's K/V in device memory and the rest in pinned host memory, mapped into the same device virtual range with CUDA VMM (cuMemCreate with CU_MEM_LOCATION_TYPE_HOST), so a context can be larger than VRAM. Attention kernels are unchanged: they only touch host pages once a sequence is that deep. Output is bit-identical to an all-device cache. When an attention op's K/V range reaches the host tail, the used host rows are first copied into a VRAM staging buffer with the copy engine (a second VMM alias maps the same VRAM head pages plus the staging buffer), and the op reads the all-VRAM alias. Prefill reads the tail once per op instead of once per query tile; decode moves it by DMA instead of SMs reading host memory from inside the kernel. GGML_CUDA_KV_TIER_STAGING=0 disables staging. Bonsai 2 27B on an RTX 4070 12 GB: q8_0 K/V at the full 262,144-token window (9.1 GB of KV), first ~94k positions in VRAM. Staging: prefill in the tail 146 -> 280 tok/s at 180k, decode +28% at 180k. Plumbed through llama_cparams to the standard KV cache and the hybrid (attention + recurrent) memory via a trailing defaulted constructor parameter; other memory types are untouched. Co-Authored-By: Claude Opus 5.5 --- common/arg.cpp | 12 ++ common/common.cpp | 1 + common/common.h | 3 + ggml/include/ggml-cuda.h | 4 + ggml/src/ggml-cuda/fattn.cu | 30 +++ ggml/src/ggml-cuda/ggml-cuda.cu | 321 +++++++++++++++++++++++++++++++- include/llama.h | 6 + src/llama-context.cpp | 2 + src/llama-cparams.h | 1 + src/llama-kv-cache.cpp | 34 +++- src/llama-kv-cache.h | 3 +- src/llama-memory-hybrid.cpp | 6 +- src/llama-memory-hybrid.h | 3 +- src/llama-model.cpp | 6 +- 14 files changed, 418 insertions(+), 14 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 94dead3ff3ba..d611bb28c154 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2457,6 +2457,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.kv_mean_center_path = value; } ).set_env("LLAMA_ARG_KV_MEAN_CENTER")); + add_opt(common_arg( + {"--kv-vram-cells"}, "N", + "tiered KV cache: keep the first N cells of each layer's K/V in VRAM and the rest in pinned system RAM\n" + "mapped into the same device range (CUDA VMM), so the context can exceed VRAM. Positions past N are\n" + "read over PCIe once a sequence is that deep; output is identical to an all-VRAM cache. 0 = off (default)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.n_kv_vram_cells = value; + } + ).set_env("LLAMA_ARG_KV_VRAM_CELLS")); add_opt(common_arg( {"--hellaswag"}, "compute HellaSwag score over random tasks from datafile supplied with -f", diff --git a/common/common.cpp b/common/common.cpp index 4045903ae1a3..eb91804c9376 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1764,6 +1764,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & // note: params (and therefore params.kv_mean_center_path) is kept alive by the caller for // at least as long as it takes to call llama_init_from_model() with the returned cparams cparams.path_kv_mean_center = params.kv_mean_center_path.empty() ? nullptr : params.kv_mean_center_path.c_str(); + cparams.n_kv_vram_cells = (uint32_t) std::max(0, params.n_kv_vram_cells); return cparams; } diff --git a/common/common.h b/common/common.h index 0ab2a3ffb3d1..bbbb76e6f29f 100644 --- a/common/common.h +++ b/common/common.h @@ -588,6 +588,9 @@ struct common_params { // only takes effect when cache_type_k == GGML_TYPE_Q4_0; see docs/kv-mean-center.md std::string kv_mean_center_path = ""; + // tiered KV cache: cells past this many live in pinned host memory (0 = all in device memory) + int32_t n_kv_vram_cells = 0; + common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO; // multimodal models (see tools/mtmd) diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 1cd81eeaebcd..b2430496b21f 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -27,6 +27,10 @@ GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); // device buffer GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device); +// device buffer whose buffers keep vram_frac of each of n_parts equal parts in VRAM and the remainder in +// pinned host memory mapped into the same device address range (CUDA VMM). nullptr without VMM. +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_tier_buffer_type(int device, double vram_frac, int n_parts, const char * tag); + // 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); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 59e6c9ba74b0..cef901f9feb2 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -595,8 +595,31 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d return f16_extra.end - (uintptr_t) dst->data; } +// ggml-cuda.cu: host-tail staging of tiered buffers (ggml_backend_cuda_tier_buffer_type) +void * ggml_cuda_tier_stage(const void * ptr, size_t nbytes, cudaStream_t stream); + void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + + // K/V in a tiered buffer whose host tail this op reaches: copy the used host rows into the VRAM staging + // buffer with the copy engine and point K/V at the all-VRAM alias of the same range for this op. Prefill + // kernels read each K/V row once per query tile, which for host rows would be PCIe traffic every time; + // decode reads each row once, but DMA moves it faster than SMs reading host memory from inside the + // kernel (RTX 4070, 180k context, 86k cells in host memory: +28% decode). Nothing is copied for ops + // that stay below the tier line. The staged bytes are the same bytes, so results are unchanged. + ggml_tensor * K = dst->src[1]; + ggml_tensor * V = dst->src[2]; + void * K_data = K ? K->data : nullptr; + void * V_data = V ? V->data : nullptr; + if (K && V && V != K) { + if (void * a = ggml_cuda_tier_stage(K->data, ggml_nbytes(K), ctx.stream())) { + K->data = a; + } + if (void * a = ggml_cuda_tier_stage(V->data, ggml_nbytes(V), ctx.stream())) { + V->data = a; + } + } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); @@ -610,6 +633,13 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst ggml_cuda_flash_attn_ext_mma_f16(ctx, dst); break; } + + if (K) { + K->data = K_data; + } + if (V) { + V->data = V_data; + } } bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 41b3be6317d8..29082fa68d58 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -82,6 +82,7 @@ #include #include #include +#include #include #include #include @@ -732,6 +733,7 @@ struct ggml_backend_cuda_buffer_context { int device; void * dev_ptr = nullptr; std::string name; + std::function release; // set by the tiered (VRAM head + host tail) buffer; default cudaFree ggml_backend_cuda_buffer_context(int device, void * dev_ptr) : device(device), dev_ptr(dev_ptr), @@ -739,7 +741,11 @@ struct ggml_backend_cuda_buffer_context { } ~ggml_backend_cuda_buffer_context() { - CUDA_CHECK(cudaFree(dev_ptr)); + if (release) { + release(); + } else { + CUDA_CHECK(cudaFree(dev_ptr)); + } } }; @@ -873,8 +879,15 @@ static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = { struct ggml_backend_cuda_buffer_type_context { int device; std::string name; + // tiered buffer (see ggml_backend_cuda_tier_buffer_type): the buffer is n_parts equal parts; the + // first vram_frac of every part is backed by VRAM, the rest by pinned host memory mapped into the + // same virtual range. vram_frac >= 1 is an ordinary device buffer. + double vram_frac = 1.0; + int n_parts = 1; }; +static ggml_backend_buffer_t ggml_backend_cuda_tier_alloc(ggml_backend_buffer_type_t buft, size_t size); + 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; @@ -890,6 +903,10 @@ static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_bac ggml_cuda_set_device(buft_ctx->device); + if (buft_ctx->vram_frac < 1.0 && size > 0) { + return ggml_backend_cuda_tier_alloc(buft, size); + } + void * dev_ptr; cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device); if (err != cudaSuccess) { @@ -963,6 +980,293 @@ ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) { return &ggml_backend_cuda_buffer_types[device]; } +// Tiered device buffer: one contiguous device virtual range whose pages are backed partly by VRAM and +// partly by pinned host memory (cuMemCreate with a host location), so kernels address it like any +// device buffer. Used for the KV cache (llama_context_params.n_kv_vram_cells): the first vram_frac of +// every K/V tensor (the early positions) stays in VRAM, positions past that live in system RAM and are +// read over PCIe only once a sequence is that deep. Nothing in the attention kernels changes. +// +// Staging: each tiered buffer also gets a second virtual range that maps the same VRAM head pages and, in +// place of each host run, a VRAM staging buffer shared by all tiered buffers (one per part: K, V). An +// attention op whose K/V range reaches the host tail copies the used host rows into staging with the copy +// engine and reads the all-VRAM alias (ggml_cuda_tier_stage). GGML_CUDA_KV_TIER_STAGING=0 disables it. +#if defined(GGML_USE_VMM) && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +struct ggml_cuda_tier_entry { + CUdeviceptr va, va2; + size_t total; + std::vector> host_runs; // (offset, length) of the host-backed runs +}; +static std::mutex g_tier_mutex; +static std::vector g_tier_entries; + +static void ggml_cuda_tier_register(CUdeviceptr va, CUdeviceptr va2, size_t total, const std::vector> & host_runs) { + std::lock_guard lock(g_tier_mutex); + g_tier_entries.push_back({ va, va2, total, host_runs }); +} + +static void ggml_cuda_tier_unregister(CUdeviceptr va) { + std::lock_guard lock(g_tier_mutex); + for (size_t i = 0; i < g_tier_entries.size(); ++i) { + if (g_tier_entries[i].va == va) { + g_tier_entries.erase(g_tier_entries.begin() + i); + return; + } + } +} + +// VRAM staging buffers shared by every tiered buffer: slot i backs the i-th host run (K tail, V tail). +// Created at the first request and kept for the life of the process; a later, longer run gets no alias. +static bool ggml_cuda_tier_staging_handle(int physical, int slot, size_t len, CUmemGenericAllocationHandle * out) { + static std::mutex m; + static std::map> handles; + std::lock_guard lock(m); + auto it = handles.find(slot); + if (it == handles.end()) { + CUmemAllocationProp pd = {}; + pd.type = CU_MEM_ALLOCATION_TYPE_PINNED; + pd.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + pd.location.id = physical; + CUmemGenericAllocationHandle h; + if (cuMemCreate(&h, len, &pd, 0) != CUDA_SUCCESS) { + return false; + } + GGML_LOG_INFO("%s: staging buffer %d: %.2f MiB VRAM\n", __func__, slot, len/1048576.0); + it = handles.emplace(slot, std::make_pair(h, len)).first; + } + if (it->second.second < len) { + return false; + } + *out = it->second.first; + return true; +} + +// For a tensor range [ptr, ptr + nbytes) inside a tiered buffer with a staging alias: copy its host-backed +// bytes into the staging buffers (stream-ordered) and return the same range in the alias, whose pages are +// all VRAM. nullptr when ptr is not in such a buffer or no byte of the range is host-backed. +void * ggml_cuda_tier_stage(const void * ptr, size_t nbytes, cudaStream_t stream) { + const CUdeviceptr p = (CUdeviceptr) ptr; + const ggml_cuda_tier_entry * e = nullptr; + { + std::lock_guard lock(g_tier_mutex); + for (auto & it : g_tier_entries) { + if (p >= it.va && p < it.va + it.total) { + e = ⁢ + break; + } + } + } + if (!e) { + return nullptr; + } + const size_t lo = p - e->va, hi = std::min(e->total, lo + nbytes); + bool any = false; + for (auto & hr : e->host_runs) { + const size_t a = std::max(lo, hr.first), b = std::min(hi, hr.first + hr.second); + if (a < b) { + CUDA_CHECK(cudaMemcpyAsync((void *) (e->va2 + a), (const void *) (e->va + a), b - a, cudaMemcpyDeviceToDevice, stream)); + any = true; + } + } + return any ? (void *) (e->va2 + lo) : nullptr; +} + +static ggml_backend_buffer_t ggml_backend_cuda_tier_alloc(ggml_backend_buffer_type_t buft, size_t size) { + ggml_backend_cuda_buffer_type_context * bctx = (ggml_backend_cuda_buffer_type_context *) buft->context; + const int device = bctx->device; + const int physical = ggml_cuda_get_physical_device(device); + + CUmemAllocationProp pd = {}; + pd.type = CU_MEM_ALLOCATION_TYPE_PINNED; + pd.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + pd.location.id = physical; + + CUmemAllocationProp ph = {}; + ph.type = CU_MEM_ALLOCATION_TYPE_PINNED; + ph.location.type = CU_MEM_LOCATION_TYPE_HOST; + ph.location.id = 0; + + size_t gd = 0, gh = 0; + CU_CHECK(cuMemGetAllocationGranularity(&gd, &pd, CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + if (cuMemGetAllocationGranularity(&gh, &ph, CU_MEM_ALLOC_GRANULARITY_MINIMUM) != CUDA_SUCCESS || gh == 0) { + gh = gd; + } + const size_t gran = std::max(gd, gh); + const size_t total = GGML_PAD(size, gran); + const int parts = std::max(1, bctx->n_parts); + const size_t part = size / parts; + + // page p is VRAM when its start lies in the head of its part + const size_t n_pages = total / gran; + std::vector in_vram(n_pages); + size_t n_vram = 0; + for (size_t p = 0; p < n_pages; ++p) { + const size_t off = p * gran; + const size_t ip = std::min(off / std::max(part, 1), parts - 1); + const size_t head = (size_t) (bctx->vram_frac * part); + in_vram[p] = off - ip*part < head; + n_vram += in_vram[p]; + } + + CUdeviceptr va = 0; + CUresult r = cuMemAddressReserve(&va, total, gran, 0, 0); + if (r != CUDA_SUCCESS) { + GGML_LOG_ERROR("%s: cuMemAddressReserve(%.1f MiB) failed: %d\n", __func__, total/1048576.0, (int) r); + return nullptr; + } + + struct run { CUdeviceptr ptr; size_t len; CUmemGenericAllocationHandle h; }; + std::vector runs; + auto undo = [runs_p = &runs, va, total]() { + for (auto & rr : *runs_p) { + cuMemUnmap(rr.ptr, rr.len); + cuMemRelease(rr.h); + } + cuMemAddressFree(va, total); + }; + + for (size_t p = 0; p < n_pages; ) { + size_t q = p; + while (q < n_pages && in_vram[q] == in_vram[p]) { + ++q; + } + const size_t len = (q - p) * gran; + CUmemGenericAllocationHandle h; + r = cuMemCreate(&h, len, in_vram[p] ? &pd : &ph, 0); + if (r != CUDA_SUCCESS) { + const char * es = nullptr; cuGetErrorString(r, &es); + GGML_LOG_ERROR("%s: cuMemCreate(%s, %.1f MiB) failed: %s\n", __func__, in_vram[p] ? "vram" : "host", len/1048576.0, es ? es : "?"); + undo(); + return nullptr; + } + r = cuMemMap(va + p*gran, len, 0, h, 0); + if (r != CUDA_SUCCESS) { + cuMemRelease(h); + GGML_LOG_ERROR("%s: cuMemMap failed: %d\n", __func__, (int) r); + undo(); + return nullptr; + } + runs.push_back({ va + p*gran, len, h }); + p = q; + } + + CUmemAccessDesc ad = {}; + ad.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + ad.location.id = physical; + ad.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + r = cuMemSetAccess(va, total, &ad, 1); + if (r != CUDA_SUCCESS) { + GGML_LOG_ERROR("%s: cuMemSetAccess failed: %d\n", __func__, (int) r); + undo(); + return nullptr; + } + + GGML_LOG_INFO("%s: %s %.2f MiB: %.2f MiB VRAM + %.2f MiB host (%zu runs, %d parts)\n", __func__, bctx->name.c_str(), + size/1048576.0, n_vram*gran/1048576.0, (n_pages - n_vram)*gran/1048576.0, runs.size(), parts); + + // staging alias (see the comment above ggml_cuda_tier_entry) + CUdeviceptr va2 = 0; + std::vector> host_runs; + { + static const bool staging = [] { const char * e = getenv("GGML_CUDA_KV_TIER_STAGING"); return !e || atoi(e) != 0; }(); + if (staging && cuMemAddressReserve(&va2, total, gran, 0, 0) == CUDA_SUCCESS) { + bool ok = true; + int ih = 0; + std::vector mapped; + for (auto & rr : runs) { + const size_t off = rr.ptr - va; + CUmemGenericAllocationHandle h = rr.h; + if (!in_vram[off / gran]) { + CUmemGenericAllocationHandle bh; + if (!ggml_cuda_tier_staging_handle(physical, ih++, rr.len, &bh)) { + ok = false; + break; + } + h = bh; + host_runs.push_back({ off, rr.len }); + } + if (cuMemMap(va2 + off, rr.len, 0, h, 0) != CUDA_SUCCESS) { + ok = false; + break; + } + mapped.push_back(va2 + off); + } + ok = ok && cuMemSetAccess(va2, total, &ad, 1) == CUDA_SUCCESS; + if (!ok) { + for (size_t i = 0; i < mapped.size(); ++i) { + cuMemUnmap(mapped[i], runs[i].len); + } + cuMemAddressFree(va2, total); + va2 = 0; + host_runs.clear(); + GGML_LOG_WARN("%s: no staging alias for %s: its host tail is read over PCIe in place\n", __func__, bctx->name.c_str()); + } else { + ggml_cuda_tier_register(va, va2, total, host_runs); + } + } + } + + ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(device, (void *) va); + ctx->release = [runs, va, va2, total, device]() { + ggml_cuda_set_device(device); + cudaDeviceSynchronize(); + if (va2) { + ggml_cuda_tier_unregister(va); + for (auto & rr : runs) { + cuMemUnmap(va2 + (rr.ptr - va), rr.len); + } + cuMemAddressFree(va2, total); + } + for (auto & rr : runs) { + cuMemUnmap(rr.ptr, rr.len); + cuMemRelease(rr.h); + } + cuMemAddressFree(va, total); + }; + return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size); +} +#else +void * ggml_cuda_tier_stage(const void * ptr, size_t nbytes, cudaStream_t stream) { + GGML_UNUSED(ptr); GGML_UNUSED(nbytes); GGML_UNUSED(stream); + return nullptr; +} + +static ggml_backend_buffer_t ggml_backend_cuda_tier_alloc(ggml_backend_buffer_type_t buft, size_t size) { + GGML_UNUSED(size); + GGML_LOG_ERROR("%s: tiered buffers need CUDA VMM (build without GGML_CUDA_NO_VMM)\n", __func__); + GGML_UNUSED(buft); + return nullptr; +} +#endif + +// A device buffer type whose buffers keep vram_frac of each of n_parts equal parts in VRAM and the rest in +// pinned host memory mapped into the same device range. `tag` makes the name unique so a caller can ask +// for one buffer per tensor group. Returns nullptr when the device has no VMM. +ggml_backend_buffer_type_t ggml_backend_cuda_tier_buffer_type(int device, double vram_frac, int n_parts, const char * tag) { + static std::mutex mutex; + std::lock_guard lock(mutex); + static std::map cache; + + if (device < 0 || device >= ggml_backend_cuda_get_device_count() || !ggml_cuda_info().devices[device].vmm) { + return nullptr; + } + const std::string name = std::string(GGML_CUDA_NAME) + std::to_string(device) + "_tier_" + (tag ? tag : "") + + "_" + std::to_string((int) (vram_frac * 1e6)) + "_" + std::to_string(n_parts); + auto it = cache.find(name); + if (it != cache.end()) { + return it->second; + } + auto * ctx = new ggml_backend_cuda_buffer_type_context{device, name}; + ctx->vram_frac = std::max(0.0, std::min(1.0, vram_frac)); + ctx->n_parts = std::max(1, n_parts); + auto * buft = new ggml_backend_buffer_type{ + /* .iface = */ ggml_backend_cuda_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device), + /* .context = */ ctx, + }; + cache[name] = buft; + return buft; +} + // Communication context for multi-GPU AllReduce during tensor parallelism. // // Created once per meta backend instance. Resources for the selected mode @@ -2447,7 +2751,7 @@ static void ggml_backend_cuda_set_tensor_async(ggml_backend_t backend, ggml_tens 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"); + GGML_ASSERT(ggml_backend_buft_is_cuda(buf->buft) && ((ggml_backend_cuda_buffer_type_context *) buf->buft->context)->device == cuda_ctx->device && "unsupported buffer type"); CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream())); } @@ -2456,7 +2760,7 @@ static void ggml_backend_cuda_get_tensor_async(ggml_backend_t backend, const ggm 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"); + GGML_ASSERT(ggml_backend_buft_is_cuda(buf->buft) && ((ggml_backend_cuda_buffer_type_context *) buf->buft->context)->device == cuda_ctx->device && "unsupported buffer type"); CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cuda_ctx->stream())); } @@ -2466,7 +2770,7 @@ static void ggml_backend_cuda_set_tensor_2d_async(ggml_backend_t backend, struct 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"); + GGML_ASSERT(ggml_backend_buft_is_cuda(buf->buft) && ((ggml_backend_cuda_buffer_type_context *) buf->buft->context)->device == 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())); @@ -2477,7 +2781,7 @@ static void ggml_backend_cuda_get_tensor_2d_async(ggml_backend_t backend, const 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"); + GGML_ASSERT(ggml_backend_buft_is_cuda(buf->buft) && ((ggml_backend_cuda_buffer_type_context *) buf->buft->context)->device == 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())); @@ -4695,12 +4999,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // On integrated GPUs (APUs, e.g. RDNA3.5) the scheduler may place a // node's output on the host-visible buffer, which the compute path // handles. Allow that here, mirroring the src-tensor check below. - assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || + assert(ggml_backend_buft_is_cuda(node->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->buffer->buft))); 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) || + assert(ggml_backend_buft_is_cuda(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); } } @@ -6041,6 +6345,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + if (strcmp(name, "ggml_backend_cuda_tier_buffer_type") == 0) { + return (void *)ggml_backend_cuda_tier_buffer_type; + } return nullptr; } diff --git a/include/llama.h b/include/llama.h index deb1e12828c6..555ff32dc51d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -396,6 +396,12 @@ extern "C" { // see tools/kv-mean-center to generate this file and docs/kv-mean-center.md for details. const char * path_kv_mean_center; + // tiered KV cache [EXPERIMENTAL]: keep only the first n_kv_vram_cells cells of each K/V tensor in + // device memory and the rest in pinned host memory mapped into the same device range (CUDA VMM), so + // the context can be larger than what fits in VRAM. Positions past the line are read over PCIe, only + // once a sequence is that deep. Output is identical to an all-device cache. 0 = all in device memory. + uint32_t n_kv_vram_cells; + // Abort callback // if it returns true, execution of llama_decode() will be aborted // currently works only with CPU execution diff --git a/src/llama-context.cpp b/src/llama-context.cpp index cb6cf1f2854a..6daf5c6745a1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -178,6 +178,7 @@ llama_context::llama_context( const auto & hparams = model.hparams; cparams.n_seq_max = std::max(1u, params.n_seq_max); + cparams.n_kv_vram_cells = params.n_kv_vram_cells; if (cparams.n_seq_max > LLAMA_MAX_SEQ) { throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); } @@ -3822,6 +3823,7 @@ llama_context_params llama_context_default_params() { /*.type_k =*/ GGML_TYPE_F16, /*.type_v =*/ GGML_TYPE_F16, /*.path_kv_mean_center =*/ nullptr, + /*.n_kv_vram_cells =*/ 0, /*.abort_callback =*/ nullptr, /*.abort_callback_data =*/ nullptr, /*.embeddings =*/ false, diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 1f013e2f0027..af3487e4e542 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -18,6 +18,7 @@ struct llama_cparams { uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback uint32_t n_outputs_max; // max outputs supported by the context uint32_t n_outputs_max_per_seq; + uint32_t n_kv_vram_cells; // tiered KV: cells past this live in host memory (0 = all in device memory) int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3df5cedd082e..723737acb7bf 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -80,7 +80,8 @@ llama_kv_cache::llama_kv_cache( llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share) : + const layer_share_cb & share, + uint32_t n_kv_vram_cells) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -163,6 +164,12 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); + const uint32_t tier_cells = GGML_PAD(n_kv_vram_cells, n_pad); + if (tier_cells > 0 && tier_cells < kv_size) { + LLAMA_LOG_INFO("%s: tiered KV: cells [0, %u) in VRAM, [%u, %u) in host memory\n", + __func__, tier_cells, tier_cells, kv_size); + } + for (uint32_t il = 0; il < n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -219,6 +226,31 @@ llama_kv_cache::llama_kv_cache( buft = ggml_backend_dev_buffer_type(dev); dev_name = ggml_backend_dev_name(dev); + + // n_kv_vram_cells = N: tiered KV. Cells [0, N) of this layer's K and V stay in VRAM, cells + // [N, kv_size) live in pinned system RAM mapped into the same device range (CUDA VMM), so the + // window can exceed VRAM and the extra positions cost PCIe reads only once a sequence reaches + // them. Kernels are unchanged and results are identical to an all-VRAM cache. Backends without + // ggml_backend_cuda_tier_buffer_type (or without VMM) keep the ordinary device buffer. + if (tier_cells > 0 && tier_cells < kv_size && n_stream == 1 && !is_mla && + n_embd_k_gqa == n_embd_v_gqa && type_k == type_v) { + using tier_fn_t = ggml_backend_buffer_type_t (*)(int, double, int, const char *); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + auto * fn = reg ? (tier_fn_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_tier_buffer_type") : nullptr; + const char * dn = ggml_backend_dev_name(dev); + int dev_idx = -1; + if (dn && strncmp(dn, "CUDA", 4) == 0) { + dev_idx = atoi(dn + 4); + } + if (fn && dev_idx >= 0) { + char tag[32]; + snprintf(tag, sizeof(tag), "kv%p_l%u", (void *) this, il); + ggml_backend_buffer_type_t tb = fn(dev_idx, (double) tier_cells / (double) kv_size, 2, tag); + if (tb) { + buft = tb; + } + } + } } LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 4d2229b90274..ee66e5c1fccc 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -112,7 +112,8 @@ class llama_kv_cache : public llama_memory_i { llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share); + const layer_share_cb & share, + uint32_t n_kv_vram_cells = 0); // tiered KV: cells past this live in host memory (0 = all VRAM) ~llama_kv_cache() = default; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..a0fe07e6b3cc 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -29,7 +29,8 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + uint32_t n_kv_vram_cells) : hparams(model.hparams), mem_attn(new llama_kv_cache( model, @@ -49,7 +50,8 @@ llama_memory_hybrid::llama_memory_hybrid( [&](int32_t il) { return !hparams.is_recr(il); } : filter_attn, nullptr, - nullptr + nullptr, + n_kv_vram_cells )), mem_recr(new llama_memory_recurrent( model, diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..52ea297abc4b 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,8 @@ class llama_memory_hybrid : public llama_memory_i { bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + uint32_t n_kv_vram_cells = 0); // tiered KV for the attention layers (see llama_context_params) ~llama_memory_hybrid() = default; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c2ac6a50b316..c607ba49a685 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2831,7 +2831,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* n_kv_vram_cells */ cparams.n_kv_vram_cells); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; @@ -2933,7 +2934,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, filter, nullptr, - nullptr); + nullptr, + cparams.n_kv_vram_cells); } } } From 50a0d87cefa90c6ed0c747c7a0ea6e1bc78ac1f1 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:54 -0500 Subject: [PATCH 2/6] cuda: quantized-KV GQA decode on the in-place MMA flash-attention kernel For up to 8 queries (single-token decode and speculative verify) with q4_0/q8_0 K/V and a GQA ratio above 4, take the MMA kernel that reads K/V in place instead of the vector kernel. The vector kernel runs one block per Q head, so every K/V row is fetched gqa_ratio times (6x for Qwen3.5 / Bonsai 2: 24 Q heads over 4 KV heads); the MMA kernel packs the GQA heads of one KV head into one tile. Same rule the F16 path uses. Bonsai 2 27B, RTX 4070, q8_0 K/V, MTP draft n-max 2: 4k 78.1 -> 86.9 tok/s, 16k 77.7 -> 111.6. GGML_CUDA_FA_MMA_DECODE_MIN_KV sets the shortest KV length that takes the route (default 256, 0 = never). Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-cuda/fattn.cu | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index cef901f9feb2..e099ec3741fc 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -475,6 +475,23 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { + // Quantized-KV decode with wide GQA (up to 8 queries: single-token decode and speculative verify): + // the vector kernel runs one block per Q head, so every K/V row is fetched gqa_ratio times (6x for + // Qwen3.5/Bonsai 2: 24 Q heads over 4 KV heads). The MMA kernel packs the GQA heads of one KV head + // into one tile and reads q4_0/q8_0 K/V in place once. Same rule the F16 path already uses. + // RTX 4070, Bonsai 2 27B, q8_0 K/V, MTP draft: 4k 78.1 -> 86.9 tok/s, 16k 77.7 -> 111.6, 32k 54.5 -> 103.6. + // GGML_CUDA_FA_MMA_DECODE_MIN_KV = shortest KV length that takes this route (default 256, 0 = never). + { + static const int min_kv = [] { + const char * e = getenv("GGML_CUDA_FA_MMA_DECODE_MIN_KV"); + return e ? atoi(e) : 256; + }(); + if (min_kv > 0 && ggml_is_quantized(K->type) && K->type == V->type && + ggml_cuda_fattn_mma_kv_native_supported(dst) && + gqa_opt_applies && gqa_ratio > 4 && Q->ne[1] <= 8 && Q->ne[3] == 1 && K->ne[1] >= min_kv) { + return BEST_FATTN_KERNEL_MMA_F16; + } + } if (can_use_vector_kernel) { // batch-invariant mode: the same (vector) kernel for 1 to 8 queries, so a token verified in a // speculative batch attends with the same arithmetic as a token decoded alone From dc0cd6c8b8b7a7f928674e8065b4309fd719c200 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:54 -0500 Subject: [PATCH 3/6] cuda: BATCH_INVARIANT: occupancy-independent FA KV split, PTQ1_0 mat-vec up to 8 columns Under GGML_CUDA_BATCH_INVARIANT: - the non-stream-k flash-attention KV split is sized from a fixed blocks-per-SM value instead of the occupancy of the template instance that runs; the 1-query and multi-query instances differ in registers and shared memory, so their splits (and combine order) differed. - PTQ1_0 batches up to MMVQ_MAX_BATCH_SIZE stay on the PT mat-vec, whose per-column arithmetic does not depend on the column count; a 5-column speculative verify on MMQ did not match the same tokens decoded alone. Measured cost: pp5 172 -> 162 tok/s, only on 5-8 column batches. GGML_CUDA_PTQ1_MMVQ_MAX overrides the mat-vec / MMQ crossover (default 4, confirmed: MMQ wins from 5). Not covered: the stream-k split of the MMA kernel follows the padded KV length (and the tile instance follows the query count), so attention in a verify batch is not bit-identical to single-token decode past ~32k, or at any depth on the MMA decode route. Measured at 4k/20k/40k: the weight path matches, a few continuations diverge at the rounding level. Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-cuda/fattn-common.cuh | 7 +++++++ ggml/src/ggml-cuda/mmvq.cu | 11 +++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index ae4318cbc881..85d34024fc68 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1112,6 +1112,13 @@ void launch_fattn( int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy. CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared)); GGML_ASSERT(max_blocks_per_sm > 0); + // Batch-invariant mode: the KV split must not depend on which template instance runs. Occupancy differs + // between the 1-query and the multi-query instances (registers, shared memory), so size the split from a + // fixed blocks-per-SM instead; that only shifts work between waves. + // (Stream-k launches keep the real occupancy: their instance is fixed per batch size range already.) + if (ggml_cuda_batch_invariant() && !stream_k) { + max_blocks_per_sm = 4; + } int parallel_blocks = max_blocks_per_sm; const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index b8279e40abb5..8b39c2fd6f2f 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -299,8 +299,15 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { if (type == GGML_TYPE_PTQ1_0 && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) { // The PT mat-vec path shares the weight decode across columns; with the branch-free PTQ1_0 // MMQ tile loader the tile path overtakes it at 5+ columns on Ada (RTX 4070, Bonsai 2 27B: - // mat-vec 155 t/s vs MMQ 244 t/s at n=8). - return ne11 <= 4; + // mat-vec 155 t/s vs MMQ 244 t/s at n=8; at n=5 172 vs 162). Under GGML_CUDA_BATCH_INVARIANT every + // batch up to MMVQ_MAX_BATCH_SIZE stays on the PT mat-vec, whose per-column arithmetic does not depend + // on the column count: a 5-column speculative verify on MMQ would not match the same tokens decoded + // alone. GGML_CUDA_PTQ1_MMVQ_MAX overrides the crossover. + static const int max_cols = [] { + const char * e = getenv("GGML_CUDA_PTQ1_MMVQ_MAX"); + return e ? atoi(e) : (ggml_cuda_batch_invariant() ? MMVQ_MAX_BATCH_SIZE : 4); + }(); + return ne11 <= max_cols; } #endif // k-quants cost more to decode and mvq redoes that per column, so MMQ wins sooner. From d2e29649b20916fbf5a03e759bc1f4fda4296c34 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:55 -0500 Subject: [PATCH 4/6] speculative: --spec-draft-window and --spec-draft-n-max-tail (MTP drafting at every depth) --spec-draft-window N: the draft (MTP) context keeps only the last N rows. The server drops older rows before feeding each batch and the draft context is sized for the window, so its cells are reused, its cache stays small and on the device, and a draft pass costs the same at any depth. The MTP head predicts the next few tokens from recent context: a 16k window accepts as many drafts as the full history at 131k. Also sizes the draft context for --spec-draft-depth-max when no window is set. --spec-draft-n-max-tail N: draft size once the sequence reaches --kv-vram-cells. Past the tiered-KV line a step is bound by reading the host tail over PCIe, and a wider verify reads it once for all columns. The drafter and the output limits are built for max(n_max, n_max_tail); each slot caps a draft by depth, and the MTP draft loop now honours that per-draft cap. Bonsai 2 27B on an RTX 4070 with the MMA decode route: drafting pays at every depth, so the --spec-draft-depth-max cutoff is no longer needed (32k 54.5 -> 103.6 tok/s, 64k 48.0 -> 90.1, same acceptance; window alone +17-18% at 131k; tail draft 4 vs 2: +26% code / +15% prose at 180k). Co-Authored-By: Claude Opus 5.5 --- common/arg.cpp | 22 +++++++++++++++++++ common/common.h | 18 +++++++++++---- common/speculative.cpp | 23 +++++++++++++++++-- tools/server/server-context.cpp | 39 ++++++++++++++++++++++++++++++++- 4 files changed, 95 insertions(+), 7 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index d611bb28c154..08979aa92b14 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4140,6 +4140,28 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_depth_max = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_DEPTH_MAX")); + add_opt(common_arg( + {"--spec-draft-window"}, "N", + string_format("draft (MTP) context keeps only the last N rows, so its cache stays small and a draft pass costs\n" + "the same at any depth; 0 = full history (default: %d)", params.speculative.draft.n_window), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.draft.n_window = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_WINDOW")); + add_opt(common_arg( + {"--spec-draft-n-max-tail"}, "N", + string_format("draft size once the sequence reaches --kv-vram-cells, where a wider verify reads the host tail\n" + "once for all columns; 0 = --spec-draft-n-max (default: %d)", params.speculative.draft.n_max_tail), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.draft.n_max_tail = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX_TAIL")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.h b/common/common.h index bbbb76e6f29f..17a5189a87a0 100644 --- a/common/common.h +++ b/common/common.h @@ -329,12 +329,22 @@ struct common_params_speculative_draft { float p_split = 0.1f; // speculative decoding split probability float p_min = 0.0f; // minimum speculative decoding probability (greedy) - // stop drafting once the sequence is this long (0 = never). Deep in the context a step is bound - // by reading the KV cache, and the draft passes plus the multi-column verify add to that read - // without shortening it: on a 4070 with Bonsai 2 27B the draft is +85% at zero depth, breaks - // even near 24k tokens and costs 30% at 64k. Past the cutoff the slot decodes one token per step. + // stop drafting once the sequence is this long (0 = never). Measured when a deep verify batch ran on the + // vector FA kernel and the draft context grew with the conversation (Bonsai 2 27B on a 4070: +85% at zero + // depth, even near 24k, -30% at 64k). With quantized-KV decode on the MMA FA kernel and n_window below, + // drafting pays at every depth (64k: 48 -> 90 tok/s), so the server default is 0. int32_t n_depth_max = 0; + // draft (MTP) context window (0 = full history): it keeps only the last n_window rows and is sized for + // them, so its cells are reused, its cache stays small and on the device, and a draft pass costs the same + // at any depth. The MTP head predicts the next few tokens from recent context: 16k rows accept as many + // drafts as the full history at 131k. + int32_t n_window = 0; + + // draft size once the sequence reaches the tiered-KV line (--kv-vram-cells; 0 = n_max). Past it a step + // is bound by reading the host tail over PCIe and a wider verify batch reads it once for all columns. + int32_t n_max_tail = 0; + bool backend_sampling = true; // offload draft sampling to the backend (default: on) common_params_model mparams; diff --git a/common/speculative.cpp b/common/speculative.cpp index 6c6c1fbd5079..2a6de502a59b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2688,7 +2688,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + // the caller's per-draft cap (the server's depth-dependent draft size) as well as the global one + if (params.n_max <= (int) result.size() || (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; continue; @@ -3440,8 +3441,26 @@ common_speculative_init_result::common_speculative_init_result( cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; } - // the draft context holds as many tokens per sequence as the target context + // the draft context holds as many tokens per sequence as the target context, unless the MTP draft only + // ever sees a bounded part of the sequence: the last --spec-draft-window rows (the server drops older ones + // before each batch), or rows up to --spec-draft-depth-max (never fed past it). Sizing it for that keeps + // its cache small; with the window its cells are reused and its attention span stays that short. cparams.n_ctx = llama_n_ctx(ctx_tgt); + if (spec_mtp) { + const int32_t window = params.speculative.draft.n_window; + const int32_t depth_max = params.speculative.draft.n_depth_max; + const int32_t span = window > 0 ? window : depth_max; + if (span > 0) { + const uint32_t n_seq = std::max(1, cparams.n_seq_max); + const uint32_t need = (uint32_t) span + 2*std::max(cparams.n_batch, cparams.n_ubatch) + 256; + const uint32_t capped = GGML_PAD(need, 256) * n_seq; + if (capped < cparams.n_ctx) { + LOG_INF("%s: draft context sized to %u tokens (%s %d, target %u)\n", __func__, capped, + window > 0 ? "spec-draft-window" : "spec-draft-depth-max", span, cparams.n_ctx); + cparams.n_ctx = capped; + } + } + } // n_rs_seq stays as common_context_params_to_llama set it: the draft context needs the same rollback window as the target, with n_rs_seq == 0 its seq_rm fails silently on partial acceptance and keeps stale positions cparams.ctx_other = ctx_tgt; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 988d9f365190..5ee7316bb617 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -207,6 +207,9 @@ struct server_slot { // speculative decoding common_speculative * spec; int32_t spec_depth_max = 0; // no drafting once the sequence is longer than this (0 = always draft) + int32_t spec_n_max = 0; // draft size (--spec-draft-n-max) below spec_tail_depth + int32_t spec_n_max_tail = 0; // draft size from spec_tail_depth on (--spec-draft-n-max-tail, 0 = spec_n_max) + int32_t spec_tail_depth = 0; // start of the tiered-KV host tail (--kv-vram-cells, 0 = none) llama_tokens spec_draft; llama_tokens spec_prompt; @@ -459,6 +462,14 @@ struct server_slot { n_draft_max = std::min(n_draft_max, n_remaining() - 1); } + // past the tiered-KV line a step is bound by reading the host tail over PCIe, and a wider verify + // batch reads it once for all columns: draft --spec-draft-n-max-tail there, --spec-draft-n-max below + const bool tail = spec_tail_depth > 0 && spec_n_max_tail > 0 && prompt.n_tokens() >= spec_tail_depth; + const int32_t cap = tail ? spec_n_max_tail : spec_n_max; + if (cap > 0) { + n_draft_max = std::min(n_draft_max, (int) cap); + } + SLT_DBG(*this, "max possible draft: %d\n", n_draft_max); return n_draft_max; @@ -850,6 +861,7 @@ struct server_context_impl { llama_context * ctx_dft = nullptr; common_speculative_init_result_ptr spec_init; + int32_t spec_n_max = 0; // --spec-draft-n-max as given (the drafter is built for max(n_max, n_max_tail)) common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -972,6 +984,12 @@ struct server_context_impl { const bool is_resume = sleeping; params_base = params; + + // the drafter and the output limits are built for the larger of the two draft sizes; each slot caps a + // draft to --spec-draft-n-max below the tiered-KV line (see server_slot::get_n_draft_max) + spec_n_max = params_base.speculative.draft.n_max; + params_base.speculative.draft.n_max = std::max(params_base.speculative.draft.n_max, params_base.speculative.draft.n_max_tail); + const auto output_limits = server_output_limits(params_base); params_base.n_outputs_max = output_limits.total; params_base.n_outputs_max_per_seq = output_limits.per_seq; @@ -1245,7 +1263,10 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.spec_depth_max = params_base.speculative.draft.n_depth_max; + slot.spec_depth_max = params_base.speculative.draft.n_depth_max; + slot.spec_n_max = spec_n_max; + slot.spec_n_max_tail = params_base.speculative.draft.n_max_tail; + slot.spec_tail_depth = params_base.n_kv_vram_cells; slot.n_ctx = n_ctx_slot; slot.stats.speculative = slot.can_speculate(); @@ -3675,6 +3696,22 @@ struct server_context_impl { } } + // --spec-draft-window: before feeding this view to the draft context, drop its rows older than the window. + // The draft context is sized for the window (common_speculative_init), so freed cells are reused and its + // attention span stays short; the MTP head predicts the next few tokens from recent context. + if (spec_process && ctx_dft && params_base.speculative.draft.n_window > 0) { + llama_pos pos_min = std::numeric_limits::max(); + for (int i = 0; i < batch_view.n_tokens; ++i) { + pos_min = std::min(pos_min, batch_view.pos[i]); + } + const llama_pos hi = pos_min - params_base.speculative.draft.n_window; + if (hi > 0) { + for (auto & slot : slots) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, 0, hi); + } + } + } + if (spec_process) { bool ok = true; queue_tasks.yield_to_queue([&]() { From a9c07107e643c7752e59ad087cb3e351e365be49 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:55 -0500 Subject: [PATCH 5/6] server: --reasoning-effort-allow/-fallback and --reasoning-max-tokens-floor --reasoning-effort-allow LIST: reasoning_effort values passed to the chat template; any other value (from the request's reasoning_effort or chat_template_kwargs) becomes --reasoning-effort-fallback instead of reaching a template that raises on it. Bonsai 2's template accepts xhigh/medium/low and raises on "high", which Cline, Kilo and Open WebUI send: every request failed with HTTP 500. --reasoning-max-tokens-floor N: with thinking on, a client max_tokens below N is raised to N. A thinking model spends its first thousands of tokens inside the thinking block; 256-4096 token app caps (PrismML's quickstart, SillyTavern, AnythingLLM, Continue) end the request there with no answer. --reasoning-budget still bounds the thinking. Both are off by default. Co-Authored-By: Claude Opus 5.5 --- common/arg.cpp | 27 +++++++++++++++++++++++++ common/common.h | 6 ++++++ tools/server/server-common.cpp | 35 +++++++++++++++++++++++++++++++++ tools/server/server-common.h | 3 +++ tools/server/server-context.cpp | 5 ++++- 5 files changed, 75 insertions(+), 1 deletion(-) diff --git a/common/arg.cpp b/common/arg.cpp index 08979aa92b14..128130f0df19 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3713,6 +3713,33 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sampling.reasoning_budget_message = value; } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); + add_opt(common_arg( + {"--reasoning-effort-allow"}, "LIST", + "comma-separated reasoning_effort values to pass to the chat template; any other value in a request (or\n" + "chat_template_kwargs) is replaced by --reasoning-effort-fallback instead of reaching a template that\n" + "raises on it (e.g. \"high\" -> HTTP 500). Empty = pass everything (default)", + [](common_params & params, const std::string & value) { + params.reasoning_effort_allow = string_split(value, ','); + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REASONING_EFFORT_ALLOW")); + add_opt(common_arg( + {"--reasoning-effort-fallback"}, "WORD", + string_format("reasoning_effort used in place of a value --reasoning-effort-allow rejects (default: %s)", params.reasoning_effort_fallback.c_str()), + [](common_params & params, const std::string & value) { + params.reasoning_effort_fallback = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REASONING_EFFORT_FALLBACK")); + add_opt(common_arg( + {"--reasoning-max-tokens-floor"}, "N", + "with thinking on, raise a client max_tokens below N to N, so a small app cap does not end the request\n" + "inside the thinking block with no answer; --reasoning-budget still bounds the thinking. 0 = off (default)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.reasoning_max_tokens_floor = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REASONING_MAX_TOKENS_FLOOR")); add_opt(common_arg( {"--reasoning-preserve"}, {"--no-reasoning-preserve"}, diff --git a/common/common.h b/common/common.h index 17a5189a87a0..506456ec80e3 100644 --- a/common/common.h +++ b/common/common.h @@ -656,6 +656,12 @@ struct common_params { bool force_pure_content_parser = false; common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; int enable_reasoning = -1; // -1 = auto, 0 = disable, 1 = enable + + // server: reasoning_effort words passed to the chat template (empty = any); others become the fallback + std::vector reasoning_effort_allow; + std::string reasoning_effort_fallback = "medium"; + // server: with thinking on, raise a client output cap below this (0 = off) + int32_t reasoning_max_tokens_floor = 0; bool prefill_assistant = true; // if true, any trailing assistant message will be prefilled into the response int sleep_idle_seconds = -1; // if >0, server will sleep after this many seconds of idle time diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index b5e052a9e84c..48466405e5ed 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1321,6 +1321,26 @@ json oaicompat_chat_params_parse( } } + // --reasoning-effort-allow: an effort word the chat template does not accept (from the request or from + // chat_template_kwargs) becomes --reasoning-effort-fallback instead of reaching the template. Bonsai 2's + // template raises on "high", which Cline, Kilo and Open WebUI send: HTTP 500 on every request. + if (!opt.reasoning_effort_allow.empty()) { + auto it = inputs.chat_template_kwargs.find("reasoning_effort"); + if (it != inputs.chat_template_kwargs.end()) { + std::string effort; + try { + effort = json::parse(it->second).get(); + } catch (...) { + effort = it->second; + } + const auto & allow = opt.reasoning_effort_allow; + if (std::find(allow.begin(), allow.end(), effort) == allow.end()) { + SRV_INF("reasoning_effort \"%s\" is not in --reasoning-effort-allow: using \"%s\"\n", effort.c_str(), opt.reasoning_effort_fallback.c_str()); + it->second = json(opt.reasoning_effort_fallback).dump(); + } + } + } + inputs.force_pure_content = opt.force_pure_content; // Apply chat template to the list of messages @@ -1388,6 +1408,21 @@ json oaicompat_chat_params_parse( } } + // --reasoning-max-tokens-floor: with thinking on, a client output cap below the floor is raised to it. A + // thinking model spends its first thousands of tokens inside the thinking block; a 256-4096 token app cap + // ends the request there with no answer. The reasoning budget still bounds the thinking. + if (opt.reasoning_max_tokens_floor > 0 && inputs.enable_thinking && !chat_params.thinking_end_tags.empty()) { + for (const char * key : {"n_predict", "max_tokens", "max_completion_tokens"}) { + if (llama_params.contains(key) && llama_params.at(key).is_number_integer()) { + const int v = llama_params.at(key).get(); + if (v > 0 && v < opt.reasoning_max_tokens_floor) { + SRV_INF("%s %d raised to %d (thinking on, --reasoning-max-tokens-floor)\n", key, v, opt.reasoning_max_tokens_floor); + llama_params[key] = opt.reasoning_max_tokens_floor; + } + } + } + } + return llama_params; } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index a624ae6a07fd..9f697ce25526 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -310,6 +310,9 @@ struct server_chat_params { std::string reasoning_budget_message; std::string media_path; bool force_pure_content = false; + std::vector reasoning_effort_allow; // --reasoning-effort-allow (empty = pass everything) + std::string reasoning_effort_fallback; // --reasoning-effort-fallback + int reasoning_max_tokens_floor = 0; // --reasoning-max-tokens-floor (0 = off) }; // used by /completions endpoint diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5ee7316bb617..5e3420509ce2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1463,7 +1463,10 @@ struct server_context_impl { /* reasoning_budget */ params_base.sampling.reasoning_budget_tokens, /* reasoning_budget_msg */ params_base.sampling.reasoning_budget_message, /* media_path */ params_base.media_path, - /* force_pure_content */ params_base.force_pure_content_parser + /* force_pure_content */ params_base.force_pure_content_parser, + /* reasoning_effort_allow */ params_base.reasoning_effort_allow, + /* reasoning_effort_fallback */ params_base.reasoning_effort_fallback, + /* reasoning_max_tok_floor */ params_base.reasoning_max_tokens_floor }; { From 31640357451f9d370f49f4284b803295399e6c5c Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:55 -0500 Subject: [PATCH 6/6] cuda: GGML_CUDA_OP_TIMING=1 per-node GPU time breakdown With CUDA graphs off (GGML_CUDA_DISABLE_GRAPHS=1), records an event before every graph node on the main stream, charges the time to the next event to that node (a fused group to its first node), aggregates by op/type/shape and prints the top 30 every 100 graphs. Launch gaps inflate the small ops; the large mat-vecs read true. Diagnostics only; off by default. Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-cuda/ggml-cuda.cu | 56 +++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 29082fa68d58..81cad7c44541 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4823,8 +4823,20 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud cuda_ctx->gdn_gathers().reset(); cuda_ctx->fwht_q8().reset(); + // GGML_CUDA_OP_TIMING=1 (with CUDA graphs off): an event on the main stream before every node; the + // time to the next event is charged to that node (a fused group to its first node). Printed below. + static const bool op_timing = getenv("GGML_CUDA_OP_TIMING") != nullptr; + const bool timing = op_timing && !use_cuda_graph; + std::vector> op_events; + for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; + if (timing) { + cudaEvent_t e; + CUDA_CHECK(cudaEventCreate(&e)); + CUDA_CHECK(cudaEventRecord(e, cuda_ctx->stream(cuda_ctx->device, 0))); + op_events.push_back({ node, e }); + } if (is_concurrent_event_active) { GGML_ASSERT(concurrent_event); @@ -5022,6 +5034,50 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud try_launch_concurrent_event(node); } } + + if (timing && !op_events.empty()) { + cudaEvent_t end; + CUDA_CHECK(cudaEventCreate(&end)); + CUDA_CHECK(cudaEventRecord(end, cuda_ctx->stream(cuda_ctx->device, 0))); + CUDA_CHECK(cudaEventSynchronize(end)); + struct acc { double ms = 0; int64_t n = 0; }; + static std::map totals; + static int64_t n_graphs = 0; + static double ms_graphs = 0; + for (size_t k = 0; k < op_events.size(); ++k) { + float ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&ms, op_events[k].second, k + 1 < op_events.size() ? op_events[k + 1].second : end)); + const ggml_tensor * t = op_events[k].first; + char key[160]; + if (t->op == GGML_OP_MUL_MAT && t->src[0]) { + snprintf(key, sizeof(key), "MUL_MAT %s %lldx%lld cols=%lld", ggml_type_name(t->src[0]->type), + (long long) t->src[0]->ne[0], (long long) t->src[0]->ne[1], (long long) t->ne[1]); + } else { + snprintf(key, sizeof(key), "%s%s%s cols=%lld", ggml_op_desc(t), t->src[0] ? " " : "", + t->src[0] ? ggml_type_name(t->src[0]->type) : "", (long long) t->ne[1]); + } + auto & a = totals[key]; + a.ms += ms; + a.n += 1; + ms_graphs += ms; + } + for (auto & pe : op_events) { + CUDA_CHECK(cudaEventDestroy(pe.second)); + } + CUDA_CHECK(cudaEventDestroy(end)); + if (++n_graphs % 100 == 0) { + std::vector> v(totals.begin(), totals.end()); + std::sort(v.begin(), v.end(), [](auto & x, auto & y) { return x.second.ms > y.second.ms; }); + GGML_LOG_INFO("op timing over %lld graphs: %.3f ms per graph\n", (long long) n_graphs, ms_graphs / n_graphs); + for (size_t k = 0; k < v.size() && k < 30; ++k) { + GGML_LOG_INFO(" %8.3f ms/graph %6.1f%% x%-5lld %s\n", v[k].second.ms / n_graphs, + 100.0 * v[k].second.ms / ms_graphs, (long long) (v[k].second.n / n_graphs), v[k].first.c_str()); + } + totals.clear(); + n_graphs = 0; + ms_graphs = 0; + } + } } #ifdef USE_CUDA_GRAPH