CuTeDSL у Perplexity

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

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

Простір конфігурації ядра для різних типів даних, розмірів голівок, констант часу компіляції та налаштувань запуску.

На високому рівні ROSE складається з рушія, який адаптує модель до інтерфейсу, з яким можуть взаємодіяти клієнти. Рушії є високопридатними для повторного використання, керуючи плануванням і пакетизацією запитів, а також отримуючи декодовані токени, ембединги або оцінки з базових моделей. Вони забезпечують ініціалізацію пристроїв, міжпроцесні та міжвузлові канали зв'язку, а також завантаження вагів. Для LLM рушій відповідає за пакетизацію, розбиття на блоки (chunking) та семплювання. Він також відповідає за виділення пам'яті під KV-кеш для повної та лінійного уваги, а також за пошук збігів префіксів. Рушії ембедингів простіші, оскільки вони виконують лише онлайнову пакетизацію. Виконання оптимізовано так, що рушій готує наступний раунд роботи паралельно з роботою моделі на прискорювачі.

ROSE надає багату колекцію спеціальних шарів для моделювання, що обгортають графічні ядра (GPU kernels), створені за допомогою CuTeDSL, Triton, CUDA, CUTLASS та cuBLAS. Для множення матриць та реалізації механізму уваги ми переважно покладаємося на ядра NVIDIA з CUTLASS та cuBLAS. Решта конвеєра, включно з ембедингами, нормуванням, MoE та ядрами активації, є частиною бібліотеки ядер ROSE. Для найкращої підтримки апаратного забезпечення, на якому ми виконуємо інференс, ROSE оснащено власними примітивами міжпристроєвого зв'язку для маршрутизації MoE, диспетчеризації та об'єднання. Хоча інші ядра є відносно простими, їх необхідно ретельно спеціалізувати для досягнення пікової продуктивності на всіх різних прихованих розмірностях і пристроях, на яких вони працюють. Крім того, ми покладаємося на спеціалізацію під час компіляції, щоб врахувати всі налаштування та нюанси різних моделей (зміщення вагів порівняно з відсутністю зміщення вагів для нормування RMS, функції активації тощо). Як наслідок, ми вирішили використовувати DSL для їх реалізації: спочатку ми створили ROSE на базі Triton. Однак відтоді ми прийняли CuTeDSL як основне середовище для програмування графічних процесорів.

З математичної точки зору ядра інференсу зазвичай є нескладними. Отже, під час розробки ядра складність зазвичай полягає не у вираженні того, які обчислення слід виконати, а в тому, як саме апаратне забезпечення повинно це зробити. Щоб створити найкращі ядра, нам потрібна мова, яка пропонує правильні абстракції над базовим апаратним забезпеченням. Правильна мова дозволяє нам визначати операції високого рівня для виконання обчислень, агресивно спеціалізуючи та оптимізуючи їх у всіх точках широкого простору конфігурацій.

CuTeDSL — це мова специфікації домену (DSL) на базі Python для графічних процесорів NVIDIA, створена на основі алгебри макетів CuTe та MLIR. Вона компілюється «на льоту» (just-in-time) в оптимізований PTX, зберігаючи при цьому контроль над низькорівневими апаратними примітивами. Вона спрямована на забезпечення таких самих можливостей і продуктивності, як CUTLASS C++, без високих витрат часу на компіляцію, які гальмують швидкість розробки. З-поміж мов, які ми досліджували або з якими працювали (зокрема CUDA, cuTile, TVM IR та Triton), ми вважаємо її найкращим вибором для розробки наших ядер.

Загалом, хоча вона не така високорівнева та лаконічна, як Triton чи cuTile, ми виявили, що алгебра макетів CuTe та її композиції (що надаються через cute.Tensor) є розумними замінниками. Крім того, CuTeDSL пропонує необмежений доступ до апаратних примітивів, забезпечуючи той рівень контролю, який нам необхідний для досягнення пікової продуктивності в усіх граничних випадках. З CUDA можливості спеціалізації на етапі компіляції були обмеженими, оскільки кожен додатковий параметр додавав невиправданих витрат часу на компіляцію або складності системі збірки. Натомість із 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 ми виявили, що стали набагато продуктивнішими під час розробки ядер.

