Perplexity жүйесіндегі CuTeDSL

Perplexity компаниясында біз эмбеддингтерден бастап триллион параметрлі Үлкен тілдік модельдерге (LLM) дейінгі модельдерге қызмет көрсету үшін меншікті Runtime-Optimized Serving Engine (ROSE) жүйесін қолданамыз. Sonar, Search және Embeddings қоса алғанда, барлық API интерфейстерінің артында жұмыс істейтін біздің орындау ортамыз зерттеушілерге ең қолайлы модельдер мен архитектураларды зерттеуге, сонымен қатар ең жақсы кандидаттарды өндіріске жылдам енгізуге мүмкіндік беру үшін жасалған. Бұған қол жеткізу үшін біз ең озық архитектураларға арналған инференсті жылдам қарқынмен іске асырып, енгізіп және оңтайландыра алуымыз керек. Осы блог жазбасында біз инференс қозғалтқышымызға біріктірілген CuTeDSL кернелдері NVIDIA GPU-ларында модельдерді ең жоғары өнімділікке жеткізу үшін қажетті GPU кернелдерін жасауға қалай мүмкіндік беретінін қарастырамыз.

Perplexity-тегі түрлі командалар әзірлеген арнайы модельдер AI Inference командасы жасаған ROSE inference engine арқылы NVIDIA Hopper және Blackwell GPU-ларында ішкі түрде орналастырылған. ROSE бастапқыда NVIDIA Triton Inference Server-мен үйлесімді интерфейс арқылы тілдік декодтау да, жіктеу де үшін теңшелген Llama модельдерін жақсырақ қызмет көрсету мақсатында жасалды. Уақыт өте келе ол Perplexity өнімдерін қуаттандыру үшін көптеген командаларға арналған елеулі LLM-дерді, сондай-ақ ranking, classification, scoring және transformer негізіндегі басқа да модельдерді қызмет көрсететін толыққанды inference engine-ге айналды. ROSE-тың жобалау мақсаты қарапайым: уақыт өте келе ең жоғары өнімділікке жету үшін күш салынады, бірақ басты басымдық — біздің модельдеріміздің үздіксіз өзгеріп отыратын архитектуралары мен use case-теріне бейімделу.

Деректер түрлері, бас өлшемдері, компиляция уақытындағы тұрақтылар және іске қосу параметрлері бойынша ядро конфигурациясының кеңістігі.

Жалпы алғанда, ROSE моделін клиенттер өзара әрекеттесе алатын интерфейске бейімдейтін қозғалтқыштан тұрады. Қозғалтқыштар сұраныстарды жоспарлау мен топтастыруды (batching) өңдеп, негізгі модельдерден декодталған токендерді, эмбеддингтерді немесе ұпайларды ала отырып, жоғары дәрежеде қайта пайдалануға жарамды. Құрылғыларды инициализациялау, процессорралық және түйінаралық байланыс арналарын орнату, сондай-ақ салмақтарды жүктеу жұмыстары солардың арқылы басқарылады. LLM үшін қозғалтқыш топтастыру, бөліктерге бөлу және іріктеу үшін жауап береді. Ол сондай-ақ толық және сызықтық назар аудару (attention) үшін, сондай-ақ префикс сәйкестігі үшін KV жадысын бөлуге жауапты. Эмбеддинг қозғалтқыштары қарапайымдау, өйткені олар тек онлайн топтастыруды орындайды. Орындау процесі модель үдеткіште параллель жұмыс істеп жатқанда, қозғалтқыш жұмыстың келесі кезеңін дайындайтындай етіп оңтайландырылған.

