Cohere、North Mini Code の高速化エンジン「デコード・メガカーネル」を公開

Cohere は本日、H100 単体で BF16 を処理するデコード用メガカーネルを中核に据えた、North Mini Code 向けのサービングエンジンを発表しました。このエンジンは vLLM と比較してエンドツーエンドで 1.25 倍から 1.41 倍の高速化を実現しています。

同社の GitHub リポジトリ cohere-megakernel では、このサービングエンジンの実装コードを公開しています。

現在、多くの LLM サービングスタックでは、フォワードパスごとにカーネルを順次実行する方式が採用されています。具体的には「QKV カーネルを実行して待機」「アテンションカーネルを実行して待機」「MoE カーネルを実行して待機」という流れです。それぞれのカーネル自体は問題なく動作しますが、課題はその間の待ち時間にあります。

バッチサイズが小さい場合、GPU はデコードステップの驚くほど長い時間を計算ではなく、待機に費やしているのです。

自己回帰的なデコード、特にバッチサイズが小さい状況では、根本的にメモリ帯域幅がボトルネックとなります。デコードステップごとに大量のデータを HBM（ハイバンド幅メモリー）から読み出しますが、実際の計算量は相対的に少なくなります。つまり、問うべきは「FLOPS（浮動小数点演算性能）をいかに引き出すか」ではなく、「メモリ帯域幅をいかに有効活用できるか」という点です。

例えば、North Mini Code は 30B パラメータを持つモデルですが、トークンごとにアクティブになるのは 3.3B パラメータです。BF16 で計算する場合、デコードステップごとに約 6.6 GB の重みデータをストリーミングし、8K コンテキストではさらに KV キャッシュとして約 0.5 GB を必要とします。

H100 は HBM を通じて 3.35 TB/s の帯域幅を提供しますが、理論上の最大速度（Speed-of-Light: SoL）は約 470 トークン/秒に達します。一方、vLLM でこのモデルをサービングすると 185 トークン/秒となり、SoL のわずか 39% に留まっています。

最近、メガカーネルが注目されています。これは、百個の小さなカーネルを個別に実行するのではなく、フォワードパス全体を一つの永続的なカーネルとして実行することで、性能ギャップを埋める手法です。

Hazy Research の「Look Ma, No Bubbles!」という先駆的な研究（設計概要は後述）を出発点に、数多くの follow-up 作品が発表され、さまざまなレベルでの高速化を実現しています。既存の取り組みは主に二つの方向に進んでいます。一つはメガカーネルを自動生成するコンパイラであり、もう一つはバッチサイズ 1 の場合のデコード速度を測定するスタンドアロンのデモです。

私たちはさらに一歩踏み出しました。本記事では、デコード用メガカーネルを中心に構築された、初めての本格的なサービングシステムをご紹介します。このシステムは、実際のサーバーに必要な機能をすべて備えています。つまり、継続的バッチ処理（continuous batching）、ページドアテンション（paged attention）、そして可変長のシーケンス長に対応し、ツール呼び出し機能も備えた OpenAI 互換エンドポイントの背後で動作します。

PointCode をこのシステムに接続すれば、コード生成が可能です。

バッチサイズ 1 の場合、当社のメガカーネルは秒間 292 トークン（tok/s）を達成し、SoL の 62% に相当します。これは vLLM よりも 1.58 倍高速です。この性能差はバッチサイズやコンテキスト長が 256K に達しても維持され、精度の低下は測定可能な範囲では確認されていません。

また、メガカーネルは評判ほど書くのが難しいものではないことも発見しました。そのため、既存のカーネルを一つのメガカーネルに移植するためのレシピも用意しています。私たちの実装では単一の CUDA ファイルで完結しており、コンパイラや新しいプログラミングパラダイム、あるいは特殊な抽象化は一切必要ありません。ただ、通常のタイル化された GEMM と通常のページドアテンションを、単一の呼び出し規約に収まるように再構成しただけです。

メガカーネルとは何か？

GPU はおよそ 100〜150 の独立したプロセッサ（SM：ストリーミングマルチプロセッサ）で構成されており、これらはすべて同じプログラム——つまり「カーネル」——を実行しつつ、異なるデータピースを処理します。メガカーネルは、フォワードパス全体を単一の永続的なカーネルとして実行するものです。各 SM に対してスレッドブロックを 1 つだけ起動し、デコードステップの全期間中そのブロックを常駐させます。

ドライバから作業を受け取るのではなく、各ブロックは「タスクリスト」を読み取ります。これはホスト側で準備され、グローバルメモリ上に置かれた、実行すべき小さな作業のリストです。データ依存関係はカーネル境界によって暗黙的に表現されるのではなく、グローバルメモリ上の明示的なカウンターとして記述されます。タスクが完了するとこのカウンターが増加し、入力が必要になるまでスピン（待機）します。

その結果、スケジューリングの単位は「操作全体」から「1 つの操作における 1 タイル」に縮小され、同期の単位も「GPU 全体」から「タスクが依存する特定の生産者」へと細分化されます。

