From a18eeb01cbba730aa4a12efe19c45af201f84137 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Thu, 28 May 2026 11:50:51 +0000 Subject: [PATCH] Qwen3.5 MTP: extract selected tokens earlier --- src/graphs/build_qwen35.cpp | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/src/graphs/build_qwen35.cpp b/src/graphs/build_qwen35.cpp index 2fd0eecd0..9dadc881f 100644 --- a/src/graphs/build_qwen35.cpp +++ b/src/graphs/build_qwen35.cpp @@ -211,14 +211,10 @@ struct ggml_tensor * llm_build_context::build_qwen35moe_mtp( const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); cur = build_std_attention(gf, mtp_layer.attn_norm, cur, - inp_pos, nullptr, nullptr, + inp_pos, inp_out_ids, nullptr, KQ_mask, nullptr, nullptr, kq_scale, 0.0f, 0, il, true, false, true, false, true, nullptr); - if (inp_out_ids) { - cur = ggml_get_rows(ctx0, cur, inp_out_ids); - } - cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur, mtp_layer.ffn_gate_inp, nullptr, mtp_layer.ffn_up_exps, nullptr, @@ -236,8 +232,6 @@ struct ggml_tensor * llm_build_context::build_qwen35moe_mtp( cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "ffn_out", il); - cb(cur, "result_norm", -1); - cur = build_output(lctx, ctx0, cur, model.output_mtp, mtp_layer.nextn.shared_head_norm, cb); cb(cur, "result_output", -1); @@ -291,14 +285,10 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp( const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); cur = build_std_attention(gf, mtp_layer.attn_norm, cur, - inp_pos, nullptr, nullptr, + inp_pos, inp_out_ids, nullptr, KQ_mask, nullptr, nullptr, kq_scale, 0.0f, 0, il, true, false, true, false, true, nullptr); - if (inp_out_ids) { - cur = ggml_get_rows(ctx0, cur, inp_out_ids); - } - // Dense FFN — optional (9B and 4B don't have FFN in MTP layer) if (mtp_layer.ffn_gate != nullptr) { cur = llm_build_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur, @@ -312,11 +302,6 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp( cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "ffn_out", il); - // As far as I can tell this was wrong. We need the FFN output, and not the normalized result. - //cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il); - cb(cur, "result_norm", -1); - - //cur = build_output(lctx, ctx0, cur, model.output, nullptr, cb); cur = build_output(lctx, ctx0, cur, model.output_mtp, mtp_layer.nextn.shared_head_norm, cb); cb(cur, "result_output", -1);