AWS EFAలో ట్రిలియన్-పారామీటర్ మోడళ్లను ప్రారంభించడం

Perplexityలో, మేము మా ఉత్పత్తి, మా APIలు మరియు మా పరిశోధనా బృందాల కోసం అత్యుత్తమ మోడళ్లను ఉపయోగిస్తాము. Kimi-K2 వంటి పెద్ద ఓపెన్-సోర్స్ మిక్చర్-ఆఫ్-ఎక్స్‌పర్ట్స్ (MoE) మోడళ్లు ప్రత్యేక సవాళ్లను ఎదుర్కొంటాయి, ఎందుకంటే 8x NVIDIA H200 GPUలతో కూడిన అతిపెద్ద ఇన్‌ఫరెన్స్ నోడ్‌లు వాటిని సమర్థవంతంగా ఉంచలేవు, దీనివల్ల మల్టీ-నోడ్ డిప్లాయ్‌మెంట్‌లు అవసరమవుతాయి. మేము ConnectX-7పై స్టేట్-ఆఫ్-ద-ఆర్ట్ లాటెన్సీలను సాధించే ఎక్స్‌పర్ట్ పారలలిజం కోసం కర్నెల్‌ల సమితిని అందిస్తున్నాము, ఇది DeepEP పనితీరును మించిపోయింది. AWS ఎలాస్టిక్ ఫాబ్రిక్ అడాప్టర్ (EFA)పై సాధ్యమయ్యే లాటెన్సీలను సాధించిన మొదటి కర్నెల్స్ ఇవే, ఇవి ట్రిలియన్-పారామీటర్ మోడల్ డిప్లాయ్‌మెంట్‌లను ప్రారంభిస్తాయి.

GitHubలో మా కర్నెల్‌లను ప్రయత్నించండి మరియు arXivలో పూర్తి పరిశోధనా పత్రాన్ని చదవండి.

పరిచయం

మిక్చర్-ఆఫ్-ఎక్స్‌పర్ట్స్ (MoE) వందల బిలియన్లు లేదా ట్రిలియన్ల పారామీటర్‌లకు స్కేల్ చేయడానికి ప్రామాణిక ఆర్కిటెక్చర్‌గా ఉద్భవించింది, అదే సమయంలో ఇన్‌ఫరెన్స్ సమయంలో సహేతుకమైన లాటెన్సీని అందిస్తుంది. MoE మోడల్స్ ట్రాన్స్‌ఫార్మర్ యొక్క డెన్స్ లేయర్‌ను ఎక్స్‌పర్ట్స్ సమితి మరియు టోకెన్‌ను దేనితో గుణించాలో నిర్ణయించే చిన్న లీనియర్ రూటింగ్ లేయర్‌తో భర్తీ చేస్తాయి. ఒకే ఎక్స్‌పర్ట్‌తో ముడిపడి ఉన్న టోకెన్‌లను గ్రూప్ చేయడం ద్వారా, సమస్య స్పార్స్‌గా మారుతుంది మరియు సులభంగా సమాంతరంగా మార్చబడుతుంది, ఎందుకంటే ఇప్పుడు వివిధ ఎక్స్‌పర్ట్‌లను బహుళ నోడ్‌లలో విస్తరించిన వివిధ GPUలలో హోస్ట్ చేయవచ్చు.

torch.distributed మరియు NCCL వంటి కలెక్టివ్ కమ్యూనికేషన్ లైబ్రరీలను ఉపయోగించి సులభంగా అమలు చేయగల టెన్సార్ పారలలిజం (TP) మరియు డేటా పారలలిజం (DP) వంటి ఇతర రకాల పారలలిజం వలె కాకుండా, MoE రూటింగ్ స్పార్స్ పీర్-టు-పీర్ కమ్యూనికేషన్‌ను కలిగి ఉంటుంది. డిస్పాచ్ కర్నెల్ ఇన్‌కమింగ్ టోకెన్‌ల సమితిని విభజించి వాటిని ఎక్స్‌పర్ట్‌లను హోస్ట్ చేసే ర్యాంక్‌లకు పంపాలి, అయితే దాని ద్వంద్వ సమస్య అయిన కంబైన్ కర్నెల్ వాటి ప్రాసెస్ చేయబడిన ప్రతిరూపాలను స్వీకరించి వాటి వెయిటెడ్ యావరేజ్‌ను లెక్కించాలి. ఈ బదిలీలు కలెక్టివ్‌లకు మ్యాప్ చేయబడవు మరియు తక్కువ లాటెన్సీని సాధించడానికి అనుకూలీకరించిన కర్నెల్స్ అవసరమవుతాయి.

Perplexityలో, మేము పోర్టబుల్ ఇంటర్-నోడ్ కర్నెల్స్ మరియు ప్రత్యేకమైన ఇంట్రా-నోడ్ కర్నెల్స్ రెండింటినీ అభివృద్ధి చేస్తున్నాము, మరియు మా పనిని మునుపటి MoE కమ్యూనికేషన్ వ్యాసంలో మరింత వివరంగా ప్రదర్శించాము. పెద్ద MoE మోడళ్ల కోసం మల్టీ-నోడ్ డిప్లాయ్‌మెంట్ తక్కువ లేటెన్సీ మరియు అధిక త్రుపుట్‌ను తీసుకురాగలదని మేము మరో బ్లాగ్ పోస్టులో కూడా చూపాము.

