Google、TPU推論の外部化を本格化へ Anthropicが最大利用者に

2026 年 9 月 7 日、SemiAnalysis ニュースレターで最初に公開されました。

業界は過去 10 年以上にわたり、Google が独自の半導体基盤の上に帝国を築く様を見守ってきました。検索、広告、YouTube、そして Gemini の全世代が TPU で動いています。他社設計の環境での性能や経済性がどうなるかについて、これほど建築的な検討や議論を呼んだアクセラレーターは他にありませんでした。Anthropic は 2029 年までに DeepMind の利用規模を超え、TPU の最大ユーザーとなっています。

Google の社内での成功自体は疑問の余地がありません。真に問われていたのは、その優位性のうちどれを他社が実際に活用できるかです。オープンウェイトモデルを採用し、馴染みのある推論エンジンで提供すれば、自社のビジネスにとって重要な経済性において NVIDIA を凌駕できるのでしょうか。

本日、InferenceX 公式プレビューにおいて、TPUv7 Ironwood の第三者による推論結果を初めて公開します。B200/B300 と同等条件で比較したところ、Ironwood はドルあたりのパフォーマンスが最大 50% 向上しています。この優位性はパレート曲線の広範囲にわたり、Google の内部総所有コスト（TCO）と実際の顧客が支払う外部 TCO の両面から経済性を検証します。

Ironwood（TPUv7）は、Google が自社クラウドで直接購入またはレンタル可能なチップを用いて、他社の推論ワークロードにも参入する初の世代です。2025 年 11 月には、Anthropic が TPU を愛用し、100 万個以上（約 40 万個以上の直接購入と GCP 経由の 60 万個以上のレンタル）をコミットしたことを既に報告しました。これらは主にトレーニングに使用されますが、推論にも活用されています。

当社のアクセラレーターモデルには、Anthropic の TPU 出荷数や Google の四半期別 TPU 出荷数の最新データに加え、TPUv8i、v8t、および各種 v9/v10 の見積もりも含まれています。

TPU 用の外部スタックである新しい TorchTPU の開発スピードに、私たちは大きな期待を抱いています。この記事の後半では、推測的デコーディングの最適化や、分離されたプレフィル、KV キャッシュのオフロード、マルチターン型エージェントワークロードなど、TPU ソフトウェアの外部化に必要な今後の取り組みについて詳しく解説します。

それでもなお、SemiAnalysis はTPU の外部化が正しい方向へ進み、全速力で推進されていると確信しています。さらに AMD がまだテストファーストなソフトウェア文化を構築する段階にあるのと対照的に、Google には数十年にわたるソフトウェアエンジニアリングの経験と、極めて確立された品質重視の文化があります。そのため、外部向け TPU ソフトウェアも急速に成熟していくと予想しています。

本記事では、オープンウェイトモデル向けの TPU カーネルおよびサービングスタックに施されたすべての最適化について取り上げます。具体的には、DP アテンションの最適化、MoE のルーティングとカーネル最適化、GDN カーネルにおけるパディング削減などです。さらに、TPU システムの詳細な解説に加え、このスタックを広く利用可能にするために、素晴らしい TPU パフォーマンスエンジニアたちが現在取り組んでいる次のステップについても掘り下げていきます。

Google は過去 10 年以上にわたり、TPU で何を実現できるかを示し続けてきました。今こそ、業界全体がこれらを活用して何を作り出せるかを測る時です。

Google の Chris Chan、Jahangir Hasan、Wangyuan Zhang、Anne Stern、Puneith Kaul、Ruizi Dong、Sangam Jindal、Qi Zhou、Madhan Jaganathan、Gang Ji、Jun Wan、Devanshu Jain、Jiaxin Cao、Srinath Mandalapu、Haowen Ning 氏ら、そして Inferact チームに感謝します。この素晴らしい TPU の基盤とパフォーマンスを実現してくれました。また、TorchTPU と SGLang の開発にも取り組んでいる RadixArk チームにも謝意を表します。

InferenceX オフィシャルプレビュー：TPUv7 Ironwood vs. Blackwell および Blackwell Ultra

現在、次世代のネイティブな TorchTPU vLLM スack が NVIDIA GPU と同条件で比較された際、すでに強力な結果を示しています。Google は初期導入モデルとして、FP8 形式の Qwen3.5 397B を採用しています。後続のセクションでは、外部 TPU 向けの vLLM/SGLang サービングにおいて、新しい TorchTPU アプローチが従来の TorchAX パスと比較してなぜ劇的な改善をもたらすのかを詳しく解説します。

この基盤が整った後、Google は Kimi K3 や GLM 5.3 など他のオープンウェイトモデルへのサポートも拡大する計画です。数々のモデルで最適化が進めば、発売直後（day-0）に幅広い人気オープンモデルの最適化サポートを追加するのは格段に容易になります。現状では vLLM と SGLang は NVIDIA に対して day-0 サポートを集中させており、AMD についてはそれなりの対応がなされていますが、TorchTPU が近い将来十分に安定すれば、vLLM や SGLang のメンテナーが TPU もその day-0 リストに追加する可能性が高いでしょう。このスタックはプライベートベータを終了し、10 月頃にオープンソース化される予定です。

