腾讯エンジニア：大規模モデルの出力を確定させるための AI インフラ手法

AI インフラの深化：大規模モデルに確定的な出力を実現する方法

腾讯程序员 2026-08-05 17:25 广东

AI インフラを一段階レベルアップさせる

著者：binnnliu

背景

大規模モデルとの対話で、同じプロンプトを入力しても毎回結果が異なることに気づいたことはありませんか？これは「大規模モデルは本質的に単語を予測する装置である」という我々の一般的な認識に合致しています。

しかし、再現性は科学発展の基盤です。そのため、大規模モデルから完全に確定的な出力を得ることは極めて重要な研究課題であり、特に強化学習においては、実験の再現性とトレーニングプロセスの安定性を保証するために、決定論的なロールアウト（Rollout）が必要です。

Reproducibility is a bedrock of scientific progress.

多くの人が最初に思い浮かべるのは、「同じハードウェアで、同じプロンプトなら」という点でしょう。Temperature を 0 に設定してランダムサンプリングを無効化すればよいのではないか？あるいは、業務要件で Temperature > 0 が必要な場合でも、乱数シード（Seed）を固定すれば毎回同じ出力が得られるのではないか？

実際に検証しました。環境は vllm serve Qwen/Qwen3-8B です。

結論は「部分的に正しい」です。

「正しい」と言える点：ランタイム実行間での確実性（Run-to-run Determinism）として、同一時間に他のリクエストがない状態で同じプロンプトを複数回送信すれば、回答は一致します。

「正しくない」と言える点：バッチ不変性（Batch Invariance / 批次不变性）の観点では、同一時間に他のリクエストが存在する場合、同じプロンプトを複数回送信しても回答が異なります。

異なるリクエスト同士が互いに影響し合うというのでしょうか？あり得ない話です。各リクエスト間のアテンション機構は物理的に厳密に隔離されており、コンテキスト汚染の問題など起こりようがありません。

実はその根源は、推論エンジン層の動的なバッチ組立てと演算子スケジューリング戦略が浮動小数点加算の順序を変化させていることにあります。そして、浮動小数点加算には結合律が成り立たないのです：(a + b) + c ≠ a + (b + c)。

システム設計の観点から見ると、結果にばらつきが生じる根本原因は情報の干渉ではなく、システムがメモリアクセス遅延を隠蔽するために、非決定論的なハードウェアレベルの並列最適化（Non-deterministic Hardware-level Parallel Optimizations）を底層で発動させている点にあります。

浮点数と定点数を比較すると、浮点数の本質は科学記数法にあります。これは限られたビット数（例えば 16 ビットや 32 ビット）の中で、表現できる数値の範囲と絶対精度の間のトレードオフを実現するための仕組みです。

指数が大きくなるほど表現可能な数値の範囲は広がりますが、尾数のビット数が固定されているため、絶対精度（つまりステップ幅）は低下します。この絶対精度（ステップ幅）は以下の式で計算されます：ULP(x) = 2E - M。ここで E はその浮動小数点数の実指数、M は尾数のビット数を表します。

IEEE 754 標準によると、FP16 は 1 ビットの符号位、5 ビットの指数位、そして 10 ビットの尾数から構成されています。例えば、数値が 2048（すなわち 2^11）の場合、その規格化された表現は 1.0000000000_2 × 2^11 となります。このとき、隣り合う表現可能な数の間のステップ幅（ULP）は 2 に増加します。つまり、FP16 では 2048 の次の数は 2050 となり、2049 を正確に表現することはできません。

もしハードウェアが 2048 + 1 を計算する場合、演算器内部ではより広い GRS 拡張ビット（保護位、舍入位、粘滞位）を用いて 2049 という正確な結果を導き出します。しかし、この結果を 10 ビットの尾数に書き戻す際、2049 はちょうど 2048 と 2050 の真ん中に位置するため、「最も近い値へ丸める（偶数に丸める）」ルール（Round to Nearest, ties to Even）が適用されます。その結果、最終的に 2048 に強制的に丸められてしまいます。

ここで注意すべき点は以下の通りです：

この例はあくまで「浮動小数点の加法は結合律を満たさない」ことを示すためのものです。実際の推論プロセスでは、GEMM（行列乗算）の入力は FP16 または BF16 ですが、累加演算は FP32 で行われます。バッチごとのばらつき（batch variance）が存在する根本原因は、FP16 の記憶精度が低いことだけでなく、リダクションのトポロジー（結合順序や構造）が変わったことにあります。

