Perplexity में CuTeDSL

Perplexity में, हम एम्बेडिंग्स से लेकर ट्रिलियन-पैरामीटर वाले LLMs तक, मॉडल्स को सर्व करने के लिए इन-हाउस रनटाइम-ऑप्टिमाइज़्ड सर्विंग इंजन (ROSE) का उपयोग करते हैं। Sonar, Search और एम्बेडिंग्स सहित सभी API के पीछे स्थित, हमारा रनटाइम शोधकर्ताओं को सबसे उपयुक्त मॉडल्स और आर्किटेक्चर खोजने के लिए सशक्त बनाने के साथ-साथ सर्वोत्तम उम्मीदवारों को उत्पादन में तेज़ी से तैनात करने के लिए बनाया गया था। इसे प्राप्त करने के लिए, हमें अत्याधुनिक आर्किटेक्चर के लिए इन्फरेंस को लागू करने, तैनात करने और ऑप्टिमाइज़ करने में सक्षम होने की आवश्यकता है। इस ब्लॉग पोस्ट में, हम बताते हैं कि कैसे हमारे इन्फरेंस इंजन में एकीकृत CuTeDSL कर्नल्स हमें NVIDIA GPUs पर मॉडल्स को चरम प्रदर्शन तक ले जाने के लिए आवश्यक GPU कर्नल्स बनाने की अनुमति देते हैं।

Perplexity में विभिन्न टीम्स (teams) द्वारा निर्मित कस्टम मॉडल्स, AI इन्फरेंस टीम द्वारा निर्मित ROSE इन्फरेंस इंजन का उपयोग करके NVIDIA Hopper और Blackwell GPUs पर इन-हाउस होस्ट किए जाते हैं। ROSE को शुरू में भाषा डिकोडिंग और वर्गीकरण, दोनों के लिए अनुकूलित Llama मॉडल्स को बेहतर ढंग से सर्व करने के लिए बनाया गया था, जिसका इंटरफ़ेस NVIDIA Triton इन्फरेंस सर्वर के साथ संगत है। समय के साथ, यह एक पूर्ण इन्फरेंस इंजन में विकसित हो गया जो Perplexity के उत्पादों को संचालित करने के लिए विभिन्न प्रकार की टीम्स (teams) के लिए बड़े LLMs, साथ ही रैंकिंग, वर्गीकरण, स्कोरिंग और अन्य ट्रांसफॉर्मर-आधारित मॉडल्स को सर्व करता है। ROSE का डिज़ाइन लक्ष्य सरल है: समय के साथ, चरम प्रदर्शन तक पहुँचने के लिए प्रयास किया जाता है, लेकिन मुख्य प्राथमिकता हमारे मॉडल्स के लिए लगातार विकसित हो रहे आर्किटेक्चर और उपयोग के मामलों के अनुकूल होना है।

डेटा प्रकारों, हेड आयामों, संकलन-समय स्थिरांकों, और लॉन्च सेटिंग्स के बीच कर्नेल कॉन्फ़िगरेशन स्पेस।

उच्च स्तर पर, ROSE में एक इंजन शामिल होता है जो एक मॉडल को ऐसे इंटरफ़ेस के अनुकूल बनाता है जिसके साथ क्लाइंट इंटरैक्ट कर सकें। ये इंजन अत्यधिक पुन: प्रयोज्य (reusable) हैं, जो रिक्वेस्ट शेड्यूलिंग और बैचिंग को संभालते हैं, और अंतर्निहित मॉडल्स से डिकोड किए गए टोकन, एम्बेडिंग्स या स्कोर प्राप्त करते हैं। डिवाइसों का आरंभीकरण, इंटर-प्रोसेस और इंटर-नोड संचार चैनल, और वेट लोडिंग, इन सभी को इनके द्वारा संभाला जाता है। LLMs के लिए, इंजन बैचिंग, चंकिंग और सैंपलिंग के लिए जिम्मेदार है। यह पूर्ण और लीनियर अटेंशन, दोनों के लिए KV स्टोरेज आवंटन, और प्रीफ़िक्स मिलान के लिए भी जिम्मेदार है। एम्बेडिंग इंजन सरल होते हैं, क्योंकि वे केवल ऑन-लाइन बैचिंग करते हैं। निष्पादन को इस तरह ऑप्टिमाइज़ किया जाता है कि जब मॉडल समानांतर में एक्सीलरेटर पर चल रहा हो, तो इंजन काम का अगला राउंड तैयार कर लेता है।