集約型推論と単一トークン予測の同等条件での比較では、FP8 を使用した集約型推論において、TPU は B200 や B300（いずれも FP8 動作）と比較して、1 ドルあたりのパフォーマンスが最大 50% 向上していることが確認されています。一方、NVIDIA GPU で FP4 モデルを推論する場合は、FP8 に比べて品質の低下が見られます。TPUv7 はネイティブな FP4 計算をサポートしていないため、現状では FP4 環境において NVIDIA GPU が依然としてリードしています。しかし、これは TPUv8i の登場で変わります。TPUv8i はネイティブな FP4 サポートを備えており、これにより「Boardfly」構成が Rubin NVL72 と競合できる水準に達すると確信しています。

Google が分散型推論とスペキュレーティブ・デコーディング（MTP 機能を含む）を両側で有効化すれば、1 ドルあたりのパフォーマンスにおける TPU の優位性は持続すると考えています。本記事の後半では、TorchTPU の基盤が整った後に Google が実装する予定の最適化手法についても触れます。初期段階では導入範囲を限定するため、8k1k のベンチマークに焦点を当てていますが、TorchTPU スタックはエージェントワークロードにおいても驚異的なパフォーマンスを発揮すると確信しており、年内にその結果も公開する予定です。さらに本記事の後半では、TPU 外部推論の今後のステップとロードマップについても解説します。

1 ドルあたりのパフォーマンスにおいて TPU が王者

ユーザーあたり 100 トークン/秒のインタラクション速度において、Ironwood のコストは総トークン数 100 万あたり約 0.181 ドルです。一方、B200 は 0.222 ドル、B300 は 0.276 ドルとなります。これは B200 より約 19%、B300 より約 34% のコスト削減を実現しつつ、ユーザーあたりの生成速度は同等に保たれています。

この比較では、ハイパースケール企業が TPU を購入する場合を想定した外部 TCO（総所有コスト）モデルと、B200/B300 の TCO を対比しています。昨年、SemiAnalysis のアクセラレーターモデルが初めて「Google がクラウド経由でのレンタルだけでなく、TPU を直接販売している」という事実を明らかにしました。TPU はスループットが低い傾向にあるものの、多くのインタラクション範囲において、その低時間単価コストがそれを上回っています。SemiAnalysis の TCO モデルでは、TPU の BoM（部品表）見積もりと TCO 推計の詳細な内訳を確認できます。

TPU は、純粋な性能曲線の大部分において NVIDIA GPU を上回るわけではありません。しかし、この比較には TPU の低い総所有コスト（TCO）が考慮されていません。TPU の購入者が最も重視するのは、収益性の高さ、つまりドルあたりのパフォーマンスやワットあたりのパフォーマンスです。

とはいえ、ユーザーあたり 20 トークン/秒の応答性を想定した場合、Ironwood は純粋なスループットでも優位に立ちます。チップあたりの総トークン数は 9,364 に達し、B200 の 8,903 や B300 の 8,925 を上回ります。これは、いずれの GPU よりも約 5% 高いスループットです。さらにモデル化された時間あたりのコストが低いことと相まって、Ironwood は B200 よりドルあたりで 50.4% 多く、B300 より 96.0% 多くのトークンを生成できます。

もし Google の内部ワークロードにおける TPU の TCO をチップあたり 1.03 ドル/時間と仮定すると、並行度 256 でのドルあたりのパフォーマンス優位性は、B200 で 76.7%、B300 で 130.2% に拡大します。ただし、これらの高並行度の結果にはレイテンシのトレードオフが伴います。例えば、並行度 256 の場合、TPU の平均 TTFT（Time to First Token）は 5.41 秒ですが、B200 は 3.75 秒、B300 は 2.40 秒です。前述の 50% から 96% の優位性は、この特定のデータポイントに適用されるものであり、すべてのレイテンシ目標に当てはまるわけではありません。

エンドツーエンドのレイテンシ（応答遅延）に注目すると、TPU は最高スループット点を超えた領域でも依然として競争力を持っています。中位値で 20 秒の応答時間を達成した場合、Pareto フロンティア（最適解の曲線）によると、TPU のコストは 100 万トークンあたり約 0.098 ドルです。一方、B200 は 0.106 ドル、B300 は 0.132 ドルとなります。つまり、TPU は B200 よりも約 8%、B300 よりも約 25% 安価に運用できることになります。

上記のすべての条件を踏まえても、B200 は 30 秒前後の中央値レスポンス時間など、特定の短い応答時間の範囲では依然として有利な立場にあります。しかし、応答時間が長くなると TPU は百万トークンあたりのコストで優位性を回復し、ここで示された重複する範囲全体において B300 に対してもコスト面で有利です。

比較対象のすり替え：GB300 NVL72 の分散型と TPUv7 の集約型

