CuTeDSL у Perplexity-ју

У Perplexity користимо сопствени систем за опслуживање оптимизован за време извршавања (ROSE) за опслуживање модела, од ембединга до ЛЛМ са билионима параметара. Смештен иза свих API-ја, укључујући Sonar, Search и Embeddings, наш рунтајм је изграђен да омогући истраживачима да истражују најпогодније моделе и архитектуре, уз брзо постављање најбољих кандидата у производњу. Да бисмо то постигли, морамо бити у стању да брзом темпом примењујемо, постављамо и оптимизујемо инференцију за најсавременије архитектуре. У овом блог посту истражујемо како нам CuTeDSL језгра интегрисана у наш мотор за инференцију омогућавају да изградимо потребна GPU језгра како бисмо довели моделе до врхунских перформанси на NVIDIA GPU-овима.

Прилагођени модели, које су изградили различити тимови у Perplexity-ју, хостују се интерно на NVIDIA Hopper и Blackwell GPU-овима користећи ROSE мотор за инференцију који је изградио тим за AI инференцију. ROSE је првобитно изграђен да боље опслужује прилагођене Llama моделе како за декодирање језика, тако и за класификацију, са интерфејсом компатибилним са NVIDIA Triton Inference Server-ом. Временом је израстао у потпуно развијен мотор за инференцију који опслужује значајне ЛЛМ-ове, као и рангирање, класификацију, бодовање и друге модове засноване на трансформаторима за широк спектар тимова који покрећу Perplexity производи. Циљ дизајна ROSE-a је да буде једноставан: временом се улаже труд да се достигну врхунске перформансе, али је главни приоритет прилагођавање архитектурама и случајевима употребе који се стално развијају за наше моделе.

Простор конфигурације језгра кроз типове података, димензије head-а, константе у време компајлирања и подешавања покретања.

На високом нивоу, ROSE се састоји од мотора који прилагођава модел интерфејсу са којим клијенти могу да интерагују. Мотори су високо реупотребљиви, управљају распоређивањем и груписањем захтева и преузимају декодиране токене, ембединге или скорове из подлежећих модела. Иницијализација уређаја, комуникациони канали између процеса и чворова и учитавање тежина су у потпуности под њиховом контролом. За ЛЛМ-ове, мотор је одговоран за груписање, дељење на делове и узорковање. Такође је одговоран за доделу KV меморије, како за пуну тако и за линеарну пажњу, као и за подударање префикса. Мотори за ембединг су једноставнији јер обављају само онлајн груписање. Извршавање је оптимизовано тако да мотор припрема следећу рунду посла док модел паралелно ради на акцелератору.

ROSE излаже богато збирку прилагођених слојева за моделирање, умотавајући GPU језгра изграђена помоћу CuTeDSL, Triton, CUDA, CUTLASS и cuBLAS. За матрично множење и имплементације пажње, првенствено се ослањамо на NVIDIA језгра из CUTLASS-а и cuBLAS-a. Остатак цевовода, укључујући ембединг, нормирање, MoE и језгра активације, чине део библиотеке језгра ROSE-a. Да би што боље подржао хардвер на којем изводимо инференцију, ROSE је опремљен сопственим примитивима за комуникацију између уређаја за MoE рутирање, диспечеровање и комбиновање. Иако су друга језгра релативно једноставна, морају бити темељно специјализована како би се постигле врхунске перформансе на свим различитим скривеним димензијама и уређајима на којима се изводе. Поред тога, ослањамо се на специјализацију у време компајлирања како бисмо прилагодили сва дугмад и подешавања различитих модела (тежина са померањем у односу на без померања за RMS норму, функције активације итд.). Сходно томе, одлучили смо да користимо DSL-ове за њихову имплементацију: првобитно смо изградили ROSE користећи Triton. Међутим, у међувремену смо усвојили CuTeDSL као primarno окружење за GPU програмирање.

Са математичке перспективе, језгра за инференцију су обично једноставна. Сходно томе, при развоју језгра, сложеност се углавном не налази у изражавању шта се рачуна, већ у томе како тачно хардвер то треба да уради. Да бисмо изградили најбоља језгра, потребан нам је језик који нуди праве апстракције над подлежећим хардвером. Прави језик нам омогућава да наведемо операције високог нивоа за извршавање прорачуна, агретивно их специјализујући и оптимизујући кроз све тачке у широком простору конфигурације.

