CuTeDSL at Perplexity

Perplexity मध्ये, आम्ही एम्बेडिंग्जपासून ते ट्रिलियन पॅरामीटर एलएलएमपर्यंतच्या मॉडेलना सेवा देण्यासाठी इन-हाउस रनटाइम-ऑप्टिमाइज्ड सर्व्हिंग इंजिन (ROSE) वापरतो. Sonar, Search आणि Embeddings सह सर्व API च्या मागे कार्यरत असलेले आमचे रनटाइम, संशोधकांना सर्वाधिक योग्य मॉडेल आणि आर्किटेक्चर शोधण्यासाठी आणि सर्वोत्तम उमेदवारांना प्रोडक्शनमध्ये त्वरीत तैनात करण्यासाठी सक्षम बनवले गेले आहे. हे साध्य करण्यासाठी, आम्हाला अत्याधुनिक आर्किटेक्चरसाठी अत्यंत वेगाने इन्फरन्स लागू करण्याची, तैनात करण्याची आणि ऑप्टिमाइझ करण्याची क्षमता असणे आवश्यक आहे. या ब्लॉग पोस्टमध्ये, इन्फरन्स इंजिनमध्ये समाकलित केलेले CuTeDSL कर्नल NVIDIA GPU वर मॉडेलना त्यांच्या सर्वोत्तम कामगिरीवर आणण्यासाठी आवश्यक जीपीयू कर्नल तयार करण्यास कसे सक्षम करतात, याचा आम्ही शोध घेतो.

Perplexity मधील विविध टीम्सद्वारे तयार केलेले कस्टम मॉडेल्स, AI इन्फरन्स टीमने तयार केलेले ROSE इन्फरन्स इंजिन वापरून NVIDIA Hopper आणि Blackwell GPU वर इन-हाउस होस्ट केले जातात. NVIDIA Triton इन्फरन्स सर्व्हरशी सुसंगत इंटरफेससह, लँग्वेज डिकोडिंग आणि क्लासिफिकेशन या दोन्हीसाठी सानुकूलित Llama मॉडेल्सना अधिक चांगल्या प्रकारे सेवा देण्यासाठी ROSE सुरुवातीला तयार केले गेले होते. कालांतराने, Perplexity च्या उत्पादनांना पॉवर देण्यासाठी विस्तृत टीम्ससाठी रँकिंग, क्लासिफिकेशन, स्कोअरिंग आणि इतर ट्रान्सफॉर्मर-आधारित मॉडेल्स तसेच महत्त्वपूर्ण LLM ला सेवा देणारे एक पूर्ण-विकसित इन्फरन्स इंजिन म्हणून ते विकसित झाले. ROSE चे डिझाइन ध्येय सोपे असणे हे आहे: कालांतराने, सर्वोच्च कार्यक्षमता गाठण्यासाठी प्रयत्न केले जातात, परंतु मुख्य प्राधान्य आमच्या मॉडेल्ससाठी सतत विकसित होणाऱ्या आर्किटेक्चर आणि युज केसेसशी जुळवून घेणे हे आहे.

डेटा प्रकार, हेड डायमेंशन्स, कंपाइल-टाइम कॉन्स्टंट्स, आणि लॉन्च सेटिंग्जमध्ये कर्नल कॉन्फिगरेशन स्पेस.

उच्च पातळीवर, ROSE मध्ये एक इंजिन असते जे मॉडेलला अशा इंटरफेसशी जुळवून घेते ज्याच्याशी क्लायंट संवाद साधू शकतात. हे इंजिन अत्यंत पुनर्वापर करण्यायोग्य (reusable) आहेत, जे रिक्वेस्ट शेड्युलिंग आणि बॅचिंग हाताळतात आणि अंतर्निहित मॉडेल्समधून डिकोड केलेले टोकन्स, एम्बेडिंग्ज किंवा स्कोअर पुनर्प्राप्त करतात. डिव्हाइसेसचे इनिशिलायझेशन, इंटर-प्रक्रिया आणि इंटर-नोड कम्युनिकेशन चॅनेल आणि वजन लोड करणे हे सर्व त्यांच्याद्वारे हाताळले जाते. LLM साठी, इंजिन बॅचिंग, चँकिंग, सॅम्पलिंगसाठी जबाबदार आहे. हे संपूर्ण आणि रेखीय लक्ष (linear attention) तसेच प्रिफिक्स मॅचिंग या दोन्हीसाठी केव्ही स्टोरेज वाटपासाठी देखील जबाबदार आहे. एम्बेडिंग इंजिन सोपी आहेत, कारण ते फक्त ऑनलाइन बॅचिंग करतात. जेव्हा मॉडेल एकाच वेळी ॲक्सिलेटरवर चालू असते तेव्हा इंजिन पुढील फेरीच्या कामाची समांतर तयारी करेल अशा प्रकारे अंमलबजावणी ऑप्टिमाइझ केली जाते.

