CuTeDSL bei Perplexity

Bei Perplexity nutzen wir die interne Runtime-Optimized Serving Engine (ROSE), um Modelle von Embeddings bis hin zu Billionen-Parameter-LLMs bereitzustellen. Unsere Runtime bildet die Basis aller APIs, einschließlich Sonar, Search und Embeddings. Sie wurde entwickelt, um Forschenden die Exploration der am besten geeigneten Modelle und Architekturen zu ermöglichen und gleichzeitig die besten Kandidaten schnell in der Produktion bereitzustellen. Um dies zu erreichen, müssen wir in der Lage sein, die Inferenz für modernste Architekturen schnell zu implementieren, bereitzustellen und zu optimieren. In diesem Blogbeitrag untersuchen wir, wie CuTeDSL-Kernels, die in unsere Inferenz-Engine integriert sind, es uns ermöglichen, die erforderlichen GPU-Kernels zu erstellen, um Modelle auf NVIDIA-GPUs auf Spitzenleistung zu bringen.

Benutzerdefinierte Modelle, die von verschiedenen Teams bei Perplexity erstellt wurden, werden intern auf NVIDIA Hopper- und Blackwell-GPUs unter Verwendung der ROSE-Inferenz-Engine gehostet, die vom AI-Inference-Team entwickelt wurde. ROSE wurde ursprünglich entwickelt, um angepasste Llama-Modelle für die Sprachdekodierung und Klassifizierung besser bedienen zu können, mit einer Schnittstelle, die mit dem NVIDIA Triton Inference Server kompatibel ist. Im Laufe der Zeit wuchs sie zu einer vollwertigen Inferenz-Engine heran, die umfangreiche LLMs sowie Ranking-, Klassifizierungs-, Scoring- und andere transformatorbasierte Modelle für eine Vielzahl von Teams zur Unterstützung der Produkte von Perplexity bereitstellt. Das Designziel von ROSE ist Einfachheit: Mit der Zeit wird Aufwand investiert, um Spitzenleistung zu erreichen, aber die Hauptpriorität liegt darauf, sich an die sich ständig weiterentwickelnden Architekturen und Anwendungsfälle für unsere Modelle anzupassen.

Der Kernel-Konfigurationsraum über Datentypen, Kopfdimensionen, Compile-Time-Konstanten und Start-Einstellungen hinweg.

Auf einer hohen Ebene besteht ROSE aus einer Engine, die ein Modell an eine Schnittstelle anpasst, mit der Clients interagieren können. Die Engines sind in hohem Maße wiederverwendbar, handhaben die Anforderungsplanung und Stapelverarbeitung (Batching) und rufen dekodierte Tokens, Embeddings oder Scores von den zugrunde liegenden Modellen ab. Die Initialisierung von Geräten, Interprozess- und Internode-Kommunikationskanälen sowie das Laden von Gewichten werden alle von ihnen abgewickelt. Für LLMs ist die Engine für das Batching, Chunking und Sampling verantwortlich. Sie ist außerdem für die KV-Speicherallokation zuständig, sowohl für vollständige als auch für lineare Aufmerksamkeit, sowie für das Prefix-Matching. Embedding-Engines sind einfacher, da sie nur Online-Batching durchführen. Die Ausführung wird so optimiert, dass die Engine die nächste Arbeitsrunde vorbereitet, während das Modell parallel auf dem Beschleuniger läuft.

ROSE stellt eine umfangreiche Sammlung benutzerdefinierter Layer für die Modellierung bereit und kapselt GPU-Kernels, die mit CuTeDSL, Triton, CUDA, CUTLASS und cuBLAS erstellt wurden. Für Matrixmultiplikations- und Aufmerksamkeitsimplementierungen verlassen wir uns primär auf NVIDIA-Kernels von CUTLASS und cuBLAS. Der Rest der Pipeline, einschließlich Embedding-, Norm-, MoE- und Aktivierungskernels, ist Teil der Kernellibrary von ROSE. Um die Hardware, auf der wir die Inferenz ausführen, optimal zu unterstützen, ist ROSE mit eigenen Inter-Device-Kommunikationsprimitiven für MoE-Routing, Dispatch und Combine ausgestattet. Während die anderen Kernels relativ einfach sind, müssen sie gründlich spezialisiert werden, um auf all den verschiedenen verborgenen Dimensionen und Geräten, auf denen sie ausgeführt werden, Spitzenleistung zu erzielen. Zusätzlich setzen wir auf Compile-Time-Spezialisierung, um alle Knöpfe und Feinheiten verschiedener Modelle (Gewicht-Bias vs. kein Gewicht-Bias für RMS-Norm, Aktivierungsfunktionen usw.) zu berücksichtigen. Folglich haben wir uns entschieden, DSLs zu deren Implementierung zu verwenden: Wir haben ROSE ursprünglich mit Triton aufgebaut. Wir haben jedoch seitdem CuTeDSL als primäre Umgebung für die GPU-Programmierung übernommen.