Google は長年にわたり、本番環境で分散型推論（disaggregated serving）を運用しており、その経路は徹底的に最適化されています。一方、外部向けの TPU 推論スタックには、まだ完全に最適化された分散型のパスが用意されていません。そのため、現状では「分散対分散」の比較において、GB200/300 NVL72 の方がパフォーマンス対価格比で競争力があります。

しかし、Google、Inferact、RadixArk、そして SemiAnalysis による最適化が進むことで、この格差は数ヶ月以内に縮まると予想されます。最終的には TPUv7 が GB200/300 NVL72 と互角の競争力を獲得するでしょう。詳細な比較データについては、追って「TPUv7 の分散型対 GB200/300 NVL72 の分散型」という記事で公開予定です。

TPUv7 ポッドは低遅延の ICI ファブリックを通じて 1,000 チップ以上へのスケーリングが可能です。これにより、大規模モデル向けの分散処理や NVIDIA NVL72 では実現できない超広帯域のエッジプロセッシング（EP）のための最適化を可能にします。

なお、「集約型対分散型」という比較対象が異なる場合でも、内部 TCO を基準にすれば、TPUv7 の集約型推論は高遅延領域や低〜中程度の遅延領域において十分に競争力があります。しかし、右側の中間帯域におけるエンドツーエンドの遅延で比較すると、GB300 NVL72 の分散型の方が TPUv7 の集約型に対して有利となります。

GB300 NVL72 の分散型サービングと TPUv7 の集約型サービングを比較すると、レイテンシが低い場合も高い場合も互角ですが、中間帯域では GB300 が TPUv7 に対してコストパフォーマンスで約 30% の優位性を示しています。しかし、TPUv7 における分散型サービングの最適化が進めば、GB300 NVL72 と比較してパレート最適曲線全体において競争力を持つようになるはずです。

新世代のネイティブ・ファーストクラスな TorchTPU バックエンドと、従来の TorchAX スタックにおける推論サービングの違い

従来の TorchAX vLLM バックエンド

Google は以前、TorchAX バックエンドを使用して、vLLM や SGLang 内の PyTorch ベースの実装を JAX に変換していました。この TorchAX は、開発者が PyTorch のモデル定義を JAX を通じて実行できるようにすることを目的としていました。しかし、このアプローチには課題が生じたため、Google と PyTorch、vLLM、SGLang のコミュニティは協力して「TorchTPU」と呼ばれる新しいアプローチを開発しました。これにより、開発者は TPUs をネイティブな PyTorch デバイスとして使用できるようになります。一方、背後では StableHLO、XLA、および TPU 最適化された Pallas カーネルが低レベルの実行を処理します。

次のセクションでは、この新しい TorchTPU バックエンドの性能を、NVIDIA Blackwell および Blackwell Ultra GPU と比較して示していきます。

vLLM の TPU サポートは、3 つの異なる段階を経て進化してきました。

• 最初の TPO プロトタイプでは PyTorch/XLA が使用されました。この遅延実行モデルは、各 PyTorch の操作を即座に実行するのではなく、操作を集めて計算グラフを作成し、それを XLA でコンパイルする仕組みを採用していました。

• 現在の公開バックエンドである tpu-inference は、利用可能な場合は TPU 最適化された JAX 実装を使用します。もし利用できない場合でも、TorchAX が元の PyTorch モデルを JAX で実行可能な操作に変換します。この目的は、すべての PyTorch モデルを書き直すことなく、成熟した TPU プリミティブや並列処理のサポートを再利用することにあります。

次期設計「TorchTPU」により、vLLM は TPUs をネイティブの PyTorch デバイスとして扱えるようになります。今後は SGLang や vLLM の TPU 向けバックエンドとして標準的に採用される見込みです。

以下の図は第 2 段階を拡大したもので、vLLM の tpu-inference バックエンドが、直接 JAX モデルと TorchAX を介して実行される PyTorch モデルをどのように統合しているかを示しています。

当初 JAX を採用したのには実用的な理由がありました。直近の 発表 で Google は、低レベル最適化やページドアテンション（paged attention）の課題、そして vLLM のワーカーモデルを TPU 実行環境に適合させる難しさを説明しています。そこで JAX が持つ成熟した TPU プリミティブと並列処理サポートに期待したのです。TorchAX を活用すれば、PyTorch モデルをすべて書き直すことなくこれらの機能を利用でき、TPU 指向の JAX 実装側でも同じカーネルやコンパイラの作業を共有できるようになります。

TorchAX を使えば、開発者は依然として PyTorch でモデルを記述できますが、実際の計算は TPU 上で JAX が実行します。TorchAX は両者のフレームワークをつなぐ翻訳層として機能します。

TorchAX のドキュメントによると、torchax.tensor.Tensor は通常の PyTorch テンソルのように動作しますが、その裏側では jax.Array がサポートしています。モデルがテンソル操作（PyTorch では ATen 操作と呼ばれる）を呼び出すと、TorchAX が __torch_dispatch__ フックを通じてこれを捕捉し、1 つ以上の JAX 操作に変換します。このため、計算を実行するのは JAX であっても、モデルのコードは PyTorch で記述したままにできます。

