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.
145 kB · 4032 lines
C++
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033#include "llama-graph.h"
#include "llama-impl.h"#include "llama-model.h"#include "llama-batch.h"#include "llama-cparams.h"#include "llama-sampler.h"
#include "llama-kv-cache.h"#include "llama-kv-cache-iswa.h"#include "llama-kv-cache-dsa.h"#include "llama-kv-cache-dsa-iswa.h"#include "llama-kv-cache-msa.h"#include "llama-kv-cache-dsv4.h"#include "llama-memory-hybrid.h"#include "llama-memory-hybrid-iswa.h"#include "llama-memory-recurrent.h"
#include <cassert>#include <cmath>#include <cstring>#include <numeric>#include <sstream>#include <string>#include <unordered_set>
// dedup helpers
static ggml_tensor * build_attn_inp_kq_mask( ggml_context * ctx, const llama_kv_cache_context * mctx, const llama_ubatch & ubatch, const llama_cparams & cparams) { const auto n_kv = mctx->get_n_kv(); const auto n_tokens = ubatch.n_tokens; const auto n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq;
// flash attention requires an f16 mask const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(res); ggml_set_name(res, "attn_inp_kq_mask");
return res;}
static bool can_reuse_kq_mask( ggml_tensor * kq_mask, const llama_kv_cache_context * mctx, const llama_ubatch & ubatch, const llama_cparams & cparams) { const auto n_kv = mctx->get_n_kv(); const auto n_tokens = ubatch.n_tokens; const auto n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq;
bool res = true;
res &= (kq_mask->ne[0] == n_kv); res &= (kq_mask->ne[1] == n_tokens/n_stream); res &= (kq_mask->ne[2] == 1); res &= (kq_mask->ne[3] == n_stream);
return res;}
// impl
void llm_graph_input_dspark_ctx::set_input(const llama_ubatch * ubatch) { if (correction_prefix_rows > 0) { if (!ctx_feat || !dctx || dctx->n_ctx_rows != ctx_feat->ne[1]) { throw std::runtime_error("DSpark correction prefix requires staged context inputs"); } for (int64_t row = 0; row < ubatch->n_tokens; ++row) { const bool expected = row >= dctx->n_ctx_rows && row < dctx->n_ctx_rows + correction_prefix_rows; if (!ubatch->output || (ubatch->output[row] != 0) != expected) { throw std::runtime_error("DSpark correction prefix requires matching output rows"); } } } if (reuse_mask_id >= 0) { if (!ubatch->token || !dctx || dctx->n_ctx_rows != ctx_feat->ne[1]) { throw std::runtime_error("DSpark mask reuse requires staged token inputs"); } for (int64_t i = dctx->n_ctx_rows + 1; i < ubatch->n_tokens; ++i) { if (ubatch->token[i] != reuse_mask_id) { throw std::runtime_error("DSpark mask reuse requires an anchor followed by MASK tokens"); } } } // ignores ubatch entirely (like llm_graph_input_cross_embd): the context // feature row count (n_ctx_rows) is independent of the current ubatch's // token count and comes purely from the staged llama_dspark_ctx. if (ctx_feat && dctx && !dctx->v_ctx_feat.empty()) { GGML_ASSERT((int64_t) dctx->v_ctx_feat.size() == ggml_nelements(ctx_feat)); ggml_backend_tensor_set(ctx_feat, dctx->v_ctx_feat.data(), 0, ggml_nbytes(ctx_feat)); }}
void llm_graph_input_dspark_logsnr::set_input(const llama_ubatch * ubatch) { // v_feat was precomputed at graph-build time and does not depend on the ubatch GGML_UNUSED(ubatch); if (feat && !v_feat.empty()) { GGML_ASSERT((int64_t) v_feat.size() == ggml_nelements(feat)); ggml_backend_tensor_set(feat, v_feat.data(), 0, ggml_nbytes(feat)); }}
bool llm_graph_input_dspark_logsnr::can_reuse(const llm_graph_params & params) { // v_feat is a function of n_tokens and n_seqs_unq, both already checked by // llm_graph_params::allow_reuse, and of min/max_log_snr, fixed per model return feat && feat->ne[1] == params.ubatch.n_tokens;}
void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { if (ubatch->token) { const int64_t n_tokens = ubatch->n_tokens;
ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens)); }
if (ubatch->embd) { GGML_ASSERT(n_embd == embd->ne[0]);
const int64_t n_tokens = ubatch->n_tokens;
ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(embd)); }}
bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) { bool res = true;
res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens);
return res;}
void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) { const int64_t n_tokens = ubatch->n_tokens;
if (ubatch->token) { ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens)); } else { // note: mtmd embedding input goes through here GGML_ASSERT(ubatch->embd); GGML_ASSERT(n_embd == embd->ne[0]);
ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h)); }
// TODO: extend llama_ubatch to differentiate between token embeddings and hidden states // for now, we assume that the hidden state is always provided as an embedding // ref: https://github.com/ggml-org/llama.cpp/pull/23643 if (ubatch->embd) { GGML_ASSERT(n_embd == h->ne[0]);
ggml_backend_tensor_set(h, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h)); }}
bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) { bool res = true;
res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens); res &= (!params.ubatch.embd) || (h && h->ne[1] == params.ubatch.n_tokens);
return res;}
void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) { if (ubatch->pos && pos) { const int64_t n_tokens = ubatch->n_tokens;
if (ubatch->token && n_pos_per_embd == 4) { // in case we're using M-RoPE with text tokens, convert the 1D positions to 4D // the 3 first dims are the same, and 4th dim is all 0 std::vector<llama_pos> pos_data(n_tokens*n_pos_per_embd); // copy the first dimension for (int i = 0; i < n_tokens; ++i) { pos_data[ i] = ubatch->pos[i]; pos_data[ n_tokens + i] = ubatch->pos[i]; pos_data[2 * n_tokens + i] = ubatch->pos[i]; pos_data[3 * n_tokens + i] = 0; // 4th dim is 0 } ggml_backend_tensor_set(pos, pos_data.data(), 0, pos_data.size()*ggml_element_size(pos)); } else { ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos)); } }}
bool llm_graph_input_pos::can_reuse(const llm_graph_params & params) { bool res = true;
res &= pos->ne[0] == params.ubatch.n_tokens*n_pos_per_embd;
return res;}
void llm_graph_input_attn_temp::set_input(const llama_ubatch * ubatch) { if (ubatch->pos && attn_scale) { const int64_t n_tokens = ubatch->n_tokens;
GGML_ASSERT(f_attn_temp_scale != 0.0f); GGML_ASSERT(n_attn_temp_floor_scale != 0);
std::vector<float> attn_scale_data(n_tokens, 0.0f); for (int i = 0; i < n_tokens; ++i) { const float pos = ubatch->pos[i]; attn_scale_data[i] = std::log( std::floor((pos + f_attn_temp_offset) / n_attn_temp_floor_scale) + 1.0 ) * f_attn_temp_scale + 1.0; }
ggml_backend_tensor_set(attn_scale, attn_scale_data.data(), 0, n_tokens*ggml_element_size(attn_scale)); }}
void llm_graph_input_pos_bucket::set_input(const llama_ubatch * ubatch) { if (pos_bucket) { const int64_t n_tokens = ubatch->n_tokens;
GGML_ASSERT(ggml_backend_buffer_is_host(pos_bucket->buffer)); GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing
int32_t * data = (int32_t *) pos_bucket->data;
for (int j = 0; j < n_tokens; ++j) { for (int i = 0; i < n_tokens; ++i) { data[j*n_tokens + i] = llama_relative_position_bucket(ubatch->pos[i], ubatch->pos[j], hparams.n_rel_attn_bkts, true); } } }}
void llm_graph_input_pos_bucket_kv::set_input(const llama_ubatch * ubatch) { if (pos_bucket) { mctx->set_input_pos_bucket(pos_bucket, ubatch); }}
void llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(out_ids);
const int64_t n_tokens = ubatch->n_tokens;
GGML_ASSERT(ggml_backend_buffer_is_host(out_ids->buffer)); int32_t * data = (int32_t *) out_ids->data;
if (n_outputs == n_tokens) { for (int i = 0; i < n_tokens; ++i) { data[i] = i; }
return; }
GGML_ASSERT(ubatch->output);
int n_outputs = 0;
for (int i = 0; i < n_tokens; ++i) { if (ubatch->output[i]) { data[n_outputs++] = i; } }}
bool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) { bool res = true;
res &= n_outputs == params.n_outputs;
return res;}
void llm_graph_input_mean::set_input(const llama_ubatch * ubatch) { if (cparams.embeddings && (cparams.pooling_type == LLAMA_POOLING_TYPE_MEAN || cparams.pooling_type == LLAMA_POOLING_TYPE_RANK )) {
const int64_t n_tokens = ubatch->n_tokens; const int64_t n_seq_tokens = ubatch->n_seq_tokens; const int64_t n_seqs_unq = ubatch->n_seqs_unq;
GGML_ASSERT(mean); GGML_ASSERT(ggml_backend_buffer_is_host(mean->buffer));
float * data = (float *) mean->data; memset(mean->data, 0, n_tokens*n_seqs_unq*ggml_element_size(mean));
std::vector<uint64_t> sums(n_seqs_unq, 0); for (int i = 0; i < n_tokens; i += n_seq_tokens) { for (int s = 0; s < ubatch->n_seq_id[i]; ++s) { const llama_seq_id seq_id = ubatch->seq_id[i][s]; const int32_t seq_idx = ubatch->seq_idx[seq_id];
sums[seq_idx] += ubatch->n_seq_tokens; } }
std::vector<float> div(n_seqs_unq, 0.0f); for (int s = 0; s < n_seqs_unq; ++s) { const uint64_t sum = sums[s]; if (sum > 0) { div[s] = 1.0f/float(sum); } }
for (int i = 0; i < n_tokens; i += n_seq_tokens) { for (int s = 0; s < ubatch->n_seq_id[i]; ++s) { const llama_seq_id seq_id = ubatch->seq_id[i][s]; const int32_t seq_idx = ubatch->seq_idx[seq_id];
for (int j = 0; j < n_seq_tokens; ++j) { data[seq_idx*n_tokens + i + j] = div[seq_idx]; } } } }}
void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) { const int64_t n_tokens = ubatch->n_tokens; const int64_t n_seqs_unq = ubatch->n_seqs_unq;
if (cparams.embeddings && ( cparams.pooling_type == LLAMA_POOLING_TYPE_CLS || cparams.pooling_type == LLAMA_POOLING_TYPE_RANK || cparams.pooling_type == LLAMA_POOLING_TYPE_LAST )) { GGML_ASSERT(cls); GGML_ASSERT(ggml_backend_buffer_is_host(cls->buffer));
uint32_t * data = (uint32_t *) cls->data; memset(cls->data, 0, n_seqs_unq*ggml_element_size(cls));
std::vector<int> target_pos(n_seqs_unq, -1); std::vector<int> target_row(n_seqs_unq, -1);
const bool last = ( cparams.pooling_type == LLAMA_POOLING_TYPE_LAST || (cparams.pooling_type == LLAMA_POOLING_TYPE_RANK && (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_QWEN3VL)) // qwen3 reranking & embedding models use last token );
for (int i = 0; i < n_tokens; ++i) { const llama_pos pos = ubatch->pos[i];
for (int s = 0; s < ubatch->n_seq_id[i]; ++s) { const llama_seq_id seq_id = ubatch->seq_id[i][s]; const int32_t seq_idx = ubatch->seq_idx[seq_id];
if ( (target_pos[seq_idx] == -1) || ( last && pos >= target_pos[seq_idx]) || (!last && pos < target_pos[seq_idx]) ) { target_pos[seq_idx] = pos; target_row[seq_idx] = i; } } }
for (int s = 0; s < n_seqs_unq; ++s) { if (target_row[s] >= 0) { data[s] = target_row[s]; } } }}
void llm_graph_input_rs::set_input_rs(const llama_memory_recurrent_context * mctx_cur, const llama_ubatch * ubatch) { const int64_t n_rs = mctx_cur->get_n_rs();
if (s_copy) { GGML_ASSERT(ggml_backend_buffer_is_host(s_copy->buffer)); int32_t * data = (int32_t *) s_copy->data;
// assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n for (uint32_t i = 0; i < n_rs; ++i) { data[i] = mctx_cur->s_copy(i); } }
if (s_write_rows) { GGML_ASSERT(ggml_backend_buffer_is_host(s_write_rows->buffer)); int64_t * data = (int64_t *) s_write_rows->data;
// destination cache row of (snapshot slot s, seq i) = s*mem_size + head + i, // matching the strided-cpy destination of the legacy gathered path const uint32_t mem_size = mctx_cur->get_size(); const uint32_t rs_head = mctx_cur->get_head(); const int64_t n_seqs = ubatch->n_seqs;
for (int64_t r = 0; r < s_write_rows->ne[0]; ++r) { data[r] = (r / n_seqs) * (int64_t) mem_size + (int64_t) rs_head + (r % n_seqs); } }}
bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mctx_cur, const llm_graph_params & params) { bool res = true;
res &= s_copy->ne[0] == mctx_cur->get_n_rs();
res &= s_copy_main->ne[0] == params.ubatch.n_seqs; res &= s_copy_extra->ne[0] == mctx_cur->get_n_rs() - params.ubatch.n_seqs;
if (s_write_rows) { res &= s_write_rows->ne[0] == std::min<int64_t>(params.ubatch.n_seq_tokens, s_write_K) * params.ubatch.n_seqs; }
res &= head == mctx_cur->get_head(); res &= rs_z == mctx_cur->get_rs_z();
res &= s_copy_main_in_place == mctx_cur->s_copy_main_in_place(params.ubatch.n_seqs);
return res;}
void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) { set_input_rs(mctx, ubatch);}
bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_memory_recurrent_context *>(params.mctx);
this->mctx = mctx;
return can_reuse_rs(mctx, params);}
void llm_graph_input_cross_embd::set_input(const llama_ubatch * ubatch) { GGML_UNUSED(ubatch);
if (cross_embd && !cross->v_embd.empty()) { assert(cross_embd->type == GGML_TYPE_F32);
ggml_backend_tensor_set(cross_embd, cross->v_embd.data(), 0, ggml_nbytes(cross_embd)); }}
template <typename T>static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n_swa, llama_swa_type swa_type) { LLAMA_LOG_DEBUG("%s: === Attention mask ===\n", __func__); const char * swa_type_str = "unknown";
switch (swa_type) { case LLAMA_SWA_TYPE_NONE: swa_type_str = "LLAMA_SWA_TYPE_NONE"; break; case LLAMA_SWA_TYPE_STANDARD: swa_type_str = "LLAMA_SWA_TYPE_STANDARD"; break; case LLAMA_SWA_TYPE_CHUNKED: swa_type_str = "LLAMA_SWA_TYPE_CHUNKED"; break; case LLAMA_SWA_TYPE_SYMMETRIC: swa_type_str = "LLAMA_SWA_TYPE_SYMMETRIC"; break; };
LLAMA_LOG_DEBUG("%s: n_swa : %d, n_kv: %d, swa_type: %s\n", __func__, (int)n_swa, (int)n_kv, swa_type_str); LLAMA_LOG_DEBUG("%s: '0' = can attend, '∞' = masked\n", __func__); LLAMA_LOG_DEBUG("%s: Rows = query tokens, Columns = key/value tokens\n\n", __func__);
LLAMA_LOG_DEBUG(" "); for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) { LLAMA_LOG_DEBUG("%2d", j); } LLAMA_LOG_DEBUG("\n");
for (int i = 0; i < std::min((int64_t)20, n_tokens); ++i) { LLAMA_LOG_DEBUG(" %2d ", i); for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) { float val = llama_cast<float>(data[i * n_kv + j]); if (val == -INFINITY) { LLAMA_LOG_DEBUG(" ∞"); } else { LLAMA_LOG_DEBUG(" 0"); } } LLAMA_LOG_DEBUG("\n"); }}
void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { const int64_t n_kv = ubatch->n_tokens; const int64_t n_tokens = ubatch->n_tokens;
const auto fill_mask = [&](auto * data, int64_t ne, int n_swa, llama_swa_type swa_type) { using T = std::remove_reference_t<decltype(*data)>; std::fill(data, data + ne, llama_cast<T>(-INFINITY));
for (int i1 = 0; i1 < n_tokens; ++i1) { const llama_seq_id s1 = ubatch->seq_id[i1][0]; const llama_pos p1 = ubatch->pos[i1];
const uint64_t idst = i1*n_kv;
for (int i0 = 0; i0 < n_tokens; ++i0) { const llama_seq_id s0 = ubatch->seq_id[i0][0]; const llama_pos p0 = ubatch->pos[i0];
// mask different sequences if (s0 != s1) { continue; }
// mask future tokens if (cparams.causal_attn && p0 > p1) { continue; }
// apply SWA if any if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) { continue; }
data[idst + i0] = llama_cast<T>(hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f); } }
if (debug) { print_mask(data, n_tokens, n_kv, n_swa, swa_type); } };
GGML_ASSERT(self_kq_mask); GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer)); if (self_kq_mask->type == GGML_TYPE_F16) { fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE); } else { fill_mask((float *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE); }
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { GGML_ASSERT(self_kq_mask_swa); GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer)); if (self_kq_mask_swa->type == GGML_TYPE_F16) { fill_mask((ggml_fp16_t *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type); } else { fill_mask((float *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type); } }}
void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch);
// the mask is left unallocated when the graph only stores K/V without attending // (e.g. DFlash's KV-injection pass) if (self_kq_mask && self_kq_mask->buffer) { mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); }
if (self_k_rot && self_k_rot->buffer) { mctx->set_input_k_rot(self_k_rot); }
if (self_v_rot && self_v_rot->buffer) { mctx->set_input_v_rot(self_v_rot); }}
bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams);
return res;}
void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch);
mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);}
bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) { mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
return can_reuse_impl(params);}
bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { bool res = true;
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams);
return res;}
llm_graph_input_attn_kv_msa::llm_graph_input_attn_kv_msa( const llama_hparams & hparams, const llama_cparams & cparams, const llama_kv_cache_msa_context * mctx) : llm_graph_input_attn_kv(hparams, cparams, mctx->get_base()), mctx_msa(mctx) {}
void llm_graph_input_attn_kv_msa::set_input(const llama_ubatch * ubatch) { llm_graph_input_attn_kv::set_input(ubatch);
if (self_k_idxs_idx) { mctx_msa->get_idx()->set_input_k_idxs(self_k_idxs_idx, ubatch); }}
bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) { mctx_msa = static_cast<const llama_kv_cache_msa_context *>(params.mctx);
// the parent class operates on the base cache context this->mctx = mctx_msa->get_base();
bool res = true;
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; if (self_k_idxs_idx) { res &= self_k_idxs_idx->ne[0] == params.ubatch.n_tokens; }
res &= can_reuse_kq_mask(self_kq_mask, this->mctx, params.ubatch, params.cparams);
return res;}
void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) { mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn);
mctx->get_lid()->set_input_k_rot(self_k_rot_lid);}
bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) { mctx = static_cast<const llama_kv_cache_dsa_context *>(params.mctx);
return can_reuse_impl(params);}
bool llm_graph_input_attn_k_dsa::can_reuse_impl(const llm_graph_params & params) { bool res = true;
res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens; res &= self_k_idxs_lid->ne[0] == params.ubatch.n_tokens;
res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams); res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams);
return res;}
void llm_graph_input_attn_k_dsa_iswa::set_input(const llama_ubatch * ubatch) { inp_dsa->set_input(ubatch); inp_swa->set_input(ubatch);}
bool llm_graph_input_attn_k_dsa_iswa::can_reuse(const llm_graph_params & params) { mctx = static_cast<const llama_kv_cache_dsa_iswa_context *>(params.mctx);
inp_dsa->mctx = mctx->get_dsa(); inp_swa->mctx = mctx->get_swa();
bool res = true;
res &= inp_dsa->can_reuse_impl(params); res &= inp_swa->can_reuse_impl(params);
return res;}
void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { // base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); if (self_v_idxs) { mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); } }
// the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live if (self_kq_mask && self_kq_mask->buffer) { mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); }
// swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); if (self_v_idxs_swa) { mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); } }
if (self_kq_mask_swa && self_kq_mask_swa->buffer) { mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); }
if (self_k_rot && self_k_rot->buffer) { mctx->get_base()->set_input_k_rot(self_k_rot); }
if (self_v_rot && self_v_rot->buffer) { mctx->get_base()->set_input_v_rot(self_v_rot); }
if (self_k_rot_swa && self_k_rot_swa->buffer) { mctx->get_swa()->set_input_k_rot(self_k_rot_swa); }
if (self_v_rot_swa && self_v_rot_swa->buffer) { mctx->get_swa()->set_input_v_rot(self_v_rot_swa); }}
bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
// base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there }
if (self_kq_mask && self_kq_mask->buffer) { res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); }
// swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there }
if (self_kq_mask_swa && self_kq_mask_swa->buffer) { res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); }
return res;}
void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) { // base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); }
// the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live if (self_kq_mask && self_kq_mask->buffer) { mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); }
// swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); }
if (self_kq_mask_swa && self_kq_mask_swa->buffer) { mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); }
if (self_k_rot && self_k_rot->buffer) { mctx->get_base()->set_input_k_rot(self_k_rot); }
if (self_k_rot_swa && self_k_rot_swa->buffer) { mctx->get_swa()->set_input_k_rot(self_k_rot_swa); }}
bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
// base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; }
if (self_kq_mask && self_kq_mask->buffer) { res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); }
// swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; }
if (self_kq_mask_swa && self_kq_mask_swa->buffer) { res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); }
return res;}
static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & src) { if (!dst || !dst->buffer) { return; }
GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst));}
static void dsv4_set_i32(ggml_tensor * dst, const std::vector<int32_t> & src) { if (!dst || !dst->buffer) { return; }
GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst));}
static void dsv4_set_kq_mask( ggml_tensor * dst, const llama_kv_cache_dsv4_context::comp_plan & plan, uint32_t n_tokens, int64_t n_stream) { if (!dst || !dst->buffer) { return; }
GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); GGML_ASSERT(dst->ne[0] == plan.n_kv); GGML_ASSERT(dst->ne[1] == (int64_t) n_tokens/n_stream); GGML_ASSERT(dst->ne[2] == 1); GGML_ASSERT(dst->ne[3] == n_stream); GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) n_tokens); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
if (dst->type == GGML_TYPE_F32) { float * data = (float *) dst->data;
for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { const int32_t n_visible = plan.n_visible[i];
for (int64_t j = 0; j < dst->ne[0]; ++j) { data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; } } } else if (dst->type == GGML_TYPE_F16) { ggml_fp16_t * data = (ggml_fp16_t *) dst->data; const ggml_fp16_t fp16_ninf = llama_cast<ggml_fp16_t>(-INFINITY); const ggml_fp16_t fp16_zero = llama_cast<ggml_fp16_t>(0.0f);
for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { const int32_t n_visible = plan.n_visible[i];
for (int64_t j = 0; j < dst->ne[0]; ++j) { data[i*dst->ne[0] + j] = j < n_visible ? fp16_zero : fp16_ninf; } } }}
static ggml_tensor * dsv4_build_raw_kq_mask( ggml_context * ctx, const llama_kv_cache_dsv4_raw_context * mctx, const llama_ubatch & ubatch, const llama_cparams & cparams, int64_t n_stream) { const auto n_kv = mctx->get_n_kv(); const auto n_tokens = ubatch.n_tokens;
GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0);
const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(res); ggml_set_name(res, "attn_inp_kq_mask");
return res;}
static bool dsv4_can_reuse_raw_kq_mask( ggml_tensor * kq_mask, const llama_kv_cache_dsv4_raw_context * mctx, const llama_ubatch & ubatch, int64_t n_stream) { const auto n_kv = mctx->get_n_kv(); const auto n_tokens = ubatch.n_tokens;
GGML_ASSERT(n_stream > 0);
bool res = true;
res &= (kq_mask->ne[0] == n_kv); res &= (kq_mask->ne[1] == n_tokens/n_stream); res &= (kq_mask->ne[2] == 1); res &= (kq_mask->ne[3] == n_stream);
return res;}
static std::string dsv4_plan_positions(const std::vector<int32_t> & values) { std::ostringstream ss; ss << "["; for (size_t i = 0; i < values.size(); ++i) { if (i > 0) { ss << ", "; } ss << values[i]; } ss << "]"; return ss.str();}
static bool dsv4_compress_debug() { static const bool debug = []() { const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG"); return env && atoi(env) > 0; }();
return debug;}
static void dsv4_set_comp_inputs( const llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, const char * name, bool debug, uint32_t n_tokens, int64_t n_stream) { dsv4_set_i32(inp.state_pos, plan.state_pos); dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs); dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs); dsv4_set_i32(inp.state_restore_src_idxs, plan.state_restore_src_idxs); dsv4_set_i32(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs); dsv4_set_i32(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs); dsv4_set_i32(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs); dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs); dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs); dsv4_set_i32(inp.state_write_pos, plan.state_write_pos); dsv4_set_kq_mask(inp.kq_mask, plan, n_tokens, n_stream);
if (debug || dsv4_compress_debug()) { LLAMA_LOG_INFO("%s: %s n_tokens=%u, n_stream=%d, state_persist_dst=%s, state_write_pos=%s\n", __func__, name, n_tokens, (int) n_stream, dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(), dsv4_plan_positions(plan.state_write_pos).c_str()); }}
static bool dsv4_can_reuse_tensor_1d(ggml_tensor * t, int64_t ne0) { return (t == nullptr && ne0 == 0) || (t != nullptr && t->ne[0] == ne0);}
static bool dsv4_can_reuse_kq_mask( ggml_tensor * t, const llama_kv_cache_dsv4_context::comp_plan & plan, uint32_t n_tokens, int64_t n_stream) { if (plan.n_kv == 0) { return t == nullptr; }
GGML_ASSERT(n_stream > 0);
return t != nullptr && t->ne[0] == plan.n_kv && t->ne[1] == (int64_t) n_tokens/n_stream && t->ne[2] == 1 && t->ne[3] == n_stream;}
static bool dsv4_can_reuse_comp_input( const llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, uint32_t n_tokens, int64_t n_stream) { bool res = true; res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_restore_src_idxs, plan.state_restore_src_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size()); res &= dsv4_can_reuse_kq_mask(inp.kq_mask, plan, n_tokens, n_stream);
return res;}
static ggml_tensor * dsv4_build_input_1d( ggml_context * ctx, ggml_type type, int64_t ne0, const std::string & name) { if (ne0 == 0) { return nullptr; }
ggml_tensor * res = ggml_new_tensor_1d(ctx, type, ne0); ggml_set_input(res); ggml_set_name(res, name.c_str());
return res;}
static void dsv4_build_comp_inputs( ggml_context * ctx, llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, const char * name, const llama_cparams & cparams, int64_t n_stream) { inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos"); inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs"); inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs"); inp.state_restore_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_src_idxs.size(), std::string("dsv4_") + name + "_state_restore_src_idxs"); inp.state_restore_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_dst_idxs.size(), std::string("dsv4_") + name + "_state_restore_dst_idxs"); inp.state_snapshot_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_src_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_src_idxs"); inp.state_snapshot_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_dst_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_dst_idxs"); inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs"); inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs"); inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos");
if (plan.n_kv > 0) { const int64_t n_tokens = (int64_t) plan.n_visible.size();
GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0);
inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); }}
void llm_graph_input_dsv4_raw::set_input(const llama_ubatch * ubatch) { if (self_k_idxs && self_k_idxs->buffer) { mctx->set_input_k_idxs(self_k_idxs); }
if (self_kq_mask && self_kq_mask->buffer) { mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); }
if (self_k_rot) { mctx->set_input_k_rot(self_k_rot); }}
void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) { const auto & plan_csa = mctx->get_csa_plan(*ubatch); const auto & plan_hca = mctx->get_hca_plan(*ubatch); const auto & plan_lid = mctx->get_lid_plan(*ubatch); const int64_t n_stream = plan_csa.n_stream;
inp_raw->mctx = mctx->get_raw(); inp_raw->set_input(ubatch);
dsv4_set_comp_inputs(inp_csa, plan_csa, "csa", debug > 0, ubatch->n_tokens, n_stream); dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream); dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream);
if (inp_csa.k_rot && inp_csa.k_rot->buffer) { mctx->get_csa()->set_input_k_rot(inp_csa.k_rot); }
if (inp_hca.k_rot && inp_hca.k_rot->buffer) { mctx->get_hca()->set_input_k_rot(inp_hca.k_rot); }
if (inp_lid.k_rot && inp_lid.k_rot->buffer) { mctx->get_lid()->set_input_k_rot(inp_lid.k_rot); }}
bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_kv_cache_dsv4_context *>(params.mctx);
this->mctx = mctx; inp_raw->mctx = mctx->get_raw();
bool res = true;
const auto & plan_csa = mctx->get_csa_plan(params.ubatch); const auto & plan_hca = mctx->get_hca_plan(params.ubatch); const auto & plan_lid = mctx->get_lid_plan(params.ubatch); const int64_t n_stream = plan_csa.n_stream;
const auto * raw_ctx = mctx->get_raw(); inp_raw->mctx = raw_ctx;
if (inp_raw->self_k_idxs && inp_raw->self_k_idxs->buffer) { res &= inp_raw->self_k_idxs->ne[0] == raw_ctx->get_n_write(); } if (inp_raw->self_kq_mask && inp_raw->self_kq_mask->buffer) { res &= dsv4_can_reuse_raw_kq_mask(inp_raw->self_kq_mask, raw_ctx, params.ubatch, n_stream); }
res &= dsv4_can_reuse_comp_input(inp_csa, plan_csa, params.ubatch.n_tokens, n_stream); res &= dsv4_can_reuse_comp_input(inp_hca, plan_hca, params.ubatch.n_tokens, n_stream); res &= dsv4_can_reuse_comp_input(inp_lid, plan_lid, params.ubatch.n_tokens, n_stream);
return res;}
void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(cross_kq_mask);
const int64_t n_enc = cross_kq_mask->ne[0]; const int64_t n_tokens = ubatch->n_tokens;
GGML_ASSERT(ggml_backend_buffer_is_host(cross_kq_mask->buffer)); GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing
const auto fill_mask = [&](auto * data) { using T = std::remove_reference_t<decltype(*data)>; for (int i = 0; i < n_tokens; ++i) { GGML_ASSERT(!cross->seq_ids_enc.empty() && "llama_encode must be called first"); for (int j = 0; j < n_enc; ++j) { float f = -INFINITY;
for (int s = 0; s < ubatch->n_seq_id[i]; ++s) { const llama_seq_id seq_id = ubatch->seq_id[i][s];
if (cross->seq_ids_enc[j].find(seq_id) != cross->seq_ids_enc[j].end()) { f = 0.0f; } }
data[i*n_enc + j] = llama_cast<T>(f); } } };
if (cross_kq_mask->type == GGML_TYPE_F16) { fill_mask((ggml_fp16_t *) cross_kq_mask->data); } else { fill_mask((float *) cross_kq_mask->data); }}
void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
if (inp_attn->self_k_rot) { mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot); }
if (inp_attn->self_v_rot) { mctx->get_attn()->set_input_v_rot(inp_attn->self_v_rot); }
inp_rs->set_input_rs(mctx->get_recr(), ubatch);}
bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams);
res &= inp_rs->can_reuse_rs(mctx->get_recr(), params);
return res;}
// TODO: Hybrid input classes are a bit redundant.// Instead of creating a hybrid input, the graph can simply create 2 separate inputs.// Refactoring is required in the future.void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
inp_rs->set_input_rs(mctx->get_recr(), ubatch);}
bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;
res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams);
res &= inp_rs->can_reuse_rs(mctx->get_recr(), params);
return res;}
void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) { const auto * attn_ctx = mctx->get_attn();
// base tensors may not be allocated if there are no non-SWA attention layers if (inp_attn->self_k_idxs && inp_attn->self_k_idxs->buffer) { attn_ctx->get_base()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); attn_ctx->get_base()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); }
if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) { attn_ctx->get_base()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); }
// swa tensors may not be allocated if there are no SWA attention layers if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) { attn_ctx->get_swa()->set_input_k_idxs(inp_attn->self_k_idxs_swa, ubatch); attn_ctx->get_swa()->set_input_v_idxs(inp_attn->self_v_idxs_swa, ubatch); }
if (inp_attn->self_kq_mask_swa && inp_attn->self_kq_mask_swa->buffer) { attn_ctx->get_swa()->set_input_kq_mask(inp_attn->self_kq_mask_swa, ubatch, cparams.causal_attn); }
if (inp_attn->self_k_rot) { attn_ctx->get_base()->set_input_k_rot(inp_attn->self_k_rot); }
if (inp_attn->self_v_rot) { attn_ctx->get_base()->set_input_v_rot(inp_attn->self_v_rot); }
if (inp_attn->self_k_rot_swa) { attn_ctx->get_swa()->set_input_k_rot(inp_attn->self_k_rot_swa); }
if (inp_attn->self_v_rot_swa) { attn_ctx->get_swa()->set_input_v_rot(inp_attn->self_v_rot_swa); }
inp_rs->set_input_rs(mctx->get_recr(), ubatch);}
bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params) { const auto * mctx = static_cast<const llama_memory_hybrid_iswa_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
const auto * attn_ctx = mctx->get_attn();
// base tensors may not be allocated if there are no non-SWA attention layers if (inp_attn->self_k_idxs && inp_attn->self_k_idxs->buffer) { res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there }
res &= can_reuse_kq_mask(inp_attn->self_kq_mask, attn_ctx->get_base(), params.ubatch, params.cparams);
// swa tensors may not be allocated if there are no SWA attention layers if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) { res &= inp_attn->self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; //res &= inp_attn->self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there }
res &= can_reuse_kq_mask(inp_attn->self_kq_mask_swa, attn_ctx->get_swa(), params.ubatch, params.cparams);
res &= inp_rs->can_reuse_rs(mctx->get_recr(), params);
return res;}
void llm_graph_input_sampling::set_input(const llama_ubatch * ubatch) { // set the inputs only for the active samplers in the current ubatch std::unordered_set<llama_seq_id> active_samplers; for (uint32_t i = 0; i < ubatch->n_tokens; i++) { if (ubatch->output[i]) { llama_seq_id seq_id = ubatch->seq_id[i][0]; active_samplers.insert(seq_id); } }
for (auto seq_id : active_samplers) { if (samplers.find(seq_id) == samplers.end()) { continue; }
auto & sampler = samplers[seq_id];
if (sampler->iface->backend_set_input) { sampler->iface->backend_set_input(sampler); } }}
bool llm_graph_input_sampling::can_reuse(const llm_graph_params & params) { if (samplers.size() != params.samplers.size()) { return false; }
for (const auto & [seq_id, sampler] : params.samplers) { if (samplers[seq_id] != sampler) { return false; } }
return true;}
//// llm_graph_result//
llm_graph_result::llm_graph_result(int64_t max_nodes) : max_nodes(max_nodes) { reset();
const char * LLAMA_GRAPH_RESULT_DEBUG = getenv("LLAMA_GRAPH_RESULT_DEBUG"); debug = LLAMA_GRAPH_RESULT_DEBUG ? atoi(LLAMA_GRAPH_RESULT_DEBUG) : 0;}
int64_t llm_graph_result::get_max_nodes() const { return max_nodes;}
void llm_graph_result::reset() { t_inp_tokens = nullptr; t_inp_embd = nullptr; t_logits = nullptr; t_embd = nullptr; t_embd_pooled = nullptr; t_h_nextn = nullptr; t_h_capture = nullptr;
t_layer_inp.resize(LLAMA_MAX_LAYERS + 1); std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr);
t_sampled.clear(); t_dspark_greedy.clear(); t_sampled_probs.clear(); t_sampled_logits.clear(); t_candidates.clear();
params = {};
inputs.clear(); fused_nodes.clear();
buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));
ggml_init_params params = { /*.mem_size =*/ buf_compute_meta.size(), /*.mem_buffer =*/ buf_compute_meta.data(), /*.no_alloc =*/ true, };
ctx_compute.reset(ggml_init(params));
gf = ggml_new_graph_custom(ctx_compute.get(), max_nodes, false);}
void llm_graph_result::set_inputs(const llama_ubatch * ubatch) { for (auto & input : inputs) { input->set_input(ubatch); }}
void llm_graph_result::set_outputs(const llm_graph_params & params) { if (t_logits != nullptr) { ggml_set_output(t_logits); } if (t_embd != nullptr) { ggml_set_output(t_embd); } if (t_embd_pooled != nullptr) { ggml_set_output(t_embd_pooled); } if (t_h_capture != nullptr) { ggml_set_output(t_h_capture); } if (t_h_nextn != nullptr) { ggml_set_output(t_h_nextn); } { const auto & embeddings_layer_inp = params.cparams.embeddings_layer_inp; for (size_t il = 0; il < embeddings_layer_inp.size(); ++il) { if (embeddings_layer_inp[il]) { GGML_ASSERT(t_layer_inp[il] != nullptr && "layer input tensor is null"); ggml_set_output(t_layer_inp[il]); } } } for (auto * tensor : t_sampled) { if (tensor != nullptr) { ggml_set_output(tensor); } } for (auto * tensor : t_sampled_probs) { if (tensor != nullptr) { ggml_set_output(tensor); } } for (auto * tensor : t_sampled_logits) { if (tensor != nullptr) { ggml_set_output(tensor); } } for (auto * tensor : t_candidates) { if (tensor != nullptr) { ggml_set_output(tensor); } }}
bool llm_graph_result::can_reuse(const llm_graph_params & params) { if (!this->params.allow_reuse(params)) { if (debug > 1) { LLAMA_LOG_DEBUG("%s: cannot reuse graph due to incompatible graph parameters\n", __func__); }
return false; }
if (debug > 1) { LLAMA_LOG_DEBUG("%s: checking compatibility of %d inputs:\n", __func__, (int) inputs.size()); }
bool res = true;
for (auto & input : inputs) { const bool cur = input->can_reuse(params);
if (debug > 1) { LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", "placeholder", cur); }
res = res && cur; }
if (debug > 0) { LLAMA_LOG_DEBUG("%s: can reuse graph = %d\n", __func__, res); }
return res;}
llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) { inputs.emplace_back(std::move(input)); return inputs.back().get();}
void llm_graph_result::add_fused_node(llm_graph_fused_node result) { fused_nodes.push_back(result);}
void llm_graph_result::set_params(const llm_graph_params & params) { this->params = params;}
//// llm_graph_context//
llm_graph_context::llm_graph_context(const llm_graph_params & params) : arch(params.arch), hparams(params.hparams), cparams(params.cparams), ubatch(params.ubatch), n_embd(hparams.n_embd), n_layer(hparams.n_layer()), n_layer_nextn(hparams.n_layer_nextn), n_rot(hparams.n_rot()), n_ctx(cparams.n_ctx), n_head(hparams.n_head()), n_head_kv(hparams.n_head_kv()), n_embd_head_k(hparams.n_embd_head_k()), n_embd_k_gqa(hparams.n_embd_k_gqa()), n_embd_head_v(hparams.n_embd_head_v()), n_embd_v_gqa(hparams.n_embd_v_gqa()), n_expert(hparams.n_expert), n_expert_used(cparams.warmup ? hparams.n_expert : hparams.n_expert_used), freq_base(cparams.rope_freq_base), freq_scale(cparams.rope_freq_scale), ext_factor(cparams.yarn_ext_factor), attn_factor(cparams.yarn_attn_factor), beta_fast(cparams.yarn_beta_fast), beta_slow(cparams.yarn_beta_slow), norm_eps(hparams.f_norm_eps), norm_rms_eps(hparams.f_norm_rms_eps), n_tokens(ubatch.n_tokens), n_outputs(params.n_outputs), n_ctx_orig(cparams.n_ctx_orig_yarn), pooling_type(cparams.pooling_type), rope_type(hparams.rope_type), sched(params.sched), backend_cpu(params.backend_cpu), cvec(params.cvec), loras(params.loras), mctx(params.mctx), cross(params.cross), dspark_ctx(params.dspark_ctx), dspark_has_context(params.dspark_has_context), dspark_ctx_rows(params.dspark_ctx_rows), dspark_ctx_width(params.dspark_ctx_width), hadamard_rotations(params.hadamard_rotations), hadamard_inverses(params.hadamard_inverses), samplers(params.samplers), cb_func(params.cb), res(params.res), ctx0(res->get_ctx()), gf(res->get_gf()) { res->set_params(params);}
void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const { if (cb_func) { cb_func(ubatch, cur, name, il); }}
ggml_tensor * llm_graph_context::build_cvec( ggml_tensor * cur, int il) const { return cvec->apply_to(ctx0, cur, il);}
ggml_tensor * llm_graph_context::build_lora_mm( ggml_tensor * w, ggml_tensor * cur, ggml_tensor * w_s) const { ggml_tensor * cur_mm = cur; if (hadamard_rotations) { const auto it = hadamard_rotations->find(w); if (it != hadamard_rotations->end()) { const auto & t = it->second; // another folded weight on this same activation already built the transform const auto memo_key = std::make_pair((const ggml_tensor *) cur, (const ggml_tensor *) t.rot); const auto memo_it = hadamard_memo.find(memo_key); if (memo_it != hadamard_memo.end()) { cur_mm = memo_it->second; } else { if (t.perm_rep > 1) { // tiled [hd, nk, rep] -> grouped [hd, rep, nk] feature order ggml_tensor * x = ggml_is_contiguous(cur_mm) ? cur_mm : ggml_cont(ctx0, cur_mm); const int64_t ne1 = x->ne[1], ne2 = x->ne[2], ne3 = x->ne[3]; x = ggml_reshape_4d(ctx0, x, t.perm_hd, t.perm_nk, t.perm_rep, ne1*ne2*ne3); x = ggml_cont(ctx0, ggml_permute(ctx0, x, 0, 2, 1, 3)); cur_mm = ggml_reshape_4d(ctx0, x, t.perm_hd*t.perm_nk*t.perm_rep, ne1, ne2, ne3); } if (t.signs) { cur_mm = ggml_mul(ctx0, cur_mm, t.signs); } cur_mm = llama_mul_mat_hadamard(ctx0, cur_mm, t.rot); hadamard_memo[memo_key] = cur_mm; } } }
ggml_tensor * res = ggml_mul_mat(ctx0, w, cur_mm);
if (w_s) { res = ggml_mul(ctx0, res, w_s); }
for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); if (lw == nullptr) { continue; }
const float adapter_scale = lora.second; const float scale = lw->get_scale(lora.first->alpha, adapter_scale);
ggml_tensor * ab_cur = ggml_mul_mat( ctx0, lw->b, ggml_mul_mat(ctx0, lw->a, cur) );
ab_cur = ggml_scale(ctx0, ab_cur, scale); res = ggml_add(ctx0, res, ab_cur); }
return res;}
ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, ggml_tensor * w_s) const { ggml_tensor * cur_mm = cur; if (hadamard_rotations) { const auto it = hadamard_rotations->find(w); if (it != hadamard_rotations->end()) { const auto & t = it->second; // another folded weight on this same activation already built the transform const auto memo_key = std::make_pair((const ggml_tensor *) cur, (const ggml_tensor *) t.rot); const auto memo_it = hadamard_memo.find(memo_key); if (memo_it != hadamard_memo.end()) { cur_mm = memo_it->second; } else { if (t.perm_rep > 1) { // tiled [hd, nk, rep] -> grouped [hd, rep, nk] feature order ggml_tensor * x = ggml_is_contiguous(cur_mm) ? cur_mm : ggml_cont(ctx0, cur_mm); const int64_t ne1 = x->ne[1], ne2 = x->ne[2], ne3 = x->ne[3]; x = ggml_reshape_4d(ctx0, x, t.perm_hd, t.perm_nk, t.perm_rep, ne1*ne2*ne3); x = ggml_cont(ctx0, ggml_permute(ctx0, x, 0, 2, 1, 3)); cur_mm = ggml_reshape_4d(ctx0, x, t.perm_hd*t.perm_nk*t.perm_rep, ne1, ne2, ne3); } if (t.signs) { cur_mm = ggml_mul(ctx0, cur_mm, t.signs); } cur_mm = llama_mul_mat_hadamard(ctx0, cur_mm, t.rot); hadamard_memo[memo_key] = cur_mm; } } }
ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur_mm, ids);
if (w_s) { const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); s = ggml_get_rows(ctx0, s, ids); res = ggml_mul(ctx0, res, s); } for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); if (lw == nullptr) { continue; }
const float alpha = lora.first->alpha; const float rank = (float) lw->b->ne[0]; const float scale = alpha ? lora.second * alpha / rank : lora.second;
ggml_tensor * ab_cur = ggml_mul_mat_id( ctx0, lw->b, ggml_mul_mat_id(ctx0, lw->a, cur, ids), ids );
ab_cur = ggml_scale(ctx0, ab_cur, scale); res = ggml_add(ctx0, res, ab_cur); }
return res;}
ggml_tensor * llm_graph_context::build_norm( ggml_tensor * cur, ggml_tensor * mw, ggml_tensor * mb, llm_norm_type type, int il) const { switch (type) { case LLM_NORM: cur = ggml_norm (ctx0, cur, hparams.f_norm_eps); break; case LLM_NORM_RMS: cur = ggml_rms_norm(ctx0, cur, hparams.f_norm_rms_eps); break; case LLM_NORM_GROUP: { cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, cur->ne[1]); cur = ggml_group_norm(ctx0, cur, hparams.n_norm_groups, hparams.f_norm_group_eps); cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], cur->ne[2]); } break; }
if (mw || mb) { cb(cur, "norm", il); }
if (mw) { cur = ggml_mul(ctx0, cur, mw); if (mb) { cb(cur, "norm_w", il); } }
if (mb) { cur = ggml_add(ctx0, cur, mb); }
return cur;}
llm_graph_qkv llm_graph_context::build_qkv( const llama_layer & layer, ggml_tensor * cur, int64_t n_embd_head, int64_t n_head, int64_t n_head_kv, int il) const { const int64_t n_embd_q = n_embd_head * n_head; const int64_t n_embd_kv = n_embd_head * n_head_kv;
ggml_tensor * Qcur, * Kcur, * Vcur;
if (layer.wqkv) { // fused QKV path ggml_tensor * qkv = build_lora_mm(layer.wqkv, cur, layer.wqkv_s); cb(qkv, "wqkv", il); if (layer.wqkv_b) { qkv = ggml_add(ctx0, qkv, layer.wqkv_b); cb(qkv, "wqkv_b", il); } if (hparams.f_clamp_kqv > 0.0f) { qkv = ggml_clamp(ctx0, qkv, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(qkv, "wqkv_clamped", il); } Qcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head, n_tokens, ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], 0); Kcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], ggml_row_size(qkv->type, n_embd_q)); Vcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], ggml_row_size(qkv->type, n_embd_q + n_embd_kv)); } else { // separate Q/K/V path Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); cb(Qcur, "Qcur", il); if (layer.wq_b) { Qcur = ggml_add(ctx0, Qcur, layer.wq_b); cb(Qcur, "Qcur", il); } if (hparams.f_clamp_kqv > 0.0f) { Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Qcur, "Qcur_clamped", il); } Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); cb(Kcur, "Kcur", il); if (layer.wk_b) { Kcur = ggml_add(ctx0, Kcur, layer.wk_b); cb(Kcur, "Kcur", il); } if (hparams.f_clamp_kqv > 0.0f) { Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Kcur, "Kcur_clamped", il); } Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); cb(Vcur, "Vcur", il); if (layer.wv_b) { Vcur = ggml_add(ctx0, Vcur, layer.wv_b); cb(Vcur, "Vcur", il); } if (hparams.f_clamp_kqv > 0.0f) { Vcur = ggml_clamp(ctx0, Vcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Vcur, "Vcur_clamped", il); } Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); }
cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il);
return { Qcur, Kcur, Vcur };}
ggml_tensor * llm_graph_context::build_ffn( ggml_tensor * cur, ggml_tensor * up, ggml_tensor * up_b, ggml_tensor * up_s, ggml_tensor * gate, ggml_tensor * gate_b, ggml_tensor * gate_s, ggml_tensor * down, ggml_tensor * down_b, ggml_tensor * down_s, ggml_tensor * act_scales, llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, int il) const { // NVFP4 support is currently restricted to // 1) LORA absence (*_s would be applied after LORA residual, which is incorrect) // 2) bias absense (*_s would be applied after bias addition, which is incorrect) // TODO: disambiguate LLM-architectural scales (which use *_s) from NVFP4 scale_2 (which also uses *_s currently) auto has_lora = [this](ggml_tensor * w) { if (!w) { return false; } for (const auto & lora : *loras) { if (lora.first->get_weight(w) != nullptr) { return true; } } return false; };
GGML_ASSERT(!up_s || !up_b || !up || up->type != GGML_TYPE_NVFP4); GGML_ASSERT(!gate_s || !gate_b || !gate || gate->type != GGML_TYPE_NVFP4); GGML_ASSERT(!down_s || !down_b || !down || down->type != GGML_TYPE_NVFP4); GGML_ASSERT(!up_s || !up || up->type != GGML_TYPE_NVFP4 || !has_lora(up)); GGML_ASSERT(!gate_s || !gate || gate->type != GGML_TYPE_NVFP4 || !has_lora(gate)); GGML_ASSERT(!down_s || !down || down->type != GGML_TYPE_NVFP4 || !has_lora(down));
ggml_tensor * tmp = up ? build_lora_mm(up, cur) : cur; cb(tmp, "ffn_up", il);
if (up_b) { tmp = ggml_add(ctx0, tmp, up_b); cb(tmp, "ffn_up_b", il); }
if (up_s) { tmp = ggml_mul(ctx0, tmp, up_s); cb(tmp, "ffn_up_s", il); }
if (gate) { switch (type_gate) { case LLM_FFN_SEQ: { cur = build_lora_mm(gate, tmp); cb(cur, "ffn_gate", il); } break; case LLM_FFN_PAR: { cur = build_lora_mm(gate, cur); cb(cur, "ffn_gate", il); } break; }
if (gate_b) { cur = ggml_add(ctx0, cur, gate_b); cb(cur, "ffn_gate_b", il); }
if (gate_s) { cur = ggml_mul(ctx0, cur, gate_s); cb(cur, "ffn_gate_s", il); }
} else { cur = tmp; }
switch (type_op) { case LLM_FFN_SILU: if (gate && type_gate == LLM_FFN_PAR) { if (il >= 0) { const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { tmp = ggml_clamp(ctx0, tmp, -limit, limit); cb(tmp, "ffn_up_clamped", il);
if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, tmp); } else { ggml_tensor * gate_act = ggml_silu(ctx0, cur); cb(gate_act, "ffn_silu", il); gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); cb(gate_act, "ffn_silu_clamped", il); cur = ggml_mul(ctx0, gate_act, tmp); } cb(cur, "ffn_swiglu_limited", il); type_gate = LLM_FFN_SEQ; break; } }
cur = ggml_swiglu_split(ctx0, cur, tmp); cb(cur, "ffn_swiglu", il); type_gate = LLM_FFN_SEQ; } else { cur = ggml_silu(ctx0, cur); cb(cur, "ffn_silu", il); } break; case LLM_FFN_GELU: if (gate && type_gate == LLM_FFN_PAR) { cur = ggml_geglu_split(ctx0, cur, tmp); cb(cur, "ffn_geglu", il); type_gate = LLM_FFN_SEQ; } else { cur = ggml_gelu(ctx0, cur); cb(cur, "ffn_gelu", il); if (act_scales != NULL) { cur = ggml_div(ctx0, cur, act_scales); cb(cur, "ffn_act", il); } } break; case LLM_FFN_RELU: if (gate && type_gate == LLM_FFN_PAR) { cur = ggml_reglu_split(ctx0, cur, tmp); cb(cur, "ffn_reglu", il); type_gate = LLM_FFN_SEQ; } else { cur = ggml_relu(ctx0, cur); cb(cur, "ffn_relu", il); } break; case LLM_FFN_RELU_SQR: { cur = ggml_relu(ctx0, cur); cb(cur, "ffn_relu", il);
cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_sqr(relu)", il); } break; case LLM_FFN_SWIGLU: { cur = ggml_swiglu(ctx0, cur); cb(cur, "ffn_swiglu", il); } break; case LLM_FFN_SWIGLU_OAI_MOE: if (gate && type_gate == LLM_FFN_PAR) { // same alpha/limit constants as gpt-oss const float alpha = 1.702f; const float limit = 7.0f; cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit); cb(cur, "ffn_swiglu_oai", il); type_gate = LLM_FFN_SEQ; } else { GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate"); } break; case LLM_FFN_GEGLU: { cur = ggml_geglu(ctx0, cur); cb(cur, "ffn_geglu", il); } break; case LLM_FFN_REGLU: { cur = ggml_reglu(ctx0, cur); cb(cur, "ffn_reglu", il); } break; case LLM_FFN_SITU: GGML_ABORT("not yet supported"); default: GGML_ABORT("fatal error"); }
if (gate && type_gate == LLM_FFN_PAR) { cur = ggml_mul(ctx0, cur, tmp); cb(cur, "ffn_gate_par", il); }
if (down) { cur = build_lora_mm(down, cur); if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators ggml_mul_mat_set_prec(cur, GGML_PREC_F32); } }
if (down_b) { cb(cur, "ffn_down", il); }
if (down_b) { cur = ggml_add(ctx0, cur, down_b); }
if (down_s) { cur = ggml_mul(ctx0, cur, down_s); cb(cur, "ffn_down_s", il); }
return cur;}
ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur, ggml_tensor * gate_inp, ggml_tensor * up_exps, ggml_tensor * gate_exps, ggml_tensor * down_exps, ggml_tensor * exp_probs_b, int64_t n_expert, int64_t n_expert_used, llm_ffn_op_type type_op, bool norm_w, float w_scale, llama_expert_gating_func_type gating_op, int il, ggml_tensor * probs_in, ggml_tensor * gate_up_exps, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, ggml_tensor * selected_experts_in) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, up_exps, /* up_exps_b */ nullptr, gate_exps, /* gate_exps_b */ nullptr, down_exps, /* down_exps_b */ nullptr, exp_probs_b, n_expert, n_expert_used, type_op, norm_w, w_scale, gating_op, il, probs_in, gate_up_exps, /* gate_up_exps_b */ nullptr, up_exps_s, gate_exps_s, down_exps_s, selected_experts_in );}
ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur, ggml_tensor * gate_inp, ggml_tensor * gate_inp_b, ggml_tensor * up_exps, ggml_tensor * up_exps_b, ggml_tensor * gate_exps, ggml_tensor * gate_exps_b, ggml_tensor * down_exps, ggml_tensor * down_exps_b, ggml_tensor * exp_probs_b, int64_t n_expert, int64_t n_expert_used, llm_ffn_op_type type_op, bool norm_w, float w_scale, llama_expert_gating_func_type gating_op, int il, ggml_tensor * probs_in, ggml_tensor * gate_up_exps, ggml_tensor * gate_up_exps_b, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, ggml_tensor * selected_experts_in) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN
ggml_tensor * logits = nullptr;
if (probs_in == nullptr) { logits = build_lora_mm(gate_inp, cur); // [n_expert, n_tokens] if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { ggml_mul_mat_set_prec(logits, GGML_PREC_F32); } cb(logits, "ffn_moe_logits", il); } else { logits = probs_in; }
if (gate_inp_b) { logits = ggml_add(ctx0, logits, gate_inp_b); cb(logits, "ffn_moe_logits_biased", il); }
ggml_tensor * probs = nullptr; switch (gating_op) { case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX: { probs = ggml_soft_max(ctx0, logits); // [n_expert, n_tokens] } break; case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID: { probs = ggml_sigmoid(ctx0, logits); // [n_expert, n_tokens] } break; case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT: { probs = logits; // [n_expert, n_tokens] } break; case LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS: { probs = ggml_sqrt(ctx0, ggml_softplus(ctx0, logits)); // [n_expert, n_tokens] } break; default: GGML_ABORT("fatal error"); } cb(probs, "ffn_moe_probs", il);
// add experts selection bias - introduced in DeepSeek V3 // leave probs unbiased as it's later used to get expert weights ggml_tensor * selection_probs = probs; if (exp_probs_b != nullptr) { selection_probs = ggml_add(ctx0, probs, exp_probs_b); cb(selection_probs, "ffn_moe_probs_biased", il); }
// llama4 doesn't have exp_probs_b, and sigmoid is only used after top_k // see: https://github.com/meta-llama/llama-models/blob/699a02993512fb36936b1b0741e13c06790bcf98/models/llama4/moe.py#L183-L198 if (arch == LLM_ARCH_LLAMA4) { selection_probs = logits; }
if (arch == LLM_ARCH_GROVEMOE) { selection_probs = ggml_sigmoid(ctx0, logits); // [n_expert, n_tokens] cb(selection_probs, "ffn_moe_probs_biased", il); }
// select top n_group_used expert groups // https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/e815299b0bcbac849fa540c768ef21845365c9eb/modeling_deepseek.py#L440-L457 if (hparams.n_expert_groups > 1 && n_tokens > 0) { const int64_t n_exp_per_group = n_expert / hparams.n_expert_groups;
// organize experts into n_expert_groups ggml_tensor * selection_groups = ggml_reshape_3d(ctx0, selection_probs, n_exp_per_group, hparams.n_expert_groups, n_tokens); // [n_exp_per_group, n_expert_groups, n_tokens]
ggml_tensor * group_scores = ggml_argsort_top_k(ctx0, selection_groups, 2); // [2, n_expert_groups, n_tokens] group_scores = ggml_get_rows(ctx0, ggml_reshape_4d(ctx0, selection_groups, 1, selection_groups->ne[0], selection_groups->ne[1], selection_groups->ne[2]), group_scores); // [1, 2, n_expert_groups, n_tokens]
// get top n_group_used expert groups group_scores = ggml_sum_rows(ctx0, ggml_reshape_3d(ctx0, group_scores, group_scores->ne[1], group_scores->ne[2], group_scores->ne[3])); // [1, n_expert_groups, n_tokens] group_scores = ggml_reshape_2d(ctx0, group_scores, group_scores->ne[1], group_scores->ne[2]); // [n_expert_groups, n_tokens]
ggml_tensor * expert_groups = ggml_argsort_top_k(ctx0, group_scores, hparams.n_group_used); // [n_group_used, n_tokens] cb(expert_groups, "ffn_moe_group_topk", il);
// mask out the other groups selection_probs = ggml_get_rows(ctx0, selection_groups, expert_groups); // [n_exp_per_group, n_group_used, n_tokens] selection_probs = ggml_set_rows(ctx0, ggml_fill(ctx0, selection_groups, -INFINITY), selection_probs, expert_groups); // [n_exp_per_group, n_expert_groups, n_tokens] selection_probs = ggml_reshape_2d(ctx0, selection_probs, n_expert, n_tokens); // [n_expert, n_tokens] cb(selection_probs, "ffn_moe_probs_masked", il); }
// select experts ggml_tensor * selected_experts = selected_experts_in; if (selected_experts == nullptr) { selected_experts = ggml_argsort_top_k(ctx0, selection_probs, n_expert_used); // [n_expert_used, n_tokens] cb(selected_experts->src[0], "ffn_moe_argsort", il); } cb(selected_experts, "ffn_moe_topk", il);
if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { // TODO: Use scalar div instead when/if implemented ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); selected_experts = ggml_cast(ctx0, ggml_scale(ctx0, f_sel, 1.0f / float(hparams.n_group_experts)), GGML_TYPE_I32); probs = ggml_reshape_3d(ctx0, probs, 1, hparams.n_expert, n_tokens); } else { probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens); }
ggml_tensor * weights = ggml_get_rows(ctx0, probs, selected_experts); // [1, n_expert_used, n_tokens] cb(weights, "ffn_moe_weights", il);
if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT) { weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens); weights = ggml_soft_max(ctx0, weights); // [n_expert_used, n_tokens] weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens); cb(weights, "ffn_moe_weights_softmax", il); }
if (norm_w) { weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens);
ggml_tensor * weights_sum = ggml_sum_rows(ctx0, weights); // [1, n_tokens] cb(weights_sum, "ffn_moe_weights_sum", il);
// Avoid division by zero, clamp to smallest number representable by F16 weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY); cb(weights_sum, "ffn_moe_weights_sum_clamped", il);
weights = ggml_div(ctx0, weights, weights_sum); // [n_expert_used, n_tokens] cb(weights, "ffn_moe_weights_norm", il);
weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens); } if (w_scale != 0.0f && w_scale != 1.0f) { weights = ggml_scale(ctx0, weights, w_scale); cb(weights, "ffn_moe_weights_scaled", il); }
//call early so that topk-moe can be used ggml_build_forward_expand(gf, weights);
cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
if (weight_before_ffn) { // repeat cur to [n_embd, n_expert_used, n_tokens] ggml_tensor * repeated = ggml_repeat_4d(ctx0, cur, n_embd, n_expert_used, n_tokens, 1); cur = ggml_mul(ctx0, repeated, weights); cb(cur, "ffn_moe_weighted", il); }
ggml_tensor * up = nullptr; ggml_tensor * experts = nullptr;
if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il);
if (up_exps_s) { cb(gate_up, "ffn_moe_gate_up_scaled", il); }
if (gate_up_exps_b) { gate_up = ggml_add_id(ctx0, gate_up, gate_up_exps_b, selected_experts); cb(gate_up, "ffn_moe_gate_up_biased", il); }
const int64_t n_ff = gate_up->ne[0] / 2; cur = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], 0); cb(cur, "ffn_moe_gate", il); up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], n_ff * gate_up->nb[0]); cb(up, "ffn_moe_up", il); } else { // separate gate and up path up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il);
if (up_exps_s) { cb(up, "ffn_moe_up_scaled", il); }
if (up_exps_b) { up = ggml_add_id(ctx0, up, up_exps_b, selected_experts); cb(up, "ffn_moe_up_biased", il); }
if (gate_exps) { cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; }
if (gate_exps_s) { cb(cur, "ffn_moe_gate_scaled", il); }
if (gate_exps_b) { cur = ggml_add_id(ctx0, cur, gate_exps_b, selected_experts); cb(cur, "ffn_moe_gate_biased", il); } }
const bool has_gate = gate_exps || gate_up_exps;
switch (type_op) { case LLM_FFN_SILU: if (gate_exps) { if (il >= 0) { const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { up = ggml_clamp(ctx0, up, -limit, limit); cb(up, "ffn_moe_up_clamped", il);
if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_moe_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, up); } else { ggml_tensor * gate_act = ggml_silu(ctx0, cur); cb(gate_act, "ffn_moe_silu", il); gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); cb(gate_act, "ffn_moe_silu_clamped", il); cur = ggml_mul(ctx0, gate_act, up); } cb(cur, "ffn_moe_swiglu_limited", il); break; } } }
if (has_gate) { cur = ggml_swiglu_split(ctx0, cur, up); cb(cur, "ffn_moe_swiglu", il); } else { cur = ggml_silu(ctx0, cur); cb(cur, "ffn_moe_silu", il); } break; case LLM_FFN_SITU: { // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb) GGML_ASSERT(has_gate); const float beta = hparams.situ_beta; const float lb = hparams.situ_linear_beta;
ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta); act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur)); if (lb > 0.0f) { up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb); } cur = ggml_mul(ctx0, act, up); cb(cur, "ffn_moe_situ", il); } break; case LLM_FFN_GELU: if (has_gate) { cur = ggml_geglu_split(ctx0, cur, up); cb(cur, "ffn_moe_geglu", il); } else { cur = ggml_gelu(ctx0, cur); cb(cur, "ffn_moe_gelu", il); } break; case LLM_FFN_SWIGLU_OAI_MOE: { // TODO: move to hparams? constexpr float alpha = 1.702f; constexpr float limit = 7.0f; cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit); cb(cur, "ffn_moe_swiglu_oai", il); } break; case LLM_FFN_RELU: if (has_gate) { cur = ggml_reglu_split(ctx0, cur, up); cb(cur, "ffn_moe_reglu", il); } else { cur = ggml_relu(ctx0, cur); cb(cur, "ffn_moe_relu", il); } break; case LLM_FFN_RELU_SQR: if (has_gate) { // TODO: add support for gated squared relu GGML_ABORT("fatal error: gated squared relu not implemented"); } else { cur = ggml_relu(ctx0, cur); cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_moe_relu_sqr", il); } break; default: GGML_ABORT("fatal error"); }
experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il);
if (down_exps_s) { cb(experts, "ffn_moe_down_scaled", il); }
if (down_exps_b) { experts = ggml_add_id(ctx0, experts, down_exps_b, selected_experts); cb(experts, "ffn_moe_down_biased", il); }
if (!weight_before_ffn) { experts = ggml_mul(ctx0, experts, weights); cb(experts, "ffn_moe_weighted", il); }
ggml_build_forward_expand(gf, experts);
ggml_tensor * cur_experts[LLAMA_MAX_EXPERTS] = { nullptr };
assert(n_expert_used > 0);
// order the views before the adds for (uint32_t i = 0; i < hparams.n_expert_used; ++i) { cur_experts[i] = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]);
ggml_build_forward_expand(gf, cur_experts[i]); }
// aggregate experts // note: here we explicitly use hparams.n_expert_used instead of n_expert_used // to avoid potentially a large number of add nodes during warmup // ref: https://github.com/ggml-org/llama.cpp/pull/14753 ggml_tensor * moe_out = cur_experts[0];
for (uint32_t i = 1; i < hparams.n_expert_used; ++i) { moe_out = ggml_add(ctx0, moe_out, cur_experts[i]);
ggml_build_forward_expand(gf, moe_out); }
if (hparams.n_expert_used == 1) { // avoid returning a non-contiguous tensor moe_out = ggml_cont(ctx0, moe_out); }
cb(moe_out, "ffn_moe_out", il);
return moe_out;}
ggml_tensor * llm_graph_context::build_embd_rows(ggml_tensor * tok_embd, ggml_tensor * ids) const { ggml_tensor * cur = ggml_get_rows(ctx0, tok_embd, ids);
// a Hadamard-latent embedding table stores rotated rows; restore the // primal basis right after the lookup: h = s * (H z) if (hadamard_inverses) { const auto it = hadamard_inverses->find(tok_embd); if (it != hadamard_inverses->end()) { cur = llama_mul_mat_hadamard(ctx0, cur, it->second.rot); if (it->second.signs) { cur = ggml_mul(ctx0, cur, it->second.signs); } } }
return cur;}
// input embeddings with optional loraggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { const int64_t n_embd_inp = hparams.n_embd_inp(); const int64_t n_embd = hparams.n_embd;
assert(n_embd_inp >= n_embd);
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens); cb(inp->tokens, "inp_tokens", -1); ggml_set_input(inp->tokens); res->t_inp_tokens = inp->tokens;
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens); cb(inp->embd, "inp_embd", -1); ggml_set_input(inp->embd);
// select one of the 2 inputs, based on the batch contents // ref: https://github.com/ggml-org/llama.cpp/pull/18550 std::array<ggml_tensor *, 2> inps;
// token embeddings path (ubatch.token != nullptr) { auto & cur = inps[0];
cur = build_embd_rows(tok_embd, inp->tokens);
// apply lora for embedding tokens if needed for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd); if (lw == nullptr) { continue; }
const float adapter_scale = lora.second; const float scale = lw->get_scale(lora.first->alpha, adapter_scale);
ggml_tensor * inpL_delta = ggml_scale(ctx0, ggml_mul_mat( ctx0, lw->b, // non-transposed lora_b ggml_get_rows(ctx0, lw->a, inp->tokens) ), scale);
cur = ggml_add(ctx0, cur, inpL_delta); }
if (n_embd_inp != n_embd) { cur = ggml_pad(ctx0, cur, hparams.n_embd_inp() - n_embd, 0, 0, 0); } }
// vector embeddings path (ubatch.embd != nullptr) { auto & cur = inps[1];
cur = inp->embd; }
assert(ggml_are_same_shape (inps[0], inps[1])); assert(ggml_are_same_stride(inps[0], inps[1]));
ggml_tensor * cur = ggml_build_forward_select(gf, inps.data(), inps.size(), ubatch.token ? 0 : 1);
if (n_embd_inp != n_embd) { cur = ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[1], 0); }
res->t_inp_embd = cur;
// For Granite architecture // NOTE: For deepstack models, only apply scale to token inputs (ie text-only input). // Raw embeddings are assumed to be multimodal inputs that should not be scaled. if (hparams.f_embedding_scale != 0.0f && (ubatch.token || hparams.n_deepstack_layers == 0)) { if (!ggml_is_contiguous(cur)) { cur = ggml_cont(ctx0, cur); } cur = ggml_scale(ctx0, cur, hparams.f_embedding_scale); }
cb(cur, "embd", -1);
res->add_input(std::move(inp));
// make sure the produced embeddings are immediately materialized in the ggml graph // ref: https://github.com/ggml-org/llama.cpp/pull/18599 ggml_build_forward_expand(gf, cur);
return cur;}
ggml_tensor * llm_graph_context::build_inp_pos() const { auto inp = std::make_unique<llm_graph_input_pos>(hparams.n_pos_per_embd());
auto & cur = inp->pos;
cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t)n_tokens*hparams.n_pos_per_embd()); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_attn_scale() const { auto inp = std::make_unique<llm_graph_input_attn_temp>(hparams.n_attn_temp_floor_scale, hparams.f_attn_temp_scale, hparams.f_attn_temp_offset);
auto & cur = inp->attn_scale;
// this need to be 1x1xN for broadcasting cur = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); ggml_set_input(cur); ggml_set_name(cur, "attn_scale");
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_out_ids() const { // note: when all tokens are output, we could skip this optimization to spare the ggml_get_rows() calls, // but this would make the graph topology depend on the number of output tokens, which can interfere with // features that require constant topology such as pipeline parallelism // ref: https://github.com/ggml-org/llama.cpp/pull/14275#issuecomment-2987424471 //if (n_outputs < n_tokens) { // return nullptr; //}
auto inp = std::make_unique<llm_graph_input_out_ids>(hparams, cparams, n_outputs);
auto & cur = inp->out_ids;
cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_outputs); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_mean() const { auto inp = std::make_unique<llm_graph_input_mean>(cparams);
auto & cur = inp->mean;
cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, ubatch.n_seqs_unq); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_cls() const { auto inp = std::make_unique<llm_graph_input_cls>(cparams, arch);
auto & cur = inp->cls;
cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_seqs_unq); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_cross_embd() const { auto inp = std::make_unique<llm_graph_input_cross_embd>(cross);
auto & cur = inp->cross_embd;
// if we have the output embeddings from the encoder, use them directly // TODO: needs more work to be correct, for now just use the tensor shape //if (cross->t_embd) { // cur = ggml_view_tensor(ctx0, cross->t_embd);
// return cur; //}
const auto n_embd = !cross->v_embd.empty() ? cross->n_embd : hparams.n_embd_inp(); const auto n_enc = !cross->v_embd.empty() ? cross->n_enc : hparams.n_ctx_train;
cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_enc); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_pos_bucket_enc() const { auto inp = std::make_unique<llm_graph_input_pos_bucket>(hparams);
auto & cur = inp->pos_bucket;
cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_tokens, n_tokens); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_inp_pos_bucket_dec() const { const auto * mctx_cur = static_cast<const llama_kv_cache_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_pos_bucket_kv>(hparams, mctx_cur);
const auto n_kv = mctx_cur->get_n_kv();
auto & cur = inp->pos_bucket;
cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens); ggml_set_input(cur);
res->add_input(std::move(inp));
return cur;}
ggml_tensor * llm_graph_context::build_pos_bias(ggml_tensor * pos_bucket, ggml_tensor * attn_rel_b) const { ggml_tensor * pos_bucket_1d = ggml_reshape_1d(ctx0, pos_bucket, pos_bucket->ne[0] * pos_bucket->ne[1]); cb(pos_bucket_1d, "pos_bucket_1d", -1);
ggml_tensor * pos_bias = ggml_get_rows(ctx0, attn_rel_b, pos_bucket_1d);
pos_bias = ggml_reshape_3d(ctx0, pos_bias, pos_bias->ne[0], pos_bucket->ne[0], pos_bucket->ne[1]); pos_bias = ggml_permute (ctx0, pos_bias, 2, 0, 1, 3); pos_bias = ggml_cont (ctx0, pos_bias);
cb(pos_bias, "pos_bias", -1);
return pos_bias;}
ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, ggml_tensor * kq_b, ggml_tensor * kq_mask, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { const bool v_trans = v->nb[1] > v->nb[2];
// split the batch into streams if needed const auto n_stream = k->ne[3];
q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[3]/n_stream, 0);
q = ggml_permute(ctx0, q, 0, 2, 1, 3); k = ggml_permute(ctx0, k, 0, 2, 1, 3); v = ggml_permute(ctx0, v, 0, 2, 1, 3);
ggml_tensor * cur;
const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr; if (use_flash_attn) { GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet");
if (v_trans) { v = ggml_transpose(ctx0, v); }
// this can happen when KV cache is not used (e.g. an embedding model with non-causal attn) if (k->type == GGML_TYPE_F32) { k = ggml_cast(ctx0, k, GGML_TYPE_F16); }
if (v->type == GGML_TYPE_F32) { v = ggml_cast(ctx0, v, GGML_TYPE_F16); }
cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
if (v_mla) {#if 0 // v_mla can be applied as a matrix-vector multiplication with broadcasting across dimension 3 == n_tokens. // However, the code is optimized for dimensions 0 and 1 being large, so this is inefficient. cur = ggml_reshape_4d(ctx0, cur, v_mla->ne[0], 1, n_head, n_tokens); cur = ggml_mul_mat(ctx0, v_mla, cur);#else // It's preferable to do the calculation as a matrix-matrix multiplication with n_tokens in dimension 1. // The permutations are noops and only change how the tensor data is interpreted. cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); cur = ggml_mul_mat(ctx0, v_mla, cur); cb(cur, "fattn_mla", il); cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); cur = ggml_cont(ctx0, cur); // Needed because ggml_reshape_2d expects contiguous inputs.#endif }
cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); } else { ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); cb(kq, "kq", il);
// note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
if (arch == LLM_ARCH_GROK) { // need to do the following: // multiply by attn_output_multiplier // and then : // kq = 30 * tanh(kq / 30) // before the softmax below
kq = ggml_tanh(ctx0, ggml_scale(ctx0, kq, hparams.f_attn_out_scale / hparams.f_attn_logit_softcapping)); cb(kq, "kq_tanh", il); kq = ggml_scale(ctx0, kq, hparams.f_attn_logit_softcapping); cb(kq, "kq_scaled", il); }
if (hparams.attn_soft_cap) { kq = ggml_scale(ctx0, kq, 1.0f / hparams.f_attn_logit_softcapping); cb(kq, "kq_scaled_1", il); kq = ggml_tanh (ctx0, kq); cb(kq, "kq_tanh", il); kq = ggml_scale(ctx0, kq, hparams.f_attn_logit_softcapping); cb(kq, "kq_scaled_2", il); }
if (kq_b) { kq = ggml_add(ctx0, kq, kq_b); cb(kq, "kq_plus_kq_b", il); }
kq = ggml_soft_max_ext(ctx0, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias); ggml_soft_max_add_sinks(kq, sinks); cb(kq, "kq_soft_max", il);
if (!v_trans) { // note: avoid this branch v = ggml_cont(ctx0, ggml_transpose(ctx0, v)); cb(v, "v_cont", il); }
ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq); cb(kqv, "kqv", il);
// for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA if (v_mla) { kqv = ggml_mul_mat(ctx0, v_mla, kqv); cb(kqv, "kqv_mla", il); }
cur = ggml_permute(ctx0, kqv, 0, 2, 1, 3);
// recombine streams cur = ggml_cont_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]);
if (!cparams.offload_kqv) { // all nodes between the KV store and the attention output are run on the CPU ggml_backend_sched_set_tensor_backend(sched, cur, backend_cpu); } }
ggml_build_forward_expand(gf, cur);
return cur;}
llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() const { auto inp = std::make_unique<llm_graph_input_attn_no_cache>(hparams, cparams);
// flash attention requires an f16 mask const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
// note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); ggml_set_input(inp->self_kq_mask);
inp->self_kq_mask_cnv = inp->self_kq_mask;
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); ggml_set_input(inp->self_kq_mask_swa);
inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; } else { inp->self_kq_mask_swa = nullptr; inp->self_kq_mask_swa_cnv = nullptr; }
return (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_no_cache * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { GGML_UNUSED(n_tokens);
// these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced ggml_build_forward_expand(gf, q_cur); ggml_build_forward_expand(gf, k_cur); ggml_build_forward_expand(gf, v_cur);
const bool is_swa = hparams.is_swa(il);
const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
// [TAG_NO_CACHE_PAD] // TODO: if ubatch.equal_seqs() == true, we can split the three tensors below into ubatch.n_seqs_unq streams // but it might not be worth it: https://github.com/ggml-org/llama.cpp/pull/15636 //assert(!ubatch.equal_seqs() || (k_cur->ne[3] == 1 && k_cur->ne[3] == ubatch.n_seqs_unq));
ggml_tensor * q = q_cur; ggml_tensor * k = k_cur; ggml_tensor * v = v_cur;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (wo) { cur = build_lora_mm(wo, cur, wo_s); }
if (wo_b) { //cb(cur, "kqv_wo", il); }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
static std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl( ggml_context * ctx0, const llama_ubatch & ubatch, const llama_hparams & hparams, const llama_cparams & cparams, const llama_kv_cache_context * mctx_cur) {
auto inp = std::make_unique<llm_graph_input_attn_kv>(hparams, cparams, mctx_cur);
{ GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; }
inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0); inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0);
return inp;}
llm_graph_input_attn_kv * llm_graph_context::build_attn_inp_kv() const { const auto * mctx_cur = static_cast<const llama_kv_cache_context *>(mctx);
auto inp = build_attn_inp_kv_impl(ctx0, ubatch, hparams, cparams, mctx_cur);
return (llm_graph_input_attn_kv *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_kv * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, // TODO: remove float kq_scale, int il) const { GGML_ASSERT(v_mla == nullptr);
if (inp->self_k_rot) { q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); }
if (inp->self_v_rot) { v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); }
// these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced // expand k later to enable rope fusion which directly writes into k-v cache ggml_build_forward_expand(gf, q_cur); ggml_build_forward_expand(gf, v_cur); ggml_build_forward_expand(gf, k_cur);
const auto * mctx_cur = inp->mctx;
// store to KV cache { const auto & k_idxs = inp->get_k_idxs(); const auto & v_idxs = inp->get_v_idxs();
// hook point for KV cache calibration tooling (e.g. tools/kv-mean-center): this is // exactly the K tensor that cpy_k() writes into the cache, after any RoPE/rotation // the architecture applies upstream cb(k_cur, "k_cache_in", il);
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); }
ggml_tensor * kq_mask = inp->get_kq_mask();
ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (inp->self_v_rot) { cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); }
if (wo) { if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators cur = build_lora_mm(wo, cur); ggml_mul_mat_set_prec(cur, GGML_PREC_F32); if (wo_s) { cur = ggml_mul(ctx0, cur, wo_s); } } else { cur = build_lora_mm(wo, cur, wo_s); } }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
static std::unique_ptr<llm_graph_input_attn_k> build_attn_inp_k_impl( ggml_context * ctx0, const llama_ubatch & ubatch, const llama_hparams & hparams, const llama_cparams & cparams, const llama_kv_cache_context * mctx_cur) {
auto inp = std::make_unique<llm_graph_input_attn_k>(hparams, cparams, mctx_cur);
{ GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; }
return inp;}
llm_graph_input_attn_k * llm_graph_context::build_attn_inp_k() const { const auto * mctx_cur = static_cast<const llama_kv_cache_context *>(mctx);
auto inp = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur);
return (llm_graph_input_attn_k *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { // these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced // expand k later to enable rope fusion which directly writes into k-v cache ggml_build_forward_expand(gf, q_cur); ggml_build_forward_expand(gf, v_cur); ggml_build_forward_expand(gf, k_cur);
const auto * mctx_cur = inp->mctx;
// store to KV cache { const auto & k_idxs = inp->get_k_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); }
const auto & kq_mask = inp->get_kq_mask();
ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (wo) { if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) { // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators cur = build_lora_mm(wo, cur); ggml_mul_mat_set_prec(cur, GGML_PREC_F32); if (wo_s) { cur = ggml_mul(ctx0, cur, wo_s); } } else { cur = build_lora_mm(wo, cur, wo_s); } }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_dsa * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, ggml_tensor * top_k, float kq_scale, int il) const { // these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced // expand k later to enable rope fusion which directly writes into k-v cache ggml_build_forward_expand(gf, q_cur); ggml_build_forward_expand(gf, v_cur); ggml_build_forward_expand(gf, k_cur);
const auto * mctx_cur = inp->mctx->get_mla();
// store to KV cache { const auto & k_idxs = inp->get_k_idxs_mla();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); }
const auto & kq_mask = inp->get_kq_mask_mla();
// prepare new kq mask - starts filled with -INFINITY ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
// reshape KQ mask into tensor with rows of size 1: // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
// reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
// prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] // this will be our source of zero values for unmasking top k mask elements ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f);
// modify KQ mask by unmasking elements that are in top_k indices // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
// reshape to restore the original shape of KQ mask: // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
// combine with the original kq mask kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (wo) { cur = build_lora_mm(wo, cur, wo_s); }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_kv_iswa * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { const bool is_swa = hparams.is_swa(il);
auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot; auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot;
if (k_rot) { q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); if (k_cur) { k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); } } if (v_rot) { if (v_cur) { v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot); } }
// these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced ggml_build_forward_expand(gf, q_cur);
if (k_cur) { ggml_build_forward_expand(gf, k_cur); }
if (v_cur) { ggml_build_forward_expand(gf, v_cur); }
const auto * mctx_iswa = inp->mctx;
const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
// optionally store to KV cache if (k_cur) { const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); }
if (v_cur) { const auto & v_idxs = is_swa ? inp->get_v_idxs_swa() : inp->get_v_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); }
const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (v_rot) { cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); }
if (wo) { cur = build_lora_mm(wo, cur, wo_s); }
if (wo_b) { //cb(cur, "kqv_wo", il); }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_iswa * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { const bool is_swa = hparams.is_swa(il);
auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
if (k_rot) { q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); if (k_cur) { k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); } }
// these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced ggml_build_forward_expand(gf, q_cur);
if (k_cur) { ggml_build_forward_expand(gf, k_cur); }
const auto * mctx_iswa = inp->mctx; const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
// optionally store to KV cache if (k_cur) { const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); }
const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
// MLA-style attention: the cached K is used as V ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (k_rot) { cur = llama_mul_mat_hadamard(ctx0, cur, k_rot); }
if (wo) { cur = build_lora_mm(wo, cur, wo_s); }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const { auto inp = std::make_unique<llm_graph_input_attn_cross>(cross);
const int32_t n_enc = !cross->v_embd.empty() ? cross->n_enc : hparams.n_ctx_train;
// flash attention requires an f16 mask const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
inp->cross_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_enc, n_tokens, 1, 1); ggml_set_input(inp->cross_kq_mask);
inp->cross_kq_mask_cnv = inp->cross_kq_mask;
return (llm_graph_input_attn_cross *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_cross * inp, ggml_tensor * wo, ggml_tensor * wo_b, ggml_tensor * wo_s, ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, int il) const { // these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced ggml_build_forward_expand(gf, q_cur); ggml_build_forward_expand(gf, k_cur); ggml_build_forward_expand(gf, v_cur);
const auto & kq_mask = inp->get_kq_mask_cross();
ggml_tensor * q = q_cur; ggml_tensor * k = k_cur; ggml_tensor * v = v_cur;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il);
if (wo) { cur = build_lora_mm(wo, cur, wo_s); }
if (wo_b) { //cb(cur, "kqv_wo", il); }
if (wo_b) { cur = ggml_add(ctx0, cur, wo_b); }
return cur;}
static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl( ggml_context * ctx0, const llama_ubatch & ubatch, const llama_hparams & hparams, const llama_cparams & cparams, const llama_kv_cache_dsa_context * mctx_cur) {
auto inp = std::make_unique<llm_graph_input_attn_k_dsa>(hparams, cparams, mctx_cur);
{ inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams); inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla; }
{ inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
// ensure that mask type matches fused lightning indexer use (requires f16 mask) auto cparams_copy = cparams; cparams_copy.flash_attn = cparams.fused_lid;
inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy); inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;
inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0); }
return inp;}
llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_context *>(mctx);
auto inp = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur);
return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));}
llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa() const { const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_iswa_context *>(mctx);
auto inp_dsa = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_dsa());
// build_attn_inp_k_impl rejects SWA caches, so construct the input directly auto inp_swa = std::make_unique<llm_graph_input_attn_k>(hparams, cparams, mctx_cur->get_swa());
inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask;
auto inp = std::make_unique<llm_graph_input_attn_k_dsa_iswa>(std::move(inp_dsa), std::move(inp_swa), mctx_cur);
return (llm_graph_input_attn_k_dsa_iswa *) res->add_input(std::move(inp));}
llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const { const auto * mctx_cur = static_cast<const llama_kv_cache_msa_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_attn_kv_msa>(hparams, cparams, mctx_cur);
const auto * mctx_base = mctx_cur->get_base(); const auto * mctx_idx = mctx_cur->get_idx();
{ GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch); inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; }
inp->self_k_rot = mctx_base->build_input_k_rot(ctx0); inp->self_v_rot = mctx_base->build_input_v_rot(ctx0);
if (msa_enabled) { inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch); }
return (llm_graph_input_attn_kv_msa *) res->add_input(std::move(inp));}
// TODO: maybe separate the inner implementation into a separate function// like with the non-sliding window equivalent// once sliding-window hybrid caches are a thing.llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const { const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_attn_kv_iswa>(hparams, cparams, mctx_cur);
{ inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; }
{ GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA");
inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch);
inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; }
inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0); inp->self_v_rot = mctx_cur->get_base()->build_input_v_rot(ctx0);
inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0); inp->self_v_rot_swa = mctx_cur->get_swa()->build_input_v_rot(ctx0);
return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));}
llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_attn_k_iswa>(hparams, cparams, mctx_cur);
{ inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; }
{ GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA");
inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; }
inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp));}
llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { const auto * mctx_cur = static_cast<const llama_kv_cache_dsv4_context *>(mctx); const auto * raw_ctx = mctx_cur->get_raw();
auto inp_raw = std::make_unique<llm_graph_input_dsv4_raw>(cparams, raw_ctx);
const int64_t n_stream = mctx_cur->get_csa_plan(ubatch).n_stream;
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "DSV4 expects SWA raw cache");
inp_raw->self_k_idxs = raw_ctx->build_input_k_idxs(ctx0, ubatch); inp_raw->self_kq_mask = dsv4_build_raw_kq_mask(ctx0, raw_ctx, ubatch, cparams, n_stream); inp_raw->self_kq_mask_cnv = inp_raw->self_kq_mask;
inp_raw->self_k_rot = raw_ctx->build_input_k_rot(ctx0); auto inp = std::make_unique<llm_graph_input_dsv4>(cparams, std::move(inp_raw), mctx_cur);
dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", cparams, n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", cparams, n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", cparams, n_stream); inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0); inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0); inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0);
return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_rs( ggml_tensor * s, ggml_tensor * state_copy_main, ggml_tensor * state_copy_extra, int32_t state_size, int32_t n_seqs, uint32_t n_rs, uint32_t rs_head, uint32_t rs_size, int32_t rs_zero, const llm_graph_get_rows_fn & get_state_rows) const {
GGML_UNUSED(rs_size); ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, s->ne[1]);
// Clear a single state which will then be copied to the other cleared states. // Note that this is a no-op when the view is zero-sized. ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0)); ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0));
// copy states // NOTE: assuming the copy destinations are ALL contained between rs_head and rs_head + n_rs // {state_size, rs_size} -> {state_size, n_seqs} ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main); ggml_build_forward_expand(gf, output_states);
// copy extra states which won't be changed further (between n_seqs and n_rs) ggml_tensor * states_extra = ggml_get_rows(ctx0, states, state_copy_extra); ggml_build_forward_expand(gf, ggml_cpy(ctx0, states_extra, ggml_view_2d(ctx0, s, state_size, (n_rs - n_seqs), s->nb[1], (rs_head + n_seqs)*s->nb[1])));
return output_states;}
static std::unique_ptr<llm_graph_input_rs> build_rs_inp_impl( ggml_context * ctx0, const llama_ubatch & ubatch, const llama_memory_recurrent_context * mctx_cur) {
auto inp = std::make_unique<llm_graph_input_rs>(mctx_cur);
const int64_t n_rs = mctx_cur->get_n_rs(); const int64_t n_seqs = ubatch.n_seqs;
inp->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs); ggml_set_input(inp->s_copy);
inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0); inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]);
inp->head = mctx_cur->get_head(); inp->rs_z = mctx_cur->get_rs_z();
inp->s_copy_main_in_place = mctx_cur->s_copy_main_in_place(n_seqs);
return inp;}
llm_graph_input_rs * llm_graph_context::build_rs_inp() const { const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);
auto inp = build_rs_inp_impl(ctx0, ubatch, mctx_cur);
return (llm_graph_input_rs *) res->add_input(std::move(inp));}
ggml_tensor * llm_graph_context::build_rs( llm_graph_input_rs * inp, ggml_tensor * s, int32_t state_size, int32_t n_seqs, const llm_graph_get_rows_fn & get_state_rows) const { const auto * kv_state = inp->mctx;
return build_rs(s, inp->s_copy_main, inp->s_copy_extra, state_size, n_seqs, kv_state->get_n_rs(), kv_state->get_head(), kv_state->get_size(), kv_state->get_rs_z(), get_state_rows);}
ggml_tensor * llm_graph_context::build_rs_cache_view( llm_graph_input_rs * inp, ggml_tensor * s, int32_t state_size, int32_t n_seqs) const { const auto * kv_state = inp->mctx;
const uint32_t n_rs = kv_state->get_n_rs(); const uint32_t rs_head = kv_state->get_head(); const int32_t rs_zero = kv_state->get_rs_z();
ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, s->ne[1]);
// same cache hygiene as build_rs, minus the main gather (the consumer reads // per-seq rows via inp->s_copy_main directly, inside the GDN op). // // build_rs gathers the main rows BEFORE this extra relocation; rows mode // defers the main read into the consumer, after it. That is only safe while // every main row lies in its own seq's cell (inp->s_copy_main_in_place): // during a cache reorder or a fresh seq reading another cell's zeroed row, // the relocation or another seq's snapshot write could clobber a row the // consumer reads later, so callers keep build_rs then. ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0)); ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0));
ggml_tensor * states_extra = ggml_get_rows(ctx0, states, inp->s_copy_extra); ggml_build_forward_expand(gf, ggml_cpy(ctx0, states_extra, ggml_view_2d(ctx0, s, state_size, (n_rs - n_seqs), s->nb[1], (rs_head + n_seqs)*s->nb[1])));
return states;}
ggml_tensor * llm_graph_context::build_rs_write_rows( llm_graph_input_rs * inp, int64_t K, int64_t n_seq_tokens, int64_t n_seqs) const { const int64_t n_write = std::min(n_seq_tokens, K);
if (inp->s_write_rows == nullptr) { inp->s_write_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_write * n_seqs); ggml_set_input(inp->s_write_rows); inp->s_write_K = K; }
GGML_ASSERT(inp->s_write_K == K); GGML_ASSERT(inp->s_write_rows->ne[0] == n_write * n_seqs);
return inp->s_write_rows;}
ggml_tensor * llm_graph_context::build_rwkv_token_shift_load( llm_graph_input_rs * inp, const llama_ubatch & ubatch, int il) const { const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);
const auto token_shift_count = hparams.token_shift_count;
const int64_t n_seqs = ubatch.n_seqs;
ggml_tensor * token_shift_all = mctx_cur->get_r_l(il);
ggml_tensor * token_shift = build_rs( inp, token_shift_all, hparams.n_embd_r(), n_seqs);
token_shift = ggml_reshape_3d(ctx0, token_shift, hparams.n_embd, token_shift_count, n_seqs);
return token_shift;}
ggml_tensor * llm_graph_context::build_rwkv_token_shift_store( ggml_tensor * token_shift, const llama_ubatch & ubatch, int il) const { const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);
const auto token_shift_count = hparams.token_shift_count; const auto n_embd = hparams.n_embd;
const int64_t n_seqs = ubatch.n_seqs;
const auto kv_head = mctx_cur->get_head();
return ggml_cpy( ctx0, ggml_view_1d(ctx0, token_shift, n_embd * n_seqs * token_shift_count, 0), ggml_view_1d(ctx0, mctx_cur->get_r_l(il), hparams.n_embd_r()*n_seqs, hparams.n_embd_r()*kv_head*ggml_element_size(mctx_cur->get_r_l(il))) );}
llm_graph_input_mem_hybrid * llm_graph_context::build_inp_mem_hybrid() const { const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx);
auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr()); auto inp_attn = build_attn_inp_kv_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn());
auto inp = std::make_unique<llm_graph_input_mem_hybrid>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
return (llm_graph_input_mem_hybrid *) res->add_input(std::move(inp));}
llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const { const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx);
auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr()); auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn());
auto inp = std::make_unique<llm_graph_input_mem_hybrid_k>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
return (llm_graph_input_mem_hybrid_k *) res->add_input(std::move(inp));}
llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() const { const auto * mctx_cur = static_cast<const llama_memory_hybrid_iswa_context *>(mctx);
auto inp_rs = build_rs_inp_impl(ctx0, ubatch, mctx_cur->get_recr());
// build iswa attention input const auto * attn_ctx = mctx_cur->get_attn();
auto inp_attn = std::make_unique<llm_graph_input_attn_kv_iswa>(hparams, cparams, attn_ctx);
{ inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch); inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch);
inp_attn->self_kq_mask = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams); inp_attn->self_kq_mask_cnv = inp_attn->self_kq_mask; }
{ inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch); inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch);
inp_attn->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams); inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa; }
auto inp = std::make_unique<llm_graph_input_mem_hybrid_iswa>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
return (llm_graph_input_mem_hybrid_iswa *) res->add_input(std::move(inp));}
void llm_graph_context::build_dense_out( ggml_tensor * dense_2, ggml_tensor * dense_2_b, ggml_tensor * dense_3) const { if (!cparams.embeddings || !(dense_2 || dense_2_b || dense_3)) { return; } ggml_tensor * cur = res->t_embd_pooled != nullptr ? res->t_embd_pooled : res->t_embd; GGML_ASSERT(cur != nullptr && "missing t_embd_pooled/t_embd");
if (dense_2) { cur = ggml_mul_mat(ctx0, dense_2, cur); } if (dense_2_b) { cur = ggml_add(ctx0, cur, dense_2_b); } if (dense_3) { cur = ggml_mul_mat(ctx0, dense_3, cur); } cb(cur, "result_embd_pooled", -1); res->t_embd_pooled = cur; ggml_build_forward_expand(gf, cur);}
void llm_graph_context::build_pooling( ggml_tensor * cls, ggml_tensor * cls_b, ggml_tensor * cls_out, ggml_tensor * cls_out_b, ggml_tensor * cls_norm) const { if (!cparams.embeddings) { return; }
ggml_tensor * inp = res->t_embd;
//// find result_norm tensor for input //for (int i = ggml_graph_n_nodes(gf) - 1; i >= 0; --i) { // inp = ggml_graph_node(gf, i); // if (strcmp(inp->name, "result_norm") == 0 || strcmp(inp->name, "result_embd") == 0) { // break; // }
// inp = nullptr; //}
GGML_ASSERT(inp != nullptr && "missing result_norm/result_embd tensor");
ggml_tensor * cur;
switch (pooling_type) { case LLAMA_POOLING_TYPE_NONE: { cur = inp; } break; case LLAMA_POOLING_TYPE_MEAN: { ggml_tensor * inp_mean = build_inp_mean(); cur = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, inp)), inp_mean); } break; case LLAMA_POOLING_TYPE_CLS: case LLAMA_POOLING_TYPE_LAST: { ggml_tensor * inp_cls = build_inp_cls(); cur = ggml_get_rows(ctx0, inp, inp_cls); } break; case LLAMA_POOLING_TYPE_RANK: { if (arch == LLM_ARCH_MODERN_BERT) { // modern bert gte reranker builds mean first then applies prediction head and classifier // https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modular_modernbert.py#L1404-1411 ggml_tensor * inp_mean = build_inp_mean(); cur = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, inp)), inp_mean); } else { ggml_tensor * inp_cls = build_inp_cls(); cur = ggml_get_rows(ctx0, inp, inp_cls); }
// classification head // https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/roberta/modeling_roberta.py#L1566 if (cls) { cur = ggml_mul_mat(ctx0, cls, cur); if (cls_b) { cur = ggml_add(ctx0, cur, cls_b); } if (arch == LLM_ARCH_MODERN_BERT) { cur = ggml_gelu(ctx0, cur); } else { cur = ggml_tanh(ctx0, cur); } if (cls_norm) { // head norm cur = build_norm(cur, cls_norm, NULL, LLM_NORM, -1); } }
// some models don't have `cls_out`, for example: https://huggingface.co/jinaai/jina-reranker-v1-tiny-en // https://huggingface.co/jinaai/jina-reranker-v1-tiny-en/blob/cb5347e43979c3084a890e3f99491952603ae1b7/modeling_bert.py#L884-L896 // Single layer classification head (direct projection) // https://github.com/huggingface/transformers/blob/f4fc42216cd56ab6b68270bf80d811614d8d59e4/src/transformers/models/bert/modeling_bert.py#L1476 if (cls_out) { cur = ggml_mul_mat(ctx0, cls_out, cur); if (cls_out_b) { cur = ggml_add(ctx0, cur, cls_out_b); } }
// softmax for qwen3 reranker if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_QWEN3VL) { cur = ggml_soft_max(ctx0, cur); } } break; default: { GGML_ABORT("unknown pooling type"); } }
cb(cur, "result_embd_pooled", -1); res->t_embd_pooled = cur;
ggml_build_forward_expand(gf, cur);}
void llm_graph_context::build_sampling() const { if (samplers.empty() || !res->t_logits) { return; }
std::array<ggml_tensor *, 2> outs; outs[0] = res->t_logits;
auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers); res->add_input(std::move(inp_sampling));
std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows; uint32_t n_rows = 0; for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { if (ubatch.output[i]) { sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++); } }
res->t_sampled.resize(n_rows, nullptr); res->t_sampled_probs.resize(n_rows, nullptr); res->t_sampled_logits.resize(n_rows, nullptr); res->t_candidates.resize(n_rows, nullptr);
// res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1) GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");
// add a dummy row to keep the single-output graph static regardless of active samplers // multi-output graphs can still vary with the number of output rows // DSpark reserve graphs can expose fewer draft rows than the output budget. const int64_t pad_rows = res->t_dspark_greedy.empty() ? 1 : std::max<int64_t>(1, int64_t(n_rows) - res->t_logits->ne[1]); ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, pad_rows, 0, 0);
for (const auto & entry : samplers) { if (entry.second->iface->backend_reset) { entry.second->iface->backend_reset(entry.second); } }
static const std::vector<uint32_t> dummy_row = { 0 };
for (const auto & [seq_id, sampler] : samplers) { const auto it = sampling_rows.find(seq_id);
// inactive samplers always work on the first row const bool active = it != sampling_rows.end(); const auto & rows = active ? it->second : dummy_row; const int i_out = active ? 1 : 0;
if (active && res->t_dspark_greedy.size() == n_rows && llama_sampler_chain_n(sampler) == 1 && std::strcmp(llama_sampler_name(llama_sampler_chain_get(sampler, 0)), "greedy") == 0) { for (uint32_t row : rows) { res->t_sampled[row] = res->t_dspark_greedy[row]; ggml_build_forward_expand(gf, res->t_sampled[row]); } continue; }
for (uint32_t i = 0; i < rows.size(); ++i) { ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]); ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
struct llama_sampler_data data = { /*.logits =*/ logits_seq, /*.probs =*/ nullptr, /*.sampled =*/ nullptr, /*.candidates =*/ nullptr, };
assert(sampler->iface->backend_apply); sampler->iface->backend_apply(sampler, ctx0, gf, &data);
if (data.sampled != nullptr) { if (active) { res->t_sampled[rows[i]] = data.sampled; } outs[1] = data.sampled; ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); }
if (data.probs != nullptr) { if (active) { res->t_sampled_probs[rows[i]] = data.probs; } outs[1] = data.probs; ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); }
if (data.logits != nullptr) { if (active) { res->t_sampled_logits[rows[i]] = data.logits; } outs[1] = data.logits; ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); }
if (data.candidates != nullptr) { if (active) { res->t_candidates[rows[i]] = data.candidates; } outs[1] = data.candidates; ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); } } }
// TODO: Call backend_accept after all samplers have been applied. /* for (const auto & [seq_id, sampler] : samplers) { const auto it = sampling_rows.find(seq_id); if (it == sampling_rows.end()) { continue; }
for (uint32_t row : it->second) { ggml_tensor * selected_token = res->t_sampled[row]; if (selected_token != nullptr && sampler->iface->backend_accept) { sampler->iface->backend_accept(sampler, ctx0, gf, selected_token); } } } */}
int32_t llama_relative_position_bucket(llama_pos x, llama_pos y, uint64_t n_buckets, bool bidirectional) { // TODO move to hparams if a T5 variant appears that uses a different value const int64_t max_distance = 128;
if (bidirectional) { n_buckets >>= 1; }
const int64_t max_exact = n_buckets >> 1;
int32_t relative_position = x - y; int32_t relative_bucket = 0;
if (bidirectional) { relative_bucket += (relative_position > 0) * n_buckets; relative_position = std::abs(relative_position); } else { relative_position = -std::min<int32_t>(relative_position, 0); }
int32_t relative_position_if_large = floorf(max_exact + logf(1.0 * relative_position / max_exact) * (n_buckets - max_exact) / log(1.0 * max_distance / max_exact)); relative_position_if_large = std::min<int32_t>(relative_position_if_large, n_buckets - 1); relative_bucket += (relative_position < max_exact ? relative_position : relative_position_if_large);
return relative_bucket;}