diff --git a/common/arg.cpp b/common/arg.cpp index 3acfbab6e544..1df1c15cf9e5 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2459,9 +2459,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. " + "Layers are selected in turn from each owning device. Unsupported cache layouts keep " + "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/include/llama.h b/include/llama.h index 386f0fd7ad16..7ec3d111dfd9 100644 --- a/include/llama.h +++ b/include/llama.h @@ -364,7 +364,7 @@ extern "C" { uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max) - uint32_t kv_gpu_layers; // with offload_kqv=false, keep this many standard or direct hybrid attention KV layers on their assigned devices + uint32_t kv_gpu_layers; // with offload_kqv=false, request up to this many independently owned attention KV layers on their assigned devices 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-context.cpp b/src/llama-context.cpp index 9aaa0f2a5cf6..a01d0b1d01e9 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..58d8ef87263a 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -287,7 +287,10 @@ llama_kv_cache::llama_kv_cache( const char * dev_name = "CPU"; - const bool layer_offload = offload || n_gpu_resident < placement.gpu_resident_layers; + // Sub-caches claim each selected layer once. + 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) { @@ -299,6 +302,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); @@ -339,9 +345,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..8997e09a2fa6 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,13 @@ class llama_io_read_i; struct llama_memory_placement_options { bool cpu_pinned = false; - uint32_t gpu_resident_layers = 0; bool recurrent_offload = false; + + // One residency set is shared by all sub-caches of this context. + std::set gpu_resident_ils; + + // The first cache that owns a selected layer claims it; auxiliary copies remain on the host. + std::shared_ptr> gpu_resident_done; }; struct llama_memory_params { @@ -121,10 +127,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 a65abfc9ca64..98436e8c2984 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -2248,15 +2249,116 @@ 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 { +// Use the same hybrid attention ownership rules for selection and cache construction. +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; + } +} + +static bool llama_layer_in_context(const llama_hparams & hparams, uint32_t il, bool is_mtp) { + return hparams.n_layer_nextn == 0 || hparams.n_layer() == 0 || hparams.router_layer >= 0 || + (il >= hparams.n_layer()) == is_mtp; +} + +// Spread the residency budget across the devices that own attention layers. +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; + + // Mixed host/device iSWA under tensor split has a known numerical discrepancy. + if (model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) { + for (uint32_t il = 0; il < hparams.n_layer_all; il++) { + if (llama_layer_in_context(hparams, il, is_mtp) && hparams.has_kv(il) && 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; + + // 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++) { + if (!llama_layer_in_context(hparams, il, is_mtp)) { + continue; + } + if (!hparams.has_kv(il)) { + continue; + } + if (filter_attn ? !filter_attn(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; + // Share one residency budget across sub-caches; full offload needs no selection. + 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); + placement.gpu_resident_done = std::make_shared>(); + } switch (arch) { // Models that need specific instantiation should be handled in the @@ -2299,7 +2401,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: @@ -2353,7 +2455,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: @@ -2397,7 +2499,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: @@ -2450,7 +2552,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: @@ -2478,7 +2580,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, @@ -2495,7 +2597,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: @@ -2520,7 +2622,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, nullptr, nullptr, - specialized_placement); + placement); break; } } @@ -2551,26 +2653,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); }; @@ -2595,7 +2690,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), @@ -2668,14 +2763,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; } - // don't filter when n_layer_nextn is repurposed for a router layer the trunk attends - // or when a model is entirely n_layer_nextn layers and has no trunk + // Router and all-nextn models keep all layers in the same context. if (hparams.n_layer_nextn > 0 && hparams.n_layer() > 0 && hparams.router_layer < 0) { - if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) { - filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; - } else { - filter = [&](uint32_t il) { return il < hparams.n_layer(); }; - } + const bool is_mtp = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP; + filter = [&, is_mtp](uint32_t il) { return llama_layer_in_context(hparams, il, is_mtp); }; } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { @@ -2710,7 +2801,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, @@ -2728,7 +2819,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()); @@ -2756,6 +2847,18 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } + // Attention compute placement follows the layers actually claimed by the caches. + 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/src/llama-model.h b/src/llama-model.h index 2cb74ea23273..fd42357977dc 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -754,7 +754,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/CMakeLists.txt b/tests/CMakeLists.txt index 04693bb0e6ed..1b5283d5ff61 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 3951dbbe8526..21c223caf34c 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -12,6 +12,7 @@ #include "../src/llama-arch.h" #include "../src/llama-context.h" #include "../src/llama-ext.h" +#include "../src/llama-model.h" #include "../src/llama-model-saver.h" #include @@ -69,7 +70,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){ @@ -1052,6 +1053,168 @@ 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; +} + +static size_t device_context_bytes(const llama_context * ctx, ggml_backend_dev_t dev) { + size_t ret = 0; + for (const auto & [buft, data] : llama_get_memory_breakdown(ctx)) { + if (!ggml_backend_buft_is_host(buft) && ggml_backend_buft_get_device(buft) == dev) { + ret += data.context; + } + } + return ret; +} + +static llama_context_params kv_residency_context_params() { + llama_context_params ret = llama_context_default_params(); + ret.n_ctx = 32; + ret.n_batch = 32; + ret.n_ubatch = 32; + ret.n_seq_max = 1; + ret.n_outputs_max = 4; + ret.n_threads = 4; + ret.n_threads_batch = 4; + return ret; +} + +// MTP owns the nextn cache, while the main context owns the trunk cache and recurrent state. +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); + + for (llama_context_type ctx_type : {LLAMA_CONTEXT_TYPE_DEFAULT, LLAMA_CONTEXT_TYPE_MTP}) { + size_t host_bytes = 0; + size_t resident_bytes = 0; + std::vector logits_host; + for (uint32_t requested : {0u, 1u, 99u}) { + llama_context_params ctx_params = kv_residency_context_params(); + ctx_params.ctx_type = ctx_type; + 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); + GGML_ASSERT(ctx->get_cparams().kv_gpu_layers == (requested > 0 ? 1u : 0u)); + const size_t host = host_context_bytes(ctx.get()); + const size_t device = device_context_bytes(ctx.get(), gpu); + if (requested == 0) { + host_bytes = host; + GGML_ASSERT(host > 0 && device == 0); + } else { + GGML_ASSERT(host < host_bytes && device > 0); + GGML_ASSERT((host == 0) == (ctx_type == LLAMA_CONTEXT_TYPE_MTP)); + if (requested == 1) { + resident_bytes = device; + } else { + GGML_ASSERT(device == resident_bytes); + } + } + if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) { + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get())); + llama_batch batch = make_mtp_batch(4, llama_model_n_embd_out(model.get()), 0, n_vocab, seed); + GGML_ASSERT(llama_decode(ctx.get(), batch) == 0); + const float * logits = llama_get_logits_ith(ctx.get(), -1); + const std::vector result(logits, logits + n_vocab); + if (requested == 0) { + logits_host = result; + } else { + const double error = nmse(logits_host, result); + GGML_ASSERT(error <= 1e-4); + printf("test_mtp_kv_residency: requested=%u actual=%u NMSE=%.2e\n", requested, ctx->get_cparams().kv_gpu_layers, error); + } + llama_batch_free(batch); + } + printf("test_mtp_kv_residency: type=%d requested=%u actual=%u host=%zu device=%zu\n", ctx_type, requested, ctx->get_cparams().kv_gpu_layers, host, device); + } + llama_context_params defaults = kv_residency_context_params(); + defaults.ctx_type = ctx_type; + GGML_ASSERT(defaults.offload_kqv && defaults.kv_gpu_layers == 0); + llama_context_ptr ctx(llama_init_from_model(model.get(), defaults)); + GGML_ASSERT(ctx); + GGML_ASSERT(host_context_bytes(ctx.get()) == 0); + GGML_ASSERT(device_context_bytes(ctx.get(), gpu) > 0); + } + printf("test_mtp_kv_residency: OK\n"); +} + +// Hybrid attention ownership differs from the recurrent layer flags and the total layer count. +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; + } + } + for (llm_arch arch : {LLM_ARCH_LLAMA, LLM_ARCH_FALCON_H1, LLM_ARCH_NEMOTRON_H, LLM_ARCH_MAMBA}) { + if (gpu == nullptr && arch != LLM_ARCH_LLAMA) { + continue; + } + gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, false); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + model_params.n_gpu_layers = arch == LLM_ARCH_LLAMA ? 0 : 99; + ggml_backend_dev_t devices[] = { arch == LLM_ARCH_LLAMA ? nullptr : 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); + + const uint32_t available = arch == LLM_ARCH_FALCON_H1 ? 2 : arch == LLM_ARCH_NEMOTRON_H ? 1 : 0; + size_t host_bytes = 0; + for (uint32_t requested : {0u, 1u, 99u}) { + llama_context_params ctx_params = kv_residency_context_params(); + 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); + const uint32_t actual = ctx->get_cparams().kv_gpu_layers; + GGML_ASSERT(actual == std::min(requested, available)); + const size_t host = host_context_bytes(ctx.get()); + const size_t device = device_context_bytes(ctx.get(), gpu); + if (requested == 0) { + host_bytes = host; + GGML_ASSERT(host > 0); + } + GGML_ASSERT((device > 0) == (actual > 0)); + GGML_ASSERT(actual > 0 ? host < host_bytes : host == host_bytes); + printf("test_kv_residency_ownership: arch=%s requested=%u actual=%u host=%zu device=%zu\n", llm_arch_name(arch), requested, actual, host, device); + } + } + 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)); @@ -1086,6 +1249,191 @@ static std::vector get_logits( return ret; } +static void test_kv_residency_devices(size_t seed) { + std::vector devices; + 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) { + devices.push_back(dev); + } + } + if (devices.empty()) { + printf("test_kv_residency_devices: skipped, GPU device required\n"); + return; + } + + const uint32_t n_devices = devices.size(); + const uint32_t n_layers = 2*n_devices; + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_LLAMA, false); + llama_model_saver ms(LLM_ARCH_LLAMA, gguf_ctx.get()); + ms.add_kv(LLM_KV_BLOCK_COUNT, n_layers); + + // Each device owns two repeating layers; the last device also owns the output layer. + std::vector tensor_split(llama_max_devices(), 0.0f); + GGML_ASSERT(n_devices <= tensor_split.size()); + for (uint32_t d = 0; d < n_devices; ++d) { + tensor_split[d] = d + 1 == n_devices ? 3.0f : 2.0f; + } + devices.push_back(nullptr); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + model_params.n_gpu_layers = n_layers + 1; + model_params.devices = devices.data(); + model_params.tensor_split = tensor_split.data(); + model_params.split_mode = LLAMA_SPLIT_MODE_LAYER; + 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); + for (uint32_t il = 0; il < n_layers; ++il) { + GGML_ASSERT(model->dev_layer(il) == devices[il/2]); + } + + const std::vector tokens = get_tokens(4, llama_vocab_n_tokens(llama_model_get_vocab(model.get())), seed); + std::vector requests = {0, 1}; + if (n_devices > 1) { + requests.push_back(n_devices); + } + requests.push_back(n_devices + 1); + requests.push_back(99); + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_Q8_0}) { + llama_context_params defaults = kv_residency_context_params(); + defaults.type_k = type; + defaults.type_v = type; + defaults.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + GGML_ASSERT(defaults.offload_kqv && defaults.kv_gpu_layers == 0); + llama_context_ptr full(llama_init_from_model(model.get(), defaults)); + GGML_ASSERT(full); + GGML_ASSERT(host_context_bytes(full.get()) == 0); + std::vector device_layer_bytes(n_devices); + for (uint32_t d = 0; d < n_devices; ++d) { + const size_t bytes = device_context_bytes(full.get(), devices[d]); + printf("test_kv_residency_devices: type=%s default device=%s layers=2 bytes=%zu\n", ggml_type_name(type), ggml_backend_dev_name(devices[d]), bytes); + fflush(stdout); + GGML_ASSERT(bytes > 0 && bytes % 2 == 0); + device_layer_bytes[d] = bytes/2; + } + const std::vector logits_full = get_logits(model.get(), full.get(), tokens); + full.reset(); + + size_t layer_bytes = 0; + std::vector logits_host; + for (uint32_t requested : requests) { + llama_context_params ctx_params = kv_residency_context_params(); + ctx_params.offload_kqv = false; + ctx_params.kv_gpu_layers = requested; + ctx_params.type_k = type; + ctx_params.type_v = type; + ctx_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + llama_context_ptr ctx(llama_init_from_model(model.get(), ctx_params)); + GGML_ASSERT(ctx); + const uint32_t actual = ctx->get_cparams().kv_gpu_layers; + const size_t host = host_context_bytes(ctx.get()); + printf("test_kv_residency_devices: type=%s requested=%u actual=%u host=%zu\n", ggml_type_name(type), requested, actual, host); + fflush(stdout); + GGML_ASSERT(actual == std::min(requested, n_layers)); + if (requested == 0) { + GGML_ASSERT(host > 0 && host % n_layers == 0); + layer_bytes = host/n_layers; + } + GGML_ASSERT(host == (n_layers - actual)*layer_bytes); + for (uint32_t d = 0; d < n_devices; ++d) { + const uint32_t resident = actual/n_devices + (d < actual % n_devices ? 1 : 0); + const size_t bytes = device_context_bytes(ctx.get(), devices[d]); + const size_t expected = resident*device_layer_bytes[d]; + printf("test_kv_residency_devices: type=%s requested=%u actual=%u device=%s layers=%u bytes=%zu expected=%zu\n", ggml_type_name(type), requested, actual, ggml_backend_dev_name(devices[d]), resident, bytes, expected); + fflush(stdout); + GGML_ASSERT(bytes == expected); + } + const std::vector logits = get_logits(model.get(), ctx.get(), tokens); + if (requested == 0) { + logits_host = logits; + } else { + const double error = nmse(logits_host, logits); + printf("test_kv_residency_devices: type=%s requested=%u host=%zu NMSE=%.2e\n", ggml_type_name(type), requested, host, error); + fflush(stdout); + GGML_ASSERT(error <= 1e-4); + } + } + const double error = nmse(logits_host, logits_full); + printf("test_kv_residency_devices: type=%s default NMSE=%.2e\n", ggml_type_name(type), error); + fflush(stdout); + GGML_ASSERT(error <= 1e-4); + } + printf("test_kv_residency_devices: OK (%u devices)\n", n_devices); +} + +static void test_kv_residency_context_filter(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_context_filter: skipped, GPU device required\n"); + return; + } + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_LLAMA, false); + for (llama_split_mode split_mode : {LLAMA_SPLIT_MODE_LAYER, LLAMA_SPLIT_MODE_TENSOR}) { + 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; + model_params.split_mode = split_mode; + 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 = kv_residency_context_params(); + ctx_params.offload_kqv = false; + llama_context_ptr ctx(llama_init_from_model(model.get(), ctx_params)); + GGML_ASSERT(ctx); + const llama_cparams resolved = ctx->get_cparams(); + ctx.reset(); + + // Test cache ownership directly; these metadata variants do not define inference graphs. + const llama_hparams original = model->hparams; + auto check = [&](const char * name, llama_context_type ctx_type, uint32_t expected) { + llama_cparams cparams = resolved; + cparams.kv_gpu_layers = 99; + cparams.ctx_type = ctx_type; + llama_memory_params mparams = {GGML_TYPE_F16, GGML_TYPE_F16, false, ctx_type, nullptr}; + std::unique_ptr memory(model->create_memory(mparams, cparams)); + GGML_ASSERT(memory); + size_t host = 0; + size_t device = 0; + for (const auto & [buft, bytes] : memory->memory_breakdown()) { + (ggml_backend_buft_is_host(buft) ? host : device) += bytes; + } + printf("test_kv_residency_context_filter: case=%s split=%d type=%d requested=99 actual=%u host=%zu device=%zu\n", name, split_mode, ctx_type, cparams.kv_gpu_layers, host, device); + GGML_ASSERT(cparams.kv_gpu_layers == expected); + GGML_ASSERT((device > 0) == (expected > 0)); + GGML_ASSERT((host > 0) == (expected == 0)); + }; + model->hparams.n_layer_nextn = 2; + check("all-nextn", LLAMA_CONTEXT_TYPE_DEFAULT, 2); + model->hparams.n_layer_nextn = 1; + model->hparams.router_layer = 1; + check("router", LLAMA_CONTEXT_TYPE_DEFAULT, 2); + model->hparams.router_layer = -1; + check("trunk", LLAMA_CONTEXT_TYPE_DEFAULT, 1); + check("mtp", LLAMA_CONTEXT_TYPE_MTP, 1); + if (split_mode == LLAMA_SPLIT_MODE_TENSOR) { + model->hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + model->hparams.n_swa = 16; + model->hparams.is_swa_impl[0] = 1; + model->hparams.is_swa_impl[1] = 0; + check("swa-trunk", LLAMA_CONTEXT_TYPE_DEFAULT, 0); + check("dense-mtp-with-swa-trunk", LLAMA_CONTEXT_TYPE_MTP, 1); + } + model->hparams = original; + } + printf("test_kv_residency_context_filter: OK\n"); +} + // decode two sequences in one batch, compare each with a decode of it alone // logits_a: logits of tokens decoded alone with the same device config static bool test_parallel_seqs(llama_model * model, const std::vector & tokens, const std::vector & logits_a, bool encode) { @@ -1541,6 +1889,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; @@ -1594,6 +1943,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); @@ -1611,6 +1964,13 @@ 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); + test_kv_residency_devices(seed); + test_kv_residency_context_filter(seed); + return 0; + } if (!out.empty()) { return save_models(arch, seed, verbosity, out); } diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index e2949bed16ac..faf8f48c919b 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -107,7 +107,7 @@ The maintained fork adds the following context parameters. The four boolean opti | Option | Behavior and limits | | --- | --- | -| `-kvgl`, `--kv-gpu-layers` | With `-nkvo 1`, requests device residency for the first N independently owned attention KV layers. Standard and direct hybrid caches support it; unsupported specialized caches ignore it. The model's layer placement and available attention layers limit what can become GPU resident. | +| `-kvgl`, `--kv-gpu-layers` | With `-nkvo 1`, requests up to N independently owned attention KV layers, selected in turn from each owning device. Sub-caches share one budget; auxiliary caches can remain host-resident. Tensor-split iSWA caches do not support partial residency. The count depends on the context's owned layers and model placement; it is not a VRAM limit. | | `-kvcp`, `--kv-cpu-pinned` | Requests pinned buffers for host KV when the backend provides them. With operation offload enabled (`-nopo 0`), it also allows attention compute on the accelerator, so a `0,1` sweep can change both storage and compute placement. | | `-rso`, `--recurrent-state-offload` | Offloads recurrent state for recurrent and hybrid models independently of host attention KV. With `-nkvo 0`, KV offload already enables recurrent state offload, so changing this option has no additional effect. | | `-paw`, `--phase-aware-workspace` | Resizes compute workspaces between prompt processing and token generation; a later prompt regrows the reservation. It can be combined with `-lcw 1`. |