ఈ కథనంలో, NVSHMEM-ఆధారిత మునుపటి మా పని పనితీరును బాగా మెరుగుపరిచే మా కొత్త ఇంటర్-నోడ్ కర్నెల్‌ల అమలులోకి మేము లోతుగా వెళ్తాము. ConnectX-7పై స్టేట్-ఆఫ్-ద-ఆర్ట్ ఫలితాలను ప్రదర్శించడంతో పాటు, మా ఇన్‌ఫరెన్స్ రన్‌టైమ్ లోపల, అతిపెద్ద ఓపెన్-సోర్స్ మోడళ్లకు సమర్థవంతంగా సేవ చేయడానికి కర్నెల్స్ ఉపయోగించవచ్చని మేము ప్రదర్శిస్తాము.

ఇన్ఫినిబ్యాండ్‌పై ఇన్‌ఫరెన్స్

NVLink మైక్రోసెకన్ల లాటెన్సీ వద్ద అధిక త్రూపుట్‌ను (900GB/s) అందిస్తుంది కాబట్టి, టెన్సార్ పారలలిజం మరియు ఎక్స్‌పర్ట్ పారలలిజం ఒకే నోడ్ లోపల సమర్థవంతంగా స్కేల్ చేయగలవు. NVIDIA GB200 మరియు GB300 సిస్టమ్‌లు ట్రిలియన్-పారామీటర్ మోడల్స్ యొక్క వెయిట్‌లు మరియు KV క్యాచెలను ఉంచడానికి తగినంత మెమరీని కలిగి ఉన్నప్పటికీ, గరిష్టంగా 8 H200 GPUలతో కూడిన AWS p5en ఇన్‌స్టాన్స్‌పై డిప్లాయ్‌మెంట్‌లు 1120GB HBMకి పరిమితం చేయబడ్డాయి, ఇది మోడల్ వెయిట్‌లు మరియు KV క్యాచెలు రెండింటి మధ్య విభజించబడాలి. ఫలితంగా, కొన్ని మోడళ్లకు ఇన్ఫినిబ్యాండ్ ద్వారా ఇంటర్-నోడ్ డిప్లాయ్‌మెంట్‌లు అవసరమవుతాయి.

ఇన్ఫినిబ్యాండ్ ఇప్పటికీ చాలా వేగంగా ఉన్నప్పటికీ, ఇది కేవలం 400Gbps (50GB/s) త్రూపుట్‌ను మాత్రమే అందిస్తుంది, MoE రూటింగ్ కోసం పది నుండి వందల మైక్రోసెకన్ల లాటెన్సీని జోడిస్తుంది. అదృష్టవశాత్తూ, పోటీ ధర వద్ద ఆచరణీయమైన త్రూపుట్‌ను అందించడానికి షేర్డ్ ఎక్స్‌పర్ట్స్, మైక్రో-బ్యాచింగ్ మరియు కంప్యూటేషన్-కమ్యూనికేషన్ ఓవర్‌లాపింగ్ వినియోగం ద్వారా జోడించబడిన లాటెన్సీలను దాచవచ్చు. DeepSeek-R1 ద్వారా డీప్‌సీక్ ఈ టెక్నిక్‌లకు మార్గదర్శకత్వం వహించింది, ఇవి H800 GPUలు మరియు NVIDIA ConnectX-7 ఇన్ఫినిబ్యాండ్ అడాప్టర్‌ల క్లస్టర్‌లపై డిప్లాయ్ చేయబడ్డాయి.

AWS క్లస్టర్లలో, వారి కస్టమ్ Elastic Fabric Adapters (EFA) వినియోగం కారణంగా పరిస్థితి భిన్నంగా ఉంటుంది. ఇవి ట్రైనింగ్ వర్క్‌లోడ్స్‌లో సాధారణంగా ఉపయోగించే కలెక్టివ్‌లపై పీక్ 400Gbps త్రుపుట్‌ను అందించినప్పటికీ, MoE డిస్పాచ్ మరియు కంబైన్ సమయంలో మార్పిడి అయ్యే సందేశ పరిమాణాల్లో ఇవి ConnectX-7 కంటే కొద్దిగా వెనుకబడతాయి. అదనంగా, EFA GPUDirect Async‌ను మద్దతు ఇవ్వదు; అందువల్ల ట్రాన్స్‌ఫర్‌లను ప్రారంభించడానికి GPU నుంచి NIC‌కు బ్రిడ్జ్‌గా పనిచేసే CPU ప్రాక్సీ థ్రెడ్ అవసరం అవుతుంది. అదనపు సంక్లిష్టతతో పాటు, అదనపు PCIe ట్రాన్సాక్షన్‌లు అన్ని ట్రాన్స్‌ఫర్‌లకు మైక్రోసెకన్ల స్థాయి ఓవర్‌హెడ్‌ను జోడిస్తాయి. NVSHMEM వంటి ఫ్రేమ్‌వర్క్‌లు కర్నెల్ రచయితలకు డివైస్-స్వతంత్ర అమలులను అందించడానికి అనుమతిస్తాయి, అయితే సాధారణ APIs వల్ల వచ్చే ఓవర్‌హెడ్‌లు ఖర్చు మరియు లేటెన్సీ రెండింటి పరంగా డిప్లాయ్‌మెంట్‌లను ప్రభావవంతం కాకుండా చేస్తాయి.

