diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index c43d984ab..374c1e913 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -36,12 +36,22 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { const auto nf = nextn_flags(ml); const int trunk_flags = nf.trunk; const int mtp_flags = nf.mtp; + const bool mtp_only = (nf.trunk & TENSOR_NOT_REQUIRED) != 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); // optional projection of the embeddings output cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), { n_embd, hparams.n_embd_out() }, TENSOR_NOT_REQUIRED); @@ -49,6 +59,7 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) { // 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); } @@ -667,5 +678,21 @@ 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; }