CuTeDSL na Perplexity

Na Perplexity, usamos o motor proprietário Runtime-Optimized Serving Engine (ROSE) para servir modelos que vão desde embeddings até LLMs de um trilhão de parâmetros. Posicionado por trás de todas as APIs, incluindo Sonar, Search e Embeddings, nosso ambiente de execução foi criado para capacitar pesquisadores a explorar os modelos e arquiteturas mais adequados, ao mesmo tempo em que implantam rapidamente os melhores candidatos em produção. Para alcançar isso, precisamos ser capazes de implementar, implantar e otimizar a inferência para arquiteturas de ponta em um ritmo acelerado. Nesta postagem do blog, exploramos como os kernels CuTeDSL integrados ao nosso motor de inferência nos permitem construir os kernels de GPU necessários para levar os modelos ao desempenho máximo em GPUs NVIDIA.

Modelos personalizados, construídos por várias equipes da Perplexity, são hospedados internamente em GPUs NVIDIA Hopper e Blackwell usando o motor de inferência ROSE construído pela equipe de Inferência de IA. O ROSE foi inicialmente construído para atender melhor a modelos Llama personalizados, tanto para decodificação de linguagem quanto para classificação, com uma interface compatível com o NVIDIA Triton Inference Server. Com o tempo, ele evoluiu para um motor de inferência completo que atende a LLMs substanciais, além de modelos de classificação, ranqueamento, pontuação e outros baseados em transformadores para uma ampla gama de equipes para impulsionar os produtos da Perplexity. O objetivo de design do ROSE é ser simples: com o tempo, o esforço é investido para atingir o desempenho máximo, mas a principal prioridade é se adaptar às arquiteturas e casos de uso em constante evolução para nossos modelos.

O espaço de configuração do kernel entre tipos de dados, dimensões de cabeçalho, constantes em tempo de compilação e configurações de lançamento.

Em alto nível, o ROSE consiste em um motor que adapta um modelo a uma interface com a qual os clientes podem interagir. Os motores são altamente reutilizáveis, gerenciando o agendamento e a criação de lotes de solicitações, e recuperando tokens decodificados, embeddings ou pontuações dos modelos subjacentes. A inicialização de dispositivos, canais de comunicação entre processos e entre nós, e o carregamento de pesos são todos gerenciados por eles. Para LLMs, o motor é responsável pelo agrupamento em lotes (batching), divisão em pedaços (chunking) e amostragem (sampling). Ele também é responsável pela alocação de armazenamento KV, tanto para atenção completa quanto linear, bem como pela correspondência de prefixos. Os motores de embedding são mais simples, pois realizam apenas agrupamento em lotes online. A execução é otimizada de forma que o motor prepara a próxima rodada de trabalho enquanto o modelo está rodando no acelerador em paralelo.

O ROSE expõe uma rica coleção de camadas personalizadas para modelagem, envolvendo kernels de GPU construídos com CuTeDSL, Triton, CUDA, CUTLASS e cuBLAS. Para implementações de multiplicação de matrizes e atenção, contamos principalmente com kernels NVIDIA do CUTLASS e cuBLAS. O restante do pipeline, incluindo embeddings, norma, MoE e kernels de ativação, faz parte da biblioteca de kernels do ROSE. Para melhor suportar o hardware no qual executamos a inferência, o ROSE está equipado com suas próprias primitivas de comunicação entre dispositivos para roteamento, despacho e combinação de MoE. Embora os outros kernels sejam relativamente simples, eles devem ser profundamente especializados para alcançar o desempenho máximo em todas as diferentes dimensões ocultas e dispositivos em que executam. Além disso, contamos com a especialização em tempo de compilação para acomodar todos os ajustes e parâmetros de vários modelos (viés de peso vs. sem viés de peso para norma RMS, funções de ativação, etc.). Consequentemente, decidimos usar DSLs para implementá-los: inicialmente construímos o ROSE usando Triton. No entanto, desde então adotamos o CuTeDSL como o ambiente principal para programação de GPU.

