ఊహాగానాల ద్వారా Sonarను వేగవంతం చేయడం
డ్రాఫ్ట్ మోడల్ ధృవీకరణ ద్వారా ఊహాజనిత డికోడింగ్ Sonar LLMలను వేగవంతం చేస్తుంది
పెర్కివులేటివ్ డికోడింగ్ (Speculative decoding) వేగవంతమైన మరియు చిన్న డ్రాఫ్ట్ మోడల్ను ఉపయోగించి పెద్ద టార్గెట్ మోడల్ ద్వారా ధృవీకరించబడే పూర్తి అభ్యర్థులను ఉత్పత్తి చేయడం ద్వారా లార్జ్ లాంగ్వేజ్ మోడల్స్ (LLMలు) యొక్క జనరేషన్ వేగాన్ని వేగవంతం చేస్తుంది. ఈ విధానంలో, ఖరీదైన టార్గెట్ యొక్క రన్ ఒక్కో టోకెన్ను ఉత్పత్తి చేయడానికి బదులుగా, ఒకే దశలో బహుళ టోకెన్లు విడుదల చేయబడతాయి. ఇక్కడ సోనార్ మోడళ్లపై టోకెన్ల మధ్య లేటెన్సీని తగ్గించడానికి Perplexity వద్ద వర్తింపజేయబడిన వివిధ రకాల ఊహాజనిత డికోడింగ్ యొక్క అమలు వివరాలను మేము అందిస్తున్నాము.
ఊహాజనిత డికోడింగ్
ఊహాజనిత డికోడింగ్ అనేది టోకెన్ జనరేషన్ను వేగవంతం చేయడానికి సహజ భాషల నిర్మాణాన్ని మరియు ట్రాన్స్ఫార్మర్ల ఆటో-రిగ్రెసివ్ స్వభావాన్ని సద్వినియోగం చేసుకుంటుంది. Llama-70B వంటి పెద్ద మోడళ్లు, Llama-1B వంటి చిన్నవాటి కంటే ఎక్కువ జ్ఞానాన్ని కలిగి ఉన్నప్పటికీ, కొన్ని సరళమైన పనులపై అవి ఒకేలా పనిచేస్తాయి. ఈ అతివ్యాప్తి కొన్ని సీక్వెన్స్లు తక్కువ ఖరీదైన మోడళ్ల ద్వారా మెరుగ్గా రూపొందించబడతాయని సూచిస్తుంది, సంక్లిష్టమైన సమస్యలను పెద్ద వాటికి వదిలివేస్తుంది. ఏ పూర్తి చేయడాలు మెరుగ్గా ఉన్నాయి మరియు చిన్న మోడల్ యొక్క జనరేషన్ పెద్ద మోడల్ వలె నాణ్యతతో ఉందా లేదా అని నిర్ణయించడంలో సవాలు ఉంది.
అదృష్టవశాత్తూ, LLMలు ఆటో-రిగ్రెసివ్ ట్రాన్స్ఫార్మర్లు: టోకెన్ల సీక్వెన్స్ ఇవ్వబడినప్పుడు, అవి తదుపరి టోకెన్ యొక్క సంభావ్యత పంపిణీని అవుట్పుట్ చేస్తాయి. అదనంగా, ఇన్పుట్ సీక్వెన్స్లోని టోకెన్లతో ముడిపడి ఉన్న మధ్యవర్తిత్వ ఫీచర్ల నుండి తీసుకోబడిన లాజిట్లు మోడల్ ఆ ఖచ్చితమైన టోకెన్లను జారీ చేయడానికి ఎంత అవకాశం ఉందో కూడా సూచిస్తాయి. ఈ గుణం ఊహాగానాలను సృష్టిస్తుంది: ఒక ఇన్పుట్ ప్రిఫిక్స్ నుండి చిన్న దాని ద్వారా టోకెన్ల సీక్వెన్స్ రూపొందించబడితే, అది టార్గెట్ మోడల్తో ఎంతవరకు సరిపోతుందో నిర్ణయించడానికి పెద్ద దాని ద్వారా అమలు చేయబడుతుంది. అభ్యర్థుల ప్రతి ప్రిఫిక్స్ సంభావ్యతతో స్కోర్ చేయబడుతుంది మరియు అంగీకార థ్రెషోల్డ్ కంటే ఎక్కువ ఉన్న పొడవైనది ఎంపిక చేయబడుతుంది. బోనస్గా, టార్గెట్ మోడల్ తదుపరి టోకెన్ను ఉచితంగా అందిస్తుంది: డ్రాఫ్ట్ మోడల్ n టోకెన్లను రూపొందిస్తే, ఒకే దశలో n + 1 వరకు విడుదల చేయవచ్చు.
ఇన్ఫరెన్స్ సమయంలో, speculative sampling processను సుమారు 4 దశలుగా విభజించవచ్చు:
ప్రిఫిల్: KV కాషి ఎంట్రీలను పాపులేట్ చేయడానికి ఇన్పుట్ సీక్వెన్స్పై టార్గెట్ మరియు డ్రాఫ్ట్ మోడల్ రెండూ అమలు చేయబడాలి. మెడుసా వంటి కొన్ని పథకాలు అంచనా కోసం సరళమైన డెన్స్ లేయర్లను ఉపయోగిస్తుండగా, ఈ పోస్ట్లో మేము వాటి స్వంత KV కాశీలు అవసరమైన ట్రాన్స్ఫార్మర్-ఆధారిత డ్రాఫ్ట్లపై దృష్టి పెడతాము.
డ్రాఫ్ట్ జనరేషన్: నిర్దిష్ట సంఖ్యలో స్థిరమైన టోకెన్లను ఉత్పత్తి చేయడానికి డ్రాఫ్ట్ మోడల్ పునరావృతమవుతుంది. డ్రాఫ్ట్ సీక్వెన్స్ లీనియర్ కావచ్చు లేదా మోడల్ ఇచ్చిన లోతు వరకు చెట్టు లాంటి నిర్మాణాన్ని అన్వేషించవచ్చు (EAGLE, Medusa). ఇక్కడ, మేము లీనియర్ సీక్వెన్స్లపై దృష్టి పెడతాము.
అంగీకారం: డ్రాఫ్ట్ సీక్వెన్స్పై టార్గెట్ మోడల్ రన్ అవుతుంది, ప్రతి డ్రాఫ్ట్ టోకెన్కు తగిన లాజిట్లను నిర్మిస్తుంది. అత్యంత ఆమోదయోగ్యమైన సీక్వెన్స్ యొక్క పొడవు నిర్ణయించబడుతుంది.
టార్గెట్ జనరేషన్: టార్గెట్ లాజిట్లను రూపొందించినందున, సరిపోలని స్థానంలో లేదా సీక్వెన్స్ చివర లాజిట్లు తదుపరి టోకెన్కు అనుగుణంగా ఉంటాయి. సీక్వెన్స్ను ముగించడానికి టార్గెట్ నుండి బలమైన టోకెన్ను అందించడానికి ఈ లాజిట్లను నమూనా చేయవచ్చు.
ఊహాజనిత డికోడింగ్ను అమలు చేయడానికి వివిధ పద్ధతులు ఉన్నాయి. ఈ పోస్ట్లో, ఇన్-హౌస్ 1B మోడల్ని ఉపయోగించి Sonar మోడళ్లను వేగవంతం చేయడానికి మేము ఉపయోగించిన పథకాలపై, అలాగే DeepSeek స్థాయిలో మోడళ్లను వేగవంతం చేయడానికి మేము నిర్మిస్తున్న అంచనా విధానాలపై దృష్టి పెడతాము.
టార్గెట్-డ్రాఫ్ట్
అభ్యర్థి సీక్వెన్స్లను రూపొందించడానికి టార్గెట్ మోడల్కు డ్రాఫ్ట్ మోడల్గా ఇప్పటికే ఉన్న చిన్న LLMని జత చేయడం ద్వారా ఊహాజనిత డికోడింగ్ను సాధించవచ్చు. ఉత్పత్తిలో, టార్గెట్ వలె అదే డేటాసెట్పై ఫైన్-ట్యూన్ చేయబడిన Llama-1B మోడల్ని ఉపయోగించి మేము Sonarను వేగవంతం చేసాము. ఈ విధానానికి మొదటి నుండి డ్రాఫ్ట్కు శిక్షణ ఇవ్వడం అవసరం లేనప్పటికీ, చిన్న మోడల్ ఇప్పటికీ గణనీయమైన KV కాషి సామర్థ్యాన్ని ఉపయోగిస్తుంది మరియు స్వల్ప ప్రిఫిల్ ఓవర్హెడ్ను పరిచయం చేస్తుంది, TTFTని పెంచుతుంది.
ఈ పద్ధతిలో, డీకోడర్ కేవలం decode-only batchesపై మాత్రమే speculate చేస్తుంది; prefill సమయంలో లేదా mixed prefill-decode batchesపై standard sampling ద్వారా tokensను ఉత్పత్తి చేస్తుంది. prefill దశలో, target logits వెంటనే sample చేయబడి, draft యొక్క KV cacheలో కొత్తగా ఉత్పత్తి చేసిన tokenను కూడా prefill చేస్తాయి. draft ఇంకా sample చేయబడలేదు, కానీ అది ఉత్పత్తి చేసే logits decode stageకు carry over చేయబడతాయి.
డీకోడ్లో, డ్రాఫ్ట్ మోడల్ ముందుకు సాగుతుంది, ప్రతి దశలో టాప్ టోకెన్ను నమూనా చేస్తుంది. కావలసిన డ్రాఫ్ట్ పొడవు చేరుకున్న తర్వాత, ఆమోదించబడిన సీక్వెన్స్ పొడవును శాంప్లర్ గుర్తించే లాజిట్లను ఉత్పత్తి చేయడానికి టోకెన్లు టార్గెట్ మోడల్ ద్వారా రన్ చేయబడతాయి. డ్రాఫ్ట్ మరియు టార్గెట్ నుండి పూర్తి సంభావ్యత పంపిణీలను పోల్చడం ద్వారా అంగీకారం నిర్ణయించబడుతుంది. టార్గెట్ ఎల్లప్పుడూ ఆమోదించబడిన డ్రాఫ్ట్ సీక్వెన్స్ తర్వాత ఒక సెట్ లాజిట్లను అవుట్పుట్ చేస్తుంది కాబట్టి, అదనపు అవుట్పుట్ను ఉత్పత్తి చేయడానికి అది నమూనా చేయబడుతుంది. డ్రాఫ్ట్ మోడల్ ఇంకా ఆ ఆమోదించబడిన టోకెన్ను చూడలేదు కాబట్టి, తదుపరి డీకోడ్ దశ కోసం తయారీలో దాని సంబంధిత KV కాషి ఎంట్రీలను పూరించడానికి అది మళ్లీ రన్ చేయబడుతుంది, లాజిట్లను మళ్లీ ముందుకు తీసుకువెళుతుంది.
EAGLE
EAGLE అనేది ఒక speculative decoding scheme, ఇది probable draft tokens యొక్క tree-like traversal ద్వారా ఉత్పత్తి చేయబడిన బహుళ draft sequencesను అన్వేషిస్తుంది. స్థిర (EAGLE) లేదా డైనమిక్గా ఆకృతీకరించిన (EAGLE-2) treeను వరుస draft token executions ద్వారా అన్వేషిస్తారు; ఇందులో రేఖీయ క్రమంలో అత్యధిక స్కోరు ఉన్న tokenను అనుసరించకుండా, ప్రతి node వద్ద Top-K candidatesను పరిగణనలోకి తీసుకుంటారు. అనంతరం sequencesకు స్కోరు ఇచ్చి, కొనసాగించడానికి తగినదానిలో అతి పొడవైన sequenceను ఎంచుకుంటారు, అలాగే target నుండి ఒక అదనపు tokenను కూడా జోడిస్తారు.
더 సరిగ్గా అంచనా వేయడానికి, EAGLE డ్రాఫ్ట్ మోడల్ టోకెన్ల ఆధారంగా మాత్రమే కాకుండా, టార్గెట్ మోడల్ యొక్క టార్గెట్ ఫీచర్లు (చివరి లేయర్ హిడెన్ స్టేట్స్) ఉపయోగించి కూడా అంచనా వేస్తుంది. EAGLE యొక్క ప్రతికూలత ఏమిటంటే, తక్కువ లేటెన్సీ బడ్జెట్లో తగిన అభ్యర్థులను ఉత్పత్తి చేయడానికి తగినంత ఖచ్చితమైన కస్టమ్, చిన్న డ్రాఫ్ట్ మోడళ్లకు శిక్షణ ఇవ్వవలసి ఉంటుంది. సాధారణంగా, డ్రాఫ్ట్ మోడల్ అనేది అసలు మోడల్ యొక్క డీకోడర్ లేయర్తో సమానంగా ఉండే ఒకే ట్రాన్స్ఫార్మర్ లేయర్, ఇది దాని ఎంబెడ్డింగ్లు మరియు lm_head ప్రొజెక్షన్లకు అనుసంధానించడం ద్వారా టార్గెట్తో గట్టిగా జత చేయబడుతుంది. దీనికి తక్కువ KV కాషి సామర్థ్యం అవసరం కాబట్టి, EAGLE తక్కువ మెమరీ ఫుట్ప్రింట్ను కలిగి ఉంటుంది.
టార్గెట్ మోడల్లో ట్రీ లాంటి సీక్వెన్స్లను ధృవీకరించడానికి, కస్టమ్ అటెన్షన్ మాస్క్లను ఉపయోగించాలి. దురదృష్టవశాత్తు, మొత్తం సీక్వెన్స్ కోసం కస్టమ్ అటెన్షన్ మాస్క్ను ఉపయోగించడం వల్ల వాస్తవిక ఇన్పుట్ పొడవుల కోసం అటెన్షన్ గణనీయంగా నెమ్మదిస్తుంది (50% వరకు), ఊహాగానాల ద్వారా సాధించగల వేగవంతంలో కొంత భాగాన్ని శూన్యం చేస్తుంది. ఈ కారణంగా మేము పూర్తి ట్రీ అన్వేషణను ఇంకా ఉత్పత్తికి అమలు చేయలేదు, దానికి బదులుగా DeepSeek-V3 సాంకేతిక నివేదికలో అందించబడిన MTP-వంటి పథకాల ద్వారా సింగిల్-టోకెన్ ప్రిడిక్షన్ యొక్క ప్రత్యేక కేస్పై దృష్టి పెట్టాము.
MTP
ఈ పథకం డ్రాఫ్ట్-టార్గెట్ డికోడింగ్ను పోలి ఉంటుంది, అంచనా కోసం టోకెన్లతో పాటు హిడెన్ స్టేట్స్ ఉపయోగించబడుతుందనే మినహాయింపు మినహా. రెగ్యులర్ డ్రాఫ్ట్-టార్గెట్ ఊహాగానాలతో పోలిస్తే ప్రిఫిల్ మరియు డీకోడ్ దశలు రెండింటిలోనూ కొద్దిగా ఎక్కువ పని చేయాలి. డ్రాఫ్ట్ మోడల్ టోకెన్లు మరియు హిడెన్ స్టేట్స్ రెండింటినీ ఉపయోగిస్తుంది: టోకెన్ t_iకి అనుగుణంగా ఉండే లాజిట్లు L_i నుండి టోకెన్ t_{i+1} నమూనా చేయబడుతుంది, ఇవి హిడెన్ స్టేట్స్ H_i నుండి తీసుకోబడతాయి. తత్ఫలితంగా, టార్గెట్ అవుట్పుట్ చేసిన హిడెన్ స్టేట్ వెక్టర్లకు సంబంధించి ఇన్పుట్ టోకెన్ బఫర్లు ఒక దశ ఎడమవైపుకు మార్చబడాలి. దిగువ చిత్రం శిక్షణ కోసం ఉపయోగించిన కరస్పాండెన్స్లను, అలాగే ఇన్ఫరెన్స్ సమయంలో షిఫ్ట్ను సూచిస్తుంది.
డీకోడింగ్ ఫ్లో డ్రాఫ్ట్-టార్గెట్ డికోడింగ్ను పోలి ఉంటుంది, హిడెన్ స్టేట్స్ మరియు లాజిట్లు రెండూ ముందుకు తీసుకెళ్లబడతాయి అనే మినహాయింపు మినహా. మా అమలు సంబంధిత శాంప్లింగ్ మరియు లాజిట్ ప్రాసెసింగ్ లాజిట్లన్నింటినీ పంచుకుంటుంది, మోడల్ ఫార్వర్డ్ కాల్లను మాత్రమే ప్రత్యేకపరుస్తుంది. బహుళ టోకెన్లు అంచనా వేయబడినప్పుడు, డ్రాఫ్ట్ మోడల్ అంచనా కోసం డ్రాఫ్ట్ హిడెన్ స్టేట్లను ఉపయోగిస్తుంది, దాని స్వంత ఫీచర్ల ఆధారంగా KV కాషి ఎంట్రీలను కూడా పాపులేట్ చేస్తుంది. దీర్ఘకాలంలో, ఇది ఖచ్చితత్త్వాన్ని తగ్గిస్తుంది. తదనంతరం, టార్గెట్ ప్రిడిక్షన్ కోసం KV కాషి ఎంట్రీని పాపులేట్ చేయడానికి డ్రాఫ్ట్ మోడల్ను రన్ చేస్తున్నప్పుడు, మేము మరింత ఖచ్చితమైన టార్గెట్ హిడెన్ స్టేట్లను ఇన్పుట్లుగా తీసుకుంటూ మొత్తం సీక్వెన్స్పై దానిని రన్ చేస్తాము. ఈ డ్రాఫ్ట్ మోడల్స్ చిన్నవి కాబట్టి, అదనపు టోకెన్లను ప్రాసెస్ చేసే అదనపు ఖర్చు చాలా తక్కువ.
MTP హెడ్స్కు శిక్షణ ఇవ్వడం
MTP నుండి ప్రయోజనం పొందేందుకు, 8xH100 పరికరాలతో కూడిన ఒక నోడ్లో రన్ అవుతున్న Perplexity యొక్క డేటాసెట్లపై మా ఫైన్-ట్యూన్ చేయబడిన మోడళ్లకు జోడించబడిన MTP హెడ్స్కు శిక్షణ ఇవ్వడానికి అవసరమైన మౌలిక సదుపాయాలను మేము నిర్మించాము. సుమారు ఒక రోజులో, మేము Llama-1B నుండి Llama-70B మరియు DeepSeek V2-Lite వరకు ఉన్న మోడళ్ల కోసం హెడ్స్ను నిర్మించగలము. పెద్ద మోడళ్ల కోసం, ఫైన్-ట్యూనింగ్ ప్రక్రియలో నిర్మించబడిన MTP హెడ్స్పై మేము ఆధారపడతాము.
MTP శిక్షణ యొక్క లక్ష్యం డ్రాఫ్ట్ హిడెన్ స్టేట్స్ మరియు టార్గెట్ హిడెన్ స్టేట్ల నుండి ఎక్స్ట్రాపోలేట్ చేయబడిన లాజిట్లను టార్గెట్ యొక్క తదుపరి టోకెన్ లాజిట్లు మరియు హిడెన్ స్టేట్లతో సరిపోల్చడం. హిడెన్ స్టేట్ల కోసం ఇన్ఫరెన్స్ ఖరీదైనది కాబట్టి, శిక్షణ సమయంలో ఉపయోగించడానికి మేము మా ఇన్ఫరెన్స్-ఆప్టిమైజ్ చేయబడిన టార్గెట్ మోడల్ అమలును ఉపయోగించి వాటిని ముందే గణిస్తాము. ఏదేమైనప్పటికీ, ఇన్ఫరెన్స్ MTP అమలును ధృవీకరించడానికి మరియు క్వాంటైజేషన్ లేదా ఆప్టిమైజేషన్ల కారణంగా సంఖ్యాపరమైన వ్యత్యాసాలు ఫలితాలకు ఆటంకం కలిగించవని నిర్ధారించడానికి, ధృవీకరణ లాస్ మరియు ఖచ్చితత్వ అంచనా కోసం మేము టార్గెట్ మరియు డ్రాఫ్ట్ మోడల్స్ రెండింటి ఇన్ఫరెన్స్ అమలును పూర్తిగా మళ్లీ ఉపయోగిస్తాము.
అసలు పేపర్లో ఉపయోగించిన ShareGPT డేటాసెట్ నుండి పెద్ద నమూనాలకు స్కేలింగ్ చేస్తున్నప్పుడు, EAGLE పేపర్లో రూపుమాపబడిన మరియు అమలు చేయబడిన MTP హెడ్ ఆర్కిటెక్చర్ 70B-పరిమాణ మోడళ్లకు శిక్షణ ఇవ్వడంలో విఫలమైందని మేము గమనించాము. పెద్ద సంఖ్యలో చిన్న సీక్వెన్స్లను కలిగి ఉన్న ShareGPT వలె కాకుండా, మేము కొద్దిగా తక్కువ సంఖ్యలో గణనీయంగా పొడవైన ప్రాంప్ట్లపై శిక్షణ ఇస్తాము. అసలు EAGLE హెడ్స్ సాధారణ ట్రాన్స్ఫార్మర్ నుండి నిర్మాణంలో కొద్దిగా వైదొలిగాయి కాబట్టి, తీసివేయబడిన కొన్ని RMS నార్మలైజేషన్ లేయర్లను మేము తిరిగి పరిచయం చేసాము. ఇది శిక్షణ కన్వర్జ్ అవ్వడానికి అనుమతించడమే కాకుండా, హెడ్స్ యొక్క ఖచ్చితత్త్వాన్ని కొన్ని శాతాల పాయింట్ల ద్వారా పెంచినట్లు మేము కనుగొన్నాము.
layer norms శిక్షణను సులభతరం చేయడమే కాదు, normsను తిరిగి ప్రవేశపెట్టడం గణితపరంగా కూడా సహజంగా ఉంటుంది. MTP heads, target model యొక్క embeddings మరియు logit projectionsను మళ్లీ ఉపయోగిస్తాయి, ఎందుకంటే అవి పరిమాణంలో గణనీయంగా ఉండవచ్చు (Llama 70B కోసం సుమారు 2 GB). శిక్షణ సమయంలో ఇవి frozenగా ఉంటాయి, మరియు MTP layer, అసలు model యొక్క projection layer శిక్షణ సమయంలో నేర్చుకున్న దానితో అదే vector spaceలో predictionsను embed చేయడం నేర్చుకుంటుందని ఆశిస్తారు. normsను తీసివేయడం వల్ల, single MLP, MLP తరువాత ఒక norm ఉన్నదానితో సమానమైన functionను నేర్చుకోవాల్సి వస్తుంది, ఇది draft మరియు target models యొక్క hidden states మధ్య matchupను అడ్డుకుంటుంది.
ఊహాజనిత డికోడింగ్తో ఇన్ఫరెన్స్
ఇన్ఫరెన్స్ ఇంజిన్లో, ఇన్పుట్ సీక్వెన్స్ల కోసం టోకెన్లను ఉత్పత్తి చేయడానికి, అవి మొదట సహేతుకమైన పరిమాణ గణాలుగా (batches) సమూహపరచబడాలి, ఆపై తదుపరి టోకెన్ల కోసం KV కాషిలో పేజీలు కేటాయించబడాలి. ఇన్పుట్ టోకెన్లు మరియు KV పేజీ సమాచారం ఆపై మోడల్ను రన్ చేస్తున్న అన్ని సమాంతర రాంక్లకు బ్రాడ్కాస్ట్ చేయబడిన బఫర్లోకి ప్యాక్ చేయబడతాయి. చివరిగా, మెటాడేటా GPU మెమరీకి కాపీ చేయబడుతుంది మరియు తదుపరి టోకెన్ నమూనా చేయబడే లాజిట్లను ఉత్పత్తి చేయడానికి మోడల్ అమలు చేయబడుతుంది.
కోణికంగా ఒక draft మరియు target inference serverను వాటి మధ్య requestsను orchestrate చేసే wrapper ద్వారా సడలుగా couple చేసే కొన్ని implementationsకు భిన్నంగా, మా draft-target pairs tightly coupledగా ఉండి generationను unisonలో step through చేస్తాయి. speculative decoding యొక్క అన్ని రూపాల కోసం batch scheduling మరియు KV page allocation models మధ్య పంచుకుంటారు: దీని ద్వారా modelను overarching inference serverతో కలిపే logic ఏకీకృతమవుతుంది, ఎందుకంటే అవన్నీ ఒకే interfaceను expose చేస్తాయి.
Perplexity వద్ద ఇన్ఫరెన్స్ రన్టైమ్ FlashInfer చుట్టూ రూపొందించబడింది, ఇది అటెన్షన్ కెర్నల్ను కాన్ఫిగర్ చేయడానికి మరియు షెడ్యూల్ చేయడానికి నిర్మించాల్సిన మెటాడేటాను నిర్ణయిస్తుంది. బ్యాచ్ను ఏర్పరిచే కొన్ని ఇన్పుట్ సీక్వెన్స్లు ఇవ్వబడినప్పుడు, ప్రిఫిల్, డీకోడ్ లేదా ధృవీకరణ కోసం, ఇంటర్మీడియట్ బఫర్లను కేటాయించడానికి మరియు అటెన్షన్లో ఉపయోగించే కొన్ని స్థిర బఫర్లను పాపులేట్ చేయడానికి CPU-వైపు పని చేయాలి. ఈ పని బ్యాచ్ షెడ్యూలింగ్ మరియు KV పేజీ కేటాయింపు యొక్క ఖర్చుతో కూడుకున్నది, ఇవి GPU వినియోగాన్ని పెంచడానికి దాచవలసిన లేటెన్సీలను కూడా కలిగి ఉంటాయి.
ఊహాగానాలు లేకుండా ఇన్ఫరెన్స్ కోసం మేము CPU-వైపు మరియు GPU-వైపు పనిని పూర్తిగా సమాంతరీకరించినప్పటికీ, ఊహాజనిత డికోడింగ్ కోసం CPU-GPU బ్యాలెన్స్ మరింత క్లిష్టమైనదని మేము కనుగొన్నాము. ఆమోదించబడిన టోకెన్ల సంఖ్య తదుపరి రన్ కోసం సీక్వెన్స్ పొడవును నిర్ణయిస్తుంది అనే వాస్తవం నుండి ప్రధాన సవాలు ఉద్భవించింది, ఇది నివారించడానికి కష్టమైన GPU-to-CPU సింక్రొనైజేషన్ పాయింట్ను పరిచయం చేస్తుంది. CPU పని యొక్క లేటెన్సీని ఉత్తమంగా దాచడానికి మేము వివిధ షెడ్యూలింగ్ పథకాలతో ప్రయోగం చేసాము.
డ్రాఫ్ట్-టార్గెట్ షెడ్యూల్
టార్గెట్ మోడల్ కంటే చిన్నది అయినప్పటికీ, మొత్తం LLM డ్రాఫ్ట్గా ఉపయోగించినప్పుడు, ఇది ఇప్పటికీ GPU పై గణనీయమైన లేటెన్సీని పరిచయం చేస్తుంది, ఖరీదైన CPU ఆపరేషన్లను దాచడానికి కొంత హెడ్రూమ్ను అందిస్తుంది. చిన్న మోడళ్లు టెన్సార్ సమాంతరత నుండి ప్రయోజనం పొందవు కాబట్టి, టార్గెట్ మరియు డ్రాఫ్ట్ షార్డ్ చేయబడిన రాంక్ల సంఖ్య మధ్య సరిపోలిక లేదు. మా అమలులో, డ్రాఫ్ట్ మోడల్ TP సమూహం యొక్క లీడర్ రాంక్ పై మాత్రమే రన్ అవుతుంది.
గమనించినట్లుగా, డీకోడ్ దశ తదుపరి రన్లోకి లాజిట్లను ముందుకు తీసుకువెళుతుంది. డ్రాఫ్ట్ మోడల్ యొక్క ఒక అమలును CPU-వైపు బ్యాచ్ షెడ్యూలింగ్ పనితో అతివ్యాప్తి చేయడానికి ఇది మనలను అనుమతిస్తుంది. బ్యాచ్ను అమర్చిన తర్వాత, శాంప్లర్ మరియు డ్రాఫ్ట్కు పునరావృత కాల్లు డ్రాఫ్ట్ టోకెన్లను ఉత్పత్తి చేస్తాయి. సమాంతరంగా, ధృవీకరణ కోసం బ్యాచ్ టార్గెట్ మోడల్ కోసం అమర్చబడుతుంది మరియు సమాంతర వర్కర్లతో సింక్రొనైజ్ చేయబడుతుంది. ఆమోదించబడిన సీక్వెన్స్ పొడవులను నిర్ణయించడానికి టార్గెట్ లాజిట్లు ధృవీకరించబడతాయి మరియు నమూనా చేయబడతాయి. ఈ సమయంలో, తదుపరి సీక్వెన్స్ పొడవులను నిర్ణయించడానికి GPU-to-CPU సింక్రొనైజేషన్ అవసరం. డ్రాఫ్ట్ మోడల్ లీడర్ నోడ్ పై మాత్రమే రన్ చేయబడుతుంది కాబట్టి, దాని బ్యాచ్ క్రమంగా సెటప్ చేయబడుతుంది మరియు టార్గెట్ ఉత్పత్తి చేసిన అదనపు టోకెన్తో దాని KV కాషి ఎంట్రీలను పాపులేట్ చేయడానికి దాని అమలు ప్రారంభించబడుతుంది. ప్రస్తుత రన్లో ఈ డ్రాఫ్ట్ రన్ ద్వారా ఉత్పత్తి చేయబడిన లాజిట్లు తదుపరి రన్లో మొదటి డ్రాఫ్ట్ టోకెన్ను నమూనా చేయడానికి ఉపయోగించబడతాయి. అత్యంత ముఖ్యమైనది ఏమిటంటే, డ్రాఫ్ట్ నడుస్తున్నప్పుడు, తదుపరి బ్యాచ్ షెడ్యూల్ చేయబడవచ్చు.
ఒక టోకెన్ కోసం MTP షెడ్యూల్
రన్టైమ్ ఇంకా ఈగల్-శైలి డ్రాఫ్ట్ ట్రీ అన్వేషణను అందించనప్పటికీ, మేము ఈ పథకం యొక్క ప్రత్యేక కేస్ను అమలు చేసాము, ఒకే ట్రాన్స్ఫార్మర్ డీకోడర్ లేయర్ పరిమాణం గల మోడల్ ద్వారా ఉత్పత్తి చేయబడిన డ్రాఫ్ట్ టోకెన్ల లీనియర్ సీక్వెన్స్ను పరిగణనలోకి తీసుకున్నాము. DeepSeek R1 యొక్క ఓపెన్ సోర్స్ వెయిట్లను ఉపయోగించి డ్రాఫ్ట్ ప్రిడిక్షన్ కోసం ఈ పథకాన్ని ఉపయోగించవచ్చు. ఒకే టోకెన్ను అంచనా వేసే ఉప-కేస్ ఆసక్తికరంగా ఉంటుంది, ఎందుకంటే పెద్ద MTP లేయర్లు వాటి ఓవర్హెడ్ను సమర్థించడానికి తగినంత అధిక అంగీకార రేట్లను సాధిస్తాయి.
MTP షెడ్యూలింగ్ కొంతవరకు మరింత సంక్లిష్టంగా ఉంటుంది, ఎందుకంటే డ్రాఫ్ట్ మోడల్ చాలా వేగంగా ఉంటుంది, తక్కువ CPU-వైపు లేటెన్సీని దాస్తుంది. అదనంగా, డ్రాఫ్ట్ టార్గెట్ మోడల్తో పాటు షార్డ్ చేయబడుతుంది, బ్యాచ్ సమాచారం కోసం షేర్డ్ మెమరీ బదిలీలు అవసరం. మునుపటి పథకం వలె కాకుండా, క్యారీ-ఓవర్ లాజిట్ల నుండి బ్యాచ్ సమాచారాన్ని బదిలీ చేయడం మరియు మొదటి టోకెన్ను నమూనా చేయడం ద్వారా రన్ ప్రారంభమవుతుంది. తదుపరి, టోకెన్లను ధృవీకరించడానికి టార్గెట్ రన్ చేయబడుతుంది, 2 * D టోకెన్లను ప్రాసెస్ చేస్తుంది, ఇక్కడ D అనేది డీకోడ్ బ్యాచ్ పరిమాణం. బ్యాచ్ సమానంగా రెండు భాగాలుగా విడిపోతుంది కాబట్టి, InfiniBand వంటి నెమ్మదైన ఇంటర్కనెక్ట్లపై మిక్చర్-ఆఫ్-ఎక్స్పర్ట్స్ (MoE) మోడళ్లలో మైక్రో-బ్యాచింగ్కు ఇది ఆదర్శంగా సరిపోతుంది. టార్గెట్ యొక్క హిడెన్ స్టేట్స్ తదుపరి డ్రాఫ్ట్ రన్కు తీసుకువెళ్లబడతాయి, అయితే లాజిట్లు ధృవీకరణ కోసం శాంప్లర్లోకి పాస్ చేయబడతాయి.
GPU పై పరిమిత మొత్తంలో అదనపు పనిని చేయడం ద్వారా, డ్రాఫ్ట్ సీక్వెన్స్ అంగీకారం తర్వాత CPU-to-CPU సింక్రొనైజేషన్ను మేము నివారిస్తాము. టార్గెట్ల ఇన్పుట్ టోకెన్లు మార్చబడిన తర్వాత, కెర్నల్ తదుపరి టార్గెట్ టోకెన్లను వాటి సంబంధిత స్థానాల్లో ప్లగ్ చేస్తుంది. టార్గెట్ వలె అదే బ్యాచ్ సమాచారంతో డ్రాఫ్ట్ మళ్లీ రన్ చేయబడుతుంది, KV కాషి ఎంట్రీలను పాపులేట్ చేస్తుంది మరియు తదుపరి రన్ కోసం లాజిట్లు మరియు హిడెన్ స్టేట్లను నిర్మిస్తుంది, ఆమోదించబించని టోకెన్లపై కొంత అనవసరమైన పని చేస్తుంది. ఈ పరిస్థితులలో, డ్రాఫ్ట్ మోడల్ యొక్క చిన్న పరిమాణం కారణంగా ఉపయోగించని పని యొక్క లేటెన్సీ అతిగా కొలవదగనిది. డ్రాఫ్ట్ రన్తో సమాంతరంగా, CPU పై సీక్వెన్స్ పొడవులు నిర్ణయించబడతాయి మరియు GPU పని ముగిసే వరకు వేచి ఉండకుండా, తదుపరి బ్యాచ్ యొక్క షెడ్యూలింగ్ ప్రారంభించబడుతుంది.
డ్రాఫ్ట్ లేయర్లో అదనపు పని యొక్క ఓవర్హెడ్ అటెన్షన్లో గమనించదగ్గది కాదు, అయినప్పటికీ MLP లేయర్లు మరింత సమస్యాత్మకమైనవి. టోకెన్ల సంఖ్య యొక్క పరిమాణం వెంబడి మ్యాట్రిక్స్ గుణకార సూచనలు 64 సరిహద్దుకు ప్యాడ్ చేయబడతాయి కాబట్టి, రెట్టింపు చేయడం గణనీయంగా ఎక్కువ బ్లాక్లు అవసరం లేకపోతే, ఓవర్హెడ్ దాచబడుతుంది. పొడవైన డ్రాఫ్ట్ సీక్వెన్స్ల కోసం ఓవర్హెడ్ మరింత ఖరీదైనది మరియు రెగ్యులర్ డ్రాఫ్ట్-టార్గెట్ మోడళ్ల కోసం ఉపయోగించే పథకం మెరుగ్గా పనిచేస్తుంది.
ఉపయోగించిన సూచనలు (References)
ఊహాజనిత డికోడింగ్ ద్వారా ట్రాన్స్ఫార్మర్ల నుండి వేగవంతమైన ఇన్ఫరెన్స్
EAGLE: ఊహాజనిత నమూనా ఫీచర్ అనూహ్యతను పునరాలోచించడం అవసరం
EAGLE-2: డైనమిక్ డ్రాఫ్ట్ ట్రీలతో భాషా మోడళ్ల వేగవంతమైన ఇన్ఫరెన్స్
Medusa: బహుళ డికోడింగ్ హెడ్స్తో సాధారణ LLM ఇన్ఫరెన్స్ త్వరణ ఫ్రేమ్వర్క్
FlashInfer: LLM ఇన్ఫరెన్స్ సర్వింగ్ కోసం సమర్థవంతమైన మరియు అనుకూలీకరించదగిన అటెన్షన్ ఇంజిన్
ఉదాహరణ
@online{perplexity2025AcceleratingSonarThroughSpeculation,
author = {{Perplexity} and Nandor Licker and Erin McKnight},
title = {{Accelerating Sonar through speculation}},
date = {2025-06-10},
url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}