Perplexityలో CuTeDSL

Perplexityలో, ఎంబెడ్డింగ్‌ల నుండి ట్రిలియన్-పారామీటర్ LLMల వరకు మోడల్‌లను అందించడానికి మేము అంతర్గత రన్‌టైమ్-ఆప్టిమైజ్డ్ సర్వింగ్ ఇంజిన్ (ROSE)ని ఉపయోగిస్తాము. Sonar, Search మరియు Embeddings సహా అన్ని APIల వెనుక ఉంటూ, మా పరిశోధకులు అత్యంత తగిన మోడల్స్ మరియు ఆర్కిటెక్చర్‌లను అన్వేషించడానికి, అలాగే ఉత్తమ అభ్యర్థులను త్వరగా ప్రొడక్షన్‌లోకి మోహరించడానికి మా రన్‌టైమ్ శక్తినిస్తుంది. దీన్ని సాధించడానికి, మేము స్టేట్-ఆఫ్-ది-ఆర్ట్ ఆర్కిటెక్చర్‌ల కోసం ఇన్‌ఫరెన్స్‌ను అమలు చేయడం, మోహరించడం మరియు ఆప్టిమైజ్ చేయగలగాలి. ఈ బ్లాగ్ పోస్ట్‌లో, మా ఇన్‌ఫరెన్స్ ఇంజిన్‌లో ఏకీకృతం చేయబడిన CuTeDSL కెర్నల్‌లు NVIDIA GPUలపై మోడల్‌లను గరిష్ట పనితీరుకు తీసుకురావడానికి అవసరమైన GPU కెర్నళ్లను రూపొందించడానికి మాకు ఎలా అనుమతిస్తాయో మేము అన్వేషిస్తాము.

Perplexityలోని వివిధ బృందాలు నిర్మించిన కస్టమ్ మోడల్‌లు, AI ఇన్‌ఫరెన్స్ బృందం నిర్మించిన ROSE ఇన్‌ఫరెన్స్ ఇంజిన్‌ను ఉపయోగించి NVIDIA హోప్పర్ మరియు బ్లాక్‌వెల్ GPUలలో అంతర్గతంగా హోస్ట్ చేయబడతాయి. NVIDIA ట్రిటన్ ఇన్‌ఫరెన్స్ సర్వర్‌తో అనుకూలమైన ఇంటర్‌ఫేస్‌తో లాంగ్వేజ్ డీకోడింగ్ మరియు వర్గీకరణ రెండింటి కోసం అనుకూలీకరించిన Llama మోడళ్లను మెరుగ్గా అందించడానికి ROSE ప్రారంభంలో నిర్మించబడింది. కాలక్రమేణా, ఇది Perplexity ఉత్పత్తులకు శక్తినివ్వడానికి విస్తృత శ్రేణి బృందాల కోసం గణనీయమైన LLMలను, అలాగే ర్యాంకింగ్, వర్గీకరణ, స్కోరింగ్ మరియు ఇతర ట్రాన్స్‌ఫార్మర్-ఆధారిత మోడళ్లను అందించే పూర్తి స్థాయి ఇన్‌ఫరెన్స్ ఇంజిన్‌గా అభివృద్ధి చెందింది. ROSE యొక్క డిజైన్ లక్ష్యం సరళంగా ఉండటం: కాలక్రమేణా, గరిష్ట పనితీరును చేరుకోవడానికి కృషి పెట్టుబడి పెట్టబడుతుంది, అయితే మా మోడల్‌ల కోసం నిరంతరం అభివృద్ధి చెందుతున్న ఆర్కిటెక్చర్‌లు మరియు ఉపయోగ సందర్భాలకు అనుగుణంగా మారడమే ప్రధాన ప్రాధాన్యత.

డేటా రకాల, హెడ్ డైమెన్షన్లు, కంపైల్-టైమ్ కాన్స్టెంట్లు, మరియు లాంచ్ సెట్టింగ్‌లలోని కర్నెల్ కాన్ఫిగరేషన్ స్పేస్.

అధిక స్థాయిలో, ROSE ఒక ఇంజిన్ను కలిగి ఉంటుంది; అది ఒక మోడల్‌ను క్లయింట్లు పరస్పర చర్య చేయగల ఇంటర్‌ఫేస్‌కు అనుకూలిస్తుంది. ఈ ఇంజిన్‌లు అత్యంత పునర్వినియోగయోగ్యమైనవి; రిక్వెస్ట్ షెడ్యూలింగ్ మరియు బ్యాచింగ్‌ను నిర్వహించడం, అలాగే అంతర్లీన మోడళ్ల నుంచి డీకోడ్ చేసిన టోకెన్లు, ఎంబెడ్డింగ్లు లేదా స్కోర్లను తిరిగి పొందడం వీటి బాధ్యత. డివైస్‌ల ప్రారంభీకరణ, ఇంటర్-ప్రాసెస్ మరియు ఇంటర్-నోడ్ కమ్యూనికేషన్ ఛానెల్‌లు, అలాగే వెయిట్ లోడింగ్ అన్నీ వీటిద్వారా నిర్వహించబడతాయి. LLMల కోసం, ఇంజిన్ బ్యాచింగ్, చంకింగ్, సాంప్లింగ్‌కు బాధ్యత వహిస్తుంది. KV స్టోరేజ్ కేటాయింపు కూడా, ఫుల్ మరియు లీనియర్ అటెన్షన్ రెండింటికీ, అలాగే ప్రిఫిక్స్ మ్యాచ్‌ింగ్‌కు కూడా ఇదే బాధ్యత వహిస్తుంది. ఎంబెడ్డింగ్ ఇంజిన్‌లు మరింత సరళమైనవి, ఎందుకంటే అవి కేవలం ఆన్‌లైన్ బ్యాచింగ్‌ను మాత్రమే నిర్వహిస్తాయి. అమలు ఈ విధంగా ఆప్టిమైజ్ చేయబడింది: మోడల్ యాక్సిలరేటర్‌పై సమాంతరంగా నడుస్తున్నప్పుడు, ఇంజిన్ తదుపరి పనుల చక్రాన్ని సిద్ధం చేస్తుంది.

