KTransformersがSGLangの混合推理を高速化

背景：スパースMoEモデルの混合推理

DeepSeek-V3などの現代的なMixture-of-Experts（MoE）言語モデルは、数千億のパラメータを含んでいますが、各トークンはエキスパートのごく一部のみをアクティベートします。このスパースアクティベーションパターンにより、MoEモデルはCPU/GPU混合推理に非常に適しています：スパースにアクティベートされたエキスパートは大容量メモリを持つCPU上で効率的に実行でき、密な計算コンポーネントは帯域幅とスループットがより高いGPU上で実行できます。

KTransformers：MoEモデルCPU/GPU混合推理の潜在能力を解放

混合推理における課題を解決するため、清華大学のMadSysとApproaching.AIはKTransformersプロジェクトを開発し、CPU/GPU協調推理をより効率的にする一連の最適化を提供しています。その改善は主に3つの側面に分かれています：

1. AMX最適化CPUカーネル

KTransformersは、Intel AMX最適化カーネルとキャッシュ階層構造向けのメモリレイアウトを通じてCPU計算を再設計しました。単一のXeonソケット上で、AMX最適化カーネルの持続スループットは21.3 TFLOPSに達し、PyTorchネイティブ実装より3.9倍高速です。

2. 効率的なデバイス協調

NUMA認識テンソル並列とCUDAグラフサポートのスケジューリングを導入することで、KTransformersはCPUとGPU間の協調コストを大幅に削減しました。NUMA認識テンソル並列は、エキスパートの重みフラグメントを各NUMAノードのローカルメモリに配置し、高価なクロスNUMAメモリトラフィックを回避することで、最大63%のデコーディングスループット向上を実現しました。

3. エキスパート遅延メカニズム

特定のエキスパートの実行を遅延させることで、KTransformersはCPUエキスパート計算とGPUアテンション処理を重複して実行できるようにし、デバイスの並行利用率を向上させ、精度の変化を0.5%以内に抑えながら最大1.45倍のデコーディングスループット向上を実現しました。

KTransformersをSGLangに統合

SGLangは現在KTransformersをバックエンドライブラリとして統合し、MoEモデルのCPU/GPU混合推理をより効率的にしています。GPUテンソル並列とCPU/GPUハイブリッドエキスパート並列を組み合わせ、ヘテロジニアスデバイス上での推論をサポートしています。

インストールガイド

SGLangのKTransformers混合推理を使用するには、SGLangとKTransformers CPUカーネル（kt-kernel）をインストールする必要があります。システムが以下の要件を満たしていることを確認してください：CUDAバージョン12.1以上、Linux x86_64オペレーティングシステム、gcc、g++ >= 11、CMake >= 3.25、Python 3.11。

ベンチマーク結果（プレビュー）

単一GPU+CPU構成では、KTransformersはすべてのプロンプト長においてベースラインを上回り、AMX最適化CPUカーネルのおかげで最大20倍の高速化を達成しました。デコーディング段階でも、KTransformersは優れた性能を示し、主にCPU/GPU協調オーバーヘッドの削減により、最大4倍の高速化を実現しました。