vLLM はまた、生成プロセス中に状態が変化する部分を管理する必要があります。特に KV キャッシュは、以前のトークンからのアテンション情報を保存して再利用する重要な役割を果たします。モデルラッパーでは、モデルの重みと KV キャッシュを JAX に対して明示的な状態として提示します。このプロセスは「関数化（functionalization）」と呼ばれ、本来 Python オブジェクト内部で起こる変化を、コンパイラが追跡可能な入力と出力に変換します。その結果、jax.jit は各推論ステップを計算グラフとして捉え、JAX のコンパイルパイプラインを通じて繰り返し実行できるようになります。

各コンポーネントには明確な役割が与えられています。Pallas はパフォーマンスが重要な処理向けに最適化された TPU カーネルを提供します。StableHLO は、コンパイラが処理できる標準形式でプログラムを表現し、XLA がそれを TPU 用の実行可能コードに変換します。

vLLM の標準的な PyTorch パスでは torch.compile がグラフのキャプチャとコンパイルを調整しますが、TorchAX パスでは JAX/XLA パイプラインがすでにこれらのタスクを実行しているため、vLLM は通常の torch.compile ルートを迂回します。これらの翻訳レイヤーは多くの課題を生んできましたが、TorchTPU vLLM スタックはその解決を目指しています。

従来の SGLang JAX バックエンド

SGLang の現在の公開 TPU スタックでは SGLang-JAX を使用しています。SGLang-JAX は vLLM の TorchAX バックエンドと同様のトレードオフを行いますが、アプローチは異なります。これは PyTorch 版 SGLang ランタイムを TorchAX を介して変換したものではなく、JAX ネイティブのサービングエンジンです。SGLang 形式のスケジューリングとプレフィックスキャッシュを JAX モデルおよび TPU 専用カーネルと組み合わせています。

Google と RadixArk は、利用可能なサービングパスを拡張するために、追加の PyTorch ネイティブバックエンドとして SGL-torchtpu を発表しました。

新しいネイティブでファーストクラスの TorchTPU バックエンドへ移行

「TorchTPU」は、フレームワークの境界を PyTorch 自体の中に移すアプローチを採用しています。これは PyTorch の PrivateUse1 バックエンド拡張ポイント を活用し、JAX アレイに依存するラッパーではなく、デバイス指定が tpu の通常の torch.Tensor を公開します。PyTorch のディスパッチャは ATen 演算を TPU バックエンドへルーティングします。

開発者は、初期動作確認やデバッグのためにコードを即時実行（eager）することもできますし、torch.compile を呼び出してグラフキャプチャを行うことも可能です。コンパイルパスでは、TorchDynamo と AOTAutograd が FX グラフを生成し、それを TorchTPU が StableHLO へ変換します。最後に XLA が TPU 向けの実行可能ファイルを生成します。

Google のカンファレンス資料では、このコンパイラの選択が明確に示されています。このパスでは Inductor や Triton ではなく、XLA が使用されます。

「ネイティブ」という用語には、特定の範囲が定義されています。具体的には PyTorch のテンソル、ディスパッチ処理、イーグル実行、コンパイレーションのエントリーポイント、そして分散 API を指します。XLA は引き続きコンパイラとして機能し、カーネル層では TPU 固有の実装が使い続けられます。TorchTPU は Pallas や JAX ベースのカスタムカーネルを呼び出すことが可能です。Helion の TPU バックエンド も同様に Pallas を出力します。ユーザーが直接触れるフレームワークは PyTorch ネイティブなものになりますが、コンパイラやパフォーマンスに直結するカーネルは TPU 対応のままです。

開発者は .to("tpu") を使用しながらも、慣れ親しんだ分散インターフェースやサービングエンジンのコードを維持できます。Google は DDP、FSDP2、DTensor、1 プロセス per デバイス、そして MPMD と SPMD の両方の実行に対応するドキュメントを提供しています。これらのインターフェースは、PyTorch サービングエンジンがすでに前提としているプロセス構成と分散モデルに適合しています。

vLLM や SGLang にとっては、PyTorch から JAX への境界をまたいで再構築する必要なく、より多くのアップストリームのモデルコードやスケジューラ、連続バッチ処理、API、機能ロジックを流用できる点が大きな機会です。最終的な TPU 実装には依然として専用カーネルが必要になる場合でも、モデルやエンジン機能の立ち上げコストを削減できるはずです。

ネイティブな PyTorch のサポートが導入されたとしても、TPU 固有の最適化が必要なくなるわけではありません。Pallas と TorchAX は異なる目的のために存在します。TorchAX は PyTorch のモデルコードを JAX を介して実行可能にするものであり、Pallas は TPU ハードウェア向けにパフォーマンスが重要な演算を実装するために使われます。TorchAX を削除したからといって、その背後にある Pallas カーネルまで消えるわけではありません。TorchTPU によりネイティブな PyTorch コードから Pallas や JAX のカーネルを呼び出せるようになったことで、以前のスタック向けに開発されたカーネルも新しいスタックへ移行できるようになります。