• 数値：規格化された実指数 / 実際の指数 (物理尾数位：10 ビット) / 累加の分解
• 0.5：1.0000000000_2 × 2^-1 / -1 / 0000000000
• 14 (01110)：1.0000000000_2 × 2^0 / 0 / 0000000000
• 2：1.0000000000_2 × 2^1 / 1 / 0000000000
• 2046：1.1111111110_2 × 2^10 / 10 / 25 (11001) → 1111111110_2 = 2^10 + 2^9 + ⋯ + 2^1
• 2047：1.1111111111_2 × 2^10 / 10 / 25 (11001) → 1111111111_2 = 2^10 + 2^9 + ⋯ + 2^0

20481.00000000002 × 21111

26 (11010)

0000000000211

20501.00000000012 × 21111

26 (11010)

0000000001211 + 21

では、具体的にどのような操作が浮動小数点加算の順序を変えてしまうのでしょうか？その答えを見つけるには、まず行列演算から見ていく必要があります。

GEMM の Batch Invariance（バッチ不変性）

なぜ GEMM 演算は浮動小数点加算の順序を変えてしまうのでしょうか？その理由は GEMM の最適化戦略にあります。

GEMM の最適化戦略

この問題を理解するには、まず「メモリウォール」の下で GPU の GEMM オペレーターがどのように進化してきたかを把握しておく必要があります。多くの人が疑問に思うでしょう。「単なる基本的な行列乗算なのに、なぜこれほど複雑になるのか？」

答えは「極めて複雑」です。実際、AI インフラの分野全体がこの細部を巡って熾烈な競争を繰り広げています。その核心にある理由は、現代 GPU の計算ユニット（Tensor Core）の演算能力が劇的に向上した一方で、それを支えるメモリ帯域幅の伸びが追いついていないからです。

この格差により、GPU は計算ユニットがデータ転送のために待たされる時間が長くなり、結果として計算ユニットは入力データを待つ間に大量にアイドル状態（Compute Bound から Memory Bound への変化）になってしまいます。これが「メモリウォール」と呼ばれる現象です。

for i from 0 to M: // 外側ループ：A の行（全 M 行）を走査

for j from 0 to N: // 中間ループ：B の列（全 N 列）を走査

for k from 0 to K: // 内側ループ：A の i 行目と B の j 列目の内積（長さ K）を計算

C[i,j] += A[i,k] * B[k,j]

Tiling - 分塊行列乗算

単純に上記の式に従って要素を一つずつ計算すると、乗算・加算演算のたびに極めて低速なグローバルメモリ（HBM）からデータを読み出す必要があります。オンチップの高速キャッシュ（SRAM/共有メモリ）やレジスタが全く利用されないため、膨大な数の HBM への重複読み出しが発生し、計算リソースを十分に活用できなくなります。

メモリ帯域のボトルネックを打破するため、業界では IO-Aware（I/O 感知）なブロック分割アルゴリズムが採用されています。その核心は「化整为零」、つまり大規模な行列をキャッシュに収まる小さなブロックに切り分けることにあります。スケジューリング時には、各 CTA（Cooperative Thread Array/Thread Block）を結果行列 C の特定の小部分 Cij に割り当て、この CTA がその小ブロックの計算を全権で担当します。

ある CTA が Cij の計算を担当する際の、完全なデータフローは以下の通りです。

累算器の初期化

レジスタを割り当て、目標とする分塊 Cij の中間累積結果を格納します。

K 次元に沿った縮約（Reduction）

共通次元である K に沿ってステップごとにスキャンを進めます。各ステップ p では以下のパイプラインを実行します。

• Load (HBM → SRAM): 計算に必要なサブブロック Aip と Bpj を、低速なグローバルメモリから高速な共有メモリ（SRAM）へ、一度にバッチ処理で転送します。
• Compute (Warp GEMM): 各 Warp が SRAM からデータをレジスタに取り出し、Tensor Core に渡して高速な行列乗算を実行します。
• Accumulate: 今回計算された積の結果を、レジスタ内の累算器に直接加算します。結果はメモリへ書き戻しません。

エピローグと書き戻し

K 次元のループが完了すると、レジスタ内の累算器には Cij の最終的な正確な値が格納されます。その後、一括してグローバルメモリ（HBM）へ書き戻します。

Split-K

これまでの整理したロジックを踏まえると、各 CTA は出力行列 C の 1 つの Tile を担当し、内部で K 次元に沿って順次ループを実行します。

しかし現実には、K 次元が極めて巨大で、出力行列 C（つまり M と N 次元）が非常に小さく、結果として切り出せる Tile がたった 4 つしか存在しないという極端な状況に直面することがあります。この場合、GPU のハードウェアスケジューラーは作業を担うために 4 つの CTA しか起動しません。NVIDIA H100 は 132 個の SM を備えているため、残りの 128 個の SM が完全にアイドル状態となり、たった 4 つの CTA が K 次元という極めて長いループを必死に実行するのを待機していることになります。これにより GPU の計算リソースは著しく浪費されてしまいます。

