CuTeDSL na Perplexity

Na Perplexity, utilizamos o mecanismo de serviço otimizado em tempo de execução (ROSE - Runtime-Optimized Serving Engine) interno para servir modelos que variam de embeddings a LLMs de trilhões de parâmetros. Por trás de todas as APIs, incluindo Sonar, Search e Embeddings, nosso tempo de execução foi criado para capacitar os pesquisadores a explorar os modelos e arquiteturas mais adequados, ao mesmo tempo em que implantamos rapidamente os melhores candidatos em produção. Para atingir esse objetivo, precisamos ser capazes de implementar, implantar e otimizar a inferência para arquiteturas de ponta em um ritmo acelerado. Nesta postagem de blog, exploramos como os kernels CuTeDSL integrados ao nosso mecanismo de inferência nos permitem criar os kernels de GPU necessários para levar os modelos ao desempenho máximo em GPUs NVIDIA.

Modelos personalizados, criados por várias equipes na Perplexity, são hospedados internamente em GPUs NVIDIA Hopper e Blackwell usando o mecanismo de inferência ROSE criado pela equipe de Inferência de IA. O ROSE foi inicialmente criado para servir melhor 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 mecanismo de inferência completo que atende a LLMs substanciais, bem como modelos de classificação, pontuação, ranking e outros baseados em transformadores para uma ampla gama de equipes que impulsionam os produtos da Perplexity. O objetivo do design do ROSE é ser simples: com o tempo, investe-se esforço para atingir o desempenho máximo, mas a prioridade principal é adaptar-se às arquiteturas e casos de uso em constante evolução para nossos modelos.

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

Em um nível elevado, o ROSE consiste em um mecanismo que adapta um modelo a uma interface com a qual os clientes podem interagir. Os mecanismos são altamente reutilizáveis, lidando com agendamento e processamento em lote 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 tratados por eles. Para LLMs, o mecanismo é responsável por processamento em lote, fragmentação e amostragem. Também é responsável pela alocação de armazenamento KV, tanto para atenção total quanto linear, bem como pela correspondência de prefixo. Os mecanismos de embedding são mais simples, pois realizam apenas processamento em lote on-line. A execução é otimizada de forma que o mecanismo prepare a próxima rodada de trabalho enquanto o modelo é executado no acelerador em paralelo.