ROSE मॉडलिंग के लिए कस्टम लेयर्स का एक समृद्ध संग्रह प्रदान करता है, जो CuTeDSL, Triton, CUDA, CUTLASS और cuBLAS के साथ निर्मित GPU कर्नल्स को रैप करता है। मैट्रिक्स गुणन और अटेंशन कार्यान्वयन के लिए, हम मुख्य रूप से CUTLASS और cuBLAS के NVIDIA कर्नल्स पर निर्भर करते हैं। पाइपलाइन का शेष भाग, जिसमें एम्बेडिंग, नॉर्म, MoE और एक्टिवेशन कर्नल्स शामिल हैं, ROSE की कर्नल लाइब्रेरी का हिस्सा हैं। हम जिस हार्डवेयर पर इन्फरेंस चलाते हैं, उसका सर्वोत्तम समर्थन करने के लिए, ROSE MoE रूटिंग, डिस्पैच और कंबाइन के लिए अपने स्वयं के इंटर-डिवाइस संचार प्रिमिटिव्स से सुसज्जित है। जबकि अन्य कर्नल्स अपेक्षाकृत सरल हैं, उन्हें उन सभी विभिन्न हिडन डाइमेंशन्स और डिवाइसों पर चरम प्रदर्शन प्राप्त करने के लिए पूरी तरह से विशेषीकृत किया जाना चाहिए, जिन पर वे चलते हैं। इसके अतिरिक्त, हम विभिन्न मॉडल्स के सभी नोब्स और ट्विक्स (जैसे RMS नॉर्म के लिए वेट बायस बनाम नो वेट बायस, एक्टिवेशन फ़ंक्शंस, आदि) को समायोजित करने के लिए कंपाइल-टाइम विशेषीकरण पर निर्भर करते हैं। परिणामस्वरूप, हमने उन्हें लागू करने के लिए DSL का उपयोग करने का निर्णय लिया: हमने शुरू में Triton का उपयोग करके ROSE बनाया था। हालाँकि, हमने तब से GPU प्रोग्रामिंग के लिए CuTeDSL को प्राथमिक वातावरण के रूप में अपना लिया है।

गणितीय दृष्टिकोण से, इन्फरेंस कर्नल्स आमतौर पर सरल होते हैं। परिणामस्वरूप, कर्नल विकसित करते समय, जटिलता आम तौर पर यह बताने में नहीं होती कि क्या गणना की जानी है, बल्कि यह होती है कि हार्डवेयर को इसे ठीक कैसे करना चाहिए। सर्वोत्तम कर्नल्स बनाने के लिए, हमें एक ऐसी भाषा की आवश्यकता होती है जो अंतर्निहित हार्डवेयर पर सही एब्स्ट्रैक्शंस प्रदान करती हो। सही भाषा हमें गणना करने के लिए उच्च-स्तरीय ऑपरेशन्स को निर्दिष्ट करने, और उन्हें एक विस्तृत कॉन्फ़िगरेशन स्पेस (स्पेस) में सभी बिंदुओं पर आक्रामक रूप से विशेषीकृत और ऑप्टिमाइज़ करने में सक्षम बनाती है।

CuTeDSL, CuTe लेआउट बीजगणित और MLIR के ऊपर निर्मित NVIDIA GPUs के लिए एक पायथन-आधारित DSL है। यह लो-लेवल हार्डवेयर प्रिमिटिव्स पर नियंत्रण बनाए रखते हुए, अत्यधिक ऑप्टिमाइज़्ड PTX के लिए जस्ट-इन-टाइम कंपाइल होता है। इसका उद्देश्य विकास की गति में बाधा डालने वाले महंगे संकलन समय के बिना, CUTLASS C++ के समान सुविधाएँ और प्रदर्शन प्रदान करना है। जिन भाषाओं का हमने सर्वेक्षण किया है या जिनके साथ काम किया है (जिनमें CUDA, cuTile, TVM IR और Triton शामिल हैं), उनमें से हम इसे अपने कर्नल्स के विकास के लिए सबसे उपयुक्त मानते हैं।

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

CUDA कर्नल्स की संकलन लागत का अधिकांश हिस्सा आमतौर पर कर्नल्स के टेम्पलेट विस्तार चरण में खर्च किया जाता है, क्योंकि PTX चरण से आगे के ऑप्टिमाइज़ेशन और संकलन को सेकंडों के भीतर पूरा किया जा सकता है। कई कंपाइल-टाइम पैरामीटर्स के साथ, संकलन का यह चरण एक गैर-तुच्छ कर्नल के लिए कई मिनट तक बढ़ सकता है। यह बाधा इस तथ्य से और बढ़ जाती है कि समानांतर संकलन के अवसर केवल बाद में, IR और PTX स्तरों पर उत्पन्न होते हैं, जबकि C++ टेम्पलेट हैंडलिंग क्रमिक रूप से होती है। हालाँकि इसे संबोधित करने के तरीके मौजूद हैं, टेम्पलेट इंस्टेंटिएशन को अलग संकलन इकाइयों में ले जाकर, उन्हें बिल्ड सिस्टम में महत्वपूर्ण बदलाव और भारी मैक्रोज़ के उपयोग की आवश्यकता होती है। इसके विपरीत, CuTeDSL कर्नल्स जिन्हें जस्ट-इन-टाइम (JIT) कंपाइल किया जाता है, उन्हें जटिल बिल्ड नियमों की आवश्यकता नहीं होती है और वे संकलन समय में बाधा नहीं डालते हैं, क्योंकि विकास चक्र को कॉन्फ़िगरेशन स्पेस (स्पेस) में एक बिंदु पर केंद्रित किया जा सकता है, जबकि बाद में सावधानीपूर्वक किए गए यूनिट टेस्ट कर्नल के पूरे डिज़ाइन स्पेस को कवर कर सकते हैं।