NVSHMEM-ఆధారిత pplx-కర్నెల్స్

మా మునుపటి కర్నెల్స్ NVSHMEM ద్వారా అంతర్లీన రవాణా అమలును నూతన పద్ధతిలో అందించాయి. సహచరులను సింక్రొనైజ్ చేయడానికి అటామిక్ కౌంటర్‌లను ఉపయోగించి, టోకెన్‌లను ఒక్కొక్కటిగా బదిలీ చేయడానికి అవి nvshmemx_putmem_signal_nbi_warpపై ఆధారపడ్డాయి. IBGDAతో కూడిన ConnectX-7 NICలపై ఇది సహేతుకమైన పనితీరును అందించినప్పటికీ, ప్రాక్సీ-ఆధారిత ConnectX-7 IBRC అమలు గణనీయంగా నెమ్మదిగా ఉంది, అయితే EFA అమలు ఒక మిలిసెకండ్ కంటే తక్కువ లాటెన్సీలను అందించడంలో విఫలమైంది.

మా మోడళ్లకు బహిర్గతం చేయడానికి అత్యంత అనుకూలమైన ఇంటర్‌ఫేస్‌ను గుర్తించడంతో పాటు, మేము అంతర్లీన హార్డ్‌వేర్ కోసం కర్నెల్స్‌ను టెస్ట్‌బెడ్‌గా ఉపయోగించాము. కర్నెల్స్ మరియు NICలు లేదా ప్రాక్సీ థ్రెడ్‌ల మధ్య తరచుగా జరిగే బదిలీలు మరియు పరస్పర చర్యలు పనితీరుకు ఆటంకం కలిగించాయని మేము ఊహించాము. అదనంగా, మా అమలు విశ్వసనీయత కాకుండా హార్డ్‌వేర్ నుండి ఎలాంటి హామీలపై ఆధారపడలేదు, సింక్రొనైజేషన్ ప్రిమిటివ్‌ల ఓవర్‌హెడ్‌ను తగ్గించడం ద్వారా మేము విషయాలను వేగవంతం చేయగలమని ఇది సూచిస్తుంది.

ఈ పరిశీలనలను అనుసరించి, గరిష్ట పనితీరు కోసం హుడ్ కింద డివైస్ కర్నెల్ - హోస్ట్ ప్రాక్సీ పరస్పర చర్యలను ఆప్టిమైజ్ చేస్తూ, అదే ఇంటర్‌ఫేస్‌ను బహిర్గతం చేసే మా EFA-ఆధారిత కర్నెల్స్‌ను మేము రూపొందించాము.

EFA ద్వారా MoE ఆల్-టు-ఆల్

విస్తరించు

కొత్త కర్నెల్స్ హైబ్రిడ్ CPU-GPU ఆర్కిటెక్చర్‌ను అనుసరిస్తాయి, ఇక్కడ GPU కర్నెల్స్ డివైస్‌లోని మోడల్‌తో పరస్పర చర్య చేస్తాయి మరియు హోస్ట్ CPUపై నడుస్తున్న ప్రాక్సీ థ్రెడ్ NICతో పరస్పర చర్యలను నిర్వహిస్తుంది. కొత్త కర్నెల్స్ మేము KV క్యాచే బదిలీలు కోసం నిర్మించిన అదే TransferEngineని తిరిగి ఉపయోగిస్తాయి.

MoE రూటింగ్ డిస్పాచ్ కర్నెల్ ద్వారా నిర్వహించబడుతుంది, ఇది మోడల్ నుండి టోకెన్‌లను మరియు వాటిని రూట్ చేయడానికి పర్-టోకెన్ ఎక్స్‌పర్ట్‌లను స్వీకరిస్తుంది, వాటిని ఎక్స్‌పర్ట్‌లను హోస్ట్ చేసే ర్యాంక్‌లకు నెట్‌వర్క్ ద్వారా పంపుతుంది. గ్రూప్ చేయబడిన GEMM కర్నెల్స్ ఎక్స్‌పర్ట్ వెయిట్‌లతో టోకెన్‌లను గుణించిన తర్వాత, కంబైన్ కర్నెల్ టోకెన్‌లను వాటి సోర్స్ ర్యాంక్‌లకు బదిలీ చేస్తుంది మరియు మోడల్‌కు తిరిగి ఇచ్చే ముందు వాటి వెయిటెడ్ యావరేజ్‌ను లెక్కిస్తుంది.

ఓవర్‌లాపింగ్ మరియు మైక్రో-బ్యాచింగ్‌ను సులభతరం చేయడానికి డిస్పాచ్ మరియు కంబైన్ కర్నెల్స్ రెండూ సెండర్ మరియు రిసీవర్ హాఫ్‌గా విభజించబడ్డాయి. సెండర్ టోకెన్‌లను సెండ్ బఫర్‌లలోకి షఫుల్ చేస్తుంది, తద్వారా వాటిని ఒకే RDMA రైట్ ఆపరేషన్ ద్వారా సహచరులకు వ్రాయవచ్చు. సెండర్‌ను అమలు చేసిన తర్వాత, నెట్‌వర్క్ ద్వారా బదిలీలు బ్యాక్‌గ్రౌండ్‌లో పెండింగ్‌లో ఉన్నప్పుడు GPU ఇతర పనిని (షేర్డ్ ఎక్స్‌పర్ట్ గుణకారం వంటివి) చేయవచ్చు. అన్ని బదిలీలు స్వీకరించబడే వరకు రిసీవర్ కర్నెల్స్ బ్లాక్ అవుతాయి. బ్యాక్‌గ్రౌండ్‌లో, హోస్ట్-సైడ్ వర్కర్ థ్రెడ్ బదిలీలను నిర్వహిస్తుంది: సెండ్ బఫర్‌లు జనాభా నింపబడినప్పుడు GPU దీనికి సిగ్నల్ ఇస్తుంది, ఆ తర్వాత అది RDMA రైట్‌లను డిస్పాచ్ చేస్తుంది. NIC పూర్తి అయినట్లు సిగ్నల్ ఇచ్చిన తర్వాత, ఇది రిసీవర్‌ను అన్‌బ్లాక్ చేస్తుంది, ఇది టోకెన్‌లను మోడల్‌కు తిరిగి ఇస్తుంది.