O ROSE expõe uma rica coleção de camadas personalizadas para modelagem, envolvendo kernels de GPU criados com CuTeDSL, Triton, CUDA, CUTLASS e cuBLAS. Para implementações de multiplicação de matrizes e atenção, dependemos principalmente de kernels da NVIDIA provenientes do CUTLASS e cuBLAS. O restante do pipeline, incluindo kernels de embedding, norm, MoE e ativação, faz parte da biblioteca de kernels do ROSE. Para oferecer o melhor suporte ao hardware no qual executamos a inferência, o ROSE é equipado com seus próprios primitivos de comunicação entre dispositivos para roteamento, despacho e combinação de MoE. Embora os outros kernels sejam relativamente simples, eles devem ser minuciosamente especializados para atingir o desempenho máximo em todas as diferentes dimensões ocultas e dispositivos nos quais são executados. Além disso, confiamos na especialização em tempo de compilação para acomodar todos os ajustes e nuances de vários modelos (peso com viés versus sem viés para RMS norm, 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 geralmente não são 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 criar os melhores kernels, precisamos de 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 e otimizando-as agressivamente em todos os pontos de um amplo espaço de configuração.

CuTeDSL é uma DSL baseada em Python para GPUs NVIDIA construída sobre a álgebra de layout CuTe e MLIR. Ela compila just-in-time para PTX altamente otimizado, mantendo o controle sobre primitivos de hardware de baixo nível. O objetivo é oferecer os mesmos recursos e desempenho que o C++ do CUTLASS, 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 ela é a melhor opção para o desenvolvimento de nossos kernels.

No geral, embora não seja tão de alto nível e sucinta 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 aos primitivos de hardware, fornecendo-nos o nível de controle de que precisamos para atingir o desempenho máximo em todos os casos extremos. Com 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 construção. Em contraste, com uma DSL compilada Just-In-Time, podemos nos especializar de forma mais agressiva sem sacrificar o ritmo de iteração. Ao contrário do TVM IR, que usava CUDA como uma etapa intermediária, descobrimos que o pipeline baseado em MLIR que reduzia para PTX era mais robusto e mais fácil de depurar do que um 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 poderíamos aproveitar totalmente 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 dos kernels, já que a otimização e a compilação do estágio PTX em diante podem ser concluídas em questão de segundos. Com muitos parâmetros de 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 para compilação paralela surgem apenas mais tarde, nos níveis IR e PTX, com o tratamento de modelos C++ ocorrendo sequencialmente. Embora existam truques para resolver isso, movendo as instanciações de modelo para unidades de compilação separadas, eles exigem ajustes significativos nos sistemas de construção e o uso de macros complexas. Em contraste, os kernels CuTeDSL que são compilados Just-In-Time (JIT) não exigem regras de construção complexas e não dificultam os tempos de compilação, já que o ciclo de desenvolvimento pode ser focado em um único ponto no espaço de configuração, enquanto testes unitários cuidadosos podem cobrir todo o espaço de projeto do kernel posteriormente.

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 na escala de segundos durante a inicialização/aquecimento do mecanismo de inferência. Em contraste, pré-compilar todas as configurações de um kernel poderia levar vários minutos. Além disso, as sobrecargas podem ser ainda mais mitigadas pelo armazenamento em cache dos artefatos compilados via JIT.

O custo de compilação pesado do CUDA é acentuado pela dificuldade de rastrear a origem dos erros quando o despacho e a especialização baseados em modelo são usados, como é o caso do CuTe e 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 declaração da DSL que o originou por meio de estágios MLIR intermediários ou confiando em informações de depuração anexadas ao PTX gerado. Graças a esses aspectos da DSL, descobrimos que somos 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. 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 de 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 desempenho quanto para depuração. Particularmente em um processador em pipeline, como é o caso das GPUs, uma ramificação inexistente é melhor do que uma ramificação sempre tomada ou nunca tomada. 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 cuda-gdb durante uma investigação.

Para um determinado tamanho de lote e configuração, o tráfego esperado para DRAM e o número de operações aritméticas podem ser estimados. Subsequentemente, decisões ideais podem ser tomadas caso a caso para alocação ideal de grade, 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 de 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. Throughput a ser maximizado tanto para embeddings quanto para os tokens de entrada de LLMs. Latência na casa das dezenas a centenas de milissegundos.

Decode: tamanhos de lote de pequenos a grandes, um (sem decodificação especulativa) a alguns tokens por sequência (especulação de alvo de rascunho, previsão de vários 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 pelo throughput 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 streaming e núcleos tensores, limitados pela largura de banda de memória do dispositivo em vez disso. Embora 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 da 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 um deles.

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 do warp de 32 tópicos por warp). É desejável tornar esse parâmetro uma constante de tempo de compilação, pois pode ajudar na alocação de registros. Em 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:

cpp

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

}

O número de registros 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 registros pode atribuir a cada tópico e pode ter sucesso ao compilar kernels otimizados de forma mais agressiva. No entanto, isso tem um custo de tempo de compilação, já que agora o kernel precisa ser recompilado para diferentes tamanhos de warp. Com CuTeDSL, isso não é mais um problema, pois não precisamos pagar o custo de compilação para 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__ em CUDA.

Além de tornar o número de warps uma constante, é desejável variá-lo entre prefill e decode. Para utilizar totalmente a largura de banda de memória de uma GPU, não é suficiente atribuir sempre um grande número de warps por bloco, pois a largura de banda total só pode ser alcançada usando o maior número possível dos Multiprocessadores de Streaming (SMs) disponíveis. Consequentemente, para um determinado número de tokens e tipos de dados e dimensões ocultas conhecidos, primeiro escolhemos uma dimensão de vetor ideal por tópico (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). Sabendo o tamanho do vetor, podemos calcular o número de warps exigidos pelo kernel. Tentamos primeiro maximizar o número de SMs: se houver menos warps do que SMs, escolhe-se um único warp por bloco. Esta é normalmente a variante decode do kernel. Se houver warps suficientes, aumentamos gradualmente o número de warps para 8, 16 ou 32, fazendo um benchmarking de um limite superior após o qual normalmente não vemos mais uma melhoria no desempenho.

