CuTeDSL в Perplexity

В Perplexity използваме собствената Runtime-Optimized Serving Engine (ROSE) за обслужване на модели – от ембединги до езикови модели (LLM) с трилиони параметри. Разположено зад всички API, включително Sonar, Search и Embeddings, нашето софтуерно време за изпълнение беше създадено, за да даде възможност на изследователите да изследват най-подходящите модели и архитектури, като същевременно бързо внедряват най-добрите кандидати в производство. За да постигнем това, трябва да можем да внедряваме, разполагаме и оптимизираме инференса за съвременни архитектури с бързи темпове. В тази публикация в блога изследваме как интегрираните в нашия инференс двигател CuTeDSL ядра ни позволяват да изградим необходимите GPU ядра, за да изведем моделите до пикова производителност на NVIDIA GPU.

Персонализираните модели, изградени от различни екипи в Perplexity, се хостват вътрешно на NVIDIA Hopper и Blackwell GPU с помощта на инференс двигателя ROSE, създаден от екипа по AI Inference. ROSE първоначално е създаден за по-добро обслужване на персонализирани модели Llama както за деобфускация на език, така и за класификация, с интерфейс, съвместим с NVIDIA Triton Inference Server. С течение на времето той се разрасна в пълноценен инференс двигател, обслужващ значителни LLM, както и модели за класиране, класификация, оценяване и други трансформерни модели за широк спектър от екипи, които да занаят продуктите на Perplexity. Проектната цел на ROSE е да бъде проста: с времето се инвестират усилия за достигане на пикова производителност, но основният приоритет е да се адаптира към непрекъснато развиващите се архитектури и случаи на употреба за нашите модели.

Пространството за конфигурация на ядрото при различни типове данни, размери на главите, константи по време на компилация и настройки на стартиране.

На високо ниво ROSE се състои от двигател, който адаптира модел към интерфейс, с който клиентите могат да взаимодействат. Двигателите са силно многократно използваеми, като управляват планирането и пакетирането на заявките и извличат деобфускирани токени, ембединги или резултати от базовите модели. Инициализацията на устройствата, междупроцесните и междувъзловите комуникационни канали и зареждането на тегла се управляват от тях. За LLM двигателите отговарят за пакетирането, разделянето на части, семплирането. Той отговаря и за разпределението на KV хранилище както за пълно, така и за линейно внимание, както и за съпоставяне на префикси. Двигателите за ембединги са по-опростени, тъй като извършват само онлайн пакетиране. Изпълнението е оптимизирано така, че двигателят подготвя следващия кръг от работа, докато моделът работи едновременно върху ускорителя.

ROSE излага богата колекция от потребителски слоеве за моделиране, обгръщащи GPU ядра, изградени с CuTeDSL, Triton, CUDA, CUTLASS и cuBLAS. За умножение на матрици и внедряване на внимание разчитаме предимно на NVIDIA ядра от CUTLASS и cuBLAS. Останалата част от пайплайна, включително ембединг, норма, MoE и активационни ядра, са част от библиотеката от ядра на ROSE. За да поддържа най-добре хардуера, на който стартираме инференс, ROSE е оборудван със собствени примитиви за комуникация между устройствата за MoE маршрутизиране, изпращане и комбиниране. Въпреки че другите ядра са относително прости, те трябва да бъдат напълно специализирани, за да постигнат пикова производителност на всички различни скрити измерения и устройства, на които работят. Освен това разчитаме на специализация по време на компилация, за да поберем всички настройки и корекции на различни модели (тегло пристрастие спрямо липса на тегло пристрастие за RMS норма, активиращи функции и т.н.). В резултат на това решихме да използваме DSL за тяхното внедряване: първоначално изградихме ROSE с помощта на Triton. Въпреки това оттогава приехме CuTeDSL като основна среда за програмиране на GPU.

От математическа гледна точка инференс ядрата обикновено са неосложнени. Следователно при разработката на ядро сложността обикновено се крие не в изразяването на това какво изчисление трябва да се извърши, а в това как точно хардуерът трябва да го направи. За да изградим най-добрите ядра, се нуждаем от език, който предлага правилните абстракции върху базовия хардуер. Правилният език ни позволява да специфицираме операциите на високо ниво за извършване на изчислението, като ги специализираме агресивно и ги оптимизираме във всички точки в широко пространство за конфигурация.

