🤖 AI

Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection

本論文は、ピボットベースの切り捨てと選択を活用して既存の GPU カーネルと比較して最大 1.4 倍のスループット向上と 50% のメモリ削減を実現する大語彙向けの高性能かつ決定論的な Top-k および Top-p サンプリングアルゴリズム「Qrita」を導入し、それが vLLM のデフォルトサンプリャとして採用されるに至ったことを述べる。

原著者： Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

公開日 2026-05-27

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは100,000種類の異なる食材（語彙）を備えたパントリーを持つ巨大なキッチン（大規模言語モデル）を運営するシェフだと想像してください。料理（単語を生成）するたびに、その巨大なパントリーから最高の食材をいくつか選び、料理が美味しくても飽きないものになるよう確保する必要があります。

AIの世界では、この選び出しのプロセスをTop-k（上位k個の最高の食材を選ぶ）とTop-p（特定の「風味スコア」に達する最小限の食材グループを選ぶ）と呼びます。

問題：遅く、散らかったキッチン

現在、ほとんどのキッチンは、100,000個の食材を巨大なカウンターにすべて出し、最高のものから最悪のものまで「風味スコア」でソートし、その後上位のものを取り出すという方法で処理しています。

問題点: 100,000個のアイテムをソートするのは遅く、散らかりやすいです。大量のカウンター空間（メモリ）を占有し、シェフの手が疲れます（計算オーバーヘッド）。
代替案: 一部のシェフは、たぶん良いだろうと推測して、ランダムに handful の食材を掴もうとします。しかし、これはリスクがあります。時には最高の食材を見逃したり、毎回料理するたびに出来上がりが異なる（非確定的）ことがあり、レシピを正確に再現する必要がある場合は問題です。

解決策：Qrita（賢いソウシェフ）

この論文は、食材を選ぶための新しい高速な方法Qritaを紹介しています。Qritaは、散らかったソート作業を完全にスキップするために2つの巧妙なトリックを使う、超賢いソウシェフだと考えてください。

トリック1: 「ガウス・シグマ・トリンケーション」（ノイズフィルター）

100,000個の食材があるパントリーの中で、99,000個は単なる「ノイズ」（塩、砂糖、小麦粉など、すべてほぼ同じような味気ないもの）だと想像してください。残りの数百個だけが「スター食材」（トリュフやサフランなど）です。

Qritaは、すべての瓶を調べる代わりに、素早い嗅ぎテストを行います。パントリーの平均的な味と「辛味」（標準偏差）を計算します。そして、あるラインを引きます。「このラインより下は単なるノイズだ。見る必要はない。」

結果: パントリーの99%を即座に捨て、興味深い瓶を200個程度だけ残します。これは単一の、雷のような速さのパスで行われます。

トリック2: 「四進ピボット探索」（4分割）

これで、シェフは200個の興味深い食材の小さな山を持っています。それでも、正確な上位50個を見つける必要があります。

古い方法: 1つずつ確認する（遅すぎる）か、山を半分に分割する（二分探索）。
Qritaの方法: 山を半分に分割する代わりに、Qritaは一度に4つのセクションに分割します。「最高の食材は、第1四半期、第2四半期、第3四半期、それとも第4四半期にあるか？」と問います。
ボーナス: また、重複する食材のための特別なルールもあります。「サフラン」の瓶が3つあり、すべて正確に同じスコアの場合、Qritaはレシピを毎回同一にするために何個保持すべきかを正確に知っています（確定的）。これにより、シェフが同一の瓶の間で決着をつけようとして無限ループに陥るのを防ぎます。

なぜこれが重要か（結果）

著者は、Triton（グラフィックカード/GPUをプログラミングするための言語）という専用ツールを使ってQritaを構築し、vLLMやSGLangなどの主要なAIエンジンが現在使用している最良の方法と比較してテストしました。

速度: 実際のサービスシナリオでは最大1.4倍、生速度テストでは最大2倍高速です。
メモリ: 100,000個のアイテムのソート済みリスト全体を保存する必要がないため、半分のメモリしか使用しません。
精度: 推測する一部の高速な方法とは異なり、Qritaは遅くても完璧なソート方法と完全に同じ結果を提供します。出力を変えることなく、それをはるかに高速に見つけ出します。

結論

