diff --git a/src/llama.cpp b/src/llama.cpp index 6710ea32b..657dfdeab 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -6949,6 +6949,10 @@ static int32_t llama_kv_cache_update_internal(struct llama_context & lctx) { // TODO: extract to a function // build worst-case graph int n_tokens = (int)std::min(lctx.cparams.n_ctx, lctx.cparams.n_ubatch); + // MTP draft generation consumes one token and one hidden-state vector per decode step. + if (lctx.cparams.mtp_op_type == MTP_OP_DRAFT_GEN) { + n_tokens = 1; + } int n_past = lctx.cparams.n_ctx - n_tokens; llama_token token = llama_token_bos(&lctx.model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph llama_batch reserve_batch = llama_batch_get_one(&token, n_tokens, n_past, 0);