CuTeDSL at Perplexity

Perplexity-ல், உட்பொதிவுகள் (embeddings) முதல் டிரில்லியன் அளவுள்ள பாராமீட்டர்களைக் கொண்ட LLM-கள் வரை பல மாதிரிகளுக்கு சேவை வழங்க, நாங்கள் எங்களின் சொந்த Runtime-Optimized Serving Engine (ROSE)-ஐப் பயன்படுத்துகிறோம். சோனார், தேடல் மற்றும் உட்பொதிவுகள் உட்பட அனைத்து API-களுக்கும் பின்னால் இயங்கும் எங்கள் ரன்டைம், மிகவும் பொருத்தமான மாதிரிகள் மற்றும் கட்டமைப்புகளை ஆராய்ச்சியாளர்கள் ஆராய்வதற்கும், சிறந்த தேர்வுகளை விரைவாக தயாரிப்பு சூழலில் (production) வரிசைப்படுத்துவதற்கும் அதிகாரம் அளிக்கும் வகையில் உருவாக்கப்பட்டது. இதை அடைய, அதிநவீன கட்டமைப்புகளுக்கான இன்ஃபரன்ஸை (inference) மிக விரைவான வேகத்தில் செயல்படுத்தவும், வரிசைப்படுத்தவும், மேம்படுத்தவும் எங்களால் முடிய வேண்டும். இந்தக் வலைப்பதிவு பதிவில், எங்கள் இன்ஃபரன்ஸ் இன்ஜினுடன் ஒருங்கிணைக்கப்பட்ட CuTeDSL கர்னல்கள் (kernels), NVIDIA GPU-களில் மாதிரிகளை உச்ச செயல்திறனுக்குக் கொண்டுவருவதற்குத் தேவையான GPU கர்னல்களை உருவாக்க நமக்கு எவ்வாறு உதவுகின்றன என்பதை ஆராய்கிறோம்.

Perplexity-ல் உள்ள பல்வேறு குழுக்களால் உருவாக்கப்பட்ட தனிப்பயன் மாதிரிகள், AI இன்ஃபரன்ஸ் குழுவால் உருவாக்கப்பட்ட ROSE இன்ஃபரன்ஸ் இன்ஜினைப் பயன்படுத்தி, NVIDIA Hopper மற்றும் Blackwell GPU-களில் இன்-ஹவுஸில் ஹோஸ்ட் செய்யப்படுகின்றன. NVIDIA Triton இன்ஃபரன்ஸ் சர்வெருடன் இணக்கமான ஒரு இடைமுகத்துடன், மொழி டிகோடிங் மற்றும் வகைப்படுத்தல் ஆகிய இரண்டிற்கும் தனிப்பயனாக்கப்பட்ட Llama மாதிரிகளுக்குச் சிறப்பாகச் சேவை செய்ய ROSE ஆரம்பத்தில் உருவாக்கப்பட்டது. காலப்போக்கில், இது குறிப்பிடத்தக்க LLM-களுக்கும், மேலும் பலவிதமான குழுக்களுக்கு Perplexity-ன் தயாரிப்புகளுக்கு சக்தியூட்ட ரேங்கிங், வகைப்படுத்தல், ஸ்கோரிங் மற்றும் பிற டிரான்ஸ்பார்மர் அடிப்படையிலான மாதிரிகளுக்கும் சேவை செய்யும் ஒரு முழுமையான இன்ஃபரன்ஸ் இன்ஜினாக வளர்ந்தது. ROSE-ன் வடிவமைப்பு இலக்கு எளிமையானதாக இருப்பதாகும்: காலப்போக்கில், உச்ச செயல்திறனை அடைய முயற்சி முதலீடு செய்யப்படுகிறது, ஆனால் மாதிரிகளுக்கான எப்போதுமான வளர்ந்து வரும் கட்டமைப்புகள் மற்றும் பயன்பாட்டு வழக்குகளுக்கு ஏற்ப தகவமைத்துக் கொள்வதே முக்கிய முன்னுரிமையாகும்.

தரவு வகைகள், head dimensions, compile-time constants, மற்றும் launch settings ஆகியவற்றின் முழுவதுமாக kernel configuration space.

உயர் மட்டத்தில், ROSE என்பது ஒரு இன்ஜினைக் கொண்டுள்ளது, இது ஒரு மாதிரியை கிளையண்டுகள் தொடர்பு கொள்ளக்கூடிய ஒரு இடைமுகத்திற்கு ஏற்ப மாற்றுகிறது. இன்ஜின்கள் மிகவும் மீண்டும் பயன்படுத்தக்கூடியவை, கோரிக்கை திட்டமிடல் மற்றும் பேட்சிங் ஆகியவற்றைக் கையாளுகின்றன, மேலும் அடிப்படை மாதிரிகளிலிருந்து டிகோட் செய்யப்பட்ட டோக்கன்கள், உட்பொதிவுகள் அல்லது மதிப்பெண்களை மீட்டெடுக்கின்றன. சாதனங்களின் துவக்கம், செயல்முறை மற்றும் முனைகளுக்கு இடையேயான தொடர்பு சேனல்கள் மற்றும் எடை ஏற்றுதல் (weight loading) ஆகியவை அனைத்தாலும் கையாளப்படுகின்றன. LLM-களுக்கு, இன்ஜின் பேட்சிங், சங்கிங், சாம்பிளிங் ஆகியவற்றிற்குப் பொறுப்பாகும். முழு மற்றும் நேரியல் கவனம் (attention) இரண்டிற்கும், அதே போல் முன்னொட்டு பொருத்தம் (prefix matching) ஆகியவற்றிற்கும் KV சேமிப்பக ஒதுக்கீட்டிற்கும் இது பொறுப்பாகும். உட்பொதிவு இன்ஜின்கள் எளிமையானவை, ஏனெனில் அவை ஆன்-லைன் பேட்சிங்கை மட்டுமே செய்கின்றன. ஆக்சலரேட்டரில் மாதிரி இணையாக இயங்கும் போது, அடுத்த சுற்றுப் பணியை இன்ஜின் தயாரிக்கும் வகையில் செயலாக்கம் உகந்ததாக்கப்படுகிறது.

