furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cleanup

Browse Source
This commit is contained in:
Kawrakow
2026-05-13 05:10:41 +00:00
parent 2050855445
commit 26591f2b57
3 changed files with 22 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-delta-net.cpp
+10 -35
View File
@@ -150,11 +150,10 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_qkvz(llama_context & lctx, ggml_context * ctx0,
ggml_tensor * wqkv, ggml_tensor * wqkv_gate,
ggml_tensor * input, int il, const llm_build_cb & cb, ggml_cgraph * gf, ggml_tensor * qkv_cpy) {
ggml_tensor * input, int il, const llm_build_cb & cb, ggml_cgraph * gf) {
const int64_t n_tok = input->ne[1];
auto qkv_mixed = qkv_cpy ? ggml_mul_mat_inplace(ctx0, wqkv, input, qkv_cpy)
: llm_build_context::llm_build_lora_mm(lctx, ctx0, wqkv, input);
auto qkv_mixed = llm_build_context::llm_build_lora_mm(lctx, ctx0, wqkv, input);
cb(qkv_mixed, "qkv_mixed", il);
ggml_tensor * z = llm_build_context::llm_build_lora_mm(lctx, ctx0, wqkv_gate, input);
cb(z, "z", il);
@@ -167,7 +166,7 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_qkvz(llama_context & lc
std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_qkvz(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_in,
int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads,
ggml_tensor * input, int il, const llm_build_cb & cb, ggml_tensor * qkv_cpy) {
ggml_tensor * input, int il, const llm_build_cb & cb) {
const int64_t n_tok = input->ne[1];
@@ -214,17 +213,17 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_qkvz(llama_context & lc
cb(value_flat, "value_flat", il);
ggml_tensor * qkv_mixed = ggml_concat(ctx0, query_flat, key_flat, 0);
qkv_mixed = qkv_cpy ? ggml_concat_inplace(ctx0, qkv_mixed, value_flat, qkv_cpy, 0) : ggml_concat(ctx0, qkv_mixed, value_flat, 0);
qkv_mixed = ggml_concat(ctx0, qkv_mixed, value_flat, 0);
cb(qkv_mixed, "qkv_mixed", il);
return { qkv_mixed, z };
}
std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_qkvz(llama_context & lctx, ggml_context * ctx0, ggml_tensor * wqkv, ggml_tensor * wqkv_gate, ggml_tensor * ssm_in,
int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, ggml_tensor * input, int il, const llm_build_cb & cb, ggml_cgraph * gf, ggml_tensor * qkv_cpy) {
int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, ggml_tensor * input, int il, const llm_build_cb & cb, ggml_cgraph * gf) {
GGML_ASSERT((wqkv && wqkv_gate) || ssm_in);
return wqkv && wqkv_gate ? build_qkvz(lctx, ctx0, wqkv, wqkv_gate, input, il, cb, gf, qkv_cpy)
: build_qkvz(lctx, ctx0, ssm_in, head_k_dim, num_k_heads, head_v_dim, num_v_heads, input, il, cb, qkv_cpy);
return wqkv && wqkv_gate ? build_qkvz(lctx, ctx0, wqkv, wqkv_gate, input, il, cb, gf)
: build_qkvz(lctx, ctx0, ssm_in, head_k_dim, num_k_heads, head_v_dim, num_v_heads, input, il, cb);
}
std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_beta_gate(llama_context & lctx, ggml_context * ctx0,
@@ -433,15 +432,6 @@ static ggml_tensor * get_input_tensor_sm_graph(ggml_context * ctx, ggml_tensor *
return cur;
}
//static void build_qkv_cpy(ggml_context * ctx0, ggml_cgraph * gf, ggml_tensor * qkv_mixed, ggml_tensor * per_step_qkv) {
// const int64_t conv_dim = qkv_mixed->ne[0];
// const int64_t n_tok_qkv = qkv_mixed->ne[1] * qkv_mixed->ne[2];
// ggml_tensor * qkv_flat = ggml_reshape_2d(ctx0, qkv_mixed, conv_dim, n_tok_qkv);
// ggml_tensor * qkv_dst = ggml_view_2d(ctx0, per_step_qkv, conv_dim, n_tok_qkv, conv_dim * sizeof(float), 0);
// auto qkv_cpy = ggml_cpy(ctx0, qkv_flat, qkv_dst);
// ggml_build_forward_expand(gf, qkv_cpy);
//}
ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_cgraph * gf,
ggml_tensor * delta_input, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids,
uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb) const {
@@ -492,21 +482,17 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
int il_cb = 1000*il + id;
int64_t num_k_heads_id, num_v_heads_id;
ggml_tensor *qkv_mixed, *z;
//auto qkv_cpy = save_per_step_states &&
// il < (int)kv_self.ckpt.per_step_qkv.size() &&
// id < (int)kv_self.ckpt.per_step_qkv[il].size() ?
// kv_self.ckpt.per_step_qkv[il][id] : nullptr;
if (split_wqkv && split_wqkv_gate) {
num_k_heads_id = split_wqkv->splits[id]->ne[1]/(head_k_dim*(2 + gqa_ratio));
num_v_heads_id = num_k_heads_id * gqa_ratio;
auto p = build_qkvz(lctx, ctx0, split_wqkv->splits[id], split_wqkv_gate->splits[id], cur, il_cb, cb, gf, nullptr); //qkv_cpy);
auto p = build_qkvz(lctx, ctx0, split_wqkv->splits[id], split_wqkv_gate->splits[id], cur, il_cb, cb, gf);
qkv_mixed = p.first;
z = p.second;
} else {
num_k_heads_id = split_smm_in->splits[id]->ne[1]/(2*head_k_dim*(1 + gqa_ratio));
num_v_heads_id = num_k_heads_id * gqa_ratio;
auto p = build_qkvz(lctx, ctx0, nullptr, nullptr, split_smm_in->splits[id],
head_k_dim, num_k_heads_id, head_v_dim, num_v_heads_id, cur, il, cb, gf, nullptr); //qkv_cpy);
head_k_dim, num_k_heads_id, head_v_dim, num_v_heads_id, cur, il, cb, gf);
qkv_mixed = p.first;
z = p.second;
}
@@ -535,11 +521,7 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
ggml_tensor * per_step_ckpt = nullptr;
if (save_per_step_states && il < (int)kv_self.ckpt.per_step_ssm.size()) {
per_step_ckpt = kv_self.ckpt.per_step_ssm[il][id];
//GGML_ASSERT(per_step_ckpt);
}
//if (save_per_step_states && il < (int)kv_self.ckpt.per_step_qkv.size() && id < (int)kv_self.ckpt.per_step_qkv[il].size()) {
// build_qkv_cpy(ctx0, gf, qkv_mixed, kv_self.ckpt.per_step_qkv[il][id]);
//}
auto per_step_conv = save_per_step_states && il < (int)kv_self.ckpt.per_step_conv.size() &&
id < (int)kv_self.ckpt.per_step_conv[il].size()
? kv_self.ckpt.per_step_conv[il][id] : nullptr;
@@ -591,10 +573,8 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
auto norm = model.layers[il].attn_norm->extra ? ((ggml_split_tensor_t *)model.layers[il].attn_norm->extra)->splits[idx] : model.layers[il].attn_norm;
auto cur = llm_build_context::llm_build_norm(ctx0, input, hparams, norm, nullptr, LLM_NORM_RMS, cb, il);
//auto qkv_cpy = save_per_step_states && il < (int)kv_self.ckpt.per_step_qkv.size() && !kv_self.ckpt.per_step_qkv[il].empty()
// ? kv_self.ckpt.per_step_qkv[il].front() : nullptr;
auto [qkv_mixed, z] = build_qkvz(lctx, ctx0, model.layers[il].wqkv, model.layers[il].wqkv_gate, model.layers[il].ssm_in,
head_k_dim, num_k_heads, head_v_dim, num_v_heads, cur, il, cb, gf, nullptr); //qkv_cpy);
head_k_dim, num_k_heads, head_v_dim, num_v_heads, cur, il, cb, gf);
auto [beta, gate] = build_beta_gate(lctx, ctx0, model.layers[il].ssm_beta_alpha, model.layers[il].ssm_beta, model.layers[il].ssm_alpha,
model.layers[il].ssm_dt, model.layers[il].ssm_a, num_k_heads, num_v_heads, n_seqs, cur, il, cb, gf);
@@ -605,10 +585,6 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
per_step_ckpt = kv_self.ckpt.per_step_ssm[il].front();
}
// Save qkv_mixed features for per-step conv state reconstruction
//if (save_per_step_states && il < (int)kv_self.ckpt.per_step_qkv.size() && !kv_self.ckpt.per_step_qkv[il].empty()) {
// build_qkv_cpy(ctx0, gf, qkv_mixed, kv_self.ckpt.per_step_qkv[il].front());
//}
auto per_step_conv = save_per_step_states && il < (int)kv_self.ckpt.per_step_conv.size() && !kv_self.ckpt.per_step_conv[il].empty()
? kv_self.ckpt.per_step_conv[il].front() : nullptr;
@@ -626,7 +602,6 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
output = ggml_add(ctx0, gated_output, input);
cb(output, "ssm_output", il);
return output;
//return build_gated_output(lctx, ctx0, model.layers[il].ssm_norm, model.layers[il].ssm_out, output, z, head_v_dim, num_v_heads, n_tok, il, cb);
}
src/llama-delta-net.h
+3 -3
View File
@@ -34,16 +34,16 @@ private:
static std::pair<ggml_tensor *, ggml_tensor *> build_qkvz(llama_context & lctx, ggml_context * ctx0,
ggml_tensor * wqkv, ggml_tensor * wqkv_gate, ggml_tensor * input, int il, const llm_build_cb & cb,
ggml_cgraph * gf, ggml_tensor * qkv_copy);
ggml_cgraph * gf);
static std::pair<ggml_tensor *, ggml_tensor *> build_qkvz(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_in,
int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, ggml_tensor * input, int il,
const llm_build_cb & cb, ggml_tensor * qkv_copy);
const llm_build_cb & cb);
static std::pair<ggml_tensor *, ggml_tensor *> build_qkvz(llama_context & lctx, ggml_context * ctx0,
ggml_tensor * wqkv, ggml_tensor * wqkv_gate, ggml_tensor * ssm_in,
int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, ggml_tensor * input,
int il, const llm_build_cb & cb, ggml_cgraph * gf, ggml_tensor * qkv_copy);
int il, const llm_build_cb & cb, ggml_cgraph * gf);
static std::pair<ggml_tensor *, ggml_tensor *> build_beta_gate(llama_context & lctx, ggml_context * ctx0,
ggml_tensor * ssm_beta_alpha, ggml_tensor * ssm_beta, ggml_tensor * ssm_alpha,
src/llama.cpp
+9 -49
View File
@@ -1671,10 +1671,8 @@ bool llama_kv_cache::per_step_alloc(const llama_model & model, int max_tokens) {
}
static void restore_recurrent_cache_tensors(int step, ggml_backend_sched_t sched,
size_t ssm_bytes, size_t conv_bytes, size_t conv_dim, int d_conv_m1,
ggml_tensor * s_l, ggml_tensor * per_step_ssm,
ggml_tensor * s_l_shadow, ggml_tensor * per_step_conv,
std::vector<float> & old_conv_buf, std::vector<float> & qkv_buf, std::vector<float> & conv_buf,
size_t ssm_bytes, size_t conv_bytes,
ggml_tensor * s_l, ggml_tensor * per_step_ssm, ggml_tensor * per_step_conv,
std::unordered_set<ggml_backend_t> & backends_to_sync) {
auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, s_l);
auto dst = *s_l;
@@ -1691,41 +1689,6 @@ static void restore_recurrent_cache_tensors(int step, ggml_backend_sched_t sched
src = dst;
src.data = (char *)per_step_conv->data + (size_t)step * conv_bytes;
ggml_backend_tensor_copy_async(dst_backend, dst_backend, &src, &dst);
//if (s_l_shadow != nullptr) {
// ggml_backend_tensor_get(s_l_shadow, old_conv_buf.data(), 0, conv_bytes);
//} else {
// memset(old_conv_buf.data(), 0, conv_bytes);
//}
//auto qkv_needed = size_t(step + 1) * conv_dim;
//if (per_step_qkv) {
// ggml_backend_tensor_get(per_step_qkv, qkv_buf.data(), 0, qkv_needed * sizeof(float));
//} else {
// memset(qkv_buf.data(), 0, qkv_needed * sizeof(float));
//}
//for (int32_t col = 0; col < d_conv_m1; col++) {
// int32_t src_token = step - (d_conv_m1 - 1) + col; // e.g., K-2, K-1, K for d_conv=4
// if (src_token >= 0) {
// for (int64_t d = 0; d < conv_dim; d++) {
// conv_buf[col + d * d_conv_m1] = qkv_buf[d + (int64_t)src_token * conv_dim];
// }
// } else {
// int32_t old_col = d_conv_m1 + src_token; // maps to 0, 1, ... for early steps
// if (old_col >= 0 && old_col < d_conv_m1) {
// for (int64_t d = 0; d < conv_dim; d++) {
// conv_buf[col + d * d_conv_m1] = old_conv_buf[old_col + d * d_conv_m1];
// }
// } else {
// for (int64_t d = 0; d < conv_dim; d++) {
// conv_buf[col + d * d_conv_m1] = 0.0f;
// }
// }
// }
//}
//ggml_backend_tensor_set(s_l, conv_buf.data(), 0, conv_bytes);
}
bool llama_kv_cache::per_step_restore(const llama_model & model, ggml_backend_sched_t sched, int step) {
@@ -1743,7 +1706,6 @@ bool llama_kv_cache::per_step_restore(const llama_model & model, ggml_backend_sc
const int64_t ssm_bytes = ssm_state_dim * sizeof(float);
const int64_t conv_bytes = conv_state_dim * sizeof(float);
const int32_t d_conv_m1 = d_conv - 1; // number of columns in conv state
std::vector<float> ssm_buf(ssm_state_dim);
std::vector<float> conv_buf(conv_state_dim); // reconstructed conv state
@@ -1759,9 +1721,9 @@ bool llama_kv_cache::per_step_restore(const llama_model & model, ggml_backend_sc
if (s_l[il] == nullptr || ckpt.per_step_ssm[il].empty()) continue;
if (!s_l[il]->extra) {
restore_recurrent_cache_tensors(step, sched, ssm_bytes, conv_bytes, conv_dim, d_conv_m1,
s_l[il], ckpt.per_step_ssm[il].front(), ckpt.s_l_shadow[il], ckpt.per_step_conv[il].front(),
old_conv_buf, qkv_buf, conv_buf, backends_to_sync);
restore_recurrent_cache_tensors(step, sched, ssm_bytes, conv_bytes,
s_l[il], ckpt.per_step_ssm[il].front(), ckpt.per_step_conv[il].front(),
backends_to_sync);
} else {
auto split_sl = (const ggml_split_tensor_t *)s_l[il]->extra;
for (int id = 0; id < split_sl->n_device; ++id) {
@@ -1773,16 +1735,14 @@ bool llama_kv_cache::per_step_restore(const llama_model & model, ggml_backend_sc
int nv = split->ne[0] / head_v_dim; // number of heads handled by this device
auto [this_conv_dim, this_ssm_dim] = model.hparams.n_embd_v_s_dims(nv);
auto this_conv_bytes = (d_conv - 1) * this_conv_dim * sizeof(float);
restore_recurrent_cache_tensors(step, sched, this_ssm_dim * sizeof(float), this_conv_bytes, this_conv_dim, d_conv_m1,
split_sl->splits[id], ckpt.per_step_ssm[il][id], ckpt.split_s_l_shadow[il][id], ckpt.per_step_conv[il][id],
old_conv_buf, qkv_buf, conv_buf, backends_to_sync);
restore_recurrent_cache_tensors(step, sched, this_ssm_dim * sizeof(float), this_conv_bytes,
split_sl->splits[id], ckpt.per_step_ssm[il][id], ckpt.per_step_conv[il][id],
backends_to_sync);
}
}
}
// Strictly speaking we shouldn't need to do this. We are doing a synchronous copy of the
// convolution state in the above loop, and that involves a ggml_backend_synchronize call.
// But just in case.
// TODO: do we need to synchronize here?
for (auto backend : backends_to_sync) {
ggml_backend_synchronize(backend);
}