CuTeDSL a Perplexitynél

A Perplexytinál a saját fejlesztésű Runtime-Optimized Serving Engine (ROSE) futtatókönynyezetet használjuk a beágyazásoktól kezdve egészen a billió paraméteres nagy nyelvi modellekig (LLM) a modellek kiszolgálására. Minden API – beleértve a Sonart, a Searcht és a beágyazásokat is – mögött ott lévő futtatókönynyezetünket azért hoztuk létre, hogy a kutatók felfedezhessék a legmegfelelőbb modelleket és architektúrákat, miközben a legjobb jelölteket gyorsan éles környezetbe is helyezhetik. Ennek eléréséhez képesnek kell lennünk a legmodernebb architektúrák következtetésének (inference) gyors ütemű implementálására, üzembe helyezésére és optimalizálására. Ebben a bejegyzésben azt vizsgáljuk meg, hogy a következtetési motorunkba integrált CuTeDSL-kernelek hogyan teszik lehetővé számunkra a szükséges GPU-kernelek felépítését, hogy a modellek csúcsteljesítményt érjenek el NVIDIA GPU-kon.

A Perplexity különböző csapatai által készített egyedi modelleket házi keretek között, NVIDIA Hopper és Blackwell GPU-kon futtatjuk az AI Inference csapat által fejlesztett ROSE következtetési motor használatával. A ROSE-t eredetileg azért építették, hogy jobban kiszolgálja a testreszabott Llama modelleket mind a nyelvi dekódolás, mind a besorolás terén, az NVIDIA Triton Inference Serverrel kompatibilis interfésszel. Idővel teljes értékű következtetési motorvá fejlődött, amely jelentős LLM-eket, valamint rangsorolási, besorolási, pontozási és egyéb transzformátoralapú modelleket szolgál ki a csapatok széles köre számára a Perplexity termékeinek támogatásához. A ROSE tervezési célja az egyszerűség: idővel erőfeszítéseket tesznek a csúcsteljesítmény elérése érdekében, de a fő prioritás az, hogy alkalmazkodjanak a modellek folyamatosan fejlődő architektúráihoz és use case-eihez.

Az adattípusok, fejdimenziók, fordítási időben megadott állandók és indítási beállítások közötti kernelkonfigurációs tér.

Magas szinten a ROSE egy olyan motorból áll, amely hozzáigazít egy modellt egy olyan interfészhez, amellyel az ügyfelek interakcióba léphetnek. A motorok nagymértékben újrafelhasználhatók, kezelik a kérések ütemezését és kötegelését, valamint lekérik a dekódolt tokeneket, beágyazásokat vagy pontszámokat az alapul szolgáló modellekből. Az eszközök inicializálását, a folyamatok és csomópontok közötti kommunikációs csatornákat, valamint a súlyok betöltését mind ezek kezelik. Az LLM-ek esetében a motor felelős a kötegelésért, darabolásért és mintavételezésért. Emellett felelős a KV-tárhely kiosztásáért is, mind a teljes, mind a lineáris figyelem (attention), valamint a előtag-illesztés (prefix matching) tekintetében. A beágyazási motorok egyszerűbbek, mivel csak online kötegelést hajtanak végre. A végrehajtás úgy van optimalizálva, hogy a motor előkészítse a következő munkamenetet, amíg a modell párhuzamosan fut a gyorsítón.

A ROSE a modellezéshez egyedi rétegek gazdag gyűjteményét teszi közzé, amelyek a CuTeDSL, a Triton, a CUDA, a CUTLASS és a cuBLAS használatával felépített GPU-kerneleket csomagolják be. A mátrixszorzási és figyelem-implementációkhoz elsősorban a CUTLASS és a cuBLAS NVIDIA-kerneleire támaszkodunk. A folyamat többi része, beleértve a beágyazási, normálási, MoE- és aktiválási kerneleket, a ROSE kerneltárának részét képezi. A futtatási hardver legjobb támogatása érdekében a ROSE saját eszközök közötti kommunikációs primitívekkel van felszerelve az MoE-útválasztáshoz, -elosztáshoz és -kombináláshoz. Bár a többi kernelek viszonylag egyszerűek, alaposan specializálni kell őket a csúcsteljesítmény eléréséhez az összes különböző rejtett dimenzión és eszközön, amelyeken futnak. Ezenkívül a fordítás idejű specializációra támaszkodunk a különböző modellek összes beállításának és finomhangolásának befogadására (súlyelfogultság vs. súlyelfogultság hiánya az RMS normánál, aktiválási függvények stb.). Ennek következtében a megvalósításukhoz a DSL-ek használata mellett döntöttünk: eredetileg a Triton használatával építettük fel a ROSE-t. Azóta azonban a CuTeDSL-t fogadtuk el a GPU-programozás elsődleges környezeteként.