CuTeDSL је DSL за NVIDIA GPU-ове заснован на Пајтону, изграђен на врху CuTe алгебре распореда и MLIR-а. Компајлира се у ходу (just-in-time) у високо оптимизовани PTX, уз задржавање контроле над хардверским примитивима ниског нивоа. Циљ му је да понуди исте функције и перформансе као CUTLASS C++, без високих времена компајлирања која успоравају брзину развоја. Од језика које смо испитали или са којима смо радили (укључујући CUDA, cuTile, TVM IR и Triton), сматрамо га најбољим избором за развој наших језгра.

Све у свему, иако није тако висок ниво и језгровит као Triton или cuTile, установили смо да су CuTe алгебра распореда и њене композиције (изложене кроз cute.Tensor) разумне замене. Поред тога, CuTeDSL нуди неограничен приступ хардверским примитивима, пружајући нам ниво контроле који нам је потребан да постигнемо врхунске перформансе у свим граничним случајевима. Са CUDA-ом, могућности за специјализацију у време компајлирања биле су ограничене, јер је сваки додатни параметар додавао неприхватљив трошак у време компајлирања или сложеност систему за градњу. Насупрот томе, са JIT компајлираним DSL-ом, можемо агресивније да се специјализујемо без жртвовања темпа итерација. За разлику од TVM IR-а, који је користио CUDA-у као посредни корак, утврдили смо да је цевовод заснован на MLIR-у који се спушта на PTX робуснији и лакши за дебаговање од оног који је емитовао CUDA-у као посредни корак. Поред тога, уз употребу прилагођених операција за умотавање PTX инструкција, утврдили смо да можемо у потпуности да искористимо све карактеристике подлежећег хардвера.

Главнина трошкова компајлирања CUDA језгра се обично троши у фази експанзије шаблона језгра, пошто се оптимизација и компајлирање од PTX фазе па надаље могу завршити за неколико секунди. Са много параметара у време компајлирања, ова фаза компајлирања може да експлодира на више минута за нетривијално језгро. Ово уска грло је наглашено чињеницом да се могућности за паралелно компајлирање појављују тек касније, на нивоима IR и PTX, при чему се руковање C++ шаблонима одвија секвенцијално. Иако постоје трикови за решавање овог проблема померањем инстанци шаблона у одвојене јединице компајлирања, они захтевају значајна дотеравања система за градњу и употребу гломазних макроа. Насупрот томе, CuTeDSL језгра која се компајлирају у ходу (JIT) не захтевају сложена правила градње и не ометају време компајлирања, пошто се циклус развоја може фокусирати на једну тачку у простору конфигурације, док пажљиви јединични тестови касније могу покрити цео простор дизајна језгра.

Иако JIT компајлирање помера трошкове са времена развоја на време извршавања, ова надземна рехија је више него прихватљива. Када се модел распореди, он активира само неколико тачака из целог простора конфигурације, што доводи до трошкова у секундама током покретања/загревања мотора за инференцију. Насупрот томе, преткомпајлирање свих конфигурација језгра могло би да траје више минута. Поред тога, надземна рехија се може додатно ублажити кеширањем JIT компајлираних артефаката.

Велики трошак компајлирања CUDA-е је наглашен тешкоћом праћења порекла грешака када се користи диспечеровање и специјализација заснована на шаблонима, као што је случај са CuTe и CUTLASS. Са DSL-ом, грешке се могу указати на линију кода која их је генерисала, што значајно олакшава дебаговање. Поред тога, много је једноставније мапирати генерисани код на DSL изјаву која га је покренула преко посредних MLIR фаза или ослањањем на информације за дебаговање приложене уз генерисани PTX. Захваљујући овим аспектима DSL-a, установили смо да смо много продуктивнији при развоју језгара.

Иако језгра деле многе сличности међу моделима, она и даље морају бити опремљена низом конфигурационих дугмади која могу имати утицај на перформансе. Релевантни параметри укључују:

Скривена димензија / Димензија главе / Број глава пажње

Типови улазних/излазних података

Шема квантизације

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

ROSE архитектура, укључујући њен мотор, модел, библиотеку језгра и услуге.

