PerplexityにおけるCuTeDSL

Perplexityでは、組み込みから数兆パラメータ規模のLLMまで、幅広いモデルのサービングに社内開発のRuntime-Optimized Serving Engine（ROSE）を使用しています。Sonar、Search、Embeddingsを含むすべてのAPIのバックエンドである当社のランタイムは、研究者が最適なモデルやアーキテクチャを探索できるようにし、同時に最良の候補を迅速に本番環境へデプロイできるように構築されました。これを実現するために、当社は最先端アーキテクチャのための推論実装、デプロイ、最適化を迅速に行う必要があります。本ブログ記事では、当社の推論エンジンに統合されたCuTeDSLカーネルが、いかにしてNVIDIA GPU上でモデルを最高性能まで引き上げるために必要なGPUカーネルを構築することを可能にするかを探ります。

Perplexityの各チームが構築したカスタムモデルは、AI推論チームが開発したROSE推論エンジンを使用し、NVIDIA HopperおよびBlackwell GPU上で社内ホストされています。ROSEは当初、NVIDIA Triton Inference Serverと互換性のあるインターフェースを備え、言語デコードと分類の両方の用途でカスタマイズされたLlamaモデルをより適切にサービングするために構築されました。時を経て、ROSEは本格的な推論エンジンへと成長し、大規模なLLMに加え、ランキング、分類、スコアリング、その他各種トランスフォーマーベースのモデルを幅広いチームに提供し、Perplexityの製品群を支えています。ROSEの設計目標はシンプルです。時間をかけて性能の極限を引き出すための投資は行いますが、最も重視しているのは、モデルの進化し続けるアーキテクチャやユースケースに適応することです。

高い視点で見ると、ROSEは、モデルをクライアントが対話可能なインターフェースに適応させるエンジンで構成されています。このエンジンは再利用性が非常に高く、リクエストのスケジューリングやバッチ処理、基盤モデルからのデコードされたトークン、埋め込み、スコアの取得を処理します。デバイスの初期化、プロセス間およびノード間の通信チャネル、重みのロードもすべてこのエンジンが担います。LLMの場合、エンジンはバッチ処理、チャンキング、サンプリングを担当します。また、フルアテンションおよびリニアアテンション双方のKVストレージ割り当てや、プレフィックスマッチングも行います。埋め込みエンジンはオンラインバッチ処理のみを行うため、よりシンプルです。実行は、モデルがアクセラレータ上で動作している間に、エンジンが並行して次の作業ラウンドを準備するように最適化されています。

ROSEは、モデリングのためのカスタムレイヤーを豊富に備えており、CuTeDSL、Triton、CUDA、CUTLASS、cuBLASで構築されたGPUカーネルをラップしています。行列乗算とアテンションの実装については、主にCUTLASSとcuBLASによるNVIDIAカーネルに依存しています。埋め込み、ノルム、MoE、活性化カーネルを含むパイプラインの残りの部分は、ROSEのカーネルライブラリの一部です。私たちが推論を実行するハードウェアを最適にサポートするために、ROSEはMoEルーティング、ディスパッチ、結合のための独自のデバイス間通信プリミティブを備えています。他のカーネルは比較的単純ですが、実行されるすべての異なる隠れ次元やデバイス上で最高性能を達成するためには、徹底的に専門化（specialized）されなければなりません。さらに、さまざまなモデルの細かな設定（RMSノルムにおける重みバイアスの有無、活性化関数など）に対応するため、コンパイル時の専門化（compile-time specialization）に依存しています。その結果、これらを実装するためにDSLを使用することに決めました。当初はROSEをTritonで構築しましたが、現在ではGPUプログラミングの主要環境としてCuTeDSLを採用しています。

数学的な観点から見ると、推論カーネルは通常、それほど複雑ではありません。したがって、カーネルを開発する際、難しさは「どのような計算を行うか」という表現ではなく、「ハードウェアがそれを正確にどう実行すべきか」にあります。最高のカーネルを構築するためには、基盤となるハードウェアに対して適切な抽象化を提供する言語が必要です。適切な言語があれば、高水準の演算を指定して計算を実行し、広範な構成スペースのすべての点において積極的に専門化と最適化を行うことが可能になります。