Хоча ядра мають багато подібностей у різних моделях, вони все одно повинні бути оснащені низкою ручок конфігурації, які можуть впливати на продуктивність. Відповідні параметри включають:

Прихована розмірність / Розмірність голови / Кількість голів уваги

Типи вхідних/вихідних даних

Схема квантування

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

Архітектура ROSE, включно з її рушієм, моделлю, бібліотекою ядра та службами.

Спеціалізація на етапі компіляції корисна як для продуктивності, так і для налагодження. Зокрема, у конвеєрному процесорі, як у випадку з графічними процесорами, неіснуюча гілка краща, ніж гілка, що завжди або ніколи не виконується. Крім того, компіляція ядра для певного набору ручок конфігурації призводить до меншої кількості інструкцій PTX/SASS, які легше аналізувати та налагоджувати за допомогою cuda-gdb під час розслідування.

Для заданого розміру пакета та конфігурації можна оцінити очікуваний трафік до DRAM та кількість арифметичних операцій. Згодом можна приймати оптимальні рішення для кожного конкретного випадку щодо оптимального виділення сітки, синхронізації та використання певних апаратних функцій. CuTeDSL дозволяє нам визначати та повторно використовувати компоненти для глибокої спеціалізації без неприпустимих накладних витрат на компіляцію.

Загалом продуктивність ядер інференсу має бути оптимізована для двох різних сценаріїв використання:

Prefill: невеликий розмір батчу (мало чанків, послідовностей), багато токенів на послідовність. Пропускну здатність потрібно максимізувати як для ембедингів, так і для вхідних токенів LLM. Затримка становить від десятків до сотень мілісекунд.

Decode: від невеликих до великих розмірів батчу, від одного (без спекулятивного декодування) до кількох токенів на послідовність (спекуляція чернетка-мета, багатотокеневе передбачення). Затримка близька до затримки запуску ядра.

Для GEMM та механізму уваги етап попереднього заповнення працює з великою кількістю токенів, а продуктивність обмежена обчислювальною пропускною здатністю пристрою. Декодування передбачає ту саму операцію для невеликих розмірів батчу, при цьому в часі виконання переважають витрати на завантаження активацій і вагів із DRAM у потокові мультипроцесори та тензорні ядра, що натомість обмежується пропускною здатністю пам'яті пристрою. Навіть попри те, що розмір батчу або загальна кількість токенів у батчі є відносно чітким індикатором, поріг перемикання між попереднім заповненням і декодуванням є більш тонким, оскільки пропускна здатність пам'яті також визначається прихованою розмірністю моделі та типом даних активації. Відповідно, ядра розділяються на реалізації попереднього заповнення та декодування, якщо алгоритмічні відмінності або призначення сітки/блоків дають перевагу будь-якому з них.

Параметром, який незмінно відрізняється між ядрами попереднього заповнення та декодування, є кількість варпів у блоці (кількість тредів, поділена на розмір варпу — 32 треди на варп). Бажано зробити цей параметр константою часу компіляції, оскільки це може допомогти у виділенні регістрів. В CUDA це зазвичай робиться шляхом вказівки кількості тредів через __launch_bounds__ та надання кількості варпів як параметра шаблону:

cpp

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

}

Кількість регістрів, доступних для всіх тредів у блоці, є апаратним обмеженням: 64K на Hopper та Blackwell. Якщо кількість тредів відома заздалегідь, компілятор точно знає, скільки регістрів він може призначити кожному треду, і може успішно компілювати більш агресивно оптимізовані ядра. Однак це супроводжується витратами часу на компіляцію, оскільки тепер ядро доводиться перекомпільовувати для різних розмірів варпів. З CuTeDSL це більше не є проблемою, оскільки нам не потрібно заздалегідь платити за витрати на компіляцію для всього простору конфігурацій. DSL також очікує, що розмір блоку буде константним виразом, неявно виконуючи оптимізацію, передбачену __latch_bounds__ в CUDA.