画像2: ホストからデバイスへのデコードステップの一例。各操作ごとにカーネルを起動するのではなく、ホストはステップをタスクに分解します（1つの操作につき1タイル）。そして、SM（ストリーミングマルチプロセッサ）に対してラウンドロビン方式で分散し、各SMがグローバルメモリ上に独自のタスクリストを受け取ります。ほとんどのタスクの順序は、スケジューラーによってホスト側で決定されます（後述の「スケジューラー」セクション参照）。ただし、フルアテンションとMoE（Mixture of Experts）はこの例外です。これらは実行中のシーケンス長やルーティング処理に依存してタスク数が変動するため、どのSMも取り出せる共有ワークキューに格納されます。詳細は以下の各セクションで解説します。

高速化の源泉

推論エンジンでは通常、RMSNorm、QKV、アテンション、MoEなど、各操作ごとに1つのカーネルを起動します。そして、最適化の大部分はこの個々のカーネルを可能な限り高速化することに注がれます。これは、計算集約型のGEMM（行列乗算）が大量に発生するトレーニングやプリフィル（初期推論）フェーズでは有効なアプローチです。これらの場合、各カーネルには十分な作業量があり、SMを飽和させることができます。

一方、デコード処理は全く異なる状況にあります。ここではメモリ帯域幅とレイテンシがボトルネックとなります。デコードステップの大部分は演算強度の低いGEMV（行列ベクトル積）であり、その速度はHBMから共有メモリへ重みをストリーミングする速さによって決まります。重みの移動を妨げるあらゆる要因は時間のロスとなり、「操作ごとに1つのカーネル」というアプローチには、重みが止まってしまう箇所が複数存在します。これらの停止時間が積み重なって、一般的な推論エンジンで利用されていない帯域幅の61%を占める主要原因となっています。

最も基本的なメリットは、起動と同期のオーバーヘッドを削減できる点にあります。連続するカーネルの間では、すべての SM（ストリーミングマルチプロセッサ）が完了するまで、どの SM も次のカーネルを開始できず、ドライバーも次のグリッドをディスパッチする必要があります。1 層あたり数十個の小さなカーネルで構成されるデコードステップの場合、これらの待ち時間が累積してしまいます。一方、メガカーネルでは、このコストは操作ごとではなく、ステップごとに一度だけ支払えば済みます。

本モデルにとってより重要となるメリットを、影響度の大きい順に 3 つ挙げます。

1. ウェーブ量子化の削減

仮に、あるカーネルで処理すべきタイルが 200 個あり、GPU に搭載されている SM が 132 個だとしましょう。最初の 132 個のタイルは並列実行されますが、残りの 68 個は第 2 ウェーブで実行され、その間 64 個の SM がアイドル状態になります。つまり、1.5 ウェーブ分の作業を行うのに、2 ウェーブ分の時間がかかってしまうのです。カーネルが小さければ小さいほど、この丸め誤差による非効率さは顕著になります。

これはタイルをより均等に分割することで解決できる問題ではありません。GEMM のタイル形状は行列の次元やカーネル設計によって制約されるため、総タイル数が SM 数のきっちり倍数になることは稀です。

メガカーネルでは、処理の境界を丸める必要はありません。入力準備が整ったタイルは、空いている任意の SM（ストリーミングマルチプロセッサ）で即座に実行されます。このアーキテクチャは「並列トランスフォーマー層」を採用しているため、他のアーキテクチャよりもメガカーネルの恩恵を大きく受けます。具体的には、アテンション処理と MoE（Mixture of Experts）のフィードフォワード計算が、同じ正規化された入力から同時に実行され、両者の出力を必要としません。これらは層の最後で、融合された残差加算と RMSNorm によってのみ結合されます。

*図 4：North Mini Code が採用する並列トランスフォーマー層の構成。

メガカーネルを用いることで、アイドル状態にある SM に「バックフィル」処理を適用できます。並列トランスフォーマー層の特性を活かせば、より積極的なバックフィルが可能になります。可能な限り、実行準備が整っていると予測されるタスクを、空いている SM へ確定的に割り当てるのです。この割り当ての詳細については、「タスクスケジューラ」のセクションで解説します。

下の図では、従来のサービングスタックで MoE デコード層を実行した場合と、メガカーネルを用いた場合を比較しています。

画像 5: MoE デコードレイヤーの 1 つ。同じ操作を同じ順序で実行し、2 つの方法でスケジューリングしています。上段：各操作ごとに 1 つのカーネル。ウェーブ量子化とハードウェアのジッターによりアテンションは同時に終了しません。カーネル境界にあるバリアによって、すべての SM が最後の SM の完了まで待機します（斜線部分）。このパターンが各カーネル境界で繰り返されます。下段：メガカーネル。各 SM がアテンション作業を終えるとすぐに MoE タイルの処理を開始するため、同じ時間間隔で待機するのではなく実際に作業が行われます。可視性を確保するために 16 個の SM と各操作あたり数十個のタイルで描画されていますが、実際のカーネルでは 132 個の SM と数千個のタスクが存在します。持続時間は例示的なものであり、時間軸は相対値です。

2. 偽の依存関係を排除する

SM は、負荷量が同じであっても必ずしも同時に作業を終えるわけではありません。カーネル境界は全グリッドにわたるバリアとして機能するため、最も遅い SM の速度がすべての SM のペースを決定します。例えば、アテンション処理が 4 つのキー/バリューグループに分割されている場合、あるグループが早く完了しても、他の 3 つのグループが追いつくまでその SM はアイドル状態になります。しかし、次の操作に必要なデータはすでにメモリ内にあるはずです。

