| |
| #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; |
| } |
|
|