Aus mathematischer Sicht sind Inferenz-Kernels in der Regel unkompliziert. Wenn man einen Kernel entwickelt, liegt die Komplexität daher im Allgemeinen nicht darin auszudrücken, welche Berechnung durchgeführt werden soll, sondern darin, wie genau die Hardware sie ausführen soll. Um die besten Kernels zu erstellen, benötigen wir eine Sprache, die die richtigen Abstraktionen für die zugrunde liegende Hardware bietet. Die richtige Sprache ermöglicht es uns, die High-Level-Operationen für die Berechnung zu spezifizieren und sie aggressiv über alle Punkte in einem weiten Konfigurationsraum zu spezialisieren und zu optimieren.

CuTeDSL ist eine Python-basierte DSL für NVIDIA-GPUs, die auf CuTe-Layout-Algebra und MLIR aufbaut. Sie kompiliert Just-in-Time zu hochoptimiertem PTX, während sie die Kontrolle über Low-Level-Hardwareprimitive behält. Sie zielt darauf ab, dieselben Funktionen und Leistungen wie CUTLASS C++ zu bieten, ohne die teuren Kompilierungszeiten, die die Entwicklungsgeschwindigkeit behindern. Von den Sprachen, die wir untersucht oder mit denen wir gearbeitet haben (einschließlich CUDA, cuTile, TVM IR und Triton), halten wir sie für die beste Wahl für die Entwicklung unserer Kernels.

Insgesamt fanden wir, obwohl sie nicht so abstrakt und prägnant wie Triton oder cuTile ist, dass die CuTe-Layout-Algebra und ihre Zusammensetzungen (bereitgestellt über cute.Tensor) angemessene Ersatzmöglichkeiten darstellen. Zusätzlich bietet CuTeDSL uneingeschränkten Zugriff auf Hardwareprimitive und gibt uns das Maß an Kontrolle, das wir benötigen, um in allen Grenzfällen Spitzenleistung zu erzielen. Mit CUDA waren die Möglichkeiten zur Compile-Time-Spezialisierung begrenzt, da jeder zusätzliche Parameter unangemessene Kosten bei den Kompilierungszeiten oder Komplexität für das Build-System verursachte. Im Gegensatz dazu können wir mit einer Just-In-Time-kompilierten DSL aggressiver spezialisieren, ohne das Iterationstempo zu opfern. Im Gegensatz zu TVM IR, das CUDA als Zwischenschritt verwendete, fanden wir die MLIR-basierte Pipeline, die zu PTX herunterstufte, robuster und einfacher zu debuggen als eine, die CUDA als Zwischenschritt ausgab. Darüber hinaus stellten wir fest, dass wir durch die Verwendung benutzerdefinierter Ops zum Kapseln von PTX-Anweisungen alle Funktionen der zugrunde liegenden Hardware vollständig nutzen konnten.

Der Großteil der Kompilierungskosten von CUDA-Kernels entfällt in der Regel auf die Template-Expansionsphase der Kernels, da die Optimierung und Kompilierung ab der PTX-Stufe in Sekundenschnelle abgeschlossen werden kann. Bei vielen Compile-Time-Parametern kann diese Kompilierungsstufe bei einem nicht-trivialen Kernel auf mehrere Minuten anwachsen. Dieser Engpass wird dadurch akzentuiert, dass Möglichkeiten zur parallelen Kompilierung erst später, auf der IR- und PTX-Ebene, entstehen, während die C++-Template-Handhabung sequentiell erfolgt. Obwohl es Tricks gibt, um dies durch Verschieben von Template-Instanziierungen in separate Kompilierungseinheiten zu adressieren, erfordern sie erhebliche Anpassungen an Build-Systemen und die Verwendung umständlicher Makros. Im Gegensatz dazu erfordern CuTeDSL-Kernels, die Just-In-Time (JIT) kompiliert werden, keine komplexen Build-Regeln und behindern nicht die Kompilierungszeiten, da sich der Entwicklungszyklus auf einen einzelnen Punkt im Konfigurationsraum konzentrieren kann, während sorgfältige Unit-Tests später den gesamten Designraum des Kernels abdecken können.