Matematikai szempontból a következtetési kernelek általában egyszerűek. Ennek következtében egy kernel fejlesztésekor a komplexitás általában nem abban rejlik, hogy kifejezzük, milyen számítást kell elvégezni, hanem abban, hogy a hardvernek pontosan hogyan kell azt megtennie. A legjobb kernelek felépítéséhez olyan nyelvre van szükségünk, amely a megfelelő absztrakciókat biztosítja az alapul szolgáló hardver felett. A megfelelő nyelv lehetővé számunkra, hogy megadjuk a számítás elvégzéséhez szükséges magas szintű műveleteket, agresszíven specializálva és optimalizálva azokat a széles konfigurációs tér minden pontján.

A CuTeDSL egy Python-alapú DSL NVIDIA GPU-khoz, amely a CuTe elrendezési algebrára és az MLIR-re épül. Just-in-time fordítást hajt végre a nagymértékben optimalizált PTX-re, miközben fenntartja az alacsony szintű hardverprimitívek feletti ellenőrzést. Célja, hogy a CUTLASS C++-szal megegyező funkciókat és teljesítményt nyújtsa, a fejlesztési sebességet hátráltató magas fordítási idők nélkül. Az általunk áttekintett vagy használt nyelvek közül (beleértve a CUDA-t, a cuTile-t, a TVM IR-t és a Tritont) ezt tartjuk a legmegfelelőbbnek a kerneleink fejlesztéséhez.

Összességében, bár nem olyan magas szintű és tömör, mint a Triton vagy a cuTile, azt találtuk, hogy a CuTe elrendezési algebra és annak kompozíciói (amelyek a cute.Tensor révén érhetők el) ésszerű helyettesítők voltak. Ezenkívül a CuTeDSL korlátozás nélküli hozzáférést biztosít a hardverprimitívekhez, biztosítva számunkra azt az ellenőrzési szintet, amelyre szükségünk van a csúcsteljesítmény eléréséhez minden élmezőbeli esetben. A CUDA esetében a fordítás idejű specializáció lehetőségei korlátozottak voltak, mivel minden egyes további paraméter indokolatlan költséget rótt a fordítási időre vagy növelte a buildrendszer komplexitásaságát. Ezzel szemben a Just-In-Time fordítású DSL-del agresszívebben specializálódhatunk az iterációs ütem feláldozása nélkül. Ellentétben a TVM IR-rel, amely a CUDA-t használta közbenső lépésként, az MLIR-alapú folyamatot, amely a PTX-re ereszkedik le, robusztusabbnak és könnyebben hibakereshetőnek találtuk, mint azt, amely a CUDA-t bocsátotta ki közbenső lépésként. Ezenkívül a PTX-utasítások becsomagolására szolgáló egyéni műveletek használatával azt találtuk, hogy teljes mértékben kiaknázhatjuk az alapul szolgáló hardver összes funkcióját.

A CUDA-kernelek fordítási költségének nagy része általában a kernelek sablonbővítési fázisában megy el, mivel a PTX-stádiumtól kezdődő optimalizálás és fordítás másodpercek alatt elvégezhető. Számos fordítás idejű paraméter mellett a fordítás ezen szakasza több percre is duzzadhat egy nem triviális kernel esetében. Ezt a szűk keresztmetszetet az a tény is fokozza, hogy a párhuzamos fordítás lehetőségei csak később, az IR- és PTX-szinteken merülnek fel, a C++ sablonkezelés pedig szekvenciálisan történik. Bár léteznek trükkök ezek kezelésére a sablonpéldányosítások külön fordítási egységekbe való áthelyezésével, azok jelentős módosításokat igényelnek a buildrendszerekben, valamint nehezen kezelhető makrók használatát. Ezzel szemben a Just-In-Time (JIT) fordított CuTeDSL-kernelek nem igényelnek összetett buildszabályokat, és nem akadályozzák a fordítási időt, mivel a fejlesztési ciklus a konfigurációs tér egyetlen pontjára összpontosítható, míg a későbbi gondos egységtesztek lefedhetik a kernel teljes tervezési terét.

