diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index a364174..bc626ea 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1193,6 +1193,7 @@ const char * llm_arch_name(llm_arch arch) { } llm_arch llm_arch_from_string(const std::string & name) { + if (name == "minimaxm31_sakura") { return LLM_ARCH_MINIMAX_M3; } for (const auto & kv : LLM_ARCH_NAMES) { // NOLINT if (kv.second == name) { return kv.first; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 719d708..efe70a2 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -307,6 +307,7 @@ struct llama_hparams { // DSA (deepseek sparse attention) uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; + uint32_t indexer_first_layer = 0; uint32_t indexer_top_k = 0; uint32_t indexer_kpool = 0; // k-pool size bool indexer_kpool_select_tail = true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4cb0409..b684aea 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2815,9 +2815,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_MINIMAX_M3: { - // sparse (MSA) layers carry an indexer key cache, but leading dense layers do not + // Only sparse attention layers carry an indexer key cache. llama_kv_cache::layer_filter_cb filter_idx = - [&](int32_t il) { return (uint32_t) il >= hparams.n_layer_dense_lead; }; + [&](int32_t il) { return (uint32_t) il >= hparams.indexer_first_layer; }; res = new llama_kv_cache_msa( *this, diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 35e7b1a..df860f6 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -13,6 +13,9 @@ 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); + hparams.indexer_first_layer = hparams.n_layer_dense_lead; + ml.get_key("minimax-m3.attention.indexer.first_layer", hparams.indexer_first_layer, false); + if (hparams.indexer_first_layer > hparams.n_layer()) { throw std::runtime_error("Invalid first sparse layer"); } ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); 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); @@ -79,7 +82,8 @@ void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { 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); - + } + if ((uint32_t) i >= hparams.indexer_first_layer) { // 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); @@ -233,6 +237,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ 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; + if (hparams.indexer_first_layer == 0 && !msa_enabled) { + throw std::runtime_error("MiniMax M3.1 requires flash attention and supported stream layout"); + } auto * inp_attn = build_attn_inp_kv_msa(msa_enabled); @@ -331,7 +338,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); - const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead; + const bool is_sparse = msa_enabled && il >= (int) hparams.indexer_first_layer; if (!is_sparse) { cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s,