diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 309dd4324..dd79513a5 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -41,14 +41,24 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + int64_t n_vocab_out = n_vocab; + const ggml_tensor * d2t_meta = ml.get_tensor_meta("d2t"); + if (mtp_only && d2t_meta) { + n_vocab_out = d2t_meta->ne[0]; + d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_out }, 0); + LLAMA_LOG_INFO("%s: QWEN35 MTP using d2t draft-vocab trim (n_vocab_out = %lld)\n", + __func__, (long long) n_vocab_out); + } + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); // output output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); - output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_out }, TENSOR_NOT_REQUIRED); // if output is NULL, init from the input tok embed if (output == NULL) { + GGML_ASSERT(!d2t && "d2t draft-vocab trim requires output.weight"); output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); } @@ -640,6 +650,22 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_lora_mm(head_w, cur, head_s); cb(cur, "result_output", -1); + if (model.d2t) { + const int64_t n_draft_vocab = cur->ne[0]; + const int64_t n_outputs = cur->ne[1]; + const int64_t n_vocab_full = (int64_t) model.vocab.n_tokens(); + + GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab); + + ggml_tensor * logits = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab_full, n_outputs), -INFINITY); + cur = ggml_set_rows(ctx0, logits, + ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + cur = ggml_reshape_2d(ctx0, cur, n_vocab_full, n_outputs); + cb(cur, "result_output_d2t", -1); + } + res->t_logits = cur; ggml_build_forward_expand(gf, cur); }