CuTeDSL, Triton, CUDA, CUTLASS மற்றும் cuBLAS ஐப் பயன்படுத்தி உருவாக்கப்பட்ட GPU கர்னல்களை ரேப் செய்யும் மாடலிங்கிற்கான தனிப்பயன் அடுக்குகளின் செழுமையான தொகுப்பை ROSE வெளிப்படுத்துகிறது. மேட்ரிக்ஸ் பெருக்கல் மற்றும் கவனச் செயலாக்கங்களுக்கு, நாங்கள் முதன்மையாக CUTLASS மற்றும் cuBLAS இலிருந்து NVIDIA கர்னல்களை நம்பியுள்ளோம். உட்பொதிவு, நார்ம் (norm), MoE மற்றும் ஆக்டிவேஷன் கர்னல்கள் உட்பட பைப்லைனின் மீதமுள்ள பகுதிகள் ROSE-ன் கர்னல் லைப்ரரியின் பகுதியாகும். நாங்கள் இன்ஃபரன்ஸ் இயக்கும் ஹோர்டுவேரைச் சிறப்பாக ஆதரிக்க, ROSE ஆனது MoE ரூட்டிங், டிஸ்பாட்ச் மற்றும் இணைப்பிற்கான (combine) சொந்த சாதனாங்களுக்கு இடையேயான தொடர்பு பிரிமிடிவ்களுடன் (primitives) பொருத்தப்பட்டுள்ளது. மற்ற கர்னல்கள் ஒப்பீட்டளவில் எளிமையானவை என்றாலும், அவை இயங்கும் பல்வேறு மறைக்கப்பட்ட பரிமாணங்கள் (hidden dimensions) மற்றும் சாதனங்கள் அனைத்திலும் உச்ச செயல்திறனை அடைய அவை முழுமையாக சிறப்புப்படுத்தப்பட வேண்டும். கூடுதலாக, பல்வேறு மாதிரிகளின் அனைத்து நாப்கள் மற்றும் ட்வீக்குகளுக்கு இடமளிக்க கம்பைல்-நேர சிறப்புப்படுத்தலை (compile-time specialization) நாங்கள் நம்பியுள்ளோம் (RMS நார்மிற்கான வெயிட் பயஸ் மற்றும் வெயிட் பயஸ் இல்லாமை, ஆக்டிவேஷன் செயல்பாடுகள் போன்றவை). இதன் விளைவாக, அவற்றைச் செயல்படுத்த 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 நிலையிலிருந்து மேற்கொண்டு ஆப்டிமைசேஷனும் கம்பைலேஷனும் சில விநாடிகளிலேயே முடிந்துவிடும். பல compile-time பராமீட்டர்கள் இருந்தால், இக் கம்பைலேஷன் கட்டம், முக்கியத்துவமுள்ள ஒரு கர்னலுக்கு பல நிமிடங்களாக நீளலாம். இந்த bottleneck, parallel கம்பைலேஷனுக்கான வாய்ப்புகள் IR மற்றும் PTX நிலைகளில்தான் பின்னரே உருவாகின்றன என்பதாலும், C++ டெம்ப்ளேட் கையாளுதல் தொடர்ச்சியாக நடைபெறுவதாலும் மேலும் தீவிரமாகிறது. இதை சமாளிக்க வழிகள் இருந்தாலும், டெம்ப்ளேட் இன்ஸ்டான்ஷியேஷன்களை தனித் கம்பைலேஷன் யூனிட்டுகளுக்கு நகர்த்துவதன் மூலம், அவை build systems-இல் குறிப்பிடத்தக்க மாற்றங்களையும் சிரமமான macro-களின் பயன்பாட்டையும் தேவைப்படுத்துகின்றன. இதற்கு மாறாக, Just-In-Time (JIT) ஆக கம்பைல் செய்யப்படும் CuTeDSL கர்னல்கள் சிக்கலான build விதிகளைத் தேவையாக்குவதில்லை, மேலும் அவை கம்பைலேஷன் நேரத்தைத் தடுக்கவும் இல்லை; காரணம், டெவலப்மென்ட் சைக்கிள் கான்ஃபிகரேஷன் ஸ்பேஸில் உள்ள ஒரு ஒற்றைப் புள்ளியில் கவனம் செலுத்த முடியும், பின்னர் கவனமாக செய்யப்படும் unit tests கர்னலின் முழு design space-ஐ கவர் செய்ய முடியும்.

JIT கம்பைலேஷன் செலவுகளை டெவலப்மென்ட் நேரத்திலிருந்து runtime-க்கு மாற்றினாலும், அந்த ஓவர்ஹெட் முற்றிலும் ஏற்றுக்கொள்ளத்தக்க அளவில்தான் உள்ளது. ஒரு மாடல் deploy செய்யப்படும்போது, அது முழு கான்ஃபிகரேஷன் ஸ்பேஸிலிருந்து சில புள்ளிகளையே செயல்படுத்துகிறது; அதனால் inference engine-இன் startup/warmup போது விநாடி அளவிலான ஓவர்ஹெட்கள் ஏற்படுகின்றன. இதற்கு மாறாக, ஒரு கர்னலின் அனைத்து கான்ஃபிகரேஷன்களையும் முன்கூட்டியே கம்பைல் செய்வது பல நிமிடங்கள் ஆகக்கூடும். மேலும், JIT-கம்பைல் செய்யப்பட்ட ஆர்ட்டிஃபேக்ட்களை cache செய்வதன் மூலம் இந்த ஓவர்ஹெட்களை மேலும் குறைக்க முடியும்.

CuTe மற்றும் CUTLASS விஷயத்தில் இருப்பது போல, டெம்ப்ளேட் அடிப்படையிலான டிஸ்பாட்ச் மற்றும் ஸ்பெஷலைசேஷன் பயன்படுத்தப்படும் போது பிழைகளின் தோற்றத்தைக் கண்டறிவதில் உள்ள சிரமத்தால் CUDA இன் அதிக கம்பைலேஷன் செலவு எடுத்துக்காட்டப்படுகிறது. ஒரு DSL உடன், அதை உருவாக்கிய குறியீட்டின் வரியை பிழைகள் சுட்டிக்காட்டலாம், இது பிழைத்திருத்தத்தை கணிசமாக எளிதாக்குகிறது. கூடுதலாக, இடைநிலை MLIR நிலைகள் வழியாக அல்லது உருவாக்கப்பட்ட PTX உடன் இணைக்கப்பட்ட பிழைத்திருத்த தகவலை நம்பி, உருவாக்கப்பட்ட குறியீட்டை அதைத் தொடங்கிய DSL அறிக்கையுடன் மேப்பிங் செய்வது மிகவும் எளிது. DSL இன் இந்த அம்சங்களுக்கு நன்றி, கர்னல்களை உருவாக்கும் போது நாங்கள் மிகவும் உற்பத்தித் திறனுடன் இருப்பதைக் கண்டறிந்தோம்.

