CuTeDSL v Perplexity

V Perplexity používáme vlastní Runtime-Optimized Serving Engine (ROSE) pro spouštění modelů od embeddings až po LLM s biliony parametrů. Náš běhový prostředí, které stojí za všemi rozhraními API včetně Sonar, Search a Embeddings, bylo vytvořeno proto, aby výzkumníkům umožnilo zkoumat nejvhodnější modely a architektury a zároveň rychle nasazovat nejlepší kandidáty do produkce. Abychom toho dosáhli, musíme být schopni rychle implementovat, nasazovat a optimalizovat inferenci pro nejmodernější architektury. V tomto příspěvku na blogu se podíváme na to, jak nám kernely CuTeDSL integrované do našeho inferenčního engine umožňují vytvářet potřebné GPU kernely, abychom dosáhli špičkového výkonu modelů na GPU NVIDIA.

Vlastní modely vytvořené různými týmy v Perplexity jsou hostovány interně na GPU NVIDIA Hopper a Blackwell pomocí inferenčního engine ROSE, který vytvořil tým pro AI inferenci. ROSE byl původně vytvořen tak, aby lépe obsluhoval přizpůsobené modely Llama pro dekódování jazyka i klasifikaci, s rozhraním kompatibilním s NVIDIA Triton Inference Server. Časem se rozrostl v plnohodnotný inferenční engine obsluhující podstatné LLM a také modely pro hodnocení, klasifikaci, skórování a další modely založené na transformerech pro širokou škálu týmů, které napájejí produkty Perplexity. Cílem návrhu ROSE je být jednoduchý: postupem času se investuje úsilí do dosažení špičkového výkonu, ale hlavní prioritou je přizpůsobit se neustále se vyvíjejícím architektuře a případům použití pro naše modely.

Prostor konfigurace kernelu napříč datovými typy, rozměry hlav, konstantami v době kompilace a nastaveními spuštění.

Na vysoké úrovni se ROSE skládá z engine, který přizpůsobuje model rozhraní, se kterým mohou klienti komunikovat. Engine jsou vysoce znovupoužitelné, zajišťují plánování a dávkování požadavků a načítání dekódovaných tokenů, embeddings nebo skóre z podkladových modelů. Zpracovávají inicializaci zařízení, meziprocesové a mezinozlové komunikační kanály a nahrávání vah. Pro LLM je engine zodpovědný za dávkování, rozdělování na bloky (chunking) a vzorkování. Je také zodpovědný za alokaci úložiště KV, a to jak pro plnou, tak pro lineární pozornost, stejně jako pro párování prefixů. Embedding enginy jsou jednodušší, protože provádějí pouze online dávkování. Provádění je optimalizováno tak, že engine připravuje další kolo práce, zatímco model běží paralelně na akcelerátoru.

ROSE poskytuje bohatou sbírku vlastních vrstev pro modelování, které obalují GPU kernely vytvořené pomocí CuTeDSL, Triton, CUDA, CUTLASS a cuBLAS. Pro implementaci maticového násobení a pozornosti spoléháme především na kernely NVIDIA z CUTLASS a cuBLAS. Zbytek pipeline, včetně embedding, norm, MoE a aktivačních kernelů, je součástí knihovny kernelů ROSE. Aby co nejlépe podpořil hardware, na kterém spouštíme inferenci, je ROSE vybaven vlastními primitivy mezizařízení komunikace pro MoE routing, dispatch a combine. Zatímco ostatní kernely jsou relativně jednoduché, musí být důkladně specializované, aby dosáhly špičkového výkonu na všech různých skrytých dimenzích a zařízeních, na kterých běží. Kromě toho spoléháme na specializaci za kompilace, abychom přizpůsobili všechny páky a úpravy různých modelů (váhová bias vs. bez váhové bias pro RMS norm, aktivační funkce atd.). V důsledku toho jsme se rozhodli použít DSL k jejich implementaci: zpočátku jsme ROSE postavili pomocí Tritonu. Od té doby jsme však přijali CuTeDSL jako primární prostředí pro programování GPU.

Z matematického hlediska jsou inferenční kernely obvykle nenáročné. V důsledku toho při vývoji kernelu složitost obecně nespočívá v vyjádření toho, jaký výpočet se má provést, ale v tom, jak přesně by jej hardware měl provést. Abychom vytvořili ty nejlepší kernely, potřebujeme jazyk, který nabízí správné abstrakce nad podkladovým hardwarem. Správný jazyk nám umožňuje specifikovat operace na vysoké úrovni k provedení výpočtu, přičemž je agresivně specializujeme a optimalizujeme ve všech bodech v širokém prostoru konfigurací.

