推論はもはや logits の「運搬作業」ではない：FlashSampling でデコード速度が最大 19% 向上

解説：大規模言語モデル（LLM）が 1 語を生成するたびに、数万語に及ぶ確率分布（logits）を算出し、その中からサンプリングを行う必要があります。この一連の中間データは大量のメモリ帯域を消費するにもかかわらず、生成後即座に破棄されてしまいます。プリンストン大学とドイツのミュンヘン・ルートヴィヒ・マクシミリアン大学（LMU Munich）の共同研究チームは、サンプリング処理を行列演算に直接「統合」する「FlashSampling」を提案しました。これにより logits テンソルを明示的に生成することなく、デコード速度を最大19%向上させることに成功。かつ、数学的には既存のサンプリング手法と完全に等価であることが証明されています。

問題の核心：サンプリングのたびに「無駄な運搬」が発生していた

LLM によるデコードの各ステップでは、主に以下の 3 つの工程が実行されます。

1. 語彙全体（10 万語以上の場合も）に対する行列演算を行い、logits を算出する。
2. Softmax 関数を用いて確率分布に変換する。
3. 確率分布に基づき、1 語をサンプリングする。

この 3 段階のプロセスにおいて、logits テンソルは「使い捨ての通過点」に過ぎません。計算が完了すれば直ちに不要となるにもかかわらず、高帯域幅メモリ（HBM）への書き込み・読み出しという形で多大なリソースを占有しています。モデルの語彙数が増大し、バッチサイズが大きくなるにつれ、このオーバーヘッドは無視できないものとなっています。

FlashSampling：サンプリングを行列演算に統合し、logits をメモリに書き込まない

中核的な洞察は、サンプリングの本質が「加重された乱数の最大値取得」であるという点です。これは「Gumbel-Max トリック」を用いて書き換えることができ、各スコアにガンベルノイズを加えた上で argmax（最大値のインデックス）を取得すれば、確率に基づくサンプリングと数学的に等価になります。

FlashSampling は、このノイズ注入の処理を行列演算（matmul）のエピローグ（後処理段階）に直接組み込むことで、サンプリングと行列演算を単一の CUDA カーネル内で完結させます。

• logits テンソルはHBM へ一度も書き込まれることなく、チップ上キャッシュ（SRAM）内で即座に生成・破棄される。
• 転送されるのは logits ベクトル全体ではなく、1 つのトークンインデックスのみとなる。
• top-k や nucleus sampling といった一般的な戦略も、モデル修正なしにサポート可能。

データが示す成果：H100 で最大 19% の高速化、1.2T パラメータ級の超大規模モデルでも恩恵

複数 GPU を用いるシナリオでは、FlashSampling により AllReduce 通信量も削減されます。従来は完全な logits の同期が必要でしたが、現在は次元削減済みの情報のみを同期すれば良いため、通信オーバーヘッドが劇的に低下します。

本技術はすでに「vLLM」へ統合済みです。数学問題の生成など実際のタスクにおいても品質の低下は確認されておらず、コードは GitHub で公開されています（https://github.com/FlashSampling/FlashSampling）。

なぜ注目すべきなのか

現在、大規模モデルの推論最適化は、主にKV Cache の圧縮（Google の TurboQuant や NVIDIA の KVTC など）や投機的デコード（Speculative Decoding）に集中しています。FlashSampling は、これまで見過ごされてきた「隠れたコスト」、すなわちサンプリングそのものがもたらすメモリ負荷に光を当てました。

これは一切の性能劣化を伴わない純粋なシステム最適化です。モデル重みの変更も、出力品質の犠牲も、追加ハードウェアの導入も不要で、プラグ＆プレイで導入可能です。本番環境で数千枚の GPU を稼働させている推論クラスター而言えば、わずか 10% のスループット向上でも、実質的な計算コストの削減に直結します。

サンプリングはこれまで推論処理のボトルネックにおける主役とは見なされてきませんでした。しかし FlashSampling は、大規模モデルシステムにおける最適化の余地は、往々にして「当然のこと」と思われている工程の裏側に潜んでいるのだと私たちに教えてくれます。

出典：arXiv:2603.15854、プリンストン大学 / LMU ミュンヘン | 2026 年 3 月 25 日