भले ही JIT संकलन लागत को विकास-समय से रनटाइम में स्थानांतरित करता है, लेकिन ओवरहेड स्वीकार्य से अधिक है। जब एक मॉडल तैनात किया जाता है, तो यह पूरे कॉन्फ़िगरेशन स्पेस (स्पेस) से केवल कुछ बिंदुओं को सक्रिय करता है, जिससे इन्फरेंस इंजन के स्टार्टअप/वार्मअप के दौरान सेकंड-स्केल ओवरहेड होता है। इसके विपरीत, कर्नल के सभी कॉन्फ़िगरेशन्स को प्री-कंपाइल करने में कई मिनट लग सकते हैं। इसके अतिरिक्त, JIT-कंपाइल किए गए आर्टिफ़ैक्ट (आर्टिफ़ैक्ट) को कैश करके ओवरहेड को और कम किया जा सकता है।

CUDA की भारी संकलन लागत त्रुटियों के मूल का पता लगाने में कठिनाई से और अधिक बढ़ जाती है जब टेम्पलेट-आधारित डिस्पैच और विशेषीकरण का उपयोग किया जाता है, जैसा कि CuTe और CUTLASS के साथ होता है। DSL के साथ, त्रुटियों को कोड की उस पंक्ति की ओर इशारा किया जा सकता है जिसने इसे उत्पन्न किया है, जिससे डीबगिंग काफी आसान हो जाती है। इसके अतिरिक्त, मध्यवर्ती MLIR चरणों के माध्यम से उत्पन्न कोड को उस DSL स्टेटमेंट से मैप करना बहुत सरल है जिससे इसकी उत्पत्ति हुई, या उत्पन्न PTX से जुड़े डीबगिंग जानकारी पर निर्भर रहना। DSL के इन पहलुओं के लिए धन्यवाद, हमने कर्नल्स विकसित करते समय खुद को बहुत अधिक उत्पादक पाया।

हालाँकि कर्नल्स मॉडल्स के बीच कई समानताएं साझा करते हैं, फिर भी उन्हें कई कॉन्फ़िगरेशन नोब्स के साथ सुसज्जित किया जाना चाहिए जिनका प्रदर्शन पर प्रभाव पड़ सकता है। प्रासंगिक पैरामीटर्स में शामिल हैं:

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

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

क्वांटाइज़ेशन स्कीम

कंपाइल-टाइम स्थिरांक (बायस, एक्टिवेशन फ़ंक्शंस, आदि)

ROSE आर्किटेक्चर, जिसमें इसका इंजन, मॉडल, कर्नेल लाइब्रेरी, और सेवाएँ शामिल हैं।

कंपाइल-टाइम विशेषीकरण प्रदर्शन और डीबगिंग दोनों के लिए उपयोगी है। विशेष रूप से एक पाइपलाइन प्रोसेसर में, जैसा कि GPUs के साथ होता है, एक गैर-मौजूद शाखा हमेशा-ली जाने वाली या कभी-न-ली जाने वाली शाखा से बेहतर होती है। इसके अतिरिक्त, विशिष्ट कॉन्फ़िगरेशन नोब्स के एक सेट के लिए कर्नल को कंपाइल करने से PTX/SASS निर्देश कम हो जाते हैं, जिनका विश्लेषण करना और जांच के दौरान cuda-gdb का उपयोग करते समय डीबग करना आसान होता है।

किसी दिए गए बैच आकार और कॉन्फ़िगरेशन के लिए, DRAM के लिए अपेक्षित ट्रैफ़िक और अंकगणितीय ऑपरेशन्स की संख्या का अनुमान लगाया जा सकता है। बाद में, इष्टतम ग्रिड आवंटन, सिंक्रनाइज़ेशन और कुछ हार्डवेयर सुविधाओं के उपयोग के लिए केस-दर-केस आधार पर इष्टतम निर्णय लिए जा सकते हैं। CuTeDSL हमें अस्वीकार्य कंपाइल-टाइम ओवरहेड के बिना गहरे विशेषीकरण के लिए घटकों को परिभाषित और पुन: उपयोग करने की अनुमति देता है।

आम तौर पर, इन्फरेंस कर्नल्स के प्रदर्शन को दो अलग-अलग उपयोग के मामलों के लिए ऑप्टिमाइज़ किया जाना चाहिए:

प्रीफ़िल: छोटा बैच आकार (कुछ चंक्स, सीक्वेंस), प्रति सीक्वेंस कई टोकन। एम्बेडिंग्स और LLMs के इनपुट टोकन, दोनों के लिए थ्रूपुट को अधिकतम करना। मिलीसेकंड के दसवें से सैकड़ों भाग में लेटेंसी।

डिकोड: छोटे से बड़े बैच आकार, प्रति सीक्वेंस एक (कोई सट्टा डिकोडिंग नहीं) से लेकर कुछ टोकन (ड्राफ्ट-टारगेट सट्टा, मल्टी-टोकन भविष्यवाणी)। लेटेंसी कर्नल लॉन्च लेटेंसी के करीब।

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। यदि थ्रेड्स (थ्रेड) की संख्या पहले से ज्ञात है, तो कंपाइलर जानता है कि वह प्रत्येक थ्रेड को कितने रजिस्टर असाइन कर सकता है और वह अधिक आक्रामक रूप से ऑप्टिमाइज़्ड कर्नल्स को कंपाइल करने में सफल हो सकता है। हालाँकि, यह संकलन-समय की लागत पर आता है, क्योंकि अब कर्नल को विभिन्न वार्प आकारों के लिए फिर से कंपाइल करना पड़ता है। CuTeDSL के साथ, यह अब कोई समस्या नहीं है, क्योंकि हमें अग्रिम रूप से पूरे कॉन्फ़िगरेशन स्पेस (स्पेस) के लिए संकलन की लागत का भुगतान करने की आवश्यकता नहीं है। DSL यह भी अपेक्षा करता है कि ब्लॉक आकार एक निरंतर अभिव्यक्ति हो, जो CUDA में __latch_bounds__ द्वारा निहित ऑप्टिमाइज़ेशन को स्पष्ट रूप से निष्पादित करे।

वार्प्स की संख्या को एक स्थिरांक बनाने के अलावा, इसे प्रीफ़िल और डिकोड के बीच बदलना वांछनीय है। GPU की मेमोरी बैंडविड्थ का पूरी तरह से उपयोग करने के लिए, प्रति ब्लॉक हमेशा वार्प्स की एक बड़ी संख्या असाइन करना पर्याप्त नहीं है, क्योंकि पूर्ण बैंडविड्थ केवल उपलब्ध स्ट्रीमिंग मल्टीप्रोसेसर (SMs) के अधिक से अधिक उपयोग करके प्राप्त की जा सकती है। परिणामस्वरूप, टोकन की एक दी गई संख्या और ज्ञात डेटा प्रकारों और हिडन डाइमेंशन्स के लिए, हम पहले एक आदर्श प्रति-थ्रेड वेक्टर डाइमेंशन (bfloat16 के लिए 8, float32 के लिए 4) चुनते हैं ताकि यह सुनिश्चित हो सके कि उत्पन्न कोड 128-बिट लोड/स्टोर निर्देशों (Blackwell पर 256-बिट) का उपयोग करे। वेक्टर आकार को जानकर, हम कर्नल द्वारा आवश्यक वार्प्स की संख्या की गणना कर सकते हैं। हम पहले SMs की संख्या को अधिकतम करने का प्रयास करते हैं: यदि SMs की तुलना में कम वार्प्स हैं, तो प्रति ब्लॉक एक एकल वार्प चुना जाता है। यह आमतौर पर कर्नल का डिकोड वेरिएंट होता है। यदि पर्याप्त वार्प्स हैं, तो हम धीरे-धीरे वार्प्स की संख्या को 8, 16 या 32 तक बढ़ाते हैं, एक ऊपरी सीमा को बेंचमार्क करते हैं जिसके बाद हम आमतौर पर प्रदर्शन में सुधार नहीं देखते हैं।

हालाँकि Triton ने हुड के नीचे इनमें से कुछ ऑप्टिमाइज़ेशन किए, हमने पाया कि हम उच्च-स्तरीय कंपाइलर की कमियों से बाधित हुए बिना, CuTe लेआउट बीजगणित का सावधानीपूर्वक उपयोग करके समान परिणाम प्राप्त कर सकते हैं। विशेष रूप से जब उप-वार्प क्षैतिज रिडक्शन शामिल होते हैं, तो हमने पाया कि CuTeDSL ने हमें इष्टतम कोड उत्पन्न करने के लिए आवश्यक फाइन-ग्रेन्ड नियंत्रण प्रदान किया।