Auch wenn die JIT-Kompilierung Kosten von der Entwicklungszeit auf die Laufzeit verlagert, ist der Overhead mehr als akzeptabel. Wenn ein Modell bereitgestellt wird, aktiviert es nur wenige Punkte aus dem gesamten Konfigurationsraum, was während des Startens/Aufwärmens der Inferenz-Engine Overheads im Sekundenbereich verursacht. Im Gegensatz dazu könnte das Vorkompilieren aller Konfigurationen eines Kernels mehrere Minuten dauern. Zudem können Overheads durch das Zwischenspeichern der JIT-kompilierten Artefakte weiter gemindert werden.

Die hohen Kompilierungskosten von CUDA werden durch die Schwierigkeit akzentuiert, die Herkunft von Fehlern nachzuverfolgen, wenn template-basiertes Dispatching und Spezialisierung verwendet werden, wie es bei CuTe und CUTLASS der Fall ist. Bei einer DSL können Fehler auf die Codezeile zurückgeführt werden, die sie generiert hat, was das Debuggen erheblich erleichtert. Außerdem ist es viel einfacher, den generierten Code über MLIR-Zwischenstufen der DSL-Anweisung zuzuordnen, die ihn verursacht hat, oder sich auf Debug-Informationen zu verlassen, die an das generierte PTX angehängt sind. Dank dieser Aspekte der DSL fanden wir uns bei der Entwicklung von Kernels wesentlich produktiver.

Obwohl Kernels zwischen Modellen viele Ähnlichkeiten aufweisen, müssen sie dennoch mit einer Reihe von Konfigurationsknöpfen ausgestattet sein, die sich auf die Leistung auswirken können. Relevante Parameter sind unter anderem:

Verborgene Dimension/Head-Dimension/Anzahl der Attention-Heads

Eingabe-/Ausgabedatentypen

Quantisierungsschema

Compile-Time-Konstanten (Biases, Aktivierungsfunktionen usw.)

ROSE-Architektur, einschließlich ihrer Engine, des Modells, der Kernelbibliothek und der Dienste.

Die Compile-Time-Spezialisierung ist sowohl für die Leistung als auch für das Debuggen nützlich. Insbesondere in einem Pipelined-Prozessor, wie bei GPUs, ist ein nicht existierender Branch besser als ein Branch, der immer oder nie genommen wird. Darüber hinaus führt das Kompilieren eines Kernels für einen spezifischen Satz an Konfigurationsknöpfen zu weniger PTX/SASS-Anweisungen, die einfacher zu analysieren und zu debuggen sind, wenn cuda-gdb während einer Untersuchung verwendet wird.

Für eine gegebene Batch-Größe und Konfiguration können der erwartete Traffic zum DRAM und die Anzahl der arithmetischen Operationen geschätzt werden. Anschließend können von Fall zu Fall optimale Entscheidungen für die optimale Grid-Allokation, Synchronisation und die Verwendung bestimmter Hardwarefunktionen getroffen werden. CuTeDSL ermöglicht es uns, Komponenten für eine tiefe Spezialisierung zu definieren und wiederzuverwenden, ohne einen inakzeptablen Kompilierungs-Overhead zu verursachen.

Im Allgemeinen muss die Leistung von Inferenz-Kernels für zwei separate Anwendungsfälle optimiert werden:

Prefill: kleine Batch-Größe (wenige Chunks, Sequenzen), viele Tokens pro Sequenz. Der Durchsatz soll sowohl für Embeddings als auch für die Eingabe-Tokens von LLMs maximiert werden. Latenz im Bereich von zehn bis hunderten von Millisekunden.

Decode: Batch-Größen von klein bis groß, ein (kein spekulatives Dekodieren) bis wenige Tokens pro Sequenz (Draft-Target-Spekulation, Multi-Token-Vorhersage). Latenz nahe an der Kernel-Startlatenz.

Für GEMM und Attention läuft die Prefill-Phase auf einer großen Anzahl von Tokens, und die Leistung ist durch den Rechen-Durchsatz des Geräts begrenzt. Das Dekodieren beinhaltet dieselbe Operation bei kleinen Batch-Größen, wobei die Ausführungszeiten von den Kosten für das Laden der Aktivierungen und Gewichte vom DRAM in die Streaming-Multiprozessoren und Tensor-Cores dominiert werden, was stattdessen durch die Speicherbandbreite des Geräts eingeschränkt wird. Obwohl die Batch-Größe oder die Gesamtzahl der Tokens in einem Batch ein relativ klarer Indikator ist, ist die Umschalt-Schwelle zwischen Prefill und Decode subtiler, da die Speicherbandbreite auch durch die verborgene Dimension des Modells und den Datentyp der Aktivierung bestimmt wird. Kernels werden daher in Prefill- und Decode-Implementierungen aufgeteilt, wenn algorithmische Unterschiede oder Grid-/Block-Zuweisungen einen Vorteil für einen von beiden bieten würden.