ROSE, మోడలింగ్ కోసం, CuTeDSL, Triton, CUDA, CUTLASS, మరియు cuBLAS‌తో నిర్మించిన GPU కర్నెల్‌లను ర్యాప్ చేసే కస్టమ్ లేయర్‌ల సమృద్ధి సేకరణను అందిస్తుంది. మ్యాట్రిక్స్ మల్టిప్లికేషన్ మరియు అటెన్షన్ అమలుల కోసం, మేము ప్రధానంగా CUTLASS మరియు cuBLAS నుండి వచ్చే NVIDIA కర్నెల్‌లపై ఆధారపడుతాము. ఎంబెడ్డింగ్, నార్మ్, MoE, మరియు యాక్టివేషన్ కర్నెల్‌లను కలిగి ఉన్న మిగిలిన పైప్‌లైన్, ROSE కర్నెల్ లైబ్రరీలో భాగం. మేము ఇన్‌ఫరెన్స్ నడిపే హార్డ్‌వేర్‌కు ఉత్తమంగా మద్దతివ్వడానికి, MoE రౌటింగ్, డిస్పాచ్, మరియు కాంబైన్ కోసం ROSE తన స్వంత ఇంటర్-డివైస్ కమ్యూనికేషన్ ప్రిమిటివ్‌లతో సన్నద్ధం చేయబడింది. మిగిలిన కర్నెల్‌లు తులనాత్మకంగా సరళమైనవే అయినప్పటికీ, అవి నడిచే వివిధ హిడెన్ డైమెన్షన్‌లు మరియు డివైస్‌లన్నింటిలో గరిష్ఠ పనితీరును సాధించేందుకు అవి పూర్తిగా స్పెషలైజ్ చేయబడాలి. అదనంగా, వివిధ మోడళ్ల అన్ని నియంత్రణలు మరియు సర్దుబాట్లను (RMS norm కోసం weight bias వర్సెస్ no weight bias, యాక్టివేషన్ ఫంక్షన్‌లు మొదలైనవి) సమర్థించేందుకు మేము compile-time specialization‌పై ఆధారపడుతాము. ఫలితంగా, వాటిని అమలు చేయడానికి DSLలను ఉపయోగించాలని మేము నిర్ణయించుకున్నాము: మొదట Triton‌ను ఉపయోగించి ROSE‌ను నిర్మించాము. అయితే, అప్పటి నుంచి GPU ప్రోగ్రామింగ్ కోసం ప్రధాన వాతావరణంగా CuTeDSL‌ను స్వీకరించాము.

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

CuTeDSL అనేది CuTe లేఅవుట్ ఆల్జీబ్రా మరియు MLIR పైన నిర్మించబడిన NVIDIA GPUల కోసం పైథాన్-ఆధారిత DSL. ఇది తక్కువ-స్థాయి హార్డ్‌వేర్ ప్రిమిటివ్‌లపై నియంత్రణను కలిగి ఉంటూనే, అత్యంత ఆప్టిమైజ్ చేయబడిన PTXకి జస్ట్-ఇన్-టైమ్‌లో కంపైల్ చేస్తుంది. ఇది అభివృద్ధి వేగాన్ని అడ్డుకునే ఖరీదైన కంపైలేషన్ సమయాలు లేకుండా, CUTLASS C++ వలె అదే లక్షణాలను మరియు పనితీరును అందించాలని లక్ష్యంగా పెట్టుకుంది. మేము సర్వే చేసిన లేదా పని చేసిన భాషలలో (CUDA, cuTile, TVM IR మరియు Tritonతో సహా), మా కెర్నళ్ల అభివృద్ధికి ఇది అత్యంత సరిపోతుందని మేము భావిస్తున్నాము.

మొత్తంమీద, Triton లేదా cuTile వలె హై-లెవెల్ మరియు సంక్షిప్తంగా కాకపోయినా, CuTe లేఅవుట్ ఆల్జీబ్రా మరియు దాని కూర్పులు (cute.Tensor ద్వారా ప్రదర్శించబడతాయి) సహేతుకమైన ప్రత్యామ్నాయాలుగా ఉన్నాయని మేము కనుగొన్నాము. అదనంగా, CuTeDSL హార్డ్‌వేర్ ప్రిమిటివ్‌లకు అనియంత్రిత యాక్సెస్‌ను అందిస్తుంది, అన్ని ఎడ్జ్ కేస్‌లలో గరిష్ట పనితీరును సాధించడానికి మాకు అవసరమైన స్థాయి నియంత్రణను అందిస్తుంది. CUDAతో, ప్రతి అదనపు పరామితి కంపైలేషన్ సమయాలకు సహేతుకం కాని ఖర్చును లేదా బిల్డ్ సిస్టమ్‌కు సంక్లిష్టతను జోడించినందున, కంపైల్-టైమ్ స్పెషలైజేషన్ అవకాశాలు పరిమితం చేయబడ్డాయి. దీనికి విరుద్ధంగా, జస్ట్-ఇన్-టైమ్ కంపైల్ చేయబడిన DSLతో, ఇటరేషన్ వేగాన్ని త్యాగం చేయకుండా మేము మరింత దూకుడుగా ప్రత్యేకతను సాధించవచ్చు. CUDAని మధ్యవర్తి దశగా ఉపయోగించిన TVM IR కాకుండా, PTXకి తగ్గిన MLIR-ఆధారిత పైప్‌లైన్ CUDAని మధ్యవర్తి దశగా విడుదల చేసిన దానికంటే మరింత బలమైనదిగా మరియు డీబగ్ చేయడానికి సులభంగా ఉంటుందని మేము కనుగొన్నాము. అదనంగా, PTX సూచనలను చుట్టడానికి కస్టమ్ ఆప్స్‌ను ఉపయోగించడంతో, అంతర్లీన హార్డ్‌వేర్ యొక్క అన్ని లక్షణాలను మేము పూర్తిగా ఉపయోగించుకోవచ్చని మేము కనుగొన్నాము.

CUDA కెర్నళ్ల కంపైలేషన్ ఖర్చులో ఎక్కువ భాగం సాధారణంగా కెర్నళ్ల టెంప్లేట్ విస్తరణ దశలో ఖర్చు చేయబడుతుంది, ఎందుకంటే PTX దశ నుండి ఆప్టిమైజేషన్ మరియు కంపైలేషన్ కొన్ని సెకన్లలో పూర్తవుతాయి. అనేక కంపైల్-టైమ్ పారామీటర్‌లతో, నాన్-ట్రివియల్ కెర్నగ్‌ కోసం కంపైలేషన్ యొక్క ఈ దశ అనేక నిమిషాలకు పేలవచ్చు. సమాంతర కంపైలేషన్ కోసం అవకాశాలు IR మరియు PTX స్థాయిలలో మాత్రమే తదుపరి దశలో ఉద్భవించడమనే వాస్తవం ద్వారా ఈ బాటిల్‌నెక్ ఉద్ఘాటించబడింది, C++ టెంప్లేట్ నిర్వహణ క్రమక్రమంగా జరుగుతుంది. దీన్ని పరిష్కరించడానికి ఉపాయాలు ఉన్నప్పటికీ, టెంప్లేట్ ఇన్‌స్టాన్షియేషన్‌లను వేరువేరు కంపైలేషన్ యూనిట్లకు మార్చడం ద్వారా, వాటికి బిల్డ్ సిస్టమ్‌లకు గణనీయమైన ట్వీక్‌లు మరియు అసాధ్యమైన మాక్రోల వినియోగం అవసరం. దీనికి విరుద్ధంగా, జస్ట్-ఇన్-టైమ్ (JIT) కంపైల్ చేయబడిన CuTeDSL కెర్నళ్లకు సంక్లిష్టమైన బిల్డ్ నిబంధనలు అవసరం లేదు మరియు కంపైలేషన్ సమయాలను అడ్డుకోవు, ఎందుకంటే అభివృద్ధి చక్రం కాన్ఫిగరేషన్ స్పేస్ లోని ఒకే పాయింట్‌పై దృష్టి పెట్టవచ్చు, అయితే తదుపరి జాగ్రత్తగా యూనిట్ పరీక్షలు కెర్నగ్‌ యొక్క మొత్తం డిజైన్ స్పేస్‌ను కవర్ చేయగలవు.

