File size: 6,166 Bytes
12eb3a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | // EPQ target with Qwen3.6 MTP draft context.
#include "llama.h"
#include "sdq_pager.h"
#include "common.h"
#include "sampling.h"
#include "speculative.h"
#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
using namespace sdq;
namespace {
struct opt_t { std::string model, prompt; int ctx=4096, threads=8, n=48, draft_n=2, ram=8192, ubatch=128; bool no_pf=false; };
bool parse(int argc, char ** argv, opt_t & o) {
auto next = [&](int & i) -> const char * { return ++i < argc ? argv[i] : nullptr; };
for (int i=1; i<argc; i++) {
const std::string a=argv[i]; const char * v=nullptr;
if (a=="-m" || a=="--model") { if (!(v=next(i))) return false; o.model=v; }
else if (a=="-p" || a=="--prompt") { if (!(v=next(i))) return false; o.prompt=v; }
else if (a=="-n") { if (!(v=next(i))) return false; o.n=atoi(v); }
else if (a=="--draft-n") { if (!(v=next(i))) return false; o.draft_n=atoi(v); }
else if (a=="-c") { if (!(v=next(i))) return false; o.ctx=atoi(v); }
else if (a=="-t") { if (!(v=next(i))) return false; o.threads=atoi(v); }
else if (a=="-b") { if (!(v=next(i))) return false; o.ubatch=atoi(v); }
else if (a=="--ram-budget-mb") { if (!(v=next(i))) return false; o.ram=atoi(v); }
else if (a=="--no-prefetch") o.no_pf=true;
else return false;
}
return !o.model.empty() && !o.prompt.empty();
}
double now_s() { return std::chrono::duration<double>(std::chrono::steady_clock::now().time_since_epoch()).count(); }
}
int main(int argc, char ** argv) {
opt_t o;
if (!parse(argc, argv, o)) return 1;
setenv("SDQ_MODEL_PATH", o.model.c_str(), 1);
setenv("SDQ_IO_THREADS", std::to_string(o.threads*3).c_str(), 1);
setenv("SDQ_RAM_BUDGET_MB", std::to_string(o.ram).c_str(), 1);
setenv("SDQ_FADV_DONTNEED", "1", 1);
if (o.no_pf) setenv("SDQ_PREFETCH", "0", 1);
llama_backend_init();
llama_numa_init(GGML_NUMA_STRATEGY_DISABLED);
llama_model_params mp=llama_model_default_params();
mp.n_gpu_layers=0; mp.load_mode=LLAMA_LOAD_MODE_MMAP; mp.load_mtp=true;
sdq_set_enabled(true); sdq_install_buft_overrides(&mp);
llama_model * model=llama_model_load_from_file(o.model.c_str(), mp);
if (!model || !sdq_init(o.model)) return 2;
llama_context_params c=llama_context_default_params();
c.n_ctx=o.ctx; c.n_seq_max=1; c.n_batch=std::max(o.ubatch,o.draft_n+1); c.n_ubatch=o.ubatch;
c.n_outputs_max=o.draft_n+1; c.n_outputs_max_per_seq=o.draft_n+1;
c.n_threads=o.threads; c.n_threads_batch=o.threads; c.ctx_type=LLAMA_CONTEXT_TYPE_DEFAULT;
llama_context * ctx=llama_init_from_model(model,c);
if (!ctx) return 3;
common_params p;
p.model.path=o.model; p.n_ctx=o.ctx; p.n_batch=c.n_batch; p.n_ubatch=o.ubatch;
p.cpuparams.n_threads=o.threads; p.cpuparams_batch.n_threads=o.threads;
p.speculative.types={COMMON_SPECULATIVE_TYPE_DRAFT_MTP};
p.speculative.draft.n_max=o.draft_n;
auto init=common_speculative_init_from_params(p,model,ctx);
llama_context * dctx=init->context();
p.speculative.draft.ctx_tgt=ctx; p.speculative.draft.ctx_dft=dctx;
if (!dctx) { fprintf(stderr,"[sdq-mtp] MTP context init failed\n"); return 4; }
common_speculative * spec=common_speculative_init(p.speculative,1);
if (!spec) { fprintf(stderr,"[sdq-mtp] speculative init failed\n"); return 5; }
common_params_sampling sp; sp.temp=0; sp.top_k=1;
common_sampler * sampler=common_sampler_init(model,sp);
auto input=common_tokenize(ctx,o.prompt,true,true);
if (input.empty()) return 6;
common_batch pb(ctx);
for (size_t i=0;i+1<input.size();i++) pb.add(input[i],(llama_pos)i,0,false);
if (pb.size() && llama_process(ctx,LLAMA_PROCESS_TYPE_DECODE,pb.get())!=0) return 7;
pager::instance().end_token(); common_speculative_process(spec,pb);
llama_tokens prompt(input.begin(),input.end()-1); common_speculative_begin(spec,0,prompt);
llama_token last=input.back(); int past=(int)prompt.size(), gen=0, drafted=0, accepted=0;
common_prompt_checkpoint ck; llama_tokens draft;
const double t0=now_s(); printf("assistant: "); fflush(stdout);
while (gen<o.n) {
ck.update_pos(prompt.size(),llama_memory_seq_pos_min(llama_get_memory(ctx),0),llama_memory_seq_pos_max(llama_get_memory(ctx),0));
ck.update_tgt(ctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ck.update_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
if (draft.empty()) { auto & dp=common_speculative_get_draft_params(spec,0); dp={}; dp.drafting=true; dp.n_max=std::min(o.draft_n,o.n-gen-1); dp.pos0=past; dp.id_last=last; dp.prompt=&prompt; dp.result=&draft; common_speculative_draft(spec); }
ck.load_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
common_batch vb(ctx); vb.add(last,past++,0,true); for(size_t i=0;i<draft.size();i++) vb.add(draft[i],past+(llama_pos)i,0,true);
if (llama_process(ctx,LLAMA_PROCESS_TYPE_DECODE,vb.get())!=0) break;
pager::instance().end_token();
if(!common_speculative_process(spec,vb)) break;
auto ids=common_sampler_sample_and_accept_n(sampler,ctx,draft); if(ids.empty()) break;
drafted+=(int)draft.size(); accepted+=(int)ids.size()-1;
if(ids.size()-1<draft.size()) { ck.load_tgt(ctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ck.load_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); prompt.resize(ck.n_tokens); past=ck.n_tokens; draft=std::move(ids); continue; }
for(auto id:ids) { if(gen>=o.n) break; printf("%s",common_token_to_piece(ctx,id).c_str()); gen++; prompt.push_back(last); last=id; }
common_speculative_accept(spec,0,(uint16_t)(ids.size()-1)); past+=(int)ids.size()-1;
}
printf("\n"); double dt=now_s(); fprintf(stderr,"[sdq-mtp] generated=%d drafted=%d accepted=%d acceptance=%.1f%% speed=%.2f tok/s\n",gen,drafted,accepted,drafted?100.0*accepted/drafted:0,dt-t0?gen/(dt-t0):0);
common_speculative_print_stats(spec); common_speculative_free(spec); common_sampler_free(sampler); sdq_shutdown(); llama_free(ctx); llama_model_free(model); llama_backend_free(); return 0;
}
|