Специјализација у време компајлирања је корисна и за перформансе и за дебаговање. Посебно у цевоводном процесору, као што је случај са GPU-овима, непостојећа грана је боља од гране која се увек узима или се никада не узима. Поред тога, компајлирање језгра за одређени скуп конфигурационих дугмади резултира мањим бројем PTX/SASS инструкција, које је лакше анализирати и дебаговати када се користи cuda-gdb током истраге.

За дату величину серије и конфигурацију, могу се проценити очекивани саобраћај ка DRAM-у и број аритметичких операција. Након тога, оптималне одлуке се могу доносити од случаја до случаја за оптималну доделу мреже, синхронизацију и употребу одређених хардверских функција. CuTeDSL нам омогућава да дефинишемо и поново користимо компоненте за дубоку специјализацију без стварања неприхватљивих трошкова у време компајлирања.

Генерално, перформансе језгара за инференцију морају бити оптимизоване за два одвојена случаја употребе:

Предпуњење (Prefill): мала величина серије (мало делова, секвенци), много токена по секвенци. Пропусни опсег који треба максимизирати и за ембединге и за улазне токене ЛЛМ-ова. Латентност у десетинама до стотинама милисекунди.

Декодирање: величина серије од мале до велике, један (без спекулативног декодирања) до неколико токена по секвенци (спекулација циљ-нацрт, предвиђање више токена). Латентност блиска латентности покретања језгра.

За GEMM и пажњу, фаза предпуњења ради на великом броју токена и перформансе су ограничене пропусним опсегом рачунања уређаја. Декодирање укључује исту операцију на малим величинама серија, при чему време извршавања доминира трошком учитавања активација и тежина из DRAM-а у стриминг мултипроцесоре и тензорске језгре, што је уместо тога ограничено пропусним опсегом меморије уређаја. Иако је величина серије или укупан број токена у серији релативно јасан показатељ, праг преласка између предпуњења и декодирања је суптилнији зато што је пропусни опсег меморије такође одређен скривеном димензијом модела и типом података активације. Језгра су сходно томе подељена на имплементације предпуњења и декодирања ако би алгоритамске разлике или доделе мреже/блокова давале предност било којој од њих.

Параметар који се универзално разликује између језгара предпуњења и декодирања је број warp-ова у блоку (број тредова, подељен са величином warp-а од 32 треда по warp-у). Пожељно је овај параметар учинити константом у време компајлирања, јер може помоћи у додели регистра. У CUDA-овима се то обично ради дефинисањем броја тредова преко __launch_bounds__ и обезбеђивањем броја warp-ова као параметара шаблона:

cpp

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

}

Број регистра доступних свим тредовима у блоку је хардверско ограничење, 64K на Hopper-у и Blackwell-у. Ако је број тредова познат унапред, компајлер тачно зна колико регистра може доделити сваком треду и може успешно да компајлира агресивније оптимизована језгра. Међутим, ово долази уз цену у време компајлирања, јер сада језгро мора бити поново компајлирано за различите величине warp-ова. Са CuTeDSL-ом ово више није проблем, јер не морамо унапред да плаћамо трошкове компајлирања за цео простор конфигурације. DSL такође очекује да величина блока буде константни израз, чиме се имплицитно извршава оптимизација коју подразумева __latch_bounds__ у CUDA-и.

Осим што се број warp-ова прави константним, пожељно је варирати га између предпуњења и декодирања. Да би се у потпуности искористио пропусни опсег меморије GPU-а, није довољно увек додељивати велики број warp-ова по блоку, јер се пун пропусни опсег може постићи само коришћењем што већег броја доступних стриминг мултипроцесора (SM). Сходно томе, за дати број токена и познате типове података и скривене димензије, прво бирамо идеалну векторску димензију по треду (8 за bfloat16, 4 за float32) како бисмо осигурали да генерисани код користи инструкције учитавања/чувања од 128 бита (256 бита на Blackwell-у). Знајући величину вектора, можемо израчунати број warp-ова који захтева језгро. Прво покушавамо да максимизирамо број SM-ова: ако има мање warp-ова него SM-ова, бира се један warp по блоку. Ово је обично варијанта декодирања језгра. Ако има довољно warp-ова, постепено повећавамо број warp-ова на 8, 16 или 32, бенчмаркујући горњу границу након које обично више не видимо побољшање перформанси.