JIT కంపైలేషన్ డెవలప్‌మెంట్-టైమ్ నుండి రన్‌టైమ్‌కు ఖర్చులను మార్చినప్పటికీ, ఓవర్‌హెడ్ ఆమోదయోగ్యమైనదానికంటే ఎక్కువ. ఒక మోడల్ మోహరించబడినప్పుడు, అది మొత్తం కాన్ఫిగరేషన్ స్పేస్ నుండి కొన్ని పాయింట్లను మాత్రమే సక్రియం చేస్తుంది, ఇన్‌ఫరెన్స్ ఇంజిన్ ప్రారంభం/వార్మప్ సమయంలో రెండవ-స్థాయి ఓవర్‌హెడ్‌లను కలిగిస్తుంది. దీనికి విరుద్ధంగా, కెర్నగ్‌ యొక్క అన్ని కాన్ఫిగరేషన్‌లను ముందుగా కంపైల్ చేయడం అనేక నిమిషాలను విస్తరించవచ్చు. అదనంగా, JIT-కంపైల్ చేయబడిన ఆర్టిఫ్యాక్ట్‌లను కాషింగ్ చేయడం ద్వారా ఓవర్‌హెడ్‌లను మరింత తగ్గించవచ్చు.

CuTe మరియు CUTLASS విషయంలో వలె, టెంప్లేట్-ఆధారిత డిస్పాచ్ మరియు స్పెషలైజేషన్ ఉపయోగించినప్పుడు లోపాల మూలాన్ని ట్రేస్ చేయడం కష్టంగా ఉండటం ద్వారా CUDA యొక్క భారీ కంపైలేషన్ ఖర్చు ఉద్ఘాటించబడుతుంది. DSLతో, లోపాలు దాన్ని oluşturిన కోడ్ లైన్‌కు చూపబడతాయి, డీబగ్గింగ్‌ను గణనీయంగా సులభతరం చేస్తాయి. అదనంగా, మధ్యవర్తి MLIR దశల ద్వారా లేదా రూపొందించబడిన PTXకి జోడించబడిన డీబగ్గింగ్ సమాచారంపై ఆధారపడటం ద్వారా రూపొందించబడిన కోడ్‌ను దాన్ని ఉద్భవించిన DSL స్టేట్‌మెంట్‌కు మ్యాప్ చేయడం చాలా సులభం. DSL యొక్క ఈ అంశాలకు ధన్యవాదాలు, కెర్నళ్లను అభివృద్ధి చేస్తున్నప్పుడు మేము చాలా ఎక్కువ ఉత్పాదకతను కలిగి ఉన్నామని కనుగొన్నాము.

కెర్నళ్లు మోడల్‌ల అంతటా అనేక సారూప్యతలను పంచుకున్నప్పటికీ, పనితీరు ప్రభావాన్ని చూపే అనేక కాన్ఫిగరేషన్ నాబ్‌లతో అవి ఇప్పటికీ అమర్చబడి ఉండాలి. సంబంధిత పారామీటర్లలో ఇవి ఉన్నాయి:

హిడెన్ డైమెన్షన్/హెడ్ డైమెన్షన్/అటెన్షన్ హెడ్‌ల సంఖ్య

ఇన్‌పుట్/అవుట్‌పుట్ డేటా రకాలు

క్వాంటైజేషన్ స్కీమ్

కంపైల్-టైమ్ స్థిరాంకాలు (బయాస్‌లు, యాక్టివేషన్ ఫంక్షన్‌లు మొదలైనవి)

ROSE ఆర్కిటెక్చర్, దాని ఇంజిన్, మోడల్, కర్నెల్ లైబ్రరీ, మరియు సేవలతో సహా.

కంపైల్-టైమ్ స్పెషలైజేషన్ పనితీరు మరియు డీబగ్గింగ్ రెండింటికీ ఉపయోగపడుతుంది. ముఖ్యంగా పైప్‌లైన్ ప్రాసెసర్‌లో, GPUల విషయంలో వలె, ఎప్పటికీ తీసుకోని లేదా ఎప్పుడూ తీసుకోని బ్రాంచ్ కంటే ఉనికిలో లేని బ్రాంచ్ మెరుగైనది. అదనంగా, నిర్దిష్ట కాన్ఫిగరేషన్ నాబ్‌ల సెట్ కోసం కెర్నగ్‌ను కంపైల్ చేయడం వలన తక్కువ PTX/SASS సూచనలు వస్తాయి, ఇవి పరిశోధన సమయంలో cuda-gdb ఉపయోగిస్తున్నప్పుడు విశ్లేషించడానికి మరియు డీబగ్ చేయడానికి సులభంగా ఉంటాయి.

ఇచ్చిన బ్యాచ్ పరిమాణం మరియు కాన్ఫిగరేషన్ కోసం, DRAMకి अपेक्षित ట్రాఫిక్ మరియు అంకగణిత కార్యకలాపాల సంఖ్యను అంచనా వేయవచ్చు. తదనంతరం, సరైన గ్రిడ్ కేటాయింపు, సింక్రొనైజేషన్ మరియు కొన్ని హార్డ్‌వేర్ ఫీచర్‌ల వినియోగం కోసం కేస్-బై-కేస్ ప్రాతిపదికన సరైన నిర్ణయాలు తీసుకోవచ్చు. ఆమోదయోగ్యం కాని కంపైలేషన్-టైమ్ ఓవర్‌హెడ్‌ను భరించకుండా లోతైన స్పెషలైజేషన్ కోసం భాగాలను నిర్వచించడానికి మరియు పునర్వినియోగించడానికి CuTeDSL మాకు అనుమతిస్తుంది.

సాధారణంగా, ఇన్‌ఫరెన్స్ కెర్నళ్ల పనితీరు రెండు వేర్వేరు ఉపయోగ సందర్భాల కోసం ఆప్టిమైజ్ చేయబడాలి:

ప్రిఫిల్: చిన్న బ్యాచ్ పరిమాణం (కొన్ని చంక్‌లు, సీక్వెన్స్‌లు), ప్రతి సీక్వెన్స్‌కు అనేక టోకెన్‌లు. ఎంబెడ్డింగ్‌లు మరియు LLMల ఇన్‌పుట్ టోకెన్‌లు రెండింటికీ త్రూపుట్ గరిష్టీకరించబడాలి. పది నుండి వందల మిల్లీసెకన్లలో లేటెన్సీ.

డీకోడ్: చిన్న-నుంచి-పెద్ద బ్యాచ్ పరిమాణాలు, ప్రతి సీక్వెన్స్‌కు ఒకటి (స్పెక్యులేటివ్ డీకోడింగ్ లేదు) నుండి కొన్ని టోకెన్‌లు (డ్రాఫ్ట్-టార్గెట్ స్పెక్యులేషన్, మల్టీ-టోకెన్ ప్రిడిక్షన్). లేటెన్సీ కర్నెల్ ప్రారంభ లేటెన్సీకి దగ్గరగా ఉంటుంది.

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