CuTeDSL е базиран на Python DSL за NVIDIA GPU, изграден върху алгебрата на оформлението CuTe и MLIR. Той се компилира just-in-time до силно оптимизиран PTX, като същевременно запазва контрола върху хардуерните примитиви на ниско ниво. Целта му е да предложи същите функции и производителност като CUTLASS C++, без скъпите времена за компилация, които възпрепятстват скоростта на разработка. От езиците, които сме проучили или с които сме работили (включително CUDA, cuTile, TVM IR и Triton), го смятаме за най-подходящия за разработването на нашите ядра.

Като цяло, макар и да не е толкова високо ниво и лаконичен като Triton или cuTile, установихме, че CuTe Layout Algebra и неговите композиции (изложени чрез cute.Tensor) са разумни заместители. Освен това CuTeDSL предлага неограничен достъп до хардуерни примитиви, осигурявайки ни нивото на контрол, от което се нуждаем, за да постигнем пикова производителност във всички крайни случаи. С CUDA възможностите за специализация по време на компилация бяха ограничени, тъй като всеки допълнителен параметър добавяше неразумна цена към времето за компилация или сложност към системата за изграждане. За разлика от това, с Just-In-Time Compiled DSL можем да специализираме по-агресивно, без да жертваме темпото на итерация. За разлика от TVM IR, който използва CUDA като междинна стъпка, установихме, че базираният на MLIR пайплайн, който се спуска до PTX, е по-здрав и по-лесен за отстраняване на грешки от този, който излъчва CUDA като междинна стъпка. Освен това с използването на потребителски операции за обвиване на PTX инструкции установихме, че можем напълно да се възползваме от всички функции на базовия хардуер.

Голяма част от разходите за компилация на CUDA ядрата обикновено се отделят за фазата на разширяване на шаблона на ядрата, тъй като оптимизацията и компилацията от етапа на PTX нататък могат да бъдат завършени за секунди. С много параметри за време на компилация този етап на компилация може да нарасне до множество минути за необикновено ядро. Тази тесен място се подчертава от факта, че възможностите за паралелна компилация възникват едва по-късно, на нивата IR и PTX, като обработката на шаблони на C++ се случва последователно. Въпреки че съществуват трикове за справяне с това чрез преместване на инстациранията на шаблони в отделни компилационни единици, те изискват значителни настройки на системите за изграждане и използване на неудобни макроси. За разлика от тях, CuTeDSL ядрата, които се компилират Just-In-Time (JIT), не изискват сложни правила за изграждане и не пречат на времето за компилация, тъй като цикълът на разработка може да бъде насочен към една точка в пространството за конфигурация, докато внимателните модулни тестове по-късно могат да покрият цялото пространство за проектиране на ядрото.

Въпреки че JIT компилацията премества разходите от времето за разработка към времето за изпълнение, надвишаването е повече от приемливо. Когато даден модел е внедрен, той активира само няколко от точките от цялото пространство за конфигурация, което води до надвишаване на мащаба на секундите по време на стартирането/загряването на инференс двигателя. За разлика от това, предварителното компилиране на всички конфигурации на ядро може да обхване множество минути. Освен това надвишаването може да бъде допълнително смекчено чрез кеширане на JIT-компилираните артефакти.

Голямата цена за компилация на CUDA се подчертава от трудността при проследяване на произхода на грешките, когато се използва базирано на шаблони разпределение и специализация, както е в случая с CuTe и CUTLASS. С DSL грешките могат да бъдат насочени към реда код, който го е генерирал, което значително улеснява отстраняването на грешки. Освен това е много по-опростено да се картографира генерираният код към DSL изявлението, което го е произвело чрез междинни MLIR етапи или като се разчита на информация за отстраняване на грешки, прикачена към генерирания PTX. Благодарение на тези аспекти на DSL, установихме, че сме много по-продуктивни при разработването на ядра.

Въпреки че ядрата споделят много сходства между моделите, те все пак трябва да бъдат оборудвани с набор от конфигурационни копчета, които могат да окажат влияние върху производителността. Съответните параметри включват:

Скрито измерение/Измерение на глава/Брой глави на внимание

Типове данни за вход/изход

Схема за квантоване

Константи за време на компилация (пристрастия, активиращи функции и др.)