CuTeDSLは、CuTeレイアウト代数とMLIRの上に構築された、NVIDIA GPU向けのPythonベースのDSLです。低レイヤーのハードウェアプリミティブを制御しつつ、高度に最適化されたPTXにジャストインタイム（JIT）コンパイルされます。開発速度を損なう長時間のコンパイルを必要とせず、CUTLASS C++と同等の機能と性能を提供することを目指しています。これまでに調査または使用してきた言語（CUDA、cuTile、TVM IR、Tritonなど）の中で、私たちはCuTeDSLがカーネル開発に最も適していると考えています。

全体として、TritonやcuTileほど高レベルで簡潔ではありませんが、CuTeレイアウト代数とその構成（cute.Tensorを介して公開）は妥当な代替手段であると判断しました。さらに、CuTeDSLはハードウェアプリミティブへの無制限のアクセスを提供し、あらゆるエッジケースで最高性能を達成するために必要なレベルの制御を可能にします。CUDAでは、パラメータを追加するたびにコンパイル時間が不当に増大したり、ビルドシステムの複雑さが増したりするため、コンパイル時の専門化の機会は限られていました。対照的に、JITコンパイルされるDSLを使用することで、反復速度を犠牲にすることなく、より積極的に専門化を行うことができます。中間ステップとしてCUDAを使用するTVM IRとは異なり、PTXへ低層化するMLIRベースのパイプラインの方が、CUDAを中間ステップとして出力するものよりも堅牢でデバッグが容易であることがわかりました。さらに、PTX命令をラップするカスタムオペレーションの使用により、基盤となるハードウェアの全機能を完全に活用できることがわかりました。

CUDAカーネルのコンパイルコストの大部分は、通常、カーネルのテンプレート展開フェーズで消費されます。PTX段階以降の最適化とコンパイルは数秒で完了するためです。コンパイル時のパラメータが多い場合、このコンパイル段階は非自明なカーネルに対しては数分にも及ぶことがあります。このボトルネックは、C++テンプレートの処理が順次行われる一方で、並列コンパイルの機会がIRおよびPTXレベルになって初めて発生するという事実によって強調されます。これに対処する手法は存在しますが、テンプレートのインスタンス化を別々のコンパイルユニットに移動させる必要があり、ビルドシステムの大幅な調整や扱いにくいマクロの使用が必要になります。対照的に、JITコンパイルされるCuTeDSLカーネルは、複雑なビルドルールを必要とせず、コンパイル時間を妨げません。開発サイクルを構成スペース内の単一の点に集中させ、後の慎重なユニットテストによってカーネルの設計スペース全体をカバーできるためです。

JITコンパイルはコストを開発時から実行時に移動させますが、そのオーバーヘッドは十分に許容範囲内です。モデルがデプロイされると、構成スペース全体のうちのいくつかの点のみが有効化され、推論エンジンの起動/ウォームアップ時に秒単位のオーバーヘッドが発生するに過ぎません。これに対し、カーネルのすべての構成を事前にコンパイルすると、数分かかる可能性があります。さらに、JITコンパイルされたアーティファクトをキャッシュすることで、オーバーヘッドをさらに軽減できます。

CUDAの重いコンパイルコストは、CuTeやCUTLASSの場合のように、テンプレートベースのディスパッチや専門化が使用される場合に、エラーの発生源を特定することが困難であるという点によって強調されます。DSLを使用すると、エラーを生成したコード行を指摘できるため、デバッグが大幅に容易になります。さらに、中間MLIRステージを経由するか、生成されたPTXに付加されたデバッグ情報に依存することで、生成されたコードを発生元のDSLステートメントにマッピングする方がはるかに簡単です。DSLのこれらの側面のおかげで、私たちはカーネル開発において生産性が大幅に向上したことを実感しています。

カーネルはモデル間で多くの類似点を共有していますが、パフォーマンスに影響を与える可能性のあるさまざまな構成ノブを備えている必要があります。関連するパラメータには以下が含まれます：

隠れ次元/ヘッド次元/アテンションヘッド数