ROSE модельдеуге арналған арнайы қабаттардың бай жиынтығын ұсынады, олар CuTeDSL, Triton, CUDA, CUTLASS және cuBLAS көмегімен жасалған GPU кернелдерін орайды. Матрицалық көбейту және назар аудару (attention) іске асырулары үшін біз негізінен CUTLASS және cuBLAS ұсынған NVIDIA кернелдеріне сүйенеміз. Конвейердің қалған бөлігі, соның ішінде эмбеддинг, норма, MoE және активация кернелдері ROSE кернелдер кітапханасына кіреді. Инференс жұмысын атқаратын аппараттық құралдарды ең жақсы түрде қолдау үшін ROSE MoE маршруттау, тарату және біріктіруге арналған өз құрылғыаралық байланыс примитивтерімен жабдықталған. Басқа кернелдер салыстырмалы түрде қарапайым болғанымен, олар өзі жұмыс істейтін барлық әртүрлі жасырын өлшемдер мен құрылғыларда ең жоғары өнімділікке қол жеткізу үшін мұқият мамандандырылуы тиіс. Сонымен қатар, біз әртүрлі модельдердің барлық параметрлері мен баптауларын (RMS нормасы үшін салмақ смещенісі (bias) мен онсыз болуы, активация функциялары және т.б.) ескеру үшін құрастыру кезіндегі (compile-time) мамандануға сүйенеміз. Сондықтан біз оларды іске асыру үшін DSL тілдерін пайдалануды шештік: бастапқыда ROSE жүйесін Triton көмегімен құрдық. Дегенмен, содан бері біз GPU бағдарламалау үшін негізгі орта ретінде CuTeDSL қабылдадық.

Математикалық тұрғыдан алғанда, инференс кернелдері әдетте күрделі емес. Демек, кернелді әзірлеу кезінде күрделілік әдетте қандай есептеу жасалу керектігін білдіруде емес, аппараттық құралдың оны дәл қалай жасауы керектігінде болады. Ең жақсы кернелдерді жасау үшін бізге негізгі аппараттық құралдардың үстінен дұрыс абстракцияларды ұсынатын тіл қажет. Дұрыс тіл бізге есептеуді орындау үшін жоғары деңгейлі операцияларды анықтауға, оларды кең конфигурациялық кеңістіктің барлық нүктелерінде агрессивті түрде мамандандыруға және оңтайландыруға мүмкіндік береді.

CuTeDSL – бұл CuTe орналасу алгебрасы мен MLIR негізінде жасалған NVIDIA GPU үшін Python тіліне негізделген DSL. Ол төмен деңгейлі аппараттық примитивтерді басқаруды сақтай отырып, жоғары деңгейде оңтайландырылған PTX форматына Just-In-Time (JIT) режимінде компиляцияланады. Оның мақсаты – әзірлеу жылдамдығына кедергі келтіретін қымбат компиляция уақытынсыз CUTLASS C++ сияқты мүмкіндіктер мен өнімділікті ұсыну. Біз зерттеген немесе жұмыс істеген тілдердің ішінде (CUDA, cuTile, TVM IR және Triton қоса алғанда) біз оны өз кернелдерімізді әзірлеу үшін ең қолайлы деп санаймыз.

Жалпы алғанда, Triton немесе cuTile сияқты жоғары деңгейлі және қысқа болмаса да, CuTe орналасу алгебрасы және оның композициялары (cute.Tensor арқылы ұсынылған) қолайлы алмастырғыштар екенін анықтадық. Сонымен қатар, CuTeDSL аппараттық примитивтерге шектеусіз қол жеткізуді ұсынады, бұл бізге барлық шекаралық жағдайларда ең жоғары өнімділікке қол жеткізу үшін қажетті басқару деңгейін қамтамасыз етеді. CUDA көмегімен құрастыру кезіндегі мамандану мүмкіндіктері шектеулі болды, өйткені әрбір қосымша параметр компиляция уақытына орынсыз шығын немесе құрастыру жүйесіне күрделілік қосты. Керісінше, Just-In-Time компиляцияланатын DSL көмегімен біз қайталау қарқынын жоғалтпастан агрессивті түрде мамандана аламыз. CUDA тілін аралық қадам ретінде пайдаланған TVM IR айырмашылығы, біз PTX-ке түсіретін MLIR негізіндегі конвейердің CUDA-ны аралық қадам ретінде шығаратын конвейерге қарағанда сенімдірек әрі жөндеуге оңайырақ екенін анықтадық. Сонымен қатар, PTX нұсқауларын орау үшін арнайы операцияларды пайдалану арқылы біз негізгі аппараттық құралдардың барлық мүмкіндіктерін толығымен пайдалана алатынымызды анықтадық.

