GTD

Header Ad Placement

TorchTPU Architecture: PrivateUse1 Eager Dispatch, Deferred StableHLO Fusion, and a Three-Tier Compilation Cache Bringing Native PyTorch to Google TPUs

ML COMPILERS & RUNTIMES

Google and Meta's open-sourced TorchTPU makes the TPU a first-class PyTorch device. The real engineering story is how it reconciles PyTorch's mutable, eager worldview with XLA's immutable, compile-time one.

By Ossam Hashim

KEY FINDINGS

What the open-source release establishes

  • TorchTPU registers a C++ backend on the PyTorch dispatcher as PrivateUse1, with TPU tensor storage defined through the standard at::Tensor and c10::StorageImpl mechanisms, so stock PyTorch code targets a tpu device with a two-line change.
  • Eager execution has three modalities: Strict (asynchronous op-by-op dispatch), Debug (synchronous, with bounds checks), and DeferAndFuse, a hybrid that batches Aten ops into StableHLO chunks and is reported by Google to run 50% or more faster than Strict eager while reducing peak HBM usage.
  • A three-tier compilation cache (process memory, host shared memory, and a persistent filesystem tier) keys compiled executables on a 192-bit fingerprint: structure, dimensions, and compiler options plus component versions.
  • PyTorch's mutable view semantics are preserved over immutable StableHLO through view decomposition and inversion, with in-place writes lowered to dynamic_update_slice-style selective updates.
  • Collectives compile to StableHLO ops instead of calling a separate NCCL-style library; distribution runs one process per TPU device, with a topology_aware_mesh() API that aligns tensor, context, expert, and data parallelism to the physical ICI interconnect.
  • The stack already plugs into vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, and Hugging Face Transformers v5, where deploying to TPU is a device_map="tpu" argument.

A Third Native Backend for PyTorch

Announced at PyTorch Conference 2026 and open-sourced under the google-pytorch GitHub organization, TorchTPU is co-developed by Google and Meta's PyTorch team. Its stated ambition is "PyTorch citizenship": no JAX dependency in the runtime, a solid eager mode, a torch.compile path built on Dynamo rather than LazyTensor, and wide support for distributed idioms like DDP, FSDP, and DTensor. Meta's framing is hardware diversity - the ability to co-develop hardware with the core ML systems team instead of being constrained to a single architecture. The launch matters because it attacks the two historical frictions of TPU adoption from the PyTorch side. The first friction was API drift: earlier bridges asked users to learn TPU-specific idioms. TorchTPU's answer is an autoload mechanism - on a TPU host with the wheel installed, standard code acquires a tpu device and a tpu_dist process-group backend without structural rewrites. The second friction was serving: Google paired the release with native TorchTPU backends in vLLM (continuous batching, quantization, multi-chip parallelism, custom Pallas kernels), SGLang (whose RadixAttention prefix-cache reuse now runs on TPU), and llm-d for Kubernetes-native fleet routing across TPU and GPU. Hugging Face committed to nightly CI of the Transformers suite on TPU infrastructure.

Eager Mode, Three Ways

The dispatcher callback is where the design gets unusual. In Strict mode each Aten op executes on the TPU asynchronously with the CPU, the GPU-like default. Debug mode (TORCH_TPU_LAUNCH_BLOCKING=1) makes execution synchronous and adds bounds checks. The interesting one is DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1): each dispatched op, say aten::mul, creates a Deferred Op instance holding a StableHLO builder for that op plus optional static parameters for cache-key computation, and returns a tensor whose storage refers, via indirection, to that deferred computation instead of a materialized buffer. Materialization happens only when it must: a .to("cpu") call, an .item(), or a similar sink forces TorchTPU to analyze the region of deferred ops feeding the requested tensor, run liveness analysis (side-effecting ops such as collectives are always preserved in order), fingerprint the resulting graph, consult the compilation cache, compile on a thread pool if missed, dispatch to the TPU, and copy back the one requested value. A Repeated Ops Heuristic detects recurring op sequences - think a stack of transformer blocks - and triggers materialization before the host-to-device boundary, so steady-state CPU execution stays ahead of TPU execution. Op fusion in this mode also shrinks temporary HBM allocations, which is where the reported memory reduction comes from.

The Compilation Cache Is the Product

XLA compilation is expensive enough that cache design determines whether eager TPU execution is viable at all. TorchTPU's tier-1 cache is a process-local in-memory store of ready-to-schedule C++ executables. Tier-2 is machine-local shared memory for cross-process sharing on multi-device hosts, holding serialized executables. Tier-3 is a persistent cache on remote filesystems such as Google Cloud Storage, surviving job restarts and shared across hosts - aimed at Colab-style and datacenter environments. Cache keys are deliberately not the StableHLO graphs. TorchTPU fingerprints each computation into three 64-bit integers: a shapeless key capturing graph structure, a dimensions key capturing tensor sizes, and a compile-options key capturing XLA options plus the versions of TorchTPU, libtpu, and the PjRT runtime. A collision across the 192-bit space is treated as a non-event in practice. The worked example below models this scheme in miniature.

Bridging Mutable PyTorch and Immutable StableHLO

