2026__MLSys2026__Dataflow Is All You Need

# Dataflow Is All You Need > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 5 (May 22 / Fri)、Grand Ballroom 1、Industry Track Oral Presentation: Compilers/HW（08:15–08:30 PDT、Moderator: Phitchaya Phothilimthana） > - **登壇者:** Darshan Gandhi（SambaNova Systems、筆頭著者） > - **全著者:** Darshan Gandhi¹, Pushkar Nandkar¹, David Koeplinger¹, Nasim Farahini¹, Romy Tsoupidi¹, Samuel Rydh¹, Matheen Musaddiq¹, Tuowen Zhao¹, Reid Goodbar¹, Nathan Sheeley¹, Leon Zhang¹, Matthew Shaffer¹, John Long¹, Han Wang¹, Angela Wang¹, Arjun Sabnis¹, Joshua Brot¹, Yun Du¹, Håkan Zeffer¹, Mingran Wang¹, Raghu Prabhakar¹（¹SambaNova Systems） > - **URL:** https://mlsys.org/virtual/2026/oral/3852 > - **OpenReview:** https://openreview.net/forum?id=7wOOhxkuN8 > - **関連:** cloud.sambanova.ai（SambaNova Cloud 本番推論基盤、SN40 Reconfigurable Dataflow Unit） > [!abstract] 概要（論文アブストラクト） > 自己回帰デコードにおけるトークン生成は、メモリ帯域律速という特性から現代の AI ワークフローで主要な性能ボトルネックとなっている。大規模コンテキストウィンドウや chain-of-thought 推論などの手法がこの問題をさらに増幅させている。一般的な GPU アーキテクチャでは、利用可能な HBM 帯域の 21% 程度しか重みと KV キャッシュの読み込みに活用できていない。これは CPU スケジューリングオーバーヘッド、カーネル同期オーバーヘッド、および演算と通信の重複不足が非同期実行の範囲を制限しているためである。先行研究はカーネル融合や GPU 上での非同期実行でこれらのオーバーヘッドに取り組んできたが、多くは単一 GPU に焦点を当て、異なるモデルアーキテクチャへの汎化が不十分である。本論文は、ピークデコード性能を達成するためにデータフローアーキテクチャ――SambaNova SN40 Reconfigurable Dataflow Unit (RDU)――とのハードウェア・ソフトウェア協調設計を記録する。データフロー実行が可能にする 3 つのコンパイラ最適化 KernelLooping・BatchStreaming・ScheduleOffloading は、小規模から大規模、密・疎、MoE、ハイブリッドを含む多様なモデルと注意機構に汎化する。これらの最適化は理論ピークルーフライン性能の 75% 超を広範なオープンソースモデルで達成し、投機的デコーディング使用時に単体モデル比で 6 倍超の高速化を示す。投機的デコーディングは 16 基の SN40 チップで DGX H100 比 1.7 倍高速であり、両システムの HBM 帯域は同等である。記載手法は cloud.sambanova.ai の本番 AI 推論クラウドにデプロイ済みである。 ## 背景と問題設定 ### GPU デコードの非効率性 - 現代の LLM 推論ではコンテキスト長の急増（システムプロンプト、推論トークン、コード・画像・動画トークン）により、デコード段階のメモリ帯域律速が深刻化している。 - DGX H100 上の Llama3.1-8B の Time-Per-Output-Token (TPOT) 内訳（Figure 2）では、GEMM が全体の 31%（8 基 H100）にとどまり、AllReduce が 19%、SDPA が 12%、RMSNorm/SiLU/Mul が 11%、非デコーダ演算が 28% を占める。スケーリングの効率も低く、H100 を 2 基→8 基へ増やしても TPOT は 4,462 μs → 2,864 μs にとどまる（1.56 倍）。 - 根本原因はカーネル境界での強制同期と演算・通信の重複不足である。MoE では AllToAll・AllGather など非重複の集団通信に約 68% の時間を費やすとの報告がある。 - メガカーネルのアプローチはカーネル起動オーバーヘッドを軽減するが、カーネル内部での HBM アクセス排除や集団通信との演算重複は不完全であり、単一 GPU に限定される傾向がある。また CUDA Graph は MoE の動的エキスパートルーティングとは非互換であり、投機的デコーディングの動的スケジュールアンローリングとも相性が悪い。 ### 投機的デコーディングの重要性 - 推論時間演算（テスト時計算）の活用が AI コミュニティで急速に進み、より多くのトークンを生成して精度を向上させるアプローチが注目されている。この手法はモデルパラメータのスケーリングよりスケーラブルであるとされる。 - 投機的デコーディングは小さなドラフトモデルで $k$ 個のトークンを生成し、ターゲットモデルで一括検証する手法である。しかし GPU ではドラフトモデルのようなパラメータの少ないモデルの性能スケールが悪い（Figure 1a: H100 を 2 基→4 基に増やしても 205→265 tokens/s で 1.29 倍のみ）。 - Figure 1b では、8B ドラフト＋70B ターゲットの投機的デコーディングにおいて、H100 は 1 ステップ 45 ms のうち 72% をドラフトモデルに費やしている。 ## SN40 アーキテクチャ概要 - SN40 は 5FF TSMC プロセスの 2 ダイソケットで、ダイあたり 2 演算コア、2 HBM モジュール、3 DDR チャネルを搭載する。 - 1 ソケットあたり BF16 638 TFLOPS、オンチップ SRAM 520 MB、HBM 64 GB（1.6 TB/s）、DDR 1.5 TB（100 GB/s 超）。 - 各コアは 1,040 基の Pattern Compute Unit (PCU) と 1,040 基の Pattern Memory Unit (PMU) で構成される。PCU はシストリックアレイとして行列積を実行し、PMU は 512 KB のプログラマブルオンチップメモリを提供する。 - PCU と PMU はプログラマブルな Reconfigurable Device Network (RDN) でデータを交換する。Address Generation and Coalescing Unit (AGCU) がローカルメモリ（HBM・DDR）、ホストメモリ、リモート SN40 コアへのブリッジとなる。 - チップ間通信はピアツーピア (P2P) プロトコルで HBM・RDMA エンジンを介さず直接オンチップメモリ間でストリームされる。これにより AllReduce 等の集団通信を HBM 帯域を消費せず非同期に構築できる。 - ハードウェアグラフオーケストレーション機構により、ホスト介入なしにカーネル起動を連鎖させることが可能である。 ## 提案手法: 3 つのデータフロー最適化 ### KernelLooping - 現代の LLM はデコーダ層を反復する構造を持つ。GPU (H100) では 1 デコーダ層が 10 個のカーネル呼び出し（K1–K10）に分割されるのに対し、SN40-16 では全デコーダ演算を単一カーネル K0 に融合する。KernelLooping はさらに、この単一カーネルを外側ループでラップし、32 層全デコーダの実行をわずか 4 回のカーネル呼び出し（Embedding、all_decoders_nosync、Classifier、Sampling）に圧縮する（Figure 4b）。H100 では同等の処理に 320 回のカーネル呼び出しが必要である。 - **実装は 2 フェーズ**で構成される: 1. **Pattern Matching**: データフロー解析で互換性のある連続カーネル呼び出しを特定する。連鎖出力テンソルの使用パターンが厳密な使い回し規約に従うことを検証する。 2. **Transformation**: 特定された連続呼び出しをループ付き単一カーネルに置換する。入力の結合・外側ループイテレータの導入・中間バッファの適切なメモリ階層（SRAM・HBM・DRAM）への配置を行う（Figure 6）。 - **データフローハードウェアの恩恵**: 同期アーキテクチャ（SIMT 等）では、カウンタベースの依存性追跡がグローバルメモリ経由の高レイテンシ同期を要し、報告されるオーバーヘッドはサブミリ秒推論のレイテンシ予算の 20% 超を占める。SN40 は軽量制御ハンドシェイク信号で非同期性をネイティブにサポートし、同期レイテンシをナノ秒オーダーに低減する。 - KernelLooping はさらにカーネルの複数インスタンスをパイプライン実行し、デコーダ間のデッドタイムを排除してメモリシステムを常時稼働させる。 - **評価** (Figure 10): 幅広いモデルアーキテクチャ・パラメータ規模・バッチサイズ・シーケンス長で高速化を達成し、幾何平均 1.6 倍である。Qwen2.5-72B のような大規模モデルでは大バッチサイズで 2 倍近くに達し、同期オーバーヘッドが無視できないことを示している。Figure 11 の HBM 帯域利用率トレースでは、KernelLooping（橙色）がベースライン（青色）に比べ同期による帯域低下を解消し、ピーク帯域に近い安定した利用率を維持している。 ### BatchStreaming - KernelLooping でカーネル呼び出しを排除しても、バッチ処理ではデコーダ層末尾で全サンプルの完了を待つ人為的同期境界が残る。この同期はパイプラインのウォームアップコストをデコーダ層ごとに発生させ、時間あたり数マイクロ秒しかかからないドラフトモデルの小規模モデルで顕著に影響する。 - **BatchStreaming** はデコーダ層間にノンブロッキング通信チャネルを導入し、サンプルが非同期にデコーダ層を流れるようにするコンパイラ最適化である（Figure 7）。BatchStreaming なしではバッチ内全サンプルが現在の層を完了してから次層に進むが、BatchStreaming ではサンプル 0 が現在の層を抜けた時点で次層の処理を開始し、パイプラインの全段を常時稼働させる。 - **実装**: オンチップの LoopBuffer が複数 PMU を結合した仮想スクラッチパッドメモリとして機能し、デコーダ層の出力を保持する。書き込みはサンプル完了ごとに適切なオフセットへ行われ、読み出しは後続デコーダが入力を受け取り可能な時点で非同期に行われる。read-after-write 依存関係が強制され、グローバル同期は存在しない。 - **評価** (Figure 12): Llama 1B・3B・8B のドラフトモデルをバッチサイズ 2, 4, 8, 16 で測定。バッチサイズ 2→4→8 で効果が増大し、8B モデル・バッチサイズ 8 で約 1.4 倍の高速化を達成。バッチサイズ 16 でも安定した高速化が持続する。 ### ScheduleOffloading - SN40 ではデコードスケジュールが Processor Execution Format (PEF) に格納される。通常の実行ではトークン生成ごとにホストへ制御が戻り、次の反復を起動する。これはデータフローアーキテクチャの利点を無効化する。 - **ScheduleOffloading** はデコードスケジュールの動的結合をサポートし、$M$ 回の反復分のスケジュールをホスト介入なしに連鎖実行する（Figure 9）。$M$ は推測ウィンドウ $k$ と I/O 特性から経験的に決定され、実用上 $k$ で上界される。ドラフトモデルとターゲットモデルのアンローリング回数 $M$ を独立に設定可能である。 - CUDA Graph も同様に CPU オーバーヘッドを削減するが、MoE の動的エキスパートルーティングには非互換であり、投機的デコーディングの動的スケジュールアンローリングとも静的グラフモデルの制約で完全互換ではない。 - **評価** (Figure 13): ドラフトモデルのみの高速化を測定（$k=9$）。小規模モデル（Qwen2.5-0.5B）でバッチサイズ 1 において約 1.9 倍、Llama 3.2 1B で約 1.5 倍の高速化を達成。ホストオーバーヘッドの比重が大きい小規模モデルほど効果が顕著である。 ## 実験評価 ### 実験設定 - **評価モデル** (Table 2): GPT-OSS 120B (FP8)、DeepSeek R1 671B (FP8)、Llama 4 17B128E (FP8)、Llama 3.2 1B/3B (BF16)、Llama 3.1 8B (BF16)、Llama 3.3 70B・3.1 405B (BF16)、Mixtral 8x7B (BF16)、Qwen2.5 72B・0.5B (BF16) を幅広いバッチサイズ・シーケンス長で評価。 - **ハードウェア** (Table 3): DGX H100 (8 ソケット、BF16 8,000 TFLOPS、HBM 24.0 TB/s) と SN40-16 (16 ソケット、BF16 10,208 TFLOPS、HBM 25.6 TB/s)。両システムの HBM 帯域は同等である。 - DGX H100 の性能は TensorRT-LLM で取得。SN40 の性能は 10 プロンプト（出力トークン 100〜1,000 超）の平均トークン/秒で計測。 ### 投機的デコーディングの受容率 - Table 1 では複数のドラフト/ターゲットモデル組み合わせ、推測ウィンドウ $k=5$ と $k=9$、および 3 つのベンチマーク（AlpacaEval・HumanEval・GSM8K）で受容率を報告。Llama 3.1 70B/8B で $k=5$ 時に AlpacaEval 0.72、HumanEval 0.95、GSM8K 0.83 と、コーディングタスクでの受容率が一般タスクより高い。$k$ を増やすと受容率は低下するがタスクとモデルの組み合わせにより変動幅が異なる。 - ルーフラインモデル（Equation 3）を導入: $\text{Perf}_{RL} = \frac{\text{samples} \times \text{HBM\_bw}}{\text{weights\_size} + \text{KV\_cache} \times \text{samples}}$。HBM 帯域律速下での理論上限スループットを与え、最適化の達成率を評価する指標として用いる。 ### SN40 対 GPU: ルーフラインとの比較 - **ターゲットモデル性能** (Figure 14a): Llama 3.1 70B・シーケンス長 4K において、ベースライン SN40-16 はルーフラインの約 45%、最適化済み SN40-16 は 60〜78% を達成。DGX H100 は OOM（Out of Memory）が発生しない構成で同等以下の達成率にとどまる。405B モデルは DGX H100 ではメモリ不足で実行不能だが、SN40-16 では実行可能かつ 45〜78% のルーフライン達成率を示す。 - **投機的デコーディング性能** (Figure 14b): Llama 3.1 70B/8B の投機的デコーディングで、最適化済み SN40-16 は DGX H100 に対し 60〜80% 高い高速化を達成。 - **時間内訳** (Figure 15a): バッチサイズの増加とともにルーフライン超過分（SN40-Excess）の比率が上がる。これは (A) HBM ベースのルーフラインが捕捉しない演算律速挙動の増加と (B) より大きな I/O オーバーヘッドによる。モデルサイズが大きくなると SN40-Excess と Host 双方の比率が低下し、HBM 利用率がルーフラインに近づく。 ### バッチ投機的デコーディング - Table 4 ではターゲットモデル単体ベースラインに対する正規化高速化を報告。ドラフト/ターゲットのサイズ比が大きいほど高速化が大きく、Llama 405B-3B で BS=4, $k=9$ 時に 6.1 倍、Llama 405B-8B で BS=1, $k=9$ 時に 6.0 倍を達成。 - バッチサイズの増加とともに高速化は減衰する。これは重みの 1 回読み込みコストがサンプル間で償却され、ドラフトモデルの KV キャッシュオーバーヘッドの比率が高まるためである。シーケンス長の増加（4K→32K）も同様にドラフトモデルの KV キャッシュが肥大化するため高速化が低下する。 ## デプロイメント - 全手法は SambaNova Cloud (cloud.sambanova.ai) の全モデルにデプロイ済みである。TTFT（最初のトークンまでの時間）、トークンスループット、コンテキスト長などの主要指標は Artificial Analysis や OpenRouter などのサードパーティベンチマークで公開されている。 - **グラフコンパイル**: KernelLooping と BatchStreaming を含むコンパイル時最適化は無視できるオーバーヘッドで追加される。テンソル並列（TP）下では全デバイスで同一の演算グラフが共有され、ループ化はデコーダ 1 層に削減される。コンパイル時間とメモリはデコーダ数にスケールせず、モデルパラメータにも依存しない。 - **ランタイムメモリ**: ScheduleOffloading の PEF スケジュールメモリはアンローリング回数 $M$ に線形にスケールするが、数 MB にとどまり、$M$ の効果は約 20 で飽和する。HBM（数百 GB）や DRAM（数 TB）に対して無視できる。 - **チェックポイント処理**: KernelLooping では、デコーダ層間で重みを高次元テンソルにパックし HBM から効率的に読み出せるようチェックポイントを再構成する。自動前処理インフラがコンパイラ生成メタデータを読み取り、適切にチェックポイント入力を変換する。 ## 結論 - デコードの非効率性を克服するには実行方式を第一原理から再考する必要がある。データフロー実行は演算・メモリ・通信をチップ内およびチップ間で本質的に重畳させる。 - SambaNova SN40 RDU 上のハードウェア・ソフトウェア協調設計により、コンパイラ駆動の 3 つの最適化――KernelLooping・BatchStreaming・ScheduleOffloading――が小規模密モデルから大規模 MoE まで多様なモデルで理論ルーフラインの 75% 超を達成した。 - 投機的デコーディングのエンドツーエンド結果は、単体モデルベースラインに対し 6 倍超のトークン生成高速化を示し、同等の HBM 帯域を持つ DGX H100 システムに対し 1.7 倍のスループットを達成した。 - これらの手法は本番環境で堅牢に稼働し、cloud.sambanova.ai で商用推論を支えている。