மாதிரிகள் முழுவதும் கர்னல்கள் பல ஒற்றுமைகளைப் பகிர்ந்து கொண்டாலும், செயல்திறன் தாக்கத்தை ஏற்படுத்தக்கூடிய பல்வேறு உள்ளமைவு நாப்களுடன் அவை இன்னும் பொருத்தப்பட்டிருக்க வேண்டும். தொடர்புடைய அளவுருக்கள் அடங்கும்:

மறைக்கப்பட்ட பரிமாணம்/தலை பரிமாணம்/கவனத் தலைகளின் எண்ணிக்கை (Hidden Dimension/Head Dimension/Number of Attention Heads)

உள்ளீடு/வெளியீடு தரவு வகைகள் (Input/Output data types)

குவாண்டைசேஷன் திட்டம் (Quantization scheme)

கம்பைல்-நேர மாறிலிகள் (பயஸ்கள், ஆக்டிவேஷன் செயல்பாடுகள் போன்றவை)

ROSE கட்டமைப்பு, அதன் engine, model, kernel library, மற்றும் services உட்பட.

செயல்திறன் மற்றும் பிழைத்திருத்தம் ஆகிய இரண்டிற்கும் கம்பைல்-நேர சிறப்புப்படுத்தல் பயனுள்ளதாக இருக்கும். குறிப்பாக ஒரு பைப்லைன் செய்யப்பட்ட ப்ராசஸரில், GPU-களில் உள்ளதைப் போல, எப்போதும் எடுக்கப்படும் அல்லது எப்போதும் எடுக்கப்படாத கிளைக்கு பதிலாக இல்லாத கிளை சிறந்தது. கூடுதலாக, ஒரு குறிப்பிட்ட அளவிலான உள்ளமைவு நாப்களுக்கு ஒரு கர்னலை கம்பைல் செய்வது குறைந்த PTX/SASS வழிமுறைகளை விளைவிக்கிறது, ஆய்வின் போது cuda-gdb ஐப் பயன்படுத்தும் போது பகுப்பாய்வு செய்யவும் பிழைத்திருத்தம் செய்யவும் எளிதானது.

கொடுக்கப்பட்ட பேட்ச் அளவு மற்றும் உள்ளமைவுக்கு, DRAM-க்கான எதிர்பார்க்கப்படும் ட்ராஃபிக் மற்றும் கணித செயல்பாடுகளின் எண்ணிக்கையை மதிப்பிடலாம். அதைத் தொடர்ந்து, உகண்ட கிரிட் ஒதுக்கீடு, ஒத்திசைவு மற்றும் சில வன்பொருள் அம்சங்களைப் பயன்படுத்துவதற்கு வழக்கு அடிப்படையில் உகந்த முடிவுகளை எடுக்கலாம். ஏற்றுக்கொள்ள முடியாத கம்பைல்-நேர மேலதிகச் செலவை ஏற்காமல், ஆழமான சிறப்புப்படுத்தலுக்கான கூறுகளை வரையறுக்கவும் மீண்டும் பயன்படுத்தவும் CuTeDSL நம்மை அனுமதிக்கிறது.

பொதுவாக, இன்ஃபரன்ஸ் கர்னல்களின் செயல்திறன் இரண்டு தனித்தனி பயன்பாட்டு வழக்குகளுக்கு உகந்ததாக்கப்பட வேண்டும்:

ப்ரீஃபில் (Prefill): சிறிய பேட்ச் அளவு (சில சங்கா்கள், சீக்வென்ஸ்கள்), ஒரு சீக்வென்ஸுக்கு பல டோக்கன்கள். உட்பொதிவுகள் மற்றும் LLM-களின் உள்ளீட்டு டோக்கன்கள் ஆகிய இரண்டிற்கும் த்ரூபுட் (throughput) அதிகப்படுத்தப்பட வேண்டும். தாமதம் (latency) பத்துகளில் இருந்து நூற்றுக்கணக்கான மில்லி விநாடிகளில் உள்ளது.

Decode: சிறியதிலிருந்து பெரியதுவரையிலான batch sizes, ஒன்று (speculative decoding இல்லை) முதல் ஒவ்வொரு sequence-க்கும் சில tokens வரை (draft-target speculation, multi-token prediction). Latency kernel launch latency-க்கு நெருக்கமாக உள்ளது.

GEMM மற்றும் கவனத்திற்கு (attention), ப்ரீஃபில் கட்டம் அதிக எண்ணிக்கையிலான டோக்கன்களில் இயங்குகிறது மற்றும் செயல்திறன் சாதனத்தின் கணக்கீட்டு த்ரூபுட்டால் பிணைக்கப்பட்டுள்ளது. டிகோடிங்கில் சிறிய பேட்ச் அளவுகளில் அதே செயல்பாடு அடங்கும், மேலும் இயக்க நேரங்கள் DRAM- இலிருந்து ஆக்டிவேஷன்கள் மற்றும் எடைகளை ஸ்ட்ரீமிங் மல்டிபிராசஸர்கள் மற்றும் டென்சார் கோர்களில் ஏற்றுவதற்கான செலவால் ஆதிக்கம் செலுத்தப்படுகின்றன, மாறாக சாதனத்தின் நினைவக அலைவரிசையால் (memory bandwidth) கட்டுப்படுத்தப்படுகின்றன. பேட்ச் அளவு அல்லது ஒரு பேட்சில் உள்ள டோக்கன்களின் ஒட்டுமொத்த எண்ணிக்கை ஒப்பீட்டளவில் தெளிவான குறிகாட்டியாக இருந்தாலும், ப்ரீஃபில் மற்றும் டிகோட்டுக்கு இடையேயான மாறுதல் வரம்பு மிகவும் நுட்பமானது, ஏனெனில் நினைவக அலைவரிசையும் மாதிரியின் மறைக்கப்பட்ட பரிமாணம் மற்றும் ஆக்டிவேஷனின் தரவு வகையால் தீர்மானிக்கப்படுகிறது. இதன் விளைவாக, அல்காரிதமிக் வேறுபாடுகள் அல்லது கிரிட்/பிளாக் ஒதுக்கீடுகள் அவற்றில் ஏதேனும் ஒன்றிற்கு நன்மை பயக்கும் என்றால், கர்னல்கள் ப்ரீஃபில் மற்றும் டிகோட் செயலாக்கங்களாகப் பிரிக்கப்படுகின்றன.

prefill மற்றும் decode கர்னல்களுக்கு இடையில் உலகளாவிய வகையில் வேறுபடும் பராமீட்டர் என்பது ஒரு block-இல் உள்ள warps-இன் எண்ணிக்கை (threads-இன் எண்ணிக்கை, அதை ஒரு warp-க்கு 32 threads என்ற warp size-ஆல் வகுத்தது). இந்த பராமீட்டரை compile-time constant ஆக மாற்றுவது விரும்பத்தக்கது, ஏனெனில் அது register allocation-க்கு உதவக்கூடும். CUDA-வில், இது பொதுவாக __launch_bounds__ மூலம் threads எண்ணிக்கையை குறிப்பிட்டு, warps எண்ணிக்கையை template parameter ஆக வழங்குவதன் மூலம் செய்யப்படுகிறது:

cpp

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

}

ஒரு block-இல் உள்ள எல்லா threads-க்கும் கிடைக்கும் registers எண்ணிக்கை என்பது ஒரு hardware constraint; Hopper மற்றும் Blackwell-இல் இது 64K. threads எண்ணிக்கை முன்கூட்டியே தெரிந்திருந்தால், compiler ஒவ்வொரு thread-க்கும் எவ்வளவு registers ஒதுக்க முடியும் என்பதைத் துல்லியமாக அறிந்து, மேலும் aggressive ஆக ஆப்டிமைஸ் செய்யப்பட்ட கர்னல்களை compile செய்வதில் வெற்றியடைய முடியும். ஆனால், இதற்கு compilation-time செலவு உண்டு; ஏனெனில் இப்போது கர்னலை வேறு வேறு warp sizes-க்கு மீண்டும் compile செய்ய வேண்டியிருக்கிறது. CuTeDSL-இல், இது இனி பிரச்சினை அல்ல, காரணம் முழு கான்ஃபிகரேஷன் ஸ்பேஸுக்கான கம்பைலேஷன் செலவை முன்கூட்டியே செலுத்த வேண்டியதில்லை. இந்த DSL, block size ஒரு constant expression ஆக இருக்க வேண்டும் என்றும் எதிர்பார்க்கிறது; இதன் மூலம் CUDA-வில் __latch_bounds__ மூலம் குறிப்பிக்கப்பட்ட optimization-ஐ மறைமுகமாக செய்கிறது.

வார்ப்புகளின் எண்ணிக்கையை ஒரு மாறிலியாக மாற்றுவதைத் தவிர, ப்ரீஃபில் மற்றும் டிகோட்டுக்கு இடையில் அதை மாற்றுவது விரும்பத்தக்கது. ஒரு GPU-ன் நினைவக அலைவரிசையை முழுமையாகப் பயன்படுத்த, ஒரு தொகுதிக்கு எப்போதும் அதிக எண்ணிக்கையிலான வார்ப்களை ஒதுக்குவது போதாது, ஏனெனில் கிடைக்கக்கூடிய அனைத்து ஸ்ட்ரீமிங் மல்டிபிராசஸர்களையும் (SMs) முடிந்தவரை பயன்படுத்துவதன் மூலம் மட்டுமே முழு அலைவரிசையை அடைய முடியும். இதன் விளைவாக, கொடுக்கப்பட்ட எண்ணிக்கையிலான டோக்கன்கள் மற்றும் அறியப்பட்ட தரவு வகைகள் மற்றும் மறைக்கப்பட்ட பரிமாணங்களுக்கு, உருவாக்கப்பட்ட குறியீடு 128-பிட் லோட்/ஸ்டோர் வழிமுறைகளைப் (Blackwell-ல் 256-பிட்) பயன்படுத்துவதை உறுதிசெய்ய, முதலில் ஒரு சிறந்த த்ரெட் வாரியான வெக்டர் பரிமாணத்தைத் தேர்ந்தெடுக்கிறோம் (bfloat16-க்கு 8, float32-க்கு 4). வெக்டர் அளவை அறிந்தால், கர்னலுக்குத் தேவையான வார்ப்களின் எண்ணிக்கையைக் கணக்கிடலாம். நாங்கள் முதலில் SM-களின் எண்ணிக்கையை அதிகரிக்க முயற்சிக்கிறோம்: SM-களை விட குறைவான வார்ப்கள் இருந்தால், ஒரு தொகுதிக்கு ஒரு வார்ப் தேர்ந்தெடுக்கப்படுகிறது. இது பொதுவாக கர்னலின் டிகோட் மாறுபாடாகும். போதுமான வார்ப்கள் இருந்தால், வார்ப்களின் எண்ணிக்கையை படிப்படியாக 8, 16 அல்லது 32 ஆக அதிகரிக்கிறோம், அதன் பிறகு செயல்திறனில் எந்த முன்னேற்றத்தையும் நாம் பொதுவாகக் காணாத ஒரு மேல் வரம்பை பெஞ்ச்மார்க் செய்கிறோம்.

Triton இந்த உகந்ததாக்கல்களில் சிலவற்றைத் தன் கட்டுப்பாட்டின் கீழ் செய்தாலும், உயர் மட்ட கம்பைலரின் குறைபாடுகளால் கட்டுப்படுத்தப்படாமல், CuTe லேஅவுட் அல்ஜீப்ராவை கவனமாகப் பயன்படுத்துவதன் மூலம் இதே போன்ற முடிவுகளை அடைய முடியும் என்பதைக் கண்டறிந்தோம். குறிப்பாக சப்-வார்ப் கிடைமட்ட குறைப்புகள் (horizontal reductions) சம்பந்தப்பட்ட போது, உகந்த குறியீட்டை உருவாக்கத் தேவையான நுண்ணிய கட்டுப்பாட்டை CuTeDSL எங்களுக்கு வழங்கியதை நாங்கள் கண்டறிந்தோம்.

சில சூழ்நிலைகளில், முடிந்தவரை பல தொகுதிகளில் பணியை விநியோகிக்க, பங்கேற்கும் அனைத்து தொகுதிகளையும் ஒத்திசைக்க cooperative_groups::this_grid().sync() பயன்படுத்தப்படலாம். கட்டுப்பாடுகள் பொருந்தும்: ஒரு தொகுதி செயல்படுத்தத் தொடங்கியவுடன் அதை ஒரு SM- இலிருந்து டி-ஷெடியூல் செய்ய முடியாது என்பதால், முன்னோக்கி முன்னேற்றத்திற்கு உத்தரவாதம் அளிக்க, தொகுதிகளின் எண்ணிக்கை சாதனத்தில் கிடைக்கும் SM-களின் எண்ணிக்கையை விட அதிகமாக இருக்கக்கூடாது. ஒத்திசைவில் ஈடுபட்டுள்ள உலகளாவிய நினைவக இருப்பிடங்களை எந்த நேரத்திலும் அனைத்து தொகுதிகளும் புதுப்பிக்கவும் மற்றும் நிழலாடவும் (poll) முடியும் என்பதை உறுதிப்படுத்த இது உதவுகிறது. இந்த கட்டுப்பாடு cudaLaunchCooperativeKernel மூலம் செயல்படுத்தப்படுகிறது. ஒரு கிரிட் தடையின் (barrier) செலவு Hopper-ல் சுமார் 3us ஆகும், மேலும் சில சூழ்நிலைகளில் மறைமுக நினைவக தடை (implicit memory barrier) மூலம் முந்தைய அனைத்து நினைவக செயல்பாடுகளையும் ஃப்ளஷ் செய்வதன் செலவு ஆகும். ஒட்டுமொத்தமாக, இது ஒரு கர்னலை இரண்டு லாஞ்ச்களாகப் பிரிப்பதன் மேலதிகச் செலவுடன் ஒப்பிடத்தக்கது, இது செயல்திறன்-உணர்திறன் வடிவமைப்பு முடிவுக்கு வழிவகுக்கிறது. Hopper இலிருந்து புரோகிராமெட்டிக் டிபென்டன்ட் லாஞ்ச் (PDL) லாஞ்ச் மேலதிகச் செலவுகளைக் குறைக்க முடியும் என்ற போதிலும், டிகோட் கர்னல்களுக்கு, மேலதிகச் செலவைத் தவிர்ப்பது மதிப்புக்குரியது.

