論文／ブログ「Defeating Nondeterminism in LLM Inference」の技術的なポイント

 

LLM 推論における非決定性を克服する

Thinking Machines Lab（以下 TML）が論じる「非決定性（nondeterminism）」は、ただの「GPUの並列＋浮動小数点誤差」の話を越えており、「バッチ不変性（batch invariance）」という概念を中心とした設計選択が鍵、というものです。以下詳細。

---

背景と問題の定義

• LLM（Large Language Models）を使って「同じプロンプトを、temperature=0 に設定 → 決定的（deterministic）応答が返るべき」状況であっても、実務上／API提供環境では毎回応答が異なることがある。 Simon Willison’s Weblog+3Tildes+3eu.36kr.com+3

• 従来の仮説 (“concurrency + floating point hypothesis”): 浮動小数点演算の非結合性 (associativity 非保持) と、GPU/並列演算における計算順序の揺らぎ（どのスレッド／コアがいつ実行を終えるか等）が、異なる実行で微小誤差を生み、それが最終出力に影響を与える、というもの。 Tildes+1

TML はこの仮説を「正しくはあるが、決定性が失われる根本原因とは言いきれない」と指摘し、「バッチサイズやバッチ処理の条件」に依存する設計が、非決定性を引き起こすもっとも重要なファクターだとした。 eu.36kr.com+3Tildes+3Analytics India Magazine+3

---

バッチ不変性 (Batch Invariance) の意味

「バッチ不変性(batch invariance)」とは：

入力（プロンプトなど）が同じであれば、バッチの構成（バッチサイズ、同時リクエスト数、他のリクエストの有無や順序など）が異なっても、各入力に対する出力はビット単位で同一であるべき、という性質。 Tildes+1

さらに、

• サーバーサイドでは、多数のユーザー／リクエストを効率よく処理するために、リクエストをある規模ごと（サーバ負荷／リクエスト到来状況）にまとめて一括（バッチ）処理することが常。これにより “バッチサイズ” が実行ごとに変動する。 Tildes+2eu.36kr.com+2

• 多くの GPU カーネル（行列乗算やアテンションなど）は、バッチサイズに応じて内部の計算戦略を切り替える（分割アルゴリズムを変える、使うユニット（TensorCoreなど）を変える、最適化パスを変える等）ように設計されており、これが「非決定性」の実際の発生源である。 eu.36kr.com+1

---

非決定性が発生する具体メカニズム

以下のようなチェーンが問題を引き起こす：

1. サーバー負荷の変動 → 同時処理すべきリクエスト数が変わる

2. バッチサイズの変動 → 単一の入力が “大きなバッチ” 内で処理されるか、“小さなバッチ” で処理されるかが変わる

3. カーネルの戦略切り替え
    - 小バッチだとコア利用率を上げたりアイドル時間を減らす戦略を使う。
    - 大バッチや特定形状 (行列の次元) の場合は別のアルゴリズムや最適化パス（たとえば分割処理、Split‑K／Split‑KV, FlashAttention モードなど）を使う。

4. 演算順序／削減 (reduction) の内部パスの変化 → 浮動小数点演算の非結合性により、順序が違えば丸め誤差等の差異が生じる

5. 差異の累積と伝播 → Transformer のような多数層を持つモデルでは、このわずかな違いが層を重ねるごとに拡大し、最終トークン生成やロジット (logit) 出力に目に見える違いをもたらす。 eu.36kr.com+1

---

対応策：バッチ不変なカーネル設計

TML は、この非決定性を抑えるために、以下のような設計変更／実装戦略を提案／検証している：

演算（Kernel）	問題点	解決策
RMSNorm	正規化操作で “Reduce” を行う部分で、バッチ要素をどう扱うか（分割するかどうか、複数コアで分担するかどうか）によって演算順序が変わる。バッチサイズが小さいと Split‑Reduction のような戦略が使われ, 演算経路が異なることがある。 eu.36kr.com+1	バッチ不変性を維持するため、「データ並列(data‑parallel)」戦略を使う。つまりバッチサイズがどうであれ、各バッチ要素を同様な扱いで処理し、削減 (reduction) の順序を一定に保つ実装を強制する。
行列乗算（MatMul）	バッチの次元やサイズ（M, N, K 等）に応じて異なるアルゴリズムを選択することがある。たとえば、小さな行列サイズでは Tensor Core を使わない、または Split‑K を使う／使わないなど。これが部分的に演算順序の変化を引き起こす。 eu.36kr.com+1	全ての形状（matrix shapes）で単一のカーネル構成を使う。アルゴリズムやタイルサイズ（tile size）を予め固定し、それに最適化されたコードを用意しておく。パフォーマンスは落ちるが、実証実験では約 20% の落ち込みでバッチ不変性を維持可能と示されている。 eu.36kr.com
アテンション (Attention)	アテンションには行列乗算の他、シーケンス長(seq length)、チャンク（chunked prefill, prefix caching 等）の扱い、KV キャッシュの扱いなどバッチの構成以外の変数が絡む。これらが異なる戦略やレイアウトを引き起こし、演算パス（特に reduction）を変える原因となる。 eu.36kr.com	- KV キャッシュを使う場合、常にキャッシュと現在のデータを一貫したメモリレイアウト (layout) に統合してからアテンション演算をする。 - チャンク／分割（Split‑KV, FlashDecode など）の戦略を固定分割サイズ (fixed‑size splits) を使うようにし、バッチやシーケンス長の違いによる分割方式のばらつきを減らす。 - 注意して設計されたアテンション・カーネルを使う。 eu.36kr.com