CUDA кернелдерінің компиляция құнының негізгі бөлігі әдетте кернелдерді шаблондық кеңейту фазасына жұмсалады, өйткені PTX кезеңінен бастап оңтайландыру мен компиляция бірнеше секунд ішінде аяқталуы мүмкін. Көптеген құрастыру кезіндегі параметрлермен компиляцияның бұл кезеңі күрделі емес кернел үшін бірнеше минутқа дейін ұлғаюы мүмкін. Бұл кедергі C++ шаблонын өңдеу тізбекті түрде жүретіндіктен, параллель компиляция мүмкіндіктері тек кейінірек, IR және PTX деңгейлерінде пайда болатындығымен айқындала түседі. Шаблонды инстанциациялауды бөлек компиляциялық бірліктерге жылжыту арқылы мұны шешу тәсілдері болғанымен, олар құрастыру жүйелеріне айтарлықтай түзетулер енгізуді және қолайсыз макростарды пайдалануды талап етеді. Керісінше, Just-In-Time (JIT) режимінде компиляцияланатын CuTeDSL кернелдері күрделі құрастыру ережелерін талап етпейді және компиляция уақытына кедергі келтірмейді, өйткені әзірлеу циклі конфигурациялық кеңістіктегі бір ғана нүктеге бағытталуы мүмкін, ал мұқият бірлік тесттер кейін кернелдің бүкіл жобалау кеңістігін қамти алады.

JIT компиляциясы шығындарды әзірлеу уақытынан орындалу уақытына жылжытқанымен, үстеме шығын (overhead) толығымен қабылдауға болады. Модель орналастырылған кезде, ол бүкіл конфигурациялық кеңістіктегі тек бірнеше нүктені іске қосады, бұл инференс қозғалтқышының іске қосылуы/қыздырылуы кезінде секундтық үстеме шығындарды тудырады. Керісінше, кернелдің барлық конфигурацияларын алдын ала компиляциялау бірнеше минутқа созылуы мүмкін. Сонымен қатар, JIT арқылы компиляцияланған артефакттарды кэштеу арқылы үстеме шығындарды одан әрі азайтуға болады.

CUDA компиляциясының ауыр шығыны CuTe және CUTLASS сияқты шаблонға негізделген тарату (dispatch) және мамандану қолданылған кезде қателердің шығу тегін қадағалаудың қиындығымен күшейе түседі. DSL көмегімен қателер оны тудырған код жолына көрсетіле алады, бұл жөндеуді (debugging) айтарлықтай жеңілдетеді. Сонымен қатар, аралық MLIR кезеңдері арқылы немесе генерленген PTX-ке қосылған жөндеу ақпаратына сүйене отырып, генерленген кодты оны тудырған DSL нұсқауына сәйкестендіру әлдеқайда қарапайым. DSL-дің осы аспектілерінің арқасында біз кернелдерді әзірлеу кезінде өзімізді әлдеқайда өнімдірек сезіндік.

Кернелдер модельдер арасында көптеген ұқсастықтарды бөліссе де, олар әлі де өнімділікке әсер ете алатын бірқатар конфигурациялық түймешіктермен (knobs) жабдықталуы тиіс. Тиісті параметрлерге мыналар жатады:

Жасырын өлшем (Hidden Dimension) / Бас өлшемі (Head Dimension) / Назар аудару бастарының саны (Number of Attention Heads)

Кіріс/шығыс дерек түрлері

Кванттау схемасы

Құрастыру кезіндегі тұрақтылар (ығысулар (biases), активация функциялары және т.б.)

ROSE архитектурасы, оның ішінде қозғалтқышы, моделі, ядро кітапханасы және қызметтері.

Құрастыру кезіндегі мамандану өнімділік пен жөндеу үшін де пайдалы. Әсіресе GPU-дағы сияқты конвейерлік процессорда әрқашан орындалатын немесе ешқашан орындалмайтын тармаққа қарағанда жоқ тармақ жақсырақ. Сонымен қатар, кернелді конфигурациялық түймешіктердің белгілі бір жиынтығы үшін компиляциялау аз PTX/SASS нұсқауларын береді, оларды зерттеу кезінде cuda-gdb пайдаланған кезде талдау және жөндеу оңайырақ болады.