பணிச்சுமைகளில் ஒன்று SM-களின் எண்ணிக்கையை விட அதிக தொகுதிகள் தேவைப்பட்டால், கர்னல் பிரிக்கப்பட வேண்டும், இது கர்னல் லாஞ்ச்களின் மறைமுக தடையை நம்பியுள்ளது. இது முதன்மையாக ப்ரீஃபி லுக்கானது. அனைத்து பணிச்சுமைகளுக்கும் ஏற்ற தொகுதிகளின் எண்ணிக்கை கிரிட்டிற்குள் பொருந்தினால், குறைந்த எண்ணிக்கையிலான டோக்கன்களில் இயங்கும் டிகோட் கர்னல்களைப் போல, ஒரு கிரிட் தடை விரும்பப்படுகிறது.

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 மாடல்களின் சுய-கவன அடுக்கில் (self-attention layer) பயன்படுத்தப்படும் QK நார்ம், ஒரு தலை அடிப்படையில், 64, 128 அல்லது 256 கூறுகளில் பயன்படுத்தப்படுகிறது. பகிரப்பட்ட நினைவகம் மற்றும் ரெஜிஸ்டர்கள் வழியாக ஒரு கிடைமட்ட குறைப்பு தேவைப்படும், ஆயிரங்களைத் தாண்டிய மறைக்கப்பட்ட பரிமாணங்களில் இயங்கும் வழக்கமான RMS நார்மைப் போலன்றி, QK நார்மிற்கு வார்ப்களில் மட்டுமே குறைப்பு தேவைப்படுகிறது, பல தலைகளை ஒரே வார்பிற்குள் சுருக்கக்கூடிய சாத்தியக்கூறுகளுடன். அனைத்து தலை பரிமாணங்கள் மற்றும் தரவு வகைகளுக்கும் சிறந்த டோக்கன்-டு-வார்ப் ஒதுக்கீடுகளைக் கண்டறிய கம்பைல் நேரத்தில் நாங்கள் ஆக்கிரமிப்புடன் உகந்ததாக்கக்கூடிய ஒரு பயன்பாட்டு நிகழ்வாக இது மாறியது.

RMS நார்ம் என்பது மிகவும் நேரடியான செயலாகும், இருப்பினும் ஒரு வெக்டரின் வர்க்க நார்மைக் கணக்கிட ஒரு கிடைமட்ட கூட்டுத்தொகை தேவைப்படுகிறது:

cpp

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

attention heads பொதுவாக சிறியவையாகவும், ஒரு thread-க்கு செய்யப்படும் operations மிகவும் vectorized ஆகவும் இருப்பதால், ஒரு warp square norm-ஐ கணக்கிட முடியும். memory bandwidth-ஐ அதிகப்படுத்த, ஒவ்வொரு thread-உம் 128 bits input-ஐ செயலாக்க வேண்டும்; இது 8 bfloat16 values-க்கு சமம். ஆகவே, ஒரு half-warp ஒரு head-ஐ செயலாக்க போதுமானது, மேலும் squared norm-ஐ shuffle_sync_bfly மூலம் செய்யப்படும் horizontal reduction பயன்படுத்தி நிர்ணயிக்க முடியும். head dimension மற்றும் data type-இலிருந்து, ஒரு warp-க்கு எத்தனை heads உள்ளன என்பதைத் தீர்மானித்து, compilation நேரத்திலேயே arguments-ஐயும் reduction rounds எண்ணிக்கையையும் குறிப்பிடலாம்.

பின்னர், auto-vectorized loads/stores-ஐ நம்பியவாறும், warp reduction-க்கு கீழ்நிலை architecture-specific operations-ஐ பயன்படுத்தியவாறும், ஒப்பிடுகையில் குறைந்த வரிகளிலான code-இல் ஒரு specialized கர்னலை உருவாக்கலாம். அந்த கர்னல், எந்த head (Q அல்லது K) மீது வேலை செய்ய வேண்டும் என்பதையும், பயன்படுத்த வேண்டிய தொடர்புடைய weights-ஐயும் தீர்மானிக்கிறது. அது head-இலிருந்து ஒரு முழு vector-ஐ load செய்து, முதலில் உள்ளூராக கூட்டுகிறது. பின்னர், ஒரு head-க்கு உள்ள threads எண்ணிக்கையின் அடிப்படையில், squared norm கணக்கிடப்படுகிறது. இறுதியாக, auto-vectorized RMS norm நிர்ணயிக்கப்பட்டு memory-க்கு மீண்டும் எழுதப்படுகிறது. WEIGHT_BIAS மற்றும் EPS போன்ற model-specific பராமீட்டர்கள் constant-fold செய்யப்படுகின்றன.

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 மாதிரிகளை நம்பியிருப்பதால், செயல்திறன் குறிப்பாக ஆர்வமாக உள்ளது. வரிசைப்படுத்தல் வகைகளின் அடிப்படையில், உகந்ததாக்க வெவ்வேறு கர்னல்களை நாங்கள் வேறுபடுத்துகிறோம்:

TP=1: ஒரே GPU-க்கான உள்ளூர் கர்னல். டிஸ்பாட்ச் டோக்கன்களை தொடர்புடைய நிபுணர்களுக்கு ஷஃபிள் செய்கிறது மற்றும் இணைப்பானது அவற்றைச் சேகரிக்கிறது. H100 இல் Qwen3 30B வரிசைப்படுத்தல்களுக்கு முதன்மையாகப் பயன்படுத்தப்படுகிறது.

