PyTorch Helion Adds a TPU Backend for Heterogeneous Kernel Authoring

Lead

PyTorch is pushing custom ML kernel development toward a higher-level, more portable model. In a new PyTorch Blog post, the team describes a TPU backend for Helion, a high-level DSL designed to let developers write kernels in a PyTorch-friendly style while the compiler generates optimized backend code.

The new path compiles Helion kernels to Pallas, Google’s lower-level DSL for TPU programming. That matters because TPU kernels can be difficult to write well: performance often depends on explicitly managing data movement between off-chip HBM and on-chip VMEM, and on overlapping those transfers with matrix and vector computation.

Key points

• Helion targets performance portability. Developers write kernels once in a higher-level form, and Helion maps them to hardware-specific implementations. The TPU backend extends that idea to Pallas-generated kernels.
• TPUs require a different mental model from GPUs. GPUs rely on massive SIMT parallelism and hardware-managed caches, while TPUs expose more explicit memory spaces and favor carefully scheduled sequential execution with wide compute units.
• Pipelining is the central optimization. Helion’s generated Pallas code aims to overlap HBM-to-VMEM transfers with computation. In a simple tensor add example, the tile loop becomes a host-side grid, and each device invocation works on VMEM-resident tiles.
• Flash Attention shows the harder case. The kernel has an outer loop over query tiles and an inner loop over key/value tiles. Helion autotunes between different inner-loop code generation strategies, including Pallas emit_pipeline and prefetching/unrolling when VMEM capacity allows.
• The reported result is notable. On the Flash Attention workload described in the source, the Helion-generated kernel reaches 838 TFLOPs on TPU v7, around 79% MFU of one tensor core.

Why it matters

The broader implication is that PyTorch is treating TPUs as a first-class target for performance-critical workloads, not merely as an alternative runtime. Current-generation TPUs and high-end GPUs are comparable on key hardware metrics such as BF16 compute and HBM bandwidth, but the developer experience is still very different. A DSL such as Helion can help bridge that gap by moving more of the scheduling and tuning burden into the compiler.

For engineering teams, this could be useful in three scenarios: optimizing critical kernels without becoming Pallas experts, maintaining a shared kernel codebase across TPU and GPU backends, and relying on autotuning to handle shape-dependent choices that are tedious to optimize manually.

The caveat is that the blog presents results for specific workloads and configurations, so it should not be read as a universal performance guarantee. Still, the direction is clear: the next layer of AI infrastructure competition is not only about faster chips, but also about tools that make those chips easier to use efficiently.

Source: PyTorch Blog