vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-07-09 17:08:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

mtp drafting added (+1 squashed commits)

Browse source 
Squashed commits:

[b0fa1a6de] wip 1
This commit is contained in:
Concedo 2026-06-11 15:11:41 +08:00
parent 10e4b6d5e3
commit 40084bb1f8
4 changed files with 156 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

1
expose.h
View file

@ -51,6 +51,7 @@ struct load_model_inputs
const int visionmaxtokens = -1;
const bool use_mmap = false;
const bool use_mlock = false;
const bool use_mtp = false;
const bool use_smartcontext = false;
const bool use_contextshift = false;
const bool use_fastforward = false;

167
gpttype_adapter.cpp
View file

@ -119,6 +119,8 @@ static llama_context * llama_ctx_v4 = nullptr;
static llama_context * draft_ctx = nullptr; //will remain null if speculative is unused
static common_speculative * draft_spec = nullptr; // llama.cpp speculative state, used for MTP draft heads
static bool draft_is_mtp = false;
static bool mtp_uses_spec_checkpoint = false;
static common_prompt_checkpoint mtp_spec_ckpt;
static llama_context * guidance_ctx = nullptr; //for classifier free guidance, will be null if unused
static mtmd_context * mtmd_ctx = nullptr; //for multimodal media
@ -657,6 +659,60 @@ const char * kcpp_print_system_info(void) {
return s.c_str();
}
static bool mtp_speculative_state_setup(llama_context * main_ctx, const llama_context_params & mtp_ctx_params, int draft_gpulayers)
{
common_params_speculative spec_params;
spec_params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
spec_params.draft.ctx_tgt = main_ctx;
spec_params.draft.ctx_dft = draft_ctx;
spec_params.draft.n_max = speculative_chunk_amt;
spec_params.draft.n_min = 0;
spec_params.draft.p_min = 0.0f;
spec_params.draft.backend_sampling = false;
spec_params.draft.n_gpu_layers = draft_gpulayers;
spec_params.draft.cache_type_k = mtp_ctx_params.type_k;
spec_params.draft.cache_type_v = mtp_ctx_params.type_v;
draft_spec = common_speculative_init(spec_params, 1);
if(draft_spec == nullptr)
{
printf("Error: failed to initialize MTP speculative decoding state.\n");
llama_free(draft_ctx);
draft_ctx = nullptr;
draft_is_mtp = false;
return false;
}
return true;
}
static void mtp_decoding_setup(llama_model * main_model, llama_context * main_ctx, const llama_context_params & base_ctx_params)
{
if(main_model == nullptr || main_model->hparams.n_layer_nextn <= 0)
{
printf("Warning: --usemtp was enabled, but this model does not expose built-in MTP layers. MTP will not be used.\n");
draft_is_mtp = false;
return;
}
llama_context_params mtp_ctx_params = base_ctx_params;
mtp_ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
mtp_ctx_params.ctx_other = main_ctx;
mtp_ctx_params.n_rs_seq = 0;
mtp_ctx_params.n_outputs_max = 1;
printf("\nAttempting to create built-in MTP context from the main model.\n");
draft_ctx = llama_init_from_model(main_model, mtp_ctx_params);
if(draft_ctx == nullptr)
{
printf("Error: failed to create built-in MTP context. MTP will not be used!\n");
draft_is_mtp = false;
return;
}
draft_is_mtp = true;
mtp_speculative_state_setup(main_ctx, mtp_ctx_params, 0);
}
//loads a model for speculative decoding.
static void speculative_decoding_setup(std::string spec_model_filename, llama_context * main_ctx, const llama_model_params & base_model_params, const llama_context_params & base_ctx_params, int base_n_vocab, const float * draft_gpusplit, int draft_gpulayers)
{
@ -755,25 +811,7 @@ static void speculative_decoding_setup(std::string spec_model_filename, llama_co
if(draft_ctx && draft_is_mtp)
{
common_params_speculative spec_params;
spec_params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
spec_params.draft.ctx_tgt = main_ctx;
spec_params.draft.ctx_dft = draft_ctx;
spec_params.draft.n_max = speculative_chunk_amt;
spec_params.draft.n_min = 0;
spec_params.draft.p_min = 0.0f;
spec_params.draft.backend_sampling = false;
spec_params.draft.n_gpu_layers = draft_gpulayers;
spec_params.draft.cache_type_k = draft_ctx_params.type_k;
spec_params.draft.cache_type_v = draft_ctx_params.type_v;
draft_spec = common_speculative_init(spec_params, 1);
if(draft_spec == nullptr)
{
printf("Error: failed to initialize MTP speculative decoding state.\n");
llama_free(draft_ctx);
draft_ctx = nullptr;
draft_is_mtp = false;
}
mtp_speculative_state_setup(main_ctx, draft_ctx_params, draft_gpulayers);
}
}
}
@ -783,6 +821,10 @@ static int32_t kcpp_decode_main_and_mtp_spec(llama_context * main_ctx, llama_bat
const int32_t decode_status = llama_decode(main_ctx, batch);
if(decode_status == 0 && draft_spec && draft_is_mtp)
{
if(draft_ctx && llama_get_ctx_other(draft_ctx) != main_ctx && batch.n_tokens > 0 && batch.n_seq_id[0] > 0)
{
llama_memory_seq_rm(llama_get_memory(draft_ctx), batch.seq_id[0][0], batch.pos[0], -1);
}
if(!common_speculative_process(draft_spec, batch))
{
printf("\nERROR: MTP speculative state update failed!\n");
@ -888,6 +930,22 @@ static speculative_draft_result speculative_decoding_eval_mtp_chunk(llama_contex
real_embd.push_back(drafted_ids[i]);
}
results.verify_tokens.assign(real_embd.begin(), real_embd.end());
results.verify_n_past = n_past;
if(mtp_uses_spec_checkpoint)
{
mtp_spec_ckpt.clear();
mtp_spec_ckpt.update_pos(n_past,
llama_memory_seq_pos_min(llama_get_memory(main_ctx), 0),
llama_memory_seq_pos_max(llama_get_memory(main_ctx), 0));
mtp_spec_ckpt.update_tgt(main_ctx, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
if(draft_ctx)
{
mtp_spec_ckpt.update_dft(draft_ctx, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
}
kcpp_embd_batch batch = kcpp_embd_batch(real_embd, n_past, use_mrope, true);
const int32_t decode_status = kcpp_decode_main_and_mtp_spec(main_ctx, batch.batch);
if(decode_status != 0)
@ -2458,6 +2516,8 @@ ModelLoadResult gpttype_load_model(const load_model_inputs inputs, FileFormat in
}
draft_ctx = nullptr;
draft_is_mtp = false;
mtp_uses_spec_checkpoint = false;
mtp_spec_ckpt.clear();
guidance_ctx = nullptr;
if(mtmd_ctx)
{
@ -3075,17 +3135,36 @@ ModelLoadResult gpttype_load_model(const load_model_inputs inputs, FileFormat in
const llama_vocab * tmpvocab = llama_model_get_vocab(llamamodel);
n_vocab = llama_vocab_n_tokens(tmpvocab);
if(draftmodel_filename !="" && file_format==FileFormat::GGUF_GENERIC)
if((draftmodel_filename != "" || inputs.use_mtp) && file_format==FileFormat::GGUF_GENERIC)
{
if(mtmd_ctx!=nullptr)
{
printf("Error: Speculative decoding cannot be used with multimodal projectors!\n");
printf("Error: Speculative decoding and MTP cannot be used with multimodal projectors!\n");
}
else
{
printf("\nAttempting to load draft model for speculative decoding. It will be fully offloaded if possible. Vocab must match the main model.\n");
speculative_chunk_amt = inputs.draft_amount;
speculative_decoding_setup(draftmodel_filename, llama_ctx_v4, model_params, llama_ctx_params, n_vocab, inputs.draft_gpusplit, inputs.draft_gpulayers);
if(draftmodel_filename != "")
{
if(inputs.use_mtp)
{
printf("\nBoth --draftmodel and --usemtp were provided. The draft model will be used for speculative decoding.\n");
}
printf("\nAttempting to load draft model for speculative decoding. It will be fully offloaded if possible. Vocab must match the main model.\n");
speculative_decoding_setup(draftmodel_filename, llama_ctx_v4, model_params, llama_ctx_params, n_vocab, inputs.draft_gpusplit, inputs.draft_gpulayers);
}
else
{
mtp_decoding_setup(llamamodel, llama_ctx_v4, llama_ctx_params);
}
}
}
if(draft_is_mtp && draft_spec)
{
mtp_uses_spec_checkpoint = common_context_can_seq_rm(llama_ctx_v4) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
if(mtp_uses_spec_checkpoint)
{
printf("\nMTP speculative decoding will use checkpoints for draft mismatch recovery.\n");
}
}
@ -5862,6 +5941,11 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
if (!startedsampling)
{
startedsampling = true;
if(draft_is_mtp && draft_spec)
{
llama_tokens prompt_tokens(current_context_tokens.begin(), current_context_tokens.end());
common_speculative_begin(draft_spec, 0, prompt_tokens);
}
process_time = timer_check();
timer_start();
if(allow_regular_prints)
@ -6226,13 +6310,48 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
logits_sampled += 1;
}
bool mtp_recovered_from_checkpoint = false;
if(draft_used && draft_is_mtp && abort_draft && mtp_uses_spec_checkpoint && !mtp_spec_ckpt.empty())
{
const size_t replay_count = std::min(draft_results.verify_tokens.size(), (size_t) draft_accepted_this_round + 1);
mtp_spec_ckpt.load_tgt(llama_ctx_v4, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
if(draft_ctx)
{
mtp_spec_ckpt.load_dft(draft_ctx, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
if(replay_count > 0)
{
std::vector<int32_t> replay_tokens(
draft_results.verify_tokens.begin(),
draft_results.verify_tokens.begin() + replay_count);
kcpp_embd_batch replay_batch = kcpp_embd_batch(replay_tokens, draft_results.verify_n_past, use_mrope, true);
const int32_t replay_status = kcpp_decode_main_and_mtp_spec(llama_ctx_v4, replay_batch.batch);
if(replay_status != 0)
{
printf("\nERROR: MTP speculative checkpoint replay failed! (code:%d)\n", replay_status);
output.text = nullptr;
output.status = 0;
output.prompt_tokens = output.completion_tokens = 0;
last_stop_reason = stop_reason::ERROR_ENCOUNTERED;
output.stopreason = last_stop_reason;
generation_finished = true;
return output;
}
n_past = draft_results.verify_n_past + replay_tokens.size();
}
mtp_recovered_from_checkpoint = true;
}
if(draft_used && draft_is_mtp && draft_spec)
{
common_speculative_accept(draft_spec, 0, draft_accepted_this_round);
}
//if we have somehow skipped ahead (e.g drafting), ensure that all tokens after npast are purged
if (file_format == FileFormat::GGUF_GENERIC && draft_used)
if (file_format == FileFormat::GGUF_GENERIC && draft_used && !mtp_recovered_from_checkpoint)
{
llama_memory_seq_rm(llama_get_memory(llama_ctx_v4), 0, n_past, -1);
if (draft_ctx) {

13
koboldcpp.py
View file

@ -265,6 +265,7 @@ class load_model_inputs(ctypes.Structure):
("visionmaxtokens", ctypes.c_int),
("use_mmap", ctypes.c_bool),
("use_mlock", ctypes.c_bool),
("use_mtp", ctypes.c_bool),
("use_smartcontext", ctypes.c_bool),
("use_contextshift", ctypes.c_bool),
("use_fastforward", ctypes.c_bool),
@ -1929,6 +1930,7 @@ def load_model(model_filename):
inputs.blasthreads = args.blasthreads
inputs.use_mmap = args.usemmap
inputs.use_mlock = args.usemlock
inputs.use_mtp = args.usemtp
inputs.lora_filename = "".encode("UTF-8")
inputs.lora_multiplier = args.loramult
if args.lora:
@ -7942,6 +7944,7 @@ def show_gui():
draftamount_var = ctk.StringVar(value=str(default_draft_amount))
draftgpulayers_var = ctk.StringVar(value=str(999))
draftgpusplit_str_vars = ctk.StringVar(value="")
usemtp_var = ctk.IntVar(value=0)
nomodel = ctk.IntVar(value=0)
download_dir_var = ctk.StringVar()
@ -8762,9 +8765,10 @@ def show_gui():
makelabelentry(model_tab, "", vision_max_tokens_var, 9, padx=(284),width=36, singleline=True, tooltip="Override the maximum tokens for the MMProj embedding (default -1).", labelpadx=(260))
makecheckbox(model_tab, "V.Force CPU", mmprojcpu_var, 9, padx=340, tooltiptxt="Force CLIP for Vision mmproj always on CPU.")
makefileentry(model_tab, "Draft Model:", "Select Speculative Text Model File", draftmodel_var, 11,width=280,singlerow=True,tooltiptxt="Select a draft text model file to use for speculative decoding.\nLeave blank to skip.")
makelabelentry(model_tab, "Draft Amount: ", draftamount_var, 13, 50,padx=(100),singleline=True,tooltip="How many tokens to draft per chunk before verifying results")
makelabelentry(model_tab, "Splits: ", draftgpusplit_str_vars, 13, 50,padx=(210),singleline=True,tooltip="Distribution of draft model layers. Leave blank to follow main model's gpu split. Only works if multi-gpu (All) selected in main model.", labelpadx=(160))
makelabelentry(model_tab, "Layers: ", draftgpulayers_var, 13, 50,padx=(320),singleline=True,tooltip="How many layers to GPU offload for the draft model", labelpadx=(270))
makelabelentry(model_tab, "Draft Amount: ", draftamount_var, 13, 40,padx=(100),singleline=True,tooltip="How many tokens to draft per chunk before verifying results")
makelabelentry(model_tab, "Splits: ", draftgpusplit_str_vars, 13, 50,padx=(190),singleline=True,tooltip="Distribution of draft model layers. Leave blank to follow main model's gpu split. Only works if multi-gpu (All) selected in main model.", labelpadx=(150))
makelabelentry(model_tab, "Layers: ", draftgpulayers_var, 13, 40,padx=(300),singleline=True,tooltip="How many layers to GPU offload for the draft model", labelpadx=(254))
makecheckbox(model_tab, "Use MTP", usemtp_var, 13, 0,padx=(364),tooltiptxt="Allows using MTP layers for drafting in MTP models.")
makefileentry(model_tab, "Embeds Model:", "Select Embeddings Model File", embeddings_model_var, 15, width=130,singlerow=True, filetypes=[("*.gguf","*.gguf")], tooltiptxt="Select an embeddings GGUF model that can be used to generate embedding vectors.")
makelabelentry(model_tab, "ECtx: ", embeddings_ctx_var, 15, 50,padx=(335),singleline=True,tooltip="If set above 0, limits max context for embedding model to save memory.", labelpadx=(302))
makecheckbox(model_tab, "GPU", embeddings_gpu_var, 15, 0,padx=(390),tooltiptxt="Uses the GPU for Embeddings.")
@ -9241,6 +9245,7 @@ def show_gui():
args.draftmodel = None if draftmodel_var.get() == "" else draftmodel_var.get()
args.draftamount = int(draftamount_var.get()) if draftamount_var.get()!="" else default_draft_amount
args.draftgpulayers = int(draftgpulayers_var.get()) if draftgpulayers_var.get()!="" else 999
args.usemtp = usemtp_var.get() == 1
args.ssl = None if (ssl_cert_var.get() == "" or ssl_key_var.get() == "") else ([ssl_cert_var.get(), ssl_key_var.get()])
args.password = None if (password_var.get() == "") else (password_var.get())
@ -9535,6 +9540,7 @@ def show_gui():
draftamount_var.set(mydict["draftamount"])
if "draftgpulayers" in mydict:
draftgpulayers_var.set(mydict["draftgpulayers"])
usemtp_var.set(1 if "usemtp" in mydict and mydict["usemtp"] else 0)
ssl_cert_var.set("")
ssl_key_var.set("")
@ -12100,6 +12106,7 @@ if __name__ == '__main__':
advparser.add_argument("--ssl", help="Allows all content to be served over SSL instead. A valid UNENCRYPTED SSL cert and key .pem files must be provided", metavar=('[cert_pem]', '[key_pem]'), nargs='+')
advparser.add_argument("--swapadding", help="How much extra to pad the SWA KV cache, this affects the rewind limit before reprocessing is forced.", type=int, default=swa_padding_default)
advparser.add_argument("--unpack", help="Extracts the file contents of the KoboldCpp binary into a target directory.", metavar=('destination'), type=str, default="")
advparser.add_argument("--usemtp", help="Enables MTP layers to be used for drafting (speculative decoding) if present", action='store_true')
advparser.add_argument("--usemlock","--mlock", help="Enables mlock, preventing the RAM used to load the model from being paged out. Not usually recommended.", action='store_true')
compatgroup3 = advparser.add_mutually_exclusive_group()
compatgroup3.add_argument("--usemmap", help="If set, uses mmap to load model.", action='store_true')

2
otherarch/otherarch.h
View file

@ -529,9 +529,11 @@ struct media_object
struct speculative_draft_result
{
std::vector<int32_t> draftids;
std::vector<int32_t> verify_tokens;
std::vector<float *> actual_logits;
bool draft_success = false;
int drafted_amount = 0;
int verify_n_past = 0;
};
struct savestate_data