De uma perspectiva matemática, os kernels de inferência são geralmente pouco sofisticados. Consequentemente, ao desenvolver um kernel, a complexidade geralmente não reside em expressar qual computação deve ser feita, mas em como exatamente o hardware deve fazê-lo. Para construir os melhores kernels, exigimos uma linguagem que ofereça as abstrações corretas sobre o hardware subjacente. A linguagem certa nos permite especificar as operações de alto nível para realizar a computação, especializando-as e otimizando-as agressivamente em todos os pontos em um amplo espaço de configuração.

O CuTeDSL é uma DSL baseada em Python para GPUs NVIDIA, construída sobre a álgebra de layout CuTe e MLIR. Ele compila just-in-time para PTX altamente otimizado, mantendo o controle sobre primitivas de hardware de baixo nível. O objetivo é oferecer os mesmos recursos e desempenho do CUTLASS C++, sem os tempos de compilação caros que dificultam a velocidade de desenvolvimento. Das linguagens que pesquisamos ou com as quais trabalhamos (incluindo CUDA, cuTile, TVM IR e Triton), consideramos que ele é o mais adequado para o desenvolvimento de nossos kernels.

No geral, embora não seja tão de alto nível e conciso quanto o Triton ou cuTile, descobrimos que a álgebra de layout CuTe e suas composições (expostas via cute.Tensor) eram substitutos razoáveis. Além disso, o CuTeDSL oferece acesso irrestrito a primitivas de hardware, fornecendo-nos o nível de controle que exigimos para alcançar o desempenho máximo em todos os casos de borda. Com o CUDA, as oportunidades de especialização em tempo de compilação eram limitadas, pois cada parâmetro adicional adicionava um custo irracional aos tempos de compilação ou complexidade ao sistema de compilação. Em contraste, com uma DSL compilada Just-In-Time, podemos nos spezializar mais agressivamente sem sacrificar o ritmo de iteração. Ao contrário do TVM IR, que usava CUDA como uma etapa intermediária, achamos que o pipeline baseado em MLIR que desce para PTX era mais robusto e mais fácil de depurar do que aquele que emitia CUDA como uma etapa intermediária. Além disso, com o uso de operações personalizadas para envolver instruções PTX, descobrimos que podíamos aproveitar plenamente todos os recursos do hardware subjacente.

A maior parte do custo de compilação dos kernels CUDA é normalmente gasta na fase de expansão de modelo (template) dos kernels, pois a otimização e a compilação a partir do estágio PTX em diante podem ser concluídas em questão de segundos. Com muitos parâmetros em tempo de compilação, esse estágio de compilação pode explodir para vários minutos para um kernel não trivial. Esse gargalo é acentuado pelo fato de que as oportunidades de compilação paralela surgem apenas mais tarde, nos níveis de IR e PTX, com o manuseio de modelos (templates) em C++ ocorrendo sequencialmente. Embora existam truques para resolver isso, movendo as instanciações de modelos para unidades de compilação separadas, eles exigem ajustes significativos nos sistemas de compilação e o uso de macros pesadas. Em contraste, os kernels CuTeDSL que são compilados Just-In-Time (JIT) não exigem regras de compilação complexas e não atrapalham os tempos de compilação, pois o ciclo de desenvolvimento pode ser focado em um único ponto no espaço de configuração, enquanto testes de unidade cuidadosos posteriores podem cobrir todo o espaço de design do kernel.

Embora a compilação JIT mova os custos do tempo de desenvolvimento para o tempo de execução, a sobrecarga é mais do que aceitável. Quando um modelo é implantado, ele ativa apenas alguns dos pontos de todo o espaço de configuração, incorrendo em sobrecargas de escala de segundos durante a inicialização/aquecimento do motor de inferência. Em contraste, pré-compilar todas as configurações de um kernel pode abranger vários minutos. Além disso, as sobrecargas podem ser mitigadas ainda mais armazenando em cache os artefatos compilados por JIT.

O pesado custo de compilação do CUDA é acentuado pela dificuldade de rastrear a origem dos erros quando o despacho baseado em modelo (template) e a especialização são usados, como é o caso do CuTe e do CUTLASS. Com uma DSL, os erros podem ser apontados para a linha de código que os gerou, facilitando significativamente a depuração. Além disso, é muito mais simples mapear o código gerado para a instrução DSL que o originou por meio de estágios intermediários de MLIR ou contando com informações de depuração anexadas ao PTX gerado. Graças a esses aspectos da DSL, nos vimos muito mais produtivos ao desenvolver kernels.

