メモリの壁を打破する:標準的なGPUで3,000 tokens/sの推論を実現する
現在のLLM推論の状況は、多くの場合、集計スループット(数百人のユーザーに対してサーバーが生成できる合計トークン数)によって測定されます。これはクラウドプロバイダーにとっては重要な指標ですが、次世代のAI、すなわち自律型エージェントにとっては誤った指標です。エージェントのワークフローは本質的に逐次的です。エージェントは、各ステップが前のステップに依存するループの中で、調査、計画、編集、テストを行わなければなりません。この領域では、唯一重要な指標はsingle-request decode speed(単一リクエストのデコード速度)です。
Kog AIは、Kog Inference Engine (KIE)のテックプレビューをリリースし、標準的なデータセンター用GPUを使用して、2Bモデルにおいてリクエストあたり3,000 tokens/s (8× AMD MI300X上)および2,100 tokens/s (8× NVIDIA H200上)の速度を達成できることを実証しました。モデルアーキテクチャ、ランタイム、および低レベルのGPUカーネルを単一の共同設計システムとして扱うことで、Kogは汎用GPUと専用の推論用シリコンの間のギャップを埋めようとしています。
ボトルネック:計算量ではなくメモリ帯域幅
単一リクエストのデコード(バッチサイズ1)において、主な制約は1秒あたりの浮動小数点演算数(FLOPS)ではなく、メモリ帯域幅です。自己回帰的なデコードにおいて、GPUは生成される各トークンごとに、すべての有効なモデルの重みをHigh Bandwidth Memory (HBM)から計算プロセッサへ移動させなければなりません。
これにより、重みをいかに速くストリーミングできるかにほぼ完全に制限される、低い演算強度のシナリオが生まれます。H200のような最新のGPUは膨大な理論的帯域幅を持っていますが、ほとんどの推論スタックはそれを効果的に活用できていません。8× H200ノードにおけるFP16の2Bモデルの理論的な「光速」上限は、およそ7,700 tokens/sです。ほとんどのフレームワークがこの値を大幅に下回って動作しているという事実は、ボトルネックがハードウェアではなく、ソフトウェアスタックにあることを示唆しています。
マイクロ秒が失われる場所
3,000 tokens/sという目標速度では、1トークンあたりの合計時間予算はわずか333マイクロ秒です。標準的な推論スタック(例:vLLMやTensorRT-LLM)では、この予算は、高いバッチサイズでは無視できるものの、単一リクエストのレイテンシに対しては致命的なオーバーヘッドによって消費されます。
- Kernel Launch Overhead: カーネルの起動とクリーンアップには約4.5µsかかることがあります。Transformerレイヤーが10個のカーネルを必要とする場合、25個のレイヤーがあれば、実際の計算が行われる前に、1トークンあたり1,100µs以上のオーバーヘッドが発生します。
- CPU Scheduling: ホスト側のロジックとGPU-CPU通信は、サンプリングと実行における遅延を導入します。
- Synchronization: GPU全体での同期とGPU間コレクティブ(テンソル並列化におけるAllReduceなど)は、操作ごとに数マイクロ秒を追加します。
- Topology Ignorance: GPUチップレットアーキテクチャ内の非一様メモリ・アクセス(NUMA)効果は、レイテンシの偏り(skew)を導入する可能性があります。
Kogの解決策:共同設計とモノカーネル
これらの失われたマイクロ秒を取り戻すために、KogはクリティカルパスにおいてPyTorchやTritonのような高レベルフレームワークを避け、代わりに手作業で作成されたGPUコード(CUDA/PTXおよびHIP/CDNA ISA)を選択しています。彼らのアプローチは、3つの核心的な革新に基づいています。
1. The Monokernel Runtime
カーネルのシーケンスを実行する代わりに、Kogはpersistent monokernelを使用します。これは、MatMul、attention、normalization、およびsamplingを行う、デコードパス全体を処理する単一のGPU常駐プログラムです。これにより、CPUに制御を戻すことなく、カーネル起動オーバーヘッドとホスト側のスケジューリングを完全に排除します。
2. KCCL (Kog Collective Communication Layer)
標準的な通信ライブラリは、スループットではなくマイクロ秒単位のレイテンシを最適化しています。KCCLは、AllReduceのレイテンシを3µs未満に抑えるためにアセンブリレベルでチューニングされたカスタムレイヤーです。これにより、GPU間通信がメモリ・ストリーミング・プロセスを停止させないようにします。
3. Laneformer Architecture
Kogは、Laneformerアーキテクチャを導入しました。これは、**Delayed Tensor Parallelism (DTP)**を利用しています。DTPはデコーディング・プロセスの依存関係構造を変更し、デバイス間通信を計算とオーバーラップさせることで、クリティカルパスをブロックしないようにします。
スケーリングをFrontier MoEモデルへ
3,000 tokens/sというベンチマークは2Bモデルを使用していますが、Kogは、同じ原理が大規模なMixture-of-Experts (MoE)モデルにも適用可能であると主張しています。MoEはトークンごとに全パラメータのわずかな一部のみを活性化させるため、帯域幅のボトルネックは全パラメータではなく、activeパラメータに依存します。
Kogは、Frontierモデル(DeepSeekやQwenなど)に対して、彼らのモノカーネルとKCCLの最適化を適用することで、標準的なデータセンター用GPUにおいて、出力速度を1,000–5,000 tokens/sの範囲に押し上げることができると見積もっています。特に、次世代ハードウェア(NVIDIA RubinやAMD MI450など)がHBM帯域幅を増大させるにつれて、その効果は顕著になります。
批判的な視点とトレードオフ
印象的なベンチマークとは裏腹に、コミュニティからは、この技術の実際の応用に関するいくつかの重要な点が指摘されています。
- Model Utility vs. Speed: 一部のユーザーは、モデルの知能がFrontierモデルに及ばない場合、高いトークン・スループットは無意味であると指摘しています。あるコメントでは、「3.4k tok/sの純粋なナンセンス」は、実用的な製品ではありません、と述べられています。
- Hardware Definitions: 「標準的なGPU」とは何を指すのかについて議論があります。批判的な意見では、8× H200ノードはエリート・データセンター用ハードウェアであり、アクセス可能な「標準的」なGPUではない、と指摘されています。
- Scaling Concerns: 懐疑的な人々は、2Bモデルで見られた線形的な利得が、メモリ圧迫や通信オーバーヘッドが異なるスケールで変化する30B+モデルにどのように移行するかを疑問視しています。
最終的に、Kogの成果は、AIインフラストラクチャの風景における根本的な変化を強調しています。チャットボットから自律型エージェントへ移行するにつれ、業界は、集計スループットの最大化から、単一の逐次的な思考の連鎖のレイテンシを最小化することへと、方向を転換しなければなりません。