GPU 環境での PyTorch vLLM と同様、高性能推論にはカスタムカーネルが依然として必要です。TPU の行列演算ユニットを効率的に利用するには、テンソルの形状やレイアウトの調整が不可欠です。既存の Pallas カーネルの多くは新しいネイティブな TorchTPU サービングスタックへ容易に移行できると考えていますが、すべてのカーネルが変更なしでそのまま使えるわけではありません。ラッパー関数、テンソルレイアウト、ランタイムとの統合など、いくつかの要素には微調整と検証が必要になる可能性があります。

Google TPU チームとの強力なエコシステム連携

vLLM と SGLang は、主要なオープンソースの生産環境向け推論サービングスタックです。Inferact、RadixArk、Red Hat は、Google と協力して TorchTPU バックエンドを両方のスタックでファーストクラスエクスペリエンスとして確立することに注力しています。このコミュニティからの支援は、TPU での推論を外部化していく上で極めて重要です。

9 月初旬現在、TorchTPU はまだベータ版（限定公開）の段階ですが、PyTorch Conference の開催を控えた 10 月中旬にオープンソース化される予定です。これにより、TPU が PyTorch、vLLM、SGLang において第一級の体験を提供できるようになることを楽しみにしています。TorchTPU vLLM と TorchTPU SGLang が公開されたら、InferenceX/AgentX 向けの TPU ベンチマークを、私たちが管理するフォークから公式リポジトリへ移行します。

TPU 推論最適化の深掘り

ネイティブな TorchTPU サービングスタックの最初のモデル立ち上げに向けた、TPU 推論の最適化 sprint（短期集中開発）には、数百時間に及ぶエンジニアリング作業と多数の PR が投入されました。後述する通り、TPU のアーキテクチャは GPU と異なるため、モデルを効率的に実行するには独自の最適化セットが必要です。以下では、Ironwood 上で Qwen3.5-397B やその他のモデルのパフォーマンス向上（あるいは一部の場合には機能実現）に寄与した最適化手法について解説します。Pallas カーネルの最適化の多くは、新しいネイティブな TorchTPU スタックへも転用可能です。

DP Attention 最適化

TPU に限らず、モデルアーキテクチャの並列性をハードウェアにどう効果的にマッピングするかが課題です。例えば Qwen3.5 の GQA アテンション層では、クエリヘッドが 32 個ある一方で共有 KV ヘッドは 2 つだけです。TP8（論理デバイス 8 台）環境では、クエリ計算を 8 台のデバイスに均等に分散でき、1 台あたり 4 個ずつ処理できます。一方、KV ヘッドは各デバイスにきれいに割り切れません。各 KV ヘッドは 16 個のクエリヘッドで共有されるため、4 台のデバイスが同じ KV データを必要としてローカルアテンション計算を行います。TP サイズが KV ヘッド数を超えると、標準的な対応策は TP ランク間で KV ヘッドを複製することです。vLLM はすでにこの機能をサポートしています。TPU バックエンドでは、既存の動作を有効化し、不要な All-to-All 通信を防ぐために互換性修正が必要でした（詳細）。

高同時実行処理のためのサービスでは、TPU バックエンドは 8 方向アテンションデータ並列化 (DP8) と 8 方向エキスパート並列化 (EP8) を組み合わせたサポートを追加しました。これは単に「DP アテンション」または DEP8 と呼ばれています。

すべてのリクエストのアテンションを 8 つのデバイス全体で分割するのではなく、各デバイスは異なるリクエストのサブセットを処理し、そのリクエストに対する KV ヘッダを両方ともローカルに保持します。アテンション重みは複製されますが、KV キャッシュには同じ履歴の繰り返しコピーではなく、異なるリクエスト履歴が格納されます。

一方、512 のルーティング済みエキスパートはデバイス全体に分散されたままとなり、より大規模なエキスパート重みプール全体の複製を回避しています。これを実現するには、サービスエンジンにおける リクエスト割り当て、再帰状態スロット、ブロックテーブルの調整 が必要でした。

コミュニケーションの最適化

DP アテンションと EP を採用した TPU の GroupedGEMM 実装では、エキスパートが動作する前にトークン活性化値やルーティングメタデータをすべて集約（all-gather）し、その後、reduce-scatter を用いて重み付けされた出力を合算して、各トークンの結果を対応するアテンションランクに返します。

当初のバックエンドでは、選択されたエキスパート ID とルーティング重みを別々に収集していました。Google はこれらを 1 つの all-gather に統合し（PR）、小さな配列における転送時間を支配しがちだった追加的な集合操作を削除しました。

この PR では、DeepSeek-V3 の測定結果に基づき、レイヤーあたり約 80 マイクロ秒の削減が報告されています。80 µs は一見小さく思えますが、DeepSeek-V3/R1 にはこの AllGather を必要とするレイヤーが 58 層あります。これらを合算すると、前向きパス（forward pass）あたり約 4.64 ミリ秒の削減効果があり（80 µs × 58）、TPOT の短縮と対話性の向上につながります。

