llama-io.cpp
llama-kv-cache.cpp
llama-kv-cache-iswa.cpp
+ llama-kv-cache-dsa.cpp
llama-memory-recurrent.cpp
llama-memory-hybrid.cpp
llama-memory-hybrid-iswa.cpp
auto it = ctx_map.find(buft);
if (it == ctx_map.end()) {
ggml_init_params params = {
- /*.mem_size =*/ hparams.n_layer*ggml_tensor_overhead(),
+ /*.mem_size =*/ hparams.n_layer()*ggml_tensor_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
};
// make tensors
- tensors.reserve(hparams.n_layer);
+ tensors.reserve(hparams.n_layer());
tensors.push_back(nullptr); // there's never a tensor for layer 0
- for (size_t il = 1; il < hparams.n_layer; il++) {
+ for (size_t il = 1; il < hparams.n_layer(); il++) {
ggml_backend_buffer_type_t buft = model.select_buft(il);
ggml_context * ctx = ctx_for_buft(buft);
if (!ctx) {
layer_start = il_start;
layer_end = il_end;
- for (size_t il = 1; il < hparams.n_layer; il++) {
+ for (size_t il = 1; il < hparams.n_layer(); il++) {
assert(tensors[il] != nullptr);
const size_t off = n_embd * (il - 1); // buffer doesn't have data for layer 0, since it's never present
{ LLM_ARCH_GEMMA3, "gemma3" },
{ LLM_ARCH_GEMMA3N, "gemma3n" },
{ LLM_ARCH_GEMMA4, "gemma4" },
+ { LLM_ARCH_GEMMA4_ASSISTANT, "gemma4-assistant" },
{ LLM_ARCH_GEMMA_EMBEDDING, "gemma-embedding" },
{ LLM_ARCH_STARCODER2, "starcoder2" },
{ LLM_ARCH_MAMBA, "mamba" },
{ LLM_ARCH_DEEPSEEK, "deepseek" },
{ LLM_ARCH_DEEPSEEK2, "deepseek2" },
{ LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" },
+ { LLM_ARCH_DEEPSEEK32, "deepseek32" },
{ LLM_ARCH_CHATGLM, "chatglm" },
{ LLM_ARCH_GLM4, "glm4" },
{ LLM_ARCH_GLM4_MOE, "glm4moe" },
{ LLM_ARCH_MAINCODER, "maincoder" },
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_TALKIE, "talkie" },
+ { LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
{ LLM_KV_MOE_LATENT_SIZE, "%s.moe_latent_size" },
{ LLM_KV_NEXTN_PREDICT_LAYERS, "%s.nextn_predict_layers" },
{ LLM_KV_NUM_DEEPSTACK_LAYERS, "%s.n_deepstack_layers" },
+ { LLM_KV_DEEPSTACK_MAPPING, "%s.deepstack_mapping" },
+ { LLM_KV_HIDDEN_ACT, "%s.hidden_activation" },
{ LLM_KV_POOLING_TYPE, "%s.pooling_type" },
{ LLM_KV_LOGIT_SCALE, "%s.logit_scale" },
{ LLM_KV_DECODER_START_TOKEN_ID, "%s.decoder_start_token_id" },
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
{ LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" },
+ { LLM_KV_ATTENTION_RECURRENT_LAYERS, "%s.attention.recurrent_layers" },
{ LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" },
{ LLM_KV_ROPE_DIMENSION_COUNT_SWA, "%s.rope.dimension_count_swa" },
{ LLM_KV_TOKENIZER_HF_JSON, "tokenizer.huggingface.json" },
{ LLM_KV_TOKENIZER_RWKV, "tokenizer.rwkv.world" },
{ LLM_KV_TOKENIZER_CHAT_TEMPLATE, "tokenizer.chat_template" },
+ { LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE, "tokenizer.ggml.normalizer.lowercase" },
{ LLM_KV_TOKENIZER_FIM_PRE_ID, "tokenizer.ggml.fim_pre_token_id" },
{ LLM_KV_TOKENIZER_FIM_SUF_ID, "tokenizer.ggml.fim_suf_token_id" },
{ LLM_KV_TOKENIZER_FIM_MID_ID, "tokenizer.ggml.fim_mid_token_id" },
{ LLM_KV_TOKENIZER_FIM_PAD_ID, "tokenizer.ggml.fim_pad_token_id" },
{ LLM_KV_TOKENIZER_FIM_REP_ID, "tokenizer.ggml.fim_rep_token_id" },
{ LLM_KV_TOKENIZER_FIM_SEP_ID, "tokenizer.ggml.fim_sep_token_id" },
+ { LLM_KV_TOKENIZER_SUPPRESS_TOKENS, "tokenizer.ggml.suppress_tokens" },
{ LLM_KV_ADAPTER_TYPE, "adapter.type" },
{ LLM_KV_ADAPTER_LORA_ALPHA, "adapter.lora.alpha" },
{ LLM_TENSOR_FFN_NORM_EXPS, "blk.%d.ffn_norm_exps" },
{ LLM_TENSOR_ATTN_K_B, "blk.%d.attn_k_b" },
{ LLM_TENSOR_ATTN_V_B, "blk.%d.attn_v_b" },
+ { LLM_TENSOR_NEXTN_PROJ_PRE, "nextn.pre_projection" },
+ { LLM_TENSOR_NEXTN_PROJ_POST, "nextn.post_projection" },
{ LLM_TENSOR_NEXTN_EH_PROJ, "blk.%d.nextn.eh_proj" },
{ LLM_TENSOR_NEXTN_EMBED_TOKENS, "blk.%d.nextn.embed_tokens" },
{ LLM_TENSOR_NEXTN_ENORM, "blk.%d.nextn.enorm" },
{LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_NEXTN_PROJ_POST, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
// NextN/MTP tensors are stored per-block (blk.%d.nextn.*) even though only the
// last nextn_predict_layers blocks carry them. Classify as LAYER_REPEATING so
// the model loader doesn't fault on the block index.
case LLM_ARCH_OLMO2:
case LLM_ARCH_OLMOE:
case LLM_ARCH_DEEPSEEK2:
+ case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_BITNET:
case LLM_ARCH_T5:
LLM_ARCH_GEMMA3,
LLM_ARCH_GEMMA3N,
LLM_ARCH_GEMMA4,
+ LLM_ARCH_GEMMA4_ASSISTANT,
LLM_ARCH_GEMMA_EMBEDDING,
LLM_ARCH_STARCODER2,
LLM_ARCH_MAMBA,
LLM_ARCH_DEEPSEEK,
LLM_ARCH_DEEPSEEK2,
LLM_ARCH_DEEPSEEK2OCR,
+ LLM_ARCH_DEEPSEEK32,
LLM_ARCH_CHATGLM,
LLM_ARCH_GLM4,
LLM_ARCH_GLM4_MOE,
LLM_ARCH_MAINCODER,
LLM_ARCH_KIMI_LINEAR,
LLM_ARCH_TALKIE,
+ LLM_ARCH_MELLUM,
LLM_ARCH_UNKNOWN,
};
LLM_KV_MOE_LATENT_SIZE,
LLM_KV_NEXTN_PREDICT_LAYERS,
LLM_KV_NUM_DEEPSTACK_LAYERS,
+ LLM_KV_DEEPSTACK_MAPPING,
+ LLM_KV_HIDDEN_ACT,
LLM_KV_POOLING_TYPE,
LLM_KV_LOGIT_SCALE,
LLM_KV_DECODER_START_TOKEN_ID,
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
LLM_KV_ATTENTION_INDEXER_TOP_K,
LLM_KV_ATTENTION_SHARED_KV_LAYERS,
+ LLM_KV_ATTENTION_RECURRENT_LAYERS,
LLM_KV_ROPE_DIMENSION_COUNT,
LLM_KV_ROPE_DIMENSION_COUNT_SWA,
LLM_KV_TOKENIZER_HF_JSON,
LLM_KV_TOKENIZER_RWKV,
LLM_KV_TOKENIZER_CHAT_TEMPLATE,
+ LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE,
LLM_KV_TOKENIZER_FIM_PRE_ID,
LLM_KV_TOKENIZER_FIM_SUF_ID,
LLM_KV_TOKENIZER_FIM_MID_ID,
LLM_KV_TOKENIZER_FIM_PAD_ID,
LLM_KV_TOKENIZER_FIM_REP_ID,
LLM_KV_TOKENIZER_FIM_SEP_ID,
+ LLM_KV_TOKENIZER_SUPPRESS_TOKENS,
LLM_KV_ADAPTER_TYPE,
LLM_KV_ADAPTER_LORA_ALPHA,
LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K,
LLM_TENSOR_INDEXER_ATTN_Q_B,
+ LLM_TENSOR_NEXTN_PROJ_PRE,
+ LLM_TENSOR_NEXTN_PROJ_POST,
LLM_TENSOR_NEXTN_EH_PROJ,
LLM_TENSOR_NEXTN_EMBED_TOKENS,
LLM_TENSOR_NEXTN_ENORM,
cparams.n_rs_seq = 0;
}
- cparams.n_threads = params.n_threads;
- cparams.n_threads_batch = params.n_threads_batch;
- cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor;
- cparams.yarn_attn_factor = params.yarn_attn_factor >= 0.0f ? params.yarn_attn_factor : hparams.yarn_attn_factor;
- cparams.yarn_beta_fast = params.yarn_beta_fast >= 0.0f ? params.yarn_beta_fast : hparams.yarn_beta_fast;
- cparams.yarn_beta_slow = params.yarn_beta_slow >= 0.0f ? params.yarn_beta_slow : hparams.yarn_beta_slow;
- cparams.embeddings = params.embeddings;
- cparams.embeddings_pre_norm = false;
- cparams.embeddings_pre_norm_masked = false;
- cparams.offload_kqv = params.offload_kqv;
- cparams.no_perf = params.no_perf;
- cparams.pooling_type = params.pooling_type;
- cparams.warmup = false;
+ cparams.n_threads = params.n_threads;
+ cparams.n_threads_batch = params.n_threads_batch;
+ cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor;
+ cparams.yarn_attn_factor = params.yarn_attn_factor >= 0.0f ? params.yarn_attn_factor : hparams.yarn_attn_factor;
+ cparams.yarn_beta_fast = params.yarn_beta_fast >= 0.0f ? params.yarn_beta_fast : hparams.yarn_beta_fast;
+ cparams.yarn_beta_slow = params.yarn_beta_slow >= 0.0f ? params.yarn_beta_slow : hparams.yarn_beta_slow;
+ cparams.embeddings = params.embeddings;
+ cparams.embeddings_nextn = false;
+ cparams.embeddings_nextn_masked = false;
+ cparams.offload_kqv = params.offload_kqv;
+ cparams.no_perf = params.no_perf;
+ cparams.warmup = false;
+
+ cparams.ctx_type = params.ctx_type;
+ cparams.pooling_type = params.pooling_type;
cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx;
cparams.rope_freq_base = params.rope_freq_base == 0.0f ? hparams.rope_freq_base_train : params.rope_freq_base;
cparams.cb_eval = params.cb_eval;
cparams.cb_eval_user_data = params.cb_eval_user_data;
- cparams.ctx_type = params.ctx_type;
+ cparams.ctx_other = nullptr;
+
+ // TODO: more generic
+ if (model.arch == LLM_ARCH_GEMMA4_ASSISTANT) {
+ if (params.ctx_other == nullptr) {
+ // TODO: change from runtime_error to llama_exception to avoid printing error message
+ throw std::runtime_error("Gemma4Assistant requires ctx_other to be set (this is normal during memory fitting)");
+ }
+
+ cparams.ctx_other = params.ctx_other;
+ }
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
+ cparams.n_outputs_max = params.n_outputs_max == 0 ? cparams.n_batch : params.n_outputs_max;
+
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);
LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);
LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);
+ LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);
if (cparams.n_ctx_seq < hparams.n_ctx_train) {
LLAMA_LOG_WARN("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n",
// init the memory module
if (!hparams.vocab_only) {
llama_memory_params params_mem = {
- /*.type_k =*/ params.type_k,
- /*.type_v =*/ params.type_v,
- /*.swa_full =*/ params.swa_full,
- /*.ctx_type= */ cparams.ctx_type,
+ /*.type_k =*/ params.type_k,
+ /*.type_v =*/ params.type_v,
+ /*.swa_full =*/ params.swa_full,
+ /*.ctx_type =*/ cparams.ctx_type,
+ /*.mem_other =*/ llama_get_memory(cparams.ctx_other),
};
memory.reset(model.create_memory(params_mem, cparams));
// enabling pipeline parallelism in the scheduler increases memory usage, so it is only done when necessary
bool pipeline_parallel =
model.n_devices() > 1 &&
- model.n_gpu_layers() > model.hparams.n_layer &&
+ model.n_gpu_layers() > model.hparams.n_layer_all &&
model.split_mode() == LLAMA_SPLIT_MODE_LAYER &&
cparams.offload_kqv &&
!model.has_tensor_overrides();
// note: n_outputs must match n_tokens for embedding models with mean/rank pooling,
// because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies
// it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens,
- // the ggml_mul_mat assertion fails. this matches the pp reservation below (line ~553).
+ // the ggml_mul_mat assertion fails.
const uint32_t n_tokens_ch = 16*n_seqs;
auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
if (!gf) {
int n_splits_tg = -1;
int n_nodes_tg = -1;
+ const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max);
+
// reserve pp (prompt processing) graph first so that buffers are only allocated once
{
- auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(),
+ auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(),
model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr);
if (!gf) {
if (cparams.pipeline_parallel) {
LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
cparams.pipeline_parallel = false;
sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));
- gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
+ gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
}
if (!gf) {
throw std::runtime_error("failed to allocate compute pp buffers");
//
// auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get());
//
- auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get(), model.hparams.no_alloc);
+ auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc);
if (!gf) {
throw std::runtime_error("failed to allocate compute pp buffers");
}
const uint32_t n_seqs = cparams.n_seq_max;
const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch);
- auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
+ const uint32_t n_outputs_max = std::min(n_tokens, cparams.n_outputs_max);
+
+ auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_max, mctx.get());
if (!gf) {
LLAMA_LOG_ERROR("%s: failed to reserve graph after the memory update\n", __func__);
}
return it->second.data();
}
-float * llama_context::get_embeddings_pre_norm() {
+float * llama_context::get_embeddings_nextn() {
output_reorder();
- return embd_pre_norm.data;
+ return embd_nextn.data;
}
-float * llama_context::get_embeddings_pre_norm_ith(int32_t i) {
+float * llama_context::get_embeddings_nextn_ith(int32_t i) {
output_reorder();
try {
- if (embd_pre_norm.data == nullptr) {
- throw std::runtime_error("no pre-norm embeddings");
+ if (embd_nextn.data == nullptr) {
+ throw std::runtime_error("no nextn embeddings");
}
- const uint32_t n_embd = model.hparams.n_embd;
+ const uint32_t n_embd = model.hparams.n_embd_out();
- if (!cparams.embeddings_pre_norm_masked) {
- // unmasked: pre-norm rows are stored densely, indexed by raw token position.
- if (i < 0 || (size_t)(i + 1) * n_embd > embd_pre_norm.size) {
- throw std::runtime_error(format("out of range [0, %zu)", embd_pre_norm.size / n_embd));
+ if (!cparams.embeddings_nextn_masked) {
+ // unmasked: nextn rows are stored densely, indexed by raw token position.
+ if (i < 0 || (size_t)(i + 1) * n_embd > embd_nextn.size) {
+ throw std::runtime_error(format("out of range [0, %zu)", embd_nextn.size / n_embd));
}
- return embd_pre_norm.data + (size_t) i * n_embd;
+ return embd_nextn.data + (size_t) i * n_embd;
}
const int64_t j = output_resolve_row(i);
- return embd_pre_norm.data + j*n_embd;
+ return embd_nextn.data + j*n_embd;
} catch (const std::exception & err) {
- LLAMA_LOG_ERROR("%s: invalid pre-norm embeddings id %d, reason: %s\n", __func__, i, err.what());
+ LLAMA_LOG_ERROR("%s: invalid nextn embeddings id %d, reason: %s\n", __func__, i, err.what());
#ifndef NDEBUG
GGML_ABORT("fatal error");
#else
//sched_need_reserve = true;
}
-void llama_context::set_embeddings_pre_norm(bool value, bool masked) {
+void llama_context::set_embeddings_nextn(bool value, bool masked) {
LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked);
- cparams.embeddings_pre_norm = value;
- cparams.embeddings_pre_norm_masked = masked;
+ cparams.embeddings_nextn = value;
+ cparams.embeddings_nextn_masked = masked;
}
void llama_context::set_causal_attn(bool value) {
}
int llama_context::encode(const llama_batch & batch_inp) {
- // MTP hook batches carry both token (next-token id) and embd (h_pre_norm row),
+ // MTP hook batches carry both token (next-token id) and embd (h_nextn row),
// so accept either present rather than requiring exactly one.
GGML_ASSERT(batch_inp.token || batch_inp.embd);
}
}
- auto * t_logits = res->get_logits();
- auto * t_embd = res->get_embd_pooled() ? res->get_embd_pooled() : res->get_embd();
- auto * t_h_pre_norm = cparams.embeddings_pre_norm ? res->get_h_pre_norm() : nullptr;
+ auto * t_logits = res->get_logits();
+ auto * t_embd = res->get_embd_pooled() ? res->get_embd_pooled() : res->get_embd();
+ auto * t_h_nextn = cparams.embeddings_nextn ? res->get_h_nextn() : nullptr;
// extract logits
if (logits.data && t_logits) {
}
}
- // extract pre-norm embeddings (hidden state before the final output norm)
- if (embd_pre_norm.data && t_h_pre_norm && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) {
- ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_pre_norm);
+ // extract nextn embeddings (hidden state before the final output norm)
+ if (embd_nextn.data && t_h_nextn && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) {
+ ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn);
GGML_ASSERT(backend_h != nullptr);
- const uint32_t n_embd = hparams.n_embd;
- GGML_ASSERT(n_tokens*n_embd <= (int64_t) embd_pre_norm.size);
- ggml_backend_tensor_get_async(backend_h, t_h_pre_norm, embd_pre_norm.data, 0, n_tokens*n_embd*sizeof(float));
+ const uint32_t n_embd = hparams.n_embd_out();
+ GGML_ASSERT(n_tokens*n_embd <= (int64_t) embd_nextn.size);
+ ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn.data, 0, n_tokens*n_embd*sizeof(float));
}
// TODO: hacky solution
}
int llama_context::decode(const llama_batch & batch_inp) {
- // MTP hook batches carry both token (next-token id) and embd (h_pre_norm row),
+ // MTP hook batches carry both token (next-token id) and embd (h_nextn row),
// so accept either present rather than requiring exactly one.
GGML_ASSERT(batch_inp.token || batch_inp.embd);
// ggml_graph_dump_dot(gf, NULL, "llama.dot");
//}
- auto * t_logits = res->get_logits();
- auto * t_embd = cparams.embeddings ? res->get_embd() : nullptr;
- auto * t_h_pre_norm = cparams.embeddings_pre_norm ? res->get_h_pre_norm() : nullptr;
+ auto * t_logits = res->get_logits();
+ auto * t_embd = cparams.embeddings ? res->get_embd() : nullptr;
+ auto * t_h_nextn = cparams.embeddings_nextn ? res->get_h_nextn() : nullptr;
if (t_embd && res->get_embd_pooled()) {
t_embd = res->get_embd_pooled();
}
}
- // extract pre-norm embeddings (hidden state before the final output norm)
+ // extract nextn embeddings before
// only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored.
{
- const bool masked = cparams.embeddings_pre_norm_masked;
+ const bool masked = cparams.embeddings_nextn_masked;
const int64_t n_rows = masked ? n_outputs : (int64_t) ubatch.n_tokens;
const int64_t offset = masked ? n_outputs_prev : n_tokens_prev;
- if (embd_pre_norm.data && t_h_pre_norm && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) {
- ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_pre_norm);
+ if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) {
+ ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn);
GGML_ASSERT(backend_h != nullptr);
- const uint32_t n_embd = hparams.n_embd;
- float * embd_pre_norm_out = embd_pre_norm.data + offset*n_embd;
+ const uint32_t n_embd = hparams.n_embd_out();
+ float * embd_nextn_out = embd_nextn.data + offset*n_embd;
- GGML_ASSERT((offset + n_rows)*n_embd <= (int64_t) embd_pre_norm.size);
- ggml_backend_tensor_get_async(backend_h, t_h_pre_norm, embd_pre_norm_out, 0, n_rows*n_embd*sizeof(float));
+ GGML_ASSERT((offset + n_rows)*n_embd <= (int64_t) embd_nextn.size);
+ ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn_out, 0, n_rows*n_embd*sizeof(float));
}
}
const auto n_batch = cparams.n_batch;
const auto n_vocab = vocab.n_tokens();
- const auto n_embd = hparams.n_embd;
const auto n_embd_out = hparams.n_embd_out();
- bool has_logits = true;
- bool has_embd = cparams.embeddings;
- bool has_embd_pre_norm = cparams.embeddings_pre_norm;
+ bool has_logits = true;
+ bool has_embd = cparams.embeddings;
+ bool has_embd_nextn = cparams.embeddings_nextn;
// TODO: hacky enc-dec support
if (model.arch == LLM_ARCH_T5) {
size_t backend_float_count = 0;
size_t backend_token_count = 0;
- logits.size = has_logits ? n_vocab*n_outputs_max : 0;
- embd.size = has_embd ? n_embd_out*n_outputs_max : 0;
- embd_pre_norm.size = has_embd_pre_norm ? n_embd*n_outputs_max : 0;
+ logits.size = has_logits ? n_vocab*n_outputs_max : 0;
+ embd.size = has_embd ? n_embd_out*n_outputs_max : 0;
+ embd_nextn.size = has_embd_nextn ? n_embd_out*n_outputs_max : 0;
- if (has_embd_pre_norm && !cparams.embeddings_pre_norm_masked) {
- // unmasked: pre-norm row exists for every token in the batch, not just
+ if (has_embd_nextn && !cparams.embeddings_nextn_masked) {
+ // unmasked: nextn row exists for every token in the batch, not just
// those flagged via batch.logits[i] -> size by token count instead.
- embd_pre_norm.size = (size_t) n_embd * n_batch;
+ embd_nextn.size = (size_t) n_embd_out * n_batch;
}
// Allocate backend sampling output buffers if there are backend samplers configured.
const size_t prev_size = buf_output ? ggml_backend_buffer_get_size(buf_output.get()) : 0;
const size_t new_size =
- (logits.size + embd.size + embd_pre_norm.size + backend_float_count) * sizeof(float) +
+ (logits.size + embd.size + embd_nextn.size + backend_float_count) * sizeof(float) +
( backend_token_count) * sizeof(llama_token);
// alloc only when more than the current capacity is required
buf_output = nullptr;
logits.data = nullptr;
embd.data = nullptr;
- embd_pre_norm.data = nullptr;
+ embd_nextn.data = nullptr;
}
auto * buft = ggml_backend_cpu_buffer_type();
embd = has_embd ? buffer_view<float>{(float *) (base + offset), embd.size} : buffer_view<float>{nullptr, 0};
offset += embd.size * sizeof(float);
- embd_pre_norm = has_embd_pre_norm ? buffer_view<float>{(float *) (base + offset), embd_pre_norm.size} : buffer_view<float>{nullptr, 0};
- offset += embd_pre_norm.size * sizeof(float);
+ embd_nextn = has_embd_nextn ? buffer_view<float>{(float *) (base + offset), embd_nextn.size} : buffer_view<float>{nullptr, 0};
+ offset += embd_nextn.size * sizeof(float);
if (has_sampling) {
sampling.logits = {(float *) (base + offset), (size_t)(n_vocab*n_outputs_max)};
this->n_outputs = 0;
+ GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max);
+
return n_outputs_max;
}
}
}
- if (embd_pre_norm.size > 0) {
+ if (embd_nextn.size > 0) {
for (uint64_t k = 0; k < n_embd; k++) {
- std::swap(embd_pre_norm.data[i0*n_embd + k], embd_pre_norm.data[i1*n_embd + k]);
+ std::swap(embd_nextn.data[i0*n_embd + k], embd_nextn.data[i1*n_embd + k]);
}
}
if (n_tokens % n_seqs != 0) {
n_tokens = ((n_tokens + (n_seqs - 1)) / n_seqs) * n_seqs; // round to next multiple of n_seqs
- n_outputs = std::max(n_outputs, n_tokens);
-
LLAMA_LOG_DEBUG("%s: making n_tokens a multiple of n_seqs - n_tokens = %u, n_seqs = %u, n_outputs = %u\n", __func__, n_tokens, n_seqs, n_outputs);
}
// norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
// FIXME: fix in ggml_backend_sched
- const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer;
+ const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all;
if (ubatch.n_tokens < 32 || full_offload) {
if (il != -1 && strcmp(name, "norm") == 0) {
const auto & dev_layer = model.dev_layer(il);
/*.n_ubatch =*/ 512,
/*.n_seq_max =*/ 1,
/*.n_rs_seq =*/ 0,
+ /*.n_outputs_max =*/ 0,
/*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
/*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS,
/*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT,
/*.kv_unified =*/ false,
/*.sampler =*/ nullptr,
/*.n_sampler =*/ 0,
+ /*.ctx_other =*/ nullptr,
};
return result;
LLAMA_LOG_ERROR("%s: SPLIT_MODE_TENSOR requires flash_attn to be enabled\n", __func__);
return nullptr;
}
- if (ggml_is_quantized(params.type_k) || ggml_is_quantized(params.type_v)) {
- LLAMA_LOG_ERROR("%s: simultaneous use of SPLIT_MODE_TENSOR and KV cache quantization not implemented\n", __func__);
- return nullptr;
- }
}
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) {
const uint32_t blck_size = ggml_blck_size(params.type_k);
- for (uint32_t il = 0; il < model->hparams.n_layer; ++il) {
+ for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
if (model->hparams.n_embd_head_k(il) % blck_size != 0) {
LLAMA_LOG_ERROR("%s: K cache type %s with block size %u does not divide n_embd_head_k=%u\n",
__func__, ggml_type_name(params.type_k), blck_size, model->hparams.n_embd_head_k(il));
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_v)) {
const uint32_t blck_size = ggml_blck_size(params.type_v);
- for (uint32_t il = 0; il < model->hparams.n_layer; ++il) {
+ for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
if (model->hparams.n_embd_head_v(il) % blck_size != 0) {
LLAMA_LOG_ERROR("%s: V cache type %s with block size %u does not divide n_embd_head_v=%u\n",
__func__, ggml_type_name(params.type_v), blck_size, model->hparams.n_embd_head_v(il));
}
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
- model->hparams.nextn_predict_layers == 0) {
+ model->hparams.n_layer_nextn == 0) {
LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
return nullptr;
}
-
try {
auto * ctx = new llama_context(*model, params);
return ctx;
return ctx->get_embeddings_seq(seq_id);
}
-void llama_set_embeddings_pre_norm(llama_context * ctx, bool value, bool masked) {
- ctx->set_embeddings_pre_norm(value, masked);
+void llama_set_embeddings_nextn(llama_context * ctx, bool value, bool masked) {
+ ctx->set_embeddings_nextn(value, masked);
}
-float * llama_get_embeddings_pre_norm(llama_context * ctx) {
+llama_memory_t llama_get_memory(const struct llama_context * ctx) {
+ if (!ctx) {
+ return nullptr;
+ }
+
+ return ctx->get_memory();
+}
+
+float * llama_get_embeddings_nextn(llama_context * ctx) {
ctx->synchronize();
- return ctx->get_embeddings_pre_norm();
+ return ctx->get_embeddings_nextn();
}
-float * llama_get_embeddings_pre_norm_ith(llama_context * ctx, int32_t i) {
+float * llama_get_embeddings_nextn_ith(llama_context * ctx, int32_t i) {
ctx->synchronize();
- return ctx->get_embeddings_pre_norm_ith(i);
+ return ctx->get_embeddings_nextn_ith(i);
}
bool llama_set_sampler(llama_context * ctx, llama_seq_id seq_id, llama_sampler * smpl) {
uint32_t n_tokens,
uint32_t n_seqs,
uint32_t n_outputs) {
- auto * memory = ctx->get_memory();
+ auto memory = ctx->get_memory();
llama_memory_context_ptr mctx;
if (memory) {
mctx = memory->init_full();
// memory
//
-llama_memory_t llama_get_memory(const struct llama_context * ctx) {
- return ctx->get_memory();
-}
-
void llama_memory_clear(llama_memory_t mem, bool data) {
if (!mem) {
return;
llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx) {
return ctx->memory_breakdown();
}
+
+llama_context * llama_get_ctx_other(struct llama_context * ctx) {
+ return ctx->get_cparams().ctx_other;
+}
#include "llama-graph.h"
#include "llama-adapter.h"
#include "llama-impl.h"
+#include "llama-memory.h"
#include "ggml-cpp.h"
#include "ggml-opt.h"
float * get_embeddings_ith(int32_t i);
float * get_embeddings_seq(llama_seq_id seq_id);
- float * get_embeddings_pre_norm();
- float * get_embeddings_pre_norm_ith(int32_t i);
+ float * get_embeddings_nextn();
+ float * get_embeddings_nextn_ith(int32_t i);
llama_token * get_sampled_tokens() const;
llama_token get_sampled_token_ith(int32_t idx);
void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data);
void set_embeddings (bool value);
- void set_embeddings_pre_norm(bool value, bool masked);
+ void set_embeddings_nextn(bool value, bool masked);
void set_causal_attn(bool value);
void set_warmup(bool value);
llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably
- std::unique_ptr<llama_memory_i> memory;
+ llama_memory_ptr memory;
// decode output (2-dimensional array: [n_outputs][n_vocab])
buffer_view<float> logits = {nullptr, 0};
// populated only when pooling_type == LLAMA_POOLING_TYPE_NONE
buffer_view<float> embd = {nullptr, 0};
- // hidden state before the final output norm (2-dimensional array: [n_outputs][n_embd])
- // populated only when cparams.embeddings_pre_norm is enabled and the model graph
- // sets llm_graph_result::t_h_pre_norm
- buffer_view<float> embd_pre_norm = {nullptr, 0};
+ // hidden state required by the nextn layers (2-dimensional array: [n_outputs][n_embd])
+ // populated only when cparams.embeddings_nextn is enabled and the model graph
+ // sets llm_graph_result::t_h_nextn
+ buffer_view<float> embd_nextn = {nullptr, 0};
struct sampling_info {
// !samplers.empty() to check if any samplers are active
uint32_t n_ubatch;
uint32_t n_seq_max;
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback
+ uint32_t n_outputs_max; // max outputs supported by the context
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
float yarn_beta_slow;
bool embeddings;
- bool embeddings_pre_norm; // also extract the hidden state before the final output norm
- bool embeddings_pre_norm_masked; // extract for only rows where batch.logits != 0
+ bool embeddings_nextn; // also extract the hidden state before the final output norm
+ bool embeddings_nextn_masked; // extract for only rows where batch.logits != 0
bool causal_attn;
bool offload_kqv;
bool flash_attn;
bool fused_gdn_ch; // use fused gated delta net (chunked)
bool auto_fgdn;
bool no_perf;
- bool warmup;
+ bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
bool op_offload;
bool kv_unified;
bool pipeline_parallel;
ggml_backend_sched_eval_callback cb_eval;
void * cb_eval_user_data;
+
+ llama_context * ctx_other;
};
LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx);
-//
-// pre-norm embeddings (hidden state before the final output norm)
-//
-
-// Set whether the context outputs pre-norm embeddings or not
+// Set whether the context outputs nextn embeddings or not
// If masked == true, output the embeddings only for the tokens with batch.logits != 0
// If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits
-LLAMA_API void llama_set_embeddings_pre_norm(struct llama_context * ctx, bool value, bool masked);
+LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
// mirrors:
// LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
-LLAMA_API float * llama_get_embeddings_pre_norm (struct llama_context * ctx);
+LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx);
// LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i);
-LLAMA_API float * llama_get_embeddings_pre_norm_ith(struct llama_context * ctx, int32_t i);
+LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i);
+
+LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
+#include "llama-kv-cache-dsa.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-recurrent.h"
const auto n_tokens = ubatch.n_tokens;
const auto n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq;
- ggml_tensor * res = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_tokens/n_stream, 1, n_stream);
+ // flash attention requires an f16 mask
+ const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
+ ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream);
ggml_set_input(res);
ggml_set_name(res, "attn_inp_kq_mask");
return res;
}
+void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) {
+ const int64_t n_tokens = ubatch->n_tokens;
+
+ if (ubatch->token) {
+ ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));
+ } else {
+ // note: mtmd embedding input goes through here
+ GGML_ASSERT(ubatch->embd);
+ GGML_ASSERT(n_embd == embd->ne[0]);
+
+ ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h));
+ }
+
+ // TODO: extend llama_ubatch to differentiate between token embeddings and hidden states
+ // for now, we assume that the hidden state is always provided as an embedding
+ // ref: https://github.com/ggml-org/llama.cpp/pull/23643
+ if (ubatch->embd) {
+ GGML_ASSERT(n_embd == h->ne[0]);
+
+ ggml_backend_tensor_set(h, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h));
+ }
+}
+
+bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) {
+ bool res = true;
+
+ res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens);
+ res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens);
+ res &= (!params.ubatch.embd) || (h && h->ne[1] == params.ubatch.n_tokens);
+
+ return res;
+}
+
void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) {
if (ubatch->pos && pos) {
const int64_t n_tokens = ubatch->n_tokens;
}
}
-static void print_mask(const float * data, int64_t n_tokens, int64_t n_kv, int64_t n_swa, llama_swa_type swa_type) {
+template <typename T>
+static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n_swa, llama_swa_type swa_type) {
LLAMA_LOG_DEBUG("%s: === Attention mask ===\n", __func__);
const char * swa_type_str = "unknown";
case LLAMA_SWA_TYPE_SYMMETRIC: swa_type_str = "LLAMA_SWA_TYPE_SYMMETRIC"; break;
};
- LLAMA_LOG_DEBUG("%s: n_swa : %d, n_kv: %d, swq_type: %s\n", __func__, (int)n_swa, (int)n_kv, swa_type_str);
+ LLAMA_LOG_DEBUG("%s: n_swa : %d, n_kv: %d, swa_type: %s\n", __func__, (int)n_swa, (int)n_kv, swa_type_str);
LLAMA_LOG_DEBUG("%s: '0' = can attend, '∞' = masked\n", __func__);
LLAMA_LOG_DEBUG("%s: Rows = query tokens, Columns = key/value tokens\n\n", __func__);
for (int i = 0; i < std::min((int64_t)20, n_tokens); ++i) {
LLAMA_LOG_DEBUG(" %2d ", i);
for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) {
- float val = data[i * n_kv + j];
+ float val = llama_cast<float>(data[i * n_kv + j]);
if (val == -INFINITY) {
LLAMA_LOG_DEBUG(" ∞");
} else {
const int64_t n_kv = ubatch->n_tokens;
const int64_t n_tokens = ubatch->n_tokens;
- const auto fill_mask = [&](float * data, int n_swa, llama_swa_type swa_type) {
+ const auto fill_mask = [&](auto * data, int64_t ne, int n_swa, llama_swa_type swa_type) {
+ using T = std::remove_reference_t<decltype(*data)>;
+ std::fill(data, data + ne, llama_cast<T>(-INFINITY));
+
for (int i1 = 0; i1 < n_tokens; ++i1) {
const llama_seq_id s1 = ubatch->seq_id[i1][0];
const llama_pos p1 = ubatch->pos[i1];
continue;
}
- data[idst + i0] = hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f;
+ data[idst + i0] = llama_cast<T>(hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f);
}
}
- };
-
- {
- GGML_ASSERT(self_kq_mask);
- GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
-
- float * data = (float *) self_kq_mask->data;
-
- std::fill(data, data + ggml_nelements(self_kq_mask), -INFINITY);
-
- fill_mask(data, 0, LLAMA_SWA_TYPE_NONE);
if (debug) {
- print_mask(data, n_tokens, n_kv, 0, LLAMA_SWA_TYPE_NONE);
+ print_mask(data, n_tokens, n_kv, n_swa, swa_type);
}
+ };
+
+ GGML_ASSERT(self_kq_mask);
+ GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
+ if (self_kq_mask->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE);
+ } else {
+ fill_mask((float *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE);
}
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
GGML_ASSERT(self_kq_mask_swa);
GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));
-
- float * data = (float *) self_kq_mask_swa->data;
-
- std::fill(data, data + ggml_nelements(self_kq_mask_swa), -INFINITY);
-
- fill_mask(data, hparams.n_swa, hparams.swa_type);
-
- if (debug) {
- print_mask(data, n_tokens, n_kv, hparams.n_swa, hparams.swa_type);
+ if (self_kq_mask_swa->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type);
+ } else {
+ fill_mask((float *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type);
}
}
}
return res;
}
+void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
+ mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
+
+ mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
+
+ mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
+
+ mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn);
+
+ mctx->get_lid()->set_input_k_rot(self_k_rot_lid);
+}
+
+bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
+ const auto * mctx = static_cast<const llama_kv_cache_dsa_context *>(params.mctx);
+
+ this->mctx = mctx;
+
+ bool res = true;
+
+ res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens;
+ res &= self_k_idxs_lid->ne[0] == params.ubatch.n_tokens;
+
+ res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams);
+ res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams);
+
+ return res;
+}
+
void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
// base tensors may not be allocated if there are no non-SWA attention layers
if (self_k_idxs && self_k_idxs->buffer) {
mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);
mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch);
-
- mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
}
+ mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
+
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);
mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch);
-
- mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
}
+ mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
+
if (self_k_rot) {
mctx->get_base()->set_input_k_rot(self_k_rot);
}
if (self_k_idxs && self_k_idxs->buffer) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
-
- res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
}
+ res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
+
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
//res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
-
- res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
}
+ res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
+
return res;
}
GGML_ASSERT(ggml_backend_buffer_is_host(cross_kq_mask->buffer));
GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing
- float * data = (float *) cross_kq_mask->data;
-
- for (int i = 0; i < n_tokens; ++i) {
- GGML_ASSERT(!cross->seq_ids_enc.empty() && "llama_encode must be called first");
- for (int j = 0; j < n_enc; ++j) {
- float f = -INFINITY;
+ const auto fill_mask = [&](auto * data) {
+ using T = std::remove_reference_t<decltype(*data)>;
+ for (int i = 0; i < n_tokens; ++i) {
+ GGML_ASSERT(!cross->seq_ids_enc.empty() && "llama_encode must be called first");
+ for (int j = 0; j < n_enc; ++j) {
+ float f = -INFINITY;
- for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {
- const llama_seq_id seq_id = ubatch->seq_id[i][s];
+ for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {
+ const llama_seq_id seq_id = ubatch->seq_id[i][s];
- if (cross->seq_ids_enc[j].find(seq_id) != cross->seq_ids_enc[j].end()) {
- f = 0.0f;
+ if (cross->seq_ids_enc[j].find(seq_id) != cross->seq_ids_enc[j].end()) {
+ f = 0.0f;
+ }
}
- }
- data[i*n_enc + j] = f;
+ data[i*n_enc + j] = llama_cast<T>(f);
+ }
}
+ };
+
+ if (cross_kq_mask->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) cross_kq_mask->data);
+ } else {
+ fill_mask((float *) cross_kq_mask->data);
}
}
if (inp_attn->self_k_idxs && inp_attn->self_k_idxs->buffer) {
attn_ctx->get_base()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
attn_ctx->get_base()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
+ }
+ if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) {
attn_ctx->get_base()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
}
if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) {
attn_ctx->get_swa()->set_input_k_idxs(inp_attn->self_k_idxs_swa, ubatch);
attn_ctx->get_swa()->set_input_v_idxs(inp_attn->self_v_idxs_swa, ubatch);
+ }
+ if (inp_attn->self_kq_mask_swa && inp_attn->self_kq_mask_swa->buffer) {
attn_ctx->get_swa()->set_input_kq_mask(inp_attn->self_kq_mask_swa, ubatch, cparams.causal_attn);
}
if (inp_attn->self_k_idxs && inp_attn->self_k_idxs->buffer) {
res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;
//res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
-
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask, attn_ctx->get_base(), params.ubatch, params.cparams);
}
+ res &= can_reuse_kq_mask(inp_attn->self_kq_mask, attn_ctx->get_base(), params.ubatch, params.cparams);
+
// swa tensors may not be allocated if there are no SWA attention layers
if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) {
res &= inp_attn->self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
//res &= inp_attn->self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
-
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask_swa, attn_ctx->get_swa(), params.ubatch, params.cparams);
}
+ res &= can_reuse_kq_mask(inp_attn->self_kq_mask_swa, attn_ctx->get_swa(), params.ubatch, params.cparams);
+
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
if (t_embd_pooled != nullptr) {
ggml_set_output(t_embd_pooled);
}
- if (t_h_pre_norm != nullptr) {
- ggml_set_output(t_h_pre_norm);
+ if (t_h_nextn != nullptr) {
+ ggml_set_output(t_h_nextn);
}
for (auto & [seq_id, t] : t_sampled) {
if (t != nullptr) {
cparams (params.cparams),
ubatch (params.ubatch),
n_embd (hparams.n_embd),
- n_layer (hparams.n_layer),
+ n_layer (hparams.n_layer()),
+ n_layer_nextn (hparams.n_layer_nextn),
n_rot (hparams.n_rot()),
n_ctx (cparams.n_ctx),
n_head (hparams.n_head()),
res->t_inp_embd = cur;
// For Granite architecture
- if (hparams.f_embedding_scale != 0.0f) {
+ // NOTE: Only apply scale to token inputs. Raw embeddings are assumed to be
+ // multimodal inputs that should not be scaled.
+ if (ubatch.token && hparams.f_embedding_scale != 0.0f) {
+ if (!ggml_is_contiguous(cur)) {
+ cur = ggml_cont(ctx0, cur);
+ }
cur = ggml_scale(ctx0, cur, hparams.f_embedding_scale);
}
llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() const {
auto inp = std::make_unique<llm_graph_input_attn_no_cache>(hparams, cparams);
+ // flash attention requires an f16 mask
+ const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
// note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens, 1, 1);
+ inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask);
- inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
- inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens, 1, 1);
+ inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask_swa);
- inp->self_kq_mask_swa_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_swa, GGML_TYPE_F16) : inp->self_kq_mask_swa;
+ inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
} else {
inp->self_kq_mask_swa = nullptr;
inp->self_kq_mask_swa_cnv = nullptr;
inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams);
- inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
}
inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0);
inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams);
- inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
}
return inp;
return cur;
}
+ggml_tensor * llm_graph_context::build_attn(
+ llm_graph_input_attn_k_dsa * inp,
+ ggml_tensor * wo,
+ ggml_tensor * wo_b,
+ ggml_tensor * wo_s,
+ ggml_tensor * q_cur,
+ ggml_tensor * k_cur,
+ ggml_tensor * v_cur,
+ ggml_tensor * kq_b,
+ ggml_tensor * sinks,
+ ggml_tensor * v_mla,
+ ggml_tensor * top_k,
+ float kq_scale,
+ int il) const {
+ // these nodes are added to the graph together so that they are not reordered
+ // by doing so, the number of splits in the graph is reduced
+ // expand k later to enable rope fusion which directly writes into k-v cache
+ ggml_build_forward_expand(gf, q_cur);
+ ggml_build_forward_expand(gf, v_cur);
+ ggml_build_forward_expand(gf, k_cur);
+
+ const auto * mctx_cur = inp->mctx->get_mla();
+
+ // store to KV cache
+ {
+ const auto & k_idxs = inp->get_k_idxs_mla();
+
+ ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
+ }
+
+ const auto & kq_mask = inp->get_kq_mask_mla();
+
+ // prepare new kq mask - starts filled with -INFINITY
+ ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
+
+ // reshape KQ mask into tensor with rows of size 1:
+ // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
+ kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
+
+ // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
+ ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
+
+ // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
+ // this will be our source of zero values for unmasking top k mask elements
+ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
+ zeros = ggml_fill(ctx0, zeros, 0.0f);
+
+ // modify KQ mask by unmasking elements that are in top_k indices
+ // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
+ ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
+
+ // reshape to restore the original shape of KQ mask:
+ // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
+ kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
+
+ // combine with the original kq mask
+ kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+
+ ggml_tensor * q = q_cur;
+ ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+ ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
+
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il);
+ cb(cur, "kqv_out", il);
+
+ if (wo) {
+ cur = build_lora_mm(wo, cur, wo_s);
+ }
+
+ if (wo_b) {
+ cur = ggml_add(ctx0, cur, wo_b);
+ }
+
+ return cur;
+}
+
ggml_tensor * llm_graph_context::build_attn(
llm_graph_input_attn_kv_iswa * inp,
ggml_tensor * wo,
const int32_t n_enc = !cross->v_embd.empty() ? cross->n_enc : hparams.n_ctx_train;
- inp->cross_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_enc, n_tokens, 1, 1);
+ // flash attention requires an f16 mask
+ const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
+ inp->cross_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_enc, n_tokens, 1, 1);
ggml_set_input(inp->cross_kq_mask);
- inp->cross_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->cross_kq_mask, GGML_TYPE_F16) : inp->cross_kq_mask;
+ inp->cross_kq_mask_cnv = inp->cross_kq_mask;
return (llm_graph_input_attn_cross *) res->add_input(std::move(inp));
}
return cur;
}
+llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
+ const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_context *>(mctx);
+
+ auto inp = std::make_unique<llm_graph_input_attn_k_dsa>(hparams, cparams, mctx_cur);
+
+ {
+ inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch);
+
+ inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams);
+ inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla;
+ }
+
+ {
+ inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
+
+ // ensure F32 mask
+ auto cparams_copy = cparams;
+ cparams_copy.flash_attn = false;
+
+ inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
+ inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;
+
+ inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0);
+ }
+
+ return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));
+}
+
// TODO: maybe separate the inner implementation into a separate function
// like with the non-sliding window equivalent
// once sliding-window hybrid caches are a thing.
inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
- inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
}
{
inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
- inp->self_kq_mask_swa_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_swa, GGML_TYPE_F16) : inp->self_kq_mask_swa;
+ inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
}
inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch);
inp_attn->self_kq_mask = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams);
- inp_attn->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp_attn->self_kq_mask, GGML_TYPE_F16) : inp_attn->self_kq_mask;
+ inp_attn->self_kq_mask_cnv = inp_attn->self_kq_mask;
}
{
inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch);
inp_attn->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams);
- inp_attn->self_kq_mask_swa_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp_attn->self_kq_mask_swa, GGML_TYPE_F16) : inp_attn->self_kq_mask_swa;
+ inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa;
}
auto inp = std::make_unique<llm_graph_input_mem_hybrid_iswa>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
struct llama_memory_context_i;
class llama_kv_cache_context;
+class llama_kv_cache_dsa_context;
class llama_kv_cache_iswa_context;
class llama_memory_recurrent_context;
class llama_memory_hybrid_context;
LLM_GRAPH_TYPE_DECODER_MTP,
};
-enum llm_ffn_op_type {
+enum llm_ffn_op_type : int {
+ LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use
LLM_FFN_SILU,
LLM_FFN_GELU,
LLM_FFN_RELU,
const int64_t n_embd = 0;
};
+// similar to llm_graph_input_embd but with an additional hidden state input
+class llm_graph_input_embd_h : public llm_graph_input_i {
+public:
+ llm_graph_input_embd_h(int64_t n_embd) : n_embd(n_embd) {}
+ virtual ~llm_graph_input_embd_h() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ ggml_tensor * tokens = nullptr; // I32 [n_batch]
+ ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch]
+ ggml_tensor * h = nullptr; // F32 [n_embd, n_batch]
+
+ const int64_t n_embd = 0;
+};
+
class llm_graph_input_pos : public llm_graph_input_i {
public:
llm_graph_input_pos(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {}
ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
// n_tokens == n_batch
- ggml_tensor * self_kq_mask = nullptr; // F32 [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32 [n_tokens, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]
const llama_hparams hparams;
const llama_cparams cparams;
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]
- ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
// note: assumes v_rot^2 == I
ggml_tensor * self_k_rot = nullptr;
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
- ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
const llama_hparams hparams;
const llama_cparams cparams;
const llama_kv_cache_context * mctx;
};
+class llm_graph_input_attn_k_dsa : public llm_graph_input_i {
+public:
+ llm_graph_input_attn_k_dsa(
+ const llama_hparams & hparams,
+ const llama_cparams & cparams,
+ const llama_kv_cache_dsa_context * mctx) :
+ hparams(hparams),
+ cparams(cparams),
+ mctx(mctx) {
+ }
+ ~llm_graph_input_attn_k_dsa() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; }
+ ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; }
+
+ ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; }
+ ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; }
+
+ ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch]
+ ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch]
+
+ ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+
+ ggml_tensor * self_k_rot_lid = nullptr;
+
+ const llama_hparams hparams;
+ const llama_cparams cparams;
+
+ const llama_kv_cache_dsa_context * mctx;
+};
+
class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {
public:
llm_graph_input_attn_kv_iswa(
ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
ggml_tensor * self_v_idxs_swa = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]
- ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+ ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_k_rot = nullptr;
ggml_tensor * self_v_rot = nullptr;
ggml_tensor * get_kq_mask_cross() const { return cross_kq_mask_cnv; }
- ggml_tensor * cross_kq_mask = nullptr; // F32 [n_outputs_enc, n_batch, 1, 1]
- ggml_tensor * cross_kq_mask_cnv = nullptr; // F32 [n_outputs_enc, n_batch, 1, 1]
+ ggml_tensor * cross_kq_mask = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]
+ ggml_tensor * cross_kq_mask_cnv = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]
const llama_cross * cross = nullptr;
};
ggml_tensor * get_logits() const { return t_logits; }
ggml_tensor * get_embd() const { return t_embd; }
ggml_tensor * get_embd_pooled() const { return t_embd_pooled; }
- ggml_tensor * get_h_pre_norm() const { return t_h_pre_norm; }
+ ggml_tensor * get_h_nextn() const { return t_h_nextn; }
ggml_cgraph * get_gf() const { return gf; }
ggml_context * get_ctx() const { return ctx_compute.get(); }
ggml_tensor * t_logits = nullptr;
ggml_tensor * t_embd = nullptr;
ggml_tensor * t_embd_pooled = nullptr;
- ggml_tensor * t_h_pre_norm = nullptr; // [n_embd, n_outputs] hidden state before final output norm
+ ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm
std::map<llama_seq_id, ggml_tensor*> t_sampled_logits;
std::map<llama_seq_id, ggml_tensor*> t_candidates;
const int64_t n_embd;
const int64_t n_layer;
+ const int64_t n_layer_nextn;
const int64_t n_rot;
const int64_t n_ctx; // user-specified context size (can be different from n_ctx_train)
const int64_t n_head;
float kq_scale,
int il) const;
+ llm_graph_input_attn_k_dsa * build_attn_inp_k_dsa() const;
+
+ ggml_tensor * build_attn(
+ llm_graph_input_attn_k_dsa * inp,
+ ggml_tensor * wo,
+ ggml_tensor * wo_b,
+ ggml_tensor * wo_s,
+ ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]
+ ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]
+ ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens]
+ ggml_tensor * kq_b,
+ ggml_tensor * sinks, // [n_head_q]
+ ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
+ ggml_tensor * top_k, // [n_indexer_top_k, n_tokens]
+ float kq_scale,
+ int il) const;
+
llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
// note: if k_cur or v_cur are not provided, they will not be stored in the memory
void llama_hparams::set_swa_pattern(uint32_t n_pattern, bool dense_first) {
if (dense_first) {
- for (uint32_t il = 0; il < n_layer; ++il) {
- swa_layers[il] = n_pattern == 0 || (il % n_pattern != 0);
+ for (uint32_t il = 0; il < n_layer(); ++il) {
+ is_swa_impl[il] = n_pattern == 0 || (il % n_pattern != 0);
}
} else {
- for (uint32_t il = 0; il < n_layer; ++il) {
- swa_layers[il] = n_pattern == 0 || (il % n_pattern < (n_pattern - 1));
+ for (uint32_t il = 0; il < n_layer(); ++il) {
+ is_swa_impl[il] = n_pattern == 0 || (il % n_pattern < (n_pattern - 1));
}
}
+
+ for (uint32_t il = n_layer(); il < n_layer_all; ++il) {
+ is_swa_impl[il] = false;
+ }
+}
+
+void llama_hparams::set_recr_pattern(uint32_t n_pattern, bool dense_first) {
+ if (dense_first) {
+ for (uint32_t il = 0; il < n_layer(); ++il) {
+ is_recr_impl[il] = n_pattern == 0 || (il % n_pattern != 0);
+ }
+ } else {
+ for (uint32_t il = 0; il < n_layer(); ++il) {
+ is_recr_impl[il] = n_pattern == 0 || (il % n_pattern < (n_pattern - 1));
+ }
+ }
+
+ for (uint32_t il = n_layer(); il < n_layer_all; ++il) {
+ is_recr_impl[il] = false;
+ }
}
bool llama_hparams::is_swa_any() const {
- for (uint32_t il = 0; il < n_layer; ++il) {
- if (swa_layers[il]) {
+ for (uint32_t il = 0; il < n_layer_all; ++il) {
+ if (is_swa_impl[il]) {
return true;
}
}
}
uint32_t llama_hparams::n_head(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return n_head_arr[il];
}
}
uint32_t llama_hparams::n_head_kv(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return n_head_kv_arr[il];
}
}
uint32_t llama_hparams::n_ff(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return n_ff_arr[il];
}
}
uint32_t llama_hparams::n_rot(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return is_swa(il) ? n_rot_swa : n_rot_full;
}
}
uint32_t llama_hparams::n_embd_inp() const {
+ if (n_embd_inp_impl > 0) {
+ return n_embd_inp_impl;
+ }
+
uint32_t n_embd_inp = n_embd;
if (n_deepstack_layers > 0) {
}
uint32_t llama_hparams::n_embd_head_k(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return is_swa(il) ? n_embd_head_k_swa : n_embd_head_k_full;
}
}
uint32_t llama_hparams::n_embd_head_v(uint32_t il) const {
- if (il < n_layer) {
+ if (il < n_layer_all) {
return is_swa(il) ? n_embd_head_v_swa : n_embd_head_v_full;
}
bool llama_hparams::is_n_embd_k_gqa_variable() const {
const uint32_t val = n_embd_k_gqa();
- for (uint32_t il = 0; il < n_layer; ++il) {
+ for (uint32_t il = 0; il < n_layer_all; ++il) {
if (val != n_embd_k_gqa(il)) {
return true;
}
bool llama_hparams::is_n_embd_v_gqa_variable() const {
const uint32_t val = n_embd_v_gqa();
- for (uint32_t il = 0; il < n_layer; ++il) {
+ for (uint32_t il = 0; il < n_layer_all; ++il) {
if (val != n_embd_v_gqa(il)) {
return true;
}
uint32_t llama_hparams::n_embd_k_gqa_max() const {
uint32_t val = n_embd_k_gqa();
- for (uint32_t il = 0; il < n_layer; ++il) {
+ for (uint32_t il = 0; il < n_layer_all; ++il) {
val = std::max(val, n_embd_k_gqa(il));
}
uint32_t llama_hparams::n_embd_v_gqa_max() const {
uint32_t val = n_embd_v_gqa();
- for (uint32_t il = 0; il < n_layer; ++il) {
+ for (uint32_t il = 0; il < n_layer_all; ++il) {
val = std::max(val, n_embd_v_gqa(il));
}
return ssm_d_state * ssm_d_inner;
}
-bool llama_hparams::is_recurrent(uint32_t il) const {
- if (il < n_layer) {
- return recurrent_layer_arr[il];
+bool llama_hparams::is_recr(uint32_t il) const {
+ if (il < n_layer_all) {
+ return is_recr_impl[il];
}
- GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer);
+ GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all);
}
uint32_t llama_hparams::n_pos_per_embd() const {
}
bool llama_hparams::is_swa(uint32_t il) const {
- if (il < n_layer) {
- return swa_layers[il];
+ if (il < n_layer_all) {
+ return is_swa_impl[il];
}
- GGML_ABORT("fatal error");
+ GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all);
}
bool llama_hparams::is_mla() const {
}
bool llama_hparams::has_kv(uint32_t il) const {
- if (kv_only_nextn) {
- // MTP head: only the trailing nextn_predict_layers blocks own a KV cache;
- // the leading trunk blocks are not executed in this graph.
- return nextn_predict_layers > 0 && il >= (n_layer - nextn_predict_layers);
- }
-
if (n_layer_kv_from_start >= 0) {
if (il < (uint32_t) n_layer_kv_from_start) {
return true;
return true;
}
-uint32_t llama_hparams::n_layer_kv() const {
- uint32_t res = 0;
-
- for (uint32_t il = 0; il < n_layer; ++il) {
- if (has_kv(il)) {
- res++;
- }
- }
-
- return res;
+uint32_t llama_hparams::n_layer() const {
+ return n_layer_all - n_layer_nextn;
}
bool llama_hparams::use_mrope() const {
LLAMA_SWA_TYPE_SYMMETRIC = 3,
};
+// forward declaration; full definition in llama-graph.h
+enum llm_ffn_op_type : int;
+
struct llama_hparams_posnet {
uint32_t n_embd;
uint32_t n_layer;
};
struct llama_hparams {
+ // note: use the `_impl` suffix to avoid name conflict between members and getters
+ // for example: n_embd_out() vs n_embd_out_impl
+
bool vocab_only;
bool no_alloc;
bool rope_finetuned;
uint32_t n_ctx_train; // context size the model was trained on
uint32_t n_embd;
- uint32_t n_layer;
- int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
+ uint32_t n_layer_all;
+ uint32_t n_layer_nextn = 0;
uint32_t n_expert = 0;
uint32_t n_expert_used = 0;
uint32_t n_rel_attn_bkts = 0;
+ // TODO: this needs to be reworked
+ int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
+
// different head size for full_attention and SWA layers
uint32_t n_embd_head_k_full; // dimension of keys (d_k). d_q is assumed to be the same, but there are n_head q heads, and only n_head_kv k-v heads
uint32_t n_embd_head_v_full; // dimension of values (d_v) aka n_embd_head
uint32_t expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_NONE;
uint32_t moe_every_n_layers = 0;
uint32_t moe_latent_size = 0;
- uint32_t nextn_predict_layers = 0;
-
- bool kv_only_nextn = false; // if true, only the last nextn_predict_layers blocks have a KV cache (MTP head arches)
float f_norm_eps;
float f_norm_rms_eps;
llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
// the size of the sliding window (0 - no SWA)
uint32_t n_swa = 0;
- // if swa_layers[il] == 1, then layer il is SWA
- // if swa_layers[il] == 0, then layer il is dense (i.e. non-SWA)
+
+ // if is_swa_impl[il] == 1, then layer il is SWA
+ // if is_swa_impl[il] == 0, then layer il is dense (i.e. non-SWA)
// by default, all layers are dense
// note: using uint32_t type for compatibility reason
- std::array<uint32_t, LLAMA_MAX_LAYERS> swa_layers;
+ std::array<uint32_t, LLAMA_MAX_LAYERS> is_swa_impl;
+
+ // for hybrid state space models
+ std::array<uint32_t, LLAMA_MAX_LAYERS> is_recr_impl;
// for State Space Models
uint32_t ssm_d_conv = 0;
// for Kimi Linear KDA
uint32_t n_embd_head_kda = 0;
- // for hybrid state space models
- std::array<bool, LLAMA_MAX_LAYERS> recurrent_layer_arr;
-
bool ssm_dt_b_c_rms = false;
float f_clamp_kqv = 0.0f;
// for Classifiers
uint32_t n_cls_out = 1;
+ // input embedding dimension (0 = use n_embd)
+ uint32_t n_embd_inp_impl = 0;
+
// output embedding dimension (0 = use n_embd)
uint32_t n_embd_out_impl = 0;
uint32_t indexer_top_k = 0;
// qwen3vl deepstack
+ // When parsed from GGUF, this implies the first N layers consume the first
+ // N deepstack embeddings. Use deepstack_mapping_arr if you need a more
+ // complex mapping. If using deepstack_mapping_arr, also make sure to set
+ // n_deepstack_layers to the number of unique deepstack layers so that
+ // n_embd_imp is accurate (see granite.cpp).
+ // TODO: can be expressed via the `new n_embd_inp_impl` and remove this param
uint32_t n_deepstack_layers = 0;
+ // deepstack layer array (Granite4 Vision)
+ // -1 => no deepstack
+ // >=0 => input embedding index for deepstack injection
+ std::array<int32_t, LLAMA_MAX_LAYERS> deepstack_mapping_arr;
+
// gemma4 per-layer embedding
uint32_t n_embd_per_layer = 0;
enum llama_rope_scaling_type rope_scaling_type_train = LLAMA_ROPE_SCALING_TYPE_NONE;
+ // Resolved FFN gated activation flavor for archs that read
+ // `<arch>.hidden_activation` from the GGUF (e.g. ModernBert derivatives).
+ // Defaults to LLM_FFN_NONE (sentinel = 0); the mapping from the GGUF
+ // string to a real op is done at hparam-load time via
+ // llm_ffn_op_type_from_string() in llama-model.cpp, mirroring how
+ // rope_scaling_type_train is handled.
+ enum llm_ffn_op_type llm_ffn_op;
+
// Step35: optional per-layer clamps for (Swi)GLU
std::array<float, LLAMA_MAX_LAYERS> swiglu_clamp_exp; // clamping for expert FFN
std::array<float, LLAMA_MAX_LAYERS> swiglu_clamp_shexp; // shared expert
// return true if one of the layers is SWA
bool is_swa_any() const;
+ bool is_swa(uint32_t il) const;
+
+ void set_recr_pattern(uint32_t n_pattern, bool dense_first = false);
+
+ // whether or not the given layer is recurrent (for hybrid models)
+ bool is_recr(uint32_t il) const;
+
uint32_t n_head(uint32_t il = 0) const;
uint32_t n_head_kv(uint32_t il = 0) const;
// dimension of the recurrent state embeddings
uint32_t n_embd_s() const;
- // whether or not the given layer is recurrent (for hybrid models)
- bool is_recurrent(uint32_t il) const;
-
uint32_t n_pos_per_embd() const;
- bool is_swa(uint32_t il) const;
-
// note: currently only support if either all or none of the layers are MLA
bool is_mla() const;
bool has_kv(uint32_t il) const;
- // number of layers for which has_kv() returns true
- uint32_t n_layer_kv() const;
+ // number of effective layers (excludes nextn layers)
+ uint32_t n_layer() const;
// note that this function uses different SWA parameters from those in the hparams
// note: inlined on purpose for performance reasons
#include "ggml.h" // for ggml_log_level
#include <string>
+#include <type_traits>
#include <vector>
#ifdef __GNUC__
no_init() = default;
};
+template <typename dst_t, typename src_t>
+static inline dst_t llama_cast(src_t v) {
+ if constexpr (std::is_same_v<src_t, dst_t>) {
+ return v;
+ } else if constexpr (std::is_same_v<src_t, ggml_fp16_t> && std::is_same_v<dst_t, float>) {
+ return ggml_fp16_to_fp32(v);
+ } else if constexpr (std::is_same_v<src_t, float> && std::is_same_v<dst_t, ggml_fp16_t>) {
+ return ggml_fp32_to_fp16(v);
+ } else {
+ static_assert(std::is_same_v<dst_t, void>, "unsupported type combination");
+ }
+}
+
struct time_meas {
time_meas(int64_t & t_acc, bool disable = false);
~time_meas();
--- /dev/null
+#include "llama-kv-cache-dsa.h"
+
+#include "llama-impl.h"
+#include "llama-batch.h"
+#include "llama-model.h"
+
+#include <algorithm>
+#include <cassert>
+
+//
+// llama_kv_cache_dsa
+//
+
+llama_kv_cache_dsa::llama_kv_cache_dsa(
+ const llama_model & model,
+ ggml_type type_k,
+ ggml_type type_v,
+ bool v_trans,
+ bool offload,
+ bool unified,
+ uint32_t kv_size,
+ uint32_t n_seq_max,
+ uint32_t n_pad,
+ uint32_t n_swa,
+ llama_swa_type swa_type,
+ const layer_filter_cb & filter,
+ const layer_reuse_cb & reuse) :
+ hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) {
+
+ LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
+
+ kv_mla = std::make_unique<llama_kv_cache>(
+ model, model.hparams, type_k, type_v,
+ v_trans, offload, unified, kv_size, n_seq_max, n_pad,
+ n_swa, swa_type, nullptr, filter, reuse, nullptr);
+
+ // we use llama_kv_cache for caching indexer keys
+ // by hand-tweaking some hparams we fool it to create
+ // indexer key cache tensors with correct dimensions
+ // https://github.com/ggml-org/llama.cpp/pull/21149#discussion_r3015940823
+
+ // DSA lightning indexer uses MQA with single key head
+ std::fill(hparams_lid.n_head_kv_arr.begin(), hparams_lid.n_head_kv_arr.end(), 1);
+ hparams_lid.n_embd_head_k_full = model.hparams.indexer_head_size;
+ hparams_lid.rope_type = LLAMA_ROPE_TYPE_NEOX;
+
+ LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
+
+ kv_lid = std::make_unique<llama_kv_cache>(
+ model, hparams_lid, type_k, type_v,
+ v_trans, offload, unified, kv_size, n_seq_max, n_pad,
+ n_swa, swa_type, nullptr, filter, reuse, nullptr);
+}
+
+void llama_kv_cache_dsa::clear(bool data) {
+ kv_mla->clear(data);
+ kv_lid->clear(data);
+}
+
+bool llama_kv_cache_dsa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
+ bool res = true;
+
+ res = res & kv_mla->seq_rm(seq_id, p0, p1);
+ res = res & kv_lid->seq_rm(seq_id, p0, p1);
+
+ return res;
+}
+
+void llama_kv_cache_dsa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ kv_mla->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+ kv_lid->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+}
+
+void llama_kv_cache_dsa::seq_keep(llama_seq_id seq_id) {
+ kv_mla->seq_keep(seq_id);
+ kv_lid->seq_keep(seq_id);
+}
+
+void llama_kv_cache_dsa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+ kv_mla->seq_add(seq_id, p0, p1, shift);
+ kv_lid->seq_add(seq_id, p0, p1, shift);
+}
+
+void llama_kv_cache_dsa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+ kv_mla->seq_div(seq_id, p0, p1, d);
+ kv_lid->seq_div(seq_id, p0, p1, d);
+}
+
+llama_pos llama_kv_cache_dsa::seq_pos_min(llama_seq_id seq_id) const {
+ return kv_mla->seq_pos_min(seq_id);
+}
+
+llama_pos llama_kv_cache_dsa::seq_pos_max(llama_seq_id seq_id) const {
+ return kv_mla->seq_pos_max(seq_id);
+}
+
+std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_dsa::memory_breakdown() const {
+ std::map<ggml_backend_buffer_type_t, size_t> mb = kv_mla->memory_breakdown();
+ for (const auto & buft_size : kv_lid->memory_breakdown()) {
+ mb[buft_size.first] += buft_size.second;
+ }
+ return mb;
+}
+
+llama_memory_context_ptr llama_kv_cache_dsa::init_batch(
+ llama_batch_allocr & balloc,
+ uint32_t n_ubatch,
+ bool embd_all) {
+ GGML_UNUSED(embd_all);
+
+ do {
+ balloc.split_reset();
+
+ std::vector<llama_ubatch> ubatches;
+ while (true) {
+ auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
+
+ if (ubatch.n_tokens == 0) {
+ break;
+ }
+
+ ubatches.push_back(std::move(ubatch)); // NOLINT
+ }
+
+ if (balloc.get_n_used() < balloc.get_n_tokens()) {
+ // failed to find a suitable split
+ break;
+ }
+
+ auto sinfos_mla = kv_mla->prepare(ubatches);
+ if (sinfos_mla.empty()) {
+ break;
+ }
+
+ auto sinfos_lid = kv_lid->prepare(ubatches);
+ if (sinfos_lid.empty()) {
+ break;
+ }
+
+ assert(sinfos_mla.size() == sinfos_lid.size());
+
+ return std::make_unique<llama_kv_cache_dsa_context>(
+ this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(ubatches));
+ } while (false);
+
+ return std::make_unique<llama_kv_cache_dsa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
+}
+
+llama_memory_context_ptr llama_kv_cache_dsa::init_full() {
+ return std::make_unique<llama_kv_cache_dsa_context>(this);
+}
+
+llama_memory_context_ptr llama_kv_cache_dsa::init_update(llama_context * lctx, bool optimize) {
+ return std::make_unique<llama_kv_cache_dsa_context>(this, lctx, optimize);
+}
+
+bool llama_kv_cache_dsa::get_can_shift() const {
+ return kv_mla->get_can_shift() &&
+ kv_lid->get_can_shift() &&
+ kv_mla->get_size() == kv_lid->get_size();
+}
+
+void llama_kv_cache_dsa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+ kv_mla->state_write(io, seq_id, flags);
+ kv_lid->state_write(io, seq_id, flags);
+}
+
+void llama_kv_cache_dsa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+ kv_mla->state_read(io, seq_id, flags);
+ kv_lid->state_read(io, seq_id, flags);
+}
+
+llama_kv_cache * llama_kv_cache_dsa::get_mla() const {
+ return kv_mla.get();
+}
+
+llama_kv_cache * llama_kv_cache_dsa::get_lid() const {
+ return kv_lid.get();
+}
+
+//
+// llama_kv_cache_dsa_context
+//
+
+llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(llama_memory_status status) : status(status) {}
+
+llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv) :
+ ctx_mla(kv->get_mla()->init_full()),
+ ctx_lid(kv->get_lid()->init_full()),
+ status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) {
+}
+
+llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv,
+ llama_context * lctx,
+ bool optimize) :
+ ctx_mla(kv->get_mla()->init_update(lctx, optimize)),
+ ctx_lid(kv->get_lid()->init_update(lctx, optimize)),
+ status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) {
+}
+
+llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv,
+ slot_info_vec_t sinfos_mla,
+ slot_info_vec_t sinfos_lid,
+ std::vector<llama_ubatch> ubatches) :
+ ubatches(std::move(ubatches)),
+ // note: here we copy the ubatches. not sure if this is ideal
+ ctx_mla(new llama_kv_cache_context(kv->get_mla(), std::move(sinfos_mla), this->ubatches)),
+ ctx_lid(new llama_kv_cache_context(kv->get_lid(), std::move(sinfos_lid), this->ubatches)),
+ status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) {
+}
+
+llama_kv_cache_dsa_context:: ~llama_kv_cache_dsa_context() = default;
+
+bool llama_kv_cache_dsa_context::next() {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ ctx_mla->next();
+ ctx_lid->next();
+
+ if (++i_next >= ubatches.size()) {
+ return false;
+ }
+
+ return true;
+}
+
+bool llama_kv_cache_dsa_context::apply() {
+ assert(!llama_memory_status_is_fail(status));
+
+ bool res = true;
+
+ res = res & ctx_mla->apply();
+ res = res & ctx_lid->apply();
+
+ return res;
+}
+
+llama_memory_status llama_kv_cache_dsa_context::get_status() const {
+ return status;
+}
+
+const llama_ubatch & llama_kv_cache_dsa_context::get_ubatch() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return ubatches[i_next];
+}
+
+const llama_kv_cache_context * llama_kv_cache_dsa_context::get_mla() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return static_cast<const llama_kv_cache_context *>(ctx_mla.get());
+}
+
+const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return static_cast<const llama_kv_cache_context *>(ctx_lid.get());
+}
--- /dev/null
+#pragma once
+
+#include "llama-kv-cache.h"
+
+#include <vector>
+
+//
+// llama_kv_cache_dsa
+//
+
+// utilizes two instances of llama_kv_cache:
+// - the first instance is for caching key tensors of the model,
+// - the second instance is for caching lightning indexer key tensors
+
+class llama_kv_cache_dsa : public llama_memory_i {
+public:
+ llama_kv_cache_dsa(
+ const llama_model & model,
+ ggml_type type_k,
+ ggml_type type_v,
+ bool v_trans,
+ bool offload,
+ bool unified,
+ uint32_t kv_size,
+ uint32_t n_seq_max,
+ uint32_t n_pad,
+ uint32_t n_swa,
+ llama_swa_type swa_type,
+ const layer_filter_cb & filter,
+ const layer_reuse_cb & reuse);
+
+ ~llama_kv_cache_dsa() = default;
+
+ //
+ // llama_memory_i
+ //
+
+ llama_memory_context_ptr init_batch(
+ llama_batch_allocr & balloc,
+ uint32_t n_ubatch,
+ bool embd_all) override;
+
+ llama_memory_context_ptr init_full() override;
+
+ llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;
+
+ bool get_can_shift() const override;
+
+ void clear(bool data) override;
+
+ bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override;
+ void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;
+ void seq_keep(llama_seq_id seq_id) override;
+ void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override;
+ void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override;
+
+ llama_pos seq_pos_min(llama_seq_id seq_id) const override;
+ llama_pos seq_pos_max(llama_seq_id seq_id) const override;
+
+ std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
+
+ // state write/load
+
+ void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
+ void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
+
+ //
+ // llama_kv_cache_dsa specific API
+ //
+
+ llama_kv_cache * get_mla() const;
+ llama_kv_cache * get_lid() const;
+
+private:
+ // we keep indexer KV cache hparams instance here as llama_kv_cache stores only reference to it
+ llama_hparams hparams_lid;
+ const uint32_t n_stream = 1;
+
+ std::unique_ptr<llama_kv_cache> kv_mla;
+ std::unique_ptr<llama_kv_cache> kv_lid;
+};
+
+class llama_kv_cache_dsa_context : public llama_memory_context_i {
+public:
+ using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
+
+ // used for errors
+ llama_kv_cache_dsa_context(llama_memory_status status);
+
+ // used to create a full-cache context
+ llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv);
+
+ // used to create an update context
+ llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv,
+ llama_context * lctx,
+ bool optimize);
+
+ // used to create a batch processing context from a batch
+ llama_kv_cache_dsa_context(
+ llama_kv_cache_dsa * kv,
+ slot_info_vec_t sinfos_base,
+ slot_info_vec_t sinfos_ik,
+ std::vector<llama_ubatch> ubatches);
+
+ virtual ~llama_kv_cache_dsa_context();
+
+ //
+ // llama_memory_context_i
+ //
+
+ bool next() override;
+ bool apply() override;
+
+ llama_memory_status get_status() const override;
+ const llama_ubatch & get_ubatch() const override;
+
+ //
+ // llama_kv_cache_dsa_context specific API
+ //
+
+ const llama_kv_cache_context * get_mla() const;
+ const llama_kv_cache_context * get_lid() const;
+
+private:
+ //llama_kv_cache_dsa * kv;
+
+ // the index of the next ubatch to process
+ size_t i_next = 0;
+
+ std::vector<llama_ubatch> ubatches;
+
+ const llama_memory_context_ptr ctx_mla;
+ const llama_memory_context_ptr ctx_lid;
+
+ const llama_memory_status status;
+};
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
+ llama_memory_t mem_other,
const layer_filter_cb & filter,
- const layer_reuse_cb & reuse) : hparams(model.hparams), unified(unified) {
+ const layer_reuse_cb & reuse,
+ const layer_share_cb & share) : hparams(model.hparams), unified(unified) {
// chain filters
const layer_filter_cb filter_base = [&](int32_t il) {
LLAMA_LOG_INFO("%s: creating non-SWA KV cache, size = %u cells\n", __func__, size_base);
+ llama_memory_t mem_other_base = nullptr;
+ if (mem_other) {
+ mem_other_base = static_cast<llama_kv_cache_iswa *>(mem_other)->get_base();
+ }
+
+ llama_memory_t mem_other_swa = nullptr;
+ if (mem_other) {
+ mem_other_swa = static_cast<llama_kv_cache_iswa *>(mem_other)->get_swa();
+ }
+
kv_base = std::make_unique<llama_kv_cache>(
- model, type_k, type_v,
+ model, hparams, type_k, type_v,
v_trans, offload, unified, size_base, n_seq_max, n_pad,
- 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse);
+ 0, LLAMA_SWA_TYPE_NONE, mem_other_base, filter_base, reuse, share);
LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa);
kv_swa = std::make_unique<llama_kv_cache>(
- model, type_k, type_v,
+ model, hparams, type_k, type_v,
v_trans, offload, unified, size_swa, n_seq_max, n_pad,
- hparams.n_swa, hparams.swa_type, filter_swa, reuse);
+ hparams.n_swa, hparams.swa_type, mem_other_swa, filter_swa, reuse, share);
}
void llama_kv_cache_iswa::clear(bool data) {
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
+ llama_memory_t mem_other,
const layer_filter_cb & filter,
- const layer_reuse_cb & reuse);
+ const layer_reuse_cb & reuse,
+ const layer_share_cb & share);
~llama_kv_cache_iswa() = default;
llama_kv_cache::llama_kv_cache(
const llama_model & model,
+ const llama_hparams & hparams,
ggml_type type_k,
ggml_type type_v,
bool v_trans,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
+ llama_memory_t mem_other,
const layer_filter_cb & filter,
- const layer_reuse_cb & reuse) :
- model(model), hparams(model.hparams), v_trans(v_trans),
- n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type) {
+ const layer_reuse_cb & reuse,
+ const layer_share_cb & share) :
+ model(model), hparams(hparams), v_trans(v_trans),
+ n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type),
+ other(static_cast<llama_kv_cache *>(mem_other)),
+ v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
+ v_cells(*v_cells_impl) {
+
+ // shared cells view the source cache's K/V tensors, so the cell count
+ // follows the source allocation: a fitted target can be smaller than the
+ // draft default and oversized views would overflow the source tensors
+ if (other) {
+ const uint32_t size_other = other->get_size();
+ if (kv_size != size_other) {
+ LLAMA_LOG_WARN("%s: kv_size = %u overridden to %u to match the shared source cache\n", __func__, kv_size, size_other);
+ kv_size = size_other;
+ }
+ }
GGML_ASSERT(kv_size % n_pad == 0);
- const uint32_t n_layer_kv = hparams.n_layer_kv();
+ const uint32_t n_layer = hparams.n_layer_all;
// define a comparator for the buft -> ctx map to ensure that the order is well-defined:
struct ggml_backend_buft_comparator {
auto it = ctx_map.find(buft);
if (it == ctx_map.end()) {
ggml_init_params params = {
- /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer_kv*ggml_tensor_overhead()),
+ /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
const bool is_mla = hparams.is_mla();
- for (uint32_t il = 0; il < hparams.n_layer; il++) {
+ for (uint32_t il = 0; il < n_layer; il++) {
if (!hparams.has_kv(il)) {
LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il);
continue;
continue;
}
+ if (share && other) {
+ const int32_t il_share = share(il);
+
+ if (il_share >= 0) {
+ const auto & layer_share = other->layers[other->map_layer_ids[il_share]];
+
+ LLAMA_LOG_WARN("%s: layer %3d: sharing with layer %d. k = %p, v = %p\n", __func__, il, il_share,
+ layer_share.k->data, layer_share.v->data);
+
+ map_layer_ids[il] = layers.size();
+
+ layers.push_back(layer_share);
+ layers.back().il = il;
+
+ continue;
+ }
+ }
+
if (n_embd_head_k_all == 0) {
n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il);
} else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) {
if (reuse) {
LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__);
- for (uint32_t il = 0; il < hparams.n_layer; il++) {
+ for (uint32_t il = 0; il < n_layer; il++) {
const int32_t il_reuse = reuse(il);
if (il_reuse < 0) {
// allocate tensors and initialize the buffers to avoid NaNs in the padding
for (auto & [buft, ctx] : ctx_map) {
ggml_backend_buffer_t buf;
- if (model.hparams.no_alloc) {
+ if (hparams.no_alloc) {
buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
}
- const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
- const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
- if (attn_rot_disable) {
- LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
- }
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ n_embd_head_k_all = other->n_embd_head_k_all;
+ n_embd_head_v_all = other->n_embd_head_v_all;
+
+ attn_rot_k = other->attn_rot_k;
+ attn_rot_v = other->attn_rot_v;
+ } else {
+ const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
+ const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
+ if (attn_rot_disable) {
+ LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
+ }
+
+ attn_rot_k =
+ !attn_rot_disable &&
+ n_embd_head_k_all > 0 &&
+ ggml_is_quantized(type_k) &&
+ hparams.n_embd_head_k() % 64 == 0;
- attn_rot_k =
- !attn_rot_disable &&
- n_embd_head_k_all > 0 &&
- ggml_is_quantized(type_k) &&
- hparams.n_embd_head_k() % 64 == 0;
+ // always create Hadamard rotation tensors for DeepSeek V3.2 DSA lightning indexer
+ if (model.arch == LLM_ARCH_DEEPSEEK32 && hparams.n_embd_head_k_full == hparams.indexer_head_size) {
+ attn_rot_k = true;
+ }
- attn_rot_v =
- !attn_rot_disable &&
- n_embd_head_v_all > 0 &&
- ggml_is_quantized(type_v) &&
- hparams.n_embd_head_v() % 64 == 0;
+ attn_rot_v =
+ !attn_rot_disable &&
+ n_embd_head_v_all > 0 &&
+ ggml_is_quantized(type_v) &&
+ hparams.n_embd_head_v() % 64 == 0;
+ }
LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all);
LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all);
}
bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return true;
+ }
+
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
if (p0 < 0) {
}
void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());
GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());
}
void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
auto & cells = v_cells[seq_to_stream[seq_id]];
}
void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1");
}
void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1");
}
llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return other->seq_pos_min(seq_id);
+ }
+
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
const auto & cells = v_cells[seq_to_stream[seq_id]];
}
llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return other->seq_pos_max(seq_id);
+ }
+
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
const auto & cells = v_cells[seq_to_stream[seq_id]];
}
bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return true;
+ }
+
bool updated = false;
auto * sched = lctx->get_sched();
}
void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
// keep track of the max sequence position that we would overwrite with this ubatch
// for non-SWA cache, this would be always empty
llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ];
int64_t n_tps;
};
-template<bool causal, bool swa, bool is_2d, bool alibi>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, float * data) {
+template<typename T, bool causal, bool swa, bool is_2d, bool alibi>
+static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
//const auto & hparams = args.hparams;
const auto & ubatch = args.ubatch;
const int64_t n_stream = args.n_stream;
const int64_t n_tps = args.n_tps;
+ const T mask_keep = llama_cast<T>(0.0f);
+ const T mask_drop = llama_cast<T>(-INFINITY);
+
// the min position in the batch for each sequence
llama_pos seq_pos_min[LLAMA_MAX_SEQ];
std::fill(seq_pos_min, seq_pos_min + LLAMA_MAX_SEQ, INT32_MAX);
}
if (alibi) {
- data[idst + j] = -std::abs(p0 - p1);
+ data[idst + j] = llama_cast<T>(static_cast<float>(-std::abs(p0 - p1)));
} else {
- data[idst + j] = 0.0f;
+ data[idst + j] = mask_keep;
}
continue;
skip:
- data[idst + j] = -INFINITY;
+ data[idst + j] = mask_drop;
}
}
}
}
-template<bool causal, bool swa, bool is_2d>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, float * data) {
+template<typename T, bool causal, bool swa, bool is_2d>
+static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
const bool alibi = args.hparams.use_alibi;
if (alibi) {
- set_input_kq_mask_impl<causal, swa, is_2d, true> (args, data);
+ set_input_kq_mask_impl<T, causal, swa, is_2d, true> (args, data);
} else {
- set_input_kq_mask_impl<causal, swa, is_2d, false>(args, data);
+ set_input_kq_mask_impl<T, causal, swa, is_2d, false>(args, data);
}
}
-template<bool causal, bool swa>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, float * data) {
+template<typename T, bool causal, bool swa>
+static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
const bool is_2d = args.ubatch->is_pos_2d();
if (is_2d) {
- set_input_kq_mask_impl<causal, swa, true> (args, data);
+ set_input_kq_mask_impl<T, causal, swa, true> (args, data);
} else {
- set_input_kq_mask_impl<causal, swa, false>(args, data);
+ set_input_kq_mask_impl<T, causal, swa, false>(args, data);
}
}
-template<bool causal>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, float * data) {
+template<typename T, bool causal>
+static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
const bool swa = args.swa_type != LLAMA_SWA_TYPE_NONE;
if (swa) {
- set_input_kq_mask_impl<causal, true> (args, data);
+ set_input_kq_mask_impl<T, causal, true> (args, data);
+ } else {
+ set_input_kq_mask_impl<T, causal, false>(args, data);
+ }
+}
+
+template<typename T>
+static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data, bool causal_attn) {
+ if (causal_attn) {
+ set_input_kq_mask_impl<T, true> (args, data);
} else {
- set_input_kq_mask_impl<causal, false>(args, data);
+ set_input_kq_mask_impl<T, false>(args, data);
}
}
const uint32_t n_tokens = ubatch->n_tokens;
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
- float * data = (float *) dst->data;
const int64_t n_kv = dst->ne[0];
const int64_t n_stream = dst->ne[3]; // num streams in the current ubatch
/*.n_tps =*/ n_tps,
};
- if (causal_attn) {
- set_input_kq_mask_impl<true> (args, data);
+ if (dst->type == GGML_TYPE_F16) {
+ set_input_kq_mask_impl<ggml_fp16_t>(args, (ggml_fp16_t *) dst->data, causal_attn);
} else {
- set_input_kq_mask_impl<false>(args, data);
+ set_input_kq_mask_impl<float>(args, (float *) dst->data, causal_attn);
}
//const int64_t t_end = ggml_time_us();
}
ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_context * lctx) const {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ GGML_ASSERT(!other);
+
auto * ctx = res->get_ctx();
auto * gf = res->get_gf();
}
void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_UNUSED(flags);
io.write(&n_stream, sizeof(n_stream));
uint32_t cell_range_begin = cells.size();
for (uint32_t i = 0; i < cells.size(); ++i) {
- if (!cells.is_empty(i) && (seq_id == -1 || cells.seq_has(i, seq_id))) {
+ bool add_cell = true;
+
+ add_cell = add_cell && !cells.is_empty(i);
+ add_cell = add_cell && (seq_id == -1 || cells.seq_has(i, seq_id));
+
+ // check the cell is not SWA-masked
+ if (add_cell && seq_id != -1) {
+ const bool is_masked = llama_hparams::is_masked_swa(n_swa, swa_type, cells.pos_get(i), cells.seq_pos_max(seq_id));
+
+ add_cell = !is_masked;
+ }
+
+ if (add_cell) {
++cell_count;
if (cell_range_begin == cells.size()) {
cell_range_begin = i;
}
void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+ // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
+ if (other) {
+ return;
+ }
+
GGML_UNUSED(flags);
GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
sinfo = find_slot(ubatch, false);
if (sinfo.empty()) {
- LLAMA_LOG_ERROR("%s: failed to find available cells in kv cache\n", __func__);
+ LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count);
return false;
}
using slot_info_vec_t = std::vector<slot_info>;
+ // TODO: refactor the memory instances to not depend on `llama_model`
+ // instead pass all necessary info (e.g. hparams, dev layers, arch, etc.) directly
+ // likely through `struct llama_memory_params`
llama_kv_cache(
const llama_model & model,
+ const llama_hparams & hparams,
ggml_type type_k,
ggml_type type_v,
bool v_trans,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
+ llama_memory_t mem_other,
const layer_filter_cb & filter,
- const layer_reuse_cb & reuse);
+ const layer_reuse_cb & reuse,
+ const layer_share_cb & share);
~llama_kv_cache() = default;
// note: this is not part of the KV state and it's only used to speed-up the find_slot() method
std::vector<uint32_t> v_heads;
- std::vector<llama_kv_cells> v_cells;
+ // TODO: temporary until we refactor to be able to share the same cells between 2 kv caches [TAG_KV_CACHE_SHARE_CELLS]
+ llama_kv_cache * other;
+
+ std::shared_ptr<llama_kv_cells_vec> v_cells_impl;
+
+ llama_kv_cells_vec & v_cells;
// maps from a sequence id to a stream id
std::vector<uint32_t> seq_to_stream;
}
}
};
+
+using llama_kv_cells_vec = std::vector<llama_kv_cells>;
n_seq_max,
n_ubatch,
n_pad,
+ nullptr,
filter_attn == nullptr ?
- [&](int32_t il) { return !hparams.is_recurrent(il); }
+ [&](int32_t il) { return !hparams.is_recr(il); }
: filter_attn,
+ nullptr,
nullptr
)),
mem_recr(new llama_memory_recurrent(
n_seq_max,
n_rs_seq,
filter_recr == nullptr ?
- [&](int32_t il) { return hparams.is_recurrent(il); }
+ [&](int32_t il) { return hparams.is_recr(il); }
: filter_recr
)) {}
hparams(model.hparams),
mem_attn(new llama_kv_cache(
model,
+ model.hparams,
type_k,
type_v,
v_trans,
n_pad,
n_swa,
swa_type,
+ nullptr,
filter_attn == nullptr ?
- [&](int32_t il) { return !hparams.is_recurrent(il); }
+ [&](int32_t il) { return !hparams.is_recr(il); }
: filter_attn,
+ nullptr,
nullptr
)),
mem_recr(new llama_memory_recurrent(
n_seq_max,
n_rs_seq,
filter_recr == nullptr ?
- [&](int32_t il) { return hparams.is_recurrent(il); }
+ [&](int32_t il) { return hparams.is_recr(il); }
: filter_recr
)) {}
uint32_t n_seq_max,
uint32_t n_rs_seq,
const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) {
- const int32_t n_layer = hparams.n_layer;
+ const int32_t n_layer = hparams.n_layer();
head = 0;
size = mem_size;
void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges) const {
const uint32_t s_trans = 0;
- const uint32_t n_layer = hparams.n_layer;
+ const uint32_t n_layer = hparams.n_layer();
io.write(&s_trans, sizeof(s_trans));
io.write(&n_layer, sizeof(n_layer));
io.read(&s_trans, sizeof(s_trans));
io.read(&n_layer, sizeof(n_layer));
- if (n_layer != hparams.n_layer) {
- LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, hparams.n_layer);
+ if (n_layer != hparams.n_layer()) {
+ LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, hparams.n_layer());
return false;
}
if (cell_count > size) {
bool swa_full;
llama_context_type ctx_type;
+
+ llama_memory_t mem_other;
};
enum llama_memory_status {
// return negative value to indicate that the layer il should not reuse memory
using layer_reuse_cb = std::function<int32_t(int32_t il)>;
+ using layer_share_cb = std::function<int32_t(int32_t il)>;
+
virtual ~llama_memory_i() = default;
// split the input batch into a set of ubatches and verify that they can fit into the cache
const enum gguf_type arr_type = gguf_get_arr_type(ctx, k);
return ArrayInfo {
arr_type,
- size_t(gguf_get_arr_n(ctx, k)),
+ gguf_get_arr_n(ctx, k),
arr_type == GGUF_TYPE_STRING ? nullptr : gguf_get_arr_data(ctx, k),
};
}
}
template bool llama_model_loader::get_arr<std::vector<std::string>>(enum llm_kv kid, std::vector<std::string> & result, bool required);
+ template bool llama_model_loader::get_arr<std::array<int32_t, 512>>(enum llm_kv kid, std::array<int32_t, 512> & result, bool required);
template<typename T>
bool llama_model_loader::get_key(const std::string & key, T & result, bool required) {
}
if (n > N_MAX) {
- throw std::runtime_error(format("n > N_MAX: %u > %u for key %s", (uint32_t) n, (uint32_t) N_MAX, key.c_str()));
+ throw std::runtime_error(format("n > N_MAX: %u > %u for key %s", n, (uint32_t) N_MAX, key.c_str()));
}
if (gguf_get_kv_type(metadata, kid) == GGUF_TYPE_ARRAY) {
}
// TODO: this is not very clever - figure out something better
- template bool llama_model_loader::get_key_or_arr<std::array<int, 4>>(enum llm_kv kid, std::array<int, 4> & result, uint32_t n, bool required);
+ template bool llama_model_loader::get_key_or_arr<std::array<int, 4>> (enum llm_kv kid, std::array<int, 4> & result, uint32_t n, bool required);
template bool llama_model_loader::get_key_or_arr<std::array<uint32_t, 512>>(enum llm_kv kid, std::array<uint32_t, 512> & result, uint32_t n, bool required);
- template bool llama_model_loader::get_key_or_arr<std::array<float, 512>>(enum llm_kv kid, std::array<float, 512> & result, uint32_t n, bool required);
+ template bool llama_model_loader::get_key_or_arr<std::array<float, 512>>(enum llm_kv kid, std::array<float, 512> & result, uint32_t n, bool required);
llama_model_loader::llama_model_loader(
if (it == ctx_map.end()) {
// one ggml context per buffer type
int max_n_tensors = n_tensors;
- max_n_tensors += 1; // duplicated output tensor
- max_n_tensors += hparams.n_layer*2; // duplicated rope freq tensors
+ max_n_tensors += 1; // duplicated output tensor
+ max_n_tensors += hparams.n_layer()*2; // duplicated rope freq tensors
if (files.empty()) {
- max_n_tensors += hparams.n_layer*256; // this should be well above what any model actually uses
+ max_n_tensors += hparams.n_layer()*256; // this should be well above what any model actually uses
}
const size_t ctx_size = ggml_tensor_overhead()*max_n_tensors;
bool llama_model_saver_supports_arch(llm_arch arch) {
switch (arch) {
- case LLM_ARCH_QWEN3NEXT:
- case LLM_ARCH_QWEN35:
- case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_PLAMO3:
case LLM_ARCH_GEMMA3:
case LLM_ARCH_GEMMA3N:
case LLM_ARCH_APERTUS:
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
+ case LLM_ARCH_MELLUM:
return false;
default:
return true;
template <typename Container>
void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, const bool per_layer) {
GGML_ASSERT(model != nullptr || !per_layer);
- const size_t n_values = per_layer ? size_t(model->hparams.n_layer) : value.size();
+ const size_t n_values = per_layer ? size_t(model->hparams.n_layer()) : value.size();
GGML_ASSERT(n_values <= value.size());
if (n_values == 0) {
gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT8, value.data(), n_values);
} else if (std::is_same<typename Container::value_type, uint32_t>::value) {
gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT32, value.data(), n_values);
+ } else if (std::is_same<typename Container::value_type, bool>::value) {
+ gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_BOOL, value.data(), n_values);
} else if (std::is_same<typename Container::value_type, int32_t>::value) {
gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT32, value.data(), n_values);
} else if (std::is_same<typename Container::value_type, float>::value) {
if (hparams.n_embd_out_impl > 0) {
add_kv(LLM_KV_EMBEDDING_LENGTH_OUT, hparams.n_embd_out_impl);
}
- add_kv(LLM_KV_BLOCK_COUNT, hparams.n_layer);
+ add_kv(LLM_KV_BLOCK_COUNT, hparams.n_layer_all);
add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true);
add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
add_kv(LLM_KV_EXPERT_GROUP_SCALE, hparams.expert_group_scale);
add_kv(LLM_KV_EXPERTS_PER_GROUP, hparams.n_group_experts);
add_kv(LLM_KV_MOE_EVERY_N_LAYERS, hparams.moe_every_n_layers);
- add_kv(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers);
+ add_kv(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn);
add_kv(LLM_KV_NUM_DEEPSTACK_LAYERS, hparams.n_deepstack_layers);
+ add_kv(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr);
add_kv(LLM_KV_POOLING_TYPE, uint32_t(hparams.pooling_type));
add_kv(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
add_kv(LLM_KV_DECODER_START_TOKEN_ID, hparams.dec_start_token_id);
add_kv(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale);
add_kv(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count);
add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step);
- // add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, ???);
+ // add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, ???); // saved as LLM_KV_ATTENTION_RECURRENT_LAYERS instead
add_kv(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, true);
add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, hparams.n_head_kv_arr, true);
add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
+ add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
+#include "llama-kv-cache-dsa.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-recurrent.h"
return new llama_model_mpt(params);
case LLM_ARCH_STABLELM:
return new llama_model_stablelm(params);
+ case LLM_ARCH_MELLUM:
+ return new llama_model_mellum(params);
case LLM_ARCH_QWEN:
return new llama_model_qwen(params);
case LLM_ARCH_QWEN2:
return new llama_model_gemma3n(params);
case LLM_ARCH_GEMMA4:
return new llama_model_gemma4(params);
+ case LLM_ARCH_GEMMA4_ASSISTANT:
+ return new llama_model_gemma4_assistant(params);
case LLM_ARCH_GEMMA_EMBEDDING:
return new llama_model_gemma_embedding(params);
case LLM_ARCH_STARCODER2:
return new llama_model_deepseek2(params);
case LLM_ARCH_DEEPSEEK2OCR:
return new llama_model_deepseek2ocr(params);
+ case LLM_ARCH_DEEPSEEK32:
+ return new llama_model_deepseek32(params);
case LLM_ARCH_GLM_DSA:
return new llama_model_glm_dsa(params);
case LLM_ARCH_MISTRAL4:
// count only the same type of previous layers to avoid this
auto get_il_eff = [&](const size_t il){
size_t ret = 0;
- const bool il_is_recurrent = hparams.is_recurrent(il);
- const bool il_is_swa = hparams.is_swa(il);
+ const bool il_is_recr = hparams.is_recr(il);
+ const bool il_is_swa = hparams.is_swa(il);
for (size_t il_prev = 0; il_prev < il; il_prev++) {
- ret += hparams.is_recurrent(il_prev) == il_is_recurrent && hparams.is_swa(il_prev) == il_is_swa;
+ ret += hparams.is_recr(il_prev) == il_is_recr && hparams.is_swa(il_prev) == il_is_swa;
}
return ret;
};
rotation = get_il_eff(il) % ud->n_devices;
} else {
il = 0;
- rotation = hparams.n_layer % ud->n_devices;
+ rotation = hparams.n_layer() % ud->n_devices;
}
const ggml_tensor * tensor_axis_0 = suffix.empty() ? tensor : ud->model->get_tensor((prefix + suffix).c_str());
if (tensor_axis_0 == nullptr) {
auto get_tensor_config = [&]() -> tensor_config {
// standard attention
if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) {
- return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight");
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
}
if (std::regex_match(tensor_name, pattern_q_bias) || std::regex_match(tensor_name, pattern_kv_bias)) {
- return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight");
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight", "ssm_out.weight");
}
if (std::regex_match(tensor_name, pattern_qkv_weight)) {
- return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
}
if ( std::regex_match(tensor_name, pattern_qkv_bias)) {
- return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight", "ssm_out.weight");
}
if (std::regex_match(tensor_name, pattern_qk_norm)) {
return get_tensor_config_impl(tensor->ne[1] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight");
}
if (std::regex_match(tensor_name, pattern_attn_gate_weight)) {
- return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1);
+ return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
}
if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a)) {
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ssm_out.weight");
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
};
- auto get_split_segments = [&](int axis, uint32_t il) -> std::vector<int64_t> {
+ auto get_split_segments = [&](int axis, uint32_t il) -> std::vector<std::pair<int64_t, uint32_t>> {
if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) {
const int64_t head_k_dim = hparams.ssm_d_state;
const int64_t head_v_dim = hparams.ssm_d_state;
if (ud->model->arch == LLM_ARCH_QWEN3NEXT) {
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
- return {key_dim, key_dim, value_dim};
+ return {{key_dim, 2}, {value_dim, 1}};
}
} else {
const int64_t head_ratio = n_v_heads / n_k_heads;
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
- return std::vector<int64_t>(2 + head_ratio, key_dim);
+ return {{key_dim, 2 + head_ratio}};
}
if (std::regex_match(tensor_name, pattern_attn_gate_weight) || std::regex_match(tensor_name, pattern_ssm_out_weight)) {
- return std::vector<int64_t>(head_ratio, key_dim);
+ return {{key_dim, head_ratio}};
}
if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) ||
std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta)) {
- return std::vector<int64_t>(head_ratio, n_k_heads);
+ return {{n_k_heads, head_ratio}};
}
if (std::regex_match(tensor_name, pattern_r_cache)) {
- return std::vector<int64_t>(2 + head_ratio, key_dim * (hparams.ssm_d_conv - 1));
+ return {{key_dim * (hparams.ssm_d_conv - 1), 2 + head_ratio}};
}
if (std::regex_match(tensor_name, pattern_s_cache)) {
- return std::vector<int64_t>(head_ratio, n_k_heads * head_v_dim * head_v_dim);
+ return {{n_k_heads * head_v_dim * head_v_dim, head_ratio}};
}
}
if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
const int64_t n_ff_exp = hparams.n_ff_exp;
GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
- return {n_ff_exp, n_ff_exp};
+ return {{n_ff_exp, 2}};
}
- return {tensor->ne[axis]};
+ return {{tensor->ne[axis], 1}};
}
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il);
GGML_ASSERT(hparams.n_embd_k_gqa() == n_embd_gqa);
GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);
- return {n_embd, n_embd_gqa, n_embd_gqa};
+ return {{n_embd, 1}, {n_embd_gqa, 2}};
}
if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
const int64_t n_ff_exp = hparams.n_ff_exp;
GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
- return {n_ff_exp, n_ff_exp};
+ return {{n_ff_exp, 2}};
}
- return {tensor->ne[axis]};
+ return {{tensor->ne[axis], 1}};
};
- auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector<int64_t> & segments) -> std::vector<int64_t> {
- if (hparams.is_recurrent(il)) {
+ auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector<std::pair<int64_t, uint32_t>> & segments) -> std::vector<int64_t> {
+ // for better performance it may make sense to round up blck_size to a higher power of 2 so that more efficient kernels can be used
+ if (hparams.is_recr(il)) {
// linear attention
- const int64_t head_dim = hparams.ssm_d_state;
- const int64_t granularity_qkv = std::lcm(blck_size, head_dim);
+ const int64_t head_dim = hparams.ssm_d_state;
+ const int64_t blck_size_perf = std::lcm(blck_size, 128);
+ const int64_t granularity_qkv = std::lcm(blck_size_perf, head_dim);
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_attn_gate_weight) ||
std::regex_match(tensor_name, pattern_ssm_conv1d) || std::regex_match(tensor_name, pattern_ssm_out_weight)) {
return std::vector<int64_t>(segments.size(), granularity_qkv);
// regular attention
const uint32_t n_gqa = hparams.n_gqa(il);
const uint32_t n_embd_q = n_gqa * hparams.n_embd_head_k(il);
+
+ // to handle head sizes like 80, only increase granularity while it doesn't cause underutilization
+ int64_t blck_size_perf = blck_size;
+ while (blck_size_perf < 128 && blck_size_perf*ud->n_devices < n_embd_q) {
+ blck_size_perf *= 2;
+ }
+
if (std::regex_match(tensor_name, pattern_attn_sinks)) {
GGML_ASSERT(segments.size() == 1);
- return {std::lcm(n_embd_q, blck_size)/n_embd_q * n_gqa};
+ return {std::lcm(n_embd_q, blck_size_perf)/n_embd_q * n_gqa};
}
- const int64_t granularity_q = std::lcm(n_embd_q, blck_size);
+ const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf);
if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) {
GGML_ASSERT(segments.size() == 1);
// some models have Q gate tensors, for those cases the granularity needs to be doubled:
if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) {
- return {std::lcm(2*n_embd_q, blck_size)};
+ return {std::lcm(2*n_embd_q, blck_size_perf)};
}
return {granularity_q};
}
return {granularity_kv};
}
if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
- GGML_ASSERT(segments.size() == 3);
- return {granularity_q, granularity_kv, granularity_kv};
+ GGML_ASSERT(segments.size() == 2);
+ return {granularity_q, granularity_kv};
}
}
// FFN
if (std::regex_match(tensor_name, pattern_ffn_up_gate_weight) || std::regex_match(tensor_name, pattern_ffn_up_gate_bias) ||
std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) {
- GGML_ASSERT(segments.size() <= 2);
- return std::vector<int64_t>(segments.size(), blck_size);
+ const int64_t blck_size_perf = std::lcm(blck_size, 128);
+ GGML_ASSERT(segments.size() == 1);
+ return {blck_size_perf};
}
// everything else
tensor_config tc = get_tensor_config();
split_state.axis = tc.axis;
if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) {
- const int64_t ne_full = tensor->ne[split_state.axis];
const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type);
const float * tensor_split = ud->model->tensor_split();
std::vector<float> tensor_split_scan;
tensor_split_scan[j] += tensor_split_scan[j - 1];
}
}
- const std::vector<int64_t> segments = get_split_segments(split_state.axis, tc.il);
+ const std::vector<std::pair<int64_t, uint32_t>> segments = get_split_segments(split_state.axis, tc.il);
const std::vector<int64_t> granularity = get_split_granularity(blck_size, tc.il, segments);
for (size_t is = 0; is < segments.size(); is++) {
- const int64_t ne_s = segments[is];
- const int64_t g_s = granularity[is];
- GGML_ASSERT(ne_full % g_s == 0);
+ const int64_t ne_s = segments[is].first;
+ const uint32_t nr_s = segments[is].second;
+ const int64_t g_s = granularity[is];
int64_t low = 0;
size_t j = 0;
for (; j < ud->n_devices - 1; j++) {
low = high;
}
split_state.ne[is*ud->n_devices + (j + tc.rotation) % ud->n_devices] = ne_s - low;
+ split_state.nr[is] = nr_s;
}
split_state.n_segments = segments.size();
} else {
memset(split_state.ne, 0, sizeof(split_state.ne));
+ split_state.nr[0] = 1;
split_state.n_segments = 1;
}
return split_state;
case LLM_TYPE_A13B: return "A13B";
case LLM_TYPE_7B_A1B: return "7B.A1B";
case LLM_TYPE_8B_A1B: return "8B.A1B";
+ case LLM_TYPE_12B_A2_5B: return "12B.A2.5B";
case LLM_TYPE_16B_A1B: return "16B.A1B";
case LLM_TYPE_21B_A3B: return "21B.A3B";
case LLM_TYPE_24B_A2B: return "24B.A2B";
case LLM_TYPE_310B_A15B: return "310B.A15B";
case LLM_TYPE_355B_A32B: return "355B.A32B";
case LLM_TYPE_397B_A17B: return "397B.A17B";
+ case LLM_TYPE_685B_A37B: return "685B.A37B";
case LLM_TYPE_744B_A40B: return "744B.A40B";
case LLM_TYPE_E2B: return "E2B";
case LLM_TYPE_E4B: return "E4B";
return LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED;
}
+// Maps the GGUF `<arch>.hidden_activation` string to the FFN op type used by the
+// graph builders. Only gated activations that map cleanly to llm_ffn_op_type are
+// listed; unrecognized values fall back to GeGLU, which matches the historical
+// default for ModernBert-style architectures.
+static const std::map<std::string, llm_ffn_op_type> LLM_FFN_OP_TYPES_FROM_STRING = {
+ { "gelu", LLM_FFN_GEGLU },
+ { "geglu", LLM_FFN_GEGLU },
+ { "silu", LLM_FFN_SWIGLU },
+ { "swish", LLM_FFN_SWIGLU },
+ { "swiglu", LLM_FFN_SWIGLU },
+ { "relu", LLM_FFN_RELU },
+ { "reglu", LLM_FFN_REGLU },
+};
+
+llm_ffn_op_type llm_ffn_op_type_from_string(const std::string & name, llm_ffn_op_type fallback) {
+ const auto it = LLM_FFN_OP_TYPES_FROM_STRING.find(name);
+ if (it != LLM_FFN_OP_TYPES_FROM_STRING.end()) {
+ return it->second;
+ }
+ return fallback;
+}
+
// CPU: ACCEL -> GPU host -> CPU extra -> CPU
static buft_list_t make_cpu_buft_list(const std::vector<llama_device> & devices, bool use_extra_bufts, bool no_host) {
buft_list_t buft_list;
ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT, hparams.n_embd_out_impl, false);
ml.get_key(LLM_KV_ATTENTION_CAUSAL, hparams.causal_attn, false);
ml.get_key(LLM_KV_POOLING_TYPE, hparams.pooling_type, false);
- ml.get_key(LLM_KV_BLOCK_COUNT, hparams.n_layer);
+ ml.get_key(LLM_KV_BLOCK_COUNT, hparams.n_layer_all);
ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert, false);
ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false);
ml.get_key(LLM_KV_EXPERT_GROUP_COUNT, hparams.n_expert_groups, false);
std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0);
std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);
std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0);
- std::fill(
- hparams.recurrent_layer_arr.begin(),
- hparams.recurrent_layer_arr.end(),
- llm_arch_is_recurrent(ml.get_arch()));
std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);
- std::fill(hparams.swa_layers.begin(), hparams.swa_layers.end(), 0);
+ std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0);
+ std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0);
std::fill(hparams.xielu_alpha_n.begin(), hparams.xielu_alpha_n.end(), 0.0f);
std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f);
- std::fill(hparams.xielu_beta.begin(), hparams.xielu_beta.end(), 0.0f);
- std::fill(hparams.xielu_eps.begin(), hparams.xielu_eps.end(), 0.0f);
+ std::fill(hparams.xielu_beta.begin(), hparams.xielu_beta.end(), 0.0f);
+ std::fill(hparams.xielu_eps.begin(), hparams.xielu_eps.end(), 0.0f);
+
std::fill(hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.end(), 0.0f);
std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f);
- ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, hparams.n_layer, false);
- ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer, false);
+ ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, hparams.n_layer(), false);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer(), false);
+
+ // Populate deepstack_mapping_arr - initialized to -1 (no deepstack)
+ std::fill(hparams.deepstack_mapping_arr.begin(), hparams.deepstack_mapping_arr.end(), -1);
// n_head_kv is optional, default to n_head
hparams.n_head_kv_arr = hparams.n_head_arr;
- ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT_KV, hparams.n_head_kv_arr, hparams.n_layer, false);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT_KV, hparams.n_head_kv_arr, hparams.n_layer(), false);
bool rope_finetuned = false;
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
const auto & use_mlock = params.use_mlock;
const auto & tensor_split = params.tensor_split;
- const int n_layer = hparams.n_layer;
+ const int n_layer_all = hparams.n_layer_all;
const int n_gpu_layers = this->n_gpu_layers();
const bool use_mmap_buffer = true;
splits[i] /= split_sum;
}
- const int i_gpu_start = std::max(int(hparams.n_layer) + 1 - n_gpu_layers, 0);
- const int act_gpu_layers = devices.empty() ? 0 : std::min(n_gpu_layers, int(n_layer) + 1);
+ const int i_gpu_start = std::max(n_layer_all + 1 - n_gpu_layers, 0);
+ const int act_gpu_layers = devices.empty() ? 0 : std::min(n_gpu_layers, n_layer_all + 1);
auto get_layer_buft_list = [&](int il) -> llama_model::impl::layer_dev {
- const bool is_swa = il < int(hparams.n_layer) && hparams.is_swa(il);
+ const bool is_swa = il < n_layer_all && hparams.is_swa(il);
if (il < i_gpu_start || (il - i_gpu_start) >= act_gpu_layers) {
LLAMA_LOG_DEBUG("load_tensors: layer %3d assigned to device %s, is_swa = %d\n", il, ggml_backend_dev_name(cpu_dev), is_swa);
return {cpu_dev, &pimpl->cpu_buft_list};
pimpl->dev_input = { cpu_dev, &pimpl->cpu_buft_list };
// assign the repeating layers to the devices according to the splits
- pimpl->dev_layer.resize(n_layer);
- for (int il = 0; il < n_layer; ++il) {
+ pimpl->dev_layer.resize(n_layer_all);
+ for (int il = 0; il < n_layer_all; ++il) {
pimpl->dev_layer[il] = get_layer_buft_list(il);
}
// assign the output layer
- pimpl->dev_output = get_layer_buft_list(n_layer);
+ pimpl->dev_output = get_layer_buft_list(n_layer_all);
const auto TENSOR_NOT_REQUIRED = llama_model_loader::TENSOR_NOT_REQUIRED;
throw std::runtime_error("model has expert layers but no expert layers are used");
}
- layers.resize(n_layer);
+ layers.resize(n_layer_all);
// call the per-model loading function
load_arch_tensors(ml);
// generic pass: load optional per-tensor/per-expert ".scale" tensors (e.g. NVFP4 scale2)
// this avoids having to add scale loading to every architecture
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
// attention weight scales (per-tensor, shape {1})
}
if (llama_supports_gpu_offload()) {
- const int n_gpu = std::min(n_gpu_layers, int(hparams.n_layer));
+ const int n_gpu = std::min(n_gpu_layers, n_layer_all);
int n_repeating = n_gpu;
if (n_repeating > 0) {
}
LLAMA_LOG_INFO("%s: offloading %d repeating layers to GPU\n", __func__, n_repeating);
- const int max_backend_supported_layers = hparams.n_layer + 1;
- const int max_offloadable_layers = hparams.n_layer + 1;
+ const int max_backend_supported_layers = n_layer_all + 1;
+ const int max_offloadable_layers = n_layer_all + 1;
LLAMA_LOG_INFO("%s: offloaded %d/%d layers to GPU\n", __func__, std::min(n_gpu_layers, max_offloadable_layers), max_backend_supported_layers);
}
}
uint32_t llama_model::n_gpu_layers() const {
- return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer + 1;
+ // note: plus 1 for the "output" layer
+ return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer_all + 1;
}
llama_split_mode llama_model::split_mode() const {
void llama_model::print_info() const {
const std::string rope_scaling_type = llama_rope_scaling_type_name(hparams.rope_scaling_type_train);
- auto print_f = [](const std::function<uint32_t(uint32_t)> & f, uint32_t n) {
+ auto print_f = [](const std::function<int32_t(uint32_t)> & f, uint32_t n) {
bool is_var = false;
- std::vector<uint32_t> v;
+ std::vector<int32_t> v;
for (uint32_t i = 0; i < n; ++i) {
v.push_back(f(i));
if (v[i] != v[0]) {
if (!hparams.vocab_only) {
LLAMA_LOG_INFO("%s: n_ctx_train = %u\n", __func__, hparams.n_ctx_train);
- LLAMA_LOG_INFO("%s: n_embd = %u\n", __func__, hparams.n_embd);
LLAMA_LOG_INFO("%s: n_embd_inp = %u\n", __func__, hparams.n_embd_inp());
- LLAMA_LOG_INFO("%s: n_layer = %u\n", __func__, hparams.n_layer);
- LLAMA_LOG_INFO("%s: n_head = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_head(il); }, hparams.n_layer).c_str());
- LLAMA_LOG_INFO("%s: n_head_kv = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_head_kv(il); }, hparams.n_layer).c_str());
+ LLAMA_LOG_INFO("%s: n_embd = %u\n", __func__, hparams.n_embd);
+ LLAMA_LOG_INFO("%s: n_embd_out = %u\n", __func__, hparams.n_embd_out());
+ LLAMA_LOG_INFO("%s: n_layer = %u\n", __func__, hparams.n_layer());
+ LLAMA_LOG_INFO("%s: n_layer_all = %u\n", __func__, hparams.n_layer_all);
+ LLAMA_LOG_INFO("%s: n_head = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_head(il); }, hparams.n_layer_all).c_str());
+ LLAMA_LOG_INFO("%s: n_head_kv = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_head_kv(il); }, hparams.n_layer_all).c_str());
LLAMA_LOG_INFO("%s: n_rot = %u\n", __func__, hparams.n_rot_full);
LLAMA_LOG_INFO("%s: n_swa = %u\n", __func__, hparams.n_swa);
LLAMA_LOG_INFO("%s: is_swa_any = %u\n", __func__, hparams.is_swa_any());
LLAMA_LOG_INFO("%s: n_embd_head_k = %u\n", __func__, hparams.n_embd_head_k_full);
LLAMA_LOG_INFO("%s: n_embd_head_v = %u\n", __func__, hparams.n_embd_head_v_full);
- LLAMA_LOG_INFO("%s: n_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_gqa(il); }, hparams.n_layer).c_str());
- LLAMA_LOG_INFO("%s: n_embd_k_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_embd_k_gqa(il); }, hparams.n_layer).c_str());
- LLAMA_LOG_INFO("%s: n_embd_v_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_embd_v_gqa(il); }, hparams.n_layer).c_str());
+ LLAMA_LOG_INFO("%s: n_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_gqa(il); }, hparams.n_layer_all).c_str());
+ LLAMA_LOG_INFO("%s: n_embd_k_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_embd_k_gqa(il); }, hparams.n_layer_all).c_str());
+ LLAMA_LOG_INFO("%s: n_embd_v_gqa = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_embd_v_gqa(il); }, hparams.n_layer_all).c_str());
LLAMA_LOG_INFO("%s: f_norm_eps = %.1e\n", __func__, hparams.f_norm_eps);
LLAMA_LOG_INFO("%s: f_norm_rms_eps = %.1e\n", __func__, hparams.f_norm_rms_eps);
LLAMA_LOG_INFO("%s: f_clamp_kqv = %.1e\n", __func__, hparams.f_clamp_kqv);
LLAMA_LOG_INFO("%s: f_logit_scale = %.1e\n", __func__, hparams.f_logit_scale);
LLAMA_LOG_INFO("%s: f_attn_scale = %.1e\n", __func__, hparams.f_attention_scale);
LLAMA_LOG_INFO("%s: f_attn_value_scale = %.4f\n", __func__, hparams.f_attn_value_scale);
- LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer).c_str());
+ LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer_all).c_str());
LLAMA_LOG_INFO("%s: n_expert = %u\n", __func__, hparams.n_expert);
LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used);
LLAMA_LOG_INFO("%s: n_expert_groups = %d\n", __func__, hparams.n_expert_groups);
LLAMA_LOG_INFO("%s: n_ctx_orig_yarn = %u\n", __func__, hparams.n_ctx_orig_yarn);
LLAMA_LOG_INFO("%s: rope_yarn_log_mul = %.4f\n", __func__, hparams.rope_yarn_log_mul);
LLAMA_LOG_INFO("%s: rope_finetuned = %s\n", __func__, hparams.rope_finetuned ? "yes" : "unknown");
+ if (arch == LLM_ARCH_GRANITE &&
+ std::any_of(hparams.deepstack_mapping_arr.begin(),
+ hparams.deepstack_mapping_arr.end(),
+ [](const auto & entry) { return entry >= 0; })) {
+ LLAMA_LOG_INFO("%s: deepstack_mapping_arr = %s\n", __func__,
+ print_f([&](uint32_t il) { return hparams.deepstack_mapping_arr[il]; },
+ hparams.n_layer_all).c_str());
+ }
// MRoPE (Multi-axis Rotary Position Embedding) sections
if (const auto & s = hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) {
LLAMA_LOG_INFO("%s: mrope sections = [%d, %d, %d, %d]\n", __func__, s[0], s[1], s[2], s[3]);
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
}
- if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4) {
+ if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4) {
LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead);
LLAMA_LOG_INFO("%s: n_lora_q = %d\n", __func__, hparams.n_lora_q);
LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv);
LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp);
}
- if (arch == LLM_ARCH_QWEN3MOE || arch == LLM_ARCH_OPENAI_MOE || arch == LLM_ARCH_QWEN3VLMOE || arch == LLM_ARCH_RND1) {
+ if (arch == LLM_ARCH_MELLUM ||
+ arch == LLM_ARCH_QWEN3MOE ||
+ arch == LLM_ARCH_OPENAI_MOE ||
+ arch == LLM_ARCH_QWEN3VLMOE ||
+ arch == LLM_ARCH_RND1) {
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
}
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm);
LLAMA_LOG_INFO("%s: expert_gating_func = %s\n", __func__, llama_expert_gating_func_name((llama_expert_gating_func_type) hparams.expert_gating_func));
- LLAMA_LOG_INFO("%s: nextn_predict_layers = %d\n", __func__, hparams.nextn_predict_layers);
+ LLAMA_LOG_INFO("%s: n_layer_nextn = %d\n", __func__, hparams.n_layer_nextn);
}
if (arch == LLM_ARCH_SMALLTHINKER || arch == LLM_ARCH_LFM2MOE) {
{
res = nullptr;
} break;
+ case LLM_ARCH_DEEPSEEK32:
+ {
+ res = new llama_kv_cache_dsa(
+ *this,
+ params.type_k,
+ params.type_v,
+ !cparams.flash_attn,
+ cparams.offload_kqv,
+ cparams.kv_unified,
+ cparams.n_ctx_seq,
+ cparams.n_seq_max,
+ 1,
+ hparams.n_swa,
+ hparams.swa_type,
+ nullptr,
+ nullptr);
+ } break;
// Models that need standard caching should rely on recurrent/hybrid
// checks
default:
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
llama_memory_hybrid::layer_filter_cb filter_recr = nullptr;
if (arch == LLM_ARCH_FALCON_H1) {
- filter_attn = [&](int32_t) { return true; };
- filter_recr = [&](int32_t) { return true; };
+ filter_attn = [&](uint32_t) { return true; };
+ filter_recr = [&](uint32_t) { return true; };
} else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
- filter_attn = [&](int32_t il) {
- return !hparams.is_recurrent(il) && hparams.n_ff(il) == 0;
+ filter_attn = [&](uint32_t il) {
+ return !hparams.is_recr(il) && hparams.n_ff(il) == 0;
};
- filter_recr = [&](int32_t il) {
- return hparams.is_recurrent(il) && hparams.n_ff(il) == 0;
+ filter_recr = [&](uint32_t il) {
+ return hparams.is_recr(il) && hparams.n_ff(il) == 0;
};
} else if (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
- const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
- filter_attn = [&, n_main](int32_t il) {
- return (uint32_t)il < n_main && !hparams.is_recurrent(il);
+ filter_attn = [&](uint32_t il) {
+ return il < hparams.n_layer() && !hparams.is_recr(il);
};
- filter_recr = [&, n_main](int32_t il) {
- return (uint32_t)il < n_main && hparams.is_recurrent(il);
+ filter_recr = [&](uint32_t il) {
+ return il < hparams.n_layer() && hparams.is_recr(il);
};
}
/* filter_recr */ std::move(filter_recr));
}
} else {
- llama_memory_i::layer_reuse_cb reuse = nullptr;
llama_kv_cache::layer_filter_cb filter = nullptr;
+ llama_memory_i::layer_reuse_cb reuse = nullptr;
+ llama_kv_cache::layer_share_cb share = nullptr;
if (arch == LLM_ARCH_GEMMA3N || arch == LLM_ARCH_GEMMA4) {
- reuse = [&](int32_t il) {
- if (il >= (int32_t) hparams.n_layer_kv_from_start) {
- return (int32_t) hparams.n_layer_kv_from_start - (hparams.is_swa(il) ? 2 : 1);
+ reuse = [&](uint32_t il) {
+ GGML_ASSERT(hparams.n_layer_kv_from_start >= 2);
+
+ if (il >= (uint32_t)hparams.n_layer_kv_from_start) {
+ return hparams.n_layer_kv_from_start - (hparams.is_swa(il) ? 2 : 1);
}
return -1;
}
if (mtp_on_hybrid_qwen35) {
- const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
- filter = [n_main](int32_t il) { return (uint32_t)il >= n_main; };
+ filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
+ }
+
+ if (arch == LLM_ARCH_STEP35 && hparams.n_layer_nextn > 0) {
+ if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
+ filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
+ } else {
+ filter = [&](uint32_t il) { return il < hparams.n_layer(); };
+ }
}
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
GGML_ASSERT(hparams.is_swa_any());
- res = new llama_kv_cache_iswa(
- *this,
- params.type_k,
- params.type_v,
- !cparams.flash_attn,
- cparams.offload_kqv,
- params.swa_full,
- cparams.kv_unified,
- cparams.n_ctx_seq,
- cparams.n_seq_max,
- cparams.n_ubatch,
- 1,
- filter,
- reuse);
+ if (arch == LLM_ARCH_GEMMA4_ASSISTANT) {
+ llama_memory_t mem_other = llama_get_memory(cparams.ctx_other);
+
+ share = [&](int32_t il) {
+ const llama_model * model_other = llama_get_model(cparams.ctx_other);
+
+ if (hparams.is_swa(il)) {
+ return llama_model_n_layer(model_other) - 2;
+ }
+
+ return llama_model_n_layer(model_other) - 1;
+ };
+
+ res = new llama_kv_cache_iswa(
+ *this,
+ params.type_k,
+ params.type_v,
+ !cparams.flash_attn,
+ cparams.offload_kqv,
+ params.swa_full,
+ cparams.kv_unified,
+ cparams.n_ctx_seq,
+ cparams.n_seq_max,
+ cparams.n_ubatch,
+ 1,
+ mem_other,
+ filter,
+ reuse,
+ share);
+ } else {
+ res = new llama_kv_cache_iswa(
+ *this,
+ params.type_k,
+ params.type_v,
+ !cparams.flash_attn,
+ cparams.offload_kqv,
+ params.swa_full,
+ cparams.kv_unified,
+ cparams.n_ctx_seq,
+ cparams.n_seq_max,
+ cparams.n_ubatch,
+ 1,
+ nullptr,
+ filter,
+ reuse,
+ share);
+ }
} else {
GGML_ASSERT(!hparams.is_swa_any());
res = new llama_kv_cache(
*this,
+ hparams,
params.type_k,
params.type_v,
!cparams.flash_attn,
1,
hparams.n_swa,
hparams.swa_type,
+ nullptr,
filter,
+ nullptr,
nullptr);
}
}
}
int32_t llama_model_n_layer(const llama_model * model) {
- return model->hparams.n_layer;
+ return model->hparams.n_layer();
}
int32_t llama_model_n_head(const llama_model * model) {
case LLM_ARCH_DEEPSEEK:
case LLM_ARCH_DEEPSEEK2:
case LLM_ARCH_DEEPSEEK2OCR:
+ case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_PLM:
case LLM_ARCH_CHATGLM:
case LLM_ARCH_GRANITE:
case LLM_ARCH_GEMMA3:
case LLM_ARCH_GEMMA3N:
case LLM_ARCH_GEMMA4:
+ case LLM_ARCH_GEMMA4_ASSISTANT:
case LLM_ARCH_GEMMA_EMBEDDING:
case LLM_ARCH_STARCODER2:
case LLM_ARCH_OPENELM:
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
case LLM_ARCH_TALKIE:
+ case LLM_ARCH_MELLUM:
return LLAMA_ROPE_TYPE_NEOX;
case LLM_ARCH_QWEN2VL:
LLM_TYPE_A13B,
LLM_TYPE_7B_A1B,
LLM_TYPE_8B_A1B, // lfm2moe
+ LLM_TYPE_12B_A2_5B,
LLM_TYPE_16B_A1B,
LLM_TYPE_21B_A3B, // Ernie MoE small
LLM_TYPE_24B_A2B, // lfm2moe
LLM_TYPE_310B_A15B, // /MiMo-V2-Flash
LLM_TYPE_355B_A32B, // GLM-4.5
LLM_TYPE_397B_A17B, // Qwen3.5
+ LLM_TYPE_685B_A37B, // DeepSeek V3.2
LLM_TYPE_744B_A40B, // GLM-5
LLM_TYPE_E2B,
LLM_TYPE_E4B,
std::string llama_rope_scaling_type_name(llama_rope_scaling_type rope_scaling_type);
+// Map a GGUF activation-name string to llm_ffn_op_type. Returns `fallback` if
+// the string is empty or not recognized.
+llm_ffn_op_type llm_ffn_op_type_from_string(const std::string & name, llm_ffn_op_type fallback);
+
struct llama_layer_posnet {
// resnet
struct ggml_tensor * norm1 = nullptr;
struct ggml_tensor * output_s = nullptr;
struct ggml_tensor * output_in_s = nullptr;
+ // NextN/MTP model-level projections
+ struct ggml_tensor * nextn_proj_pre = nullptr;
+ struct ggml_tensor * nextn_proj_post = nullptr;
+
// classifier
struct ggml_tensor * cls = nullptr;
struct ggml_tensor * cls_b = nullptr;
// convenience macro for loading local variables for load_tensors() in llama_model_base
// note: cast to int64_t since we will use these for the tensor dimensions
#define LLAMA_LOAD_LOCALS \
- const int n_layer = hparams.n_layer; GGML_UNUSED(n_layer); \
+ const int n_layer = hparams.n_layer(); GGML_UNUSED(n_layer); \
+ const int n_layer_all = hparams.n_layer_all; GGML_UNUSED(n_layer_all); \
+ const int n_layer_nextn = hparams.n_layer_nextn; GGML_UNUSED(n_layer_nextn); \
const int64_t n_head = hparams.n_head(); GGML_UNUSED(n_head); \
const int64_t n_head_kv = hparams.n_head_kv(); GGML_UNUSED(n_head_kv); \
const int64_t n_embd = hparams.n_embd; GGML_UNUSED(n_embd); \
qs.has_tied_embeddings = false;
}
}
- qs.n_ffn_down = qs.n_ffn_gate = qs.n_ffn_up = (int)qs.model.hparams.n_layer;
+ qs.n_ffn_down = qs.n_ffn_gate = qs.n_ffn_up = (int)qs.model.hparams.n_layer();
}
//
model->hparams.n_embd = desc->n_embd;
model->hparams.n_embd_head_k_full = desc->n_embd_head_k;
model->hparams.n_embd_head_v_full = desc->n_embd_head_v;
- model->hparams.n_layer = desc->n_layer;
+ model->hparams.n_layer_all = desc->n_layer;
model->hparams.n_expert = desc->n_expert;
for (uint32_t i = 0; i < desc->n_layer; i++) {
case LLAMA_VOCAB_PRE_TYPE_CODESHELL:
case LLAMA_VOCAB_PRE_TYPE_EXAONE:
case LLAMA_VOCAB_PRE_TYPE_MINERVA:
+ case LLAMA_VOCAB_PRE_TYPE_MELLUM2:
regex_exprs = {
"\\p{N}",
"'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)",
"[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
};
break;
+ case LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI:
+ // Same lookaheads as GPT4O but with \p{M} added so combining marks
+ // (diacritics) attach to their base letters. Avoids excessive
+ // backtracking on scripts that use them heavily (Bengali, Hindi,
+ // Telugu, Thai, ...). See PR #22716 for benchmarks.
+ regex_exprs = {
+ "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+ };
+ break;
case LLAMA_VOCAB_PRE_TYPE_TINY_AYA:
regex_exprs = {
// original regex from tokenizer.json: "\\d{1,3}(?=(?:\\d{3})*\\b)"
"(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}+| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
};
break;
+ case LLAMA_VOCAB_PRE_TYPE_WHITESPACE:
+ // whitespace pre-tokenizer (jinaai/jina-embeddings-v2-base-zh)
+ regex_exprs = {
+ "\\S+",
+ };
+ byte_encode = false;
+ break;
default:
// default regex for BPE tokenization pre-processing
regex_exprs = {
void tokenize(const std::string & text, std::vector<llama_token> & output) {
// normalize and split by whitespace
- std::vector<std::string> words = preprocess(text);
+ std::vector<std::string> words = preprocess(text, vocab.get_normalizer_lowercase());
// bos token prepended already
// find the longest tokens that form the words
}
// TODO: reduce string copies by using cpts_offs array
- static std::vector<std::string> preprocess(const std::string & text) {
+ static std::vector<std::string> preprocess(const std::string & text, bool lowercase) {
const std::vector<uint32_t> cpts_nfd = unicode_cpts_normalize_nfd(unicode_cpts_from_utf8(text));
std::vector<std::string> words(1, "");
continue;
}
- const std::string s = unicode_cpt_to_utf8(unicode_tolower(cpt));
+ const std::string s = unicode_cpt_to_utf8(lowercase ? unicode_tolower(cpt) : cpt);
if (flags.is_punctuation || ( cpt < 0x7F && flags.is_symbol ) || is_chinese_char(cpt)) {
if (words.back().size()) { // finish previous word if any
words.emplace_back();
const llama_vocab & vocab;
};
+struct llm_tokenizer_whitespace_session : llm_tokenizer_bpe_session {
+ llm_tokenizer_whitespace_session(const llama_vocab & vocab, const llm_tokenizer_bpe & tokenizer) : llm_tokenizer_bpe_session{vocab, tokenizer}, vocab{vocab} {}
+
+ void tokenize(const std::string & text, std::vector<llama_token> & output) override {
+ const bool lowercase = vocab.get_normalizer_lowercase();
+
+ std::string segment;
+ auto flush = [&]() {
+ if (!segment.empty()) {
+ llm_tokenizer_bpe_session::tokenize(segment, output);
+ segment.clear();
+ }
+ };
+
+ for (uint32_t cpt : unicode_cpts_from_utf8(text)) {
+ // drop whitespace
+ if (unicode_cpt_flags_from_cpt(cpt).is_whitespace) {
+ flush();
+ } else {
+ segment += unicode_cpt_to_utf8(lowercase ? unicode_tolower(cpt) : cpt);
+ }
+ }
+ flush();
+ }
+
+private:
+ const llama_vocab & vocab;
+};
+
//
// impl
//
bool remove_extra_whitespaces = false;
bool escape_whitespaces = true;
bool treat_whitespace_as_suffix = false;
+ bool normalizer_lowercase = true; // Lowercase normalizer (tokenizer.json)
std::unordered_map<std::string, llama_token> token_to_id;
std::vector<token_data> id_to_token;
// set of all tokens that cause "end of generation"
std::set<llama_token> special_eog_ids;
+ std::vector<llama_token> suppress_tokens;
+
std::unique_ptr<llm_tokenizer> tokenizer;
std::vector<char> precompiled_charsmap;
special_mask_id = 103;
add_sep = true;
- } else if (tokenizer_model == "gpt2" || tokenizer_model == "hybriddna") {
+ } else if (tokenizer_model == "gpt2" || tokenizer_model == "hybriddna" || tokenizer_model == "whitespace") {
type = LLAMA_VOCAB_TYPE_BPE;
// read bpe merges and populate bpe ranks
tokenizer_pre == "jais-2") {
pre_type = LLAMA_VOCAB_PRE_TYPE_JAIS2;
} else if (
- tokenizer_pre == "gemma4") {
+ tokenizer_pre == "gemma4" ||
+ tokenizer_pre == "granite-embed-multi-311m") {
pre_type = LLAMA_VOCAB_PRE_TYPE_GEMMA4;
escape_whitespaces = true;
} else if (
tokenizer_pre == "roberta-bpe") {
pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2;
add_sep = true;
+ } else if (
+ tokenizer_pre == "whitespace") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_WHITESPACE;
+ normalizer_lowercase = false;
} else if (
tokenizer_pre == "refact") {
pre_type = LLAMA_VOCAB_PRE_TYPE_REFACT;
tokenizer_pre == "talkie") {
pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O;
clean_spaces = false;
+ } else if (
+ tokenizer_pre == "granite-embed-multi-97m") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI;
+ clean_spaces = false;
+ ignore_merges = true;
} else if (
tokenizer_pre == "tiny_aya") {
pre_type = LLAMA_VOCAB_PRE_TYPE_TINY_AYA;
tokenizer_pre == "solar-open") {
pre_type = LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN;
clean_spaces = false;
+ } else if (
+ tokenizer_pre == "mellum2") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_MELLUM2;
} else {
throw std::runtime_error(format("unknown pre-tokenizer type: '%s'", tokenizer_pre.c_str()));
}
}
}
+ // Lowercase normalizer flag (consulted by WPM / whitespace BPE)
+ ml.get_key(LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE, normalizer_lowercase, false);
+
+ // suppress tokens
+ {
+ const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str());
+ if (suppress_idx != -1) {
+ const int n = gguf_get_arr_n(ctx, suppress_idx);
+ const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
+ suppress_tokens.assign(data, data + n);
+ }
+ }
+
// auto-detect special tokens by text
// TODO: convert scripts should provide these tokens through the KV metadata LLM_KV_TOKENIZER_...
// for now, we apply this workaround to find the tokens based on their text
std::unique_ptr<llm_tokenizer_bpe_session> session;
if (vocab.get_tokenizer_model() == "hybriddna") {
session = std::make_unique<llm_tokenizer_hybriddna_session>(vocab, *tok_bpe);
+ } else if (vocab.get_tokenizer_model() == "whitespace") {
+ session = std::make_unique<llm_tokenizer_whitespace_session>(vocab, *tok_bpe);
} else {
session = std::make_unique<llm_tokenizer_bpe_session>(vocab, *tok_bpe);
}
return pimpl->treat_whitespace_as_suffix;
}
+bool llama_vocab::get_normalizer_lowercase() const {
+ return pimpl->normalizer_lowercase;
+}
+
+const std::vector<llama_token> & llama_vocab::get_suppress_tokens() const {
+ return pimpl->suppress_tokens;
+}
+
int llama_vocab::max_token_len() const {
return pimpl->max_token_len;
}
// pre-tokenization types
enum llama_vocab_pre_type {
- LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0,
- LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1,
- LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2,
- LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3,
- LLAMA_VOCAB_PRE_TYPE_FALCON = 4,
- LLAMA_VOCAB_PRE_TYPE_MPT = 5,
- LLAMA_VOCAB_PRE_TYPE_STARCODER = 6,
- LLAMA_VOCAB_PRE_TYPE_GPT2 = 7,
- LLAMA_VOCAB_PRE_TYPE_REFACT = 8,
- LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9,
- LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10,
- LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11,
- LLAMA_VOCAB_PRE_TYPE_OLMO = 12,
- LLAMA_VOCAB_PRE_TYPE_DBRX = 13,
- LLAMA_VOCAB_PRE_TYPE_SMAUG = 14,
- LLAMA_VOCAB_PRE_TYPE_PORO = 15,
- LLAMA_VOCAB_PRE_TYPE_CHATGLM3 = 16,
- LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17,
- LLAMA_VOCAB_PRE_TYPE_VIKING = 18,
- LLAMA_VOCAB_PRE_TYPE_JAIS = 19,
- LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20,
- LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21,
- LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22,
- LLAMA_VOCAB_PRE_TYPE_BLOOM = 23,
- LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24,
- LLAMA_VOCAB_PRE_TYPE_EXAONE = 25,
- LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26,
- LLAMA_VOCAB_PRE_TYPE_MINERVA = 27,
- LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28,
- LLAMA_VOCAB_PRE_TYPE_GPT4O = 29,
- LLAMA_VOCAB_PRE_TYPE_SUPERBPE = 30,
- LLAMA_VOCAB_PRE_TYPE_TRILLION = 31,
- LLAMA_VOCAB_PRE_TYPE_BAILINGMOE = 32,
- LLAMA_VOCAB_PRE_TYPE_LLAMA4 = 33,
- LLAMA_VOCAB_PRE_TYPE_PIXTRAL = 34,
- LLAMA_VOCAB_PRE_TYPE_SEED_CODER = 35,
- LLAMA_VOCAB_PRE_TYPE_HUNYUAN = 36,
- LLAMA_VOCAB_PRE_TYPE_KIMI_K2 = 37,
- LLAMA_VOCAB_PRE_TYPE_HUNYUAN_DENSE = 38,
- LLAMA_VOCAB_PRE_TYPE_GROK_2 = 39,
- LLAMA_VOCAB_PRE_TYPE_GRANITE_DOCLING = 40,
- LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2 = 41,
- LLAMA_VOCAB_PRE_TYPE_AFMOE = 42,
- LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN = 43,
- LLAMA_VOCAB_PRE_TYPE_YOUTU = 44,
- LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE = 45,
- LLAMA_VOCAB_PRE_TYPE_QWEN35 = 46,
- LLAMA_VOCAB_PRE_TYPE_TINY_AYA = 47,
- LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM = 48,
- LLAMA_VOCAB_PRE_TYPE_JAIS2 = 49,
- LLAMA_VOCAB_PRE_TYPE_GEMMA4 = 50,
- LLAMA_VOCAB_PRE_TYPE_SARVAM_MOE = 51,
- LLAMA_VOCAB_PRE_TYPE_MINICPM5 = 52,
+ LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0,
+ LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1,
+ LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2,
+ LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3,
+ LLAMA_VOCAB_PRE_TYPE_FALCON = 4,
+ LLAMA_VOCAB_PRE_TYPE_MPT = 5,
+ LLAMA_VOCAB_PRE_TYPE_STARCODER = 6,
+ LLAMA_VOCAB_PRE_TYPE_GPT2 = 7,
+ LLAMA_VOCAB_PRE_TYPE_REFACT = 8,
+ LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9,
+ LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10,
+ LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11,
+ LLAMA_VOCAB_PRE_TYPE_OLMO = 12,
+ LLAMA_VOCAB_PRE_TYPE_DBRX = 13,
+ LLAMA_VOCAB_PRE_TYPE_SMAUG = 14,
+ LLAMA_VOCAB_PRE_TYPE_PORO = 15,
+ LLAMA_VOCAB_PRE_TYPE_CHATGLM3 = 16,
+ LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17,
+ LLAMA_VOCAB_PRE_TYPE_VIKING = 18,
+ LLAMA_VOCAB_PRE_TYPE_JAIS = 19,
+ LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20,
+ LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21,
+ LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22,
+ LLAMA_VOCAB_PRE_TYPE_BLOOM = 23,
+ LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24,
+ LLAMA_VOCAB_PRE_TYPE_EXAONE = 25,
+ LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26,
+ LLAMA_VOCAB_PRE_TYPE_MINERVA = 27,
+ LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28,
+ LLAMA_VOCAB_PRE_TYPE_GPT4O = 29,
+ LLAMA_VOCAB_PRE_TYPE_SUPERBPE = 30,
+ LLAMA_VOCAB_PRE_TYPE_TRILLION = 31,
+ LLAMA_VOCAB_PRE_TYPE_BAILINGMOE = 32,
+ LLAMA_VOCAB_PRE_TYPE_LLAMA4 = 33,
+ LLAMA_VOCAB_PRE_TYPE_PIXTRAL = 34,
+ LLAMA_VOCAB_PRE_TYPE_SEED_CODER = 35,
+ LLAMA_VOCAB_PRE_TYPE_HUNYUAN = 36,
+ LLAMA_VOCAB_PRE_TYPE_KIMI_K2 = 37,
+ LLAMA_VOCAB_PRE_TYPE_HUNYUAN_DENSE = 38,
+ LLAMA_VOCAB_PRE_TYPE_GROK_2 = 39,
+ LLAMA_VOCAB_PRE_TYPE_GRANITE_DOCLING = 40,
+ LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2 = 41,
+ LLAMA_VOCAB_PRE_TYPE_AFMOE = 42,
+ LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN = 43,
+ LLAMA_VOCAB_PRE_TYPE_YOUTU = 44,
+ LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE = 45,
+ LLAMA_VOCAB_PRE_TYPE_QWEN35 = 46,
+ LLAMA_VOCAB_PRE_TYPE_TINY_AYA = 47,
+ LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM = 48,
+ LLAMA_VOCAB_PRE_TYPE_JAIS2 = 49,
+ LLAMA_VOCAB_PRE_TYPE_GEMMA4 = 50,
+ LLAMA_VOCAB_PRE_TYPE_SARVAM_MOE = 51,
+ LLAMA_VOCAB_PRE_TYPE_MINICPM5 = 52,
+ LLAMA_VOCAB_PRE_TYPE_WHITESPACE = 53,
+ LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54,
+ LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55,
};
struct LLM_KV;
bool get_remove_extra_whitespaces () const;
bool get_escape_whitespaces () const;
bool get_treat_whitespace_as_suffix() const;
+ bool get_normalizer_lowercase () const;
+
+ const std::vector<llama_token> & get_suppress_tokens() const;
int max_token_len() const;
}
case GGML_BACKEND_DEVICE_TYPE_IGPU:
- igpus.push_back({false, dev});
+ if (igpus.empty()) {
+ igpus.push_back({false, dev});
+ }
break;
case GGML_BACKEND_DEVICE_TYPE_META:
GGML_ABORT("fatal error");
// add GPUs
model->devices.insert(model->devices.end(), gpus.begin(), gpus.end());
- // add integrated GPUs only if no other devices were found
- if (model->devices.empty()) {
+ // add integrated GPUs only if no discrete GPUs were found
+ // (RPC servers do not count, otherwise the local iGPU would be dropped on iGPU+RPC setups)
+ if (gpus.empty()) {
model->devices.insert(model->devices.end(), igpus.begin(), igpus.end());
}
}
uint32_t n_ubatch; // physical maximum batch size
uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
+ uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch)
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
// note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init)
struct llama_sampler_seq_config * samplers;
size_t n_samplers;
+
+ // a source/target/parent context
+ // can be utilized in various ways, for example by sharing results or llama_memory between 2 contexts
+ struct llama_context * ctx_other;
};
struct llama_model_tensor_override {
// Set whether the model is in warmup mode or not
// If true, all model tensors are activated during llama_decode() to load and cache their weights.
- LLAMA_API void llama_set_warmup(struct llama_context * ctx, bool warmup);
+ //
+ // note: using this can cause extra graph reallocations because it changes the graph topology with MoE models,
+ // so it is generally not recommended to use in practice. will be removed in the future
+ DEPRECATED(LLAMA_API void llama_set_warmup(struct llama_context * ctx, bool warmup),
+ "user code should do warmup runs manually [TAG_LLAMA_GRAPH_NO_WARMUP]");
// Set abort callback
LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data);
hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 56: type = LLM_TYPE_6B; break;
case 32: type = LLM_TYPE_26B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_apertus::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- ml.get_key_or_arr(LLM_KV_XIELU_ALPHA_N, hparams.xielu_alpha_n, hparams.n_layer);
- ml.get_key_or_arr(LLM_KV_XIELU_ALPHA_P, hparams.xielu_alpha_p, hparams.n_layer);
- ml.get_key_or_arr(LLM_KV_XIELU_BETA, hparams.xielu_beta, hparams.n_layer);
- ml.get_key_or_arr(LLM_KV_XIELU_EPS, hparams.xielu_eps, hparams.n_layer);
- switch (hparams.n_layer) {
+ ml.get_key_or_arr(LLM_KV_XIELU_ALPHA_N, hparams.xielu_alpha_n, hparams.n_layer());
+ ml.get_key_or_arr(LLM_KV_XIELU_ALPHA_P, hparams.xielu_alpha_p, hparams.n_layer());
+ ml.get_key_or_arr(LLM_KV_XIELU_BETA, hparams.xielu_beta, hparams.n_layer());
+ ml.get_key_or_arr(LLM_KV_XIELU_EPS, hparams.xielu_eps, hparams.n_layer());
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_8B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
// Arcee uses the same structure as Llama
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 36: type = LLM_TYPE_4B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
if (hparams.n_expert == 128) {
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 35: type = LLM_TYPE_10B_128x3_66B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_ATTENTION_GATE_LORA_RANK, hparams.n_lora_gate, false);
ml.get_key(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 12:
switch (hparams.n_embd) {
case 768: type = LLM_TYPE_190M; break;
void llama_model_baichuan::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 40: type = LLM_TYPE_13B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 28: type = LLM_TYPE_16B; break;
case 88: type = LLM_TYPE_290B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
- // TODO: when MTP is implemented, this should probably be updated if needed
- hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
+ GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_impl");
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 20: type = LLM_TYPE_16B_A1B; break;
- case 21: type = LLM_TYPE_16B_A1B; break;
case 32: type = LLM_TYPE_100B_A6B; break;
- case 33: type = LLM_TYPE_100B_A6B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
GGML_ASSERT(n_expert > 0 && "n_expert must be > 0 for bailingmoe2");
GGML_ASSERT(n_expert_used > 0 && "n_expert_used must be > 0 for bailingmoe2");
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
// skip all tensors in the NextN layers
flags |= TENSOR_SKIP;
}
}
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
ggml_tensor * inp_out_ids = build_inp_out_ids();
- const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
// norm
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
}
- if (il == n_transformer_layers - 1 && inp_out_ids) {
+ 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);
}
#include "models.h"
void llama_model_bert::load_arch_hparams(llama_model_loader & ml) {
- ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 3:
type = LLM_TYPE_17M; break; // bge-micro
case 6:
void llama_model_bitnet::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 26: type = LLM_TYPE_3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_bloom::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1B; break;
case 30:
switch (hparams.n_embd) {
hparams.f_norm_eps = 1e-5; // eps for qk-norm, torch default
ml.get_key(LLM_KV_SWIN_NORM, hparams.swin_norm, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 48: type = LLM_TYPE_34B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_chatglm::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 28: {
if (hparams.n_head(0) == 16) {
type = LLM_TYPE_1_5B;
void llama_model_codeshell::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 42: type = LLM_TYPE_7B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_cogvlm::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_13B; break;
default: type = LLM_TYPE_UNKNOWN;
}
uint32_t swa_period = 4;
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
hparams.set_swa_pattern(swa_period);
+
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_8B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_command_r::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 40: type = LLM_TYPE_35B; break;
default: type = LLM_TYPE_UNKNOWN;
}
#include "models.h"
void llama_model_dbrx::load_arch_hparams(llama_model_loader & ml) {
-ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
-ml.get_key(LLM_KV_ATTENTION_CLAMP_KQV, hparams.f_clamp_kqv);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
+ ml.get_key(LLM_KV_ATTENTION_CLAMP_KQV, hparams.f_clamp_kqv);
-switch (hparams.n_layer) {
- case 40: type = LLM_TYPE_16x12B; break;
- default: type = LLM_TYPE_UNKNOWN;
+ switch (hparams.n_layer()) {
+ case 40: type = LLM_TYPE_16x12B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
}
- }
void llama_model_dbrx::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
void llama_model_deci::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 80: type = LLM_TYPE_70B; break;
case 162: type = LLM_TYPE_405B; break;
ml.get_key(LLM_KV_VOCAB_SIZE, n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, n_vocab, false);
// lite variants include DeepSeek-V2-Lite, GigaChat3-10B-A1.8B, Kanana-2-30B-A3B
- const bool is_lite = (hparams.n_layer == 27 || hparams.n_layer == 26 || (hparams.n_layer == 48 && n_vocab == 128256));
+ const bool is_lite = (hparams.n_layer() == 27 || hparams.n_layer() == 26 || (hparams.n_layer() == 48 && n_vocab == 128256));
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
// for compatibility with existing DeepSeek V2 and V2.5 GGUFs
// that have no expert_gating_func model parameter set
- if ((hparams.n_layer == 47 || hparams.n_layer == 48) && n_vocab == 154880) {
+ if ((hparams.n_layer() == 47 || hparams.n_layer() == 48) && n_vocab == 154880) {
// GLM 4.7 Lite
hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
} else {
hparams.f_attn_temp_offset = 0.0f;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 27: type = LLM_TYPE_16B; break;
case 47: type = LLM_TYPE_30B_A3B; break;
case 60: type = LLM_TYPE_236B; break;
ggml_tensor * inp_out_ids = build_inp_out_ids();
- int effective_n_layers = hparams.n_layer - hparams.nextn_predict_layers;
- for (int il = 0; il < effective_n_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
// norm
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
}
}
- if (il == effective_n_layers - 1 && inp_out_ids) {
+ 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);
}
hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 12: type = LLM_TYPE_3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
--- /dev/null
+#include "models.h"
+
+#include "llama-kv-cache.h"
+#include "llama-kv-cache-dsa.h"
+
+void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ hparams.f_norm_eps = 1e-6; // eps for layer norm
+ ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
+
+ // MoE parameters
+ ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert);
+ ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used);
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+ ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
+
+ // deepseek MLA parameters
+ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
+ ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl, false);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl, false);
+ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+
+ // DSA parameters
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
+
+ // Expert gating function
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
+
+ if (ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, 0.0f)) {
+ // [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]
+ // cancel the factor from the convert script
+ hparams.rope_yarn_log_mul /= 0.1f;
+ }
+
+ // NextN/MTP parameters
+ 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 62: type = LLM_TYPE_685B_A37B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+ const bool is_mla = hparams.is_mla();
+ if (!is_mla) {
+ throw std::runtime_error("DEEPSEEK32 architecture requires MLA");
+ }
+
+ // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+
+ const int64_t n_embd_head_qk_rope = hparams.n_rot();
+ const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;
+
+ const int64_t q_lora_rank = hparams.n_lora_q;
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ // output
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ // try to load output.weight, if not found, use token_embd (tied embeddings)
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+ if (!output) {
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+ }
+
+ for (int i = 0; i < n_layer_all; ++i) {
+ int flags = 0;
+ if (i >= n_layer) {
+ // skip all tensors in the NextN layers
+ // TODO @ngxson : TENSOR_NOT_REQUIRED was a hack, need to remove it later
+ flags |= TENSOR_SKIP | TENSOR_NOT_REQUIRED;
+ }
+
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags);
+
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);
+ layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, flags);
+
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags);
+
+ // note: only old legacy GGUF files will have the unsplit wkv_b tensor in
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags);
+
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
+
+ // DSA indexer
+ layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags);
+ layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {hparams.indexer_head_size}, flags);
+ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags);
+ layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags);
+ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size}, flags);
+ if (i < (int) hparams.n_layer_dense_lead) {
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);
+ } else {
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+
+ if (n_expert == 0) {
+ throw std::runtime_error("n_expert must be > 0");
+ }
+ if (n_expert_used == 0) {
+ throw std::runtime_error("n_expert_used must be > 0");
+ }
+
+ // MoE branch
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);
+
+ // Shared expert branch
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
+ }
+
+ // NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
+ if (i >= n_layer) {
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);
+
+ // Optional tensors
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags | TENSOR_NOT_REQUIRED);
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_deepseek32::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context(params) {
+ const bool is_mla = hparams.is_mla();
+ GGML_ASSERT(is_mla);
+
+ // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
+ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();
+ GGML_UNUSED(n_embd_head_v);
+
+ const int64_t n_embd_head_qk_rope = hparams.n_rot();
+ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
+
+ const int64_t n_indexer_head = hparams.indexer_n_head;
+ const int64_t n_embd_indexer_head = hparams.indexer_head_size;
+ const int64_t n_embd_indexer_head_rope = hparams.n_rot();
+ const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope;
+ const uint32_t n_indexer_top_k = hparams.indexer_top_k;
+
+ const uint32_t kv_lora_rank = hparams.n_lora_kv;
+
+ // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
+ // See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.
+ // And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]
+
+ // first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor
+ GGML_ASSERT(ext_factor >= 0.0f);
+ const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));
+
+ // use the original attn_factor to pre-scale the kq_scale
+ const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
+ const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ // {n_embd, n_tokens}
+ inpL = build_inp_embd(model.tok_embd);
+
+ // inp_pos - contains the positions
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa();
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * inpSA = inpL;
+
+ // norm
+ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // self_attention
+ {
+ ggml_tensor * qr = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);
+ cb(qr, "qr", il);
+
+ qr = build_norm(qr, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(qr, "qr", il);
+
+ ggml_tensor * top_k = nullptr;
+
+ // lightning indexer
+ {
+ ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr);
+ cb(indexer_q, "indexer_q", il);
+
+ // split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens}
+ ggml_tensor * indexer_q_pe =
+ ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens,
+ ggml_row_size(indexer_q->type, n_embd_indexer_head),
+ ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0);
+ cb(indexer_q_pe, "indexer_q_pe", il);
+
+ // and {n_embd_indexer_head_nope, n_indexer_head, n_tokens}
+ ggml_tensor * indexer_q_nope =
+ ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens,
+ ggml_row_size(indexer_q->type, n_embd_indexer_head),
+ ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head,
+ ggml_row_size(indexer_q->type, n_embd_indexer_head_nope));
+ cb(indexer_q_nope, "indexer_q_nope", il);
+
+ indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot,
+ LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(indexer_q_pe, "indexer_q_pe", il);
+
+ // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens}
+ indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0);
+ cb(indexer_q, "indexer_q", il);
+
+ ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur);
+ cb(indexer_k, "indexer_k", il);
+
+ indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il);
+ cb(indexer_k, "indexer_k", il);
+
+ // split into {n_embd_indexer_head_rope, 1, n_tokens}
+ ggml_tensor * indexer_k_pe =
+ ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens,
+ ggml_row_size(indexer_k->type, n_embd_indexer_head),
+ ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0);
+ cb(indexer_k_pe, "indexer_k_pe", il);
+
+ // and {n_embd_indexer_head_nope, 1, n_tokens}
+ ggml_tensor * indexer_k_nope =
+ ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens,
+ ggml_row_size(indexer_k->type, n_embd_indexer_head),
+ ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1,
+ ggml_row_size(indexer_k->type, n_embd_indexer_head_nope));
+ cb(indexer_k_nope, "indexer_k_nope", il);
+
+ indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot,
+ LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(indexer_k_pe, "indexer_k_pe", il);
+
+ // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens}
+ indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0);
+ cb(indexer_k, "indexer_k", il);
+
+ // perform Hadamard transform on indexer q and k
+ indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q);
+ cb(indexer_q, "indexer_q", il);
+ indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k);
+ cb(indexer_k, "indexer_k", il);
+
+ // store indexer keys to KV cache
+ const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();
+ const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();
+ ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il));
+
+ // prepare indexer weights
+ ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur);
+ cb(indexer_weights, "indexer_weights", il);
+
+ // get cached indexer keys
+ indexer_k = mctx_lid->get_k(ctx0, il);
+
+ // split the batch into streams if needed
+ const auto n_stream = indexer_k->ne[3];
+ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);
+ indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
+
+ // calculate indexer kq
+ indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
+ cb(indexer_q, "indexer_q", il);
+ indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
+ cb(indexer_k, "indexer_k", il);
+
+ ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
+ cb(indexer_kq, "indexer_kq", il);
+
+ // ReLU requires contiguous tensors
+ indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
+ cb(indexer_kq, "indexer_kq", il);
+
+ // apply ReLU
+ ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
+ cb(indexer_score, "indexer_score", il);
+
+ // pre-scale weights to avoid scaling operations on huge indexer_score tensor
+ indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));
+ cb(indexer_weights, "indexer_weights", il);
+
+ // multiply scores by indexer weights
+ indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
+ cb(indexer_score, "indexer_score", il);
+
+ // sum by q n_indexer_head dimension
+ indexer_score = ggml_sum_rows(ctx0, indexer_score);
+ cb(indexer_score, "indexer_score", il);
+
+ // permute result to match KQ mask
+ indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
+ cb(indexer_score, "indexer_score", il);
+
+ // mask indexer scores
+ ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+ indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
+ cb(indexer_score, "indexer_score", il);
+
+ // get indices of top k indexer scores
+ uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;
+ top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
+ cb(top_k, "top_k", il);
+ }
+
+ ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_b, qr);
+ cb(q, "q", il);
+
+ // split into {n_embd_head_qk_nope, n_head, n_tokens}
+ ggml_tensor * q_nope =
+ ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
+ ggml_row_size(q->type, n_embd_head_k) * n_head, 0);
+ cb(q_nope, "q_nope", il);
+
+ // and {n_embd_head_qk_rope, n_head, n_tokens}
+ ggml_tensor * q_pe = ggml_view_3d(
+ ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
+ ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));
+ cb(q_pe, "q_pe", il);
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);
+ cb(kv_cmpr_pe, "kv_cmpr_pe", il);
+
+ // split into {kv_lora_rank, n_tokens}
+ ggml_tensor * kv_cmpr =
+ ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ // and {n_embd_head_qk_rope, 1, n_tokens}
+ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+ cb(k_pe, "k_pe", il);
+
+ q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(q_pe, "q_pe", il);
+
+ k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(k_pe, "k_pe", il);
+
+ kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ // MLA attention
+ {
+ // {n_embd_head_qk_nope, n_tokens, n_head}
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "q_nope_perm", il);
+
+ // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope);
+ cb(q_nope_absorbed, "q_nope_absorbed", il);
+
+ // {kv_lora_rank, n_head, n_tokens}
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "q_nope_absorbed_perm", il);
+
+ // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}
+ // note: rope must go first for in-place context shifting in build_rope_shift()
+ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ cb(Qcur, "Qcur", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr_reshape", il);
+
+ // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
+ cb(Kcur, "Kcur", il);
+
+ // {kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Vcur = kv_cmpr;
+ cb(Vcur, "Vcur", il);
+
+ // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
+ cur = build_attn(inp_attn_dsa,
+ model.layers[il].wo, NULL, model.layers[il].wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, 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);
+
+ cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ if ((uint32_t) il < hparams.n_layer_dense_lead) {
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,
+ model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+ model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ } else {
+ // MoE branch
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ model.layers[il].ffn_gate_inp,
+ model.layers[il].ffn_up_exps,
+ model.layers[il].ffn_gate_exps,
+ model.layers[il].ffn_down_exps,
+ model.layers[il].ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il,
+ nullptr,
+ model.layers[il].ffn_gate_up_exps,
+ model.layers[il].ffn_up_exps_s,
+ model.layers[il].ffn_gate_exps_s,
+ model.layers[il].ffn_down_exps_s);
+ cb(moe_out, "ffn_moe_out", il);
+
+ // FFN shared expert
+ {
+ ggml_tensor * ffn_shexp =
+ build_ffn(cur,
+ model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
+ model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
+ model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ 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 = inpL;
+
+ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ // lm_head
+ cur = ggml_mul_mat(ctx0, model.output, cur);
+
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 62: type = LLM_TYPE_142B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_dream::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
// Dream models are primarily 7B with 28 layers
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 28:
type = LLM_TYPE_7B;
break;
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 18: type = LLM_TYPE_0_3B; break;
case 28: type = LLM_TYPE_21B_A3B; break;
case 54: type = LLM_TYPE_300B_A47B; break;
void llama_model_eurobert::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- if (hparams.n_layer == 12) {
+ if (hparams.n_layer() == 12) {
type = LLM_TYPE_SMALL; // 0.2B
}
}
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_30B_A3B; break;
- case 48:
- case 49: type = LLM_TYPE_235B_A22B; break;
+ case 48: type = LLM_TYPE_235B_A22B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
// skip all tensors in the NextN layers
flags |= TENSOR_SKIP;
}
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
// dense layers for first n_layer_dense_lead layers or nextn_predict_layers layers at the end
- if (i < (int) hparams.n_layer_dense_lead || (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers)) {
+ if (i < (int) hparams.n_layer_dense_lead || (i >= n_layer)) {
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);
}
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags);
ggml_tensor * inp_out_ids = build_inp_out_ids();
- const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
// use RoPE for SWA layers
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
cb(cur, "attn_out", il);
}
- if (il == n_transformer_layers - 1 && inp_out_ids) {
+ 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);
}
void llama_model_exaone::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_8B; break;
default: type = LLM_TYPE_UNKNOWN;
}
#include "models.h"
void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) {
- if (hparams.n_layer == 64) { // 32B
+ if (hparams.n_layer() == 64) { // 32B
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.n_swa = 4096;
uint32_t swa_period = 4;
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);
- switch (hparams.n_layer) {
+ 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;
case 64: type = LLM_TYPE_32B; break;
default: type = LLM_TYPE_UNKNOWN;
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
+ const bool is_nextn = i >= n_layer;
+ int flags = 0;
+ if (is_nextn) {
+ // NextN/MTP layers are preserved in GGUF but are not executed yet.
+ flags |= TENSOR_SKIP;
+ }
+
auto & layer = layers[i];
- 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, n_embd}, 0);
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, flags);
- layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+ if (!is_nextn) {
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+ }
+
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, flags);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, flags);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, flags);
- layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
- 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);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);
+ layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, flags);
- 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);
- layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
+ if (is_nextn) {
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, flags | TENSOR_NOT_REQUIRED);
+ }
}
}
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);
- std::fill(hparams.recurrent_layer_arr.begin(), hparams.recurrent_layer_arr.end(), true);
+ std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), true);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 36:
type = LLM_TYPE_0_5B; break;
case 24:
void llama_model_falcon::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 60: type = LLM_TYPE_40B; break;
default: type = LLM_TYPE_UNKNOWN;
GGML_ASSERT((hparams.dense_2_feat_in == 0 || hparams.dense_2_feat_in == hparams.n_embd) && "dense_2_feat_in must be equal to n_embd");
GGML_ASSERT((hparams.dense_3_feat_out == 0 || hparams.dense_3_feat_out == hparams.n_embd) && "dense_3_feat_out must be equal to n_embd");
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_0_3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_gemma::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 18: type = LLM_TYPE_2B; break;
case 28: type = LLM_TYPE_7B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_ATTN_LOGIT_SOFTCAPPING, hparams.f_attn_logit_softcapping, false);
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 26: type = LLM_TYPE_2B; break;
case 42: type = LLM_TYPE_9B; break;
case 46: type = LLM_TYPE_27B; break;
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 18: type = LLM_TYPE_270M; break;
case 26: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_8B; break; // Rnj-1
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.set_swa_pattern(swa_period);
- hparams.n_layer_kv_from_start = 20;
- hparams.f_attention_scale = 1.0f;
+ hparams.n_layer_kv_from_start = 20;
+ hparams.f_attention_scale = 1.0f;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 30: type = LLM_TYPE_E2B; break;
case 35: type = LLM_TYPE_E4B; break;
default: type = LLM_TYPE_UNKNOWN;
--- /dev/null
+#include "models.h"
+
+void llama_model_gemma4_assistant::load_arch_hparams(llama_model_loader & ml) {
+ hparams.n_embd_inp_impl = hparams.n_embd_out();
+
+ 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());
+
+ uint32_t n_kv_shared_layers = 0;
+ ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false);
+
+ hparams.f_attention_scale = 1.0f;
+
+ 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_impl");
+
+ ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa);
+}
+
+void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+
+ if (n_embd_head_k != n_embd_head_v) {
+ throw std::runtime_error("Gemma 4 assistant requires n_embd_head_k == n_embd_head_v");
+ }
+ if (hparams.n_embd_head_k_swa != hparams.n_embd_head_v_swa) {
+ throw std::runtime_error("Gemma 4 assistant requires n_embd_head_k_swa == n_embd_head_v_swa");
+ }
+ if (hparams.n_embd_out() == n_embd) {
+ throw std::runtime_error("Gemma 4 assistant requires embedding_length_out to carry the target hidden size");
+ }
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
+
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);
+
+ const int64_t n_embd_backbone = hparams.n_embd_inp();
+ nextn_proj_post = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_POST, "weight"), { n_embd, n_embd_backbone }, 0);
+
+ int rope_freqs_flag = 0;
+
+ for (int i = 0; i < n_layer_nextn; ++i) {
+ auto & layer = layers[i];
+
+ const int64_t n_head = hparams.n_head(i);
+ const int64_t n_embd_head = hparams.n_embd_head_k(i);
+ const int64_t n_ff = hparams.n_ff(i);
+
+ if (i == 0) {
+ nextn_proj_pre = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_PRE, "weight", i), { 2*n_embd_backbone, n_embd }, 0);
+ }
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head*n_head }, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head*n_head, n_embd }, 0);
+
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head }, 0);
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
+
+ layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), { 1u }, 0);
+
+ if (!hparams.is_swa(i)) {
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), { n_embd_head/2 }, rope_freqs_flag);
+ rope_freqs_flag = TENSOR_DUPLICATED;
+ }
+
+ 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_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "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_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, 0);
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_gemma4_assistant::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_gemma4_assistant::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context(params) {
+ const int64_t n_embd_backbone = hparams.n_embd_inp();
+
+ ggml_tensor * inp_tokens;
+ ggml_tensor * inp_h;
+ {
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd_backbone);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);
+ cb(inp->tokens, "inp_tokens", -1);
+ ggml_set_input(inp->tokens);
+ inp_tokens = inp->tokens;
+ res->t_inp_tokens = inp->tokens;
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_backbone, ubatch.n_tokens);
+ cb(inp->embd, "inp_h", -1);
+ ggml_set_input(inp->embd);
+ inp_h = inp->embd;
+ res->t_inp_embd = inp->embd;
+
+ res->add_input(std::move(inp));
+ }
+
+ GGML_ASSERT(cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(cparams.ctx_other);
+
+ ggml_tensor * x = ggml_get_rows(ctx0, model_other->tok_embd, inp_tokens);
+ x = ggml_scale(ctx0, x, sqrtf((float) n_embd_backbone));
+ cb(x, "inp_embd_target", -1);
+
+ ggml_tensor * xh = ggml_concat(ctx0, x, inp_h, 0);
+ cb(xh, "inp_xh", -1);
+
+ ggml_tensor * cur = ggml_mul_mat(ctx0, model.nextn_proj_pre, xh);
+ cb(cur, "pre_proj", -1);
+
+ auto * inp_attn = build_attn_inp_kv_iswa();
+ ggml_tensor * inp_pos = build_inp_pos();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ ggml_tensor * inpL = cur;
+
+ for (int il = 0; il < n_layer_nextn; ++il) {
+ const bool is_swa = hparams.is_swa(il);
+
+ const int64_t n_embd_head = hparams.n_embd_head_k(il);
+ const int64_t n_head = hparams.n_head(il);
+
+ const float freq_base_l = model.get_rope_freq_base(cparams, il);
+ const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
+ const int n_rot_l = hparams.n_rot(il);
+
+ ggml_tensor * cur_norm = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur_norm, "attn_norm", il);
+
+ ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur_norm);
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+
+ ggml_tensor * freq_factors = is_swa ? nullptr : model.layers[il].rope_freqs;
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig,
+ freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(Qcur, "Qcur_pos", il);
+
+ cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr,
+ Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);
+
+ if (il == n_layer_nextn - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+ }
+
+ cur = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_post_norm", il);
+
+ ggml_tensor * attn_out = ggml_add(ctx0, cur, inpL);
+ cb(attn_out, "attn_out", il);
+
+ cur = build_norm(attn_out, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, nullptr, nullptr,
+ model.layers[il].ffn_gate, nullptr, nullptr,
+ model.layers[il].ffn_down, nullptr, nullptr,
+ nullptr,
+ LLM_FFN_GELU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+
+ cur = build_norm(cur, model.layers[il].ffn_post_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "ffn_post_norm", il);
+
+ cur = ggml_add(ctx0, cur, attn_out);
+
+ cur = ggml_mul(ctx0, cur, model.layers[il].out_scale);
+ cb(cur, "out_scaled", il);
+
+ inpL = cur;
+ }
+ cur = inpL;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+
+ ggml_tensor * logits = build_lora_mm(model.output, cur);
+ cb(logits, "result_output", -1);
+ res->t_logits = logits;
+
+ ggml_tensor * h_next = ggml_mul_mat(ctx0, model.nextn_proj_post, cur);
+ cb(h_next, "h_nextn", -1);
+ res->t_h_nextn = h_next;
+
+ ggml_build_forward_expand(gf, logits);
+ ggml_build_forward_expand(gf, h_next);
+}
void llama_model_gemma4::load_arch_hparams(llama_model_loader & ml) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
- ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
uint32_t n_kv_shared_layers = 0;
ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false);
- hparams.n_layer_kv_from_start = hparams.n_layer - (int32_t)n_kv_shared_layers;
+ hparams.n_layer_kv_from_start = hparams.n_layer_all - (int32_t)n_kv_shared_layers;
hparams.f_attention_scale = 1.0f; // Gemma4 uses self.scaling = 1.0 (no pre-attn scaling)
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa);
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 30: type = LLM_TYPE_26B_A4B; break;
case 35: type = LLM_TYPE_E2B; break;
case 42: type = LLM_TYPE_E4B; break;
idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
}
+// TODO @ngxson : maybe improve this in the future
+class llm_graph_input_logits_bias : public llm_graph_input_i {
+public:
+ llm_graph_input_logits_bias(const llama_vocab & vocab) {
+ arr.resize(vocab.n_tokens(), 0.0f);
+ for (llama_token id : vocab.get_suppress_tokens()) {
+ if (0 <= id && id < (int32_t)vocab.n_tokens()) {
+ arr[id] = -INFINITY;
+ }
+ }
+ }
+ virtual ~llm_graph_input_logits_bias() = default;
+
+ void set_input(const llama_ubatch * /*ubatch*/) override {
+ const int64_t n_vocab = arr.size();
+ ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias));
+ }
+
+ bool can_reuse(const llm_graph_params & /*params*/) override {
+ return true;
+ }
+
+ ggml_tensor * logits_bias = nullptr; // F32 [n_vocab]
+
+ std::vector<float> arr;
+};
+
llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_graph_context(params),
model(model),
}
// TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
- if (il == n_layer - 1 && inp_out_ids) {
+ // keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
// TODO @ngxson : improve this
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
}
model.output_norm, nullptr,
LLM_NORM_RMS, -1);
+ // Expose the post-output-norm hidden state (the LM-head input feature) so that
+ // MTP draft contexts can read it via llama_get_embeddings_nextn_ith() as the
+ // recurrent h input. This matches the reference (transformers/vLLM/SGLang),
+ // which feeds the drafter the target's post-final-norm hidden state.
+ 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;
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
}
+ // apply logits bias if needed (e.g. for gemma4_unified patch)
+ // this is to mirror the suppress_tokens patch on transformers, to avoid model from outputing <image|> and <audio|> tokens (which is a known issue related to the checkpoint)
+ // TODO: maybe handle this inside the sampling system in the future
+ if (!model.vocab.get_suppress_tokens().empty()) {
+ auto inp_bias = std::make_unique<llm_graph_input_logits_bias>(model.vocab);
+ inp_bias->logits_bias = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, inp_bias->arr.size());
+ cur = ggml_add(ctx0, cur, inp_bias->logits_bias);
+ res->add_input(std::move(inp_bias));
+ }
+
cb(cur, "result_output", -1);
res->t_logits = cur;
}
// NextN/MTP parameters
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
- // TODO: when MTP is implemented, this should probably be updated if needed
- hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
-
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 79: type = LLM_TYPE_744B_A40B; break;
default: type = LLM_TYPE_UNKNOWN;
}
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
// skip all tensors in the NextN layers
// TODO @ngxson : TENSOR_NOT_REQUIRED was a hack, need to remove it later
flags |= TENSOR_SKIP | TENSOR_NOT_REQUIRED;
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
}
- // NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ // NextN/MTP tensors (preserved but unused) - conditionally load for last n_layer_nextn
+ if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);
}
// NextN/MTP parameters
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
- // TODO: when MTP is implemented, this should probably be updated if needed
- hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
-
- switch (hparams.n_layer) {
- case 47: type = LLM_TYPE_106B_A12B; break; // GLM-4.5-Air (46 layers + 1 NextN layer)
+ switch (hparams.n_layer()) {
+ case 46: type = LLM_TYPE_106B_A12B; break; // GLM-4.5-Air
case 48: type = LLM_TYPE_102B_A12B; break; // Solar Open
- case 93: type = LLM_TYPE_355B_A32B; break; // GLM-4.5 (92 layers + 1 NextN layer)
+ case 92: type = LLM_TYPE_355B_A32B; break; // GLM-4.5
default: type = LLM_TYPE_UNKNOWN;
}
}
// Load ALL tensors including NextN layer to satisfy total tensor count
// but only PROCESS up to last layer (skipping final NextN layer) in forward pass
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
// skip all tensors in the NextN layers
flags |= TENSOR_SKIP;
}
}
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);
// Only process up to last layer (skip final NextN layer)
// Final layer tensors are loaded but not processed in forward pass
- const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
// Pre-attention norm
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_transformer_layers - 1 && inp_out_ids) {
+ 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);
}
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
// NextN/MTP parameters (GLM-OCR)
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
- // TODO: when MTP is implemented, this should probably be updated if needed
- hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
-
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 17: type = LLM_TYPE_1B; break; // GLM-OCR
case 40: type = LLM_TYPE_9B; break;
case 61: type = LLM_TYPE_32B; break;
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
// skip all tensors in the NextN layers
flags |= TENSOR_SKIP;
}
layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, flags);
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
- if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers) {
+ if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);
// Only process up to last layer (skip final NextN layer)
// Final layer tensors are loaded but not processed in forward pass
- const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
// Pre-attention norm
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_transformer_layers - 1 && inp_out_ids) {
+ 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);
}
void llama_model_gpt2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 12: type = LLM_TYPE_SMALL; break;
case 24: type = LLM_TYPE_MEDIUM; break;
case 36: type = LLM_TYPE_LARGE; break;
void llama_model_gptneox::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
ml.get_key(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 6:
switch (hparams.n_ff()) {
case 512: type = LLM_TYPE_14M; break;
hparams.rope_finetuned = rope_finetuned;
// A layer is recurrent IFF the n_head_kv value is set to 0
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0;
+ for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
+ hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
// norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- if (hparams.is_recurrent(i)) {
+ 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);
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// ssm layer //
cur = build_mamba2_layer(inp->get_recr(), cur, model, ubatch, il);
} else {
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
hparams.rope_finetuned = rope_finetuned;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_3B; break;
case 40: type = LLM_TYPE_3B; break;
// Add additional layer/vocab/etc checks here for other model sizes
#include "models.h"
+#include <sstream>
+
void llama_model_granite::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_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false);
+ // Granite4 Vision uses array deepstack_mapping
+ ml.get_arr(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr, false);
+
+ // Count the unique deepstack input indices
+ std::unordered_set<uint32_t> unique_deepstack_idxs;
+ for (const auto val : hparams.deepstack_mapping_arr) {
+ if (val >= 0) {
+ unique_deepstack_idxs.insert(val);
+ }
+ }
+ hparams.n_deepstack_layers = unique_deepstack_idxs.size();
+
+ // Ensure all values are valid (avoid overflow attacks)
+ for (const auto val : unique_deepstack_idxs) {
+ if (val > hparams.n_deepstack_layers) {
+ std::stringstream ss;
+ ss << "Invalid deepstack index: " << val << " > " << hparams.n_deepstack_layers;
+ throw std::runtime_error(ss.str());
+ }
+ }
+
// Granite uses rope_finetuned as a switch for rope, so default to true
bool rope_finetuned = true;
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
hparams.rope_finetuned = rope_finetuned;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_3B; break;
case 40: type = LLM_TYPE_3B; break;
// Add additional layer/vocab/etc checks here for other model sizes
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
+
+ // Granite Vision 4.1 deepstack: inject the projector stream that
+ // targets decoder layer `il` before the decoder runs.
+ // NOTE: skip the first deepstack layer since that's inpL
+ const auto & deepstack_emb_idx = hparams.deepstack_mapping_arr[il];
+ if (il > 0 && deepstack_emb_idx >= 0) {
+ ggml_tensor * ds = ggml_view_2d(ctx0,
+ res->t_inp_embd, n_embd, n_tokens,
+ res->t_inp_embd->nb[1],
+ deepstack_emb_idx * n_embd * sizeof(float));
+ inpL = ggml_add(ctx0, inpL, ds);
+ cb(inpL, "deepstack_in", il);
+ }
+
ggml_tensor * inpSA = inpL;
// norm
ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_FAST, hparams.yarn_beta_fast, false);
ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_SLOW, hparams.yarn_beta_slow, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 64: type = LLM_TYPE_314B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERTS_PER_GROUP, hparams.n_group_experts);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_30B_A3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_A13B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_internlm2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 48: type = LLM_TYPE_20B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
ml.get_key(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, hparams.f_max_alibi_bias, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1_3B; break;
case 40: type = LLM_TYPE_13B; break;
/* TODO: add variants */
void llama_model_jais2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_8B; break;
case 68: type = LLM_TYPE_70B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0;
+ for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
+ hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
// TODO: Jamba layers are a bit heterogeneous, so naming this is hard.
case 12: // 900M 8x???M
case 32: // 51B 16x?B
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
hparams.f_max_alibi_bias = 8.0f;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 4: type = LLM_TYPE_33M; break; // jina-embeddings-small
case 12: type = LLM_TYPE_137M; break; // jina-embeddings-base
default: type = LLM_TYPE_UNKNOWN;
void llama_model_jina_bert_v3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24:
type = LLM_TYPE_558M; break;
default: type = LLM_TYPE_UNKNOWN;
// Mark KDA layers as recurrent using n_head_kv pattern (like Jamba)
// Set n_head_kv = 0 for KDA layers (recurrent), n_head_kv = n_head for MLA layers (attention)
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0; // KDA layers are recurrent
+ for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
+ hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; // KDA layers are recurrent
}
// MoE parameters - Kimi uses moe_intermediate_size = 1024
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 27: type = LLM_TYPE_48B_A3B; break; // Kimi-Linear-48B-A3B
default: type = LLM_TYPE_UNKNOWN;
}
const int64_t n_embd_head_v_kda = hparams.n_embd_head_kda;
const int64_t ssm_d_conv = hparams.ssm_d_conv;
- if (hparams.is_recurrent(i)) {
+ if (hparams.is_recr(i)) {
// Conv1d weights: try 4D first, then 3D (quantization may remove trailing 1)
// 4D: [d_conv, 1, d_inner, 1], 3D: [d_conv, 1, d_inner]
layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head, 1}, TENSOR_NOT_REQUIRED);
ggml_build_forward_expand(gf, cur);
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// === KDA Layer (Kimi Delta Attention) with Recurrent State ===
// Reference: vLLM kda.py
const auto * mctx_cur = inp_rs->mctx;
void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SHORTCONV_L_CACHE, hparams.n_shortconv_l_cache);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- for (uint32_t il = 0; il < hparams.n_layer; ++il) {
- hparams.recurrent_layer_arr[il] = hparams.n_head_kv(il) == 0;
+
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0;
}
- hparams.n_layer_dense_lead = hparams.n_layer;
+
+ hparams.n_layer_dense_lead = hparams.n_layer();
+
switch (hparams.n_ff()) {
case 4608: type = LLM_TYPE_350M; break;
case 6912: type = LLM_TYPE_700M; break;
case 10752: type = LLM_TYPE_2_6B; break;
default: type = LLM_TYPE_UNKNOWN;
}
+
if (const auto is_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); is_swa && hparams.n_swa > 0) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
- for (uint32_t il = 0; il < hparams.n_layer; ++il) {
- hparams.swa_layers[il] = !hparams.recurrent_layer_arr[il];
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ hparams.is_swa_impl[il] = !hparams.is_recr_impl[il];
}
}
}
// for operator_norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- if (!hparams.is_recurrent(i)) {
+ if (!hparams.is_recr(i)) {
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);
GGML_ASSERT(n_embd_v_gqa == n_embd_k_gqa);
cur = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "model.layers.{}.operator_norm", il);
- cur = hparams.is_recurrent(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
- build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
+ cur = hparams.is_recr(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
+ build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
- for (uint32_t il = 0; il < hparams.n_layer; ++il) {
- hparams.recurrent_layer_arr[il] = hparams.n_head_kv(il) == 0;
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_8B_A1B; break;
case 40: type = LLM_TYPE_24B_A2B; break;
default: type = LLM_TYPE_UNKNOWN;
// for operator_norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- if (!hparams.is_recurrent(i)) {
+ if (!hparams.is_recr(i)) {
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);
GGML_ASSERT(n_embd_v_gqa == n_embd_k_gqa);
void llama_model_llada_moe::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
-
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
// diffusion language model uses non-causal attention
hparams.causal_attn = false;
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 16: type = LLM_TYPE_A1_7B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_llada::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
// LLaDA-8B has 32 layers, similar to LLaMA but for diffusion
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32:
type = LLM_TYPE_8B;
break;
default:
type = LLM_TYPE_UNKNOWN;
}
+
// Set non-causal attention for diffusion models
hparams.causal_attn = false;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
if (hparams.n_expert == 8) {
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_8x7B; break;
case 56: type = LLM_TYPE_8x22B; break;
default: type = LLM_TYPE_UNKNOWN;
}
} else {
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 16: type = LLM_TYPE_1B; break; // Llama 3.2 1B
case 22: type = LLM_TYPE_1B; break;
case 26: type = LLM_TYPE_3B; break;
const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
if (found_swa && hparams.n_swa == 0) {
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
- hparams.n_no_rope_layer_step = hparams.n_layer; // always use rope
+ hparams.n_no_rope_layer_step = hparams.n_layer(); // always use rope
} else {
hparams.swa_type = LLAMA_SWA_TYPE_CHUNKED;
hparams.n_swa = 8192;
hparams.n_attn_temp_floor_scale = 8192;
hparams.f_attn_temp_scale = 0.1f;
hparams.f_attn_temp_offset = 1.0f;
- uint32_t swa_period = 4; // pattern: 3 chunked - 1 full
+
+ uint32_t swa_period = 4; // pattern: 3 chunked - 1 full
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
hparams.set_swa_pattern(swa_period);
void llama_model_maincoder::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_1B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24:
switch (hparams.n_embd) {
case 768: type = LLM_TYPE_SMALL; break;
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24:
switch (hparams.n_embd) {
case 768: type = LLM_TYPE_SMALL; break;
--- /dev/null
+#include "models.h"
+
+void llama_model_mellum::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_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
+
+ if (hparams.n_swa > 0) {
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+
+ uint32_t swa_period = 4;
+ const auto res = ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
+ if (res) {
+ 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());
+ }
+
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+
+ ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
+ } else {
+ hparams.swa_type = LLAMA_SWA_TYPE_NONE;
+ }
+
+ switch (hparams.n_layer()) {
+ case 28: type = LLM_TYPE_12B_A2_5B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_mellum::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
+ 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];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+
+ if (n_expert == 0) {
+ throw std::runtime_error("n_expert must be > 0 for Mellum");
+ }
+ if (n_expert_used == 0) {
+ throw std::runtime_error("n_expert_used must be > 0 for Mellum");
+ }
+
+ const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
+
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { 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, n_embd, n_expert}, 0);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0);
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_mellum::build_arch_graph(const llm_graph_params & params) const {
+ if (hparams.swa_type == LLAMA_SWA_TYPE_STANDARD) {
+ return std::make_unique<graph<true>>(*this, params);
+ }
+ return std::make_unique<graph<false>>(*this, params);
+}
+
+template <bool iswa>
+llama_model_mellum::graph<iswa>::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);
+
+ // inp_pos - contains the positions
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
+ inp_attn_type * inp_attn = nullptr;
+
+ if constexpr (iswa) {
+ inp_attn = build_attn_inp_kv_iswa();
+ } else {
+ inp_attn = build_attn_inp_kv();
+ }
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * inpSA = inpL;
+
+ // norm
+ cur = build_norm(inpL,
+ model.layers[il].attn_norm, nullptr,
+ LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // self_attention
+ {
+ // compute Q and K and RoPE them
+ auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+ n_embd_head, n_head, n_head_kv, il);
+
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+
+ Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Kcur, "Kcur_normed", il);
+
+ const bool is_swa = hparams.is_swa(il);
+
+ if (is_swa) {
+ // For sliding window layers, use regular rope with no yarn rope scaling.
+ // This is achieved here by setting freq_scale and attn_factor to 1.
+ // We also set ext_factor to 0 to avoid a few unnecessary computations.
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base, 1.0,
+ 0.0, 1.0, beta_fast, beta_slow
+ );
+
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base, 1.0,
+ 0.0, 1.0, beta_fast, beta_slow
+ );
+ } else {
+ 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, model.layers[il].wo_b, 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);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "ffn_inp", il);
+
+ // MoE
+ cur = build_norm(ffn_inp,
+ model.layers[il].ffn_norm, nullptr,
+ LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ ggml_tensor * moe_out =
+ build_moe_ffn(cur,
+ model.layers[il].ffn_gate_inp,
+ model.layers[il].ffn_up_exps,
+ model.layers[il].ffn_gate_exps,
+ model.layers[il].ffn_down_exps,
+ nullptr,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, true,
+ hparams.expert_weights_scale,
+ LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,
+ il,
+ nullptr, nullptr,
+ model.layers[il].ffn_up_exps_s,
+ model.layers[il].ffn_gate_exps_s,
+ model.layers[il].ffn_down_exps_s);
+ cb(moe_out, "ffn_moe_out", il);
+ cur = moe_out;
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "ffn_out", il);
+
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ // input for next layer
+ inpL = cur;
+ }
+ cur = inpL;
+
+ cur = build_norm(cur,
+ model.output_norm, nullptr,
+ LLM_NORM_RMS, -1);
+
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ // lm_head
+ 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);
+}
+
+template struct llama_model_mellum::graph<false>;
+template struct llama_model_mellum::graph<true>;
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
- ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
+
+ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
float value_scale = 0.0f;
if (ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, value_scale, false) && value_scale != 1.0f) {
hparams.f_attn_value_scale = value_scale;
}
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
- hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
+ 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_impl");
- switch (hparams.n_layer - hparams.nextn_predict_layers) {
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_310B_A15B; break;
default: type = LLM_TYPE_UNKNOWN;
}
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);
- const uint32_t n_nextn = hparams.nextn_predict_layers;
-
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
uint32_t n_head = hparams.n_head(i);
// NextN/MTP layers (the last n_nextn blocks) are preserved but disabled pending support
- const bool is_nextn = (n_nextn > 0) && (static_cast<uint32_t>(i) >= n_layer - n_nextn);
+ const bool is_nextn = i >= n_layer;
const int skip = is_nextn ? TENSOR_SKIP : 0;
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip);
const float v_scale = hparams.f_attn_value_scale;
- // The last hparams.nextn_predict_layers blocks are MTP heads, currently inactive
- const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
-
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
uint32_t n_head_l = hparams.n_head(il);
}
}
- if (il == n_transformer_layers - 1 && inp_out_ids) {
+ 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);
}
void llama_model_minicpm::load_arch_hparams(llama_model_loader & ml) {
// Backward-compatible defaults for older MiniCPM GGUFs
hparams.f_embedding_scale = 12.0f;
- hparams.f_residual_scale = 1.4f / sqrtf(float(hparams.n_layer));
+ hparams.f_residual_scale = 1.4f / sqrtf(float(hparams.n_layer()));
hparams.f_logit_scale = hparams.n_embd ? (256.0f / float(hparams.n_embd)) : 1.0f;
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
// MiniCPM uses rope by default, unlike Granite which uses it as a switch
hparams.rope_finetuned = true;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 52: type = LLM_TYPE_1B; break;
case 40: type = LLM_TYPE_2B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 62: type = LLM_TYPE_4B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 62: type = LLM_TYPE_230B_A10B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 26: type = LLM_TYPE_3B; break;
case 34: type = LLM_TYPE_8B; break;
case 40: type = LLM_TYPE_14B; break;
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
+struct llama_model_mellum : public llama_model_base {
+ llama_model_mellum(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;
+
+ template <bool iswa>
+ 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_qwen : public llama_model_base {
llama_model_qwen(const struct llama_model_params & params) : llama_model_base(params) {}
};
+struct llama_model_gemma4_assistant : public llama_model_base {
+ llama_model_gemma4_assistant(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_gemma_embedding : public llama_model_base {
llama_model_gemma_embedding(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
};
+struct llama_model_deepseek32 : public llama_model_base {
+ llama_model_deepseek32(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_deepseek2ocr : public llama_model_base {
llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
graph(const llama_model & model, const llm_graph_params & params);
};
+ 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;
};
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ // Some ModernBert derivatives (e.g. IBM Granite Embedding 97m R2) use
+ // SiLU/SwiGLU in the FFN instead of the default GELU/GeGLU.
+ hparams.llm_ffn_op = LLM_FFN_GEGLU;
+ std::string hidden_act;
+ if (ml.get_key(LLM_KV_HIDDEN_ACT, hidden_act, false)) {
+ hparams.llm_ffn_op = llm_ffn_op_type_from_string(hidden_act, LLM_FFN_GEGLU);
+ }
+
+ switch (hparams.n_layer()) {
case 12:
type = LLM_TYPE_47M; break; // granite-embedding-small
case 22:
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
- LLM_FFN_GEGLU, LLM_FFN_SEQ, il);
+ hparams.llm_ffn_op,
+ LLM_FFN_SEQ, il);
// attentions bypass the intermediate layer
cur = ggml_add(ctx0, cur, ffn_inp);
ml.get_key(LLM_KV_ATTENTION_CLAMP_KQV, hparams.f_clamp_kqv, false);
ml.get_key(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, hparams.f_max_alibi_bias, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 48: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
// 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.recurrent_layer_arr[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 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);
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_MOE_LATENT_SIZE, hparams.moe_latent_size, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 52: type = LLM_TYPE_31B_A3_5B; break; // Nemotron-H_MOE 31B
case 56: type = LLM_TYPE_9B; break;
case 88: type = LLM_TYPE_120B_A12B; break;
// all blocks use the attn norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- if (hparams.is_recurrent(i)) {
+ 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);
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// ssm layer //
cur = build_mamba2_layer(inp->get_recr(), cur, model, ubatch, il);
} else if (hparams.n_ff(il) == 0) {
void llama_model_nemotron::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_4B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_neo_bert::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- if (hparams.n_layer == 28) {
+ if (hparams.n_layer() == 28) {
type = LLM_TYPE_250M;
}
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
ml.get_key(LLM_KV_MOE_EVERY_N_LAYERS, hparams.moe_every_n_layers, 0);
- if (hparams.n_layer == 12 && hparams.n_embd == 768) {
+ if (hparams.n_layer() == 12 && hparams.n_embd == 768) {
if (arch == LLM_ARCH_NOMIC_BERT) {
type = LLM_TYPE_137M;
} else if (arch == LLM_ARCH_NOMIC_BERT_MOE && hparams.moe_every_n_layers == 2) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
ml.get_key(LLM_KV_MOE_EVERY_N_LAYERS, hparams.moe_every_n_layers, 0);
- if (hparams.n_layer == 12 && hparams.n_embd == 768) {
+ if (hparams.n_layer() == 12 && hparams.n_embd == 768) {
if (arch == LLM_ARCH_NOMIC_BERT) {
type = LLM_TYPE_137M;
} else if (arch == LLM_ARCH_NOMIC_BERT_MOE && hparams.moe_every_n_layers == 2) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
ml.get_key(LLM_KV_ATTENTION_CLAMP_KQV, hparams.f_clamp_kqv, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 22: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_7B; break;
case 80: type = LLM_TYPE_70B; break;
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 16: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_7B; break;
case 40: type = LLM_TYPE_13B; break;
void llama_model_olmoe::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 16: type = LLM_TYPE_A1_7B; break;
default: type = LLM_TYPE_UNKNOWN;
}
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_20B; break;
case 36: type = LLM_TYPE_120B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_openelm::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
- case 16: type = LLM_TYPE_270M; break;
- case 20: type = LLM_TYPE_450M; break;
- case 28: type = LLM_TYPE_1B; break;
- case 36: type = LLM_TYPE_3B; break;
- default: type = LLM_TYPE_UNKNOWN;
+ switch (hparams.n_layer()) {
+ case 16: type = LLM_TYPE_270M; break;
+ case 20: type = LLM_TYPE_450M; break;
+ case 28: type = LLM_TYPE_1B; break;
+ case 36: type = LLM_TYPE_3B; break;
+ default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_orion::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 40: type = LLM_TYPE_14B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_pangu_embed::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 26: type = LLM_TYPE_1B; break; // openPangu-Embedded-1B-V1.1
case 34: type = LLM_TYPE_7B; break; // openPangu-Embedded-7B-V1.1
default: type = LLM_TYPE_UNKNOWN;
void llama_model_phi2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_3B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_phi3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_3B; break;
case 40: type = LLM_TYPE_14B; break;
void llama_model_phimoe::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_16x3_8B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_plamo::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 40: type = LLM_TYPE_13B; break;
default: type = LLM_TYPE_UNKNOWN;
}
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);
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0;
+ for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
+ hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 16: type = LLM_TYPE_1B; break;
case 32:
if (hparams.n_embd == 2048) {
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
- bool is_mamba_layer = hparams.is_recurrent(i);
+ bool is_mamba_layer = hparams.is_recr(i);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
// check if this layer is Mamba or Attention
- const bool is_mamba_layer = hparams.is_recurrent(il);
+ const bool is_mamba_layer = hparams.is_recr(il);
if (is_mamba_layer) {
// PLaMo-2 Mamba layer
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_2B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_plm::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_KV_LORA_RANK, hparams.n_lora_kv);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_1_8B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_qwen::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 40: type = LLM_TYPE_13B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_qwen2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 24: type = hparams.n_embd == 1024 ? LLM_TYPE_0_5B : LLM_TYPE_1B; break;
case 28: type = hparams.n_embd == 1536 ? LLM_TYPE_1_5B : LLM_TYPE_7B; break;
case 32: type = LLM_TYPE_7B; break;
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_A2_7B; break;
case 28: type = LLM_TYPE_57B_A14B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_qwen3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 28: type = hparams.n_embd == 1024 ? LLM_TYPE_0_6B : LLM_TYPE_1_7B; break;
case 36: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_8B; break;
case 40: type = LLM_TYPE_14B; break;
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// NextN/MTP (Qwen3.5/3.6): extra decoder block appended beyond the main stack
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
// Mark recurrent layers (linear attention layers). MTP layers are dense
// attention-only and must be flagged non-recurrent.
- {
- const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
+ if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = (i < n_main) && ((i + 1) % full_attn_interval != 0);
+ for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+ hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
}
}
- switch (hparams.n_layer - hparams.nextn_predict_layers) {
+ switch (hparams.n_layer()) {
case 24: type = hparams.n_embd == 1024 ? LLM_TYPE_0_8B : LLM_TYPE_2B; break;
case 32: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_9B; break;
case 64: type = LLM_TYPE_27B; break;
void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
- const uint32_t n_main = n_layer - hparams.nextn_predict_layers;
- const bool mtp_only = (hparams.nextn_predict_layers > 0) &&
- (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ 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;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
- if (!hparams.is_recurrent(il)) {
+ if (!hparams.is_recr(il)) {
// Attention layers
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
};
- for (int i = 0; i < (int) n_main; ++i) {
+ for (int i = 0; i < n_layer; ++i) {
load_block_trunk(i, trunk_flags);
}
- for (int i = (int) n_main; i < n_layer; ++i) {
+ for (int i = n_layer; i < n_layer_all; ++i) {
load_block_mtp(i);
}
}
ggml_tensor * inp_out_ids = build_inp_out_ids();
// MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
- const int n_transformer_layers = n_layer - (int) hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
ggml_build_forward_expand(gf, cur);
// Determine layer type and build appropriate attention mechanism
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// Linear attention layer (gated delta net)
cur = build_layer_attn_linear(inp->get_recr(), cur, il);
} else {
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
}
- if (il == n_transformer_layers - 1 && inp_out_ids && cparams.embeddings_pre_norm_masked) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
}
cur = inpL;
- cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
- if (!cparams.embeddings_pre_norm_masked && inp_out_ids) {
+ if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
- // Final norm
- cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
-
cb(cur, "result_norm", -1);
res->t_embd = cur;
// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3.5/3.6 dense series
llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
- GGML_ASSERT(hparams.nextn_predict_layers > 0 && "QWEN35 MTP requires nextn_predict_layers > 0");
- GGML_ASSERT(hparams.nextn_predict_layers == 1 && "QWEN35 MTP currently only supports a single MTP block");
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN35 MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN35 MTP currently only 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());
// hparams.n_layer includes both main model layers and MTP layers. The MTP
// layer is stored immediately after the main layers in model.layers[].
- const int il = (int) hparams.n_layer - (int) hparams.nextn_predict_layers;
+ const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
int sections[4];
std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);
- auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
+ // TODO: extract in a common llm_graph_context::build_inp_embd_h()
+ 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, n_tokens);
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- ggml_set_name(inp->embd, "mtp_h_input");
- ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+ // TODO: make static using `ggml_build_forward_select()`
+ // see llm_graph_context::build_inp_embd() for reference
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
- ggml_tensor * h_input = inp->embd;
- ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ 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_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
- auto * inp_attn = build_attn_inp_kv();
- ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ 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);
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
- // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
- // (In the trunk graph this is `t_h_pre_norm`; the MTP head reuses the same slot.)
- cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
-
- cur = ggml_get_rows(ctx0, cur, inp_out_ids);
-
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
? layer.nextn.shared_head_norm
: model.output_norm;
GGML_ASSERT(head_norm_w && "QWEN35 MTP: missing both nextn.shared_head_norm and output_norm");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "mtp_shared_head_norm", -1);
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// NextN/MTP (Qwen3.5/3.6): extra decoder block appended beyond the main stack
- ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
- GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
+ 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_impl");
// Mark recurrent layers (linear attention layers). MTP layers are dense
// attention-only and must be flagged non-recurrent.
- {
- const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
+ if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = (i < n_main) && ((i + 1) % full_attn_interval != 0);
+ for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+ hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
}
}
- switch (hparams.n_layer - hparams.nextn_predict_layers) {
+ switch (hparams.n_layer()) {
case 40: type = LLM_TYPE_35B_A3B; break;
case 48: type = LLM_TYPE_122B_A10B; break;
case 60: type = LLM_TYPE_397B_A17B; break;
void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
- const uint32_t n_main = n_layer - hparams.nextn_predict_layers;
- const bool mtp_only = (hparams.nextn_predict_layers > 0) &&
- (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ 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;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
- if (!hparams.is_recurrent(il)) {
+ if (!hparams.is_recr(il)) {
// Attention layers
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
};
- for (int i = 0; i < (int) n_main; ++i) {
+ for (int i = 0; i < n_layer; ++i) {
load_block_trunk(i, trunk_flags);
}
- for (int i = (int) n_main; i < n_layer; ++i) {
+ for (int i = n_layer; i < n_layer_all; ++i) {
load_block_mtp(i);
}
}
ggml_tensor * inp_out_ids = build_inp_out_ids();
// MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
- const int n_transformer_layers = n_layer - (int) hparams.nextn_predict_layers;
- for (int il = 0; il < n_transformer_layers; ++il) {
+ for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
ggml_build_forward_expand(gf, cur);
// Determine layer type and build appropriate attention mechanism
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// Linear attention layer (gated delta net)
cur = build_layer_attn_linear(inp->get_recr(), cur, il);
} else {
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
}
- if (il == n_transformer_layers - 1 && inp_out_ids && cparams.embeddings_pre_norm_masked) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
}
cur = inpL;
- cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
+ // post-norm hidden state feeds both the LM head and the MTP seed below
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
- if (!cparams.embeddings_pre_norm_masked && inp_out_ids) {
+ if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
- // Final norm
- cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
-
cb(cur, "result_norm", -1);
res->t_embd = cur;
// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3.5/3.6 MoE
llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
- GGML_ASSERT(hparams.nextn_predict_layers > 0 && "QWEN35MOE MTP requires nextn_predict_layers > 0");
- GGML_ASSERT(hparams.nextn_predict_layers == 1 && "QWEN35MOE MTP currently only supports a single MTP block");
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN35MOE MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN35MOE MTP currently only 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 = (int) hparams.n_layer - (int) hparams.nextn_predict_layers;
+ const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
int sections[4];
std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);
- auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
+ // TODO: extract in a common llm_graph_context::build_inp_embd_h()
+ 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, n_tokens);
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- ggml_set_name(inp->embd, "mtp_h_input");
- ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+ // TODO: make static using `ggml_build_forward_select()`
+ // see llm_graph_context::build_inp_embd() for reference
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
- ggml_tensor * h_input = inp->embd;
- ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ 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_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
- auto * inp_attn = build_attn_inp_kv();
+ auto * inp_attn = build_attn_inp_kv();
- ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ 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);
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
- // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
- cb(cur, "h_pre_norm", -1);
- res->t_h_pre_norm = cur;
-
- cur = ggml_get_rows(ctx0, cur, inp_out_ids);
-
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
? layer.nextn.shared_head_norm
: model.output_norm;
GGML_ASSERT(head_norm_w && "QWEN35MOE MTP: missing both nextn.shared_head_norm and output_norm");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn= cur;
+
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "mtp_shared_head_norm", -1);
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
#include "models.h"
void llama_model_qwen3moe::load_arch_hparams(llama_model_loader & ml) {
- ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
-
+ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_30B_A3B; break;
case 94: type = LLM_TYPE_235B_A22B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// Mark recurrent layers (linear attention layers)
- {
+ if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
- for (uint32_t i = 0; i < hparams.n_layer; ++i) {
- hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0);
+ for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+ hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
}
}
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_80B_A3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
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);
- if (!hparams.is_recurrent(i)) {
+ if (!hparams.is_recr(i)) {
// Attention layers
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, 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);
ggml_build_forward_expand(gf, cur);
// Determine layer type and build appropriate attention mechanism
- if (hparams.is_recurrent(il)) {
+ if (hparams.is_recr(il)) {
// Linear attention layer (gated delta net)
cur = build_layer_attn_linear(inp->get_recr(), cur, il);
} else {
ml.get_key(LLM_KV_NUM_DEEPSTACK_LAYERS, hparams.n_deepstack_layers, false);
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 28: type = LLM_TYPE_1_7B; break;
case 36: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_8B; break;
case 64: type = LLM_TYPE_32B; break;
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_30B_A3B; break;
case 94: type = LLM_TYPE_235B_A22B; break;
default: type = LLM_TYPE_UNKNOWN;
void llama_model_refact::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_1B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_rnd1::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
-
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_30B_A3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
+
// Set non-causal attention for diffusion models
hparams.causal_attn = false;
}
ml.get_key(LLM_KV_RESCALE_EVERY_N_LAYERS, hparams.rescale_every_n_layers, false);
ml.get_key(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1_6B; break;
case 32:
switch (hparams.n_embd) {
ml.get_key(LLM_KV_RESCALE_EVERY_N_LAYERS, hparams.rescale_every_n_layers, false);
ml.get_key(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1_6B; break;
case 32:
switch (hparams.n_embd) {
ml.get_key(LLM_KV_ATTENTION_GATE_LORA_RANK, hparams.n_lora_gate, false);
ml.get_key(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 12:
switch (hparams.n_embd) {
case 768: type = LLM_TYPE_190M; break;
void llama_model_seed_oss::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 64: type = LLM_TYPE_36B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
} else {
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
- hparams.n_no_rope_layer_step = hparams.n_layer;
+ hparams.n_no_rope_layer_step = hparams.n_layer();
}
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_4B; break;
case 52: type = LLM_TYPE_20B; break;
default: type = LLM_TYPE_UNKNOWN;
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
hparams.n_no_rope_layer_step = 4;
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 36: type = LLM_TYPE_3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_stablelm::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1B; break;
case 32: type = LLM_TYPE_3B; break;
case 40: type = LLM_TYPE_12B; break;
void llama_model_starcoder::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 24: type = LLM_TYPE_1B; break;
case 36: type = LLM_TYPE_3B; break;
case 42: type = LLM_TYPE_7B; break;
void llama_model_starcoder2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 30: type = LLM_TYPE_3B; break;
case 32: type = LLM_TYPE_7B; break;
case 40: type = LLM_TYPE_15B; break;
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
- ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
- ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer, false);
- ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer, false);
- switch (hparams.n_layer) {
+ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
+
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer(), false);
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer(), false);
+
+ // NextN/MTP (Step3p5): extra decoder block appended beyond the main stack.
+ 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_impl");
+
+ switch (hparams.n_layer()) {
case 45: type = LLM_TYPE_196B_A11B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
-void llama_model_step35::load_arch_tensors(llama_model_loader &) {
+void llama_model_step35::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);
+ // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
+ // tensors live in a separate file (e.g. user split target/draft). Mark
+ // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
+ const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
+ const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
+ const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+ const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// output
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);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, trunk_flags);
// STEP35 supports per-layer partial RoPE dims; rope factors are stored as a single shared tensor
// ("rope_freqs.weight") and ggml uses only the first (n_rot_l/2) entries per layer.
n_rot_max = n_rot;
}
- for (int i = 0; i < n_layer; ++i) {
+ auto load_block_trunk = [&](int i, int flags) {
auto & layer = layers[i];
const uint32_t n_head_l = hparams.n_head(i);
const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
- 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}, flags);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
}
- create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_l, n_embd_k_gqa, n_embd_v_gqa, 0);
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, 0);
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_l, n_embd_k_gqa, n_embd_v_gqa, flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, flags);
// head-wise attention gate (Step35 self_attn.g_proj)
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, TENSOR_NOT_REQUIRED);
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
// dense MLP (leading dense blocks)
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
+ };
+
+ auto load_block_mtp = [&](int i, bool is_first_mtp) {
+ auto & layer = layers[i];
+
+ const uint32_t n_head_l = hparams.n_head(i);
+ const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
+ const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
+
+ // The MTP block is a full Step3p5 decoder layer (mtp_block) plus the
+ // NextN-specific wiring (enorm/hnorm/eh_proj + optional shared head).
+ // `mtp_flags` becomes NOT_REQUIRED when the GGUF is trunk-only.
+ //
+ // Only the FIRST MTP block (i == n_main) is required for the
+ // single-block MTP runtime; trailing MTP blocks are always tolerated
+ // as missing so pruned GGUFs (block 0 only) load cleanly. Override
+ // mtp_flags to NOT_REQUIRED for those.
+ const int eff_mtp_flags = is_first_mtp ? mtp_flags : (mtp_flags | TENSOR_NOT_REQUIRED);
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, eff_mtp_flags);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
+
+ if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
+ layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
+ layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
+ } else {
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
+ }
+
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_l, n_embd_k_gqa, n_embd_v_gqa, eff_mtp_flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, eff_mtp_flags);
+
+ layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, TENSOR_NOT_REQUIRED);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, eff_mtp_flags);
+
+ // dense MLP (leading dense blocks) — present if the MTP block isn't MoE
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
+
+ // MoE routed experts + selection bias (router_bias)
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
+
+ // NextN-specific tensors that define the MTP block.
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, eff_mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, eff_mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, eff_mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED);
+ };
+
+ for (int i = 0; i < n_layer; ++i) {
+ load_block_trunk(i, trunk_flags);
+ }
+ // Only the first MTP block (i == n_main) is required at runtime — the
+ // single-block-MTP graph in build_arch_graph always uses that one.
+ // Trailing MTP blocks are loaded if present (so an un-pruned GGUF with
+ // all MTP layers still works) but tolerated when absent via the pruning
+ // path. See scripts/prune_step35_extra_mtp.py for the pruner.
+ for (int i = n_layer; i < n_layer_all; ++i) {
+ load_block_mtp(i, /*is_first_mtp=*/ i == n_layer);
}
}
std::unique_ptr<llm_graph_context> llama_model_step35::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);
}
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * inp_out_ids = build_inp_out_ids();
+ // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cb(cur, "attn_proj", il);
}
- if (il == n_layer - 1 && inp_out_ids) {
- cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = inpL;
+ 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);
+ }
+
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
ggml_build_forward_expand(gf, cur);
}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for Step3p5 (MoE)
+llama_model_step35::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "STEP35 MTP requires n_layer_nextn > 0");
+
+ // Single-block MTP only: always run the first trained MTP block (Qwen
+ // MTP / vLLM single-MTP-layer style). Multi-block round-robin proved to
+ // be a much deeper refactor than this PR justifies; the trailing MTP
+ // blocks are loaded with TENSOR_NOT_REQUIRED so pruned GGUFs (with just
+ // block 0) also work — see load_arch_tensors below and
+ // scripts/prune_step35_extra_mtp.py.
+ const int il = hparams.n_layer();
+ const auto & layer = model.layers[il];
+
+ GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
+ GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
+ GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
+
+ const uint32_t n_head_l = hparams.n_head(il);
+ const uint32_t n_head_kv_l = hparams.n_head_kv(il);
+
+ const float freq_base_l = model.get_rope_freq_base(cparams, il);
+ const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
+
+ auto inp = std::make_unique<llm_graph_input_embd>(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, n_tokens);
+ ggml_set_input(inp->embd);
+ ggml_set_name(inp->embd, "mtp_h_input");
+
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+
+ ggml_tensor * h_input = inp->embd;
+ ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ cb(tok_embd, "mtp_tok_embd", il);
+
+ res->add_input(std::move(inp));
+
+ ggml_tensor * inp_pos = build_inp_pos();
+ auto * inp_attn = build_attn_inp_kv_iswa();
+
+ ggml_tensor * h_norm = build_norm(h_input, 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);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ // mtp_block: full Step3p5 decoder layer (attention with optional head-wise gate, then MoE/dense FFN)
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ ggml_tensor * Qcur = build_lora_mm(layer.wq, cur, layer.wq_s);
+ ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s);
+ ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s);
+ cb(Qcur, "mtp_Qcur", il);
+ cb(Kcur, "mtp_Kcur", il);
+ cb(Vcur, "mtp_Vcur", il);
+
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
+ Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
+
+ if (layer.attn_q_norm) {
+ Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "mtp_Qcur_normed", il);
+ }
+ if (layer.attn_k_norm) {
+ Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Kcur, "mtp_Kcur_normed", il);
+ }
+
+ const bool is_swa = hparams.is_swa(il);
+ ggml_tensor * rope_factors = is_swa ? nullptr : model.get_rope_factors(cparams, il);
+ const int64_t n_rot_l = hparams.n_rot(il);
+
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, rope_factors,
+ n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, rope_factors,
+ n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(Qcur, "mtp_Qcur_pos", il);
+ cb(Kcur, "mtp_Kcur_pos", il);
+
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
+ ggml_tensor * attn_out = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(attn_out, "mtp_attn_out", il);
+
+ // head-wise attention gate: sigmoid(g_proj(x))
+ if (layer.wqkv_gate) {
+ ggml_tensor * gate = build_lora_mm(layer.wqkv_gate, cur); // [n_head_l, n_tokens]
+ cb(gate, "mtp_attn_gate", il);
+
+ gate = ggml_sigmoid(ctx0, gate);
+ cb(gate, "mtp_attn_gate_sigmoid", il);
+
+ ggml_tensor * attn_3d = ggml_reshape_3d(ctx0, attn_out, n_embd_head_v, n_head_l, n_tokens);
+ ggml_tensor * gate_3d = ggml_reshape_3d(ctx0, gate, 1, n_head_l, n_tokens);
+ cb(gate_3d, "mtp_attn_gate_3d", il);
+
+ attn_3d = ggml_mul(ctx0, attn_3d, gate_3d);
+ cb(attn_3d, "mtp_attn_gated_3d", il);
+
+ attn_out = ggml_reshape_2d(ctx0, attn_3d, n_embd_head_v * n_head_l, n_tokens);
+ cb(attn_out, "mtp_attn_gated", il);
+ }
+
+ cur = build_lora_mm(layer.wo, attn_out, layer.wo_s);
+ cb(cur, "mtp_attn_proj", il);
+
+ cur = ggml_add(ctx0, cur, inpSA);
+ cb(cur, "mtp_attn_residual", il);
+
+ ggml_tensor * ffn_inp = cur;
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ // FFN: dense MLP or MoE (mirrors trunk path)
+ if (layer.ffn_gate_inp == nullptr) {
+ cur = build_ffn(cur,
+ layer.ffn_up, layer.ffn_up_b, nullptr,
+ layer.ffn_gate, layer.ffn_gate_b, nullptr,
+ layer.ffn_down, layer.ffn_down_b, nullptr,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "mtp_ffn_out", il);
+ } else {
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ ggml_tensor * sh_out = build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, nullptr,
+ layer.ffn_gate_shexp, nullptr, nullptr,
+ layer.ffn_down_shexp, nullptr, nullptr,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(sh_out, "mtp_ffn_shared_out", il);
+
+ cur = ggml_add(ctx0, moe_out, sh_out);
+ cb(cur, "mtp_ffn_out", il);
+ }
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "mtp_post_ffn", il);
+
+ // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "STEP35 MTP: missing both nextn.shared_head_norm and output_norm");
+ cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "mtp_shared_head_norm", -1);
+
+ ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+ GGML_ASSERT(head_w && "STEP35 MTP: missing LM head (nextn.shared_head_head or model.output)");
+ cur = build_lora_mm(head_w, cur);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
hparams.dec_start_token_id = dec_start_token_id;
}
- hparams.dec_n_layer = hparams.n_layer;
+ hparams.dec_n_layer = hparams.n_layer();
ml.get_key(LLM_KV_DECODER_BLOCK_COUNT, hparams.dec_n_layer, false);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_60M; break; // t5-small
case 8: type = LLM_TYPE_80M; break; // flan-t5-small
case 12:
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);
- switch (hparams.n_layer) {
+ switch (hparams.n_layer()) {
case 40: type = LLM_TYPE_13B; break;
default: type = LLM_TYPE_UNKNOWN;
}
void llama_model_xverse::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
- switch (hparams.n_layer) {
+
+ switch (hparams.n_layer()) {
case 32: type = LLM_TYPE_7B; break;
case 40: type = LLM_TYPE_13B; break;
case 80: type = LLM_TYPE_65B; break;
overview / pkg / ggml / sources / whisper.cpp / commitdiff