TPU Inference Externalization Full Steam Ahead - InferenceX

What it claims

SemiAnalysis benchmarked TPUv7 Ironwood against B200 and B300 on Qwen3.5-397B in FP8, using 8k1k aggregated serving with single-token prediction (no MTP or speculative decoding) on Google's pre-release TorchTPU vLLM stack. At 100 tok/s/user, Ironwood costs about $0.181 per million total tokens. B200 costs $0.222 and B300 costs $0.276, so Ironwood is 19% cheaper than B200 and 34% cheaper than B300. That comparison uses an external TPU TCO, meaning the price a hyperscaler lab pays to buy TPUs. On raw throughput, TPU trails Nvidia across most of the curve, and its lead comes from a lower hourly cost. On FP4, Nvidia leads because Ironwood has no native FP4 compute.

How it works

TorchTPU replaces the TorchAX path. In TorchAX, PyTorch ATen ops were intercepted via __torch_dispatch__ and translated to JAX, and vLLM's usual torch.compile route was bypassed. TorchTPU registers TPU through PyTorch's PrivateUse1 extension point, so tensors are ordinary torch.Tensor on device="tpu". XLA remains the compiler (not Inductor or Triton), and existing Pallas kernels can still be called. Google documents support for DDP, FSDP2, DTensor, and both MPMD and SPMD execution.

flowchart LR
  M[PyTorch model in vLLM/SGLang] --> A{Backend}
  A -->|old: TorchAX| J[__torch_dispatch__ to JAX, jax.jit]
  A -->|new: TorchTPU| T[PrivateUse1 device, eager or torch.compile]
  T --> F[Dynamo/AOTAutograd FX graph]
  J --> S[StableHLO]
  F --> S
  S --> X[XLA to TPU executable]
  P[Pallas kernels] --> X

The kernel work most relevant to anyone porting models:

• Combined all-gather of expert IDs and routing weights saves about 80 µs per layer, or about 4.64 ms per forward pass on DeepSeek-V3.
• ReduceScatter moved to SparseCore, with die-to-die reduction first, raises 8k1k throughput by 4.1 to 14.2% (concurrency 64 to 512) and 1k8k throughput by 26.1% at concurrency 512. Small collectives stay on TensorCore below a threshold derived from VMEM.
• A packed sort key for MoE routing cuts sort time from 106.6 µs to 21.7 µs.
• The GDN output-projection rewrite computes Sq while the VPU builds the next state. GDN v3 fuses Conv1D and GDN, with kernel-only speedups of 1.41× decode, 1.60× prefill and 2.14× mixed.
• Compact recurrent-state slots reclaim about 76 GiB of HBM and add 18% to 1k8k throughput at concurrency 64. BF16 state storage with FP32 math adds 15% at concurrency 512.
• A sequence-on-lane KV layout doubles usable KV pages (5,141 to 10,283) and lowers the head-dim requirement to a multiple of 32. It lifts 8k1k throughput by 16.5% at concurrency 128 and cuts median TTFT by 95%, at about 3% extra per-token latency at low concurrency.
• Splitting the RPA v3 fetch block (16k) from the compute block (4k) raises decode throughput by 49%, from 64.9k to 96.3k tok/s, on Qwen3-0.6B.

Evidence

At 20 tok/s/user (concurrency 256), Ironwood reaches 9,364 tok/s/chip, against 8,903 for B200 and 8,925 for B300. That gives 50.4% more tokens per dollar than B200 and 96.0% more than B300. At Google's internal TCO of $1.03/chip-hour, the advantage rises to 76.7% and 130.2%. The cost is latency: mean TTFT at that point is 5.41 s on TPU, 3.75 s on B200 and 2.40 s on B300. At a 20 s median response time, costs are $0.098, $0.106 and $0.132 per million tokens. B200 wins around a 30 s median response time. Against GB300 NVL72 disaggregated serving, TPU aggregated serving is about 30% behind on performance per dollar in the middle of the latency range.

How to use it

As of early September, TorchTPU was in private beta. It was slated for open-sourcing around mid-October at the PyTorch Conference, with SGLang support (SGL-torchtpu) coming via RadixArk. The InferenceX TPU runs live in a SemiAnalysis fork until then. The public path today is vLLM tpu-inference, which uses TorchAX and is to be deprecated. Disaggregated KV transfer and DRAM offload use the open-sourced TPU-Sync library (formerly TPU-raiden), with llm-d and Mooncake Store support in progress.

Limits

Only one model was measured (Qwen3.5-397B), on one workload (8k1k random input), with no MTP, no disaggregation and no prefix-cache reuse. Agentic, Kimi K3, GLM5.3 and Gemma4 results are future work. Some sub-results come from DeepSeek-V3, Qwen3-32B or Qwen3-0.6B, and some are kernel-only. One low-concurrency change lowered 8k1k throughput by 5.3% at concurrency 8. TPU's 256×256 MXU penalizes small shapes: head dim 64 (as in gpt-oss) caps attention matmuls at 25% utilization, and MLA's 192 also tiles poorly. Such models need new kernels. Performance-per-dollar figures depend on SemiAnalysis's own TCO model. Its forecasts that the disaggregation gap will close within months and that TPUv8i will compete with Rubin NVL72 are SemiAnalysis's opinions.