Иако је Triton извршио неке од ових оптимизација испод хаубе, установили смо да сличне резултате можемо постићи пажљивим коришћењем CuTe алгебре распореда, без ограничења недостацима компилатора високог нивоа. Посебно када су укључене хоризонталне редукције под-warp-а, установили смо да нам је CuTeDSL понудио фино подешену контролу која је неопходна за генерисање оптималног кода.

У неким ситуацијама, да би се посао распоредио на што више блокова, cooperative_groups::this_grid().sync() се може користити за синхронизацију свих блокова који учествују. Важе ограничења: пошто се блок не може уклонити са распореда са SM-а када једном почне са извршавањем, да би се гарантовао напредак напред, број блокова не може прећи број SM-ова доступних на уређају, како би се осигурало да у било ком тренутку сви блокови могу ажурирати и анкетирати локалне меморијске локације укључене у синхронизацију. Ово ограничење намеће cudaLaunchCooperativeKernel. Цена мрежне баријере је око 3us на Hopper-у, плус цена чишћења свих претходних меморијских операција путем имплицитне меморијске баријере у неким ситуацијама. Све у свему, ово је упоредиво са надземном рехијом дељења језгра на два покретања, што доводи до доношења одлуке о дизајну осетљивој на перформансе. Упркос чињеници да програмски зависно покретање (PDL) може смањити надземну рехију покретања од Hopper-а надаље, за језгра декодирања избегавање ове надземне рехије се исплати.

Ако један од радних оптерећења захтијева више блокова него што је број SM-ова, кернел треба подијелити, ослањајући се на имплицитну баријеру покретања кернела. Ово је прије свега случај за prefill. Ако број блокова који је идеалан за сва радна оптерећења стаје у grid, предност има grid баријера, као у случају decode кернела који раде над мањим бројем токена.

У 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 норма захтева редукцију само кроз warp-ове, уз могућност сабијања више глава у један warp. Ово је постало случај употребе у којем смо могли агретивно да се оптимизујемо у време компајлирања како бисмо пронашли најбоље доделе токена за warp за све димензије глава и типове података.

RMS норма је прилично једноставна операција, иако захтева хоризонтални збир за израчунавање квадратне норме вектора:

cpp

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

Пошто су главе пажње обично мале и операције по треду су веома векторизоване, један warp може израчунати квадратну норму. Да би се максимизирао пропусни опсег меморије, сваки тред треба да обради 128 бита улаза, што се преводи у 8 bfloat16 вредности. Сходно томе, пола 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 norm кернелом, ова специјализована варијанта испоручује 2–3x бољи проток и латенцију за и prefill и decode.

MoE диспечеровање и комбиновање мешају улазне токене са стручњацима којима су рутирани, обављајући пондерисану акумулацију након што MoE GEMM језгра изврше релевантна матрична множења. Перформансе су од посебног интереса, јер се различити производи ослањају на MoE моделе различитих величина. На основу типова распоређивања, разликујемо различите језгра за оптимизацију:

TP=1: Локално језгро за један GPU. Диспечеровање меша токене до релевантних стручњака, а комбиновање их акумулира. Користи се првенствено за Qwen3 30B распоређивања на H100.

DP=1, TP ≠ 1: Тензорски паралелно језгро за један чвор. Улаз је реплициран, диспечеровање бира токене за тренутни ранг, а комбиновање их акумулира кроз све рангове. Акумулација преко NVLink-a или NVLS-a. Користи се за велике уређаје за предпуњење.

DP ≠ 1, NVLink: Дата-паралелно језгро са опционом тензорском паралелношћу. Сваки ранг шаље токене сваком другом рангу који хостује одређене стручњаке. Информације о рутирању се агрегирају, токени се преносе преко NVIDIA NVLink-a. Користи се за дата-паралелна H200 или GB200 распоређивања.

DP ≠ 1, InfiniBand: Peer-to-Peer дата-паралелно и тензорски паралелно језгро које ради преко NVLink-а и NVIDIA Quantum InfiniBand-а. Користи се за вишечворна Hopper распоређивања.

Локално језгро броји број токена додељених сваком стручњаку на основу табеле рутирања и додељује јединствени индекс сваком токену унутар листе токена који припадају сваком стручњаку. Након корака рутирања, токени се копирају у опционо попуњене бафере који се прослеђују MoE GEMM језгру. Користећи локације токена, језгро за комбиновање затим врши акумулацију у регистрама, уписујући коначне токене у излазни бафер.

