diff --git a/common/arg.cpp b/common/arg.cpp index 94dead3ff3ba..128130f0df19 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", @@ -3701,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"}, @@ -4128,6 +4167,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.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..506456ec80e3 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; @@ -588,6 +598,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) @@ -643,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/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/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-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/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 59e6c9ba74b0..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 @@ -595,8 +612,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 +650,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..81cad7c44541 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())); @@ -4519,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); @@ -4695,12 +5011,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))); } } @@ -4718,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 @@ -6041,6 +6401,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/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. 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); } } } 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 988d9f365190..5e3420509ce2 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(); @@ -1442,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 }; { @@ -3675,6 +3699,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([&]() {