Qritaは、世界中のすべてのスパイスを手作業でソートするシェフから、退屈なものを即座に無視し、興味深いものを一度に4つの山に分割し、重複を完璧に処理するスマートなアシスタントへアップグレードするようなものです。これにより、回答の品質を変えることなく、AI生成をより高速かつ効率的にします。

注：この論文では、QritaがAIモデルを実行するための人気ツールであるvLLMのGPUパスにおけるデフォルト方法となり、コードは他の人が使用できるよう公開されていると述べています。

技術概要：Qrita - ピボットベースの切り捨てと選択を用いた高性能 Top-k および Top-p

問題提起

大規模言語モデル（LLM）の巨大な語彙（多くの場合 10 万トークンを超える）に対する Top-k および Top-p サンプリングアルゴリズムの効率的な実装は、GPU 上で依然として大きな課題です。既存のアプローチは、計算効率と出力の正確性の間でトレードオフに直面しています：

ソートベースの手法： vLLM などのエンジンで使用されるこれらの手法は、正確で決定論的な結果を提供しますが、全語彙のソートとストライドメモリアクセスにより、高い計算およびメモリオーバーヘッドを伴います。
確率的な手法： FlashInfer や SGLang などのライブラリは、ソートを回避するために棄却サンプリングを使用し、高いスループットを提供します。しかし、これらは完全な候補セットではなく、単一のサンプリングされたトークンのみを返すため、スペキュレーティブデコーディングの検証や強化学習など、決定論的な候補セットを必要とするタスクとは互換性がありません。
統計的カットオフ手法： Top-nσや Min-p などのアプローチは、統計的閾値（平均や標準偏差など）を使用してソートを排除します。効率的である一方で、サンプリング分布と出力を変更するため、生成の品質や一貫性に関する議論を招いています。

Top-k と Top-p の正確なセマンティクスを犠牲にすることなく、メモリ効率と厳密な決定論的動作（完全な候補セットの返却）を実現するアルゴリズムに対する重要なギャップが存在します。

手法

著者は、ピボットベースの切り捨てと選択に基づく GPU 向けの高性能 Top-k および Top-p アルゴリズム「Qrita」を提案します。Qrita は、元々グラフニューラルネットワーク向けに設計された RTop-k のピボットベース戦略を拡張し、決定論性を保証しながら大規模な LLM 語彙を処理できるようにします。この手法は、以下の 2 つの主要な技術に依存しています：

1. ガウスベースの $\sigma$ -切り捨て

LLM のロジットは通常、高確率トークンが「ノイズ」の背景に対して外れ値を形成する準ガウス分布に従うことを認識し、Qrita は事前フィルタリングステップを採用します。

メカニズム： アルゴリズムはロジット分布を $Z \sim \mathcal{N}(\mu, \sigma^2)$ として近似します。サンプルブロックから平均（ $\mu$ ）と標準偏差（ $\sigma$ ）を計算し、切り捨て閾値 $t = \mu + \delta \cdot \sigma$ を算出します。
ルックアップテーブル： 係数 $\delta$ は、正規分布の分位点関数から導出された事前計算されたルックアップテーブルによって決定され、ターゲットの $k$ または $p$ の値に合わせて調整されます。
結果： このステップにより、語彙の巨大な「ノイズ」領域が単一パスで切り捨てられます。切り捨てられたセットのサイズ $|Z_t|$ がターゲット $k$ を超える場合（「切り捨てヒット」）、検索はこのより小さな部分集合のみで続行されます。そうでない場合は、全ロジットセットにフォールバックします。これにより、後続のピボット検索のための検索空間が大幅に削減されます。

2. 重複処理を伴う四分探索

正確なピボット値 $\tau$ （ちょうど $k$ 個の要素が $\ge \tau$ である値）を効率的かつ決定論的に見つけるために：

四分探索： 二分探索の代わりに、Qrita は各反復で現在の探索範囲の 25 パーセンタイル、50 パーセンタイル、75 パーセンタイルの 3 つのピボットを評価します。これにより、二分探索と比較して探索反復数が半分に削減されます。
重複処理： 標準的なピボット探索は、重複したロジット値が存在する場合（低精度または長尾分布で一般的）、失敗したり非決定論的になったりする可能性があります。Qrita は、ピボットより大きい最小ロジット値 $z_{mn}$ とその出現回数 $N_{zmn}$ を追跡します。
終了条件： 探索は、 $\tau_n$ に対して、 $\tau_n$ より厳密に大きい要素の数が $k$ 未満であるが、最小重複値 $z_{mn}$ 以上である要素の数が $k$ 以上であるようなピボット $\tau_n$ が見つかったときに終了します。これにより、Top-k 制約を満たすために境界値のどの程度の重複インスタンスを保持すべきかを正確に決定でき、再現性のある決定論的出力を保証します。