Фаза диспечеровања варира између предпуњења и декодирања: за мало токена има мало информација о рутирању за читање, па их обрађујемо у једном блоку и сигнализирамо остатак мреже, при чему појединачни блокови копирају делове токена. Додела токена се одређује атомским повећањем бројача у дељеној меморији. Пошто предпуњење мора да обради више информација о рутирању, рутирање је подељено на више блокова у одвојеном језгру. Индекси се прво локално додељују коришћењем бројача у дељеној меморији, а затим се агрегирају кроз све блокове у глобалној меморији да би се одредила позиција сваког токена. Мрежа овог језгра је ограничена на број SM-ова како би се могли користити мрежне баријере. Наредно језгро врши копирање, покренуто преко много већег броја блокова и warp-ова како би се у потпуности искористио расположиви пропусни опсег меморије.

Језгро за комбиновање је релативно једноставно, јер врши векторизовану акумулацију, ослањајући се на информације о рутирању и на оффсете које израчунава диспечеровање. CuTeDSL се овде показао вредним, јер се прорачун могао изразити кроз једноставне CuTe тензоре који се компајлирају у ефикасан код за све типове података и скривене димензије.

У тензорски паралелном распоређивању, информације о рутирању и улазни токени се реплицирају. Локално језгро за диспечеровање се може у потпуности поново користити уз мању промену за рутирање и одабир само токена додељених стручњацима који се налазе на тренутном рангу. Са неким параметрима и специјализацијом у време компајлирања, локална језгра за диспечеровање се поново користе.

Пошто је број рангова обично мањи од броја стручњака којима је токен рутиран, у просечном случају се може претпоставити да сваки токен има бар једну од својих реплика рутирану на сваки ранг. У овој ситуацији, локално језгро за комбиновање се може поново користити за агрегирање локалних копија токена, након чега следи all-reduce операција за сабирање кроз рангове који учествују. У случају језгра предпуњења, покрећемо одвојена језгра, међутим декодирање спаја обе операције у једно покретање.

За all-reduce користимо обичан NVLink ако учествују само два ранга, прелазећи на NVLS када је укључено 4 или више рангова за повећање пропусног опсега од 10-20%. Захваљујући DSL-у, можемо да делимо већи део логике редукције и размене бафера вршњака између MoE и редовних all-reduce језгра. Ово је ситуација у којој је блиска контрола над хардвером корисна, јер су DSL операције над multimem.ld_reduce и multimem.st неопходне за NVLS. Сама имплементација редукције углавном прати all_reduce пример из NCCL тестова.

Дата-паралелна језгра се ослањају на NVLink за агрегирање токена за више уређаја. У фази диспечеровања, пре размене токена, вршњаци прво локално рутирају да би идентификовали циљни оффсет на сваком рангу за сваки токен који шаљу. Након тога, почињу да копирају токене у пријемне бафере додељене за сваког пошиљаоца како би се избегли сукоби. Писања са ослобађањем и складиштењем по токену се користе за сигнализацију завршетка. На крају пријемника, информације о рутирању се агрегирају како би се сви токени могли сабити у један бафер за MoE GEMM. Пријемници анкетирају заставицу коју су написали вршњаци са аквизицијом и учитавањем како би сачекали да се писања заврше.

На страни комбиновања, токени се копирају преко NVLink-а до изворних рангова, са заставицом по токену која указује на пријем. Језгро за комбиновање је слично локалним и једночворним варијантама, уз мање разлике у адресирању и додатну логику синхронизације.

Изградили смо библиотеку PTX омотача да изложимо учитавања и складиштења са ослобађањем и аквизицијом са опсегом на нивоу система и уређаја како бисмо имплементирали ово језгро. Побољшањем грануларности заставица синхронизације и додељивањем више блокова за слање и пријем токена преко NVLink-а, постигли смо 2 пута боља побољшања у односу на наша претходна CUDA језгра.

