TPUはなぜ流行らないのか — トークン単価が安いのに誰も使わない理由

スペックシートだけ見るとTPUの話は圧倒的だ。トークン単価が安い、ワットあたりの処理量が桁違いに良い、レイテンシが決定論的。Trainiumも同じ物語を語る。にもかかわらず、業界の相当部分は——ChatGPTのような一般向けチャットUIの裏側の推論トラフィックの多くを含めて——いまだにNVIDIAで動いている。

「紙の上では安い」と「実際にデプロイされているもの」のこの乖離は、マーケティングの失敗ではない。シストリックアレイ型のシリコンが、**コード・パイプライン・組織構造の3つに課してくる「税金」**そのものだ。この記事では、その税金がどこから来るのか、そしてなぜごく一握りの会社しかそれを払えないのかを掘り下げる。

すべてを説明する1つの事実：静的シェイプ

NVIDIA GPUはSIMT（Single Instruction, Multiple Threads）プロセッサだ。実行時にスレッドを動的にスケジューリングし、メモリをオンデマンドでページングする。一方、TPUやAWS TrainiumはGPUではない。シストリックアレイ——隣接する演算器（MAC）同士が物理的な銅線で直結された格子状の構造で、事前コンパイラ（TPUならXLA、TrainiumならNeuronコンパイラ）がデータを流し込む。

シストリックアレイがピーク性能を出すのは、流れるデータの形状（シェイプ）がコンパイル時に固定されているときだけだ。重み（Weight）は最初に一度ロードされたら演算器内に静止（Stationary）し、入力（Activation）だけがバケツリレーのように隣へスライドしていく。シーケンス長やバッチサイズが1トークンでも変わると、データの経路とメモリアドレスを全部再計算する必要がある——つまり、コンパイラが新しいバイナリを生成し直すことになる。

このたった1つの制約が、下流のすべての苦しみの源だ。推論時に何が起きるか：

実行時の入力	NVIDIA（動的）	TPU / Trainium（静的）
コンパイル済みバケットより大きい	動的確保で処理	Shape Mismatchでクラッシュ
バケットより小さい	無駄なく処理	JIT再コンパイルでストール（数分）またはゼロ埋めで空転
未知の長さ	そのまま動く	対応バイナリが存在しないとストール

トークンがチップに到達する前に、「これはどんな形状で、どのコンパイル済みバイナリにルーティングするのか」に答えなければならない。NVIDIAではこの問いを立てる必要すらない。

動的 vs 静的のアナロジー：Python vs Java

一番きれいなメンタルモデルはこれだ。NVIDIAはPython、TPU/TrainiumはJava。

• NVIDIA = Python。 動的型付け ≒ 動的シェイプ。ランタイムがカオスを吸収する。100トークンのプロンプトだろうが50,000トークンだろうが、同じforwardに放り込めば「それなりに速く」動く。コンパイルの工程が目の前に出てこない。
• TPU/Trainium = Java。 静的型付け ≒ 静的シェイプ。固定バイナリ（NeuronなNEFF、TPUならXLA実行ファイル）にコンパイルされるまで1行も動かない。ボイラープレートと厳格な規律と引き換えに、すべてが「契約」に収まったときの極限の実行効率を手に入れる。

ちなみにAMDのInstinct系（CDNA / ROCm）は完全にNVIDIA（Python）側だ。SIMT、動的シェイプ、PagedAttention対応、そして既存のCUDAコードをそのまま動かすためのHIPIFYツールチェーン。本当の断層はベンダーのロゴではなく、静的か動的かにある。

「静的ハードで動的入力を捌く」がコード上で何を意味するか

3人のユーザーが同時に来たとする。3,000 / 4,000 / 1,000トークン。NVIDIAならパディングもマスク生成もしない。フラットな8,000トークンのバッファに連結して、境界を示すcu_seqlensインデックスをFlashAttentionに渡すだけだ。

# NVIDIA: 可変長アテンション。パディングなし、マスク行列なし。
# フラットなデータ + 累積シーケンス長 [0, 3000, 7000, 8000] を渡すだけ。
outputs = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q, cu_seqlens_k,
    max_seqlen_q, max_seqlen_k,
)

カーネルが境界インデックスを見て、各ユーザーの文脈をハードウェアレベルで隔離する。ユーザー間アテンションの無駄なFLOPsはゼロ。コードは「ただのモデルロジック」だ。

TPUではシストリックアレイの形を変えられないので、逆をやる。すべてを固定の[batch, STATIC_SEQ_LEN]の四角形に押し込み、計算したくない部分を数式で消す。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_xla.core.xla_model as xm

class StaticShapeAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads, self.d_k = n_heads, d_model // n_heads
        self.q = nn.Linear(d_model, d_model)
        self.k = nn.Linear(d_model, d_model)
        self.v = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, attention_mask):
        # x は常に [batch, STATIC_SEQ_LEN, d_model]。形状は絶対に変わらない。
        b, s, _ = x.size()
        q = self.q(x).view(b, s, self.n_heads, self.d_k).transpose(1, 2)
        k = self.k(x).view(b, s, self.n_heads, self.d_k).transpose(1, 2)
        v = self.v(x).view(b, s, self.n_heads, self.d_k).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)

        # シストリックアレイは全セルを既に計算済み。パディングも他人の領域も含めて。
        # それを後から消す: e^(-1e9) -> 0
        scores = scores.masked_fill(attention_mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)

        ctx = torch.matmul(attn, v).transpose(1, 2).contiguous().view(b, s, -1)
        return self.out(ctx)

このコードをXLA上で動かす際の2点は、静的シリコンの純粋な帰結だ。

1. xm.mark_step() こそが実際の実行トリガー。 冒頭のimport torch_xlaは飾りではない。CUDAのeager modeと違い、XLAでmodel(x)を叩いてもグラフが蓄積されるだけ。mark_step()——これはforwardの中ではなくサービングのループ側で呼ぶ——が溜まったグラフを1つの固定バイナリにコンパイルして送り込むまで、チップ上では何も走らない。新しい形状が来れば新しいコンパイルだ。（近年のPyTorch/XLAはこれを隠すeager modeを追加しているが、「形状ごとにコンパイル」という根っこのモデルは変わらない。）
2. masked_fill(..., -1e9) は最適化ではなくハック。 NVIDIAのvarlenはユーザー間の掛け算をそもそもスキップする。シストリックアレイはスキップできない。四角形の全セル（ゼロも含む）を律儀に掛け算してから、softmaxで数学的に潰す。電力を消費してから結果を捨てている。

「入力が小さすぎる」という罠

オーバーフロー側（はみ出しクラッシュ）は直感的だ。1,024用にコンパイルしたバイナリに1,025トークンを流せばShape Mismatchで落ちる。厄介なのはアンダーフロー——1,024のシステムに100トークンのリクエストが来た場合だ。

• そのまま流す: XLAが新しい形状と認識してJIT再コンパイルを起動。本番では数分のフリーズになる。ストール。
• 1,024にパディング: アレイは演算器の約9割で0 × 0 + 0を律儀に実行し、何もしないために全電力を消費する。効率が暴落する。

逃げ道がパッキングだ。1バケット1ユーザーではなく、複数ユーザーのリクエストをテトリスのように固定の四角形に詰め込み、ユーザー間でアテンションが漏れないようセグメントIDのマスクを生成する。

固定バケット [ 8192 トークン ]
├─ ユーザーA クエリ (3000)
├─ ユーザーB クエリ (4000)
├─ ユーザーC クエリ (1000)
└─ パディング        (192)   <-- 無駄はこれだけ

「この四角形」が物理的に何なのかを具体的にするとわかりやすい。BATCH_SIZE = 4, STATIC_SEQ_LEN = 8192でコンパイルすると、XLAはTPUのHBM上に1つの地続きの[4, 8192]静的領域を予約する。「8192の部屋が4つ独立して並ぶ」のではなく、コンパイラが演算回路のルートを固定する1枚の巨大なシートだ。1ユーザーで8,192レーンを1本使い切ることすら稀なので、サービング層はこの4本のレーンに複数ユーザーを同時に詰め込む。

【1基のTPUプロセッサ = 1枚の静的な [4 x 8192] シート】

レーン[0] (8192): [ A(2000) + B(5000) + C(1000) + 補間(192) ]
レーン[1] (8192): [ D(8000)                      + 補間(192) ]
レーン[2] (8192): [ E(3000) + F(3000) + G(2100)  + 補間(92)  ]
レーン[3] (8192): [ H(4000) + I(4000)            + 補間(192) ]

物理的には4本のレーン（計32K空間）だが、論理的にはプロキシが**9人のラギッドなユーザー（A〜I）**をそこに圧入しただけだ。アプリ側から見ると「1基のTPUが多数の小さなリクエストを並列で同時に捌いている」ように見えるが、実態は1枚の硬直したシートにセグメントマスクを被せているだけ。ハードが「最初から小部屋を切る」のではなく「1枚の太いシート」を欲しがる理由はシストリックアレイの物理だ。行列がデカいほどアレイの充填率が上がり、データ供給の合間の空転が減る。