DP=1, TP ≠ 1: டென்சார்-இணை, ஒற்றை முனை கர்னல். உள்ளீடு நகலெடுக்கப்படுகிறது, டிஸ்பாட்ச் தற்போதைய ரேங்கிற்கான டோக்கன்களைத் தேர்ந்தெடுக்கிறது, மேலும் இணைப்பானது அனைத்து ரேங்குகளிலும் அவற்றைச் சேகரிக்கிறது. NVLink அல்லது NVLS முழுவதும் குவிப்பு. பெரிய ப்ரீஃபில்லர்களுக்குப் பயன்படுத்தப்படுகிறது.

DP ≠ 1, NVLink: விருப்பத்தேர்வு டென்சார் இணையுடன் தரவு-இணை கர்னல். குறிப்பிட்ட நிபுணர்களை ஹோஸ்ட் செய்யும் மற்ற ஒவ்வொரு ரேங்கிற்கும் ஒவ்வொரு ரேங்கும் டோக்கன்களை அனுப்புகிறது. ரூட்டிங் தகவல் திரட்டப்படுகிறது, NVIDIA NVLink முழுவதும் டோக்கன்கள் மாற்றப்படுகின்றன. தரவு-இணை H200 அல்லது GB200 வரிசைப்படுத்தல்களுக்குப் பயன்படுத்தப்படுகிறது.

DP ≠ 1, InfiniBand: NVLink மற்றும் NVIDIA Quantum InfiniBand இரண்டிலும் இயங்கும் பியர்-டு-பியர் தரவு-இணை மற்றும் டென்சார்-இணை கர்னல். பல-முனை Hopper வரிசைப்படுத்தல்களுக்குப் பயன்படுத்தப்படுகிறது.

உள்ளூர் கர்னல் ரூட்டிங் அட்டவணையின் அடிப்படையில் ஒவ்வொரு நிபுணருக்கும் ஒதுக்கப்பட்ட டோக்கன்களின் எண்ணிக்கையைக் கணக்கிடுகிறது மற்றும் ஒவ்வொரு நிபுணருக்கும் சொந்தமான டோக்கன்களின் பட்டியலுக்குள் ஒவ்வொரு டோக்கனுக்கும் ஒரு தனித்துவமான குறியீட்டை (index) ஒதுக்குகிறது. ரூட்டிங் படிநிலைக்குப் பிறகு, MoE GEMM கர்னலுக்கு அனுப்ப விருப்பத்தேர்வுடன் பேட் செய்யப்பட்ட பஃபர்களில் டோக்கன்கள் நகலெடுக்கப்படுகின்றன. டோக்கன் இருப்பிடங்களைப் பயன்படுத்தி, இணைப்புக் கர்னல் ரெஜிஸ்டர்களில் ஒரு குவிப்பைச் செய்கிறது, இறுதி டோக்கன்களை வெளியீட்டு பஃபரில் எழுதுகிறது.

ப்ரீஃபில் மற்றும் டிகோட்டுக்கு இடையில் டிஸ்பாட்ச் கட்டம் மாறுபடுகிறது: சில டோக்கன்களுக்கு, படிக்க ரூட்டிங் தகவல் மிகக் குறைவாகவே உள்ளது, எனவே அதை ஒரே தொகுதியில் செயலாக்கி, கிரிட்டின் மீதமுள்ள பகுதியைக் குறிக்கிறோம், தனிப்பட்ட தொகுதிகள் டோக்கன்களின் டாங்குகளை நகலெடுக்கின்றன. பகிரப்பட்ட நினைவகத்தில் உள்ள கவுன்ட்டர்களை அணு முறையில் அதிகரிப்பதன் மூலம் டோக்கன் ஒதுக்கீடு தீர்மானிக்கப்படுகிறது. ப்ரீஃபில் அதிக ரூட்டிங் தகவலைச் செயலாக்க வேண்டும் என்பதால், ரூட்டிங் ஒரு தனி கர்னலில் பல தொகுதிகளாகப் பிரிக்கப்படுகிறது. பகிரப்பட்ட நினைவகத்தில் உள்ள கவுன்ட்டர்களைப் பயன்படுத்தி குறியீடுகள் முதலில் உள்ளூரில் ஒதுக்கப்படுகின்றன, பின்னர் ஒவ்வொரு டோக்கனின் நிலையையும் தீர்மானிக்க உலகளாவிய நினைவகத்தில் உள்ள அனைத்து தொகுதிகளிலும் திரட்டப்படுகின்றன. கிரிட்-வாரைஸ் தடைகளைப் பயன்படுத்த முடியும் என்பதற்காக இந்த கர்னலின் கிரிட் SM-களின் எண்ணிக்கைக்குக் கிம்ப் செய்யப்படுகிறது. கிடைக்கக்கூடிய நினைவக அலைவரிசையை முழுமையாகப் பயன்படுத்த மிகவும் பெரிய எண்ணிக்கையிலான தொகுதிகள் மற்றும் வார்ப்களில் லாஞ்ச் செய்யப்பட்ட ஒரு அடுத்தடுத்த கர்னல் நகல்களைச் செய்கிறது.

இணைப்புக் கர்னல் ஒப்பீட்டளவில் எளிமையானது, ஏனெனில் இது ரூட்டிங் தகவலையும் டிஸ்பாட்சால் கணக்கிடப்பட்ட ஆஃப்செட்களையும் நம்பி, வெக்டரைஸ் செய்யப்பட்ட குவிப்பைச் செய்கிறது. CuTeDSL இங்கு மதிப்புமிக்கது என்பதை நிரூபித்தது, ஏனெனில் அனைத்து தரவு வகைகள் மற்றும் மறைக்கப்பட்ட பரிமாணங்களுக்கான திறமையான குறியீட்டாக கம்பைல் செய்யும் எளிய CuTe டென்சர்கள் மூலம் கணக்கீட்டை வெளிப்படுத்த முடியும்.

டென்சார்-இணை வரிசைப்படுத்தலில், ரூட்டிங் தகவல் மற்றும் உள்ளீட்டு டோக்கன்கள் நகலெடுக்கப்படுகின்றன. தற்போதைய ரேங்கில் அமைந்துள்ள நிபுணர்களுக்கு ஒதுக்கப்பட்ட டோக்கன்களை மட்டும் ரூட் செய்து தேர்ந்தெடுக்க ஒரு சிறிய மாற்றத்துடன் உள்ளூர் டிஸ்பாட்ச் கர்னலை முழுமையாக மீண்டும் பயன்படுத்தலாம். சில அளவுருக்கள் மற்றும் கம்பைல்-நேர சிறப்புப்படுத்தலுடன், உள்ளூர் டிஸ்பாட்ச் கர்னல்கள் மீண்டும் பயன்படுத்தப்படுகின்றன.