入出力データ型

量子化スキーム

コンパイル時の定数（バイアス、活性化関数など）

ROSEアーキテクチャ。エンジン、モデル、カーネルライブラリ、サービスを含みます。

コンパイル時の専門化は、パフォーマンスとデバッグの両方に役立ちます。特に、GPUのようなパイプライン化されたプロセッサにおいて、常に実行される、あるいは決して実行されない分岐よりも、存在しない分岐の方が優れています。さらに、特定の構成ノブのセットに対してカーネルをコンパイルすると、PTX/SASS命令の数が減り、調査中にcuda-gdbを使用する際の分析とデバッグが容易になります。

特定のバッチサイズと構成に対して、DRAMへの期待トラフィックと算術演算数を推定できます。その後、最適なグリッド割り当て、同期、および特定のハードウェア機能の使用について、ケースバイケースで最適な決定を下すことができます。CuTeDSLを使用すると、容認できないコンパイル時のオーバーヘッドを招くことなく、詳細な専門化のためのコンポーネントを定義および再利用できます。

一般に、推論カーネルのパフォーマンスは、2つの別々のユースケースに対して最適化する必要があります：

プリフィル（Prefill）: 小さなバッチサイズ（少数のチャンク、シーケンス）、シーケンスごとのトークン数が多い。埋め込みとLLMの入力トークンの両方でスループットを最大化する必要がある。レイテンシは数十から数百ミリ秒。

デコード（Decode）: 小から大のバッチサイズ、シーケンスごとに1トークン（推測デコードなし）から少数トークン（ドラフトターゲット推測、マルチトークン予測）。レイテンシはカーネル起動レイテンシに近い。

GEMMとアテンションの場合、プリフィル段階は多数のトークン上で実行され、パフォーマンスはデバイスの計算スループットに制約されます。デコードでは小規模バッチサイズで同じ操作が行われますが、実行時間はDRAMからストリーミングマルチプロセッサ（SM）およびテンソルコアへの活性化値と重みのロードコストに支配され、代わりにデバイスのメモリ帯域幅に制約されます。バッチサイズやバッチ内のトークンの総数は比較的明確な指標ですが、メモリ帯域幅はモデルの隠れ次元や活性化値のデータ型によっても決定されるため、プリフィルとデコードの切り替えしきい値はより微妙です。したがって、アルゴリズムの違いやグリッド/ブロックの割り当てによって利点が得られる場合、カーネルはプリフィルとデコードの実装に分割されます。

プリフィルカーネルとデコードカーネルの間で普遍的に異なるパラメータは、ブロック内のワープ数（スレッド数をワープサイズである32で割った数）です。このパラメータをコンパイル時の定数にすることは、レジスタ割り当てを支援できるため望ましいです。CUDAでは、これは通常、__launch_bounds__を介してスレッド数を指定し、テンプレートパラメータとしてワープ数を提供することで行われます：

cpp

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

}

ブロック内のすべてのスレッドが利用可能なレジスタ数は、HopperとBlackwellで64Kというハードウェア制約があります。スレッド数が事前にわかっていれば、コンパイラは各スレッドに割り当て可能なレジスタ数を正確に把握でき、より積極的に最適化されたカーネルをコンパイルできます。しかし、これにはコンパイル時間のコストが伴います。カーネルを異なるワープサイズに対して再コンパイルする必要があるためです。CuTeDSLを使用すると、構成スペース全体に対して事前にコンパイルコストを支払う必要がないため、これはもはや問題ではありません。DSLはブロックサイズが定数式であることも想定しており、CUDAの__launch_bounds__が意味する最適化を暗黙的に実行します。