Embora o Triton realizasse algumas dessas otimizações nos bastidores, descobrimos que podemos obter resultados semelhantes usando cuidadosamente a álgebra de layout CuTe, sem sermos restringidos pelas deficiências de um compilador de alto nível. Particularmente quando reduções horizontais sub-warp estão envolvidas, descobrimos que o CuTeDSL nos ofereceu o controle de granulação fina necessário para gerar código ideal.

Em algumas situações, para distribuir o trabalho entre o 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 desprogramado de um SM assim que inicia a execução, para garantir o progresso, 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 pesquisar os locais de memória global envolvidos na sincronização. Essa restrição é aplicada por cudaLaunchCooperativeKernel. O custo de uma barreira de grade é de cerca de 3us no Hopper, mais o custo de liberar todas as operações de memória anteriores por meio 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, 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, confiando na barreira implícita dos lançamentos de kernel. Este é principalmente o caso para 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 com os kernels decode 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 essa 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 throughput máximo. Ao contrário das soluções baseadas em CUDA, essa lógica pode ser definida de forma mais sucinta e com menos sobrecarga de 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 de nossos kernels MoE all-to-all, conforme detalhado abaixo, ter a flexibilidade de construir kernels com ou sem barreiras de grade e otimizá-los no tempo de compilação nos ajudou a obter ganhos substanciais de desempenho por meio do CuTeDSL.

O primeiro kernel que construímos em 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 que excedem milhares, exigindo uma redução horizontal através de memória compartilhada e registros, a norma QK precisa de redução apenas entre warps, com a possibilidade de espremer múltiplas cabeças em um único warp. Este tornou-se um caso de uso onde poderíamos otimizar agressivamente no 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 calcular a norma ao quadrado 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 calcular a norma ao quadrado. Para maximizar a largura de banda da memória, cada tópico deve processar 128 bits de entrada, o que se traduz em 8 valores bfloat16. Consequentemente, meio warp é suficiente para processar uma cabeça, e a norma ao quadrado pode ser determinada usando uma redução horizontal realizada por shuffle_sync_bfly. Começando 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, confiando em cargas/armazenamentos automaticamente vetorizados e usando operações específicas de arquitetura de nível inferior para a redução de warp. O kernel determina em qual cabeça (Q ou K) trabalhar e identifica os pesos correspondentes a aplicar. 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 ao quadrado é calculada. Finalmente, a norma RMS, automaticamente vetorizada, é determinada e gravada de volta na memória. Parâmetros específicos do modelo, como WEIGHT_BIAS e EPS, são dobrados como constantes.

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 oferece throughput e latência de 2 a 3 vezes melhores para prefill e decode.

O despacho e combinação de MoE embaralham os tokens de entrada para os especialistas para os quais são roteados, realizando um acúmulo ponderado após os kernels MoE GEMM realizarem as multiplicações de matrizes relevantes. O desempenho é de particular interesse, já que 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 por tensor, de nó único. A entrada é replicada, o despacho seleciona os tokens para a classificação atual e a combinação os acumula em todas as classificações. Acúmulo via NVLink ou NVLS. Usado para pré-preenchedores grandes.

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

DP ≠ 1, InfiniBand: Kernel de paralelismo de dados e paralelismo de tensor Peer-to-Peer operando em NVLink e NVIDIA Quantum InfiniBand. Usado para implantações Hopper multi-nó.

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 para serem passados para um kernel MoE GEMM. Usando as localizações dos tokens, o kernel de combinação então realiza um acúmulo em registros, 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 processamos em um único bloco e sinalizamos o restante da grade, com blocos individuais copiando pedaços de tokens. A atribuição de token é 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 depois são agregados em todos os blocos na memória global para determinar a posição de cada token. A grade desse kernel é fixada ao número de SMs para poder usar barreiras de nível de grade. Um kernel subsequente realiza as cópias, lançado em um número muito maior de blocos e warps para utilizar totalmente a largura de banda de memória disponível.