空間次元（M, N）でこれ以上のタスクを切り出せない場合、例えば LLM 推論の Decoding 段階では Batch Size が極小となり M=1 です。もし強制的に空間分割のみを行うと、GPU のコアの多くが作業待ちとなってしまいます。そこで K 次元をさらに細かく分割し、異なる CTA に計算を割り当てることはできないでしょうか？これが「Split-K」最適化戦略です。

その核心となる考え方は、単一の CTA が K 次元全体を独自に計算する従来のロジックを打破することにあります。K 次元を複数のセグメント（Split_K）に分割し、異なる複数の CTA に同時に同じ Cij ブロックの異なる K セグメントにおける局所的な乗算・累積処理を担当させます。

（注意：通常の Tiling 計算でも K 次元に沿って分段して計算することはあります。しかし両者には本質的な違いがあります。BLOCK_K は時間軸上の順次処理であり、同じ CTA がグローバルメモリから BLOCK_K サイズのデータを搬送し、レジスタ内で固定された順序で累積していきます。一方、SPLIT_K は空間軸上の並列処理であり、物理的に独立した複数の CTA に強制的にタスクを分散させて同時に計算を行います。）

図のように Split-K = 2 と設定すると、異なる色の Tile が 2 つの独立した CTA によって並列処理され、計算リソースの利用効率が即座に倍増します。最終的には atomic_add を用いて結果を行列 C の同一位置に累積加算します。ここで留意すべき点は、SPLIT_K の値が大きすぎると atomic_add の競合が増大し、その効果は次第に減衰していくことです。

原子加（atomic_add）以外にも、Workspace Reduction という実装手法があります。その具体的なワークフローは以下の通りです。

Workspace の割り当て（メモリ確保）

HBM に [split_k, M, N] サイズの一時バッファ（Workspace）を別途確保します。

Partial GEMM Kernel の実行

主となる行列乗算カーネルを開始します。各スレッドブロックは K 次元の断片（スライス）を担当し、計算が完了すると原子加を使わずに、その結果を Workspace 内の自分専用の領域（workspace[k_idx, m, n]）へ直接書き込みます。

Reduction Kernel の実行

主カーネルの実行と同期が完了した後、独立した軽量な規約カーネルを起動します。このカーネルは split_k 次元に沿って Workspace 内の局所結果を加算し、最終的な合計値を出力行列 C に書き込みます。

GROUP_M - Swizzle L2 Cache

CTA（Cooperative Thread Array）のスケジューリング順序を制約することで、広範囲に散らばったメモリアクセスによる L2 キャッシュの競合（キャッシュスラッシング）やデータ頻繁な追放を防ぎます。これは論理的に複数の独立した CTA を、GROUP_SIZE_M × N のサイズを持つマクロスケジューリング行列へと再構成するものです。

この限定された連続実行区間では、スケジューリングされた CTA が空間的に隣接する出力ブロックを集中的に処理するため、L2 キャッシュにロード済みの行列 A と B のデータを共有できます。これによりデータの時間的局所性と L2 キャッシュの再利用率が最大化されます。

これは別の次元におけるタイリング（Tiling）と言えるでしょう。目的は L2 の再利用です。前述した従来のタイリングが SRAM やレジスタの再利用を目的としたのに対し、さらに高度なテンソル並列化（Column/Row Parallel）もまた、より高次元のタイリングと捉えることができます。

演算子チューニング

前文探讨了 Tiling、Split-K、Swizzle 等旨在缓解访存瓶颈的核心优化策略。但在实际的 GPU 硬件执行层面，针对不同规模的输入张量（Tensor Shape），系统必须确定具体的并行化切分配置：即矩阵分块的具体维度（BLOCK_M / BLOCK_N / BLOCK_K）以及 K 维度的切分段数（SPLIT_K）。这些切分参数的选择，不仅直接决定了算子的访存效率与硬件利用率（Occupancy），更关键的是——分块参数的动态变化会重塑底层浮点数累加的归约拓扑（Reduction Tree），这是导致大模型推理在不同批次下失去 Batch Invariance 的根本原因。

为了解析推理引擎底层为何会根据输入特征动态变更这些调度参数，我们需要先回顾 GPU 算子的编程范式演进，并深入探讨现代编译器引入的自动调优（AutoTune）机制。

CUDA VS Triton

在之前的《AI Infra 入门：GPU 是如何工作的》中，我们探讨了 CUDA 编程模型与 GPU 硬件执行模型：Grid 定义了内部 Thread Block 的组织形式 (gridDim)；Thread Block 定义了内部 thread 的组织形式 (blockDim)。