Архитектурата ROSE, включително нейния engine, модел, библиотека на ядрото и услуги.

Специализацията по време на компилация е полезна както за производителността, така и за отстраняването на грешки. Особено в конвейерен процесор, какъвто е случаят с GPU, несъществуващият клон е по-добър от винаги взет или никога невзет клон. Освен това компирането на ядро за определен набор от конфигурационни копчета води до по-малко PTX/SASS инструкции, които са по-лесни за анализиране и отстраняване на грешки при използване на cuda-gdb по време на разследване.

За даден размер на партидата и конфигурация могат да бъдат оценени очакваният трафик към DRAM и броят на аритметичните операции. Впоследствие могат да бъдат взети оптимални решения за всеки отделен случай за оптимално разпределение на мрежата, синхронизация и използване на определени хардуерни функции. CuTeDSL ни позволява да дефинираме и използваме повторно компоненти за дълбока специализация, без да предизвикваме неприемливи режийни разходи за време на компилация.

Като цяло производителността на инференс ядрата трябва да бъде оптимизирана за два отделни случая на употреба:

Предварително зареждане: малък размер на партидата (малко парчета, последователности), много токени на последователност. Пропуск способността трябва да бъде максимизирана както за ембединги, така и за входните токени на LLMs. Латентност в десетки до стотици милисекунди.

Декодиране: размери на пакета от малки към големи, един (без спекулативно декодиране) до няколко токена на последователност (спекулация между чернова и цел, прогнозиране на множество токени). Закъснение, близко до латентността при стартиране на kernel.

За GEMM и внимание предварително зареденият етап работи върху голям брой токени и производителността е ограничена от изчислителната производителност на устройството. Декодирането включва същата операция при малки размери на партидите, като времето за изпълнение е доминирано от цената на зареждане на активациите и теглата от DRAM в стрийминг мултипроцесорите и тензорните ядра, ограничено от честотната лента на паметта на устройството вместо това. Въпреки че размерът на партидата или общият брой токени в партидата е относително ясен индикатор, прагът на превключване между предварително зареждане и декодиране е по-фин, тъй като честотната лента на паметта се определя и от скритото измерение на модела и типа данни на активацията. Следователно ядрата се разделят на реализации за предварително зареждане и декодиране, ако алгоритмичните разлики или назначенията на мрежа/блок биха дали предимство на някое от тях.

Параметърът, който е универсално различен между предварително заредените и декодиращите ядра, е броят на усукванията в блок (броят на темите, разделен на размера на усукване от 32 теми на усукване). Желателно е този параметър да бъде константа за време на компилация, тъй като може да помогне за разпределението на регистрите. В CUDA това обикновено се прави чрез указване на броя на темите чрез __launch_bounds__ и предоставяне на броя на усукванията като параметър на шаблон:

cpp

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

}

Броят на регистрите, достъпни за всички теми в блок, е хардуерно ограничение – 64K на Hopper и Blackwell. Ако броят на темите е известен предварително, компилаторът знае точно колко регистъра може да назначи на всяка тема и може да успее да компилира по-агресивно оптимизирани ядра. Това обаче е свързано с разходи за време на компилация, тъй като сега ядрото трябва да бъде прекомпилирано за различни размери на усукванията. С CuTeDSL това вече не е проблем, тъй като не е нужно да плащаме разходите за компилация за цялото пространство за конфигурация предварително. DSL също така очаква размерът на блока да бъде константен израз, като имплицитно извършва оптимизацията, заложена от __latch_bounds__ в CUDA.

Освен че прави броя на усукванията константа, е желателно да се променя между предварително зареждане и декодиране. За да се използва напълно честотната лента на паметта на GPU, не е достатъчно винаги да се назначава висок брой усуквания на блок, тъй като пълната честотна лента може да бъде постигната само чрез използване на колкото се може повече от наличните Streaming Multiprocessors (SMs). Следователно за даден брой токени и известни типове данни и скрити измерения първо избираме идеален вектор на тема размер (8 за bfloat16, 4 за float32), за да гарантираме, че генерираният код използва 128-битовите инструкции за зареждане/запис (256-битови на Blackwell). Познавайки размера на вектора, можем да изчислим броя на усукванията, изисквани от ядрото. Първо се опитваме да максимизираме броя на SM: ако има по-малко усуквания от SM, се избира едно усукване на блок. Това обикновено е дефинираният вариант на декодиране на ядрото. Ако има достатъчно усуквания, постепенно увеличаваме броя на усукванията до 8, 16 или 32, като сравняваме горна граница, след която обикновено вече не виждаме подобрение в производителността.