ప్రిఫిల్ మరియు డీకోడ్ కెర్నళ్ల మధ్య సార్వత్రికంగా భిన్నంగా ఉండే పరామితి బ్లాక్‌లోని వార్ప్‌ల సంఖ్య (థ్రెడ్‌ల సంఖ్య, ప్రతి వార్ప్‌కు 32 థ్రెడ్‌ల వార్ప్ పరిమాణంతో భాగించబడుతుంది). ఈ పరామితిని కంపైల్-టైమ్ స్థిరాంకంగా చేయడం వాంఛనీయమైనది, ఎందుకంటే ఇది రిజిస్టర్ కేటాయింపులో సహాయపడుతుంది. CUDAలో, ఇది సాధారణంగా __launch_bounds__ ద్వారా థ్రెడ్‌ల సంఖ్యను పేర్కొనడం ద్వారా మరియు వార్ప్‌ల సంఖ్యను టెంప్లేట్ పరామితిగా అందించడం ద్వారా చేయబడుతుంది:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

ఒక బ్లాక్‌లోని అన్ని థ్రెడ్‌లకు అందుబాటులో ఉన్న రిజిస్టర్‌ల సంఖ్య హార్డ్‌వేర్ పరిమితి; Hopper మరియు Blackwellలో ఇది 64K. థ్రెడ్‌ల సంఖ్య ముందుగానే తెలిసి ఉంటే, కంపైలర్ ప్రతి థ్రెడ్‌కు ఎంత రిజిస్టర్‌లను కేటాయించగలదో ఖచ్చితంగా తెలుసుకుంటుంది మరియు మరింత దూకుడుగా ఆప్టిమైజ్ చేసిన కర్నెల్‌లను కంపైల్ చేయడంలో విజయవంతమవుతుంది. అయితే, దీనివల్ల కంపైలేషన్-సమయ ఖర్చు పెరుగుతుంది, ఎందుకంటే ఇప్పుడు వివిధ warp పరిమాణాల కోసం కర్నెల్‌ను మళ్లీ కంపైల్ చేయాల్సి వస్తుంది. CuTeDSL‌తో, ఇది ఇకపై సమస్య కాదు, ఎందుకంటే మొత్తం కాన్ఫిగరేషన్ స్థలానికి ముందుగానే కంపైలేషన్ ఖర్చును చెల్లించాల్సిన అవసరం లేదు. DSL, CUDAలో __latch_bounds__ సూచించే ఆప్టిమైజేషన్‌ను పరోక్షంగా నిర్వహిస్తూ, బ్లాక్ పరిమాణం ఒక స్థిర వ్యక్తీకరణగా ఉండాలని కూడా ఆశిస్తుంది.

వార్ప్‌ల సంఖ్యను స్థిరాంకంగా చేయడమే కాకుండా, ప్రిఫిల్ మరియు డీకోడ్ మధ్య దాన్ని మార్చడం వాంఛనీయమైనది. GPU యొక్క మెమరీ బ్యాండ్‌విడ్త్‌ను పూర్తిగా ఉపయోగించుకోవడానికి, ప్రతి బ్లాక్‌కు ఎల్లప్పుడూ ఎక్కువ సంఖ్యలో వార్ప్‌లను కేటాయించడం సరిపోదు, ఎందుకంటే అందుబాటులో ఉన్న స్ట్రీమింగ్ మల్టీప్రాసెసర్‌లలో (SMలు) వీలైనన్ని ఎక్కువ ఉపయోగించడం ద్వారా మాత్రమే పూర్తి బ్యాండ్‌విడ్త్‌ను సాధించవచ్చు. ఫలితంగా, ఇచ్చిన టోకెన్‌ల సంఖ్య మరియు తెలిసిన డేటా రకాలు మరియు హిడెన్ డైమెన్షన్‌ల కోసం, రూపొందించబడిన కోడ్ 128-బిట్ లోడ్/స్టోర్ సూచనలను (బ్లాక్‌వెల్‌పై 256-బిట్) ఉపయోగించేలా చూసుకోవడానికి మేము మొదట ఆదర్శవంతమైన పర్-థ్రెడ్ వెక్టర్ డైమెన్షన్‌ను (bfloat16 కోసం 8, float32 కోసం 4) ఎంచుకుంటాము. వెక్టర్ పరిమాణం తెలిసిన తర్వాత, కెర్నల్ ద్వారా అవసరమైన వార్ప్‌ల సంఖ్యను మేము లెక్కించవచ్చు. మేము మొదట SMల సంఖ్యను గరిష్టీకరించడానికి ప్రయత్నిస్తాము: SMల కంటే తక్కువ వార్ప్‌లు ఉంటే, ప్రతి బ్లాక్‌కు ఒకే వార్ప్ ఎంపిక చేయబడుతుంది. ఇది సాధారణంగా కెర్నగ్‌ యొక్క డీకోడ్ వేరియంట్. తగినన్ని వార్ప్‌లు ఉంటే, పనితీరులో మెరుగుదల ఇకపై కనిపించని అప్పర్ బౌండ్‌ను బెం్‌మార్క్ చేస్తూ, మేము క్రమంగా వార్ప్‌ల సంఖ్యను 8, 16 లేదా 32కి పెంచుతాము.

Triton ఈ ఆప్టిమైజేషన్‌లలో కొన్నింటిని అంతర్గతంగా నిర్వహించినప్పటికీ, ఉన్నత-స్థాయి కంపైలర్ యొక్క లోపాల వల్ల పరిమితం కాకుండా, CuTe లేఅవుట్ ఆల్జెబ్రాను జాగ్రత్తగా ఉపయోగించడం ద్వారా మేము ఇలాంటి ఫలితాలను సాధించగలమని కనుగొన్నాము. ముఖ్యంగా sub-warp హారిజాంటల్ రిడక్షన్‌లు ఉన్నప్పుడు, ఉత్తమ కోడ్‌ను రూపొందించడానికి అవసరమైన సున్నితమైన నియంత్రణను CuTeDSL మాకు అందించిందని మేము కనుగొన్నాము.

కొన్ని సందర్భాల్లో, వీలైనన్ని ఎక్కువ బ్లాక్‌లలో పనిని పంపిణీ చేయడానికి, పాల్గొనే అన్ని బ్లాక్‌లను సింక్రొనైజ్ చేయడానికి cooperative_groups::this_grid().sync() ఉపయోగించవచ్చు. పరిమితులు వర్తిస్తాయి: ఎగ్జిక్యూషన్ ప్రారంభమైన తర్వాత ఒక బ్లాక్ SM నుండి డి-షెడ్యూల్ చేయబడదు కాబట్టి, ఫార్వర్డ్ ప్రోగ్రెస్‌కు హామీ ఇవ్వడానికి, ఏ సమయంలోనైనా అన్ని బ్లాక్‌లు సింక్రొనైజేషన్‌లో పాల్గొన్న గ్లోబల్ మెమరీ స్థానాలను అప్‌డేట్ చేయగలవని మరియు పోల్ చేయగలవని నిర్ధారించడానికి, బ్లాక్‌ల సంఖ్య పరికరంలో అందుబాటులో ఉన్న SMల సంఖ్యను మించకూడదు. ఈ పరిమితి cudaLaunchCooperativeKernel ద్వారా అమలు చేయబడుతుంది. హోప్పర్‌పై గ్రిడ్ బారియర్ ఖర్చు సుమారు 3us, కొన్ని సందర్భాల్లో అంతర్లీన మెమరీ బారియర్ ద్వారా మునుపటి మెమరీ ఆపరేషన్‌లన్నింటినీ ఫ్లష్ చేసే ఖర్చుతో కూడుకున్నది. మొత్తంమీద, ఇది కెర్నగ్‌ను రెండు లాంచ్‌లుగా విభజించడం యొక్క ఓవర్‌హెడ్‌తో పోల్చదగినది, ఇది పనితీరు-సున్నితమైన డిజైన్ నిర్ణయానికి దారితీస్తుంది. ప్రోగ్రామాటిక్ డిపెండెంట్ లాంచ్ (PDL) హోప్పర్ నుండి లాంచ్ ఓవర్‌హెడ్‌లను తగ్గించగలదనే వాస్తవం ఉన్నప్పటికీ, డీకోడ్ కెర్నళ్ల కోసం, ఓవర్‌హెడ్‌ను నివారించడం విలువైనది.

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