ROSE CuTeDSL, Triton, CUDA, CUTLASS, आणि cuBLAS सह तयार केलेले GPU कर्नल रॅप करून मॉडेलिंगसाठी सानुकूल थरांची एक समृद्ध श्रेणी सादर करते. मॅट्रिक्स मल्टिप्लिकेशन आणि अटेंशन इम्प्लीमेंटेशनसाठी, आम्ही प्रामुख्याने CUTLASS आणि cuBLAS मधील NVIDIA कर्नलवर अवलंबून असतो. एम्बेडिंग, नॉर्म, MoE आणि ॲक्टिव्हेशन कर्नलसह पाइपलाइनचे उर्वरित भाग, ROSE च्या कर्नल लायब्ररीचा भाग आहेत. आम्ही ज्या हार्डवेअरवर इन्फरन्स चालवतो त्याला सर्वोत्तम समर्थन देण्यासाठी, ROSE MoE रुटिंग, डिस्पॅच आणि कम्बाइनसाठी स्वतःच्या इंटर-डिव्हाइस कम्युनिकेशन प्रिमिटिव्ह्ससह सुसज्जित आहे. इतर कर्नल तुलनेने सोपे असले तरी, ते चालत असलेल्या सर्व वेगवेगळ्या लपविलेल्या परिमाणांवर (hidden dimensions) आणि डिव्हाइसेसवर सर्वोच्च कार्यक्षमता साध्य करण्यासाठी त्यांना सखोलपणे स्पेशलाइज्ड केले जाणे आवश्यक आहे. याव्यतिरिक्त, विविध मॉडेल्सचे सर्व नॉब्स आणि ट्विक्स सामावून घेण्यासाठी आम्ही कंपाईल-टाइम स्पेशलायझेशनवर अवलंबून असतो (RMS नॉर्मसाठी वजन बायस विरुद्ध कोणताही वजन बायस नाही, ॲक्टिव्हेशन फंक्शन्स इ.). परिणामी, आम्ही त्यांना लागू करण्यासाठी DSL वापरण्याचा निर्णय घेतला: आम्ही सुरुवातीला Triton वापरून ROSE तयार केले. तथापि, आम्ही तेव्हापासून GPU प्रोग्रामिंगसाठी प्राथमिक वातावरण म्हणून CuTeDSL स्वीकारले आहे.

गणितीय दृष्टिकोनातून, इन्फरन्स कर्नल सामान्यत: अनपेक्षित असतात. परिणामी, कर्नल विकसित करताना, गुंतागुंत सामान्यत: कोणती गणना करायची हे व्यक्त करण्यात नसते, तर हार्डवेअरने ते नेमके कसे करावे यात असते. सर्वोत्तम कर्नल तयार करण्यासाठी, आम्हाला अशा भाषेची आवश्यकता आहे जी अंतर्निहित हार्डवेअरवर योग्य ॲब्स्ट्रॅक्शन प्रदान करते. योग्य भाषा आम्हाला गणना करण्यासाठी उच्च-स्तरीय ऑपरेशन निर्दिष्ट करण्यास सक्षम करते, त्यांना विस्तृत कॉन्फिगरेशन स्पेस मधील सर्व बिंदूंवर आक्रमकपणे स्पेशलाइज्ड आणि ऑप्टिमाइझ करते.

CuTeDSL हे CuTe लेआउट बीजगणित आणि MLIR च्या शीर्षस्थानी तयार केलेल्या NVIDIA GPU साठी Python-आधारित DSL आहे. लो-लेव्हल हार्डवेअर प्रिमिटिव्ह्सवर नियंत्रण राखून ते अत्यंत ऑप्टिमाइझ केलेल्या PTX वर जस्ट-इन-टाइम संकलित (compiles just-in-time) होते. विकसन गतीला बाधा आणणाऱ्या महागड्या कंपायलेशन वेळेशिवाय, CUTLASS C++ प्रमाणेच वैशिष्ट्ये आणि कार्यक्षमता προσφέρει ऑफर करण्याचे त्याचे उद्दिष्ट आहे. आम्ही सर्वेक्षण केलेल्या किंवा काम केलेल्या भाषांपैकी (CUDA, cuTile, TVM IR आणि Triton सह), आम्ही आमच्या कर्नलच्या विकासासाठी हे सर्वात योग्य मानतो.

एकंदरीत, Triton किंवा cuTile इतके उच्च-स्तरीय आणि संक्षिप्त नसले तरी, आम्हाला आढळले की CuTe लेआउट बीजगणित आणि त्याचे कंपोझिशन्स (cute.Tensor द्वारे प्रदर्शित) हे वाजवी पर्याय होते. याव्यतिरिक्त, CuTeDSL हार्डवेअर प्रिमिटिव्ह्समध्ये विनाअट प्रवेश प्रदान करते, ज्यामुळे आम्हाला सर्व एज केसेसमध्ये सर्वोच्च कार्यक्षमता साध्य करण्यासाठी आवश्यक असलेले नियंत्रणाचा स्तर प्रदान होतो. CUDA सह, कंपाईल-टाइम स्पेशलायझेशनच्या संधी मर्यादित होत्या, कारण प्रत्येक अतिरिक्त पॅरामीटरने कंपायलेशन वेळेत अवाजवी खर्च किंवा बिल्ड सिस्टममध्ये गुंतागुंत वाढवली. याउलट, जस्ट-इन-टाइम कंपाइल केलेल्या DSL सह, आपण इटेशनच्या गतीचा बळी न देता अधिक आक्रमकपणे स्पेशलाइज्ड करू शकतो. TVM IR च्या विपरीत, ज्याने CUDA ला मध्यवर्ती पायरी म्हणून वापरले, आम्हाला PTX कडे लोअर होणारी MLIR-आधारित पाइपलाइन ही CUDA ला मध्यवर्ती पायरी म्हणून उत्सर्जित करण्यापेक्षा अधिक मजबूत आणि डीबग करण्यास सोपी असल्याचे आढळले. याव्यतिरिक्त, PTX सूचना रॅप करण्यासाठी कस्टम ऑप्सच्या वापरामुळे, आम्हाला आढळले की आम्ही अंतर्निहित हार्डवेअरची सर्व वैशिष्ट्ये पूर्णपणे वापरू शकतो.

CUDA kernels च्या compilation cost पैकी बहुतेक खर्च सामान्यतः kernels च्या template expansion टप्प्यात होतो, कारण optimization आणि PTX टप्प्यापासूनची compilation काही सेकंदांत पूर्ण होऊ शकते. अनेक compile-time parameters असल्यास, compilation चा हा टप्पा non-trivial kernel साठी अनेक मिनिटांपर्यंत वाढू शकतो. या bottleneck चे कारण म्हणजे parallel compilation साठीच्या संधी नंतरच, IR आणि PTX स्तरांवर, निर्माण होतात; तर C++ template handling क्रमाने (sequentially) होते. हे हाताळण्यासाठी काही युक्त्या असल्या तरी, template instantiations वेगळ्या compilation units मध्ये हलवून, त्यांना build systems मध्ये लक्षणीय बदल आणि अवजड macros चा वापर करावा लागतो. याउलट, Just-In-Time (JIT) ने compiled केलेल्या CuTeDSL kernels ना जटिल build rules ची आवश्यकता नसते आणि compilation times मध्ये अडथळा येत नाही, कारण development cycle configuration space मधील एका एकमेव बिंदूवर केंद्रित करता येते, तर नंतर केलेल्या काळजीपूर्वक unit tests kernels च्या संपूर्ण design space ला कव्हर करू शकतात.

