Helion is PyTorch's high-level DSL for writing performance-portable ML kernels. A new TPU backend compiles Helion kernels to Pallas (Google's low-level TPU DSL), enabling PyTorch-style kernel authoring for TPUs without deep Pallas expertise. On a flash attention workload, the Helion-generated kernel achieves 838 TFLOPs (~79% MFU) on TPU v7. The key innovation is an autotuner that selects between two pipelining strategies — `emit_pipeline` (memory-efficient, scales to long sequences) and `unroll` (pre-fetches K/V into VMEM, eliminates compute bubbles but uses more memory) — based on input shape. Across a broader benchmark suite, Helion achieves a geometric mean 1.55x speedup over TorchTPU eager and 1.12x over torch.compile with XLA, with the largest gains on fusion-heavy kernels like attention that XLA struggles to optimize automatically.