CuTeDSLలో, ఈ నిర్ణయాలను host-side launch codeలో తీసుకోవచ్చు, మరియు అవసరమైనప్పుడు kernels‌ను సంయోజించి grid barriers‌తో మధ్యలో కలిపి అమలు చేయవచ్చు, క్రింది ఉదాహరణలో చూపినట్లు. ఇలాంటి హైబ్రిడ్ వ్యూహంతో, డీకోడ్ kernels‌ను కనిష్ట లేటెన్సీకి చేరేలా ఆప్టిమైజ్ చేయవచ్చు, మరియు prefill kernels‌ను గరిష్ఠ throughput కోసం ఆప్టిమైజ్ చేయవచ్చు. CUDA-ఆధారిత పరిష్కారాల కంటే, ఈ లాజిక్‌ను మరింత సంక్షిప్తంగా మరియు తక్కువ కంపైలేషన్-సమయ ఓవర్‌హెడ్‌తో నిర్వచించవచ్చు.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

క్రింద వివరించిన మా MoE all-to-all kernels సందర్భంలో, grid barriers‌తో లేదా వాటి లేకుండా kernels‌ను నిర్మించి, వాటిని compilation సమయంలో ఆప్టిమైజ్ చేసే లవచికత ఉండడం వల్ల CuTeDSL ద్వారా గణనీయమైన పనితీరు మెరుగుదలలను సాధించగలిగాము.

మేము CuTeDSLలో నిర్మించిన మొదటి kernel చిన్న డైమెన్షన్‌ల కోసం ఆప్టిమైజ్ చేసిన ఒక ప్రత్యేక RMS norm kernel. తాజా Qwen మోడళ్ల self-attention layerలో ఉపయోగించే QK norm, ప్రతి head ఆధారంగా 64, 128 లేదా 256 ఎలిమెంట్‌లపై వర్తించబడుతుంది. వేలల్లో ఉన్న hidden dimensions‌పై పనిచేస్తూ, shared memory మరియు registers రెండింటి ద్వారా horizontal reduction అవసరమయ్యే సాధారణ RMS norm‌కు విరుద్ధంగా, QK norm‌కు రిడక్షన్ warps అంతటా మాత్రమే అవసరం, అలాగే బహుళ heads‌ను ఒకే warpలో ఒదిగించే అవకాశం కూడా ఉంటుంది. ఇది అన్ని head dimensions మరియు data types కోసం ఉత్తమ token-to-warp assignments‌ను కనుగొనడానికి compile timeలో దూకుడుగా ఆప్టిమైజ్ చేయగల వినియోగ సందర్భంగా మారింది.

RMS నార్మ్ అనేది చాలా సులభమైన ఆపరేషన్, అయినప్పటికీ వెక్టర్ యొక్క వర్గీకరించబడిన నార్మ్‌ను లెక్కించడానికి దీనికి హారిజాంటెల్ మొత్తం అవసరం:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

అటెన్షన్ హెడ్‌లు సాధారణంగా చిన్నవిగా ఉంటాయి మరియు పర్-థ్రెడ్ ఆపరేషన్‌లు తీవ్రంగా వెక్టరైజ్ చేయబడతాయి కాబట్టి, ఒకే వార్ప్ స్క్వేర్ నార్మ్‌ను లెక్కించవచ్చు. మెమరీ బ్యాండ్‌విడ్త్‌ను గరిష్టీకరించడానికి, ప్రతి థ్రెడ్ 128 బిట్ల ఇన్‌పుట్‌ను ప్రాసెస్ చేయాలి, ఇది 8 bfloat16 విలువలకు అనువదిస్తుంది. ఫలితంగా, ఒక హెడ్‌ను ప్రాసెస్ చేయడానికి హాఫ్-వార్ప్ సరిపోతుంది, మరియు shuffle_sync_bfly ద్వారా నిర్వహించబడే హారిజాంటెల్ తగ్గింపును ఉపయోగించి స్క్వేర్డ్ నార్మ్‌ను నిర్ణయించవచ్చు. హెడ్ డైమెన్షన్ మరియు డేటా రకం నుండి ప్రారంభించి, కంపైలేషన్ సమయంలో ప్రతి వార్ప్‌కు హెడ్‌ల సంఖ్యను మేము నిర్ణయించవచ్చు మరియు ఆర్గుమెంట్‌లు మరియు తగ్గింపు రౌండ్‌ల సంఖ్యను పేర్కొనవచ్చు.

ఆటోమేటిక్‌గా వెక్టరైజ్ చేయబడిన లోడ్స్/స్టోర్స్‌పై ఆధారపడి మరియు వార్ప్ తగ్గింపు కోసం లోయర్-లెవెల్ ఆర్కిటెక్చర్-నిర్దిష్ట ఆపరేషన్‌లను ఉపయోగించి మేము సాపేక్షంగా కొన్ని కోడ్ లైన్‌లతో ప్రత్యేకమైన కెర్నగ్‌ను నిర్మించవచ్చు. కెర్నల్ ఏ హెడ్ (Q లేదా K)పై పనిచేయాలో నిర్ణయిస్తుంది మరియు వర్తింపజేయడానికి తగిన వెయిట్‌లను గుర్తిస్తుంది. ఇది హెడ్ నుండి మొత్తం వెక్టర్‌ను లోడ్ చేస్తుంది మరియు మొదట స్థానికంగా మొత్తం చేస్తుంది. తదనంతరం, ప్రతి హెడ్‌కు థ్రెడ్‌ల సంఖ్య ఆధారంగా, స్క్వేర్డ్ నార్మ్ లెక్కించబడుతుంది. చివరిగా, ఆటోమేటిక్‌గా వెక్టరైజ్ చేయబడిన RMS నార్మ్ నిర్ణయించబడుతుంది మరియు మెమరీకి తిరిగి వ్రాయబడుతుంది. WEIGHT_BIAS మరియు EPS వంటి మోడల్-నిర్దిష్ట పరామితులు స్థిరంగా మడవబడతాయి.

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