---

実装上の工夫：ツールとライブラリ

• TML は batch_invariant_ops というライブラリを発表していて、既存の PyTorch モデルを比較的容易に “batch‑invariant モード” に切り替えることができるようになっている。 GitHub

  • このライブラリは、既存のカーネルを置き換える形で、「バッチ不変なカーネル」を導入する。

  • torch.Library を使って既存の演算（matmul, addmm, log_softmax など）を batch-invariant なものに差し替える仕組みを含む。 GitHub

• 実験例：vLLM を使った推論で、1000 回の同一プロンプト／条件下でサンプリングをしたところ、非決定性を抑えたモードではすべて同じ出力が得られた。 Tildes+2eu.36kr.com+2

---

トレードオフ・限界

技術的に有望ではあるものの、以下のような注意点や制限が存在する：

1. 性能低下（スループット／レイテンシ）
   固定戦略を使う、最適化パスを限定する、分割 (split) の自由度を下げるなど、演算パスの選択肢を狭めるため、特定の入力サイズ／バッチ構成では効率が落ちる。実験で報告されている数字は MatMul などで約 20% の性能低下。 eu.36kr.com

2. メモリ使用量やコア利用効率の影響
   特に小さなバッチやチャンク処理を固定戦略にする場合、GPU コアがアイドルになる時間が増えるかもしれない。リソースの最適利用 vs 決定性の確保、というトレードオフ。

3. システム全体の制御の難しさ
   本番環境では、リクエスト到来の変動、サーバー負荷、他ユーザーリクエストとの干渉など、バッチサイズを意図的に制御できない要因が多い。完全な決定性を保証するには、そのような外部要因を含めたシステム設計／運用が必要。

4. 用途による要求の違い

   • 研究・監査・法規制用途では「ビット単位での決定性」が重要。

   • 対話型チャットなどでは「多様性」「創造性」の方が価値を持つ場合もあり、あえて多少の揺らぎを残すことが望ましいケースもある。

   • RL（強化学習）での「オンポリシー vs オフポリシー」の整合性にも関連するため、トレーニングおよび推論の一貫性を維持する設計が重要。

---

実用に向けた設計上のチェックリスト

LLM を提供・構築しようとするエンジニア／研究者として、非決定性を抑えて決定性を確保するために確認すべきポイント：

1. バッチサイズの制御可能性と一貫性
   　- 推論エンドポイントで負荷によってバッチが動的に変わるかを把握する。
   　- 必要なら「batch invariant モード／設定」を導入。

2. 主要カーネルの設計・実装
   　- RMSNorm、MatMul、Attention の reduced operations 部分がバッチ不変かどうかを確認。
   　- 使用するライブラリ／フレームワーク（PyTorch, cuBLAS, FlashAttention 等）のカーネルがどのような戦略を取っているか調べる。
   　たとえば、小行列の場合にはどの戦略が選ばれているか、TensorCore の切り替え条件、split‑K や split‑KV の使用時の挙動など。

3. KV キャッシュ／チャンク分割（Prefill vs Decode 等）の扱い
   　- キャッシュの内容と現在の入力の扱いが常に同じ内部レイアウトで処理されるか。
   　異なるレイアウトや異なるマスク (mask/full/partial) の扱いが、計算パスを変えないようにすること。

4. ライブラリとフレームワークのサポート
   　- TML の batch_invariant_ops のような補助ライブラリを使う or それに類する取り組みがなされているか。 GitHub
   　- フレームワーク自体が決定性モードを提供しているかどうか。

5. 性能計測と妥協点の評価

   • 決定性モードでの速度・スループット低下を測る。

   • バッチ構成・シーケンス長による違いを実際の運用ワークロードで確認する。

   • 出力の決定性 vs レイテンシ／コストのバランスをどう設計するかを決める。

---

総括

TML の研究は、LLM の「応答の揺らぎ」を単に避けられない副作用とみなすのではなく、システム設計／カーネル設計の選択次第で大きく制御できるということを示した点で技術的に非常に重要です。エンジニアリング視点では、「バッチ不変性」を設計パラメータの一つと位置付け、そのための実装戦略をあらかじめ設計に組み込むことが、今後の LLM 推論サービスの品質向上（再現性・信頼性・監査可能性等）にとって鍵になります。