然而，原生 CUDA 编程的门槛依然极高，开发者需要手动实现各种底层且繁琐的优化逻辑。OpenAI Triton 的出现，彻底改变了这一现状。它大幅降低了编写高性能 GPU 算子（Kernel）的门槛，让开发者能以接近 Python 的生产力，写出接近 CUDA C++ 专家级性能的代码。借助内核融合（Kernel Fusion）、IO-Aware 等底层优化技术，Triton 不仅保证了极致性能，还兼顾了极佳的硬件无关性。

两者核心思想的差异在于编程范式：

CUDA（以线程为中心）：开发者需要显式定义底层网格，例如 <<>>，由底层生成 128 个线程，硬件再将其划分为 4 个 Warp 调度执行。

Triton（以 Tiling/Block 为中心与自动调优）：开发者只需关注数据分块，并可通过 @triton.autotune 装饰器提供一个配置搜索空间（如不同的分块大小、测试 4 个或 8 个 Warp）。Triton 编译器不仅会在底层自动将需求（如 4 个 Warp）精准映射为 4×32=128 个线程以屏蔽硬件细节（即 <<>>），更会在运行时自动进行基准测试，智能选出当前硬件下的最优参数组合。

在 PyTorch 2.0 及更高版本中，引入了 OpenAI Triton 作为编译器：

前端分析：当用户调用 torch.compile(model) 时，PyTorch 的前端（如 TorchDynamo）会捕获计算图。

后端优化与代码生成：默认后端 TorchInductor 会分析图中可以被 Kernel Fusion 的操作，生成 OpenAI Triton 代码。

JIT コンパイルと実行：最後に、Triton の JIT（Just-In-Time）コンパイラーがこのコードを引き継ぎ、特定の GPU ハードウェア向けに高度に最適化された単一の融合カーネル（Fused Kernel）としてコンパイルし、最終的に効率的に実行します。

このエコシステムにおいて、PyTorch（Inductor と Triton を活用）は、従来 CUDA の専門家が多額の時間を費やして行っていた性能チューニングを自動的に担当するようになりました。自動調優機能を持つ Inductor+Triton は、ほとんどのシナリオで劇的な性能向上をもたらしますが、本質的には汎用的な自動化ソリューションです。極めてクリティカルで、性能を限界まで引き出す必要がある特殊な演算子においては、依然として手動による低レベルの最適化が不可欠です。

いずれにせよ、コンパイラーによる自動調優であれ、人手による手動調優であれ、抽象化の奥にある GPU 底辺の物理的な制約とチューニングの論理を理解することが必須となります。

ハードウェア実行層では、複数の Thread Block が同じストリームマルチプロセッサ（SM）を共有できます。SM の物理リソース（レジスタや共有メモリなど）がまだ満杯でない限り、ハードウェアスケジューラーはより多くの Block を同一の SM 内に配置して並列実行します。ただし、各 SM には並行処理の硬性上限（例：最大スレッド数 2048、最大 Block 数 32）が存在します。

同じ Thread Block 内の Thread は、共有メモリ（Shared Memory）を介してデータ交換を行い、同期（__syncthreads()）を行うことができます。

レジスタ制限

A100/H100 の各 SM には、物理的なレジスタファイルサイズが固定されており、65,536 個の 32 ビットレジスタ（256 KB）です。ハードウェアスレッドごとに割り当てられる最大レジスタ数は 255 です。一方、単一のスレッドが必要とするレジスタ数は以下の 2 つの部分で構成されます。

• ベースオーバーヘッド（Base Overhead）：各スレッド固有のもので、メモリポインタ、ループカウンタ、TMA 状態などを格納するために使用されます。
• アキュムレータの配分（Accumulator Share）：出力ブロックサイズを仮定し、アキュムレータが BLOCK_M × BLOCK_N の場合、このレジスタは全スレッドで均等に割り当てられます。各プロセスのレジスタ使用量は (BLOCK_M × BLOCK_N) / NUM_WARPS / 32 で計算されます。Tiling サイズが大きくなれば、num_warps を同期して増やさなければならず、そうしないとレジスタスパリング（Register Spilling）が発生します。

関連する式：(BLOCK_M * BLOCK_N) / (num_warps * 32) <= 255

通常は 128 程度に抑える必要があります。もし 255 を超えた場合は、num_warps を大きくするか、BLOCK サイズを小さくする必要があります。

実行パターンの制約と共有メモリの制限

従来の CUDA 最適化において、メモリアクセスの待ち時間を隠すために Occupancy（占用率）を高めることは基本原則でした。これはスレッドレベルの並列処理（TLP）による遅延隠蔽です。SM（Streaming Multiprocessor）上に常駐するアクティブな Warp の数、つまり Occupancy が高いほど、ハードウェアスケジューラは現在の Warp がメモリアクセスで停止した際に、すぐに切り替え可能な他の待機状態の Warp を見つけやすくなります。このゼロオーバーヘッドのコンテキストスイッチにより、計算命令とメモリアクセスが時間軸上で重なり合い、グローバルメモリの物理的な遅延を効果的に隠すことができます。