微細粒度のバリアを導入することで、この偽の依存関係を解消できます。特定の KV グループに対する O-proj（出力投影）処理は、*そのグループの*アテンション出力が到着した時点で即座に開始されます。同様に、MoE の下流投影タスクは、対応するエキスパートの上流投影タスクが完了すればいつでも開始でき、すべてのエキスパートの上流投影を待機する必要はありません。

3. 重みのプリフェッチ

重みは不変であり、このステップでの活性化値に一切依存しません。そのため、タスクは活性化の依存関係が満たされる前に、HBM から共有メモリへ重みのタイルをストリーミングし始めることができます。これはカーネル境界では許されない行為です。

特にルーターや QKV 投影層に対してはこの手法を積極的に採用しています。これらの層は、前のレイヤーの O-投影演算が終了する直後、RMSNorm が実行される前というタイミングで重みをプリフェッチします。これにより、本来なら無駄になる帯域幅を有効活用できるのです。

設計の原点

私たちの設計は、前述した Hazy Research の画期的なポストから多くの知見を借用しています。同研究では Llama-3.2-1B の順伝播処理を単一のカーネルに統合し、バッチサイズ 1 の条件下で H100 のメモリ帯域幅の 78% を達成しました。これは vLLM や SGLang が同等の条件で達成できた値（約半分）と比較して圧倒的な成果です。

彼らのアイデアのうち、以下の 3 つが私たちの設計において特に重要でした：

GPU 上での「タスクインタプリター」パターン。各 SM は、ホスト上で用意され、順次実行 across で再利用されるタスク記述子のリストを走査します。コントローラーワープが記述子を読み取り、それぞれが特定の操作を実装するオンデバイス関数へタスクをディスパッチします。このアイデアの実装については、図 2 と図 6 で説明しています。

カウンターベースの同期。依存関係バリアはグローバルメモリ上の単純な整数であり、各ステップ前にゼロ化されます。タスクが完了するとカウンターをインクリメントし、開始する前に該当するカウンターでスピンします。詳細な実装については「バリアーズ」セクションおよび図 7 で解説しています。

タスク境界を超えたオーバーラップ。ある SM 上の前のタスクが結果の保存を行っている間に、次のタスクはすでにウェイトの読み込みを開始できます。

独自の取り組み

当社の実装は、同様のメガカーネルとはいくつかの点で異なります。GEMM 実装では、バッチサイズが 1 の場合でもテンサコア命令を積極的に活用しています。これは、wgmma が CUDA コアよりもわずかに高速であり、レジスタ圧力を低減できることを確認したためです。

もう一つの違いは、重みのプリフェッチ実装に共有メモリページングを使用していない点です。当初は共有メモリページングの実験も行いましたが、これは前のタスクがバッファを解放する前にメモリアクセスを開始できるという利点があります。しかし実際には、管理処理が複雑になり、バグの恒久的な原因となり、そのオーバーヘッドが得られるメリットを上回ってしまいました。

そのため、各オペコードは独自のワープ特化パイプラインを持ち、共有メモリレイアウトはコンパイル時に静的に決定されます。重なり合い（オーバーラップ）は、よりコストの低い2つの箇所から実現しています。

同じタイプの連続する GEMM タスクの間。MoE の GEMM タスクはタイルのリストを走査しますが、パイプラインはそのステージフェーズを各タイル境界で空にして再充填するのではなく、リスト全体にわたって維持します。MoE GEMM 内では、現在のタイルがまだテンサーコアまたはエピローグ処理中であっても、次のタイルの重みはすでに転送中です。同じ操作であり共有メモリレイアウトも同一であるため、共有メモリページングは単なるマルチステージパイプラインに帰着します。この種のプリフェッチは、永続化されたグループ化 GEMM カーネルと同じ精神に基づいています。

GEMM パイプラインの内部構造について解説します。重みは不変であるため、プロデューサーワープはクロス SM 間のアクティベーション待機前に重みタイルの読み出しを開始します。後ほど本文で示す擬似コード（リスト 1）でも、prefetch_weight_tiles が wait_input_bars よりも上位に配置されていることが確認できます。MoE の下段投影演算がその典型例です。専門家の重みは HBM からストリーミングされ始める一方で、上段やゲートはまだ下段で消費される隠れ状態を計算中です。つまり、入力待ちのタスクであってもデータ転送は継続して行われます。

もう一つの例として、RMS 正規化が完了する前に QKV とルーターのプリフェッチを行うケースがあります。

スケジューリングにも多大な労力を注ぎました。ほとんどのタスクは、私たちが精密にチューニング可能なホスト構築型の静的スケジュールに従います。一方、アテンションと MoE では、連続バッチ処理やルーティングに伴うランタイム依存のワークロードをバランスさせるため、ローカルワーステールを採用しています。これについては後ほど詳しく触れます。

すべての演算に共通の ABI

このメガカーネル全体は、共通の呼び出し規約によって結合された複数の小さなカーネルとみなすことができます。各小さなカーネルは、必ず 3 つのワープグループ（各グループは 4 つのワープから構成）で実装されなければなりません。具体的には、8 つのコンシューマー・ワープ、1 つのコントローラー・ワープ、1 つのプロデューサー・ワープ、そして 1 つのストアー・ワープを含みます。各ワープは独自のレジスタ要件を満たす必要があります。さらに、各小さなカーネルは固定サイズのタスク記述子から「パラメータ」を読み取る必要があります。