Annak ellenére, hogy a JIT fordítás a fejlesztési időről a futási időre helyezi át a költségeket, a többletmunka több mint elfogadható. Amikor egy modellt üzembe helyeznek, az a teljes konfigurációs térnek csak néhány pontját aktiválja, ami másodperces nagyságrendű többletköltséget okoz a következtetési motor indítása/bemelegítése során. Ezzel szemben egy kernel összes konfigurációjának előfordítása több percet is igénybe vehet. Ezenkívül a többletmunka tovább mérsékelhető a JIT-fordított műtermékek gyorsítótárazásával.

A CUDA magas fordítási költségét fokozza a hibák eredetének felkutatásának nehézsége, amikor sablonalapú elosztást és specializációt használnak, mint a CuTe és CUTLASS esetében. Egy DSL-lel a hibák rámutathatnak arra a kódsorra, amely létrehozta azokat, jelentősen megkönnyítve a hibakeresést. Ezenkívül sokkal egyszerűbb a generált kódot a létrehozó DSL-utasításhozképezni a közbenső MLIR-szakaszokon keresztül, vagy a generált PTX-hez csatolt hibakeresési információkra támaszkodva. A DSL ezen aspektusainak köszönhetően sokkal produktívabbnak találtuk magunkat a kernelek fejlesztése során.

Bár a kernelek sok hasonlóságot mutatnak a modellek között, mégis fel kell őket szerelni egy sor konfigurációs gombbal, amelyek hatással lehetnek a teljesítményre. A vonatkozó paraméterek a következők:

Rejtett dimenzió / Fejdimenzió / Figyelemfejek száma

Bemeneti/kimeneti adattípusok

Kvantálási séma

Fordítás idejű konstansok (torzítások, aktiválási függvények stb.)

A ROSE architektúra, beleértve az engine-jét, a modelljét, a kernelkönyvtárát és a szolgáltatásait.

A fordítás idejű specializáció hasznos a teljesítmény és a hibakeresés szempontjából egyaránt. Különösen egy csővezetékes (pipelined) processzorban, mint amilyen a GPU is, a nem létező ág jobb, mint a mindig vett vagy soha nem vett ág. Ezenkívül a kernel fordítása a konfigurációs gombok egy adott készletére kevesebb PTX-/SASS-utasítást eredményez, amelyeket könnyebb elemezni és hibakeresni a cuda-gdb használatával egy vizsgálat során.

Adott kötegméret és konfiguráció esetén megbecsülhető a DRAM felé irányuló várható forgalom és a aritmetikai műveletek száma. Ezt követően eseti alapon optimális döntések hozhatók az optimális rácselosztásért, szinkronizációért és bizonyos hardverfunkciók használatáért. A CuTeDSL lehetővé teszi számunkra, hogy komponenseket definiáljunk és újrahasznosítsunk a mély specializációhoz anélkül, hogy elfogadhatatlan fordítás idejű többletköltséggel járna.

Általában a következtetési kernelek teljesítményét két különálló use case-re kell optimalizálni:

Előtöltés (Prefill): kis kötegméret (kevés darab, szekvencia), sok token szekvenciánként. Az átviteli sebesség maximalizálása mind a beágyazások, mind az LLM-ek bemeneti tokenjei számára. Késleltetés a tucattól a századik ezredmásodpercig.

Dekódolás: kis- és nagy kötegméretek, egy (spekulatív dekódolás nélkül) vagy néhány token szekvenciánként (vázlat-cél spekuláció, többtokonos előrejelzés). A kernelindítási késleltetéshez közeli késleltetés.

A GEMM és a figyelem (attention) esetében az előtöltési szakasz nagy számú tokenen fut, és a teljesítményt az eszköz számítási átviteli sebessége határozza meg. A dekódolás ugyanezt a műveletet foglalja magában kis kötegméreteken, ahol a végrehajtási időt az dominálja, hogy a DRAM-ból betöltik az aktiválásokat és a súlyokat a folyamatos multiprocesszorokba és tenzormagokba, amit ehelyett az eszköz memóriasávszélessége korlátoz. Annak ellenére, hogy a kötegméret vagy egy kötegben lévő tokenek összesített száma viszonylag világos mutató, az előtöltés és a dekódolás közötti váltási küszöb finomabb, mivel a memóriasávszélességet a modell rejtett dimenziója és az aktiválás adattípusa is meghatározza. Ennek következtében a kernelek előtöltési és dekódolási megvalósításokra oszlanak, ha az algoritmikus különbségek vagy a rács-/blokk-hozzárendelések bármelyiküknek előnyére válnának.

