Qwen3.5 MTP: extract selected tokens earlier

This commit is contained in:
Kawrakow
2026-05-28 11:50:51 +00:00
parent 3bf7e836c2
commit a18eeb01cb
+2 -17
View File
@@ -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);