このような規約は、コンパイラやオペレーティングシステムにおけるアプリケーション・バイナリ・インタフェース（ABI）に似ていると考えられます。本稿では、この規約を指すために"ABI"という用語を使用します。

ABI の仕組みを理解するために、GEMM（行列乗算）の例を見てみましょう。

手書きでも記述可能に保たれている理由は、GEMM（行列乗算）に限らずすべての演算が同じ ABI（アプリケーション・バイナリ・インタフェース）に従っているからです。具体的には、タスクの定義、実行時の形状、完了の合図方法が統一されています。すべての演算がこの ABI に従うことで、それらをマーカーカーネルに組み立てる作業は非常に管理しやすくなります。このセクションでは、これら 3 つの要素について詳しく解説します。

タスク

マーカーカーネルが新しい演算を考案しているわけではありません。通常のデコードグラフ（QKV、アテンション、O プロジェクション、ルーター、MoE、リジューアル/RMSNorm、LM ヘッド）を引き継ぎ、それを小さなタイル化された「タスク」のリストに落とし込みます。デコードステップ全体をカバーする命令語は 16 種類です。

Dense GEMMs

QKV_PROJ, O_PROJ, FFN_UPGATE_ACT, FFN_DOWN, LM_HEAD

これらは行列乗算の出力タイル 1 つ、あるいは分割 K 還元（split-K reduction）のスライス 1 つを扱います。QKV_PROJ はエピローグ内で RoPE を適用し、FFN_UPGATE_ACT はエピローグで SiLU と乗算を融合します。

Attention

ATTN_DECODE, ATTN_COMBINE, ATTN_DRAIN

ATTN_DECODE は部分的な KV スライスに対してアテンションを計算し、ATTN_COMBINE はそれらの結果を実際の出力に統合します。ATTN_DRAIN はフルアテンション層でのみ使用され、動的に生成された作業キューに対するクレーマー（要求者）として機能します。詳細は後述のスケジューリングセクションをご覧ください。

MoE routing

ROUTER_GEMM, ROUTER_TOPK, ROUTE_FINALIZE, MOE_GATHER

128 個のエキスパートに対するスコア付けを行い、トークンごとに上位 k 個のエキスパートを計算して、MoE GEMM の作業キューを生成します。

MoE GEMMs

MOE_UPGATE_ACT_DRAIN, MOE_DOWN_DRAIN, MOE_COMBINE

ROUT_FINALIZE が生成した作業キューを管理するエキスパート FFN において、MOE_COMBINE は各トークンの 8 つのエキスパート出力を、ルーターが算出したスコアに基づいて重み付けして合計します。

すべての GEMM オペコードは単一のパイプラインを共有しています。 O-proj（出力投影）、ルーター、密な FFN、LM ヘッド、そして両方の MoE エキスパート GEMM は、同じワープ特化型の本体を実行します。プロデューサーが重みと活性化のタイルを共有メモリ上のステージリングにストリーミングし、コンシューマーがテンソルコア上で累積し、ストアラーが出力タイルを書き込んでバリアで同期します。各オペコードごとに異なるのは、対象となるテンソルや待機・シグナルするバリア、エピローグで SiLU と乗算を融合させるかどうかなど、短い詳細リストのみです。また、ストア処理が split-K 還元を行うかもこれに含まれます。QKV は HBM への書き戻し前にレジスタ内で RoPE を適用した点以外は同じパイプラインです。MoE のドレイン（排出）も同様のパイプラインですが、次のタイルの選択方法が異なります。座標を記述子から読み取るのではなく、キューから作業を引き受けるのです。

その結果、新しい GEMM タスクを追加する際は、カーネル自体を書き換える必要はなく、小さな編集で済みます。ロード・計算・ストアのループを再構築するのではなく、必要な詳細情報だけを埋め込むだけです。

タスク記述子

各タスク記述子は 32 の int32 フィールドから構成され、ホストエンコーダーによって各 SM（ストリーミングマルチプロセッサ）のバッファに書き込まれます。フィールド 0 はオペコードを示し、残りのフィールドはオペ固有の情報です。具体的には、どのレイヤーか、どの出力タイルか、どの split-K スライスか、何個の K タイルを処理するか、そして最も重要な点としてどのバリアで待ち、何をカウントして待機し、完了時にどのバリアをシグナルするかが記述されます。

デバイス側の実装はシンプルです。ブロックごとに 1 つのワープがコントローラーとして機能し、先読みしたディスクリプタを共有メモリのリングバッファに格納します。これにより、SM（ストリーミングマルチプロセッサ）がタスクディスクリプタの待機でストールすることを防ぎます。残りのワープは、このリングバッファからタスクを取り出して実行します。

スレッドブロックの形状

すべてのタスクは、オペコードの種類に関わらず、同じスレッドブロック形状で実行されます。12 個のワープが 3 つのワーググループに編成されています。

WG0 warp 0: コントローラー

共有メモリリングバッファへ次のタスクディスクリプタを先読みする

WG0 warp 1: プロデューサー

TMA ロードを発行し、セマフォを初期化して入力バリアを待機する