しかし、Ampere や Hopper アーキテクチャでは、Tensor Core が極めて高い浮動小数点演算スループットを提供する一方で、HBM（High Bandwidth Memory）の物理的なメモリアクセス遅延はそれに見合うほど短縮されていません。計算能力がメモリ帯域を圧倒的に上回るこの状況下では、高 Occupancy 戦略の効果が急激に低下します。SM 上に多くの Warp が常駐していても、それらはすぐに現在の計算タスクを完了し、一斉にメモリアクセス要求を発行するため、すべての Warp が同時に停止してしまうのです。この場合、スレッドレベルのコンテキストスイッチだけでは、メモリレベルの物理的な遅延を隠すことができません。

Tensor Core の高いスループットを維持するために、現代の GPU ではハードウェアレベルの非同期メモリアクセス転送機構（Ampere アーキテクチャでは cp.async、Hopper アーキテクチャでは TMA エンジン）が導入されました。num_stages の核心となる考え方は、マルチバッファリング（Multi-Buffering）を活用して計算命令とデータ転送命令を非同期で並行実行することです。つまり、GEMM の最適化方針は「高並列処理によるメモリアクセス遅延の隠蔽」から、「非同期パイプライン（Asynchronous Pipelining）による遅延の隠蔽」へとシフトしています。

num_stages の設定が小さすぎると（デフォルト値の 2 など）、計算プロセスがデータ準備完了を待ってパイプラインに隙間（Pipeline Bubble）が生じ、全体のスループットが制限されてしまいます。逆に設定が大きすぎると、SRAM が溢れてコンパイルエラー（Out of Shared Memory）が発生するだけでなく、マルチステージパイプラインの状態を管理するためのポインタが追加のレジスタリソースを消費し、性能が低下することもあります。SRAM 容量がより大きい H100 などの異なるハードウェアアーキテクチャで高度なチューニングを行う際の本質的な課題は、リソース溢れを引き起こさずにオーバーラップ率を最大化する最適な境界条件を見つけることです。

(BLOCK_M * BLOCK_K + BLOCK_N * BLOCK_K) * 1 バイト * num_stages <= SMEM の物理上限（H100 では 228KB）

チューニングのトレードオフ

Llama-3-8B の Q/K/V Linear Proj - Fused QKV を例に挙げます。hidden_size は 4096、qkv_proj_size は 6144 です。

• 項目：詳細
• 目標行列：Fused QKV
• 計算式：QKVout = X · Wqkv
• 入力 X の次元：[num_sched_tokens, hidden_size]
• 重み行列 Wqkv の次元：[hidden_size, qkv_proj_size]
• 出力結果の次元：[num_sched_tokens, qkv_proj_size]
• 備考：重みを列方向に結合し、1 回の広幅 GEMM で実行。4096 (Q) + 1024 (K) + 1024 (V) = 6144

つまり、M=num_sched_tokens、N=6144、K=4096 となります。ここで BLOCK_K は通常、16、32、または 64 に設定されます。

BLOCK_M / BLOCK_N の設定：

M が大きい場合（M > 4096）は、原則として計算強度を高め、BLOCK_M × BLOCK_N を可能な限り大きくして SMEM の利用効率を最大化する必要があります。この際、num_warps の適切な値を見つけることが重要です。

num_warps が小さすぎると、実際に作業を行うスレッド数が不足し、1 つのスレッドに割り当てられる累加器用レジスタが多くなりすぎてしまいます。その結果、スレッドあたりの最大 255 個という物理的なレジスタ数の上限を超え、「Register Spilling（レジスタの溢れ）」が発生します。

一方、num_warps が大きすぎると、基礎的なオーバーヘッドの割合が大きくなり、1 つの CTA が多くの SM リソースを占有してしまいます。その結果、GPU の Occupancy（稼働率）が低下します。したがって、BLOCK_M × BLOCK_N を十分に大きく保つために、適切な num_warps の値を選ぶ必要があります。

M が小さい場合（M < 64）は、BLOCK_M を M に合わせて縮小し、SM の計算能力を最大限に活用するために、以下の 2 つの対策が有効です。1. BLOCK_N を適度に小さくしてより多くの Tile を生成し、異なる SM に分散する。2. Split-K を有効にして、K 次元方向により多くの CTA を並列処理として起動する。

指標と最適化目標