కర్నెల్‌లు హోస్ట్‌తో గట్టిగా అనుసంధానమై ఉంటాయి, unified memory మరియు GDRCopy రెండింటి ద్వారా పరస్పరం కమ్యూనికేట్ అవుతాయి. CUDA unified memory ద్వారా device memory లేదా host memory రెండింటినీ రెండింటి address spaces‌లో మ్యాప్ చేయవచ్చు, PCIe ద్వారా కమ్యూనికేషన్‌ను పారదర్శకంగా నిర్వహిస్తూ.

బల్క్ బదిలీలకు యూనిఫైడ్ మెమరీ అనుకూలంగా ఉన్నప్పటికీ, మైక్రోసెకండ్ లాటెన్సీ వద్ద GPU మెమరీ నుండి చదవడానికి మరియు వ్రాయడానికి GDRCopy GPUDirect RDMAని ఉపయోగిస్తుంది.

RDMAతో పాటు, ఒక నోడ్ లోపల కర్నెల్స్ టోకెన్‌లను మార్పిడి చేయడానికి NVLinkని కూడా సద్వినియోగం చేసుకుంటాయి. ఇది EP32 కోసం నాల్గవ వంతు వరకు మరియు EP64 కోసం ఎనిమిదో వంతు వరకు ట్రాఫిక్‌లో గణనీయమైన భాగాన్ని ఆఫ్‌లోడ్ చేయగలదు, తద్వారా మొత్తం తక్కువ లాటెన్సీలను సాధించవచ్చు.

సెండ్ మరియు రిసీవ్ బఫర్లు

హోస్ట్ ప్రాక్సీని ಒಳಗೊಂಡ ఆచరణీయమైన లాటెన్సీలను సాధించడంలో డిస్పాచ్ కర్నెల్ కోసం సెండ్ బఫర్ల డిజైన్ చాలా కీలకం. ప్రతి ఒక్కటి T టోకెన్‌లను డిస్పాచ్ చేసే N ర్యాంక్‌లు ఉంటే, ప్రతి ఒక్కటి Eలో R ఎక్స్‌పర్ట్‌లకు రూట్ చేయబడితే, అన్నింటినీ దీనికి రూట్ చేస్తే, చెత్త సాధ్యమైన సందర్భంలో ఒక ర్యాంక్ గరిష్టంగా N * T * max(R, E / N) టోకెన్‌లను స్వీకరించగలదు. అంతేకాకుండా, ప్రైవేట్ రిసీవ్ బఫర్‌లు చాలా ఎక్కువ మెమరీని వినియోగిస్తాయి కాబట్టి, సెండర్‌లు తమ రైట్‌లను వరుసగా ఉంచినప్పుడు మాత్రమే ఈ కనిష్ట బౌండ్ సాధించబడుతుంది. సెండర్‌లు వైరుధ్యాలు లేకుండా వరుసగా తమ రైట్‌లను ఉంచగలగడానికి, రూటింగ్ సమాచారం గురించిన పూర్తి జ్ఞానం ప్రతి ర్యాంక్‌కు పంపిణీ చేయబడాలి. ఫలితంగా, మా డిస్పాచ్ కర్నెల్స్ మొదట పర్-ఎక్స్‌పర్ట్ టోకెన్ గణనలను మార్పిడి చేసుకుంటాయి, ఆ తర్వాత అవి తమ రైట్‌లు ఎక్కడికి వెళ్తాయి మరియు ఏ ఆఫ్‌సెట్ నుండి ప్రతి ర్యాంక్ నుండి టోకెన్‌లను స్వీకరిస్తాయో ప్రతి ఒక్కటి నిర్ణయించగలవు.

Routing సమాచారాన్ని కేవలం మార్పిడి చేయడం వల్ల అనవసరంగా అధిక లేటెన్సీలు వస్తాయి, ఎందుకంటే ప్రతి rank తన స్వంత writes‌ను dispatch చేయడానికి ముందు writes వచ్చే వరకు వేచి ఉండాలి. బదులుగా, మేము ప్రతి receiver rank‌పై ప్రతి sender‌కు చెందిన కొద్దిపాటి private space‌ను రిజర్వ్ చేసి, routing సమాచారంలోని writes‌తో పాటు కొన్ని tokens‌ను dispatch చేస్తాము. ఇది dispatch ప్రక్రియలో ఎక్కువ భాగం boyunca పూర్తి bandwidth వినియోగాన్ని నిర్ధారిస్తుంది. routes మార్పిడి అయిన తర్వాత, మిగిలిన tokens ప్రతి remote peer‌కు ఒకే write ద్వారా బదిలీ అవుతాయి. routing సమాచార మార్పిడి మరియు fixed-sized buffers, వాటిని నింపే transfers‌తో కలిసి, క్రింది చిత్రంలో చూపబడ్డాయి.