Наша Peer-to-Peer језгра су претходно постигла перформансе на врхунском нивоу на Hopper GPU-овима упареним са ConnectX-7 InfiniBand адаптерима. У међувремену смо пренели језгра за комбиновање на CuTeDSL ради даљих добитака у перформансама. Пошто су CUDA верзије већ у великој мери зависле од специјализације у време компајлирања, показале су се као уска грла у време компајлирања у нашем складишту, приморавајући нас да изоставимо неке параметре. Миграцијом на JIT-компиловани DSL специјализовали смо се за више параметара и дотерали доделе мреже, уводећи додатне разлике између предпуњења и декодирања за око 3% повећања латентности.

Из перспективе ROSE језгра, подршка за Blackwell спада у две категорије: архитектурно агностична и архитектурно специфична језгра. Архитектурно специфична језгра, која се ослањају на тензорске језгре, захтевају експлицитан рад за прилагођавање. Међутим, захваљујући употреби DSL-а, архитектурно агностична језгра се прилагођавају тривијално.

Blackwell GPU има до 160 SM-ова, за разлику од Hopper-ових 132. За већину језгра ово је константа у време компајлирања, а логика у коду за покретање на страни хоста то узима у обзир за одређивање оптималног језгра за покретање (предпуњење наспрам декодирања). Поред тога, Blackwell такође уводи инструкције учитавања од 256 бита, које се могу лако емитовати повећањем ширине вектора (последња димензија CuTe тензора) ради побољшања искоришћења пропусног опсега меморије.

Blackwell уводи хардверску подршку за значајан број шема квантизације (MXFP8, MXFP4, NVFP4). Подршка за њих је имплементирана првенствено подешавањем постојећих кернела за FP8 block-wise scaling. JIT подршка је овдје кључна, јер се нови режими могу подржати без експлозије простора конфигурације у вријеме компилације. У прошлости смо сматрали да је Triton незграпан за изражавање хоризонталних редукција над малим бројем елемената, преко половине или четвртине warp-а. Финија контрола над редукцијама била је кључна за постизање оптималних перформанси при рачунању апсолутне максималне вриједности над 16 или 32 елемента за NVFP4, односно MXFP8/MXFP4.

Као релативно нов језик, ЛЛМ-ови су имали мало излагања CuTeDSL коду током свог предтренирања. Поред тога, стил нашег садашњег кода, са језгрима оптимизованим за активације и норме, мало се разликовао од постојећих CUTLASS или Quack примера, који су били веома оптимизовани за TMA и архитектурно специфичне инструкције матричног множења. Достизање тачке у којој ЛЛМ-ови могу помоћи у писању језгра био је постепен процес.

Ручно смо изградили почетни корпус кернела и услужних функција, пратећи досљедне обрасце, прије него што смо могли да се ослонимо на агенте за конверзију кернела. Како је корпус напредовао, могли смо да вршимо грубе преводе са CUDA на CuTeDSL, након чега је слиједила значајна ручна дорада. Посебно смо установили да се модели тешко сналазе са конверзијама pointer-to-tensor и векторизацијом, запињући док су покушавали да исправе грешке у типовима како би тест успио да се компајлира. Показали су мало разумијевања за grid-базирани карактер кернела, мијешајући синхронизацију треда и grid-а. Најфрустрирајуће је било искушење да се халуцинирају непостојећи преводи CUDA-специфичних библиотека умјесто да се ради уз наше сопствене CuTeDSL замјене. Временом, како је корпус CuTeDSL кернела растао, установили смо да LLM-ови могу да изграде већину skeleton кода и попуне празнине, што нам омогућава да се фокусирамо на дизајн на високом нивоу.

У поређењу са CUDA-ом, сматрамо да CuTeDSL пружа много боље перформансе у време компајлирања, омогућавајући нам да се специјализујемо за више параметара и постигнемо побољшања перформанси. Иако су језгра мало опширнија због недостатка одређених апстракција у поређењу са Тритоном, сматрамо да је сложеност оправдана јер нам фино подешена контрола над хардвером омогућава да изградимо бржа језгра.

У време писања овог текста, јаз у перформансама између CUTLASS C++ и CuTeDSL језгра је минималан, отварајући узбудљиве могућности за састављање и спајање високоперформансних GEMM језгра са комуникацијом између уређаја, нормализацијом или активацијама. Опвезивањем за овај DSL, моћи ћемо да извучемо додатне перформансе из наших уређаја како бисмо побољшали пропусни опсег кроз наша радна оптерећења, од ембединга до ЛЛМ-ова.

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)

Закључак и будући рад

Референце

Цитат