うまくやればMFU（Model FLOPs Utilization）は、よく調整されたLLMサービングが実際に到達する50〜60%台まで上がる（PyTorch/XLAはLlama 2 70BのTPU学習で約53%のMFUを報告している）——1バケット1ユーザーの素朴な方式が一桁まで暴落するのと対照的だ。100%は誰も触れない天井で、要点は「パッキングが損失の大半を取り戻す」ことだ。だが、ここで何を作ったかに注目してほしい。クラスターの手前に、ラギッドな入力をキャッチしてリアルタイムで四角形に詰め込むためだけの、高速なGo/C++製プロキシだ。NVIDIAではこのレイヤーがそもそも存在しない。

1つの関数じゃない——パイプライン全体が分岐する

xm.xla_device()が共有のOpenXLA/PJRTランタイム（TPUならlibtpu.so、Neuronならlibneuronpjrt.so）のおかげでTPUにもTrainiumにも透過的に対応するので、「torch_xlaがハードウェアを抽象化してくれる」と思いがちだ。model.to(device)や基本演算については本当だ。だが、肝心な部分については完全に嘘になる。

forwardのシグネチャからして分岐する。

# NVIDIA forward: ラギッドなデータ + 境界インデックス。長さは毎回任意。
def forward(self, input_ids, cu_seqlens, max_seqlen):
    return self.flash_attn_func(input_ids, cu_seqlens, max_seqlen)

# 静的 forward: 固定の四角形 + 自前で組み立てるマスク行列。
def forward(self, input_ids, attention_mask):  # input_ids は [batch, FixedSeqLen]
    return self.static_attn_func(input_ids, attention_mask)

そしてこれは最下層までカスケードしていく。

コンポーネント	NVIDIA パイプライン	Trainium パイプライン
推論エンジン	vLLM (CUDA), TensorRT-LLM	NxD / vllm-neuron
カスタムカーネル	Triton, CUDA C++ (FlashAttention)	NKI (Neuron Kernel Interface)、ゼロから書き直し
ベースイメージ	nvcr.io/nvidia/pytorch	AWS Neuron DLC
CIビルド成果物	weight + CUDA/Tritonバイナリ	weight + バケットごとのNEFF静的バイナリ
デプロイ先	g5 / p5 インスタンス	trn1 / inf2 インスタンス
監視	nvidia-smi, DCGM exporter	neuron-top, Neuron exporter

完全に並行する2つの世界だ。CUDAコンテナも、evalスクリプトも、オートスケーリングのトリガーも、何ひとつ流用できない。vLLMのハードウェアプラグイン機構がビジネスロジック層に「皮一枚」の共通化を与えてくれるが、その下のエンジンは100%別物のコードで、別物のバグを持っている。

データ型がさらに状況を悪くする

データ型の話も対称ではない。BF16（Google初代TPUが提唱した）は両者で安定する。FP32と同じ指数の範囲を持つので、-1e9のマスク値を食らってもNaN化しない。だがFP8——現在のスループットの主力——はNVIDIA有利だ。FP8のアテンションスコアは激しく振れるので、クリッピングを防ぐために実行時の動的スケーリングが要る。静的コンパイラはコンパイル時に固定のスケール係数を焼き込むしかないので、TPU/Trainiumで攻めたFP8アテンションを動かすとクリッピングでモデルの賢さ（Perplexity）が劣化するリスクが高まる。「FP8に切り替えよう」はNVIDIAなら一行、静的シリコンなら研究プロジェクトだ。

隠れたコスト：組織が壊れる

ここが導入を殺す部分で、誰もスライドに書かない。NVIDIAにはきれいな抽象化の境界がある。

[ AIエンジニア / データサイエンティスト ]
   アーキテクチャ、ハイパーパラメータ、Eval
        │
        ▼  境界線: Hugging Face形式のweight / 標準PyTorch
        │
[ MLOps / LLMOps エンジニア ]
   vLLMに載せ、PagedAttentionを設定し、スケールアウト

データサイエンティストはメモリ配置を考えなくていい。MLOpsエンジニアはアテンションの数式を読まなくていい。きれいなインターフェースを介して成果物を受け渡す。

TPUに移った瞬間、この壁が消える。モデル構造が物理制約に直結するからだ。