విస్తరించు

డిస్పాచ్ దశ నుండి రూటింగ్ సమాచారం ఇప్పటికే అందుబాటులో ఉన్నందున కంబైన్ కర్నెల్స్ చాలా సులభం. ఇవి పీర్‌కు ఒక్క రైట్ చొప్పున టోకెన్‌లను కాపీ చేయగలవు, ఆ తర్వాత అవుట్‌పుట్ టెన్సార్‌లలోకి టోకెపాల వెయిటెడ్ యావరేజ్‌ను లెక్కించడానికి ఇవి రిసీవ్ బఫర్ నుండి చదువుతాయి.

NVLink కోసం, ఒక నోడ్ లోపల ఉండే సెండ్ మరియు రిసీవ్ బఫర్లు పీర్‌లలోకి మ్యాప్ చేయబడతాయి, తద్వారా అవి ఇరుగుపొరుగు ర్యాంక్‌ల నుండి నేరుగా టోకెన్‌లను పుష్ చేయగలవు లేదా పుల్ చేయగలవు.

డిస్పాచ్ మరియు కంబైన్

క్రింది చిత్రాలు GPU, CPU మరియు NIC మధ్య సమయం మరియు పరస్పర చర్యలను వివరంగా తెలియజేస్తాయి, మోడల్ కర్నెల్స్‌తో ఇంటర్‌లీవ్ చేయబడిన డిస్పాచ్ మరియు కంబైన్ యొక్క ఒక పునరావృతాన్ని చూపుతాయి.

పని dispatch send‌లో ప్రారంభమవుతుంది, అక్కడ కర్నెల్ routing సమాచారాన్ని సమీకరించి, single block‌లోని shared memoryలో atomic accumulation ద్వారా ప్రతి expert‌కు token countలను గణిస్తుంది. token countలు అందుబాటులోకి వచ్చిన తర్వాత, అవి host memoryకు కాపీ చేయబడతాయి మరియు proxy వెంటనే ఇతర peer ranks అన్నింటికీ routing సమాచారాన్ని dispatch చేయడం ప్రారంభించడానికి తెలియజేయబడుతుంది. సమాంతరంగా, dispatch token countల cumulative sumను గణించడం కొనసాగిస్తుంది, contiguous per-peer send buffersలో tokens‌ను pack చేయడానికి అవసరమైన offsets‌ను గుర్తిస్తుంది. offsets అందుబాటులోకి వచ్చిన తర్వాత, EFA ద్వారా పంపాల్సిన అన్ని tokens send buffersలో pack చేయబడతాయి. inter-node peers యొక్క send buffers‌లో writes‌ను post చేయడానికి proxyకి తెలియజేయబడుతుంది. intra-nodeలో, కర్నెల్‌లు అదే nodeలోని private buffers‌లోకి NVLink ద్వారా tokens‌ను వెంటనే push చేస్తాయి.

విస్తరించు

విస్తరించు

ఒక ర్యాంక్ అన్ని పీర్ల నుండి రూటింగ్ సమాచారాన్ని స్వీకరించిన తర్వాత, అది దానిని కేంద్రీకరిస్తుంది మరియు డిస్పాచ్ మరియు కంబైన్ కోసం టోకెన్‌లను గుర్తించడానికి అవసరమైన వివిధ ఆఫ్‌సెట్‌లను లెక్కిస్తుంది. బదిలీలు ఇప్పటికే పెండింగ్‌లో ఉన్నప్పుడు ఈ పని CPUపై చేయబడుతుంది, కాబట్టి ఇది లాటెన్సీ క్రిటికల్ కాదు. తదనంతరం, అన్ని ఆఫ్‌సెట్‌లు తెలిసిన తర్వాత, ప్రైవేట్ బఫర్ పరిమాణాన్ని మించిన మిగిలిన టోకెన్‌ల కోసం రైట్‌లు పోస్ట్ చేయబడతాయి మరియు డిస్పాచ్ యొక్క సెండ్ దశ పూర్తవుతుంది.

receiver వైపున, కర్నెల్ మొదట అదే node‌లోని private buffers సామర్థ్యాన్ని మించిన tokens‌ను NVLink ద్వారా pull చేస్తుంది. తరువాత, అన్ని transfers పూర్తైనట్లు host proxy సంకేతం ఇచ్చే వరకు అది block అవుతుంది. సంకేతం అందిన తర్వాత, tokens receive buffers నుంచి group GEMM‌కు పంపిన output tensor‌లోకి shuffle చేయబడతాయి. dispatch మరియు combine మధ్య buffers మళ్లీ ఉపయోగించబడుతున్నందున, తదుపరి కర్నెల్ వాటిని overwrite చేయడానికి అనుమతించే ముందు receiver buffers చదివివేయబడ్డాయని నిర్ధారించడానికి NVLink మరియు RDMA barriers రెండూ ఉపయోగించబడతాయి.

