File size: 4,600 Bytes
99127b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | 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,
|