From 944da4ef585ba0e585af04470a82e7d23dbbc36a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 10:46:56 +0200 Subject: [PATCH 01/22] docs : link CLAUDE.md to AGENTS.md Assisted-by: Claude Opus 5 --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) mode change 100644 => 120000 CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index 302cdeab99cf..000000000000 --- a/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -IMPORTANT: Ensure you’ve thoroughly reviewed the [AGENTS.md](AGENTS.md) file before beginning any work. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 000000000000..47dc3e3d863c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file From 9b7ac978c66fa446a000822131860b02dd4813dc Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 10:55:06 +0200 Subject: [PATCH 02/22] ggml-meta : split a host-resident KV cache by head A KV cache in host memory reaches attention as a scheduler copy, which is a leaf in the compute buffer. Such a leaf never reached the device split-state callback and fell through to MIRRORED, while the queries stayed split by head: each device then attended heads whose keys live on the other device. With more than one KV head that aborts in the FlashAttention kernel, or returns wrong output where the query split happens to stay a multiple of the KV head count. Offer a copied-in leaf to the callback under the name the graph gave it, and let the callback recognise the cache there. The cache folds its heads into one flat axis but the copy arrives permuted, with the heads on an axis of their own, so its segments and granularity are rescaled to whole heads. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 71 +++++++++++++++++++++++++++++++++- src/llama-model.cpp | 55 ++++++++++++++++++++++---- 2 files changed, 117 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index fe58ea3bb7a6..4074bc3147c6 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -485,6 +485,27 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct return it->second[index]; } +// The scheduler names a tensor it copies in from another backend "##", where +// is the graph name of the source plus the suffixes ggml appends for views. Undo both, so +// that the copy can be recognised under the name the graph gave it. +static std::string ggml_backend_meta_copy_source_name(const char * name) { + std::string decorated = name; + const size_t first = decorated.find('#'); + const size_t last = decorated.rfind('#'); + if (first == std::string::npos || last <= first) { + return decorated; + } + std::string ret = decorated.substr(first + 1, last - first - 1); + while (!ret.empty() && ret.back() == ')') { + const size_t paren = ret.rfind(" ("); + if (paren == std::string::npos) { + break; + } + ret.erase(paren); + } + return ret; +} + static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( @@ -831,10 +852,28 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { + // A tensor the scheduler copied in from another backend is a leaf in the compute buffer, + // so it never reached the device callback and fell through to MIRRORED. A host-resident + // KV cache arrives this way, and mirroring it while the queries stay split by head makes + // each device attend the wrong heads. Ask the callback; it still answers MIRRORED for + // names it does not know. + const bool copied_in_leaf = + ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + tensor->op == GGML_OP_NONE && (tensor->flags & GGML_TENSOR_FLAG_INPUT) == 0; + + if ((ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE || copied_in_leaf) && + tensor->view_src == nullptr) { ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); + // the callback classifies by name, so offer a copy under the name it was copied from + const ggml_tensor * tensor_query = tensor; + ggml_tensor tensor_named; + if (copied_in_leaf) { + tensor_named = *tensor; + ggml_set_name(&tensor_named, ggml_backend_meta_copy_source_name(tensor->name).c_str()); + tensor_query = &tensor_named; + } + ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor_query, dev_ctx->get_split_state_ud); if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; int64_t ne_sum = 0; @@ -1410,6 +1449,34 @@ static void ggml_backend_meta_buffer_memset_tensor( static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. + // The heads are the split axis but they are interleaved inside each cell, so the chunk splice + // below cannot express the write. Each device's heads are one contiguous run per cell, which + // is a strided copy: ne[1] cells, from the full cell stride to the device's own. + const bool strided_head_split = + !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->ne[3] == 1 && tensor->nb[1] > tensor->nb[2] && + offset == 0 && size == ggml_nbytes(tensor); + + if (strided_head_split) { + size_t offset_data = 0; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[j] * tensor->nb[2]; + if (nbytes == 0) { + continue; + } + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, 0, nbytes, + tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + } + GGML_ASSERT(offset_data == (size_t) tensor->ne[2] * tensor->nb[2]); + return; + } + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 781db00c7fb5..bbffd04b7bd3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -461,7 +461,36 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return {axis, tensor_axis_0, il, rotation}; }; + // An attention cache lives in host memory when it is not offloaded, so it reaches attention as + // a scheduler copy rather than as the cache itself. The cache folds all heads into one flat + // axis, the copy arrives permuted, with the heads on an axis of their own. The split follows + // the cache either way, so record the axis and the unit - one head - the copy is split in. + struct host_cache_copy { + bool valid = false; + ggml_backend_meta_split_axis axis = GGML_BACKEND_SPLIT_AXIS_0; + int64_t unit = 1; // cache elements per element of axis + }; + + const host_cache_copy cache_copy = [&]() -> host_cache_copy { + if (is_dsv4 || !std::regex_match(tensor_name, pattern_kv_cache)) { + return {}; + } + // [head_dim, n_kv, n_head_kv, n_stream] + const uint32_t il = std::stoul(tensor_name.substr(tensor_name.find("_l", 6) + 2)); + const int64_t head_dim = tensor_name[6] == 'k' ? hparams.n_embd_head_k(il) : hparams.n_embd_head_v(il); + if (hparams.n_head_kv(il) > 1 && tensor->ne[0] == head_dim && + tensor->ne[2] == (int64_t) hparams.n_head_kv(il)) { + return {true, GGML_BACKEND_SPLIT_AXIS_2, head_dim}; + } + return {}; + }(); + auto get_tensor_config = [&]() -> tensor_config { + // a graph tensor only reaches here as a scheduler copy, and only a copy of a host-resident + // cache is split - everything else the graph copies in keeps the mirrored fallback + if (!cache_copy.valid && ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } if (is_dsv4) { if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_dsv4_state)) { @@ -575,7 +604,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); }; - auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { + // ne_axis is the extent of the split axis as the cache sees it, which for a host-resident + // cache copy is not the extent of the axis the copy is split on + auto get_split_segments = [&](int axis, uint32_t il, int64_t ne_axis) -> std::vector> { if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; @@ -622,7 +653,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp); return {{n_ff_exp, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { @@ -638,14 +669,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (tensor->ne[axis] == 2*n_ff) { return {{n_ff, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; } if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) { const int64_t n_ff_exp = hparams.n_ff_exp; GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp); return {{n_ff_exp, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; }; auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector> & segments) -> std::vector { @@ -767,6 +798,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_config tc = get_tensor_config(); split_state.axis = tc.axis; if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { + // a host-resident cache copy is split on the axis its heads arrive on, in whole heads, + // so the segments and the granularity of the cache are rescaled to that unit + const bool unfolded = cache_copy.valid && split_state.axis == GGML_BACKEND_SPLIT_AXIS_0; + const int64_t unit = unfolded ? cache_copy.unit : 1; + const int64_t ne_axis = unfolded ? tensor->ne[cache_copy.axis]*unit : tensor->ne[split_state.axis]; const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); std::vector tensor_split_scan; @@ -777,12 +813,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_split_scan[j] += tensor_split_scan[j - 1]; } } - const std::vector> segments = get_split_segments(split_state.axis, tc.il); + const std::vector> segments = get_split_segments(split_state.axis, tc.il, ne_axis); const std::vector granularity = get_split_granularity(blck_size, tc.il, segments); + if (unfolded) { + split_state.axis = cache_copy.axis; + } for (size_t is = 0; is < segments.size(); is++) { - const int64_t ne_s = segments[is].first; + GGML_ASSERT(segments[is].first % unit == 0); + GGML_ASSERT(granularity[is] % unit == 0); + const int64_t ne_s = segments[is].first / unit; const uint32_t nr_s = segments[is].second; - const int64_t g_s = granularity[is]; + const int64_t g_s = granularity[is] / unit; int64_t low = 0; size_t j = 0; for (; j < ud->n_devices - 1; j++) { From a59d6852b59a2ad01b9aa11470b36076fe02b3e4 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 11:29:31 +0200 Subject: [PATCH 03/22] llama : keep the recurrent state on device under split mode tensor A linear-attention op packs the state it writes back together with its output, so the split of that state does not line up with the one the cache expects. On device that stays consistent, since each device reads back its own slice, but a host cache is laid out globally and the two orders disagree. Keep the recurrent state device-resident when the model is split by tensor. Next to the attention cache it is small, and --no-kv-offload keeps working on hybrid models instead of aborting. Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c11a68a34d22..9b70bd699bf7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -144,6 +144,16 @@ llama_context::llama_context( cparams.offload_kqv = params.offload_kqv; cparams.kv_cpu_pinned = params.kv_cpu_pinned; cparams.recurrent_state_offload = params.recurrent_state_offload; + + // A linear-attention op packs the state it writes back together with its output, so a split of + // that state does not line up with the split the cache expects. That only matters once the + // state has to pass through host memory, where it is laid out globally rather than per device. + // Keep it on its device instead; next to the attention cache it is small. + if (!cparams.recurrent_state_offload && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR && + (llm_arch_is_recurrent(model.arch) || llm_arch_is_hybrid(model.arch))) { + LLAMA_LOG_INFO("%s: split mode tensor: keeping the recurrent state device-resident\n", __func__); + cparams.recurrent_state_offload = true; + } cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && params.kv_cpu_pinned); cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.phase_aware_workspace = params.phase_aware_workspace; From ab26f4b88e212f6e5672099daf7b4ed601f3c60e Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 11:42:02 +0200 Subject: [PATCH 04/22] kv-cache : spread partial GPU residency across devices --kv-gpu-layers keeps the first N attention KV layers device-resident. Taking them in layer order fills the device that owns the first layers and leaves the free memory of the others unused, so on two GPUs only half the budget is reachable. Take one layer per owning device in turn instead. A single device, and a model split by tensor, see one device group and are unaffected. Assisted-by: Claude Opus 5 --- src/llama-kv-cache.cpp | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 397cf7f92526..151a8a435660 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include static bool ggml_is_power_of_2(int n) { @@ -180,6 +181,45 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); uint32_t n_gpu_resident = 0; + // Pick the layers that stay device-resident when the cache is otherwise host-resident. + // Taking them in layer order fills the device that owns the first layers and leaves the free + // memory of the others unused, so take one layer per device in turn instead. + std::set gpu_resident_ils; + if (!offload && placement.gpu_resident_layers > 0) { + std::vector devs; + std::vector> devs_ils; + + for (uint32_t il = 0; il < n_layer; il++) { + if (!hparams.has_kv(il) || (filter && !filter(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; gpu_resident_ils.size() < placement.gpu_resident_layers; i++) { + bool any = false; + for (size_t d = 0; d < devs.size() && gpu_resident_ils.size() < placement.gpu_resident_layers; d++) { + if (i < devs_ils[d].size()) { + gpu_resident_ils.insert(devs_ils[d][i]); + any = true; + } + } + if (!any) { + break; + } + } + } + enum class kv_store_side { key, value, @@ -285,7 +325,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 || 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) { From c3c5a57f111b170f3e765141b070b205226cbff9 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 12:23:12 +0200 Subject: [PATCH 05/22] kv-cache : apply partial GPU residency to every cache layout --kv-gpu-layers was ignored by every cache made of several sub-caches - iSWA and the DSA, DSV4 and MSA variants among them - because each sub-cache would have taken the full budget and kept up to N layers of its own. Resolve the device-resident layers once for the whole model instead. Each cache then honours a set of layer indices, so one budget covers all of its sub-caches and the per-layout opt-in is no longer needed. An iSWA cache split by tensor still keeps the whole cache in host memory: a mix of device- and host-resident layers stays coherent there but costs about 5% perplexity, which is not understood yet. Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 6 +-- src/llama-kv-cache.cpp | 47 ++------------------ src/llama-kv-cache.h | 4 -- src/llama-memory-hybrid.h | 4 -- src/llama-memory.h | 11 ++--- src/llama-model.cpp | 91 +++++++++++++++++++++++++++++++++------ 6 files changed, 88 insertions(+), 75 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9b70bd699bf7..666305d5603e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -446,11 +446,7 @@ llama_context::llama_context( } 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 151a8a435660..6ceba91b72e4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,7 +11,6 @@ #include #include #include -#include #include static bool ggml_is_power_of_2(int n) { @@ -181,45 +180,6 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); uint32_t n_gpu_resident = 0; - // Pick the layers that stay device-resident when the cache is otherwise host-resident. - // Taking them in layer order fills the device that owns the first layers and leaves the free - // memory of the others unused, so take one layer per device in turn instead. - std::set gpu_resident_ils; - if (!offload && placement.gpu_resident_layers > 0) { - std::vector devs; - std::vector> devs_ils; - - for (uint32_t il = 0; il < n_layer; il++) { - if (!hparams.has_kv(il) || (filter && !filter(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; gpu_resident_ils.size() < placement.gpu_resident_layers; i++) { - bool any = false; - for (size_t d = 0; d < devs.size() && gpu_resident_ils.size() < placement.gpu_resident_layers; d++) { - if (i < devs_ils[d].size()) { - gpu_resident_ils.insert(devs_ils[d][i]); - any = true; - } - } - if (!any) { - break; - } - } - } - enum class kv_store_side { key, value, @@ -325,7 +285,7 @@ llama_kv_cache::llama_kv_cache( const char * dev_name = "CPU"; - const bool layer_offload = offload || gpu_resident_ils.count(il) > 0; + 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) { @@ -377,9 +337,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 5b23b591a35d..cd1aa4c88694 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -136,10 +136,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 bbffd04b7bd3..3d10b57735f4 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -2230,15 +2231,79 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } +// 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 one layer per owning device in turn. +static std::set llama_pick_gpu_resident_layers(const llama_model & model, uint32_t n_layers) { + 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(); il++) { + if (!hparams.has_kv(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++) { + if (i < devs_ils[d].size()) { + ret.insert(devs_ils[d][i]); + any = true; + } + } + if (!any) { + break; + } + } + + return ret; +} + llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { llama_memory_i * res; - const llama_memory_placement_options placement = { + 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; + // 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 + placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers); + if (!placement.gpu_resident_ils.empty()) { + 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); + } switch (arch) { // Models that need specific instantiation should be handled in the @@ -2281,7 +2346,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: @@ -2335,7 +2400,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_DOTS3NOTE: @@ -2388,7 +2453,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: @@ -2416,7 +2481,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, @@ -2433,7 +2498,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: @@ -2458,7 +2523,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, nullptr, nullptr, - specialized_placement); + placement); break; } } @@ -2522,7 +2587,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), @@ -2615,7 +2680,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, @@ -2633,7 +2698,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()); From fa537d123f164bee026077b56332743279145edc Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 17:10:51 +0200 Subject: [PATCH 06/22] kv-cache : keep a recurrent layer out of the partial residency budget A hybrid model counts every layer, so most of the budget went to layers the attention cache does not own. Qwen3.5 spent 3 of 4 requested layers that way. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3d10b57735f4..fe00a472743c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2259,7 +2259,8 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod std::vector> devs_ils; for (uint32_t il = 0; il < hparams.n_layer(); il++) { - if (!hparams.has_kv(il)) { + // 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); From 1fa5c5b7fb5a33beb2438a51b11ff7652ed2d74d Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 17:11:32 +0200 Subject: [PATCH 07/22] kv-cache : spend the partial residency budget on the slowest link first A device-resident layer saves the transfer it would cost, which is worth the most where the link is slowest. Two cards in different slots differ by several times, so measure the host-to-device bandwidth once and take the layers of the slowest device before the others. Devices of similar speed stay one group and keep the previous round-robin. Assisted-by: Claude Opus 5 --- common/arg.cpp | 6 +-- src/llama-model.cpp | 98 ++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 92 insertions(+), 12 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 6cc1e412465d..52f0859a0543 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2451,9 +2451,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 spent on the devices with the slowest host link first. Unsupported " + "specialized caches ignore this option (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-model.cpp b/src/llama-model.cpp index fe00a472743c..008551eed247 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2231,9 +2231,58 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } +// Measure the host-to-device bandwidth of a device, in bytes per microsecond. Returns 0 when it +// cannot be measured. +static double llama_dev_h2d_bandwidth(ggml_backend_dev_t dev) { + const size_t n_bytes = 16u*1024*1024; + const int n_reps = 3; + + ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(dev); + if (!buft) { + return 0.0; + } + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx { ggml_init(ip) }; + if (!ctx) { + return 0.0; + } + + ggml_tensor * dst = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I8, n_bytes); + ggml_backend_buffer_ptr buf { ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft) }; + if (!buf) { + return 0.0; + } + + // the cache is delivered from a pinned buffer where the backend has one, so measure from there + ggml_backend_buffer_type_t host_buft = ggml_backend_dev_host_buffer_type(dev); + ggml_backend_buffer_ptr host_buf { host_buft ? ggml_backend_buft_alloc_buffer(host_buft, n_bytes) : nullptr }; + std::vector host_fallback; + void * src = host_buf ? ggml_backend_buffer_get_base(host_buf.get()) : nullptr; + if (!src) { + host_fallback.resize(n_bytes); + src = host_fallback.data(); + } + memset(src, 0, n_bytes); + + ggml_backend_tensor_set(dst, src, 0, n_bytes); + + const int64_t t_start = ggml_time_us(); + for (int i = 0; i < n_reps; i++) { + ggml_backend_tensor_set(dst, src, 0, n_bytes); + } + const int64_t t_us = ggml_time_us() - t_start; + + return t_us > 0 ? (double) n_bytes*n_reps/t_us : 0.0; +} + // 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 one layer per owning device in turn. +// 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) { std::set ret; if (n_layers == 0) { @@ -2276,17 +2325,48 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod 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++) { - if (i < devs_ils[d].size()) { - ret.insert(devs_ils[d][i]); - any = true; + // A device-resident layer saves the transfer it would otherwise cost, so it is worth the most + // on the slowest link. Slots differ by several times in a mixed setup, so rank the devices by + // what a transfer costs on each and spend the budget on the slowest first. Devices of similar + // speed stay one group, so the budget still spreads over them and their memory use stays even. + std::vector bw(devs.size(), 1.0); + if (devs.size() > 1) { + std::vector measured(devs.size()); + for (size_t d = 0; d < devs.size(); d++) { + measured[d] = llama_dev_h2d_bandwidth(devs[d]); + } + if (std::all_of(measured.begin(), measured.end(), [](double b) { return b > 0.0; })) { + bw = measured; + for (size_t d = 0; d < devs.size(); d++) { + LLAMA_LOG_INFO("%s: %s: host-to-device %.1f GB/s\n", __func__, + ggml_backend_dev_name(devs[d]), bw[d]/1000.0); } } - if (!any) { - break; + } + + std::vector order(devs.size()); + std::iota(order.begin(), order.end(), 0); + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return bw[a] < bw[b]; }); + + for (size_t first = 0; first < order.size() && ret.size() < n_layers; ) { + size_t last = first + 1; + while (last < order.size() && bw[order[last]] <= bw[order[first]]*1.15) { + last++; + } + for (size_t i = 0; ret.size() < n_layers; i++) { + bool any = false; + for (size_t d = first; d < last && ret.size() < n_layers; d++) { + const std::vector & ils = devs_ils[order[d]]; + if (i < ils.size()) { + ret.insert(ils[i]); + any = true; + } + } + if (!any) { + break; + } } + first = last; } return ret; From 553a9341a4091aef31418bd3436c572af8ede5ca Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 17:46:05 +0200 Subject: [PATCH 08/22] llama : add an attention split separate from the tensor split Tensor parallelism splits attention by head, and the share a device gets decides both how much of a host-resident cache it receives and how much attention work it does. Neither has to follow the memory split, so give the heads a ratio of their own. Unset, --attn-split follows --tensor-split. Assisted-by: Claude Opus 5 --- common/arg.cpp | 20 ++++++++++++++++++++ common/common.cpp | 2 ++ common/common.h | 1 + include/llama.h | 4 ++++ src/llama-model.cpp | 21 +++++++++++++++++++++ src/llama-model.h | 1 + 6 files changed, 49 insertions(+) diff --git a/common/arg.cpp b/common/arg.cpp index 52f0859a0543..3bd78079c1db 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2882,6 +2882,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_TENSOR_SPLIT")); + add_opt(common_arg( + {"-as", "--attn-split"}, "N0,N1,N2,...", + "fraction of the attention heads to give each GPU under --split-mode tensor, comma-separated, " + "e.g. 3,1. Only the heads move, the rest of the model still follows --tensor-split. The share is " + "rounded to whole heads, so a ratio the head count cannot express takes the nearest one it can. " + "Useful when the GPUs differ in host bandwidth or in speed (default: follow --tensor-split)", + [](common_params & params, const std::string & value) { + const std::regex regex{ R"([,/]+)" }; + std::sregex_token_iterator it{ value.begin(), value.end(), regex, -1 }; + std::vector split_arg{ it, {} }; + if (split_arg.size() >= llama_max_devices()) { + throw std::invalid_argument( + string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) + ); + } + for (size_t i = 0; i < llama_max_devices(); ++i) { + params.attn_split[i] = i < split_arg.size() ? std::stof(split_arg[i]) : 0.0f; + } + } + ).set_env("LLAMA_ARG_ATTN_SPLIT")); add_opt(common_arg( {"-mg", "--main-gpu"}, "INDEX", string_format("the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: %d)", params.main_gpu), diff --git a/common/common.cpp b/common/common.cpp index 6aaec977e589..57ba3bf70bf3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1689,6 +1689,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; mparams.tensor_split = params.tensor_split; + mparams.attn_split = std::any_of(params.attn_split, params.attn_split + llama_max_devices(), + [](float f) { return f != 0.0f; }) ? params.attn_split : nullptr; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; diff --git a/common/common.h b/common/common.h index 8ca190c2df75..f72f158f8734 100644 --- a/common/common.h +++ b/common/common.h @@ -478,6 +478,7 @@ struct common_params { int32_t n_gpu_layers = -1; // number of layers to store in VRAM, -1 is auto, <= -2 is all int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs + float attn_split[128] = {0}; // how attention heads should be distributed across GPUs bool fit_params = true; // whether to fit unset model/context parameters to free device memory bool fit_params_print = false; // print the estimated required memory to run the model int32_t fit_params_min_ctx = 4096; // minimum context size to set when trying to reduce memory use diff --git a/include/llama.h b/include/llama.h index f4aa9a9f2944..14d363f80563 100644 --- a/include/llama.h +++ b/include/llama.h @@ -321,6 +321,10 @@ extern "C" { // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; + // proportion of the attention heads to give each GPU under split mode tensor, size: llama_max_devices() + // NULL follows tensor_split. Only the heads move - the rest of the model still follows tensor_split + const float * attn_split; + // Called with a progress value between 0.0 and 1.0. Pass NULL to disable. // If the provided progress_callback returns true, model loading continues. // If it returns false, model loading is immediately aborted. diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 008551eed247..b6deda9d5b58 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -806,6 +806,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t ne_axis = unfolded ? tensor->ne[cache_copy.axis]*unit : tensor->ne[split_state.axis]; const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); + + // Tensor parallelism splits attention by head, and the share a device gets decides both how much + // of the cache it receives and how much attention work it does. Neither has to follow the memory + // split, so an attention split of its own can follow the host link or the speed of each device. + // Only the heads move - a linear-attention layer is a different mechanism and keeps tensor_split. + if (ud->model->attn_split() != nullptr && !hparams.is_recr(tc.il) && + (std::regex_match(tensor_name, pattern_kv_cache) || + (tensor_name.substr(0, 4) == "blk." && tensor_name.find(".attn_") != std::string::npos))) { + tensor_split = ud->model->attn_split(); + } + std::vector tensor_split_scan; tensor_split_scan.reserve(ud->n_devices); for (size_t j = 0; j < ud->n_devices; j++) { @@ -1194,6 +1205,7 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + std::vector attn_split_owned; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -1203,6 +1215,10 @@ llama_model::llama_model(const llama_model_params & params) : params(params), pi pimpl->tensor_split_owned.assign(params.tensor_split, params.tensor_split + llama_max_devices()); this->params.tensor_split = pimpl->tensor_split_owned.data(); } + if (params.attn_split != nullptr) { + pimpl->attn_split_owned.assign(params.attn_split, params.attn_split + llama_max_devices()); + this->params.attn_split = pimpl->attn_split_owned.data(); + } pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern; } @@ -1883,6 +1899,10 @@ const float * llama_model::tensor_split() const { return params.tensor_split; } +const float * llama_model::attn_split() const { + return params.attn_split; +} + uint32_t llama_model::n_gpu_layers() const { // note: plus 1 for the "output" layer return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer_all + 1; @@ -2844,6 +2864,7 @@ llama_model_params llama_model_default_params() { /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, + /*.attn_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, diff --git a/src/llama-model.h b/src/llama-model.h index f44c896dc9f6..8843661e3af6 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -693,6 +693,7 @@ struct llama_model { size_t n_tensors() const; size_t n_devices() const; const float * tensor_split() const; + const float * attn_split() const; uint32_t n_gpu_layers() const; llama_split_mode split_mode() const; From 4ea6cd69cbc4178b1978185757330a659a829297 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 31 Aug 2026 00:32:00 +0200 Subject: [PATCH 09/22] ggml : report allocation failure from the meta buffer type An out of memory in a simple buft aborted the process instead of returning nullptr, so the caller could not fall back. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 6fed68dfeae3..76c358c29170 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1767,9 +1767,15 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_bac std::vector bufs; bufs.reserve(n_simple_bufts); for (size_t i = 0; i < n_simple_bufts; i++) { - bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size)); - GGML_ASSERT(bufs.back() != nullptr); - max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back())); + ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size); + if (buf == nullptr) { + for (ggml_backend_buffer_t b : bufs) { + ggml_backend_buffer_free(b); + } + return nullptr; + } + bufs.push_back(buf); + max_size = std::max(max_size, ggml_backend_buffer_get_size(buf)); } ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); @@ -1824,7 +1830,10 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc t->buffer = meta_buf_ctx->bufs[i].get(); } } - GGML_ASSERT(meta_buf_ctx->bufs[i]); + if (meta_buf_ctx->bufs[i] == nullptr) { + ggml_backend_buffer_free(meta_buf); + return nullptr; + } meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get())); } return meta_buf; From ca580ef2e0df899cb4fa855d7eb4507d904a23ac Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 16:51:09 +0200 Subject: [PATCH 10/22] llama : split a tied output projection under split mode tensor A model without output.weight reuses the embedding table as the output projection, under the name of the table. Only that copy reaches the meta device, so the name is unambiguous there, but it fell through to MIRRORED and every device held the whole table and ran the whole projection. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 93203770d518..5a8090b38048 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -413,7 +413,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ffn_gate_shexp_weight ("blk\\.\\d*\\.ffn_gate_shexp.weight"); static const std::regex pattern_ffn_down_shexp_weight ("blk\\.\\d*\\.ffn_down_shexp.weight"); - static const std::regex pattern_output_weight("output\\.weight"); + // a model with tied embeddings has no output.weight - its output projection is a copy of the + // embedding table under the name of the table. Only that copy reaches the meta device, the + // table itself stays on the input device, so the name is unambiguous here. + static const std::regex pattern_output_weight("(output|token_embd)\\.weight"); static const std::regex pattern_output_bias ("output\\.bias"); struct tensor_config { From a1cd839ab98a6fb2f8bc2aa419d7d43f5d2ccc84 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 21:54:46 +0200 Subject: [PATCH 11/22] ggml : parse a truncated scheduler copy name The scheduler builds "##" into GGML_MAX_NAME, and the meta backend name grows with the device count, so from five devices on the name is cut and the trailing part is gone. The old parse then returned the whole mangled string, the cache copy was not recognised and fell back to MIRRORED - each device attended the wrong heads. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 76c358c29170..d7c3c70331df 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -486,22 +486,25 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct } // The scheduler names a tensor it copies in from another backend "##", where -// is the graph name of the source plus the suffixes ggml appends for views. Undo both, so +// is the graph name of the source plus the suffixes ggml appends for views. The name is cut +// to GGML_MAX_NAME, so what follows the source can be missing or half there. Undo what is left, so // that the copy can be recognised under the name the graph gave it. static std::string ggml_backend_meta_copy_source_name(const char * name) { - std::string decorated = name; - const size_t first = decorated.find('#'); - const size_t last = decorated.rfind('#'); - if (first == std::string::npos || last <= first) { - return decorated; + std::string ret = name; + const size_t first = ret.find('#'); + if (first == std::string::npos) { + return ret; } - std::string ret = decorated.substr(first + 1, last - first - 1); - while (!ret.empty() && ret.back() == ')') { - const size_t paren = ret.rfind(" ("); - if (paren == std::string::npos) { - break; - } - ret.erase(paren); + ret.erase(0, first + 1); + // ggml writes a view suffix as " (...)", a graph name has no spaces + const size_t suffix = ret.find(" ("); + if (suffix != std::string::npos) { + ret.erase(suffix); + return ret; + } + const size_t copy = ret.rfind('#'); + if (copy != std::string::npos && ret.find_first_not_of("0123456789", copy + 1) == std::string::npos) { + ret.erase(copy); } return ret; } From 86bc2f1cb9a031adb9b83d1fcf5bd9a4876271f2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 21:55:06 +0200 Subject: [PATCH 12/22] ggml : clear the tensors when a meta allocation fails The tensors of the context already point at the meta buffer when a simple buffer fails, so freeing it leaves them dangling. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index d7c3c70331df..2ba9f78373fb 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1813,27 +1813,32 @@ struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struc t->data = (void *) 0x2000000000000000; // FIXME } for (size_t i = 0; i < n_simple_bufts; i++) { - ggml_context * ctx = meta_buf_ctx->stc_static.ctxs[i].get(); + ggml_context * ctx_simple = meta_buf_ctx->stc_static.ctxs[i].get(); ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, i); // If a ggml_context only has zero-sized tensors, ggml_backend_alloc_ctx_tensors_from_buft returns NULL. // For those edge cases, allocate a dummy buffer instead. bool any_nonzero_slice = false; - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) { if (ggml_nelements(t) != 0) { any_nonzero_slice = true; break; } } if (any_nonzero_slice) { - meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, simple_buft)); + meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_simple, simple_buft)); } else { meta_buf_ctx->bufs[i].reset(ggml_backend_buft_alloc_buffer(simple_buft, 0)); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx_simple); t != nullptr; t = ggml_get_next_tensor(ctx_simple, t)) { t->buffer = meta_buf_ctx->bufs[i].get(); } } if (meta_buf_ctx->bufs[i] == nullptr) { + // the loop above already pointed the tensors at the buffer that is freed here + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = nullptr; + t->data = nullptr; + } ggml_backend_buffer_free(meta_buf); return nullptr; } From e89286d74d208e1077d837e2bb058d1727fecaa9 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 21:55:06 +0200 Subject: [PATCH 13/22] llama : split a V tensor on its own head size The KV granularity came from n_embd_head_k, so a model with a value length of its own put the device boundaries between heads of the V cache. Under split mode tensor with a host-resident cache it aborted on the head assert. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5a8090b38048..ffa94231cf3c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -783,12 +783,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return {granularity_q}; } - const int64_t granularity_kv = granularity_q / n_gqa; + const int64_t n_head_gran = granularity_q / n_embd_q; // KV heads per granule + const int64_t granularity_kv = n_head_gran * hparams.n_embd_head_k(il); if (std::regex_match(tensor_name, pattern_kv_weight) || std::regex_match(tensor_name, pattern_kv_bias) || std::regex_match(tensor_name, pattern_kv_cache)) { GGML_ASSERT(segments.size() == 1); - return {granularity_kv}; + // the V side can have a head size of its own, the split still lands on head boundaries + const bool is_v = tensor_name.compare(0, 8, "cache_v_") == 0 || + tensor_name.find(".attn_v.") != std::string::npos; + return {is_v ? n_head_gran * hparams.n_embd_head_v(il) : granularity_kv}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { GGML_ASSERT(segments.size() == 2); From cf91605fe54d279e5f4ffa341ca22d572961a8bc Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 21:55:13 +0200 Subject: [PATCH 14/22] llama : resolve the KV residency set only where it applies The set was resolved even with the cache fully offloaded, which measured every host link for nothing. The attention compute then moved off the CPU on the request alone, so a layout that keeps the whole cache in host memory streamed all of it to the device each ubatch. Report back what the cache got, and let the compute follow that. Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 1 + src/llama-model.cpp | 21 ++++++++++++++------- src/llama-model.h | 3 ++- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c8c6add4759e..4de9bea288a6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -445,6 +445,7 @@ 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) { cparams.offload_attn_compute = cparams.offload_attn_compute || cparams.op_offload; } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index ffa94231cf3c..0c8bb6eebbf9 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2417,18 +2417,25 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod return ret; } -llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { +llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_cparams & cparams) const { llama_memory_i * res; llama_memory_placement_options placement = { cparams.kv_cpu_pinned, 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 - placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers); - if (!placement.gpu_resident_ils.empty()) { - 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); + // 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 and must not pay for measuring the links. + if (!cparams.offload_kqv && cparams.kv_gpu_layers > 0) { + placement.gpu_resident_ils = llama_pick_gpu_resident_layers(*this, cparams.kv_gpu_layers); + 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) { diff --git a/src/llama-model.h b/src/llama-model.h index 5b7a60e86066..21cf1529c068 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -750,7 +750,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; From 52b7265c2d47f6c36d10558b7d05b13b5b63ab37 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:55:20 +0200 Subject: [PATCH 15/22] ggml : keep the source name of a scheduler copy The copy name is "##" in a field of GGML_MAX_NAME, and the meta backend name grows with the device count, so at eight devices the source name is the part that is cut. Nothing can recover the identity of the copy after that. Cut the backend label instead - it is only a label, the source name is what identifies the copy. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8138028ee8cf..88e9dc2cf3ab 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1033,6 +1033,19 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str } } +// Name a copy of an input "##". The name field has a fixed size, so cut the +// backend label rather than the source name - the source name is what identifies the copy. +static void ggml_backend_sched_name_copy( + struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c) { + const int n_tail = snprintf(NULL, 0, "#%s#%d", src->name, c); + const int n_max = GGML_MAX_NAME - 1 - n_tail; + int n_head = (int) strlen(backend_name); + if (n_head > n_max) { + n_head = n_max > 0 ? n_max : 0; + } + ggml_format_name(copy, "%.*s#%s#%d", n_head, backend_name, src->name, c); +} + static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; ggml_backend_buffer_type_t buft = NULL; @@ -1391,7 +1404,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra tensor_copy = src; // use the original tensor as the current copy } else { tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c); } ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor @@ -1412,7 +1425,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra ggml_backend_t backend = sched->backends[cur_backend_id]; for (int c = 0; c < sched->n_copies; c++) { struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c); if (sched->n_copies > 1) { ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor From 9b864740f6cef355e3a61de98ab454f85596831a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:55:28 +0200 Subject: [PATCH 16/22] llama : require a layer index in the cache patterns The patterns accepted a name with no digits, so a name that lost its layer index still matched and the layer parse threw on the empty string. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0c8bb6eebbf9..97963efa24b7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -379,8 +379,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias"); static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); - static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); - static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); + static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d+"); + static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d+"); static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); @@ -395,9 +395,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight"); static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); - static const std::regex pattern_r_cache ("cache_r_l\\d*"); - static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d*"); - static const std::regex pattern_s_cache ("cache_s_l\\d*"); + static const std::regex pattern_r_cache ("cache_r_l\\d+"); + static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d+"); + static const std::regex pattern_s_cache ("cache_s_l\\d+"); static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); From 47d9f4f501770960e9b30eeb7b352249a473cb6e Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:55:40 +0200 Subject: [PATCH 17/22] llama : pick the resident KV layers of the context that owns them An MTP cache holds only the nextn layers, which sit above the layers of the main context. The picker looked at the trunk alone, so it never named a layer the MTP cache owns while still reporting the requested count as delivered - spec_draft_kv_gpu_layers kept the whole draft cache in host memory. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 97963efa24b7..23f11ebffa83 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2328,7 +2328,7 @@ static double llama_dev_h2d_bandwidth(ggml_backend_dev_t dev) { // 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) { +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; @@ -2352,7 +2352,11 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod std::vector devs; std::vector> devs_ils; - for (uint32_t il = 0; il < hparams.n_layer(); il++) { + 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; @@ -2427,7 +2431,8 @@ 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 and must not pay for measuring the links. 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 { From 49256242be6173ec17f83b4aa5469e9abdf8f21d Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:55:47 +0200 Subject: [PATCH 18/22] llama : initialize every field of the placement options The aggregate left the residency set out, which fails a build with -Werror=missing-field-initializers. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 23f11ebffa83..532c72114e9b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2423,10 +2423,9 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_cparams & cparams) const { llama_memory_i * res; - llama_memory_placement_options placement = { - cparams.kv_cpu_pinned, - cparams.offload_kqv || cparams.recurrent_state_offload, - }; + 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 and must not pay for measuring the links. From 5786654a3bc855eac3f10237aab8036ac525f096 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:55:53 +0200 Subject: [PATCH 19/22] common : allow an attention split entry per device The bound rejected exactly as many entries as the public array holds. Assisted-by: Claude Opus 5 --- common/arg.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/arg.cpp b/common/arg.cpp index fff49d07773f..6b4383573995 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2940,7 +2940,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex const std::regex regex{ R"([,/]+)" }; std::sregex_token_iterator it{ value.begin(), value.end(), regex, -1 }; std::vector split_arg{ it, {} }; - if (split_arg.size() >= llama_max_devices()) { + if (split_arg.size() > llama_max_devices()) { throw std::invalid_argument( string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) ); From 62ec35348d89eb66049da43df5c72d65cadfc665 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:56:14 +0200 Subject: [PATCH 20/22] llama : shorten the comments of the host KV cache Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 21 ++++++++------------- src/llama-context.cpp | 6 ++---- src/llama-model.cpp | 20 ++++++++------------ 3 files changed, 18 insertions(+), 29 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 2ba9f78373fb..44e62a7b9ccd 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -485,10 +485,8 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct return it->second[index]; } -// The scheduler names a tensor it copies in from another backend "##", where -// is the graph name of the source plus the suffixes ggml appends for views. The name is cut -// to GGML_MAX_NAME, so what follows the source can be missing or half there. Undo what is left, so -// that the copy can be recognised under the name the graph gave it. +// A scheduler copy is named "##", where also carries the suffixes +// ggml appends for views. Recover the graph name the copy was made from. static std::string ggml_backend_meta_copy_source_name(const char * name) { std::string ret = name; const size_t first = ret.find('#'); @@ -855,11 +853,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } - // A tensor the scheduler copied in from another backend is a leaf in the compute buffer, - // so it never reached the device callback and fell through to MIRRORED. A host-resident - // KV cache arrives this way, and mirroring it while the queries stay split by head makes - // each device attend the wrong heads. Ask the callback; it still answers MIRRORED for - // names it does not know. + // A host-resident KV cache reaches the graph as a copied-in leaf of the compute buffer. + // Mirroring it while the queries stay split by head makes each device attend the wrong + // heads, so ask the callback; it still answers MIRRORED for names it does not know. const bool copied_in_leaf = ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->op == GGML_OP_NONE && (tensor->flags & GGML_TENSOR_FLAG_INPUT) == 0; @@ -1459,10 +1455,9 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. - // The heads are the split axis but they are interleaved inside each cell, so the chunk splice - // below cannot express the write. Each device's heads are one contiguous run per cell, which - // is a strided copy: ne[1] cells, from the full cell stride to the device's own. + // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. The + // heads split, but interleaved per cell, so the chunk splice below cannot express the write. + // Each device's heads are one contiguous run per cell: ne[1] cells, from one stride to another. const bool strided_head_split = !ggml_is_contiguous(tensor) && split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4de9bea288a6..271d8a8dd054 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -145,10 +145,8 @@ llama_context::llama_context( cparams.kv_cpu_pinned = params.kv_cpu_pinned; cparams.recurrent_state_offload = params.recurrent_state_offload; - // A linear-attention op packs the state it writes back together with its output, so a split of - // that state does not line up with the split the cache expects. That only matters once the - // state has to pass through host memory, where it is laid out globally rather than per device. - // Keep it on its device instead; next to the attention cache it is small. + // A linear-attention op packs the state it writes back together with its output, so that split + // does not line up with the one a host-resident state expects. Keep the state on its device. if (!cparams.recurrent_state_offload && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR && (llm_arch_is_recurrent(model.arch) || llm_arch_is_hybrid(model.arch))) { LLAMA_LOG_INFO("%s: split mode tensor: keeping the recurrent state device-resident\n", __func__); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 532c72114e9b..6c555602d175 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -470,10 +470,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return {axis, tensor_axis_0, il, rotation}; }; - // An attention cache lives in host memory when it is not offloaded, so it reaches attention as - // a scheduler copy rather than as the cache itself. The cache folds all heads into one flat - // axis, the copy arrives permuted, with the heads on an axis of their own. The split follows - // the cache either way, so record the axis and the unit - one head - the copy is split in. + // A host-resident cache reaches attention as a scheduler copy, permuted, with the heads on an + // axis of their own. The split follows the cache, so record that axis and its unit - one head. struct host_cache_copy { bool valid = false; ggml_backend_meta_split_axis axis = GGML_BACKEND_SPLIT_AXIS_0; @@ -831,10 +829,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); - // Tensor parallelism splits attention by head, and the share a device gets decides both how much - // of the cache it receives and how much attention work it does. Neither has to follow the memory - // split, so an attention split of its own can follow the host link or the speed of each device. - // Only the heads move - a linear-attention layer is a different mechanism and keeps tensor_split. + // The share of the attention heads decides how much cache a device gets and how much attention + // work it does, neither of which has to follow the memory split. A linear-attention layer is a + // different mechanism and keeps tensor_split. if (ud->model->attn_split() != nullptr && !hparams.is_recr(tc.il) && (std::regex_match(tensor_name, pattern_kv_cache) || (tensor_name.substr(0, 4) == "blk." && tensor_name.find(".attn_") != std::string::npos))) { @@ -2374,10 +2371,9 @@ static std::set llama_pick_gpu_resident_layers(const llama_model & mod devs_ils[it - devs.begin()].push_back(il); } - // A device-resident layer saves the transfer it would otherwise cost, so it is worth the most - // on the slowest link. Slots differ by several times in a mixed setup, so rank the devices by - // what a transfer costs on each and spend the budget on the slowest first. Devices of similar - // speed stay one group, so the budget still spreads over them and their memory use stays even. + // A device-resident layer saves the transfer it would cost, which is worth most on the slowest + // link, so spend the budget there first. Devices of similar speed stay one group and keep the + // round-robin, so their memory use stays even. std::vector bw(devs.size(), 1.0); if (devs.size() > 1) { std::vector measured(devs.size()); From 655d998f90b2b1dee6807d8ff16f892d25e34afd Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:56:29 +0200 Subject: [PATCH 21/22] ggml : read back a strided row split from a meta buffer A fused QKV makes Kcur and Vcur strided views, so a host-resident cache reads them back from a tensor that is not contiguous. Only the write side handled that; the read hit the contiguity assert. On falcon, split mode tensor with no_kv_offload aborted. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 44e62a7b9ccd..33e844678f94 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1609,6 +1609,33 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back here. + // The rows split, so the chunk splice below cannot express the read. Each device's part of a row + // is contiguous: ne[1] rows, from the device's own stride to the view's. + const bool strided_rows = + !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_0 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->ne[2] == 1 && tensor->ne[3] == 1 && + offset == 0 && size == ggml_nbytes(tensor); + + if (strided_rows) { + size_t offset_data = 0; + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = ggml_row_size(tensor->type, split_state.ne[j]); + if (nbytes == 0) { + continue; + } + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, 0, nbytes, + tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + } + GGML_ASSERT(offset_data == ggml_row_size(tensor->type, tensor->ne[0])); + return; + } + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { From bd62eb6f753b486ec9f78ba6ffae8782acdf41e6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 09:56:38 +0200 Subject: [PATCH 22/22] tests : cover a host-resident KV cache split by tensor The matrix built every context with the cache offloaded, so it never ran the path this branch fixes. Add two columns to the meta config - the cache in host memory, and the same with all heads on the first device and part of the cache kept resident - and a check that an MTP context spends its residency budget on the layers it owns. The copy name only fills the name field from eight devices on, so run one arch at that count in the CI. Assisted-by: Claude Opus 5 --- ci/run.sh | 4 +- tests/test-llama-archs.cpp | 106 +++++++++++++++++++++++++++++++++++-- 2 files changed, 104 insertions(+), 6 deletions(-) diff --git a/ci/run.sh b/ci/run.sh index 1701bc7ed058..af36eadbe240 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -312,6 +312,8 @@ function gg_run_test_llama_archs_tensor_split { GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + # the scheduler names a copy after its source, and 8 device names fill the name field + GGML_CUDA_DEVICES=8 ./build-ci-release/bin/test-llama-archs -s 1 -a llama 2>&1 fi if [ ! -z ${GG_BUILD_METAL} ]; then @@ -327,7 +329,7 @@ function gg_run_test_llama_archs_tensor_split { function gg_sum_test_llama_archs_tensor_split { gg_printf '### %s\n\n' "${ci}" - gg_printf 'Runs test-llama-archs with 1 to 4 devices\n' + gg_printf 'Runs test-llama-archs with 1 to 4 and 8 devices\n' gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)" gg_printf '```\n' gg_printf '%s\n' "$(cat $OUT/${ci}.log)" diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index c27440d57ac5..9ebc887f9d1a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -354,9 +354,17 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } +// with offload_kqv=false the cache lives in host memory; attn_split gives the heads a share of +// their own and kv_gpu_layers keeps some of the layers on their devices +struct kv_config { + bool offload_kqv = true; + uint32_t kv_gpu_layers = 0; + std::vector attn_split; +}; + static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, const kv_config & kvc = {}) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -364,11 +372,18 @@ static std::pair get_model_and_ctx( devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; + std::vector attn_split = kvc.attn_split; + if (!attn_split.empty()) { + attn_split.resize(llama_max_devices(), 0.0f); + model_params.attn_split = attn_split.data(); + } llama_context_params ctx_params = llama_context_default_params(); ctx_params.n_ctx = 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; + ctx_params.offload_kqv = kvc.offload_kqv; + ctx_params.kv_gpu_layers = kvc.kv_gpu_layers; if (!encode) { ctx_params.n_ubatch = 64; } @@ -1008,6 +1023,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)); @@ -1242,9 +1317,10 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg std::vector devs; std::string label; llama_split_mode split_mode; + kv_config kvc; - device_config(std::vector devs, std::string name, llama_split_mode split_mode) - : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {} + device_config(std::vector devs, std::string name, llama_split_mode split_mode, kv_config kvc = {}) + : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), kvc(std::move(kvc)) {} }; std::vector dev_configs; @@ -1266,6 +1342,25 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg } dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR); + + // a host-resident cache reaches attention as a scheduler copy that is split by head + kv_config kvc_host; + kvc_host.offload_kqv = false; + dev_configs.emplace_back(devices_meta, "Meta -nkvo", LLAMA_SPLIT_MODE_TENSOR, kvc_host); + + // the same, with all heads on the first device and part of the cache kept device-resident + if (devices_meta.size() > 1) { + kv_config kvc_attn; + kvc_attn.offload_kqv = false; + kvc_attn.kv_gpu_layers = 2; + kvc_attn.attn_split.assign(devices_meta.size(), 0.0f); + kvc_attn.attn_split[0] = 1.0f; + dev_configs.emplace_back(devices_meta, "Meta -nkvo -as", LLAMA_SPLIT_MODE_TENSOR, kvc_attn); + } + + for (const device_config & dc : dev_configs) { + max_device_label_length = std::max(max_device_label_length, dc.label.length()); + } } size_t max_arch_name_length = 0; @@ -1336,7 +1431,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode); } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { - model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode); + model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode, dc.kvc); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); @@ -1358,7 +1453,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg ms.save(file); rewind(file); - auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode); + auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode, dc.kvc); const std::vector logits_roundtrip = get_logits( model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode); status_roundtrip = "\033[1;32mOK\033[0m"; @@ -1461,6 +1556,7 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, log_level, out); } + test_mtp_kv_residency(seed); return test_backends(arch, seed, log_level); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what());