public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-10-17 20:03:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

talk-llama : fix build (#0)

Browse Source
This commit is contained in:
Georgi Gerganov
2026-08-14 22:16:06 +03:00
parent 733f281987
commit 1fe009caed
29 changed files with 2010 additions and 414 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/talk-llama/models/clip.cpp
+18
View File
@@ -0,0 +1,18 @@
#include "models.h"
// Stub to allow llama-quantize to open mmproj GGUFs
[[noreturn]]
void llama_model_clip::load_arch_hparams(llama_model_loader &) {
GGML_ABORT("CLIP is a quant-only stub; load_arch_hparams should not be called");
}
[[noreturn]]
void llama_model_clip::load_arch_tensors(llama_model_loader &) {
GGML_ABORT("CLIP is a quant-only stub; load_arch_tensors should not be called");
}
[[noreturn]]
std::unique_ptr<llm_graph_context> llama_model_clip::build_arch_graph(const llm_graph_params &) const {
GGML_ABORT("CLIP has no inference graph via llama_model dispatch; runtime lives in tools/mtmd/clip.cpp");
}
examples/talk-llama/models/dflash.cpp
+24 -15
View File
@@ -14,11 +14,14 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd;
LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__);
for (size_t i = 0; i < target_layer_ids.size(); ++i) {
LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : "");
std::string layers;
const char * sep = "";
for (const auto id : target_layer_ids) {
layers += sep;
layers += std::to_string(id);
sep = ", ";
}
LLAMA_LOG_INFO("]\n");
LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());
// DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
@@ -66,7 +69,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
// DFlash has a single rope, so the SWA rope == main rope.
if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
}
@@ -79,6 +82,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
//
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
@@ -97,6 +101,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
}
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
@@ -205,7 +210,7 @@ template <>
llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
ggml_tensor * cur = build_inp_embd_enc();
cur = build_lora_mm(model.fc, cur);
cur = build_lora_mm(model.fc, cur, model.fc_s);
cb(cur, "fc_out", -1);
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
@@ -460,9 +465,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
layer.ffn_up, NULL, NULL,
layer.ffn_gate, NULL, NULL,
layer.ffn_down, NULL, NULL,
layer.ffn_up, NULL, layer.ffn_up_s,
layer.ffn_gate, NULL, layer.ffn_gate_s,
layer.ffn_down, NULL, layer.ffn_down_s,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
@@ -479,15 +484,17 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
res->t_embd = cur;
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
@@ -655,15 +662,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
cb(cur, "result_norm", -1);
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
examples/talk-llama/models/exaone4.cpp
+3 -3
View File
@@ -1,6 +1,9 @@
#include "models.h"
void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
if (hparams.n_layer() == 64) { // 32B
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.n_swa = 4096;
@@ -15,9 +18,6 @@ void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
switch (hparams.n_layer()) {
case 30: type = LLM_TYPE_1_2B; break;
examples/talk-llama/models/granite-switch.cpp
+426
View File
@@ -0,0 +1,426 @@
#include "models.h"
#include <cmath>
void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false);
ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false);
bool rope_finetuned = true;
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
hparams.rope_finetuned = rope_finetuned;
switch (hparams.n_layer()) {
case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
case 64: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
ml.get_key(LLM_KV_ADAPTER_COUNT, n_adapters);
ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank);
ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false);
// bound counts that size tensors
if (n_adapters > 4096) {
throw std::runtime_error(format("graniteswitch: invalid adapter count %u", n_adapters));
}
if (max_lora_rank > 4096) {
throw std::runtime_error(format("graniteswitch: invalid lora rank %u", max_lora_rank));
}
std::vector<llama_token> token_ids;
std::vector<llama_token> substitute_ids;
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, token_ids);
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids);
if (token_ids.size() != n_adapters || substitute_ids.size() != n_adapters) {
throw std::runtime_error(format(
"graniteswitch: adapter token id arrays (%zu activate, %zu substitute) do not match adapter count %u",
token_ids.size(), substitute_ids.size(), n_adapters));
}
adapter_token_to_slot.clear();
adapter_token_to_substitute.clear();
for (uint32_t i = 0; i < n_adapters; ++i) {
// adapter i -> stacked slot i+1 (slot 0 is the base/zero delta)
adapter_token_to_slot[token_ids[i]] = (int32_t) (i + 1);
adapter_token_to_substitute[token_ids[i]] = substitute_ids[i];
}
// extra single-head attention layer at the END (index n_real) holds the router
// K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers
// keep their indices and the KV cache shift/defrag skips the router layer.
// n_layer_nextn is repurposed here (no MTP): it leaks as 1 into the
// llama_model_n_layer_nextn() getter and a re-saved nextn_predict_layers
const uint32_t n_real = hparams.n_layer();
if (n_real >= LLAMA_MAX_LAYERS) {
throw std::runtime_error(format("graniteswitch: block count %u exceeds LLAMA_MAX_LAYERS", n_real));
}
hparams.router_layer = (int32_t) n_real;
hparams.n_layer_all = n_real + 1;
hparams.n_layer_nextn = 1;
hparams.n_head_arr[n_real] = 1;
hparams.n_head_kv_arr[n_real] = 1;
hparams.n_ff_arr[n_real] = 0;
}
void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
const int64_t n_slots = (int64_t) n_adapters + 1; // slot 0 = base/zero delta
const int64_t n_rank = (int64_t) max_lora_rank;
const int64_t n_embd_q = n_embd_head_k * n_head;
const int64_t n_embd_kv = n_embd_k_gqa;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// substitute ids index tok_embd rows directly; range-check against n_vocab
for (const auto & kv : adapter_token_to_substitute) {
const llama_token sub = kv.second;
if (sub < 0 || (int64_t) sub >= n_vocab) {
throw std::runtime_error(format(
"graniteswitch: substitute token id %d out of range [0, %d)", sub, (int) n_vocab));
}
}
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
auto & sl = layer.switch_lora;
sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0);
sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0);
sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots}, 0);
sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
}
}
class llm_graph_input_switch : public llm_graph_input_i {
public:
llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {}
virtual ~llm_graph_input_switch() = default;
void set_input(const llama_ubatch * ubatch) override;
ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] adapter-substituted token ids
ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (+/-gain)
ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0)
ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0)
const llama_model_granite_switch & smodel;
};
// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then
// lets a single visible adapter token dominate so the readback recovers its slot.
void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) {
if (!ubatch->token) {
return;
}
const int64_t n_tokens = ubatch->n_tokens;
std::vector<int32_t> sub (n_tokens);
std::vector<float> ksig(n_tokens);
std::vector<float> vval(n_tokens);
std::vector<float> q (n_tokens, 1.0f);
for (int64_t i = 0; i < n_tokens; ++i) {
const llama_token tok = ubatch->token[i];
const auto it = smodel.adapter_token_to_slot.find(tok);
if (it != smodel.adapter_token_to_slot.end()) {
ksig[i] = +smodel.router_gain;
vval[i] = (float) it->second;
} else {
ksig[i] = -smodel.router_gain;
vval[i] = 0.0f;
}
const auto sit = smodel.adapter_token_to_substitute.find(tok);
sub[i] = (sit != smodel.adapter_token_to_substitute.end())
? (int32_t) sit->second
: (int32_t) tok;
}
ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens));
ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig));
ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval));
ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q));
}
std::unique_ptr<llm_graph_context> llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids.
// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
const int64_t n_in = cur->ne[0];
const int64_t n_tokens = cur->ne[1];
ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens);
ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens);
ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens}
ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens}
return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens);
}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
ggml_tensor * base = ggml_mul_mat(ctx0, w, cur);
ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids);
return ggml_add(ctx0, base, delta);
}
llama_model_granite_switch::graph::graph(
const llama_model & model,
const llm_graph_params & params)
: llm_graph_context(params) {
const auto & smodel = static_cast<const llama_model_granite_switch &>(model);
// TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed
GGML_ASSERT(ubatch.token && "granite-switch requires token input");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
auto inp_switch = std::make_unique<llm_graph_input_switch>(smodel);
inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_q = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
ggml_set_input(inp_switch->sub_tokens);
ggml_set_input(inp_switch->router_ksig);
ggml_set_input(inp_switch->router_vval);
ggml_set_input(inp_switch->router_q);
ggml_tensor * sub_tokens = inp_switch->sub_tokens;
ggml_tensor * router_ksig = inp_switch->router_ksig;
ggml_tensor * router_vval = inp_switch->router_vval;
ggml_tensor * router_q = inp_switch->router_q;
res->add_input(std::move(inp_switch));
// embed the substituted ids directly; build_inp_embd would embed the raw tokens
ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens);
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
cb(inpL, "inp_embd", -1);
ggml_tensor * inp_pos = nullptr;
if (hparams.rope_finetuned) {
inp_pos = build_inp_pos();
}
auto * inp_attn = build_attn_inp_kv();
// single causal head at layer R recovers the adapter index in-graph: only dim 0
// carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded.
const int R = hparams.router_layer;
GGML_ASSERT(R >= 0);
auto router_lane = [&](ggml_tensor * sig1d) {
ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens);
return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0);
};
ggml_tensor * Qr = router_lane(router_q);
ggml_tensor * Kr = router_lane(router_ksig);
ggml_tensor * Vr = router_lane(router_vval);
ggml_tensor * router_out = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R);
cb(router_out, "router_out", R);
// row 0 of router_out is the attended slot; clamp+round to an I32 index
ggml_tensor * slot_f = ggml_cont(ctx0,
ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0));
slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens);
slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters);
slot_f = ggml_round(ctx0, slot_f);
ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32);
cb(adapter_ids, "adapter_ids", -1);
ggml_tensor * inp_out_ids = build_inp_out_ids();
ggml_tensor * cur;
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// keep adapter_ids aligned to the kept rows (2D round-trip for get_rows)
const int64_t n_out = inp_out_ids->ne[0];
adapter_ids = ggml_get_rows(ctx0,
ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids);
adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out);
}
cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il);
inpL = cur;
}
cur = inpL;
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
const int64_t n_head = hparams.n_head(il);
const int64_t n_head_kv = hparams.n_head_kv(il);
ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur);
cb(qkv, "wqkv", il);
const int64_t n_embd_q = n_embd_head * n_head;
const int64_t n_embd_kv = n_embd_head * n_head_kv;
// slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added
ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0));
ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv)));
ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv)));
Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids));
Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids));
Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids));
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
if (hparams.rope_finetuned) {
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
}
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
// wo = nullptr so build_attn returns concatenated heads; o-proj is switched below
ggml_tensor * attn = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(attn, "attn_pre_o", il);
cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids);
cb(cur, "attn_out", il);
return cur;
}
ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids);
ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids);
g = ggml_silu(ctx0, g);
ggml_tensor * gu = ggml_mul(ctx0, g, u);
cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids);
cb(cur, "ffn_out", il);
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
return cur;
}
examples/talk-llama/models/mamba-base.cpp
+32 -16
View File
@@ -2,6 +2,8 @@
#include "llama-memory-recurrent.h"
#include <algorithm>
llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}
ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
@@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
@@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
int il) const {
const auto * mctx_cur = inp->mctx;
const auto kv_head = mctx_cur->get_head();
const auto kv_head = mctx_cur->get_head();
const auto mem_size = mctx_cur->get_size();
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = hparams.ssm_d_inner;
@@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
@@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
const int64_t state_slots = ssm_states_all->ne[1];
ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
@@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
// copy last (d_conv - 1) columns back into the state cache
ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0]));
const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
const size_t row_size = ggml_row_size(conv_states_all->type, row_count);
const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
ggml_view_1d(ctx0, conv_states_all,
(d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
ggml_element_size(conv_states_all))));
for (int64_t slot = 0; slot < n_written; ++slot) {
ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
conv_states_all->nb[1],
((size_t) slot * mem_size + kv_head) * row_size)));
}
// 1D convolution
// The equivalent is to make a self-overlapping view of conv_x
@@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// (this is necessary in order to properly use the states before they are overwritten,
// while avoiding to make unnecessary copies of the states)
auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);
// TODO: use semistructured matrices to implement state-space duality
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
// K > 1 asks the backend to return rollback snapshots in addition to the final state.
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
const int64_t D = d_state * d_inner;
const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
const size_t row_size = ggml_row_size(ssm_states_all->type, D);
const size_t y_row_size = ggml_row_size(y_ssm->type, D);
const size_t state_offset = ggml_nelements(x) * ggml_element_size(x);
// store last states
ggml_build_forward_expand(
gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]),
ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,
kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));
gf, ggml_cpy(ctx0,
ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,
y_row_size, y_row_size * n_seqs, state_offset),
ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,
ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));
ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],
n_seq_tokens * n_head * x->nb[1], 0);
examples/talk-llama/models/models.h
+96
View File
@@ -386,6 +386,22 @@ struct llama_model_bloom : public llama_model_base {
};
// Quant-only stub for mmproj GGUFs
// none of these are ever called, they only exist to satisfy the llama_model_base interface
struct llama_model_clip : public llama_model_base {
llama_model_clip(const struct llama_model_params & params) : llama_model_base(params) {}
[[noreturn]]
void load_arch_hparams(llama_model_loader & ml) override;
[[noreturn]]
void load_arch_tensors(llama_model_loader & ml) override;
[[noreturn]]
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_mpt : public llama_model_base {
llama_model_mpt(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@@ -697,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base {
};
struct llama_model_pockettts : public llama_model_base {
llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@@ -1028,6 +1057,19 @@ struct llama_model_olmoe : public llama_model_base {
};
struct llama_model_muse_glimmer : public llama_model_base {
llama_model_muse_glimmer(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_openelm : public llama_model_base {
llama_model_openelm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@@ -1461,6 +1503,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h {
using graph = llama_model_nemotron_h::graph;
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@@ -1596,6 +1642,56 @@ struct llama_model_granite_moe : public llama_model_base {
};
struct llama_model_granite_switch : public llama_model_base {
llama_model_granite_switch(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
uint32_t n_adapters = 0;
uint32_t max_lora_rank = 0;
float router_gain = 15.0f;
std::unordered_map<llama_token, int32_t> adapter_token_to_slot;
std::unordered_map<llama_token, llama_token> adapter_token_to_substitute;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
private:
ggml_tensor * build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il);
ggml_tensor * build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_minicpm : public llama_model_base {
llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
examples/talk-llama/models/muse-glimmer.cpp
+208
View File
@@ -0,0 +1,208 @@
#include "models.h"
void llama_model_muse_glimmer::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
uint32_t swa_period = 4;
if (ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false)) {
hparams.set_swa_pattern(swa_period);
} else {
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
}
switch (hparams.n_layer()) {
case 52: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_muse_glimmer::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
// Pre/post-attention norms (Muse Glimmer's `weight + 1` applied at conversion time).
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
// Q/K/V/O projections.
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
// QK-norm. Weights are synthesized at conversion time to absorb `qk_scale_factor`.
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
// Attention output gate: sigmoid(gate) * attn_out before o_proj (same as afmoe).
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
// Pre/post-FFN norms (FFN_PRE_NORM is aliased to LLM_TENSOR_FFN_NORM).
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
// Dense FFN (unlike afmoe, no MoE branches).
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
llama_model_muse_glimmer::graph::graph(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
// Different to f_norm_rms_eps for post-attn / post-FFN norms
const float post_norm_eps = 1e-8f;
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1);
cb(inpL, "embd_norm", -1);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
for (int il = 0; il < n_layer; ++il) {
// expose per-layer residual for speculative drafts (see LLM_KV_TARGET_LAYERS).
res->t_layer_inp[il] = inpL;
const float freq_base_l = model.get_rope_freq_base (cparams, il);
const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
ggml_tensor * inpSA = inpL;
// RoPE runs on the SWA layers, NoPE on full ones.
const bool use_rope = hparams.is_swa(il);
// pre-attention norm (weight+1 folded at conversion time)
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
// self-attention: attention output gate around SDPA (afmoe.cpp:147-191)
{
ggml_tensor * attn_inp = cur; // save input for gate computation
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
// gate = wqkv_gate @ attn_inp (from pre-attn hidden state)
ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
cb(gate, "attn_gate_proj", il);
// QK-norm. attn_q_norm weight was synthesized at conversion to broadcast
// qk_scale_factor across head_dim; attn_k_norm is identity (ones).
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
cb(Qcur, "Qcur_normed", il);
cb(Kcur, "Kcur_normed", il);
if (use_rope) {
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Qcur, "Qcur_rope", il);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Kcur, "Kcur_rope", il);
}
// SDPA. wo is deferred; the gate goes between attn_out and o_proj.
cur = build_attn(inp_attn,
NULL, NULL, NULL,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "attn_out", il);
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "attn_gate_sig", il);
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
cb(cur, "attn_o_proj", il);
}
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].attn_post_norm);
cb(cur, "attn_post_norm", il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
// pre-FFN norm
cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
// SwiGLU dense FFN
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
model.layers[il].ffn_gate, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].ffn_post_norm);
cb(cur, "ffn_post_norm", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
inpL = cur;
}
cur = inpL;
// final norm
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
// lm_head, followed by output multiplier
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
// Final logit tanh softcap (from gemma3.cpp).
if (hparams.f_final_logit_softcapping) {
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
cur = ggml_tanh(ctx0, cur);
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
}
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
std::unique_ptr<llm_graph_context> llama_model_muse_glimmer::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
examples/talk-llama/models/nemotron-h-moe.cpp
+150
View File
@@ -1,6 +1,156 @@
#include "models.h"
std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
// MTP draft head for Nemotron-H MoE
llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);
GGML_ASSERT(layer.ffn_gate_inp);
// token embedding weights
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * tok_embd;
if (ubatch.token) {
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
} else {
tok_embd = inp->embd;
}
cb(tok_embd, "mtp_tok_embd", il);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * h_embd = inp->h;
res->add_input(std::move(inp));
ggml_tensor * inp_out_ids = build_inp_out_ids();
// attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
// it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
const bool emit_h_nextn = cparams.embeddings_nextn;
const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
cb(concat, "mtp_concat", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
// dense NoPE attention sub-layer (mtp.layers.0)
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
{
auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "mtp_attn_out", il);
}
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "mtp_attn_residual", il);
// gather the output rows here so the MoE FFN below only runs on the positions we keep
if (crop_before_ffn) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// MoE FFN sub-layer (mtp.layers.1)
ggml_tensor * ffn_residual = cur;
cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_post_norm", il);
{
ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);
cb(router_logits, "mtp_ffn_moe_logits", il);
ggml_tensor * moe_out =
build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
nullptr, // no gate
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_RELU_SQR, hparams.expert_weights_norm,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
il,
router_logits, nullptr,
layer.ffn_up_exps_s,
nullptr, // no gate
layer.ffn_down_exps_s);
cb(moe_out, "mtp_ffn_moe_out", il);
ggml_tensor * ffn_shexp = build_ffn(cur,
layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
NULL, NULL, NULL,
layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
NULL,
LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
// final head norm: the MTP head has its own LayerNorm
GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");
cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!crop_before_ffn && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// LM head
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
examples/talk-llama/models/nemotron-h.cpp
+83 -23
View File
@@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// NextN/MTP: optional draft head appended as extra trailing block(s)
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
// A layer is recurrent IFF the n_head_kv value is set to 0 and
// the n_ff value is set to 0
for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0);
// the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent)
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
@@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
// mamba2 Mixer SSM params
// NOTE: int64_t for tensor dimensions
const int64_t d_conv = hparams.ssm_d_conv;
@@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
auto & layer = layers[i];
// all blocks use the attn norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags);
if (hparams.is_recr(i)) {
// ssm layers
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags);
layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags);
// no "weight" suffix for these
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags);
// out_proj
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags);
} else if (hparams.n_ff(i) == 0) {
// attention layers (with optional bias)
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
} else {
if (n_expert != 0) {
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags);
// MoE branch
layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags);
// Shared expert branch
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags);
} else {
// mlp layers
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags);
layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED);
}
}
}
// NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE
// sub-layer into a single trailing block
for (int i = n_layer; i < n_layer_all; ++i) {
auto & layer = layers[i];
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
// NextN input-fusion tensors
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags);
// attention sub-layer
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED);
// MoE sub-layer
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags);
}
}
std::unique_ptr<llm_graph_context> llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const {
@@ -135,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
auto * inp = build_inp_mem_hybrid();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
struct ggml_tensor * inpSA = inpL;
// norm
@@ -153,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_ffn_layer(cur, model, il);
}
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -167,9 +212,24 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
}
cur = inpL;
if (extract_final_inp) {
res->t_layer_inp[n_layer] = cur;
if (inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
// seed for the MTP/NextN draft head
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;
examples/talk-llama/models/plamo2.cpp
+1 -1
View File
@@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
examples/talk-llama/models/pockettts.cpp
+146
View File
@@ -0,0 +1,146 @@
#include "models.h"
// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_109M; break;
case 24: type = LLM_TYPE_335M; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
// no output head, the logits are unused; reuse the embedding table so a sampler can still run
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
cur = build_norm(inpL,
model.layers[il].attn_norm,
model.layers[il].attn_norm_b,
LLM_NORM, il);
cb(cur, "attn_norm", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
// FF
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm,
model.layers[il].ffn_norm_b,
LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_GELU, LLM_FFN_SEQ, il);
cb(cur, "ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = build_norm(inpL,
model.output_norm,
model.output_norm_b,
LLM_NORM, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}