graph : fix unused input tensors in minimax m3 graph (#26519)

diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index fdab7b8dde679b068b6b024e2c9049373bb92c92..9dde345df406bfe749b14fd0dc2b711762f01b50 100644 (file)
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -530,7 +530,9 @@ llm_graph_input_attn_kv_msa::llm_graph_input_attn_kv_msa(
 void llm_graph_input_attn_kv_msa::set_input(const llama_ubatch * ubatch) {
     llm_graph_input_attn_kv::set_input(ubatch);
 
-    mctx_msa->get_idx()->set_input_k_idxs(self_k_idxs_idx, 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) {
@@ -541,8 +543,10 @@ bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
 
     bool res = true;
 
-    res &= self_k_idxs    ->ne[0] == params.ubatch.n_tokens;
-    res &= self_k_idxs_idx->ne[0] == params.ubatch.n_tokens;
+    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);
 
@@ -3218,7 +3222,7 @@ 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() const {
+llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const {
     const auto * mctx_cur = static_cast<const llama_kv_cache_msa_context *>(mctx);
 
     auto inp = std::make_unique<llm_graph_input_attn_kv_msa>(hparams, cparams, mctx_cur);
@@ -3239,7 +3243,9 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa() const {
     inp->self_k_rot = mctx_base->build_input_k_rot(ctx0);
     inp->self_v_rot = mctx_base->build_input_v_rot(ctx0);
 
-    inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch);
+    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));
 }

diff --git a/src/llama-graph.h b/src/llama-graph.h
index ff216302dbb11f7b4aa404dab890e335e97ec0d5..32d8d395aa4546ed7e90e7d26d24218fcd37547a 100644 (file)
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1190,7 +1190,7 @@ 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() 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,

diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp
index 8bd6a4298e0d7bfb2d8553fc27b6caa0c6a36214..854d5aed0f822df145801d13790872275dfa5cb0 100644 (file)
--- a/src/models/minimax-m3.cpp
+++ b/src/models/minimax-m3.cpp
@@ -213,7 +213,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
     inpL = build_inp_embd(model.tok_embd);
 
     ggml_tensor * inp_pos = build_inp_pos();
-    auto inp_attn = build_attn_inp_kv_msa();
+
+    // ==========================================
+    // 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
@@ -225,6 +227,8 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
     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 "
@@ -237,6 +241,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
                        "-> 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;