set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
+ file(GLOB SRC_KV_CACHE llama-kv-cache-*.cpp)
+ file(GLOB SRC_MEMORY llama-memory-*.cpp)
file(GLOB SRC_MODELS models/*.cpp)
set(TARGET whisper-talk-llama)
llama-impl.cpp
llama-io.cpp
llama-kv-cache.cpp
- llama-kv-cache-iswa.cpp
- llama-kv-cache-dsa.cpp
- llama-kv-cache-dsv4.cpp
- llama-memory-recurrent.cpp
- llama-memory-hybrid.cpp
- llama-memory-hybrid-iswa.cpp
+ ${SRC_KV_CACHE}
llama-memory.cpp
+ ${SRC_MEMORY}
llama-mmap.cpp
llama-model-loader.cpp
llama-model-saver.cpp
{ LLM_ARCH_DOTS1, "dots1" },
{ LLM_ARCH_ARCEE, "arcee" },
{ LLM_ARCH_AFMOE, "afmoe" },
+ { LLM_ARCH_LAGUNA, "laguna" },
{ LLM_ARCH_ERNIE4_5, "ernie4_5" },
{ LLM_ARCH_ERNIE4_5_MOE, "ernie4_5-moe" },
{ LLM_ARCH_HUNYUAN_MOE, "hunyuan-moe" },
{ LLM_ARCH_HUNYUAN_DENSE, "hunyuan-dense" },
{ LLM_ARCH_HUNYUAN_VL, "hunyuan_vl" },
+ { LLM_ARCH_HY_V3, "hy_v3" },
{ LLM_ARCH_SMOLLM3, "smollm3" },
{ LLM_ARCH_OPENAI_MOE, "gpt-oss" },
{ LLM_ARCH_LFM2, "lfm2" },
{ LLM_ARCH_GROVEMOE, "grovemoe" },
{ LLM_ARCH_APERTUS, "apertus" },
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
+ { LLM_ARCH_MINIMAX_M3, "minimax-m3" },
{ LLM_ARCH_COGVLM, "cogvlm" },
{ LLM_ARCH_RND1, "rnd1" },
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_TALKIE, "talkie" },
{ LLM_ARCH_MELLUM, "mellum" },
+ { LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
{ LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" },
{ LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" },
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
+ { LLM_KV_NUM_LOOPS, "%s.num_loops" },
+ { LLM_KV_SKIP_LOOP_FINAL_NORM, "%s.skip_loop_final_norm" },
{ LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" },
{ LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" },
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
{ 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_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" },
+ { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" },
+ { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
{ LLM_KV_TARGET_LAYERS, "%s.target_layers" },
{ LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" },
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
+ { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },
{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
// sentence-transformers dense modules feature dims
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
+ { LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
+ { LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
+ { LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
{ LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" },
{ LLM_TENSOR_FC, "fc" },
{ LLM_TENSOR_D2T, "d2t" },
+ { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
+ { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
+ { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
};
// declare information about the model weight tensors:
{LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
- {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+ {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_ATTN_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_ATTN_K_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_V_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{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_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
- {LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+ {LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_INDEXER_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_FFN_GATE_TID2EID, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
// eagle3
{LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
{LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
+ // dspark
+ {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
+ {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
};
LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
+ case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
switch (arch) {
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
+ case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
case LLM_ARCH_LFM2:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_MINIMAX_M2:
+ case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
return false;
LLM_ARCH_DOTS1,
LLM_ARCH_ARCEE,
LLM_ARCH_AFMOE,
+ LLM_ARCH_LAGUNA,
LLM_ARCH_ERNIE4_5,
LLM_ARCH_ERNIE4_5_MOE,
LLM_ARCH_HUNYUAN_MOE,
LLM_ARCH_HUNYUAN_DENSE,
LLM_ARCH_HUNYUAN_VL,
+ LLM_ARCH_HY_V3,
LLM_ARCH_SMOLLM3,
LLM_ARCH_OPENAI_MOE,
LLM_ARCH_LFM2,
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
+ LLM_ARCH_MINIMAX_M3,
LLM_ARCH_DFLASH,
+ LLM_ARCH_NANBEIGE,
LLM_ARCH_UNKNOWN,
};
LLM_KV_TOKEN_SHIFT_COUNT,
LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
LLM_KV_FULL_ATTENTION_INTERVAL,
+ LLM_KV_NUM_LOOPS,
+ LLM_KV_SKIP_LOOP_FINAL_NORM,
LLM_KV_ATTENTION_HEAD_COUNT,
LLM_KV_ATTENTION_HEAD_COUNT_KV,
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
LLM_KV_ATTENTION_INDEXER_TOP_K,
+ LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,
+ LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,
+ LLM_KV_ATTENTION_INDEXER_TYPES,
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
LLM_KV_TARGET_LAYERS,
LLM_KV_TARGET_HIDDEN_SIZE,
LLM_KV_NORM_BEFORE_RESIDUAL,
+ LLM_KV_NORM_BEFORE_FC,
LLM_KV_SHORTCONV_L_CACHE,
LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K,
LLM_TENSOR_INDEXER_ATTN_Q_B,
+ LLM_TENSOR_INDEXER_Q_PROJ,
+ LLM_TENSOR_INDEXER_K_PROJ,
+ LLM_TENSOR_INDEXER_Q_NORM,
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
LLM_TENSOR_MASKED_EMBD_ORDERING,
LLM_TENSOR_FC,
LLM_TENSOR_D2T,
+ LLM_TENSOR_DSPARK_MARKOV_W1,
+ LLM_TENSOR_DSPARK_MARKOV_W2,
+ LLM_TENSOR_DSPARK_CONF_PROJ,
};
/*.n_tokens_per_seq =*/ 16,
};
+static const llm_fused_op_probe llm_fused_op_lid_probe = {
+ /*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER,
+ /*.name =*/ "Lightning Indexer",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = {
+ /*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE,
+ /*.name =*/ "fused DeepSeek V4 HC pre",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_dsv4_hc_comb_probe = {
+ /*.op =*/ LLM_FUSED_OP_DSV4_HC_COMB,
+ /*.name =*/ "fused DeepSeek V4 HC comb",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {
+ /*.op =*/ LLM_FUSED_OP_DSV4_HC_POST,
+ /*.name =*/ "fused DeepSeek V4 HC post",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
llama_context::llama_context(
const llama_model & model,
llama_context_params params) :
cparams.no_perf = params.no_perf;
cparams.warmup = false;
- cparams.embeddings_layer_inp.resize(hparams.n_layer(), false);
- embd_layer_inp.resize(hparams.n_layer());
+ // +1: id n_layer() taps the output of the last layer ("input" of the head)
+ cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);
+ embd_layer_inp.resize(hparams.n_layer() + 1);
cparams.ctx_type = params.ctx_type;
cparams.pooling_type = params.pooling_type;
cparams.fused_gdn_ch = true;
cparams.auto_fgdn = true;
+ cparams.fused_lid = true;
+ cparams.auto_flid = true;
+
+ cparams.fused_dsv4_hc_pre = true;
+ cparams.fused_dsv4_hc_comb = true;
+ cparams.fused_dsv4_hc_post = true;
+ cparams.auto_fhc = true;
+
// with causal attention, the batch size is limited by the context size
cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
}
llama_context::~llama_context() {
+ // wait for any pending asynchronous copies into the output buffers before they are freed
+ synchronize();
+
if (!model.hparams.no_alloc) {
for (size_t i = 0; i < backend_ptrs.size(); ++i) {
ggml_backend_t backend = backend_ptrs[i];
resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
cparams.auto_fgdn = false;
}
+
+ if (cparams.auto_flid) {
+ LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func);
+ resolve(llm_fused_op_lid_probe, cparams.fused_lid);
+ cparams.auto_flid = false;
+ }
+
+ if (cparams.auto_fhc) {
+ LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 HC support:\n", func);
+ resolve(llm_fused_op_dsv4_hc_pre_probe, cparams.fused_dsv4_hc_pre);
+ resolve(llm_fused_op_dsv4_hc_comb_probe, cparams.fused_dsv4_hc_comb);
+ resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);
+ cparams.auto_fhc = false;
+ }
}
void llama_context::sched_reserve() {
void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {
LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable);
- GGML_ASSERT(lid < model.hparams.n_layer());
+ GGML_ASSERT(lid <= model.hparams.n_layer());
cparams.embeddings_layer_inp[lid] = enable;
// micro-batching is not possible for non-causal encoding, so we process the batch in a single shot
GGML_ASSERT(cparams.n_ubatch >= n_tokens && "encoder requires n_ubatch >= n_tokens");
+ // TODO: this clear of the buffer can easily be forgotten - need something better
+ // sync first so any in-flight async copies into embd_seq complete before it is freed
+ if (!embd_seq.empty()) {
+ synchronize();
+ }
+ embd_seq.clear();
+
if (t_compute_start_us == 0) {
t_compute_start_us = ggml_time_us();
}
- // TODO: this clear of the buffer can easily be forgotten - need something better
- embd_seq.clear();
-
sched_reserve();
n_queued_tokens += n_tokens;
const auto & hparams = model.hparams;
const int64_t n_vocab = vocab.n_tokens();
- const int64_t n_embd = hparams.n_embd_inp();
+ const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;
+ const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
// when computing embeddings, all tokens are output
const bool output_all = cparams.embeddings;
GGML_ASSERT((cparams.causal_attn || cparams.n_ubatch >= n_tokens_all) && "non-causal attention requires n_ubatch >= n_tokens");
+ // TODO: this clear of the buffer can easily be forgotten - need something better
+ // sync first so any in-flight async copies into embd_seq complete before it is freed
+ if (!embd_seq.empty()) {
+ synchronize();
+ }
+ embd_seq.clear();
+
if (t_compute_start_us == 0) {
t_compute_start_us = ggml_time_us();
}
n_queued_tokens += n_tokens_all;
- // TODO: this clear of the buffer can easily be forgotten - need something better
- embd_seq.clear();
output_swaps.clear();
sched_reserve();
}
void llama_context::output_reorder() {
- const uint64_t n_vocab = model.vocab.n_tokens();
- const uint64_t n_embd = model.hparams.n_embd;
+ const uint64_t n_vocab = model.vocab.n_tokens();
+ const uint64_t n_embd = model.hparams.n_embd;
+ const uint64_t n_embd_out = model.hparams.n_embd_out();
for (size_t s = 0; s < output_swaps.size(); ++s) {
const uint64_t i0 = output_swaps[s].i0;
}
if (embd.size > 0) {
- for (uint64_t k = 0; k < n_embd; k++) {
- std::swap(embd.data[i0*n_embd + k], embd.data[i1*n_embd + k]);
+ for (uint64_t k = 0; k < n_embd_out; k++) {
+ std::swap(embd.data[i0*n_embd_out + k], embd.data[i1*n_embd_out + k]);
}
}
if (embd_nextn.size > 0) {
- for (uint64_t k = 0; k < n_embd; k++) {
- std::swap(embd_nextn.data[i0*n_embd + k], embd_nextn.data[i1*n_embd + k]);
+ for (uint64_t k = 0; k < n_embd_out; k++) {
+ std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]);
}
}
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
- model.arch == LLM_ARCH_DEEPSEEK4) {
+ model.arch == LLM_ARCH_DEEPSEEK4 ||
+ (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
+ model.arch == LLM_ARCH_NANBEIGE ||
+ model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
ggml_set_name(cur, name);
}
- // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
+ // - norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
+ // - force the last op of the layer on the specified backend to avoid running it on the backend of the next layer due to scheduling
// FIXME: fix in ggml_backend_sched
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) {
+ if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) {
const auto & dev_layer = model.dev_layer(il);
for (const auto & backend : backends) {
if (ggml_backend_get_device(backend.get()) == dev_layer) {
}
}
+ if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) {
+ LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v));
+ return nullptr;
+ }
+
+ if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
+ if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
+ LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
+ params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+ }
+ if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
+ LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\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) {
}
}
- if (ggml_is_quantized(params.type_v) && params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
- LLAMA_LOG_ERROR("%s: V cache quantization requires flash_attn\n", __func__);
- return nullptr;
- }
-
if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED &&
params.pooling_type != model->hparams.pooling_type) {
//user-specified pooling-type is different from the model default
bool fused_gdn_ar; // use fused gated delta net (autoregressive)
bool fused_gdn_ch; // use fused gated delta net (chunked)
bool auto_fgdn;
+ bool fused_lid; // use fused lightning indexer
+ bool auto_flid;
+ bool fused_dsv4_hc_pre;
+ bool fused_dsv4_hc_comb;
+ bool fused_dsv4_hc_post;
+ bool auto_fhc;
bool no_perf;
bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
bool op_offload;
vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
}
+ // Validate that all rule references point to valid rules
+ for (size_t i = 0; i < n_rules; i++) {
+ for (const auto & elem : vec_rules[i]) {
+ if (elem.type == LLAMA_GRETYPE_RULE_REF) {
+ if (elem.value >= n_rules || vec_rules[elem.value].empty()) {
+ LLAMA_LOG_ERROR("invalid grammar: rule %zu references undefined rule %u\n", i, elem.value);
+ return nullptr;
+ }
+ }
+ }
+ }
+
// Check for left recursion
std::vector<bool> rules_visited(n_rules);
std::vector<bool> rules_in_progress(n_rules);
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
#include "llama-kv-cache-dsa.h"
+#include "llama-kv-cache-msa.h"
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
return res;
}
+llm_graph_input_attn_kv_msa::llm_graph_input_attn_kv_msa(
+ const llama_hparams & hparams,
+ const llama_cparams & cparams,
+ const llama_kv_cache_msa_context * mctx) :
+ llm_graph_input_attn_kv(hparams, cparams, mctx->get_base()),
+ mctx_msa(mctx) {
+}
+
+void llm_graph_input_attn_kv_msa::set_input(const llama_ubatch * ubatch) {
+ llm_graph_input_attn_kv::set_input(ubatch);
+
+ if (self_k_idxs_idx) {
+ mctx_msa->get_idx()->set_input_k_idxs(self_k_idxs_idx, ubatch);
+ }
+}
+
+bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
+ mctx_msa = static_cast<const llama_kv_cache_msa_context *>(params.mctx);
+
+ // the parent class operates on the base cache context
+ this->mctx = mctx_msa->get_base();
+
+ bool res = true;
+
+ res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
+ if (self_k_idxs_idx) {
+ res &= self_k_idxs_idx->ne[0] == params.ubatch.n_tokens;
+ }
+
+ res &= can_reuse_kq_mask(self_kq_mask, this->mctx, params.ubatch, params.cparams);
+
+ 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);
return res;
}
+void llm_graph_input_attn_k_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);
+ }
+
+ // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live
+ if (self_kq_mask && self_kq_mask->buffer) {
+ 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);
+ }
+
+ if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
+ mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
+ }
+
+ if (self_k_rot && self_k_rot->buffer) {
+ mctx->get_base()->set_input_k_rot(self_k_rot);
+ }
+
+ if (self_k_rot_swa && self_k_rot_swa->buffer) {
+ mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
+ }
+}
+
+bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
+ const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);
+
+ this->mctx = mctx;
+
+ bool res = true;
+
+ // base tensors may not be allocated if there are no non-SWA attention layers
+ if (self_k_idxs && self_k_idxs->buffer) {
+ res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
+ }
+
+ if (self_kq_mask && self_kq_mask->buffer) {
+ 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;
+ }
+
+ if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
+ res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
+ }
+
+ return res;
+}
+
static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & src) {
if (!dst || !dst->buffer) {
return;
dsv4_set_i32(inp.state_pos, plan.state_pos);
dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs);
dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs);
+ dsv4_set_i32(inp.state_restore_src_idxs, plan.state_restore_src_idxs);
+ dsv4_set_i32(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs);
+ dsv4_set_i32(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs);
+ dsv4_set_i32(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs);
dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs);
dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs);
dsv4_set_i32(inp.state_write_pos, plan.state_write_pos);
res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size());
+ res &= dsv4_can_reuse_tensor_1d(inp.state_restore_src_idxs, plan.state_restore_src_idxs.size());
+ res &= dsv4_can_reuse_tensor_1d(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs.size());
+ res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs.size());
+ res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size());
inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos");
inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs");
inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs");
+ inp.state_restore_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_src_idxs.size(), std::string("dsv4_") + name + "_state_restore_src_idxs");
+ inp.state_restore_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_dst_idxs.size(), std::string("dsv4_") + name + "_state_restore_dst_idxs");
+ inp.state_snapshot_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_src_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_src_idxs");
+ inp.state_snapshot_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_dst_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_dst_idxs");
inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs");
inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs");
inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos");
GGML_ASSERT(n_stream > 0);
GGML_ASSERT(n_tokens%n_stream == 0);
- inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
+ inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
ggml_set_input(inp.kq_mask);
ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str());
}
t_embd_pooled = nullptr;
t_h_nextn = nullptr;
- t_layer_inp.resize(LLAMA_MAX_LAYERS);
+ t_layer_inp.resize(LLAMA_MAX_LAYERS + 1);
std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr);
t_sampled.clear();
tmp = ggml_clamp(ctx0, tmp, -limit, limit);
cb(tmp, "ffn_up_clamped", il);
- if (arch == LLM_ARCH_DEEPSEEK4) {
+ if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
cb(cur, "ffn_gate_clamped", il);
cur = ggml_swiglu_split(ctx0, cur, tmp);
cur = ggml_swiglu(ctx0, cur);
cb(cur, "ffn_swiglu", il);
} break;
+ case LLM_FFN_SWIGLU_OAI_MOE:
+ if (gate && type_gate == LLM_FFN_PAR) {
+ // same alpha/limit constants as gpt-oss
+ const float alpha = 1.702f;
+ const float limit = 7.0f;
+ cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit);
+ cb(cur, "ffn_swiglu_oai", il);
+ type_gate = LLM_FFN_SEQ;
+ } else {
+ GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate");
+ } break;
case LLM_FFN_GEGLU:
{
cur = ggml_geglu(ctx0, cur);
up = ggml_clamp(ctx0, up, -limit, limit);
cb(up, "ffn_moe_up_clamped", il);
- if (arch == LLM_ARCH_DEEPSEEK4) {
+ if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
cb(cur, "ffn_moe_gate_clamped", il);
cur = ggml_swiglu_split(ctx0, cur, up);
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
- const auto & kq_mask = inp->get_kq_mask();
+ ggml_tensor * kq_mask = inp->get_kq_mask();
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
return cur;
}
+ggml_tensor * llm_graph_context::build_attn(
+ llm_graph_input_attn_k_iswa * 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,
+ float kq_scale,
+ int il) const {
+ const bool is_swa = hparams.is_swa(il);
+
+ GGML_UNUSED(v_cur);
+
+ auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
+
+ if (k_rot) {
+ q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
+ if (k_cur) {
+ k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot);
+ }
+ }
+
+ // 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
+ ggml_build_forward_expand(gf, q_cur);
+
+ if (k_cur) {
+ ggml_build_forward_expand(gf, k_cur);
+ }
+
+ const auto * mctx_iswa = inp->mctx;
+ const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
+
+ // optionally store to KV cache
+ if (k_cur) {
+ const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
+
+ ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
+ }
+
+ const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
+
+ // MLA-style attention: the cached K is used as V
+ ggml_tensor * q = q_cur;
+ ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+ ggml_tensor * v = k;
+
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
+ cb(cur, "kqv_out", il);
+
+ if (k_rot) {
+ cur = llama_mul_mat_hadamard(ctx0, cur, k_rot);
+ }
+
+ if (wo) {
+ cur = build_lora_mm(wo, cur, wo_s);
+ }
+
+ if (wo_b) {
+ cur = ggml_add(ctx0, cur, wo_b);
+ }
+
+ return cur;
+}
+
llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const {
auto inp = std::make_unique<llm_graph_input_attn_cross>(cross);
{
inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
- // ensure F32 mask
+ // ensure that mask type matches fused lightning indexer use (requires f16 mask)
auto cparams_copy = cparams;
- cparams_copy.flash_attn = false;
+ cparams_copy.flash_attn = cparams.fused_lid;
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;
return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));
}
+llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const {
+ const auto * mctx_cur = static_cast<const llama_kv_cache_msa_context *>(mctx);
+
+ auto inp = std::make_unique<llm_graph_input_attn_kv_msa>(hparams, cparams, mctx_cur);
+
+ const auto * mctx_base = mctx_cur->get_base();
+ const auto * mctx_idx = mctx_cur->get_idx();
+
+ {
+ GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
+
+ inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch);
+ inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch);
+
+ inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams);
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
+ }
+
+ inp->self_k_rot = mctx_base->build_input_k_rot(ctx0);
+ inp->self_v_rot = mctx_base->build_input_v_rot(ctx0);
+
+ if (msa_enabled) {
+ inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch);
+ }
+
+ return (llm_graph_input_attn_kv_msa *) 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.
return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));
}
+llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
+ const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
+
+ auto inp = std::make_unique<llm_graph_input_attn_k_iswa>(hparams, cparams, mctx_cur);
+
+ {
+ inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
+
+ inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
+ inp->self_kq_mask_cnv = inp->self_kq_mask;
+ }
+
+ {
+ GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA");
+
+ inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_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 = inp->self_kq_mask_swa;
+ }
+
+ inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
+
+ inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
+
+ return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp));
+}
+
llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_dsv4_context *>(mctx);
const auto * raw_ctx = mctx_cur->get_raw();
class llama_kv_cache_context;
class llama_kv_cache_dsa_context;
+class llama_kv_cache_msa_context;
class llama_kv_cache_dsv4_raw_context;
class llama_kv_cache_dsv4_context;
class llama_kv_cache_iswa_context;
LLM_FUSED_OP_FLASH_ATTN,
LLM_FUSED_OP_GDN_AR,
LLM_FUSED_OP_GDN_CH,
+ LLM_FUSED_OP_LIGHTNING_INDEXER,
+ LLM_FUSED_OP_DSV4_HC_PRE,
+ LLM_FUSED_OP_DSV4_HC_COMB,
+ LLM_FUSED_OP_DSV4_HC_POST,
};
enum llm_ffn_op_type : int {
const llama_kv_cache_dsa_context * mctx;
};
+// standard K/V attention input against the base cache, plus destination indices for the indexer key cache
+class llm_graph_input_attn_kv_msa : public llm_graph_input_attn_kv {
+public:
+ llm_graph_input_attn_kv_msa(
+ const llama_hparams & hparams,
+ const llama_cparams & cparams,
+ const llama_kv_cache_msa_context * mctx);
+ ~llm_graph_input_attn_kv_msa() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ ggml_tensor * get_k_idxs_idx() const { return self_k_idxs_idx; }
+
+ ggml_tensor * self_k_idxs_idx = nullptr; // I64 [n_batch]
+
+ const llama_kv_cache_msa_context * mctx_msa;
+};
+
class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {
public:
llm_graph_input_attn_kv_iswa(
const llama_kv_cache_iswa_context * mctx;
};
+class llm_graph_input_attn_k_iswa : public llm_graph_input_i {
+public:
+ llm_graph_input_attn_k_iswa(
+ const llama_hparams & hparams,
+ const llama_cparams & cparams,
+ const llama_kv_cache_iswa_context * mctx) :
+ hparams(hparams),
+ cparams(cparams),
+ mctx(mctx) {
+ }
+ ~llm_graph_input_attn_k_iswa() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ ggml_tensor * get_k_idxs() const { return self_k_idxs; }
+ ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
+
+ ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
+ ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
+
+ ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
+ ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
+
+ 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_k_rot_swa = nullptr;
+
+ const llama_hparams hparams;
+ const llama_cparams cparams;
+
+ const llama_kv_cache_iswa_context * mctx;
+};
+
// DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped
// so raw K can be concatenated with DSV4 compressed K in one attention op.
class llm_graph_input_dsv4_raw {
ggml_tensor * state_pos = nullptr; // I32 [n_state]
ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist]
ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist]
+ ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore]
+ ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore]
+ ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot]
+ ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot]
ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write]
ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]
ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write]
ggml_tensor * build_attn_mha(
ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens]
ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens]
- ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false)
+ ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
ggml_tensor * sinks, // [n_head_q]
llm_graph_input_attn_k_dsa * build_attn_inp_k_dsa() const;
+ llm_graph_input_attn_kv_msa * build_attn_inp_kv_msa(bool msa_enabled) const;
+
ggml_tensor * build_attn(
llm_graph_input_attn_k_dsa * inp,
ggml_tensor * wo,
float kq_scale,
int il) const;
+ llm_graph_input_attn_k_iswa * build_attn_inp_k_iswa() const;
+
+ // note: if k_cur is not provided, it will not be stored in the memory
+ // note: the K cache is used as V (MLA-style attention)
+ ggml_tensor * build_attn(
+ llm_graph_input_attn_k_iswa * 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] optional
+ ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional
+ 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]
+ float kq_scale,
+ int il) const;
+
llm_graph_input_attn_cross * build_attn_inp_cross() const;
ggml_tensor * build_attn(
return n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0;
}
+bool llama_hparams::is_indexer_full(uint32_t il) const {
+ if (il < n_layer()) {
+ return is_indexer_full_impl[il];
+ }
+
+ GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer());
+}
+
uint32_t llama_hparams::n_embd_head_k_mla() const {
return is_mla() ? n_embd_head_k_mla_impl : n_embd_head_k();
}
bool use_par_res;
bool swin_norm;
bool norm_before_residual = false;
+ bool norm_before_fc = false;
uint32_t n_ctx_train; // context size the model was trained on
uint32_t n_embd;
uint32_t indexer_n_head = 0;
uint32_t indexer_head_size = 0;
uint32_t indexer_top_k = 0;
+ // MSA
+ uint32_t indexer_block_size = 0;
+ uint32_t indexer_local_blocks = 0;
+
+ // Indexer is "full" (1) or "shared" (0)
+ // Shared indexers reuse top-k from previous full layer
+ std::array<uint32_t, LLAMA_MAX_LAYERS> is_indexer_full_impl;
// DeepSeek-V4
uint32_t dsv4_o_group_count = 0;
bool is_swa(uint32_t il) const;
+ bool is_indexer_full(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)
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
- const layer_filter_cb & filter,
+ const layer_filter_cb & filter_mla,
+ const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse) :
hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) {
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);
+ n_swa, swa_type, nullptr, filter_mla, reuse, nullptr);
// we use llama_kv_cache for caching indexer keys
// by hand-tweaking some hparams we fool it to create
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);
+ n_swa, swa_type, nullptr, filter_lid, reuse, nullptr);
}
void llama_kv_cache_dsa::clear(bool data) {
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
- const layer_filter_cb & filter,
+ const layer_filter_cb & filter_mla,
+ const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse);
~llama_kv_cache_dsa() = default;
static constexpr uint32_t DSV4_STATE_VERSION = 1;
static constexpr uint32_t DSV4_STATE_MODE_FULL = 0;
static constexpr uint32_t DSV4_STATE_MODE_PARTIAL = 1;
-static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 1;
+static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 2;
static constexpr uint32_t DSV4_COMP_STATE_VER = 1;
static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
}
+static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {
+ GGML_ASSERT(ggml_is_contiguous(tensor));
+ GGML_ASSERT(tensor->ne[3] == 1);
+ GGML_ASSERT(stream < (uint32_t) tensor->ne[2]);
+
+ const size_t stream_size = tensor->nb[2];
+ ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);
+}
+
+static uint32_t dsv4_state_n_used_k_rows(llama_pos pos_max, uint32_t ratio, uint32_t kv_size) {
+ if (pos_max < 0) {
+ return 0;
+ }
+
+ const uint64_t n_rows = ((uint64_t) pos_max + 1)/ratio;
+
+ return (uint32_t) std::min<uint64_t>(kv_size, n_rows);
+}
+
static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
if (n_stream <= 1) {
return 0;
static void dsv4_state_write_tensor_streams(
llama_io_write_i & io,
ggml_tensor * tensor,
+ uint32_t tensor_rows,
uint32_t n_rows,
uint32_t s0,
- uint32_t ns) {
+ uint32_t ns,
+ const std::vector<uint32_t> * stream_ids = nullptr) {
const int32_t type_i = (int32_t) tensor->type;
const uint64_t ne0 = tensor->ne[0];
const uint64_t rows = n_rows;
const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]);
+ if (n_rows > tensor_rows) {
+ throw std::runtime_error("DSV4 state tensor row count exceeds storage");
+ }
+
io.write(&type_i, sizeof(type_i));
io.write(&ne0, sizeof(ne0));
io.write(&rows, sizeof(rows));
io.write(&row_size, sizeof(row_size));
- const size_t offset = (size_t) s0*n_rows*row_size;
- const size_t size = (size_t) ns*n_rows*row_size;
+ const size_t stream_stride = (size_t) tensor_rows*row_size;
+ const size_t size = (size_t) n_rows*row_size;
+ if (size == 0) {
+ return;
+ }
+
+ if (stream_ids && stream_ids->size() != ns) {
+ throw std::runtime_error("DSV4 state tensor stream map size mismatch");
+ }
- io.write_tensor(tensor, offset, size);
+ for (uint32_t s = 0; s < ns; ++s) {
+ const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s;
+ if ((int64_t) stream >= tensor->ne[2]) {
+ throw std::runtime_error("DSV4 state tensor stream out of range");
+ }
+ const size_t offset = (size_t) stream*stream_stride;
+ io.write_tensor(tensor, offset, size);
+ }
}
static void dsv4_state_read_tensor_streams(
llama_io_read_i & io,
ggml_tensor * tensor,
+ uint32_t tensor_rows,
uint32_t n_rows,
uint32_t s0,
uint32_t ns) {
if (type_i != type_i_ref || ne0 != ne0_ref || rows != rows_ref || row_size != row_size_ref) {
throw std::runtime_error("DSV4 state tensor metadata mismatch");
}
+ if (n_rows > tensor_rows) {
+ throw std::runtime_error("DSV4 state tensor row count exceeds storage");
+ }
- const size_t offset = (size_t) s0*n_rows*row_size;
- const size_t size = (size_t) ns*n_rows*row_size;
+ const size_t stream_stride = (size_t) tensor_rows*row_size;
+ const size_t size = (size_t) n_rows*row_size;
+ if (size == 0) {
+ return;
+ }
- io.read_tensor(tensor, offset, size);
+ for (uint32_t s = 0; s < ns; ++s) {
+ const size_t offset = (size_t) (s0 + s)*stream_stride;
+ io.read_tensor(tensor, offset, size);
+ }
}
static void dsv4_state_write_k_cache(
llama_io_write_i & io,
const llama_kv_cache * kv,
llama_seq_id seq_id,
- llama_state_seq_flags flags) {
+ llama_state_seq_flags flags,
+ uint32_t n_rows) {
GGML_UNUSED(flags);
uint32_t s0;
const auto layer_ids = kv->get_layer_ids();
const uint32_t n_layer = layer_ids.size();
+ if (n_rows > kv_size) {
+ throw std::runtime_error("DSV4 K-cache state row count exceeds cache size");
+ }
+
io.write(&version, sizeof(version));
- io.write(&kv_size, sizeof(kv_size));
+ io.write(&n_rows, sizeof(n_rows));
io.write(&ns, sizeof(ns));
io.write(&n_layer, sizeof(n_layer));
for (uint32_t il : layer_ids) {
io.write(&il, sizeof(il));
- dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, s0, ns);
+ dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows, s0, ns);
}
}
GGML_UNUSED(flags);
uint32_t version;
- uint32_t kv_size_ref;
+ uint32_t n_rows_ref;
uint32_t ns;
uint32_t n_layer_ref;
io.read(&version, sizeof(version));
- io.read(&kv_size_ref, sizeof(kv_size_ref));
+ io.read(&n_rows_ref, sizeof(n_rows_ref));
io.read(&ns, sizeof(ns));
io.read(&n_layer_ref, sizeof(n_layer_ref));
- if (version != DSV4_K_CACHE_STATE_VER) {
+ if (version != 1 && version != DSV4_K_CACHE_STATE_VER) {
throw std::runtime_error("DSV4 K-cache state version mismatch");
}
- if (kv_size_ref != kv->get_size()) {
+
+ const uint32_t kv_size = kv->get_size();
+ if (version == 1 && n_rows_ref != kv_size) {
+ LLAMA_LOG_INFO("kv size ref %d kv %d\n", n_rows_ref, kv_size);
+ throw std::runtime_error("DSV4 K-cache state size mismatch");
+ }
+ if (n_rows_ref > kv_size) {
+ LLAMA_LOG_INFO("kv rows ref %d kv %d\n", n_rows_ref, kv_size);
throw std::runtime_error("DSV4 K-cache state size mismatch");
}
throw std::runtime_error("DSV4 K-cache layer id mismatch");
}
- dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv->get_size(), s0, ns);
+ dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
}
}
bool overlap,
uint32_t state_size,
uint32_t kv_size,
- uint32_t n_stream) {
+ uint32_t n_stream,
+ uint32_t n_rs_seq,
+ const std::vector<uint32_t> & rs_idx) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
std::vector<int32_t> overlap_cur_reads;
std::map<std::pair<llama_seq_id, llama_pos>, int64_t> curr_token_idx_map;
+ std::map<llama_seq_id, uint32_t> state_write_counts;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
plan.state_write_idxs.push_back(cache_off + pos/ratio);
plan.state_write_pos.push_back((int32_t) source_start);
+ ++state_write_counts[seq_id];
if (overlap) {
const llama_pos prev_start = source_start - ratio;
}
}
- if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) {
- // Non-boundary CSA steps still need a write op so their graph matches
- // boundary steps. Use a padded scratch row that is masked from attention.
+ if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
assert(kv_size > 0);
- uint32_t i = 0;
- while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
- ++i;
- }
- assert(i < ubatch.n_tokens);
+ // Pad each stream to the reserve plan's block count.
+ const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) {
+ const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
+ const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);
- const llama_pos pos = ubatch.pos[i];
- const llama_seq_id seq_id = ubatch.seq_id[i][0];
- const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
- const int32_t source_idx = state_source_idx(seq_id, pos);
+ plan.state_write_idxs.push_back(cache_off + kv_size - 1);
+ plan.state_write_pos .push_back(0);
- plan.state_write_idxs.push_back(cache_off + kv_size - 1);
- plan.state_write_pos .push_back(0);
+ if (overlap) {
+ for (uint32_t j = 0; j < ratio; ++j) {
+ overlap_prev_reads.push_back(source_idx);
+ overlap_cur_reads .push_back(source_idx);
+ }
+ } else {
+ for (uint32_t j = 0; j < ratio; ++j) {
+ plan.state_read_idxs.push_back(source_idx);
+ }
+ }
+ };
- if (overlap) {
- for (uint32_t j = 0; j < ratio; ++j) {
- overlap_prev_reads.push_back(source_idx);
- overlap_cur_reads .push_back(source_idx);
+ if (dsv4_ubatch_has_coupled(ubatch)) {
+ if (plan.state_write_idxs.empty()) {
+ uint32_t i = 0;
+ while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
+ ++i;
+ }
+ assert(i < ubatch.n_tokens);
+ append_dummy_block(ubatch.seq_id[i][0], i);
}
} else {
- for (uint32_t j = 0; j < ratio; ++j) {
- plan.state_read_idxs.push_back(source_idx);
+ const uint32_t n_blocks = (std::max<uint32_t>(1, ubatch.n_seq_tokens) + ratio - 1)/ratio;
+
+ for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+ const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+ const uint32_t n_writes = state_write_counts[seq_id];
+ if (n_writes >= n_blocks) {
+ continue;
+ }
+ if (n_writes + 1 != n_blocks) {
+ throw std::runtime_error("DSV4 CSA sequence positions are not contiguous");
+ }
+
+ uint32_t i = 0;
+ while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) {
+ ++i;
+ }
+ assert(i < ubatch.n_tokens);
+ append_dummy_block(seq_id, i);
}
}
}
plan.state_persist_dst_idxs.push_back(row.dst);
}
+
+ if (n_rs_seq > 0) {
+ for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+ const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+ if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
+ continue;
+ }
+
+ const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
+ const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
+ // Keep the restore graph fixed-width when no rollback is pending.
+ const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;
+ for (uint32_t r = 0; r < state_size; ++r) {
+ plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r));
+ plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r));
+ }
+
+ std::vector<uint32_t> token_idxs;
+ token_idxs.reserve(ubatch.n_tokens);
+ for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+ if (dsv4_token_has_seq(ubatch, i, seq_id)) {
+ token_idxs.push_back(i);
+ }
+ }
+ if (token_idxs.empty()) {
+ continue;
+ }
+
+ const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();
+ const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);
+ for (uint32_t d = 1; d <= n_rs_seq; ++d) {
+ const int64_t dst_plane = (int64_t) d*state_rows;
+
+ for (uint32_t r = 0; r < state_size; ++r) {
+ int32_t src;
+ if (d <= n_seq_tokens) {
+ const uint32_t prefix = n_seq_tokens - d;
+ src = (int32_t) (stream_off + r);
+
+ for (uint32_t j = 0; j < prefix; ++j) {
+ const uint32_t i_tok = token_idxs[j];
+ if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
+ src = (int32_t) (scratch_off + i_tok);
+ }
+ }
+ } else {
+ const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows;
+ src = (int32_t) (src_plane + stream_off + r);
+ }
+
+ plan.state_snapshot_src_idxs.push_back(src);
+ plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r));
+ }
+ }
+ }
+ }
+
static const bool debug = []() {
const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");
return env && atoi(env) > 0;
bool overlap,
uint32_t state_size,
uint32_t kv_size,
- uint32_t n_stream) {
+ uint32_t n_stream,
+ uint32_t n_rs_seq,
+ const std::vector<uint32_t> & rs_idx) {
std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
plans.reserve(ubatches.size());
for (const llama_ubatch & ubatch : ubatches) {
- plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream));
+ plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
}
return plans;
bool overlap,
uint32_t state_size,
uint32_t kv_size,
- uint32_t n_stream) {
+ uint32_t n_stream,
+ uint32_t n_rs_seq) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
const uint64_t state_rows = (uint64_t) state_size*n_stream;
const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);
+ const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq) : 0;
+ const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq);
plan.state_pos .resize(ubatch.n_tokens);
plan.state_persist_src_idxs.resize(n_persist);
plan.state_persist_dst_idxs.resize(n_persist);
+ plan.state_restore_src_idxs.resize(n_restore);
+ plan.state_restore_dst_idxs.resize(n_restore);
+ plan.state_snapshot_src_idxs.resize(n_snapshot);
+ plan.state_snapshot_dst_idxs.resize(n_snapshot);
plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks);
plan.state_write_idxs.resize(n_blocks);
plan.state_write_pos .resize(n_blocks);
uint32_t ratio,
uint32_t state_size,
uint32_t n_embd_state,
+ uint32_t n_rs_seq,
const char * name,
const llama_memory_i::layer_filter_cb & filter) :
ratio(ratio),
state_size(state_size),
n_embd_state(n_embd_state),
- n_stream(unified ? 1 : n_seq_max) {
+ n_stream(unified ? 1 : n_seq_max),
+ n_rs_seq(n_rs_seq) {
const llama_hparams & hparams = model.hparams;
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*hparams.n_layer()*ggml_tensor_overhead()),
+ /*.mem_size =*/ size_t(2u*(1 + n_stream)*hparams.n_layer()*ggml_tensor_overhead()),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
throw std::runtime_error("failed to create ggml context for DSV4 compressor state");
}
- ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
- ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
+ const uint32_t n_planes = n_stream*(1 + n_rs_seq);
+ ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
+ ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il);
ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
+ std::vector<ggml_tensor *> kv_stream;
+ std::vector<ggml_tensor *> score_stream;
+
+ for (uint32_t s = 0; s < n_stream; ++s) {
+ kv_stream.push_back(ggml_view_2d(ctx, kv, n_embd_state, state_size, kv->nb[1], s*kv->nb[2]));
+ score_stream.push_back(ggml_view_2d(ctx, score, n_embd_state, state_size, score->nb[1], s*score->nb[2]));
+ }
+
map_layer_ids[il] = layers.size();
- layers.push_back({ il, kv, score });
+ layers.push_back({ il, kv, score, std::move(kv_stream), std::move(score_stream) });
}
for (auto & [buft, ctx] : ctx_map) {
ctxs_bufs.emplace_back(std::move(ctx), buf);
}
- LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n",
- __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
+ LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",
+ __func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0);
}
-void llama_dsv4_comp_state::clear(bool data) {
+void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
if (!data) {
return;
}
+ if (seq_id >= 0) {
+ GGML_ASSERT((uint32_t) seq_id < n_stream);
+
+ for (const auto & layer : layers) {
+ for (uint32_t d = 0; d <= n_rs_seq; ++d) {
+ const uint32_t stream = d*n_stream + (uint32_t) seq_id;
+ dsv4_clear_tensor_stream(layer.kv, stream);
+ dsv4_clear_tensor_stream(layer.score, stream);
+ }
+ }
+ return;
+ }
+
for (auto & [_, buf] : ctxs_bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
}
}
+void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) {
+ GGML_ASSERT(seq_id_src >= 0 && (uint32_t) seq_id_src < n_stream);
+ GGML_ASSERT(seq_id_dst >= 0 && (uint32_t) seq_id_dst < n_stream);
+
+ if (seq_id_src == seq_id_dst) {
+ return;
+ }
+
+ clear(seq_id_dst, true);
+
+ sc_info.ssrc.push_back((uint32_t) seq_id_src);
+ sc_info.sdst.push_back((uint32_t) seq_id_dst);
+}
+
+void llama_dsv4_comp_state::apply_copies(const stream_copy_info & sc_info) const {
+ for (size_t i = 0; i < sc_info.ssrc.size(); ++i) {
+ const uint32_t ssrc = sc_info.ssrc[i];
+ const uint32_t sdst = sc_info.sdst[i];
+
+ for (const auto & layer : layers) {
+ ggml_backend_tensor_copy(layer.kv_stream[ssrc], layer.kv_stream[sdst]);
+ ggml_backend_tensor_copy(layer.score_stream[ssrc], layer.score_stream[sdst]);
+ }
+ }
+}
+
uint32_t llama_dsv4_comp_state::get_ratio() const {
return ratio;
}
return n_stream;
}
+uint32_t llama_dsv4_comp_state::get_n_rs_seq() const {
+ return n_rs_seq;
+}
+
+uint32_t llama_dsv4_comp_state::get_n_rows() const {
+ return state_size*n_stream;
+}
+
std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const auto & [_, buf] : ctxs_bufs) {
return ret;
}
-void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+void llama_dsv4_comp_state::state_write(
+ llama_io_write_i & io,
+ llama_seq_id seq_id,
+ llama_state_seq_flags flags,
+ const std::vector<uint32_t> & rs_idx) const {
GGML_UNUSED(flags);
uint32_t s0;
uint32_t ns;
dsv4_state_src_stream_range(n_stream, seq_id, s0, ns);
+ std::vector<uint32_t> stream_ids(ns);
+ for (uint32_t s = 0; s < ns; ++s) {
+ const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s;
+ if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) {
+ throw std::runtime_error("DSV4 recurrent state rollback index out of range");
+ }
+ stream_ids[s] = rs_idx[seq]*n_stream + s0 + s;
+ }
+
const uint32_t version = DSV4_COMP_STATE_VER;
const uint32_t n_layer = layers.size();
for (const auto & layer : layers) {
io.write(&layer.il, sizeof(layer.il));
- dsv4_state_write_tensor_streams(io, layer.kv, state_size, s0, ns);
- dsv4_state_write_tensor_streams(io, layer.score, state_size, s0, ns);
+ dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns, &stream_ids);
+ dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids);
}
}
throw std::runtime_error("DSV4 compressor state layer id mismatch");
}
- dsv4_state_read_tensor_streams(io, layer.kv, state_size, s0, ns);
- dsv4_state_read_tensor_streams(io, layer.score, state_size, s0, ns);
+ dsv4_state_read_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);
+ dsv4_state_read_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
}
}
-ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
+ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
-
ggml_tensor * state = layers[ids].kv;
- return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
+ return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
}
-ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {
+ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
-
ggml_tensor * state = layers[ids].score;
- return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
+ return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
+}
+
+ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
+ ggml_tensor * state = get_kv_all(ctx, il);
+ const size_t row_size = ggml_row_size(state->type, state->ne[0]);
+
+ return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
+}
+
+ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {
+ ggml_tensor * state = get_score_all(ctx, il);
+ const size_t row_size = ggml_row_size(state->type, state->ne[0]);
+
+ return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
}
ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
- return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs);
+ return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs);
}
ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
- return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs);
+ return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs);
}
size_t llama_dsv4_comp_state::total_size() const {
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
+ uint32_t n_rs_seq,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse) :
hparams_raw(model.hparams),
hparams_csa(model.hparams),
hparams_hca(model.hparams),
hparams_lid(model.hparams),
- n_seq_max(n_seq_max) {
+ n_seq_max(n_seq_max),
+ n_rs_seq(n_rs_seq),
+ rs_idx(n_seq_max, 0) {
const layer_filter_cb filter_raw = [&](int32_t il) {
if (filter && !filter(il)) {
// Keep DSV4 KV/state streams per sequence even when public KV mode is unified.
const bool unified_raw = false;
+ hparams_raw.n_layer_nextn = 0;
+ hparams_csa.n_layer_nextn = 0;
+ hparams_hca.n_layer_nextn = 0;
+ hparams_lid.n_layer_nextn = 0;
+
LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__);
dsv4_make_k_only(hparams_raw);
csa_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
- 2*model.hparams.n_embd_head_k(), "csa", filter_csa);
+ 2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
hca_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO,
- model.hparams.n_embd_head_k(), "hca", filter_hca);
+ model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
lid_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
- 2*model.hparams.indexer_head_size, "lid", filter_csa);
+ 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
// DSV4 attention reads compressed-K / compressor-state rows that the current
// graph does not necessarily overwrite; uninitialized buffer contents would
// otherwise leak in (instance-specific garbage) and corrupt recall. Zero all
// compressed buffers up front so reads of un-written rows are deterministic.
- clear_compressed(true);
+ clear_compressed(-1, true);
}
llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
}
llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) {
- return std::make_unique<llama_kv_cache_dsv4_context>(this, lctx, optimize);
+ return std::make_unique<llama_kv_cache_dsv4_context>(
+ this,
+ lctx,
+ optimize,
+ std::move(csa_state->sc_info),
+ std::move(hca_state->sc_info),
+ std::move(lid_state->sc_info));
}
bool llama_kv_cache_dsv4::get_can_shift() const {
void llama_kv_cache_dsv4::clear(bool data) {
kv_raw->clear(data);
- clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows
+ clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows
}
bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
}
if (p0 > 0) {
- // DSV4 compressed cache rows are derived from running compressor state,
- // so arbitrary rollback is not reconstructible from the raw cache alone.
- // Allow the common prompt-cache cleanup no-op: remove [end, infinity).
- if (seq_id >= 0 && p0 > kv_raw->seq_pos_max(seq_id)) {
- return true;
+ if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
+ return false;
}
- return false;
+ const llama_pos pos_max = kv_raw->seq_pos_max(seq_id);
+ if (p0 > pos_max) {
+ bool res = true;
+
+ res = res & kv_raw->seq_rm(seq_id, p0, -1);
+ res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
+ res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
+ res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
+
+ return res;
+ }
+
+ if (n_rs_seq == 0) {
+ return false;
+ }
+
+ const llama_pos rollback = pos_max - (p0 - 1);
+ if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
+ return false;
+ }
+
+ const bool res = kv_raw->seq_rm(seq_id, p0, p1);
+ if (res) {
+ rs_idx[seq_id] = (uint32_t) rollback;
+ }
+
+ return res;
}
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
if (res) {
- clear_compressed(true);
+ clear_compressed(seq_id, true);
}
return res;
}
void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ GGML_ASSERT(p0 <= 0 && p1 < 0 && "DSV4 only supports full sequence copies");
+
kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
- clear_compressed(true);
+ kv_csa->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+ kv_hca->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+ kv_lid->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+
+ csa_state->seq_cp(seq_id_src, seq_id_dst);
+ hca_state->seq_cp(seq_id_src, seq_id_dst);
+ lid_state->seq_cp(seq_id_src, seq_id_dst);
+
+ if (seq_id_src != seq_id_dst) {
+ rs_idx[seq_id_dst] = 0;
+ }
}
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
+ GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max);
+
kv_raw->seq_keep(seq_id);
- clear_compressed(true);
+
+ for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) {
+ if (id == seq_id) {
+ continue;
+ }
+
+ kv_raw->seq_rm(id, -1, -1);
+ clear_compressed(id, true);
+ }
}
void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
kv_raw->seq_add(seq_id, p0, p1, shift);
- clear_compressed(true);
}
void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
kv_raw->seq_div(seq_id, p0, p1, d);
- clear_compressed(true);
}
llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const {
kv_raw->state_write(io, seq_id, flags);
if (!partial_only) {
- dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags);
- dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags);
- dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags);
+ const llama_pos pos_max = seq_id >= 0 ? kv_raw->seq_pos_max(seq_id) : -1;
+
+ //FIXME : note that we conflate token positions with rows, which is not true for multi-modal case.
+ const uint32_t n_rows_csa = seq_id >= 0 ?
+ dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_csa->get_size()) : kv_csa->get_size();
+ const uint32_t n_rows_hca = seq_id >= 0 ?
+ dsv4_state_n_used_k_rows(pos_max, DSV4_HCA_RATIO, kv_hca->get_size()) : kv_hca->get_size();
+ const uint32_t n_rows_lid = seq_id >= 0 ?
+ dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_lid->get_size()) : kv_lid->get_size();
+
+ dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa);
+ dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca);
+ dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
}
- csa_state->state_write(io, seq_id, flags);
- hca_state->state_write(io, seq_id, flags);
- lid_state->state_write(io, seq_id, flags);
+ csa_state->state_write(io, seq_id, flags, rs_idx);
+ hca_state->state_write(io, seq_id, flags, rs_idx);
+ lid_state->state_write(io, seq_id, flags, rs_idx);
}
void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
kv_raw->state_read(io, seq_id, flags);
if (!partial_only) {
+ kv_csa->clear(true);
+ kv_hca->clear(true);
+ kv_lid->clear(true);
+
dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);
dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags);
hca_state->state_read(io, seq_id, flags);
lid_state->state_read(io, seq_id, flags);
+ if (seq_id >= 0) {
+ GGML_ASSERT((uint32_t) seq_id < n_seq_max);
+ rs_idx[seq_id] = 0;
+ } else {
+ std::fill(rs_idx.begin(), rs_idx.end(), 0);
+ }
}
llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const {
return lid_state.get();
}
-void llama_kv_cache_dsv4::clear_compressed(bool data) {
- kv_csa->clear(data);
- kv_hca->clear(data);
- kv_lid->clear(data);
- csa_state->clear(data);
- hca_state->clear(data);
- lid_state->clear(data);
+uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
+ return n_rs_seq;
+}
+
+const std::vector<uint32_t> & llama_kv_cache_dsv4::get_rs_idx() const {
+ return rs_idx;
+}
+
+void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
+ if (n_rs_seq == 0) {
+ return;
+ }
+
+ for (const llama_ubatch & ubatch : ubatches) {
+ for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+ for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
+ const llama_seq_id seq_id = ubatch.seq_id[i][s];
+ if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) {
+ rs_idx[seq_id] = 0;
+ }
+ }
+ }
+ }
+}
+
+void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
+ if (seq_id < 0) {
+ kv_csa->clear(data);
+ kv_hca->clear(data);
+ kv_lid->clear(data);
+ } else {
+ GGML_ASSERT((uint32_t) seq_id < n_seq_max);
+
+ const auto clear_seq = [seq_id, data](llama_kv_cache * kv) {
+ kv->seq_rm(seq_id, -1, -1);
+
+ if (data) {
+ for (uint32_t il : kv->get_layer_ids()) {
+ dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id);
+ }
+ }
+ };
+
+ clear_seq(kv_csa.get());
+ clear_seq(kv_hca.get());
+ clear_seq(kv_lid.get());
+ }
+
+ csa_state->clear(seq_id, data);
+ hca_state->clear(seq_id, data);
+ lid_state->clear(seq_id, data);
+
+ if (seq_id >= 0) {
+ rs_idx[seq_id] = 0;
+ } else {
+ std::fill(rs_idx.begin(), rs_idx.end(), 0);
+ }
}
//
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
llama_kv_cache_dsv4 * kv,
llama_context * lctx,
- bool optimize) :
+ bool optimize,
+ stream_copy_info sc_info_csa,
+ stream_copy_info sc_info_hca,
+ stream_copy_info sc_info_lid) :
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(kv->get_raw(), lctx, optimize)),
ctx_csa_mem(kv->get_csa()->init_update(lctx, optimize)),
ctx_hca_mem(kv->get_hca()->init_update(lctx, optimize)),
ctx_lid_mem(kv->get_lid()->init_update(lctx, optimize)),
- ctx_csa(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_csa())),
- ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_hca())),
- ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_lid())),
csa_state(kv->get_csa_state()),
hca_state(kv->get_hca_state()),
lid_state(kv->get_lid_state()),
+ sc_info_csa(std::move(sc_info_csa)),
+ sc_info_hca(std::move(sc_info_hca)),
+ sc_info_lid(std::move(sc_info_lid)),
status(llama_memory_status_combine(
- llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
- llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) {
+ llama_memory_status_combine(
+ llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
+ llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status())),
+ this->sc_info_csa.empty() && this->sc_info_hca.empty() && this->sc_info_lid.empty() ?
+ LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)) {
}
llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
std::vector<llama_ubatch> ubatches_raw) :
ubatches(std::move(ubatches)),
plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
- kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())),
+ kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(),
+ kv->get_n_rs_seq(), kv->get_rs_idx())),
plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false,
- kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())),
- plans_lid(plans_csa),
+ kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(),
+ kv->get_n_rs_seq(), kv->get_rs_idx())),
+ plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
+ kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(),
+ kv->get_n_rs_seq(), kv->get_rs_idx())),
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(
kv->get_raw(),
std::move(sinfos_raw_base_write),
hca_state(kv->get_hca_state()),
lid_state(kv->get_lid_state()),
status(ctx_raw->get_status()) {
+ kv->reset_rs_idx_for_ubatches(this->ubatches);
}
llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default;
res = res & ctx_raw->apply();
+ if (ctx_csa_mem) {
+ res = res & ctx_csa_mem->apply();
+ res = res & ctx_hca_mem->apply();
+ res = res & ctx_lid_mem->apply();
+ }
+
+ if (ubatches.empty()) {
+ csa_state->apply_copies(sc_info_csa);
+ hca_state->apply_copies(sc_info_hca);
+ lid_state->apply_copies(sc_info_lid);
+ }
+
return res;
}
reserve_plan_csa = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
- csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream());
+ csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq());
return reserve_plan_csa;
}
reserve_plan_hca = dsv4_build_reserve_comp_plan(
ubatch, DSV4_HCA_RATIO, false,
- hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream());
+ hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq());
return reserve_plan_hca;
}
reserve_plan_lid = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
- lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream());
+ lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq());
return reserve_plan_lid;
}
class llama_dsv4_comp_state {
public:
+ using stream_copy_info = llama_kv_cache::stream_copy_info;
+
+ stream_copy_info sc_info;
+
llama_dsv4_comp_state(
const llama_model & model,
bool offload,
uint32_t ratio,
uint32_t state_size,
uint32_t n_embd_state,
+ uint32_t n_rs_seq,
const char * name,
const llama_memory_i::layer_filter_cb & filter);
- void clear(bool data);
+ void clear(llama_seq_id seq_id, bool data);
+ void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
+ void apply_copies(const stream_copy_info & sc_info) const;
- uint32_t get_ratio() const;
+ uint32_t get_ratio() const;
uint32_t get_state_size() const;
- uint32_t get_n_stream() const;
+ uint32_t get_n_stream() const;
+ uint32_t get_n_rs_seq() const;
+ uint32_t get_n_rows() const;
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
- void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const;
+ void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, const std::vector<uint32_t> & rs_idx) const;
void state_read (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags);
- ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const;
- ggml_tensor * get_score(ggml_context * ctx, int32_t il) const;
+ ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const;
+ ggml_tensor * get_score (ggml_context * ctx, int32_t il) const;
+ ggml_tensor * get_kv_all (ggml_context * ctx, int32_t il) const;
+ ggml_tensor * get_score_all(ggml_context * ctx, int32_t il) const;
ggml_tensor * cpy_kv (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
ggml_tensor * kv;
ggml_tensor * score;
+
+ std::vector<ggml_tensor *> kv_stream;
+ std::vector<ggml_tensor *> score_stream;
};
const uint32_t ratio;
const uint32_t state_size;
const uint32_t n_embd_state;
const uint32_t n_stream;
+ const uint32_t n_rs_seq;
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
// DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches.
// The compressed caches are storage only; DSV4-specific visibility and block
// planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4.
+// FIXME: currently the cache only supports non-unified mode even if unified flag is passed
+// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819
class llama_kv_cache_dsv4 : public llama_memory_i {
public:
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
+ uint32_t n_rs_seq,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse);
llama_dsv4_comp_state * get_hca_state() const;
llama_dsv4_comp_state * get_lid_state() const;
+ uint32_t get_n_rs_seq() const;
+ const std::vector<uint32_t> & get_rs_idx() const;
+ void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
+
private:
llama_hparams hparams_raw;
llama_hparams hparams_csa;
llama_hparams hparams_lid;
const uint32_t n_seq_max;
+ const uint32_t n_rs_seq;
+
+ std::vector<uint32_t> rs_idx;
std::unique_ptr<llama_kv_cache_iswa> kv_raw;
std::unique_ptr<llama_kv_cache> kv_csa;
std::unique_ptr<llama_dsv4_comp_state> hca_state;
std::unique_ptr<llama_dsv4_comp_state> lid_state;
- void clear_compressed(bool data);
+ void clear_compressed(llama_seq_id seq_id, bool data);
};
// DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept
class llama_kv_cache_dsv4_context : public llama_memory_context_i {
public:
using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
+ using stream_copy_info = llama_kv_cache::stream_copy_info;
struct comp_plan {
// Per-ubatch recipe for updating compressor state, committing completed
std::vector<int32_t> state_persist_src_idxs;
std::vector<int32_t> state_persist_dst_idxs;
+ // Device-side rollback restore copies snapshot planes back to the
+ // current compressor-state plane before the graph reads it.
+ std::vector<int32_t> state_restore_src_idxs;
+ std::vector<int32_t> state_restore_dst_idxs;
+
+ // Device-side rollback snapshots copy rows from the graph-local
+ // [persistent_state | current_ubatch_scratch] tensor into rollback
+ // planes after the graph has computed current-token compressor state.
+ std::vector<int32_t> state_snapshot_src_idxs;
+ std::vector<int32_t> state_snapshot_dst_idxs;
+
// Flattened source row ids used for state-backed commits. Source rows
// index the graph-local [persistent_state | current_ubatch_scratch]
// tensor. For overlapped compression the first half is previous rows
llama_kv_cache_dsv4_context(
llama_kv_cache_dsv4 * kv,
llama_context * lctx,
- bool optimize);
+ bool optimize,
+ stream_copy_info sc_info_csa,
+ stream_copy_info sc_info_hca,
+ stream_copy_info sc_info_lid);
llama_kv_cache_dsv4_context(
llama_kv_cache_dsv4 * kv,
const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_hca;
const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_lid;
- const llama_dsv4_comp_state * csa_state = nullptr;
- const llama_dsv4_comp_state * hca_state = nullptr;
- const llama_dsv4_comp_state * lid_state = nullptr;
+ llama_dsv4_comp_state * csa_state = nullptr;
+ llama_dsv4_comp_state * hca_state = nullptr;
+ llama_dsv4_comp_state * lid_state = nullptr;
+
+ stream_copy_info sc_info_csa;
+ stream_copy_info sc_info_hca;
+ stream_copy_info sc_info_lid;
bool reserve_plans = false;
mutable comp_plan reserve_plan_csa;
--- /dev/null
+#include "llama-kv-cache-msa.h"
+
+#include "llama-impl.h"
+#include "llama-batch.h"
+#include "llama-model.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+
+// llama_kv_cache_msa
+
+llama_kv_cache_msa::llama_kv_cache_msa(
+ 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_filter_cb & filter_idx,
+ const layer_reuse_cb & reuse) :
+ hparams_idx(model.hparams),
+ n_stream(unified ? 1 : n_seq_max), n_seq_max(n_seq_max), n_pad(n_pad),
+ n_swa(n_swa), swa_type(swa_type) {
+
+ LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
+
+ kv_base = 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);
+
+ // the MSA indexer uses a single key head per layer
+ std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
+ hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
+ // the rope parameters are kept identical to the main cache
+
+ LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
+
+ kv_idx = std::make_unique<llama_kv_cache>(
+ model, hparams_idx, type_k, type_v,
+ v_trans, offload, unified, kv_size, n_seq_max, n_pad,
+ n_swa, swa_type, nullptr, filter_idx, reuse, nullptr);
+}
+
+void llama_kv_cache_msa::clear(bool data) {
+ kv_base->clear(data);
+ kv_idx ->clear(data);
+}
+
+bool llama_kv_cache_msa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
+ bool res = true;
+
+ res = res & kv_base->seq_rm(seq_id, p0, p1);
+ res = res & kv_idx ->seq_rm(seq_id, p0, p1);
+
+ return res;
+}
+
+void llama_kv_cache_msa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+ kv_base->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+ kv_idx ->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+}
+
+void llama_kv_cache_msa::seq_keep(llama_seq_id seq_id) {
+ kv_base->seq_keep(seq_id);
+ kv_idx ->seq_keep(seq_id);
+}
+
+void llama_kv_cache_msa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+ kv_base->seq_add(seq_id, p0, p1, shift);
+ kv_idx ->seq_add(seq_id, p0, p1, shift);
+}
+
+void llama_kv_cache_msa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+ kv_base->seq_div(seq_id, p0, p1, d);
+ kv_idx ->seq_div(seq_id, p0, p1, d);
+}
+
+llama_pos llama_kv_cache_msa::seq_pos_min(llama_seq_id seq_id) const {
+ return kv_base->seq_pos_min(seq_id);
+}
+
+llama_pos llama_kv_cache_msa::seq_pos_max(llama_seq_id seq_id) const {
+ return kv_base->seq_pos_max(seq_id);
+}
+
+std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_msa::memory_breakdown() const {
+ std::map<ggml_backend_buffer_type_t, size_t> mb = kv_base->memory_breakdown();
+ for (const auto & buft_size : kv_idx->memory_breakdown()) {
+ mb[buft_size.first] += buft_size.second;
+ }
+ return mb;
+}
+
+llama_memory_context_ptr llama_kv_cache_msa::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, 0);
+
+ if (ubatch.n_tokens == 0) {
+ break;
+ }
+
+ ubatches.push_back(std::move(ubatch));
+ }
+
+ if (balloc.get_n_used() < balloc.get_n_tokens()) {
+ // failed to find a suitable split
+ break;
+ }
+
+ auto sinfos_base = kv_base->prepare(ubatches);
+ if (sinfos_base.empty()) {
+ break;
+ }
+
+ auto sinfos_idx = kv_idx->prepare(ubatches);
+ if (sinfos_idx.empty()) {
+ break;
+ }
+
+ assert(sinfos_base.size() == sinfos_idx.size());
+
+ return std::make_unique<llama_kv_cache_msa_context>(
+ this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches));
+ } while (false);
+
+ return std::make_unique<llama_kv_cache_msa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
+}
+
+llama_memory_context_ptr llama_kv_cache_msa::init_full() {
+ return std::make_unique<llama_kv_cache_msa_context>(this);
+}
+
+llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) {
+ return std::make_unique<llama_kv_cache_msa_context>(this, lctx, optimize);
+}
+
+bool llama_kv_cache_msa::get_can_shift() const {
+ return kv_base->get_can_shift() &&
+ kv_idx ->get_can_shift() &&
+ kv_base->get_size() == kv_idx->get_size();
+}
+
+void llama_kv_cache_msa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+ kv_base->state_write(io, seq_id, flags);
+ kv_idx ->state_write(io, seq_id, flags);
+}
+
+void llama_kv_cache_msa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+ kv_base->state_read(io, seq_id, flags);
+ kv_idx ->state_read(io, seq_id, flags);
+}
+
+llama_kv_cache * llama_kv_cache_msa::get_base() const {
+ return kv_base.get();
+}
+
+llama_kv_cache * llama_kv_cache_msa::get_idx() const {
+ return kv_idx.get();
+}
+
+// llama_kv_cache_msa_context
+
+llama_kv_cache_msa_context::llama_kv_cache_msa_context(llama_memory_status status) :
+ kv(nullptr), status(status) {}
+
+llama_kv_cache_msa_context::llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv) :
+ kv(kv),
+ ctx_base(kv->get_base()->init_full()),
+ ctx_idx (kv->get_idx ()->init_full()),
+ status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
+}
+
+llama_kv_cache_msa_context::llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv,
+ llama_context * lctx,
+ bool optimize) :
+ kv(kv),
+ ctx_base(kv->get_base()->init_update(lctx, optimize)),
+ ctx_idx (kv->get_idx ()->init_update(lctx, optimize)),
+ status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
+}
+
+llama_kv_cache_msa_context::llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv,
+ slot_info_vec_t sinfos_base,
+ slot_info_vec_t sinfos_idx,
+ std::vector<llama_ubatch> ubatches) :
+ kv(kv),
+ ubatches(std::move(ubatches)),
+ // here we copy the ubatches. not sure if this is ideal
+ ctx_base(new llama_kv_cache_context(kv->get_base(), std::move(sinfos_base), this->ubatches)),
+ ctx_idx (new llama_kv_cache_context(kv->get_idx (), std::move(sinfos_idx), this->ubatches)),
+ status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
+}
+
+llama_kv_cache_msa_context::~llama_kv_cache_msa_context() = default;
+
+bool llama_kv_cache_msa_context::next() {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ ctx_base->next();
+ ctx_idx ->next();
+
+ if (++i_next >= ubatches.size()) {
+ return false;
+ }
+
+ return true;
+}
+
+bool llama_kv_cache_msa_context::apply() {
+ assert(!llama_memory_status_is_fail(status));
+
+ bool res = true;
+
+ res = res & ctx_base->apply();
+ res = res & ctx_idx ->apply();
+
+ return res;
+}
+
+llama_memory_status llama_kv_cache_msa_context::get_status() const {
+ return status;
+}
+
+const llama_ubatch & llama_kv_cache_msa_context::get_ubatch() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return ubatches[i_next];
+}
+
+const llama_kv_cache_context * llama_kv_cache_msa_context::get_base() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return static_cast<const llama_kv_cache_context *>(ctx_base.get());
+}
+
+const llama_kv_cache_context * llama_kv_cache_msa_context::get_idx() const {
+ assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
+
+ return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
+}
+
+uint32_t llama_kv_cache_msa_context::get_n_pos() const {
+ // pad the value so that the graph remains constant across batches and can be reused
+ const uint32_t n_pad_cur = std::max(kv->get_n_pad(), 256u);
+
+ llama_pos pos_max = -1;
+
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) kv->get_n_seq_max(); ++seq_id) {
+ pos_max = std::max(pos_max, kv->seq_pos_max(seq_id));
+ }
+
+ return std::max(n_pad_cur, GGML_PAD((uint32_t) (pos_max + 1), n_pad_cur));
+}
+
+void llama_kv_cache_msa_context::set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const {
+ GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
+ GGML_ASSERT(dst->type == GGML_TYPE_I32);
+ GGML_ASSERT(div > 0);
+
+ const int64_t n_tokens = ubatch->n_tokens;
+ const int64_t n_kv = dst->ne[0];
+ const int64_t n_stream_ub = dst->ne[1];
+
+ GGML_ASSERT(n_tokens % n_stream_ub == 0);
+ const int64_t n_tps = n_tokens/n_stream_ub;
+
+ int32_t * data = (int32_t *) dst->data;
+
+ for (int64_t s = 0; s < n_stream_ub; ++s) {
+ const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
+
+ const auto & cells = kv->get_base()->get_cells(seq_id);
+
+ for (int64_t j = 0; j < n_kv; ++j) {
+ // the value for empty or other-sequence cells is irrelevant as consumers mask them
+ data[s*n_kv + j] =
+ cells.is_empty(j) || !cells.seq_has(j, seq_id)
+ ? 0
+ : (int32_t) (cells.pos_get(j)/div);
+ }
+ }
+}
+
+void llama_kv_cache_msa_context::set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const {
+ GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
+ GGML_ASSERT(dst->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_F32);
+
+ const int64_t n_tokens = ubatch->n_tokens;
+ const int64_t n_pos = dst->ne[0];
+ const int64_t n_stream_ub = dst->ne[1];
+
+ GGML_ASSERT(n_tokens % n_stream_ub == 0);
+ const int64_t n_tps = n_tokens/n_stream_ub;
+
+ for (int64_t s = 0; s < n_stream_ub; ++s) {
+ const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
+
+ const auto & cells = kv->get_base()->get_cells(seq_id);
+
+ std::vector<int32_t> map(n_pos, 0);
+
+ for (uint32_t j = 0; j < cells.size(); ++j) {
+ if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
+ continue;
+ }
+
+ const llama_pos p0 = cells.pos_get(j);
+
+ if (p0 < 0 || p0 >= n_pos) {
+ continue;
+ }
+
+ map[p0] = (int32_t) j;
+ }
+
+ if (dst->type == GGML_TYPE_I32) {
+ int32_t * data = (int32_t *) dst->data + s*n_pos;
+ std::copy(map.begin(), map.end(), data);
+ } else {
+ float * data = (float *) dst->data + s*n_pos;
+ for (int64_t p = 0; p < n_pos; ++p) {
+ data[p] = (float) map[p];
+ }
+ }
+ }
+}
+
+void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const {
+ GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ const int64_t n_tokens = ubatch->n_tokens;
+ const int64_t n_pos = dst->ne[0];
+
+ GGML_ASSERT(dst->ne[1] == n_tokens);
+
+ const uint32_t n_swa = kv->get_n_swa();
+ const llama_swa_type swa_type = kv->get_swa_type();
+
+ float * data = (float *) dst->data;
+
+ std::fill(data, data + n_pos*n_tokens, -INFINITY);
+
+ for (int64_t i = 0; i < n_tokens; ++i) {
+ const llama_seq_id seq_id = ubatch->seq_id[i][0];
+
+ const auto & cells = kv->get_base()->get_cells(seq_id);
+
+ const llama_pos p1 = ubatch->pos[i];
+
+ for (uint32_t j = 0; j < cells.size(); ++j) {
+ if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
+ continue;
+ }
+
+ const llama_pos p0 = cells.pos_get(j);
+
+ if (p0 < 0 || p0 >= n_pos) {
+ continue;
+ }
+
+ // causal mask
+ if (p0 > p1) {
+ continue;
+ }
+
+ // apply SWA if any
+ if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
+ continue;
+ }
+
+ data[i*n_pos + p0] = 0.0f;
+ }
+ }
+}
--- /dev/null
+#pragma once
+
+#include "llama-kv-cache.h"
+
+#include <vector>
+
+// llama_kv_cache_msa
+
+// uses two instances of llama_kv_cache, one for K/V tensors, and one for the MSA indexer tensors
+// both receive identical sequence operations and identical ubatches, so their cell layouts stay in synced.
+// the context also exposes per-ubatch pos - cell translation maps populated from llama_kv_cells via
+// llama_kv_cache::get_cells(), which the model graph uses to run MSA block selection in position space
+
+class llama_kv_cache_msa : public llama_memory_i {
+public:
+ llama_kv_cache_msa(
+ 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_filter_cb & filter_idx,
+ const layer_reuse_cb & reuse);
+
+ ~llama_kv_cache_msa() = 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_msa specific API
+
+ llama_kv_cache * get_base() const;
+ llama_kv_cache * get_idx () const;
+
+ uint32_t get_n_pad() const { return n_pad; }
+ uint32_t get_n_seq_max() const { return n_seq_max; }
+ uint32_t get_n_swa() const { return n_swa; }
+ llama_swa_type get_swa_type() const { return swa_type; }
+
+private:
+ // keep the indexer KV cache hparams instance here as llama_kv_cache stores only a reference
+ llama_hparams hparams_idx;
+
+ const uint32_t n_stream = 1;
+ const uint32_t n_seq_max = 1;
+ const uint32_t n_pad = 1;
+
+ const uint32_t n_swa = 0;
+ const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
+
+ std::unique_ptr<llama_kv_cache> kv_base;
+ std::unique_ptr<llama_kv_cache> kv_idx;
+};
+
+class llama_kv_cache_msa_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_msa_context(llama_memory_status status);
+
+ // used to create a full-cache context
+ llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv);
+
+ // used to create an update context
+ llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv,
+ llama_context * lctx,
+ bool optimize);
+
+ // used to create a batch processing context from a batch
+ llama_kv_cache_msa_context(
+ llama_kv_cache_msa * kv,
+ slot_info_vec_t sinfos_base,
+ slot_info_vec_t sinfos_idx,
+ std::vector<llama_ubatch> ubatches);
+
+ virtual ~llama_kv_cache_msa_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_msa_context specific API
+
+ const llama_kv_cache_context * get_base() const;
+ const llama_kv_cache_context * get_idx () const;
+
+ // max position currently present in the cache plus one, padded MSA blocks are defined over token positions
+ // so the block-selection tensors are sized by this value rather than by the number of cells
+ uint32_t get_n_pos() const;
+
+ // position <-> cell translation maps, populated from the base cache cells
+ // the model graph relates cache contents to token positions only through these per ubatch inputs
+ // value for empty or other-sequence cells is 0 so consumers must mask them
+ void set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const;
+ // positions without a cell map to cell 0, consumers must mask them assumes one sequence per stream
+ void set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const;
+ void set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const;
+
+private:
+ llama_kv_cache_msa * 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_base;
+ const llama_memory_context_ptr ctx_idx;
+
+ const llama_memory_status status;
+};
hparams.n_embd_head_k() % 64 == 0;
// always create Hadamard rotation tensors for DeepSeek lightning indexers
- if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) &&
+ if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || model.arch == LLM_ARCH_GLM_DSA) &&
hparams.n_embd_head_k_full == hparams.indexer_head_size) {
attn_rot_k = true;
}
return layers[ikv].k;
}
+const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const {
+ GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
+
+ return v_cells[seq_to_stream[seq_id]];
+}
+
uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
uint32_t result = 0;
bool res = true;
res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
- res = res && state_read_data(io, strm, cell_count, sinfo);
+
+ try {
+ res = res && state_read_data(io, strm, cell_count, sinfo);
+ } catch (...) {
+ res = false;
+ }
if (!res) {
if (seq_id == -1) {
std::vector<uint32_t> get_layer_ids() const;
ggml_tensor * get_k_storage(int32_t il) const;
+ const llama_kv_cells & get_cells(llama_seq_id seq_id) const;
+
//
// graph_build API
//
bool res = true;
res = res && state_read_meta(io, cell_count, seq_id);
- res = res && state_read_data(io, cell_count);
+
+ try {
+ res = res && state_read_data(io, cell_count);
+ } catch (...) {
+ res = false;
+ }
if (!res) {
if (seq_id == -1) {
#include "ggml.h"
#include "gguf.h"
#include "llama-hparams.h"
+#include "llama.h"
#include <algorithm>
#include <array>
const std::string & fname,
std::vector<std::string> & splits,
FILE * file,
- bool use_mmap,
- bool use_direct_io,
+ llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
+ bool load_mtp,
const llama_model_kv_override * param_overrides_p,
const llama_model_tensor_buft_override * param_tensor_buft_overrides_p)
: metadata(meta), set_tensor_data(set_tensor_data), set_tensor_data_ud(set_tensor_data_ud) {
tensor_buft_overrides = param_tensor_buft_overrides_p;
+ this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK;
+ this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO;
+
if (!fname.empty()) {
// Load the main GGUF
struct ggml_context * ctx = NULL;
files.emplace_back(new llama_file(fname.c_str(), "rb", use_direct_io));
contexts.emplace_back(ctx);
- if (use_mmap && use_direct_io) {
- if (files.back()->has_direct_io()) {
- LLAMA_LOG_WARN("%s: direct I/O is enabled, disabling mmap\n", __func__);
- use_mmap = false;
- } else {
- LLAMA_LOG_WARN("%s: direct I/O is not available, using mmap\n", __func__);
- use_direct_io = false;
-
- // reopen file using std::fopen for mmap
- files.pop_back();
- files.emplace_back(new llama_file(fname.c_str(), "rb", false));
- }
- }
-
// Save tensors data offset of the main file.
// For subsidiary files, `meta` tensor data offset must not be used,
// so we build a unified tensors index for weights.
}
}
- if (!llama_mmap::SUPPORTED) {
+ if (this->use_mmap && !llama_mmap::SUPPORTED) {
LLAMA_LOG_WARN("%s: mmap is not supported on this platform\n", __func__);
- use_mmap = false;
+ this->use_mmap = false;
}
- this->use_mmap = use_mmap;
- this->use_direct_io = use_direct_io;
this->check_tensors = check_tensors;
this->no_alloc = no_alloc;
+ this->load_mtp = load_mtp;
}
std::string llama_model_loader::get_arch_name() const {
return tensor;
}
-const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::string & name, const std::vector<int64_t> & ne, bool required) const {
+const struct ggml_tensor * llama_model_loader::check_tensor_dims(
+ const std::string & name,
+ const std::vector<int64_t> & ne,
+ bool required,
+ bool allow_reshape) const {
const struct ggml_tensor * cur = get_tensor_meta(name.c_str());
if (cur == NULL) {
throw std::runtime_error(format("%s: tensor '%s' not found", __func__, name.c_str()));
}
- {
- bool is_ok = true;
+ bool is_ok = true;
+
+ if (allow_reshape) {
+ // check total number of elements only
+ const int64_t ncur = ggml_nelements(cur);
+ int64_t nexp = 1;
+ for (size_t i = 0; i < ne.size(); ++i) {
+ nexp *= ne[i];
+ }
+ if (ncur != nexp) {
+ is_ok = false;
+ }
+ } else {
for (size_t i = 0; i < GGML_MAX_DIMS; ++i) {
if ((i < ne.size() && ne[i] != cur->ne[i]) || (i >= ne.size() && cur->ne[i] != 1)) {
is_ok = false;
break;
}
}
- if (!is_ok) {
- throw std::runtime_error(
- format("%s: tensor '%s' has wrong shape; expected %s, got %s",
- __func__, name.c_str(),
- llama_format_tensor_shape(ne).c_str(),
- llama_format_tensor_shape(cur).c_str()));
- }
+ }
+
+ if (!is_ok) {
+ throw std::runtime_error(
+ format("%s: tensor '%s' has wrong shape; expected %s, got %s",
+ __func__, name.c_str(),
+ llama_format_tensor_shape(ne).c_str(),
+ llama_format_tensor_shape(cur).c_str()));
}
return cur;
return ret;
}
- ggml_tensor * t_meta = get_tensor_meta(tn.str().c_str());
- ggml_backend_buffer_type_t buft = buft_for_tensor(t_meta);
+ LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str());
+ const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED), flags & TENSOR_ALLOW_RESHAPE);
+ if (cur == NULL) {
+ return NULL;
+ }
+
+ ggml_tensor t_meta = *cur;
+ if (flags & TENSOR_ALLOW_RESHAPE) {
+ for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
+ t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
+ t_meta.nb[dim] = dim == 0 ? ggml_type_size(t_meta.type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
+ }
+ }
+
+ ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
if (buft == nullptr) {
- return nullptr; // return type is ggml_tensor *
+ return nullptr;
}
+
ggml_context * ctx = ctx_for_buft(buft);
// if duplicated, check if the original tensor was allocated in the same buffer type context and avoid creating a new one
}
}
- LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str());
- const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED));
-
- if (cur == NULL) {
- return NULL;
- }
-
const bool duplicated = flags & TENSOR_DUPLICATED;
- struct ggml_tensor * tensor = ggml_dup_tensor(ctx, cur);
- ggml_set_name(tensor, ggml_get_name(cur));
+ struct ggml_tensor * tensor = ggml_dup_tensor(ctx, &t_meta);
+ ggml_set_name(tensor, ggml_get_name(&t_meta));
if (duplicated) {
- size_data += ggml_nbytes(cur);
+ size_data += ggml_nbytes(&t_meta);
} else {
n_created++;
}
return tensor;
}
-struct ggml_tensor * llama_model_loader::create_tensor_as_view(struct ggml_context * ctx, struct ggml_tensor * base, const std::string & name, const std::initializer_list<int64_t> & ne, size_t offset, bool required) {
- const struct ggml_tensor * cur = check_tensor_dims(name, ne, required);
-
- if (cur == NULL) {
- return NULL;
- }
-
- if (cur->type != base->type) {
- throw std::runtime_error(format("%s: tensor '%s' has wrong type; expected %s, got %s", __func__, name.c_str(), ggml_type_name(base->type), ggml_type_name(cur->type)));
- }
-
- std::array<int64_t, GGML_MAX_DIMS> dims;
- for (size_t i = 0; i < GGML_MAX_DIMS; ++i) {
- dims[i] = i < ne.size() ? ne.begin()[i] : 1;
- }
-
- struct ggml_tensor * tensor = ggml_view_4d(ctx, base,
- dims[0], dims[1], dims[2], dims[3],
- cur->nb[1], cur->nb[2], cur->nb[3],
- offset);
-
- ggml_set_name(tensor, name.c_str());
-
- n_created++;
-
- return tensor;
-}
-
void llama_model_loader::done_getting_tensors(bool partial) const {
if (n_created > n_tensors) {
throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created));
static const int TENSOR_DUPLICATED = 1 << 1;
static const int TENSOR_SKIP = 1 << 2;
static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3;
+ static const int TENSOR_ALLOW_RESHAPE = 1 << 4;
int n_kv = 0;
int n_tensors = 0;
bool use_direct_io = false;
bool check_tensors;
bool no_alloc;
+ bool load_mtp;
llama_files files;
llama_ftype ftype;
const std::string & fname,
std::vector<std::string> & splits, // optional, only need if the split does not follow naming scheme
FILE * file,
- bool use_mmap,
- bool use_direct_io,
+ llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
+ bool load_mtp,
const llama_model_kv_override * param_overrides_p,
const llama_model_tensor_buft_override * param_tensor_buft_overrides_p);
struct ggml_tensor * require_tensor_meta(const std::string & name) const;
- const struct ggml_tensor * check_tensor_dims(const std::string & name, const std::vector<int64_t> & ne, bool required) const;
+ const struct ggml_tensor * check_tensor_dims(
+ const std::string & name,
+ const std::vector<int64_t> & ne,
+ bool required,
+ bool allow_reshape) const;
struct ggml_tensor * create_tensor(
const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output,
const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags);
- struct ggml_tensor * create_tensor_as_view(struct ggml_context * ctx, struct ggml_tensor * base, const std::string & name, const std::initializer_list<int64_t> & ne, size_t offset, bool required = true);
-
void done_getting_tensors(bool partial = false) const;
void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr);
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
case LLM_ARCH_MELLUM:
+ case LLM_ARCH_LAGUNA:
return false;
default:
return 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_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);
+ add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
+ add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
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-kv-cache-msa.h"
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-recurrent.h"
+#include "llama.h"
#include "models/models.h"
#include "ggml.h"
return new llama_model_stablelm(params);
case LLM_ARCH_MELLUM:
return new llama_model_mellum(params);
+ case LLM_ARCH_NANBEIGE:
+ return new llama_model_nanbeige(params);
case LLM_ARCH_QWEN:
return new llama_model_qwen(params);
case LLM_ARCH_QWEN2:
return new llama_model_arcee(params);
case LLM_ARCH_AFMOE:
return new llama_model_afmoe(params);
+ case LLM_ARCH_LAGUNA:
+ return new llama_model_laguna(params);
case LLM_ARCH_ERNIE4_5:
return new llama_model_ernie4_5(params);
case LLM_ARCH_ERNIE4_5_MOE:
return new llama_model_hunyuan_vl(params);
case LLM_ARCH_HUNYUAN_DENSE:
return new llama_model_hunyuan_dense(params);
+ case LLM_ARCH_HY_V3:
+ return new llama_model_hy_v3(params);
case LLM_ARCH_SMOLLM3:
return new llama_model_smollm3(params);
case LLM_ARCH_OPENAI_MOE:
return new llama_model_apertus(params);
case LLM_ARCH_MINIMAX_M2:
return new llama_model_minimax_m2(params);
+ case LLM_ARCH_MINIMAX_M3:
+ return new llama_model_minimax_m3(params);
case LLM_ARCH_COGVLM:
return new llama_model_cogvlm(params);
case LLM_ARCH_PANGU_EMBED:
if (model != nullptr) {
model->arch = arch;
- auto & devices = model->devices;
- if (!devices.empty() && devices[0].is_meta && !llm_arch_supports_sm_tensor(arch)) {
+ if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && !llm_arch_supports_sm_tensor(arch)) {
throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'");
}
}
const llama_hparams & hparams = ud->model->hparams;
const std::string tensor_name = tensor->name;
- const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
- const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
- const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight");
- const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias");
- const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias");
- const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
- const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
- const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
- const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
- const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
- const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
- const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
-
- const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
- const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a");
- const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight");
- const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight");
- const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight");
- const std::regex pattern_r_cache ("cache_r_l\\d*");
- const std::regex pattern_s_cache ("cache_s_l\\d*");
- const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight");
- const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight");
-
- const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
- const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
- const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
- const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
- const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
- const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
-
- const std::regex pattern_output_weight("output\\.weight");
- const std::regex pattern_output_bias ("output\\.bias");
+ static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight");
+ static const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight");
+ static const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight");
+ static const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias");
+ static const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias");
+ static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias");
+ static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
+ static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*");
+ static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight");
+ static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
+ static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias");
+ static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
+
+ static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias");
+ static const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a");
+ static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight");
+ static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight");
+ static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight");
+ static const std::regex pattern_r_cache ("cache_r_l\\d*");
+ static const std::regex pattern_s_cache ("cache_s_l\\d*");
+ static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight");
+ static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight");
+
+ static const std::regex pattern_ffn_up_weight ("blk\\.\\d*\\.ffn_up(_exps)?.weight");
+ static const std::regex pattern_ffn_up_bias ("blk\\.\\d*\\.ffn_up(_exps)?.bias");
+ static const std::regex pattern_ffn_gate_weight ("blk\\.\\d*\\.ffn_gate(_exps)?.weight");
+ static const std::regex pattern_ffn_gate_bias ("blk\\.\\d*\\.ffn_gate(_exps)?.bias");
+ static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
+ static const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
+ static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias");
+ static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
+
+ static const std::regex pattern_output_weight("output\\.weight");
+ static const std::regex pattern_output_bias ("output\\.bias");
struct tensor_config {
ggml_backend_meta_split_axis axis;
}
// FFN
- if (std::regex_match(tensor_name, pattern_ffn_up_gate_weight)) {
+ if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_gate_weight)) {
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down.weight", "ffn_down_exps.weight");
}
- if (std::regex_match(tensor_name, pattern_ffn_up_gate_bias)) {
+ if (std::regex_match(tensor_name, pattern_ffn_up_bias) || std::regex_match(tensor_name, pattern_ffn_gate_bias)) {
return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down.weight", "ffn_down_exps.weight");
}
if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);
return {{n_embd, 1}, {n_embd_gqa, 2}};
}
+ if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias)) {
+ const int64_t n_ff = hparams.n_ff(il);
+ // some models such as Phi 3 have fused up + gate tensors named "up" tensors, which need to be segmented
+ if (tensor->ne[axis] == 2*n_ff) {
+ return {{n_ff, 2}};
+ }
+ return {{tensor->ne[axis], 1}};
+ }
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);
}
// FFN
- if (std::regex_match(tensor_name, pattern_ffn_up_gate_weight) || std::regex_match(tensor_name, pattern_ffn_up_gate_bias) ||
+ if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias) ||
+ std::regex_match(tensor_name, pattern_ffn_gate_weight) || std::regex_match(tensor_name, pattern_ffn_gate_bias) ||
std::regex_match(tensor_name, pattern_ffn_gate_up_weight) || std::regex_match(tensor_name, pattern_ffn_down_weight)) {
const int64_t blck_size_perf = std::lcm(blck_size, 128);
GGML_ASSERT(segments.size() == 1);
case LLM_TYPE_100B_A6B: return "100B.A6B";
case LLM_TYPE_102B_A12B: return "102B.A12B";
case LLM_TYPE_106B_A12B: return "106B.A12B";
+ case LLM_TYPE_118B_A8B: return "118B.A8B";
case LLM_TYPE_120B_A12B: return "120B.A12B";
case LLM_TYPE_122B_A10B: return "122B.A10B";
case LLM_TYPE_196B_A11B: return "196B.A11B";
case LLM_TYPE_230B_A10B: return "230B.A10B";
+ case LLM_TYPE_428B_A23B: return "428B.A23B";
case LLM_TYPE_235B_A22B: return "235B.A22B";
case LLM_TYPE_300B_A47B: return "300B.A47B";
case LLM_TYPE_310B_A15B: return "310B.A15B";
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_all);
+ GGML_ASSERT(hparams.n_layer_all > 0 && hparams.n_layer_all <= LLAMA_MAX_LAYERS);
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.rope_sections.begin(), hparams.rope_sections.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.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 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);
bool llama_model_base::load_tensors(llama_model_loader & ml) {
const auto & split_mode = params.split_mode;
- const auto & use_mlock = params.use_mlock;
+ const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK || params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK;
const auto & tensor_split = params.tensor_split;
const int n_layer_all = hparams.n_layer_all;
this->ml = &ml; // to be used by create_tensor() and load_arch_tensors()
- LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (mmap = %s, direct_io = %s)\n",
- __func__, ml.use_mmap ? "true" : "false", ml.use_direct_io ? "true" : "false");
+ LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n",
+ __func__, llama_load_mode_name(params.load_mode));
// build a list of buffer types for the CPU and GPU devices
pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host);
{
res = nullptr;
} break;
- case LLM_ARCH_DEEPSEEK32:
+ case LLM_ARCH_MINIMAX_M3:
{
- res = new llama_kv_cache_dsa(
+ // sparse (MSA) layers carry an indexer key cache, but leading dense layers do not
+ llama_kv_cache::layer_filter_cb filter_idx =
+ [&](int32_t il) { return (uint32_t) il >= hparams.n_layer_dense_lead; };
+
+ res = new llama_kv_cache_msa(
*this,
params.type_k,
params.type_v,
hparams.n_swa,
hparams.swa_type,
nullptr,
+ filter_idx,
nullptr);
} break;
+ case LLM_ARCH_GLM_DSA:
+ case LLM_ARCH_DEEPSEEK32:
+ {
+ if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) {
+ // The NextN/MTP draft head runs dense MLA (no DSA indexer), so the
+ // MTP context uses a plain attention KV cache holding only the
+ // nextn layer(s) - same pattern as the hybrid Qwen3.5 MTP context.
+ llama_kv_cache::layer_filter_cb filter =
+ [&](uint32_t il) { return il >= hparams.n_layer(); };
+
+ res = new llama_kv_cache(
+ *this,
+ hparams,
+ 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,
+ filter,
+ nullptr,
+ nullptr);
+ } else {
+ // Main context: DSA cache for the trunk layers only - the nextn
+ // layer(s) are never attended by the trunk graph.
+ llama_kv_cache::layer_filter_cb filter_mla = nullptr;
+ if (hparams.n_layer_nextn > 0) {
+ filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); };
+ }
+ llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); };
+
+ 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,
+ filter_mla,
+ filter_lid,
+ nullptr);
+ }
+ } break;
+ case LLM_ARCH_DEEPSEEK4:
+ {
+ GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
+
+ if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
+ const llama_memory_i::layer_filter_cb filter_mtp = [&](int32_t il) {
+ return il >= (int32_t) hparams.n_layer();
+ };
+
+ 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_mtp,
+ nullptr,
+ nullptr);
+ } else {
+ res = new llama_kv_cache_dsv4(
+ *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,
+ cparams.n_rs_seq,
+ nullptr,
+ nullptr);
+ }
+ } break;
+ case LLM_ARCH_DFLASH:
+ {
+ // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring)
+ if (hparams.dsv4_hc_mult > 0) {
+ GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
+
+ 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,
+ nullptr,
+ nullptr,
+ nullptr);
+ break;
+ }
+ }
+ [[fallthrough]];
// Models that need standard caching should rely on recurrent/hybrid
// checks
default:
{
- // The MTP head is dense-attention only on hybrid Qwen3.5/3.6, so use a plain
+ // The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain
// attention KV cache for the MTP context instead of the hybrid wrapper.
- const bool mtp_on_hybrid_qwen35 =
+ const bool mtp_on_hybrid_qwen =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
- (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
+ (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
if (llm_arch_is_recurrent(arch)) {
res = new llama_memory_recurrent(
cparams.n_seq_max,
cparams.n_rs_seq,
nullptr);
- } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen35) {
+ } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
// The main difference between hybrid architectures is the
// layer filters, so pick the right one here
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
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) {
+ } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
filter_attn = [&](uint32_t il) {
return il < hparams.n_layer() && !hparams.is_recr(il);
};
};
}
- if (mtp_on_hybrid_qwen35) {
+ if (mtp_on_hybrid_qwen) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
- if (arch == LLM_ARCH_STEP35 && hparams.n_layer_nextn > 0) {
+ if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
+ arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_DEEPSEEK32) &&
+ hparams.n_layer_nextn > 0) {
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
} else {
}
}
- if (arch == LLM_ARCH_DEEPSEEK4) {
- GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
-
- res = new llama_kv_cache_dsv4(
- *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);
- } else if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
+ if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
GGML_ASSERT(hparams.is_swa_any());
if (arch == LLM_ARCH_GEMMA4_ASSISTANT) {
/*.tensor_buft_overrides =*/ nullptr,
/*.n_gpu_layers =*/ -1,
/*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER,
+ /*.load_mode =*/ LLAMA_LOAD_MODE_MMAP,
/*.main_gpu =*/ 0,
/*.tensor_split =*/ nullptr,
/*.progress_callback =*/ nullptr,
/*.progress_callback_user_data =*/ nullptr,
/*.kv_overrides =*/ nullptr,
/*.vocab_only =*/ false,
- /*.use_mmap =*/ true,
- /*.use_direct_io =*/ false,
- /*.use_mlock =*/ false,
/*.check_tensors =*/ false,
/*.use_extra_bufts =*/ true,
/*.no_host =*/ false,
/*.no_alloc =*/ false,
+ /*.load_mtp =*/ false,
};
return result;
case LLM_ARCH_LLAMA_EMBED:
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
+ case LLM_ARCH_NANBEIGE:
return LLAMA_ROPE_TYPE_NORM;
// the pairs of head values are offset by n_rot/2
case LLM_ARCH_JAIS2:
case LLM_ARCH_OPENAI_MOE:
case LLM_ARCH_HUNYUAN_DENSE:
+ case LLM_ARCH_HY_V3:
case LLM_ARCH_LFM2:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_SMALLTHINKER:
case LLM_ARCH_GROVEMOE:
case LLM_ARCH_APERTUS:
case LLM_ARCH_MINIMAX_M2:
+ case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_COGVLM:
case LLM_ARCH_PANGU_EMBED:
case LLM_ARCH_AFMOE:
+ case LLM_ARCH_LAGUNA:
case LLM_ARCH_QWEN3NEXT:
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
case LLM_ARCH_TALKIE:
case LLM_ARCH_MELLUM:
- case LLM_ARCH_DFLASH:
return LLAMA_ROPE_TYPE_NEOX;
+ case LLM_ARCH_DFLASH:
+ // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX
+ return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
+
case LLM_ARCH_QWEN2VL:
case LLM_ARCH_PADDLEOCR:
return LLAMA_ROPE_TYPE_MROPE;
TENSOR_DUPLICATED (llama_model_loader::TENSOR_DUPLICATED),
TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED),
TENSOR_SKIP (llama_model_loader::TENSOR_SKIP),
- TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL) {}
+ TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL),
+ TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE) {}
ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
GGML_ASSERT(ml != nullptr);
LLM_TYPE_100B_A6B,
LLM_TYPE_102B_A12B, // Solar-Open
LLM_TYPE_106B_A12B, // GLM-4.5-Air
+ LLM_TYPE_118B_A8B, // Laguna-S-2
LLM_TYPE_120B_A12B, // Nemotron 3 Super
LLM_TYPE_122B_A10B, // Qwen3.5
LLM_TYPE_196B_A11B, // Step3.5-Flash
LLM_TYPE_230B_A10B, // Minimax M2
+ LLM_TYPE_428B_A23B, // Minimax M3
LLM_TYPE_235B_A22B,
LLM_TYPE_300B_A47B, // Ernie MoE big
LLM_TYPE_310B_A15B, // /MiMo-V2-Flash
struct ggml_tensor * indexer_attn_k = nullptr;
struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias
+ // MSA
+ struct ggml_tensor * index_q_proj = nullptr;
+ struct ggml_tensor * index_k_proj = nullptr;
+ struct ggml_tensor * index_q_norm = nullptr;
+ struct ggml_tensor * index_k_norm = nullptr;
+
// gemma4 layer output scale, reused for talkie embedding skip scale
struct ggml_tensor * out_scale = nullptr;
struct ggml_tensor * fc = nullptr; // feature fusion layer
struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping
+ // dspark
+ struct ggml_tensor * dspark_markov_w1 = nullptr;
+ struct ggml_tensor * dspark_markov_w2 = nullptr;
+ struct ggml_tensor * dspark_conf_proj = nullptr;
+ struct ggml_tensor * dspark_conf_proj_b = nullptr;
+
// unified vector to store target-model extracted layer ids in eagle3, dflash, etc.
std::vector<int32_t> target_layer_ids;
const int TENSOR_NOT_REQUIRED;
const int TENSOR_SKIP;
const int TENSOR_SKIP_IF_VIRTUAL;
+ const int TENSOR_ALLOW_RESHAPE;
explicit llama_model_base(const llama_model_params & params);
virtual ~llama_model_base() = default;
#include "llama-model.h"
#include "llama-model-loader.h"
#include "llama-ext.h"
+#include "llama.h"
#include <algorithm>
#include <cmath>
// NOTE: can't use LLM_TN here because the layer number is not known
quantize &= name.find("ffn_gate_inp.weight") == std::string::npos;
+ // do not quantize the i32 token-id -> expert-id routing table (DeepSeek-V4)
+ quantize &= name.find("ffn_gate_tid2eid.weight") == std::string::npos;
+
// these are very small (e.g. 4x4)
quantize &= name.find("altup") == std::string::npos;
quantize &= name.find("laurel") == std::string::npos;
quantize &= name.find("ssm_conv1d") == std::string::npos;
quantize &= name.find("shortconv.conv.weight") == std::string::npos;
+ // do not quantize MiniMax's indexer projection weights, they are tiny
+ quantize &= name.find("indexer.k_proj.weight") == std::string::npos;
+ quantize &= name.find("indexer.q_proj.weight") == std::string::npos;
+
// do not quantize RWKV's small yet 2D weights
quantize &= name.find("time_mix_first.weight") == std::string::npos;
quantize &= name.find("time_mix_w0.weight") == std::string::npos;
quantize &= name.find(".patch_embd") == std::string::npos;
quantize &= name.find(".patch_merger") == std::string::npos;
+ // audio codebook
+ quantize &= name.find("a.rvq.codebook") == std::string::npos;
+ quantize &= name.find("mm.a.code_embd") == std::string::npos;
+
return quantize;
}
ggml_type new_type = default_type;
// get more optimal quantization type based on the tensor shape, layer, etc.
- if (!params->pure && ggml_is_quantized(default_type)) {
+ if (ggml_is_quantized(default_type)) {
// if the user provided tensor types - use those
bool manual = false;
if (!qs.tensor_type_patterns.empty()) {
}
// if not manual - use the standard logic for choosing the quantization type based on the selected mixture
- if (!manual) {
+ if (!manual && !params->pure) {
new_type = llama_tensor_get_type_impl(qs, new_type, tensor, params->ftype, tm.category);
}
// mmap consistently increases speed on Linux, and also increases speed on Windows with
// hot cache. It may cause a slowdown on macOS, possibly related to free memory.
#if defined(__linux__) || defined(_WIN32)
- constexpr bool use_mmap = true;
+ constexpr llama_load_mode load_mode = LLAMA_LOAD_MODE_MMAP;
#else
- constexpr bool use_mmap = false;
+ constexpr llama_load_mode load_mode = LLAMA_LOAD_MODE_NONE;
#endif
const llama_model_kv_override * kv_overrides = params->kv_overrides;
std::vector<std::string> splits = {};
llama_model_loader ml(/*metadata*/ nullptr, /*set_tensor_data*/ nullptr, /*set_tensor_data_ud*/ nullptr,
- fname_inp, splits, /*file*/ nullptr, use_mmap, /*use_direct_io*/ false, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
+ fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, /*load_mtp*/ true, kv_overrides, nullptr);
ml.init_mappings(false); // no prefetching
auto mparams = llama_model_default_params();
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_all = desc->n_layer;
+ GGML_ASSERT(desc->n_layer > 0 && desc->n_layer <= LLAMA_MAX_LAYERS);
model->hparams.n_expert = desc->n_expert;
for (uint32_t i = 0; i < desc->n_layer; i++) {
*/
static void llama_sampler_temp_impl(llama_token_data_array * cur_p, float temp) {
+ if (cur_p->size == 0) {
+ return;
+ }
+
if (temp <= 0.0f) {
// find the token with the highest logit and set the rest to -inf
size_t max_i = 0;
GGML_UNUSED(gf);
GGML_UNUSED(smpl);
- struct ggml_tensor * curl = ggml_argmax(ctx, data->logits);
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
+ struct ggml_tensor * curl = ggml_argmax(ctx, logits);
ggml_set_name(curl, "greedy_argmax");
data->sampled = curl;
ggml_set_name (sctx->inp_uniform, "uniform");
ggml_set_input(sctx->inp_uniform);
- struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
+ // flatten
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
+ struct ggml_tensor * probs = ggml_soft_max(ctx, logits);
ggml_set_name(probs, "dist_probs");
struct ggml_tensor * cumsum = ggml_cumsum(ctx, probs);
struct llama_sampler_data * data) {
auto * sctx = (llama_sampler_top_k *) smpl->ctx;
- struct ggml_tensor * top_k = ggml_top_k(ctx, data->logits, sctx->k);
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
+ struct ggml_tensor * top_k = ggml_top_k(ctx, logits, sctx->k);
ggml_set_name(top_k, "top_k");
if (data->candidates) {
struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
data->candidates = ggml_get_rows(ctx, candidates_rows, top_k);
- data->candidates = ggml_reshape_1d(ctx, data->candidates, sctx->k);
ggml_set_name(data->candidates, "top_k_candidates");
} else {
data->candidates = top_k;
}
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
- struct ggml_tensor * top_k_rows = ggml_get_rows(ctx, logits_rows, top_k);
- data->logits = ggml_reshape_1d(ctx, top_k_rows, sctx->k);
- ggml_set_name(top_k_rows, "top_k_rows");
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]);
+ data->logits = ggml_get_rows(ctx, logits_rows, top_k);
+ ggml_set_name(data->logits, "top_k_rows");
GGML_UNUSED(gf);
}
struct llama_sampler_data * data) {
auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+ // flatten
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
auto ggml_sort = [ctx](struct ggml_tensor * a, struct ggml_tensor * b) {
GGML_ASSERT(ggml_nrows(a) == 1);
struct ggml_tensor * a_reshaped = ggml_reshape_2d(ctx, a, 1, a->ne[0]);
struct ggml_tensor * a_sorted = ggml_get_rows(ctx, a_reshaped, b);
- return ggml_reshape_1d(ctx, a_sorted, a->ne[0]);
+ return a_sorted;
};
// Get the sorted logits in descending order.
- struct ggml_tensor * sorted_idx = ggml_argsort(ctx, data->logits, GGML_SORT_ORDER_DESC);
+ struct ggml_tensor * sorted_idx = ggml_argsort(ctx, logits, GGML_SORT_ORDER_DESC);
ggml_set_name(sorted_idx, "top_p_sorted_idx");
// Do the sorting via reshape + get_rows
- struct ggml_tensor * sorted_logits = ggml_sort(data->logits, sorted_idx);
+ struct ggml_tensor * sorted_logits = ggml_sort(logits, sorted_idx);
ggml_set_name(sorted_logits, "top_p_sorted_logits");
+ sorted_logits = ggml_reshape_1d(ctx, sorted_logits, ggml_nelements(sorted_logits));
struct ggml_tensor * softmax = ggml_soft_max(ctx, sorted_logits);
ggml_set_name(softmax, "top_p_softmax");
struct llama_sampler_data * data) {
auto * sctx = (llama_sampler_min_p *) smpl->ctx;
- struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
+ struct ggml_tensor * max_idx = ggml_argmax(ctx, logits);
ggml_set_name(max_idx, "max_idx");
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]);
ggml_set_name(logits_rows, "logits_rows");
struct ggml_tensor * max_logit = ggml_get_rows(ctx, logits_rows, max_idx);
ggml_set_name(threshold, "min_p_threshold");
// Subtract the threshold from logits.
- struct ggml_tensor * sub = ggml_sub(ctx, data->logits, threshold);
+ struct ggml_tensor * sub = ggml_sub(ctx, logits, threshold);
// Create a mask where logits below the threshold are 0 (discard),
// and others are 1 (keep).
struct ggml_tensor * min_p_bias = ggml_log(ctx, mask);
ggml_set_name(min_p_bias, "min_p_bias");
- data->logits = ggml_add(ctx, data->logits, min_p_bias);
+ data->logits = ggml_add(ctx, logits, min_p_bias);
ggml_set_name(data->logits, "min_p_logits");
GGML_UNUSED(gf);
struct llama_sampler_data * data,
float temp) {
if (temp <= 0.0f) {
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
// Find the most probable token index.
- struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits);
+ struct ggml_tensor * max_idx = ggml_argmax(ctx, logits);
ggml_set_name(max_idx, "temp_max_idx");
if (data->candidates) {
- struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]);
+ struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
data->candidates = ggml_get_rows(ctx, candidates_rows, max_idx);
} else {
data->candidates = max_idx;
}
- struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]);
+ struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
data->logits = ggml_get_rows(ctx, logits_rows, max_idx);
return;
return;
}
+ struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+
// Calculate min_temp, max_temp, and max_entropy.
const float min_temp = std::max(0.0f, sctx->temp - sctx->delta);
const float max_temp = sctx->temp + sctx->delta;
- const float max_entropy = logf(data->logits->ne[0]);
+ const float max_entropy = logf(logits->ne[0]);
// Calculate the probabilities.
- struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits);
+ struct ggml_tensor * probs = ggml_soft_max(ctx, logits);
ggml_set_name(probs, "temp_ext_softmax_probs");
// Clamp probabilities to avoid log(0) which would give -inf
ggml_set_name(dyn_temp, "temp_ext_dyn_temp");
// Scale the logits by the dynamic temperature
- struct ggml_tensor * scaled_logits = ggml_div(ctx, data->logits, dyn_temp);
+ struct ggml_tensor * scaled_logits = ggml_div(ctx, logits, dyn_temp);
ggml_set_name(scaled_logits, "temp_ext_scaled_logits");
data->logits = scaled_logits;
// penalties
-struct llama_sampler_penalties {
+struct llama_sampler_penalties : public llama_sampler_backend {
+ const int32_t n_vocab;
const int32_t penalty_last_n;
const float penalty_repeat;
const float penalty_freq;
// a frequency map to count token occurrences
std::unordered_map<llama_token, int> token_count;
+
+ // backend graph inputs
+ ggml_tensor * inp_token_ids = nullptr;
+ ggml_tensor * inp_counts = nullptr;
+
+ // backend helpers
+ int32_t n_max = 0;
+ bool has_candidates = false;
+
+ std::vector<int32_t> host_token_ids;
+ std::vector<int32_t> host_counts;
+
+ static bool is_disabled(
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present) {
+ return penalty_last_n == 0 ||
+ (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f);
+ }
+
+ bool is_disabled() const {
+ return is_disabled(penalty_last_n, penalty_repeat, penalty_freq, penalty_present);
+ }
+
+ llama_sampler_penalties(
+ int32_t n_vocab,
+ int32_t penalty_last_n,
+ float penalty_repeat,
+ float penalty_freq,
+ float penalty_present)
+ : llama_sampler_backend("penalties")
+ , n_vocab (n_vocab)
+ , penalty_last_n (penalty_last_n)
+ , penalty_repeat (penalty_repeat)
+ , penalty_freq (penalty_freq)
+ , penalty_present (penalty_present)
+ , prev (penalty_last_n) {
+ }
};
-static const char * llama_sampler_penalties_name(const struct llama_sampler * /*smpl*/) {
- return "penalties";
+static const char * llama_sampler_penalties_name(const struct llama_sampler * smpl) {
+ auto * ctx = (llama_sampler_penalties *) smpl->ctx;
+ return ctx->get_name();
}
static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_token token) {
static void llama_sampler_penalties_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_penalties *) smpl->ctx;
- if ((ctx->penalty_last_n == 0) ||
- (ctx->penalty_repeat == 1.0f && ctx->penalty_freq == 0.0f && ctx->penalty_present == 0.0f)) {
+ if (ctx->is_disabled()) {
return;
}
static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_sampler * smpl) {
const auto * ctx = (const llama_sampler_penalties *) smpl->ctx;
auto * result = llama_sampler_init_penalties(
+ ctx->n_vocab,
ctx->penalty_last_n,
ctx->penalty_repeat,
ctx->penalty_freq,
{
auto * result_ctx = (llama_sampler_penalties *) result->ctx;
- result_ctx->prev = ctx->prev;
+ result_ctx->prev = ctx->prev;
+ result_ctx->token_count = ctx->token_count;
}
return result;
delete (llama_sampler_penalties *) smpl->ctx;
}
+static bool llama_sampler_penalties_backend_init(
+ struct llama_sampler * smpl,
+ ggml_backend_buffer_type_t buft) {
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ const bool res = llama_sampler_backend_support(smpl, buft);
+
+ sctx->init(res);
+
+ return res;
+}
+
+static void llama_sampler_penalties_backend_apply(
+ struct llama_sampler * smpl,
+ struct ggml_context * ctx,
+ struct ggml_cgraph * gf,
+ struct llama_sampler_data * data) {
+ GGML_UNUSED(gf);
+
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ if (sctx->is_disabled()) {
+ return;
+ }
+
+ GGML_ASSERT(sctx->n_vocab > 0);
+
+ sctx->has_candidates = data->candidates != nullptr;
+ sctx->n_max = std::min(sctx->penalty_last_n, sctx->n_vocab);
+
+ sctx->inp_token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
+ ggml_set_name(sctx->inp_token_ids, "penalties_token_ids");
+ ggml_set_input(sctx->inp_token_ids);
+
+ sctx->inp_counts = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max);
+ ggml_set_name(sctx->inp_counts, "penalties_counts");
+ ggml_set_input(sctx->inp_counts);
+
+ if ((int32_t) sctx->host_token_ids.size() != sctx->n_max) {
+ sctx->host_token_ids.assign(sctx->n_max, 0);
+ sctx->host_counts.assign(sctx->n_max, 0);
+ }
+
+ // flatten
+ ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
+ ggml_tensor * gathered = logits;
+ ggml_tensor * counts_f32 = ggml_cast(ctx, sctx->inp_counts, GGML_TYPE_F32);
+
+ if (sctx->has_candidates) {
+ ggml_tensor * candidates = ggml_reshape_1d(
+ ctx, data->candidates, ggml_nelements(data->candidates));
+ const int64_t n_candidates = candidates->ne[0];
+ GGML_ASSERT(n_candidates == ggml_nelements(logits));
+
+ ggml_tensor * counts_rows = ggml_fill(
+ ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, sctx->n_vocab), 0.0f);
+ ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, counts_f32, 1, sctx->n_max);
+ counts_rows = ggml_set_rows(ctx, counts_rows, scatter_rows, sctx->inp_token_ids);
+ counts_f32 = ggml_get_rows(ctx, counts_rows, candidates);
+ counts_f32 = ggml_reshape_1d(ctx, counts_f32, n_candidates);
+ } else {
+ ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
+ gathered = ggml_get_rows(ctx, logits_rows, sctx->inp_token_ids);
+ gathered = ggml_reshape_1d(ctx, gathered, sctx->n_max);
+ }
+
+ ggml_tensor * active_mask = ggml_step(ctx, counts_f32);
+ ggml_tensor * inactive_mask = ggml_sub(ctx, ggml_fill(ctx, active_mask, 1.0f), active_mask);
+
+ ggml_tensor * penalized = gathered;
+
+ if (sctx->penalty_repeat != 1.0f) {
+ ggml_tensor * pos_mask = ggml_step(ctx, penalized);
+ ggml_tensor * neg_mask = ggml_sub(ctx, ggml_fill(ctx, pos_mask, 1.0f), pos_mask);
+
+ ggml_tensor * pos_scale = ggml_scale(ctx, pos_mask, 1.0f/sctx->penalty_repeat);
+ ggml_tensor * neg_scale = ggml_scale(ctx, neg_mask, sctx->penalty_repeat);
+ ggml_tensor * repeat_scale = ggml_add(ctx, pos_scale, neg_scale);
+
+ // scale inactive entries with 1 to avoid -INF * 0 = NaN for values masked by top-p
+ repeat_scale = ggml_mul(ctx, repeat_scale, active_mask);
+ repeat_scale = ggml_add(ctx, repeat_scale, inactive_mask);
+ penalized = ggml_mul(ctx, gathered, repeat_scale);
+ }
+
+ if (sctx->penalty_freq != 0.0f) {
+ ggml_tensor * penalty_freq = ggml_scale(ctx, counts_f32, sctx->penalty_freq);
+ penalized = ggml_sub(ctx, penalized, penalty_freq);
+ }
+
+ if (sctx->penalty_present != 0.0f) {
+ ggml_tensor * penalty_present = ggml_scale(ctx, active_mask, sctx->penalty_present);
+ penalized = ggml_sub(ctx, penalized, penalty_present);
+ }
+
+ if (sctx->has_candidates) {
+ data->logits = penalized;
+ } else {
+ ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits));
+ ggml_tensor * scatter_rows = ggml_reshape_2d(ctx, penalized, 1, sctx->n_max);
+ logits_rows = ggml_set_rows(ctx, logits_rows, scatter_rows, sctx->inp_token_ids);
+ data->logits = ggml_reshape_1d(ctx, logits_rows, ggml_nelements(logits));
+ }
+}
+
+static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smpl) {
+ auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+
+ if (!sctx->inp_token_ids || !sctx->inp_counts || sctx->n_max <= 0 || sctx->n_vocab <= 0) {
+ return;
+ }
+
+ if (sctx->is_disabled()) {
+ return;
+ }
+
+ // fill active entries from the map
+ int32_t n_active = 0;
+
+ for (const auto & it : sctx->token_count) {
+ GGML_ASSERT(n_active < sctx->n_max);
+ sctx->host_token_ids[n_active] = it.first;
+ sctx->host_counts [n_active] = it.second;
+ ++n_active;
+ }
+
+ // Sorting is required because backend_apply uses ggml_set_rows (a scatter-back operation)
+ std::vector<std::pair<int32_t, int32_t>> entries;
+ entries.reserve(n_active);
+ for (int32_t i = 0; i < n_active; ++i) {
+ entries.emplace_back(sctx->host_token_ids[i], sctx->host_counts[i]);
+ }
+ std::sort(entries.begin(), entries.end(), [](const auto & a, const auto & b) {
+ return a.first < b.first;
+ });
+ for (int32_t i = 0; i < n_active; ++i) {
+ sctx->host_token_ids[i] = entries[i].first;
+ sctx->host_counts [i] = entries[i].second;
+ }
+
+ // Padding: Finds a filler token id that is not present in token_count.
+ // Use it to do padding for the arrays, it avoids resizing every time.
+ // The arrays must always have exactly n_max entries (the GPU tensor is a fixed size).
+ int32_t filler = 0;
+ if (n_active < sctx->n_max) {
+ while (sctx->token_count.find(filler) != sctx->token_count.end()) {
+ ++filler;
+ }
+ GGML_ASSERT(filler < sctx->n_vocab);
+ }
+
+ // Fill the rest of the arrays with the filler token id and count 0.
+ // Inactive slots are padded with a unique dummy token ID (count = 0).
+ // The uniqueness matters because ggml_set_rows with duplicate indices can produce non-deterministic or incorrect results.
+ // Using a filler token with count 0 that isn't in the active set is safe, because the active_mask step in backend_apply filters them out via ggml_step(counts_f32)
+ for (int32_t i = n_active; i < sctx->n_max; ++i) {
+ sctx->host_token_ids[i] = filler;
+ sctx->host_counts [i] = 0;
+ }
+
+ ggml_backend_tensor_set(sctx->inp_token_ids, sctx->host_token_ids.data(), 0, sctx->n_max * sizeof(int32_t));
+ ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t));
+}
+
static struct llama_sampler_i llama_sampler_penalties_i = {
/* .name = */ llama_sampler_penalties_name,
/* .accept = */ llama_sampler_penalties_accept,
/* .reset = */ llama_sampler_penalties_reset,
/* .clone = */ llama_sampler_penalties_clone,
/* .free = */ llama_sampler_penalties_free,
- /* .backend_init = */ nullptr,
+ /* .backend_init = */ llama_sampler_penalties_backend_init,
/* .backend_accept = */ nullptr,
- /* .backend_apply = */ nullptr,
- /* .backend_set_input = */ nullptr,
+ /* .backend_apply = */ llama_sampler_penalties_backend_apply,
+ /* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
};
struct llama_sampler * llama_sampler_init_penalties(
+ int32_t n_vocab,
int32_t penalty_last_n,
float penalty_repeat,
float penalty_freq,
float penalty_present) {
penalty_last_n = std::max(penalty_last_n, 0);
- const bool is_empty = (penalty_last_n == 0 || (penalty_repeat == 1.0f && penalty_freq == 0.0f && penalty_present == 0.0f));
-
- if (is_empty) {
+ if (llama_sampler_penalties::is_disabled(
+ penalty_last_n, penalty_repeat, penalty_freq, penalty_present)) {
return llama_sampler_init_empty("?penalties");
}
return llama_sampler_init(
/* .iface = */ &llama_sampler_penalties_i,
- /* .ctx = */ new llama_sampler_penalties {
- /* .penalty_last_n = */ penalty_last_n,
- /* .penalty_repeat = */ penalty_repeat,
- /* .penalty_freq = */ penalty_freq,
- /* .penalty_present = */ penalty_present,
- /* .prev = */ ring_buffer<llama_token>(penalty_last_n),
- /* .token_count = */ {},
- }
+ /* .ctx = */ new llama_sampler_penalties(
+ n_vocab,
+ penalty_last_n,
+ penalty_repeat,
+ penalty_freq,
+ penalty_present)
);
}
"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\\r\\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
};
break;
+ case LLAMA_VOCAB_PRE_TYPE_LAGUNA:
+ regex_exprs = {
+ "[^\\n]+|[\\n]+",
+ "(?:'[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_EXAONE_MOE:
regex_exprs = {
// original regex from tokenizer.json
token_id = node->value;
token_length = position + 1;
}
+ if (position + 1 >= text.size()) {
+ break;
+ }
node = node->traverse(text[++position]);
}
tokenizer_pre == "afmoe") {
pre_type = LLAMA_VOCAB_PRE_TYPE_AFMOE;
clean_spaces = false;
+ } else if (
+ tokenizer_pre == "laguna") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_LAGUNA;
+ clean_spaces = false;
} else if (
tokenizer_pre == "minimax-m2") {
pre_type = LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2;
const std::string & key = kv(std::get<0>(it));
int32_t & id = std::get<1>(it);
+ if (id >= 0 && static_cast<size_t>(id) >= id_to_token.size()) {
+ LLAMA_LOG_WARN("%s: default special token '%s' = %d out of vocab range, disabling\n",
+ __func__, key.c_str(), id);
+ id = LLAMA_TOKEN_NULL;
+ }
+
uint32_t new_id;
if (!ml.get_key(std::get<0>(it), new_id, false)) {
continue;
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);
+ // drop out-of-range ids
+ suppress_tokens.reserve(n);
+ for (int i = 0; i < n; ++i) {
+ const int32_t id = data[i];
+ if (id >= 0 && id < (int) id_to_token.size()) {
+ suppress_tokens.push_back(id);
+ }
+ }
}
}
|| t.first == "<turn|>" // gemma4
|| t.first == "<|tool_response>" // gemma4
|| t.first == "<｜end▁of▁sentence｜>" // deepseek-ocr
+ || t.first == "[e~[" // minimax-m2/m3
) {
special_eog_ids.insert(t.second);
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
LLAMA_LOG_INFO("%s: printing all EOG tokens:\n", __func__);
for (auto tid : special_eog_ids) {
+ if (tid < 0 || tid >= (llama_token) id_to_token.size()) {
+ LLAMA_LOG_WARN("%s: EOG token id %d is out of range (vocab size %zu), skipping\n",
+ __func__, tid, id_to_token.size());
+ continue;
+ }
auto & text = id_to_token[tid].text;
LLAMA_LOG_INFO("%s: - %d ('%s')\n", __func__, tid, text.c_str());
llama_token s_id = LLAMA_TOKEN_NULL;
for (auto tid : special_eog_ids) {
+ if (tid < 0 || tid >= (llama_token) id_to_token.size()) {
+ continue;
+ }
const auto & text = id_to_token[tid].text;
if (text == "<|tool_response>") {
has_tool_response = true;
}
std::vector<std::string> llama_vocab::get_bpe_merges() const {
- std::vector<std::string> result(pimpl->bpe_ranks.size());
+ int max_rank = -1;
+ for (const auto & pair : pimpl->bpe_ranks) {
+ max_rank = std::max(max_rank, pair.second);
+ }
+ std::vector<std::string> result(max_rank + 1);
for (const auto & pair : pimpl->bpe_ranks) {
result[pair.second] = pair.first.first + " " + pair.first.second;
return vocab->get_add_sep();
}
+const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens) {
+ const std::vector<llama_token> & tokens = vocab->get_suppress_tokens();
+ if (n_suppress_tokens) {
+ *n_suppress_tokens = (int32_t) tokens.size();
+ }
+ return tokens.data();
+}
+
llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab) {
return vocab->token_fim_pre();
}
LLAMA_VOCAB_PRE_TYPE_WHITESPACE = 53,
LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54,
LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55,
+ LLAMA_VOCAB_PRE_TYPE_LAGUNA = 56,
};
struct LLM_KV;
GGML_ABORT("fatal error");
}
+const char * llama_load_mode_name(enum llama_load_mode load_mode) {
+ switch (load_mode) {
+ case LLAMA_LOAD_MODE_NONE:
+ return "none";
+ case LLAMA_LOAD_MODE_MMAP:
+ return "mmap";
+ case LLAMA_LOAD_MODE_MLOCK:
+ return "mlock";
+ case LLAMA_LOAD_MODE_MMAP_MLOCK:
+ return "mmap+mlock";
+ case LLAMA_LOAD_MODE_DIRECT_IO:
+ return "dio";
+ }
+ GGML_ABORT("fatal error");
+}
+
+enum llama_load_mode llama_load_mode_from_str(const char * str) {
+ if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
+ if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
+ if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
+ if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; }
+ if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
+ throw std::invalid_argument(std::string("unknown load mode: ") + str);
+}
+
struct llama_sampler_chain_params llama_sampler_chain_default_params() {
struct llama_sampler_chain_params result = {
/*.no_perf =*/ true,
static std::pair<int, llama_model *> llama_model_load(struct gguf_context * metadata, llama_model_set_tensor_data_t set_tensor_data, void * set_tensor_data_ud,
const std::string & fname, std::vector<std::string> & splits, FILE * file, llama_model_params & params) {
try {
- llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.use_mmap, params.use_direct_io,
- params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides);
+ llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode,
+ params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides);
ml.print_info();
std::unique_ptr<llama_model> model_ptr(llama_model_create(ml, params));
GGML_ASSERT(metadata != nullptr);
std::string path_model;
std::vector<std::string> splits = {};
- params.use_mmap = false;
+ params.load_mode = LLAMA_LOAD_MODE_NONE;
params.use_extra_bufts = false;
return llama_model_load_from_file_impl(metadata, set_tensor_data, set_tensor_data_ud, path_model, splits, /*file*/ nullptr, params);
}
LLAMA_SPLIT_MODE_TENSOR = 3,
};
+ enum llama_load_mode {
+ LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
+ LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
+ LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing
+ LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
+ LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available
+ };
+
+ LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
+ LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str);
+
enum llama_context_type {
LLAMA_CONTEXT_TYPE_DEFAULT = 0,
LLAMA_CONTEXT_TYPE_MTP = 1,
int32_t n_gpu_layers; // number of layers to store in VRAM, a negative value means all layers
enum llama_split_mode split_mode; // how to split the model across multiple GPUs
+ enum llama_load_mode load_mode; // how to load the model
// the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
int32_t main_gpu;
// Keep the booleans together to avoid misalignment during copy-by-value.
bool vocab_only; // only load the vocabulary, no weights
- bool use_mmap; // use mmap if possible
- bool use_direct_io; // use direct io, takes precedence over use_mmap when supported
- bool use_mlock; // force system to keep model in RAM
bool check_tensors; // validate model tensor data
bool use_extra_bufts; // use extra buffer types (used for weight repacking)
bool no_host; // bypass host buffer allowing extra buffers to be used
bool no_alloc; // only load metadata and simulate memory allocations
+ bool load_mtp; // whether to load MTP layers
};
struct llama_sampler_seq_config {
LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab);
LLAMA_API bool llama_vocab_get_add_sep(const struct llama_vocab * vocab);
+ // model-specific suppress tokens (gguf key: tokenizer.ggml.suppress_tokens)
+ LLAMA_API const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens);
+
LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab);
LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab);
LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab);
/// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first.
LLAMA_API struct llama_sampler * llama_sampler_init_penalties(
+ int32_t n_vocab,
int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size)
- float penalty_repeat, // 1.0 = disabled
- float penalty_freq, // 0.0 = disabled
- float penalty_present); // 0.0 = disabled
+ float penalty_repeat, // must be > 0.0, 1.0 = disabled
+ float penalty_freq, // must be finite, 0.0 = disabled
+ float penalty_present); // must be finite, 0.0 = disabled
/// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
LLAMA_API struct llama_sampler * llama_sampler_init_dry(
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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
hparams.rope_yarn_log_mul /= 0.1f;
}
+ // NextN/MTP
+ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
+ GGML_ASSERT(hparams.n_layer_nextn == 0 ||
+ hparams.n_layer() + hparams.n_layer_nextn == hparams.n_layer_all);
+
// (optional) temperature tuning - used by mistral-large
ml.get_key(LLM_KV_ATTENTION_TEMPERATURE_SCALE, hparams.f_attn_temp_scale, false);
ml.get_key(LLM_KV_ATTENTION_TEMPERATURE_LENGTH, hparams.n_attn_temp_floor_scale, false); // FIXME why not use temperature_length?
}
}
-void llama_model_deepseek2::load_arch_tensors(llama_model_loader &) {
+void llama_model_deepseek2::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t n_expert_shared = hparams.n_expert_shared;
+ const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ 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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
+
const bool is_mla = hparams.is_mla();
// note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
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) {
auto & layer = layers[i];
+ const int flags = i < n_layer ? trunk_flags : mtp_flags;
- 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);
if (q_lora_rank > 0) {
- layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
+ 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}, 0);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags);
if (q_lora_rank > 0) {
- layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
- layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, 0);
+ 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);
} else {
- layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k_mla}, 0);
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, 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}, 0);
+ 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
if (is_mla) {
- layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, 0);
- layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, 0);
+ 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);
} else {
- layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v_mla)}, 0);
+ layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v_mla)}, flags);
}
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0);
+ 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}, 0);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 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}, 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_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}, 0);
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+ 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 | flags);
if (n_expert == 0) {
throw std::runtime_error("n_expert must be > 0");
}
// MoE branch
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
- create_tensor_gate_up_exps(layer, 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}, flags);
+ create_tensor_gate_up_exps(layer, 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}, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ 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
+ if (i >= n_layer) {
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED | flags);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_deepseek2::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);
}
+llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM4 MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM4 MTP currently only supports a single MTP block");
+ GGML_ASSERT(hparams.is_mla() && "GLM4 MTP requires MLA");
+ GGML_ASSERT(hparams.f_attn_temp_scale == 0.0f && "GLM4 MTP does not support attention temperature scaling");
+
+ // The appended MTP block is stored immediately after the main decoder 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");
+ GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
+ GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
+
+ GGML_ASSERT((uint32_t) il >= hparams.n_layer_dense_lead && "GLM4 MTP block expected to use MoE FFN");
+
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_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 kv_lora_rank = hparams.n_lora_kv;
+
+ GGML_ASSERT(n_embd_head_qk_nope >= 1);
+ GGML_ASSERT(hparams.n_lora_q > 0);
+ GGML_ASSERT(layer.wq_a);
+ GGML_ASSERT(layer.attn_q_a_norm);
+ GGML_ASSERT(layer.wq_b);
+ GGML_ASSERT(layer.wkv_a_mqa);
+ GGML_ASSERT(layer.attn_kv_a_norm);
+ GGML_ASSERT(layer.wk_b);
+
+ const bool has_split_exps =
+ layer.ffn_up_exps != nullptr &&
+ layer.ffn_gate_exps != nullptr;
+
+ const bool has_fused_exps = layer.ffn_gate_up_exps != nullptr;
+
+ GGML_ASSERT(has_split_exps || has_fused_exps);
+ GGML_ASSERT(layer.ffn_norm);
+ GGML_ASSERT(layer.ffn_gate_inp);
+ GGML_ASSERT(layer.ffn_down_exps);
+ GGML_ASSERT(layer.ffn_gate_shexp);
+ GGML_ASSERT(layer.ffn_down_shexp);
+ GGML_ASSERT(layer.ffn_up_shexp);
+
+ auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens
+ ? layer.nextn.embed_tokens
+ : model.tok_embd;
+
+ 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_k = build_attn_inp_k();
+
+ ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);
+ cb(q, "mtp_q_a", il);
+
+ q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(q, "mtp_q_a_norm", il);
+
+ q = ggml_mul_mat(ctx0, layer.wq_b, q);
+ cb(q, "mtp_q_b", il);
+
+ 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_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head, 0);
+ cb(q_nope, "mtp_q_nope", il);
+
+ 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_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head,
+ ggml_row_size(q->type, n_embd_head_qk_nope));
+ cb(q_pe, "mtp_q_pe", il);
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+ cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il);
+
+ 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, "mtp_kv_cmpr", il);
+
+ 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, "mtp_k_pe", il);
+
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(kv_cmpr, "mtp_kv_cmpr_norm", il);
+
+ GGML_ASSERT(ext_factor >= 0.0f);
+
+ const float attn_factor_org =
+ attn_factor * (1.0f + 0.1f * logf(1.0f / freq_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_mla));
+
+ 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, "mtp_q_pe_rope", 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, "mtp_k_pe_rope", il);
+
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "mtp_q_nope_perm", il);
+
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+ cb(q_nope_absorbed, "mtp_q_nope_absorbed", il);
+
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il);
+
+ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ cb(Qcur, "mtp_Qcur", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, hparams.n_lora_kv, 1, n_tokens);
+ cb(kv_cmpr, "mtp_kv_cmpr_reshape", il);
+
+ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
+ cb(Kcur, "mtp_Kcur", il);
+
+ ggml_tensor * Vcur = kv_cmpr;
+ cb(Vcur, "mtp_Vcur", il);
+
+ cur = build_attn(inp_attn_k,
+ layer.wo, nullptr, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);
+ cb(cur, "mtp_attn_out", il);
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "mtp_ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ 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,
+ nullptr,
+ layer.ffn_gate_up_exps);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = 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(ffn_shexp, "mtp_ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "mtp_ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "mtp_post_ffn", il);
+
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "GLM4 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;
+
+ if (inp_out_ids) {
+ 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;
+
+ ggml_tensor * head_s = layer.nextn.shared_head_head
+ ? layer.nextn.shared_head_head_s
+ : model.output_s;
+
+ GGML_ASSERT(head_w && "GLM4 MTP: missing LM head (nextn.shared_head_head or model.output)");
+
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
+
llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_graph_context(params) {
// lite variants include DeepSeek-V2-Lite, GigaChat3-10B-A1.8B
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
}
}
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (cparams.embeddings_nextn && !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;
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;
+ case 61: type = LLM_TYPE_685B_A37B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
-void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
+void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
+
+ const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ const 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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
+
const bool is_mla = hparams.is_mla();
if (!is_mla) {
throw std::runtime_error("DEEPSEEK32 architecture requires MLA");
}
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;
- }
+ const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;
auto & layer = layers[i];
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
+ // NextN/MTP tensors - 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);
}
std::unique_ptr<llm_graph_context> llama_model_deepseek32::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);
}
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);
+ ggml_tensor * indexer_score = nullptr;
+ if (cparams.fused_lid) {
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ cb(indexer_score, "indexer_score", il);
+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
+ } else {
+ // 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
+ indexer_score = ggml_relu(ctx0, indexer_kq);
+ cb(indexer_score, "indexer_score", 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;
Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il);
}
}
- if (il == n_layer - 1 && inp_out_ids) {
+ // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
+ // so the early output masking has to be skipped (it is applied after the final norm instead)
+ if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+ // post-norm hidden state feeds the NextN/MTP draft head
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (cparams.embeddings_nextn && !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;
ggml_build_forward_expand(gf, cur);
}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).
+// Semantics mirror the deepseek-family NextN/MTP layer:
+// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->
+// full deepseek32 decoder block (dense MLA attention + sigmoid-gated MoE FFN
+// with shared expert, exactly as the trunk deepseek2 graph builds it) ->
+// shared_head_norm (fallback output_norm) -> shared LM head.
+// The DSA indexer is not used at runtime.
+llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK32 MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK32 MTP currently only supports a single MTP block");
+ GGML_ASSERT(hparams.is_mla() && "DEEPSEEK32 MTP requires MLA");
+
+ const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+ GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+ "nextn_layer_offset out of range [0, n_layer_nextn)");
+ 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");
+ GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
+
+ // 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_qk_rope = hparams.n_rot();
+ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
+
+ 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 the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.
+ GGML_ASSERT(ext_factor >= 0.0f);
+ const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_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));
+
+ // 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_inp(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+
+ 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();
+
+ // MLA with the absorption optimization uses a K-only cache (V is a view of K)
+ auto * inp_attn = build_attn_inp_k();
+
+ ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ // self-attention: dense MLA, same construction as the deepseek2 trunk graph
+ {
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);
+ cb(q, "mtp_q", il);
+
+ q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(q, "mtp_q", il);
+
+ q = ggml_mul_mat(ctx0, layer.wq_b, q);
+ cb(q, "mtp_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, "mtp_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, "mtp_q_pe", il);
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+ cb(kv_cmpr_pe, "mtp_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, "mtp_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, "mtp_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, "mtp_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, "mtp_k_pe", il);
+
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(kv_cmpr, "mtp_kv_cmpr", il);
+
+ // {n_embd_head_qk_nope, n_tokens, n_head}
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "mtp_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, layer.wk_b, q_nope);
+ cb(q_nope_absorbed, "mtp_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, "mtp_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, "mtp_Qcur", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "mtp_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, "mtp_Kcur", il);
+
+ // {kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Vcur = kv_cmpr;
+ cb(Vcur, "mtp_Vcur", il);
+
+ // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
+ cur = build_attn(inp_attn,
+ layer.wo, NULL, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);
+ cb(cur, "mtp_attn_out", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "mtp_ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ // MoE FFN with shared expert - same construction as the deepseek2 trunk graph
+ 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,
+ nullptr,
+ layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ // FFN shared expert
+ ggml_tensor * ffn_shexp =
+ build_ffn(cur,
+ layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "mtp_ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "mtp_post_ffn", il);
+
+ // shared_head_norm applied after the decoder block, before the shared LM head.
+ // The post-norm hidden state seeds the next MTP step.
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "DEEPSEEK32 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;
+ ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w && "DEEPSEEK32 MTP: missing LM head (nextn.shared_head_head or model.output)");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
+
}
void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
+ if (hparams.n_layer_nextn > 0 && hparams.n_layer_nextn < hparams.n_layer_all) {
+ const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;
+ const std::string mtp_probe = "blk." + std::to_string(n_layer_main) + ".nextn.eh_proj.weight";
+ if (ml.get_weight(mtp_probe.c_str()) == nullptr) {
+ hparams.n_layer_nextn = 0;
+ }
+ }
+ GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count");
+
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
- ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer());
- if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer(), 0)) {
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);
+ if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
}
ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
ml.get_key(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count);
+ hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
+
uint32_t n_compress_ratios = 0;
ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios);
- if (n_compress_ratios < hparams.n_layer()) {
+ if (n_compress_ratios < hparams.n_layer_all) {
throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count");
}
ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios);
}
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.set_swa_pattern(0);
+ for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) {
+ hparams.is_swa_impl[il] = true;
+ }
switch (hparams.n_layer()) {
case 43: type = LLM_TYPE_UNKNOWN; break;
}
}
-void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
+void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t hc_dim = hc_mult * n_embd;
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
+ const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+ const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP;
+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
- for (int i = 0; i < n_layer; ++i) {
+ for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
-
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
- layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
- layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
- layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
- layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
- layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
- layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0);
- layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
-
- layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
- layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
- layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
- layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
- layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
- layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
+ const int flags = i < n_layer ? trunk_flags : mtp_flags;
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
+ layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, flags);
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {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}, flags);
+ layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, flags);
+ layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, flags);
+ // for wo_a, the shape in the file is (n_head * n_embd_head / o_groups, o_lora_rank*o_groups)
+ // so we reshape here, to avoid reshaping the tensor in the graph
+ layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, flags | TENSOR_ALLOW_RESHAPE);
+ layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, flags);
+
+ layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags);
+ layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, flags);
+ layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
+ layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags);
+ layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, flags);
+ layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, flags);
const int64_t ratio = hparams.dsv4_compress_ratios[i];
if (ratio != 0) {
const int64_t coff = ratio == 4 ? 2 : 1;
- layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, 0);
- layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, 0);
- layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, 0);
- layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0);
+ layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, flags);
+ layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, flags);
+ layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, flags);
+ layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, flags);
if (ratio == 4) {
const int64_t n_embd_indexer = hparams.indexer_head_size;
- layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, 0);
- layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, 0);
+ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags);
+ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags);
- layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
- layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
- layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, 0);
- layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, 0);
+ layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
+ layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
+ layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, flags);
+ layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, flags);
} else if (ratio != 128) {
throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128");
}
}
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);
if ((uint32_t) i < hparams.dsv4_hash_layer_count) {
- layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, 0);
+ layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, flags);
} else {
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags);
+ }
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
+
+ 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);
+
+ 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);
+
+ 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);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
}
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
-
- 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);
-
- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_deepseek4::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);
}
return ggml_concat(ctx, t, row, 1);
}
-static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) {
- if (dep == nullptr) {
- return t;
+struct dsv4_state_tensors {
+ ggml_tensor * kv;
+ ggml_tensor * score;
+};
+
+static dsv4_state_tensors dsv4_build_state_restore(
+ ggml_context * ctx,
+ const llm_graph_input_dsv4::comp_input & inp,
+ const llama_dsv4_comp_state * state,
+ int32_t il) {
+ dsv4_state_tensors restored = {
+ state->get_kv_all(ctx, il),
+ state->get_score_all(ctx, il),
+ };
+
+ if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) {
+ return restored;
}
- ggml_tensor * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f);
- return ggml_add(ctx, t, zero);
+ ggml_tensor * kv_rows = ggml_get_rows(ctx, restored.kv, inp.state_restore_src_idxs);
+ restored.kv = state->cpy_kv(ctx, kv_rows, inp.state_restore_dst_idxs, il);
+
+ ggml_tensor * score_rows = ggml_get_rows(ctx, restored.score, inp.state_restore_src_idxs);
+ restored.score = state->cpy_score(ctx, score_rows, inp.state_restore_dst_idxs, il);
+
+ return restored;
+}
+
+static dsv4_state_tensors dsv4_build_state_snapshot(
+ ggml_context * ctx,
+ const llm_graph_input_dsv4::comp_input & inp,
+ const llama_dsv4_comp_state * state,
+ ggml_tensor * source_kv,
+ ggml_tensor * source_score,
+ int32_t il) {
+ if (inp.state_snapshot_src_idxs == nullptr || inp.state_snapshot_dst_idxs == nullptr ||
+ source_kv == nullptr || source_score == nullptr) {
+ return {};
+ }
+
+ ggml_tensor * kv_rows = ggml_get_rows(ctx, source_kv, inp.state_snapshot_src_idxs);
+ ggml_tensor * kv = state->cpy_kv(ctx, kv_rows, inp.state_snapshot_dst_idxs, il);
+
+ ggml_tensor * score_rows = ggml_get_rows(ctx, source_score, inp.state_snapshot_src_idxs);
+ ggml_tensor * score = state->cpy_score(ctx, score_rows, inp.state_snapshot_dst_idxs, il);
+
+ return { kv, score };
}
static constexpr int64_t DSV4_CSA_RATIO = 4;
static constexpr int64_t DSV4_HCA_RATIO = 128;
+// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
+static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) {
+ const int64_t hc = x->ne[1];
+
+ ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0);
+ for (int64_t s = 1; s < hc; ++s) {
+ acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
+ }
+ return ggml_scale(ctx, acc, 1.0f/hc);
+}
+
static ggml_tensor * dsv4_hc_affine(
ggml_context * ctx,
ggml_tensor * x,
return x;
}
-ggml_tensor * llama_model_deepseek4::graph::build_hc_weighted_sum(
+ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
ggml_tensor * x,
- ggml_tensor * weights) const {
+ ggml_tensor * weights,
+ int il) const {
+ GGML_ASSERT(x->ne[0] == n_embd);
+ GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult);
+
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t nt = x->ne[2];
- ggml_tensor * acc = nullptr;
+ if (cparams.fused_dsv4_hc_pre && il >= 0) {
+ ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il});
+ return result;
+ }
+
+ ggml_tensor * result = nullptr;
for (int64_t ih = 0; ih < hc; ++ih) {
ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]);
ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]);
-
ggml_tensor * cur = ggml_mul(ctx0, xh, wh);
- acc = acc ? ggml_add(ctx0, acc, cur) : cur;
+ result = result ? ggml_add(ctx0, result, cur) : cur;
}
- return acc;
+ return result;
}
ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn(
ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0);
ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1);
- ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0);
ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc);
- ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0);
pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre);
*post = ggml_scale(ctx0, *post, 2.0f);
cb(*post, "hc_post", il);
- *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
- *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
- *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
- *comb = build_hc_sinkhorn(*comb, il);
+ if (cparams.fused_dsv4_hc_comb) {
+ *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
+ (int32_t) hparams.dsv4_hc_sinkhorn_iters);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});
+ } else {
+ ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
+ ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
+
+ *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
+ *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
+ *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
+ *comb = build_hc_sinkhorn(*comb, il);
+ }
cb(*comb, "hc_comb", il);
- return build_hc_weighted_sum(x, pre);
+ ggml_tensor * result = build_hc_pre(x, pre, il);
+ return result;
}
ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
ggml_tensor * post,
ggml_tensor * comb,
int il) const {
- GGML_UNUSED(il);
+ GGML_ASSERT(x->ne[0] == n_embd);
+ GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult);
+
+ if (cparams.fused_dsv4_hc_post) {
+ ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il});
+ return result;
+ }
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t nt = x->ne[1];
for (int64_t src = 0; src < hc; ++src) {
ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]);
- ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], dst*comb->nb[0] + src*comb->nb[1]);
+ ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2],
+ dst*comb->nb[0] + src*comb->nb[1]);
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst));
}
pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
cb(pre, "hc_head_pre", -1);
- return build_hc_weighted_sum(x, pre);
+ return build_hc_pre(x, pre, -1);
}
ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state(
kv_state = dsv4_append_zero_row(ctx0, kv_state, false);
score_state = dsv4_append_zero_row(ctx0, score_state, true);
- ggml_tensor * prev_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, 0);
- ggml_tensor * cur_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, ratio*n_blocks);
+ const int64_t n_read = ratio*n_blocks;
- ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state, prev_idxs);
- kv_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_prev, n_embd_head, ratio*n_blocks, kv_prev->nb[1], 0));
+ ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs);
+ ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs);
+
+ ggml_tensor * kv_prev = ggml_cont(ctx0,
+ ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0));
kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks);
cb(kv_prev, name, il);
- ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state, prev_idxs);
- score_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, score_prev, n_embd_head, ratio*n_blocks, score_prev->nb[1], 0));
+ ggml_tensor * score_prev = ggml_cont(ctx0,
+ ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0));
score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks);
cb(score_prev, name, il);
- ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state, cur_idxs);
- kv_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_cur, n_embd_head, ratio*n_blocks, kv_cur->nb[1],
- ggml_row_size(kv_cur->type, n_embd_head)));
+ ggml_tensor * kv_cur = ggml_cont(ctx0,
+ ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1],
+ n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head)));
kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks);
- ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state, cur_idxs);
- score_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, score_cur, n_embd_head, ratio*n_blocks, score_cur->nb[1],
- ggml_row_size(score_cur->type, n_embd_head)));
+ ggml_tensor * score_cur = ggml_cont(ctx0,
+ ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1],
+ n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head)));
score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);
ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1);
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);
- indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
- cb(indexer_q, "lid_q", il);
- indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
- cb(indexer_k, "lid_k", il);
+ ggml_tensor * indexer_score = nullptr;
+ if (cparams.fused_lid) {
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask);
+ cb(indexer_score, "lid_score_masked", il);
+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
+ } else {
+ indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
+ cb(indexer_q, "lid_q", il);
+ indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
+ cb(indexer_k, "lid_k", il);
- ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
- cb(indexer_kq, "lid_kq", il);
+ ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
+ cb(indexer_kq, "lid_kq", il);
- indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
- cb(indexer_kq, "lid_kq", il);
+ indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
+ cb(indexer_kq, "lid_kq", il);
- ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
- indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
- indexer_score = ggml_sum_rows(ctx0, indexer_score);
- indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
- cb(indexer_score, "lid_score", il);
+ indexer_score = ggml_relu(ctx0, indexer_kq);
+ indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
+ indexer_score = ggml_sum_rows(ctx0, indexer_score);
+ indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
+ cb(indexer_score, "lid_score", il);
- indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
- cb(indexer_score, "lid_score_masked", il);
+ indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
+ cb(indexer_score, "lid_score_masked", il);
+ }
const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k;
ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const {
+ return build_attention_impl(model, inp_dsv4, nullptr, cur, inp_pos, il);
+}
+
+ggml_tensor * llama_model_deepseek4::graph::build_attention(
+ const llama_model & model,
+ llm_graph_input_attn_k_iswa * inp_mtp,
+ ggml_tensor * cur,
+ ggml_tensor * inp_pos,
+ int il) const {
+ return build_attention_impl(model, nullptr, inp_mtp, cur, inp_pos, il);
+}
+
+ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
+ const llama_model & model,
+ llm_graph_input_dsv4 * inp_dsv4,
+ llm_graph_input_attn_k_iswa * inp_mtp,
+ ggml_tensor * cur,
+ ggml_tensor * inp_pos,
+ int il) const {
+ GGML_ASSERT((inp_dsv4 == nullptr) != (inp_mtp == nullptr));
+
const auto & layer = model.layers[il];
- llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
+ llm_graph_input_dsv4_raw * inp_attn = inp_dsv4 ? inp_dsv4->get_raw() : nullptr;
const int64_t n_embd_head = hparams.n_embd_head_k();
const int64_t n_embd_head_rope = hparams.n_rot();
cb(kv, "kv", il);
const int64_t ratio = hparams.dsv4_compress_ratios[il];
+ GGML_ASSERT(inp_dsv4 || ratio == 0);
ggml_tensor * hca_state_kv = nullptr;
ggml_tensor * hca_state_score = nullptr;
+ ggml_tensor * hca_source_kv = nullptr;
+ ggml_tensor * hca_source_score = nullptr;
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur);
cb(hca_state_kv, "hca_state_kv", il);
GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs);
- ggml_tensor * csa_source_kv = ggml_concat(ctx0,
- inp_dsv4->mctx->get_csa_state()->get_kv(ctx0, il), csa_state_kv, 1);
- ggml_tensor * csa_source_score = ggml_concat(ctx0,
- inp_dsv4->mctx->get_csa_state()->get_score(ctx0, il), csa_state_score, 1);
+ const auto * csa_state = inp_dsv4->mctx->get_csa_state();
+ const dsv4_state_tensors csa_restored = dsv4_build_state_restore(
+ ctx0, inp_dsv4->get_csa(), csa_state, il);
+ ggml_tensor * csa_base_kv = dsv4_view_2d(
+ ctx0, csa_restored.kv, csa_restored.kv->ne[0], csa_state->get_n_rows(), 0);
+ ggml_tensor * csa_base_score = dsv4_view_2d(
+ ctx0, csa_restored.score, csa_restored.score->ne[0], csa_state->get_n_rows(), 0);
+
+ ggml_tensor * csa_source_kv = ggml_concat(ctx0, csa_base_kv, csa_state_kv, 1);
+ ggml_tensor * csa_source_score = ggml_concat(ctx0, csa_base_score, csa_state_score, 1);
ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state(
csa_source_kv,
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0,
kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il));
- csa_state_kv = dsv4_with_zero_dep(ctx0, csa_state_kv, kv_comp_csa_state);
- csa_state_score = dsv4_with_zero_dep(ctx0, csa_state_score, kv_comp_csa_state);
+ ggml_tensor * csa_snapshot_source_kv = ggml_concat(ctx0,
+ csa_restored.kv, csa_state_kv, 1);
+ ggml_tensor * csa_snapshot_source_score = ggml_concat(ctx0,
+ csa_restored.score, csa_state_score, 1);
+
+ const dsv4_state_tensors csa_snapshot = dsv4_build_state_snapshot(
+ ctx0, inp_dsv4->get_csa(), csa_state, csa_snapshot_source_kv, csa_snapshot_source_score, il);
+ if (csa_snapshot.kv != nullptr) {
+ ggml_build_forward_expand(gf, csa_snapshot.kv);
+ }
+ if (csa_snapshot.score != nullptr) {
+ ggml_build_forward_expand(gf, csa_snapshot.score);
+ }
ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs);
ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs);
GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs);
- ggml_tensor * lid_source_kv = ggml_concat(ctx0,
- inp_dsv4->mctx->get_lid_state()->get_kv(ctx0, il), lid_state_kv, 1);
- ggml_tensor * lid_source_score = ggml_concat(ctx0,
- inp_dsv4->mctx->get_lid_state()->get_score(ctx0, il), lid_state_score, 1);
+ const auto * lid_state = inp_dsv4->mctx->get_lid_state();
+ const dsv4_state_tensors lid_restored = dsv4_build_state_restore(
+ ctx0, inp_dsv4->get_lid(), lid_state, il);
+ ggml_tensor * lid_base_kv = dsv4_view_2d(
+ ctx0, lid_restored.kv, lid_restored.kv->ne[0], lid_state->get_n_rows(), 0);
+ ggml_tensor * lid_base_score = dsv4_view_2d(
+ ctx0, lid_restored.score, lid_restored.score->ne[0], lid_state->get_n_rows(), 0);
+
+ ggml_tensor * lid_source_kv = ggml_concat(ctx0, lid_base_kv, lid_state_kv, 1);
+ ggml_tensor * lid_source_score = ggml_concat(ctx0, lid_base_score, lid_state_score, 1);
ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state(
lid_source_kv,
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0,
kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il));
- lid_state_kv = dsv4_with_zero_dep(ctx0, lid_state_kv, kv_comp_lid_state);
- lid_state_score = dsv4_with_zero_dep(ctx0, lid_state_score, kv_comp_lid_state);
+ ggml_tensor * lid_snapshot_source_kv = ggml_concat(ctx0,
+ lid_restored.kv, lid_state_kv, 1);
+ ggml_tensor * lid_snapshot_source_score = ggml_concat(ctx0,
+ lid_restored.score, lid_state_score, 1);
+
+ const dsv4_state_tensors lid_snapshot = dsv4_build_state_snapshot(
+ ctx0, inp_dsv4->get_lid(), lid_state, lid_snapshot_source_kv, lid_snapshot_source_score, il);
+ if (lid_snapshot.kv != nullptr) {
+ ggml_build_forward_expand(gf, lid_snapshot.kv);
+ }
+ if (lid_snapshot.score != nullptr) {
+ ggml_build_forward_expand(gf, lid_snapshot.score);
+ }
ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs);
ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs);
ggml_build_forward_expand(gf, lid_state_score);
}
- ggml_tensor * hca_state_dep = nullptr;
+ const llama_dsv4_comp_state * hca_state = nullptr;
+ dsv4_state_tensors hca_restored = {};
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) {
GGML_ASSERT(hca_state_kv);
GGML_ASSERT(hca_state_score);
- ggml_tensor * hca_source_kv = ggml_concat(ctx0,
- inp_dsv4->mctx->get_hca_state()->get_kv(ctx0, il), hca_state_kv, 1);
- ggml_tensor * hca_source_score = ggml_concat(ctx0,
- inp_dsv4->mctx->get_hca_state()->get_score(ctx0, il), hca_state_score, 1);
+ hca_state = inp_dsv4->mctx->get_hca_state();
+ hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
+ ggml_tensor * hca_base_kv = dsv4_view_2d(
+ ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
+ ggml_tensor * hca_base_score = dsv4_view_2d(
+ ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
+
+ hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
+ hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state(
hca_source_kv,
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0,
kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il));
- hca_state_dep = kv_comp_hca;
}
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
GGML_ASSERT(hca_state_kv);
GGML_ASSERT(hca_state_score);
- hca_state_kv = dsv4_with_zero_dep(ctx0, hca_state_kv, hca_state_dep);
- hca_state_score = dsv4_with_zero_dep(ctx0, hca_state_score, hca_state_dep);
+ if (hca_state == nullptr) {
+ hca_state = inp_dsv4->mctx->get_hca_state();
+ }
+ if (hca_restored.kv == nullptr) {
+ hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
+ }
+ if (hca_source_kv == nullptr || hca_source_score == nullptr) {
+ ggml_tensor * hca_base_kv = dsv4_view_2d(
+ ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
+ ggml_tensor * hca_base_score = dsv4_view_2d(
+ ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
+
+ hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
+ hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
+ }
+
+ ggml_tensor * hca_snapshot_source_kv = ggml_concat(ctx0,
+ hca_restored.kv, hca_state_kv, 1);
+ ggml_tensor * hca_snapshot_source_score = ggml_concat(ctx0,
+ hca_restored.score, hca_state_score, 1);
+
+ const dsv4_state_tensors hca_snapshot = dsv4_build_state_snapshot(
+ ctx0, inp_dsv4->get_hca(), hca_state, hca_snapshot_source_kv, hca_snapshot_source_score, il);
+ if (hca_snapshot.kv != nullptr) {
+ ggml_build_forward_expand(gf, hca_snapshot.kv);
+ }
+ if (hca_snapshot.score != nullptr) {
+ ggml_build_forward_expand(gf, hca_snapshot.score);
+ }
ggml_tensor * hca_persist_kv = ggml_get_rows(ctx0, hca_state_kv, inp_dsv4->get_hca().state_persist_src_idxs);
ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, inp_dsv4->get_hca().state_persist_src_idxs);
}
ggml_tensor * out = nullptr;
- if (ratio == DSV4_CSA_RATIO &&
+ if (inp_mtp) {
+ out = build_attn(inp_mtp,
+ nullptr, nullptr, nullptr,
+ q, kv, nullptr,
+ nullptr, layer.attn_sinks, nullptr,
+ 1.0f/sqrtf(float(n_embd_head)), il);
+ cb(out, "attn_raw", il);
+ } else if (ratio == DSV4_CSA_RATIO &&
inp_dsv4->get_csa().kq_mask &&
inp_dsv4->get_lid().kq_mask &&
inp_dsv4->get_lid().k_rot) {
out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt);
out = ggml_permute(ctx0, out, 0, 2, 1, 3);
- ggml_tensor * oa = ggml_mul_mat(ctx0, ggml_reshape_3d(ctx0, layer.wo_a, layer.wo_a->ne[0], o_lora_rank, n_groups), out);
+ ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
cb(oa, "attn_wo_a", il);
oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, nt);
cb(inpL, "hc_init", -1);
for (int il = 0; il < n_layer; ++il) {
+ if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) {
+ res->t_layer_inp[il] = dsv4_hc_mean(ctx0, inpL);
+ cb(res->t_layer_inp[il], "layer_inp", il);
+ ggml_build_forward_expand(gf, res->t_layer_inp[il]);
+ }
+
ggml_tensor * residual = inpL;
ggml_tensor * post = nullptr;
ggml_tensor * comb = nullptr;
&post, &comb, il);
cb(cur, "hc_ffn_pre", il);
+ ggml_build_forward_expand(gf, residual);
+ ggml_build_forward_expand(gf, post);
+ ggml_build_forward_expand(gf, comb);
+
cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
inpL = build_hc_post(cur, residual, post, comb, il);
inpL = build_cvec(inpL, il);
- cb(inpL, "l_out", il);
+ cb(inpL, "l_last", il);
+ }
+
+ if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) {
+ res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL);
+ cb(res->t_layer_inp[n_layer], "layer_inp", n_layer);
+ ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]);
+ }
+
+ ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+ ggml_tensor * flat_out = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
+
+ if (cparams.embeddings_nextn) {
+ ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL;
+ cb(h_nextn, "h_nextn", -1);
+ res->t_h_nextn = h_nextn;
}
if (inp_out_ids) {
- ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
- flat = ggml_get_rows(ctx0, flat, inp_out_ids);
- inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);
+ inpL = ggml_reshape_3d(ctx0, flat_out, n_embd, hc, n_outputs);
}
cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
ggml_build_forward_expand(gf, cur);
}
+
+
+llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
+ graph(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK4 MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK4 MTP currently only supports a single MTP block");
+ GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+ "nextn_layer_offset out of range [0, n_layer_nextn)");
+ GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input");
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "DEEPSEEK4 MTP hidden width mismatch");
+
+ const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+ 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");
+
+ auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd_out());
+
+ 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_out(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
+ ggml_set_input(inp->h);
+ ggml_set_name(inp->h, "mtp_h_input");
+
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+ ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ cb(tok_embd, "mtp_tok_embd", il);
+
+ ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens);
+ cb(h_state, "mtp_h_state", il);
+
+ res->add_input(std::move(inp));
+
+ ggml_tensor * inp_pos = build_inp_pos();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ ggml_tensor * h_norm = build_norm(h_state, 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);
+ e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens);
+ e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * inpL = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(inpL, "mtp_eh_proj", il);
+
+ ggml_tensor * residual = inpL;
+ ggml_tensor * post = nullptr;
+ ggml_tensor * comb = nullptr;
+
+ ggml_tensor * cur = build_hc_pre(inpL,
+ layer.hc_attn_fn,
+ layer.hc_attn_scale,
+ layer.hc_attn_base,
+ &post, &comb, il);
+ cb(cur, "mtp_hc_attn_pre", il);
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ cur = build_attention(model, inp_attn, cur, inp_pos, il);
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ cb(inpL, "mtp_hc_attn_post", il);
+
+ residual = inpL;
+ cur = build_hc_pre(inpL,
+ layer.hc_ffn_fn,
+ layer.hc_ffn_scale,
+ layer.hc_ffn_base,
+ &post, &comb, il);
+ cb(cur, "mtp_hc_ffn_pre", il);
+
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ GGML_ASSERT((uint32_t) il >= hparams.dsv4_hash_layer_count && "DEEPSEEK4 MTP does not support hash-routed MTP blocks");
+ 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, hparams.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 * ffn_shexp = 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(ffn_shexp, "mtp_ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "mtp_ffn_out", il);
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ inpL = build_cvec(inpL, il);
+ cb(inpL, "mtp_l_out", il);
+
+ ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+ ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids);
+ cb(h_nextn, "h_nextn", -1);
+ res->t_h_nextn = h_nextn;
+
+ inpL = ggml_reshape_3d(ctx0, h_nextn, n_embd, hc, n_outputs);
+
+ cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+ cb(cur, "mtp_hc_head", -1);
+
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
+ GGML_ASSERT(head_norm_w && "DEEPSEEK4 MTP missing shared head norm");
+ cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "mtp_shared_head_norm", -1);
+ res->t_embd = cur;
+
+ ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+ GGML_ASSERT(head_w && "DEEPSEEK4 MTP missing LM head");
+ cur = ggml_mul_mat(ctx0, head_w, cur);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
#include "models.h"
+#include "llama-impl.h"
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
}
LLAMA_LOG_INFO("]\n");
+ // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
+ ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
+ if (hparams.dsv4_hc_mult > 0) {
+ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+ 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);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);
+ if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
+ hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
+ }
+ ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);
+ ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
+ ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
+
+ if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
+ throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
+ }
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ if (hparams.dsv4_compress_ratios[il] != 0) {
+ throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");
+ }
+ }
+
+ GGML_ASSERT(hparams.n_swa > 0);
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+ hparams.set_swa_pattern(0);
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ hparams.is_swa_impl[il] = true;
+ }
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+
+ type = LLM_TYPE_UNKNOWN;
+ return;
+ }
+
// optional interleaved sliding-window attention with per-layer pattern array.
// DFlash has a single rope, so the SWA rope == main rope.
if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
+ // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
+ //
+ // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
+ // need their own conversion path and graph tweaks
+ const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight");
+ if (markov_meta) {
+ const int64_t dspark_markov_rank = markov_meta->ne[0];
+
+ dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
+ dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0);
+
+ dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0);
+ dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
+
+ LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
+ }
+
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
+ if (hparams.dsv4_hc_mult > 0) {
+ const int64_t q_lora_rank = hparams.n_lora_q;
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+ const int64_t n_embd_head = hparams.n_embd_head_k();
+ const int64_t o_groups = hparams.dsv4_o_group_count;
+ const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
+ const int64_t hc_mult = hparams.dsv4_hc_mult;
+ const int64_t hc_dim = hc_mult * n_embd;
+ const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
+
+ hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
+ hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
+ hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 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);
+ layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
+ layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
+ layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
+ layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
+ layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0);
+ layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
+
+ layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
+ layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
+ layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
+ layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
+
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ 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);
+
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ }
+ return;
+ }
+
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
+ if (hparams.dsv4_hc_mult > 0) {
+ return std::make_unique<graph_dsv4>(*this, params);
+ }
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
ggml_build_forward_expand(gf, cur);
}
+// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position
+static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
+ ggml_context * ctx0 = g.ctx0;
+ auto & res = g.res;
+
+ ggml_tensor * w1 = model.dspark_markov_w1;
+ ggml_tensor * w2 = model.dspark_markov_w2;
+ GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded");
+
+ ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
+ const int64_t n_vocab = base->ne[0];
+ const int64_t n_tok = base->ne[1];
+
+ const auto it = model.gguf_kv.find("dflash.block_size");
+ GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata");
+ const int64_t block_size = std::stoi(it->second);
+ GGML_ASSERT(block_size > 0);
+
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");
+ // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size
+ const int64_t block_drafts = n_tok / n_blocks;
+ if (block_drafts > block_size) {
+ return;
+ }
+
+ // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view
+ const size_t token_stride = (size_t) block_drafts * tokens->nb[0];
+ const size_t base_stride = (size_t) block_drafts * base->nb[1];
+
+ ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);
+ prev = ggml_cont_1d(ctx0, prev, n_blocks);
+
+ // confidence head input: predicts per-position acceptance
+ ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]
+
+ ggml_tensor * cat = nullptr;
+ ggml_tensor * cat_conf = nullptr;
+
+ // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
+ // token pick, not the Markov conditioning path
+ for (int64_t i = 0; i < block_drafts; ++i) {
+ ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
+ ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks]
+
+ // position i of every block: strided view [n_vocab, n_blocks]
+ ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);
+ ggml_tensor * col = ggml_add(ctx0, base_i, bias);
+
+ cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;
+
+ // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]
+ ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
+ (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
+ ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);
+ ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
+ if (model.dspark_conf_proj_b) {
+ conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
+ }
+ conf = ggml_sigmoid(ctx0, conf);
+
+ cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
+
+ if (i + 1 < block_drafts) {
+ prev = ggml_argmax(ctx0, col);
+ }
+ }
+
+ // cat is position-major; restore ubatch block-major order
+ ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts);
+ out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]
+ out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);
+
+ {
+ ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);
+ conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));
+ conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);
+
+ // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn`
+ conf = ggml_repeat(ctx0, conf, res->t_embd);
+ res->t_h_nextn = conf;
+ ggml_build_forward_expand(g.gf, conf);
+ }
+
+ res->t_logits = out;
+ ggml_build_forward_expand(g.gf, out);
+}
+
// DFlash decoder, dual-mode by batch type:
// * embd batch -> fused target features: project + inject K/V into the cache.
// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base();
ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs();
ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs();
+ // rotate K/V into the cache's rotated space
+ ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa : inp_attn_iswa->self_k_rot;
+ ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa : inp_attn_iswa->self_v_rot;
+ if (k_rot) {
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);
+ }
+ if (v_rot) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot);
+ }
ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il));
ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il));
} else {
+ // rotate K/V into the cache's rotated space
+ if (inp_attn->self_k_rot) {
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
+ }
+ if (inp_attn->self_v_rot) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+ }
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
}
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
+ ggml_tensor * inp_tokens = inp->tokens;
+
ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
cb(inpL, "inp_noise_embd", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
+
+ // DSpark: bias the draft logits with the Markov head
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens);
+ }
+}
+
+// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
+// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache
+// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
+llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :
+ llama_model_deepseek4::graph(params) {
+ const int64_t n_embd_head = hparams.n_embd_head_k();
+ const int64_t n_embd_head_rope = hparams.n_rot();
+ const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ // KV cache injection: fused target features from the encoder
+ if (ubatch.embd) {
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * inp_g = inp->embd;
+ cb(inp_g, "inp_g_embeddings", -1);
+
+ res->add_input(std::move(inp));
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same
+ // rope parameters as the uncompressed layers in build_attention_impl
+ ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
+ kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+ kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+
+ ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, n_tokens,
+ ggml_row_size(kv->type, n_embd_head),
+ ggml_row_size(kv->type, n_embd_head),
+ 0);
+ ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, n_tokens,
+ ggml_row_size(kv->type, n_embd_head),
+ ggml_row_size(kv->type, n_embd_head),
+ ggml_row_size(kv->type, n_embd_head_nope));
+ kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+ freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ kv = ggml_concat(ctx0, kv_nope, kv_pe, 0);
+ cb(kv, "kv_injected", il);
+
+ if (inp_attn->self_k_rot_swa) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+ }
+
+ res->t_embd = inp_g;
+
+ ggml_build_forward_expand(gf, inp_g);
+ return;
+ }
+
+ // tok_embd from the target model (shared via ctx_other)
+ auto * tok_embd = model.tok_embd;
+ if (tok_embd == nullptr) {
+ GGML_ASSERT(cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(cparams.ctx_other);
+
+ GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
+ tok_embd = model_other->tok_embd;
+ }
+
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ ggml_tensor * inp_tokens = inp->tokens;
+
+ ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+ cb(inpL, "inp_noise_embd", -1);
+
+ res->add_input(std::move(inp));
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
+ inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
+ cb(inpL, "hc_init", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * residual = inpL;
+ ggml_tensor * post = nullptr;
+ ggml_tensor * comb = nullptr;
+
+ ggml_tensor * cur = build_hc_pre(inpL,
+ layer.hc_attn_fn,
+ layer.hc_attn_scale,
+ layer.hc_attn_base,
+ &post, &comb, il);
+ cb(cur, "hc_attn_pre", il);
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ cur = build_attention(model, inp_attn, cur, inp_pos, il);
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ cb(inpL, "hc_attn_post", il);
+
+ residual = inpL;
+ cur = build_hc_pre(inpL,
+ layer.hc_ffn_fn,
+ layer.hc_ffn_scale,
+ layer.hc_ffn_base,
+ &post, &comb, il);
+ cb(cur, "hc_ffn_pre", il);
+
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ 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, hparams.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, "ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = 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(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ cb(inpL, "l_out", il);
+ }
+
+ ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+ cb(cur, "hc_head", -1);
+
+ // confidence head input: the reference scores the pre-norm collapsed hidden state
+ res->t_embd = cur;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+
+ // lm_head from the target model (shared via ctx_other)
+ auto * output = model.output;
+ if (output == nullptr) {
+ GGML_ASSERT(cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(cparams.ctx_other);
+ GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
+ output = model_other->output;
+ }
+
+ cur = build_lora_mm(output, cur);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens);
+ }
}
LLAMA_LOG_INFO("%s: EAGLE3gnorm_before_residual = true\n", __func__);
}
+ // eagle3 norm_before_fc (optional, default false)
+ // compatible with eagle3.1 (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
+ ml.get_key(LLM_KV_NORM_BEFORE_FC, hparams.norm_before_fc, false);
+
type = LLM_TYPE_UNKNOWN;
}
// Feature fusion layer: projects 3 target layers to draft hidden size
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), {n_embd_inp, n_embd}, 0);
+ // RMSNorm on the fused target features (input to fc), only when norm_before_fc is set.
+ if (hparams.norm_before_fc) {
+ output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), {n_embd_inp}, 0);
+ }
+
// Output layer (uses draft vocab size)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_draft_vocab}, TENSOR_NOT_REQUIRED);
cur = build_inp_embd_enc();
+ // RMSNorm on the fused target features before fc
+ if (hparams.norm_before_fc) {
+ cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
+ cb(cur, "enc_input_norm", -1);
+ }
+
// Feature fusion layer
cur = build_lora_mm(model.fc, cur);
cb(cur, "fc_out", -1);
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),
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;
#include "models.h"
+#include "llama-kv-cache-dsa.h"
+
+// https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L26
+const std::array<uint32_t, LLAMA_MAX_LAYERS> GLM_5_2_DEFAULT_INDEXER_TYPES = {
+ 1, 1,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+ 1, 0, 0, 0,
+};
+
void llama_model_glm_dsa::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);
// 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_impl");
+ GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
+
+ // BC for GLM 5, 5.1 (full indexers) without indexer_types metadata
+ const bool is_pre_5_2 = hparams.n_ctx_train < 1048576;
+ if (is_pre_5_2) {
+ std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);
+ } else {
+ hparams.is_indexer_full_impl = GLM_5_2_DEFAULT_INDEXER_TYPES;
+ }
+ ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);
switch (hparams.n_layer()) {
- case 79: type = LLM_TYPE_744B_A40B; break;
+ case 78: // GGUF with NextN/MTP metadata: n_layer() excludes the nextn layer
+ case 79:
+ type = LLM_TYPE_744B_A40B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
-void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) {
+void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t n_expert_shared = hparams.n_expert_shared;
+ // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft).
+ 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 (or were stripped at conversion). 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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
+
const bool is_mla = hparams.is_mla();
if (!is_mla) {
throw std::runtime_error("GLM_DSA architecture requires MLA");
}
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;
- }
+ // NextN/MTP layers (i >= n_layer) are full decoder blocks used by the
+ // LLM_GRAPH_TYPE_DECODER_MTP draft head; load them like qwen35moe/step35/hy_v3.
+ const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;
auto & layer = layers[i];
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 n_layer_nextn
+ // NextN/MTP tensors - the NextN-specific wiring around the extra decoder block
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);
}
std::unique_ptr<llm_graph_context> llama_model_glm_dsa::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);
}
+llama_model_glm_dsa::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();
+
+ // Difference vs Deepseek 3.2: shared indexer layers reuse the top_k from the previous full indexer layers
+ // See https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L30
+ ggml_tensor * prev_top_k = nullptr;
+ 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
+ if (hparams.is_indexer_full(il)) {
+ // "full" layer
+ 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_NORM, 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_NORM, 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);
+
+ // 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);
+
+ ggml_tensor * indexer_score = nullptr;
+ if (cparams.fused_lid) {
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ cb(indexer_score, "indexer_score", il);
+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
+ } else {
+ // 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
+ indexer_score = ggml_relu(ctx0, indexer_kq);
+ cb(indexer_score, "indexer_score", 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));
+ prev_top_k = top_k;
+ cb(top_k, "top_k", il);
+ } else {
+ // "shared" indexer layer - reuse top-k from a previous full layer
+ GGML_ASSERT(prev_top_k != nullptr && "shared indexer layer must follow a previous full indexer layer");
+ top_k = prev_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);
+ }
+ }
+ // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
+ // so the early output masking has to be skipped (it is applied after the final norm instead)
+ if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ 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);
+
+ // post-norm hidden state feeds the NextN/MTP draft head
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (cparams.embeddings_nextn && !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;
+
+ // 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);
+}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for GLM-5.2 (GLM_DSA).
+// Semantics mirror the deepseek-family NextN/MTP layer:
+// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->
+// full glm_dsa decoder block (dense MLA attention + sigmoid-gated MoE FFN
+// with shared expert, exactly as the trunk deepseek2 graph builds it) ->
+// shared_head_norm (fallback output_norm) -> shared LM head.
+// The DSA indexer is not used at runtime (same as the trunk graph).
+llama_model_glm_dsa::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM_DSA MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM_DSA MTP currently only supports a single MTP block");
+ GGML_ASSERT(hparams.is_mla() && "GLM_DSA MTP requires MLA");
+
+ const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+ GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+ "nextn_layer_offset out of range [0, n_layer_nextn)");
+ 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");
+ GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
+
+ // 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_qk_rope = hparams.n_rot();
+ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
+
+ 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 the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.
+ GGML_ASSERT(ext_factor >= 0.0f);
+ const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_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));
+
+ // 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_inp(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+
+ 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();
+
+ // MLA with the absorption optimization uses a K-only cache (V is a view of K)
+ auto * inp_attn = build_attn_inp_k();
+
+ ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ // self-attention: dense MLA, same construction as the deepseek2 trunk graph
+ {
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);
+ cb(q, "mtp_q", il);
+
+ q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(q, "mtp_q", il);
+
+ q = ggml_mul_mat(ctx0, layer.wq_b, q);
+ cb(q, "mtp_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, "mtp_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, "mtp_q_pe", il);
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+ cb(kv_cmpr_pe, "mtp_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, "mtp_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, "mtp_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, "mtp_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, "mtp_k_pe", il);
+
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(kv_cmpr, "mtp_kv_cmpr", il);
+
+ // {n_embd_head_qk_nope, n_tokens, n_head}
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "mtp_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, layer.wk_b, q_nope);
+ cb(q_nope_absorbed, "mtp_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, "mtp_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, "mtp_Qcur", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "mtp_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, "mtp_Kcur", il);
+
+ // {kv_lora_rank, 1, n_tokens}
+ ggml_tensor * Vcur = kv_cmpr;
+ cb(Vcur, "mtp_Vcur", il);
+
+ // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
+ cur = build_attn(inp_attn,
+ layer.wo, NULL, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);
+ cb(cur, "mtp_attn_out", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "mtp_ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ // MoE FFN with shared expert - same construction as the deepseek2 trunk graph
+ 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,
+ nullptr,
+ layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ // FFN shared expert
+ ggml_tensor * ffn_shexp =
+ build_ffn(cur,
+ layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "mtp_ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "mtp_post_ffn", il);
+
+ // shared_head_norm applied after the decoder block, before the shared LM head.
+ // The post-norm hidden state seeds the next MTP step.
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "GLM_DSA 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;
+ ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w && "GLM_DSA MTP: missing LM head (nextn.shared_head_head or model.output)");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
--- /dev/null
+#include "models.h"
+
+void llama_model_hy_v3::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_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, 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);
+
+ // HY V3 uses a sigmoid router with expert selection bias by default
+ if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
+ hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
+ }
+
+ // NextN/MTP (HY V3): extra decoder block(s) 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_all");
+
+ switch (hparams.n_layer()) {
+ case 48: type = LLM_TYPE_30B_A3B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_hy_v3::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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+ if (output == NULL) {
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+ }
+
+ auto load_block = [&](int i, int flags) {
+ auto & layer = layers[i];
+ const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / (n_expert_used > 0 ? n_expert_used : 1);
+ const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp;
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
+
+ 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_head_k * n_head, 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.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
+
+ // dense FFN (leading dense blocks, first_k_dense_replace)
+ 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 (sigmoid router + expert selection bias)
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {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);
+ create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, TENSOR_NOT_REQUIRED);
+
+ // shared expert (always active, no gate)
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
+ };
+
+ for (int i = 0; i < n_layer; ++i) {
+ load_block(i, trunk_flags);
+ }
+
+ // NextN/MTP block(s): a full hy_v3 decoder block plus the NextN projections.
+ for (int i = n_layer; i < n_layer_all; ++i) {
+ auto & layer = layers[i];
+
+ load_block(i, mtp_flags);
+
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, mtp_flags);
+ layer.nextn.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);
+ // hy_v3 stores the MTP block's trailing final_layernorm here (applied
+ // after the decoder block, before the shared LM head).
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED);
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_hy_v3::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);
+}
+
+llama_model_hy_v3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+ GGML_ASSERT(n_embd_head == n_rot);
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+ ggml_tensor * inp_pos = build_inp_pos();
+ auto * inp_attn = build_attn_inp_kv();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
+
+ // 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;
+
+ cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // self-attention
+ {
+ ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+
+ 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);
+ Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);
+
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ 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, kq_scale, il);
+ cb(cur, "attn_out", il);
+ }
+
+ 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);
+ }
+
+ 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, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ if (model.layers[il].ffn_gate_inp == nullptr) {
+ // dense FFN (leading dense blocks)
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s,
+ model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s,
+ model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_dense_out", il);
+ } else {
+ // MoE routed experts (sigmoid gating + expert selection bias)
+ 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);
+
+ // shared expert (always active, no gate)
+ ggml_tensor * sh_out = build_ffn(cur,
+ model.layers[il].ffn_up_shexp, nullptr, model.layers[il].ffn_up_shexp_s,
+ model.layers[il].ffn_gate_shexp, nullptr, model.layers[il].ffn_gate_shexp_s,
+ model.layers[il].ffn_down_shexp, nullptr, model.layers[il].ffn_down_shexp_s,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(sh_out, "ffn_shared_out", il);
+
+ cur = ggml_add(ctx0, moe_out, sh_out);
+ cb(cur, "ffn_out", il);
+ }
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+ }
+
+ cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+
+ // Post-final-norm hidden state: what the MTP draft head's hnorm consumes.
+ // vLLM feeds the target model's normed output states, and the MTP layer
+ // itself returns final_layernorm(h), so the chained state is post-norm.
+ 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 = build_lora_mm(model.output, cur, model.output_s);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for HY V3 (MoE).
+// Semantics mirror vLLM's HYV3MultiTokenPredictorLayer (hy_v3_mtp.py):
+// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->
+// hy_v3 decoder block -> final_layernorm (stored as nextn.shared_head_norm) ->
+// shared LM head (the main model's lm_head; the checkpoint has no separate
+// MTP head or MTP embeddings).
+llama_model_hy_v3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "HY_V3 MTP requires n_layer_nextn > 0");
+
+ 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);
+
+ const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+ GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+ "nextn_layer_offset out of range [0, n_layer_nextn)");
+ 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");
+
+ 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();
+ 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);
+ 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: a full hy_v3 decoder layer (mirrors the trunk graph)
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ {
+ ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+
+ auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il);
+
+ Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
+
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
+
+ cur = build_attn(inp_attn,
+ layer.wo, layer.wo_b, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(cur, "mtp_attn_out", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "mtp_ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ if (layer.ffn_gate_inp == nullptr) {
+ cur = build_ffn(cur,
+ layer.ffn_up, layer.ffn_up_b, layer.ffn_up_s,
+ layer.ffn_gate, layer.ffn_gate_b, layer.ffn_gate_s,
+ layer.ffn_down, layer.ffn_down_b, layer.ffn_down_s,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "mtp_ffn_dense_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,
+ nullptr, layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ ggml_tensor * sh_out = build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,
+ 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);
+
+ // final_layernorm applied after the decoder block, before the shared head.
+ // The post-norm hidden state seeds the next MTP step (matches vLLM, where
+ // HYV3MultiTokenPredictorLayer returns final_layernorm(h)).
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "HY_V3 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;
+ ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w && "HY_V3 MTP: missing LM head (nextn.shared_head_head or model.output)");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
--- /dev/null
+// Laguna (poolside): sigmoid-routed MoE with a score-correction bias, one shared
+// expert, a softplus attention output gate, QK-norm, and per-layer-type RoPE
+// (YaRN on full-attention layers, plain RoPE on sliding-window layers). XS.2 is
+// hybrid full/SWA with a per-head gate; M.1 is full-attention with a per-element
+// gate. Shares the MoE/gate structure with afmoe.
+
+#include "models.h"
+
+void llama_model_laguna::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_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
+ 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);
+ 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);
+
+ // Laguna ships one shared expert and stores its size directly (routed and
+ // shared experts may differ), so read the size from expert_shared_feed_forward_length.
+ // The count is not in the config; default to 1 but read the key if present.
+ hparams.n_expert_shared = 1;
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared, false);
+ ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
+ if (hparams.n_ff_shexp == 0) {
+ // Weightless fixtures (test-llama-archs) omit this key; derive a nonzero
+ // size so the shared expert is still built. Real GGUFs always carry the
+ // exact value (routed and shared FF lengths may differ).
+ hparams.n_ff_shexp = hparams.n_ff_exp * hparams.n_expert_shared;
+ }
+
+ // Sliding-window attention is OPTIONAL. XS.2 is hybrid (full / SWA / SWA /
+ // SWA repeating, period 4 starting with full); M.1 has no sliding window
+ // (all layers full attention). When sliding_window is absent or zero we
+ // leave swa_type = NONE and skip the SWA-specific per-layer-type RoPE.
+ hparams.n_swa = 0;
+ 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;
+ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
+ hparams.set_swa_pattern(swa_period, /*dense_first=*/true); // XS.2: FULL at il%4==0
+
+ // Per-layer-type RoPE: full layers use YaRN θ=500000 over 64 dims;
+ // SWA layers use default RoPE θ=10000 over 128 dims. Base load_hparams
+ // already reads ROPE_FREQ_BASE and ROPE_DIMENSION_COUNT into the
+ // non-SWA fields; we explicitly pull the SWA mirrors here.
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = 1.0f; // SWA uses plain RoPE (no YaRN scaling); do NOT inherit full layers 1/factor
+ ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
+ ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT_SWA, hparams.n_rot_swa, false);
+ }
+
+ // Default the expert gating function to SIGMOID when the key is absent
+ // (matches the HF reference).
+ if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
+ hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
+ }
+
+ switch (hparams.n_layer()) {
+ case 40: type = LLM_TYPE_30B_A3B; break; // Laguna-XS.2
+ case 48: type = LLM_TYPE_118B_A8B; break; // Laguna-S.2
+ case 70: type = LLM_TYPE_230B_A10B; break; // Laguna-M.1
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_laguna::load_arch_tensors(llama_model_loader & ml) {
+ LLAMA_LOAD_LOCALS;
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+ if (output == NULL) {
+ // tied embeddings fallback
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+ }
+
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+ const int64_t n_ff_shexp = hparams.n_ff_shexp;
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ // Per-layer head count — Laguna varies n_head between full and SWA
+ // layers (48 vs 64 in XS.2). KV head count is uniform.
+ const int64_t n_head_il = hparams.n_head(i);
+ const int64_t n_head_kv_il = hparams.n_head_kv(i);
+ const int64_t n_embd_q_il = n_embd_head_k * n_head_il;
+ const int64_t n_embd_k_il = n_embd_head_k * n_head_kv_il;
+ const int64_t n_embd_v_il = n_embd_head_v * n_head_kv_il;
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+ create_tensor_qkv(layer, i, n_embd, n_embd_q_il, n_embd_k_il, n_embd_v_il, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q_il, 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);
+
+ // Attention output gate. XS.2 is per-head (g_proj -> n_head, one scalar
+ // per head broadcast over head_dim at multiply time); M.1 is per-element
+ // (g_proj -> n_head*head_dim, like afmoe). Detect from the stored tensor
+ // shape so a single arch handles both; the graph mirrors this check.
+ // Gate width selects per-head vs per-element. Real GGUFs always carry the
+ // gate tensor, so read the width from it and require EXACTLY one of the two
+ // valid widths -- never guess between them. Weightless fixtures
+ // (test-llama-archs) have no gate tensor; fall back to the per-head layout so
+ // the per-head reshape path is still exercised.
+ const int64_t n_gate_per_head = n_head_il;
+ const int64_t n_gate_per_elem = n_embd_head_k * n_head_il;
+ const ggml_tensor * gate_meta = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_GATE, "weight", i).str().c_str());
+ int64_t n_gate_out;
+ if (gate_meta != nullptr) {
+ n_gate_out = gate_meta->ne[1];
+ if (n_gate_out != n_gate_per_head && n_gate_out != n_gate_per_elem) {
+ GGML_ABORT("Laguna: unexpected attention gate width %lld at layer %d "
+ "(expected %lld per-head or %lld per-element)",
+ (long long) n_gate_out, i, (long long) n_gate_per_head, (long long) n_gate_per_elem);
+ }
+ } else {
+ n_gate_out = n_gate_per_head;
+ }
+ layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_gate_out}, 0);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ if ((uint32_t)i >= hparams.n_layer_dense_lead) {
+ // MoE layer
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, 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);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
+
+ // Always-on shared expert.
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
+ } else {
+ // Dense layer (the leading n_layer_dense_lead layers)
+ 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);
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_laguna::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+ // No MuP embedding scale (laguna omits this; afmoe scales by sqrt(hidden)).
+
+ ggml_tensor * inp_pos = build_inp_pos();
+ // XS.2 is hybrid SWA -> interleaved-SWA KV input; M.1 is all-full -> plain
+ // KV input. Pick the matching input (and build_attn overload) per swa_type.
+ const bool has_swa = hparams.swa_type != LLAMA_SWA_TYPE_NONE;
+ llm_graph_input_attn_kv * inp_attn_kv = has_swa ? nullptr : build_attn_inp_kv();
+ llm_graph_input_attn_kv_iswa * inp_attn_iswa = has_swa ? build_attn_inp_kv_iswa() : nullptr;
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
+
+ for (int il = 0; il < n_layer; ++il) {
+ const bool is_swa_il = hparams.is_swa(il);
+ const int64_t n_head_il = hparams.n_head(il);
+ const int64_t n_head_kv_il = hparams.n_head_kv(il);
+
+ // Per-layer-type RoPE config. SWA layers run plain rope (no YaRN),
+ // achieved by zeroing the YaRN ext/beta params for those layers.
+ const int n_rot_l = is_swa_il ? hparams.n_rot_swa : n_rot;
+ const float freq_base_l = is_swa_il ? hparams.rope_freq_base_train_swa : freq_base;
+ const float freq_scale_l = is_swa_il ? hparams.rope_freq_scale_train_swa : freq_scale;
+ const float ext_factor_l = is_swa_il ? 0.0f : ext_factor;
+ // YaRN magnitude scaling (mscale) is already handled by the framework:
+ // llama_context pre-divides cparams.yarn_attn_factor by (1 + 0.1*ln(factor))
+ // to cancel ggml rope_yarn's internal mscale *= 1 + 0.1*ln(1/freq_scale).
+ // Pass attn_factor straight through (like every other arch); SWA layers run
+ // plain RoPE (ext_factor 0, no mscale) so force 1.0 there.
+ const float attn_factor_l = is_swa_il ? 1.0f : attn_factor;
+ const float beta_fast_l = is_swa_il ? 0.0f : beta_fast;
+ const float beta_slow_l = is_swa_il ? 0.0f : beta_slow;
+ const int n_ctx_orig_l = is_swa_il ? hparams.n_ctx_train : n_ctx_orig;
+
+ ggml_tensor * inpSA = inpL;
+
+ // Pre-norm
+ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // Self-attention
+ {
+ ggml_tensor * attn_inp = cur; // saved for the gate projection
+
+ auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+ n_embd_head, n_head_il, n_head_kv_il, il);
+
+ // g_proj on the *pre-attention* hidden state (matches HF
+ // reference: gate is computed from the same `hidden_states`
+ // input as q/k/v, not from the attn output).
+ ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
+ cb(gate, "attn_gate_proj", il);
+
+ // QK RMSNorm at head_dim level (Qwen3 style)
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
+ Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+ cb(Kcur, "Kcur_normed", il);
+
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
+ n_rot_l, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l,
+ ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+ Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
+ n_rot_l, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l,
+ ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+ cb(Qcur, "Qcur_rope", il);
+ cb(Kcur, "Kcur_rope", il);
+
+ cur = has_swa
+ ? build_attn(inp_attn_iswa,
+ NULL, NULL, NULL, // o_proj deferred until after gating
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
+ : build_attn(inp_attn_kv,
+ NULL, NULL, NULL,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(cur, "attn_out", il);
+
+ // Softplus output gate (the unary kernel computes softplus in fp32
+ // and casts back). Two shapes, distinguished by the g_proj output
+ // dim (matching the load-time detection):
+ // XS.2 per-head : gate [n_head_il, n_tokens] -> reshape to
+ // [1, n_head_il, n_tokens] and broadcast over
+ // head_dim against cur [head_dim, n_head, T].
+ // M.1 per-element : gate [n_head_il*head_dim, n_tokens] spans the
+ // full attention output -> direct ggml_mul.
+ gate = ggml_softplus(ctx0, gate);
+ cb(gate, "attn_gate_softplus", il);
+
+ const int64_t n_tokens = cur->ne[1];
+ if (model.layers[il].wqkv_gate->ne[1] == n_head_il) {
+ cur = ggml_reshape_3d(ctx0, cur, n_embd_head, n_head_il, n_tokens);
+ gate = ggml_reshape_3d(ctx0, gate, 1, n_head_il, n_tokens);
+ cur = ggml_mul(ctx0, cur, gate);
+ cur = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head_il, n_tokens);
+ } else {
+ cur = ggml_mul(ctx0, cur, gate);
+ }
+ cb(cur, "attn_gated", il);
+
+ cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
+ cb(cur, "attn_o_proj", il);
+ }
+
+ if (il == n_layer - 1 && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "ffn_inp", il);
+
+ // Pre-norm only (no post-attn norm)
+ 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) {
+ // MoE: sigmoid routing + score-correction bias + sum-norm +
+ // routed_scaling_factor (all handled by build_moe_ffn).
+ 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);
+ cb(moe_out, "ffn_moe_out", il);
+
+ // Always-on shared expert, summed in parallel.
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ model.layers[il].ffn_up_shexp, NULL, NULL,
+ model.layers[il].ffn_gate_shexp, NULL, NULL,
+ model.layers[il].ffn_down_shexp, NULL, NULL,
+ 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);
+ } else {
+ // Dense FFN for the leading n_layer_dense_lead layers (XS.2: 1, M.1: 3)
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, NULL, NULL,
+ model.layers[il].ffn_gate, NULL, NULL,
+ model.layers[il].ffn_down, NULL, NULL,
+ NULL,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ }
+
+ // No post-ffn norm
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+ }
+
+ cur = inpL;
+ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ cur = build_lora_mm(model.output, cur);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
}
}
-void llama_model_mimo2::load_arch_tensors(llama_model_loader &) {
+void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
+ 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);
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// output
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 = i >= n_layer;
- const int skip = is_nextn ? TENSOR_SKIP : 0;
+ const int flags = is_nextn ? mtp_flags : 0;
- create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip);
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, skip);
+ 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_head_v * n_head, n_embd }, flags);
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, skip);
- layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | skip);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
+ layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | flags);
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, skip);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
// non-MoE branch
- layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | skip);
- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | flags);
// MoE branch
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 | skip);
- layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
- layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | skip);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | flags);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | flags);
if (is_nextn) {
- layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, skip);
- layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, skip);
- layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, skip);
- layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, skip);
+ 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.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
+ layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_mimo2::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);
}
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float v_scale = hparams.f_attn_value_scale;
+ const bool emit_h_nextn = cparams.embeddings_nextn;
+ const bool crop_last_layer = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
}
}
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_layer - 1 && crop_last_layer) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = inpL;
+ if (emit_h_nextn) {
+ 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, NULL,
LLM_NORM_RMS, -1);
ggml_build_forward_expand(gf, cur);
}
+
+// Mirrors MiMo's appended NextN block: normalize and fuse token and hidden inputs, run the decoder block,
+// expose its pre-head-norm state to the next draft step, then apply the shared output norm and LM head.
+// Converted checkpoints may store that shared norm as layer_out_norm, so it remains in the fallback chain.
+llama_model_mimo2::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "MIMO2 MTP requires n_layer_nextn > 0");
+
+ const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+ GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+ "nextn_layer_offset out of range [0, n_layer_nextn)");
+
+ const auto & layer = model.layers[il];
+ GGML_ASSERT(layer.nextn.eh_proj && "MIMO2 MTP block missing nextn.eh_proj");
+ GGML_ASSERT(layer.nextn.enorm && "MIMO2 MTP block missing nextn.enorm");
+ GGML_ASSERT(layer.nextn.hnorm && "MIMO2 MTP block missing nextn.hnorm");
+ GGML_ASSERT(layer.wqkv && "MIMO2 MTP requires fused attn_qkv");
+
+ 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);
+ const float v_scale = hparams.f_attn_value_scale;
+
+ 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();
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+ 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, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ ggml_tensor * qkv = build_lora_mm(layer.wqkv, cur, layer.wqkv_s);
+ cb(qkv, "mtp_wqkv", il);
+
+ const size_t row_k = ggml_row_size(qkv->type, n_embd_head_k);
+ const size_t row_v = ggml_row_size(qkv->type, n_embd_head_v);
+ const size_t row_full = qkv->nb[1];
+ const size_t k_off = row_k * n_head_l;
+ const size_t v_off = k_off + row_k * n_head_kv_l;
+
+ ggml_tensor * Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_l, n_tokens, row_k, row_full, 0);
+ ggml_tensor * Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_kv_l, n_tokens, row_k, row_full, k_off);
+ ggml_tensor * Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_kv_l, n_tokens, row_v, row_full, v_off);
+
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ cb(Qcur, "mtp_Qcur", il);
+ cb(Kcur, "mtp_Kcur", il);
+ cb(Vcur, "mtp_Vcur", il);
+
+ cur = build_attn(inp_attn,
+ layer.wo, nullptr, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr,
+ 1.0f / sqrtf(float(n_embd_head_k)), il);
+ cb(cur, "mtp_attn_out", il);
+
+ if (v_scale) {
+ cur = ggml_scale(ctx0, cur, v_scale);
+ cb(cur, "mtp_attn_out_scaled", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "mtp_ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_ffn_norm", il);
+
+ GGML_ASSERT(layer.ffn_gate && layer.ffn_down && layer.ffn_up && "MIMO2 MTP requires dense FFN tensors");
+ 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);
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "mtp_post_ffn", il);
+
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+
+ 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
+ : (layer.layer_out_norm ? layer.layer_out_norm : model.output_norm);
+ GGML_ASSERT(head_norm_w && "MIMO2 MTP missing head norm fallback");
+ 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_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w && "MIMO2 MTP missing LM head fallback");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
+ res->t_layer_inp[il] = inpL;
+
ggml_tensor * inpSA = inpL;
cur = inpL;
--- /dev/null
+#include "models.h"
+#include "llama-kv-cache-msa.h"
+#include <cmath>
+#include <vector>
+#include <cstdint>
+
+// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with
+// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling),
+// swigluoai activation, and MiniMax Sparse Attention (MSA). MTP is not in released model weights.
+// MSA blocks are defined over token positions. The graph translates between position space (block
+// selection) and cell space (K/V/indexer storage) via per-ubatch pos<->cell maps populated from llama_kv_cells
+
+void llama_model_minimax_m3::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_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, 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);
+ 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_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);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
+ msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };
+
+ switch (hparams.n_layer()) {
+ case 60: type = LLM_TYPE_428B_A23B; break;
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+ const int64_t n_ff_exp = hparams.n_ff_exp;
+
+ 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];
+
+ 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_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+ // per-head QK-norm: a single head_dim vector applied to every head
+ 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_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ if (i < (int) hparams.n_layer_dense_lead) {
+ // leading dense layers
+ 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);
+ } else {
+ // routed experts
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+ layer.ffn_gate_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);
+
+ // shared expert
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+
+ // indexer
+ layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, 0);
+ layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0);
+ layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0);
+ layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0);
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+class llm_graph_input_msa : public llm_graph_input_i {
+public:
+ llm_graph_input_msa(const llama_kv_cache_msa_context * mctx, int blk, int local) :
+ mctx(mctx), blk(blk), local(local) {}
+
+ void set_input(const llama_ubatch * ubatch) override {
+ if (pos_slot_i) { mctx->set_input_pos_slot(pos_slot_i, ubatch); }
+ if (pos_slot_f) { mctx->set_input_pos_slot(pos_slot_f, ubatch); }
+ if (cell_blk) { mctx->set_input_cell_pos(cell_blk, ubatch, blk); }
+ if (pos_mask) { mctx->set_input_pos_mask(pos_mask, ubatch); }
+
+ // local-force bias over position blocks
+ if (bias && ubatch->pos) {
+ const int64_t n_tokens = ubatch->n_tokens;
+ const int64_t nblk = bias->ne[0];
+ std::vector<float> data((size_t) nblk * n_tokens, 0.0f);
+ for (int64_t i = 0; i < n_tokens; ++i) {
+ const int64_t L = ubatch->pos[i] / blk;
+ for (int l = 0; l < local && L - l >= 0; ++l) {
+ if (L - l < nblk) {
+ data[(size_t) i * nblk + (L - l)] = 1e30f;
+ }
+ }
+ }
+ ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));
+ }
+ }
+
+ // valid as long as the tensor dims still match the new ubatch/cache window and the
+ // ubatch is in the same regime (decode graphs have pos_slot_f, batch graphs cell_blk)
+ bool can_reuse(const llm_graph_params & params) override {
+ const auto * mctx_new = static_cast<const llama_kv_cache_msa_context *>(params.mctx);
+
+ this->mctx = mctx_new;
+
+ const int64_t n_ps = GGML_PAD((int64_t) mctx_new->get_n_pos(), blk);
+ const int64_t ns = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq;
+
+ const bool decode = params.ubatch.n_tokens == ns; // one token per stream
+
+ bool res = true;
+
+ res &= bias->ne[0] * blk == n_ps;
+ res &= bias->ne[1] == params.ubatch.n_tokens;
+
+ res &= pos_mask->ne[0] == n_ps;
+ res &= pos_mask->ne[1] == params.ubatch.n_tokens;
+
+ res &= pos_slot_i->ne[0] == n_ps;
+ res &= pos_slot_i->ne[1] == ns;
+
+ res &= decode == (pos_slot_f != nullptr);
+ res &= decode == (cell_blk == nullptr);
+
+ if (pos_slot_f) {
+ res &= pos_slot_f->ne[0] == n_ps;
+ res &= pos_slot_f->ne[1] == ns;
+ }
+
+ if (cell_blk) {
+ res &= cell_blk->ne[0] == (int64_t) mctx_new->get_base()->get_n_kv();
+ res &= cell_blk->ne[1] == ns;
+ }
+
+ return res;
+ }
+
+ ggml_tensor * bias = nullptr; // F32 [nblk, n_tokens] local-force bias (position blocks)
+ ggml_tensor * pos_mask = nullptr; // F32 [n_ps, n_tokens] 0/-inf visibility, by position
+ ggml_tensor * pos_slot_i = nullptr; // I32 [n_ps, ns] pos -> cell (get_rows index)
+ ggml_tensor * pos_slot_f = nullptr; // F32 [n_ps, ns] pos -> cell (gatherable values, decode)
+ ggml_tensor * cell_blk = nullptr; // I32 [n_kv, ns] cell -> position block (batch)
+
+ const llama_kv_cache_msa_context * mctx;
+
+ int blk;
+ int local;
+};
+
+// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])
+ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa(
+ ggml_tensor * q_cur, // [D, HQ, T]
+ ggml_tensor * k, // [D, n_keys, 1, C]
+ ggml_tensor * v, // [D, n_keys, 1, C]
+ ggml_tensor * mask, // [n_keys, R, 1, C] f16, contiguous
+ int64_t Gp, float kq_scale, int il) const {
+
+ const int64_t D = q_cur->ne[0];
+ const int64_t HQ = q_cur->ne[1];
+ const int64_t T = q_cur->ne[2];
+ const int64_t C = k->ne[3];
+ const int64_t R = HQ*T/(Gp*C);
+ GGML_ASSERT(Gp*C*R == HQ*T);
+ GGML_ASSERT(mask->type == GGML_TYPE_F16);
+
+ // [D, HQ, T] -> [D, Gp, C, R] -> [D, R, Gp, C]
+ // batch (C=HKV, R=T): channel = group
+ // decode (C=HKV*ns, R=1): channel = (group, stream), group innermost
+ ggml_tensor * q = ggml_reshape_4d(ctx0, q_cur, D, Gp, C, R);
+ q = ggml_permute(ctx0, q, 0, 2, 3, 1);
+
+ ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale,
+ hparams.f_max_alibi_bias, 0.0f);
+ ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32);
+ cb(o, "msa_fattn", il);
+
+ // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T]
+ o = ggml_permute(ctx0, o, 0, 1, 3, 2);
+ if (!ggml_is_contiguous(o)) {
+ o = ggml_cont(ctx0, o); // no-op layout at decode (R == 1), copy at batch
+ }
+ return ggml_reshape_2d(ctx0, o, D*HQ, T);
+}
+
+llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+ const auto & mm = static_cast<const llama_model_minimax_m3 &>(model);
+
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+ // partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ // ==========================================
+ // TODO: avoid such kind of complexity in the model graphs
+
+ // MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that
+ // llama.cpp only provides when flash attention is enabled. Block selection is anchored
+ // to absolute KV cache slots, which equal positions only for append-only per-stream
+ // caches either a single sequence, or multiple sequences with kv_unified == false (each
+ // stream then has its own slot space). A unified cache with multiple sequences
+ // interleaves slots and would silently break block anchoring so it falls back to dense.
+ const bool fa_on = cparams.flash_attn;
+ const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified;
+ const bool msa_enabled = fa_on && streams_ok;
+
+ auto * inp_attn = build_attn_inp_kv_msa(msa_enabled);
+
+ static bool warned_no_fa = false;
+ if (!fa_on && !warned_no_fa) {
+ LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention "
+ "(output may be degraded). Enable flash attention for MSA.\n", __func__);
+ warned_no_fa = true;
+ }
+ static bool warned_unified = false;
+ if (fa_on && !streams_ok && !warned_unified) {
+ LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams "
+ "-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);
+ warned_unified = true;
+ }
+ // ==========================================
+
+ // hoisted per-graph MSA state (shared by every sparse layer)
+ llm_graph_input_msa * msa = nullptr;
+ ggml_tensor * msa_kqm = nullptr;
+ ggml_tensor * msa_mf = nullptr; // F32 copy of the FA mask for the final mask add
+ int64_t n_kv = 0, n_ps = 0, nblk = 0, ns = 1, n_tps = 0;
+ bool msa_decode = false; // gather (1 token per stream) vs mask
+ const int blk = mm.msa_p.blk;
+ const int64_t Hd = hparams.indexer_n_head; // one indexer head per GQA group
+
+ if (msa_enabled) {
+ const auto * mctx_msa = static_cast<const llama_kv_cache_msa_context *>(mctx);
+
+ msa_kqm = inp_attn->get_kq_mask();
+ n_kv = msa_kqm->ne[0];
+ n_tps = msa_kqm->ne[1]; // tokens per stream
+ ns = msa_kqm->ne[3]; // streams in this ubatch
+ GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask");
+ GGML_ASSERT(n_tps*ns == n_tokens);
+
+ // the position axis covers every position currently in the cache and is padded to whole blocks
+ n_ps = GGML_PAD((int64_t) mctx_msa->get_n_pos(), blk);
+ nblk = n_ps / blk;
+ msa_decode = n_tps == 1;
+
+ auto inp = std::make_unique<llm_graph_input_msa>(mctx_msa, blk, mm.msa_p.local);
+
+ inp->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens
+ ggml_set_input(inp->bias);
+
+ inp->pos_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, n_tokens);
+ ggml_set_input(inp->pos_mask);
+
+ inp->pos_slot_i = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_ps, ns);
+ ggml_set_input(inp->pos_slot_i);
+
+ if (msa_decode) {
+ inp->pos_slot_f = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, ns);
+ ggml_set_input(inp->pos_slot_f);
+ } else {
+ inp->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns);
+ ggml_set_input(inp->cell_blk);
+
+ msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);
+ }
+
+ msa = (llm_graph_input_msa *) res->add_input(std::move(inp));
+ }
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * inpSA = inpL;
+
+ // self-attention
+ {
+ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+ n_embd_head, n_head, n_head_kv, il);
+
+ // per-head QK RMSNorm (weights already include Gemma's +1)
+ Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
+ cb(Qcur, "Qcur_normed", il);
+ Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
+ cb(Kcur, "Kcur_normed", il);
+
+ // partial rotary: only the first n_rot dims are rotated
+ 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);
+
+ const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead;
+
+ if (!is_sparse) {
+ cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,
+ 1.0f/sqrtf(float(n_embd_head)), il);
+ } else {
+ const int64_t n_idx_dim = hparams.indexer_head_size; // 128
+
+ // Index Branch, project, norm, partial RoPE, cache
+ ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);
+ ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);
+ iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, Hd, n_tokens);
+ ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens);
+ iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il); // +1 baked
+ ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il);
+ iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
+ freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
+ ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
+ freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
+
+ const auto * mctx_msa_l = static_cast<const llama_kv_cache_msa_context *>(mctx);
+ const auto * mctx_cur = mctx_msa_l->get_base();
+ const auto * mctx_idx = mctx_msa_l->get_idx();
+ ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs_idx(), il));
+ ggml_tensor * ik_kv = mctx_idx->get_k(ctx0, il);
+
+ if (inp_attn->self_k_rot) {
+ Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
+ }
+ if (inp_attn->self_v_rot) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+ }
+
+ // Main branch: store K/V, take cache views
+ ggml_build_forward_expand(gf, Qcur);
+ ggml_build_forward_expand(gf, Kcur);
+ ggml_build_forward_expand(gf, Vcur);
+ ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
+ ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
+ ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+ ggml_tensor * v = mctx_cur->get_v(ctx0, il);
+ GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA assumes v_trans=false (FA on)");
+
+ const int64_t D = k->ne[0];
+ const int64_t HKV = k->ne[1];
+ const int64_t Gp = n_head/HKV;
+ GGML_ASSERT(HKV == Hd && "MSA: one indexer head per GQA group");
+ GGML_ASSERT(k->ne[3] == ns);
+ const int K = mm.msa_p.topk_blocks < (int) nblk ? mm.msa_p.topk_blocks : (int) nblk;
+
+ const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
+
+ if (msa_decode) {
+ // decode: batched over streams top-k + gather, one grouped FA
+ // gather the indexer keys through the pos -> cell map
+ ggml_tensor * ik3 = ggml_view_3d(ctx0, ik_kv, n_idx_dim, n_kv, ns,
+ ik_kv->nb[2], ik_kv->nb[3], 0);
+ ggml_tensor * ikp = ggml_get_rows(ctx0, ik3, msa->pos_slot_i); // [n_idx_dim, n_ps, ns]
+ ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);
+ ggml_tensor * sc = ggml_mul_mat(ctx0,
+ ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4);
+ ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
+ // unmapped positions come out -inf, so they can never rank into the top-k
+ sc = ggml_add_inplace(ctx0, sc,
+ ggml_reshape_4d(ctx0, msa->pos_mask, n_ps, 1, 1, ns));
+ ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
+ cb(bs, "msa_bs", il);
+
+ ggml_tensor * bsf = ggml_add(ctx0, bs,
+ ggml_reshape_4d(ctx0, msa->bias, nblk, 1, 1, ns));
+ ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // position blocks
+
+ // pos idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (positions - mask gather)
+ // cell idx: cs[t,k,h,s] = pos_slot[tj] (pos -> cell translation)
+ // row idx: tr[t,k,h,s] = cs*HKV + h (per-stream K/V gather)
+ ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk);
+ a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns);
+ ggml_tensor * tj = ggml_add(ctx0,
+ ggml_repeat_4d(ctx0, a, blk, K, Hd, ns),
+ ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1));
+
+ ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
+
+ ggml_tensor * cs = ggml_get_rows(ctx0,
+ ggml_reshape_3d(ctx0, msa->pos_slot_f, 1, n_ps, ns), tokj); // [1, blk*K*Hd, ns]
+ cs = ggml_reshape_4d(ctx0, cs, blk, K, Hd, ns);
+
+ ggml_tensor * tr = ggml_add(ctx0,
+ ggml_scale(ctx0, cs, (float) HKV),
+ ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));
+
+ ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
+
+ ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0);
+ ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0);
+ ggml_tensor * mp = ggml_reshape_3d(ctx0, msa->pos_mask, 1, n_ps, ns);
+
+ ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr);
+ ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr);
+ ggml_tensor * mg = ggml_get_rows(ctx0, mp, tokj);
+
+ // fold (group, stream) onto the FA channel dim
+ const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type;
+ const ggml_type vt = ggml_is_quantized(v->type) ? GGML_TYPE_F16 : v->type;
+ ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, D, (int64_t) blk*K, 1, Hd*ns);
+ ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, D, (int64_t) blk*K, 1, Hd*ns);
+ if (kfa->type != kt) { kfa = ggml_cast(ctx0, kfa, kt); }
+ if (vfa->type != vt) { vfa = ggml_cast(ctx0, vfa, vt); }
+ // the FA mask must be F16
+ ggml_tensor * mfa = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk*K, 1, 1, Hd*ns), GGML_TYPE_F16);
+
+ cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il);
+ } else {
+ // batch: per-stream loop
+ std::vector<ggml_tensor *> outs(ns);
+ for (int64_t st = 0; st < ns; ++st) {
+ ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps,
+ iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]);
+ ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv,
+ ik_kv->nb[2], st*ik_kv->nb[3]);
+ ggml_tensor * psl_s = ggml_view_1d(ctx0, msa->pos_slot_i, n_ps,
+ st*msa->pos_slot_i->nb[1]);
+ ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, n_tps,
+ msa->pos_mask->nb[1], msa->pos_mask->nb[1], st*n_tps*msa->pos_mask->nb[1]);
+ ggml_tensor * cb_s = ggml_view_1d(ctx0, msa->cell_blk, n_kv,
+ st*msa->cell_blk->nb[1]);
+ ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, n_tps, 1,
+ msa_mf->nb[1], msa_mf->nb[3], st*msa_mf->nb[3]);
+ ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, n_tps,
+ msa->bias->nb[1], msa->bias->nb[1], st*n_tps*msa->bias->nb[1]);
+ ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,
+ Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);
+ ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,
+ k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]);
+ ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1,
+ v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]);
+
+ // block scores: the indexer keys are gathered through the pos -> cell map first
+ // scores are unscaled, only the top-k ordering matters
+ ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s); // [n_idx_dim, n_ps]
+ ggml_tensor * sc = ggml_mul_mat(ctx0, ikp,
+ ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));
+ // indexer scores run in F32
+ ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
+ sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps);
+ // unmapped positions (holes, padding, empty cells) come out -inf
+ sc = ggml_add_inplace(ctx0, sc, pm_s);
+ ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
+ cb(bs, "msa_bs", il);
+
+ // bias the scores so locally-forced blocks always rank first
+ ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s); // [nblk, Hd, n_tps]
+ cb(bsf, "msa_bsf", il);
+
+ ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // [K, Hd, n_tps] i32
+
+ ggml_tensor * ninf = ggml_cast(ctx0,
+ ggml_scale_bias(ctx0, bias_s, 0.0f, -1e30f),
+ GGML_TYPE_F16); // [nblk, 1, n_tps]
+ ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, n_tps, 1);
+ ggml_tensor * zero = ggml_scale(ctx0,
+ ggml_cast(ctx0, idx, GGML_TYPE_F32), 0.0f);
+ ggml_tensor * bm = ggml_set_rows(ctx0,
+ ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*n_tps),
+ ggml_reshape_3d(ctx0, zero, 1, K, Hd*n_tps),
+ ggml_reshape_2d(ctx0, idx, K, Hd*n_tps));
+ bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, n_tps);
+ bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd]
+ cb(bm, "msa_block_mask", il);
+
+ // expand block -> cell granularity through the cell -> position block
+ // map, then combine with the causal mask. empty cells are masked by the causal mask.
+ ggml_tensor * bm2 = ggml_cont(ctx0, ggml_transpose(ctx0,
+ ggml_reshape_2d(ctx0, bm, nblk, n_tps*Hd))); // [n_tps*Hd, nblk]
+ ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s); // [n_tps*Hd, n_kv] F32
+ ggml_tensor * bmx = ggml_cont(ctx0, ggml_transpose(ctx0, bmc));
+ bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd);
+ ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, mf_s);
+ mask4 = ggml_cast(ctx0,
+ ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd), GGML_TYPE_F16);
+ cb(mask4, "msa_mask4", il);
+
+ // cache views with groups on ne[3];
+ ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2);
+ ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2);
+
+ outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il);
+ }
+ cur = outs[0];
+ for (int64_t st = 1; st < ns; ++st) {
+ cur = ggml_concat(ctx0, cur, outs[st], 1);
+ }
+ }
+ if (inp_attn->self_v_rot) {
+ cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);
+ }
+ cb(cur, "kqv_out", il);
+ if (model.layers[il].wo) {
+ cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
+ }
+ }
+ }
+
+ 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) {
+ // leading dense FFN (swigluoai)
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, NULL, NULL,
+ model.layers[il].ffn_gate, NULL, NULL,
+ model.layers[il].ffn_down, NULL, NULL,
+ NULL,
+ LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ } else {
+ // routed experts (swigluoai MoE)
+ 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_SWIGLU_OAI_MOE, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ // shared expert (swigluoai)
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ model.layers[il].ffn_up_shexp, NULL, NULL,
+ model.layers[il].ffn_gate_shexp, NULL, NULL,
+ model.layers[il].ffn_down_shexp, NULL, NULL,
+ NULL,
+ LLM_FFN_SWIGLU_OAI_MOE, 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 = build_lora_mm(model.output, cur, model.output_s);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
+struct llama_model_nanbeige : public llama_model_base {
+ llama_model_nanbeige(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;
+
+ int n_loops = 1;
+ int n_layer_phys = 0;
+ bool skip_loop_final_norm = false;
+
+ 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) {}
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;
};
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;
};
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
+ graph(const llm_graph_params & params) : llm_graph_context(params) {}
graph(const llama_model & model, const llm_graph_params & params);
ggml_tensor * build_hc_pre(
ggml_tensor * inp_pos,
int il) const;
+ ggml_tensor * build_attention(
+ const llama_model & model,
+ llm_graph_input_attn_k_iswa * inp_mtp,
+ ggml_tensor * cur,
+ ggml_tensor * inp_pos,
+ int il) const;
+
+ ggml_tensor * build_attention_impl(
+ const llama_model & model,
+ llm_graph_input_dsv4 * inp_dsv4,
+ llm_graph_input_attn_k_iswa * inp_mtp,
+ ggml_tensor * cur,
+ ggml_tensor * inp_pos,
+ int il) const;
+
ggml_tensor * build_hca_compressed_kv_from_state(
ggml_tensor * kv_state,
ggml_tensor * score_state,
float kq_scale,
int il) const;
- ggml_tensor * build_hc_weighted_sum(
+ ggml_tensor * build_hc_pre(
ggml_tensor * x,
- ggml_tensor * weights) const;
+ ggml_tensor * weights,
+ int il) const;
ggml_tensor * build_hc_sinkhorn(
ggml_tensor * comb,
int il) const;
};
+ struct graph_mtp : public graph {
+ 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;
};
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
- using graph = llama_model_deepseek2::graph;
+ struct graph : public llm_graph_context {
+ 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;
};
ggml_tensor * build_inp_embd_enc() const;
};
+ struct graph_dsv4 : public llama_model_deepseek4::graph {
+ graph_dsv4(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_laguna : public llama_model_base {
+ llama_model_laguna(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_ernie4_5 : public llama_model_base {
llama_model_ernie4_5(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
+struct llama_model_hy_v3 : public llama_model_base {
+ llama_model_hy_v3(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);
+ };
+
+ 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;
+};
+
struct llama_model_hunyuan_vl : public llama_model_base {
llama_model_hunyuan_vl(const struct llama_model_params & params) : llama_model_base(params) {}
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
+struct msa_params {
+ int blk;
+ int topk_blocks;
+ int local;
+};
+
+struct llama_model_minimax_m3 : public llama_model_base {
+ llama_model_minimax_m3(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;
+ msa_params msa_p;
+ struct graph : public llm_graph_context {
+ graph(const llama_model & model, const llm_graph_params & params);
+
+ ggml_tensor * build_attn_msa_fa(
+ ggml_tensor * q_cur, // [D, HQ, S] f32
+ ggml_tensor * k, // [D, n_keys, 1, C] C = HKV or HKV*n_stream
+ ggml_tensor * v, // [D, n_keys, 1, C]
+ ggml_tensor * mask, // [n_keys, R, 1, C] f16, R = HQ*T/(Gp*C)
+ int64_t Gp, float kq_scale, int il) const;
+ };
+ std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
struct llama_model_cogvlm : public llama_model_base {
llama_model_cogvlm(const struct llama_model_params & params) : llama_model_base(params) {}
const llama_model & model;
};
+ 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;
};
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;
};
--- /dev/null
+#include "models.h"
+
+void llama_model_nanbeige::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+
+ uint32_t n_loops_u = 1;
+ ml.get_key(LLM_KV_NUM_LOOPS, n_loops_u, false);
+ GGML_ASSERT(n_loops_u >= 1);
+
+ skip_loop_final_norm = false;
+ ml.get_key(LLM_KV_SKIP_LOOP_FINAL_NORM, skip_loop_final_norm, false);
+
+ n_layer_phys = (int) hparams.n_layer();
+
+ // Bound-check before casting: signed int mul can overflow and bypass the guard.
+ GGML_ASSERT((size_t) n_layer_phys * (size_t) n_loops_u <= (size_t) LLAMA_MAX_LAYERS);
+ n_loops = (int) n_loops_u;
+
+ // Expand logical layer count before load_tensors() allocates layers / KV.
+ if (n_loops > 1) {
+ for (int j = 1; j < n_loops; ++j) {
+ for (int i = 0; i < n_layer_phys; ++i) {
+ const int dst = i + j * n_layer_phys;
+ hparams.n_head_arr[dst] = hparams.n_head_arr[i];
+ hparams.n_head_kv_arr[dst] = hparams.n_head_kv_arr[i];
+ hparams.n_ff_arr[dst] = hparams.n_ff_arr[i];
+ hparams.is_swa_impl[dst] = hparams.is_swa_impl[i];
+ hparams.is_recr_impl[dst] = hparams.is_recr_impl[i];
+ }
+ }
+ hparams.n_layer_all = (uint32_t) ((size_t) n_layer_phys * (size_t) n_loops);
+ }
+
+ type = LLM_TYPE_UNKNOWN;
+}
+
+void llama_model_nanbeige::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+ if (output == NULL) {
+ output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+ }
+
+ const int n_phys = n_layer_phys > 0 ? n_layer_phys : n_layer;
+ for (int i = 0; i < n_phys; ++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_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);
+
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2},
+ TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
+ }
+
+ // Share physical weights across loops; each slot still has its own KV index.
+ if (n_loops > 1) {
+ for (int j = 1; j < n_loops; ++j) {
+ for (int i = 0; i < n_phys; ++i) {
+ layers[i + j * n_phys] = layers[i];
+ }
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_nanbeige::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context(params) {
+ const auto & nb = static_cast<const llama_model_nanbeige &>(model);
+
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+ const int n_phys = nb.n_layer_phys > 0 ? nb.n_layer_phys : (int) n_layer;
+ const int n_loops = nb.n_loops > 0 ? nb.n_loops : 1;
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ inpL = build_inp_embd(model.tok_embd);
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ auto * inp_attn = build_attn_inp_kv();
+
+ const float kq_scale = hparams.f_attention_scale == 0.0f
+ ? 1.0f / sqrtf(float(n_embd_head))
+ : hparams.f_attention_scale;
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * inpSA = inpL;
+
+ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ {
+ ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+
+ auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+ n_embd_head, n_head, n_head_kv, il);
+
+ Qcur = ggml_rope_ext(
+ ctx0, Qcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ Kcur = ggml_rope_ext(
+ ctx0, Kcur, inp_pos, rope_factors,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+
+ cb(Qcur, "Qcur", il);
+ cb(Kcur, "Kcur", il);
+ cb(Vcur, "Vcur", il);
+
+ 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, kq_scale, il);
+ cb(cur, "attn_out", 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);
+
+ cur = build_ffn(cur,
+ model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s,
+ model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s,
+ model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+ cb(cur, "ffn_out", il);
+
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+
+ if (n_loops > 1 &&
+ ((il + 1) % n_phys) == 0 &&
+ (il + 1) < n_layer &&
+ !nb.skip_loop_final_norm) {
+ cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "loop_norm", il);
+ inpL = cur;
+ }
+ }
+
+ cur = inpL;
+
+ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ cur = build_lora_mm(model.output, cur, model.output_s);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ ggml_build_forward_expand(gf, cur);
+}
cb(cur, "attn_out", il);
}
- if (il == n_layer - 1) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
// skip computing output for unused tokens
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = inpL;
+ 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, NULL,
LLM_NORM_RMS, -1);
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;
+ int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
auto & layer = layers[il];
// MTP block looks like a full-attention Qwen3.5 decoder block.
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0);
- layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
- create_tensor_qkv(layer, il, 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", il), { n_embd_head_k * n_head, n_embd }, 0);
- layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
- layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
+ create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
- layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, 0);
- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, 0);
- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, 0);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, mtp_flags);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, mtp_flags);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, mtp_flags);
// NextN-specific tensors that define the MTP block.
- layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, 0);
- layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, 0);
- layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, 0);
- layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
- layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
- layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);
};
for (int i = 0; i < n_layer; ++i) {
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;
+ int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;
// MTP block looks like a full-attention Qwen3.5 decoder block with MoE FFN.
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0);
- layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
- create_tensor_qkv(layer, il, 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", il), { n_embd_head_k * n_head, n_embd }, 0);
- layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
- layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
+ create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
+ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
// Routed experts
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
- create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, mtp_flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, mtp_flags);
+ create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, mtp_flags);
// Shared experts
- layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);
- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0);
+ layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, mtp_flags);
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, mtp_flags);
// NextN-specific tensors that define the MTP block.
- layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, 0);
- layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, 0);
- layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, 0);
- layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
- layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
- layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);
};
for (int i = 0; i < n_layer; ++i) {
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);
- // Mark recurrent layers (linear attention layers)
+ // NextN/MTP: 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_all");
+
+ // 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);
}
}
-void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
+void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
if (n_expert == 0) {
throw std::runtime_error(arch_name() + " model cannot have zero experts");
}
+ 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;
+ int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
// output
const int64_t qkvz_dim = key_dim * 2 + value_dim * 2;
const int64_t ba_dim = n_v_heads * 2;
- for (int i = 0; i < n_layer; ++i) {
- auto & layer = layers[i];
- const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(i);
+ auto load_block_trunk = [&](int il, int flags) {
+ auto & layer = layers[il];
+ const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(il);
- 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);
+ 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_recr(i)) {
+ if (!hparams.is_recr(il)) {
// 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);
-
+ 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);
// Q/K normalization for attention layers
- 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.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags);
+ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags);
} else {
// Linear attention (gated delta net) specific tensors
// Create tensors with calculated dimensions
// note: ssm_in is used by legacy GGUF
- layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED);
- layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
- layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
- layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
- layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
- layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
- layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", i), { n_embd, ba_dim }, 0);
- layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
- layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
+ layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", il), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED | flags);
+ layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | flags);
+ layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | flags);
+ layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
+ layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);
+ layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);
+ layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", il), { n_embd, ba_dim }, flags);
+ layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);
+ layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);
}
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, 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);
- create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, 0);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags);
+ create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags);
// Shared experts
- layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0);
- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0);
+ layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags);
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags);
+ };
+
+ auto load_block_mtp = [&](int il) {
+ // MTP head is identical to the trunk block (full attention + FFN)
+ load_block_trunk(il, mtp_flags);
+
+ auto & layer = layers[il];
+
+ // NextN-specific tensors that define the MTP block.
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags | TENSOR_NOT_REQUIRED);
+ };
+
+ for (int i = 0; i < n_layer; i++) {
+ load_block_trunk(i, trunk_flags);
+ }
+ for (int i = n_layer; i < n_layer_all; i++) {
+ load_block_mtp(i);
}
}
std::unique_ptr<llm_graph_context> llama_model_qwen3next::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);
}
ggml_tensor * inp_pos = build_inp_pos();
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) {
res->t_layer_inp[il] = inpL;
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il);
}
- if (il == n_layer - 1 && 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;
- // Final norm
+ // post-norm hidden state is input to both the LM head and the MTP head
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_nextn_masked && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
cb(cur, "result_norm", -1);
res->t_embd = cur;
ggml_build_forward_expand(gf, cur);
}
-// utility to get one slice from the third dimension
-// input dim: [x, y, c, b]
-// output dim: [x, y, 1, b]
-static ggml_tensor * get_slice_2d(ggml_context * ctx0, ggml_tensor * t, int64_t c) {
- return ggml_view_4d(ctx0, t, t->ne[0], t->ne[1], 1, t->ne[3],
- t->nb[1], t->nb[2], t->nb[3], t->nb[2] * c);
-}
-
ggml_tensor * llama_model_qwen3next::graph::build_norm_gated(
ggml_tensor * input,
ggml_tensor * weights,
// Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention
// Qwen3Next uses a single Q projection that outputs query + gate
- ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur);
+ ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
cb(Qcur_full, "Qcur_full", il);
Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1);
Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full));
cb(gate, "gate", il);
- ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
+ ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);
cb(Kcur, "Kcur", il);
- ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
+ ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s);
cb(Vcur, "Vcur", il);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "gate_sigmoid", il);
- gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
-
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
LLM_FFN_SILU, true,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
- nullptr, model.layers[il].ffn_gate_up_exps);
+ 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);
// Add shared experts if present - following Qwen3Next reference implementation
if (model.layers[il].ffn_up_shexp != nullptr) {
ggml_tensor * ffn_shexp =
build_ffn(cur,
- model.layers[il].ffn_up_shexp, NULL, NULL,
- model.layers[il].ffn_gate_shexp, NULL, NULL,
- model.layers[il].ffn_down_shexp, NULL, NULL,
+ 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);
}
return cur;
}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3-Next
+llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN3NEXT MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN3NEXT 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 = 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");
+ GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
+
+ // 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_inp(), n_tokens);
+ ggml_set_input(inp->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;
+
+ 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_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s);
+ cb(Qcur_full, "mtp_Qcur_full", il);
+
+ ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full,
+ n_embd_head, n_head, n_tokens,
+ ggml_element_size(Qcur_full) * n_embd_head * 2,
+ ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
+ 0);
+ Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Qcur, "mtp_Qcur_normed", il);
+
+ ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
+ Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
+ cb(Kcur, "mtp_Kcur_normed", il);
+
+ ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s);
+ Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
+
+ 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, "mtp_Qcur", il);
+ cb(Kcur, "mtp_Kcur", il);
+ cb(Vcur, "mtp_Vcur", il);
+
+ const float kq_scale = hparams.f_attention_scale == 0.0f
+ ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
+
+ cur = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(cur, "mtp_attn_pregate", il);
+
+ ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full,
+ n_embd_head, n_head, n_tokens,
+ ggml_element_size(Qcur_full) * n_embd_head * 2,
+ ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
+ ggml_element_size(Qcur_full) * n_embd_head);
+
+ // TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont
+ gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
+ cb(gate, "mtp_gate", il);
+
+ cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
+ cur = build_lora_mm(layer.wo, cur, layer.wo_s);
+ cb(cur, "mtp_attn_out", il);
+
+ if (inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ cur = ggml_add(ctx0, cur, inpSA);
+ cb(cur, "mtp_attn_residual", il);
+
+ ggml_tensor * ffn_residual = cur;
+ cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_post_norm", il);
+
+ // MoE FFN — routed experts plus gated shared expert (mirrors the trunk).
+ ggml_tensor * moe_out =
+ build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.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, layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ if (layer.ffn_up_shexp != nullptr) {
+ ggml_tensor * ffn_shexp =
+ build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+ ggml_tensor * shared_gate = build_lora_mm(layer.ffn_gate_inp_shexp, cur);
+ shared_gate = ggml_sigmoid(ctx0, shared_gate);
+ cb(shared_gate, "mtp_shared_expert_gate_sigmoid", il);
+
+ ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);
+ cb(ffn_shexp, "mtp_ffn_shexp_gated", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ } else {
+ cur = moe_out;
+ }
+ cb(cur, "mtp_ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_residual);
+ cb(cur, "mtp_post_ffn", il);
+
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "QWEN3NEXT 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;
+
+ ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+ ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w && "QWEN3NEXT MTP: missing LM head (nextn.shared_head_head or model.output)");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
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;
+ int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+ if (!ml.load_mtp) {
+ mtp_flags |= TENSOR_SKIP;
+ }
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
overview / pkg / ggml / sources / whisper.cpp / commitdiff