• 項目：関連ハードウェアリソース / 役割と定義 / Trade-Off（トレードオフ）
• BLOCK_M / BLOCK_N：共有メモリ (SMEM)、L1 キャッシュ / 行列の行方向と列方向におけるサブブロックのサイズ（例：128, 256）。これにより、共有メモリに読み込まれるデータ量が決定されます。 / サイズが大きすぎると、各スレッドが消費するレジスタ数が増えすぎて「Register Spilling」を引き起こします。これはレジスタが極端に遅いローカルメモリへ溢れ出す現象であり、性能の急激な低下（雪崩）を招きます。
• BLOCK_K：共有メモリ (SMEM) / K 次元方向での累加処理の長さです。 / 上記と同様の理由で、通常は 32 または 64 に設定されます。M/N/num_stages と組み合わせて計算し、総 SMEM 使用量を調整する必要があります。
• num_warps：レジスタ、Occupancy（並行度） / 各 Thread Block に割り当てられる Warp の数です（1 Warp = 32 スレッド）。本質的には計算負荷を分散させるための分母となります。 / 小さすぎると累加器がスレッドあたりの最大 255 個のレジスタ上限を超え、致命的な Register Spilling が発生します。大きすぎると、1 つの Block が多くの Warp を占有してしまい、SM 内に収められる Block 数が減り、並行度（Occupancy）が急激に低下します。
• num_stages：グローバルメモリ (HBM)、レイテンシのカバー / ソフトウェアによるパイプライン処理の段数です。複数のバッファを確保し、計算ユニットが現在のブロックを処理している間に、バックグラウンドで次のブロックの読み込みを非同期で行います。 / 設定値が大きすぎると SMEM が溢れる原因となります。
• GROUP_M (Swizzle)：L2 キャッシュ / 連続する複数の行をグループ化してパックし、デフォルトの逐行スキャンを打破します。L2 キャッシュに既にロードされている行列 A と B のデータを共有することで、データの時間的局所性と L2 キャッシュの再利用率を最大化します。
• SPLIT_K：ストリーミングマルチプロセッサ (SM)、計算利用率 / 非常に長い K 次元を複数の Block に分割して同時に処理し、最後に原子加算（Atomic Add）で結果を結合します。M と N が極小で K が極大の場合（Flash-Decoding など）にのみ有効化され、SM の利用効率を向上させます。

Batch Invariance の由来

前節では多くの最適化戦略について議論しましたが、具体的にどのパラメータの変更が浮動小数点加算の順序を変えているのでしょうか？

数学的な本質を振り返ると、行列乗算 C = A × B（A は [M, K] 行列、B は [K, N] 行列、C は [M, N] 行列）において、任意の要素 Ci,j の計算ロジックは一定です。つまり、A の i 行目と B の j 列目の内積を計算するだけです。

Ci,j = Σk=0K-1 Ai,k × Bk,j

Triton や CUDA の底层优化において、各種スケジューリングパラメータの動的調整がもたらす「決定性（Determinism）」への影響には、本質的な違いがあります。

BLOCK_M と BLOCK_N（影響なし）

分塊のロジックは、Tile_A [BLOCK_M, BLOCK_K] と Tile_B [BLOCK_K, BLOCK_N] を読み込み、出力として Tile_C [BLOCK_M, BLOCK_N] を計算することです。これは空間次元におけるタスクマッピングに過ぎず、どの要素をまとめて独立して計算するか、そして特定の CTA (Thread Block) が行列のどの領域を担当するかを決定するだけです。

出力行列 C の特定要素 Ci,j について考えてみましょう。それがどの CTA に割り当てられ、どのような隣接要素とともに計算されようとも、その根底にある K 次元の点積ロジックは変わりません。したがって、この 2 つのパラメータを変更しても、浮動小数点加算の順序には影響しません。

GROUP_M / Swizzle（影響なし）

Swizzle の本質は、複数の CTA がグリッド (Grid) レベルで実行される順序を変えることで、L2 Cache のヒット率を向上させることにあります。これは「どの空間ブロックを先に計算し、その後に行うか」を決めるだけで、特定の分塊内部における乗算と加算の過程には一切干渉しません。したがって、Swizzle も Batch Invariance を破壊することはありません。

BLOCK_K（確定的な誤差を導入）

これは縮約次元 (Reduction Dimension) におけるステップサイズを定義するパラメータです。BLOCK_K の値を変更すると、1 回のループでレジスタに読み込まれるデータ量が直接変化し、結果として Tensor Core 内部の MMA 命令が構成する累加ツリーのトポロジーが変わります。この変更は浮動小数点加算の順序を変化させます（つまり、異なる BLOCK_K では異なる結果が得られますが、同じ BLOCK_K の値であれば結果は一定になります）。

SPLIT_K（非確定的な誤差を導入）

これは K 次元を複数の CTA に強制的に分割し、並列実行させる手法です。Atomic Add を用いて結果を統合する場合、GPU ハードウェアがどのスレッドブロックを先に実行するかは完全にランダムであるため、加算の順序も制御不能なランダムなものになります。Workspace Reduction を採用した場合でも、SPLIT_K の段数が動的に変化すれば、累加ツリーのトポロジー自体が変わってしまいます。