WG0 warp 2: ストアラー

TMA ストアを発行し、出力バリアにシグナルを送る

WG0 warp 3: アイドル

未使用

この構成の大部分は、スタンドアロンの Hopper GEMM で標準的なプロデューサー/コンシューマーのワープ特化パターンに従っています。メガカーネルが独自に貢献している点は、アテンションや RMSNorm を含むすべてのオペコードがこの同じ形状を使用することです。つまり、ある SM が QKV タイルの処理を終えた直後に、ワープ数や各ワープの役割を変更せずに、すぐに次のアテンションスライスを実行できます。

これらの役割はコンパイル時のタグであり、ランタイムでの分岐ではありません。各操作本体は役割に対してテンプレート化されており、if constexpr (role == PRODUCER) によって到達不可能なコードがコンパイル時に削除されます。これにより、通常の GEMM のようにプロデューサー、コンシューマー、ストアラーを一つの関数内で記述できつつ、各ワープは実際に実行するパスのみを保持します。

コントローラーは決してその関数には入りません。コントローラーには独自のループがあり、その役割は次のディスクリプタを小さな共有メモリリングにプリフェッチすることと、ワーカーに対してスロットが準備できたことを通知することの 2 つだけです。

ただし、注意すべき点があります。ワーカー同士で同期を行う際に __syncthreads() を使用することはできません。これは、ワーカーとは独立して動作するコントローラー・ワープを待機させてしまうためです。そのため、ワーカーたちはコントローラーを除外した名前付きバリア（worker_sync）で互いに合流します。

ここでは、メガカーネルで実装された GEMM（行列乗算）タスクの概要を示します。ワープ特化型のホッパー用カーネルを記述した経験があれば、この構造は馴染み深いものとなるでしょう。MoE ドレインも同様のタスクであり、キューから取得したタイルに対してループ処理を行う形で呼び出されます。

def gemm_task(task):  # compute Y = X @ W
    if constexpr (role == producer):
        prefetch_weight_tiles(task)      # optional, before the wait
        wait_input_bars(task)            # cross-SM: my activations ready?
        for k in k_tiles(task):
            tma_load_A_B(k)              # async copy HBM -> shared memory
            signal_stage_ready(k)        # intra-block: stage k is loaded
    elif constexpr (role == consumer):
        for k in k_tiles(task):
            wait_stage_ready(k)
            wgmma(k)                     # tensor-core MMA
        epilogue_to_smem()
    elif constexpr (role == storer):
        tma_store_Y()                    # async copy shared memory -> HBM
        arrive_output_bar(task)          # cross-SM: my tile is visible
    worker_sync()                        # all worker warps; excluding controller

def moe_drain_task(task):
    while True:
        tile_id = atomicAdd(n_claimed_tiles, 1)
        if tile_id >= len(moe_workqueue)：
            break
        tile = moe_workqueue[tile_id]
        gemm_task(tile)

def controller_loop(tasks):
    for i in range(len(tasks)):
        slot = ring[i % RING]
        if i >= RING:
            wait(slot.done)              # workers finished with this slot
        slot.task = load(tasks[i])   # prefetch the next descriptor
        arrive(slot.ready)               # signal workers: you can start

def worker_loop():
    for i in range(len(tasks)):
        slot = ring[i % RING]
        wait(slot.ready)                 # descriptor is in shared memory
        gemm_task(slot.task)             # or switch(opcode) onto another body
        if storer:
            arrive(slot.done)            # slot is free for the next prefetch

リスト 1：メガカーネルにおける GEMM タスクの実装方法を示します。全体構造は、ホッパー向けに標準的なワープ特化型 GEMM の構成に従っています。

if constexpr ブランチは、上記のコンパイル時のタグであり、実行時の切り替えではありません。単独のカーネルでは現れない 2 つの行、すなわち wait_input_bars と arrive_output_bar は、SM（ストリーミングマルチプロセッサ）間を跨ぐカウンタです。

記述子に格納される。待機処理の上方にある prefetch_weight_tiles は、メガカーネル固有の別のオプション操作である。

コントローラーはワーカーたちとバリア（同期点）に参加することなく、常に 1 つ先の記述子を保持します。アテンション、ルーティング、MoE の処理は、同じプロデューサー/コンシューマー/ストレージの役割に割り当てられます。worker_loop 内の switch (opcode) のみが変更されます。既存のカーネルをメガカーネルに組み込む際の摩擦となるのは、このラップループと 2 つのバリア呼び出しだけであり、新しいプログラミングモデルが必要になるわけではありません。

バリアの実装

Hazy Research に着想を得て、当社のバリアはグローバルメモリ上のカウンタとして実装されています。

// wait: spin until enough upstream tasks have arrived
while (*(volatile const uint32_t*)bar < target) {
    __nanosleep(20);
}
__threadfence();

// arrive: publish my tile, then signal downstream tasks
fence.proxy.async;     // make async (TMA) stores visible first
__threadfence();       // make data computed by this SM visible to others
atomicAdd(bar, 1);

リスト 2: バリアをグローバルメモリのカウンタとして実装しています。待機中のスレッドは依存関係が満たされるまでスピンプール（空回し）を行います。

