diff --git a/examples/talk-llama/CMakeLists.txt b/examples/talk-llama/CMakeLists.txt index 59643c390..f6901f920 100644 --- a/examples/talk-llama/CMakeLists.txt +++ b/examples/talk-llama/CMakeLists.txt @@ -2,6 +2,8 @@ if (WHISPER_SDL2) 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) @@ -19,13 +21,9 @@ if (WHISPER_SDL2) 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 diff --git a/examples/talk-llama/llama-arch.cpp b/examples/talk-llama/llama-arch.cpp index b890e66fc..ea0ddd114 100644 --- a/examples/talk-llama/llama-arch.cpp +++ b/examples/talk-llama/llama-arch.cpp @@ -108,11 +108,13 @@ static const std::map LLM_ARCH_NAMES = { { 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" }, @@ -125,6 +127,7 @@ static const std::map LLM_ARCH_NAMES = { { 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" }, @@ -140,6 +143,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, + { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -218,6 +222,8 @@ static const std::map LLM_KV_NAMES = { { 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" }, @@ -251,6 +257,9 @@ static const std::map LLM_KV_NAMES = { { 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" }, @@ -308,6 +317,7 @@ static const std::map LLM_KV_NAMES = { { 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 @@ -594,6 +604,9 @@ static const std::map LLM_TENSOR_NAMES = { { 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" }, @@ -603,6 +616,9 @@ static const std::map LLM_TENSOR_NAMES = { { 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: @@ -664,7 +680,7 @@ static const std::map LLM_TENSOR_INFOS = { {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}}, @@ -829,9 +845,12 @@ static const std::map LLM_TENSOR_INFOS = { {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}}, @@ -854,6 +873,10 @@ static const std::map LLM_TENSOR_INFOS = { // 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) {} @@ -945,6 +968,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_DEEPSEEK4: return true; default: return false; @@ -967,6 +991,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { switch (arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_DEEPSEEK4: return true; default: return false; @@ -998,6 +1023,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { 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; diff --git a/examples/talk-llama/llama-arch.h b/examples/talk-llama/llama-arch.h index a4f5091e7..cbc97085e 100644 --- a/examples/talk-llama/llama-arch.h +++ b/examples/talk-llama/llama-arch.h @@ -113,11 +113,13 @@ enum llm_arch { 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, @@ -144,7 +146,9 @@ enum llm_arch { LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, + LLM_ARCH_MINIMAX_M3, LLM_ARCH_DFLASH, + LLM_ARCH_NANBEIGE, LLM_ARCH_UNKNOWN, }; @@ -223,6 +227,8 @@ enum llm_kv { 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, @@ -256,6 +262,9 @@ enum llm_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, @@ -354,6 +363,7 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, LLM_KV_NORM_BEFORE_RESIDUAL, + LLM_KV_NORM_BEFORE_FC, LLM_KV_SHORTCONV_L_CACHE, @@ -594,6 +604,9 @@ enum llm_tensor { 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, @@ -611,6 +624,9 @@ enum llm_tensor { 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, }; diff --git a/examples/talk-llama/llama-context.cpp b/examples/talk-llama/llama-context.cpp index 5edfc85ab..19cca7df1 100644 --- a/examples/talk-llama/llama-context.cpp +++ b/examples/talk-llama/llama-context.cpp @@ -55,6 +55,30 @@ static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { /*.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) : @@ -96,8 +120,9 @@ llama_context::llama_context( 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; @@ -226,6 +251,14 @@ llama_context::llama_context( 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; @@ -442,6 +475,9 @@ llama_context::llama_context( } 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]; @@ -522,6 +558,20 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 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() { @@ -1115,7 +1165,7 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) { 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; @@ -1371,13 +1421,17 @@ int llama_context::encode(const llama_batch & batch_inp) { // 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; @@ -1663,7 +1717,8 @@ int llama_context::decode(const llama_batch & batch_inp) { 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; @@ -1716,13 +1771,18 @@ int llama_context::decode(const llama_batch & batch_inp) { 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(); @@ -2217,8 +2277,9 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to } 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; @@ -2231,14 +2292,14 @@ void llama_context::output_reorder() { } 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]); } } @@ -2292,7 +2353,10 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { 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(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); @@ -2426,11 +2490,12 @@ llm_graph_cb llama_context::graph_get_cb() const { 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) { @@ -3492,6 +3557,22 @@ llama_context * llama_init_from_model( } } + 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) { @@ -3514,11 +3595,6 @@ llama_context * llama_init_from_model( } } - 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 diff --git a/examples/talk-llama/llama-cparams.h b/examples/talk-llama/llama-cparams.h index 546ae1e2c..5018170ed 100644 --- a/examples/talk-llama/llama-cparams.h +++ b/examples/talk-llama/llama-cparams.h @@ -41,6 +41,12 @@ struct llama_cparams { 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; diff --git a/examples/talk-llama/llama-grammar.cpp b/examples/talk-llama/llama-grammar.cpp index badcbfd0f..363644464 100644 --- a/examples/talk-llama/llama-grammar.cpp +++ b/examples/talk-llama/llama-grammar.cpp @@ -1139,6 +1139,18 @@ struct llama_grammar * llama_grammar_init_impl( 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 rules_visited(n_rules); std::vector rules_in_progress(n_rules); diff --git a/examples/talk-llama/llama-graph.cpp b/examples/talk-llama/llama-graph.cpp index a8fd11ebc..2be3b75fb 100644 --- a/examples/talk-llama/llama-graph.cpp +++ b/examples/talk-llama/llama-graph.cpp @@ -8,6 +8,7 @@ #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" @@ -518,6 +519,40 @@ bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) { 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(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); @@ -619,6 +654,63 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { 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(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 & src) { if (!dst || !dst->buffer) { return; @@ -754,6 +846,10 @@ static void dsv4_set_comp_inputs( 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); @@ -798,6 +894,10 @@ static bool dsv4_can_reuse_comp_input( 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()); @@ -832,6 +932,10 @@ static void dsv4_build_comp_inputs( 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"); @@ -842,7 +946,7 @@ static void dsv4_build_comp_inputs( 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()); } @@ -1195,7 +1299,7 @@ void llm_graph_result::reset() { 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(); @@ -1650,7 +1754,7 @@ ggml_tensor * llm_graph_context::build_ffn( 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); @@ -1709,6 +1813,17 @@ ggml_tensor * llm_graph_context::build_ffn( 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); @@ -2034,7 +2149,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( 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); @@ -2668,7 +2783,7 @@ ggml_tensor * llm_graph_context::build_attn( 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); @@ -2951,6 +3066,75 @@ ggml_tensor * llm_graph_context::build_attn( 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(cross); @@ -3025,9 +3209,9 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { { 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; @@ -3038,6 +3222,34 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { 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(mctx); + + auto inp = std::make_unique(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. @@ -3073,6 +3285,34 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const 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(mctx); + + auto inp = std::make_unique(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(mctx); const auto * raw_ctx = mctx_cur->get_raw(); diff --git a/examples/talk-llama/llama-graph.h b/examples/talk-llama/llama-graph.h index 97141ef93..32d8d395a 100644 --- a/examples/talk-llama/llama-graph.h +++ b/examples/talk-llama/llama-graph.h @@ -23,6 +23,7 @@ struct llama_memory_context_i; 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; @@ -42,6 +43,10 @@ enum llm_fused_op { 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 { @@ -421,6 +426,26 @@ public: 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( @@ -467,6 +492,45 @@ public: 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 { @@ -501,6 +565,10 @@ public: 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] @@ -1064,7 +1132,7 @@ struct llm_graph_context { 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] @@ -1122,6 +1190,8 @@ struct llm_graph_context { 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, @@ -1156,6 +1226,24 @@ struct llm_graph_context { 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( diff --git a/examples/talk-llama/llama-hparams.cpp b/examples/talk-llama/llama-hparams.cpp index 9d0683d2f..846d4c69a 100644 --- a/examples/talk-llama/llama-hparams.cpp +++ b/examples/talk-llama/llama-hparams.cpp @@ -248,6 +248,14 @@ bool llama_hparams::is_mla() const { 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(); } diff --git a/examples/talk-llama/llama-hparams.h b/examples/talk-llama/llama-hparams.h index 8be5f28f3..6e8336c98 100644 --- a/examples/talk-llama/llama-hparams.h +++ b/examples/talk-llama/llama-hparams.h @@ -47,6 +47,7 @@ struct llama_hparams { 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; @@ -226,6 +227,13 @@ struct llama_hparams { 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 is_indexer_full_impl; // DeepSeek-V4 uint32_t dsv4_o_group_count = 0; @@ -302,6 +310,8 @@ struct llama_hparams { 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) diff --git a/examples/talk-llama/llama-kv-cache-dsa.cpp b/examples/talk-llama/llama-kv-cache-dsa.cpp index 241c50365..96cb045d2 100644 --- a/examples/talk-llama/llama-kv-cache-dsa.cpp +++ b/examples/talk-llama/llama-kv-cache-dsa.cpp @@ -23,7 +23,8 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( 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) { @@ -32,7 +33,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( kv_mla = std::make_unique( 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 @@ -49,7 +50,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( kv_lid = std::make_unique( 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) { diff --git a/examples/talk-llama/llama-kv-cache-dsa.h b/examples/talk-llama/llama-kv-cache-dsa.h index e2b330993..e74fc4d91 100644 --- a/examples/talk-llama/llama-kv-cache-dsa.h +++ b/examples/talk-llama/llama-kv-cache-dsa.h @@ -26,7 +26,8 @@ public: 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; diff --git a/examples/talk-llama/llama-kv-cache-dsv4.cpp b/examples/talk-llama/llama-kv-cache-dsv4.cpp index 9fccf347e..5caa05e8b 100644 --- a/examples/talk-llama/llama-kv-cache-dsv4.cpp +++ b/examples/talk-llama/llama-kv-cache-dsv4.cpp @@ -22,13 +22,32 @@ static constexpr uint32_t DSV4_STATE_MAGIC = 0x34565344; // DSV4 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(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(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; @@ -230,28 +249,49 @@ static void dsv4_state_dst_stream_range( 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 * 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; + } - io.write_tensor(tensor, offset, size); + if (stream_ids && stream_ids->size() != ns) { + throw std::runtime_error("DSV4 state tensor stream map size mismatch"); + } + + 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) { @@ -273,18 +313,28 @@ static void dsv4_state_read_tensor_streams( 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; @@ -296,14 +346,18 @@ static void dsv4_state_write_k_cache( 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); } } @@ -315,19 +369,26 @@ static void dsv4_state_read_k_cache( 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"); } @@ -346,7 +407,7 @@ static void dsv4_state_read_k_cache( 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); } } @@ -369,7 +430,9 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( 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 & 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); @@ -399,6 +462,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( std::vector overlap_cur_reads; std::map, int64_t> curr_token_idx_map; + std::map 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) { @@ -461,6 +525,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( 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; @@ -479,33 +544,57 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( } } - 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(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); } } } @@ -531,6 +620,63 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( 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 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; @@ -552,12 +698,14 @@ static std::vector dsv4_build_comp_plans 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 & rs_idx) { std::vector 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; @@ -644,7 +792,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( 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); @@ -662,10 +811,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( const uint64_t state_rows = (uint64_t) state_size*n_stream; const size_t n_persist = (size_t) std::min(ubatch.n_tokens, state_rows); + const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max(1, ubatch.n_seqs_unq) : 0; + const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max(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); @@ -691,12 +846,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( 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 { @@ -711,7 +868,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( 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, }; @@ -752,15 +909,24 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( 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 kv_stream; + std::vector 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) { @@ -777,20 +943,59 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( 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; } @@ -803,6 +1008,14 @@ uint32_t llama_dsv4_comp_state::get_n_stream() const { 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 llama_dsv4_comp_state::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { @@ -812,13 +1025,26 @@ std::map llama_dsv4_comp_state::memory_break 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 & 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 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(); @@ -832,8 +1058,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_ 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); } } @@ -874,33 +1100,45 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id 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_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_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 { - const int32_t ids = map_layer_ids.at(il); + ggml_tensor * state = get_score_all(ctx, il); + const size_t row_size = ggml_row_size(state->type, state->ne[0]); - 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(), 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 { @@ -929,13 +1167,16 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( 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)) { @@ -950,6 +1191,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( // 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); @@ -1016,25 +1262,25 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( csa_state = std::make_unique( 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( 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( 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( @@ -1136,7 +1382,13 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_full() { } llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) { - return std::make_unique(this, lctx, optimize); + return std::make_unique( + 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 { @@ -1147,7 +1399,7 @@ 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) { @@ -1156,43 +1408,86 @@ 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 { @@ -1251,14 +1546,24 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id 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) { @@ -1289,6 +1594,10 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, 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); @@ -1298,6 +1607,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, 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 { @@ -1328,13 +1643,63 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() 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 & llama_kv_cache_dsv4::get_rs_idx() const { + return rs_idx; +} + +void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector & 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); + } } // @@ -1595,20 +1960,26 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( 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(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(kv->get_csa())), - ctx_hca(std::make_unique(kv->get_hca())), - ctx_lid(std::make_unique(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( @@ -1620,10 +1991,14 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( std::vector 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( kv->get_raw(), std::move(sinfos_raw_base_write), @@ -1650,6 +2025,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( 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; @@ -1676,6 +2052,18 @@ bool llama_kv_cache_dsv4_context::apply() { 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; } @@ -1773,7 +2161,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ 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; } @@ -1787,7 +2175,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ 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; } @@ -1801,7 +2189,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ 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; } diff --git a/examples/talk-llama/llama-kv-cache-dsv4.h b/examples/talk-llama/llama-kv-cache-dsv4.h index 772b428cd..ce39867c0 100644 --- a/examples/talk-llama/llama-kv-cache-dsv4.h +++ b/examples/talk-llama/llama-kv-cache-dsv4.h @@ -10,6 +10,10 @@ 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, @@ -18,22 +22,29 @@ public: 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 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 & 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; @@ -44,12 +55,16 @@ private: ggml_tensor * kv; ggml_tensor * score; + + std::vector kv_stream; + std::vector 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> ctxs_bufs; @@ -67,6 +82,8 @@ private: // 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: @@ -82,6 +99,7 @@ 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); @@ -130,6 +148,10 @@ public: 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 & get_rs_idx() const; + void reset_rs_idx_for_ubatches(const std::vector & ubatches); + private: llama_hparams hparams_raw; llama_hparams hparams_csa; @@ -137,6 +159,9 @@ private: llama_hparams hparams_lid; const uint32_t n_seq_max; + const uint32_t n_rs_seq; + + std::vector rs_idx; std::unique_ptr kv_raw; std::unique_ptr kv_csa; @@ -146,7 +171,7 @@ private: std::unique_ptr hca_state; std::unique_ptr 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 @@ -243,6 +268,7 @@ private: 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 @@ -256,6 +282,17 @@ public: std::vector state_persist_src_idxs; std::vector 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 state_restore_src_idxs; + std::vector 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 state_snapshot_src_idxs; + std::vector 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 @@ -289,7 +326,10 @@ public: 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, @@ -349,9 +389,13 @@ private: const std::unique_ptr ctx_hca; const std::unique_ptr 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; diff --git a/examples/talk-llama/llama-kv-cache-msa.cpp b/examples/talk-llama/llama-kv-cache-msa.cpp new file mode 100644 index 000000000..55ef286ca --- /dev/null +++ b/examples/talk-llama/llama-kv-cache-msa.cpp @@ -0,0 +1,395 @@ +#include "llama-kv-cache-msa.h" + +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-model.h" + +#include +#include +#include + +// 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( + 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( + 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 llama_kv_cache_msa::memory_breakdown() const { + std::map 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 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( + this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches)); + } while (false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_kv_cache_msa::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(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 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(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(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 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; + } + } +} diff --git a/examples/talk-llama/llama-kv-cache-msa.h b/examples/talk-llama/llama-kv-cache-msa.h new file mode 100644 index 000000000..f09b6d32b --- /dev/null +++ b/examples/talk-llama/llama-kv-cache-msa.h @@ -0,0 +1,153 @@ +#pragma once + +#include "llama-kv-cache.h" + +#include + +// 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 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 kv_base; + std::unique_ptr 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 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 ubatches; + + const llama_memory_context_ptr ctx_base; + const llama_memory_context_ptr ctx_idx; + + const llama_memory_status status; +}; diff --git a/examples/talk-llama/llama-kv-cache.cpp b/examples/talk-llama/llama-kv-cache.cpp index e70583e64..8678a326d 100644 --- a/examples/talk-llama/llama-kv-cache.cpp +++ b/examples/talk-llama/llama-kv-cache.cpp @@ -323,7 +323,7 @@ llama_kv_cache::llama_kv_cache( 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; } @@ -1224,6 +1224,12 @@ ggml_tensor * llama_kv_cache::get_k_storage(int32_t il) const { 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; @@ -2054,7 +2060,12 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama 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) { diff --git a/examples/talk-llama/llama-kv-cache.h b/examples/talk-llama/llama-kv-cache.h index 531d99dbd..6cb6dbd2f 100644 --- a/examples/talk-llama/llama-kv-cache.h +++ b/examples/talk-llama/llama-kv-cache.h @@ -164,6 +164,8 @@ public: std::vector 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 // diff --git a/examples/talk-llama/llama-memory-recurrent.cpp b/examples/talk-llama/llama-memory-recurrent.cpp index 3d6c6db87..ef82eb976 100644 --- a/examples/talk-llama/llama-memory-recurrent.cpp +++ b/examples/talk-llama/llama-memory-recurrent.cpp @@ -819,7 +819,12 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i 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) { diff --git a/examples/talk-llama/llama-model-loader.cpp b/examples/talk-llama/llama-model-loader.cpp index 28f8bb793..b31e92e2d 100644 --- a/examples/talk-llama/llama-model-loader.cpp +++ b/examples/talk-llama/llama-model-loader.cpp @@ -4,6 +4,7 @@ #include "ggml.h" #include "gguf.h" #include "llama-hparams.h" +#include "llama.h" #include #include @@ -522,10 +523,10 @@ llama_model_loader::llama_model_loader( const std::string & fname, std::vector & 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) { @@ -542,6 +543,9 @@ llama_model_loader::llama_model_loader( 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; @@ -562,20 +566,6 @@ llama_model_loader::llama_model_loader( 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. @@ -816,15 +806,14 @@ llama_model_loader::llama_model_loader( } } - 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 { @@ -868,7 +857,11 @@ struct ggml_tensor * llama_model_loader::require_tensor_meta(const std::string & return tensor; } -const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::string & name, const std::vector & ne, bool required) const { +const struct ggml_tensor * llama_model_loader::check_tensor_dims( + const std::string & name, + const std::vector & ne, + bool required, + bool allow_reshape) const { const struct ggml_tensor * cur = get_tensor_meta(name.c_str()); if (cur == NULL) { @@ -878,21 +871,33 @@ const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::stri 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; @@ -1257,11 +1262,25 @@ struct ggml_tensor * llama_model_loader::create_tensor( return ret; } - ggml_tensor * t_meta = get_tensor_meta(tn.str().c_str()); - ggml_backend_buffer_type_t buft = buft_for_tensor(t_meta); - if (buft == nullptr) { - return nullptr; // return type is ggml_tensor * + 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; + } + 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 @@ -1272,20 +1291,13 @@ struct ggml_tensor * llama_model_loader::create_tensor( } } - 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++; } @@ -1293,34 +1305,6 @@ struct ggml_tensor * llama_model_loader::create_tensor( 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 & 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 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)); diff --git a/examples/talk-llama/llama-model-loader.h b/examples/talk-llama/llama-model-loader.h index c476026d3..d6b31c231 100644 --- a/examples/talk-llama/llama-model-loader.h +++ b/examples/talk-llama/llama-model-loader.h @@ -67,6 +67,7 @@ struct llama_model_loader { 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; @@ -79,6 +80,7 @@ struct llama_model_loader { bool use_direct_io = false; bool check_tensors; bool no_alloc; + bool load_mtp; llama_files files; llama_ftype ftype; @@ -126,10 +128,10 @@ struct llama_model_loader { const std::string & fname, std::vector & 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); @@ -176,14 +178,16 @@ struct llama_model_loader { 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 & ne, bool required) const; + const struct ggml_tensor * check_tensor_dims( + const std::string & name, + const std::vector & 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 & 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 & 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); diff --git a/examples/talk-llama/llama-model-saver.cpp b/examples/talk-llama/llama-model-saver.cpp index a3928523b..3812c594e 100644 --- a/examples/talk-llama/llama-model-saver.cpp +++ b/examples/talk-llama/llama-model-saver.cpp @@ -28,6 +28,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: case LLM_ARCH_MELLUM: + case LLM_ARCH_LAGUNA: return false; default: return true; @@ -280,6 +281,9 @@ void llama_model_saver::add_kv_from_model() { 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; diff --git a/examples/talk-llama/llama-model.cpp b/examples/talk-llama/llama-model.cpp index adacf702d..333f506de 100644 --- a/examples/talk-llama/llama-model.cpp +++ b/examples/talk-llama/llama-model.cpp @@ -11,11 +11,13 @@ #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" @@ -84,6 +86,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params 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: @@ -250,6 +254,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params 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: @@ -262,6 +268,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params 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: @@ -280,6 +288,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params 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: @@ -313,8 +323,7 @@ llama_model * llama_model_create(llm_arch arch, const llama_model_params & param 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) + "'"); } } @@ -336,38 +345,40 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str 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"); + 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"); - 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"); + 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"); - 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"); + 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"); - const std::regex pattern_output_weight("output\\.weight"); - const std::regex pattern_output_bias ("output\\.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; @@ -468,10 +479,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } // 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)) { @@ -555,6 +566,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str 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); @@ -631,7 +650,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } // 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); @@ -799,10 +819,12 @@ const char * llm_type_name(llm_type type) { 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"; @@ -1068,6 +1090,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { 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); @@ -1113,6 +1136,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { 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); @@ -1229,7 +1253,7 @@ void llama_model_base::load_vocab(llama_model_loader & ml) { 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; @@ -1239,8 +1263,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { 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); @@ -2048,9 +2072,13 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, { 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, @@ -2063,17 +2091,140 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 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( @@ -2085,7 +2236,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 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; @@ -2100,7 +2251,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 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); }; @@ -2166,11 +2317,13 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; } - 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 { @@ -2178,24 +2331,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } - 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) { @@ -2304,19 +2440,18 @@ llama_model_params llama_model_default_params() { /*.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; @@ -2472,6 +2607,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { 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 @@ -2526,6 +2662,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { 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: @@ -2533,17 +2670,22 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { 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; @@ -2725,7 +2867,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l 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 & ne, int flags) { GGML_ASSERT(ml != nullptr); diff --git a/examples/talk-llama/llama-model.h b/examples/talk-llama/llama-model.h index 45b054ced..6b9e94a0a 100644 --- a/examples/talk-llama/llama-model.h +++ b/examples/talk-llama/llama-model.h @@ -130,10 +130,12 @@ enum llm_type { 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 @@ -515,6 +517,12 @@ struct llama_layer { 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; @@ -599,6 +607,12 @@ struct llama_model { 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 target_layer_ids; @@ -705,6 +719,7 @@ struct llama_model_base : public llama_model { 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; diff --git a/examples/talk-llama/llama-quant.cpp b/examples/talk-llama/llama-quant.cpp index aebbc1ffb..fd6e787bd 100644 --- a/examples/talk-llama/llama-quant.cpp +++ b/examples/talk-llama/llama-quant.cpp @@ -2,6 +2,7 @@ #include "llama-model.h" #include "llama-model-loader.h" #include "llama-ext.h" +#include "llama.h" #include #include @@ -306,6 +307,9 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // 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; @@ -322,6 +326,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param 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; @@ -351,6 +359,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param 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; } @@ -673,7 +685,7 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod 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()) { @@ -692,7 +704,7 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod } // 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); } @@ -873,15 +885,15 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // 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 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(); @@ -1351,6 +1363,7 @@ llama_model * llama_quant_model_from_metadata(const llama_quant_model_desc * des 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++) { diff --git a/examples/talk-llama/llama-sampler.cpp b/examples/talk-llama/llama-sampler.cpp index 2370e91a1..6cf2d27cf 100644 --- a/examples/talk-llama/llama-sampler.cpp +++ b/examples/talk-llama/llama-sampler.cpp @@ -263,6 +263,10 @@ static void llama_log_softmax(float * array, size_t size) { */ 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; @@ -989,7 +993,9 @@ static void llama_sampler_greedy_backend_apply( 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; @@ -1154,7 +1160,10 @@ static void llama_sampler_dist_backend_apply( 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); @@ -1285,22 +1294,22 @@ static void llama_sampler_top_k_backend_apply( 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); } @@ -1431,21 +1440,25 @@ static void llama_sampler_top_p_backend_apply( 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"); @@ -1622,10 +1635,12 @@ static void llama_sampler_min_p_backend_apply( 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); @@ -1636,7 +1651,7 @@ static void llama_sampler_min_p_backend_apply( 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). @@ -1648,7 +1663,7 @@ static void llama_sampler_min_p_backend_apply( 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); @@ -1825,18 +1840,20 @@ static void llama_sampler_backend_temp_sampling( 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; @@ -2015,13 +2032,15 @@ static void llama_sampler_temp_ext_backend_apply( 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 @@ -2059,7 +2078,7 @@ static void llama_sampler_temp_ext_backend_apply( 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; @@ -2619,7 +2638,8 @@ struct llama_sampler * llama_sampler_init_grammar_lazy_patterns( // 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; @@ -2629,10 +2649,50 @@ struct llama_sampler_penalties { // a frequency map to count token occurrences std::unordered_map 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 host_token_ids; + std::vector 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) { @@ -2669,8 +2729,7 @@ static void llama_sampler_penalties_accept(struct llama_sampler * smpl, llama_to 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; } @@ -2708,6 +2767,7 @@ static void llama_sampler_penalties_reset(struct llama_sampler * smpl) { 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, @@ -2717,7 +2777,8 @@ static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_s { 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; @@ -2727,6 +2788,170 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) { 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> 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, @@ -2734,35 +2959,33 @@ static struct llama_sampler_i llama_sampler_penalties_i = { /* .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(penalty_last_n), - /* .token_count = */ {}, - } + /* .ctx = */ new llama_sampler_penalties( + n_vocab, + penalty_last_n, + penalty_repeat, + penalty_freq, + penalty_present) ); } diff --git a/examples/talk-llama/llama-vocab.cpp b/examples/talk-llama/llama-vocab.cpp index fdd447147..10032a8c6 100644 --- a/examples/talk-llama/llama-vocab.cpp +++ b/examples/talk-llama/llama-vocab.cpp @@ -496,6 +496,12 @@ struct llm_tokenizer_bpe : llm_tokenizer { "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][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 @@ -1325,6 +1331,9 @@ struct llm_tokenizer_rwkv_session { token_id = node->value; token_length = position + 1; } + if (position + 1 >= text.size()) { + break; + } node = node->traverse(text[++position]); } @@ -2342,6 +2351,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { 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; @@ -2519,6 +2532,12 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const std::string & key = kv(std::get<0>(it)); int32_t & id = std::get<1>(it); + if (id >= 0 && static_cast(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; @@ -2565,7 +2584,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { 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); + } + } } } @@ -2796,6 +2822,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "" // 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) { @@ -2855,6 +2882,11 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { 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()); @@ -2889,6 +2921,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { 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; @@ -4018,7 +4053,11 @@ int llama_vocab::find_bpe_rank(const std::string & token_left, const std::string } std::vector llama_vocab::get_bpe_merges() const { - std::vector 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 result(max_rank + 1); for (const auto & pair : pimpl->bpe_ranks) { result[pair.second] = pair.first.first + " " + pair.first.second; @@ -4179,6 +4218,14 @@ bool llama_vocab_get_add_sep(const struct llama_vocab * vocab) { 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 & 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(); } diff --git a/examples/talk-llama/llama-vocab.h b/examples/talk-llama/llama-vocab.h index 707cd4bac..b7c289263 100644 --- a/examples/talk-llama/llama-vocab.h +++ b/examples/talk-llama/llama-vocab.h @@ -64,6 +64,7 @@ enum llama_vocab_pre_type { 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; diff --git a/examples/talk-llama/llama.cpp b/examples/talk-llama/llama.cpp index 0de6048f2..d6e0bbfef 100644 --- a/examples/talk-llama/llama.cpp +++ b/examples/talk-llama/llama.cpp @@ -46,6 +46,31 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty 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, @@ -279,8 +304,8 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama static std::pair 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 & 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 model_ptr(llama_model_create(ml, params)); @@ -412,7 +437,7 @@ struct llama_model * llama_model_init_from_user( GGML_ASSERT(metadata != nullptr); std::string path_model; std::vector 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); } diff --git a/examples/talk-llama/llama.h b/examples/talk-llama/llama.h index a311ac202..fb2ca38ce 100644 --- a/examples/talk-llama/llama.h +++ b/examples/talk-llama/llama.h @@ -202,6 +202,17 @@ extern "C" { 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, @@ -301,6 +312,7 @@ extern "C" { 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; @@ -321,13 +333,11 @@ extern "C" { // 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 { @@ -1093,6 +1103,9 @@ extern "C" { 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); @@ -1411,10 +1424,11 @@ extern "C" { /// 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( diff --git a/examples/talk-llama/models/cohere2moe.cpp b/examples/talk-llama/models/cohere2moe.cpp index 499c73a1c..3acb7e77a 100644 --- a/examples/talk-llama/models/cohere2moe.cpp +++ b/examples/talk-llama/models/cohere2moe.cpp @@ -55,7 +55,11 @@ void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) { 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); diff --git a/examples/talk-llama/models/deepseek2.cpp b/examples/talk-llama/models/deepseek2.cpp index a9e8bc514..ba90c0d07 100644 --- a/examples/talk-llama/models/deepseek2.cpp +++ b/examples/talk-llama/models/deepseek2.cpp @@ -37,6 +37,11 @@ void llama_model_deepseek2::load_arch_hparams(llama_model_loader & ml) { 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? @@ -52,10 +57,20 @@ void llama_model_deepseek2::load_arch_hparams(llama_model_loader & ml) { } } -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 @@ -81,44 +96,45 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader &) { 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"); @@ -128,21 +144,281 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader &) { } // 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 llama_model_deepseek2::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*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(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 @@ -365,7 +641,7 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p 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); } @@ -425,6 +701,13 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p 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; diff --git a/examples/talk-llama/models/deepseek32.cpp b/examples/talk-llama/models/deepseek32.cpp index 9a20e2ce9..8a07a0b71 100644 --- a/examples/talk-llama/models/deepseek32.cpp +++ b/examples/talk-llama/models/deepseek32.cpp @@ -44,13 +44,24 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) { 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"); @@ -80,12 +91,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) { } 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]; @@ -138,7 +144,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) { 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); @@ -153,6 +159,9 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) { } std::unique_ptr llama_model_deepseek32::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -301,43 +310,50 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ 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); + 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); - // sum by q n_indexer_head dimension - indexer_score = ggml_sum_rows(ctx0, indexer_score); - cb(indexer_score, "indexer_score", il); + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "indexer_kq", 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); + // ReLU requires contiguous tensors + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "indexer_kq", 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); + // 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; @@ -423,7 +439,9 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ 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); } @@ -486,6 +504,14 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ 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; @@ -497,3 +523,243 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ 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(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); +} + diff --git a/examples/talk-llama/models/deepseek4.cpp b/examples/talk-llama/models/deepseek4.cpp index 07aa477e1..89cd46176 100644 --- a/examples/talk-llama/models/deepseek4.cpp +++ b/examples/talk-llama/models/deepseek4.cpp @@ -16,6 +16,16 @@ static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) { } 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); @@ -24,8 +34,8 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { 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; } @@ -41,9 +51,11 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { 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); @@ -54,6 +66,9 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { } 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; @@ -61,7 +76,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { } } -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; @@ -75,6 +90,10 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { 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); @@ -84,69 +103,84 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { 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]; + 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_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.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}, 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.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}, 0); + 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}, 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_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}, 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); + + 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); + } } } std::unique_ptr llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -175,18 +209,69 @@ static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, b 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, @@ -197,22 +282,31 @@ static ggml_tensor * dsv4_hc_affine( 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( @@ -275,11 +369,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( 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); @@ -293,13 +385,23 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_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( @@ -308,7 +410,14 @@ 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]; @@ -320,7 +429,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_post( 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)); } @@ -350,7 +460,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_head( 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( @@ -435,27 +545,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta 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); @@ -556,25 +668,32 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( 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)); @@ -770,8 +889,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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(); @@ -839,9 +979,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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); @@ -872,10 +1015,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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, @@ -896,8 +1045,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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); @@ -924,10 +1084,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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, @@ -948,8 +1114,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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); @@ -963,15 +1140,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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, @@ -990,15 +1173,41 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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); @@ -1013,7 +1222,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( } 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) { @@ -1044,7 +1260,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( 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); @@ -1072,6 +1288,12 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p 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; @@ -1099,6 +1321,10 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p &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); @@ -1141,13 +1367,26 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p 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); @@ -1163,3 +1402,145 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p 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(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); +} diff --git a/examples/talk-llama/models/dflash.cpp b/examples/talk-llama/models/dflash.cpp index a7b4f4435..6c82ab3da 100644 --- a/examples/talk-llama/models/dflash.cpp +++ b/examples/talk-llama/models/dflash.cpp @@ -1,5 +1,6 @@ #include "models.h" +#include "llama-impl.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -19,6 +20,48 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { } 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) { @@ -36,10 +79,77 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { 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]; @@ -66,6 +176,9 @@ std::unique_ptr llama_model_dflash::build_arch_graph(const ll return std::make_unique>(*this, params); case LLM_GRAPH_TYPE_DEFAULT: case LLM_GRAPH_TYPE_DECODER: + if (hparams.dsv4_hc_mult > 0) { + return std::make_unique(*this, params); + } return std::make_unique>(*this, params); default: GGML_ABORT("invalid graph type"); @@ -104,6 +217,94 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_grap 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 @@ -164,9 +365,25 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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)); } @@ -193,6 +410,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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); @@ -273,4 +492,184 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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(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(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); + } } diff --git a/examples/talk-llama/models/eagle3.cpp b/examples/talk-llama/models/eagle3.cpp index 9d96fae59..be466056d 100644 --- a/examples/talk-llama/models/eagle3.cpp +++ b/examples/talk-llama/models/eagle3.cpp @@ -28,6 +28,10 @@ void llama_model_eagle3::load_arch_hparams(llama_model_loader & ml) { 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; } @@ -53,6 +57,11 @@ void llama_model_eagle3::load_arch_tensors(llama_model_loader &) { // 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); @@ -130,6 +139,12 @@ llama_model_eagle3::graph::graph(const llama_model & model, const llm_grap 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); diff --git a/examples/talk-llama/models/gemma4.cpp b/examples/talk-llama/models/gemma4.cpp index 6a96979ce..e44f423bd 100644 --- a/examples/talk-llama/models/gemma4.cpp +++ b/examples/talk-llama/models/gemma4.cpp @@ -142,33 +142,6 @@ static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, in 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 arr; -}; - llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model), @@ -429,16 +402,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para 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 and 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(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; diff --git a/examples/talk-llama/models/glm-dsa.cpp b/examples/talk-llama/models/glm-dsa.cpp index 32fe6def6..360c2ee77 100644 --- a/examples/talk-llama/models/glm-dsa.cpp +++ b/examples/talk-llama/models/glm-dsa.cpp @@ -1,5 +1,31 @@ #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 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); @@ -34,18 +60,43 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { // 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"); @@ -74,12 +125,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { } 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]; @@ -132,7 +180,7 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { 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); @@ -147,6 +195,616 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { } std::unique_ptr 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(*this, params); + } return std::make_unique(*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(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); +} diff --git a/examples/talk-llama/models/hy-v3.cpp b/examples/talk-llama/models/hy-v3.cpp new file mode 100644 index 000000000..61db93af8 --- /dev/null +++ b/examples/talk-llama/models/hy-v3.cpp @@ -0,0 +1,394 @@ +#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 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(*this, params); + } + return std::make_unique(*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(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); +} diff --git a/examples/talk-llama/models/laguna.cpp b/examples/talk-llama/models/laguna.cpp new file mode 100644 index 000000000..82c9a9538 --- /dev/null +++ b/examples/talk-llama/models/laguna.cpp @@ -0,0 +1,333 @@ +// 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 llama_model_laguna::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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); +} diff --git a/examples/talk-llama/models/mimo2.cpp b/examples/talk-llama/models/mimo2.cpp index 889891605..d50e186cc 100644 --- a/examples/talk-llama/models/mimo2.cpp +++ b/examples/talk-llama/models/mimo2.cpp @@ -25,9 +25,17 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) { } } -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 @@ -40,41 +48,46 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader &) { 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 llama_model_mimo2::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -89,6 +102,8 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param 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; @@ -168,7 +183,7 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param } } - 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); } @@ -218,6 +233,15 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param 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); @@ -233,3 +257,143 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param 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(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); +} diff --git a/examples/talk-llama/models/minimax-m2.cpp b/examples/talk-llama/models/minimax-m2.cpp index b25435e4d..86a8ae2b1 100644 --- a/examples/talk-llama/models/minimax-m2.cpp +++ b/examples/talk-llama/models/minimax-m2.cpp @@ -60,6 +60,8 @@ llama_model_minimax_m2::graph::graph(const llama_model & model, const llm_graph_ 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; diff --git a/examples/talk-llama/models/minimax-m3.cpp b/examples/talk-llama/models/minimax-m3.cpp new file mode 100644 index 000000000..854d5aed0 --- /dev/null +++ b/examples/talk-llama/models/minimax-m3.cpp @@ -0,0 +1,601 @@ +#include "models.h" +#include "llama-kv-cache-msa.h" +#include +#include +#include + +// 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 llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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 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(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(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(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(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(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 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); +} diff --git a/examples/talk-llama/models/models.h b/examples/talk-llama/models/models.h index 7a52e7bc1..5f206621d 100644 --- a/examples/talk-llama/models/models.h +++ b/examples/talk-llama/models/models.h @@ -424,6 +424,22 @@ struct llama_model_mellum : public llama_model_base { std::unique_ptr 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 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; @@ -1068,6 +1084,10 @@ struct llama_model_deepseek2 : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -1081,6 +1101,10 @@ struct llama_model_deepseek32 : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -1091,6 +1115,7 @@ struct llama_model_deepseek4 : public llama_model_base { 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( @@ -1122,6 +1147,21 @@ struct llama_model_deepseek4 : public llama_model_base { 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, @@ -1187,15 +1227,20 @@ struct llama_model_deepseek4 : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -1216,7 +1261,13 @@ struct llama_model_glm_dsa : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -1249,6 +1300,10 @@ struct llama_model_dflash : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -1680,6 +1735,19 @@ struct llama_model_afmoe : public llama_model_base { }; +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 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; @@ -1729,6 +1797,22 @@ struct llama_model_hunyuan_moe : public llama_model_base { std::unique_ptr 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 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) {} @@ -1870,6 +1954,29 @@ struct llama_model_minimax_m2 : public llama_model_base { std::unique_ptr 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 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) {} @@ -1934,6 +2041,10 @@ struct llama_model_qwen3next : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; @@ -2052,6 +2163,10 @@ struct llama_model_mimo2 : public llama_model_base { 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 build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/examples/talk-llama/models/nanbeige.cpp b/examples/talk-llama/models/nanbeige.cpp new file mode 100644 index 000000000..3a546600f --- /dev/null +++ b/examples/talk-llama/models/nanbeige.cpp @@ -0,0 +1,184 @@ +#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 llama_model_nanbeige::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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(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); +} diff --git a/examples/talk-llama/models/openai-moe.cpp b/examples/talk-llama/models/openai-moe.cpp index 6d74f9c7e..c91bae1c3 100644 --- a/examples/talk-llama/models/openai-moe.cpp +++ b/examples/talk-llama/models/openai-moe.cpp @@ -116,7 +116,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_ 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); @@ -154,6 +154,12 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_ } 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); diff --git a/examples/talk-llama/models/qwen35.cpp b/examples/talk-llama/models/qwen35.cpp index d8ffe43ae..309dd4324 100644 --- a/examples/talk-llama/models/qwen35.cpp +++ b/examples/talk-llama/models/qwen35.cpp @@ -39,6 +39,7 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { 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); @@ -97,25 +98,25 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { 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) { diff --git a/examples/talk-llama/models/qwen35moe.cpp b/examples/talk-llama/models/qwen35moe.cpp index 7b0876cbb..38f2a5798 100644 --- a/examples/talk-llama/models/qwen35moe.cpp +++ b/examples/talk-llama/models/qwen35moe.cpp @@ -42,6 +42,7 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { 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); @@ -113,32 +114,32 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { 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) { diff --git a/examples/talk-llama/models/qwen3next.cpp b/examples/talk-llama/models/qwen3next.cpp index 09b66423d..0808fd87a 100644 --- a/examples/talk-llama/models/qwen3next.cpp +++ b/examples/talk-llama/models/qwen3next.cpp @@ -13,7 +13,11 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); - // 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); @@ -28,13 +32,17 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) { } } -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 @@ -61,49 +69,73 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) { 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 llama_model_qwen3next::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -120,6 +152,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p 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; @@ -139,7 +172,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p 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); } @@ -171,9 +204,16 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p } 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; @@ -186,14 +226,6 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p 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, @@ -216,7 +248,7 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( // 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); @@ -232,10 +264,10 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( 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); @@ -274,8 +306,6 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( 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); @@ -550,16 +580,19 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c 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); @@ -593,3 +626,198 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c } 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(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); +} diff --git a/examples/talk-llama/models/step35.cpp b/examples/talk-llama/models/step35.cpp index 9b7b18a36..5b1d90258 100644 --- a/examples/talk-llama/models/step35.cpp +++ b/examples/talk-llama/models/step35.cpp @@ -48,7 +48,11 @@ void llama_model_step35::load_arch_tensors(llama_model_loader & ml) { 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);