Окрім того, щоб зробити кількість варпів константою, бажано змінювати її між попереднім заповненням та декодуванням. Для повного використання пропускної здатності пам'яті графічного процесора недостатньо завжди призначати велику кількість варпів на блок, оскільки повної пропускної здатності можна досягти лише за умови використання якомога більшої кількості доступних потокових мультипроцесорів (SM). Отже, для заданої кількості токенів та відомих типів даних і прихованих розмірностей ми спочатку вибираємо ідеальну векторну розмірність на тред (8 для bfloat16, 4 для float32), щоб гарантувати, що згенерований код використовує 128-бітні інструкції завантаження/збереження (256-бітні на Blackwell). Знаючи розмір вектора, ми можемо обчислити кількість варпів, необхідних для ядра. Ми спочатку намагаємося максимізувати кількість SM: якщо варпів менше, ніж SM, вибирається один варп на блок. Зазвичай це варіант ядра для декодування. Якщо варпів достатньо, ми поступово збільшуємо їх кількість до 8, 16 або 32, проводячи бенчмаркінг верхньої межі, після якої ми зазвичай більше не спостерігаємо покращення продуктивності.

Хоча Triton виконував деякі з цих оптимізацій «під капотом», ми виявили, що можемо досягти подібних результатів, обережно використовуючи алгебру макетів CuTe, не будучи обмеженими недоліками високорівневого компілятора. Зокрема, коли задіяні субварпові горизонтальні редукції, ми виявили, що CuTeDSL запропонував нам детальний контроль, необхідний для генерації оптимального коду.

У деяких ситуаціях, щоб розподілити роботу між якомога більшою кількістю блоків, можна використовувати cooperative_groups::this_grid().sync() для синхронізації всіх блоків, що беруть участь. Існують обмеження: оскільки блок не може бути депланований (de-scheduled) з SM після початку його виконання, щоб гарантувати поступальний прогрес, кількість блоків не може перевищувати кількість SM, доступних на пристрої, щоб гарантувати, що в будь-який момент усі блоки можуть оновлювати та опитувати місця глобальної пам'яті, задіяні в синхронізації. Це обмеження примусово застосовується cudaLaunchCooperativeKernel. Вартість бар'єру сітки становить близько 3 мкс на Hopper плюс вартість скидання (flushing) всіх попередніх операцій з пам'яті через неявний бар'єр пам'яті в деяких ситуаціях. Загалом це можна порівняти з накладними витратами на розділення ядра на два запуски, що призводить до прийняття проєктного рішення, чутливого до продуктивності. Попри те, що програмований залежний запуск (Programmatic Dependent Launch, PDL) може зменшити накладні витрати на запуск починаючи з архітектури Hopper, для ядер декодування уникнення цих накладних витрат є виправданим.

Якщо одне з робочих навантажень вимагає більше блоків, ніж кількість SM, ядро слід розділити, покладаючись на неявний бар'єр запусків ядер. В основному це стосується попереднього заповнення. Якщо кількість блоків, яка є ідеальною для всіх робочих навантажень, вміщується в сітку, перевага віддається бар'єру сітки, як у випадку з ядрами декодування, що працюють із меншою кількістю токенів.

У CuTeDSL ці рішення можуть бути прийняті в коді запуску на стороні хоста, а ядра можуть компонуватися та чергуватися з бар'єрами сітки за потреби, як показано на прикладі нижче. Завдяки такій гібридній стратегії ядра декодування можна оптимізувати для досягнення мінімальних затримок, а ядра попереднього заповнення — для максимальної пропускної здатності. На відміну від рішень на базі CUDA, цю логіку можна визначити більш лаконічно та з меншими накладними витратами часу на компіляцію.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

У випадку наших ядер MoE all-to-all, як детально описано нижче, наявність гнучкості для створення ядер з бар'єрами сітки або без них та їх оптимізація на етапі компіляції допомогли нам досягти суттєвого приросту продуктивності завдяки CuTeDSL.

