diff --git a/docs/development/on-demand-tensor-reload.md b/docs/development/on-demand-tensor-reload.md index 607a01049..69189b440 100644 --- a/docs/development/on-demand-tensor-reload.md +++ b/docs/development/on-demand-tensor-reload.md @@ -186,6 +186,23 @@ Driver loop in a nutshell: This mode works with any of the evaluation flavours (`perplexity`, `--kl-divergence`, HellaSwag, ...); KL divergence is fully re-entrant since the base-logits file is re-read on every iteration. +**Load-time-derived tensors (MLA).** `llm_prepare_mla` derives tensors at load time: when a GGUF ships `attn_k_b`/`attn_v_b` but no `attn_kv_b`, a combined `attn_kv_b` is computed from them - and the `mla>1` prompt-processing path consumes the *derived* tensor, not the originals. Reloading `attn_k_b`/`attn_v_b` therefore also re-derives the affected layer's `attn_kv_b` in place (`llm_refresh_computed_wkv_b`), logging `hotswap: refreshed derived tensor 'blk.N.attn_kv_b.weight'`. If a refresh is impossible (per-rank split tensors under `-sm graph/attn`, or the reverse derivation where `attn_k_b`/`attn_v_b` were computed from a reloaded `attn_kv_b`), a `hotswap: failed to refresh derived tensor ...` warning is emitted so external drivers can quarantine the affected evaluation instead of recording stale results. + +**Other load-time weight transforms (refused or propagated).** Several mechanisms transform, merge or fuse weights at load/context-creation time, making a raw-byte reload wrong or ineffective. The reload machinery handles them as follows: + +* **Refused with a loud `hotswap: tensor '...' cannot be hot-swapped: ` warning** (`reload_tensor` returns false, so drivers see no `reloaded tensor` line and quarantine the round): + * views into merged tensors created by **`-mqkv`** (merged `wqkv`) and **`-muge`** (merged `ffn_gate_up_exps`, which is additionally re-interleaved per expert at context creation); + * `attn_k_b`/`attn_v_b`/`attn_kv_b` when the **`-khad`** Hadamard transform was folded into the MLA weights (`khad_pretransformed`); + * `ffn_gate_inp_s` (scaled in place at load by `llm_scale_gate_inp_s`); + * BitNet `*.scale` tensors (fused into the paired weight's `op_params` at load); + * OpenPangu `param_sink*` and `attn_kv_a_norm` (feed load-time-derived parameter sinks); + * same-dtype swaps of tensors whose original data is **mmap-backed** (read-only mapping cannot be refreshed in place; use `--no-mmap` for hot-swap workflows). +* **Propagated automatically**: tensors physically duplicated at load under the same name (tied lm-head copy of `token_embd`, per-layer `rope_freqs`/`rope_factors` copies, expert-bias `*_dup` copies) - the reloaded data is written to every duplicate instance, or a `failed to refresh derived tensor` warning is emitted when that is impossible. +* **Warned as stale**: the pre-transposed `wk_b_pp` under `-sm graph/attn`, and the runtime-requantized MTP head (`output_extra.weight`) after an `output.weight` reload. +* **`-rtr` (run-time repacking)**: a warning is emitted at registration - restores write the plain file type and cannot reproduce the repacked state (mathematically equivalent for lossless repacks, but `F16 -> BF16_R16` is lossy), so benchmark results may be inconsistent. + +**llama-server specifics.** The `/health` hot-swap hook now only attempts a reload when no slot is processing (best effort - a request arriving concurrently can still race it), and after a successful reload it clears the KV cache and all cached prompts, since those were computed with the previous weights. Note that a draft/speculative model and an mtmd projector are separate models with their own files and are not covered by the hot-swap registry. + ### `examples/server/server.cpp` On every health-check (`/health`) request, if `LLAMA_HOTSWAP_ENABLED` is set, the server calls `llama_reload_changed_tensors()`. This provides a convenient, external trigger: simply `touch` or overwrite a tensor’s source GGUF file and poll `/health` to apply the change. diff --git a/examples/server/server.cpp b/examples/server/server.cpp index 1e33a8be5..1975d0639 100644 --- a/examples/server/server.cpp +++ b/examples/server/server.cpp @@ -767,8 +767,19 @@ int main(int argc, char ** argv) { const char * hotswap_env = std::getenv("LLAMA_HOTSWAP_ENABLED"); if (hotswap_env) { // WARNING: llama_reload_changed_tensors is NOT thread-safe with active inference. - // Only enable this when you can guarantee the server is idle during health checks. - llama_reload_changed_tensors(ctx_server.ctx); + // Best effort: only attempt the reload when no slot is processing (a request + // arriving between the metrics snapshot and the reload can still race it). + if (n_processing_slots > 0) { + LOG_INFO("hotswap: skipping tensor reload, slots are processing", {{"processing", n_processing_slots}}); + } else if (llama_reload_changed_tensors(ctx_server.ctx)) { + // KV cache entries and cached prompts were computed with the + // previous weights; drop them so they cannot be reused. + ctx_server.kv_cache_clear(); + for (auto & slot : ctx_server.slots) { + slot.cache_tokens.clear(); + } + LOG_INFO("hotswap: tensors reloaded; KV cache and cached prompts cleared", {}); + } } break; diff --git a/src/llama-model.h b/src/llama-model.h index 8de5b220e..f3f05bbf1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -606,6 +606,11 @@ struct llama_model { std::unique_ptr reload; }; +// Recompute the load-time-derived combined wkv_b (computed_wkv_b) of layer il +// from the current wk_b/wv_b tensor data (hot-swap support). Defined in +// llama.cpp next to llm_prepare_mla. +bool llm_refresh_computed_wkv_b(llama_model & model, int il); + struct llama_lora_weight { struct ggml_tensor * a = nullptr; struct ggml_tensor * b = nullptr; diff --git a/src/llama-reload.cpp b/src/llama-reload.cpp index d8ef6f2c2..c4f01cb5b 100644 --- a/src/llama-reload.cpp +++ b/src/llama-reload.cpp @@ -12,6 +12,7 @@ #include #include #include +#include // ------------------------------------------------------------------ @@ -144,6 +145,49 @@ static bool is_original_snapshot_buffer(llama_model & model, ggml_backend_buffer return false; } +// ------------------------------------------------------------------ +// Guard: tensors that were transformed, merged or fused at load time cannot +// be hot-swapped by writing raw file bytes over them - the compute graph +// either consumes transformed data (so raw bytes would be WRONG) or consumes +// something else entirely (so the swap would be silently ineffective). +// Refuse loudly so external drivers can quarantine the attempt. +// ------------------------------------------------------------------ +static const char * hotswap_unsupported_reason(const llama_model & model, const char * name, const struct ggml_tensor * tensor) { + if (tensor->view_src) { + // -mqkv / -muge merge wq/wk/wv (or ffn_gate/up_exps) into one parent + // tensor consumed by the graph; the parent may additionally be + // re-interleaved per expert at context creation, so writing raw bytes + // through the view offset can corrupt the merged layout + return "it is a view into a merged tensor (e.g. -mqkv / -muge)"; + } + std::string n(name); + if (model.khad_pretransformed && + (n.find(".attn_k_b.weight") != std::string::npos || + n.find(".attn_v_b.weight") != std::string::npos || + n.find(".attn_kv_b.weight") != std::string::npos)) { + // llm_apply_khad_pretransform folded the 64-block Hadamard into + // wv_b/wk_b_pp in place and the graph permanently skips the runtime + // un-Hadamard; raw file bytes would produce garbage attention output + return "the Hadamard transform was folded into the MLA weights at load time (khad)"; + } + if (n.find(".ffn_gate_inp_s.") != std::string::npos) { + // llm_scale_gate_inp_s scales this tensor in place at load time + return "it was scaled in place at load time (llm_scale_gate_inp_s)"; + } + if (model.arch == LLM_ARCH_BITNET && n.size() > 6 && n.compare(n.size() - 6, 6, ".scale") == 0) { + // the scalar was copied into the paired weight's op_params at load + // time and the graph reads it from there + return "its value was fused into the paired weight's op_params at load time (BitNet)"; + } + if (model.arch == LLM_ARCH_OPENPANGU && + (n.find("param_sink") != std::string::npos || n.find(".attn_kv_a_norm.") != std::string::npos)) { + // llm_prepare_openpangu_param_sinks derives sink tensors from these + // at load time; the derived tensors would not be refreshed + return "it feeds load-time-derived parameter sink tensors (OpenPangu)"; + } + return nullptr; +} + // ------------------------------------------------------------------ // mmap guard: tensors whose data lives inside a (read-only) file mapping // must not be written to in place @@ -288,6 +332,13 @@ static void snapshot_tensor_source(struct ggml_tensor * tensor, // Constructor // ------------------------------------------------------------------ reload_info::reload_info(const llama_model_loader & ml) { + if (ml.repack_tensors) { + // -rtr repacks host tensors in place (type changes to the interleaved + // _R4/_R8/_R16 variants); a later restore writes the plain file type, + // which is mathematically equivalent for lossless repacks but cannot + // reproduce the repacked state (and F16 -> BF16_R16 is lossy) + LLAMA_LOG_WARN("hotswap: run-time repacking (-rtr) is enabled; restored tensors cannot reproduce the repacked state and results may be inconsistent\n"); + } for (const auto & w : ml.weights) { if (!w.tensor || w.idx >= (int)ml.files.size()) continue; @@ -885,9 +936,23 @@ bool reload_info::reload_tensor(const char * name, llama_model & model) { } } + // Refuse tensors that were transformed/merged/fused at load time + if (const char * reason = hotswap_unsupported_reason(model, name, tensor)) { + LLAMA_LOG_WARN("hotswap: tensor '%s' cannot be hot-swapped: %s; skipping\n", name, reason); + return false; + } + ggml_backend_buffer_t old_buf = tensor->buffer; bool returning = (curr_type == src.original_type); + // A same-dtype swap reattaches the original buffer and refreshes its data; + // that is impossible when the original data lives inside a read-only file + // mapping, so refuse instead of silently keeping the old weights + if (returning && tensor_data_in_mmap(model, src.original_data)) { + LLAMA_LOG_WARN("hotswap: tensor '%s' cannot be hot-swapped: its original data is mmap-backed (read-only); skipping\n", name); + return false; + } + std::vector host_buf; if (!returning) { if (curr_type != tensor->type) { @@ -918,6 +983,27 @@ bool reload_info::reload_tensor(const char * name, llama_model & model) { } if (ok) { + // Some tensors are physically DUPLICATED at load time under the same + // name (tied lm_head copy of token_embd, per-layer rope_freqs / + // rope_factors copies, expert-bias *_dup copies): propagate the new + // data to every other instance, or fail loudly when that is + // impossible so external drivers can quarantine the result. + for (auto & p : model.tensors_by_name) { + if (p.second == tensor || p.first != name) continue; + auto dup = p.second; + bool same_meta = dup && dup->type == tensor->type; + for (int i = 0; i < GGML_MAX_DIMS && same_meta; ++i) { + same_meta = dup->ne[i] == tensor->ne[i]; + } + if (same_meta && !dup->view_src && dup->buffer && !host_buf.empty() && + !tensor_data_in_mmap(model, dup->data)) { + ggml_backend_tensor_set(dup, host_buf.data(), 0, host_buf.size()); + LLAMA_LOG_INFO("hotswap: propagated reloaded data to a duplicate instance of '%s'\n", name); + } else { + LLAMA_LOG_WARN("hotswap: failed to refresh derived tensor '%s' (duplicate instance); evaluation results would be stale\n", name); + } + } + src.last_mtime = st.st_mtime; #ifdef __linux__ src.last_mtime_ns = st.st_mtim.tv_nsec; @@ -977,10 +1063,56 @@ bool reload_info::reload_changed_tensors(llama_model & model) { }); bool r = false; + std::vector mla_refresh_layers; for (auto & j : jobs) { if (reload_tensor(j.name, model)) { r = true; LLAMA_LOG_INFO("reloaded tensor '%s'\n", j.name); + + // MLA models may derive tensors at load time (llm_prepare_mla): + // when the GGUF ships attn_k_b/attn_v_b, a combined attn_kv_b is + // computed from them, and the mla>1 prompt-processing path consumes + // the DERIVED tensor. It must be refreshed after a reload, or + // evaluations would keep using weights derived from the old data. + int il = -1; + std::string n(j.name); + if (sscanf(j.name, "blk.%d.", &il) == 1 && il >= 0 && il < (int) model.layers.size()) { + if (n.find(".attn_k_b.weight") != std::string::npos || + n.find(".attn_v_b.weight") != std::string::npos) { + if (model.layers[il].computed_wkv_b && + std::find(mla_refresh_layers.begin(), mla_refresh_layers.end(), il) == mla_refresh_layers.end()) { + mla_refresh_layers.push_back(il); + } + // Under -sm graph/attn a pre-transposed wk_b_pp (and + // per-rank replicas) is derived from wk_b for the + // prompt-processing path; it cannot be refreshed here. + if (model.layers[il].wk_b_pp || model.layers[il].computed_wk_b_pp) { + LLAMA_LOG_WARN("hotswap: failed to refresh derived tensor 'blk.%d.attn_kv_b.weight' (pre-transposed wk_b_pp); evaluation results would be stale\n", il); + } + } else if (n.find(".attn_kv_b.weight") != std::string::npos) { + // The reverse derivation (attn_k_b/attn_v_b computed from + // attn_kv_b, e.g. mainline-converted models) is not + // refreshable here; make the staleness loud so external + // drivers can quarantine the benchmark. + if (model.layers[il].computed_wk_b || model.layers[il].computed_wv_b) { + LLAMA_LOG_WARN("hotswap: failed to refresh derived tensor 'blk.%d.attn_k_b/attn_v_b.weight' after attn_kv_b reload; evaluation results would be stale\n", il); + } + } + } + + // The runtime-requantized MTP head (output_extra.weight) is + // derived from output.weight at load time and is not refreshed + if (n == "output.weight" && model.output_mtp && model.output_mtp != model.output) { + LLAMA_LOG_WARN("hotswap: failed to refresh derived tensor 'output_extra.weight' (requantized MTP head); MTP evaluation results would be stale\n"); + } + } + } + + for (int il : mla_refresh_layers) { + if (llm_refresh_computed_wkv_b(model, il)) { + LLAMA_LOG_INFO("hotswap: refreshed derived tensor 'blk.%d.attn_kv_b.weight' from the reloaded attn_k_b/attn_v_b\n", il); + } else { + LLAMA_LOG_WARN("hotswap: failed to refresh derived tensor 'blk.%d.attn_kv_b.weight'; evaluation results would be stale\n", il); } } diff --git a/src/llama.cpp b/src/llama.cpp index ac0843a7d..6710ea32b 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -2609,6 +2609,98 @@ static void llm_requantize_output_tensor(llama_model & model, ggml_type new_type } +// Compute the combined wkv_b tensor of one layer from the current wk_b/wv_b +// tensor data, on the CPU. Non-host source tensors are staged through +// wk_buffer/wv_buffer; intermediate and final results live in tmp_buffer, so +// the returned tensor (data pointing into tmp_buffer) is only valid until the +// buffers are reused, and only its type/shape/data may be used: its src chain +// and the populated graph reference stack locals of this function, so do not +// walk the srcs or recompute the graph — clear it with ggml_graph_clear before +// reuse. ctx/graph must be able to hold the ~8 nodes of the derivation. +// Returns null if the graph computation fails. +static ggml_tensor * llm_compute_wkv_b(ggml_context * ctx, ggml_cgraph * graph, + const ggml_tensor * wk_b_src, const ggml_tensor * wv_b_src, + std::vector & wk_buffer, std::vector & wv_buffer, + std::vector & tmp_buffer, std::vector & work_data) { + auto wk_b = *wk_b_src; + auto wv_b = *wv_b_src; + wk_b.extra = nullptr; + wv_b.extra = nullptr; + if (!ggml_backend_buffer_is_host(wk_b_src->buffer)) { + auto nbytes = ggml_nbytes(wk_b_src); + if (wk_buffer.size() < nbytes) wk_buffer.resize(nbytes); + ggml_backend_tensor_get(wk_b_src, wk_buffer.data(), 0, nbytes); + wk_b.data = wk_buffer.data(); + } + if (!ggml_backend_buffer_is_host(wv_b_src->buffer)) { + auto nbytes = ggml_nbytes(wv_b_src); + if (wv_buffer.size() < nbytes) wv_buffer.resize(nbytes); + ggml_backend_tensor_get(wv_b_src, wv_buffer.data(), 0, nbytes); + wv_b.data = wv_buffer.data(); + } + + auto n_wk = ggml_nelements(&wk_b); + auto n_wv = ggml_nelements(&wv_b); + + size_t tot_size = 0; + if (wk_b.type != GGML_TYPE_F32) { + tot_size += n_wk*sizeof(float); + } + tot_size += n_wk*sizeof(float); // ggml_cont(ctx, ggml_transpose(ctx, wk_b_used)); + if (wv_b.type != GGML_TYPE_F32) { + tot_size += n_wv*sizeof(float); + } + tot_size += (n_wk + n_wv)*sizeof(float); // ggml_concat(ctx, wk_b_transposed, wv_b_used, 0); + tot_size += (n_wk + n_wv)*sizeof(float); // ggml_cast(ctx, wkv_b_f32, new_type); + + if (tmp_buffer.size() < tot_size) tmp_buffer.resize(tot_size); + + auto ptr = tmp_buffer.data(); + + auto wk_b_used = &wk_b; + if (wk_b.type != GGML_TYPE_F32) { + wk_b_used = ggml_cast(ctx, &wk_b, GGML_TYPE_F32); + wk_b_used->data = ptr; + ptr += ggml_nbytes(wk_b_used); + } + auto wk_b_transposed = ggml_cont(ctx, ggml_transpose(ctx, wk_b_used)); + wk_b_transposed->data = ptr; + ptr += ggml_nbytes(wk_b_transposed); + + auto wv_b_used = &wv_b; + if (wv_b.type != GGML_TYPE_F32) { + wv_b_used = ggml_cast(ctx, &wv_b, GGML_TYPE_F32); + wv_b_used->data = ptr; + ptr += ggml_nbytes(wv_b_used); + } + + auto wkv_b_f32_3d = ggml_concat(ctx, wk_b_transposed, wv_b_used, 1); + wkv_b_f32_3d->data = ptr; + ptr += ggml_nbytes(wkv_b_f32_3d); + + auto wkv_b_f32 = ggml_view_2d(ctx, wkv_b_f32_3d, wkv_b_f32_3d->ne[0], wkv_b_f32_3d->ne[1]*wkv_b_f32_3d->ne[2], + wkv_b_f32_3d->nb[1], 0); + + auto new_type = wk_b.type == GGML_TYPE_BF16 && wv_b.type == GGML_TYPE_BF16 ? GGML_TYPE_BF16 + : wk_b.type == GGML_TYPE_F16 && wv_b.type == GGML_TYPE_F16 ? GGML_TYPE_F16 + : GGML_TYPE_Q8_0; + + auto wkv_b = ggml_cast(ctx, wkv_b_f32, new_type); + wkv_b->data = ptr; + ptr += ggml_nbytes(wkv_b); + + ggml_build_forward_expand(graph, wkv_b); + + auto plan = ggml_graph_plan(graph, std::thread::hardware_concurrency()/2); + if (plan.work_size > work_data.size()) work_data.resize(plan.work_size); + plan.work_data = work_data.data(); + + auto status = ggml_graph_compute(graph, &plan); + if (status != GGML_STATUS_SUCCESS) return nullptr; + + return wkv_b; +} + static void llm_prepare_mla(llama_model & model, int mla) { if (!model.is_mla_model()) return; const auto& hparams = model.hparams; @@ -3082,79 +3174,9 @@ static void llm_prepare_mla(llama_model & model, int mla) { for (int il = 0; il < n_layer; ++il) { auto& l = model.layers[il]; if (l.wkv_b || !l.wk_b || !l.wv_b || (l.wo && l.wo->extra)) continue; - auto wk_b = *l.wk_b; - auto wv_b = *l.wv_b; - if (!ggml_backend_buffer_is_host(l.wk_b->buffer)) { - auto nbytes = ggml_nbytes(l.wk_b); - if (wk_buffer.size() < nbytes) wk_buffer.resize(nbytes); - ggml_backend_tensor_get(l.wk_b, wk_buffer.data(), 0, nbytes); - wk_b.data = wk_buffer.data(); - } - if (!ggml_backend_buffer_is_host(l.wv_b->buffer)) { - auto nbytes = ggml_nbytes(l.wv_b); - if (wv_buffer.size() < nbytes) wv_buffer.resize(nbytes); - ggml_backend_tensor_get(l.wv_b, wv_buffer.data(), 0, nbytes); - wv_b.data = wv_buffer.data(); - } - auto n_wk = ggml_nelements(&wk_b); - auto n_wv = ggml_nelements(&wv_b); - - size_t tot_size = 0; - if (wk_b.type != GGML_TYPE_F32) { - tot_size += n_wk*sizeof(float); - } - tot_size += n_wk*sizeof(float); // ggml_cont(ctx, ggml_transpose(ctx, wk_b_used)); - if (wv_b.type != GGML_TYPE_F32) { - tot_size += n_wv*sizeof(float); - } - tot_size += (n_wk + n_wv)*sizeof(float); // ggml_concat(ctx, wk_b_transposed, wv_b_used, 0); - tot_size += (n_wk + n_wv)*sizeof(float); // ggml_cast(ctx, wkv_b_f32, new_type); - - if (tmp_buffer.size() < tot_size) tmp_buffer.resize(tot_size); - - auto ptr = tmp_buffer.data(); - - auto wk_b_used = &wk_b; - if (wk_b.type != GGML_TYPE_F32) { - wk_b_used = ggml_cast(ctx, &wk_b, GGML_TYPE_F32); - wk_b_used->data = ptr; - ptr += ggml_nbytes(wk_b_used); - } - auto wk_b_transposed = ggml_cont(ctx, ggml_transpose(ctx, wk_b_used)); - wk_b_transposed->data = ptr; - ptr += ggml_nbytes(wk_b_transposed); - - auto wv_b_used = &wv_b; - if (wv_b.type != GGML_TYPE_F32) { - wv_b_used = ggml_cast(ctx, &wv_b, GGML_TYPE_F32); - wv_b_used->data = ptr; - ptr += ggml_nbytes(wv_b_used); - } - - auto wkv_b_f32_3d = ggml_concat(ctx, wk_b_transposed, wv_b_used, 1); - wkv_b_f32_3d->data = ptr; - ptr += ggml_nbytes(wkv_b_f32_3d); - - auto wkv_b_f32 = ggml_view_2d(ctx, wkv_b_f32_3d, wkv_b_f32_3d->ne[0], wkv_b_f32_3d->ne[1]*wkv_b_f32_3d->ne[2], - wkv_b_f32_3d->nb[1], 0); - - auto new_type = wk_b.type == GGML_TYPE_BF16 && wv_b.type == GGML_TYPE_BF16 ? GGML_TYPE_BF16 - : wk_b.type == GGML_TYPE_F16 && wv_b.type == GGML_TYPE_F16 ? GGML_TYPE_F16 - : GGML_TYPE_Q8_0; - - auto wkv_b = ggml_cast(ctx, wkv_b_f32, new_type); - wkv_b->data = ptr; - ptr += ggml_nbytes(wkv_b); - - ggml_build_forward_expand(graph, wkv_b); - - auto plan = ggml_graph_plan(graph, std::thread::hardware_concurrency()/2); - if (plan.work_size > work_data.size()) work_data.resize(plan.work_size); - plan.work_data = work_data.data(); - - auto status = ggml_graph_compute(graph, &plan); - if (status != GGML_STATUS_SUCCESS) throw std::runtime_error("Failed to compute wkv_b"); + auto wkv_b = llm_compute_wkv_b(ctx, graph, l.wk_b, l.wv_b, wk_buffer, wv_buffer, tmp_buffer, work_data); + if (!wkv_b) throw std::runtime_error("Failed to compute wkv_b"); auto name = std::string{"blk."} + std::to_string(il) + ".attn_kv_b.weight"; @@ -3182,6 +3204,66 @@ static void llm_prepare_mla(llama_model & model, int mla) { ggml_free(ctx); } +// Recompute the load-time-derived combined wkv_b (computed_wkv_b) of one layer +// from the current wk_b/wv_b tensor data. Used by the hot-swap reload: the +// MLA>1 prompt-processing path consumes the derived wkv_b, so after attn_k_b / +// attn_v_b are reloaded from disk the derived tensor must be refreshed as well, +// or evaluations would keep using weights derived from the previously loaded +// data. Returns true when the layer has no derived wkv_b (nothing to do) or +// when the refresh succeeded. +bool llm_refresh_computed_wkv_b(llama_model & model, int il) { + if (il < 0 || il >= (int) model.layers.size()) return false; + auto & l = model.layers[il]; + if (!l.computed_wkv_b || l.wkv_b != l.computed_wkv_b.get()) return true; // wkv_b is not derived + if (!l.wk_b || !l.wv_b) return false; + if (l.wk_b->extra || l.wv_b->extra) return false; // per-device split tensors are not supported here + + ggml_init_params params{ggml_tensor_overhead()*32 + ggml_graph_overhead_custom(8, false), nullptr, true}; + auto ctx = ggml_init(params); + if (!ctx) return false; + auto graph = ggml_new_graph_custom(ctx, 8, false); + + std::vector wk_buffer, wv_buffer, tmp_buffer; + std::vector work_data; + auto wkv_b = llm_compute_wkv_b(ctx, graph, l.wk_b, l.wv_b, wk_buffer, wv_buffer, tmp_buffer, work_data); + if (!wkv_b) { + ggml_free(ctx); + return false; + } + + auto dst = l.computed_wkv_b.get(); + if (dst->type == wkv_b->type && ggml_nbytes(dst) == ggml_nbytes(wkv_b)) { + ggml_backend_tensor_set(dst, wkv_b->data, 0, ggml_nbytes(wkv_b)); + } else { + // The derived type/size changed (e.g. BF16 -> Q8_0 after a swap to a + // quantized wk_b/wv_b, or the load-time in-place host repack changed + // the type). Reallocate the destination buffer; the tensor object + // itself must be preserved (it is registered in tensors_by_name and + // referenced as l.wkv_b by the compute graph builder). + auto buft = ggml_backend_buffer_get_type(dst->buffer); + auto new_buf = ggml_backend_buft_alloc_buffer(buft, ggml_nbytes(wkv_b)); + if (!new_buf) { + ggml_free(ctx); + return false; + } + ggml_backend_buffer_set_usage(new_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_free(dst->buffer); + dst->type = wkv_b->type; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + dst->ne[i] = wkv_b->ne[i]; + dst->nb[i] = wkv_b->nb[i]; + } + dst->buffer = new_buf; + dst->data = ggml_backend_buffer_get_base(new_buf); + ggml_backend_tensor_set(dst, wkv_b->data, 0, ggml_nbytes(wkv_b)); + } + if (ggml_backend_buffer_is_host(dst->buffer)) { + iqk_modify_tensor(dst); + } + ggml_free(ctx); + return true; +} + static void llm_prepare_openpangu_param_sinks(llama_model & model) { if (model.arch != LLM_ARCH_OPENPANGU) return;