कुछ स्थितियों में, काम को अधिक से अधिक ब्लॉक में वितरित करने के लिए, cooperative_groups::this_grid().sync() का उपयोग सभी सहभागी ब्लॉकों को सिंक्रनाइज़ करने के लिए किया जा सकता है। प्रतिबंध लागू होते हैं: चूंकि एक ब्लॉक निष्पादन शुरू होने के बाद SM से डी-शेड्यूल नहीं किया जा सकता है, इसलिए फॉरवर्ड प्रगति की गारंटी देने के लिए, ब्लॉकों की संख्या डिवाइस पर उपलब्ध 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 नॉर्म कर्नल था, जिसे छोटे डाइमेंशन्स के लिए ऑप्टिमाइज़ किया गया था। QK नॉर्म, जिसका उपयोग हाल के Qwen मॉडल्स की सेल्फ-अटेंशन लेयर में किया जाता है, को प्रति-हेड आधार पर 64, 128 या 256 तत्वों पर लागू किया जाता है। नियमित RMS नॉर्म के विपरीत, जो हजारों से अधिक हिडन डाइमेंशन्स पर काम करता है, जिसके लिए शेयर्ड मेमोरी और रजिस्टर दोनों के माध्यम से क्षैतिज रिडक्शन की आवश्यकता होती है, QK नॉर्म को केवल वार्प्स के बीच रिडक्शन की आवश्यकता होती है, जिसमें एक ही वार्प में कई हेड्स को निचोड़ने की संभावना होती है। यह एक ऐसा उपयोग मामला बन गया जहाँ हम सभी हेड डाइमेंशन्स और डेटा प्रकारों के लिए सर्वोत्तम टोकन-टू-वार्प असाइनमेंट खोजने के लिए कंपाइल समय पर आक्रामक रूप से ऑप्टिमाइज़ कर सकते थे।

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 कर्नल्स द्वारा प्रासंगिक मैट्रिक्स गुणन करने के बाद एक वेटेड संचय (weighted accumulation) करते हैं। प्रदर्शन विशेष रुचि का है, क्योंकि विभिन्न उत्पाद विभिन्न आकारों के MoE मॉडल्स पर निर्भर करते हैं। तैनाती के प्रकारों के आधार पर, हम ऑप्टिमाइज़ करने के लिए अलग-अलग कर्नल्स को अलग करते हैं:

TP=1: एक एकल GPU के लिए स्थानीय कर्नल। डिस्पैच टोकन को संबंधित विशेषज्ञों के पास शफल करता है और कंबाइन उन्हें संचित करता है। मुख्य रूप से H100 पर Qwen3 30B तैनाती के लिए उपयोग किया जाता है।

DP=1, TP ≠ 1: टेंसर-समानांतर, सिंगल-नोड कर्नल। इनपुट प्रतिकृत (replicated) है, डिस्पैच वर्तमान रैंक के लिए टोकन का चयन करता है, और कंबाइन उन्हें सभी रैंकों पर संचित करता है। NVLink या NVLS के पार संचय। बड़े प्रीफ़िलर्स के लिए उपयोग किया जाता है।

DP ≠ 1, NVLink: वैकल्पिक टेंसर समानता के साथ डेटा-समानांतर कर्नल। प्रत्येक रैंक विशिष्ट विशेषज्ञों को होस्ट करने वाले प्रत्येक अन्य रैंक को टोकन भेजता है। रूटिंग जानकारी एकत्रित की जाती है, टोकन NVIDIA NVLink के माध्यम से स्थानांतरित किए जाते हैं। डेटा-समानांतर H200 या GB200 तैनाती के लिए उपयोग किया जाता है।

DP ≠ 1, InfiniBand: पीयर-टू-पीयर डेटा-समानांतर और टेंसर-समानांतर कर्नल, NVLink और NVIDIA Quantum InfiniBand, दोनों के पार काम कर रहा है। मल्टी-नोड Hopper तैनाती के लिए उपयोग किया जाता है।

स्थानीय कर्नल रूटिंग तालिका के आधार पर प्रत्येक विशेषज्ञ को असाइन किए गए टोकन की संख्या की गणना करता है और प्रत्येक विशेषज्ञ से संबंधित टोकन की सूची के भीतर प्रत्येक टोकन को एक अद्वितीय इंडेक्स असाइन करता है। रूटिंग चरण के बाद, टोकन को MoE GEMM कर्नल को पास करने के लिए वैकल्पिक रूप से पैडेड बफ़र्स में कॉपी किया जाता है। टोकन स्थानों का उपयोग करके, कंबाइन कर्नल फिर रजिस्टरों में एक संचय करता है, अंतिम टोकन को एक आउटपुट बफ़र में लिखता है।

