diff --git a/gpttype_adapter.cpp b/gpttype_adapter.cpp index 63fc58685..5ea5b14df 100644 --- a/gpttype_adapter.cpp +++ b/gpttype_adapter.cpp @@ -6657,7 +6657,7 @@ size_t gpttype_save_state_kv(int slot) savestates[slot].current_savestate_size = newsize; savestates[slot].savestate_context_tokens = current_context_tokens; savestates[slot].media_signature = media_composite_image_signature; - float * lgptr = llama_get_logits(llama_ctx_v4); + float * lgptr = (draft_is_mtp ? llama_get_logits_ith(llama_ctx_v4, -1) : llama_get_logits(llama_ctx_v4)); savestates[slot].latest_logits.assign(lgptr,lgptr+n_vocab); int maxedpos = llama_memory_seq_pos_max(llama_get_memory(llama_ctx_v4),0); //kcpp: so maxedpos appears to always be equal to ctx tokens - 2, if savestate_ctx_tokens > maxedpos + 2 then trim excess