ரேங்குகளின் எண்ணிக்கை பொதுவாக ஒரு டோக்கன் ரூட் செய்யப்படும் நிபுணர்களின் எண்ணிக்கையை விட குறைவாக இருப்பதால், சராசரி வழக்கில் ஒவ்வொரு டோக்கனும் அதன் பிரதிகளில் குறைந்தது ஒன்றைக் கொண்டு ஒவ்வொரு ரேங்கிற்கும் ரூட் செய்யப்படுவதாகக் கருதலாம். இந்த சூழ்நிலையில், பங்கேற்கும் ரேங்குகளில் கூட்டுத்தொகை பெற ஒரு ஆல்-ரெடியூஸைத் தொடர்ந்து, ஒரு டோக்கனின் உள்ளூர் நகல்களைத் திரட்ட உள்ளூர் இணைப்புக் கர்னலை மீண்டும் பயன்படுத்தலாம். ப்ரீஃபில் கர்னலின் விஷயத்தில், நாங்கள் தனித்தனி கர்னல்களை லாஞ்ச் செய்கிறோம், இருப்பினும் டிகோட் இரண்டு செயல்பாடுகளையும் ஒரே லாஞ்சாக இணைக்கிறது.

ஆல்-ரெடியூஸுக்கு, இரண்டு ரேங்குகள் மட்டுமே பங்கேற்றால் எளிய NVLink-ஐப் பயன்படுத்துகிறோம், 4 அல்லது அதற்கு மேற்பட்ட ரேங்குகள் சம்பந்தப்பட்ட போது த்ரூபுட்டில் 10-20% அதிகரிப்புக்கு NVLS-க்கு மாறுகிறோம். DSL-க்கு நன்றி, MoE மற்றும் வழக்கமான ஆல்-ரெடியூஸ் கர்னல்களுக்கு இடையே பெரும்பாலான குறைப்பு தர்க்கம் மற்றும் பியர் பஃபர் பரிமாற்றங்களைப் பகிரலாம். NVLS-க்கு multimem.ld_reduce மற்றும் multimem.st மீதான DSL செயல்பாடுகள் தேவைப்படுவதால், வன்பொருள் மீது நெருக்கமான கட்டுப்பாடு பயனுள்ளதாக இருக்கும் ஒரு சூழ்நிலை இதுவாகும். குறைப்பின் செயலாக்கமானது NCCL சோதனைகளிலிருந்து all_reduce உதாரணத்தை பெரிதும் பின்பற்றுகிறது.

Data Parallel கர்னல்கள், பல devices-க்கான tokens-ஐ தொகுக்க NVLink-ஐ நம்புகின்றன. dispatch கட்டத்தில், tokens-ஐ பரிமாறுவதற்கு முன், peers முதலில் உள்ளூராக route செய்து, அவர்கள் அனுப்பும் ஒவ்வொரு token-க்கும் ஒவ்வொரு rank-இல் உள்ள target offset-ஐ அடையாளம் காண்கின்றனர். அதன்பின், conflicts-ஐ தவிர்க்க ஒவ்வொரு sender-க்கும் ஒதுக்கப்பட்ட receive buffers-க்குள் tokens-ஐ நகலெடுக்கத் தொடங்குகின்றனர். ஒவ்வொரு token-க்கும் release-store writes completion-ஐ signal செய்ய பயன்படுத்தப்படுகின்றன. receiver புறத்தில், routing தகவல் aggregate செய்யப்பட்டு, MoE GEMM-க்காக அனைத்து tokens-ஐயும் ஒரு single buffer-ஆக condense செய்ய முடியும். எழுதல்கள் முடிவடைந்துள்ளதை காத்திருக்க receivers, peer எழுதப்பட்ட flag-ஐ acquire-load மூலம் poll செய்கின்றனர்.

இணைப்புப் பக்கத்தில், டோக்கன்கள் NVLink வழியாக மூல ரேங்குகளுக்கு நகலெடுக்கப்படுகின்றன, ரசீதைக் குறிக்கும் டோக்கன் வாரியான கொடியுடன். இணைப்புக் கர்னல் உள்ளூர் மற்றும் ஒற்றை முனை மாறுபாடுகளைப் போன்றது, முகவரிப்படுத்தலில் சிறிய வேறுபாடுகள் மற்றும் சில கூடுதல் ஒத்திசைவு தர்க்கத்துடன்.

இந்த கர்னலைச் செயல்படுத்த சிஸ்டம் மற்றும் டிவைஸ் மட்டத்தில் ஸ்கோப் செய்யப்பட்ட ரிலீஸ்-அக்வைர் லோடுகள் மற்றும் ஸ்டோர்களை வெளிப்படுத்த PTX ரேப்பர்களின் நூலகத்தை நாங்கள் உருவாக்கினோம். ஒத்திசைவு கொடிகளின் கிரானுலாரிட்டியை மேம்படுத்துவதன் மூலமும், NVLink வழியாக டோக்கன்களை அனுப்பவும் பெறவும் கூடுதல் தொகுதிகளை ஒதுக்கீடு செய்வதன் மூலமும், எங்கள் முந்தைய CUDA கர்னல்களை விட 2 மடங்கு முன்னேற்றங்களை அடைந்தோம்.

எங்கள் பியர்-டு-பியர் கர்னல்கள் முன்பு ConnectX-7 InfiniBand அடாப்டர்களுடன் இணைக்கப்பட்ட Hopper GPU-களில் அதிநவீன செயல்திறனை அடைந்தன. செயல்திறனில் மேலும் லாபத்திற்காக இணைப்புக் கர்னல்களை CuTeDSL-க்கு மாற்றியுள்ளோம். CUDA பதிப்புகள் ஏற்கனவே கம்பைல்-நேர சிறப்புப்படுத்தலை பெரிதும் நம்பியிருந்ததால், அவை எங்கள் ரிபாசிட்டரியில் கம்பைலேஷன்-நேர தடைகளாக நிரூபிக்கப்பட்டன, இதனால் சில அளவுருக்களை நாங்கள் தவிர்க்க வேண்டிய கட்டாயம் ஏற்பட்டது. JIT-கம்பைல் செய்யப்பட்ட DSL-க்கு மாறுவதன் மூலம், தாமதத்தில் தோராயமாக 3% ஊக்கத்திற்காக ப்ரீஃபில் மற்றும் டிகோட்டுக்கு இடையே கூடுதல் வேறுபாடுகளை அறிமுகப்படுத்தி, கூடுதல் அளவுருக்களுக்காக சிறப்புப்படுத்தினோம் மற்றும் கிரிட் ஒதுக்கீடுகளை ட்வீக் செய்தோம்.

