【ITニュース解説】From a Rust Inference Engine to Heterogeneous Serving

大規模言語モデル（LLM）の活用が広がる中、その推論、つまり学習済みのモデルを使って新しいテキストを生成する過程をいかに効率的に行うかが、システムエンジニアリングの重要な課題になっている。この記事では、筆者がこの課題に取り組むため、Rustというプログラミング言語を使ってmacOS上で動作する「mini-vllm-rs」というLLM推論エンジンを開発した経験から得られた知見を紹介する。このプロジェクトは、LLMの推論を高速化し、より多くのリクエストを効率的に処理するための技術的な工夫と、その実現過程で直面した具体的な課題を明らかにしている。

mini-vllm-rsは、最新のLLM推論技術を多数取り入れている。例えば、「連続バッチ処理」という技術は、複数のリクエストを同時に処理することで全体の処理能力を高める。また、「ページングKVキャッシュ」は、モデルが過去の情報を記憶しておくためのメモリ領域（KVキャッシュ）を効率的に管理し、無駄なく再利用できるようにする。「プレフィックスキャッシュ」は、共通の入力部分を持つリクエストの計算を省略し、処理を高速化する。「推測デコーディング」は、小さなモデルを使って先行してテキストを生成し、その予測を大きなモデルで検証することで、生成速度を向上させる。さらに、このエンジンは、macOSのCPUとMetal GPUという異なるハードウェアを同時に活用し、それぞれの得意な処理を分担させることで、全体の性能を高めようと試みている。

このような高度な機能を実装する過程で、筆者はいくつかの予期せぬ課題に直面した。その一つが、連続バッチ処理の導入によって処理速度が低下するという現象である。具体的には、Apple GPU上で同じ量（各512トークン）のテキストを生成する2つのリクエストを連続して処理すると、毎秒16.8トークンの速度が出た。しかし、これらを連続バッチ処理でまとめて実行すると、速度は毎秒8.7トークンにまで落ち込んでしまった。この原因は、バッチ処理によってデータ（行列）の形状が変わり、バックエンドの数値計算フレームワークであるCandleが、より小規模な計算に適した「GEMV（Generalized Matrix-Vector multiplication）」カーネルから、大規模な並列計算に適した「GEMM（Generalized Matrix-Matrix multiplication）」カーネルに切り替わってしまったためだった。このGEMMカーネルは、小さいバッチサイズではGEMVよりも効率が悪く、結果として性能が低下したのである。筆者はこの問題を解決するため、小さなバッチサイズの場合にはGEMVカーネルに戻るようなフォールバック機構を実装し、両者の性能が逆転するしきい値をベンチマークで特定して、最適な切り替えポイントを設定した。

この経験は、推論エンジンの設計において、スケジューリング（リクエストの処理順序や方法）、モデルの実行、そして各リクエストの状態管理という三つの要素を明確に分離することの重要性を浮き彫りにした。連続バッチ処理では、各リクエストのKV状態をモデルの実行とは独立して管理する必要があり、さらにプレフィックスキャッシュを導入すると、いつそのメモリを再利用したり解放したりできるかが複雑になる。しかし、これらの要素を最初から分離して設計したことで、後からCPUとGPUを並行して動作させるワーカーを追加する際に、大規模な設計変更をすることなくスムーズに対応できた。また、開発過程でAI（人工知能）をコーディングの補助として活用したが、その際にはAIが生成したコードを盲目的に採用するのではなく、厳密な設計レビューと実験を重ねて、本当に効果的な変更だけを採用することが重要であると筆者は述べている。

もう一つの重要な課題は、「なぜGPUだけでモデル全体を実行できるのに、CPUとGPUで仕事を分ける必要があるのか」という点だった。筆者のプロジェクトでは、Apple SiliconのCPUとGPUが共有メモリを使っているため、デバイス間のデータ転送が高速であるという利点があった。しかし、Candleのようなフレームワークがデバイスごとにメモリを管理するモデルを採用していると、CPUとGPUの間で状態を共有するのが複雑になる。さらに、リクエストの同時実行数が少ない場合、両方のデバイスを常に最大限に活用し続けるのが難しく、結局、CPUとGPUで仕事を分割する実装は行われなかった。

一方で、より大規模なシステムでは、この「異種混合推論」と呼ばれるアプローチが積極的に探求されている。例えば、NVIDIAは特定のGPUと別の種類の高速な処理装置を組み合わせて推論を行う構成を発表しているし、AWSも特定の種類のアクセラレータ（Trainium）でモデルの最初の部分（プレフィル）を処理し、別のアクセラレータ（Cerebras）で残りの部分（デコード）を行うという方法を導入している。これは、それぞれのハードウェアの得意な処理に特化させることで、全体の効率を高めようという考え方に基づいている。

専門化されたハードウェアの利用は魅力的だが、そこには「通信コスト」という課題が伴う。異なるデバイス間でデータをやり取りするには時間がかかり、この通信コストが、専門化によって得られる計算速度の向上を上回ってしまう可能性がある。筆者は具体的な例を挙げてこの点を解説している。例えば、2ギガバイトのKVキャッシュを毎秒25ギガバイトの速度で転送する場合、約86ミリ秒の時間がかかる。もし計算と通信を同時に行えないと仮定すると、128トークンを生成する出力の場合、専門化されたアクセラレータによるデコードがGPU単独のデコードよりも1トークンあたり0.67ミリ秒以上速くないと、このデータ転送にかかるコストを回収できない。512トークンのような長い出力の場合、このしきい値は1トークンあたり0.17ミリ秒に低下する。このことから、専門アクセラレータは、短い回答よりも長い回答を生成する際に、より大きなメリットを発揮する可能性があることがわかる。

記事では、プレフィルとデコードの分割、アテンションとフィードフォワードネットワークの分割、ターゲットモデルとドラフトモデルの分離といった様々な分割方法を比較検討している。そして、なぜ本番環境では、単一のリクエストの処理速度が必ずしも最速ではないパスを選ぶことがあるのか、という問いにも答えている。それは、KVキャッシュの容量を増やしたり、より大きなバッチサイズを扱ったり、あるいはCPUやGPUなどのハードウェアプールをそれぞれ独立して拡張できるようにすることで、個々のリクエストの完了時間が多少遅くなったとしても、システム全体で処理できるリクエストの総数（サービス容量）を大幅に向上させることができるためである。

このプロジェクトとそこから得られた知見は、LLM推論エンジンの開発が、単にコードを書くだけでなく、ハードウェアの特性、ソフトウェアのアーキテクチャ、そしてシステム全体の効率性という多角的な視点から、様々なトレードオフを慎重に検討する複雑なエンジニアリング課題であることを示している。特にシステムエンジニアを目指す初心者にとって、これらの具体的な課題と解決策の検討は、将来のシステム設計において非常に貴重な学びとなるだろう。