ワープ数を定数にすることに加え、プリフィルとデコードの間でそれを変化させることも望ましいです。GPUのメモリ帯域幅を完全に活用するためには、常にブロックあたりのワープ数を高く割り当てるだけでは不十分です。利用可能なすべてのストリーミングマルチプロセッサ（SM）を最大限に使用して初めて、フルの帯域幅を達成できるためです。したがって、特定のトークン数、既知のデータ型、隠れ次元に対して、生成されたコードが128ビットのロード/ストア命令（Blackwellでは256ビット）を使用するように、最初に理想的なスレッドごとのベクトル次元（bfloat16では8、float32では4）を選択します。ベクトルサイズがわかれば、カーネルに必要なワープ数を計算できます。まずSMの数を最大化しようとします。もしワープ数がSMより少なければ、ブロックあたり1ワープが選択されます。これは通常、カーネルのデコードバリアントです。十分なワープ数がある場合は、ワープ数を8、16、または32へと段階的に増やし、パフォーマンスの向上が見られなくなる上限をベンチマークで測定します。

Tritonはこれらの最適化のいくつかを内部で実行していましたが、高レベルコンパイラの欠点に縛られることなく、CuTeレイアウト代数を慎重に使用することで同等の結果を達成できることがわかりました。特にワープ内の水平方向のリダクションが関与する場合、CuTeDSLは最適なコードを生成するために必要な微細な制御を提供してくれることがわかりました。

状況によっては、可能な限り多くのブロックに作業を分散させるために、cooperative_groups::this_grid().sync()を使用して、参加しているすべてのブロックを同期できます。制限が適用されます。ブロックが一度実行を開始するとSMからデスケジュールできないため、前方進行を保証するには、ブロック数がデバイス上の利用可能なSM数を超えてはなりません。これは、どの時点においても、すべてのブロックが同期に関与するグローバルメモリ位置を更新およびポーリングできることを保証するためです。この制限はcudaLaunchCooperativeKernelによって強制されます。グリッドバリアのコストはHopper上で約3μsであり、状況によっては暗黙のメモリバリアを介して以前のすべてのメモリ操作をフラッシュするコストが加わります。全体として、これはカーネルを2つの起動に分割するオーバーヘッドと同等であり、パフォーマンスに敏感な設計決定を導きます。Programmatic Dependent Launch（PDL）によりHopper以降は起動オーバーヘッドを削減できますが、デコードカーネルの場合、オーバーヘッドを回避する価値はあります。

ワークロードのいずれかがSM数よりも多くのブロックを必要とする場合、カーネル起動の暗黙的なバリアに依存して、カーネルを分割する必要があります。これは主にプリフィルに当てはまります。すべてのワークロードにとって理想的なブロック数がグリッドに収まる場合、グリッドバリアが推奨されます。これは、より少ないトークンで動作するデコードカーネルの場合と同様です。

CuTeDSLでは、これらの決定をホスト側の起動コードで行うことができ、カーネルを必要に応じてグリッドバリアとともに構成およびインターリーブできます（下の例を参照）。このようなハイブリッド戦略により、デコードカーネルを最小のレイテンシで到達するように最適化し、プリフィルカーネルを最大のスループットのために最適化できます。CUDAベースのソリューションとは異なり、このロジックをより簡潔に、かつコンパイル時のオーバーヘッドを抑えて定義できます。

python

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

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

以下で詳しく説明するMoE all-to-allカーネルの場合、グリッドバリアの有無にかかわらずカーネルを構築し、コンパイル時に最適化する柔軟性を持つことで、CuTeDSLを通じて大幅なパフォーマンス向上を達成できました。

私たちがCuTeDSLで最初に構築したカーネルは、小さな次元に対して最適化された専門化RMSノルムカーネルでした。最近のQwenモデルのセルフアテンションレイヤーで使用されるQKノルムは、ヘッドごとに適用され、64、128、または256要素に対して行われます。数千を超える隠れ次元で動作し、共有メモリとレジスタの両方を介した水平方向のリダクションを必要とする通常のRMSノルムとは異なり、QKノルムはワープ間のリダクションのみを必要とし、複数のヘッドを1つのワープに詰め込む可能性もあります。これは、コンパイル時に積極的に最適化して、すべてのヘッド次元とデータ型に対して最適なトークン対ワープの割り当てを見つけることができるユースケースとなりました。

RMSノルムはかなり単純な操作ですが、ベクトルの二乗ノルムを計算するために水平方向の合計が必要です：

cpp

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