Берілген топтама өлшемі (batch size) мен конфигурация үшін DRAM-ға болатын трафик пен арифметикалық операциялардың санын бағалауға болады. Содан кейін то торды оңтайлы бөлу, синхрондау және кейбір аппараттық мүмкіндіктерді пайдалану үшін әр жағдай бойынша оңтайлы шешімдер қабылдауға болады. CuTeDSL компиляция уақытындағы қабылданбайтын үстеме шығынға ұшырамай, терең мамандану үшін компоненттерді анықтауға және қайта пайдалануға мүмкіндік береді.

Жалпы алғанда, инференс кернелдерінің өнімділігі екі бөлек қолдану жағдайы үшін оңтайландырылуы тиіс:

Префил (Prefill): кішкентай топтама өлшемі (аз бөліктер, тізбектер), әр тізбек үшін көп токендер. Эмбеддингтер мен LLM кіріс токендері үшін өткізу қабілеттілігін барынша арттыру керек. Кідіріс ондаған мен жүздеген миллисекунд аралығында.

Decode: шағыннан үлкенге дейінгі batch өлшемдері, бір (no speculative decoding) токеннен бастап әр sequence үшін бірнеше токенге дейін (draft-target speculation, multi-token prediction). Latency kernel іске қосу latency-іне жақын.

GEMM және назар аудару (attention) үшін префил кезеңі көптеген токендерде жұмыс істейді және өнімділік құрылғының есептеу өткізу қабілетімен шектеледі. Декодтау кішкентай топтама өлшемдеріндегі ұқсас операцияны қамтиды, мұнда орындалу уақыты DRAM-нан активациялар мен салмақтарды ағындық мультипроцессорлар мен тензорлық ядроларға жүктеу құнымен басым болады, оның орнына құрылғының жад өткізу қабілетімен шектеледі. Топтама өлшемі немесе топтамадағы токендердің жалпы саны салыстырмалы түрде айқын көрсеткіш болғанымен, префил мен декод арасындағы ауысу шегі нәзіктерек, себебі жад өткізу қабілеті сонымен қатар модельдің жасырын өлшемімен және активацияның дерек түрімен анықталады. Демек, егер алгоритмдік айырмашылықтар немесе торды/блокты тағайындаулар олардың кез келгеніне артықшылық берсе, кернелдер префил және декод іске асыруларына бөлінеді.

Префил және декод кернелдері арасында әмбебап түрде ерекшеленетін параметр – блоктағы тредтер (warps) саны (тредтер саны әр тред үшін 32 тред болатын warp өлшеміне бөлінеді). Бұл параметрді компиляция кезіндегі тұрақты ету ұсынылады, өйткені ол регистрді бөлуге көмектесе алады. CUDA-да бұл әдетте __launch_bounds__ арқылы тредтер санын көрсету және warp санын шаблон параметрі ретінде қамтамасыз ету арқылы жасалады:

cpp

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

}

Блоктағы барлық тредтерге қолжетімді регистрлер саны аппараттық шектеу болып табылады (Hopper және Blackwell жүйелерінде 64К). Егер тредтер саны алдын ала белгілі болса, компилятор әр тредке қанша регистр тағайындай алатынын дәл біледі және агрессивті түрде оңтайландырылған кернелдерді сәтті компиляциялай алады. Дегенмен, бұл компиляция кезіндегі шығындармен келеді, өйткені енді кернел әртүрлі warp өлшемдері үшін қайта компиляциялануы керек. CuTeDSL көмегімен бұл енді мәселе емес, өйткені біз барлық конфигурациялық кеңістік үшін компиляция құнын алдын ала төлеудің қажеті жоқ. DSL сонымен қатар блок өлшемінің тұрақты өрнек болуын күтеді, бұл CUDA-дағы __latch_bounds__ білдіретін оңтайландыруды жанама түрде орындайды.