combine కర్నెల్‌లు సుమారు ఇలాంటి నమూనాను అనుసరిస్తాయి, tokens‌ను shuffle చేయడం, proxyకి signal ఇవ్వడం మరియు completion కోసం వేచి ఉండడం. అయితే, వాటి writes‌ను post చేయడానికి అవసరమైన metadata సులభంగా అందుబాటులో ఉంటుంది మరియు అన్ని remote offsets కూడా తెలిసి ఉంటాయి; అందువల్ల intra-nodeగా బదిలీ అయ్యే అన్ని tokens NVLink ద్వారా receivers‌కు push చేయబడతాయి మరియు inter-node tokens ఒకే write ద్వారా బదిలీ అవుతాయి. ఫలితాలను తిరిగి ఇచ్చే ముందు, కర్నెల్ receiver buffers నుంచి tokens‌ను సమీకరించి, వాటిని output tensors‌లోకి రాస్తుంది.

MoE రూటింగ్ కోసం TransferEngine

మా అమలును ప్రోటోటైప్ చేయడానికి, RDMA రైట్‌లను జారీ చేయడానికి మరియు ట్రాక్ చేయడానికి మేము మొదట డిస్గ్రిగేట్ ప్రిఫిల్ మరియు డికోడ్ కోసం నిర్మించిన Transfer Engineని తిరిగి ఉపయోగించాము. ఏదేమైనప్పటికీ, ప్రతి నోడ్ అదే పీర్‌కు పెద్ద సంఖ్యలో రైట్‌లను పోస్ట్ చేసే డిస్గ్రిగేటెడ్ డికోడింగ్ వలె కాకుండా, MoE డిస్పాచ్ అన్ని పీర్‌లకు ఒకే ఆపరేషన్‌లో చిన్న బదిలీలను పోస్ట్ చేస్తుంది. ఆటోమేటిక్‌గా, పనితీరు ఆప్టిమల్‌గా లేదు మరియు మా అమలు ConnectX-7పై DeepEP కంటే దాదాపు 20us వెనుకబడి ఉంది. ఫలితంగా, పెద్ద సంఖ్యలో పీర్‌లతో పరస్పర చర్యలను ఆప్టిమైజ్ చేయడానికి స్క్యాటర్ మరియు బ్యారియర్ అనే రెండు ఆపరేషన్‌లకు మద్దతు ఇవ్వడానికి మేము TransferEngineని ప్రత్యేకతను సంతరించుకున్నాము.

స్క్యాటర్ మరియు బ్యారియర్ ఆపరేషన్‌లు రెండూ అంతర్లీన రవాణా అమలు (వర్క్ రిక్వెస్ట్ (WR) టెంప్లేటింగ్) ద్వారా ఉపయోగించబడే డేటా స్ట్రక్చర్‌లను ముందే జనాభా నింపడానికి ముందే రిజిస్టర్ చేయగల పీర్ గ్రూప్‌లను బహిర్గతం చేస్తాయి. స్క్యాటర్ అదే సోర్స్ బఫర్ నుండి వేర్వేరు స్లైస్‌లను పీర్ గ్రూప్ లోని సభ్యులందరిపై వేర్వేరు ఆఫ్‌సెట్‌లలోకి కాపీ చేస్తుంది. బ్యారియర్ ఇమ్మీడియట్ ద్వారా అన్ని పీర్‌లకు సిగ్నల్ ఇస్తుంది. ఈ ఆపరేషన్‌లను బల్క్‌లో ఎన్‌క్యూ చేయడం మరియు నిర్వహించడం వలన క్రిటికల్ పాథ్‌పై విలువైన మైక్రోసెకండ్‌లు ఆదా అవుతాయి. రెండు 200Gbps NICలు 400Gbps బ్యాండ్‌విడ్త్‌ను అందించడానికి సేకరించే EFAపై, ప్రతి బదిలీ యొక్క బైట్‌లను షార్డింగ్ చేయడానికి బదులుగా మేము రెండు NICలలో గ్రూప్ పీర్‌లను షార్డ్ చేస్తాము.

ప్రారంభంలో మేము libfabricని ఉపయోగించి EFAపై మాత్రమే TransferEngineని అభివృద్ధి చేసినప్పటికీ, libibverbsని ఉపయోగించి TransferEngineలో ConnectX-7 మద్దతును జోడించిన తర్వాత MoE కర్నెల్స్ ఒక్క కోడ్ లైన్‌ను కూడా మార్చకుండా ConnectX-7పై పని చేస్తాయి. స్కేలబుల్ రిలయబుల్ డేటాగ్రామ్ (SRD) ఆధారిత EFAతో పోలిస్తే, ConnectX-7 కనెక్షన్ సెటప్ మరియు పీర్ మేనేజ్‌మెంట్ కోసం అదనపు పని అవసరం. మేము రిలయబుల్ కనెక్షన్ (RC)ని బూట్‌స్ట్రాప్ చేయడానికి అన్‌లయబుల్ డేటాగ్రామ్ (UD) క్యూ జతను ఉపయోగిస్తాము. మేము ప్రతి రిమోట్ పీర్ కోసం రెండు క్యూ జతలను సృష్టిస్తాము, ఒకటి టూ-సైడెడ్ SEND/RECV కోసం మరియు మరొకటి వన్-సైడెడ్ WRITE_IMM కోసం, ఎందుకంటే RECV మరియు WRITE_IMM రెండూ క్రమంలో RECV ఆపరేషన్‌ను వినియోగిస్తాయి అయితే రెండు ఆపరేషన్‌లకు వేర్వేరు బఫర్ పరిమాణాలు అవసరమవుతాయి. అంతేకాకుండా, మేము ConnectX-7 మరియు GPU మెమరీ మధ్య PCIe ట్రాన్సాక్షన్ ఆర్డరింగ్‌ను సడలిస్తాము. EFA వలె కాకుండా, మేము WR టెంప్లేటింగ్ ఆప్టిమైజేషన్‌ను అమలు చేస్తాము. అదనంగా, ఒకే పీర్‌కు బహుళ చిన్న సందేశాలను పంపేటప్పుడు, సందేశ రేటును మెరుగుపరచడానికి మేము WR చైనింగ్‌ను ఉపయోగిస్తాము.