O kernel de combinação é relativamente simples, pois realiza acúmulo vetorizado, confiando nas informações de roteamento e nos offsets calculados pelo despacho. O CuTeDSL provou ser valioso aqui, pois a computação poderia ser expressa 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 por tensor, as informações de roteamento e 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 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. Nessa 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 entre as 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 duas classificações participarem, mudando para NVLS quando 4 ou mais classificações estiverem envolvidas para um aumento de 10-20% no throughput. Graças à DSL, podemos compartilhar a maior parte da lógica de redução e as trocas de buffer de peer entre os kernels MoE e os kernels all-reduce regulares. Esta é uma situação em que o controle próximo sobre o hardware é benéfico, pois operações de DSL sobre multimem.ld_reduce e multimem.st são necessárias para NVLS. A implementação da redução em si segue amplamente o exemplo de all_reduce dos testes NCCL.

Os kernels de Paralelo de Dados dependem do NVLink para agregar tokens para vários dispositivos. Na fase de despacho, antes de trocar tokens, os peers primeiro roteiam localmente para identificar o offset alvo em cada classificação para cada token que estão enviando. Subsequentemente, eles começam a copiar os tokens para buffers de recepção alocados para cada remetente para evitar conflitos. Gravações de 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 MoE GEMM. Os receptores pesquisam o sinalizador escrito pelo peer com um acquire-load para aguardar a conclusão das gravações.

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

Construímos uma biblioteca de wrappers PTX para expor carregamentos e armazenamentos de release-acquire com escopo no nível do sistema e do dispositivo para implementar este kernel. Ao melhorar a granularidade dos sinalizadores de sincronização e alocar mais blocos para enviar e receber tokens via NVLink, alcançamos melhorias de 2x sobre nossos kernels CUDA anteriores.

Nossos kernels Peer-to-Peer alcançaram anteriormente um 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 obter mais ganhos 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 via JIT, especializamo-nos para mais parâmetros e ajustamos as alocações de grade, introduzindo mais diferenças entre prefill e decode para cerca de 3% de aumento na latência.

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

Uma GPU Blackwell tem até 160 SMs, em contraste com os 132 do Hopper. Para a maioria dos kernels, esta é uma constante de tempo de compilação e a lógica no código do lançador do lado do host a considera para determinar o kernel (prefill vs decode) ideal a ser executado. Além disso, 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.

Blackwell introduz suporte de hardware para um número significativo de esquemas de quantização (MXFP8, MXFP4, NVFP4). O suporte para eles foi implementado principalmente ajustando os kernels existentes para escalonamento por bloco FP8. O suporte JIT é crucial aqui, já que novos modos podem ser suportados sem explodir o espaço de configuração de tempo de compilação. No passado, achamos o Triton estranho ao expressar reduções horizontais em poucos elementos, mais da metade ou um quarto de um warp. Um controle mais refinado sobre as reduções foi crucial para atingir o desempenho ideal ao calcular 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 TMA e instruções de multiplicação de matrizes específicas da arquitetura. Chegar a um ponto em que os LLMs pudessem ajudar na escrita do kernel foi um processo gradual.

Construímos manualmente um corpus inicial de kernels e funções utilitárias, seguindo padrões consistentes, antes que pudéssemos confiar em agentes para a conversão de kernel. À medida que o corpus progredia, podíamos realizar traduções aproximadas de CUDA para CuTeDSL, com um trabalho manual significativo seguindo. 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 pouco entendimento da natureza baseada em grade dos kernels, confundindo a sincronização de tópico e grade. Mais frustrantemente, havia uma tentação de alucinar traduções inexistentes de bibliotecas específicas de CUDA em vez de construir contra 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, descobrimos que o CuTeDSL oferece um desempenho de tempo de compilação muito melhor, permitindo-nos especializar para mais parâmetros e obter melhorias de desempenho. Embora os kernels sejam um pouco mais detalhados devido à falta de certas abstrações em comparação com o Triton, descobrimos que a complexidade é justificada, pois um controle mais refinado sobre o hardware nos permite construir kernels mais rápidos.

No momento da redação, 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 o throughput em nossas cargas de trabalho, de embeddings a 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 CuTeDSL?

Experiência do desenvolvedor

Grande espaço de configuração

Pré-preenchimento e especialização de decodificação

Sincronização de grade

Kernels CuTeDSL

Norma QK

Despacho/Combinação de Mistura de Especialistas (MoE)

Local

Nó único

Paralelo de Dados

Peer-to-Peer

Suporte a Blackwell

Geração de Kernel Assistida por IA

Conclusão e Trabalho Futuro

Referências

Citação