डिस्पैच चरण प्रीफ़िल और डिकोड के बीच भिन्न होता है: कुछ टोकन के लिए, पढ़ने के लिए बहुत कम रूटिंग जानकारी होती है, इसलिए हम इसे एक एकल ब्लॉक में प्रोसेस करते हैं और ग्रिड के शेष भाग को सिग्नल करते हैं, जिसमें व्यक्तिगत ब्लॉक टोकन के चंक्स को कॉपी करते हैं। टोकन असाइनमेंट शेयर्ड मेमोरी में काउंटर को परमाणु रूप से बढ़ाकर निर्धारित किया जाता है। चूंकि प्रीफ़िल को अधिक रूटिंग जानकारी प्रोसेस करनी चाहिए, रूटिंग को एक अलग कर्नल में कई ब्लॉकों में विभाजित किया जाता है। इंडेक्स पहले शेयर्ड मेमोरी में काउंटर का उपयोग करके स्थानीय रूप से असाइन किए जाते हैं और फिर प्रत्येक टोकन की स्थिति निर्धारित करने के लिए वैश्विक मेमोरी में सभी ब्लॉकों में एकत्रित किए जाते हैं। ग्रिड-वार बैरियर का उपयोग करने में सक्षम होने के लिए इस कर्नल के ग्रिड को SMs की संख्या तक क्लैंप किया जाता है। एक बाद का कर्नल कॉपी करता है, जो उपलब्ध मेमोरी बैंडविड्थ का पूरी तरह से उपयोग करने के लिए ब्लॉकों और वार्प्स की एक बहुत बड़ी संख्या में लॉन्च किया जाता है।

कंबाइन कर्नल अपेक्षाकृत सरल है, क्योंकि यह रूटिंग जानकारी और डिस्पैच द्वारा गणना किए गए ऑफ़सेट्स पर निर्भर करते हुए, वेक्टरकृत संचय करता है। CuTeDSL यहाँ मूल्यवान साबित हुआ, क्योंकि गणना को सरल CuTe टेंसर के माध्यम से व्यक्त किया जा सकता है जो सभी डेटा प्रकारों और हिडन डाइमेंशन्स के लिए कुशल कोड में कंपाइल होते हैं।

टेंसर-समानांतर तैनाती में, रूटिंग जानकारी और इनपुट टोकन प्रतिकृत होते हैं। स्थानीय डिस्पैच कर्नल को वर्तमान रैंक पर स्थित विशेषज्ञों को असाइन किए गए टोकन को रूट करने और चुनने के लिए थोड़ा बदलाव के साथ पूरी तरह से पुन: उपयोग किया जा सकता है। कुछ पैरामीटर्स और कंपाइल-टाइम विशेषीकरण के साथ, स्थानीय डिस्पैच कर्नल्स को पुन: उपयोग किया जाता है।

चूंकि रैंकों की संख्या आमतौर पर उन विशेषज्ञों की संख्या से कम होती है जिनके पास टोकन रूट किया जाता है, औसत मामले में यह माना जा सकता है कि प्रत्येक टोकन की कम से कम एक प्रतिकृति प्रत्येक रैंक पर रूट की जाती है। इस स्थिति में, स्थानीय कंबाइन कर्नल को एक टोकन की स्थानीय प्रतियों को संचित करने के लिए पुन: उपयोग किया जा सकता है, जिसके बाद सहभागी रैंकों में योग करने के लिए एक ऑल-रिड्यूस किया जाता है। प्रीफ़िल कर्नल के मामले में, हम अलग कर्नल्स लॉन्च करते हैं, हालाँकि डिकोड दोनों ऑपरेशन्स को एक ही लॉन्च में फ्यूज़ कर देता है।

ऑल-रिड्यूस के लिए, हम सादे NVLink का उपयोग करते हैं यदि केवल दो रैंक भाग लेते हैं, थ्रूपुट में 10-20% की वृद्धि के लिए 4 या अधिक रैंक शामिल होने पर NVLS पर स्विच करते हैं। DSL के लिए धन्यवाद, हम MoE और नियमित ऑल-रिड्यूस कर्नल्स के बीच अधिकांश रिडक्शन तर्क और पीयर बफ़र एक्सचेंजों को साझा कर सकते हैं। यह एक ऐसी स्थिति है जहाँ हार्डवेयर पर करीबी नियंत्रण फायदेमंद है, क्योंकि NVLS के लिए multimem.ld_reduce और multimem.st पर DSL ऑपरेशन्स की आवश्यकता होती है। रिडक्शन का कार्यान्वयन स्वयं काफी हद तक NCCL टेस्ट्स से all_reduce उदाहरण का पालन करता है।

डेटा समानांतर कर्नल्स कई डिवाइसों के लिए टोकन को एकत्रित करने के लिए NVLink पर निर्भर करते हैं। डिस्पैच चरण में, टोकन का आदान-प्रदान करने से पहले, पीयर्स पहले यह पहचानने के लिए स्थानीय रूप से रूट करते हैं कि वे जिस टोकन को भेज रहे हैं उसके लिए प्रत्येक रैंक पर टारगेट ऑफ़सेट क्या है। बाद में, वे संघर्षों से बचने के लिए प्रत्येक प्रेषक के लिए आवंटित प्राप्त बफ़र्स में टोकन को कॉपी करना शुरू करते हैं। पूर्णता का संकेत देने के लिए प्रति-टोकन रिलीज़-स्टोर राइट्स का उपयोग किया जाता है। रिसीवर छोर पर, रूटिंग जानकारी को MoE GEMM के लिए सभी टोकन को एक एकल बफ़र में संघनित (condense) करने में सक्षम होने के लिए एकत्रित किया जाता है। रिसीवर्स राइट्स के पूरा होने की प्रतीक्षा करने के लिए एक्वायर-लोड के साथ पीयर-लिखित ध्वज (peer-written flag) को पोल करते हैं।

