PyTorchで実行される自己回帰型デコーダーモデルのトークン生成性能を最適化する手法として、CUDAストリームのインタリーブ（交互実行）を活用するアプローチが提案されている。この技術は、GPUのアイドル時間を削減し、特に「早期停止（EOS）チェック」に伴うCPU-GPU同期によるブロッキングを回避することで、実行効率を大幅に向上させる。研究対象はHugging FaceのGPT-2モデル（v5.1.0）で、NVIDIA L40S GPUとPyTorch 2.10.0環境下で実験が行われた。 まず、基本的な生成関数は各トークンごとにモデルを再実行するため、O(N²)の計算量となり、メモリ断片化も深刻だった。これを改善するため、KVキャッシュを導入することで計算量をO(N)に削減。さらに、PyTorchの実験的機能「expandable_segments」を有効化することで、メモリ断片化を軽減。さらに、固定サイズのキャッシュ「StaticCache」を採用することで、メモリ管理の負荷を低下させたが、長文生成では計算の無駄が生じるという課題も指摘された。 さらに、モデルコンパイル（torch.compile）を適用することで、計算グラフの最適化が進み、性能が向上。しかし、最も顕著な改善は、CUDAストリームのインタリーブによるものだった。2つのストリームを交互に使い、CPUがトークンN+1の生成をGPUに送信するタイミングで、CPUは前回のトークンNのEOSチェックを非同期に処理できるようにした。これにより、GPUのアイドル時間が解消され、GPU利用効率が安定化。結果として、バッチサイズ32、最大シーケンス長100の条件下で、ベースライン比約5倍の性能向上を達成。 ただし、この手法には注意点がある。CUDAストリームの使用には明示的な同期制御が必要で、データ転送のミスやメモリエラー、データ破損のリスクがある。また、複数ストリームの使用はメモリ予約量の増加や断片化を引き起こす可能性があるため、実運用ではベンチマークによる検証が不可欠。 結論として、PyTorchネイティブな推論環境においては、KVキャッシュ、モデルコンパイル、そしてCUDAストリームのインタリーブを組み合わせることで、自己回帰生成の性能を劇的に改善できる。ただし、環境依存性が高いため、自社のワークロードに合わせた検証が必須である。