From eacf736f9c973c8a8e38a04fec2621117f1b61e3 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:31:41 +0200 Subject: [PATCH 1/3] kv-cache : resolve the partial KV residency set once for the model --kv-gpu-layers was resolved inside each cache, so a cache built from several sub-caches (iSWA, DSA, DSV4, MSA) would have given each of them the full budget and was disabled for all of them. It also counted layers the attention cache does not own, so on a hybrid model most of the budget went to recurrent layers, and it took layers in layer order, which fills the device owning the first layers and leaves the free memory of the others unused. Resolve one set of layer indices for the whole model instead, from the layers the requesting context actually owns, and take them per owning device. An MTP context owns the nextn layers, which sit above the layers of the main context. Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 7 +-- src/llama-kv-cache.cpp | 7 +-- src/llama-kv-cache.h | 4 -- src/llama-memory-hybrid.h | 4 -- src/llama-memory.h | 11 ++-- src/llama-model.cpp | 114 +++++++++++++++++++++++++++++++------ src/llama-model.h | 3 +- tests/test-llama-archs.cpp | 61 ++++++++++++++++++++ 8 files changed, 170 insertions(+), 41 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 233dfdcf790a..c58bd61ae24e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -435,12 +435,9 @@ llama_context::llama_context( cparams.live_context_workspace = false; } + // create_memory lowered kv_gpu_layers to the layers that became device-resident if (!cparams.offload_kqv && cparams.kv_gpu_layers > 0) { - if (memory && memory->get_supports_partial_kv()) { - cparams.offload_attn_compute = cparams.offload_attn_compute || cparams.op_offload; - } else { - LLAMA_LOG_WARN("%s: partial GPU KV residency is not supported for this memory layout; ignoring kv_gpu_layers\n", __func__); - } + cparams.offload_attn_compute = cparams.offload_attn_compute || cparams.op_offload; } } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index b526af616a74..843ff0bf3207 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -287,7 +287,7 @@ llama_kv_cache::llama_kv_cache( const char * dev_name = "CPU"; - const bool layer_offload = offload || n_gpu_resident < placement.gpu_resident_layers; + const bool layer_offload = offload || placement.gpu_resident_ils.count(il) > 0; ggml_backend_buffer_type_t buft = llama_kv_cache_get_host_buft(model, il, placement.cpu_pinned); if (layer_offload) { @@ -339,9 +339,8 @@ llama_kv_cache::llama_kv_cache( layers.push_back({ il, k, v, k_store_quantize, v_store_quantize, k_stream, v_stream, }); } - if (!offload && placement.gpu_resident_layers > 0) { - LLAMA_LOG_INFO("%s: partial GPU KV residency: %u of %u requested owned attention layers device-resident\n", - __func__, n_gpu_resident, placement.gpu_resident_layers); + if (!offload && n_gpu_resident > 0) { + LLAMA_LOG_DEBUG("%s: %u attention layers device-resident\n", __func__, n_gpu_resident); } if (reuse) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index af765830d01a..dab1c0095b30 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -138,10 +138,6 @@ class llama_kv_cache : public llama_memory_i { bool get_can_shift() const override; - bool get_supports_partial_kv() const override { - return true; - } - void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 393e12c44de6..7c36f5f3dc55 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -66,10 +66,6 @@ class llama_memory_hybrid : public llama_memory_i { bool recurrent_sparse_snapshots_supported() const override; bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) override; - bool get_supports_partial_kv() const override { - return true; - } - void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory.h b/src/llama-memory.h index 72ec7340d93e..1d5f36731ab7 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -5,6 +5,7 @@ #include #include +#include #include struct llama_ubatch; @@ -16,8 +17,12 @@ class llama_io_read_i; struct llama_memory_placement_options { bool cpu_pinned = false; - uint32_t gpu_resident_layers = 0; bool recurrent_offload = false; + + // layers whose attention KV stays device-resident while the rest of the cache is in host + // memory, resolved once for the whole model so that a cache built from several sub-caches + // shares one budget instead of giving each of them the full count + std::set gpu_resident_ils; }; struct llama_memory_params { @@ -121,10 +126,6 @@ struct llama_memory_i { virtual bool recurrent_sparse_snapshots_supported() const { return false; } virtual bool recurrent_set_sparse_snapshot_mode(bool, int32_t) { return false; } - virtual bool get_supports_partial_kv() const { - return false; - } - // // ops // diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 541a99efca26..3201df2fb811 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -2231,15 +2232,92 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } -llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { +// Pick the layers whose attention KV stays device-resident while the rest of the cache is in host +// memory. Taking them in layer order fills the device that owns the first layers and leaves the +// free memory of the others unused, so take them per owning device instead. +static std::set llama_pick_gpu_resident_layers(const llama_model & model, uint32_t n_layers, bool is_mtp) { + std::set ret; + if (n_layers == 0) { + return ret; + } + + const llama_hparams & hparams = model.hparams; + + // A model split by tensor mishandles an iSWA cache whose layers are partly device- and partly + // host-resident: the result stays coherent but measurably degrades. Leave the whole cache in + // host memory there until the mix is understood. + if (model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) { + for (uint32_t il = 0; il < hparams.n_layer(); il++) { + if (hparams.is_swa(il)) { + LLAMA_LOG_WARN("%s: kv_gpu_layers is not supported for an iSWA cache split by tensor, ignoring it\n", + __func__); + return ret; + } + } + } + + std::vector devs; + std::vector> devs_ils; + + for (uint32_t il = 0; il < hparams.n_layer_all; il++) { + // a nextn layer is cached by the MTP context, the rest of the model by the main one + if (hparams.n_layer_nextn > 0 && (il >= hparams.n_layer()) != is_mtp) { + continue; + } + // a recurrent layer keeps its state elsewhere, so it must not take from the budget + if (!hparams.has_kv(il) || hparams.is_recr(il)) { + continue; + } + auto * dev = model.dev_layer(il); + if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + continue; + } + auto it = std::find(devs.begin(), devs.end(), dev); + if (it == devs.end()) { + devs.push_back(dev); + devs_ils.emplace_back(); + it = devs.end() - 1; + } + devs_ils[it - devs.begin()].push_back(il); + } + + for (size_t i = 0; ret.size() < n_layers; i++) { + bool any = false; + for (size_t d = 0; d < devs.size() && ret.size() < n_layers; d++) { + const std::vector & ils = devs_ils[d]; + if (i < ils.size()) { + ret.insert(ils[i]); + any = true; + } + } + if (!any) { + break; + } + } + + return ret; +} + +llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_cparams & cparams) const { llama_memory_i * res; - const llama_memory_placement_options placement = { - cparams.kv_cpu_pinned, - cparams.kv_gpu_layers, - cparams.offload_kqv || cparams.recurrent_state_offload, - }; - llama_memory_placement_options specialized_placement = placement; - specialized_placement.gpu_resident_layers = 0; + llama_memory_placement_options placement; + placement.cpu_pinned = cparams.kv_cpu_pinned; + placement.recurrent_offload = cparams.offload_kqv || cparams.recurrent_state_offload; + // Resolved here rather than per cache, so that a cache built from several sub-caches shares one + // budget instead of giving each of them the full count. An offloaded cache is device-resident + // already, so it needs no set. + if (!cparams.offload_kqv && cparams.kv_gpu_layers > 0) { + placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers, + params.ctx_type == LLAMA_CONTEXT_TYPE_MTP); + if (placement.gpu_resident_ils.empty()) { + LLAMA_LOG_WARN("%s: no attention layer can be kept device-resident; ignoring kv_gpu_layers\n", __func__); + } else { + LLAMA_LOG_INFO("%s: partial GPU KV residency: %zu of %u requested attention layers device-resident\n", + __func__, placement.gpu_resident_ils.size(), cparams.kv_gpu_layers); + } + // the attention compute follows the cache, so report back what the cache got + cparams.kv_gpu_layers = (uint32_t) placement.gpu_resident_ils.size(); + } switch (arch) { // Models that need specific instantiation should be handled in the @@ -2282,7 +2360,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, filter_idx, nullptr, - specialized_placement); + placement); } break; case LLM_ARCH_GLM_DSA: case LLM_ARCH_DEEPSEEK32: @@ -2336,7 +2414,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_mla, filter_lid, nullptr, - specialized_placement); + placement); } } break; case LLM_ARCH_HY_V4: @@ -2380,7 +2458,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, filter_lid, nullptr, - specialized_placement); + placement); } } break; case LLM_ARCH_DOTS3NOTE: @@ -2433,7 +2511,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_mla, filter_lid, nullptr, - specialized_placement); + placement); } } break; case LLM_ARCH_DEEPSEEK4: @@ -2461,7 +2539,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_mtp, nullptr, nullptr, - specialized_placement); + placement); } else { res = new llama_kv_cache_dsv4( *this, @@ -2478,7 +2556,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.n_rs_seq, nullptr, nullptr, - specialized_placement); + placement); } } break; case LLM_ARCH_DFLASH: @@ -2503,7 +2581,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, nullptr, nullptr, - specialized_placement); + placement); break; } } @@ -2578,7 +2656,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* attn_n_ubatch */ cparams.n_ubatch, /* attn_n_pad */ 1, /* attn_offload */ cparams.offload_kqv, - /* placement */ specialized_placement, + /* placement */ placement, /* recurrent_type_r */ GGML_TYPE_F32, /* recurrent_type_s */ GGML_TYPE_F32, /* recurrent_rs_size */ std::max((uint32_t) 1, cparams.n_seq_max), @@ -2693,7 +2771,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter, reuse, share, - specialized_placement); + placement); } else { res = new llama_kv_cache_iswa( *this, @@ -2711,7 +2789,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter, reuse, share, - specialized_placement); + placement); } } else { GGML_ASSERT(!hparams.is_swa_any()); diff --git a/src/llama-model.h b/src/llama-model.h index 934301389fc9..5a3049dee686 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -751,7 +751,8 @@ struct llama_model { ggml_tensor * get_rope_factors(const llama_cparams & cparams, int il) const; - llama_memory_i * create_memory(const llama_memory_params & params, const llama_cparams & cparams) const; + // cparams.kv_gpu_layers is updated to the number of layers that became device-resident + llama_memory_i * create_memory(const llama_memory_params & params, llama_cparams & cparams) const; ggml_cgraph * build_graph(const llm_graph_params & params) const; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3bae33a0873e..010f00870886 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1050,6 +1050,66 @@ static void test_phase_workspace_mismatched_placement(size_t seed) { GGML_ASSERT(llama_contexts_share_workspace(target.get(), draft.get()) == (status == 1)); } +static size_t host_context_bytes(const llama_context * ctx) { + size_t ret = 0; + for (const auto & [buft, data] : llama_get_memory_breakdown(ctx)) { + if (ggml_backend_buft_is_host(buft)) { + ret += data.context; + } + } + return ret; +} + +// An MTP context owns only the nextn layers, which sit above the layers of the main context. A +// budget that only looks at the main context's layers leaves the whole MTP cache in host memory. +static void test_mtp_kv_residency(size_t seed) { + ggml_backend_dev_t gpu = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) { + gpu = dev; + break; + } + } + if (gpu == nullptr) { + printf("test_mtp_kv_residency: skipped, GPU device required\n"); + return; + } + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_QWEN35, false, true); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + model_params.n_gpu_layers = 99; + model_params.load_mtp = true; + ggml_backend_dev_t devices[] = { gpu, nullptr }; + model_params.devices = devices; + + size_t tensor_seed = seed; + llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tensor_seed, model_params)); + GGML_ASSERT(model); + + size_t host_bytes[2] = {0, 0}; + for (uint32_t kv_gpu_layers : {0u, 1u}) { + llama_context_params ctx_params = llama_context_default_params(); + ctx_params.n_ctx = 32; + ctx_params.n_batch = 32; + ctx_params.n_ubatch = 32; + ctx_params.n_seq_max = 1; + ctx_params.n_outputs_max = 1; + ctx_params.n_threads = 4; + ctx_params.n_threads_batch = 4; + ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + ctx_params.offload_kqv = false; + ctx_params.kv_gpu_layers = kv_gpu_layers; + llama_context_ptr ctx(llama_init_from_model(model.get(), ctx_params)); + GGML_ASSERT(ctx); + host_bytes[kv_gpu_layers] = host_context_bytes(ctx.get()); + } + GGML_ASSERT(host_bytes[0] > 0); + GGML_ASSERT(host_bytes[1] < host_bytes[0]); + printf("test_mtp_kv_residency: OK\n"); +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -1522,6 +1582,7 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, verbosity, out); } + test_mtp_kv_residency(seed); return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); From fe498b7f326b500973d2b815db76628943bc570d Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:34:01 +0200 Subject: [PATCH 2/3] kv-cache : pick the residency set from what the caches own and keep The picker excluded every recurrent layer and counted its own choice as the result. Falcon H1 marks all of its layers recurrent yet caches all of them, so a request placed nothing; Nemotron H caches only the non-recurrent layers without an FFN, so part of the budget went to layers the cache then dropped, and the count still said they were resident, which can enable the attention compute offload for nothing. Ask the same ownership filter the hybrid cache uses, and report the layers the caches did place. Assisted-by: Claude Opus 5 --- common/arg.cpp | 6 ++-- src/llama-kv-cache.cpp | 3 ++ src/llama-memory.h | 5 +++ src/llama-model.cpp | 72 ++++++++++++++++++++++++++++---------- tests/test-llama-archs.cpp | 54 ++++++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 22 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 04f660a5be2e..47b84a5b566d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2464,9 +2464,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_LIVE_CONTEXT_WORKSPACE")); add_opt(common_arg( {"--kv-gpu-layers"}, "N", - string_format("with --no-kv-offload, keep the first N independently owned attention KV layers device-resident " - "for standard and direct hybrid caches. Unsupported specialized caches ignore this option " - "(default: %d)", params.kv_gpu_layers), + string_format("with --no-kv-offload, keep N independently owned attention KV layers device-resident. " + "The layers are taken per owning device. A model whose cache cannot hold them keeps " + "fewer, or none (default: %d)", params.kv_gpu_layers), [](common_params & params, int value) { if (value < 0) { throw std::invalid_argument("--kv-gpu-layers must not be negative"); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 843ff0bf3207..be04da99b53e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -299,6 +299,9 @@ llama_kv_cache::llama_kv_cache( if (!offload && layer_offload) { ++n_gpu_resident; + if (placement.gpu_resident_done) { + placement.gpu_resident_done->insert(il); + } } LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); diff --git a/src/llama-memory.h b/src/llama-memory.h index 1d5f36731ab7..32841c7e118b 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -23,6 +23,11 @@ struct llama_memory_placement_options { // memory, resolved once for the whole model so that a cache built from several sub-caches // shares one budget instead of giving each of them the full count std::set gpu_resident_ils; + + // the layers a cache did keep device-resident, which can be fewer than the set above: a cache + // filter can drop a picked layer, and a model without a standard attention cache takes none. + // Shared by the sub-caches of one model, so the count is over the whole model. + std::shared_ptr> gpu_resident_done; }; struct llama_memory_params { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3201df2fb811..947928b2ce5e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2232,6 +2232,35 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } +// The layers whose attention KV a hybrid cache of this model owns. A null filter means the model is +// not hybrid and every layer with a KV state belongs to the attention cache. +static llama_memory_hybrid::layer_filter_cb llama_hybrid_filter_attn(const llama_model & model, bool is_mtp) { + const llama_hparams & hparams = model.hparams; + if (is_mtp) { + // an MTP head caches the nextn layers with a plain attention cache, not the hybrid wrapper + return nullptr; + } + switch (model.arch) { + case LLM_ARCH_FALCON_H1: + return [](uint32_t) { return true; }; + case LLM_ARCH_NEMOTRON_H: + case LLM_ARCH_NEMOTRON_H_MOE: + return [&hparams](uint32_t il) { + return !hparams.is_recr(il) && hparams.n_ff(il) == 0; + }; + case LLM_ARCH_QWEN3NEXT: + case LLM_ARCH_QWEN35: + case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: + case LLM_ARCH_MINIMAX_01: + return [&hparams](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + default: + return nullptr; + } +} + // Pick the layers whose attention KV stays device-resident while the rest of the cache is in host // memory. Taking them in layer order fills the device that owns the first layers and leaves the // free memory of the others unused, so take them per owning device instead. @@ -2259,13 +2288,18 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod std::vector devs; std::vector> devs_ils; + // a layer the attention cache does not own must not take from the budget, so ask the same filter + const llama_memory_hybrid::layer_filter_cb filter_attn = llama_hybrid_filter_attn(model, is_mtp); + for (uint32_t il = 0; il < hparams.n_layer_all; il++) { // a nextn layer is cached by the MTP context, the rest of the model by the main one if (hparams.n_layer_nextn > 0 && (il >= hparams.n_layer()) != is_mtp) { continue; } - // a recurrent layer keeps its state elsewhere, so it must not take from the budget - if (!hparams.has_kv(il) || hparams.is_recr(il)) { + if (!hparams.has_kv(il)) { + continue; + } + if (filter_attn ? !filter_attn(il) : hparams.is_recr(il)) { continue; } auto * dev = model.dev_layer(il); @@ -2307,16 +2341,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // budget instead of giving each of them the full count. An offloaded cache is device-resident // already, so it needs no set. if (!cparams.offload_kqv && cparams.kv_gpu_layers > 0) { - placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers, + placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers, params.ctx_type == LLAMA_CONTEXT_TYPE_MTP); - if (placement.gpu_resident_ils.empty()) { - LLAMA_LOG_WARN("%s: no attention layer can be kept device-resident; ignoring kv_gpu_layers\n", __func__); - } else { - LLAMA_LOG_INFO("%s: partial GPU KV residency: %zu of %u requested attention layers device-resident\n", - __func__, placement.gpu_resident_ils.size(), cparams.kv_gpu_layers); - } - // the attention compute follows the cache, so report back what the cache got - cparams.kv_gpu_layers = (uint32_t) placement.gpu_resident_ils.size(); + placement.gpu_resident_done = std::make_shared>(); } switch (arch) { @@ -2612,26 +2639,19 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) { // The main difference between hybrid architectures is the // layer filters, so pick the right one here - llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; + llama_memory_hybrid::layer_filter_cb filter_attn = llama_hybrid_filter_attn(*this, false); llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; // only the sparse-attention architectures use llama_memory_hybrid_idx // a null filter_idx means the GGUF has no indexer tensors llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); if (arch == LLM_ARCH_FALCON_H1) { - filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { - filter_attn = [&](uint32_t il) { - return !hparams.is_recr(il) && hparams.n_ff(il) == 0; - }; filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01) { - filter_attn = [&](uint32_t il) { - return il < hparams.n_layer() && !hparams.is_recr(il); - }; filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; @@ -2817,6 +2837,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } + // Report what the caches did rather than what was picked: a cache filter can drop a picked + // layer, and a model without a standard attention cache keeps none. The attention compute + // follows the cache, so cparams must carry the count the caches reached. + if (placement.gpu_resident_done) { + const uint32_t n_resident = (uint32_t) placement.gpu_resident_done->size(); + if (n_resident == 0) { + LLAMA_LOG_WARN("%s: no attention layer can be kept device-resident; ignoring kv_gpu_layers\n", __func__); + } else { + LLAMA_LOG_INFO("%s: partial GPU KV residency: %u of %u requested attention layers device-resident\n", + __func__, n_resident, cparams.kv_gpu_layers); + } + cparams.kv_gpu_layers = n_resident; + } + return res; } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 010f00870886..a2e119ee5cc5 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1110,6 +1110,59 @@ static void test_mtp_kv_residency(size_t seed) { printf("test_mtp_kv_residency: OK\n"); } +// The picked set must follow the ownership filter of the attention cache and the reported count must +// be what the caches did place: Falcon H1 marks every layer recurrent yet caches all of them, Nemotron +// H caches only the non-recurrent layers without an FFN, and a recurrent model caches none. +static void test_kv_residency_ownership(size_t seed) { + ggml_backend_dev_t gpu = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) { + gpu = dev; + break; + } + } + if (gpu == nullptr) { + printf("test_kv_residency_ownership: skipped, GPU device required\n"); + return; + } + + auto resident_layers = [&](llm_arch arch, bool moe, uint32_t requested) { + gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + model_params.n_gpu_layers = 99; + ggml_backend_dev_t devices[] = { gpu, nullptr }; + model_params.devices = devices; + + size_t tensor_seed = seed; + llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tensor_seed, model_params)); + GGML_ASSERT(model); + + llama_context_params ctx_params = llama_context_default_params(); + ctx_params.n_ctx = 32; + ctx_params.n_batch = 32; + ctx_params.n_ubatch = 32; + ctx_params.n_seq_max = 1; + ctx_params.n_outputs_max = 1; + ctx_params.n_threads = 4; + ctx_params.n_threads_batch = 4; + ctx_params.offload_kqv = false; + ctx_params.kv_gpu_layers = requested; + llama_context_ptr ctx(llama_init_from_model(model.get(), ctx_params)); + GGML_ASSERT(ctx); + return ctx->get_cparams().kv_gpu_layers; + }; + + // every layer is recurrent, but the attention cache owns all of them + GGML_ASSERT(resident_layers(LLM_ARCH_FALCON_H1, false, 1) == 1); + // the attention cache owns only layer 0, so the second requested layer cannot be placed + GGML_ASSERT(resident_layers(LLM_ARCH_NEMOTRON_H, false, 2) == 1); + // a recurrent model builds no attention cache at all + GGML_ASSERT(resident_layers(LLM_ARCH_MAMBA, false, 1) == 0); + printf("test_kv_residency_ownership: OK\n"); +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -1583,6 +1636,7 @@ int main(int argc, char ** argv) { return save_models(arch, seed, verbosity, out); } test_mtp_kv_residency(seed); + test_kv_residency_ownership(seed); return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); From 5c9a74f7656b13e98acc6753de74005254398a08 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 7 Sep 2026 21:50:07 +0200 Subject: [PATCH 3/3] kv-cache : claim a picked residency layer once A cache built from several sub-caches placed the same layer in each of them. The residency tests get their own invocation instead of running in the arch sweep. Assisted-by: Claude Opus 5 --- src/llama-kv-cache.cpp | 6 +++++- tests/CMakeLists.txt | 7 +++++++ tests/test-llama-archs.cpp | 14 +++++++++++--- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index be04da99b53e..5a8ceebed6e5 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -287,7 +287,11 @@ llama_kv_cache::llama_kv_cache( const char * dev_name = "CPU"; - const bool layer_offload = offload || placement.gpu_resident_ils.count(il) > 0; + // a layer picked for residency is claimed by the first cache that owns it, so that the + // sub-caches of one model share the budget instead of each spending it again + const bool layer_picked = placement.gpu_resident_ils.count(il) > 0 && + (!placement.gpu_resident_done || placement.gpu_resident_done->count(il) == 0); + const bool layer_offload = offload || layer_picked; ggml_backend_buffer_type_t buft = llama_kv_cache_get_host_buft(model, il, placement.cpu_pinned); if (layer_offload) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 23b5d764f604..2d5de7bf5686 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -214,6 +214,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) ARGS --test-live-context-workspace ) + llama_test( + test-llama-archs + NAME test-kv-residency + LABEL main + ARGS --test-kv-residency + ) + set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index a2e119ee5cc5..b3cc41f8d527 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -68,7 +68,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace] [--test-kv-residency]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -1562,6 +1562,7 @@ int main(int argc, char ** argv) { std::string out; bool test_phase_workspace = false; bool test_live_context_workspace = false; + bool test_kv_residency = false; int verbosity = LOG_LEVEL_ERROR; @@ -1615,6 +1616,10 @@ int main(int argc, char ** argv) { test_live_context_workspace = true; continue; } + if (strcmp(argv[i], "--test-kv-residency") == 0) { + test_kv_residency = true; + continue; + } } printf("%s: using seed %zu\n", __func__, seed); @@ -1632,11 +1637,14 @@ int main(int argc, char ** argv) { test_live_context_workspace_unsupported(seed); return 0; } + if (test_kv_residency) { + test_mtp_kv_residency(seed); + test_kv_residency_ownership(seed); + return 0; + } if (!out.empty()) { return save_models(arch, seed, verbosity, out); } - test_mtp_kv_residency(seed); - test_kv_residency_ownership(seed); return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what());