graph : fix unused input tensors in minimax m3 graph (#26519)
This commit is contained in:
+11
-5
@@ -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));
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user