CuTeDSL в Perplexity

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

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

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

На верхнем уровне ROSE состоит из движка, который адаптирует модель к интерфейсу, с которым могут взаимодействовать клиенты. Движки обладают высокой степенью переиспользования, обрабатывая планирование и пакетирование запросов, а также извлекая декодированные токены, эмбеддинги или оценки из базовых моделей. Инициализация устройств, каналов межпроцессного и межпроцессорного взаимодействия, а также загрузка весов — всё это обрабатывается ими. Для LLM движок отвечает за пакетирование, фрагментирование и сэмплирование. Он также отвечает за распределение KV-хранилища как для полной, так и для линейной аттеншн, а также за префиксное сопоставление. Движки эмбеддингов проще, так как они выполняют только пакетную обработку в режиме онлайн. Выполнение оптимизировано таким образом, что движок подготавливает следующую порцию работы параллельно с тем, как модель выполняется на ускорителе.

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

С математической точки зрения ядра инференса обычно несложны. Следовательно, при разработке ядра сложность обычно заключается не в том, чтобы выразить, какие вычисления должны быть выполнены, а в том, как именно оборудование должно их выполнить. Чтобы создавать лучшие ядра, нам требуется язык, который предлагает правильные абстракции поверх базового оборудования. Правильный язык позволяет нам специфицировать высокоуровневые операции для выполнения вычислений, агрессивно специализируя и оптимизируя их во всех точках широкого пространства конфигураций.

CuTeDSL — это DSL на базе Python для GPU NVIDIA, построенный поверх алгебры раскладки CuTe и MLIR. Он компилируется «на лету» (JIT) в высокооптимизированный 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, включая ее движок, модель, библиотеку ядра и службы.

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

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

Как правило, производительность ядер инференса должна быть оптимизирована для двух отдельных вариантов использования:

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

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

Для GEMM и аттеншн этап prefill выполняется на большом количестве токенов, а производительность ограничена вычислительной пропускной способностью устройства. Decode включает ту же операцию на малых размерах пакетов, где время выполнения определяется стоимостью загрузки активаций и весов из DRAM в потоковые мультипроцессоры и тензорные ядра, ограничиваясь вместо этого пропускной способностью памяти устройства. Хотя размер пакета или общее количество токенов в пакете является относительно четким индикатором, порог переключения между prefill и decode более тонкий, поскольку пропускная способность памяти также определяется скрытой размерностью модели и типом данных активации. Следовательно, ядра разделяются на реализации prefill и decode, если алгоритмические различия или назначения сетки/блоков дали бы преимущество любому из них.