Google はまた、データ移動のような不規則な操作に適した「ReduceScatter」集合演算を SparseCore 上で実装しました。これにより、各チップ内で Ironwood の高速なダイ間リンクを活用して寄与値を結合し、チップ間の部分和の交換を行うことが可能になります。この 2 デバイス構成とチップ間ネットワーク（ICI）の詳細は後述します。

ダブルバッファリングを採用することで、カーネルは一つのチャンクの転送中に別のチャンクを蓄積できます。これにより、ローカルな集約処理やダイ間の転送が、より遅いチップ間通信と重なるように最適化されます。また、SparseCore で集合演算を実行することで、TensorCore の実行リソースを他の操作に解放することも可能です。

上記の図に示すように、通信ステージが重なることでエンドツーエンドの実行時間が短縮されます。例えば、t₁という時点では、MB（マイクロバッチ）1 に対するチップ内 DMA ScatterReduce (P1) が完了しています。これにより、次の MB に対する P1 と並行して、ICI dim 0 を跨ぐ ScatterReduce (P2.0) が実行可能になります。

この最適化をベースラインと比較すると、並列度 64 から 512 の範囲でスループットが 4.1%〜14.2%向上 します。具体的には、並列度 256 で 8.5%、並列度 512 の 1k8k モデルで 26.1%の向上が確認されています。

この集約処理を TensorCore から SparseCore へ移動させることで、集約処理実行中に TensorCore で追加の作業をパイプライン化できるようになります。ただし、すべての集約処理を移動させることが常に最適とは限りません。場合によっては、SparseCore へのオフロードが逆にパフォーマンスを低下させることもあります。

例えば Qwen3.5 では、すべてのアロリデュースとオールギャザー操作を SparseCore にオフロードするタイミングを決定する閾値 が導入されています。これは、VMEM に収まる小規模な集約処理は TensorCore 上で実行し、SparseCore へのオフロードオーバーヘッドによって遅くなるのを防ぐためです。デフォルトの閾値は VMEM の容量に基づいて設定されています。

この PR では、並列度 64 でスループットが 2.7%、並列度 128 で 5.7%向上したと報告されています。

MoE ルーティングとエキスパートカーネルの最適化

Mixture-of-Experts (MoE) レイヤーは、各エキスパートに対して不規則なトークンのグループを生成します。これらの「ラギッド（不均一）」なグループを、TPU の行列演算ユニットが効率的に処理できる形式に変形する必要があります。

以下に示す改善の一部は、MoE モデル全般のルーティングやグループ化された行列乗算 (grouped-matmul) バックエンドに対する汎用的な改良です。一方で、Qwen3.5 に対して直接測定された効果も含まれています。

grouped matmul の第 2 版 では、エキスパート入力を MXU (Matrix eXecution Unit) に供給する方法を根本から再設計しました。これにより、冗長なタイル計算が排除され、転送サイズがパディングされた最大値ではなく有効な行数に合わせて最適化されます。また、エキスパート重みはトリプルバッファリングされるため、現在のグループの演算中に次のグループの重みが既に転送経路に載っています。さらに、グループメタデータの生成もカーネル内に統合されました。

不規則なエキスパート入力の再配置処理は、その後 SparseCore へ移管されました。SparseCore がデータ転送を担当し、各エキスパートのトークンを連続したグループに集約する一方、TensorCore はエキスパート行列乗算の実行に専念します。

その後の書き換え ではメモリアクセスのパイプライン化が改善され、結合処理 (combine work) をトークンと隠れ次元 (hidden dimensions) に分割して並列化しました。この PR によると、元の SparseCore カーネルと比較して、8k1k のサービススループットは 12% 向上し、TTFT (Time To First Token) と TPOT (Token Per Output Token) も低下しています。

さらに、不規則な gather-reduce パス内の top-k 重み収集を SparseCore に移動する という最適化が行われました。ここで言う「gather（収集）」とは、1 つのデバイス上のメモリから特定のエントリを読み込むことを指し、デバイス間での全結合 gather を意味しません。従来は、TensorCore が前処理段階でルーティング重みとソースインデックスを収集していたため、TensorCore と SparseCore の並行実行が制限されていました。これらの収集処理を SparseCore カーネル内に移すことで、DeepSeek-V3 のマイクロベンチマーク（バッチサイズ 2k、16 分割の専門家並列）において、TensorCore のオーバーヘッドは 29 µs から 14 µs に、全体の操作レイテンシは 146 µs から 137 µs に短縮されました。

小規模バッチの場合、一般的な不規則パスでは処理自体よりもオーバヘッドの方が大きくなる傾向があります。そこで、専用の小規模バッチ用置換 が導入され、ワンホット行列を構築して通常の行列乗算を用いてトークンを専門家へ配置し、結果を逆変換する処理が行われています。8k1k のワークロードでは、並行度 64 でスループットが 7.3% 向上し、並行度 128 では 5.1% の改善が見られました。