Der Parameter, der sich zwischen Prefill- und Decode-Kernels universell unterscheidet, ist die Anzahl der Warps in einem Block (die Anzahl der Threads, geteilt durch die Warp-Größe von 32 Threads pro Warp). Es ist wünschenswert, diesen Parameter zu einer Compile-Time-Konstante zu machen, da er bei der Registerallokation helfen kann. In CUDA geschieht dies typischerweise durch die Angabe der Anzahl der Threads über __launch_bounds__ und die Bereitstellung der Anzahl der Warps als Template-Parameter:

cpp

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

}

Die Anzahl der für alle Threads in einem Block verfügbaren Register ist eine Hardware-Beschränkung, 64K auf Hopper und Blackwell. Wenn die Anzahl der Threads im Voraus bekannt ist, weiß der Compiler genau, wie viele Register er jedem Thread zuweisen kann und kann erfolgreich aggressiver optimierte Kernels kompilieren. Dies geht jedoch mit Kompilierungskosten einher, da der Kernel nun für verschiedene Warp-Größen neu kompiliert werden muss. Mit CuTeDSL ist dies kein Problem mehr, da wir die Kosten für die Kompilierung für den gesamten Konfigurationsraum nicht im Voraus zahlen müssen. Die DSL erwartet auch, dass die Block-Größe ein konstanter Ausdruck ist, wodurch implizit die durch __latch_bounds__ in CUDA implizierte Optimierung durchgeführt wird.

Neben der Tatsache, dass die Anzahl der Warps eine Konstante ist, ist es wünschenswert, sie zwischen Prefill und Decode zu variieren. Um die Speicherbandbreite einer GPU vollständig zu nutzen, reicht es nicht aus, immer eine hohe Anzahl von Warps pro Block zuzuweisen, da die volle Bandbreite nur durch die Nutzung so vieler verfügbarer Streaming-Multiprozessoren (SMs) wie möglich erreicht werden kann. Daher wählen wir für eine gegebene Anzahl von Tokens und bekannte Datentypen und verborgene Dimensionen zunächst eine ideale Vektor-Dimension pro Thread (8 für bfloat16, 4 für float32), um sicherzustellen, dass der generierte Code die 128-Bit-Lade-/Speicheranweisungen (256-Bit auf Blackwell) verwendet. Da wir die Vektorgröße kennen, können wir die Anzahl der für den Kernel benötigten Warps berechnen. Wir versuchen zunächst, die Anzahl der SMs zu maximieren: Wenn es weniger Warps als SMs gibt, wird ein einzelner Warp pro Block gewählt. Dies ist typischerweise die Decode-Variante des Kernels. Wenn genügend Warps vorhanden sind, erhöhen wir die Anzahl der Warps schrittweise auf 8, 16 oder 32 und bewerten eine Obergrenze, nach der wir normalerweise keine Leistungssteigerung mehr sehen.

Während Triton einige dieser Optimierungen unter der Haube durchführte, stellten wir fest, dass wir ähnliche Ergebnisse durch den sorgfältigen Einsatz der CuTe-Layout-Algebra erzielen können, ohne durch die Mängel eines High-Level-Compilers eingeschränkt zu sein. Insbesondere wenn horizontale Sub-Warp-Reduktionen involviert sind, bot uns CuTeDSL die fein abgestimmte Kontrolle, die erforderlich ist, um optimalen Code zu generieren.

In einigen Situationen, um die Arbeit auf so viele Blöcke wie möglich zu verteilen, kann cooperative_groups::this_grid().sync() verwendet werden, um alle beteiligten Blöcke zu synchronisieren. Es gelten Einschränkungen: Da ein Block nicht von einem SM de-scheduled werden kann, sobald er mit der Ausführung beginnt, darf die Anzahl der Blöcke die Anzahl der auf dem Gerät verfügbaren SMs nicht überschreiten, um einen Fortschritt zu gewährleisten und sicherzustellen, dass zu jedem Zeitpunkt alle Blöcke die beteiligten globalen Speicherorte aktualisieren und abfragen können, die in der Synchronisation involviert sind. Diese Einschränkung wird durch cudaLaunchCooperativeKernel erzwungen. Die Kosten einer Grid-Barriere betragen auf Hopper etwa 3us, zuzüglich der Kosten für das Flushen aller vorherigen Speicheroperationen über eine implizite Speicherbarriere in einigen Situationen. Insgesamt ist dies mit dem Overhead vergleichbar, einen Kernel in zwei Starts aufzuteilen, was zu einer leistungssensitiven Designentscheidung führt. Trotz der Tatsache, dass Programmatic Dependent Launch (PDL) ab Hopper die Start-Overheads reduzieren kann, ist es für Decode-Kernels lohnenswert, den Overhead zu vermeiden.