Въпреки че Triton извърши някои от тези оптимизации под капака, установихме, че можем да постигнем подобни резултати чрез внимателно използване на CuTe layout алгебра, без да бъдем ограничени от недостатъците на компилатор с високо ниво. Особено когато са включени хоризонтални редукции под усукването, установихме, че CuTeDSL ни предложи фино зърнестия контрол, необходим за генериране на оптимален код.

В някои ситуации, за да се разпредели работата между възможно най-много блокировки, cooperative_groups::this_grid().sync() може да се използва за синхронизиране на всички участващи блокировки. Прилагат се ограничения: тъй като блокът не може да бъде отписан от SM, след като започне изпълнението, за да се гарантира напредък в движение, броят на блоковете не може да надвишава броя на наличните SM на устройството, за да се гарантира, че във всеки момент всички блокове могат да актуализират и анкетират глобалните местоположения в паметта, участващи в синхронизацията. Това ограничение се налага от cudaLaunchCooperativeKernel. Цената на мрежова бариера е около 3us на Hopper, плюс цената на изчистване на всички предишни операции в паметта чрез имплицитна бариера в паметта в някои ситуации. Като цяло това е сравнимо с режийните разходи за разделяне на ядро на две стартирания, което води до чувствително към производителността решение за дизайн. Въпреки факта, че 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 норма ядро, оптимизирано за малки размери. 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-3 пъти по-добра пропуск способност и латентност както за предварително зареждане, така и за декодиране.

MoE изпращането и комбинирането разбъркват входните токени към експертите, към които са маршрутизирани, извършвайки претеглено натрупване, след като MoE GEMM ядрата извършват съответните умножения на матрици. Производителността е от особен интерес, тъй като различните продукти разчитат на MoE модели с различни размери. Въз основа на типовете внедряване отличаваме различни ядра за оптимизиране:

TP=1: Локално ядро за един GPU. Изпращането разбърква токените към съответните експерти, а комбинирането ги натрупва. Използва се основно за внедрявания на Qwen3 30B на H100.

DP=1, TP ≠ 1: Тензорно-паралелно ядро с един възел. Входът е репликиран, изпращането избира токените за текущия ранг, а комбинирането ги натрупва във всички рангове. Натрупване през NVLink или NVLS. Използва се за големи предварителни запълвачи.

DP ≠ 1, NVLink: Данно-паралелно ядро с незадължителен тензорен паралелизъм. Всеки ранг изпраща токени към всеки друг ранг, хостващ специфични експерти. Информацията за маршрутизиране се агрегира, токените се прехвърлят през NVIDIA NVLink. Използва се за данно-паралелни внедрявания на H200 или GB200.

DP ≠ 1, InfiniBand: Peer-to-Peer данно-паралелно и тензорно-паралелно ядро, работещо както през NVLink, така и през NVIDIA Quantum InfiniBand. Използва се за многовъзлови внедрявания на Hopper.

Локалното ядро преброява броя на токените, назначени на всеки експерт въз основа на таблицата за маршрутизиране, и задава уникален индекс на всеки токен в списъка с токени, принадлежащи на всеки експерт. След стъпката на маршрутизиране токените се копират в незадължително подплатени буфери, които да бъдат предадени на MoE GEMM ядро. Използвайки местоположенията на токените, ядрото за комбиниране след това извършва натрупване в регистрите, като записва крайните токени в изходен буфер.

Фазата на изпращане варира между предварително зареждане и декодиране: за малко токени има малко информация за маршрутизиране за четене, така че я обработваме в един блок и сигнализираме останалата част от мрежата, като отделни блокове копират парчета токени. Назначаването на токени се определя чрез атомно инкрементиране на броячи в споделената памет. Тъй като предварителното зареждане трябва да обработи повече информация за маршрутизиране, маршрутизирането се разделя на множество блокове в отделно ядро. Индексите първо се задават локално с помощта на броячи в споделената памет и след това се агрегират в цялата мрежа в глобалната памет, за да се определи позицията на всеки токен. Мрежата на това ядро е затегната до броя на SM, за да може да се използват мрежови бариери. Последващо ядро извършва копиранията, стартирани в много по-голям брой блокове и усуквания, за да се използва напълно наличната честотна лента на паметта.

