mirror of
https://github.com/ikawrakow/ik_llama.cpp.git
synced 2026-08-12 22:29:39 +04:00
hotswap: keep load-time-derived and transformed tensors coherent after reloads (follow-up to #2131) (#2163)
* hotswap: keep load-time-derived and transformed tensors coherent after reloads - Re-derive the MLA combined attn_kv_b (computed_wkv_b) in place when a layer's attn_k_b/attn_v_b are hot-swapped: the mla>1 prompt-processing path consumes the derived tensor, so swaps of the source tensors previously had no effect (KLD stayed exactly 0 in per-tensor benchmarks, e.g. GLM-5.2 attn_k_b/v_b). - Refuse (loudly) hot-swaps that cannot be correct: views into -mqkv/-muge merged tensors, khad-folded MLA weights, in-place-scaled ffn_gate_inp_s, BitNet fused scales, OpenPangu parameter-sink sources, and same-dtype swaps of mmap-backed tensors. A refused reload produces no 'reloaded tensor' line, so benchmark drivers quarantine the round instead of recording wrong data. - Propagate reloaded data to same-name duplicate instances (tied lm head copy of token_embd, per-layer rope_freqs/rope_factors copies, expert-bias dups), warning when a duplicate cannot be refreshed. - Warn that derived state stays stale where a refresh is not possible: pre-transposed wk_b_pp under -sm graph/attn, requantized MTP head (output_extra.weight), and k_b/v_b derived from a reloaded attn_kv_b. - Warn at registration time when -rtr is enabled (restores cannot reproduce the run-time-repacked state; F16 -> BF16_R16 is lossy). - server: only attempt the /health hot-swap reload when no slot is processing, and clear the KV cache + cached prompts after a successful reload (they were computed with the previous weights). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Deduplicate llm_compute_wkv_b --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thireus ☠
co-authored by
Claude Fable 5
Claude Fable 5
parent
3861e045fc
commit
e5357286c0
@@ -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: <reason>` 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.
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -606,6 +606,11 @@ struct llama_model {
|
||||
std::unique_ptr<reload_info> 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;
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
#include <vector>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <cstdio>
|
||||
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
@@ -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<char> 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<int> 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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+154
-72
@@ -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<char> & wk_buffer, std::vector<char> & wv_buffer,
|
||||
std::vector<char> & tmp_buffer, std::vector<uint8_t> & 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<char> wk_buffer, wv_buffer, tmp_buffer;
|
||||
std::vector<uint8_t> 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;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user