రెగ్యులర్ RMS నార్మ్ కెర్నగ్‌తో పోలిస్తే, ఈ ప్రత్యేకమైన వేరియంట్ ప్రిఫిల్ మరియు డీకోడ్ రెండింటికీ 2-3x మెరుగైన త్రూపుట్ మరియు లేటెన్సీని అందిస్తుంది.

MoE డిస్పాచ్ మరియు కంబైన్ ఇన్‌పుట్ టోకెన్‌లను అవి రూట్ చేయబడిన నిపుణులకు షఫుల్ చేస్తాయి, MoE GEMM కెర్నగ్‌లు సంబంధిత మ్యాట్రిక్స్ గుణకారాలను నిర్వహించిన తర్వాత వెయిటెడ్ సంచితాన్ని నిర్వహిస్తాయి. వివిధ ఉత్పత్తులు వివిధ పరిమాణాల MoE మోడళ్లపై ఆధారపడటం వలన, పనితీరు ప్రత్యేక ఆసక్తిని కలిగి ఉంది. మోహరింపు రకాల ఆధారంగా, ఆప్టిమైజ్ చేయడానికి మేము వివిధ కెర్నళ్లను వేరు చేస్తాము:

TP=1: ఒకే GPU కోసం స్థానిక కెర్నల్. డిస్పాచ్ టోకెన్‌లను సంబంధిత నిపుణులకు షఫుల్ చేస్తుంది మరియు కంబైన్ వాటిని సంచితం చేస్తుంది. H100పై Qwen3 30B మోహరింపుల కోసం ప్రాథమికంగా ఉపయోగించబడుతుంది.

DP=1, TP ≠ 1: టెన్సర్-పారలల్, సింగిల్‌నోడ్ కెర్నల్. ఇన్‌పుట్ ప్రతిరూపం చేయబడుతుంది, డిస్పాచ్ ప్రస్తుత ర్యాంక్ కోసం టోకెన్‌లను ఎంచుకుంటుంది మరియు కంబైన్ వాటిని అన్ని ర్యాంక్‌లలో సంచితం చేస్తుంది. NVLink లేదా NVLS అంతటా సంచితం. పెద్ద ప్రిఫిల్లర్ల కోసం ఉపయోగించబడుతుంది.

DP ≠ 1, NVLink: ఐచ్ఛిక టెన్సర్ సమాంతరతతో డేటా-పారలల్ కెర్నల్. ప్రతి ర్యాంక్ నిర్దిష్ట నిపుణులను హోస్ట్ చేసే ప్రతి ఇతర ర్యాంక్‌కు టోకెన్‌లను పంపుతుంది. రూటింగ్ సమాచారం సమగ్రపరచబడుతుంది, టోకెన్లు NVIDIA NVLink అంతటా బదిలీ చేయబడతాయి. డేటా-పారలల్ H200 లేదా GB200 మోహరింపుల కోసం ఉపయోగించబడుతుంది.

DP ≠ 1, InfiniBand: NVLink మరియు NVIDIA క్వాంటమ్ ఇన్‌ఫినిబ్యాండ్ రెండింటిలోనూ పనిచేసే పీర్-టు-పీర్ డేటా-పారలల్ మరియు టెన్సర్-పారలల్ కెర్నల్. మల్టీ-నోడ్ హోప్పర్ మోహరింపుల కోసం ఉపయోగించబడుతుంది.

స్థానిక kernel routing table ఆధారంగా ప్రతి expert‌కు కేటాయించబడిన tokenల సంఖ్యను లెక్కించి, ప్రతి expert‌కు చెందిన tokenల జాబితాలోని ప్రతి tokenకు ప్రత్యేక index‌ను కేటాయిస్తుంది. రౌటింగ్ దశ తర్వాత, tokenలు MoE GEMM kernel‌కు పంపించడానికి ఐచ్చికంగా padded buffers‌లోకి కాపీ చేయబడతాయి. token స్థానాలను ఉపయోగించి, combine kernel తరువాత registers‌లో accumulation నిర్వహించి, తుది tokenలను output buffer‌లో వ్రాస్తుంది.

ప్రిఫిల్ మరియు డీకోడ్ మధ్య డిస్పాచ్ దశ మారుతూ ఉంటుంది: కొన్ని టోకెన్‌ల కోసం, చదవడానికి కొద్దిగా రూటింగ్ సమాచారం ఉంటుంది, కాబట్టి మేము దీన్ని ఒకే బ్లాక్‌లో ప్రాసెస్ చేస్తాము మరియు గ్రిడ్ మిగిలిన భాగానికి సిగ్నల్ చేస్తాము, వ్యక్తిగత బ్లాక్‌లు టోకెన్‌ల చంక్‌లను కాపీ చేస్తాయి. షేర్డ్ మెమరీలో కౌంటర్‌లను అటామిక్‌గా ఇంక్రిమెంట్ చేయడం ద్వారా టోకెన్ అసైన్‌మెంట్ నిర్ణయించబడుతుంది. ప్రిఫిల్ మరింత రూటింగ్ సమాచారాన్ని ప్రాసెస్ చేయాలి కాబట్టి, వేర్వేరు కెర్నల్‌లో బహుళ బ్లాక్‌లలో రూటింగ్ విభజించబడుతుంది. షేర్డ్ మెమరీలోని కౌంటర్‌లను ఉపయోగించి సూచికలు మొదట స్థానికంగా కేటాయించబడతాయి మరియు ప్రతి టోకెన్ స్థానాన్ని నిర్ణయించడానికి గ్లోబల్ మెమరీలోని అన్ని బ్లాక్‌లలో తదనంతరం సమగ్రపరచబడతాయి. గ్రిడ్-వైజ్ బారియర్‌లను ఉపయోగించగలగడానికి ఈ కెర్నల్ గ్రిడ్ SMల సంఖ్యకు బిగించబడుతుంది. తదుపరి కెర్నల్ కాపీలను నిర్వహిస్తుంది, అందుబాటులో ఉన్న మెమరీ బ్యాండ్‌విడ్త్‌ను పూర్తిగా ఉపయోగించుకోవడానికి చాలా ఎక్కువ సంఖ్యలో బ్లాక్‌లు మరియు వార్ప్‌లలో లాంచ్ చేయబడుతుంది.

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

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

ranks సంఖ్య సాధారణంగా token రూట్ చేయబడే experts సంఖ్య కంటే తక్కువగా ఉన్నందున, సగటు సందర్భంలో ప్రతి token యొక్క ప్రతులలో కనీసం ఒకటి ప్రతి rank‌కు రూట్ చేయబడుతుందని అనుకోవచ్చు. ఈ పరిస్థితిలో, స్థానిక combine kernel‌ను token యొక్క స్థానిక కాపీలను సమీకరించడానికి మళ్లీ ఉపయోగించవచ్చు; ఆ తర్వాత పాల్గొనే ranks అంతటా మొత్తాన్ని కలిపేందుకు ఒక all-reduce జరుగుతుంది. prefill kernel సందర్భంలో, మేము వేర్వేరు kernels‌ను ప్రారంభిస్తాము; అయితే decode ఈ రెండు ఆపరేషన్‌లను ఒకే launchలో విలీనం చేస్తుంది.