CuTeDSL je jazyk DSL pro GPU NVIDIA založený na Pythonu, který je postaven nad algebrou rozvržení CuTe a MLIR. Kompiluje se just-in-time do vysoce optimalizovaného PTX při zachování kontroly nad nízkoúrovňovými hardwarovými primitivy. Jeho cílem je nabídnout stejné funkce a výkon jako CUTLASS C++, aniž by docházelo k nákladným časům kompilace, které brzdí rychlost vývoje. Z jazyků, které jsme zkoumali nebo se kterými jsme pracovali (včetně CUDA, cuTile, TVM IR a Triton), ho považujeme za nejvhodnější pro vývoj našich kernelů.

Celkově, i když není tak vysokoúrovňový a stručný jako Triton nebo cuTile, zjistili jsme, že algebra rozvržení CuTe a její kompozice (poskytované prostřednictvím cute.Tensor) byly přiměřenými náhradami. CuTeDSL navíc nabízí neomezený přístup k hardwarovým primitivům, což nám poskytuje úroveň kontroly, kterou potřebujeme k dosažení špičkového výkonu ve všech okrajových případech. S CUDA byly možnosti specializace za kompilace omezené, protože každý další parametr přidal nepřiměřené náklady na časy kompilace nebo složitost do systému sestavení. Naproti tomu s JIT (Just-In-Time) kompilovaným DSL můžeme specializovat agresivněji bez obětování tempa iterací. Na rozdíl od TVM IR, který používal CUDA jako mezikrok, jsme zjistili, že pipeline založená na MLIR, která sestupuje do PTX, je robustnější a snáze se ladí než ta, která generovala CUDA jako mezikrok. Navíc s použitím vlastních operací k zabalení instrukcí PTX jsme zjistili, že dokážeme plně využít všechny funkce podkladového hardwaru.

Většina nákladů na kompilaci kernelů CUDA se obvykle utratí ve fázi expanze šablon kernelů, protože optimalizace a kompilace od fáze PTX dále mohou být dokončeny během několika sekund. S mnoha parametry v době kompilace může tato fáze kompilace explodovat na několik minut u netriviálního kernelu. Toto úzké hrdlo je umocněno tím, že příležitosti pro paralelní kompilaci vznikají až později, na úrovních IR a PTX, přičemž zpracování šablon v C++ probíhá sekvenčně. I když existují triky, jak to vyřešit – přesunutím instancí šablon do samostatných kompilačních jednotek – vyžadují významné úpravy systémů sestavení a použití neohrabaných maker. Naproti tomu kernely CuTeDSL zkompilované metodou JIT (Just-In-Time) nevyžadují složitá pravidla sestavení a nebrzdí časy kompilace, protože vývojový cyklus se může soustředit na jediný bod v prostoru konfigurací, zatímco pečlivé jednotkové testy mohou později pokrýt celý designový prostor kernelu.

I když kompilace JIT přesouvá náklady z doby vývoje do doby běhu (runtime), režie je více než přijatelná. Když je model nasazen, aktivuje pouze několik bodů z celého prostoru konfigurací, což způsobuje režii v řádu sekund během spouštění/zahřívání inferenčního engine. Naproti tomu předkompilace všech konfigurací kernelu by mohla trvat několik minut. Režie lze navíc dále zmírnit ukládáním JIT zkompilovaných artefaktů do mezipaměti.

Vysoké náklady na kompilaci CUDA jsou umocněny obtížností sledování původu chyb, když se používá odesílání a specializace založené na šablonách, což je případ CuTe a CUTLASS. S DSL lze chyby ukázat na konkrétní řádek kódu, který je vygeneroval, což výrazně usnadňuje ladění. Navíc je mnohem jednodušší mapovat vygenerovaný kód na příkaz DSL, který jej vytvořil, prostřednictvím středních fází MLIR nebo spoléháním se na ladičky připevněné k vygenerovanému PTX. Díky těmto aspektům DSL jsme zjistili, že jsme při vývoji kernelů mnohem produktivnější.

Zatímco kernely sdílejí mnoho podobností napříč modely, stále musí být vybaveny řadou konfiguračních ovládacích prvků (knobs), které mohou mít vliv na výkon. Mezi relevantní parametry patří:

Skrytá dimenze / Dimenze hlavy / Počet hlav pozornosti

Vstupní/výstupní datové typy

Schéma kvantování

Konstanty doby kompilace (bias, aktivační funkce atd.)

Architektura ROSE, včetně jejího enginu, modelu, knihovny kernelů a služeb.

Specializace za kompilace je užitečná jak pro výkon, tak pro ladění. Zejména v pipeline procesoru, jako je tomu u GPU, je neexistující větev lepší než větev, která se bere vždy nebo nikdy. Kompilace kernelu pro specifickou sadu konfiguračních prvků navíc vede k menšímu počtu instrukcí PTX/SASS, které se snáze analyzují a ladí při použití cuda-gdb během vyšetřování.

Pro danou velikost dávky a konfiguraci lze odhadnout očekávaný provoz do DRAM a počet aritmetických operací. Následně lze případ od případu provádět optimální rozhodnutí pro optimální alokaci mřížky (grid), synchronizaci a používání určitých hardwarových funkcí. CuTeDSL nám umožňuje definovat a znovu použít komponenty pro hlubokou specializaci, aniž by došlo k nepřijatelné režii v době kompilace.

Výkon inferenčních kernelů musí být obecně optimalizován pro dva samostatné případy použití:

Prefill: malá velikost dávky (málo bloků, sekvencí), mnoho tokenů na sekvenci. Propustnost se maximalizuje jak pro embedding, tak pro vstupní tokeny LLM. Latence v desítkách až stovkách milisekund.

Decode: malá až velká velikost dávky, jeden (bez spekulativního dekódování) až několik tokenů na sekvenci (draft-target spekulace, predikce více tokenů). Latence blízká latenci spuštění kernelu.

Pro GEMM a pozornost běží fáze prefill na velkém počtu tokenů a výkon je omezen výpočetní propustností zařízení. Dekódování zahrnuje stejnou operaci na malých velikostech dávek, přičemž době běhu dominuje cena za načtení aktivací a vah z DRAM do streamovacích multiprocesorů a tenzorových jader, přičemž je místo toho omezena paměťovou šířkou pásma zařízení. I když je velikost dávky nebo celkový počet tokenů v dávce relativně jasným ukazatelem, práh přepnutí mezi prefill a decode je jemnější, protože šířka pásma paměti je určena také skrytou dimenzí modelu a datovým typem aktivace. Kernely jsou v důsledku toho rozděleny na implementace prefill a decode, pokud by algoritmické rozdíly nebo přiřazení mřížky/bloku přinesly výhodu některému z nich.

Parametr, který se univerzálně liší mezi kernely prefill a decode, je počet warpů v bloku (počet vláken vydělený velikostí warpu 32 vláken na warp). Je žádoucí učinit tento parametr konstantou doby kompilace, protože může pomoci při alokaci registrů. V CUDA se toho obvykle dosahuje specifikací počtu vláken pomocí __launch_bounds__ a poskytnutím počtu warpů jako parametru šablony:

cpp

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

}

Počet registrů dostupných pro všechna vlákna v bloku je hardwarové omezení – 64K na Hopper a Blackwell. Pokud je počet vláken znám předem, kompilátor přesně ví, kolik registrů může přiřadit každému vláknu, a může úspěšně kompilovat agresivněji optimalizované kernely. To však přináší náklady v době kompilace, protože nyní je nutné kernel znovu zkompilovat pro různé velikosti warpů. S CuTeDSL to již není problém, protože nemusíme platit náklady na kompilaci pro celý prostor konfigurací předem. DSL také očekává, že velikost bloku bude konstantním výrazem, čímž implicitně provádí optimalizaci implikovanou __latch_bounds__ v CUDA.

Kromě toho, že je počet warpů konstantní, je žádoucí jej měnit mezi prefill a decode. K plnému využití paměťové šířki pásma GPU nestačí vždy přiřadit vysoký počet warpů na blok, protože plné šířky pásma lze dosáhnout pouze použitím co největšího počtu dostupných streamovacích multiprocesorů (SM). Pro daný počet tokenů a známé datové typy a skryté dimenze proto nejprve vybereme ideální dimenzi vektoru na vlákno (8 pro bfloat16, 4 pro float32), abychom zajistili, že generovaný kód používá 128bitové instrukce načtení/uložení (256bitové na Blackwellu). Když známe velikost vektoru, můžeme spočítat počet warpů vyžadovaných kernelem. Nejprve se pokusíme maximalizovat počet SM: pokud je méně warpů než SM, zvolí se jeden warp na blok. Toto je obvykle dekódovací varianta kernelu. Pokud je dostatek warpů, postupně zvyšujeme jejich počet na 8, 16 nebo 32 a provádíme srovnávací test (benchmarking) horní hranice, po jejímž překročení již obvykle nezaznamenáváme zlepšení výkonu.