これが依存関係のメカニズム全体です。タスクは特定の上位タスクではなく、単一の「カウント」を待ちます。これにより、ファンインやファンアウトの規模に関係なく、シグナル送信と待機の両方が O(1) の時間で完了します。SM（ストリーミングマルチプロセッサ）間でデータ競合が発生し、結果としてデータ破損につながるのを防ぐために、フェンス処理は不可欠です。

メガカーネルの書き方

メガカーネルとは、まるで完全な書き直しのように聞こえるかもしれません。しかし実際には、ABI（アプリケーションバイナリインタフェース）によって作業範囲が制限されています。一度、ある操作が共通のスレッドブロック形状と記述子フォーマットに従えば、他のタスクと同様にメガカーネルに組み込まれます。パフォーマンスが重要なロジックは、既存のカーネルから再利用可能です。

既存のカーネルライブラリから始める場合、私たちは以下のような手順を提案します：

まず、すでに単体で十分な競争力を持つ GEMM やアテンションカーネルから始めましょう。メガカーネルは演算間の「接着剤」を除去するものであり、遅い演算を高速化するものではありません。私たちの実装では、TMA マルチキャストやピンポンスケジューリングを使わず、単純なワープ特化パイプラインを用いて、cuBLAS や FlashAttention-3 と同等の性能を出しています。

次に、カーネルを ABI に適合させます。コンシューマー用ワープは正確に 8 つ、プロデューサー/ストレージ用ワープは最大 3 つとします。標準的なホッパーアーキテクチャ向けワープ特化カーネルでは、これは主にプロデューサー側とコンシューマー側のコードを別関数に分けることを意味し、これにより各ワーグルのレジスタ制限が有効になります。

次に、データ依存関係を表現するバリアを追加します。1 つの GEMM タイルまたは 1 つの KV グループが 1 つのタスクとなります。下流の演算は、データの整合性を確保するためにこれらのバリアを慎重に待機する必要があります。

最後に、グローバルメモリのタスクリストへタスク記述子を出力します。最適なスループットを得るために、タスクの実行順序と配置を調整します。

この中で最も難しいのはステップ 3 です。カウンティングバリアは、どのスレッドが到着したかを確認せずに N 個のスレッドが揃うと解放されます。そのため、間違った回数で到達するオペコードがあると、ワープが異なるタスクに逸脱し、後ほどデッドロックを引き起こす可能性があります。バリアの管理には、明確な不変条件と慎重なテストが必要です。

記述子フォーマット、ブロック形状、およびカウンタープロトコルが、統合契約全体を形成します。これにより、新しい演算を追加する作業は、新たなカーネルアーキテクチャを開発するプロジェクトではなく、境界が明確なエンジニアリングタスクとなります。

タスクスケジューラ：SM へのタスク割り当て

メガカーネルの実装が完了したら、次は各 SM（ストリーミングマルチプロセッサ）にどのタスクを割り当てるか、またその実行順序をどう決めるかを設計する必要があります。この割り当てのことを「タスクスケジューリング」と呼びます。

タスクスケジューラは、タスクグラフを入力として受け取り、各 SM が順次実行すべきタスクの具体的なリストを出力します。実は、スケジュールの設計には非常に大きな自由度があります。細かいバリア（同期ポイント）を設定するのは、メガカーネルの正しさを保証するためであり、あるタスクがその入力データをすべて取得するまで開始されないように制御するためのものです。

スケジュールが循環待ちを引き起こさない限り、タスクの順序を自由に並べ替えることができます。異なるスケジュールでも最終的な出力は同じになりますが、スループットには大きな差が生じます。

当社のスケジューラは基本的に静的ですが、一部の操作については動的なワークスチールを採用しています。

静的な部分：ウェーブ順序と配置。ホストは各レイヤーに対して名前付きのウェーブを構築します（qkv、router、attn、moe_up_drain、oproj、rmsnorm など）。

順序を選択し、それらを単一のタスクリストに統合します。

タスク k を SM (k mod 132) に割り当てる。現在の順序では、ルーターとルーティング設定がアテンションよりも先に配置されています。**

qkv → router → top-k → route setup → MoE gather → attention → MoE up/down → O-proj → RMSNorm

**

この順序付けにより、MoE ブランチには先行するメリットが生まれ、計算負荷の低いルーティング処理を、より重い QKV GEMM 処理と組み合わせてスケジューリングできます。これによって SM（ストリーミングマルチプロセッサ）の利用効率が向上します。

しかし、この単純な図式を超えると、スケジュールの影響を単独で切り離して評価するのは困難です。あるウェーブの順序を変更すると、その依存関係にあるタスクがいつ準備完了するか、どのタスクが同じ SM を共有するかが変わり、さらに HBM バンド幅が各タスク間でどのように配分されるかも影響を受けます。これらの選択は、レイヤー全体にわたって相互に影響し合います。

私たちの実験では、スケジュールの重要性が明らかになりました。入力長 8K の均一ルーティングにおけるアブレーション実験を 2 つ行い、ウェーブの順序のみを変更しました。

Interleaved（交互実行）

アテンション処理開始後、MoE のアップ/ダウン処理をアテンション結合の前に実行します。このアプローチの狙いは、両方のブランチから準備完了したタスクを常に利用可能に保ち、SM が待機する時間を減らすことです。

qkv → router → top-k → route setup → MoE gather
→ attention → MoE up/down → attention combine → MoE combine → O-proj → RMSNorm