• パッキングの方式（MLOps側）とforward内のセグメントマスクのロジック（AIエンジニア側）は、1つの設計の両面だ。バッチの組み方を変えれば数式も同時に変えなければならない。仕様書ベースで別々の人間に分業させるのは不可能。
• AIエンジニアが気軽にif分岐を足したりレイヤー数を変えたりすると、コンパイル後のグラフトポロジーが変わり、本番でJITストールやOOMを引き起こす。デバッグにはXLA HLOグラフのダンプ解析が必要で、AIエンジニアを「インフラ」障害に巻き込むことになる。
• 「BF16 → FP8でスループット2倍」（MLOps）と「FP8の静的スケールだと特定タスクでハルシネーションが起きる」（データサイエンティスト）が真っ向から衝突する。NVIDIAならランタイムが交渉してくれる。TPUでは人間2人が膝を突き合わせて交渉するしかない。

TPU/Trainiumで最も先行している組織——GoogleのGeminiチーム（チップまで自社で垂直統合）、AnthropicのClaudeチーム、そして2026年にGoogle TPUをレンタルしてLlamaの学習・推論を試し始めたMeta——は、横割りの「データサイエンス部 / インフラ部」という分業から離れる方向にある。代わりに、アテンションの数式とコンパイラの内部の両方に精通した、垂直統合型のワンチームを据えている。大半の会社はそんなチームを編成できない。古い分業を維持しようとしたプロジェクトから順に、コンパイルエラーとOOMの山に埋もれて死んでいく。

では誰が使うのか？ 入力がロックできる者だ

入力チャネルを自分で支配して形状が予測可能になった瞬間、すべての計算がひっくり返る。きれいな例が2つある。

• Google / YouTube要約。 内部の正確なパイプラインは公開されていないが、制約から導かれる形はこうだ。Googleは動画を見直さない。アップロード時に（余剰TPUリソースで）非同期バッチがASRを走らせ、タイムスタンプ付きテキストをBigtable等のストレージに格納する。要約を求められたとき、正確なテキスト長はトークン単位ですでに判明している——だからルーターはジャストサイズのバケットを選べ、パッキングの無駄はほぼゼロ、Gemini Flashのような軽量モデルが事前生成済みテキストをスキャンする。「2時間の動画を一瞬で要約」の魔法の正体は、「数ヶ月前にほぼ無料で作っておいた小さなテキストインデックスをスキャンしただけ」だ。
• Anthropic / Claude Code。 CLIのコーディングエージェントは入力がほぼ確定している。リポジトリ構造、ツール定義、Git差分、システムプロンプト。コンテキストの最初の約90%が不変——これはまさに静的コンパイルとPrompt Cachingが大好物とする形だ。実際Anthropicは、Trainium・TPU・NVIDIAを混在させ、ワークロードに応じて最適なチップに振り分けてClaudeをサービングしており、Trainiumを大規模に運用している（neuronx-distributed）。リアルタイムパッキングをやる高スループットなGo/C++プロキシは静的パスの自然なフロントエンドだが、製品ごとの正確な内訳は公表されていない。Claude Codeは——皮肉に読めば——Java型シリコンを苦労に見合うものにするための、ほぼ完璧な「入力ロックチャネル」だ。長文ワークロードも追い風になる。200Kトークンのprefillは多数のバケットを隙間なく連結して埋めるので、パディングの相対的な無駄がゼロに近づく——静的アレイの弱点が、Claudeの最も強い領域でちょうど薄れる。

逆もまた論理的で、チャットUIが動的SIMTハードに最も強く依存する理由を説明する。ChatGPTやClaude.aiのWebフロントは任意のテキストを受け付け、突然の画像アップロードや会話途中の話題転換が来る。送信ボタンが押されるまでシステムは形状を予測できない。そのカオスこそ、動的SIMT + PagedAttentionが作られた理由そのものだ。

まとめ

• TPUが流行らないのは遅いからでも高いからでもない——トークン単価はむしろ安い。 流行らないのは、その安さが「全テンソル形状をコンパイル時に固定する」という、大半のチームが守れない規律を条件にしているからだ。
• コストは消えたのではなく移動した。 静的シリコンは不確実性をすべてハードから追い出し、ソフトウェア（パッキング、マスキング、バケットルーティング）と人間（崩壊した開発/運用の境界）に押し付ける。CapEx（シリコン、電力）をOpEx（ハック層を保守するトップエンジニア）にトレードしている。
• 意思決定のルールはチップではなくチャネルにある。 入力を支配しているなら——CLI、固定の業務ワークフロー、自社のストレージパイプライン——TPU/Trainiumは武器だ。入力がフリーフォームのチャットボックスやサードパーティAPI連携なら、NVIDIA（かAMD）が唯一まともな選択肢で、EC2のカタログ価格の安さだけでTPUに突っ込むのはMFUが静かに一桁まで暴落する道だ。

スペックシートはトークン単価について嘘をついていなかった。ただ、その前に買わなければならないエンジニア、分岐したパイプライン、組織の再設計の値段が載っていなかっただけだ。