Az a paraméter, amely univerzálisan eltér az előtöltési és a dekódolási kernelek között, a blokkban lévő warp-ok száma (a szálak száma osztva a warp-onkénti 32 szálból álló warp-mérettel). Kívánatos ezt a paramétert fordítás idejű konstanssá tenni, mivel ez segíthet a regiszterek lefoglalásában. A CUDA-ban ezt általában úgy érik el, hogy megadják a szálak számát a __launch_bounds__ segítségével, és a warp-ok számát sablonparaméterként adják meg:

cpp

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

}

A blokk összes szálja számára elérhető regiszterek száma hardverkorlát, 64K a Hopper és Blackwell esetében. Ha a szálak száma előre ismert, a fordító pontosan tudja, hogy hány regisztert rendelhet egyes szálakhoz, és sikerrel fordíthat agresszívabban optimalizált kerneleket. Ez azonban fordítás idejű költséggel jár, mivel most a kernelt újra kell fordítani a különböző warp-méretekhez. A CuTeDSL-lel ez már nem probléma, mivel nem kell előre megfizetnünk a teljes konfigurációs tér fordítási költségét. A DSL azt is elvárja, hogy a blokkméret konstans kifejezés legyen, implicit módon végrehajtva a CUDA-ban szereplő __latch_bounds__ által sugallt optimalizálást.

A warp-ok konstanstová tétele mellett kívánatos azokat az előtöltés és a dekódolás között változtatni. A GPU memóriasávszélességének teljes kihasználásához nem elegendő mindig magas számú warp-ot rendelni blokkonként, mivel a teljes sávszélesség csak az elérhető folyamatos multiprocesszorok (SM-ek) minél nagyobb számának használatával érhető el. Következésképpen egy adott számú tokenhez, valamint az ismert adattípusokhoz és rejtett dimenziókhoz először kiválasztunk egy ideális szálankénti vektordimenziót (8 a bfloat16 esetében, 4 a float32 esetében), hogy biztosítsuk a generált kód 128 bites betöltési/tárolási utasításainak (Blackwellen 256 bites) használatát. Ismerve a vektorméretet, kiszámíthatjuk a kernel által igényelt warp-ok számát. Először megpróbáljuk maximalizálni az SM-ek számát: ha kevesebb warp van, mint SM, akkor blokkonként egyetlen warp-ot választunk. Ez tipikusan a kernel dekódolási változata. Ha elegendő warp van, fokozatosan növeljük a warp-ok számát 8-ra, 16-ra vagy 32-re, bevizsgálva (benchmarking) egy felső határt, amely után tipikusan már nem látunk javulást a teljesítményben.

Bár a Triton a motorháztető alatt végrehajtott néhány ilyen optimalizálást, azt találtuk, hogy a CuTe elrendezési algebra gondos használatával hasonló eredményeket érhetünk el anélkül, hogy a magas szintű fordító hiányosságai korlátoznának minket. Különösen akkor, amikor al-warp vízszintes redukciók szerepelnek, azt találtuk, hogy a CuTeDSL biztosította számunkra az optimális kód generálásához szükséges finom szemcsézetű vezérlést.

Bizonyos helyzetekben a munka minél több blokk közötti elosztásához a cooperative_groups::this_grid().sync() használható az összes részt vevő blokk szinkronizálására. Korlátozások érvényesek: mivel a végrehajtás megkezdése után egy blokkot nem lehet ütemezésből kivenni (de-schedule) egy SM-ből, az előrehaladás garantálása érdekében a blokkok száma nem haladhatja meg az eszközön elérhető SM-ek számát, biztosítva, hogy a szinkronizációban részt vevő globális memóriacímeket bármely ponton az összes blokk frissíthesse és lekérdezhesse. Ezt a korlátozást a cudaLaunchCooperativeKernel kényszeríti ki. A rácsgát (grid barrier) költsége Hopporon körülbelül 3 us, plusz bizonyos helyzetekben az összes korábbi memíriaművelet kiürítésének (flushing) költsége egy implicit memóriagáton keresztül. Összességében ez összehasonlítható egy kernel két indításra való felosztásának többletköltségével, ami teljesítményérzékeny tervezési döntéshez vezet. Annak ellenére, hogy a Programmatic Dependent Launch (PDL) csökkentheti az indítási többletköltségeket Hopportól kezdve, a dekódolási kernelek esetében megéri elkerülni a többletköltséget.

