Cohereは本日、North Mini Code向けのサービスエンジンを発表しました。このエンジンはデコード処理をメガカーネルとして統一的に実行し、H100をBF16で使用した際にvLLMと比較してエンドツーエンドで1.25〜1.41倍のスループットを実現します。コードはGitHubで公開されています。
従来のLLM推論スタックは、各フォワードパスを一連の小さなカーネルとして扱います。QKVを起動して待ち、アテンションを起動して待ち、MoEを起動して待つ。個々のカーネルは最適化されていても、カーネル間の待ち時間により、特にバッチサイズが小さいデコード時にはGPUの帯域を生かしきれません。自回归デコードは、メモリ帯域ではなく計算量ではなくメモリ帯域によって制限される小さなバッチで特に顕著です。
North Mini Codeは30Bモデルで、トークンごとにアクティブなパラメータは3.3Bです。BF16では各デコードステップに約6.6GBの重みと、8Kコンテキストで約0.5GBのKVキャッシュをHBMから読み出します。H100の帯域3.35TB/sで単純計算すると理論限界は約470tok/sです。vLLMは185tok/sで、これは理論限界の39%にすぎません。
メガカーネルでは、フォワードパス全体を1つの常駐カーネルとして実行します。GPUの各SMに1スレッドブロックが常駐し、ホストが用意したタスクリストを読み込んで処理します。データ依存関係はグローバルメモリ上のカウンタによって明示され、カーネル境界による全Grid同期を回避します。これにより、スケジューリング単位は「オペレーション」ではなく「タイル」になり、同期も依存先のSMだけに限定できます。
高速化の要因としては、起動/同期オーバーヘッドの削減、波形量子化の解消、誤った依存関係の排除、重みのプリフェッチが挙げられます。North Mini Codeは並列Transformer層を採用しており、アテンションとMoEの出力を待ち合わせずに済むため、空いたSMへのバックフィルがしやすくなっています。
バッチサイズ1では292tok/s、理論限界の62%、vLLMの1.58倍に達しました。この利点はバッチサイズを変えても、コンテキスト長が256Kまで伸びても維持され、精度の低下も測定されていません。連続バッチ処理、ページ化アテンション、可変長シーケンス、OpenAI互換APIとツール呼び出しを備えた本格的なサービスエンジンとして動作します。実装はCUDAファイル1つで完結し、コンパイラや特別な抽象化は不要。既存カーネルを移植するためのレシピも紹介しています。
設計ではHazy Researchの「Look Ma, No Bubbles!」の3つのアイデア、タスクインタプリタ、カウンタ同期、タスク間オーバーラップを活用しています。一方で、Cohereはbatch size 1でもテンソルコアのwgmma命令を使い、共有メモリページングによるプリフェッチは採用せず、代わりに各オペコードが専用のwarp-specializedパイプラインを持ち、同種GEMM同士のパイプライン段を重ねたり、GEMM内部で入力待ちの間に重みをロードすることで帯域を活用しています。