Warp санын тұрақты етуден басқа, оны префил мен декод арасында өзгерту ұсынылады. GPU жадының өткізу қабілетін толығымен пайдалану үшін әрқашан блоктағы warp санын жоғары етіп тағайындау жеткіліксіз, өйткені толық өткізу қабілетіне қол жетімді ағындық мультипроцессорлардың (SM) мүмкіндігінше көбін пайдалану арқылы ғана қол жеткізуге болады. Демек, берілген токендер саны пен белгілі дерек түрлері мен жасырын өлшемдер үшін біз алдымен генерленген кодтың 128 биттік жүктеу/сақтау нұсқауларын (Blackwell жүйесінде 256 бит) пайдалануын қамтамасыз ету үшін өр-тредтік (per-thread) оңтайлы вектор өлшемін (bfloat16 үшін 8, float32 үшін 4) таңдаймыз. Вектор өлшемін біле отырып, біз кернелге қажетті warp санын есептей аламыз. Біз алдымен SM санын барынша арттыруға тырысамыз: егер warp саны SM-нен аз болса, әр блок үшін жалғыз warp таңдалады. Бұл әдетте кернелдің декод нұсқасы болып табылады. Егер warp жеткілікті болса, біз өнімділіктің жақсарғанын көруді тоқтататын жоғарғы шекті бенчмарктей отырып, warp санын біртіндеп 8, 16 немесе 32-ге дейін арттырамыз.

Triton бұл оңтайландырулардың кейбірін өз бетінше орындағанымен, біз жоғары деңгейлі компилятордың кемшіліктерімен шектелмей, CuTe орналасу алгебрасын мұқият пайдалану арқылы ұқсас нәтижелерге қол жеткізе алатынымызды анықтадық. Әсіресе warp-ішілік көлденең қысқартулар (sub-warp horizontal reductions) қатысқан кезде, CuTeDSL бізге оңтайлы кодты шығару үшін қажетті дәл басқаруды ұсынғанын байқадық.

Кейбір жағдайларда жұмысты мүмкіндігінше көп блоктарға бөлу үшін барлық қатысушы блоктарды синхрондау үшін cooperative_groups::this_grid().sync() пайдаланылуы мүмкін. Шектеулер қолданылады: блок орындала бастағаннан кейін SM-нан кестеден тыс қала алмайтындықтан, алға жылжуды қамтамасыз ету үшін блоктар саны құрылғыда қолжетімді SM санынан аспауы керек, бұл кез келген уақытта барлық блоктар синхрондауда қатысатын жаһандық жад орындарын жаңартып, сұрау ала алатынын қамтамасыз ету үшін. Бұл шектеуді cudaLaunchCooperativeKernel мәжбүрлейді. Тор барьерінің құны Hopper жүйесінде шамамен 3 мкс құрайды, сонымен қатар кейбір жағдайларда жадтың жанама барьері арқылы барлық алдыңғы жад операцияларын тазарту құны қосылады. Жалпы алғанда, бұл кернелді екі іске қосуға бөлудің үстеме шығынымен салыстырмалы, бұл өнімділікке сезімтал жобалау шешіміне әкеледі. Programmatic Dependent Launch (PDL) Hopper-ден бастап іске қосу үстеме шығындарын азайта алатынына қарамастан, декод кернелдері үшін үстеме шығынды болдырмау тұрарлық.

Егер жұмыс жүктемелерінің біріне 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 all-to-all кернелдеріміздің жағдайында, төменде егжей-тегжейлі көрсетілгендей, тор барьерлерімен немесе онсыз кернелдерді жасау және оларды компиляция кезінде оңтайландыру икемділігіне ие болу бізге CuTeDSL арқылы айтарлықтай өнімділік жетістіктеріне қол жеткізуге көмектесті.

CuTeDSL арқылы жасаған бірінші кернел шағын өлшемдер үшін оңтайландырылған арнайы RMS норма кернелі болды. Соңғы Qwen модельдерінің өзін-өзі назар аудару (self-attention) қабатында қолданылатын QK нормасы 64, 128 немесе 256 элементте бас (head) негізінде қолданылады. Мыңнан асатын жасырын өлшемдерде жұмыс істейтін және ортақ жад пен регистрлер арқылы көлденең қысқартуды қажет ететін кәдімгі RMS нормасынан айырмашылығы, QK нормасына тек warp-тар арасында қысқарту қажет, бұл бірнеше бастарды бір warp-қа сығымдау мүмкіндігі бар. Бұл біз барлық бас өлшемдері мен дерек түрлері үшін ең жақсы токен-warp тағайындауларын табу үшін компиляция кезінде агрессивті түрде оңтайландыра алатын қолдану жағдайына айналды.