कंबाइन पक्ष पर, टोकन को NVLink के माध्यम से स्रोत रैंकों पर कॉपी किया जाता है, जिसमें प्रति-टोकन ध्वज प्राप्ति का संकेत देता है। कंबाइन कर्नल स्थानीय और सिंगल-नोड वेरिएंट के समान है, जिसमें एड्रेसिंग और कुछ अतिरिक्त सिंक्रनाइज़ेशन तर्क में मामूली अंतर है।

हमने इस कर्नल को लागू करने के लिए सिस्टम और डिवाइस स्तर, दोनों पर स्कोप किए गए रिलीज़-एक्वायर लोड और स्टोर को एक्सपोज़ करने के लिए PTX रैपर्स की एक लाइब्रेरी बनाई। सिंक्रनाइज़ेशन ध्वजों की ग्रैन्युलैरिटी (granularity) में सुधार करके और NVLink पर टोकन भेजने और प्राप्त करने के लिए अधिक ब्लॉक आवंटित करके, हमने अपने पिछले CUDA कर्नल्स की तुलना में 2x सुधार प्राप्त किया।

हमारे पीयर-टू-पीयर कर्नल्स ने पहले ConnectX-7 InfiniBand एडेप्टर के साथ जोड़े गए Hopper GPUs पर अत्याधुनिक प्रदर्शन प्राप्त किया। हमने तब से प्रदर्शन में और लाभ के लिए कंबाइन कर्नल्स को CuTeDSL में पोर्ट कर लिया है। चूंकि CUDA संस्करण पहले से ही कंपाइल-टाइम विशेषीकरण पर बहुत अधिक निर्भर थे, वे हमारे रिपॉजिटरी में संकलन-समय की बाधाएं साबित हुए, जिससे हमें कुछ पैरामीटर्स को छोड़ने के लिए मजबूर होना पड़ा। JIT-कंपाइल किए गए DSL में माइग्रेट करके, हमने अधिक पैरामीटर्स के लिए विशेषीकरण किया और ग्रिड आवंटन में बदलाव किया, जिससे लेटेंसी में लगभग 3% की वृद्धि के लिए प्रीफ़िल और डिकोड के बीच और अंतर पैदा हुआ।

ROSE कर्नल्स के दृष्टिकोण से, Blackwell समर्थन दो श्रेणियों में आता है: आर्किटेक्चर-अज्ञेयवादी (agnostic) और आर्किटेक्चर-विशिष्ट कर्नल्स। आर्किटेक्चर-विशिष्ट कर्नल्स, टेंसर कोर पर निर्भर होने के कारण, अनुकूलन के लिए स्पष्ट काम की आवश्यकता होती है। हालाँकि, DSL के उपयोग के लिए धन्यवाद, आर्किटेक्चर-अज्ञेयवादी कर्नल्स आसानी से अनुकूलित हो जाते हैं।

एक Blackwell GPU में Hopper के 132 के विपरीत, 160 SMs तक होते हैं। अधिकांश कर्नल्स के लिए, यह एक कंपाइल-टाइम स्थिरांक है और होस्ट-साइड लॉन्चर कोड में तर्क चलाने के लिए इष्टतम (प्रीफ़िल बनाम डिकोड) कर्नल निर्धारित करने के लिए इसका कारक है। इसके अतिरिक्त, Blackwell 256-बिट लोड निर्देश भी पेश करता है, जिन्हें मेमोरी बैंडविड्थ उपयोग में सुधार करने के लिए वेक्टर चौड़ाई (CuTe टेंसर की अंतिम डाइमेंशन) को बढ़ाकर आसानी से उत्सर्जित किया जा सकता है।

Blackwell बड़ी संख्या में क्वांटाइज़ेशन योजनाओं (MXFP8, MXFP4, NVFP4) के लिए हार्डवेयर समर्थन पेश करता है। इनके लिए समर्थन मुख्य रूप से FP8 ब्लॉक-वार स्केलिंग के लिए मौजूदा कर्नल्स को ट्विक करके लागू किया गया था। JIT समर्थन यहाँ महत्वपूर्ण है, क्योंकि कंपाइल-टाइम कॉन्फ़िगरेशन स्पेस (स्पेस) को बढ़ाए बिना नए मोड्स का समर्थन किया जा सकता है। अतीत में, हमने पाया कि Triton कुछ तत्वों पर, एक वार्प के आधे या चौथाई हिस्से पर क्षैतिज रिडक्शन को व्यक्त करने में अजीब था। NVFP4 और MXFP8/MXFP4 के लिए क्रमशः 16 या 32 तत्वों पर एक निरपेक्ष अधिकतम मूल्य की गणना करते समय इष्टतम प्रदर्शन प्राप्त करने के लिए रिडक्शन पर फाइन-ग्रेन्ड नियंत्रण महत्वपूर्ण था।