Zatímco Triton prováděl některé z těchto optimalizací pod kapotou, zjistili jsme, že můžeme dosáhnout podobných výsledků pečlivým použitím algebry rozvržení CuTe, aniž bychom byli omezeni nedostatky vysokoúrovňového kompilátoru. Zejména když se jedná o sub-warp horizontální redukce, zjistili jsme, že CuTeDSL nám nabídlo jemnozrnnou kontrolu potřebnou ke generování optimálního kódu.

V některých situacích lze k distribuci práce mezi co nejvíce bloků použít cooperative_groups::this_grid().sync() k synchronizaci všech zúčastněných bloků. Platí omezení: jelikož blok nelze odplánovat ze SM, jakmile začne provádění, aby byl zaručen pokrok vpřed, počet bloků nesmí překročit počet SM dostupných na zařízení. To zajišťuje, že v libovolném bodě mohou všechny bloky aktualizovat a dotazovat globální paměťová umístění zapojená do synchronizace. Toto omezení vynucuje cudaLaunchCooperativeKernel. Náklady na bariéru mřížky jsou na Hopperu přibližně 3 us plus náklady na vyprázdnění všech předchozích paměťových operací prostřednictvím implicitní paměťové bariéry v některých situacích. Celkově je to srovnatelné s režijními náklady na rozdělení kernelu na dvě spuštění, což vede k rozhodnutí o návrhu citlivému na výkon. Nehledě na to, že Programmatic Dependent Launch (PDL) může snížit režii spuštění od architektury Hopper dále, u dekódovacích kernelů se vyplatí se této režii vyhnout.

Pokud jedno z pracovních zatížení vyžaduje více bloků, než je počet SM, měl by být kernel rozdělen a spoléhat na implicitní bariéru spuštění kernelů. To je primárně případ prefill. Pokud se počet bloků, které jsou ideální pro všechna pracovní zatížení, vejde do mřížky, upřednostňuje se bariéra mřížky, což je případ dekódovacích kernelů pracujících s menším počtem tokenů.

V CuTeDSL lze tato rozhodnutí učinit ve spouštěcím kódu na straně hostitele a kernely lze skládat a prokládat s mřížkovými bariérami podle potřeby, jak ukazuje příklad níže. S takovou hybridní strategií lze dekódovací kernely optimalizovat tak, aby dosahovaly minimální latence, a prefill kernely lze optimalizovat pro maximální propustnost. Na rozdíl od řešení založených na CUDA lze tuto logiku definovat stručněji a s menší režií v době kompilace.

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)
		)

V případě našich MoE all-to-all kernelů, jak je podrobně popsáno níže, nám flexibilita při vytváření kernelů s mřížkovými bariérami nebo bez nich a jejich optimalizace v době kompilace pomohla dosáhnout výrazného nárůstu výkonu prostřednictvím CuTeDSL.

První kernel, který jsme vytvořili v CuTeDSL, byl specializovaný RMS norm kernel optimalizovaný pro malé dimenze. QK norm, který se používá ve vrstvě self-attention nedávných modelů Qwen, se aplikuje na bázi jednotlivých hlav (per-head) na 64, 128 nebo 256 prvků. Na rozdíl od běžného RMS norm, který pracuje nad skrytými dimenzemi přesahujícími tisíce a vyžaduje horizontální redukci jak přes sdílenou paměť, tak registry, QK norm vyžaduje redukci pouze napříč warpy s možností vtěsnat více hlav do jediného warpu. To se stalo případem použití, kdy jsme mohli agresivně optimalizovat v době kompilace, abychom našli nejlepší přiřazení tokenů k warpům pro všechny dimenze hlav a datové typy.

RMS norm je poměrně přímočará operace, i když vyžaduje horizontální součet pro výpočet odmocniny normy vektoru:

cpp

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