StableHLO is SSA and purely functional: values are immutable, layouts are compiler-chosen, and there is no stride metadata to poke. PyTorch assumes the opposite - tensors as pointers into shared byte arrays, cheap views, and in-place writes visible across aliases. TorchTPU bridges the gap with view decomposition and view inversion layered on deferred ops, lowering selective writes to dynamic_update_slice-style operations. The payoff is concrete: concatenated parameter arrays for distributed training, a pattern earlier bridges struggled with, work natively. Two more mismatches get explicit treatment. Shapes: XLA compiles statically, so torch.compile on TPU currently forces dynamic=False while bounded dynamism - allocate to an upper bound, track logical sizes, pad and slice automatically - is under active development. Precision: the MXU's native format is BF16 with FP32 accumulation, so TorchTPU's default "medium" mode maps float32 matmuls to single-pass BF16 multiplication with FP32 accumulation, while "high" and "highest" use three-pass and six-pass decompositions at roughly 3x and 6x MXU cost, with a TORCH_WARN_ONCE guard. Elementwise FP32 work runs on the vector unit with full IEEE-754 semantics.

Distribution Without NCCL

Collectives lower to StableHLO ops rather than an external library, and because a TPU device admits a single user-space process, TorchTPU adopts one process per device - a deliberate contrast with JAX's one-process-per-host model, and the reason RL colocation is explicitly called out as unsupported for now. The PyTorch Stream API exists, but with TPU semantics: a TensorCore is a single-threaded processor, so streams overlap compute with DMA-engine transfers and SparseCore offload (sparse math and collectives), and a stream can be pinned to a SparseCore explicitly. Topology is treated as a first-class scheduling input. ICI links deliver tens to hundreds of GB/s inside a slice; DCN Ethernet between slices is slower and higher-latency. The topology_aware_mesh() API computes rank layouts that place the most communication-heavy parallelism - tensor parallel all-reduces every layer, then context and expert parallel all-to-alls, then FSDP, then replicated data parallel - onto the fastest contiguous ICI rings. For custom communication, distributed Pallas kernels issue hardware RDMA over ICI directly, and torch.tpu.pallas.jax_op exposes them as ordinary PyTorch functions. For production inference planning on this kind of stack, our production inference benchmark guide for serving engines remains the companion read, as does the Humming quantized-kernel library analysis for what the kernel layer looks like when it is rebuilt rather than bridged. Readers who followed our earlier dissection of TorchDynamo and TorchInductor will recognize the frontend: TorchTPU registers a Dynamo backend, traces the FX graph with all ops deferred, and compiles one StableHLO program per graph.

Assumptions and limitations: the example is a didactic model of the fingerprinting scheme, not TorchTPU's internal key derivation; real keys also fold in compiler options and component versions. All performance figures cited here - including the 50% DeferAndFuse speedup and the 3x/6x precision-mode costs - are vendor-reported from the launch documentation and have not been independently benchmarked. Bounded dynamism and RL colocation remain open limitations per the same sources.

Worked Example: Modeling the Compilation Fingerprint

TorchTPU's cache correctness rests on the claim that three 64-bit fingerprints - structure, sizes, options - uniquely identify a compiled graph. The simulation below builds a tiny deferred-op graph, computes analogous keys, and shows the behavior the design predicts: a repeated computation hits the cache, while the same op sequence at a new shape misses because the dimensions key changes.

import hashlib

def key64(s):
    h = hashlib.sha256(s.encode())
    d = h.digest()[:8]
    return int.from_bytes(d, "big")

class Node:
    def __init__(s, name, ins, shape):
        s.name = name
        s.ins = ins
        s.shape = shape

def leaf(shape):
    return Node("leaf", [], shape)

def mul(a, b):
    return Node("mul", [a, b], a.shape)

def fingerprint(root):
    names, dims = [], []
    def walk(n):
        for i in n.ins:
            walk(i)
        names.append(n.name)
        dims.append(n.shape)
    walk(root)
    shapeless = key64("|".join(names))
    dim_key = key64(str(dims))
    opt_key = key64("ttpu|libtpu|pjrt")
    return shapeless, dim_key, opt_key

cache = {}

def materialize(root):
    k = fingerprint(root)
    hit = k in cache
    cache[k] = "exe"
    return hit

a = leaf((4, 16))
b = leaf((4, 16))
c = leaf((8, 16))
print("first:  ", materialize(mul(a, b)))
print("repeat: ", materialize(mul(a, b)))
print("reshape:", materialize(mul(c, b)))

Assumptions and limitations: the example is a didactic model of the fingerprinting scheme, not TorchTPU's internal key derivation; real keys also fold in compiler options and component versions. All performance figures cited here - including the 50% DeferAndFuse speedup and the 3x/6x precision-mode costs - are vendor-reported from the launch documentation and have not been independently benchmarked. Bounded dynamism and RL colocation remain open limitations per the same sources.

Sources & Provenance

Analysis is based on the launch-week primary documentation listed above; no hands-on TPU run was performed for this piece.

Comments

Post a Comment

Subscribe to: Post Comments (Atom)

Kernel Engineering Under Physical Non-Idealities: Floorsweeping, Wave-Equalized Launches, and the DVFS Confounder Behind H100 Clustering Gains

GPU SYSTEMS An NUS team shows the H100 you programmed is not the H100 you have: fabrication binning quietly removes whole SMs from cluster reach, and power limits rewrite the speedup math. Accounting for both recovers up to 5.30% on FlashAttention 4. By Ossam Hashim KEY FINDINGS Floorsweeping, the binning step that disables defective silicon to salvage a die, is invisible to the ISA but visible to cluster placement: on the H100 SXM5 it removes GPC-level clustered execution capacity, cutting SMs reachable by clusters from 132 to 124 (paper's measurement). Cluster placement also drifts at runtime: the paper's microbenchmark records only 120 unique SM IDs per kernel iteration but 124 across iterations, so the effective SM budget oscillates launch to launch. Fewer usable SMs can add a full execution wave: a launch that needs one wave at 132 SMs can need two at 120, erasing the data-movement savings DSMEM clustering was supposed to buy. The fix is arithmetic,...

Editorial Governance & Pages