Wenn eine der Workloads mehr Blöcke als die Anzahl der SMs erfordert, sollte der Kernel aufgeteilt werden, wobei auf die implizite Barriere von Kernel-Starts zurückgegriffen wird. Dies ist primär für Prefill der Fall. Wenn die Anzahl der Blöcke, die für alle Workloads ideal ist, in das Grid passt, wird eine Grid-Barriere bevorzugt, wie es bei Decode-Kernels der Fall ist, die mit einer kleineren Anzahl von Tokens arbeiten.

In CuTeDSL können diese Entscheidungen im Host-seitigen Launch-Code getroffen werden, und die Kernels können nach Bedarf mit Grid-Barrieren zusammengesetzt und verschachtelt werden, wie im folgenden Beispiel veranschaulicht. Mit einer solchen hybriden Strategie können Decode-Kernels für minimale Latenzen und Prefill-Kernels für maximalen Durchsatz optimiert werden. Im Gegensatz zu CUDA-basierten Lösungen kann diese Logik prägnanter und mit weniger Kompilierungs-Overhead definiert werden.

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

Im Falle unserer MoE-All-to-All-Kernels, wie unten detailliert beschrieben, half uns die Flexibilität, Kernels mit oder ohne Grid-Barrieren zu erstellen und sie zur Kompilierungszeit zu optimieren, dabei, durch CuTeDSL erhebliche Leistungssteigerungen zu erzielen.

Der erste Kernel, den wir in CuTeDSL gebaut haben, war ein spezialisierter RMS-Norm-Kernel, der für kleine Dimensionen optimiert wurde. Die QK-Norm, die in der Self-Attention-Schicht der aktuellen Qwen-Modelle verwendet wird, wird pro Head auf 64, 128 oder 256 Elemente angewendet. Im Gegensatz zur regulären RMS-Norm, die auf verborgenen Dimensionen operiert, die in die Tausende gehen, und eine horizontale Reduktion über Shared Memory und Register erfordert, benötigt die QK-Norm die Reduktion nur über Warps hinweg, mit der Möglichkeit, mehrere Heads in einen einzigen Warp zu quetschen. Dies wurde zu einem Anwendungsfall, bei dem wir aggressiv zur Kompilierungszeit optimieren konnten, um die besten Token-zu-Warp-Zuweisungen für alle Head-Dimensionen und Datentypen zu finden.

Die RMS-Norm ist eine ziemlich unkomplizierte Operation, obwohl sie eine horizontale Summe erfordert, um die quadratische Norm eines Vektors zu berechnen:

cpp

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

Da Attention-Heads typischerweise klein sind und Operationen pro Thread stark vektorisiert sind, kann ein einzelner Warp die quadratische Norm berechnen. Um die Speicherbandbreite zu maximieren, sollte jeder Thread 128 Bits an Eingabe verarbeiten, was 8 bfloat16-Werten entspricht. Folglich reicht ein Half-Warp aus, um einen Head zu verarbeiten, und die quadratische Norm kann unter Verwendung einer horizontalen Reduktion durch shuffle_sync_bfly bestimmt werden. Ausgehend von der Head-Dimension und dem Datentyp können wir die Anzahl der Heads pro Warp bestimmen und die Argumente sowie die Anzahl der Reduktionsrunden zum Zeitpunkt der Kompilierung festlegen.

Wir können dann einen spezialisierten Kernel mit relativ wenigen Codezeilen erstellen, der sich auf automatisch vektorisierte Lade-/Speichervorgänge verlässt und Low-Level-architekturspezifische Operationen für die Warp-Reduktion verwendet. Der Kernel bestimmt, an welchem Head (Q oder K) gearbeitet werden soll, und identifiziert die entsprechenden Gewichte, die angewendet werden sollen. Er lädt einen gesamten Vektor aus dem Head und summiert zunächst lokal auf. Anschließend wird basierend auf der Anzahl der Threads pro Head die quadratische Norm berechnet. Schließlich wird die RMS-Norm, automatisch vektorisiert, bestimmt und zurück in den Speicher geschrieben. Modellspezifische Parameter wie WEIGHT_BIAS und EPS werden constant-gefaltet.

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)