アテンションヘッドは通常小さく、スレッドごとの操作は高度にベクトル化されているため、単一のワープで二乗ノルムを計算できます。メモリ帯域幅を最大化するために、各スレッドは128ビットの入力を処理する必要があり、これは8つのbfloat16値に変換されます。したがって、ハーフワープでヘッドを処理するのに十分であり、二乗ノルムはshuffle_sync_bflyによって実行される水平方向のリダクションを使用して決定できます。ヘッド次元とデータ型から始めて、ワープあたりのヘッド数を決定し、コンパイル時に引数とリダクションラウンド数を指定できます。

その後、自動的にベクトル化されたロード/ストアに依存し、ワープリダクションには低レイヤーのアーキテクチャ固有の操作を使用して、比較的少数のコード行で専門化されたカーネルを構築できます。カーネルはどのヘッド（QまたはK）に対して作業するかを決定し、適用する対応する重みを識別します。ヘッドからベクトル全体をロードし、最初にローカルで合計します。その後、ヘッドあたりのスレッド数に基づいて、二乗ノルムが計算されます。最後に、自動的にベクトル化されたRMSノルムが決定され、メモリに書き戻されます。WEIGHT_BIASやEPSなどのモデル固有のパラメータは、定数畳み込み（constant-folded）されます。

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

通常のRMSノルムカーネルと比較して、この専門化されたバリアントは、プリフィルとデコードの両方で2〜3倍優れたスループットとレイテンシを実現します。

MoEディスパッチと結合は、入力トークンをルーティング先の専門家（エキスパート）にシャッフルし、MoE GEMMカーネルが関連する行列乗算を実行した後に加重累積を実行します。さまざまな製品がさまざまなサイズのMoEモデルに依存しているため、パフォーマンスは特に重要です。デプロイメントタイプに基づいて、最適化するために異なるカーネルを区別します：

TP=1：単一GPU用のローカルカーネル。ディスパッチはトークンを関連する専門家にシャッフルし、結合はそれらを累積します。主にH100でのQwen3 30Bデプロイメントに使用されます。

DP=1, TP ≠ 1：テンソル並列、シングルノードカーネル。入力は複製され、ディスパッチは現在のランクのトークンを選択し、結合はすべてのランクにわたってそれらを累積します。NVLinkまたはNVLSのいずれかを介した累積。大規模なプリフィラーに使用されます。

DP ≠ 1, NVLink：オプションのテンソル並列処理を備えたデータ並列カーネル。各ランクは特定の専門家をホストする各ランクにトークンを送信します。ルーティング情報が集約され、トークンはNVIDIA NVLinkを介して転送されます。データ並列のH200またはGB200デプロイメントに使用されます。

DP ≠ 1, InfiniBand：NVLinkとNVIDIA Quantum InfiniBandの両方で動作するピア・ツー・ピアのデータ並列およびテンソル並列カーネル。マルチノードのHopperデプロイメントに使用されます。

ローカルカーネルは、ルーティングテーブルに基づいて各専門家に割り当てられたトークン数をカウントし、各専門家に属するトークンリスト内の各トークンに一意のインデックスを割り当てます。ルーティングステップの後、トークンは必要に応じてパディングされたバッファにコピーされ、MoE GEMMカーネルに渡されます。トークン位置を使用して、結合カーネルはレジスタ内で累積を実行し、最終的なトークンを出力バッファに書き込みます。

ディスパッチフェーズはプリフィルとデコードの間で異なります。トークンが少ない場合は読み取るルーティング情報がほとんどないため、単一のブロックで処理し、グリッドの残りの部分に信号を送り、個々のブロックがトークンのチャンクをコピーします。トークンの割り当ては、共有メモリ内のカウンタをアトミックにインクリメントすることで決定されます。プリフィルはより多くのルーティング情報を処理する必要があるため、ルーティングは別のカーネルで複数のブロックに分割されます。インデックスはまず共有メモリ内のカウンタを使用してローカルに割り当てられ、次にすべてのブロック間でグローバルメモリ内で集約され、各トークンの位置が決定されます。このカーネルのグリッドは、グリッド単位のバリアを使用できるように、SMの数にクランプ（制限）されます。後続のカーネルはコピーを実行し、利用可能なメモリ帯域幅を最大限に活用するために、より多くのブロックとワープにわたって起動されます。