num_warps と num_stages の本質は、ハードウェアリソースの割り当てとパイプラインスケジューリング（どの CTA を使って計算するか、非同期転送のためにいくつのキャッシュプールを確保するか）にあります。これらはマクロな行列点積の計算ロジックや Reduction Tree のトポロジーを変えるものではありません。

以上より、推論エンジンには性能と確実性の間に本質的なアーキテクチャ上の衝突が存在します。極限のメモリアクセス再利用（BLOCK_K の動的調整）を追求し、かつ SM (Stream Multiprocessor) の並行利用率を高めるために SPLIT_K を動的に有効化すると、その結果として基盤となるヒューリスティックなスケジューリング戦略が浮動小数点演算の Reduction Tree トポロジーを不可避的に変更してしまいます。これが、大規模モデルが動的 Batch 下で Batch Invariance を失う根本的な理由です。

では、実際の推論シナリオにおいて、このメカニズムはどのように修復されるのでしょうか？次は vLLM の実装に深く入り込み、それがどうやって底層のハードウェア特性と上層のスケジューリングロジックの間で Trade-off を行っているかを見ていきましょう。

vLLM における GEMM の Batch Invariance サポート

解決策はどうあるべきでしょうか？直感的には、推論エンジン内で Split-K をグローバルに無効化し、BLOCK_K を固定すればよいように思えます。しかし、現実のエンジニアリング実装はこれほど単純ではありません。

現代の推論エンジン（vLLM など）における GEMM 操作の実行パスは、複数の抽象層と多様なバックエンドをまたぐ複雑な経路をたどります。厳密な Batch Invariance を実現するには、単一の設定変更で済む話ではなく、特定のハードウェア・ソフトウェア実行環境に応じて、多次元の執行パスのルーティングとパラメータ制約が必要です。

具体的な GEMM 実行パスは、以下の 3 つの次元の組み合わせによって決定されます：

ハードウェアアーキテクチャの違い：SM80（Ampere）と SM90/SM100（Hopper/Blackwell）では、基盤となる GEMM 実行のパラダイムが根本的に異なります。前者はワープレベルの並行処理でメモリアクセス遅延を隠蔽しますが、後者は TMA と WGMMA を活用した非同期パイプラインに依存しています。この実行パラダイムの違いにより、Batch Invariance（バッチ不変性）を実現するロジックも変わります。

データ精度の違い：bf16 や fp16 といった精度は、生態系が最も成熟している cuBLASLt が担当します。一方、fp8 や fp4 などの低精度では CUTLASS（さらには DeepGEMM などの専用カーネル）へと切り替わります。これに伴い、異なるバックエンドの制御ロジックも適用されます。

オペレーター入口（API Backend）の違い：nn.Linear は vLLM 独自のディスパッチを経由し、カスタム高性能カーネルへ直接ルーティングできます。一方、裸の torch.mm や bmm は PyTorch のディスパッチャー → cuBLAS/cuBLASLt を経由するため、调度権限はフレームワークとクローズドソースライブラリ側にあります。

SM8x と SM90/100 の比較

SM80：ワープレベルのテンソルコア（mma.sync）スケジューリング。最適化のポイントは利用効率を高めることで、ワープがデータ転送と計算の両方を行います。

4096 * 4096 → CTA Tile（例：128*128）→ Warp Tile（32 * 64）→ mma 命令レベル (16 * 8 * 16)

SM90/100：TMA ハードウェアユニットの導入と WGMMA 命令（ワープグループ 128 スレッド）により、パラダイムが転換しました。非同期パイプラインを活用して遅延を隠蔽します。

4096 × 4096 の大行列では、CTA（Cooperative Thread Array）に割り当てられるメモリ空間が十分確保されるため、各 Warp に広めの専用領域（例：32×64 など）を割り当てることができます。この大きな Tile をさらに小さな単位に分割して Warp に配分し、複数の Warp が交互にメモリアクセスと計算を行うことで、メモリの待ち時間を隠蔽します。K 次元方向では各 Warp が独立して実行され、互いに干渉することなくシリアル処理を完遂するため、実行間（run-to-run）で結果の安定性は保たれます。

しかし、バッチサイズが変わる場合、BLOCK_K の値も変動し、浮動小数点加算の順序が揺らぐ可能性があります。ただし、必ずしもそうなるわけではありません。これは Warp レベルでの K 処理が一定かどうかにかかっています。例えば、ヒューリスティックな自動チューニングにより BLOCK_K が 32 から 64 に変わっても、実際の基盤実行では K=16 をステップとして Tile を順次処理するため、バッチによる変動は底層で解消されます。つまり、BLOCK_K の影響を無視して Split-K に注目すればよいのでしょうか？大行列においては、この考え方は概ね妥当です。