Im Vergleich zu einem regulären RMS-Norm-Kernel liefert diese spezialisierte Variante einen 2-3-fach besseren Durchsatz und eine geringere Latenz für Prefill und Decode.

MoE-Dispatch und -Combine mischen Eingabe-Tokens zu den Experten, zu denen sie geroutet werden, und führen eine gewichtete Akkumulation durch, nachdem MoE-GEMM-Kernels die relevanten Matrixmultiplikationen durchgeführt haben. Leistung ist von besonderem Interesse, da verschiedene Produkte auf MoE-Modellen unterschiedlicher Größe basieren. Basierend auf den Bereitstellungstypen unterscheiden wir verschiedene Kernels zur Optimierung:

TP=1: Lokaler Kernel für eine einzelne GPU. Dispatch mischt Tokens zu den relevanten Experten, und Combine akkumuliert sie. Wird primär für Qwen3 30B-Bereitstellungen auf H100 verwendet.

DP=1, TP ≠ 1: Tensor-parallel, Singlenode-Kernel. Die Eingabe wird repliziert, Dispatch wählt die Tokens für den aktuellen Rang aus, und Combine akkumuliert sie über alle Ränge hinweg. Akkumulation über NVLink oder NVLS. Wird für große Prefiller verwendet.

DP ≠ 1, NVLink: Data-parallel Kernel mit optionaler Tensor-Parallelität. Jeder Rang sendet Tokens an jeden anderen Rang, der spezifische Experten hostet. Routing-Informationen werden aggregiert, Tokens werden über NVIDIA NVLink übertragen. Wird für Data-Parallel H200 oder GB200-Bereitstellungen verwendet.

DP ≠ 1, InfiniBand: Peer-to-Peer Data-Parallel- und Tensor-Parallel-Kernel, der über sowohl NVLink als auch NVIDIA Quantum InfiniBand operiert. Wird für Multi-Node-Hopper-Bereitstellungen verwendet.

Der lokale Kernel zählt die Anzahl der Tokens, die jedem Experten basierend auf der Routing-Tabelle zugewiesen sind, und weist jedem Token innerhalb der Liste der zu jedem Experten gehörenden Tokens einen eindeutigen Index zu. Nach dem Routing-Schritt werden die Tokens in optional gepolsterte Buffer kopiert, um an einen MoE-GEMM-Kernel weitergegeben zu werden. Unter Verwendung der Token-Speicherorte führt der Combine-Kernel dann eine Akkumulation in Registern durch und schreibt die finalen Tokens in einen Ausgabe-Buffer.

Die Dispatch-Phase variiert zwischen Prefill und Decode: Bei wenigen Tokens gibt es wenig Routing-Informationen auszulesen, daher verarbeiten wir sie in einem einzigen Block und signalisieren dem Rest des Grids, wobei einzelne Blöcke Chunks von Tokens kopieren. Die Token-Zuweisung wird durch atomares Inkrementieren von Zählern im Shared Memory bestimmt. Da Prefill mehr Routing-Informationen verarbeiten muss, wird das Routing über mehrere Blöcke in einem separaten Kernel aufgeteilt. Indizes werden zunächst lokal unter Verwendung von Zählern im Shared Memory zugewiesen und dann über alle Blöcke im globalen Speicher aggregiert, um die Position jedes Tokens zu bestimmen. Das Grid dieses Kernels ist auf die Anzahl der SMs begrenzt, um Grid-weise Barrieren verwenden zu können. Ein nachfolgender Kernel führt die Kopien durch, der über eine viel größere Anzahl von Blöcken und Warps gestartet wird, um die verfügbare Speicherbandbreite vollständig zu nutzen.

Der Combine-Kernel ist relativ einfach, da er eine vektorisierte Akkumulation durchführt und sich auf die Routing-Informationen und die von Dispatch berechneten Offsets verlässt. CuTeDSL erwies sich hier als wertvoll, da die Berechnung durch einfache CuTe-Tensoren ausgedrückt werden konnte, die zu effizientem Code für alle Datentypen und verborgenen Dimensionen kompiliert werden.

Bei einer Tensor-parallelen Bereitstellung werden die Routing-Informationen und Eingabe-Tokens repliziert. Der lokale Dispatch-Kernel kann vollständig wiederverwendet werden, mit einer leichten Änderung, um nur die Tokens zu routen und auszuwählen, die den auf dem aktuellen Rang befindlichen Experten zugewiesen sind. Mit einigen Parametern und Compile-Time-Spezialisierung werden die lokalen Dispatch-Kernels wiederverwendet.