Jelikož jsou hlavy pozornosti obvykle malé a operace na vláknech jsou vysoce vektorizované, jeden warp dokáže spočítat čtvercovou normu. Aby se maximalizovala šířka pásma paměti, každé vlákno by mělo zpracovat 128 bitů vstupu, což se převádí na 8 hodnot bfloat16. Poloviční warp (half-warp) tedy stačí ke zpracování hlavy a čtvercovou normu lze určit pomocí horizontální redukce prováděné funkcí shuffle_sync_bfly. Počínaje dimenzí hlavy a datovým typem můžeme určit počet hlav na warp a specifikovat argumenty a počet redukčních kol v době kompilace.

Poté můžeme vytvořit specializovaný kernel s relativně malým počtem řádků kódu, který se spoléhá na automaticky vektorizované načítání/ukládání a používá nízkoúrovňové operace specifické pro architekturu pro redukci warpu. Kernel určí, na které hlavě (Q nebo K) má pracovat, a identifikuje odpovídající váhy, které se mají aplikovat. Načte celý vektor z hlavy a nejprve provede lokální součet. Následně se na základě počtu vláken na hlavu vypočítá čtvercová norma. Nakonec se automaticky vektorizovaný RMS norm určí a zapíše zpět do paměti. Parametry specifické pro model, jako jsou WEIGHT_BIAS a EPS, jsou konstantně složené (constant-folded).

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)

Ve srovnání s běžným RMS norm kernelem přináší tato specializovaná varianta 2–3× lepší propustnost a latenci pro prefill i decode.

MoE dispatch a combine přehazují vstupní tokeny na experty, ke kterým jsou směrovány, a provádějí váženou akumulaci poté, co MoE GEMM kernely provedou příslušné maticové násobení. Výkon je obzvláště zajímavý, protože různé produkty spoléhají na modely MoE různých velikostí. Na základě typů nasazení rozlišujeme různé kernely k optimalizaci:

TP=1: Lokální kernel pro jediné GPU. Dispatch přehazuje tokeny na příslušné experty a combine je akumuluje. Používá se primárně pro nasazení Qwen3 30B na H100.

DP=1, TP ≠ 1: Tenzorově paralelní, jednounozlový kernel. Vstup je replikován, dispatch vybírá tokeny pro aktuální hodnost (rank) a combine je akumuluje napříč všemi hodnostmi. Akumulace buď přes NVLink, nebo NVLS. Používá se pro velké prefillery.

DP ≠ 1, NVLink: Datově paralelní kernel s volitelnou tenzorovou paralelou. Každá hodnost odesílá tokeny všem ostatním hodnostem hostujícím konkrétní experty. Informace o routování se agregují, tokeny se přenášejí přes NVIDIA NVLink. Používá se pro datově paralelní nasazení H200 nebo GB200.

DP ≠ 1, InfiniBand: Peer-to-Peer datově paralelní a tenzorově paralelní kernel fungující napříč NVLink i NVIDIA Quantum InfiniBand. Používá se pro víceuhelná nasazení Hopper.

Lokální kernel počítá počet tokenů přiřazených každému expertovi na základě směrovací tabulky a přiřazuje každému tokenu unikátní index v rámci seznamu tokenů patřících ke každému expertovi. Po kroku routování jsou tokeny zkopírovány do volitelně zarovnaných vyrovnávacích pamětí (buffers), které mají být předány MoE GEMM kernelu. Pomocí umístění tokenů pak combine kernel provede akumulaci v registrech a zapíše finální tokeny do výstupního bufferu.

Fáze dispatch se liší mezi prefill a decode: pro málo tokenů je k dispozici málo informací o routování ke čtení, takže je zpracováváme v jediném bloku a signalizujeme zbytek mřížky, přičemž jednotlivé bloky kopírují části tokenů. Přiřazení tokenů je určeno atomickým inkrementováním čítačů ve sdílené paměti. Vzhledem k tomu, že prefill musí zpracovat více informací o routování, je routování rozděleno do několika bloků v samostatném kernelu. Indexy jsou nejprve přiřazeny lokálně pomocí čítačů ve sdílené paměti a poté jsou agregovány napříč všemi bloky v globální paměti, aby se určila pozice každého tokenu. Mřížka tohoto kernelu je omezena na počet SM, aby bylo možné používat mřížkové bariéry. Následný kernel provádí kopie a spouští se napříč mnohem větším počtem bloků a warpů, aby plně využil dostupnou šířku pásma paměti.