一方、32×32 といった小行列の場合、1 つの CTA に割り当てられ（K=4096）、各 Warp が得られる Tile が極めて小さくなります。その結果、十分な数の Tile を確保できず、M/N 方向の並列度が不足します。利用可能な Warp の数が減り、SM の利用率が低下するため、K 次元でのシリアル処理だけではメモリの待ち時間を十分に隠すことができません。そこでシステムは、すべての Warp が同一の 32×32 行列上で重なるようにし、時間（K）方向で強制的に分割する「Warp-level K-Slicing」を採用します。複数の Warp がそれぞれ K の一部を計算し、最後に共有メモリ内の Reduction Tree で結果を統合します。バッチパラメータの変化によって K の分割数が変わると、Reduction Tree の形状も変化し、浮動小数点加算の順序が揺らぐ原因となります。

小行列の場合、GEMM ライブラリはより小さな Tile や GEMV/SIMT 方式、Persistent Kernel、Warp-level K-Slicing、CTA-level Split-K、Stream-K など、状況に応じた複数のアルゴリズムから最適なものを選択します。Warp-level K-Slicing は、同一の出力 Tile を複数の Warp がそれぞれ K の断片ごとに計算し、CTA 内で部分結果を統合する手法です。これにより単一 CTA 内の有効な Warp 数は増えますが、Grid 全体の CTA 数が増えるわけではありません。SM の利用率をさらに高めるには、通常は CTA-level Split-K が併用されます。

しかし、SM90 以降のアーキテクチャでは、従来の高 Occupancy を利用してメモリアクセス待ち時間を隠す依存度が低下しています。標準的な WGMMA mainloop では、1 つの Warpgroup が分散された Accumulator を共同で管理し、K の Tile を順次処理します。これにより、複数の独立した Warp による部分結果を生成してから CTA 内で統合する手順は不要になります。

図：大行列処理における CTA、Warp Group Tile、および指令レベルのデータフロー

コードコメント「Hopper (SM90) と Blackwell (SM100): バッチのばらつきが生じる唯一の原因は Split-K である」からわかるように、現在サポートされテストされている FP16/BF16 の PyTorch/cuBLASLt パスでは、SM90 Hopper および SM100 Blackwell アーキテクチャ上で Split-K を禁止することでバッチ不変性（batch invariance）が保証されます。

ここで注意が必要なのは、この結論は cuBLAS や cuBLASLt がクローズドソースであるためテスト結果に基づいている点です。SM90/SM100 では Split-K を固定するだけでバッチ不変性が得られますが、SM8x アーキテクチャではバッチ不変性を確保するために、BLOCK-K と Split-K の両方を同時に固定する必要があります。

また、「Split-k で複数のワークスペースを使用すると、収縮（reduction）ツリーは順序通りに処理されるため結果が一定になる」と誤解する人もいますが、これは実行間での安定性（run-to-run 次元）に過ぎません。K の値が変われば結果も変化します。

ここまで説明した通り、最終的な浮動小数点加算の順序は、メインループにおける BLOCK_K のステージング（CTA Tiling）だけでなく、より低次元の Warp Tiling や Warp Group Tiling、さらにハードウェア固有のアトミック命令の粒度にも影響されます。ただし、この問題を解決するアプローチとしては CTA Tiling の観点から検討する必要があります。

SM80 における GEMM のバッチ不変性サポートについて

SM80 アーキテクチャでは SPLIT_K と BLOCK_K の両方に注意を払う必要があります。Linear レイヤーの GEMM 演算においては、SM のバージョンや eager モード・compile モードのいずれであっても、必ず linear_batch_invariant から matmul_persistent が呼び出されるフローになります。

class UnquantizedLinearMethod(LinearMethodBase):

def apply(

self,

layer: torch.nn.Module,

x: torch.Tensor,

bias: torch.Tensor | None = None,

) -> torch.Tensor:

if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():

return linear_batch_invariant(x, layer.weight, bias)

return dispatch_unquantized_gemm()(layer, x, layer.weight, bias) 而对于非线性层的计算，比如代码中直接调用 torch.mm，则还需要额外考虑。由于 SM80 在 cuBLASLt 上不支持禁用 BLOCK_K，只能更换 Triton 实现，最终也会走到 matmul_persistent。其中 matmul_persistent 是一个 persistent kernel 即 griddim=num_sm，固定 K 顺序、固定 tile、不切 SPLIT_K。

if current_platform.is_device_capability_family(80):

# SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the

# triton persistent matmul overrides for mm/addmm/matmul/linear.

_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")

_batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")

_batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")

_batch_invariant_LIB.impl("aten::