Embora os kernels compartilhem muitas semelhanças entre os modelos, eles ainda devem ser equipados com uma gama de botões de configuração que podem ter um impacto no desempenho. Os parâmetros relevantes incluem:

Dimensão Oculta/Dimensão de Cabeça/Número de Cabeças de Atenção

Tipos de dados de Entrada/Saída

Esquema de quantização

Constantes em tempo de compilação (vieses, funções de ativação, etc.)

Arquitetura ROSE, incluindo seu mecanismo, modelo, biblioteca de kernel e serviços.

A especialização em tempo de compilação é útil tanto para o desempenho quanto para a depuração. Particularmente em um processador em pipeline, como é o caso das GPUs, um ramo inexistente é melhor do que um ramo sempre tomado ou nunca tomado. Além disso, compilar um kernel para um conjunto específico de botões de configuração resulta em menos instruções PTX/SASS, que são mais fáceis de analisar e depurar ao usar o cuda-gdb durante uma investigação.

Para um tamanho de lote e configuração específicos, o tráfego esperado para a DRAM e o número de operações aritméticas podem ser estimados. Posteriormente, decisões ótimas podem ser tomadas caso a caso para a alocação ideal de grades, sincronização e o uso de determinados recursos de hardware. O CuTeDSL nos permite definir e reutilizar componentes para especialização profunda sem incorrer em uma sobrecarga inaceitável em tempo de compilação.

Geralmente, o desempenho dos kernels de inferência deve ser otimizado para dois casos de uso separados:

Prefill: tamanho de lote pequeno (poucos pedaços, sequências), muitos tokens por sequência. A vazão (throughput) deve ser maximizada tanto para embeddings quanto para os tokens de entrada de LLMs. Latência na faixa de dezenas a centenas de milissegundos.

Decode: tamanhos de lote pequenos a grandes, um (sem decodificação especulativa) a alguns tokens por sequência (especulação de alvo de rascunho, previsão de múltiplos tokens). Latência próxima à latência de lançamento do kernel.

Para GEMM e atenção, o estágio de prefill é executado em um grande número de tokens e o desempenho é limitado pela vazão de computação do dispositivo. A decodificação envolve a mesma operação em tamanhos de lote pequenos, com tempos de execução dominados pelo custo de carregar as ativações e pesos da DRAM para os multiprocessadores de fluxo (streaming multiprocessors) e núcleos de tensor, limitados em vez disso pela largura de banda de memória do dispositivo. Mesmo que o tamanho do lote ou o número total de tokens em um lote seja um indicador relativamente claro, o limite de transição entre prefill e decode é mais sutil porque a largura de banda de memória também é determinada pela dimensão oculta do modelo e pelo tipo de dado da ativação. Consequentemente, os kernels são divididos em implementações de prefill e decode se diferenças algorítmicas ou atribuições de grade/bloco beneficiarem qualquer uma delas.

O parâmetro que é universalmente diferente entre os kernels de prefill e decode é o número de warps em um bloco (o número de tópicos, dividido pelo tamanho de warp de 32 tópicos por warp). É desejável tornar este parâmetro uma constante em tempo de compilação, pois pode ajudar na alocação de registradores. No CUDA, isso é normalmente feito especificando o número de tópicos via __launch_bounds__ e fornecendo o número de warps como um parâmetro de modelo (template):

cpp

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

}

O número de registradores disponíveis para todos os tópicos em um bloco é uma restrição de hardware, 64K no Hopper e Blackwell. Se o número de tópicos for conhecido com antecedência, o compilador sabe exatamente quantos registradores pode atribuir a cada tópico e pode ter sucesso na compilação de kernels otimizados de forma mais agressiva. No entanto, isso tem um custo em tempo de compilação, pois agora o kernel precisa ser recompilado para diferentes tamanhos de warp. Com o CuTeDSL, isso não é mais um problema, pois não precisamos pagar o custo de compilação por todo o espaço de configuração antecipadamente. A DSL também espera que o tamanho do bloco seja uma expressão constante, realizando implicitamente a otimização implicada por __latch_bounds__ no CUDA.