Da die Anzahl der Ränge normalerweise niedriger ist als die Anzahl der Experten, zu denen ein Token geroutet wird, kann im Durchschnitt angenommen werden, dass jedes Token mindestens eines seiner Replikate zu jedem Rang geroutet hat. In dieser Situation kann der lokale Combine-Kernel wiederverwendet werden, um die lokalen Kopien eines Tokens zu aggregieren, gefolgt von einem All-Reduce, um über die teilnehmenden Ränge hinweg zu summieren. Im Fall des Prefill-Kernels starten wir separate Kernels, jedoch fusioniert Decode beide Operationen in einen einzigen Start.

Für All-Reduce verwenden wir normales NVLink, wenn nur zwei Ränge teilnehmen, und wechseln zu NVLS, wenn 4 oder mehr Ränge beteiligt sind, was zu einer 10-20%igen Durchsatzsteigerung führt. Dank der DSL können wir den Großteil der Reduktionslogik und den Peer-Buffer-Austausch zwischen den MoE- und den regulären All-Reduce-Kernels teilen. Dies ist eine Situation, in der eine enge Kontrolle über die Hardware vorteilhaft ist, da für NVLS DSL-Operationen über multimem.ld_reduce und multimem.st erforderlich sind. Die Implementierung der Reduktion selbst folgt weitgehend dem all_reduce-Beispiel aus NCCL-Tests.

Die Data-Parallel-Kernels verlassen sich auf NVLink, um Tokens für mehrere Geräte zu aggregieren. In der Dispatch-Phase routen Peers vor dem Austausch von Tokens zunächst lokal, um den Ziel-Offset auf jedem Rang für jedes Token zu identifizieren, das sie senden. Anschließend beginnen sie, die Tokens in Empfangs-Buffer zu kopieren, die für jeden Sender allokiert wurden, um Konflikte zu vermeiden. Pro-Token-Release-Store-Writes werden verwendet, um den Abschluss zu signalisieren. Auf der Empfängerseite werden Routing-Informationen aggregiert, um alle Tokens in einem einzigen Buffer für MoE-GEMM kondensieren zu können. Empfänger fragen das vom Peer geschriebene Flag mit einem Acquire-Load ab, um auf den Abschluss der Schreibvorgänge zu warten.

Auf der Combine-Seite werden Tokens über NVLink zu den Quellrängen kopiert, wobei ein Pro-Token-Flag den Empfang anzeigt. Der Combine-Kernel ist den lokalen und Singlenode-Varianten ähnlich, mit leichten Unterschieden bei der Adressierung und etwas zusätzlicher Synchronisationslogik.

Wir haben eine Bibliothek von PTX-Wrappern erstellt, um Release-Acquire-Loads und Stores auf System- und Geräteebene bereitzustellen, um diesen Kernel zu implementieren. Durch die Verbesserung der Granularität von Synchronisationsflags und die Allokation von mehr Blöcken zum Senden und Empfangen von Tokens über NVLink erzielten wir 2-fache Verbesserungen gegenüber unseren vorherigen CUDA-Kernels.

Unsere Peer-to-Peer-Kernels erzielten zuvor Spitzenleistung auf Hopper-GPUs, die mit ConnectX-7 InfiniBand-Adaptern gekoppelt waren. Wir haben seitdem die Combine-Kernels zu CuTeDSL portiert, um weitere Leistungsgewinne zu erzielen. Da die CUDA-Versionen bereits stark auf Compile-Time-Spezialisierung beruhten, erwiesen sie sich als Engpässe bei der Kompilierungszeit in unserem Repository, was uns zwang, einige Parameter wegzulassen. Durch die Migration zur JIT-kompilierten DSL spezialisierten wir auf mehr Parameter und passten Grid-Allokationen an, was weitere Unterschiede zwischen Prefill und Decode für etwa 3% Latenzverbesserung einführte.

Aus der Perspektive von ROSE-Kernels fällt die Blackwell-Unterstützung in zwei Kategorien: architekturagnostische und architekturspezifische Kernels. Architekturspezifische Kernels, die auf Tensor-Cores beruhen, erfordern explizite Arbeit zur Anpassung. Dank der Verwendung der DSL passen sich die architekturagnostischen Kernels jedoch trivial an.

Eine Blackwell-GPU hat bis zu 160 SMs, im Gegensatz zu Hoppers 132. Für die meisten Kernels ist dies eine Compile-Time-Konstante, und die Logik im Host-seitigen Launcher-Code berücksichtigt dies bei der Bestimmung des optimalen (Prefill vs Decode) Kernels, der ausgeführt werden soll. Zusätzlich führt Blackwell auch 256-Bit-Ladeanweisungen ein, die leicht durch Erhöhung der Vektorbreite (letzte Dimension der CuTe-Tensoren) ausgegeben werden können, um die Speicherauslastung zu verbessern.