Ha a terhelések egyike több blokkot igényel, mint az SM-ek száma, a kernel felosztása javasolt, a kernelindítások implicit gátjára támaszkodva. Ez elsősorban az előtöltés esetére igaz. Ha a rácsba illeszkedik az összes terheléshez ideális blokkok száma, a rácsgát előnyösebb, mint a kisebb számú tokenen működő dekódolási kernelek esetében.

A CuTeDSL-ben ezek a döntések meghozhatók a hoszt oldali indítási kódban, és a kernelek szükség szerint összeállíthatók és rácsgátakkal összefűzhetők, ahogy az alábbi példa is mutatja. Egy ilyen hibrid stratégiával a dekódolási kernelek optimalizálhatók a minimális késleltetés elérése érdekében, az előtöltési kernelek pedig a maximális átviteli sebességre optimalizálhatók. A CUDA-alapú megoldásokkal ellentétben ez a logika tömörebben és kevesebb fordítás idejű többletköltséggel definiálható.

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

Az MoE all-to-all kerneleink esetében – amint az alábbiakban részletezzük – az a rugalmasság, hogy rácsgáttal vagy anélkül építhetünk kerneleket, és a fordítás időpontjában optimalizálhatjuk őket, segített jelentős teljesítménynövekedést elérni a CuTeDSL révén.

Az első kernel, amelyet a CuTeDSL-ben építettünk, egy specializált RMS-norm kernel volt, amelyet kis dimenziókra optimalizáltak. A QK-norma, amelyet a legújabb Qwen modellek önfigyelési (self-attention) rétegében használnak, fej-alapon (per-head), 64, 128 vagy 256 elemen alkalmazható. Ellentétben a szokásos RMS-normával, amely ezreket meghaladó rejtett dimenziókon működik, és mind a megosztott memórián, mind a regisztereken keresztül történő vízszintes redukciót igényel, a QK-normának csak warp-ok között van szüksége redukcióra, felajánlva annak a lehetőségét, hogy több fejet préseljenek be egyetlen warp-ba. Ez olyan use case-vé vált, ahol agresszíven optimalizálhatunk a fordítás idejében, hogy megtaláljuk a legjobb token-warp hozzárendeléseket az összes fejdimenzióhoz és adattípushoz.

Az RMS-norma meglehetősen egyértelmű művelet, bár vízszintes összeget igényel egy vektor négyzetes normájának kiszámításához:

cpp

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

Mivel a figyelemfejek általában kicsik, és a szálankénti műveletek erősen vektorizáltak, egyetlen warp ki tudja számítani a négyzetes normát. A memóriasávszélesség maximalizálása érdekében minden egyes szálnak 128 bitnyi bemenetet kell feldolgoznia, ami 8 bfloat16 értéknek felel meg. Következésképpen a fél-warp (half-warp) elegendő egy fej feldolgozásához, és a négyzetes norma meghatározható egy a shuffle_sync_bfly által végzett vízszintes redukció használatával. A fejdimenzióból és az adattípusból kiindulva meghatározhatjuk a warp-onkénti fejek számát, és megadhatjuk az argumentumokat és a redukciós körök számát a fordítás idején.

Ezután felépíthetünk egy specializált kernelt viszonylag kevés kódsorral, az automatikusan vektorizált betöltésekre/tárolásokra támaszkodva, és alacsonyabb szintű, architektúra-specifikus műveleteket használva a warp-redukcióhoz. A kernel meghatározza, hogy melyik fejen (Q vagy K) kell dolgozni, és azonosítja a alkalmazandó megfelelő súlyokat. Betölt egy teljes vektort a fejből, és először helyben összeadja. Ezután a fejenkénti szálak száma alapján kiszámításra kerül a négyzetes norma. Végül a wautomatikusan vektorizált RMS-norma meghatározásra kerül, és visszaíródik a memóriába. A modellspecifikus paraméterek, mint például a WEIGHT_BIAS és az EPS, konstanst behajtva (constant-folded) szerepelnek.

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)

Egy szokásos RMS-norm kernelhez képest ez a specializált változat 2-3-szor jobb átviteli sebességet és késleltetést biztosít mind az előtöltés, mind a dekódolás terén.