Першим ядром, яке ми створили в CuTeDSL, було спеціалізоване ядро нормалізації RMS, оптимізоване для невеликих розмірностей. Нормування QK, що використовується в шарі само-уваги (self-attention) останніх моделей Qwen, застосовується для кожної голови окремо, на 64, 128 або 256 елементах. На відміну від звичайного нормування RMS, яке працює з прихованими розмірностями, що перевищують тисячі, вимагаючи горизонтальної редукції як через розділювану пам'ять, так і через регістри, нормування QK вимагає редукції лише по варпах із можливістю втиснути кілька голів в один варп. Це стало сценарієм використання, де ми могли агресивно оптимізувати на етапі компіляції, щоб знайти найкращі призначення токенів до варпів для всіх розмірностей голів і типів даних.

Нормування RMS є досить прямолінійною операцією, хоча воно вимагає горизонтальної суми для обчислення квадрата норми вектора:

cpp

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

Оскільки голови уваги зазвичай є невеликими, а операції для кожного треду сильно векторизовані, один варп може обчислити квадрат норми. Щоб максимізувати пропускну здатність пам'яті, кожен тред повинен обробляти 128 бітів вхідних даних, що відповідає 8 значенням bfloat16. Отже, напівварпа достатньо для обробки однієї голови, а квадрат норми можна визначити за допомогою горизонтальної редукції, виконаної за допомогою shuffle_sync_bfly. Починаючи з розмірності голови та типу даних, ми можемо визначити кількість голів на варп і вказати аргументи та кількість раундів редукції під час компіляції.

Потім ми можемо побудувати спеціалізоване ядро з відносно невеликою кількістю рядків коду, покладаючись на автоматично векторизовані завантаження/збереження та використовуючи низькорівневі операції, специфічні для архітектури, для редукції у варпі. Ядро визначає, над якою головою (Q або K) працювати, і визначає відповідні ваги для застосування. Воно завантажує весь вектор із голови та спочатку підсумовує локально. Згодом, залежно від кількості тредів на голову, обчислюється квадрат норми. Нарешті, нормування RMS, автоматично векторизоване, визначається та записується назад у пам'ять. Параметри, специфічні для моделі, такі як WEIGHT_BIAS та EPS, згортаються на етапі компіляції (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)

Порівняно зі звичайним ядром нормалізації RMS, цей спеціалізований варіант забезпечує у 2–3 рази кращу пропускну здатність та затримку як для попереднього заповнення, так і для декодування.

Диспетчеризація та об'єднання MoE перетасовують вхідні токени до експертів, до яких вони спрямовані, виконуючи зважене накопичення після того, як ядра GEMM MoE виконають відповідні множення матриць. Продуктивність становить особливий інтерес, оскільки різні продукти покладаються на моделі MoE різного розміру. Залежно від типів розгортання ми виділяємо різні ядра для оптимізації:

TP=1: Локальне ядро для одного графічного процесора. Диспетчеризація перетасовує токени до відповідних експертів, а об'єднання накопичує їх. Використовується переважно для розгортань Qwen3 30B на H100.

DP=1, TP ≠ 1: Тензорно-паралельне однооузлове ядро. Вхідні дані реплікуються, диспетчеризація вибирає токени для поточного рангу, а об'єднання накопичує їх по всіх рангах. Накопичування через NVLink або NVLS. Використовується для великих модулів попереднього заповнення.

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

DP ≠ 1, InfiniBand: Рівноправне (Peer-to-Peer) дата-паралельне та тензорно-паралельне ядро, що працює як через NVLink, так і через NVIDIA Quantum InfiniBand. Використовується для багатовузлових розгортань Hopper.

Локальне ядро підраховує кількість токенів, призначених кожному експерту на основі таблиці маршрутизації, і призначає унікальний індекс кожному токену в межах списку токенів, що належать кожному експерту. Після етапу маршрутизації токени копіюються в буфери з можливим доповненням (padded buffers) для передачі в ядро GEMM MoE. Використовуючи розташування токенів, ядро об'єднання потім виконує накопичення в регістрах, записуючи кінцеві токени у вихідний буфер.

