The wall between CUDA-flavored kernel code and Google’s TPU stack just got a door. PyTorch announced on July 23 that Helion, its high-level DSL for ML kernels, now has a TPU backend built in partnership with Google. The backend compiles Helion kernels down to Pallas — the low-level TPU kernel language — and the headline number is competitive: a Helion-generated flash attention kernel hits 838 TFLOPs on TPU v7, roughly 79% MFU of a single tensor core.
The pitch is that you write PyTorch-style Python once and let the compiler fight the hardware. Helion kernels look like tiled PyTorch loops — for tile in hl.tile(out.size()) — and until now they targeted GPUs by lowering to Triton. The new backend gives that same source a second target, which matters because the alternative on TPU is writing Pallas by hand, a low-level DSL the PyTorch team politely describes as having “a steep learning curve and code complexity.”
Why TPUs are worth the compiler engineering
The timing tracks the hardware. Google’s TPU v7, Ironwood, delivers performance PyTorch describes as comparable to NVIDIA’s B200 — similar BF16 TFLOPS and HBM bandwidth, the two metrics that dominate modern ML workloads — with a potentially lower total cost of ownership. That combination has made TPUs a serious second platform for large-scale training and inference, and it’s exactly the audience Helion names: teams that need autotuned performance, engineers who aren’t Pallas experts, and cross-hardware shops that want to maintain one set of kernels across TPU and GPU rather than two divergent codebases.
The architectural gap Helion has to paper over is real. A TPU is essentially a sequential machine with wide vector registers and an explicit memory hierarchy — the kernel author is responsible for orchestrating when data moves from off-chip HBM into the fast on-chip VMEM scratchpad, and a performant kernel must overlap those transfers with compute in the matrix (MXU) and vector units. GPUs hide most of that behind hardware-managed caches and massively parallel SIMT execution. Same math, very different programming model.
Pipelining is the whole game
Helion’s Pallas codegen strategy is built around maximizing software pipelining. The outer loop of a Helion kernel becomes a host-side grid driven through pallas_call, which pipelines device invocations automatically — while one tile computes, the next tile’s data streams from HBM into VMEM. For kernels with inner loops, like flash attention’s tiled sweep across the K and V sequences, the compiler autotunes between two strategies: Pallas’ device-side emit_pipeline API, or prefetching everything into VMEM and unrolling when it fits. Buffer sizes are autotuned too, keyed on a pallas_loop_type config.
That autotuning is where the 838 TFLOPs comes from. On different input shapes, Helion explores different code-generation strategies and pipeline schedules to pick whichever makes best use of available VMEM and compute — the kind of shape-by-shape tuning that hand-written kernels rarely get because nobody has time to re-derive the optimal schedule for every sequence length.
The bigger bet: hardware-heterogeneous kernels
The blog title’s phrase — “hardware heterogeneous kernel authoring” — is the strategic tell. PyTorch is positioning Helion as the write-once layer above vendor kernel languages: Triton on NVIDIA and AMD, Pallas on TPU, whatever comes next underneath. For teams currently maintaining parallel CUDA and TPU kernel implementations, a single Helion source that autotunes per-target is a genuine maintenance win — if the generated code keeps landing within striking distance of hand-tuned performance.
One flash attention benchmark at 79% MFU doesn’t prove that across the kernel zoo, and Helion’s TPU story is explicitly “towards” heterogeneous authoring, not “arrived.” But a first-party PyTorch path onto Ironwood-class hardware, co-built with Google, is the clearest signal yet that the kernel layer of the stack is consolidating above the vendor line — and open source, at that.