all-reduce కోసం, కేవలం రెండు ranks పాల్గొంటే మేము సాధారణ NVLink‌ను ఉపయోగిస్తాము; 4 లేదా అంతకంటే ఎక్కువ ranks పాల్గొంటే, throughputలో 10-20% పెరుగుదల కోసం NVLS‌కు మారుతాము. DSLకు ధన్యవాదాలు, MoE మరియు సాధారణ all-reduce kernels మధ్య reduction logicలో ఎక్కువ భాగాన్ని, అలాగే peer buffer exchanges‌ను పంచుకోగలుగుతాము. ఇది హార్డ్‌వేర్‌పై సన్నిహిత నియంత్రణ ప్రయోజనకరమైన ఒక పరిస్థితి, ఎందుకంటే NVLS కోసం multimem.ld_reduce మరియు multimem.st పై DSL ఆపరేషన్‌లు అవసరం. reduction యొక్క అమలు ప్రధానంగా NCCL tests‌లోని all_reduce ఉదాహరణను అనుసరిస్తుంది.

డేటా పారలల్ కెర్నళ్లు బహుళ పరికరాల కోసం టోకెన్‌లను సమగ్రపరచడానికి NVLinkపై ఆధారపడతాయి. డిస్పాచ్ దశలో, టోకెన్‌లను మార్పిడి చేయడానికి ముందు, పీర్లు తాము పంపుతున్న ప్రతి టోకెన్ కోసం ప్రతి ర్యాంక్‌పై లక్ష్య ఆఫ్‌సెట్‌ను గుర్తించడానికి మొదట స్థానికంగా రూట్ చేస్తాయి. తదనంతరం, వైరుధ్యాలను నివారించడానికి ప్రతి పంపేవారికి కేటాయించబడిన రిసీవ్ బఫర్‌లలోకి టోకెన్‌లను కాపీ చేయడం ప్రారంభిస్తాయి. పూర్తి కావడాన్ని సిగ్నల్ చేయడానికి పర్-టోకెన్ రిలీజ్-స్టోర్ రైట్‌లు ఉపయోగించబడతాయి. రిసీవర్ ముగింపులో, MoE GEMM కోసం అన్ని టోకెన్‌లను ఒకే బఫర్‌లోకి కండెన్స్ చేయగలిగేలా రూటింగ్ సమాచారం సమగ్రపరచబడుతుంది. రైట్‌లు పూర్తి కావడానికి వేచి ఉండటానికి రిసీవర్‌లు అక్వైర్-లోడ్‌తో పీర్-రాయబడిన ఫ్లాగ్‌ను పోల్ చేస్తాయి.

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

ఈ కెర్నగ్‌ను అమలు చేయడానికి సిస్టమ్ మరియు డివైజ్ స్థాయిలో స్కోప్ చేయబడిన రిలీజ్-అక్వైర్ లోడ్‌లు మరియు స్టోర్‌లను ప్రదర్శించడానికి మేము PTX మోస్తర్ల లైబ్రరీని నిర్మించాము. సింక్రొనైజేషన్ ఫ్లాగ్‌ల గ్రాన్యులారిటీని మెరుగుపరచడం ద్వారా మరియు NVLink ద్వారా టోకెన్‌లను పంపడానికి మరియు స్వీకరించడానికి మరిన్ని బ్లాక్‌లను కేటాయించడం ద్వారా, మా మునుపటి CUDA కెర్నళ్లతో పోలిస్తే మేము 2x మెరుగుదలలను సాధించాము.

మా పీర్-టు-పీర్ కెర్నళ్లు గతంలో ConnectX-7 ఇన్‌ఫినిబ్యాండ్ అడాప్టర్‌లతో జత చేయబడిన హోప్పర్ GPUలపై స్టేట్-ఆఫ్-ది-ఆర్ట్ పనితీరును సాధించాయి. పనితీరులో మరింత లాభాల కోసం మేము అప్పటి నుండి కంబైన్ కెర్నళ్లను CuTeDSLకి పోర్ట్ చేసాము. CUDA వెర్షన్‌లు ఇప్పటికే కంపైల్-టైమ్ స్పెషలైజేషన్‌పై తీవ్రంగా ఆధారపడ్డాయి కాబట్టి, అవి మా రిపోజిటరీలో కంపైలేషన్-టైమ్ బాటిల్‌నెక్స్ అని నిరూపించబడ్డాయి, కొన్ని పరామితులను మినహాయించమని మాకు బలవంతం చేశాయి. JIT-కంపైల్ చేయబడిన DSLకి మైగ్రేట్ చేయడం ద్వారా, మేము మరిన్ని పరామితుల కోసం ప్రత్యేకతను సాధించాము మరియు గ్రిడ్ కేటాయింపులను ట్వీక్ చేసాము, లేటెన్సీకి దాదాపు 3% బూస్ట్ కోసం ప్రిఫిల్ మరియు డీకోడ్ మధ్య తదుపరి వ్యత్యాసాలను పరిచయం చేసాము.

ROSE కెర్నళ్ల కోణం నుండి, బ్లాక్‌వెల్ మద్దతు రెండు వర్గాలుగా వస్తుంది: ఆర్కిటెక్చర్-అజ్ఞాస్టిక్ మరియు ఆర్కిటెక్చర్-నిర్దిష్ట కెర్నళ్లు. టెన్సర్ కోర్లపై ఆధారపడే ఆర్కిటెక్చర్-నిర్దిష్ట కెర్నళ్లకు అనుగుణంగా స్పష్టమైన పని అవసరం. అయినప్పటికీ, DSL వాడకానికి ధన్యవాదాలు, ఆర్కిటెక్చర్-అజ్ఞాస్టిక్ కెర్నళ్లు సులభంగా అనుగుణంగా ఉంటాయి.

బ్లాక్‌వెల్ GPU హోప్పర్ యొక్క 132తో పోలిస్తే గరిష్టంగా 160 SMలను కలిగి ఉంది. చాలా కెర్నళ్లకు, ఇది కంపైల్-టైమ్ స్థిరాంకం మరియు రన్ చేయడానికి ఆప్టిమైజడ్ (ప్రిఫిల్ vs డీకోడ్) కెర్నగ్‌ను నిర్ణయించడానికి హోస్ట్-సైడ్ లాంచర్ కోడ్‌లోని లాజిక్ దీన్ని కారకం చేస్తుంది. అదనంగా, బ్లాక్‌వెల్ 256-బిట్ లోడ్ సూచనలను కూడా పరిచయం చేస్తుంది, మెమరీ బ్యాండ్‌విడ్త్ వినియోగాన్ని మెరుగుపరచడానికి వెక్టర్ వెడల్పును (CuTe టెన్సర్‌ల చివరి డైమెన్షన్) పెంచడం ద్వారా దీన్ని సులభంగా విడుదల చేయవచ్చు.

