nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: adjust token count for MTP draft generation in kv cache update (#2181)

This commit is contained in:
Samuel Oliveira Alves 2026-07-25 12:07:48 -03:00 committed by GitHub
parent 3e2f5696a1
commit de55d9e2f6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/llama.cpp
View File

@ -6949,6 +6949,10 @@ static int32_t llama_kv_cache_update_internal(struct llama_context & lctx) {
// TODO: extract to a function // TODO: extract to a function
// build worst-case graph // build worst-case graph
int n_tokens = (int)std::min(lctx.cparams.n_ctx, lctx.cparams.n_ubatch); 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; 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_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); llama_batch reserve_batch = llama_batch_get_one(&token, n_tokens, n_past, 0);