LlamaIndexのチームは、最新のドキュメント解析モデル「OCRia-1」を展開する際、推論のレイテンシが顧客にとって最大の関心事の一つであることに気づきました。一連の最適化の結果、シングルB200 GPU上でのトークン生成速度を約285トークン/秒から約2800トークン/秒へとほぼ10倍に向上させることに成功しました。これにより、ベンチマークの反復速度が上がるだけでなく、エンドユーザーにもより迅速に結果を届けられるようになります。
OCRiaの推論は主にプレフィルとデコードの2段階で構成されます。プレフィルは計算負荷が高く、デコードはメモリ帯域幅に依存します。B200とこのワークロードでは、デコード段階がメモリ帯域幅によって大きく制限されており、GPUは実際の計算よりもモデルの重みや状態の移動に多くの時間を費やしています。そのため、最適化の方向性は、デコードの各ステップで移動するデータ量を減らすこと、ターゲットモデルの1回のフォワードパスでより多くのトークンを生成すること、そして推論エンジンのオーバーヘッドを可能な限り排除することに焦点を当てました。
まず、ワークロードに応じた量子化を採用しました。すべてのパラメータを一律に量子化するのではなく、層ごとの量子化感度に基づいて混合精度を使用します。敏感な層は高精度に保ち、それ以外の層はより積極的に圧縮しました。実際の入力分布を用いたアクティベーション認識量子化キャリブレーションを行い、さまざまな量子化の組み合わせを精度と速度の両面で評価しました。その結果、ビジョンエンコーダは視覚精度を保つために量子化せず、ターゲットタスクでは事実上ロスレスであることを確認しました。重みにはBF16、FP8、NVFP4を組み合わせて使用しています。
次に、投機的デコーディングがスループット向上の鍵となりました。従来の自己回帰モデルは各ステップで1トークンしか生成せず、毎回すべてのモデル重みを読み込む必要があるため、小さなバッチサイズでは計算リソースが大幅に遊休します。チームは小型のドラフトモデル「DFlash」を導入しました。これはターゲットモデルの隠れ状態特徴に基づき、1回のフォワードパスで複数の未来トークンを並行して予測します。ターゲットモデルはその後、1回のフォワードパスでこれらのトークンを並行検証し、受け入れられた結果がターゲットモデルの分布と一致することを数学的に保証します。これにより損失なしの加速が実現します。初期のDFlashの受け入れ率は9トークンでしたが、ローカル畳み込みと候補選択器を追加したDFlash2では12トークンを超えました。この時点で、バッチサイズ8〜12において、ユーザーあたりのスループットは約200トークン/秒から1400トークン/秒以上に向上しています。
さらなる性能向上のため、チームは投機的検証の部分を深く最適化しました。OCRiaのアーキテクチャはゲート付きデルタネットワーク（GDN）を使用しており、全注意層のようにKVキャッシュを保持するのではなく、固定サイズの再帰状態で履歴を要約します。投機的デコーディングでドラフトトークンが拒否されると、従来の方法では完全な状態を復元する必要があり、大量のメモリ転送が発生します。チームはReplaySSM技術を適応させました。再帰状態のチェックポイントと最近の入力バッファを保持し、定期的なフラッシュ時のみ完全なチェックポイントを更新します。拒否が発生した場合は、バッファポインタを調整して受け入れられたプレフィックスを保持するだけで、完全な状態復元を回避できます。さらにカーネル自体を最適化し、上流のTriton実装では層あたり72〜80マイクロ秒かかっていた処理を、カスタムCUDA実装で23〜24マイクロ秒まで短縮しました。これは同じベンチマークで約3.3倍の高速化です。
また、演算融合によって不要なメモリ転送を削減しました。例えば、MLPのSiLU活性化と乗算、およびそのFP8量子化を融合し、ゲート付きRMSNormとFP8量子化も融合しました。さらにNVFP4アクティベーション量子化を活性化を生成するカーネルに統合し、GDN層向けに高速な因果畳み込みを実装しました。SiLUと量子化の融合だけで、測定で2.3〜2.6%のスループット向上が得られました。
CPU側の準備もレイテンシの重要な要因であることが判明しました。推論エンジンは各デコードステップの前に、バッチを記述するメタデータ（シーケンス長、トークンオフセット、キャッシュ位置など）を構築する必要があります。vLLMはKVキャッシュグループごとに層を整理するため、ターゲットモデルとドラフトモデルで4種類のキャッシュレイアウト（24のGDN層、8のターゲット注意層、5のドラフトスライディングウィンドウ層、1のドラフト全注意層）が存在し、グループ化ルールによって38グループが生成され、それぞれがメタデータ構築をトリガーしていました。しかし、多くの情報は共有可能です。チームは、互換性のあるビルダークラスとキャッシュ仕様ごとに共有メタデータを一度だけ構築し、他のグループには独自のブロックテーブルとスロットマッピングを持つ浅いコピーを提供するように変更しました。GDNメタデータについては、グループに依存しない計算をNumPyに移し、ドラフトリクエストがバッチの先頭行を占める特性を利用してマスクを単純化しました。これにより、ステップあたり約290回の小規模なGPU起動をわずか1回に削減しました。
最後に、チームはモデルアーキテクチャに特化したFused Triton CUDAカーネルや、特定の既存カーネルおよびcuDNNコンポーネントの選択など、さらに17の最適化を実装しました。個々のマイクロ最適化は2〜5%のスループット向上にとどまりますが、積み重ねることで大きな効果を発揮します。これらの最適化はOCRiaの効率的な推論スタックの基盤を形成しており、一部の改善（ReplaySSMの適応など）は現在もドラフトPRとしてvLLMコミュニティに貢献されています。