Além de tornar o número de warps uma constante, é desejável variá-lo entre prefill e decode. Para utilizar plenamente a largura de banda de memória de uma GPU, não é suficiente atribuir sempre um número alto de warps por bloco, pois a largura de banda total só pode ser alcançada usando o máximo possível de Multiprocessadores de Fluxo (SMs) disponíveis. Consequentemente, para um número dado de tokens e tipos de dados e dimensões ocultas conhecidos, escolhemos primeiro uma dimensão vetorial por tópico ideal (8 para bfloat16, 4 para float32) para garantir que o código gerado use as instruções de carga/armazenamento de 128 bits (256 bits no Blackwell). Conhecendo o tamanho do vetor, podemos computar o número de warps exigidos pelo kernel. Tentamos primeiro maximizar o número de SMs: se houver menos warps do que SMs, um único warp por bloco é escolhido. Esta é tipicamente a variante de decodificação (decode) do kernel. Se houver warps suficientes, aumentamos gradualmente o número de warps para 8, 16 ou 32, avaliando (benchmarking) um limite superior após o qual normalmente não vemos mais melhoria no desempenho.

Embora o Triton tenha realizado algumas dessas otimizações por baixo dos panos, descobrimos que podemos alcançar resultados semelhantes usando cuidadosamente a álgebra de layout CuTe, sem sermos limitados pelas deficiências de um compilador de alto nível. Particularmente quando reduções horizontais de sub-warp estão envolvidas, descobrimos que o CuTeDSL nos ofereceu o controle granular necessário para gerar código ideal.

Em algumas situações, para distribuir o trabalho pelo maior número possível de blocos, cooperative_groups::this_grid().sync() pode ser usado para sincronizar todos os blocos participantes. Aplicam-se restrições: como um bloco não pode ser des-agendado de um SM uma vez que inicia a execução, para garantir o progresso adiante, o número de blocos não pode exceder o número de SMs disponíveis no dispositivo, para garantir que a qualquer momento todos os blocos possam atualizar e sondar os locais de memória global envolvidos na sincronização. Esta restrição é aplicada por cudaLaunchCooperativeKernel. O custo de uma barreira de grade é de cerca de 3us no Hopper, mais o custo de descarregar (flushing) todas as operações de memória anteriores através de uma barreira de memória implícita em algumas situações. No geral, isso é comparável à sobrecarga de dividir um kernel em dois lançamentos, levando a uma decisão de design sensível ao desempenho. Apesar do fato de que o Programmatic Dependent Launch (PDL) pode reduzir as sobrecargas de lançamento a partir do Hopper em diante, para kernels de decodificação, evitar a sobrecarga vale a pena.

Se uma das cargas de trabalho exigir mais blocos do que o número de SMs, o kernel deve ser dividido, contando com a barreira implícita dos lançamentos de kernel. Este é principalmente o caso do prefill. Se o número de blocos que são ideais para todas as cargas de trabalho couber na grade, uma barreira de grade é preferida, como é o caso dos kernels de decodificação operando em um número menor de tokens.

No CuTeDSL, essas decisões podem ser tomadas no código de lançamento do lado do host e os kernels podem ser compostos e intercalados com barreiras de grade conforme necessário, conforme ilustrado pelo exemplo abaixo. Com tal estratégia híbrida, os kernels de decodificação podem ser otimizados para atingir latências mínimas e os de prefill podem ser otimizados para a vazão máxima. Ao contrário das soluções baseadas em CUDA, esta lógica pode ser definida de forma mais concisa e com menos sobrecarga em tempo de compilação.

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

No caso dos nossos kernels MoE all-to-all, conforme detalhado abaixo, ter a flexibilidade de construir kernels com ou sem barreiras de grade e otimizá-los em tempo de compilação nos ajudou a alcançar ganhos substanciais de desempenho através do CuTeDSL.

O primeiro kernel que construímos no CuTeDSL foi um kernel de norma RMS especializado, otimizado para pequenas dimensões. A norma QK, usada na camada de auto-atenção dos recentes modelos Qwen, é aplicada por cabeça, em 64, 128 ou 256 elementos. Ao contrário da norma RMS regular, que opera em dimensões ocultas superiores a milhares, exigindo uma redução horizontal através de memória compartilhada e registradores, a norma QK precisa de redução apenas entre warps, com a possibilidade de espremer várias cabeças em um único warp. Este se tornou um caso de uso onde pudemos otimizar agressivamente em tempo de compilação para encontrar as melhores atribuições de token para warp para todas as dimensões de cabeça e tipos de dados.