Параметр, который повсеместно различается между ядрами prefill и decode, — это количество варпов в блоке (количество тредов, деленное на размер варпа в 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.

Помимо того, что количество варпов является константой, желательно варьировать его между prefill и decode. Чтобы в полной мере использовать пропускную способность памяти GPU, недостаточно всегда назначать большое количество варпов на блок, так как полная пропускная способность может быть достигнута только при использовании как можно большего количества доступных потоковых мультипроцессоров (SM). Следовательно, для заданного количества токенов и известных типов данных и скрытых размерностей мы сначала выбираем идеальную векторную размерность на тред (8 для bfloat16, 4 для float32), чтобы гарантировать, что сгенерированный код использует 128-битные инструкции загрузки/сохранения (256-битные на Blackwell). Зная размер вектора, мы можем вычислить количество варпов, требуемое ядром. Мы сначала пытаемся максимизировать количество SM: если варпов меньше, чем SM, выбирается один варп на блок. Обычно это вариант ядра decode. Если варпов достаточно, мы постепенно увеличиваем количество варпов до 8, 16 или 32, проводя бенчмаркинг верхней границы, после которой мы обычно больше не наблюдаем улучшения производительности.

Хотя Triton выполнял некоторые из этих оптимизаций «под капотом», мы обнаружили, что можем достичь аналогичных результатов, осторожно используя алгебру раскладки CuTe, не будучи ограниченными недостатками высокоуровневого компилятора. Особенно когда дело доходит до суб-варпных горизонтальных редукций, мы обнаружили, что CuTeDSL предложил нам тонкозернистый контроль, необходимый для генерации оптимального кода.

В некоторых ситуациях, чтобы распределить работу между как можно большим количеством блоков, можно использовать cooperative_groups::this_grid().sync() для синхронизации всех участвующих блоков. Применяются ограничения: поскольку блок не может быть снят с планирования из SM после того, как он начал выполнение, чтобы гарантировать прогресс вперед, количество блоков не может превышать количество SM, доступных на устройстве, чтобы гарантировать, что в любой момент времени все блоки могут обновлять и опрашивать глобальные ячейки памяти, участвующие в синхронизации. Это ограничение обеспечивается cudaLaunchCooperativeKernel. Стоимость барьера сетки составляет около 3 мкс на Hopper, плюс стоимость сброса всех предыдущих операций с памятью через неявный барьер памяти в некоторых ситуациях. В целом, это сопоставимо с накладными расходами на разделение ядра на два запуска, что приводит к критическому с точки зрения производительности проектному решению. Несмотря на то, что Программный зависимый запуск (PDL) может снизить накладные расходы на запуск начиная с Hopper, для ядер decode избегание этих накладных расходов является целесообразным.

Если одна из рабочих нагрузок требует больше блоков, чем количество SM, ядро должно быть разделено, полагаясь на неявный барьер запусков ядер. В основном это касается prefill. Если количество блоков, которое идеально подходит для всех рабочих нагрузок, помещается в сетку, предпочтительнее использовать барьер сетки, как это происходит с ядрами decode, работающими с меньшим количеством токенов.

В CuTeDSL эти решения могут быть приняты в коде запуска на стороне хоста, а ядра могут быть составлены и перемежены барьерами сетки по мере необходимости, как проиллюстрировано в примере ниже. Благодаря такой гибридной стратегии ядра decode могут быть оптимизированы для достижения минимальных задержек, а ядра prefill — для максимальной пропускной способности. В отличие от решений на основе 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, сворачиваются как константы.

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 раза лучшую пропускную способность и задержку как для prefill, так и для decode.

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

TP=1: Локальное ядро для одного GPU. Диспетчеризация перетасовывает токены к соответствующим экспертам, а комбинирование накапливает их. Используется в основном для развертываний Qwen3 30B на H100.

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

DP ≠ 1, NVLink: Ядро параллелизма по данным с опциональным тензорным параллелизмом. Каждый ранг отправляет токены на каждый другой ранг, на котором размещены специфические эксперты. Информация о маршрутизации агрегируется, токены передаются через NVIDIA NVLink. Используется для развертываний с параллелизмом по данным H200 или GB200.

DP ≠ 1, InfiniBand: Одноранговое ядро параллелизма по данным и тензорного параллелизма, работающее как через NVLink, так и через NVIDIA Quantum InfiniBand. Используется для многоузловых развертываний Hopper.

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

Фаза диспетчеризации варьируется между prefill и decode: для небольшого количества токенов мало информации для маршрутизации, поэтому мы обрабатываем её в одном блоке и сигнализируем остальной части сетки, при этом отдельные блоки копируют фрагменты токенов. Назначение токенов определяется атомарным увеличением счетчиков в общей памяти. Поскольку prefill должен обрабатывать больше информации о маршрутизации, маршрутизация разделяется на несколько блоков в отдельном ядре. Индексы сначала назначаются локально с использованием счетчиков в общей памяти, а затем агрегируются по всем блокам в глобальной памяти для определения позиции каждого токена. Сетка этого ядра привязывается к количеству SM, чтобы иметь возможность использовать сеточные барьеры. Последующее ядро выполняет копирование, запускаемое по гораздо большему количеству блоков и варпов, чтобы в полной мере использовать доступную пропускную способность памяти.

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

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

Поскольку количество рангов обычно меньше количества экспертов, к которым маршрутизируется токен, в среднем случае можно предположить, что хотя бы одна из реплик каждого токена маршрутизируется на каждый ранг. В этой ситуации локальное ядро комбинирования может быть переиспользовано для агрегирования локальных копий токена, за которым следует all-reduce для суммирования по всем участвующим рангам. В случае ядра prefill мы запускаем отдельные ядра, однако decode объединяет обе операции в один запуск.

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

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

На стороне комбинирования токены копируются через NVLink на ранги-источники, с флагом на каждый токен, указывающим на получение. Ядро комбинирования похоже на локальный и одноузловой варианты, с небольшими различиями в адресации и некоторой дополнительной логикой синхронизации.

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

Наши P2P-ядра ранее достигали передовой производительности на GPU Hopper в паре с адаптерами InfiniBand ConnectX-7. С тех пор мы портировали ядра комбинирования на CuTeDSL для дальнейшего повышения производительности. Поскольку версии CUDA уже сильно полагались на специализацию во время компиляции, они оказались «узкими местами» времени компиляции в нашем репозитории, заставляя нас опускать некоторые параметры. Перейдя на DSL, компилируемый «на лету», мы специализировались на большем количестве параметров и скорректировали распределение сетки, внеся дополнительные различия между prefill и decode для увеличения задержки примерно на 3%.

С точки зрения ядер ROSE поддержка Blackwell делится на две категории: ядра, не зависящие от архитектуры, и ядра, специфичные для архитектуры. Ядра, специфичные для архитектуры, полагающиеся на тензорные ядра, требуют явной работы для адаптации. Однако благодаря использованию DSL ядра, не зависящие от архитектуры, адаптируются тривиально.

GPU Blackwell имеет до 160 SM, в отличие от 132 у Hopper. Для большинства ядер это константа времени компиляции, и логика в коде запуска на стороне хоста учитывает это при определении оптимального ядра (prefill или decode) для запуска. Кроме того, Blackwell также вводит 256-битные инструкции загрузки, которые могут быть легко сгенерированы путем увеличения ширины вектора (последняя размерность тензоров CuTe) для улучшения использования пропускной способности памяти.

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

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

Мы вручную создали начальный корпус ядер и служебных функций, следуя согласованным шаблонам, прежде чем смогли положиться на агентов для конвертации ядер. По мере развития корпуса мы могли выполнять грубый перевод с CUDA на CuTeDSL со значительной последующей ручной доработкой. В частности, мы обнаружили, что модели испытывали трудности с преобразованием указателей в тензоры и векторизацией, застревая при попытке исправить ошибки типов в стремлении добиться компиляции теста. Они демонстрировали мало понимания сеточной природы ядер, путая синхронизацию тредов и сетки. Самое неприятное — присутствовало искушение галлюцинировать несуществующие переводы библиотек, специфичных для CUDA, вместо того чтобы строить их на наших собственных заменах CuTeDSL. Со временем, по мере роста корпуса ядер CuTeDSL, мы обнаружили, что LLM могут создавать большую часть каркасного кода и заполнять пробелы, позволяя нам сосредоточиться на высокоуровневом дизайне.

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

На момент написания этой статьи разрыв в производительности между ядрами CUTLASS C++ и CuTeDSL минимален, что открывает захватывающие возможности для компоновки и слияния высокопроизводительных ядер GEMM с межпроцессорным взаимодействием, нормализацией или активациями. Сделав ставку на этот DSL, мы сможем извлечь дополнительную производительность из наших устройств для улучшения пропускной способности всех наших рабочих нагрузок, от эмбеддингов до LLM.

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

https://github.com/NVIDIA/cutlass

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

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

BibTeX

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

Инференс в Perplexity AI

Почему CuTeDSL?

Опыт разработчика

Большое пространство конфигураций

Предзаполнение и декодирование специализации

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

Ядра CuTeDSL

QK Norm

Диспетчеризация/Комбинирование Mixture-of-Experts

Локальное

Одноузловое

Параллелизм по данным

Одноранговое (Peer-to-Peer)

Поддержка Blackwell

Генерация ядер с помощью ИИ

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

Ссылки

Цитирование