GPU-సైడ్ పని రెండు అడాప్టర్‌లలో ఒకే విధంగా ఉన్నందున, EFA మరియు ConnectX-7 MoE రూటింగ్ మధ్య అంతరం అంతర్లీన హార్డ్‌వేర్ సామర్ధ్యాల ద్వారా నిర్ణయించబడుతుందని హైలైట్ చేయడానికి DeepEP పనితీరును మించిపోవడానికి మేము అంతర్లీన రవాణా అమలును దాని పరిమితులకు నెట్టాము.

మూల్యాంకనం

డిస్పాచ్/కంబైన్ లాటెన్సీ

7168 హిడెన్ డైమెన్షన్, బ్లాక్-స్కేల్డ్ fp8 డిస్పాచ్ మరియు bf16 కంబైన్‌తో ప్రామాణిక DeepSeek-V3 మరియు Kimi-K2 కాన్ఫిగరేషన్‌పై ConnectX-7 మరియు EFA అడాప్టర్‌లతో కూడిన H200 సిస్టమ్‌లపై మేము డిస్పాచ్ మరియు కంబైన్ లాటెన్సీలను మూల్యాంకనం చేసాము. మేము డిస్పాచ్ కోసం 128 మరియు కంబైన్ కోసం 4096 బ్యాచ్ పరిమాణాన్ని పరిగణలోకి తీసుకుంటాము. ConnectX-7పై, మేము GPUDirect Async (IBGDA) మరియు హోస్ట్ ప్రాక్సీ (IBRC) రెండింటినీ ఉపయోగించి DeepEP మరియు మా మునుపటి NVSHMEM-ఆధారిత pplx-కర్నెల్స్‌తో పోల్చాము. మేము EP16, EP32 మరియు EP64పై వరుసగా 459us, 582us మరియు 692us యొక్క మిళిత లాటెన్సీలను సాధిస్తాము. మేము 519us, 966us మరియు 1159us సాధించే UCCL-EP వంటి సాధారణ ప్రాక్సీలపై ఆధారపడే ఇతర EFA-ఆధారిత కర్నెల్స్‌ను అధిగమిస్తాము.

విస్తరించు

డికోడ్‌పై, మా కర్నెల్స్ DeepEP కంటే కొన్ని మైక్రోసెకండ్లు నెమ్మదిగా ఉన్నప్పటికీ, మా కంబైన్ కర్నెల్స్ గణనీయంగా వేగంగా ఉండటం వలన మేము మొత్తం తక్కువ లాటెన్సీలను సాధిస్తాము. మేము హోస్ట్ ప్రాక్సీని ఉపయోగించినప్పటికీ, మా బల్క్ బదిలీల కారణంగా మేము అన్ని IBGDA-ఆధారిత ప్రాక్సీలెస్ అమలుల పనితీరును మించిపోతాము. మా అత్యంత ప్రత్యేకమైన ప్రాక్సీ సాధారణ NVSHMEM అమలు కంటే పది రెట్లు వేగంగా ఉంటుంది కాబట్టి, ప్రాక్సీ అమలు ముఖ్యం అని మేము కూడా చూపుతాము. MoE రూటింగ్ ద్వారా మార్పిడి చేయబడిన సాధారణ ప్యాకెట్ పరిమాణానికి దగ్గరగా ఉండే 256KB ప్యాకెట్లను బదిలీ చేసేటప్పుడు EFA ConnectX-7 కంటే రెండు రెట్లు నెమ్మదిగా ఉన్నప్పటికీ, రెండు అడాప్టర్‌ల మధ్య అంతరం మొత్తం మీద తక్కువ గణనీయంగా ఉంటుంది.

విస్తరించు

మా డికోడ్-ఆప్టిమైజ్ చేయబడిన కర్నెల్స్ ఎలాంటి అదనపు ట్వీక్‌లు లేకుండా ప్రిఫిల్ బ్యాచ్ పరిమాణాలకు స్కేల్ చేయగలప్పటికీ, పనితీరు DeepEP వెనుకబడి ఉంది. మేము కంబైన్ యొక్క రిసీవర్ వైపు మాత్రమే టోకెన్‌లను సేకరిస్తాము, అయితే బదిలీ చేయబడిన డేటా మొత్తాన్ని తగ్గించడానికి DeepEP సెండర్‌లో సేకరిస్తుంది అనే వాస్తవం దీనికి ప్రధాన కారణం. ఏదేమైనప్పటికీ, ఈ ఆప్టిమైజేషన్ ఖచ్చితత్వ ప్రభావాలను కలిగి ఉంటుంది. అదనంగా, కర్నెల్స్ ఒకే రైట్‌లో డేటాను బదిలీ చేస్తాయి కాబట్టి, మేము సెండ్ మరియు రిసీవ్ బఫర్‌ల కోసం ఎక్కువ మెమరీని ఉపయోగిస్తాము.