A norma RMS é uma operação bastante direta, embora exija uma soma horizontal para computar a norma quadrada de um vetor:

cpp

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

Como as cabeças de atenção são tipicamente pequenas e as operações por tópico são fortemente vetorizadas, um único warp pode computar a norma quadrada. Para maximizar a largura de banda de memória, cada tópico deve processar 128 bits de entrada, o que se traduz em 8 valores de bfloat16. Consequentemente, um meio-warp é suficiente para processar uma cabeça, e a norma quadrada pode ser determinada usando uma redução horizontal realizada por shuffle_sync_bfly. A partir da dimensão da cabeça e do tipo de dado, podemos determinar o número de cabeças por warp e especificar os argumentos e o número de rodadas de redução no momento da compilação.

Podemos então construir um kernel especializado com relativamente poucas linhas de código, contando com cargas/armazenamentos vetorizados automaticamente e usando operações de baixo nível específicas da arquitetura para a redução de warp. O kernel determina em qual cabeça (Q ou K) trabalhar e identifica os pesos correspondentes a serem aplicados. Ele carrega um vetor inteiro da cabeça e primeiro soma localmente. Subsequentemente, com base no número de tópicos por cabeça, a norma quadrada é computada. Finalmente, a norma RMS, vetorizada automaticamente, é determinada e gravada de volta na memória. Parâmetros específicos do modelo, como WEIGHT_BIAS e EPS, são dobrados em constantes (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)

Comparado a um kernel de norma RMS regular, esta variante especializada entrega uma vazão e latência de 2 a 3 vezes melhores para prefill e decode.

O despacho e a combinação de MoE embaralham os tokens de entrada para os especialistas para os quais são roteados, realizando uma acumulação ponderada após os kernels MoE GEMM realizarem as multiplicações de matrizes relevantes. O desempenho é de particular interesse, pois vários produtos dependem de modelos MoE de vários tamanhos. Com base nos tipos de implantação, distinguimos diferentes kernels para otimizar:

TP=1: Kernel local para uma única GPU. O despacho embaralha os tokens para os especialistas relevantes e a combinação os acumula. Usado principalmente para implantações do Qwen3 30B no H100.

DP=1, TP ≠ 1: Kernel paralelo a tensor, de nó único. A entrada é replicada, o despacho seleciona os tokens para a classificação (rank) atual e a combinação os acumula em todas as classificações. Acumulação através de NVLink ou NVLS. Usado para prefillers grandes.

DP ≠ 1, NVLink: Kernel paralelo a dados com paralelismo de tensor opcional. Cada classificação envia tokens para cada outra classificação que hospeda especialistas específicos. As informações de roteamento são agregadas, os tokens são transferidos através do NVIDIA NVLink. Usado para implantações de H200 ou GB200 paralelas a dados.

DP ≠ 1, InfiniBand: Kernel paralelo a dados e paralelo a tensor Ponto a Ponto operando em ambos os NVLink e NVIDIA Quantum InfiniBand. Usado para implantações Hopper em vários nós.

O kernel local conta o número de tokens atribuídos a cada especialista com base na tabela de roteamento e atribui um índice exclusivo a cada token dentro da lista de tokens pertencentes a cada especialista. Após a etapa de roteamento, os tokens são copiados em buffers opcionalmente preenchidos (padded) para serem repassados a um kernel MoE GEMM. Usando os locais dos tokens, o kernel de combinação realiza então uma acumulação em registradores, gravando os tokens finais em um buffer de saída.

A fase de despacho varia entre prefill e decode: para poucos tokens, há pouca informação de roteamento para ler, então a processamos em um único bloco e sinalizamos o restante da grade, com blocos individuais copiando pedaços de tokens. A atribuição de tokens é determinada incrementando atomicamente contadores na memória compartilhada. Como o prefill deve processar mais informações de roteamento, o roteamento é dividido em vários blocos em um kernel separado. Os índices são primeiro atribuídos localmente usando contadores na memória compartilhada e são então agregados em todos os blocos na memória global para determinar a posição de cada token. A grade deste kernel é fixada (clamped) ao número de SMs para poder usar barreiras a nível de grade. Um kernel subsequente executa as cópias, lançado em um número muito maior de blocos e warps para utilizar plenamente a largura de banda de memória disponível.