जरी JIT कंपायलेशन खर्च विकास-वेळेपासून रनटाइमपर्यंत हलवते, तरी ओव्हरहेड स्वीकारण्यापेक्षा खूप जास्त आहे. जेव्हा एखादे मॉडेल तैनात केले जाते, तेव्हा ते संपूर्ण कॉन्फिगरेशन स्पेस मधील काहीच बिंदू सक्रिय करते, ज्यामुळे इन्फरन्स इंजिनच्या स्टार्टअप/वॉर्मअप दरम्यान सेकंद-स्तरीय ओव्हरहेड उद्भवतात. याउलट, कर्नलच्या सर्व कॉन्फिगरेशन आधीपासून कंपाइल केल्याने अनेक मिनिटे लागू शकतात. याव्यतिरिक्त, JIT-कंपाइल केलेले आर्टिफॅक्ट कॅश करून ओव्हरहेड आणखी कमी केले जाऊ शकते.

CuTe आणि CUTLASS च्या बाबतीत असल्याप्रमाणे, टेम्पलेट-आधारित डिस्पॅच आणि स्पेशलायझेशन वापरले जाते तेव्हा त्रुटींच्या मुळाचा शोध घेण्याच्या अडीचमुळे CUDA चा मोठा कंपायलेशन खर्च अधिक तीव्र होतो. DSL सह, त्रुटी निर्माण करणाऱ्या कोडच्या ओळीकडे त्रुटी निर्देशित केल्या जाऊ शकतात, ज्यामुळे डीबगिंग கணிசமாக सोपे होते. याव्यतिरिक्त, मध्यवर्ती MLIR टप्प्यांद्वारे किंवा जनरेट केलेल्या PTX शी जोडलेल्या डीबगिंग माहितीवर अवलंबून राहून जनरेट केलेल्या कोडला तो तयार करणाऱ्या DSL विधानाशी मॅप करणे बरेच सोपे आहे. DSL च्या या पैलूंमुळे, कर्नल विकसित करताना आम्हाला स्वतःला खूप अधिक उत्पादनक्षम असल्याचे आढळले.

मॉडेल्समध्ये कर्नल अनेक साम्य सामायिक करत असताना, त्यांना अद्याप कॉन्फिगरेशन नॉब्सच्या श्रेणीने सुसज्जित करणे आवश्यक आहे ज्याचा कार्यक्षमतेवर प्रभाव पडू शकतो. संबंधित पॅरामीटर्समध्ये हे समाविष्ट आहे:

हिडन डायमेंशन/हेड डायमेंशन/अटेन्शन हेड्सची संख्या

इनपुट/आउटपुट डेटा प्रकार

क्वांटायझेशन योजना

कंपाईल-टाइम स्थिरांक (बायस, ॲक्टिव्हेशन फंक्शन्स इ.)

ROSE आर्किटेक्चर, ज्यामध्ये त्याचे इंजिन, मॉडेल, कर्नल लायब्ररी आणि सेवा समाविष्ट आहेत.

कंपाईल-टाइम स्पेशलायझेशन कार्यक्षमता आणि डीबगिंग दोन्हीसाठी उपयुक्त आहे. विशेषतः पाइपलाइन केलेल्या प्रोसेसरमध्ये, जसे की GPU च्या बाबतीत, नेहमी-घेतलेल्या किंवा कधीही न घेतलेल्या शाखेपेक्षा नसलेली शाखा चांगली असते. याव्यतिरिक्त, विशिष्ट कॉन्फिगरेशन नॉब्सच्या संचासाठी कर्नल कंपाइल केल्यामुळे कमी PTX/SASS सूचना मिळतात, ज्या तपासणीदरम्यान cuda-gdb वापरताना विश्लेषण करणे आणि डीबग करणे सोपे असते.

दिलेल्या बॅच आकारासाठी आणि कॉन्फिगरेशनसाठी, DRAM कडील अपेक्षित रहदारी (traffic) आणि अंकगणितीय ऑपरेशनची संख्या अंदाजित केली जाऊ शकते. त्यानंतर, ऑप्टिमायझ्ड ग्रीड वाटप, सिंक्रोनायझेशन आणि काही हार्डवेअर वैशिष्ट्यांच्या वापरासाठी केस-बाय-केस आधारावर इष्टतम निर्णय घेतले जाऊ शकतात. CuTeDSL आम्हाला अस्वीकार्य कंपाईल-टाइम ओव्हरहेड न उचलता सखोल स्पेशलायझेशनसाठी घटक परिभाषित आणि पुन्हा वापरण्याची परवानगी देतो.

सामान्यतः, इन्फरन्स कर्नलची कार्यक्षमता दोन वेगवेगळ्या युज केसेससाठी ऑप्टिमाइझ केली गेली पाहिजे:

प्रिफिल: लहान बॅच आकार (काही चंक्स, सिक्वेन्सेस), प्रति सिक्वेन्स बरेच टोकन्स. एम्बेडिंग्ज आणि LLM च्या इनपुट टोकन्स दोन्हीसाठी थ्रूपुट जास्तीत जास्त असणे आवश्यक आहे. विलंब (latency) दसरा ते शेकडो मिलिसेकंदांमध्ये.

डिकोड: लहान ते मोठे बॅच आकार, एक (कोणतेही स्पेक्युलेटिव्ह डिकोडिंग नाही) ते प्रति सिक्वेन्स काही टोकन्स (ड्राफ्ट-टार्गेट स्पेक्युलेशन, मल्टी-टोकन प्रेडिक्शन). कर्नल लाँच लेटेन्सीच्या जवळ लेटेन्सी.