結合カーネルは、ベクトル化された累積を実行するため、比較的単純であり、ディスパッチによって計算されたルーティング情報とオフセットに依存します。計算はすべてのデータ型と隠れ次元に対して効率的なコードにコンパイルされる単純なCuTeテンソルを通じて表現できるため、CuTeDSLはここで価値があることが証明されました。

テンソル並列デプロイメントでは、ルーティング情報と入力トークンが複製されます。ローカルディスパッチカーネルは、現在のランクにある専門家に割り当てられたトークンのみをルーティングおよび選択するように若干変更することで、完全に再利用できます。いくつかのパラメータとコンパイル時の専門化により、ローカルディスパッチカーネルが再利用されます。

ランク数は通常、トークンがルーティングされる専門家の数よりも少ないため、平均的なケースでは、各トークンには少なくとも1つのレプリカが各ランクにルーティングされていると想定できます。この状況では、ローカル結合カーネルを再利用してトークンのローカルコピーを集約し、続いてオールリデュース（all-reduce）を実行して参加ランク間で合計できます。プリフィルカーネルの場合は別々のカーネルを起動しますが、デコードでは両方の操作を単一の起動に融合します。

オールリデュースには、2つのランクのみが参加する場合は単純なNVLinkを使用し、4つ以上のランクが関与する場合はNVLSに切り替えてスループットを10〜20%向上させます。DSLのおかげで、MoEカーネルと通常のオールリデュースカーネルの間でリダクションロジックとピアバッファ交換の大部分を共有できます。これはハードウェアを厳密に制御することが有益なケースであり、NVLSにはmultimem.ld_reduceとmultimem.stに対するDSL操作が必要だからです。リダクション自体の実装は、主にNCCLテストのall_reduceの例に従っています。

データ並列カーネルは、NVLinkを使用して複数のデバイスのトークンを集約します。ディスパッチフェーズでは、トークンを交換する前に、ピアはまずローカルにルーティングして、送信する各トークンの各ランクでのターゲットオフセットを特定します。その後、競合を避けるために、送信者ごとに割り当てられた受信バッファにトークンをコピーし始めます。トークンごとのリリースストア書き込みを使用して、完了を通知します。受信側では、MoE GEMM用にすべてのトークンを単一のバッファに凝縮できるように、ルーティング情報が集約されます。受信者は、書き込みが完了するのを待つために、獲得ロード（acquire-load）を使用してピアが書き込んだフラグをポーリングします。

結合側では、トークンは受信を示すトークンごとのフラグとともに、ソースランクへNVLink経由でコピーされます。結合カーネルはローカルおよびシングルノードのバリアントと似ていますが、アドレス指定への若干の差異と追加の同期ロジックがあります。

このカーネルを実装するために、システムレベルとデバイスレベルの両方でスコープされたリリース獲得（release-acquire）ロードとストアを公開するPTXラッパーのライブラリを構築しました。同期フラグの粒度を改善し、NVLink経由でトークンを送受信するためにブロックを増やすことで、以前のCUDAカーネルに対して2倍の改善を達成しました。

私たちのピア・ツー・ピアカーネルは、以前はConnectX-7 InfiniBandアダプタとペアになったHopper GPU上で最高のパフォーマンスを達成していました。その後、パフォーマンスをさらに向上させるために、結合カーネルをCuTeDSLに移植しました。CUDAバージョンはすでにコンパイル時の専門化に大きく依存していたため、リポジトリ内でコンパイル時のボトルネックであることが判明し、いくつかのパラメータを省略せざるを得ませんでした。JITコンパイルされるDSLに移行することで、より多くのパラメータに対して専門化し、グリッド割り当てを調整し、プリフィルとデコードの間にさらなる差異を導入することで、レイテンシを約3%向上させました。

ROSEカーネルの観点から見ると、Blackwellのサポートは、アーキテクチャに依存しないカーネルと、アーキテクチャ固有のカーネルの2つのカテゴリに分類されます。テンソルコアに依存するアーキテクチャ固有のカーネルには、適応のための明示的な作業が必要です。しかし、DSLの使用のおかげで、アーキテクチャに依存しないカーネルは簡単に適応します。