O kernel de combinação é relativamente simples, pois realiza acumulação vetorizada, contando com as informações de roteamento e com os deslocamentos (offsets) computados pelo despacho. O CuTeDSL provou ser valioso aqui, pois a computação podia ser expressada por meio de tensores CuTe simples que compilam para código eficiente para todos os tipos de dados e dimensões ocultas.

Em uma implantação paralela a tensor, as informações de roteamento e os tokens de entrada são replicados. O kernel de despacho local pode ser totalmente reutilizado com uma pequena alteração para rotear e selecionar apenas os tokens atribuídos aos especialistas localizados na classificação (rank) atual. Com alguns parâmetros e especialização em tempo de compilação, os kernels de despacho local são reutilizados.

Como o número de classificações é geralmente menor do que o número de especialistas para os quais um token é roteado, no caso médio pode-se assumir que cada token tem pelo menos uma de suas réplicas roteada para cada classificação. Nesta situação, o kernel de combinação local pode ser reutilizado para agregar as cópias locais de um token, seguido por um all-reduce para somar através das classificações participantes. No caso do kernel de prefill, lançamos kernels separados, no entanto o decode funde ambas as operações em um único lançamento.

Para all-reduce, usamos NVLink simples se apenas dois ranks participam, mudando para NVLS quando 4 ou mais ranks estão envolvidos, para um aumento de 10-20% na vazão. Graças ao DSL, podemos compartilhar a maior parte da lógica de redução e as trocas de buffer entre pares entre os kernels MoE e os kernels all-reduce regulares. Trata-se de uma situação em que um controle rigoroso sobre o hardware é benéfico, já que são necessárias operações DSL em multimem.ld_reduce e multimem.st para NVLS. A implementação da própria redução segue em grande parte o exemplo de all_reduce dos testes do NCCL.

Os kernels paralelos a dados dependem do NVLink para agregar tokens para múltiplos dispositivos. Na fase de despacho, antes de trocar tokens, os pares primeiro roteiam localmente para identificar o deslocamento de destino em cada classificação para cada token que estão enviando. Subsequentemente, eles começam a copiar os tokens para buffers de recebimento alocados para cada remetente para evitar conflitos. Escritas release-store por token são usadas para sinalizar a conclusão. No lado do receptor, as informações de roteamento são agregadas para poder condensar todos os tokens em um único buffer para o MoE GEMM. Os receptores sondam a flag escrita pelo par com um acquire-load para aguardar a conclusão das escritas.

No lado da combinação, os tokens são copiados via NVLink para as classificações de origem, com uma flag por token indicando o recebimento. O kernel de combinação é semelhante às variantes local e de nó único, com pequenas diferenças no endereçamento e alguma lógica de sincronização adicional.

Construímos uma biblioteca de wrappers PTX para expor cargas e armazenamentos release-acquire escopados tanto no nível do sistema quanto do dispositivo para implementar este kernel. Ao melhorar a granularidade das flags de sincronização e alocar mais blocos para enviar e receber tokens via NVLink, alcançamos melhorias de 2x em relação aos nossos kernels CUDA anteriores.

Nossos kernels Ponto a Ponto alcançavam anteriormente desempenho de ponta em GPUs Hopper emparelhadas com adaptadores InfiniBand ConnectX-7. Desde então, portamos os kernels de combinação para o CuTeDSL para ganhos adicionais de desempenho. Como as versões CUDA já dependiam fortemente da especialização em tempo de compilação, elas provaram ser os gargalos de tempo de compilação em nosso repositório, forçando-nos a omitir alguns parâmetros. Ao migrar para a DSL compilada por JIT, nos especializamos para mais parâmetros e ajustamos alocações de grade, introduzindo mais diferenças entre prefill e decode para um aumento de cerca de 3% na latência.

Da perspectiva dos kernels ROSE, o suporte ao Blackwell se enquadra em duas categorias: kernels agnósticos de arquitetura e específicos de arquitetura. Os kernels específicos de arquitetura, dependendo de núcleos de Tensor, exigem trabalho explícito para se adaptarem. No entanto, graças ao uso da DSL, os kernels agnósticos de arquitetura se adaptam trivialmente.