Combine kernel je relativně jednoduchý, protože provádí vektorizovanou akumulaci, přičemž spoléhá na informace o routování a offsety vypočítané dispatchem. CuTeDSL se zde ukázal jako cenný, protože výpočet mohl být vyjádřen pomocí jednoduchých tenzorů CuTe, které se zkompilují do efektivního kódu pro všechny datové typy a skryté dimenze.

V tenzorově paralelním nasazení jsou informace o routování a vstupní tokeny replikovány. Lokální dispatch kernel lze plně znovu použít s drobnou změnou pro routování a výběr pouze těch tokenů, které jsou přiřazeny expertům umístěným na aktuální hodnosti. S některými parametry a specializací za kompilace se lokální dispatch kernely znovu použijí.

Jelikož je počet hodností obvykle menší než počet expertů, na které je token směrován, lze v průměru předpokládat, že každý token má alespoň jednu ze svých replik směrovanou na každou hodnost. V této situaci lze lokální combine kernel znovu použít k agregaci lokálních kopií tokenu, následovanou operací all-reduce pro součet napříč zúčastněnými hodnostmi. V případě prefill kernelu spouštíme samostatné kernely, avšak dekódování spojuje obě operace do jediného spuštění.

Pro all-reduce používáme obyčejný NVLink, pokud se účastní pouze dvě hodnosti, a přepínáme na NVLS, když jsou zapojeny 4 nebo více hodností, což zajišťuje 10–20% nárůst propustnosti. Díky DSL můžeme sdílet většinu redukční logiky a výměn peer bufferů mezi MoE a běžnými all-reduce kernely. Toto je situace, kdy je úzká kontrola nad hardwarem přínosná, protože operace DSL nad multimem.ld_reduce a multimem.st jsou vyžadovány pro NVLS. Samotná implementace redukce do značné míry následuje příklad all_reduce z testů NCCL.

Datově paralelní kernely spoléhají na NVLink pro agregaci tokenů pro více zařízení. Ve fázi dispatch, před výměnou tokenů, peers nejprve lokálně směrují, aby identifikovali cílový offset na každé hodnosti pro každý token, který odesílají. Následně začnou kopírovat tokeny do přijímacích bufferů alokovaných pro každého odesílatele, aby se předešlo konfliktům. K signalizaci dokončení se používají zápisy release-store na jeden token. Na straně přijímače se informace o routování agregují, aby bylo možné kondenzovat všechny tokeny do jediného bufferu pro MoE GEMM. Přijímače se dotazují vlajky zapsané peerem pomocí acquire-load, aby počkaly na dokončení zápisů.

Na straně combine jsou tokeny kopírovány přes NVLink do zdrojových hodností, přičemž vlajka pro jeden token indikuje příjem. Combine kernel je podobný lokálním a jednounozlovým variantám s drobnými rozdíly v adresování a dodatečnou synchronizační logikou.

K implementaci tohoto kernelu jsme vytvořili knihovnu obalů PTX (wrappers) pro zpřístupnění načtení a uložení typu release-acquire v rozsahu jak na úrovni systému, tak na úrovni zařízení. Zlepšením granularity synchronizačních vlajek a alokací více bloků pro odesílání a přijímání tokenů přes NVLink jsme dosáhli 2× zlepšení oproti našim předchozím kernelům CUDA.

Naše Peer-to-Peer kernely dříve dosahovaly špičkového výkonu na GPU Hopper spárovaných s adaptéry ConnectX-7 InfiniBand. Od té doby jsme portovali combine kernely do CuTeDSL pro další nárůst výkonu. Vzhledem k tomu, že verze v CUDA již silně spoléhaly na specializaci za kompilace, ukázaly se jako úzká hrdla doby kompilace v našem repozitáři, což nás nutilo vynechat některé parametry. Migrací na DSL kompilovaný metodou JIT jsme provedli specializaci pro více parametrů a vyladili alokace mřížky, což přineslo další rozdíly mezi prefill a decode pro zhruba 3% nárůst latence.

Z pohledu kernelů ROSE spadá podpora pro Blackwell do dvou kategorií: kernely agnostické k architektuře a kernely specifické pro architekturu. Kernely specifické pro architekturu, spoléhající na tenzorová jádra, vyžadují explicitní práci k přizpůsobení. Díky použití DSL se však kernely agnostické k architektuře přizpůsobují triviálně.