GEMM आणि attention साठी, prefill टप्पा मोठ्या संख्येतील tokens वर चालतो आणि performance हे device च्या compute throughput ने मर्यादित असते. Decoding मध्ये तीच operation लहान batch sizes वर होते, आणि execution times प्रामुख्याने activations आणि weights DRAM मधून streaming multiprocessors आणि tensor cores मध्ये लोड करण्याच्या खर्चामुळे ठरतात; इथे मर्यादा device च्या memory bandwidth मुळे येते. Batch size किंवा batch मधील tokens ची एकूण संख्या हा तुलनेने स्पष्ट indicator असला तरी, prefill आणि decode यांच्यातील switchover threshold अधिक सूक्ष्म असतो, कारण memory bandwidth हे model च्या hidden dimension आणि activation च्या data type वरही अवलंबून असते. त्यामुळे, algorithmic differences किंवा grid/block assignments यापैकी कुठल्याही कारणामुळे जर एखाद्याला फायदा होत असेल, तर kernels ना prefill आणि decode implementations मध्ये विभागले जाते.

प्रिफिल आणि डिकोड कर्नल दरम्यान सार्वत्रिकपणे भिन्न असलेला पॅरामीटर म्हणजे ब्लॉकमधील वॉर्प्सची संख्या (थ्रेड्सची संख्या, प्रति वॉर्प 32 थ्रेड्सच्या वॉर्प आकाराने भागले जाते). हा पॅरामीटर कंपाईल-टाइम स्थिरांक बनवणे इष्ट आहे, कारण ते रजिस्टर वाटपात मदत करू शकते. CUDA मध्ये, हे सामान्यत: __launch_bounds__ द्वारे थ्रेड्सची संख्या निर्दिष्ट करून आणि टेम्पलेट पॅरामीटर म्हणून वॉर्प्सची संख्या प्रदान करून केले जाते:

cpp

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

}

ब्लॉकमधील सर्व थ्रेड्ससाठी उपलब्ध रजिस्टरांची संख्या ही हार्डवेअर मर्यादा आहे, Hopper आणि Blackwell वर 64K. थ्रेड्सची संख्या अगोदरच माहित असल्यास, कंपाइलरला प्रत्येक थ्रेडला किती रजिस्टर नेमू शकतो हे ακριβώς ठाऊक असते आणि अधिक आक्रमकपणे ऑप्टिमाइझ केलेले कर्नल कंपाइल करण्यात यशस्वी होऊ शकते. तथापि, हे कंपायलेशन-वेळेच्या खर्चासह येते, कारण आता कर्नलला वेगवेगळ्या वॉर्प आकारांसाठी रीकंपाइल करावे लागेल. CuTeDSL सह, ही आता कोणतीही समस्या नाही, कारण आम्हाला संपूर्ण कॉन्फिगरेशन स्पेससाठी कंपायलेशनचा खर्च अगोदरच मोजावा लागत नाही. DSL ला ब्लॉकचा आकार एक स्थिर अभिव्यक्ती (constant expression) असणे देखील अपेक्षित आहे, जे CUDA मधील __latch_bounds__ द्वारे सूचित ऑप्टिमायझेशन अंतर्निहितपणे करते.

Warps ची संख्या constant ठेवण्याबरोबरच, prefill आणि decode यांच्यात ती बदलणेही desirable आहे. GPU च्या memory bandwidth चा पूर्ण उपयोग करण्यासाठी, प्रत्येक block साठी नेहमी उच्च number of warps assign करणे पुरेसे नसते, कारण full bandwidth फक्त उपलब्ध Streaming Multiprocessors (SMs) पैकी शक्य तितक्यांचा वापर करूनच साध्य होऊ शकते. म्हणून, दिलेल्या tokens च्या संख्येसाठी आणि ज्ञात data types तसेच hidden dimensions साठी, generated code 128-bit load/store instructions (Blackwell वर 256-bit) वापरत आहे याची खात्री करण्यासाठी आम्ही प्रथम आदर्श per-thread vector dimension निवडतो (bfloat16 साठी 8, float32 साठी 4). Vector size माहिती असल्यावर, kernel साठी आवश्यक warps ची संख्या आपण काढू शकतो. सर्वप्रथम आम्ही SMs ची संख्या कमाल करण्याचा प्रयत्न करतो: जर warps, SMs पेक्षा कमी असतील, तर block मागे एकच warp निवडला जातो. हे सामान्यतः kernel चा decode variant असतो. जर पुरेसे warps असतील, तर आम्ही warps ची संख्या हळूहळू 8, 16 किंवा 32 पर्यंत वाढवतो, आणि ज्या upper bound नंतर performance मध्ये साधारणपणे आणखी सुधारणा दिसत नाही, त्याचे benchmarking करतो.

Triton ने यापैकी काही ऑप्टिमायझेशन हुडखाली केले असताना, उच्च-स्तरीय कंपाइलरच्या त्रुटींमुळे मर्यादित न राहता, CuTe लेआउट बीजगणित काळजीपूर्वक वापरून आपण असेच परिणाम साध्य करू शकतो असे आम्हाला आढळले. विशेषतः जेव्हा सब-वॉर्प क्षैतिज घट (horizontal reductions) समाविष्ट असतात, तेव्हा आम्हाला आढळले की CuTeDSL ने आम्हाला इष्टतम कोड तयार करण्यासाठी आवश्यक असलेले बारीक-दाणेदार (fine-grained) नियंत्रण प्रदान केले.