Комбиниращото ядро е относително просто, тъй като извършва векторизирано натрупване, разчитайки на информацията за маршрутизиране и на отместванията, изчислени чрез изпращане. Тук CuTeDSL се оказа ценен, тъй като изчислението може да бъде изразено чрез прости CuTe тензори, които се компилират в ефективен код за всички типове данни и скрити измерения.

При тензорно-паралелно внедряване информацията за маршрутизиране и входните токени се репликират. Локалното ядро за изпращане може да бъде напълно използвано повторно с лека промяна за маршрутизиране и избор само на токените, назначени на експертите, разположени в текущия ранг. С някои параметри и специализация по време на компилация локалните ядра за изпращане се използват повторно.

Тъй като броят на ранговете обикновено е по-малък от броя на експертите, към които е маршрутизиран токен, в средния случай може да се приеме, че всеки токен има поне една от своите реплики, маршрутизирани към всеки ранг. В тази ситуация локалното комбинирано ядро може да бъде използвано повторно за агрегиране на локалните копия на токен, последвано от all-reduce за сумиране през участващите рангове. В случая на ядрото за предварително зареждане стартираме отделни ядра, но декодирането обединява двете операции в едно стартиране.

За all-reduce използваме обикновен NVLink, ако участват само два ранга, превключвайки към NVLS, когато са ангажирани 4 или повече ранга за 10-20% увеличение на пропуск способността. Благодарение на DSL можем да споделим голяма част от логиката на редукция и обмен на буфери между равни между MoE и обикновените all-reduce ядра. Това е ситуация, в която тесният контрол върху хардуера е от полза, тъй като за NVLS се изискват DSL операции над multimem.ld_reduce и multimem.st. Внедряването на самата редукция до голяма степен следва примера all_reduce от тестовете на NCCL.

Данно-паралелните ядра разчитат на NVLink за агрегиране на токени за множество устройства. Във фазата на изпращане, преди обмен на токени, колегите първо маршрутизират локално, за да идентифицират целевото отместване на всеки ранг за всеки токен, който изпращат. Впоследствие те започват да копират токените в получените буфери, разпределени за всеки подател, за да се избегнат конфликти. Записите за освобождаване-съхранение на токен се използват за сигнализиране на завършване. В края на приемника информацията за маршрутизиране се агрегира, за да може да се кондензират всички токени в един буфер за MoE GEMM. Приемниците анкетират флага, написан от връстници, с acquire-load, за да изчакат завършването на записите.

От страната на комбиниране токените се копират през NVLink към изходните рангове, с флаг за всеки токен, указващ получаването. Комбиниращото ядро е подобно на локалните варианти и вариантите с един възел, с малки разлики в адресирането и малко допълнителна логика за синхронизация.

Създадохме библиотека от PTX обвивки за излагане на освобождаване-придобиване натоварвания и магазини, обхванати както на ниво система, така и на ниво устройство, за да се реализира това ядро. Чрез подобряване на гранулирането на флаговете за синхронизация и разпределение на повече блокове за изпращане и получаване на токени през NVLink, постигнахме 2 пъти подобрение спрямо предишните ни CUDA ядра.

Нашите Peer-to-Peer ядра преди това постигнаха най-съвременна производителност на Hopper GPU, сдвоени с ConnectX-7 InfiniBand адаптери. Оттогава портнахме комбинираните ядра към CuTeDSL за допълнителни печалби в производителността. Тъй като версиите на CUDA вече разчитаха силно на специализация по време на компилация, те се оказаха тесните места за време на компилация в нашия хранилище, което ни принуди да пропуснем някои параметри. Чрез мигриране към компилирания с JIT DSL ние специализирахме за повече параметри и настроихме мрежовите разпределения, въвеждайки допълнителни разлики между предварително зареждане и декодиране за около 3% тласък на латентността.

От гледна точка на ROSE ядрата, поддръжката на Blackwell попада в две категории: архитектурно-агностични и архитектурно-специфични ядра. Архитектурно специфичните ядра, разчитащи на тензорни ядра, изискват изрична работа за адаптиране. Въпреки това, благодарение на използването на DSL, архитектурно-агностичните ядра се адаптират незначително.