Фаза диспетчеризації відрізняється між попереднім заповненням та декодуванням: для невеликої кількості токенів є мало інформації про маршрутизацію для зчитування, тому ми обробляємо її в одному блоці та сигналізуємо решті сітки, причому окремі блоки копіюють чанки токенів. Призначення токенів визначається шляхом атомарного інкрементування лічильників у розділюваній пам'яті. Оскільки попереднє заповнення повинно обробляти більше інформації про маршрутизацію, маршрутизація розділена між кількома блоками в окремому ядрі. Індекси спочатку призначаються локально за допомогою лічильників у розділюваній пам'яті, а потім агрегуються по всіх блоках у глобальній пам'яті для визначення позиції кожного токену. Сітка цього ядра обмежена кількістю SM, щоб мати можливість використовувати бар'єри на рівні сітки. Наступне ядро виконує копіювання, запущені на значно більшій кількості блоків і варпів для повного використання доступної пропускної здатності пам'яті.

Ядро об'єднання є відносно простим, оскільки воно виконує векторизоване накопичення, спираючись на інформацію про маршрутизацію та зсуви, обчислені диспетчеризацією. CuTeDSL виявився тут цінним, оскільки обчислення можна було виразити за допомогою простих тензорів CuTe, які компілюються в ефективний код для всіх типів даних і прихованих розмірностей.

У тензорно-паралельному розгортанні інформація про маршрутизацію та вхідні токени реплікуються. Локальне ядро диспетчеризації може бути повністю використане повторно з невеликою зміною для маршрутизації та вибору лише тих токенів, які призначені експертам, розташованим на поточному ранзі. Завдяки деяким параметрам і спеціалізації на етапі компіляції локальні ядра диспетчеризації використовуються повторно.

Оскільки кількість рангів зазвичай менша за кількість експертів, до яких спрямовується токен, у середньому випадку можна припустити, що принаймні одна з копій кожного токену спрямовується на кожен ранг. У цій ситуації локальне ядро об'єднання може бути використане повторно для агрегації локальних копій токену з подальшим використанням операції all-reduce для підсумовування по всіх рангах, що беруть участь. У випадку ядра попереднього заповнення ми запускаємо окремі ядра, проте декодування об'єднує обидві операції в один запуск.

Для all-reduce ми використовуємо звичайний NVLink, якщо беруть участь лише два ранги, перемикаючись на NVLS, коли задіяно 4 або більше рангів, для збільшення пропускної здатності на 10-20%. Завдяки DSL ми можемо спільно використовувати більшу частину логіки редукції та обміну буферами між одноранговими вузлами (peer buffer exchanges) між ядрами MoE та звичайними ядрами all-reduce. Це ситуація, коли тісний контроль над апаратним забезпеченням є корисним, оскільки для NVLS потрібні операції DSL над multimem.ld_reduce та multimem.st. Реалізація самої редукції значною мірою слідує прикладу all_reduce із тестів NCCL.

Дата-паралельні ядра покладаються на NVLink для агрегації токенів для кількох пристроїв. На фазі диспетчеризації перед обміном токенами однорангові вузли спочатку здійснюють локальну маршрутизацію, щоб визначити цільовий зсув на кожному ранзі для кожного токену, який вони надсилають. Згодом вони починають копіювати токени у буфери отримання, виділені для кожного відправника, щоб уникнути конфліктів. Записи release-store для кожного токену використовуються для сигналізації про завершення. На стороні приймача інформація про маршрутизація агрегується, щоб мати можливість згустити всі токени в один буфер для GEMM MoE. Одержувачі опитують прапорець, записаний одноранговим вузлом, за допомогою acquire-load, щоб дочекатися завершення записів.

На стороні об'єднання токени копіюються через NVLink на вихідні ранги з прапорцем для кожного токену, що підтверджує отримання. Ядро об'єднання подібне до локального та одновузлового варіантів, з невеликими відмінностями в адресації та деякою додатковою логікою синхронізації.

Ми створили бібліотеку обгорток PTX, щоб виставити операції завантаження та збереження release-acquire на рівні системи та пристрою для реалізації цього ядра. Покращивши деталізацію прапорців синхронізації та виділивши більше блоків для надсилання й отримання токенів через NVLink, ми досягли 2-кратного покращення порівняно з нашими попередніми ядрами CUDA.

