The news broke a few hours ago: Helion, PyTorch’s dedicated language for writing performant and portable ML kernels, now has a mature TPU backend, the result of joint work by Meta and Google. On a flash attention workload—a central operation for modern transformers—the kernel automatically generated by Helion hit 838 TFLOPs on TPU v7, roughly 79% of a tensor core’s maximum theoretical utilization. The number grabs headlines, but it is not the most important part.

How the autotuner chooses between memory and computation

The real challenge on TPUs is orchestrating data movement from HBM to fast on-chip VMEM and overlapping it with matrix and vector unit computations. Thanks to its autotuner, Helion automatically decides between two radically different pipelining strategies.

The first, emit_pipeline, loads K and V tiles from HBM anew for each Q tile. It is flexible and scales to long sequences, but it introduces compute bubbles because every time you switch Q you stall waiting for the first KV tile. The second, unroll, pre-loads the entire K and V into VMEM when the shapes allow it, eliminating dead time entirely: computations proceed without interruption. The downside is VMEM consumption, which grows linearly with sequence length and makes this approach infeasible beyond certain thresholds.

The numbers are telling: on a flash attention head with batch 8, 32 heads, dimension 256, at S=8k the emit_pipeline approach yields 653 TFLOPs, while unroll jumps to 892. At S=32k, unroll runs out of memory and the system falls back to emit_pipeline, which still reaches 695 TFLOPs. The autotuner picks according to the input shape, no developer thought required.

Hardware portability: the end of lock-in?

The real stakes go beyond teraflops. Helion was born as a portable DSL: the same kernel written for GPU is compiled for TPU as well. In a market where Google is pushing Ironwood (TPU v7) with performance comparable to NVIDIA B200 but potentially lower TCO, the ability to write a kernel once and use it on both silicon trees becomes a weapon for ML teams. Those with on-premise GPU clusters can experiment with cloud TPUs without rewriting core compute logic, and anyone designing hybrid architectures can decide where to run inference or fine-tuning based on cost and availability, not on code compatibility.

This is no small detail. The abstraction introduced by Helion (and by parallel efforts like Triton) shifts bargaining power from chip vendors to users, because it drastically cuts switching costs. If a compiler can automatically extract peak performance from every accelerator, hardware choice becomes a matter of TCO, latency, and geographic location—not of deep expertise in Pallas or CUDA. The winners are organizations with on-premise stacks that can now open a channel to the cloud without re-engineering kernels from scratch.

Tests on common operations show Helion beating TorchTPU eager (1.55x geometric mean speedup) and slightly edging out torch.compile via XLA (1.12x). The biggest gains appear on kernels with fusion patterns that are hard for compilers to discover automatically, like flash attention. For matmul and layer norm, where XLA already does well, performance is comparable. The difference comes from the autotuning logic, which on memory-intensive operations chooses the right pipeline without manual tuning.

For those evaluating hybrid deployments, AI-RADAR offers at /llm-onpremise an analytical framework to weigh these trade-offs. Meanwhile, Helion’s TPU backend still depends on the TorchTPU library, which is not yet public (expected later this year). But the direction is clear: hardware heterogeneity managed at the framework level, with a compiler smart enough to select the best pipelining strategy. This is not just a technical note; it is a signal that the software ecosystem is becoming the true fulcrum.