AI 推論與編譯器
PyTorch Helion Adds TPU Backend, Bridging Triton and Pallas with the Same High-Level Kernels
PyTorch and Google have built a TPU backend for Helion, allowing PyTorch-style kernels to compile to Pallas and reducing the need to write separate low-level TPU kernels. The team measured up to 838 TFLOPs for Flash Attention on TPU v7, though the feature remains under active development.
The PyTorch team announced on July 23 that the machine learning kernel DSL Helion can now generate TPU programs using Google Pallas as its backend. Helion originally focused on compiling Python-embedded tensor and blocked programs with PyTorch-like semantics to Triton. With the new path, engineers can potentially maintain a single high-level kernel and select a GPU or TPU backend based on the hardware, without having to directly rewrite large amounts of Pallas code for TPUs.
The work involves more than simply replacing the code generator. TPUs use a small number of large work units and explicit HBM and VMEM memory hierarchies, with performance relying on software-scheduled asynchronous data movement—a model distinct from GPU SIMT. Helion therefore converts outer tiling into a `pallas_call` pipeline and automatically evaluates `emit_pipeline` versus `unroll` for the inner loop. The former moves K and V block by block from HBM into VMEM and can handle longer sequences. The latter keeps data in VMEM ahead of time, eliminating compute bubbles but increasing memory requirements as sequence length grows.
In an official Flash Attention benchmark with B=8, H=32, and D=256, an 8K sequence achieved 892 TFLOPs using unroll, while a 32K sequence could not use that strategy because of insufficient VMEM. The autotuner can select a different pipeline based on the tensor shape. A separate measurement on TPU v7 reached 838 TFLOPs, equivalent to about 79% MFU on a single tensor core. Across a broader kernel suite, the team reported that Helion was 1.55× faster than TorchTPU eager and 1.12× faster than `torch.compile` by geometric mean.
This brings cross-accelerator kernel portability closer to practical requirements, but the figures come from kernels, shapes, and hardware selected by the project team and cannot be directly extrapolated to end-to-end model training or inference. Key areas to watch next include the TPU backend’s coverage, compilation and autotuning costs, and whether the same Helion kernels can consistently approach the performance of hand-written kernels across different generations of TPUs and GPUs.