Az MoE elosztási és kombinálási műveletei azokra a szakértőkre keverik át (shuffle) a bemeneti tokeneket, amelyekhez irányítva vannak, és súlyozott akkumulációt hajtanak végre azután, hogy az MoE GEMM kernelek elvégzik a vonatkozó mátrixszorzásokat. A teljesítmény különösen érdekes, mivel a különböző termékek különböző méretű MoE modellekre támaszkodnak. Az üzembehelyezési típusok alapján különböző optimalizálandó kerneleket különböztetünk meg:

TP=1: Helyi kernel egyetlen GPU-hoz. Az elosztás a megfelelő szakértőkhöz keveri a tokeneket, a kombinálás pedig akkumulálja őket. Elsősorban H100-on futó Qwen3 30B üzembe helyezésekhez használatos.

DP=1, TP ≠ 1: Tenzor-párhuzamos, egynodális (singlenode) kernel. A bemenet replikálva van, az elosztás kiválasztja az aktuális ranghoz tartozó tokeneket, és a kombinálás az összes rangon keresztül akkumulálja őket. Akkumuláció NVLinken vagy NVLS-en keresztül. Nagy előtöltőkhöz (prefiller) használatos.

DP ≠ 1, NVLink: Adatpárhuzamos kernel opcionális tenzorpárhuzamossággal. Minden egyes rang tokeneket küld minden más rangnak, amely sajátos szakértőket ad otthont. Az útválasztási információk aggregálódnak, a tokenek átvitelre kerülnek az NVIDIA NVLinken keresztül. Adatpárhuzamos H200-as vagy GB200-as üzembe helyezésekhez használatos.

DP ≠ 1, InfiniBand: Peer-to-Peer adatpárhuzamos és tenzorpárhuzamos kernel, amely mind az NVLink, mind az NVIDIA Quantum InfiniBand hálózaton keresztül működik. Többcsomópontos Hopper üzembe helyezésekhez használatos.

A helyi kernel az útválasztási táblázat alapján megszámolja az egyes szakértőkhöz rendelt tokenek számát, és egyedi indexet rendel minden egyes tokenhez az egyes szakértőkhöz tartozó tokenek listáján belül. Az útválasztási lépés után a tokenek opcionálisan kitöltött (padded) pufferekbe másolódnak át, hogy továbbítsdák őket egy MoE GEMM kernelnek. A tokenek helyét használva a kombinálási kernel ezután akkumulációt hajt végre a regiszterekben, a végső tokeneket egy kimeneti pufferbe írva.

Az elosztási fázis eltér az előtöltés és a dekódolás között: kevés token esetén kevés az olvasandó útválasztási információ, ezért egyetlen blokkban dolgozzuk fel, és jelezzük a rács fennmaradó részét, miközben az egyes blokkok tokenek darabjait (chunks) másolják. A token-hozzárendelést a megosztott memóriában lévő számlálók atomi növelésével határozzuk meg. Mivel az előtöltésnek több útválasztási információt kell feldolgoznia, az útválasztás több blokkra oszlik egy külön kernelben. Az indexek először helyben kerülnek hozzárendelésre a megosztott memóriában lévő számlálók használatával, majd a globális memóriában lévő összes blokkon keresztül aggregálódnak az egyes tokenek pozíciójának meghatározásához. Ennek a kernelnek a rácsa az SM-ek számára van korlátozva (clamped), hogy rács-szintű gátakat lehessen használni. Egy utólagos kernel hajtja végre a másolásokat, sokkal nagyobb számú blokkon és warpon indítva a rendelkezésre álló memóriasávszélesség teljes kihasználásához.

A kombinálási kernel viszonylag egyszerű, mivel vektorizált akkumulációt hajt végre, az útválasztási információkra és az elosztás által kiszámított eltolásokra (offsets) támaszkodva. A CuTeDSL itt értékesnek bizonyult, mivel a számítás kifejezhető volt egyszerű CuTe-tenzorokon keresztül, amelyek hatékony kóddá fordulnak le az összes adattípushoz és rejtett dimenzióhoz.

Egy tenzorpárhuzamos üzembe helyezésben az útválasztási információk és a bemeneti tokenek replikálva vannak. A helyi elosztási kernel teljes mértékben újrahasznosítható azzal a kis módosítással, hogy csak az aktuális rangon található szakértőkhöz rendelt tokeneket irányítsa és válassza ki. Néhány paraméterrel és a fordítás idejű specializációval a helyi elosztási kernelek újrahasznosításra kerülnek.