Attention first（アテンション優先）

MoE の GEMM 処理を開始する前に、まずアテンション結合と O-proj を起動します。ただし、ルーターの早期開始は可能にします。

qkv → router → top-k → route setup → MoE gather
  → attention → attention combine → O-proj → MoE up/down → MoE combine → RMSNorm

1

291

282

3%

236

19%

2

423

406

4%

364

14%

4

553

531

4%

509

8%

これらのスケジュールはすべて、正しい出力を生成します。また、依存関係に配慮した配置（dependency-affinity placement）も試みました。これは、タスクをその依存元プロデューサーと同じ SM 上に配置し、キャッシュの再利用や SM 間のハンドオフ削減を図る手法です。しかし、この方法ではカーネルの実行速度が 1〜2% 低下したため、リリース版では単純なラウンドロビン方式を採用しています。

表に示されたのは測定結果ですが、なぜ特定の順序や配置が優位になるのかを説明する完全な因果モデルは、まだ解明されていません。

動的な部分：可変サイズの作業に対するローカル・ワークスチール

フルアテンションでは、各ライブリクエストごとに読み込む KV の長さが大きく異なり、ルーターが各 MoE エキスパートに割り当てるトークン数も決定します。これらのタイル数は、ステップ実行時に初めて判明します。

この分割は意図的なものです。静的スケジュールはタスクの順序と配置の大部分を固定し、パフォーマンスチューニングに必要な制御を提供します（上記のウェーブオーダーのアブレーション実験が示す通り）。ローカル・ワークスチールは、ホスト側で事前に予測できない負荷不均衡のみを処理します。

静的なタスクリストには、一定数の小さな「クレーマー」タスクが確保されています。ATTN_DRAIN または MOE_*_DRAIN クレーマーは、原子操作によって各ステージの共有キューから次のアイテムを奪い取り、キューが空になるまで実行し続けます。アテンション用のクレーマーはアテンションタイルを、MoE 用のクレーマーは MoE タイルをそれぞれ奪います。

クレーマーの数で並列性を制御し、現在のアテンションキューにはバッチ内のライブリクエストが反映されます。これにより、静的スケジュールの安定性が保たれつつ、バッチ変更ごとにリスト全体を再構築することなく、可変な作業負荷をバランスさせることが可能になります。

以下のタスクリストグラフでは、静的なウェーブ、微細な依存関係、そしてローカルで獲得されたステージを一つのビューに統合して示しています：

画像 7：MoE のデコード層をカーネルがどのように見ているか。O-proj の分割された K が KV ヘッド数と等しいため、アテンションは 4 つの KV ヘッドとして描画されます。つまり、各 O-proj の分割はちょうど 1 つのヘッドのアテンション出力に依存します。そのため、あるヘッドが QKV を処理している間でも、別の準備完了のヘッドはすぐに O-proj に進入できます。破線のカードは、前述のローカルで確保されたステージを示しています。

試した他のスケジューラ

プロジェクト初期には、密度モデルを対象に、より洗練された複数のスケジューラを実装しました。その一つは、メモリ読み込みから各タスクのコストを見積もる貪欲型トポロジ対応スケジューラでした。また、数百のランダム候補の中から最適なスケジュールを選出する全探索アプローチも試みました。いずれも、当時のワークロードでは単純なラウンドロビン方式よりもスループットが約 10% 向上しました。しかし、MoE の登場で状況は一変します。ルーティングにより作業量と配置が動的に変化するため、密度モデル向けのスケジューリング手法はそのまま適用できませんでした。最終的に私たちは、調整済みのウェーブ順序を採用したシンプルなラウンドロビン方式と、前述のローカルワークスチールを組み合わせる方針に落ち着きました。

ただし、Blackwell GPU 上で動作するメガカーネルや、NVFP4 や FP8 といった狭帯域量子化タイプにおいて単純なラウンドロビンが不十分となった場合、より複雑なスケジューラの再検討を行う予定があります。

カーネルを取り巻くサービングエンジン

サーバーには、2 つの長期間稼働するホスト側のスレッドが存在します。1 つは Python スレッドで、制御プレーンとして機能し、リクエストの受け付け、プレフィル処理の実行、KV キャパシティとバッチの管理を担当します。もう 1 つはネイティブ C++ スレッドで、最大限のパフォーマンスを発揮するためにデコード処理を担います。

この 2 つのスレッドは交互に動作します。Python は、プレフィル処理やアクティブなバッチサイズの変更が必要な際に C++ を一時停止させます。

この引き継ぎが存在する理由は、メガカーネルが事前に構築されたタスクリストを実行しているからです。タイル数、バリアインデックス、および各タイルの所属スロットは、すでに組み込まれています。そのため、Python はバッチの受け入れ、終了、または形状変更の前にデコードを一時停止し、C++ が新しいバッチサイズとコンテキストに一致するスケジュールで再開します。

下の図は、それぞれの側が管理する状態と、スケジュールがどのように変化するかを示しています：

Python と C++ が交代で担当

C++ が連続的なデコードループを担います。KV 位置の更新、ステップごとの状態クリア、スケジュールの選択、メガカーネルの実行、トークンのサンプリング、そして生成されたトークンのストリーミングを行います。アクティブなリクエストが存在する限り、このループは繰り返されます。