काही परिस्थितींमध्ये, शक्य तितक्या ब्लॉकमध्ये कार्य वितरित करण्यासाठी, सर्व सहभागी ब्लॉक सिंक्रोनाइझ करण्यासाठी cooperative_groups::this_grid().sync() वापरले जाऊ शकते. निर्बंध लागू होतात: एकदा ब्लॉकने अंमलबजावणी सुरू केल्यावर तो SM मधून डी-शेड्युल केला जाऊ शकत नाही, त्यामुळे अग्रेषित प्रगती (forward progress) निश्चित करण्यासाठी, ब्लॉकांची संख्या डिव्हाइसवर उपलब्ध SMs च्या संख्येपेक्षा जास्त असू नये, जेणेकरून कोणत्याही वेळी सर्व ब्लॉक सिंक्रोनायझेशनमध्ये गुंतलेली जागतिक मेमरी स्थाने अपडेट आणि पोल करू शकतील. हा निर्बंध cudaLaunchCooperativeKernel द्वारे लागू केला जातो. ग्रीड बॅरियरचा खर्च Hopper वर सुमारे 3us आहे, तसेच काही परिस्थितींमध्ये अंतर्निहित मेमरी बॅरियरद्वारे मागील सर्व मेमरी ऑपरेशन फ्लश करण्याचा खर्च आहे. एकंदरीत, हे कर्नलला दोन लाँचमध्ये विभाजित करण्याच्या ओव्हरहेडच्या तुलनेत आहे, ज्यामुळे कार्यक्षमता-संवेदनशील डिझाइन निर्णय घेतला जातो. प्रोग्रामेटिक डिपेंडंट लाँच (PDL) Hopper पासून लाँच ओव्हरहेड कमी करू शकत असले तरी, डिकोड कर्नलसाठी, ओव्हरहेड टाळणे फायद्याचे आहे.

वर्कलोडपैकी एकाला SMs च्या संख्येपेक्षा जास्त ब्लॉकची आवश्यकता असल्यास, कर्नल लाँचच्या अंतर्निहित बॅरियरवर अवलंबून राहून, कर्नल विभाजित केला जावा. प्रिफिलसाठी प्रामुख्याने हेच आहे. सर्व वर्कलोडसाठी आदर्श असलेले ब्लॉकांची संख्या ग्रीडमध्ये बसत असल्यास, ग्रीड बॅरियरला प्राधान्य दिले जाते, जसे की लहान टोकन्सवर चालणाऱ्या डिकोड कर्नलच्या बाबतीत.

CuTeDSL मध्ये, हे निर्णय होस्ट-साइड लाँच कोडमध्ये घेतले जाऊ शकतात आणि खालील उदाहरणाप्रमाणे दर्शविल्याप्रमाणे कर्नल कंपोझ केले जाऊ शकतात आणि आवश्यकतेनुसार ग्रीड बॅरियरसह इंटरलीव्ह केले जाऊ शकतात. अशा हायब्रिड धोरणासह, किमान विलंब गाठण्यासाठी डिकोड कर्नल ऑप्टिमाइझ केले जाऊ शकतात आणि जास्तीत जास्त थ्रूपुटसाठी प्रिफिल ऑप्टिमाइझ केले जाऊ शकतात. 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 ऑल-टू-ऑल कर्नलच्या बाबतीत, खाली तपशीलवार सांगितल्याप्रमाणे, ग्रीड बॅरियरसह किंवा त्याशिवाय कर्नल तयार करण्याची लवचिकता असणे आणि कंपायलेशनच्या वेळी त्यांना ऑप्टिमाइझ केल्यामुळे आम्हाला CuTeDSL द्वारे महत्त्वपूर्ण कार्यक्षमता लाभ मिळण्यास मदत झाली.

CuTeDSL मध्ये आम्ही तयार केलेले पहिले कर्नल एक विशेष RMS नॉर्म कर्नल होते, जे लहान परिमाणांसाठी ऑप्टिमाइझ केलेले होते. अलीकडील Qwen मॉडेलच्या सेल्फ-अटेन्शन थरात वापरले जाणारे QK नॉर्म, प्रति-हेड आधारावर, 64, 128 किंवा 256 घटकांवर लागू केले जाते. हजारो ओलांडणाऱ्या लपविलेल्या परिमाणांवर (hidden dimensions) चालणाऱ्या, ज्यासाठी सामायिक मेमरी आणि रजिस्टर दोन्हीद्वारे क्षैतिज घट (horizontal reduction) आवश्यक असते, नियमित RMS नॉर्मच्या विपरीत, QK नॉर्मला केवळ वॉर्प्समध्ये घट (reduction) आवश्यक असते, ज्यामध्ये एकाच वॉर्पमध्ये एकाधिक हेड दाबण्याची शक्यता असते. सर्व हेड परिमाणांसाठी आणि डेटा प्रकारांसाठी सर्वोत्तम टोकन-टू-वॉर्प असाइनमेंट शोधण्यासाठी आम्ही कंपायलेशनच्या वेळी आक्रमकपणे ऑप्टिमाइझ करू शकलो अशी ही एक युज केस बनली.

RMS नॉर्म हे अगदी सोपे ऑपरेशन आहे, जरी व्हेरिएक्टरचा वर्ग नॉर्म (squared norm) मोजण्यासाठी त्याला क्षैतिज बेरीज (horizontal sum) आवश्यक असते:

cpp

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

अटेन्शन हेड सामान्यत: लहान असतात आणि प्रति-थ्रेड ऑपरेशन मोठ्या प्रमाणात व्हेक्टरइज्ड असतात, त्यामुळे एकच वॉर्प वर्ग नॉर्म (square norm) मोजू शकतो. मेमरी बँडविड्थ जास्तीत जास्त करण्यासाठी, प्रत्येक थ्रेडने 128 बिट इनपुट प्रक्रिया केली पाहिजे, जी 8 bfloat16 मूल्यांमध्ये रूपांतरित होते. परिणामी, हेडवर प्रक्रिया करण्यासाठी अर्धा-वॉर्प (half-warp) पुरेसा आहे, आणि वर्ग नॉर्म shuffle_sync_bfly द्वारे केलेल्या क्षैतिज वजावटीचा (horizontal reduction) वापर करून निश्चित केला जाऊ शकतो. हेड परिमाण आणि डेटा प्रकारापासून सुरुवात करून, आम्ही प्रति वॉर्प हेडची संख्या निश्चित करू शकतो आणि कंपायलेशनच्या वेळी युक्तिवाद आणि घट फेऱ्यांची संख्या निर्दिष्ट करू शकतो.