Blackwell GPUは、Hopperの132基に対して最大160基のSMを備えています。ほとんどのカーネルにとって、これはコンパイル時の定数であり、ホスト側のランチャーコード内のロジックが、実行する最適なカーネル（プリフィルかデコードか）を決定するために考慮に入れます。さらに、Blackwellでは256ビットのロード命令も導入されており、メモリ帯域幅の使用率を向上させるためにベクトル幅（CuTeテンソルの最後の次元）を増やすことで簡単に発行できます。

Blackwellは、かなりの数の量子化スキーム（MXFP8、MXFP4、NVFP4）のハードウェアサポートを導入しています。これらのサポートは主に、FP8のブロック単位のスケーリングのために既存のカーネルを微調整することで実装されました。コンパイル時の構成スペースを爆発させることなく新しいモードをサポートできるため、ここでもJITサポートが重要です。過去には、Tritonはワープの半分または4分の1にわたる少数の要素間の水平方向のリダクションを表現するのが扱いにくいと感じました。NVFP4とMXFP8/MXFP4について、それぞれ16または32要素にわたって絶対最大値を計算する場合、最適なパフォーマンスを達成するにはリダクションに対するより細かい制御が重要でした。

比較的新しい言語として、LLMは事前トレーニング中にCuTeDSLコードに触れる機会がほとんどありませんでした。さらに、活性化とノルム用に最適化された現在の私たちのコードスタイルは、TMAとアーキテクチャ固有の行列乗算命令用に高度に最適化された既存のCUTLASSやQuackの例とは少し異なっていました。LLMがカーネル作成を支援できる段階に達することは、段階的なプロセスでした。

エージェントによるカーネル変換に頼る前に、一貫したパターンに従ってカーネルとユーティリティ関数の初期コーパスを手動で構築しました。コーパスが発展するにつれ、CUDAからCuTeDSLへの大まかな翻訳を実行でき、その後に大幅な手動の手直しを行うことができました。特に、モデルはポインタからテンソルへの変換やベクトル化に苦労し、テストをコンパイルしようとするときに型エラーの修正で立ち往生してしまうことがわかりました。モデルはカーネルのグリッドベースの性質に対する理解がほとんど示されておらず、スレッドの同期とグリッドの同期を混同していました。最も苛立たしかったのは、独自のCuTeDSL代替品に対して構築するのではなく、存在しないCUDA固有のライブラリの翻訳を幻覚（hallucinate）させようとする誘惑があったことです。時が経ち、CuTeDSLカーネルのコーパスが成長するにつれて、LLMはスケルトンコードの大部分を構築し、空白を埋めることができ、私たちが高レベルな設計に集中できるようになりました。

CUDAと比較して、CuTeDSLはコンパイル時のパフォーマンスがはるかに優れており、より多くのパラメータに対して専門化し、パフォーマンスの向上を達成できることがわかりました。Tritonと比較して特定の抽象化が欠けているため、カーネルは若干冗長ですが、ハードウェアに対するより細かい制御が可能になるため、複雑さはより高速なカーネルを構築するための正当な理由であると考えています。

執筆時点では、CUTLASS C++とCuTeDSLカーネルのパフォーマンスギャップは最小限であり、高性能GEMMカーネルとデバイス間通信、正規化、活性化を構成および融合するためのエキサイティングな機会が広がっています。このDSLにコミットすることで、デバイスからさらなるパフォーマンスを引き出し、埋め込みからLLMに至るまで、ワークロード全体のスループットを向上させることができるようになります。

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

https://github.com/NVIDIA/cutlass

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

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

BibTeX

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

Perplexity AIにおける推論

なぜCuTeDSLなのか？

開発者体験

広大な構成スペース

プリフィルおよびデコードの特化

グリッドの同期

CuTeDSLカーネル

QKノルム

Mixture-of-Expertsのディスパッチ/結合

ローカル

シングルノード

データ並列

ピア・ツー・ピア

Blackwellのサポート

AI支援によるカーネル生成

結論と今後の展望

参考文献

引用