Наші ядра Peer-to-Peer раніше досягали найсучаснішої (state-of-the-art) продуктивності на графічних процесорах Hopper, з'єднаних з адаптерами ConnectX-7 InfiniBand. Відтоді ми перенесли ядра об'єднання на CuTeDSL для подальшого приросту продуктивності. Оскільки версії на CUDA вже сильно покладалися на спеціалізацію під час компіляції, вони виявилися «вузькими горлечками» часу компіляції в нашому репозиторії, змушуючи нас опускати деякі параметри. Здійснивши міграцію на DSL із JIT-компіляцією, ми виконали спеціалізацію для більшої кількості параметрів і налаштували виділення сітки, запровадивши додаткові відмінності між попереднім заповненням та декодуванням для приблизно 3% приросту затримки.

З погляду ядер ROSE підтримка Blackwell поділяється на дві категорії: архітектурно-агностичні та архітектурно-специфічні ядра. Архітектурно-специфічні ядра, що спираються на тензорні ядра, вимагають явної роботи для адаптації. Однак завдяки використанню DSL архітектурно-агностичні ядра адаптуються тривіально.

Графічний процесор Blackwell має до 160 SM, на відміну від 132 у Hopper. Для більшості ядер це константа часу компіляції, і логіка в коді ініціалізатора на стороні хоста враховує її для визначення оптимального ядра (попереднього заповнення чи декодування) для запуску. Крім того, Blackwell також представляє 256-бітні інструкції завантаження, які можна легко згенерувати шляхом збільшення ширини вектора (останньої розмірності тензорів CuTe) для покращення використання пропускної здатності пам'яті.

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

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

Ми вручну створили початковий корпус ядер і допоміжних функцій, дотримуючись узгоджених патернів, перш ніж змогли покладатися на агентів для конвертації ядер. У міру розвитку корпусу ми могли виконувати приблизні переклади з CUDA на CuTeDSL із подальшою значною ручною доробкою. Зокрема, ми виявили, що моделі намагалися впоратися з перетвореннями покажчиків у тензори (pointer-to-tensor) та векторизацією, застрягаючи у спробах виправити помилки типів, коли прагнули домогтися компіляції тесту. Вони продемонстрували мінімальне розуміння сіткової природи ядер, плутаючи синхронізацію тредів і сіток. Найбільш розчаровуючим було спокушання галюцинувати неіснуючі переклади бібліотек, специфічних для CUDA, замість того, щоб збиратися проти наших власних замінників CuTeDSL. З часом, у міру зростання корпусу ядер CuTeDSL, ми виявили, що LLM можуть створювати більшу частину каркасного коду (skeleton code) та заповнювати прогалини, дозволяючи нам зосередитися на проєктуванні високого рівня.

Порівняно з CUDA, ми виявили, що CuTeDSL забезпечує набагато кращу продуктивність часу компіляції, дозволяючи нам спеціалізуватися на більшій кількості параметрів та досягати покращень продуктивності. Хоча ядра є дещо більш багатослівними через відсутність певних абстракцій порівняно з Triton, ми вважаємо цю складність виправданою, оскільки детальніший контроль над апаратним забезпеченням дозволяє нам створювати швидші ядра.

На момент написання цієї статті розрив у продуктивності між ядрами CUTLASS C++ та CuTeDSL є мінімальним, що відкриває захоплюючі можливості для композиції та злиття високоефективних ядер GEMM з міжпристроєвим зв'язком, нормалізацією чи активаціями. Узявши на себе зобов'язання використовувати цей DSL, ми зможемо вичавити додаткову продуктивність з наших пристроїв, щоб покращити пропускну здатність для всіх наших робочих навантажень — від ембедингів до LLM.

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

https://github.com/NVIDIA/cutlass

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

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

BibTeX

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

Інференс у Perplexity AI

Чому CuTeDSL?

Досвід розробника

Великий простір конфігурацій

Попереднє заповнення та декодування спеціалізації

Синхронізація сітки

Ядра CuTeDSL

Нормування QK (QK Norm)

Диспетчеризація / об'єднання суміші експертів (Mixture-of-Experts Dispatch/Combine)

Локальне (Local)

Одновузлове (Singlenode)

Дата-паралельне (Data Parallel)

Рівноправні (Peer-to-Peer)

Підтримка Blackwell

Генерація ядер за допомогою штучного інтелекту

Висновок та майбутня робота

Список літератури

Цитування