GPU на Blackwell има до 160 SM, за разлика от 132 на Hopper. За повечето ядра това е константа за време на компилация и логиката в кода на стартера от страна на хоста я отчита при определяне на оптималното (преди зареждане срещу декодиране) ядро за изпълнение. Освен това Blackwell въвежда и 256-битови инструкции за зареждане, които могат лесно да бъдат излъчени чрез увеличаване на ширината на вектора (последното измерение на CuTe тензорите), за да се подобри използването на честотната лента на паметта.

Blackwell въвежда хардуерна поддръжка за значителен брой схеми за квантоване (MXFP8, MXFP4, NVFP4). Поддръжката за тях беше внедрена основно чрез настройване на съществуващи ядра за мащабиране по блокове FP8. JIT поддръжката е от решаващо значение тук, тъй като новите режими могат да се поддържат без избухване на пространството за конфигурация по време на компилация. В миналото открихме, че Triton е неудобен при изразяването на хоризонтални редукции през малко елементи, над половин или четвърт от усукване. По-финният контрол върху редукциите беше от решаващо значение за постигане на оптимална производителност при изчисляване на абсолютна максимална стойност съответно над 16 или 32 елемента за NVFP4 и MXFP8/MXFP4.

Като относително нов език, LLM имаха малко излагане на CuTeDSL код по време на тяхното предварително обучение. Освен това стилът на настоящия ни код, с ядра, оптимизирани за активации и норми, беше малко по-различен от съществуващите примери CUTLASS или Quack, които бяха силно оптимизирани за TMA и архитектурно специфични инструкции за умножение на матрици. Достигането до момент, в който LLM могат да помогнат при писането на ядро, беше постепенен процес.

Ръчно изградихме първоначален корпус от ядра и помощни функции, следвайки последователни модели, преди да можем да разчитаме на агенти за преобразуване на ядра. С напредването на корпуса бихме могли да извършваме груби преводи от CUDA към CuTeDSL, последвани от значителна ръчна преработка. По-специално установихме, че моделите се борят с преобразуванията на показалец към тензор и векторизацията, като се забиват в опити да коригират грешки в типа, когато се стремят да накарат тест да се компилира. Те показват малко разбиране за базираната на мрежа природа на ядрата, обърквайки темата и мрежовата синхронизация. Най-разочароващото беше изкушението да се халюцинират несъществуващи преводи на библиотеки, специфични за CUDA, вместо да се изгражда срещу нашите собствени заместители на CuTeDSL. С течение на времето, с нарастването на корпуса от CuTeDSL ядра, установихме, че LLM могат да изградят по-голямата част от каркасния код и да попълнят празните полета, което ни позволява да се съсредоточим върху дизайна на високо ниво.

В сравнение с CUDA, установяваме, че CuTeDSL осигурява много по-добра производителност по време на компилация, което ни позволява да специализираме за повече параметри и да постигнем подобрения в производителността. Въпреки че ядрата са малко по-многословни поради липсата на определени абстракции в сравнение с Triton, смятаме, че сложността е оправдана, тъй като по-финният контрол върху хардуера ни позволява да изграждаме по-бързи ядра.

По време на написването празнината в производителността между CUTLASS C++ и CuTeDSL ядрата е минимална, което отваря вълнуващи възможности за съставяне и сливане на високоефективни GEMM ядра с комуникация между устройствата, нормализация или активации. С ангажирането с този DSL ще можем да извлечем допълнителна производителност от нашите устройства, за да подобрим пропускателната способност на нашите работни натоварвания – от ембединги до LLM.

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

https://github.com/NVIDIA/cutlass

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

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

BibTeX

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

Инференс в Perplexity AI

Защо CuTeDSL?

Опит за разработчици

Голямо пространство за конфигурация

Предварително попълване и дешифриране на специализация

Мрежова синхронизация

CuTeDSL ядра

QK норма

Mixture-of-Experts изпращане/комбиниране

Локално

Един възел

Даннов паралелизъм

Peer-to-Peer

Поддръжка на Blackwell

Генериране на ядро, поддържано от AI

Заключение и бъдеща работа

Препратки

Цитиране