Something went wrong. Try again.
llama.cpp fork for Ternary Bonsai 2 27B: MTP speculation, suffix cache reuse (SCR), q8 KV, two-slot fairness, fused DeltaNet kernels
Something went wrong. Try again.
166 kB · 4624 lines
C++
at main
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305330633073308330933103311331233133314331533163317331833193320332133223323332433253326332733283329333033313332333333343335333633373338333933403341334233433344334533463347334833493350335133523353335433553356335733583359336033613362336333643365336633673368336933703371337233733374337533763377337833793380338133823383338433853386338733883389339033913392339333943395339633973398339934003401340234033404340534063407340834093410341134123413341434153416341734183419342034213422342334243425342634273428342934303431343234333434343534363437343834393440344134423443344434453446344734483449345034513452345334543455345634573458345934603461346234633464346534663467346834693470347134723473347434753476347734783479348034813482348334843485348634873488348934903491349234933494349534963497349834993500350135023503350435053506350735083509351035113512351335143515351635173518351935203521352235233524352535263527352835293530353135323533353435353536353735383539354035413542354335443545354635473548354935503551355235533554355535563557355835593560356135623563356435653566356735683569357035713572357335743575357635773578357935803581358235833584358535863587358835893590359135923593359435953596359735983599360036013602360336043605360636073608360936103611361236133614361536163617361836193620362136223623362436253626362736283629363036313632363336343635363636373638363936403641364236433644364536463647364836493650365136523653365436553656365736583659366036613662366336643665366636673668366936703671367236733674367536763677367836793680368136823683368436853686368736883689369036913692369336943695369636973698369937003701370237033704370537063707370837093710371137123713371437153716371737183719372037213722372337243725372637273728372937303731373237333734373537363737373837393740374137423743374437453746374737483749375037513752375337543755375637573758375937603761376237633764376537663767376837693770377137723773377437753776377737783779378037813782378337843785378637873788378937903791379237933794379537963797379837993800380138023803380438053806380738083809381038113812381338143815381638173818381938203821382238233824382538263827382838293830383138323833383438353836383738383839384038413842384338443845384638473848384938503851385238533854385538563857385838593860386138623863386438653866386738683869387038713872387338743875387638773878387938803881388238833884388538863887388838893890389138923893389438953896389738983899390039013902390339043905390639073908390939103911391239133914391539163917391839193920392139223923392439253926392739283929393039313932393339343935393639373938393939403941394239433944394539463947394839493950395139523953395439553956395739583959396039613962396339643965396639673968396939703971397239733974397539763977397839793980398139823983398439853986398739883989399039913992399339943995399639973998399940004001400240034004400540064007400840094010401140124013401440154016401740184019402040214022402340244025402640274028402940304031403240334034403540364037403840394040404140424043404440454046404740484049405040514052405340544055405640574058405940604061406240634064406540664067406840694070407140724073407440754076407740784079408040814082408340844085408640874088408940904091409240934094409540964097409840994100410141024103410441054106410741084109411041114112411341144115411641174118411941204121412241234124412541264127412841294130413141324133413441354136413741384139414041414142414341444145414641474148414941504151415241534154415541564157415841594160416141624163416441654166416741684169417041714172417341744175417641774178417941804181418241834184418541864187418841894190419141924193419441954196419741984199420042014202420342044205420642074208420942104211421242134214421542164217421842194220422142224223422442254226422742284229423042314232423342344235423642374238423942404241424242434244424542464247424842494250425142524253425442554256425742584259426042614262426342644265426642674268426942704271427242734274427542764277427842794280428142824283428442854286428742884289429042914292429342944295429642974298429943004301430243034304430543064307430843094310431143124313431443154316431743184319432043214322432343244325432643274328432943304331433243334334433543364337433843394340434143424343434443454346434743484349435043514352435343544355435643574358435943604361436243634364436543664367436843694370437143724373437443754376437743784379438043814382438343844385438643874388438943904391439243934394439543964397439843994400440144024403440444054406440744084409441044114412441344144415441644174418441944204421442244234424442544264427442844294430443144324433443444354436443744384439444044414442444344444445444644474448444944504451445244534454445544564457445844594460446144624463446444654466446744684469447044714472447344744475447644774478447944804481448244834484448544864487448844894490449144924493449444954496449744984499450045014502450345044505450645074508450945104511451245134514451545164517451845194520452145224523452445254526452745284529453045314532453345344535453645374538453945404541454245434544454545464547454845494550455145524553455445554556455745584559456045614562456345644565456645674568456945704571457245734574457545764577457845794580458145824583458445854586458745884589459045914592459345944595459645974598459946004601460246034604460546064607460846094610461146124613461446154616461746184619462046214622462346244625#include "llama-context.h"
#include "ggml.h"#include "llama-arch.h"#include "llama-batch.h"#include "llama-ext.h"#include "llama-graph.h"#include "llama-impl.h"#include "llama-io.h"#include "llama-kv-cache-iswa.h"#include "llama-kv-cache.h"#include "llama-memory-hybrid-iswa.h"#include "llama-memory-hybrid.h"#include "llama-memory.h"#include "llama-mmap.h"#include "llama-model.h"#include "llama-sampler.h"#include "llama.h"
#include <cinttypes>#include <cmath>#include <cstring>#include <limits>#include <numeric>#include <stdexcept>#include <string>
//// llama_context//
// Verify that every Hadamard-folded weight consumed by the graph receives its// activation-side transform, and every latent lookup table gets the inverse.// An architecture whose matmul path bypasses the transform helpers would// otherwise load cleanly and silently compute wrong results.static void llama_verify_hadamard_graph( ggml_cgraph * gf, const llama_hadamard_rotations & rotations, const llama_hadamard_rotations & inverses) { auto unwrap = [](const ggml_tensor * t) { while (t && (t->op == GGML_OP_RESHAPE || t->op == GGML_OP_VIEW)) { t = t->src[0]; } return t; };
std::map<const ggml_tensor *, bool> lookups; // get_rows results of latent tables
for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { const ggml_tensor * node = ggml_graph_node(gf, i);
if (node->op == GGML_OP_GET_ROWS && inverses.count(node->src[0])) { lookups.emplace(node, false); continue; }
if (node->op != GGML_OP_MUL_MAT && node->op != GGML_OP_MUL_MAT_ID) { continue; }
if (node->op == GGML_OP_MUL_MAT && ((const int32_t *) node->op_params)[1] == GGML_HINT_SRC0_IS_HADAMARD) { const auto lk = lookups.find(unwrap(node->src[1])); if (lk != lookups.end()) { lk->second = true; } continue; }
const auto it = rotations.find(node->src[0]); if (it == rotations.end()) { if (inverses.count(node->src[0])) { throw std::runtime_error(format( "Hadamard-latent table '%s' is used as a head without a forward transform", node->src[0]->name)); } continue; } const ggml_tensor * src = unwrap(node->src[1]); const bool transformed = src && src->op == GGML_OP_MUL_MAT && ((const int32_t *) src->op_params)[1] == GGML_HINT_SRC0_IS_HADAMARD && src->src[0] == it->second.rot; if (!transformed) { throw std::runtime_error(format( "Hadamard-folded weight '%s' is consumed without its activation transform; " "this graph's matmul path does not support prism.hadamard folding", node->src[0]->name)); } }
for (const auto & [node, ok] : lookups) { if (!ok) { for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { const ggml_tensor * n2 = ggml_graph_node(gf, i); for (int s = 0; s < GGML_MAX_SRC && n2->src[s]; ++s) { if (unwrap(n2->src[s]) == node) { LLAMA_LOG_WARN("%s: latent lookup '%s' consumed by op=%s name='%s' src%d hint=%d\n", __func__, node->name, ggml_op_name(n2->op), n2->name, s, ((const int32_t *) n2->op_params)[1]); } } } throw std::runtime_error(format( "Hadamard-latent table '%s' is read without the inverse transform", node->src[0]->name)); } }}
static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { switch (ctx_type) { case LLAMA_CONTEXT_TYPE_DEFAULT: return LLM_GRAPH_TYPE_DEFAULT; case LLAMA_CONTEXT_TYPE_MTP : return LLM_GRAPH_TYPE_DECODER_MTP; } throw std::runtime_error("Unsupported ctx type");}
struct llm_fused_op_probe { llm_fused_op op; const char * name; uint32_t n_tokens_per_seq;};
static const llm_fused_op_probe llm_fused_op_flash_attn_probe = { /*.op =*/ LLM_FUSED_OP_FLASH_ATTN, /*.name =*/ "Flash Attention", /*.n_tokens_per_seq =*/ 1,};
static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = { /*.op =*/ LLM_FUSED_OP_GDN_AR, /*.name =*/ "fused Gated Delta Net (autoregressive)", /*.n_tokens_per_seq =*/ 1,};
static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { /*.op =*/ LLM_FUSED_OP_GDN_CH, /*.name =*/ "fused Gated Delta Net (chunked)", /*.n_tokens_per_seq =*/ 16,};
static const llm_fused_op_probe llm_fused_op_lid_probe = { /*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER, /*.name =*/ "Lightning Indexer", /*.n_tokens_per_seq =*/ 1,};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = { /*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE, /*.name =*/ "fused DeepSeek V4 HC pre", /*.n_tokens_per_seq =*/ 1,};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_comb_probe = { /*.op =*/ LLM_FUSED_OP_DSV4_HC_COMB, /*.name =*/ "fused DeepSeek V4 HC comb", /*.n_tokens_per_seq =*/ 1,};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = { /*.op =*/ LLM_FUSED_OP_DSV4_HC_POST, /*.name =*/ "fused DeepSeek V4 HC post", /*.n_tokens_per_seq =*/ 1,};
llama_context::llama_context( const llama_model & model, llama_context_params params) : model(model), cvec(std::make_unique<llama_adapter_cvec>()), loras(std::make_unique<llama_adapter_loras>()), balloc(std::make_unique<llama_batch_allocr>(model.hparams.n_pos_per_embd())) { // TODO warning when creating llama_context with awkward ctx size that is not a power of 2, // may need to be backend-dependent LLAMA_LOG_INFO("%s: constructing llama_context\n", __func__);
t_start_us = model.t_start_us; t_load_us = model.t_load_us;
const auto & hparams = model.hparams;
cparams.n_seq_max = std::max(1u, params.n_seq_max); if (cparams.n_seq_max > LLAMA_MAX_SEQ) { throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); }
cparams.n_rs_seq = params.n_rs_seq; if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", __func__, cparams.n_rs_seq); cparams.n_rs_seq = 0; }
cparams.n_threads = params.n_threads; cparams.n_threads_batch = params.n_threads_batch; cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor; cparams.yarn_attn_factor = params.yarn_attn_factor >= 0.0f ? params.yarn_attn_factor : hparams.yarn_attn_factor; cparams.yarn_beta_fast = params.yarn_beta_fast >= 0.0f ? params.yarn_beta_fast : hparams.yarn_beta_fast; cparams.yarn_beta_slow = params.yarn_beta_slow >= 0.0f ? params.yarn_beta_slow : hparams.yarn_beta_slow; cparams.embeddings = params.embeddings; cparams.embeddings_nextn = false; cparams.embeddings_nextn_masked = false; cparams.offload_kqv = params.offload_kqv; cparams.no_perf = params.no_perf; cparams.warmup = false;
// +1: id n_layer() taps the output of the last layer ("input" of the head) cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false); embd_layer_inp.resize(hparams.n_layer() + 1);
cparams.ctx_type = params.ctx_type; cparams.pooling_type = params.pooling_type;
cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx; cparams.rope_freq_base = params.rope_freq_base == 0.0f ? hparams.rope_freq_base_train : params.rope_freq_base; cparams.rope_freq_scale = params.rope_freq_scale == 0.0f ? hparams.rope_freq_scale_train : params.rope_freq_scale;
cparams.n_ctx_orig_yarn = params.yarn_orig_ctx != 0 ? params.yarn_orig_ctx : hparams.n_ctx_orig_yarn != 0 ? hparams.n_ctx_orig_yarn : hparams.n_ctx_train;
cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data;
cparams.ctx_other = nullptr;
// TODO: more generic if (model.arch == LLM_ARCH_GEMMA4_ASSISTANT) { if (params.ctx_other == nullptr) { // TODO: change from runtime_error to llama_exception to avoid printing error message throw std::runtime_error("Gemma4Assistant requires ctx_other to be set (this warning is normal during memory fitting)"); }
cparams.ctx_other = params.ctx_other; }
if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) { if (model.tok_embd == nullptr || model.output == nullptr) { if (params.ctx_other == nullptr) { throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)"); } cparams.ctx_other = params.ctx_other; } }
hadamard_rotations = model.hadamard_rotations; hadamard_inverses = model.hadamard_inverses; if (cparams.ctx_other) { // Transform entries are keyed by tensor pointer, so borrowed target tensors remain distinct. const auto & other = cparams.ctx_other->model; hadamard_rotations.insert(other.hadamard_rotations.begin(), other.hadamard_rotations.end()); hadamard_inverses .insert(other.hadamard_inverses .begin(), other.hadamard_inverses .end()); }
auto rope_scaling_type = params.rope_scaling_type; if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) { rope_scaling_type = hparams.rope_scaling_type_train; }
if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_NONE) { cparams.rope_freq_scale = 1.0f; // never scale if scaling type is none }
if (cparams.yarn_ext_factor < 0.0f) { // negative indicates 'not set' cparams.yarn_ext_factor = rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_YARN ? 1.0f : 0.0f; }
if (cparams.yarn_ext_factor != 0) { static auto get_mscale = [](float scale, float mscale) { return scale <= 1.0f ? 1.0f : (0.1f * mscale * logf(scale) + 1.0f); };
const float factor = 1.0f / cparams.rope_freq_scale;
// ref: https://github.com/huggingface/transformers/blob/6d00f6b0a5679c36510f203e4226e36f517c3032/src/transformers/modeling_rope_utils.py#L336-L348 if (hparams.rope_yarn_log_mul != 0.0f) { // note: here we assume `mscale == 1.0f` // TODO: start reading the actual value of mscale and handle the case where it is not 1.0f float mscale = 1.0f; const float mscale_all_dims = hparams.rope_yarn_log_mul;
// [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX] // special-case DEEPSEEK v2: // https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite-Chat/blob/main/config.json#L42-L43 if (model.arch == LLM_ARCH_DEEPSEEK2 && mscale_all_dims != 1.0f) { mscale = mscale_all_dims; }
cparams.yarn_attn_factor = get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dims);
LLAMA_LOG_WARN("%s: setting new yarn_attn_factor = %.4f (mscale == %.1f, mscale_all_dim = %.1f)\n", __func__, cparams.yarn_attn_factor, mscale, mscale_all_dims); } else { cparams.yarn_attn_factor = get_mscale(factor, 1.0f); }
// when YARN is applied with yarn_ext_factor != 0.0f, we need to cancel this factor: // https://github.com/ggml-org/llama.cpp/blob/a81a569577cc38b32558958b048228150be63eae/ggml/src/ggml-cpu/ops.cpp#L5541-L5544 // // ref: https://github.com/ggml-org/llama.cpp/discussions/7416 // https://github.com/ggml-org/llama.cpp/pull/17945 cparams.yarn_attn_factor *= 1.0f / (1.0f + 0.1f * logf(factor)); }
cparams.yarn_attn_factor *= hparams.rope_attn_factor;
if (cparams.pooling_type == LLAMA_POOLING_TYPE_UNSPECIFIED) { if (hparams.pooling_type == LLAMA_POOLING_TYPE_UNSPECIFIED) { cparams.pooling_type = LLAMA_POOLING_TYPE_NONE; } else { cparams.pooling_type = hparams.pooling_type; } }
if (params.attention_type == LLAMA_ATTENTION_TYPE_UNSPECIFIED) { cparams.causal_attn = hparams.causal_attn; } else { cparams.causal_attn = params.attention_type == LLAMA_ATTENTION_TYPE_CAUSAL; }
cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO;
cparams.fused_gdn_ar = true; cparams.fused_gdn_ch = true; cparams.auto_fgdn = true;
cparams.fused_lid = true; cparams.auto_flid = true;
cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; cparams.fused_dsv4_hc_post = true; cparams.auto_fhc = true;
// with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max; cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ? cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max);
// Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later // re-reserve when graph nodes change. if (params.samplers != nullptr && params.n_samplers > 0) { for (size_t i = 0; i < params.n_samplers; ++i) { const auto & config = params.samplers[i];
if (llama_sampler_chain_get(config.sampler, -1) == nullptr) { throw std::runtime_error("the backend samplers must be of type llama_sampler_chain"); }
if (set_sampler(config.seq_id, config.sampler)) { const int n_samplers = llama_sampler_chain_n(config.sampler);
LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers); } } }
cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified;
// initialized later cparams.pipeline_parallel = false;
{ const char * LLAMA_GRAPH_REUSE_DISABLE = getenv("LLAMA_GRAPH_REUSE_DISABLE"); graph_reuse_disable = LLAMA_GRAPH_REUSE_DISABLE ? (atoi(LLAMA_GRAPH_REUSE_DISABLE) != 0) : graph_reuse_disable;
if (graph_reuse_disable) { LLAMA_LOG_WARN("%s: graph reuse disabled\n", __func__); } }
// ref: https://github.com/ggml-org/llama.cpp/pull/17046#discussion_r2503085732 cparams.n_ctx = GGML_PAD(cparams.n_ctx, 256);
if (cparams.kv_unified) { cparams.n_ctx_seq = cparams.n_ctx; } else { cparams.n_ctx_seq = cparams.n_ctx / cparams.n_seq_max; cparams.n_ctx_seq = GGML_PAD(cparams.n_ctx_seq, 256);
if (cparams.n_ctx_seq == 0) { throw std::runtime_error("n_ctx_seq == 0"); }
if (cparams.n_ctx != cparams.n_ctx_seq * cparams.n_seq_max) { cparams.n_ctx = cparams.n_ctx_seq * cparams.n_seq_max; LLAMA_LOG_WARN("%s: n_ctx is not divisible by n_seq_max - rounding down to %u\n", __func__, cparams.n_ctx); } }
LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max); LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq);
if (cparams.n_ctx_seq < hparams.n_ctx_train) { LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", __func__, cparams.n_ctx_seq, hparams.n_ctx_train); }
if (cparams.n_ctx_seq > hparams.n_ctx_train) { LLAMA_LOG_WARN("%s: n_ctx_seq (%u) > n_ctx_train (%u) -- possible training context overflow\n", __func__, cparams.n_ctx_seq, hparams.n_ctx_train); }
if (!hparams.vocab_only) { // GPU backends for (const auto & dev : model.devices) { ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr); if (backend == nullptr) { throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev))); } backends.emplace_back(backend); }
// add ACCEL backends (such as BLAS) for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_ACCEL) { ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); if (backend == nullptr) { throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev))); } backends.emplace_back(backend); } }
// add CPU backend backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); if (backend_cpu == nullptr) { throw std::runtime_error("failed to initialize CPU backend"); } backends.emplace_back(backend_cpu);
// create a list of the set_n_threads functions in the backends for (auto & backend : backends) { ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; if (reg) { auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads"); if (ggml_backend_set_n_threads_fn) { set_n_threads_fns.emplace_back(backend.get(), ggml_backend_set_n_threads_fn); } } }
llama_set_abort_callback(this, params.abort_callback, params.abort_callback_data);
// graph outputs buffer { if (output_reserve(params.n_seq_max) < params.n_seq_max) { throw std::runtime_error("failed to reserve initial output buffer"); }
LLAMA_LOG_INFO("%s: %10s output buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name (buf_output.get()), ggml_backend_buffer_get_size(buf_output.get()) / 1024.0 / 1024.0); } }
// init the memory module if (!hparams.vocab_only) { llama_memory_params params_mem = { /*.type_k =*/ params.type_k, /*.type_v =*/ params.type_v, /*.swa_full =*/ params.swa_full, /*.ctx_type =*/ cparams.ctx_type, /*.mem_other =*/ llama_get_memory(cparams.ctx_other), };
memory.reset(model.create_memory(params_mem, cparams));
if (params.path_kv_mean_center != nullptr) { // collect every standard llama_kv_cache the memory module keeps for attention // layers. hybrid (recurrent + attention) models keep one for their attention // sublayers; SWA variants keep a base/SWA pair. bias tensors are matched by // model layer id, so layers absent from a given cache are simply skipped. std::vector<llama_kv_cache *> kvs; if (auto * kv = dynamic_cast<llama_kv_cache *>(memory.get())) { kvs.push_back(kv); } else if (auto * kv_iswa = dynamic_cast<llama_kv_cache_iswa *>(memory.get())) { kvs.push_back(kv_iswa->get_base()); kvs.push_back(kv_iswa->get_swa()); } else if (auto * hyb = dynamic_cast<llama_memory_hybrid *>(memory.get())) { kvs.push_back(hyb->get_mem_attn()); } else if (auto * hyb_iswa = dynamic_cast<llama_memory_hybrid_iswa *>(memory.get())) { kvs.push_back(hyb_iswa->get_mem_attn()->get_base()); kvs.push_back(hyb_iswa->get_mem_attn()->get_swa()); } if (kvs.empty()) { throw std::runtime_error("path_kv_mean_center is only supported for standard KV caches " "(not recurrent-only or MLA/DSA memory types)"); } for (auto * kv : kvs) { if (!kv->load_kv_mean_center(params.path_kv_mean_center)) { throw std::runtime_error("failed to load K-cache mean-centering bias file"); } } } }
// init backends if (!hparams.vocab_only) { LLAMA_LOG_DEBUG("%s: enumerating backends\n", __func__);
backend_buft.clear(); backend_ptrs.clear(); backend_buf_exp_size.clear();
for (auto & backend : backends) { auto * buft = ggml_backend_get_default_buffer_type(backend.get()); auto backend_type = ggml_backend_dev_type(ggml_backend_get_device(backend.get()));
if (backend_type == GGML_BACKEND_DEVICE_TYPE_CPU && !model.devices.empty()) { // use the host buffer of the first device CPU for faster transfer of the intermediate state const auto & dev = model.devices[0]; auto * host_buft = ggml_backend_dev_host_buffer_type(dev.dev); if (host_buft) { buft = host_buft; } }
backend_buft.push_back(buft); backend_ptrs.push_back(backend.get()); backend_buf_exp_size.push_back(0); }
LLAMA_LOG_DEBUG("%s: backend_ptrs.size() = %zu\n", __func__, backend_ptrs.size());
// TODO: move these checks to ggml_backend_sched // enabling pipeline parallelism in the scheduler increases memory usage, so it is only done when necessary bool pipeline_parallel = model.n_devices() > 1 && model.n_gpu_layers() > model.hparams.n_layer_all && model.split_mode() == LLAMA_SPLIT_MODE_LAYER && cparams.offload_kqv && !model.has_tensor_overrides();
// pipeline parallelism requires support for async compute and events in all devices if (pipeline_parallel) { for (auto & backend : backends) { auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backend.get())); if (dev_type == GGML_BACKEND_DEVICE_TYPE_CPU) { // ignore CPU backend // TODO: should we ignore ACCEL types too? continue; } auto * dev = ggml_backend_get_device(backend.get()); ggml_backend_dev_props props; ggml_backend_dev_get_props(dev, &props); if (!props.caps.async || !props.caps.events) { // device does not support async compute or events pipeline_parallel = false; break; } } }
cparams.pipeline_parallel = pipeline_parallel;
if (cparams.pipeline_parallel) { LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__); }
sched_reserve();
if (!cparams.flash_attn) { if (ggml_is_quantized(params.type_v)) { throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention"); } } }
// Initialize the full vocabulary token ids for backend samplers. { const int n_vocab = model.vocab.n_tokens();
sampling.token_ids_full_vocab.resize(n_vocab); for (int i = 0; i < n_vocab; ++i) { sampling.token_ids_full_vocab[i] = i; } }}
llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize();
if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; ggml_backend_buffer_type_t buft = backend_buft[i];
const size_t size_exp = backend_buf_exp_size[i]; const size_t size_act = ggml_backend_sched_get_buffer_size(sched.get(), backend); if (size_exp == size_act) { LLAMA_LOG_DEBUG("%s: %10s compute buffer size is %8.4f MiB, matches expectation of %8.4f MiB\n", __func__, ggml_backend_buft_name(buft), size_act / (1024.0*1024.0), size_exp / (1024.0*1024.0)); } else { LLAMA_LOG_WARN("%s: %10s compute buffer size of %8.4f MiB, does not match expectation of %8.4f MiB\n", __func__, ggml_backend_buft_name(buft), size_act / (1024.0*1024.0), size_exp / (1024.0*1024.0)); } } } ggml_opt_free(opt_ctx);}
void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) { const char * func = __func__; auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) { if (!enabled) { return; }
const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs;
auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true); if (!gf) { throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check"); }
bool device_mismatch = false; for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) { if (node.op != probe.op) { continue; }
GGML_ASSERT(node.il >= 0);
ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor); ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr;
// TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior, // but is still wrong for cases like --no-kv-offload. ggml_backend_dev_t device_layer = model.dev_layer(node.il);
if (device_fused != device_layer) { LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s " "is assigned to device %s (usually due to missing support)\n", func, node.il, device_layer ? ggml_backend_dev_name(device_layer) : "none", probe.name, device_fused ? ggml_backend_dev_name(device_fused) : "none"); device_mismatch = true; break; } }
if (device_mismatch) { enabled = false; LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name); } else { enabled = true; LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name); } };
if (cparams.auto_fa) { resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn); cparams.auto_fa = false; }
if (cparams.auto_fgdn) { LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func); resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar); resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch); cparams.auto_fgdn = false; }
if (cparams.auto_flid) { LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func); resolve(llm_fused_op_lid_probe, cparams.fused_lid); cparams.auto_flid = false; }
if (cparams.auto_fhc) { LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 HC support:\n", func); resolve(llm_fused_op_dsv4_hc_pre_probe, cparams.fused_dsv4_hc_pre); resolve(llm_fused_op_dsv4_hc_comb_probe, cparams.fused_dsv4_hc_comb); resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post); cparams.auto_fhc = false; }}
// names of the samplers in a chain: two chains with the same names add the same ops to the graphstatic std::string llama_sampler_chain_sig(llama_sampler * chain) { std::string sig; for (int i = 0; i < llama_sampler_chain_n(chain); ++i) { sig += llama_sampler_name(llama_sampler_chain_get(chain, i)); sig += ';'; } return sig;}
void llama_context::sched_reserve() { if (!sched_need_reserve) { return; }
sched_need_reserve = false;
for (const auto & [seq_id, sampler] : sampling.samplers) { sampler_sigs_reserved[seq_id] = llama_sampler_chain_sig(sampler); }
LLAMA_LOG_INFO("%s: reserving ...\n", __func__);
synchronize();
const int64_t t_start_us = ggml_time_us();
const uint32_t n_seqs = cparams.n_seq_max; const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch);
const size_t max_nodes = this->graph_max_nodes(n_tokens);
LLAMA_LOG_DEBUG("%s: max_nodes = %zu\n", __func__, max_nodes);
gf_res_prev.reset(new llm_graph_result(max_nodes)); gf_res_reserve.reset(new llm_graph_result(max_nodes));
sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
llama_memory_context_ptr mctx; if (memory) { LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__); mctx = memory->init_full(); if (!mctx) { throw std::runtime_error("failed to initialize memory module"); } }
// avoid reserving graphs with zero outputs - assume one output per sequence const int n_outputs = n_seqs;
LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs);
resolve_fused_ops(mctx.get(), n_seqs);
// reserve worst-case graph int n_splits_pp = -1; int n_nodes_pp = -1;
int n_splits_tg = -1; int n_nodes_tg = -1;
const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max);
// reserve pp (prompt processing) graph first so that buffers are only allocated once { auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr); if (!gf) { if (cparams.pipeline_parallel) { LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__); cparams.pipeline_parallel = false; sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload)); gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get()); } if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } }
n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); n_nodes_pp = ggml_graph_n_nodes(gf); }
// reserve with tg (token generation) graph to get the number of splits and nodes { auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc); if (!gf) { throw std::runtime_error("failed to allocate compute tg buffers"); }
n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); n_nodes_tg = ggml_graph_n_nodes(gf); }
// reserve again with pp graph to avoid ggml-alloc reallocations during inference { // TODO: not sure if the following graph would be worst case for multi-stream KV caches: // // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); // auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } }
for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; ggml_backend_buffer_type_t buft = backend_buft[i]; if (!model.hparams.no_alloc) { backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); } if (backend_buf_exp_size[i] > 1) { LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, ggml_backend_buft_name(buft), backend_buf_exp_size[i] / 1024.0 / 1024.0); } }
if (n_nodes_pp == n_nodes_tg) { LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); } else { LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg); }
if (n_splits_pp == n_splits_tg) { LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp); } else { LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg); }
const int64_t t_end_us = ggml_time_us();
LLAMA_LOG_INFO("%s: reserve took %.2f ms, sched copies = %d\n", __func__, (t_end_us - t_start_us)/1000.0, ggml_backend_sched_get_n_copies(sched.get()));}
void llama_context::synchronize() { if (!sched) { return; }
ggml_backend_sched_synchronize(sched.get());
// FIXME: if multiple single tokens are evaluated without a synchronization, // the stats will be added to the prompt evaluation stats // this should only happen when using batch size 1 to evaluate a batch
// add the evaluation to the stats if (n_queued_tokens == 1) { if (!cparams.no_perf) { t_eval_us += ggml_time_us() - t_compute_start_us; } n_eval++; } else if (n_queued_tokens > 1) { if (!cparams.no_perf) { t_p_eval_us += ggml_time_us() - t_compute_start_us; } n_p_eval += n_queued_tokens; }
// get a more accurate load time, upon first eval if (n_queued_tokens > 0 && !has_evaluated_once) { t_load_us = ggml_time_us() - t_start_us; has_evaluated_once = true; }
n_queued_tokens = 0; t_compute_start_us = 0;}
const llama_model & llama_context::get_model() const { return model;}
const llama_cparams & llama_context::get_cparams() const { return cparams;}
ggml_backend_sched_t llama_context::get_sched() const { return sched.get();}
uint32_t llama_context::n_ctx() const { return cparams.n_ctx;}
uint32_t llama_context::n_ctx_seq() const { return cparams.n_ctx_seq;}
uint32_t llama_context::n_batch() const { return cparams.n_batch;}
uint32_t llama_context::n_ubatch() const { return cparams.n_ubatch;}
uint32_t llama_context::n_seq_max() const { return cparams.n_seq_max;}
uint32_t llama_context::n_threads() const { return cparams.n_threads;}
uint32_t llama_context::n_threads_batch() const { return cparams.n_threads_batch;}
llama_memory_t llama_context::get_memory() const { return memory.get();}
bool llama_context::memory_update(bool optimize) { if (!memory) { return false; }
{ const auto mctx = memory->init_update(this, optimize); switch (mctx->get_status()) { case LLAMA_MEMORY_STATUS_SUCCESS: { // noop } break; case LLAMA_MEMORY_STATUS_NO_UPDATE: { // no updates need to be performed return false; } case LLAMA_MEMORY_STATUS_FAILED_PREPARE: case LLAMA_MEMORY_STATUS_FAILED_COMPUTE: { LLAMA_LOG_ERROR("%s: failed to prepare memory update\n", __func__); return false; } }
// reset the previous graph result to make sure that it won't be reused // TODO: change the mctx->apply() to return information if a graph reserve is needed // reset the graph result only if the memory module did reset the scheduler gf_res_prev->reset();
if (!mctx->apply()) { LLAMA_LOG_ERROR("%s: failed to apply memory update\n", __func__); } }
// if the memory module did any computation, we have to reserve a new worst-case graph { const auto mctx = memory->init_full(); if (!mctx) { throw std::runtime_error("failed to initialize memory context"); }
const uint32_t n_seqs = cparams.n_seq_max; const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch);
const uint32_t n_outputs_max = std::min(n_tokens, cparams.n_outputs_max);
auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_max, mctx.get()); if (!gf) { LLAMA_LOG_ERROR("%s: failed to reserve graph after the memory update\n", __func__); } }
return true;}
enum llama_pooling_type llama_context::pooling_type() const { return cparams.pooling_type;}
float * llama_context::get_logits() { output_reorder();
return logits.data;}
int64_t llama_context::output_resolve_row(int32_t i) const { int64_t j = -1;
// support negative indices (last output row) if (i < 0) { j = n_outputs + i; if (j < 0) { throw std::runtime_error(format("negative index out of range [0, %d)", n_outputs)); } } else if ((size_t) i >= output_ids.size()) { throw std::runtime_error(format("out of range [0, %zu)", output_ids.size())); } else { // use output_ids to translate the batch token index into a row number // that holds this token's data. j = output_ids[i]; }
if (j < 0) { // the batch token was not configured to output anything throw std::runtime_error(format("batch.logits[%d] != true", i)); }
if (j >= n_outputs) { throw std::runtime_error(format("corrupt output buffer (j=%" PRId64 ", n_outputs=%d)", j, n_outputs)); }
return j;}
float * llama_context::get_logits_ith(int32_t i) { output_reorder();
try { if (logits.data == nullptr) { throw std::runtime_error("no logits"); }
const int64_t j = output_resolve_row(i); return logits.data + j*model.vocab.n_tokens(); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid logits id %d, reason: %s\n", __func__, i, err.what());#ifndef NDEBUG GGML_ABORT("fatal error");#else return nullptr;#endif }}
float * llama_context::get_embeddings() { output_reorder();
return embd.data;}
llama_token * llama_context::get_sampled_tokens() const{ return sampling.sampled.data;}
float * llama_context::get_embeddings_ith(int32_t i) { output_reorder();
try { if (embd.data == nullptr) { throw std::runtime_error("no embeddings"); }
const int64_t j = output_resolve_row(i); const uint32_t n_embd_out = model.hparams.n_embd_out(); return embd.data + j*n_embd_out; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid embeddings id %d, reason: %s\n", __func__, i, err.what());#ifndef NDEBUG GGML_ABORT("fatal error");#else return nullptr;#endif }}
float * llama_context::get_embeddings_seq(llama_seq_id seq_id) { auto it = embd_seq.find(seq_id); if (it == embd_seq.end()) { return nullptr; }
return it->second.data();}
float * llama_context::get_embeddings_nextn() { output_reorder();
return embd_nextn.data;}
float * llama_context::get_embeddings_nextn_ith(int32_t i) { output_reorder();
try { if (embd_nextn.data == nullptr) { throw std::runtime_error("no nextn embeddings"); }
const uint32_t n_embd = model.hparams.n_embd_out();
if (!cparams.embeddings_nextn_masked) { // unmasked: nextn rows are stored densely, indexed by raw token position. if (i < 0 || (size_t)(i + 1) * n_embd > embd_nextn.size) { throw std::runtime_error(format("out of range [0, %zu)", embd_nextn.size / n_embd)); } return embd_nextn.data + (size_t) i * n_embd; }
const int64_t j = output_resolve_row(i); return embd_nextn.data + j*n_embd; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid nextn embeddings id %d, reason: %s\n", __func__, i, err.what());#ifndef NDEBUG GGML_ABORT("fatal error");#else return nullptr;#endif }}
float * llama_context::get_embeddings_layer_inp(uint32_t lid) { output_reorder();
GGML_ASSERT(lid < embd_layer_inp.size() && embd_layer_inp[lid].has_data());
return embd_layer_inp[lid].data;}
uint32_t llama_context::get_n_capture() const { return cparams.n_capture_layers;}
float * llama_context::get_embeddings_capture() { output_reorder();
return embd_capture.data;}
float * llama_context::get_embeddings_capture_ith(int32_t i) { output_reorder();
try { if (embd_capture.data == nullptr) { throw std::runtime_error("no capture embeddings"); }
const uint32_t n_cap = cparams.n_capture_layers; const uint32_t n_embd = model.hparams.n_embd; const uint32_t row = n_cap * n_embd; // width of one concatenated row
// capture rows always follow the masked (output-row) layout, mirroring the // pre-norm masked path: the buffer holds one row per output position. const int64_t j = output_resolve_row(i); if (j < 0 || (size_t) (j + 1) * row > embd_capture.size) { throw std::runtime_error(format("out of range [0, %zu)", embd_capture.size / row)); } return embd_capture.data + (size_t) j * row; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid capture embeddings id %d, reason: %s\n", __func__, i, err.what()); return nullptr; }}
llama_token llama_context::get_sampled_token_ith(int32_t idx) { output_reorder();
if (!sampling.sampled.has_data()) { return LLAMA_TOKEN_NULL; }
try { const int64_t row = output_resolve_row(idx); GGML_ASSERT(row < (int64_t) sampling.sampled.size); return sampling.sampled.data[row]; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled token id %d, reason: %s\n", __func__, idx, err.what()); return LLAMA_TOKEN_NULL; }}
float * llama_context::get_sampled_probs_ith(int32_t idx) { output_reorder();
if (!sampling.probs.has_data()) { return nullptr; }
try { const int64_t row = output_resolve_row(idx); if ((size_t) row >= sampling.probs_count.size() || sampling.probs_count[row] == 0) { return nullptr; } return sampling.probs.data + row*model.vocab.n_tokens(); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled probs id %d, reason: %s\n", __func__, idx, err.what()); return nullptr; }}
float * llama_context::get_sampled_logits_ith(int32_t idx) { output_reorder();
if (!sampling.logits.has_data()) { return nullptr; }
try { const int64_t row = output_resolve_row(idx); if ((size_t) row >= sampling.logits_count.size() || sampling.logits_count[row] == 0) { return nullptr; } return sampling.logits.data + row*model.vocab.n_tokens(); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled logits id %d, reason: %s\n", __func__, idx, err.what()); return nullptr; }}
const llama_token * llama_context::get_sampled_candidates_ith(int32_t idx) { output_reorder();
try { const int64_t row = output_resolve_row(idx); if (sampling.candidates.has_data() && (size_t) row < sampling.candidates_count.size() && sampling.candidates_count[row] > 0) { return sampling.candidates.data + row*model.vocab.n_tokens(); } } catch (const std::exception & err) { // fallback to full vocab list GGML_UNUSED(err); }
return sampling.token_ids_full_vocab.data();}
size_t llama_context::get_sampled_candidates_count(int32_t idx) { output_reorder();
if (!sampling.candidates.has_data()) { return 0; }
try { const int64_t row = output_resolve_row(idx); if ((size_t) row >= sampling.candidates_count.size()) { return 0; } return sampling.candidates_count[row]; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled candidates count id %d, reason: %s\n", __func__, idx, err.what()); return 0; }}
size_t llama_context::get_sampled_logits_count(int32_t idx) { output_reorder();
if (!sampling.logits.has_data()) { return model.vocab.n_tokens(); }
try { const int64_t row = output_resolve_row(idx); if ((size_t) row >= sampling.logits_count.size()) { return 0; } return sampling.logits_count[row]; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled logits count id %d, reason: %s\n", __func__, idx, err.what()); return 0; }}
size_t llama_context::get_sampled_probs_count(int32_t idx) { output_reorder();
if (!sampling.probs.has_data()) { return 0; }
try { const int64_t row = output_resolve_row(idx); if ((size_t) row >= sampling.probs_count.size()) { return 0; } return sampling.probs_count[row]; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: invalid backend sampled probs count id %d, reason: %s\n", __func__, idx, err.what()); return 0; }}
void llama_context::attach_threadpool( ggml_threadpool_t threadpool, ggml_threadpool_t threadpool_batch) { LLAMA_LOG_DEBUG("%s: call\n", __func__);
this->threadpool = threadpool; this->threadpool_batch = threadpool_batch ? threadpool_batch : threadpool;}
void llama_context::detach_threadpool() { LLAMA_LOG_DEBUG("%s: call\n", __func__);
this->threadpool = nullptr; this->threadpool_batch = nullptr;}
void llama_context::set_n_threads(int32_t n_threads, int32_t n_threads_batch) { LLAMA_LOG_DEBUG("%s: n_threads = %d, n_threads_batch = %d\n", __func__, n_threads, n_threads_batch);
cparams.n_threads = n_threads; cparams.n_threads_batch = n_threads_batch;}
void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data) { LLAMA_LOG_DEBUG("%s: call\n", __func__);
this->abort_callback = abort_callback; this->abort_callback_data = abort_callback_data;
for (auto & backend : backends) { auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); if (reg) { auto * set_abort_callback_fn = (ggml_backend_set_abort_callback_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_abort_callback"); if (set_abort_callback_fn) { set_abort_callback_fn(backend.get(), this->abort_callback, this->abort_callback_data); } } }}
void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
cparams.embeddings = value;
// TODO: not sure yet if we want to reserve here //sched_need_reserve = true;}
void llama_context::set_embeddings_nextn(bool value, bool masked) { LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked);
cparams.embeddings_nextn = value; cparams.embeddings_nextn_masked = masked;}
void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) { LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable);
GGML_ASSERT(lid <= model.hparams.n_layer());
cparams.embeddings_layer_inp[lid] = enable;
// note: without this reserve, the draft acceptance drops to zero. not sure why - this is unexpected sched_need_reserve = true;}
void llama_context::set_nextn_layer_offset(int32_t offset) { cparams.nextn_layer_offset = offset;}
void llama_context::set_capture_layers(const std::vector<int32_t> & layer_ids) { const int32_t n_layer = (int32_t) model.hparams.n_layer(); if (layer_ids.size() > LLAMA_MAX_LAYERS) { throw std::invalid_argument("Too many capture layers"); } for (int32_t il : layer_ids) { if (il < 0 || il >= n_layer || il >= LLAMA_MAX_LAYERS) { throw std::invalid_argument("Capture layer index out of range"); } } // reset cparams.embeddings_capture = false; cparams.n_capture_layers = 0; cparams.capture_layer_idx = {};
// enabling/disabling capture adds/removes the t_h_capture node from the // graph (see llm_graph_result::set_outputs()), so the scheduler's // backend-assignment table -- built against whatever topology was live // at the last reserve -- must be re-derived before the next decode. // Without this, ggml_backend_sched_get_tensor_backend() on the newly // introduced t_h_capture tensor correctly reports "unknown" (nullptr), // since the scheduler never split a graph that contained it. sched_need_reserve = true;
if (layer_ids.empty()) { return; }
uint32_t n = 0; for (int32_t il : layer_ids) { cparams.capture_layer_idx[n++] = il; }
cparams.n_capture_layers = n; cparams.embeddings_capture = n > 0; // capture rows reuse the masked output-row layout; force masked extraction on. cparams.embeddings_nextn_masked = true;}
void llama_context::set_dspark_ctx(const float * feat, int64_t n_ctx_rows, int64_t n_embd_cap, const int32_t * pos) { if (n_ctx_rows <= 0 || n_embd_cap <= 0 || feat == nullptr) { // reset: no staged context (e.g. before the very first drafter round, // where the whole prompt still needs to go through as context on the // first call, or between unrelated decodes). dspark_ctx.n_ctx_rows = 0; dspark_ctx.n_embd_cap = 0; dspark_ctx.v_ctx_feat.clear(); dspark_ctx.v_ctx_pos.clear(); return; }
dspark_ctx.n_ctx_rows = n_ctx_rows; dspark_ctx.n_embd_cap = n_embd_cap;
dspark_ctx.v_ctx_feat.assign(feat, feat + (size_t) n_ctx_rows * (size_t) n_embd_cap);
dspark_ctx.v_ctx_pos.resize((size_t) n_ctx_rows); if (pos != nullptr) { std::copy(pos, pos + n_ctx_rows, dspark_ctx.v_ctx_pos.begin()); } else { // caller didn't supply explicit positions: assume a contiguous run // ending just before the current staged sequence length. this is a // convenience default; callers doing real multi-round decoding should // pass explicit positions since the growing-cache bookkeeping (Phase 2) // owns the authoritative position numbering. std::iota(dspark_ctx.v_ctx_pos.begin(), dspark_ctx.v_ctx_pos.end(), 0); }}
void llama_context::set_causal_attn(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
if (cparams.causal_attn == value) { return; }
cparams.causal_attn = value;
sched_need_reserve = true;}
void llama_context::set_warmup(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
if (cparams.warmup == value) { return; }
cparams.warmup = value;
// warmups are usually with small batches, so no need to reserve //sched_need_reserve = true;}
bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) { if (!sampler && sampling.samplers.count(seq_id) == 0) { return true; }
LLAMA_LOG_DEBUG("%s: seq_id = %d, sampler = %p\n", __func__, (int) seq_id, (void *) sampler);
if (sampler && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) { static bool warned = false; if (!warned) { LLAMA_LOG_WARN("%s: backend sampling not supported with SPLIT_MODE_TENSOR; using CPU\n", __func__); warned = true; } if (sampling.samplers.count(seq_id) > 0) { sched_need_reserve = true; } sampling.samplers.erase(seq_id); return false; }
const bool can_offload = sampler && sampler->iface->backend_init && sampler->iface->backend_apply && llama_sampler_chain_n(sampler) > 0;
if (sampler && can_offload) { auto * buft = ggml_backend_dev_buffer_type(model.dev_output());
sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq);
sampling.samplers[seq_id] = sampler;
// a server sets a fresh chain for every request; reserving again for the same chain costs ~160 ms const auto it = sampler_sigs_reserved.find(seq_id); if (it == sampler_sigs_reserved.end() || it->second != llama_sampler_chain_sig(sampler)) { sched_need_reserve = true; }
return true; }
if (sampler && !can_offload) { LLAMA_LOG_WARN("%s: sampler '%s' for seq_id = %d, cannot be offloaded to the backend\n", __func__, llama_sampler_name(sampler), seq_id);
if (sampling.samplers.count(seq_id) > 0) { sched_need_reserve = true; }
sampling.samplers.erase(seq_id);
return false; }
// removing a sampler only shrinks the graph, the reserved buffers still fit it sampling.samplers.erase(seq_id);
return true;}
void llama_context::set_adapters_lora(llama_adapter_lora ** adapters, size_t n_adapters, float * scales) { LLAMA_LOG_DEBUG("%s: adapters = %p\n", __func__, (void *) adapters);
if (adapters_lora_are_same(adapters, n_adapters, scales)) { return; }
loras.reset(new llama_adapter_loras());
for (size_t i = 0; i < n_adapters; i ++) { if (scales[i] != 0.0f) { loras->insert({adapters[i], scales[i]}); } }
sched_need_reserve = true;}
bool llama_context::adapters_lora_are_same(llama_adapter_lora ** adapters, size_t n_adapters, float * scales) { LLAMA_LOG_DEBUG("%s: adapters = %p\n", __func__, (void *) adapters);
// Adapters with a zero scale are never added to `loras`, so also ignore them for the comparison. size_t n_non_zero = 0;
for (size_t i = 0; i < n_adapters; i ++) { if (scales[i] == 0.0f) { continue; } n_non_zero++;
auto it = loras->find(adapters[i]);
if (it == loras->end() || it->second != scales[i]) { return false; } }
if (n_non_zero != loras->size()) { return false; }
return true;}
bool llama_context::set_adapter_cvec( const float * data, size_t len, int32_t n_embd, int32_t il_start, int32_t il_end) { LLAMA_LOG_DEBUG("%s: il_start = %d, il_end = %d\n", __func__, il_start, il_end);
bool res = cvec->apply(model, data, len, n_embd, il_start, il_end);
sched_need_reserve = true;
return res;}
llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret) { if (mctx && !mctx->apply()) { LLAMA_LOG_ERROR("%s: failed to apply memory context\n", __func__); ret = GGML_STATUS_FAILED; return nullptr; }
auto * res = gf_res_prev.get(); auto * gf = res->get_gf();
// the new graph parameters // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype);
if (!graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__);
// with pipeline parallelism, the previous graph_compute_async may still be running // on the GPU. we must synchronize before set_inputs to avoid overwriting input tensors // that the previous compute is still reading. if (cparams.pipeline_parallel) { ggml_backend_sched_synchronize(sched.get()); }
n_reused++; } else { res->reset();
ggml_backend_sched_reset(sched.get()); ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data);
//const auto t_start_us = ggml_time_us();
gf = model.build_graph(gparams);
//LLAMA_LOG_INFO("graph build time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0);
if (!gf) { LLAMA_LOG_ERROR("%s: failed to initialize graph\n", __func__); ret = GGML_STATUS_FAILED; return nullptr; }
if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; return nullptr; } }
// set the input data for the input tensors { //const auto t_start_us = ggml_time_us();
// FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated res->set_inputs(&ubatch);
//LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); }
const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; return nullptr; }
ret = GGML_STATUS_SUCCESS;
return res;}
int llama_context::encode(const llama_batch & batch_inp) { // MTP hook batches carry both token (next-token id) and embd (h_nextn row), // so accept either present rather than requiring exactly one. GGML_ASSERT(batch_inp.token || batch_inp.embd);
if (batch_inp.n_tokens == 0) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; }
const auto & hparams = model.hparams;
// eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim const int64_t n_embd = hparams.n_embd_inp_enc(); const int64_t n_vocab = model.vocab.n_tokens();
// note: during encode, we always pass the full sequence starting from pos = 0 if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; }
const uint32_t n_tokens = balloc->get_n_tokens();
// [TAG_NO_CACHE_PAD] // TODO: add new split mode where we pad the input sequences so that ubatch.equal_seqs == true const llama_ubatch ubatch = balloc->split_simple(n_tokens);
// micro-batching is not possible for non-causal encoding, so we process the batch in a single shot GGML_ASSERT(cparams.n_ubatch >= n_tokens && "encoder requires n_ubatch >= n_tokens");
// TODO: this clear of the buffer can easily be forgotten - need something better // sync first so any in-flight async copies into embd_seq complete before it is freed if (!embd_seq.empty()) { synchronize(); } embd_seq.clear();
if (t_compute_start_us == 0) { t_compute_start_us = ggml_time_us(); }
sched_reserve();
n_queued_tokens += n_tokens;
// reserve output buffer if (output_reserve(n_tokens) < n_tokens) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %u outputs\n", __func__, n_tokens); return -2; };
for (uint32_t i = 0; i < n_tokens; ++i) { output_ids[i] = i; }
n_outputs = n_tokens;
const auto causal_attn_org = cparams.causal_attn;
// always use non-causal attention for encoder graphs // TODO: this is a tmp solution until we have a proper way to support enc-dec models // ref: https://github.com/ggml-org/llama.cpp/pull/12181#issuecomment-2730451223 cparams.causal_attn = false;
ggml_status status; const auto * res = process_ubatch(ubatch, LLM_GRAPH_TYPE_ENCODER, nullptr, status);
cparams.causal_attn = causal_attn_org;
if (!res) { switch (status) { case GGML_STATUS_ABORTED: return 2; case GGML_STATUS_ALLOC_FAILED: return -2; case GGML_STATUS_FAILED: return -3; case GGML_STATUS_SUCCESS: GGML_ABORT("should not happen"); } }
auto * t_logits = res->get_logits(); auto * t_embd = res->get_embd_pooled() ? res->get_embd_pooled() : res->get_embd(); auto * t_h_nextn = cparams.embeddings_nextn ? res->get_h_nextn() : nullptr;
// extract logits if (logits.data && t_logits) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr);
ggml_backend_tensor_get_async(backend_res, t_logits, logits.data, 0, n_tokens*n_vocab*sizeof(float)); }
// extract embeddings if (embd.data && t_embd) { ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(sched.get(), t_embd); GGML_ASSERT(backend_embd != nullptr);
switch (cparams.pooling_type) { case LLAMA_POOLING_TYPE_NONE: { // extract token embeddings GGML_ASSERT(embd.data != nullptr); const uint32_t n_embd_out = hparams.n_embd_out();
GGML_ASSERT(n_tokens*n_embd_out <= (int64_t) embd.size); ggml_backend_tensor_get_async(backend_embd, t_embd, embd.data, 0, n_tokens*n_embd_out*sizeof(float)); } break; case LLAMA_POOLING_TYPE_MEAN: case LLAMA_POOLING_TYPE_CLS: case LLAMA_POOLING_TYPE_LAST: { // extract sequence embeddings auto & embd_seq_out = embd_seq;
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { const llama_seq_id seq_id = ubatch.seq_id_unq[s]; const int32_t seq_idx = ubatch.seq_idx[seq_id];
// use n_embd_out (not n_embd_inp) - the pooled embedding has the model's // output dimension, which differs from input dimension for deepstack models (e.g. qwen3vl) const uint32_t n_embd_out = hparams.n_embd_out(); embd_seq_out[seq_id].resize(n_embd_out); ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_out*seq_idx)*sizeof(float), n_embd_out*sizeof(float)); } } break; case LLAMA_POOLING_TYPE_RANK: { // extract the rerank score - n_cls_out floats per sequence auto & embd_seq_out = embd_seq;
const uint32_t n_cls_out = hparams.n_cls_out;
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { const llama_seq_id seq_id = ubatch.seq_id_unq[s]; const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_cls_out); ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_cls_out*seq_idx)*sizeof(float), n_cls_out*sizeof(float)); } } break; case LLAMA_POOLING_TYPE_UNSPECIFIED: { GGML_ABORT("unknown pooling type"); } } }
// extract nextn embeddings (hidden state before the final output norm) if (embd_nextn.data && t_h_nextn && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn); GGML_ASSERT(backend_h != nullptr);
const uint32_t n_embd = hparams.n_embd_out(); GGML_ASSERT(n_tokens*n_embd <= (int64_t) embd_nextn.size); ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn.data, 0, n_tokens*n_embd*sizeof(float)); }
// extract multi-layer capture embeddings (concatenated per position). // single bulk copy: t_h_capture is already [n_capture * n_embd, n_tokens]. if (embd_capture.data && cparams.n_capture_layers > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_tensor * t_cap = res->get_h_capture(); if (!t_cap) { LLAMA_LOG_ERROR("%s: target graph did not produce requested capture layers\n", __func__); return -1; } ggml_backend_t backend_c = ggml_backend_sched_get_tensor_backend(sched.get(), t_cap); GGML_ASSERT(backend_c != nullptr); const uint32_t row = cparams.n_capture_layers * hparams.n_embd; GGML_ASSERT(n_tokens * (int64_t) row <= (int64_t) embd_capture.size); ggml_backend_tensor_get_async(backend_c, t_cap, embd_capture.data, 0, (size_t) n_tokens * row * sizeof(float)); }
// TODO: hacky solution if (model.arch == LLM_ARCH_T5 && t_embd) { //cross.t_embd = t_embd;
synchronize();
cross.n_embd = t_embd->ne[0]; cross.n_enc = t_embd->ne[1]; cross.v_embd.resize(cross.n_embd*cross.n_enc); memcpy(cross.v_embd.data(), embd.data, ggml_nbytes(t_embd));
const auto & batch = balloc->get_batch();
// remember the sequence ids used during the encoding - needed for cross attention later cross.seq_ids_enc.resize(n_tokens); for (uint32_t i = 0; i < n_tokens; i++) { cross.seq_ids_enc[i].clear();
for (int s = 0; s < batch.n_seq_id[i]; s++) { const llama_seq_id seq_id = batch.seq_id[i][s];
cross.seq_ids_enc[i].insert(seq_id); } } }
return 0;}
template<typename T>static void copy_tensor_async_rows( const std::vector<ggml_tensor *> & tensors, const buffer_view<T> & dst, size_t stride, uint32_t row_offset, ggml_backend_sched_t sched, std::vector<uint32_t> * counts = nullptr) { if (!dst.has_data()) { return; }
for (size_t i = 0; i < tensors.size(); ++i) { auto * tensor = tensors[i]; if (tensor == nullptr) { continue; }
const uint32_t row = row_offset + i; const size_t n_elements = ggml_nelements(tensor); GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy"); GGML_ASSERT(n_elements <= stride); GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor); T * row_ptr = dst.data + (size_t) row * stride; ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
if (counts) { GGML_ASSERT(row < counts->size()); (*counts)[row] = n_elements; } }}
static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map<llama_seq_id, llama_sampler *> & samplers) { for (uint32_t i = 0; i < ubatch.n_tokens; i++) { if (!ubatch.output[i]) { continue; }
// Check if the output token has at least one sequence without a backend sampler. for (int32_t j = 0; j < ubatch.n_seq_id[i]; ++j) { llama_seq_id seq_id = ubatch.seq_id[i][j]; if (samplers.find(seq_id) == samplers.end()) { return true; } } } return false; // all sequences use backend sampling}
int llama_context::decode(const llama_batch & batch_inp) { // MTP hook batches carry both token (next-token id) and embd (h_nextn row), // so accept either present rather than requiring exactly one. GGML_ASSERT(batch_inp.token || batch_inp.embd);
if (!memory) { LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__); return encode(batch_inp); }
if (batch_inp.n_tokens == 0) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; }
const auto & vocab = model.vocab; const auto & hparams = model.hparams;
const int64_t n_vocab = vocab.n_tokens(); const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd; const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
// when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; const bool has_samplers = !sampling.samplers.empty();
const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
// embedding contexts output every token even when batch.logits is not set if (has_samplers && (output_all || batch_inp.logits)) { std::vector<int32_t> seq_output_count(n_seq_max, 0);
for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { if (!output_all && batch_inp.logits[i] == 0) { continue; }
const int ns = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : 1;
for (int32_t s = 0; s < ns; ++s) { const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { continue; }
seq_output_count[seq_id]++; auto sampler = sampling.samplers.find(seq_id); if (sampler != sampling.samplers.end() && seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) { LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence " "(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq, seq_id, seq_output_count[seq_id]); return -1; } } } }
if (!balloc->init(batch_inp, vocab, memory.get(), n_embd, n_seq_max, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; }
const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs();
if (output_all) { // require that all tokens are output if (n_outputs_all != n_tokens_all) { LLAMA_LOG_ERROR("%s: pooled embedding requires that all tokens are output (n_outputs_all = %d, n_tokens_all = %d)\n", __func__, n_outputs_all, n_tokens_all); return -1; } }
GGML_ASSERT(n_tokens_all <= cparams.n_batch);
GGML_ASSERT((cparams.causal_attn || cparams.n_ubatch >= n_tokens_all) && "non-causal attention requires n_ubatch >= n_tokens");
// TODO: this clear of the buffer can easily be forgotten - need something better // sync first so any in-flight async copies into embd_seq complete before it is freed if (!embd_seq.empty()) { synchronize(); } embd_seq.clear();
if (t_compute_start_us == 0) { t_compute_start_us = ggml_time_us(); } n_queued_tokens += n_tokens_all;
output_swaps.clear();
sched_reserve();
bool did_optimize = false;
// handle any pending shifts/copies memory_update(false);
llama_memory_context_ptr mctx;
while (true) { mctx = memory->init_batch(*balloc, cparams.n_ubatch, output_all); if (!mctx) { return -2; }
switch (mctx->get_status()) { case LLAMA_MEMORY_STATUS_SUCCESS: { } break; case LLAMA_MEMORY_STATUS_NO_UPDATE: { LLAMA_LOG_ERROR("%s: unexpected memory context status: %d\n", __func__, mctx->get_status());
return -2; } case LLAMA_MEMORY_STATUS_FAILED_PREPARE: { if (!did_optimize) { did_optimize = true;
if (memory_update(true)) { LLAMA_LOG_DEBUG("%s: retrying batch size %d after cache optimization\n", __func__, balloc->get_n_tokens());
continue; } }
LLAMA_LOG_WARN("%s: failed to find a memory slot for batch of size %d\n", __func__, balloc->get_n_tokens());
return 1; } case LLAMA_MEMORY_STATUS_FAILED_COMPUTE: { LLAMA_LOG_ERROR("%s: compute failed while preparing batch of size %d\n", __func__, balloc->get_n_tokens());
return -2; } }
break; }
// reserve output buffer if (output_reserve(n_outputs_all) < n_outputs_all) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); return -2; };
// start a new sampling transaction for this logical batch for (const auto & entry : sampling.samplers) { llama_sampler_backend_begin(entry.second); }
int64_t n_outputs_prev = 0; int64_t n_tokens_prev = 0;
do { const auto & ubatch = mctx->get_ubatch();
// count the outputs in this ubatch { int32_t n_outputs_new = 0;
if (n_outputs_all == n_tokens_all) { n_outputs_new = ubatch.n_tokens; } else { for (uint32_t i = 0; i < ubatch.n_tokens; i++) { n_outputs_new += (int32_t) (ubatch.output[i] != 0); } }
// needs to happen before the graph is built n_outputs = n_outputs_new; }
ggml_status status;
const auto * res = process_ubatch(ubatch, ctx_type_to_graph_type(cparams.ctx_type), mctx.get(), status);
if (!res) { // the last ubatch failed or was aborted -> remove all positions of that ubatch from the memory module llama_pos pos_min[LLAMA_MAX_SEQ]; for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { pos_min[s] = std::numeric_limits<llama_pos>::max(); }
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { const auto & seq_id = ubatch.seq_id[i][0];
pos_min[seq_id] = std::min(pos_min[seq_id], ubatch.pos[i]); }
for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { if (pos_min[s] == std::numeric_limits<llama_pos>::max()) { continue; }
LLAMA_LOG_WARN("%s: removing memory module entries for seq_id = %d, pos = [%d, +inf)\n", __func__, s, pos_min[s]);
memory->seq_rm(s, pos_min[s], -1); }
switch (status) { case GGML_STATUS_ABORTED: return 2; case GGML_STATUS_ALLOC_FAILED: return -2; case GGML_STATUS_FAILED: return -3; case GGML_STATUS_SUCCESS: GGML_ABORT("should not happen"); } }
// plot the computation graph in dot format (for debugging purposes) //if (n_past%100 == 0) { // ggml_graph_dump_dot(gf, NULL, "llama.dot"); //}
auto * t_logits = res->get_logits(); auto * t_embd = cparams.embeddings ? res->get_embd() : nullptr; auto * t_h_nextn = cparams.embeddings_nextn ? res->get_h_nextn() : nullptr;
if (t_embd && res->get_embd_pooled()) { t_embd = res->get_embd_pooled(); }
// extract multi-layer capture embeddings, concatenated per output position. // capture always uses the masked (output-row) layout, so t_h_capture is // [n_capture * n_embd, n_outputs]; copy in one shot per ubatch. if (embd_capture.data && cparams.n_capture_layers > 0 && n_outputs > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_tensor * t_cap = res->get_h_capture(); if (!t_cap) { LLAMA_LOG_ERROR("%s: target graph did not produce requested capture layers\n", __func__); return -1; } ggml_backend_t backend_c = ggml_backend_sched_get_tensor_backend(sched.get(), t_cap); GGML_ASSERT(backend_c != nullptr); const uint32_t row = cparams.n_capture_layers * hparams.n_embd; float * embd_capture_out = embd_capture.data + (size_t) n_outputs_prev * row; GGML_ASSERT((n_outputs_prev + n_outputs) * (int64_t) row <= (int64_t) embd_capture.size); ggml_backend_tensor_get_async(backend_c, t_cap, embd_capture_out, 0, (size_t) n_outputs * row * sizeof(float)); }
// extract logits if (logits.data && t_logits && n_outputs > 0 && needs_raw_logits(ubatch, sampling.samplers)) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr);
float * logits_out = logits.data + n_outputs_prev*n_vocab;
if (n_outputs) { GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all); GGML_ASSERT((n_outputs_prev + n_outputs)*n_vocab <= (int64_t) logits.size); ggml_backend_tensor_get_async(backend_res, t_logits, logits_out, 0, n_outputs*n_vocab*sizeof(float)); } }
// extract embeddings if (embd.data && t_embd && n_outputs > 0) { ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(sched.get(), t_embd); GGML_ASSERT(backend_embd != nullptr);
switch (cparams.pooling_type) { case LLAMA_POOLING_TYPE_NONE: { // extract token embeddings GGML_ASSERT(embd.data != nullptr); const uint32_t n_embd_out = hparams.n_embd_out(); float * embd_out = embd.data + n_outputs_prev*n_embd_out;
if (n_outputs) { GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all); GGML_ASSERT((n_outputs_prev + n_outputs)*n_embd_out <= (int64_t) embd.size); ggml_backend_tensor_get_async(backend_embd, t_embd, embd_out, 0, n_outputs*n_embd_out*sizeof(float)); } } break; case LLAMA_POOLING_TYPE_MEAN: case LLAMA_POOLING_TYPE_CLS: case LLAMA_POOLING_TYPE_LAST: { // extract sequence embeddings (cleared before processing each batch) auto & embd_seq_out = embd_seq;
// use n_embd_out (not n_embd_inp) - the pooled embedding has the model's // output dimension, which differs from input dimension for deepstack models (e.g. qwen3vl) const uint32_t n_embd_out = hparams.n_embd_out();
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { const llama_seq_id seq_id = ubatch.seq_id_unq[s]; const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_embd_out); ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_out*seq_idx)*sizeof(float), n_embd_out*sizeof(float)); } } break; case LLAMA_POOLING_TYPE_RANK: { // extract the rerank score - n_cls_out floats per sequence auto & embd_seq_out = embd_seq;
const uint32_t n_cls_out = hparams.n_cls_out;
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { const llama_seq_id seq_id = ubatch.seq_id_unq[s]; const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_cls_out); ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_cls_out*seq_idx)*sizeof(float), n_cls_out*sizeof(float)); } } break; case LLAMA_POOLING_TYPE_UNSPECIFIED: { GGML_ABORT("unknown pooling type"); } } }
extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens);
// extract nextn embeddings before // only meaningful in LLAMA_POOLING_TYPE_NONE (per-token); other pooling modes are ignored. { const bool masked = cparams.embeddings_nextn_masked; const int64_t n_rows = masked ? n_outputs : (int64_t) ubatch.n_tokens; const int64_t offset = masked ? n_outputs_prev : n_tokens_prev;
if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn); GGML_ASSERT(backend_h != nullptr);
const uint32_t n_embd = hparams.n_embd_out(); float * embd_nextn_out = embd_nextn.data + offset*n_embd;
GGML_ASSERT((offset + n_rows)*n_embd <= (int64_t) embd_nextn.size); ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn_out, 0, n_rows*n_embd*sizeof(float)); } }
if (has_samplers) { const auto stride = n_vocab;
// async copy the sampling data from the backend to the host copy_tensor_async_rows(res->t_sampled, sampling.sampled, 1, n_outputs_prev, sched.get()); copy_tensor_async_rows(res->t_sampled_logits, sampling.logits, stride, n_outputs_prev, sched.get(), &sampling.logits_count); copy_tensor_async_rows(res->t_sampled_probs, sampling.probs, stride, n_outputs_prev, sched.get(), &sampling.probs_count); copy_tensor_async_rows(res->t_candidates, sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count); }
n_outputs_prev += n_outputs; n_tokens_prev += ubatch.n_tokens; } while (mctx->next());
// set to total number of outputs in the batch, for use in llama_get_logits_ith n_outputs = n_outputs_all;
// set output mappings if (n_outputs > 0) { bool sorted_output = true;
auto & out_ids = balloc->get_out_ids();
GGML_ASSERT(out_ids.size() == (size_t) n_outputs);
for (int64_t i = 0; i < n_outputs; ++i) { int64_t out_id = out_ids[i]; output_ids[out_id] = i; if (out_id != i) { sorted_output = false; } }
// make the outputs have the same order they had in the user-provided batch // note: this is mostly relevant for recurrent models atm if (!sorted_output && n_outputs > 1) { GGML_ASSERT((size_t) n_outputs == out_ids.size());
// TODO: is there something more efficient which also minimizes swaps? // selection sort, to minimize swaps (from https://en.wikipedia.org/wiki/Selection_sort) for (uint32_t i = 0; i < n_outputs - 1; ++i) { uint32_t j_min = i; for (uint32_t j = i + 1; j < n_outputs; ++j) { if (out_ids[j] < out_ids[j_min]) { j_min = j; } } if (j_min == i) { continue; } std::swap(out_ids[i], out_ids[j_min]);
// remember the swaps and apply them lazily upon logits/embeddings access output_swaps.push_back({ i, j_min }); }
std::fill(output_ids.begin(), output_ids.end(), -1);
for (uint32_t i = 0; i < n_outputs; ++i) { output_ids[out_ids[i]] = i; } } }
// wait for the computation to finish (automatically done when obtaining the model output) //synchronize();
return 0;}
//// output//
uint32_t llama_context::output_reserve(int32_t n_outputs) { const auto & hparams = model.hparams; const auto & vocab = model.vocab;
const int64_t n_outputs_max = std::max<int64_t>(n_outputs, n_seq_max());
const auto n_batch = cparams.n_batch; const auto n_vocab = vocab.n_tokens(); const auto n_embd = hparams.n_embd; const auto n_embd_out = hparams.n_embd_out();
bool has_logits = true; bool has_embd = cparams.embeddings; bool has_embd_nextn = cparams.embeddings_nextn; bool has_embd_capture = cparams.n_capture_layers > 0;
// TODO: hacky enc-dec support if (model.arch == LLM_ARCH_T5) { has_logits = true; has_embd = true; }
size_t backend_float_count = 0; size_t backend_token_count = 0; size_t embd_layer_inp_float_count = 0;
logits.size = has_logits ? n_vocab*n_outputs_max : 0; embd.size = has_embd ? n_embd_out*n_outputs_max : 0; embd_nextn.size = has_embd_nextn ? n_embd_out*n_outputs_max : 0;
if (has_embd_nextn && !cparams.embeddings_nextn_masked) { // unmasked: nextn row exists for every token in the batch, not just // those flagged via batch.logits[i] -> size by token count instead. embd_nextn.size = (size_t) n_embd_out * n_batch; }
for (bool enabled : cparams.embeddings_layer_inp) { if (enabled) { embd_layer_inp_float_count += (size_t) n_embd * n_batch; } }
// Allocate backend sampling output buffers if there are backend samplers configured. const bool has_sampling = !sampling.samplers.empty(); if (has_sampling) { backend_float_count = 2 * n_vocab * n_outputs_max; // logits + probs backend_token_count = (1 + n_vocab) * n_outputs_max; // sampled + candidates }
if (output_ids.empty()) { // init, never resized afterwards output_ids.resize(n_batch); }
const size_t prev_size = buf_output ? ggml_backend_buffer_get_size(buf_output.get()) : 0; embd_capture.size = has_embd_capture ? (size_t) cparams.n_capture_layers * hparams.n_embd * n_outputs_max : 0; const size_t new_size = (logits.size + embd.size + embd_nextn.size + embd_capture.size + embd_layer_inp_float_count + backend_float_count) * sizeof(float) + (backend_token_count) * sizeof(llama_token);
// alloc only when more than the current capacity is required // TODO: also consider shrinking the buffer if (!buf_output || prev_size < new_size) { if (buf_output) {#ifndef NDEBUG // This doesn't happen often, but may be annoying in some cases (like the HellaSwag benchmark) LLAMA_LOG_DEBUG("%s: reallocating output buffer from size %.02f MiB to %.02f MiB\n", __func__, prev_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0);#endif synchronize();
// TODO: not needed? buf_output = nullptr; logits.data = nullptr; embd.data = nullptr; embd_nextn.data = nullptr; embd_capture.data = nullptr; for (auto & layer_inp : embd_layer_inp) { layer_inp = {nullptr, 0}; } }
auto * buft = ggml_backend_cpu_buffer_type(); // try to use the host buffer of the device where the output tensor is allocated for faster transfer to system memory auto * output_dev = model.dev_output(); auto * output_dev_host_buft = output_dev ? ggml_backend_dev_host_buffer_type(output_dev) : nullptr; if (output_dev_host_buft) { buft = output_dev_host_buft; } buf_output.reset(ggml_backend_buft_alloc_buffer(buft, new_size)); if (buf_output == nullptr) { LLAMA_LOG_ERROR("%s: failed to allocate output buffer of size %.2f MiB\n", __func__, new_size / (1024.0 * 1024.0)); return 0; } ggml_backend_buffer_clear(buf_output.get(), 0); }
float * output_base = (float *) ggml_backend_buffer_get_base(buf_output.get());
size_t offset = 0; uint8_t * base = (uint8_t *) output_base;
logits = has_logits ? buffer_view<float>{output_base, logits.size} : buffer_view<float>{nullptr, 0}; offset += logits.size * sizeof(float);
embd = has_embd ? buffer_view<float>{(float *) (base + offset), embd.size} : buffer_view<float>{nullptr, 0}; offset += embd.size * sizeof(float);
embd_nextn = has_embd_nextn ? buffer_view<float>{(float *) (base + offset), embd_nextn.size} : buffer_view<float>{nullptr, 0}; offset += embd_nextn.size * sizeof(float);
for (uint32_t il = 0; il < embd_layer_inp.size(); ++il) { if (cparams.embeddings_layer_inp[il]) { embd_layer_inp[il] = buffer_view<float>{(float *) (base + offset), (size_t) n_embd * n_batch}; offset += embd_layer_inp[il].size * sizeof(float); } else { embd_layer_inp[il] = buffer_view<float>{nullptr, 0}; } }
embd_capture = has_embd_capture ? buffer_view<float>{ (float *) (base + offset), embd_capture.size } : buffer_view<float>{ nullptr, 0 }; offset += embd_capture.size * sizeof(float);
if (has_sampling) { sampling.logits = {(float *) (base + offset), (size_t)(n_vocab*n_outputs_max)}; offset += sampling.logits.size * sizeof(float);
sampling.probs = {(float *) (base + offset), (size_t)(n_vocab*n_outputs_max)}; offset += sampling.probs.size * sizeof(float);
sampling.sampled = {(llama_token *) (base + offset), (size_t)n_outputs_max}; offset += sampling.sampled.size * sizeof(llama_token);
sampling.candidates = {(llama_token *) (base + offset), (size_t)(n_vocab*n_outputs_max)}; offset += sampling.candidates.size * sizeof(llama_token);
// The count vectors keep track of the actual number of logits/probs/candidates // copied from the backend for each output row.
sampling.logits_count.resize(n_outputs_max); sampling.probs_count.resize(n_outputs_max); sampling.candidates_count.resize(n_outputs_max);
std::fill(sampling.logits_count.begin(), sampling.logits_count.end(), 0); std::fill(sampling.probs_count.begin(), sampling.probs_count.end(), 0); std::fill(sampling.candidates_count.begin(), sampling.candidates_count.end(), 0);
std::fill_n(sampling.sampled.data, sampling.sampled.size, LLAMA_TOKEN_NULL); } else { sampling.logits = {nullptr, 0}; sampling.probs = {nullptr, 0}; sampling.sampled = {nullptr, 0}; sampling.candidates = {nullptr, 0};
sampling.logits_count.clear(); sampling.probs_count.clear(); sampling.candidates_count.clear(); }
// set all ids as invalid (negative) std::fill(output_ids.begin(), output_ids.end(), -1);
this->n_outputs = 0;
GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max);
return n_outputs_max;}
void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens) { for (uint32_t il = 0; il < cparams.embeddings_layer_inp.size(); ++il) { if (!cparams.embeddings_layer_inp[il]) { continue; } if (!embd_layer_inp[il].has_data()) { GGML_ABORT("output layer input buffer not allocated"); } ggml_tensor * t = res->get_layer_inp((int) il); if (!t) { GGML_ABORT("layer input tensor not found"); }
const size_t nbytes = ggml_nbytes(t); const size_t nfloats = nbytes / sizeof(float); GGML_ASSERT(n_tokens > 0); GGML_ASSERT(nfloats % n_tokens == 0);
const size_t row_floats = nfloats / n_tokens; const size_t dst_offset = token_offset * row_floats; GGML_ASSERT(dst_offset + nfloats <= embd_layer_inp[il].size);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched.get(), t); GGML_ASSERT(backend != nullptr); ggml_backend_tensor_get_async(backend, t, embd_layer_inp[il].data + dst_offset, 0, nbytes); }}
void llama_context::output_reorder() { const uint64_t n_vocab = model.vocab.n_tokens(); const uint64_t n_embd = model.hparams.n_embd; const uint64_t n_embd_out = model.hparams.n_embd_out();
for (size_t s = 0; s < output_swaps.size(); ++s) { const uint64_t i0 = output_swaps[s].i0; const uint64_t i1 = output_swaps[s].i1;
if (logits.size > 0) { for (uint64_t k = 0; k < n_vocab; k++) { std::swap(logits.data[i0*n_vocab + k], logits.data[i1*n_vocab + k]); } }
if (embd.size > 0) { for (uint64_t k = 0; k < n_embd_out; k++) { std::swap(embd.data[i0*n_embd_out + k], embd.data[i1*n_embd_out + k]); } }
if (embd_capture.size > 0) { const uint64_t row = (uint64_t) cparams.n_capture_layers * n_embd; for (uint64_t k = 0; k < row; k++) { std::swap(embd_capture.data[i0 * row + k], embd_capture.data[i1 * row + k]); } } if (embd_nextn.size > 0) { for (uint64_t k = 0; k < n_embd_out; k++) { std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]); } }
if (embd_layer_inp.size() > 0) { for (int lid = 0; lid < (int) embd_layer_inp.size(); ++lid) { if (embd_layer_inp[lid].size > 0) { for (uint64_t k = 0; k < n_embd; ++k) { std::swap(embd_layer_inp[lid].data[i0*n_embd + k], embd_layer_inp[lid].data[i1*n_embd + k]); } } } }
if (!sampling.samplers.empty()) { assert(sampling.logits.size > 0); assert(sampling.probs.size > 0); assert(sampling.candidates.size > 0); assert(sampling.sampled.size > 0); assert(sampling.logits_count.size() > 0); assert(sampling.probs_count.size() > 0); assert(sampling.candidates_count.size() > 0);
for (uint64_t k = 0; k < n_vocab; ++k) { std::swap(sampling.logits.data[i0*n_vocab + k], sampling.logits.data[i1*n_vocab + k]); }
for (uint64_t k = 0; k < n_vocab; ++k) { std::swap(sampling.probs.data[i0*n_vocab + k], sampling.probs.data[i1*n_vocab + k]); }
for (uint64_t k = 0; k < n_vocab; ++k) { std::swap(sampling.candidates.data[i0*n_vocab + k], sampling.candidates.data[i1*n_vocab + k]); }
std::swap(sampling.sampled.data[i0], sampling.sampled.data[i1]); std::swap(sampling.logits_count[i0], sampling.logits_count[i1]); std::swap(sampling.probs_count[i0], sampling.probs_count[i1]); std::swap(sampling.candidates_count[i0], sampling.candidates_count[i1]); } }
output_swaps.clear();}
//// graph//
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; if (model.arch == LLM_ARCH_KIMI_K3) { // the n_tokens*40 budget below is exhausted at ubatch 3840 res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors()); } else if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { // DFlash2's convolutions and selector are shape work rather than matmuls, // so they cost ~8.6 nodes per tensor against ~5.9 for a plain DFlash draft res = std::max<uint32_t>(1024u, 12u*model.n_tensors()); } else { res = std::max<uint32_t>(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { res += lora->get_n_nodes(); } }
uint32_t n_sampling_nodes = 0; uint32_t n_sampling_nodes_max = 0; for (const auto & [seq_id, sampler] : sampling.samplers) { const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler); n_sampling_nodes += n_nodes; if (cparams.n_outputs_max_per_seq > 1) { n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes); } }
const uint32_t n_sampling_outputs_max = std::min<uint64_t>( std::min(n_tokens, cparams.n_outputs_max), (uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq);
res += n_sampling_nodes; if (n_sampling_outputs_max > 1) { res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max; } return res;}
llm_graph_result * llama_context::get_gf_res_reserve() const { return static_cast<llm_graph_result *>(gf_res_reserve.get());}
// pack sampler outputs into as few sequences as possible before using sequences without samplersstatic void ubatch_prepare_reserve( llama_ubatch & ubatch, uint32_t n_outputs, const std::map<llama_seq_id, llama_sampler *> & samplers, uint32_t n_outputs_max_per_seq) { const uint32_t n_seqs = ubatch.n_seqs; const uint32_t n_seq_tokens = ubatch.n_seq_tokens;
for (uint32_t s = 0; s < n_seqs; ++s) { for (uint32_t t = 0; t < n_seq_tokens; ++t) { const uint32_t i = s * n_seq_tokens + t; ubatch.n_seq_id[i] = 1; ubatch.seq_id[i] = &ubatch.seq_id_unq[s]; } }
// sequences with a sampler that fit in this ubatch std::vector<uint32_t> sampler_seqs; std::vector<bool> has_sampler(n_seqs, false); for (const auto & entry : samplers) { const llama_seq_id seq_id = entry.first; if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) { continue; }
sampler_seqs.push_back(seq_id); has_sampler[seq_id] = true; }
uint32_t n_outputs_set = 0;
const uint32_t n_outputs_per_seq = std::min(n_seq_tokens, n_outputs_max_per_seq); for (uint32_t s : sampler_seqs) { if (n_outputs_set >= n_outputs) { break; }
for (uint32_t t = 0; t < n_outputs_per_seq && n_outputs_set < n_outputs; ++t) { ubatch.output[s * n_seq_tokens + t] = true; ++n_outputs_set; } }
// use sequences without samplers for any remaining outputs for (uint32_t t = 0; t < n_seq_tokens && n_outputs_set < n_outputs; ++t) { for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) { if (has_sampler[s]) { continue; }
ubatch.output[s * n_seq_tokens + t] = true; ++n_outputs_set; } }}
ggml_cgraph * llama_context::graph_reserve( uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) { LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs); GGML_ASSERT(n_outputs >= 1);
if (n_tokens % n_seqs != 0) { n_tokens = ((n_tokens + (n_seqs - 1)) / n_seqs) * n_seqs; // round to next multiple of n_seqs LLAMA_LOG_DEBUG("%s: making n_tokens a multiple of n_seqs - n_tokens = %u, n_seqs = %u, n_outputs = %u\n", __func__, n_tokens, n_seqs, n_outputs); }
ggml_backend_sched_reset(sched.get());
// when the scheduler is reset, we cannot reuse the old graph, so we reset the previous graph result to prevent that gf_res_prev->reset();
// store the n_outputs as it is, and restore it afterwards // TODO: not sure if needed, might simplify in the future by removing this const auto save_n_outputs = this->n_outputs;
this->n_outputs = n_outputs;
llama_batch_allocr balloc(model.hparams.n_pos_per_embd()); llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs);
ubatch_prepare_reserve(ubatch, n_outputs, sampling.samplers, cparams.n_outputs_max_per_seq);
auto * res = gf_res_reserve.get();
const auto gparams = graph_params(res, ubatch, mctx, ctx_type_to_graph_type(cparams.ctx_type));
res->reset();
auto * gf = model.build_graph(gparams);
// verify transform coverage on the pristine graph: after scheduling, // cross-backend copies break the producer chain the check follows if (!hadamard_verified && gf && (!hadamard_rotations.empty() || !hadamard_inverses.empty())) { llama_verify_hadamard_graph(gf, hadamard_rotations, hadamard_inverses); hadamard_verified = true; }
this->n_outputs = save_n_outputs;
// initialize scheduler with the specified graph if (split_only) { if (sizes) { ggml_backend_sched_reserve_size(sched.get(), gf, sizes); } else { ggml_backend_sched_split_graph(sched.get(), gf); } } else if (!ggml_backend_sched_reserve(sched.get(), gf)) { GGML_ASSERT(!sizes); LLAMA_LOG_ERROR("%s: failed to allocate compute buffers\n", __func__); return nullptr; }
return gf;}
llm_graph_params llama_context::graph_params( llm_graph_result * res, const llama_ubatch & ubatch, const llama_memory_context_i * mctx, llm_graph_type gtype) const { return { /*.arch =*/model.arch, /*.hparams =*/model.hparams, /*.cparams =*/cparams, /*.ubatch =*/ubatch, /*.gtype =*/gtype, /*.sched =*/sched.get(), /*.backend_cpu =*/backend_cpu, /*.cvec =*/cvec.get(), /*.loras =*/loras.get(), /*.mctx =*/mctx, /*.cross =*/&cross, /*.dspark_ctx =*/&dspark_ctx, /*.dspark_has_context =*/!dspark_ctx.v_ctx_feat.empty(), /*.dspark_ctx_rows =*/dspark_ctx.n_ctx_rows, /*.dspark_ctx_width =*/dspark_ctx.n_embd_cap, /*.hadamard_rotations =*/&hadamard_rotations, /*.hadamard_inverses =*/&hadamard_inverses, /*.samplers =*/sampling.samplers, /*.n_outputs =*/n_outputs, /*.cb =*/graph_get_cb(), /*.res =*/res, };}
ggml_status llama_context::graph_compute( ggml_cgraph * gf, bool batched) { int n_threads = batched ? cparams.n_threads_batch : cparams.n_threads; ggml_threadpool_t tp = batched ? threadpool_batch : threadpool;
if (backend_cpu != nullptr) { auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_cpu)); auto * set_threadpool_fn = (decltype(ggml_backend_cpu_set_threadpool) *) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_threadpool"); if (set_threadpool_fn) { set_threadpool_fn(backend_cpu, tp); } }
// set the number of threads for all the backends for (const auto & set_n_threads_fn : set_n_threads_fns) { set_n_threads_fn.second(set_n_threads_fn.first, n_threads); }
auto status = ggml_backend_sched_graph_compute_async(sched.get(), gf); if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, status); }
// fprintf(stderr, "splits: %d\n", ggml_backend_sched_get_n_splits(sched));
return status;}
llm_graph_cb llama_context::graph_get_cb() const { return [&](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) { if (il >= 0) { ggml_format_name(cur, "%s-%d", name, il); } else { ggml_set_name(cur, name); }
// - norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends // - force the last op of the layer on the specified backend to avoid running it on the backend of the next layer due to scheduling // FIXME: fix in ggml_backend_sched const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all; if (ubatch.n_tokens < 32 || full_offload) { if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) { const auto & dev_layer = model.dev_layer(il); for (const auto & backend : backends) { if (ggml_backend_get_device(backend.get()) == dev_layer) { if (ggml_backend_supports_op(backend.get(), cur)) { ggml_backend_sched_set_tensor_backend(sched.get(), cur, backend.get()); } } } } } };}
//// state save/load//
class llama_io_write_dummy : public llama_io_write_i {public: llama_io_write_dummy(bool skip_tensors) : skip_tensors(skip_tensors) {}
void write(const void * /* src */, size_t size) override { size_written += size; }
void write_tensor(ggml_tensor * /* tensor */, size_t /* offset */, size_t size) override { if (skip_tensors) { return; }
size_written += size; }
size_t n_bytes() override { return size_written; }
private: const bool skip_tensors;
size_t size_written = 0;};
class llama_io_write_host : public llama_io_write_i {public: llama_io_write_host( uint8_t * p, size_t len) : ptr(p), buf_size(len) {}
~llama_io_write_host() { // TODO: add backend support to batch tensor_get? or some other way to speed this up for (const auto & winfo : winfos) { ggml_backend_tensor_get(winfo.tensor, winfo.ptr, winfo.offset, winfo.size); } }
void write(const void * src, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); } memcpy(ptr, src, size); ptr += size; size_written += size; buf_size -= size; }
void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); }
// save the write for later during destruction winfos.push_back({tensor, ptr, size, offset});
ptr += size; size_written += size; buf_size -= size; }
size_t n_bytes() override { return size_written; }
private: uint8_t * ptr; size_t buf_size = 0; size_t size_written = 0;
struct write_info { ggml_tensor * tensor; uint8_t * ptr; size_t size; size_t offset; }; std::vector<write_info> winfos;};
class llama_io_read_host : public llama_io_read_i {public: llama_io_read_host(const uint8_t * p, size_t len) : ptr(p), buf_size(len) {}
~llama_io_read_host() { // flush the reads for (const auto & rinfo : rinfos) { ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); } }
void read(void * dst, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); } memcpy(dst, ptr, size); ptr += size; size_read += size; buf_size -= size; }
void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); }
// save for later during destruction rinfos.push_back({tensor, ptr, size, offset});
ptr += size; size_read += size; buf_size -= size; }
size_t n_bytes() override { return size_read; }
private: const uint8_t * ptr; size_t buf_size = 0; size_t size_read = 0;
struct read_info { ggml_tensor * tensor; const uint8_t * ptr; size_t size; size_t offset; }; std::vector<read_info> rinfos;};
class llama_io_write_file : public llama_io_write_i {public: llama_io_write_file(llama_file * f) : file(f) {}
void write(const void * src, size_t size) override { file->write_raw(src, size); size_written += size; }
void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { temp_buffer.resize(size); ggml_backend_tensor_get(tensor, temp_buffer.data(), offset, size); write(temp_buffer.data(), temp_buffer.size()); }
size_t n_bytes() override { return size_written; }
private: llama_file * file; size_t size_written = 0; std::vector<uint8_t> temp_buffer;};
class llama_io_read_file : public llama_io_read_i {public: llama_io_read_file(llama_file * f) : file(f) {}
void read(void * dst, size_t size) override { file->read_raw(dst, size); size_read += size; }
void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { temp_buffer.resize(size); read(temp_buffer.data(), size); ggml_backend_tensor_set(tensor, temp_buffer.data(), offset, size); }
size_t n_bytes() override { return size_read; }
private: llama_file * file; size_t size_read = 0; std::vector<uint8_t> temp_buffer;};
class llama_io_write_device : public llama_io_write_i {public: llama_io_write_device(uint8_t * p, size_t len, llama_memory_buffers & mbufs) : ptr(p), buf_size(len), mbufs(mbufs) { }
~llama_io_write_device() { llama_memory_buffers mbufs_new;
for (const auto & winfo : winfos) { auto * buft = ggml_backend_buffer_get_type(winfo.tensor->buffer);
mbufs_new[buft].n_tensors++; mbufs_new[buft].total_size += winfo.size; }
for (auto & [buft, mbuf] : mbufs_new) { ggml_init_params params = { /*.mem_size =*/ 2*mbuf.n_tensors*ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, };
mbuf.ctx.reset(ggml_init(params));
mbuf.org.reserve(mbuf.n_tensors); mbuf.cpy.reserve(mbuf.n_tensors); }
for (const auto & winfo : winfos) { auto * buft = ggml_backend_buffer_get_type(winfo.tensor->buffer);
const int64_t n = winfo.size/ggml_element_size(winfo.tensor);
auto & mbuf = mbufs_new[buft];
mbuf.org.push_back(ggml_view_1d (mbuf.ctx.get(), winfo.tensor, n, winfo.offset)); mbuf.cpy.push_back(ggml_new_tensor_1d(mbuf.ctx.get(), winfo.tensor->type, n)); }
for (auto & [buft, mbuf] : mbufs_new) { auto & mbuf_cur = mbufs[buft];
bool need_alloc = false;
need_alloc = need_alloc || (!mbuf_cur.buf); need_alloc = need_alloc || (mbuf_cur.org.size() != mbuf.org.size()); need_alloc = need_alloc || (mbuf_cur.total_size != mbuf.total_size);
if (!need_alloc) { for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { auto * org0 = mbuf_cur.org[i]; auto * org1 = mbuf.org[i];
if (!ggml_are_same_shape(org0, org1)) { need_alloc = true; break; }
if (org0->view_src != org1->view_src || org0->view_offs != org1->view_offs) { need_alloc = true; break; } } }
if (need_alloc) { if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { mbuf_cur = std::move(mbuf);
mbuf_cur.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(mbuf_cur.ctx.get(), buft));
LLAMA_LOG_INFO("%s: allocated '%s' buffer %.3f MiB\n", __func__, ggml_backend_buft_name(buft), mbuf.total_size/1024.0/1024.0); } else { //LLAMA_LOG_INFO("%s: reallocating tensors in '%s' buffer %.3f MiB\n", __func__, ggml_backend_buft_name(buft), mbuf.total_size/1024.0/1024.0);
// save the old buffer and allocate the new tensors in it auto buf = std::move(mbuf_cur.buf);
mbuf_cur = std::move(mbuf);
ggml_tallocr talloc = ggml_tallocr_new(buf.get());
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { ggml_backend_view_init(mbuf_cur.org[i]); ggml_tallocr_alloc(&talloc, mbuf_cur.cpy[i]); }
mbuf_cur.buf = std::move(buf); } }
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { ggml_backend_tensor_copy(mbuf_cur.org[i], mbuf_cur.cpy[i]); } } }
void write(const void * src, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); } memcpy(ptr, src, size); ptr += size; size_written += size; buf_size -= size; }
void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { // save the write for later during destruction winfos.push_back({tensor, ptr, size, offset}); }
size_t n_bytes() override { return size_written; }
private: uint8_t * ptr; size_t buf_size = 0; size_t size_written = 0;
struct write_info { ggml_tensor * tensor; uint8_t * ptr; size_t size; size_t offset; }; std::vector<write_info> winfos;
llama_memory_buffers & mbufs;};
class llama_io_read_device : public llama_io_read_i {public: llama_io_read_device(const uint8_t * p, size_t len, const llama_memory_buffers & mbufs) : ptr(p), buf_size(len), mbufs(mbufs) { }
~llama_io_read_device() { llama_memory_buffers mbufs_new;
for (const auto & rinfo : rinfos) { auto * buft = ggml_backend_buffer_get_type(rinfo.tensor->buffer);
mbufs_new[buft].n_tensors++; mbufs_new[buft].total_size += rinfo.size; }
for (auto & [buft, mbuf] : mbufs_new) { ggml_init_params params = { /*.mem_size =*/ mbuf.n_tensors*ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, };
mbuf.ctx.reset(ggml_init(params));
mbuf.org.reserve(mbuf.n_tensors); }
for (const auto & rinfo : rinfos) { auto * buft = ggml_backend_buffer_get_type(rinfo.tensor->buffer);
const int64_t n = rinfo.size/ggml_element_size(rinfo.tensor);
auto & mbuf = mbufs_new[buft];
mbuf.org.push_back(ggml_view_1d(mbuf.ctx.get(), rinfo.tensor, n, rinfo.offset));
ggml_backend_view_init(mbuf.org.back()); }
for (auto & [buft, mbuf] : mbufs_new) { const auto & mbuf_cur = mbufs.at(buft);
if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) { GGML_ABORT("%s: memory buffer mismatch\n", __func__); }
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); } }
GGML_ASSERT(buf_size == 0); }
void read(void * dst, size_t size) override { if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); } memcpy(dst, ptr, size); ptr += size; size_read += size; buf_size -= size; }
void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { // save for later during destruction rinfos.push_back({tensor, ptr, size, offset}); }
size_t n_bytes() override { return size_read; }
private: const uint8_t * ptr; size_t buf_size = 0; size_t size_read = 0;
struct read_info { ggml_tensor * tensor; const uint8_t * ptr; size_t size; size_t offset; }; std::vector<read_info> rinfos;
const llama_memory_buffers & mbufs;};
size_t llama_context::state_get_size() { llama_io_write_dummy io(false); try { return state_write_data(io); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error getting state size: %s\n", __func__, err.what()); return 0; }}
size_t llama_context::state_get_data(uint8_t * dst, size_t size) { llama_io_write_host io(dst, size); try { return state_write_data(io); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; }}
size_t llama_context::state_set_data(const uint8_t * src, size_t size) { llama_io_read_host io(src, size); try { return state_read_data(io); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; }}
static constexpr uint32_t io_magic = 0xaf143cd8;
size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_flags flags) { llama_io_write_dummy io(flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); try { io.write(&io_magic, sizeof(io_magic)); io.write(&seq_id, sizeof(seq_id));
return state_seq_write_data(io, seq_id, flags); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error getting state size: %s\n", __func__, err.what()); return 0; }}
size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags) { std::unique_ptr<llama_io_write_i> io; if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { io = std::make_unique<llama_io_write_device>(dst, size, mem_storage[seq_id]); } else { io = std::make_unique<llama_io_write_host>(dst, size); }
try { io->write(&io_magic, sizeof(io_magic)); io->write(&seq_id, sizeof(seq_id));
return state_seq_write_data(*io, seq_id, flags); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; }}
size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size, llama_state_seq_flags flags) { std::unique_ptr<llama_io_read_i> io; if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { // create a temporary io to read the magic and the src seq_id io = std::make_unique<llama_io_read_host>(src, size);
uint32_t magic_read; io->read(&magic_read, sizeof(magic_read)); if (io_magic != magic_read) { throw std::runtime_error("wrong sequence state magic"); }
llama_seq_id seq_id_read; io->read(&seq_id_read, sizeof(seq_id_read));
GGML_ASSERT(mem_storage.find(seq_id_read) != mem_storage.end());
io = std::make_unique<llama_io_read_device>(src, size, mem_storage[seq_id_read]); } else { io = std::make_unique<llama_io_read_host>(src, size); }
try { uint32_t magic_read; io->read(&magic_read, sizeof(magic_read)); if (io_magic != magic_read) { throw std::runtime_error("wrong sequence state magic"); }
llama_seq_id seq_id_read; io->read(&seq_id_read, sizeof(seq_id_read));
return state_seq_read_data(*io, seq_id, flags); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; }}
bool llama_context::state_load_file(const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { llama_file file(filepath, "rb");
// sanity checks { const uint32_t magic = file.read_u32(); const uint32_t version = file.read_u32();
if (magic != LLAMA_SESSION_MAGIC || version != LLAMA_SESSION_VERSION) { LLAMA_LOG_ERROR("%s: unknown (magic, version) for session file: %08x, %08x\n", __func__, magic, version); return false; } }
// load the prompt { const uint32_t n_token_count = file.read_u32();
if (n_token_count > n_token_capacity) { LLAMA_LOG_ERROR("%s: token count in session file exceeded capacity! %u > %zu\n", __func__, n_token_count, n_token_capacity); return false; }
file.read_raw(tokens_out, sizeof(llama_token) * n_token_count); *n_token_count_out = n_token_count; }
// restore the context state { const size_t n_state_size_cur = file.size() - file.tell();
llama_io_read_file io( &file); const size_t n_read = state_read_data(io);
if (n_read != n_state_size_cur) { LLAMA_LOG_ERROR("%s: did not read all of the session file data! size %zu, got %zu\n", __func__, n_state_size_cur, n_read); return false; } }
return true;}
bool llama_context::state_save_file(const char * filepath, const llama_token * tokens, size_t n_token_count) { llama_file file(filepath, "wb");
file.write_u32(LLAMA_SESSION_MAGIC); file.write_u32(LLAMA_SESSION_VERSION);
// save the prompt file.write_u32((uint32_t) n_token_count); file.write_raw(tokens, sizeof(llama_token) * n_token_count);
// save the context state using stream saving llama_io_write_file io(&file); state_write_data(io);
return true;}
size_t llama_context::state_seq_load_file(llama_seq_id seq_id, const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { llama_file file(filepath, "rb");
// version checks { const uint32_t magic = file.read_u32(); const uint32_t version = file.read_u32();
if (magic != LLAMA_STATE_SEQ_MAGIC || version != LLAMA_STATE_SEQ_VERSION) { LLAMA_LOG_ERROR("%s: unknown (magic, version) for sequence state file: %08x, %08x\n", __func__, magic, version); return 0; } }
// load the prompt { const uint32_t n_token_count = file.read_u32();
if (tokens_out == nullptr) { const size_t n_token_max = (file.size() - file.tell()) / sizeof(llama_token); if (n_token_count > n_token_max) { LLAMA_LOG_ERROR("%s: token count in sequence state file exceeds the file size! %u > %zu\n", __func__, n_token_count, n_token_max); return 0; }
*n_token_count_out = n_token_count; return file.tell(); }
if (n_token_count > n_token_capacity) { LLAMA_LOG_ERROR("%s: token count in sequence state file exceeded capacity! %u > %zu\n", __func__, n_token_count, n_token_capacity); return 0; }
file.read_raw(tokens_out, sizeof(llama_token) * n_token_count); *n_token_count_out = n_token_count; }
// restore the context state { const size_t state_size = file.size() - file.tell(); llama_io_read_file io(&file); const size_t nread = state_seq_read_data(io, seq_id, 0); if (!nread) { LLAMA_LOG_ERROR("%s: failed to restore sequence state\n", __func__); return 0; } GGML_ASSERT(nread <= state_size); GGML_ASSERT(nread + sizeof(uint32_t) * 3 + sizeof(llama_token) * *n_token_count_out == file.tell()); }
return file.tell();}
size_t llama_context::state_seq_save_file(llama_seq_id seq_id, const char * filepath, const llama_token * tokens, size_t n_token_count) { llama_file file(filepath, "wb");
file.write_u32(LLAMA_STATE_SEQ_MAGIC); file.write_u32(LLAMA_STATE_SEQ_VERSION);
// save the prompt file.write_u32((uint32_t) n_token_count); file.write_raw(tokens, sizeof(llama_token) * n_token_count);
// save the context state using stream saving llama_io_write_file io(&file); state_seq_write_data(io, seq_id, 0);
const size_t res = file.tell(); GGML_ASSERT(res == sizeof(uint32_t) * 3 + sizeof(llama_token) * n_token_count + io.n_bytes());
return res;}
size_t llama_context::state_write_data(llama_io_write_i & io) { LLAMA_LOG_DEBUG("%s: writing state\n", __func__);
// write model info { LLAMA_LOG_DEBUG("%s: - writing model info\n", __func__);
const std::string arch_str = llm_arch_name(model.arch); io.write_string(arch_str); // TODO: add more model-specific info which should prevent loading the session file if not identical }
if (memory != nullptr) { LLAMA_LOG_DEBUG("%s: - writing memory module\n", __func__); memory->state_write(io); }
return io.n_bytes();}
size_t llama_context::state_read_data(llama_io_read_i & io) { LLAMA_LOG_DEBUG("%s: reading state\n", __func__);
// read model info { LLAMA_LOG_DEBUG("%s: - reading model info\n", __func__);
const std::string cur_arch_str = llm_arch_name(model.arch);
std::string arch_str; io.read_string(arch_str); if (cur_arch_str != arch_str) { throw std::runtime_error(format("wrong model arch: '%s' instead of '%s'", arch_str.c_str(), cur_arch_str.c_str())); } // TODO: add more info which needs to be identical but which is not verified otherwise }
if (memory) { LLAMA_LOG_DEBUG("%s: - reading memory module\n", __func__);
memory->state_read(io); }
return io.n_bytes();}
size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { if (memory) { memory->state_write(io, seq_id, flags); }
return io.n_bytes();}
size_t llama_context::state_seq_read_data(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { if (memory) { memory->state_read(io, seq_id, flags); }
return io.n_bytes();}
//// perf//
llama_perf_context_data llama_context::perf_get_data() const { llama_perf_context_data data = {};
data.t_start_ms = 1e-3 * t_start_us; data.t_load_ms = 1e-3 * t_load_us; data.t_p_eval_ms = 1e-3 * t_p_eval_us; data.t_eval_ms = 1e-3 * t_eval_us; data.n_p_eval = std::max(1, n_p_eval); data.n_eval = std::max(1, n_eval); data.n_reused = std::max(0, n_reused);
return data;}
void llama_context::perf_reset() { t_start_us = ggml_time_us(); t_eval_us = n_eval = 0; t_p_eval_us = n_p_eval = 0; n_reused = 0;}
llama_memory_breakdown llama_context::memory_breakdown() const { std::map<ggml_backend_buffer_type_t, llama_memory_breakdown_data> ret; for (const auto & [buft, size] : model.memory_breakdown()) { ret[buft].model += size; } if (memory) { for (const auto & [buft, size] : memory->memory_breakdown()) { ret[buft].context += size; } } if (model.hparams.no_alloc) { for (size_t i = 0; i < backends.size(); ++i) { ggml_backend_t backend = backends[i].get(); ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend); ret[buft].compute += backend_buf_exp_size[i]; } } else { for (const auto & backend_ptr : backends) { ggml_backend_t backend = backend_ptr.get(); ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend); ret[buft].compute += ggml_backend_sched_get_buffer_size(sched.get(), backend); } } return ret;}
//// training//
static void llama_set_param(struct ggml_tensor * tensor, llama_opt_param_filter param_filter, void * userdata) { if (!tensor || tensor->type != GGML_TYPE_F32) { return; } if (!param_filter(tensor, userdata)) { return; } if (strcmp(tensor->name, "token_embd.weight") == 0) { return; // FIXME } if (strcmp(tensor->name, "rope_freqs.weight") == 0) { return; // FIXME } ggml_set_param(tensor);}
void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { GGML_ASSERT(!opt_ctx); model->hparams.n_ctx_train = lopt_params.n_ctx_train > 0 ? lopt_params.n_ctx_train : n_ctx(); const uint32_t n_batch = std::min(this->n_batch(), model->hparams.n_ctx_train); const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch); GGML_ASSERT(model->hparams.n_ctx_train % n_batch == 0); GGML_ASSERT(n_batch % n_ubatch == 0);
ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); opt_params.opt_period = n_batch / n_ubatch; opt_params.get_opt_pars = lopt_params.get_opt_pars; opt_params.get_opt_pars_ud = lopt_params.get_opt_pars_ud; opt_params.optimizer = lopt_params.optimizer_type; opt_ctx = ggml_opt_init(opt_params);
llama_opt_param_filter param_filter = lopt_params.param_filter; void * param_filter_ud = lopt_params.param_filter_ud;
//llama_set_param(model->tok_embd, param_filter, param_filter_ud); // FIXME llama_set_param(model->type_embd, param_filter, param_filter_ud); llama_set_param(model->pos_embd, param_filter, param_filter_ud); llama_set_param(model->tok_norm, param_filter, param_filter_ud); llama_set_param(model->tok_norm_b, param_filter, param_filter_ud); llama_set_param(model->output_norm, param_filter, param_filter_ud); llama_set_param(model->output_norm_b, param_filter, param_filter_ud); llama_set_param(model->output, param_filter, param_filter_ud); llama_set_param(model->output_b, param_filter, param_filter_ud); llama_set_param(model->output_norm_enc, param_filter, param_filter_ud); llama_set_param(model->cls, param_filter, param_filter_ud); llama_set_param(model->cls_b, param_filter, param_filter_ud); llama_set_param(model->cls_out, param_filter, param_filter_ud); llama_set_param(model->cls_out_b, param_filter, param_filter_ud); llama_set_param(model->cls_norm, param_filter, param_filter_ud);
for (struct llama_layer & layer : model->layers) { for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { llama_set_param(reinterpret_cast<struct ggml_tensor **>(&layer)[i], param_filter, param_filter_ud); } }}
void llama_context::opt_epoch_iter( ggml_opt_dataset_t dataset, ggml_opt_result_t result, const std::vector<llama_token> & tokens, const std::vector<llama_token> & labels_sparse, llama_batch & batch, ggml_opt_epoch_callback callback, bool train, int64_t idata_in_loop, int64_t ndata_in_loop, int64_t t_loop_start) { GGML_ASSERT(opt_ctx); const uint32_t n_ctx = llama_model_n_ctx_train(&model); const uint32_t n_batch = std::min(this->n_batch(), n_ctx); const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch);
memory->clear(true);
for (uint32_t pos_ctx = 0; pos_ctx < n_ctx; pos_ctx += n_batch) { batch.n_tokens = n_batch; for (uint32_t pos_batch = 0; pos_batch < n_batch; ++pos_batch) { batch.token [pos_batch] = tokens[pos_ctx + pos_batch]; batch.pos [pos_batch] = pos_ctx + pos_batch; batch.n_seq_id[pos_batch] = 1; batch.seq_id [pos_batch][0] = 0; batch.logits [pos_batch] = true; }
if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return; }
const uint32_t n_tokens_all = balloc->get_n_tokens();
n_queued_tokens += n_tokens_all;
embd_seq.clear();
uint32_t n_outputs_all = n_tokens_all;
auto mctx = memory->init_batch(*balloc, cparams.n_ubatch, true); if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); break; }
// reserve output buffer if (output_reserve(n_outputs_all) < n_outputs_all) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); GGML_ABORT("TODO: handle this error"); };
uint32_t pos_batch = 0; do { const auto & ubatch = mctx->get_ubatch();
n_outputs = ubatch.n_tokens;
if (!mctx->apply()) { LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); break; }
auto * res = gf_res_prev.get();
const auto gparams = graph_params(res, ubatch, mctx.get(), ctx_type_to_graph_type(cparams.ctx_type));
res->reset();
auto * gf = model.build_graph(gparams);
struct ggml_context * ctx_compute_opt; { const size_t size_gf = ggml_graph_size(gf); const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); struct ggml_init_params params = { /*.mem_size =*/ size_meta, /*.mem_buffer =*/ nullptr, /*.no_alloc =*/ true, }; ctx_compute_opt = ggml_init(params); } ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); ggml_opt_alloc(opt_ctx, train);
res->set_inputs(&ubatch); { struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); GGML_ASSERT(labels->ne[1] == n_ubatch); ggml_set_zero(labels); const float onef = 1.0f; for (uint32_t pos_ubatch = 0; pos_ubatch < n_ubatch; ++pos_ubatch) { const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; GGML_ASSERT(labels_sparse[ilabel] < labels->ne[0]); ggml_backend_tensor_set(labels, &onef, (pos_ubatch*labels->ne[0] + labels_sparse[ilabel])*sizeof(float), sizeof(float)); } } ggml_opt_eval(opt_ctx, result); if (callback) { callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_ubatch + 1, ndata_in_loop, t_loop_start); } ggml_free(ctx_compute_opt);
pos_batch += ubatch.n_tokens; } while (mctx->next()); }}
void llama_context::opt_epoch( ggml_opt_dataset_t dataset, ggml_opt_result_t result_train, ggml_opt_result_t result_eval, int64_t idata_split, ggml_opt_epoch_callback callback_train, ggml_opt_epoch_callback callback_eval) { const uint32_t n_ctx = this->n_ctx(); const uint32_t n_batch = std::min(cparams.n_batch, n_ctx); const uint32_t n_ubatch = std::min(cparams.n_ubatch, n_batch); const int64_t ndata = ggml_opt_dataset_ndata(dataset);
GGML_ASSERT(idata_split >= 0); GGML_ASSERT(idata_split <= ndata);
const uint32_t ubatch_per_ctx = n_ctx / n_ubatch;
struct llama_batch batch = llama_batch_init(n_batch, 0, 1); std::vector<llama_token> tokens(n_ctx); std::vector<llama_token> labels_sparse(n_ctx);
int64_t idata = 0;
int64_t t_loop_start = ggml_time_us(); int64_t ndata_in_loop = idata_split*ubatch_per_ctx; for (; idata < idata_split; ++idata) { constexpr bool train = true; const int64_t idata_in_loop = idata*ubatch_per_ctx;
ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); opt_epoch_iter(dataset, result_train, tokens, labels_sparse, batch, callback_train, train, idata_in_loop, ndata_in_loop, t_loop_start); }
t_loop_start = ggml_time_us(); ndata_in_loop = (ndata - idata_split)*ubatch_per_ctx; for (; idata < ndata; ++idata) { constexpr bool train = false; const int64_t idata_in_loop = (idata - idata_split)*ubatch_per_ctx;
ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); opt_epoch_iter(dataset, result_eval, tokens, labels_sparse, batch, callback_eval, train, idata_in_loop, ndata_in_loop, t_loop_start); }
llama_batch_free(batch);}
//// interface implementation//
llama_context_params llama_context_default_params() { llama_context_params result = { /*.n_ctx =*/ 512, /*.n_batch =*/ 2048, /*.n_ubatch =*/ 512, /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, /*.n_outputs_max =*/ 0, /*.n_outputs_max_per_seq =*/ 1, /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, /*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT, /*.rope_scaling_type =*/ LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED, /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, /*.flash_attn_type =*/ LLAMA_FLASH_ATTN_TYPE_AUTO, /*.rope_freq_base =*/ 0.0f, /*.rope_freq_scale =*/ 0.0f, /*.yarn_ext_factor =*/ -1.0f, /*.yarn_attn_factor =*/ -1.0f, /*.yarn_beta_fast =*/ -1.0f, /*.yarn_beta_slow =*/ -1.0f, /*.yarn_orig_ctx =*/ 0, /*.defrag_thold =*/ -1.0f, /*.cb_eval =*/ nullptr, /*.cb_eval_user_data =*/ nullptr, /*.type_k =*/ GGML_TYPE_F16, /*.type_v =*/ GGML_TYPE_F16, /*.path_kv_mean_center =*/ nullptr, /*.abort_callback =*/ nullptr, /*.abort_callback_data =*/ nullptr, /*.embeddings =*/ false, /*.offload_kqv =*/ true, /*.no_perf =*/ true, /*.op_offload =*/ true, /*.swa_full =*/ true, /*.kv_unified =*/ false, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, };
return result;}
llama_context * llama_init_from_model( llama_model * model, llama_context_params params) { if (!model) { LLAMA_LOG_ERROR("%s: model cannot be NULL\n", __func__); return nullptr; }
if (params.n_batch == 0 && params.n_ubatch == 0) { LLAMA_LOG_ERROR("%s: n_batch and n_ubatch cannot both be zero\n", __func__); return nullptr; }
if (params.n_ctx == 0 && model->hparams.n_ctx_train == 0) { LLAMA_LOG_ERROR("%s: n_ctx and model->hparams.n_ctx_train cannot both be zero\n", __func__); return nullptr; }
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && model->arch == LLM_ARCH_GROK) { LLAMA_LOG_WARN("%s: flash_attn is not compatible with Grok - forcing off\n", __func__); params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; }
if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR) { if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) { LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for SPLIT_MODE_TENSOR\n", __func__); params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; } if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) { LLAMA_LOG_ERROR("%s: SPLIT_MODE_TENSOR requires flash_attn to be enabled\n", __func__); return nullptr; } }
if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) { LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v)); return nullptr; }
if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) { if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) { LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__); params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; } if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) { LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\n", __func__); return nullptr; } }
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) { const uint32_t blck_size = ggml_blck_size(params.type_k); for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) { if (model->hparams.n_embd_head_k(il) % blck_size != 0) { LLAMA_LOG_ERROR("%s: K cache type %s with block size %u does not divide n_embd_head_k=%u\n", __func__, ggml_type_name(params.type_k), blck_size, model->hparams.n_embd_head_k(il)); return nullptr; } } }
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_v)) { const uint32_t blck_size = ggml_blck_size(params.type_v); for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) { if (model->hparams.n_embd_head_v(il) % blck_size != 0) { LLAMA_LOG_ERROR("%s: V cache type %s with block size %u does not divide n_embd_head_v=%u\n", __func__, ggml_type_name(params.type_v), blck_size, model->hparams.n_embd_head_v(il)); return nullptr; } } }
if (params.path_kv_mean_center != nullptr && params.type_k != GGML_TYPE_Q4_0) { LLAMA_LOG_ERROR("%s: path_kv_mean_center requires the K cache type to be Q4_0 (got %s)\n", __func__, ggml_type_name(params.type_k)); return nullptr; }
if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != model->hparams.pooling_type) { //user-specified pooling-type is different from the model default LLAMA_LOG_WARN("%s: model default pooling_type is [%d], but [%d] was specified\n", __func__, model->hparams.pooling_type, params.pooling_type); }
// router_layer >= 0 means n_layer_nextn is repurposed for a router layer, not real MTP if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && (model->hparams.n_layer_nextn == 0 || model->hparams.router_layer >= 0)) { LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__); return nullptr; }
try { auto * ctx = new llama_context(*model, params); return ctx; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: failed to initialize the context: %s\n", __func__, err.what()); }
return nullptr;}
// deprecatedllama_context * llama_new_context_with_model( llama_model * model, llama_context_params params) { return llama_init_from_model(model, params);}
void llama_free(llama_context * ctx) { delete ctx;}
uint32_t llama_n_ctx(const llama_context * ctx) { return ctx->n_ctx();}
uint32_t llama_n_ctx_seq(const llama_context * ctx) { return ctx->n_ctx_seq();}
uint32_t llama_n_batch(const llama_context * ctx) { return ctx->n_batch();}
uint32_t llama_n_ubatch(const llama_context * ctx) { return ctx->n_ubatch();}
uint32_t llama_n_seq_max(const llama_context * ctx) { return ctx->n_seq_max();}
uint32_t llama_n_rs_seq(const llama_context * ctx) { return ctx->get_cparams().n_rs_seq;}
const llama_model * llama_get_model(const llama_context * ctx) { return &ctx->get_model();}
enum llama_pooling_type llama_pooling_type(const llama_context * ctx) { return ctx->pooling_type();}
void llama_attach_threadpool( llama_context * ctx, ggml_threadpool_t threadpool, ggml_threadpool_t threadpool_batch) { ctx->attach_threadpool(threadpool, threadpool_batch);}
void llama_detach_threadpool(llama_context * ctx) { ctx->detach_threadpool();}
void llama_set_n_threads(llama_context * ctx, int32_t n_threads, int32_t n_threads_batch) { ctx->set_n_threads(n_threads, n_threads_batch);}
int32_t llama_n_threads(llama_context * ctx) { return ctx->n_threads();}
int32_t llama_n_threads_batch(llama_context * ctx) { return ctx->n_threads_batch();}
void llama_set_abort_callback(llama_context * ctx, bool (*abort_callback)(void * data), void * abort_callback_data) { ctx->set_abort_callback(abort_callback, abort_callback_data);}
void llama_set_embeddings(llama_context * ctx, bool embeddings) { ctx->set_embeddings(embeddings);}
void llama_set_causal_attn(llama_context * ctx, bool causal_attn) { ctx->set_causal_attn(causal_attn);}
void llama_set_warmup(llama_context * ctx, bool warmup) { ctx->set_warmup(warmup);}
void llama_synchronize(llama_context * ctx) { ctx->synchronize();}
float * llama_get_logits(llama_context * ctx) { ctx->synchronize();
return ctx->get_logits();}
float * llama_get_logits_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
float * res = nullptr;
res = ctx->get_sampled_logits_ith(i);
if (!res) { res = ctx->get_logits_ith(i); }
return res;}
float * llama_get_embeddings(llama_context * ctx) { ctx->synchronize();
return ctx->get_embeddings();}
float * llama_get_embeddings_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return ctx->get_embeddings_ith(i);}
float * llama_get_embeddings_seq(llama_context * ctx, llama_seq_id seq_id) { ctx->synchronize();
return ctx->get_embeddings_seq(seq_id);}
// the attention part of a context's memory, or nullptrstatic llama_kv_cache * scr_kv(llama_memory_i * mem) { if (auto * hybrid = dynamic_cast<llama_memory_hybrid *>(mem)) { return hybrid->get_mem_attn(); } return dynamic_cast<llama_kv_cache *>(mem);}
void llama_context::set_seq_spliced(llama_seq_id seq_id, bool value) { GGML_ASSERT(seq_id >= 0 && seq_id < LLAMA_MAX_SEQ); balloc->seq_spliced[seq_id] = value;}
bool llama_context::memory_relocate(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { auto * kv = scr_kv(memory.get()); if (kv == nullptr) { return false; }
synchronize();
// the relocation graph resets the scheduler, so the previous graph cannot be reused gf_res_prev->reset();
return kv->relocate(this, seq_id, p0, p1, delta);}
bool llama_scr_kv_relocate(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { return ctx->memory_relocate(seq_id, p0, p1, delta);}
void llama_scr_set_spliced(llama_context * ctx, llama_seq_id seq_id, bool value) { ctx->set_seq_spliced(seq_id, value);}
bool llama_scr_kv_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { auto * kv = scr_kv(ctx->get_memory()); return kv != nullptr && kv->seq_rm(seq_id, p0, p1);}
bool llama_scr_recr_set_pos(llama_context * ctx, llama_seq_id seq_id, llama_pos pos) { auto * hybrid = dynamic_cast<llama_memory_hybrid *>(ctx->get_memory()); if (hybrid == nullptr) { return false; }
auto * recr = hybrid->get_mem_recr();
if (pos < 0) { return recr->seq_rm(seq_id, -1, -1); }
const llama_pos cur = recr->seq_pos_max(seq_id); if (cur < 0) { return false; }
recr->seq_add(seq_id, cur, cur + 1, pos - cur);
return recr->seq_pos_max(seq_id) == pos;}
void llama_set_embeddings_nextn(llama_context * ctx, bool value, bool masked) { ctx->set_embeddings_nextn(value, masked);}
void llama_set_embeddings_layer_inp(llama_context * ctx, uint32_t lid, bool value) { ctx->set_embeddings_layer_inp(lid, value);}
void llama_set_nextn_layer_offset(llama_context * ctx, int32_t offset) { ctx->set_nextn_layer_offset(offset);}
llama_memory_t llama_get_memory(const struct llama_context * ctx) { if (!ctx) { return nullptr; }
return ctx->get_memory();}
float * llama_get_embeddings_nextn(llama_context * ctx) { ctx->synchronize();
return ctx->get_embeddings_nextn();}
float * llama_get_embeddings_nextn_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return ctx->get_embeddings_nextn_ith(i);}
float * llama_get_embeddings_layer_inp(llama_context * ctx, uint32_t lid) { ctx->synchronize();
return ctx->get_embeddings_layer_inp(lid);}
// multi-layer hidden-state tap C API (staging) -------------------------------
void llama_set_capture_layers(llama_context * ctx, const int32_t * layer_ids, size_t n_layers) { std::vector<int32_t> ids; ids.reserve(n_layers); for (size_t i = 0; i < n_layers; ++i) { ids.push_back(layer_ids[i]); } ctx->set_capture_layers(ids);}
uint32_t llama_get_n_capture(llama_context * ctx) { return ctx->get_n_capture();}
float * llama_get_embeddings_capture(llama_context * ctx) { ctx->synchronize(); return ctx->get_embeddings_capture();}
float * llama_get_embeddings_capture_ith(llama_context * ctx, int32_t i) { ctx->synchronize(); return ctx->get_embeddings_capture_ith(i);}
// dspark drafter target-context staging C API --------------------------------
void llama_set_dspark_ctx(llama_context * ctx, const float * feat, int64_t n_ctx_rows, int64_t n_embd_cap, const int32_t * pos) { ctx->set_dspark_ctx(feat, n_ctx_rows, n_embd_cap, pos);}
bool llama_set_sampler(llama_context * ctx, llama_seq_id seq_id, llama_sampler * smpl) { return ctx->set_sampler(seq_id, smpl);}
llama_token llama_get_sampled_token_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return ctx->get_sampled_token_ith(i);}
float * llama_get_sampled_probs_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return ctx->get_sampled_probs_ith(i);}
float * llama_get_sampled_logits_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return ctx->get_sampled_logits_ith(i);}
llama_token * llama_get_sampled_candidates_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return const_cast<llama_token *>(ctx->get_sampled_candidates_ith(i));}
uint32_t llama_get_sampled_candidates_count_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return static_cast<uint32_t>(ctx->get_sampled_candidates_count(i));}
uint32_t llama_get_sampled_logits_count_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return static_cast<uint32_t>(ctx->get_sampled_logits_count(i));}
uint32_t llama_get_sampled_probs_count_ith(llama_context * ctx, int32_t i) { ctx->synchronize();
return static_cast<uint32_t>(ctx->get_sampled_probs_count(i));}
struct ggml_cgraph * llama_graph_reserve( struct llama_context * ctx, uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs) { auto memory = ctx->get_memory(); llama_memory_context_ptr mctx; if (memory) { mctx = memory->init_full(); } return ctx->graph_reserve(n_tokens, n_seqs, n_outputs, mctx.get());}
// llama adapter API
int32_t llama_set_adapters_lora( llama_context * ctx, llama_adapter_lora ** adapters, size_t n_adapters, float * scales) { if (adapters == nullptr || scales == nullptr) { GGML_ASSERT(n_adapters == 0 && "invalid llama_set_adapters_lora call"); }
ctx->set_adapters_lora(adapters, n_adapters, scales);
return 0;}
int32_t llama_set_adapter_cvec( llama_context * ctx, const float * data, size_t len, int32_t n_embd, int32_t il_start, int32_t il_end) { bool res = ctx->set_adapter_cvec(data, len, n_embd, il_start, il_end);
return res ? 0 : -1;}
//// memory//
void llama_memory_clear(llama_memory_t mem, bool data) { if (!mem) { return; }
mem->clear(data);}
bool llama_memory_seq_rm( llama_memory_t mem, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { if (!mem) { return true; }
return mem->seq_rm(seq_id, p0, p1);}
void llama_memory_seq_cp( llama_memory_t mem, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { if (!mem) { return; }
mem->seq_cp(seq_id_src, seq_id_dst, p0, p1);}
void llama_memory_seq_keep( llama_memory_t mem, llama_seq_id seq_id) { if (!mem) { return; }
mem->seq_keep(seq_id);}
void llama_memory_seq_add( llama_memory_t mem, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { if (!mem) { return; }
mem->seq_add(seq_id, p0, p1, delta);}
void llama_memory_seq_div( llama_memory_t mem, llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { if (!mem) { return; }
mem->seq_div(seq_id, p0, p1, d);}
llama_pos llama_memory_seq_pos_min( llama_memory_t mem, llama_seq_id seq_id) { if (!mem) { return -1; }
return mem->seq_pos_min(seq_id);}
llama_pos llama_memory_seq_pos_max( llama_memory_t mem, llama_seq_id seq_id) { if (!mem) { return -1; }
return mem->seq_pos_max(seq_id);}
bool llama_memory_can_shift(llama_memory_t mem) { if (!mem) { return false; }
return mem->get_can_shift();}
// llama state API
// deprecatedsize_t llama_get_state_size(llama_context * ctx) { return llama_state_get_size(ctx);}
// deprecatedsize_t llama_copy_state_data(llama_context * ctx, uint8_t * dst) { return llama_state_get_data(ctx, dst, -1);}
// deprecatedsize_t llama_set_state_data(llama_context * ctx, const uint8_t * src) { return llama_state_set_data(ctx, src, -1);}
// deprecatedbool llama_load_session_file(llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { return llama_state_load_file(ctx, path_session, tokens_out, n_token_capacity, n_token_count_out);}
// deprecatedbool llama_save_session_file(llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count) { return llama_state_save_file(ctx, path_session, tokens, n_token_count);}
// Returns the *actual* size of the state.// Intended to be used when saving to state to a buffer.size_t llama_state_get_size(llama_context * ctx) { return ctx->state_get_size();}
size_t llama_state_get_data(llama_context * ctx, uint8_t * dst, size_t size) { ctx->synchronize();
return ctx->state_get_data(dst, size);}
// Sets the state reading from the specified source addresssize_t llama_state_set_data(llama_context * ctx, const uint8_t * src, size_t size) { ctx->synchronize();
return ctx->state_set_data(src, size);}
bool llama_state_load_file(llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { ctx->synchronize();
try { return ctx->state_load_file(path_session, tokens_out, n_token_capacity, n_token_count_out); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading session file: %s\n", __func__, err.what()); return false; }}
bool llama_state_save_file(llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count) { ctx->synchronize();
try { return ctx->state_save_file(path_session, tokens, n_token_count); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving session file: %s\n", __func__, err.what()); return false; }}
size_t llama_state_seq_get_size(llama_context * ctx, llama_seq_id seq_id) { return llama_state_seq_get_size_ext(ctx, seq_id, 0);}
size_t llama_state_seq_get_data(llama_context * ctx, uint8_t * dst, size_t size, llama_seq_id seq_id) { return llama_state_seq_get_data_ext(ctx, dst, size, seq_id, 0);}
size_t llama_state_seq_set_data(llama_context * ctx, const uint8_t * src, size_t size, llama_seq_id seq_id) { return llama_state_seq_set_data_ext(ctx, src, size, seq_id, 0);}
size_t llama_state_seq_get_size_ext(llama_context * ctx, llama_seq_id seq_id, llama_state_seq_flags flags) { return ctx->state_seq_get_size(seq_id, flags);}
size_t llama_state_seq_get_data_ext(llama_context * ctx, uint8_t * dst, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { ctx->synchronize();
return ctx->state_seq_get_data(seq_id, dst, size, flags);}size_t llama_state_seq_set_data_ext(llama_context * ctx, const uint8_t * src, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { ctx->synchronize();
return ctx->state_seq_set_data(seq_id, src, size, flags);}
size_t llama_state_seq_save_file(llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count) { ctx->synchronize();
try { return ctx->state_seq_save_file(seq_id, filepath, tokens, n_token_count); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving sequence state file: %s\n", __func__, err.what()); return 0; }}
size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { ctx->synchronize();
try { return ctx->state_seq_load_file(dest_seq_id, filepath, tokens_out, n_token_capacity, n_token_count_out); } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading sequence state file: %s\n", __func__, err.what()); return 0; }}
///
int32_t llama_encode( llama_context * ctx, llama_batch batch) { const int ret = ctx->encode(batch); if (ret != 0) { LLAMA_LOG_ERROR("%s: failed to encode, ret = %d\n", __func__, ret); }
return ret;}
int32_t llama_decode( llama_context * ctx, llama_batch batch) { const int ret = ctx->decode(batch); if (ret != 0 && ret != 1) { LLAMA_LOG_ERROR("%s: failed to decode, ret = %d\n", __func__, ret); }
return ret;}
//// perf//
llama_perf_context_data llama_perf_context(const llama_context * ctx) { llama_perf_context_data data = {};
if (ctx == nullptr) { return data; }
data = ctx->perf_get_data();
return data;}
void llama_perf_context_print(const llama_context * ctx) { const auto data = llama_perf_context(ctx);
const double t_end_ms = 1e-3 * ggml_time_us();
LLAMA_LOG_INFO("%s: load time = %10.2f ms\n", __func__, data.t_load_ms); LLAMA_LOG_INFO("%s: prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n", __func__, data.t_p_eval_ms, data.n_p_eval, data.t_p_eval_ms / data.n_p_eval, 1e3 / data.t_p_eval_ms * data.n_p_eval); LLAMA_LOG_INFO("%s: eval time = %10.2f ms / %5d runs (%8.2f ms per token, %8.2f tokens per second)\n", __func__, data.t_eval_ms, data.n_eval, data.t_eval_ms / data.n_eval, 1e3 / data.t_eval_ms * data.n_eval); LLAMA_LOG_INFO("%s: total time = %10.2f ms / %5d tokens\n", __func__, (t_end_ms - data.t_start_ms), (data.n_p_eval + data.n_eval)); LLAMA_LOG_INFO("%s: graphs reused = %10d\n", __func__, data.n_reused);}
void llama_perf_context_reset(llama_context * ctx) { ctx->perf_reset();}
//// training//
bool llama_opt_param_filter_all(const struct ggml_tensor * tensor, void * userdata) { GGML_UNUSED(tensor); GGML_UNUSED(userdata); return true;}
void llama_opt_init(struct llama_context * ctx, struct llama_model * model, struct llama_opt_params lopt_params) { ctx->opt_init(model, lopt_params);}
void llama_opt_epoch( struct llama_context * ctx, ggml_opt_dataset_t dataset, ggml_opt_result_t result_train, ggml_opt_result_t result_eval, int64_t idata_split, ggml_opt_epoch_callback callback_train, ggml_opt_epoch_callback callback_eval) { ctx->opt_epoch( dataset, result_train, result_eval, idata_split, callback_train, callback_eval);}
//// ext//
llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx) { return ctx->memory_breakdown();}
llama_context * llama_get_ctx_other(struct llama_context * ctx) { return ctx->get_cparams().ctx_other;}