RMS нормасы вектораның квадратты нормасын есептеу үшін көлденең қосындыны қажет етсе де, бұл біршама қарапайым операция:

cpp

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

Назар аудару бастары әдетте кішкентай болғандықтан және өр-тредтік операциялар ауқымды түрде векториалданғандықтан, бір warp квадратты норманы есептей алады. Жад өткізу қабілетін барынша арттыру үшін әр тред кірістің 128 битін өңдеуі тиіс, бұл 8 bfloat16 мәніне сәйкес келеді. Демек, басты өңдеу үшін жартылай warp (half-warp) жеткілікті, ал квадратты норманы shuffle_sync_bfly арқылы орындалатын көлденең қысқарту көмегімен анықтауға болады. Бас өлшемі мен дерек түрінен бастап біз әр warp үшін бастар санын анықтап, компиляция кезінде аргументтер мен қысқарту раундтарының санын көрсете аламыз.

Содан кейін біз автоматты түрде векториалданған жүктеулерге/сақтауларға сүйене отырып және warp-ты қысқарту үшін төмен деңгейлі архитектураға тән операцияларды пайдалана отырып, салыстырмалы түрде аз код жолдары бар арнайы кернелді құрастыра аламыз. Кернел қай баспен (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 кернелдері тиісті матрицалық көбейтулерді орындағаннан кейін салмақталған жинақтауды орындайды. Әртүрлі өнімдер әртүрлі өлшемдегі 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 арқылы жұмыс істейтін біррангіден-біррангіге (Peer-to-Peer) деректер-параллель және тензорлық-параллель кернел. Көп түйінді Hopper орналастырулары үшін қолданылады.

Жергілікті кернел маршруттау кестесіне негізделген әрбір сарапшыға тағайындалған токендер санын санайды және әрбір сарапшыға тиесілі токендер тізіміндегі әрбір токенге бірегей индекс тағайындайды. Маршруттау қадамынан кейін токендер MoE GEMM кернеліне жіберілу үшін қосымша толтырылған буферлерге көшіріледі. Токен орындарын пайдалана отырып, біріктіру кернелі содан кейін регистрлерде жинақтауды орындайды, соңғы токендерді шығыс буферіне жазады.

Тарату фазасы префил мен декод арасында өзгереді: аз токендер үшін оқып шығуға болатын маршруттау ақпараты аз, сондықтан біз оны бір блокта өңдейміз және тордың қалған бөлігіне белгі береміз, ал жеке блоктар токендердің бөліктерін көшіреді. Токенді тағайындау ортақ жадтағы санауыштарды атомдық түрде ұлғайту арқылы анықталады. Префил көбірек маршруттау ақпаратын өңдеуі тиіс болғандықтан, маршруттау бөлек кернелде бірнеше блоктарға бөлінеді. Индекстер алдымен ортақ жадтағы санауыштарды пайдаланып жергілікті түрде тағайындалады, содан кейін әрбір токеннің орнын анықтау үшін жаһандық жадтағы барлық блоктар бойынша жинақталады. Тор деңгейіндегі барьерлерді пайдалана алу үшін бұл кернелдің торы құрылғыдағы SM санына дейін қысылады. Кейінгі кернел қолжетімді жад өткізу қабілетін толығымен пайдалану үшін әлдеқайда көп блоктар мен warp-тар арасында іске қосылған көшірмелерді орындайды.

Біріктіру кернелі салыстырмалы түрде қарапайым, өйткені ол маршруттау ақпаратына және тарату арқылы есептелген ығысуларға сүйене отырып, векториалданған жинақтауды орындайды. CuTeDSL мұнда құнды екенін көрсетті, өйткені есептеу барлық дерек түрлері мен жасырын өлшемдер үшін тиімді кодқа компиляцияланатын қарапайым CuTe тензорлары арқылы көрсетіле алады.

Тензорлық-параллель орналастыруда маршруттау ақпараты мен кіріс токендері репликаланады. Жергілікті тарату кернелі тек ағымдағы рангте орналасқан сарапшыларға тағайындалған токендерді бағыттау және таңдау үшін шағын өзгеріспен толығымен қайта пайдаланылуы мүмкін. Кейбір параметрлермен және құрастыру кезіндегі маманданумен жергілікті тарату кернелдері қайта пайдаланылады.

Рангтер саны әдетте токен бағытталған сарапшылар санынан аз болатындықтан, орташа жағдайда әрбір токеннің кем дегенде бір репликасы әр рангке бағытталған деп санауға болады. Бұл жағдайда жергілікті біріктіру кернелі токеннің жергілікті көшірмелерін жинақтау үшін қайта пайдаланылуы мүмкін, соңынан қатысушы рангтер бойынша қосынды табу үшін барлығын азайту (all-reduce) орындалады. Префил кернелі жағдайында біз бөлек кернелдерді іске қосамыз, алайда декод екі операцияны да бір іске қосуға біріктіреді.

Барлығын азайту үшін, егер тек екі ранг қатысса, біз қарапайым NVLink қолданамыз, өткізу қабілетінің 10-20% өсуі үшін 4 немесе одан da көп ранг қатысқан кезде NVLS-ке ауысамыз. DSL арқасында біз MoE және кәдімгі all-reduce кернелдері арасында қысқарту логикасының көп бөлігін және рангтер аралық буфер алмасуларын бөлісе аламыз. Бұл аппаратты қатаң басқару тиімді болатын жағдай, өйткені NVLS үшін multimem.ld_reduce және multimem.st арқылы DSL операциялары қажет. Қысқартудың өзін іске асыру негізінен NCCL тесттерінен all_reduce мысалына еліктейді.

Деректер параллель кернелдері бірнеше құрылғылар үшін токендерді жинақтау үшін NVLink-ке сүйенеді. Тарату фазасында токендерді алмастырмас бұрын рангтер әр жіберіліп жатқан токен үшін әр рангтегі нысана ығысуды анықтау үшін алдымен жергілікті түрде бағыттайды. Содан кейін олар қақтығыстарды болдырмау үшін әр жіберуші үшін бөлінген қабылдау буферлеріне токендерді көшіре бастайды. Аяқталғанын белгілеу үшін токен бойынша шығару-сақтау (release-store) жазбалары қолданылады. Қабылдаушы жағында MoE GEMM үшін барлық токендерді бір буферге жинақтау үшін маршруттау ақпараты жинақталады. Қабылдаушылар жазулардың аяқталуын күту үшін сатып алу-жүктеу (acquire-load) арқылы ранг жазған жалаушаны сұрайды.

Біріктіру жағында токендер қабылдағанын көрсететін токен бойынша жалаушасы бар NVLink арқылы бастапқы рангтерге көшіріледі. Біріктіру кернелі адресациядағы шағын айырмашылықтармен және кейбір қосымша синхрондау логикасымен жергілікті және біртүйінді нұсқаларына ұқсас.

Бұл кернелді іске асыру үшін біз жүйелік және құрылғылық деңгейде де ауқымдалған шығару-сатып алу жүктеулері мен сақтауларын ашу үшін PTX орауыштарының кітапханасын құрдық. Синхрондау жалаушаларының түйіршіктілігін жақсарту және NVLink арқылы токендерді жіберу және қабылдау үшін көбірек блоктарды бөлу арқылы біз алдыңғы CUDA кернелдерімізге қарағанда 2 есе жақсартуларға қол жеткіздік.

Біздің Peer-to-Peer кернелдеріміз бұрын ConnectX-7 InfiniBand адаптерлерімен жұптастырылған Hopper GPU-ларында ең озық өнімділікке қол жеткізген болатын. Содан бері біз өнімділікті одан әрі арттыру үшін біріктіру кернелдерін CuTeDSL-ге көшірдік. CUDA нұсқалары бұрыннан да құрастыру кезіндегі мамандануға қатты сүйенгендіктен, олар біздің репозиторийде компиляция уақытының кедергілері екенін көрсетті, бұл бізді кейбір параметрлерді алып тастауға мәжбүр етті. JIT компиляцияланатын DSL-ге көшу арқылы біз көбірек параметрлер үшін мамандандық және торды бөлуді түзеттік, бұл кідірісті шамамен 3%-ға арттыру үшін префил мен декод арасындағы қосымша айырмашылықтарды енгізді.

ROSE кернелдері тұрғысынан Blackwell қолдауы екі санатқа бөлінеді: архитектураға тәуелсіз және архитектураға тән кернелдер. Тензорлық ядроларға сүйенетін архитектураға тән кернелдер бейімделу үшін нақты жұмысты талап етеді. Дегенмен, DSL қолданылуының арқасында архитектураға тәуелсіз кернелдер оңай бейімделеді.

Blackwell GPU-сында Hopper-дің 132 SM санына қарағанда 160-қа дейін SM бар. Көптеген кернелдер үшін бұл компиляция кезіндегі тұрақты болып табылады және хост жағындағы іске қосу кодындағы логика іске қосылатын оңтайлы (префил немесе декод) кернелді анықтау үшін оны ескереді. Сонымен қатар, Blackwell сонымен қатар жад өткізу қабілетін пайдалануды жақсарту үшін вектор енін (CuTe тензорларының соңғы өлшемі) арттыру арқылы оңай шығарыла алатын 256 биттік жүктеу нұсқауларын ұсынады.

Blackwell кванттау схемаларының айтарлықтай саны үшін аппараттық қолдауды ұсынады (MXFP8, MXFP4, NVFP4). Бұларға қолдау көрсету негізінен FP8 блок-тық масштабтау үшін бар кернелдерді түзету арқылы іске асырылды. JIT қолдауы мұнда өте маңызды, өйткені жаңа режимдер конфигурациялық кеңістікті шамадан тыс арттырмай қолдау тап алады. Бұрын біз Triton-ның аз элементтер бойынша, warp-тың жартысы немесе ширегі бойынша көлденең қысқартуларды көрсетуде ыңғайсыз екенін байқадық. NVFP4 және MXFP8/MXFP4 үшін сәйкесінше 16 немесе 32 элемент бойынша абсолютті максималды мәнді есептеу кезінде оңтайлы өнімділікке қол жеткізу үшін қысқартуларды дәлірек басқару өте маңызды болды.

Салыстырмалы түрде жаңа тіл ретінде LLM модельдері алдын ала оқыту кезінде CuTeDSL кодын аз көрген. Сонымен қатар, біздің қазіргі кодымыздың стилі, активациялар мен нормалар үшін оңтайландырылған кернелдері бар, TMA және архитектураға тән матрицалық көбейту нұсқаулары үшін қатты оңтайландырылған бар CUTLASS немесе Quack мысалдарынан сәл өзгеше болды. LLM модельдері кернел жазуға көмектесе алатын деңгейге жету біртіндеп жүретін процеске айналды.

Біз агенттерге кернелді түрлендіру үшін сенім артудан бұрын, дәйекті үлгілерді сақтай отырып, кернелдер мен утилиталық функциялардың бастапқы жиынтығын қолмен құрдық. Жиын дамыған сайын біз CUDA-дан CuTeDSL-ге өрескел аудармаларды орындай алдық, содан кейін айтарлықтай қолмен қайта өңдеулер жүргізілді. Атап айтқанда, біз модельдердің көрсеткіштен тензорға (pointer-to-tensor) түрлендірулермен және векторизациямен қиналатынын, тестті компиляциялауға тырысқанда типтік қателерді түзетуге тырысып қалатынын байқадық. Олар кернелдердің торға негізделген табиғатын нашар түсінетінін көрсетіп, тред пен торды синхрондауды шатастырды. Ең бастысы, өздеріміздің 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 Нормы (QK Norm)

Mixture-of-Experts Тарату/Біріктіру (Dispatch/Combine)

Жергілікті (Local)

Біртүйінді (Singlenode)

Деректер паралеллизмі (Data Parallel)

Біррангіден-біррангіге (Peer-to-Peer)

Blackwell қолдауы

ЖИ көмегімен кернелдерді жасау (AI-Assisted Kernel Generation)

Қорытынды және болашақ жұмыстар

Әдебиеттер тізімі

Дәйексөз