GPU Blackwell má až 160 SM, na rozdíl od 132 u Hoppera. Pro většinu kernelů je to konstanta doby kompilace a logika ve spouštěcím kódu na straně hostitele ji bere v úvahu při určování optimálního kernelu ke spuštění (prefill vs. decode). Blackwell navíc zavádí 256bitové instrukce načtení, které lze snadno vygenerovat zvětšením šířky vektoru (poslední dimenze tenzorů CuTe), aby se zlepšilo využití šířky pásma paměti.

Blackwell zavádí hardwarovou podporu pro významný počet schémat kvantování (MXFP8, MXFP4, NVFP4). Podpora pro ně byla implementována primárně doladěním stávajících kernelů pro škálování FP8 po blocích. Podpora JIT je zde klíčová, protože lze podporovat nové režimy, aniž by došlo k explozi prostoru konfigurací v době kompilace. V minulosti jsme zjistili, že Triton je neohrabaný při vyjadřování horizontálních redukcí napříč několika prvky, nad polovinou nebo čtvrtinou warpu. Jemnější kontrola nad redukcemi byla klíčová pro dosažení optimálního výkonu při výpočtu absolutní maximální hodnoty napříč 16 nebo 32 prvky pro NVFP4, resp. MXFP8/MXFP4.

Jako relativně nový jazyk měly LLM během svého předtrénování jen malou zkušenost s kódem CuTeDSL. Styl našeho současného kódu, s kernely optimalizovanými pro aktivace a normy, se navíc mírně lišil od stávajících příkladů CUTLASS nebo Quack, které byly silně optimalizovány pro TMA a instrukce maticového násobení specifické pro architekturu. Dostat se do bodu, kdy by LLM mohly pomáhat při psaní kernelů, byl pozvolný proces.

Ručně jsme vybudovali počáteční korpus kernelů a užitečných funkcí, které dodržovaly konzistentní vzory, než jsme se mohli spolehnout na agenty pro konverzi kernelů. Jak korpus postupoval, mohli jsme provádět hrubé překlady z CUDA do CuTeDSL, po nichž následovalo významné ruční přepracování. Zejména jsme zjistili, že modely bojovaly s konverzemi ukazatelů na tenzor (pointer-to-tensor) a vektorizací a uvízly při pokusech opravit chybné typy, když se snažily přimět test ke kompilaci. Vykazovaly jen malé pochopení mřížkové povahy kernelů, přičemž zaměňovaly synchronizaci vláken a mřížky. Co bylo nejvíce frustrující, objevilo se pokušení halucinovat neexistující překlady knihoven specifických pro CUDA namísto sestavování proti našim vlastním náhradám CuTeDSL. Postupem času, jak korpus kernelů CuTeDSL rostl, jsme zjistili, že LLM dokážou vytvořit většinu kosterního kódu a vyplnit prázdná místa, což nám umožňuje soustředit se na návrh na vysoké úrovni.

Ve srovnání s CUDA zjišťujeme, že CuTeDSL přináší mnohem lepší výkon v době kompilace, což nám umožňuje specializovat se na více parametrů a dosáhnout vylepšení výkonu. Zatímco kernely jsou o něco podrobnější kvůli absenci určitých abstrakcí ve srovnání s Tritonem, zjišťujeme, že složitost je opodstatněná, protože jemnější kontrola nad hardwarem nám umožňuje vytvářet rychlejší kernely.

V době psaní tohoto článku je propast ve výkonu mezi kernely CUTLASS C++ a CuTeDSL minimální, což otevírá vzrušující příležitosti pro skládání a slučování vysoce výkonných GEMM kernelů s komunikací mezi zařízeními, normalizací nebo aktivacemi. Tím, že se zavážeme k tomuto DSL, budeme schopni vytěžit z našich zařízení další výkon ke zlepšení propustnosti napříč našimi pracovními zatíženími, od embeddingů až po 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}
}

Inference v Perplexity AI

Proč CuTeDSL?

Vývojářská zkušenost

Velký prostor konfigurací

Předvyplnění a dekódování specializace

Synchronizace mřížky (Grid synchronization)

CuTeDSL kernely

QK Norm

Mixture-of-Experts Dispatch/Combine

Lokální

Jednounozlový

Datová paralelní (Data Parallel)

Peer-to-Peer

Podpora pro Blackwell

Generování kernelů asistované AI

Závěr a budoucí práce

Reference

Citace