WIP 変更により、専門家の ID とトークンインデックスが単一のソートキーにパックされるようになりました (https://github.com/vllm-project/tpu-inference/pull/3488)。これにより XLA は必要な順序を保ちつつ、より単純なソート処理を実行できるようになり、ソートの遅延は 106.6 µs から 21.7 µs に短縮されました。この PR では 8k1k のサービングで 0.6% から 8.5% の性能向上が報告されていますが、同様の改善により FP8 all-gather も可能になります。

Gated DeltaNet Pallas カーネルの最適化

Gated DeltaNet (GDN) は再帰計算です。つまり、各ステップで実行状態が減衰し、ランク 1 の項によって更新された後、出力を生成するために射影されます。GDN などの線形アテンションメカニズムについて詳しく知りたい場合は、以下の記事をご覧ください。

Kimi K3: The Manos, The Mythos, The Legendos — Kimi K3 は発表と同時に世界を席巻し、リーダーボードを席捲してオープンフロンティアモデルとしての地位を確立しました。コミュニティは Kimi K3 の仕組みを理解することに熱心ですが、その性能を支える非伝統的な技術の数々には驚きの声も上がっています。この記事は、Kimi K3 モデルアーキテクチャの中核となる技術を理解するための入門編です。

以下の変更は、行列演算、ベクトル更新、状態転送が MXU、VPU、VMEM、HBM 上でどのようにスケジューリングされるかを示しています。初期の Qwen3.5 サポート では、因果関係を持つ Conv1D と GDN の純粋な JAX 実装が追加され、TPU オペレータディスパッチに接続されるとともに、再帰状態のキャッシュ機能が有効化されました。その後の PR でこれらの実装はさらに最適化されています。

もう一つの即効性のある最適化は、出力投影計算の代数を再構成して MXU と VPU の処理を重畳させる ことです。従来、VPU がまず減衰した状態にランク 1 更新を適用し、その後 MXU が更新された状態とクエリを乗算して出力を生成していました。この依存関係により、MXU は待機させられていました。

減衰した状態を S、補正ベクトルを Δ、キーとクエリをそれぞれ k と q とすると、出力は以下のように展開できます：

MXU は、減衰した状態から直接 Sq を計算できるようになり、VPU が次のトークン用の更新された状態を構築します。現在の更新が出力に与える寄与は、kᵀq（ヘッドごとのスカラー内積）を用いて別々に計算され、その後、小さなスケーリングベクトル加算が行われます。これにより、MXU の依存パスから状態更新が外されます。報告された 8k1k のスループット向上率は、並列度 64 で 2.79%、並列度 512 で 4.48% です。

次の変更では、ベクトルレジスタのスパイルを削減するために、デコードループ内で Q と K をスライスし、同時に生きている値の数を減らします。これにより decode-64 カーネルは約 20% 高速化されますが、デコードステップ全体の一部に過ぎないため、エンドツーエンドでの向上率は小さくなります。8k1k では 0.8%、1k8k では並列度 512 で 3.8% です。

非同期状態転送 Asynchronous state transfers は、二重バッファリングを活用して DMA と計算処理を並行実行します。当初、第 2 のバッファセットによる VMEM 使用量の増加がデータ並列アテンションの性能低下を招きましたが、既存のスクラッチバッファの再利用と一時変数の寿命短縮によって容量を回復し、この問題も解消されました。その結果、並行度 512 でスループットが 11.3% 向上しました。

GDN v3 は、Conv1D と GDN を単一のカーネルに統合して HBM の往復回数を削減し、プリフィルレイアウトを改善するとともに、混合されたプリフィルとデコードの実行パスを統一しました。報告されているカーネルレベルの高速化率は、デコードで 1.41 倍、プリフィルで 1.60 倍、混合バッチでは 2.14 倍です。ただし、これらの数値はカーネル単体の測定であり、エンドツーエンドのサービス性能向上を保証するものではありません。

カーネル内部の処理間の重なりを改善するだけで、大きな性能向上が実現できます！

ハイブリッド状態とページドアテンションの管理

Qwen3.5 は、GQA レイヤーがトークンごとに蓄積される KV ヒストリを持ち、GDN レイヤーが各リクエストに対して固定サイズの再帰状態を保持するという 2 つの状態を持つハイブリッドモデルです。HBM の実質的な利用可能量やアテンションの効率的な実行方法は、状態割り当て、ストレージ精度、アテンションページサイズ、物理データレイアウトによって決定されます。

Batched ragged paged attention はシーケンスをバッチ化し、ページメタデータを事前計算して 3 バッファリングを行うことでパイプライン処理の改善とパディング削減を実現します。この変更は、共有アテンションバックエンドの基盤を提供するものです。PR のワークロード例では Qwen3-32B が使用されているため、その測定結果は Qwen3.5-397B の結果とは区別して扱う必要があります。

再帰状態は現在 コンパクトに割り当てられる ようになりました。以前は各レイヤーグループに対して num_blocks スロットが必要でしたが、現在はアクティブなリクエストごとに約 1 つのスロットで済むようになります。報告された構成ではこれにより HBM が約 76 GiB 節約され、アテンションブロックプールが 71% 拡大しました。その結果、コンカレンシー 64 で 1k8k の出力スループットが 18% 向上しています。

再帰状態を BF16 で保存する ことで、HBM 上のフットプリントは半減しつつ、VMEM 内での演算は FP32 のまま維持されます。これによりメモリ容量と転送帯域幅を節約できながら、FP32 演算の精度も保たれます。報告された 1k8k スループットの向上率は、コンカレンシー 512 で 15% です。

古いハイブリッドページサイズのアラインメント制約を撤廃することで、バッチ処理されたアテンション計算において適切な 2 のべき乗のページサイズ（256 トークンなど）を利用できるようになりました。これにより、並行度 512 でスループットが約 7% 向上します。この効果はプレフィックスキャッシュを使用しないパスに適用され、後述するアラインメント済みチェックポイントモードとは区別されます。

不要な KV レイアウトのコピーは、他の形状にのみ適用されるレイアウト制約を課すのではなく、KV ヘッド数に基づいてリシェープパスを選択することで回避されています。これにより、Qwen3.5 8k1k のスループットが並行度 512 で約 4.1% 向上します。

パディングの削減と低並行時のオーバーヘッド軽減

実行中のリクエスト数が 4 または 8 台だけの場合、数百の並行リクエスト向けに調整された設定では作業が無駄になります。コンパイル済みの形状バケットが大きすぎたり、メタデータが最大値に合わせてサイズ設定されすぎたり、パディングトークンによってダミーのエキスパート処理が発生したりします。InferenceX の運用ポイントはまさにこうした低並行領域に位置しており、これが以下の改善策を導いた背景です。

リクエストメタデータのバケッティングが、常に設定された最大値を使用するのではなく、アクティブなリクエスト数に基づいて行われるようになりました。8k1k 構成で並列度 64 のテストでは、GDN スケジューリングのオーバーヘッドが 283 マイクロ秒から 97 マイクロ秒に低下し、スループットはチップあたり秒間 2,328 トークンから 2,516 トークンへと向上しました。

Explicit Qwen3.5 tuning for InferenceX により、回転行列テーブルが縮小され、DP ランクごとにシーケンス制限サイズが調整されました。また、並列度 4 に特化したアテンションバケットも追加されています。これにより、並列度 4 における総合的な性能向上率は、8k1k で 13.3%、1k1k で 15.5% となりました。

さらに低並列度向けのチューニング round of low-concurrency tuning では、TP8 アテンションとエキスパート並列処理へ切り替えられ、最小トークンバケットが削減されました。パディングトークンはエキスパートゼロにルーティングされるため、追加のエキスパート重み読み込みトリガーを回避できます。これにより、1k1k 構成では並列度 4 で 22.9%、並列度 8 で 18.1% の性能向上が実現しました。一方、8k1k 構成では並列度 4 でスループットが 9.2% 向上しますが、並列度 8 では 5.3% 低下します。

ハイブリッドモデルでのプレフィックスキャッシングの有効化

上記の性能向上は、リクエスト間で共有されるデータがないランダム入力ベンチマークで測定されたものです。一方、プリフィックスキャッシングは、長いシステムプロンプトや会話履歴を再利用するエージェント型や複数回の対話ワークロードにおいて不可欠です。

ハイブリッドモデルでは、キャッシュされたプリフィックスには KV ブロックに加え、その末尾にある GDN の再帰状態も保持されなければなりません。通常、リクエストが継続するとこの再帰状態はすぐに上書きされてしまうため、「通常の」アテンションモデルに比べてプリフィックスキャッシングの実装は困難を極めます。

DP サポート付きのハイブリッドプリフィックスキャッシング は、GDN にチェックポイントの読み込み用と生状態の書き込み用の別々のスロットを与えることでこの課題を解決します。これにより、リクエストを続行してもキャッシュされたプリフィックスに対応するチェックポイントが上書きされることはなくなります。状態アドレスは、すでに KV ブロックの位置特定に使用されているブロックテーブルと同じ構造から導出されます。また、チェックポイントはアライメントされたキャッシュ粒度で取得されるため、保存された状態は常に KV ブロックの境界と一致します。

このモードでは、前述のようなコンパクトなリクエストごとの割り当てではなく、フルサイズのチェックポイントプールが必要となります。これは、リクエスト間での再利用のために HBM の一部を犠牲にするトレードオフですが、プリフィックスが実際に共有される場合にその恩恵が発揮されます。

ページドアテンションにおけるレイアウトとパイプライン深度

TPU 上で KV キャッシュをどのように配置するかは、2 つのハードウェア特性によって決定されます。以下のセクションで TPU ハードウェアの詳細について解説しますが、まず重要な点は、ベクトルユニットが処理するタイルの最後尾次元が 128 ラーン幅であることです。そのため、後続次元がこれより小さい配列はパディングが必要となります。