एक अपेक्षाकृत नई भाषा के रूप में, LLMs को अपने प्री-ट्रेनिंग के दौरान CuTeDSL कोड का बहुत कम अनुभव था। इसके अतिरिक्त, हमारे वर्तमान कोड की शैली, जिसमें एक्टिवेशन और नॉर्म्स के लिए ऑप्टिमाइज़्ड कर्नल्स थे, मौजूदा CUTLASS या Quack उदाहरणों से थोड़ी अलग थी, जो TMA और आर्किटेक्चर-विशिष्ट मैट्रिक्स गुणन निर्देशों के लिए भारी रूप से ऑप्टिमाइज़ किए गए थे। एक ऐसे बिंदु तक पहुँचना जहाँ LLMs कर्नल लेखन में सहायता कर सकें, एक क्रमिक प्रक्रिया थी।

एजेंटों पर कर्नल रूपांतरण के लिए निर्भर होने से पहले, हमने सुसंगत पैटर्न का पालन करते हुए, कर्नल्स और उपयोगिता फ़ंक्शंस का एक प्रारंभिक कॉर्पस (corpus) मैन्युअल रूप से बनाया। जैसे-जैसे कॉर्पस आगे बढ़ा, हम CUDA से CuTeDSL में मोटे तौर पर अनुवाद करने में सक्षम थे, जिसके बाद महत्वपूर्ण मैन्युअल सुधार किए गए। विशेष रूप से, हमने पाया कि मॉडल्स पॉइंटर-टू-टेंसर रूपांतरण और वेक्टरकरण के साथ संघर्ष करते थे, और जब वे टेस्ट को कंपाइल करने का प्रयास करते थे तो टाइप त्रुटियों को ठीक करने की कोशिश में फंस जाते थे। उन्होंने कर्नल्स की ग्रिड-आधारित प्रकृति की बहुत कम समझ दिखाई, और थ्रेड (थ्रेड) और ग्रिड सिंक्रनाइज़ेशन को भ्रमित किया। सबसे निराशाजनक बात यह थी कि हमारे अपने CuTeDSL प्रतिस्थापनों के खिलाफ निर्माण करने के बजाय CUDA-विशिष्ट लाइब्रेरीज़ के गैर-मौजूद अनुवादों को मतिभ्रम (hallucinate) करने का प्रलोभन था। समय के साथ, जैसे-जैसे CuTeDSL कर्नल्स का कॉर्पस बढ़ा, हमने पाया कि LLMs अधिकांश कंकाल कोड बना सकते हैं और रिक्त स्थान भर सकते हैं, जिससे हमें उच्च-स्तरीय डिज़ाइन पर ध्यान केंद्रित करने की अनुमति मिलती है।

CUDA की तुलना में, हम पाते हैं कि CuTeDSL बहुत बेहतर कंपाइल-टाइम प्रदर्शन प्रदान करता है, जिससे हमें अधिक पैरामीटर्स के लिए विशेषीकरण करने और प्रदर्शन में सुधार प्राप्त करने की अनुमति मिलती है। हालाँकि Triton की तुलना में कुछ एब्स्ट्रैक्शंस की कमी के कारण कर्नल्स थोड़े अधिक विस्तृत हैं, हम पाते हैं कि जटिलता उचित है क्योंकि हार्डवेयर पर फाइन-ग्रेन्ड नियंत्रण हमें तेज़ कर्नल्स बनाने की अनुमति देता है।

लेखन के समय, CUTLASS C++ और CuTeDSL कर्नल्स के बीच प्रदर्शन का अंतर न्यूनतम है, जो इंटर-डिवाइस संचार, सामान्यीकरण या एक्टिवेशन के साथ उच्च-प्रदर्शन GEMM कर्नल्स को कंपोज़ और फ्यूज़ करने के रोमांचक अवसर खोलता है। इस DSL के लिए प्रतिबद्ध होकर, हम एम्बेडिंग्स से लेकर LLMs तक, अपने वर्कलोड के थ्रूपुट में सुधार करने के लिए अपने डिवाइसों से अतिरिक्त प्रदर्शन निकालने में सक्षम होंगे।

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 नॉर्म

Mixture-of-Experts डिस्पैच/कंबाइन

स्थानीय

सिंगल-नोड

डेटा समानांतर

पीयर-टू-पीयर

Blackwell समर्थन

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

निष्कर्ष और भविष्य का कार्य

संदर्भ

उद्धरण