బ్లాక్‌వెల్ గణనీయమైన సంఖ్యలో క్వాంటైజేషన్ స్కీమ్‌ల (MXFP8, MXFP4, NVFP4) కోసం హార్డ్‌వేర్ మద్దతును పరిచయం చేస్తుంది. FP8 బ్లాక్-వైజ్ స్కేలింగ్ కోసం ఇప్పటికే ఉన్న కెర్నళ్లను ట్వీక్ చేయడం ద్వారా వీటి కోసం మద్దతు ప్రాథమికంగా అమలు చేయబడింది. కంపైలేషన్-టైమ్ కాన్ఫిగరేషన్ స్పేస్‌ను పేల్చకుండా కొత్త మోడ్‌లకు మద్దతు ఇవ్వగలిగేలా JIT మద్దతు చాలా కీలకం. గతంలో, కొన్ని మూలకాల అంతటా, వార్ప్‌లో సగం లేదా కాల్ భాగం కంటే ఎక్కువ హారిజాంటెల్ తగ్గింపులను వ్యక్తపరచడంలో Triton ఇబ్బందికరంగా ఉందని మేము కనుగొన్నాము. NVFP4 మరియు MXFP8/MXFP4 కోసం వరుసగా 16 లేదా 32 మూలకాల అంతటా సంపూర్ణ గరిష్ట విలువను లెక్కించేటప్పుడు సరైన పనితీరును సాధించడానికి తగ్గింపులపై మరింత చక్కటి నియంత్రణ చాలా కీలకం.

సాపేక్షంగా కొత్త భాషగా, LLMలు తమ ప్రీ-ట్రైనింగ్ సమయంలో CuTeDSL కోడ్‌కు తక్కువ బహిర్గతం అయ్యాయి. అదనంగా, TMA మరియు ఆర్కిటెక్చర్-నిర్దిష్ట మ్యాట్రిక్స్ గుణకార సూచనల కోసం తీవ్రంగా ఆప్టిమైజ్ చేయబడిన ప్రస్తుత CUTLASS లేదా Quack ఉదాహరణల నుండి, యాక్టివేషన్‌లు మరియు నార్మ్‌ల కోసం ఆప్టిమైజ్ చేయబడిన కెర్నళ్లతో మా ప్రస్తుత కోడ్ శైలి కొద్దిగా భిన్నంగా ఉంది. కెర్నల్ రాయడంలో LLMలు సహాయపడే బిందువును చేరుకోవడం క్రమమైన ప్రక్రియ.

ఏజెంట్‌లపై కెర్నల్ మార్పిడి కోసం ఆధారపడే ముందు, స్థిరమైన నమూనాలను అనుసరించి, కెర్నళ్లు మరియు యుటిలిటీ ఫంక్షన్‌ల ప్రారంభ కార్పస్‌ను మేము మాన్యువల్‌గా నిర్మించాము. కార్పస్ పురోగమిస్తున్న కొద్ది, ముఖ్యమైన మాన్యువల్ రీవర్క్ అనుసరిస్తూ, CUDA నుండి CuTeDSLకి రఫ్ అనువాదాలను మేము నిర్వహించగలిగాము. ప్రత్యేకించి, మోడళ్లు పాయింటర్-టు-టెన్సర్ మార్పిడిలు మరియు వెక్టరైజేషన్‌తో కష్టపడ్డాయని, పరీక్షను కంపైల్ చేయడానికి ప్రయత్నిస్తున్నప్పుడు టైప్ లోపాలను పరిష్కరించడానికి ప్రయత్నిస్తూ ఇరుక్కుపోయాయని మేము కనుగొన్నాము. కెర్నళ్ల గ్రిడ్-ఆధారిత స్వభావం గురించి అవి చాలా తక్కువ అవగాహనను చూపించాయి, థ్రెడ్ మరియు గ్రిడ్ సింక్రొనైజేషన్‌ను గందరగోళానికి గురిచేశాయి. అత్యంత నిరాశపరిచే విషయం ఏమిటంటే, మా స్వంత CuTeDSL భర్తీలకు వ్యతిరేకంగా బిల్డ్ చేయడానికి బదులుగా CUDA-నిర్దిష్ట లైబ్రరీల ఉనికిలో లేని అనువాదాలను భ్రమపడే ప్రలోభం ఉంది. కాలక్రమేణా, CuTeDSL కెర్నళ్ల కార్పస్ పెరిగిన కొద్ది, LLMలు చాలా వరకు స్కెలిటన్ కోడ్‌ను నిర్మించగలవని మరియు ఖాళీలను పూరించగలవని మేము కనుగొన్నాము, హై-లెవెల్ డిజైన్‌పై దృష్టి పెట్టడానికి మాకు అనుమతిస్తుంది.

CUDAతో పోలిస్తే, CuTeDSL చాలా మెరుగైన కంపైలేషన్-టైమ్ పనితీరును అందిస్తుందని మేము కనుగొన్నాము, ఇది మరిన్ని పరామితుల కోసం ప్రత్యేకతను సాధించడానికి మరియు పనితీరు మెరుగుదలలను సాధించడానికి మాకు అనుమతిస్తుంది. Tritonతో పోలిస్తే కొన్ని అబ్‌స్ట్రాక్షన్‌లు లేకపోవడం వల్ల కెర్నళ్లు కొద్దిగా ఎక్కువ వెర్బోస్‌గా ఉన్నప్పటికీ, హార్డ్‌వేర్‌పై మరింత చక్కటి నియంత్రణ వేగవంతమైన కెర్నళ్లను నిర్మించడానికి మాకు అనుమతిస్తుంది కాబట్టి సంక్లిష్టత సమర్థించబడిందని మేము కనుగొన్నాము.

రాసే సమయంలో, CUTLASS C++ మరియు CuTeDSL కెర్నళ్ల మధ్య పనితీరు అంతరం కనిష్టంగా ఉంది, ఇంటర్‌డెవైస్ కమ్యూనికేషన్, నార్మలైజేషన్ లేదా యాక్టివేషన్‌లతో హై-పెర్ఫార్మెన్స్ GEMM కెర్నళ్లను కూర్చడానికి మరియు విలీనం చేయడానికి ఉత్తేజకరమైన అవకాశాలను తెరుస్తుంది. ఈ DSLకి కట్టుబడి ఉండటం ద్వారా, ఎంబెడ్డింగ్‌ల నుండి LLMల వరకు మా వర్క్‌లోడ్‌ల అంతటా త్రూపుట్‌ను మెరుగుపరచడానికి మా పరికరాల నుండి మేము మరింత పనితీరును వెలికితీయగలుగుతాము.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AIలో ఇన్‌ఫరెన్స్

CuTeDSL ఎందుకు?

డెవలపర్ అనుభవం

పెద్ద కాన్ఫిగరేషన్ స్పేస్

ప్రిఫిల్ మరియు డికోడ్ స్పెషలైజేషన్

గ్రిడ్ సింక్రొనైజేషన్

CuTeDSL కెర్నళ్లు

QK నార్మ్

మిక్చర్-ఆఫ్-ఎక్స్‌పర్ట్స్ డిస్పాచ్/కంబైన్

స్థానిక

సింగిల్‌నోడ్

డేటా పారలల్

పీర్-టు-పీర్

బ్లాక్‌వెల్ మద్దతు

AI-సహాయక కెర్నల్ జనరేషన్

ముగింపు మరియు భవిష్యత్తు పని

ప్రస్తావనలు

ఉదాహరణ