Mivel a rangok száma általában alacsonyabb, mint azon szakértők száma, amelyekhez egy tokent irányítanak, az átlagos esetben feltételezhető, hogy minden egyes tokennek legalább az egyik replikája az egyes rangokhoz van irányítva. Ebben a helyzetben a helyi kombinálási kernel újrahasznosítható egy token helyi másolatainak aggregálására, amit egy all-reduce követ a részt vevő rangok közötti összeadáshoz. Az előtöltési kernel esetében külön kerneleket indítunk, a dekódolás azonban egyetlen indításba olvasztja mindkét műveletet.

Az all-reduce-hoz sima NVLinket használunk, ha csak két rang vesz részt, és NVLS-re váltunk át, amikor 4 vagy több rang érintett, 10-20%-os átviteli sebességnövekedés érdekében. A DSL-nek köszönhetően megoszthatjuk a redukciós logika nagy részét és a peer-puffer cseréket az MoE és a szokásos all-reduce kernelek között. Ez olyan helyzet, amikor a hardver feletti szoros ellenőrzés előnyös, mivel az NVLS-hez a multimem.ld_reduce és a multimem.st feletti DSL-műveletekre van szükség. Magának a redukciónak a megvalósítása nagyjából követi a NCCL tesztek all_reduce példáját.

Az adatpárhuzamos kernelek az NVLinkre támaszkodnak a tokenek több eszközhöz való aggregálásához. Az elosztási fázisban, a tokenek cseréje előtt, a peer-ek először helyben irányítanak, hogy azonosítsák az egyes küldött tokenek cél-eltolását az egyes rangokon. Ezt követően elkezdik átmásolni a tokeneket az egyes küldők számára lefoglalt fogadó pufferekbe. Tokenenkénti release-store írások szolgálnak a befejezés jelzésére. A fogadó oldalon az útválasztási információk aggregálódnak, hogy minden tokent egyetlen pufferbe lehessen sűríteni az MoE GEMM számára. A fogadók egy acquire-load segítségével kérdezik le a peer által írt jelölőt, hogy megvárják az írások befejeződését.

A kombinálási oldalon a tokenek átmásolódnak az NVLinken keresztül a forrásrangokra, egy tokenenkénti jelölővel jelölve a kézhezvételt. A kombinálási kernel hasonló a helyi és az egynodális változatokhoz, kisebb különbségekkel a Címezhetőségben és némi további szinkronizációs logikával.

Létrehoztunk egy PTX-burkolók library-t, amely kiteszi a rendszer- és eszközszinten is hatályos release-acquire betöltéseket és tárolásokat ezen kernel megvalósításához. A szinkronizációs jelölők részletességének javításával és több blokk lefoglalásával a tokenek NVLinken keresztüli küldésére és fogadására 2-szeres javulást értünk el a korábbi CUDA-kerneleinkhez képest.

A Peer-to-Peer kerneleink korábban csúcsteljesítményt értek el a ConnectX-7 InfiniBand adapterekkel párosított Hopper GPU-kon. Azóta a teljesítmény további növelése érdekében portoltuk a kombinálási kerneleket a CuTeDSL-re. Mivel a CUDA-verziók már eleve nagyban támaszkodtak a fordítás idejű specializációra, ezek bizonyultak a tárolónk fordítás idejű szűk keresztmetszeteinek, ami kényszerített minket néhány paraméter kihagyására. A JIT-fordított DSL-re való áttéréssel több paraméterre specializálódtunk, és finomhangoltuk a rácskiosztásokat, további különbségeket bevezetve az előtöltés és a dekódolás között, ami körülbelül 3%-os késleltetési növekedést eredményezett.

A ROSE kernelek szemszögéből a Blackwell-támogatás két kategóriába sorolható: architektúra-független és architektúra-specifikus kernelek. Az architektúra-specifikus kernelek, amelyek tenzormagokra támaszkodnak, kifejezett munkát igényelnek az adaptáláshoz. A DSL használatának köszönhetően azonban az architektúra-független kernelek triviálisan adaptálódnak.

Egy Blackwell GPU legfeljebb 160 SM-mel rendelkezik, szemben a Hopper 132-jével. A legtöbb kernel számára ez egy fordítás idejű konstans, és a hoszt oldali indító kód logikája figyelembe veszi ezt a futtatandó optimális (prefill vs decode) kernel meghatározásához. Ezenkívül a Blackwell 256 bites betöltési utasításokat is bevezet, amelyek könnyen kibocsáthatók a vektorszélesség (a CuTe-tenzorok utolsó dimenziója) növelésével a memóriasávszélesség-kihasználás javítása érdekében.