ROSE கர்னல்களின் கண்ணோட்டத்தில், Blackwell ஆதரவு இரண்டு வகைகளாகப் பிரிந்துள்ளது: கட்டமைப்பு-அக்னாஸ்டிக் (architecture-agnostic) மற்றும் கட்டமைப்பு-குறிப்பிட்ட கர்னல்கள். டென்சார் கோர்களை நம்பியிருக்கும் கட்டமைப்பு-குறிப்பிட்ட கர்னல்களுக்கு ஏற்பாடு செய்ய வெளிப்படையான வேலை தேவைப்படுகிறது. இருப்பினும், DSL இன் பயன்பாட்டிற்கு நன்றி, கட்டமைப்பு-அக்னாஸ்டிக் கர்னல்கள் எளிமையாக ஏற்படுகின்றன.

Hopper இன் 132 ஐப் போலன்றி, ஒரு Blackwell GPU ஆனது 160 SM-களை வரை கொண்டுள்ளது. பெரும்பாலான கர்னல்களுக்கு, இது ஒரு கம்பைல்-நேர மாறிலி மற்றும் இயக்க வேண்டிய உகந்த (ப்ரீஃபில் vs டிகோட்) கர்னலைத் தீர்மானிக்க ஹோஸ்ட்-சைடு லாஞ்சர் குறியீட்டில் உள்ள தர்க்கம் அதைக் காரணியாக்குகிறது. கூடுதலாக, Blackwell 256-பிட் லோட் வழிமுறைகளையும் அறிமுகப்படுத்துகிறது, இது நினைவக அலைவரிசை பயன்பாட்டை மேம்படுத்த வெக்டர் அகலத்தை (CuTe டென்சர்களின் கடைசி பரிமாணம்) அதிகரிப்பதன் மூலம் எளிதில் வெளியிடப்படலாம்.

Blackwell குறிப்பிடத்தக்க எண்ணிக்கையிலான quantization schemes-க்கு hardware support-ஐ அறிமுகப்படுத்துகிறது (MXFP8, MXFP4, NVFP4). இவைக்கான support முதன்மையாக FP8 block-wise scaling-க்காக ஏற்கனவே உள்ள கர்னல்களைச் சீரமைப்பதன் மூலம் செயல்படுத்தப்பட்டது. JIT support இங்கே மிக முக்கியமானது, ஏனெனில் compile-time configuration space-ஐ வெடித்தளவில் பெரிதாக்காமல் புதிய modes-ஐ ஆதரிக்க முடியும். கடந்த காலத்தில், warp-இன் அரை அல்லது கால் பகுதிக்கு மேல், சில elements-களில் கிடைமட்ட reductions-ஐ வெளிப்படுத்துவதில் Triton எங்களுக்கு சிரமமானதாக இருந்தது. NVFP4 மற்றும் MXFP8/MXFP4-க்கு முறையே 16 அல்லது 32 elements-க்குள் ஒரு absolute maximum value-ஐ கணக்கிடும்போது, optimal performance-ஐ அடைய reductions மீது நுணுக்கமான control மிக முக்கியமாக இருந்தது.

ஒப்பீட்டளவில் புதிய மொழியாக இருப்பதால், LLM-கள் தங்கள் ப்ரீ-டிரைனிங்கின் போது CuTeDSL குறியீட்டிற்கான வெளிப்பாடு குறைவாகவே இருந்தன. கூடுதலாக, ஆக்டிவேஷன்கள் மற்றும் நார்ம்களுக்காக உகந்ததாக்கப்பட்ட கர்னல்களுடன் எங்கள் தற்போதைய குறியீட்டின் பாணி, TMA மற்றும் கட்டமைப்பு சார்ந்த மேட்ரிக்ஸ் பெருக்கல் வழிமுறைகளுக்காக பெரிதும் உகந்ததாக்கப்பட்டிருந்த தற்போதைய CUTLASS அல்லது Quack எடுத்துக்காட்டுகளிலிருந்து சற்று வித்தியாசமாக இருந்தது. கர்னல் எழுதுவதில் LLM-கள் உதவக்கூடிய ஒரு புள்ளியை அடைவது ஒரு படிப்படியான செயல்முறையாக இருந்தது.

ஏஜெண்டுகளை கர்னல் மாற்றத்திற்கு நம்புவதற்கு முன்பு, நிலையான முறைகளைப் பின்பற்றி, கர்னல்கள் மற்றும் யூட்டிலிட்டி செயல்பாடுகளின் ஆரம்ப கார்பஸை நாங்கள் கைமுறையாக உருவாக்கினோம். கார்பஸ் முன்னேறும்போது, குறிப்பிடத்தக்க கைமுறை மறுவேலைகளைத் தொடர்ந்து, CUDA இலிருந்து CuTeDSL-க்கு தோராயமான மொழிபெயர்ப்புகளை நம்மால் செய்ய முடிந்தது. குறிப்பாக, பாயிண்டர்-டூ-டென்சார் மாற்றங்கள் மற்றும் வெக்டரைசேஷனுடன் மாதிரிகள் போராடுவதைக் கண்டறிந்தோம், ஒரு சோதனையைக் கம்பைல் செய்ய இலக்கு வைக்கும்போது வகை பிழைகளை (type errors) சரிசெய்ய முயற்சித்து சிக்கிக் கொண்டன. கர்னல்களின் கிரிட் அடிப்படையிலான தன்மையைப் பற்றிய சிறிதளவு புரிதலை மட்டுமே அவை காட்டியுள்ளன, த்ரெட் மற்றும் கிரிட் ஒத்திசைவைக் குழப்புகின்றன. மிகவும் ஏமாற்றமளிக்கும் வகையில், எங்கள் சொந்த 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 ஏன்?

டெவலப்பர் அனுபவம்

பெரிய கான்ஃபிகரேஷன் ஸ்பேஸ்

முன் நிரப்புதல் மற்றும் டிகோடிங் சிறப்பு

கிரிட் ஒத்திசைவு (Grid synchronization)

CuTeDSL கர்னல்கள்

QK நார்ம்

Mixture-of-Experts டிஸ்பாட்ச்/இணைப்பு (Combine)

உள்ளூர் (Local)

ஒற்றை முனை (Singlenode)

தரவு இணை (Data Parallel)

பியர்-டு-பியர் (Peer-to-Peer)

Blackwell ஆதரவு

AI-உதவி கர்னல் உருவாக்கம்

முடிவு மற்றும் எதிர்கால வேலை

குறிப்புகள்

மேற்கோள்