स्वयंचलितपणे व्हेक्टरइज्ड लोड/स्टोअरवर अवलंबून राहून आणि वॉर्प घट करण्यासाठी लोअर-लेव्हल आर्किटेक्चर-विशिष्ट ऑपरेशन वापरून, आम्ही नंतर अगदी कमी ओळींच्या कोडसह एक विशेष कर्नल तयार करू शकतो. कर्नल कोणते हेड (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-3 पट चांगली थ्रूपुट आणि लेटेन्सी प्रदान करतो.

MoE डिस्पॅच आणि कम्बाइन इनपुट टोकन्सना ते ज्या एक्सपर्ट्सकडे रुट केले आहेत त्यांच्याकडे शफल करतात, MoE GEMM कर्नल संबंधित मॅट्रिक्स मल्टिप्लिकेशन केल्यानंतर भारित संचयन (weighted accumulation) करतात. विविध उत्पादने विविध आकारांच्या MoE मॉडेलवर अवलंबून असल्यामुळे कार्यक्षमता विशेष स्वारस्याची आहे. तैनाती प्रकारांवर (deployment types) आधारावर, आम्ही ऑप्टिमाइझ करण्यासाठी विविध कर्नल ओळखतो:

TP=1: एकाच GPU साठी स्थानिक कर्नल. डिस्पॅच टोकन्सना संबंधित एक्सपर्ट्सकडे शफल करतो आणि कम्बाइन त्यांना जमा करतो. H100 वर Qwen3 30B Deployments साठी प्रामुख्याने वापरले जाते.

DP=1, TP ≠ 1: टेन्सर-पॅरलल, सिंगल-नोड कर्नल. इनपुट प्रतिकृती (replicated) बनते, डिस्पॅच वर्तमान रँकसाठी टोकन्स निवडतो आणि कम्बाइन त्यांना सर्व रँकमध्ये जमा करतो. NVLink किंवा NVLS ओलांडून संचयन (Accumulation). मोठ्या प्रिफिलर्ससाठी वापरले जाते.

DP ≠ 1, NVLink: पर्यायी टेन्सर पॅरालिझमसह डेटा-पॅरलल कर्नल. प्रत्येक रँक विशिष्ट एक्सपर्ट्स होस्ट करणाऱ्या इतर प्रत्येक रँकला टोकन्स पाठवतो. रुटिंग माहिती एकत्रित केली जाते, NVIDIA NVLink ओलांडून टोकन्स हस्तांतरित केले जातात. डेटा-पॅरलल H200 किंवा GB200 deployments साठी वापरले जाते.

DP ≠ 1, InfiniBand: NVLink आणि NVIDIA क्वांटम InfiniBand दोन्हीवर चालणारे पीअर-टू-पीअर डेटा-पॅरलल आणि टेन्सर-पॅरलल कर्नल. मल्टि-नोड Hopper deployments साठी वापरले जाते.

स्थानिक कर्नल रुटिंग टेबलच्या आधारावर प्रत्येक एक्सपर्टला नियुक्त केलेल्या टोकन्सची संख्या मोजतो आणि प्रत्येक एक्सपर्टच्या संबंधित टोकन्सच्या सूचीमध्ये प्रत्येक टोकनला एक अद्वितीय अनुक्रमणिका (index) नियुक्त करतो. रुटिंग पायरी नंतर, MoE GEMM कर्नलकडे पाठवण्यासाठी टोकन्स पर्यायी पॅडेड बफरमध्ये कॉपी केले जातात. टोकन स्थानांचा वापर करून, कम्बाइन कर्नल नंतर रजिस्टर्समध्ये संचयन (accumulation) करते, अंतिम टोकन्स आउटपुट बफरमध्ये लिहितो.

प्रिफिल आणि डिकोड दरम्यान डिस्पॅच टप्पा बदलतो: काही टोकन्ससाठी, वाचण्यासाठी फारशी रुटिंग माहिती नसते, म्हणून आम्ही एकाच ब्लॉकमध्ये त्यावर प्रक्रिया करतो आणि ग्रीडच्या बाकीच्या भागाला सिग्नल देतो, वैयक्तिक ब्लॉक टोकन्सचे चंक्स कॉपी करतात. सामायिक मेमरीमध्ये काउंटर ॲटोमिकली वाढवून टोकन असाइनमेंट निर्धारित केली जाते. प्रिफिलने अधिक रुटिंग माहितीवर प्रक्रिया करणे आवश्यक असल्याने, रुटिंग एका वेगळ्या कर्नलमध्ये अनेक ब्लॉकमध्ये विभागली जाते. सामायिक मेमरीमधील काउंटर वापरून इंडेक्स प्रथम स्थानिक पातळीवर नियुक्त केले जातात आणि नंतर प्रत्येक टोकनची स्थिती निर्धारित करण्यासाठी जागतिक मेमरीमधील सर्व ब्लॉकमध्ये एकत्रित केले जातात. ग्रीड-वाइज बॅरिएट वापरण्यास सक्षम होण्यासाठी या कर्नलची ग्रीड SMs च्या संख्येपर्यंत मर्यादित (clamped) केली जाते. त्यानंतरचे कर्नल उपलब्ध मेमरी बँडविड्थचा पूर्ण वापर करण्यासाठी मोठ्या संख्येने ब्लॉक आणि वॉर्प्समध्ये लाँच केलेल्या कॉपी करते.

कम्बाइन कर्नल तुलनेने सोपे आहे, कारण ते रुटिंग माहितीवर आणि डिस्पॅचद्वारे संगणित ऑफसेटवर अवलंबून राहून व्हेक्टरइज्ड संचयन करते. CuTeDSL येथे मौल्यवान सिद्ध झाले, कारण गणना साध्या CuTe टेन्सरद्वारे व्यक्त केली जाऊ शकते जे सर्व डेटा प्रकारांसाठी आणि लपविलेल्या परिमाणांसाठी कार्यक्षम कोडमध्ये संकलित होतात.

टेन्सर-पॅरलल तैनातीमध्ये, रुटिंग माहिती आणि इनपुट टोकन्सची प्रतिकृती (replicated) केली जाते. वर्तमान रँकवर असलेल्या एक्सपर्ट्सना नियुक्त केलेले टोकन्स रुट आणि निवडण्यासाठी किंचित बदलासह स्थानिक डिस्पॅच कर्नल पूर्णपणे पुन्हा वापरला जाऊ शकतो. काही पॅरामीटर्स आणि कंपाईल-टाइम स्पेशलायझेशनसह, स्थानिक डिस्पॅच कर्नल पुन्हा वापरले जातात.

रँकची संख्या सामान्यत: टोकन रुट केलेल्या एक्सपर्ट्सच्या संख्येपेक्षा कमी असल्यामुळे, सरासरी केसमध्ये असे गृहीत धरले जाऊ शकते की प्रत्येक टोकनच्या प्रत्येक रँकवर रुट केलेल्या त्याच्या प्रतिकृतींपैकी किमान एक आहे. या स्थितीत, टोकनच्या स्थानिक प्रती एकत्रित करण्यासाठी स्थानिक कम्बाइन कर्नल पुन्हा वापरला जाऊ शकतो, त्यानंतर सहभागी रँकमध्ये बेरीज करण्यासाठी ऑल-रिड्यूस केला जातो. प्रिफिल कर्नलच्या बाबतीत, आम्ही स्वतंत्र कर्नल लाँच करतो, तथापि डिकोड दोन्ही ऑपरेशन्स एकाच लाँचमध्ये एकत्र करते.

ऑल-रिड्यूससाठी, जर फक्त दोन रँक सहभागी असतील तर आम्ही साधे NVLink वापरतो, थ्रूपुटमध्ये 10-20% वाढीसाठी 4 किंवा अधिक रँक गुंतलेले असताना NVLS वर स्विच करतो. DSL मुळे, आपण MoE आणि नियमित ऑल-रिड्यूस कर्नल दरम्यान बहुतेक घट तर्कशास्त्र (reduction logic) आणि पीअर बफर एक्सचेंज सामायिक करू शकतो. ही अशी परिस्थिती आहे যেখানে हार्डवेअरवर जवळचे नियंत्रण फायदेशीर आहे, कारण NVLS साठी multimem.ld_reduce आणि multimem.st वरील DSL ऑपरेशन आवश्यक आहेत. घट स्वतःची अंमलबजावणी मोठ्या प्रमाणात NCCL चाचण्यांमधील all_reduce उदाहरणाचे अनुसरण करते.

डेटा पॅरलल कर्नल एकाधिक डिव्हाइसेससाठी टोकन्स एकत्रित करण्यासाठी NVLink वर अवलंबून असतात. डिस्पॅच टप्प्यात, टोकन्सची देवाणघेवाण करण्यापूर्वी, ते पाठवत असलेल्या प्रत्येक टोकनसाठी प्रत्येक रँकवर लक्ष्य ऑफसेट ओळखण्यासाठी पीअर्स प्रथम स्थानिक पातळीवर रुट करतात. त्यानंतर, ते संघर्ष टाळण्यासाठी प्रत्येक प्रेषकासाठी वाटप केलेल्या रिसिव्ह बफरमध्ये टोकन्स कॉपी करणे सुरू करतात. पूर्ण झाल्याचे सिग्नल देण्यासाठी प्रति-टोकन रिलीज-स्टोअर लेखनाचा वापर केला जातो. प्राप्तकर्त्याच्या टोकाला, MoE GEMM साठी सर्व टोकन्स एकाच बफरमध्ये घनीभूत करण्यास सक्षम होण्यासाठी रुटिंग माहिती एकत्रित केली जाती. लेखन पूर्ण होण्याची प्रतीक्षा करण्यासाठी प्राप्तकर्ते ॲक्वायर-लोडसह पीअर-लिखित फ्लॅग पोल करतात.

Combine side वर, tokens NVLink द्वारे source ranks कडे कॉपी केले जातात, आणि receipt दर्शवण्यासाठी प्रति-token flag असतो. Combine kernel स्थानिक आणि singlenode variants सारखाच आहे, पण addressing मध्ये थोडे फरक आणि काही अतिरिक्त synchronization logic आहेत.

हे कर्नल लागू करण्यासाठी आम्ही सिस्टम आणि डिव्हाइस स्तरावर रिलीज-ॲक्वायर लोड आणि स्टोअर उघड करण्यासाठी PTX रॅपची लायब्ररी तयार केली. सिंक्रोनायझेशन फ्लॅगची ग्रॅन्युलेटी सुधारून आणि NVLink वर टोकन्स पाठवण्यासाठी आणि प्राप्त करण्यासाठी अधिक ब्लॉक वाटप करून, आम्ही आमच्या मागील CUDA कर्नलपेक्षा 2 पट सुधारणा साध्य केली.

आमच्या पीअर-टू-पीअर कर्नलने यापूर्वी ConnectX-7 InfiniBand अडॅप्टरसह जोडलेल्या Hopper GPU वर अत्याधुनिक कार्यक्षमता साध्य केली होती. कार्यक्षमतेत आणखी वाढ करण्यासाठी आम्ही तेव्हापासून कम्बाइन कर्नल CuTeDSL वर पोर्ट केले आहेत. CUDA आवृत्त्या आधीच कंपाईल-टाइम स्पेशलायझेशनवर मोठ्या प्रमाणात अवलंबून असल्याने, त्यांनी आमच्या रिपॉझिटरीमध्ये कंपायलेशन-वेळ बॉटलनेक असल्याचे सिद्ध केले, ज्यामुळे आम्हाला काही पॅरामीटर्स वगळण्यास भाग पाडले. JIT-कंपाइल केलेल्या DSL मध्ये स्थलांतर करून, आम्ही अधिक पॅरामीटर्ससाठी स्पेशलाइज्ड केले आणि ग्रीड वाटप ट्विक केले, ज्यामुळे लेटेन्सीमध्ये सुमारे 3% वाढीसाठी प्रिफिल आणि डिकोडमध्ये आणखी फरक निर्माण झाले.

ROSE कर्नलच्या दृष्टीकोनातून, Blackwell समर्थन दोन श्रेणींमध्ये येते: आर्किटेक्चर-अॅग्नोस्टिक आणि आर्किटेक्चर-विशिष्ट कर्नल. टेन्सर कोअर्सवर अवलंबून असलेले आर्किटेक्चर-विशिष्ट कर्नलना जुळवून घेण्यासाठी स्पष्ट कामाची आवश्यकता असते. तथापि, DSL च्या वापरामुळे, आर्किटेक्चर-अॅग्नोस्टिक कर्नल क्षुल्लकपणे जुळवून घेतात.

Blackwell GPU मध्ये जास्तीत जास्त 160 SMs असू शकतात, तर Hopper मध्ये 132 आहेत. बहुतांश kernels साठी हे compile-time constant असते, आणि host-side launcher code मधील logic योग्य (prefill विरुद्ध decode) kernel चालवायचा आहे हे ठरवण्यासाठी त्याचा विचार करते. याव्यतिरिक्त, Blackwell 256-bit load instructions देखील आणतो, आणि vector width (CuTe tensors चे last dimension) वाढवून memory bandwidth utilization सुधारण्यासाठी त्या सहजपणे emit केल्या जाऊ शकतात.

Blackwell अनेक quantization schemes साठी (MXFP8, MXFP4, NVFP4) hardware support आणतो. यांचा support प्रामुख्याने FP8 block-wise scaling साठी असलेल्या विद्यमान kernels मध्ये बदल करून लागू करण्यात आला. येथे JIT support अत्यंत महत्त्वाचा आहे, कारण नवीन modes compile-time configuration space न फुगवता support करता येतात. पूर्वी, Triton ला warp च्या अर्ध्या किंवा चौथ्या भागावर काही elements across horizontal reductions व्यक्त करणे अवघड वाटत असे. NVFP4 आणि MXFP8/MXFP4 साठी अनुक्रमे 16 किंवा 32 elements across absolute maximum value मोजताना optimal performance मिळवण्यासाठी reductions वर अधिक सूक्ष्म नियंत्रण अत्यंत आवश्यक होते.

एक तुलनेने नवीन भाषा म्हणून, पूर्व-प्रशिक्षणादरम्यान LLM चा CuTeDSL कोडशी फारसा संपर्क आला नव्हता. याव्यतिरिक्त, ॲक्टिव्हेशन आणि नॉर्मसाठी ऑप्टिमाइझ केलेल्या कर्नलसह आमच्या वर्तमान कोडची शैली, विद्यमान CUTLASS किंवा Quack उदाहरणांपेक्षा किंचित भिन्न होती, जे TMA आणि आर्किटेक्चर-विशिष्ट मॅट्रिक्स मल्टिप्लिकेशन सूचनेसाठी भारी ऑप्टिमाइझ केले गेले होते. LLM कर्नल लेखनात मदत करू शकतील अशा टप्प्यावर पोहोचणे हा एक हळूहळू प्रक्रिया होती.

एजंटांवर कर्नल रूपांतरणासाठी अवलंबून राहण्यापूर्वी, आम्ही सुसंगत नमुन्यांचे अनुसरण करून, कर्नल आणि युटिलिटी फंक्शन्सचा एक प्रारंभिक कॉर्पस व्यक्तिशः तयार केला. कॉर्पस पुढे जात असताना, आम्ही CUDA मधून CuTeDSL मध्ये खडबडीत भाषांतरे करू शकलो, त्यानंतर महत्त्वपूर्ण मॅन्युअल रीवर्क केले. विशेषतः, आम्हाला आढळले की मॉडेल्सना पॉइंटर-टू-टेन्सर रूपांतरणे आणि वेक्टरयझेशनमध्ये संघर्ष करावा लागला, चाचणी कंपाइल करण्यासाठी प्रयत्न करताना टाइप त्रुटी दूर करण्याचा प्रयत्न करण्यात अडकले. त्यांनी कर्नलच्या ग्रीड-आधारित स्वरूपाची फारशी समज दर्शवली नाही, थ्रेड आणि ग्रीड सिंक्रोनायझेशनची गल्लत केली. सर्वात निराशाजनक म्हणजे, आमच्या स्वतःच्या CuTeDSL बदलांविरुद्ध तयार करण्याऐवजी CUDA-विशिष्ट लायब्ररींची अस्तित्वात नसलेली भाषांतरे हॉल्युसिनेट करण्याचा मोह होता. कालांतराने, CuTeDSL कर्नलचा कॉर्पस जसजसा वाढला, तसतसा आम्हाला आढळले की LLM बहुतेक सांगाडा कोड (skeleton code) तयार करू शकतात आणि रिक्त जागा भरू शकतात, ज्यामुळे आम्हाला उच्च-स्तरीय डिझाइनवर लक्ष केंद्रित करता येते.

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 का?

विकसक अनुभव (Developer experience)

मोठी कॉन्फिगरेशन स्पेस

प्रिफिल आणि डिकोड स्पेशलायझेशन

ग्रीड सिंक्रोनायझेशन

CuTeDSL कर्नल

QK नॉर्म

मिक्सचर-ऑफ-एक्सपर्ट्स डिस्पॅच/कम्बाइन

स्थानिक (Local)

सिंगलनोड

डेटा पॅरलल

पीअर-टू-पीअर

Blackwell समर्थन

AI-सहाय्यित कर्नल जनरेशन

निष्कर्ष आणि भविष्यातील कार्य

संदर्भ

उद्धरण