一方、バッチの変更に関わる処理はすべて Python が担当します。具体的には、リクエストの受け付け、プレフィル実行、スロットの割り当てまたは退避、そしてより小さなバッチサイズへの切り替えです。

park / resume の矢印は、この 2 つの所有者間の引き渡しを表しています。

# Python control plane
while server_is_running:
    request = wait_for_new_request()

    decode_service.request_pause()       # C++ finishes its current decode step.
    decode_service.wait_until_parked()   # GPU is now idle; C++ will not read batch state.

    prefill(request)                     # Ordinary prefill kernels populate its KV pages.
    decode_service.admit_or_evict(request)
    decode_service.switch_batch_size_if_needed()

    decode_service.resume()              # C++ snapshots the new state and keeps decoding.

// C++ decode service
while (!stop) {
    if (pause_requested || active_requests == 0) {
        signal_parked();                  // Python may now mutate the batch.
        wait_for_resume_or_admission();
        continue;
    }

    update_kv_page_tables();
    clear_scratch_and_barriers();
    select_schedule(max_live_context());
    launch_megakernel();
    sample_and_stream_tokens();
    retire_finished_requests();
}

パージング（停車）は、変更可能なバッチ状態の所有権境界です。Python は、リクエストごとのメタデータ、トークンバッファ、KV ページテーブル、および新しいバッチサイズに対応するポインタを変更します。その間、C++ は待機します。C++ は resume（再開）後にのみこの状態をスナップショットし、次のデコードステップの所有権を取得します。これにより、メガカーネルが半更新されたバッチを観測することを防ぎます。

現在のトレードオフはシンプルです。プレフィル中はデコードが一時停止するため、エンジンはまだ GPU 上でプレフィルとデコードを混在させていません。アクティブなリクエストがない場合、C++ は独自にパージングし、Python が admitting（受け入れ）後にそれを起動します。

外部からは、ストリーミング、プリフィックスキャッシング、ツール呼び出し機能を備えた OpenAI 互換サーバーとして見えます。

既知の制限事項

これらは現在の実装における制限であり、設計自体の限界ではありません。これらの機能については今後追加する予定です。

プリフィルとデコードを混在させない。プリフィルはすべてのデコードリクエストを一時停止するため、非常に短いリクエストが多数あるワークロードではパフォーマンスがやや低下します。

バッチサイズの最大値は 8 です。これはメガカーネル自体のアーキテクチャ上の制限ではなく設定によるものです。サーバーが必要であればより大きなバッチサイズをサポート可能ですが、現時点では大規模なバッチサイズに対する性能調整はまだ行われていません。

MK はデコード専用です。プリフィル処理は通常の PyTorch カーネルとして実行されます。

パフォーマンス

デコードスループット

設定。すべての数値は、North Mini Code を実行している単一の H100（SM 数 132）上の結果です。ベースラインは vLLM v0.24 で、FA3 アテンションバックエンドと Triton MoE バックエンドを使用し、プリフィルは無効化しています。両方のエンジンとも合成された KV キャッシュに対してデコードを実行するため、比較対象からプリフィルの計算を除外し、純粋なデコード処理に焦点を当てています。測定は 1,000 トークンの出力トークンに対するデコードスループットです。

2 つのベンチマーク設定を報告しており、その違いからは非常に多くの示唆が得られます。

• 実モデルチェックポイント。両方のエンジンとも、実際の North Mini Code の重みを使用します。これにより、MoE エキスパートの分布はモデル本来の相関を持つルーティングになります。
• 均一ルーティング。vLLM はシミュレーションされた均一ランダムルーティングで動作し、メガカーネルはランダムな重みを使用して約均等なエキスパート分布を生成します。

コンテキスト長 8K の場合、各バッチサイズにおける vLLM を基準に正規化した結果：

*図 9: 8K コンテキストにおける megakernel と vLLM のデコードスループットを比較。実際の専門家（エキスパート）の分布条件下では、megakernel がより大きな高速化効果をもたらします。

まず、均等ルーティングの場合、パイプラインのバブル（アイドル時間）が最も多くなる時に速度向上効果が最大になります。バッチサイズが 1 の場合、1 つのリクエストあたりの処理量は最小となるため、GPU 時間の大部分がパイプラインバブルやカーネル境界間の待機時間に費やされます。バッチサイズが大きくなると、パイプラインバブルの割合は全体に対する時間比率として減少し、速度差も縮小します。

次に、速度向上効果はエキスパートの分布に依存します。バッチサイズ 8 の場合、実際の専門家分布では megakernel は 1.32 倍高速で動作し、均等ルーティング下では 1.14 倍高速です。実際のリクエストでは同じエキスパートが選択されることが多く、結果としてアクティブなエキスパートのセットは疎になります。これにより MoE（Mixture of Experts）モデル全体の処理量が減り、パイプラインバブルがステップ全体に占める割合が大きくなります。megakernel はまさにこの隙間を埋めて高速化を実現します。一方、均等ルーティングではトークンが多数のエキスパートに分散されるため、MoE の処理量が増え、回復可能なバブルが減ります。したがって、合成された均等ルーティングは実際のトラフィックにおける megakernel の速度向上効果を過小評価することになり、ここではより厳しいケースとして報告しています。

この利点はコンテキスト...