Kimi-K2 / DeepSeek-V3 త్రూపుట్

మా కర్నెల్స్ యొక్క ఆచరణీయతను నిరూపించడానికి, ఎలాస్టిక్ ఫాబ్రిక్ అడాప్టర్లు (EFA) మరియు H200 GPUలతో కూడిన బహుళ p5en ఇన్‌స్టాన్స్‌లపై DeepSeek-V3 (671B పారామీటర్లు) మరియు Kimi-K2 (1T పారామీటర్లు) సేవలను అందించే మా ఇన్‌హౌస్ ఇన్‌ఫరెన్స్ ఇంజిన్ యొక్క ఎండ్-టు-ఎండ్ పనితీరును మేము బంచ్‌మార్క్ చేస్తాము. మేము మోడల్‌ను డ్రాఫ్ట్ వెరిఫికేషన్ మోడ్‌లో రన్ చేస్తాము మరియు ఒకే MTP లేయర్‌ని ఉపయోగించి 80% అక్సెప్టెన్స్ రేట్‌ను ఊహిస్తూ లాటెన్సీని నివేదిస్తాము.

DeepSeek-V3పై, సమర్థవంతమైన NVLink-మాత్రమే డిస్పాచ్ మరియు కంబైన్ కర్నెల్స్‌ను రన్ చేసే EP=8 ఇంట్రానోడ్ బేస్‌లైన్‌తో మేము పోల్చుతాము. ఇంటర్-నోడ్ డిప్లాయ్‌మెంట్‌లు మొత్తం మీద సింగిల్-నోడ్ డిప్లాయ్‌మెంట్ యొక్క పనితీరుతో సరిపోతాయి లేదా మించిపోతాయి, ప్రత్యేకించి మీడియం-సైజ్డ్ బ్యాచ్‌లపై, నోడ్‌లలో స్కేలింగ్ చేయడం వల్ల మొత్తం త్రూపుట్‌కు మెరుగుదలలు వస్తాయని ఇది ప్రదర్శిస్తుంది.

Kimi-K2పై, కర్నెల్స్ సర్వ్‌ను ప్రారంభిస్తాయి. దాని పరిమాణం కారణంగా, H200 నోడ్‌లో డేటా పారలలిజం (DP)ని ఉపయోగించి మోడల్ ఒకే నోడ్ లోపల సరిపోదు, అందువల్ల మోడల్‌కు సేవ చేయడానికి ఇంటర్-నోడ్ డిప్లాయ్‌మెంట్ అవసరం. మోడల్‌కు తక్కువ అటెన్షన్ హెడ్‌లు ఉన్నందున, మేము DeepSeek-V3 కంటే మెరుగైన లాటెన్సీలను సాధిస్తాము మరియు మేము ఆచరణీయమైన లాటెన్సీ వద్ద EFAపై మోడల్‌కు సేవ చేయవచ్చు.

విస్తరించు

విస్తరించు

మేము మైక్రో-బ్యాచింగ్‌ను కూడా పరిగణనలోకి తీసుకున్నాము, అయినప్పటికీ ఇది చాలా ఎక్కువ బ్యాచ్ పరిమాణాల వద్ద లేదా ఒకటి కంటే ఎక్కువ డ్రాఫ్ట్ టోకెన్‌లను ఉపయోగించినప్పుడు మాత్రమే లాటెన్సీలను మెరుగుపరిచింది. గ్రూప్ చేయబడిన GEMM సమస్యలను రెండుగా విభజించడం ద్వారా ప్రయోజనం పొందేందుకు, టోకెన్ బ్యాచ్ పరిమాణాలు తగినంత దట్టంగా ఉండాలి. ఎక్కువ సమయం ఇప్పటికీ MoE డిస్పాచ్ మరియు కంబైన్ పైనే గడపబడుతుంది, కర్నెల్స్‌కు భవిష్యత్తులో చేసే ఏవైనా మెరుగుదలలు తదనంతరం లాటెన్సీలను మెరుగుపరుస్తాయని మరియు మైక్రో-బ్యాచింగ్‌ను సులభతరం చేస్తాయని ఇది సూచిస్తుంది.

భవిష్యత్ పని

EFAపై పనితీరును మెరుగుపరచడానికి మేము AWS ఇంజనీర్‌లతో కలిసి పని చేస్తున్నాము. మా వైపు, డేటా ప్లేన్ ఓవర్‌హెడ్‌ను తగ్గించే libfabric అప్‌డేట్‌లను మేము అనుసరిస్తాము. అదనంగా, యూజర్‌స్పేస్ ఓవర్‌హెడ్‌లను మరింత తగ్గించడానికి మరియు హార్డ్‌వేర్‌తో నేరుగా ఇంటర్‌ఫేస్ చేయడానికి మేము efa-directతో ప్రయోగాలు చేస్తాము.

ప్రస్తావనలు

DeepSeek R1 టెక్నికల్ నివేదిక

Kimi K2 టెక్నికల్ నివేదిక

DeepEP

UCCL-EP

NVIDIA Magnum IO NVSHMEM మరియు GPUDirect Async ఉపయోగించి HPC సిస్టమ్‌ల నెట్‌వర్క్ పనితీరును మెరుగుపరచడం

ఉద్ధృతి

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}