Uma GPU Blackwell possui até 160 SMs, em contraste com os 132 do Hopper. Para a maioria dos kernels, esta é uma constante em tempo de compilação e a lógica no código do lançador do lado do host a considera para determinar o kernel ideal (prefill vs decode) a ser executado. Além disso, o Blackwell também introduz instruções de carga de 256 bits, que podem ser facilmente emitidas aumentando a largura do vetor (última dimensão dos tensores CuTe) para melhorar a utilização da largura de banda de memória.

O Blackwell introduz suporte de hardware para um número significativo de esquemas de quantização (MXFP8, MXFP4, NVFP4). O suporte a estes foi implementado principalmente ajustando kernels existentes para escalonamento por bloco FP8. O suporte JIT é crucial aqui, pois novos modos podem ser suportados sem explodir o espaço de configuração em tempo de compilação. No passado, achamos o Triton estranho ao expressar reduções horizontais em poucos elementos, sobre metade ou um quarto de um warp. O controle mais refinado sobre as reduções foi crucial para alcançar o desempenho ideal ao computar um valor máximo absoluto em 16 ou 32 elementos para NVFP4 e MXFP8/MXFP4, respectivamente.

Como uma linguagem relativamente nova, os LLMs tiveram pouca exposição ao código CuTeDSL durante seu pré-treinamento. Além disso, o estilo do nosso código atual, com kernels otimizados para ativações e normas, era ligeiramente diferente dos exemplos existentes do CUTLASS ou Quack, que eram fortemente otimizados para instruções de multiplicação de matrizes específicas de arquitetura e TMA. Atingir um ponto onde os LLMs pudessem ajudar na escrita de kernels foi um processo gradual.

Construímos manualmente um corpus inicial de kernels e funções utilitárias, seguindo padrões consistentes, antes de podermos contar com agentes para a conversão de kernels. Conforme o corpus progredia, podíamos realizar traduções brutas de CUDA para CuTeDSL, com retrabalho manual significativo em seguida. Em particular, descobrimos que os modelos lutavam com conversões de ponteiro para tensor e vetorização, ficando presos tentando corrigir erros de tipo ao tentar fazer um teste compilar. Eles mostraram pouca compreensão da natureza baseada em grade dos kernels, confundindo a sincronização de tópicos e grades. Mais frustrantemente, havia uma tentação de alucinar traduções inexistentes de bibliotecas específicas do CUDA em vez de construir em cima de nossas próprias substituições CuTeDSL. Com o tempo, à medida que o corpus de kernels CuTeDSL crescia, descobrimos que os LLMs podem construir a maior parte do código esqueleto e preencher as lacunas, permitindo-nos focar no design de alto nível.

Comparado ao CUDA, achamos que o CuTeDSL entrega um desempenho em tempo de compilação muito melhor, permitindo-nos especializar para mais parâmetros e alcançar melhorias de desempenho. Embora os kernels sejam ligeiramente mais verbosos devido à falta de certas abstrações em comparação com o Triton, achamos que a complexidade é justificada, pois o controle mais refinado sobre o hardware nos permite construir kernels mais rápidos.

No momento da escrita, a lacuna de desempenho entre os kernels CUTLASS C++ e CuTeDSL é mínima, abrindo oportunidades empolgantes para compor e fundir kernels GEMM de alto desempenho com comunicação entre dispositivos, normalização ou ativações. Ao nos comprometermos com esta DSL, seremos capazes de extrair mais desempenho de nossos dispositivos para melhorar a vazão em nossas cargas de trabalho, desde embeddings até LLMs.

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

Inferência na Perplexity AI

Por que o CuTeDSL?

Experiência do desenvolvedor

Grande espaço de configuração

Especialização em pré-preenchimento e decodificação

Sincronização de grade

Kernels CuTeDSL

Norma QK

Despacho/Combinação de Mistura de Especialistas (Mixture-of-Experts)

Local

Nó único (Singlenode)

Paralelo a Dados (Data Parallel)

Ponto a Ponto (Peer-to-Peer)

Suporte ao Blackwell

Geração de Kernel Assistida por IA

Conclusão e Trabalho Futuro

Referências

Citação