A Blackwell hardveres támogatást vezet be a kvantálási sémák jelentős számához (MXFP8, MXFP4, NVFP4). Ezek támogatása elsősorban a meglévő kernelek FP8 blokk-alapú skálázásra való finomhangolásával lett megvalósítva. A JIT-támogatás itt kulcsfontosságú, mivel az új módok a fordítás idejű konfigurációs tér felrobbanása nélkül támogathatók. A múltban azt találtuk, hogy a Triton kényelmetlenül fejezi ki a kis elemeken, a warp felén vagy negyedén átívelő vízszintes redukciókat. A redukciók feletti finomabb szemcsézetű vezérlés kulcsfontosságú volt az optimális teljesítmény eléréséhez, amikor abszolút maximum értéket számítunk ki 16 vagy 32 elemen az NVFP4, illetve az MXFP8/MXFP4 esetében.

Viszonylag új nyelvként az LLM-ek kevés kitettséget kaptak a CuTeDSL-kódhoz az előtanításuk során. Ezenkívül a jelenlegi kódunk stílusa, az aktiválásokra és normákra optimalizált kernelekkel, kissé eltért a meglévő CUTLASS- vagy Quack-példáktól, amelyek erősen optimalizálva voltak a TMA-ra és az architektúra-specifikus mátrixszorzási utasításokra. Elérni azt a pontot, ahol az LLM-ek segíthetnek a kernelírásban, fokozatos folyamat volt.

Kézzel építettünk fel egy kezdeti kernel- és segédfüggvény-gyűjteményt, következetes mintákat követve, mielőtt a kernelkonverzióhoz ügynökre támaszkodhattunk volna. Ahogy a gyűjtemény előrehaladt, durva fordításokat hajthattunk végre a CUDA-ról CuTeDSL-re, amit jelentős manuális utómunka követett. Különösen azt találtuk, hogy a modellek küzdöttek a pointer-tenzor konverziókkal és a vektorizációval, és elakadtak a típushibák javítási kísérleteiben, amikor egy teszt lefordítására törekedtek. Kevés megértést mutattak a kernelek rácsalapú jellegéről, összetévesztve a szál- és rácsszinkronizációt. A legfrusztrálóbb az volt a kísértésben, hogy a saját CuTeDSL-helyettesítőinkkel való építkezés helyett a CUDA-specifikus könyvtárak nem létező fordításait hallucinálják. Idővel, ahogy a CuTeDSL-kernelek gyűjteménye növekedett, azt találtuk, hogy az LLM-ek fel tudják építeni a vázkód nagy részét, és ki tudják tölteni a hiányosságokat, lehetővé téve számunkra, hogy a magas szintű tervezésre összpontosítsunk.

A CUDA-hoz képest azt találjuk, hogy a CuTeDSL sokkal jobb fordítás idejű teljesítményt nyújt, lehetővé téve számunkra, hogy több paraméterre specializálódjunk, és teljesítménynövekedést érjünk el. Bár a kernelek kissé bőbeszédűbbek a Tritonnál bizonyos absztrakciók hiánya miatt, azt találjuk, hogy a komplexitás indokolt, mivel a hardver feletti finomabb szemcsézetű vezérlés lehetővé teszi számunkra a gyorsabb kernelek felépítését.

A cikk írásakor a CUTLASS C++ és a CuTeDSL-kernelek közötti teljesítménykülönbség minimális, izgalmas lehetőségeket nyitva meg a nagy teljesítményű GEMM-kernelek eszközök közötti kommunikációval, normalizációval vagy aktiválásokkal való összetételéhez és fúziójához. Ezen a DSL-en elkötelezve magunkat képesek leszünk további teljesítményt kisajtolni az eszközeinkből a munkaterheléseink átviteli sebességének javítására, a beágyazásoktól az LLM-ekig.

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

Következtetés a Perplexity AI-nál

Miért a CuTeDSL?

Fejlesztői élmény

Nagy konfigurációs tér

Előtöltés és dekódolási specializáció

Rácsszinkronizáció

CuTeDSL-kernelek

QK Norm

Szakértők keveréke (Mixture-of-Experts) elosztás/kombinálás

Helyi (Local)

Egynodális (Singlenode)

Adatpárhuzamos (Data Parallel)

Peer-to-Peer

Blackwell-támogatás

MI-támogatású kernelgenerálás

Következtetés és jövőbeli munka

Hivatkozások

Idézet