Blackwell führt Hardwareunterstützung für eine signifikante Anzahl von Quantisierungsschemata (MXFP8, MXFP4, NVFP4) ein. Die Unterstützung hierfür wurde primär durch die Anpassung existierender Kernels für FP8 Block-weise Skalierung implementiert. JIT-Unterstützung ist hier entscheidend, da neue Modi unterstützt werden können, ohne den Compile-Time-Konfigurationsraum explodieren zu lassen. In der Vergangenheit fanden wir Triton umständlich darin, horizontale Reduktionen über wenige Elemente, über die Hälfte oder ein Viertel eines Warps auszudrücken. Feinere Kontrolle über Reduktionen war entscheidend, um eine optimale Leistung beim Berechnen eines absoluten Maximalwerts über 16 bzw. 32 Elemente für NVFP4 und MXFP8/MXFP4 zu erzielen.

Als relativ neue Sprache hatten LLMs während ihres Vortrainings wenig Kontakt mit CuTeDSL-Code. Zusätzlich war der Stil unseres aktuellen Codes, mit Kernels, die für Aktivierungen und Normen optimiert sind, etwas anders als die bestehenden CUTLASS- oder Quack-Beispiele, die stark auf TMA und architekturspezifische Matrixmultiplikationsanweisungen optimiert waren. Einen Punkt zu erreichen, an dem LLMs beim Schreiben von Kernels assistieren konnten, war ein gradueller Prozess.

Wir bauten manuell ein initiales Korpus von Kernels und Dienstfunktionen nach konsistenten Mustern auf, bevor wir uns bei der Kernel-Konvertierung auf Agenten verlassen konnten. Als das Korpus voranschritt, konnten wir grobe Übersetzungen von CUDA nach CuTeDSL durchführen, wobei anschließend eine signifikante manuelle Nacharbeit folgte. Insbesondere stellten wir fest, dass Modelle Schwierigkeiten mit Zeiger-zu-Tensor-Konvertierungen und Vektorisierung hatten und beim Versuch, Typfehler zu beheben, stecken blieben, wenn ein Test kompiliert werden sollte. Sie zeigten wenig Verständnis für die Grid-basierte Natur der Kernels und verwechselten Thread- und Grid-Synchronisation. Am frustrierendsten war die Versuchung, nicht existierende Übersetzungen von CUDA-spezifischen Bibliotheken zu halluzinieren, anstatt gegen unsere eigenen CuTeDSL-Ersatzlösungen zu bauen. Im Laufe der Zeit, als das Korpus der CuTeDSL-Kernels wuchs, stellten wir fest, dass LLMs den Großteil des Skelettcodes erstellen und Lücken füllen können, was es uns ermöglicht, uns auf das High-Level-Design zu konzentrieren.

Im Vergleich zu CUDA stellen wir fest, dass CuTeDSL eine viel bessere Kompilierungsleistung liefert, was es uns ermöglicht, auf mehr Parameter zu spezialisieren und Leistungsverbesserungen zu erzielen. Während Kernels aufgrund des Mangels bestimmter Abstraktionen im Vergleich zu Triton etwas ausführlicher sind, finden wir, dass die Komplexität gerechtfertigt ist, da feinere Kontrolle über die Hardware es uns ermöglicht, schnellere Kernels zu bauen.

Zum Zeitpunkt des Schreibens ist der Leistungsunterschied zwischen CUTLASS C++ und CuTeDSL-Kernels minimal, was spannende Möglichkeiten eröffnet, High-Performance-GEMM-Kernels mit Inter-Device-Kommunikation, Normalisierung oder Aktivierungen zu komponieren und zu fusionieren. Durch das Commitment zu dieser DSL werden wir in der Lage sein, weitere Leistung aus unseren Geräten zu extrahieren, um den Durchsatz über unsere Workloads hinweg zu verbessern, von Embeddings bis hin zu 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}
}

Inferenz bei Perplexity AI

Warum CuTeDSL?

Entwicklererfahrung

Großer Konfigurationsraum

Spezialisierung für Prefill und Decodierung vorab ausfüllen

Grid-Synchronisation

CuTeDSL Kernels

QK-Norm

Mixture-of-Experts Dispatch/Combine

Lokal

Singlenode

Data Parallel

Peer-to-Peer

Blackwell-Unterstützung

KI-gestützte Kernel-Generierung

Fazit und zukünftige Arbeiten

Referenzen

Zitieren