実装詳細

フレームワーク： GPU プログラミング言語であるTritonを使用して実装されています。
オートチューニング： 「ティアード Triton ブロックサイズオートチューン」戦略を採用しています。ベクトルサイズは全ロジットセットと切り捨てられたセットの間で劇的に変動するため、コンパイラはハードウェア利用率を最大化するために最適なブロックサイズ（例：全セットにはより大きなブロック、切り捨てられたセットにはより小さなブロック）を選択できるようにしています。
パイプライン： アルゴリズムはバッチを並列に処理し、ガウス切り捨てを適用した後、四分探索を行い、最後に重複処理マスクを生成して最終出力を生成します。

主要な貢献

決定論的かつ高性能なサンプリング： Qrita は、FlashInfer/SGLang とは異なり、完全な Top-k/Top-p 候補セットを提供しつつ、ソートベースアルゴリズムの正確なセマンティクスを維持するため、スペキュレーティブデコーディングや RL に適しています。
検索空間の削減： ロジットの統計的性質を活用することで、ガウス $\sigma$ -切り捨てが検索空間を劇的に削減し、全語彙のソートを回避します。
堅牢なピボット探索： 明示的な重複処理を組み合わせた四分探索の導入により、決定論的動作を保証し反復回数を削減し、大規模語彙における標準的なピボットベース手法の非効率性を解決します。
統合： Qrita は、vLLMの GPU 実行パスにおけるデフォルトの Top-k および Top-p サンプリング器として実装されています。

結果

著者は、NVIDIA H100、RTX4090、AMD MI355X GPU 上で、Llama、Qwen、GPT-OSS、Gemma などの各種モデルを使用して、Qrita を素の PyTorch ソート、vLLM のデフォルトサンプリング器、FlashInfer、SGLang などのベースラインと比較評価しました。

スループット： Qrita は H100 上で vLLM ベースラインに対して最大1.4 $\times$ のエンドツーエンドサービススループット向上を達成し、一部の構成では RTX4090 上で最大2 $\times$ の向上を示しました。ソートベースのベースラインを大幅に上回る性能を発揮しつつ、ソートベースアルゴリズムの出力と一致しています。
メモリ使用量： Qrita は、ソートに必要とされる大規模な中間バッファの割り当てを回避することにより、ソートベースの手法と比較してメモリ使用量を約**50%**削減します。
アブレーション研究：
- ガウス切り捨て： 最大の速度向上を提供し、全ロジットセットを検索する場合と比較してレイテンシを約 74% 削減します。
- 四分探索： 「切り捨てヒット」の場合、検索空間が減少するため二分探索よりもわずかに遅いものの、全語彙を検索しなければならないフォールバックの場合、大幅なレイテンシ削減（約 2.3ms）を提供します。
- 重複処理： オーバーヘッドは最小限（レイテンシ増加約 3.5%）ですが、正確性と決定論性にとって不可欠です。
- オートチューニング： 単一の固定ブロックサイズを使用する場合と比較して、実行レイテンシをほぼ 30% 削減します。

意義と主張

本論文は、Qrita が大規模語彙デコーディングにおけるメモリ効率と厳密な決定論的動作の間のギャップを埋めると主張しています。全語彙のソートを回避し、統計的事前フィルタリングを伴うピボットベースの選択を活用することで、Qrita は完全な候補セットを必要とする高度なタスク（例：スペキュレーティブデコーディング）と互換性のある、LLM デコーディングのための高性能ソリューションを提供すると述べています。著者は、Qrita を Top-k/Top-p の統計的セマンティクスを置き換えるものではなく、ヒューリスティックな統計的カットオフ手法や確率的単一トークンサンプリング器とは異なり、それらを正確に保持する効率的な実装として位置付けています。vLLM へのデフォルトサンプリング器としての統合は、生産環境におけるその実用性を強調しています。