SGLang 高度CUDA Graphテクノロジー解析

TL;DR

CUDA Graphはカーネル起動のオーバーヘッドを排除することを目的としているが、実際の推論エンジンでこのメリットを得るには、互換性・起動時間・メモリを犠牲にすることなく、できるだけ多くのワークロードをキャプチャする必要がある。SGLangでは、汎用的なrunner/backendインターフェースを中心にCUDA Graphサポートをリファクタリングし、異なるキャプチャ戦略を実行パス間で再利用できるようにした。

Breakable CUDA Graph（BCG）はSGLangが独自に開発したサービス技術であり、プリフィル段階においてデフォルトの手法となっている。BCGのコード量はtorch.compileの分割バックエンドの4分の1にとどまり、グラフ構築速度は3.8〜5.2倍速く、複雑な機能への対応も広い。プリフィル単独のテストでは、BCGはeager実行比1.70倍、フルキャプチャでは1.93倍の高速化を達成している。

背景

現代のLLM推論エンジンでは、1回の推論ステップに大量のGPU操作が含まれる。CPUからこれらの操作を繰り返し起動することは、特にレイテンシが重視される場面で大きなオーバーヘッドをもたらす。CUDA Graphは、GPUの処理を記録して再生することで、起動オーバーヘッドを大幅に削減できる。

SGLangにおけるCUDA Graph：Runner/Backend分離と柔軟な組み合わせ

リファクタリング後、RunnerはキャプチャおよびリプレイHに必要な実行状態を管理し、Backendはキャプチャ方式（フルグラフ・分割・コンパイラ生成）を決定する。プリフィルとデコードは独立したrunnerを使用し、投機的デコードによってさらに拡張される。

Breakable CUDA Graph

BCGはキャプチャ処理中に明示的なeagerブレークポイントを挿入することで、互換性のない操作をグラフセグメント間で実行できるようにし、コンパイラを使わずにセグメント分割を実現する。

プリフィル向けフルCUDA Graph

リクエストパディングにより、SGLangは動的なプリフィルワークロードの完全なキャプチャを実現した。

CUDA Graphのメモリ使用量

SGLangは、形状やグラフセグメントをまたいだメモリ再利用により、CUDA Graphがもたらすメモリオーバーヘッドを効果的に抑制している。