# Helion on TPU: High-Performance ML Kernel Authoring PyTorch-based Helion simplifies writing optimized kernels for TPUs. - **Source:** [PyTorch Blog](https://pytorch.org/blog) - **Published:** 2026-07-23 - **Topic:** [Data & AI](https://brifing.si/c/data-ai) - **Original article:** - **On Brifing:** https://brifing.si/s/1919edb6-9450-426b-95ce-85cf6c02ceac ## Introduction Helion is a high-level domain-specific language developed within the PyTorch ecosystem for writing machine learning kernels. Through a partnership with Google, a backend was developed that creates optimized kernels for TPUs, allowing developers to leverage TPU power through PyTorch's familiar interface. ## Importance of TPUs Google's TPU v4 (Ironwood) processors offer performance comparable to NVIDIA B200 while reducing total cost of ownership. TPUs, optimized specifically for machine learning workloads, have become an attractive alternative for large-scale training and inference tasks. ## Developer Challenges Traditionally, writing TPU kernels required learning Pallas, a low-level language with a steep learning curve and complex code structure. Helion removes this barrier by bringing PyTorch-style coding to TPUs. ## Helion Advantages Helion focuses on three main use cases: performance-critical scenarios requiring autotuning, quick onboarding for non-Pallas experts, and cross-hardware users wanting to maintain the same kernel set across TPUs and GPUs. ## TPU Architecture Unlike GPUs, TPUs are sequential machines with wide vector registers and compute units. Their memory hierarchy requires kernel authors to understand data loading timing and methods. A performant kernel must overlap HBM-to-VMEM data transfers with MXU computations. ## Performance Comparison Current TPU7x and NVIDIA B200 offer similar BF16 compute TFLOPS and HBM bandwidth. Helion's TPU code generation focuses on creating software pipelines to maximize performance from these hardware. ## Codegen Strategy Helion's Pallas codegen follows a three-layer strategy: pipeline creation via pallas_call in outer loops, autotunable pipelined loops or VMEM prefetching in inner loops, and automatically tuned pipeline buffer sizes. ## Example: Addition Operation A simple tensor addition Helion kernel transforms into a host-side tiling launcher that invokes the device function via pipeline, and a VMEM-resident addition operation. The autotuner selects buffer sizes that best overlap memory transfers with computation. ## Example: Flash Attention Attention, a key operation in modern language models, can be optimized for TPUs using Helion. The Flash Attention pattern computes in tiles to avoid materializing full S×S attention matrices. Helion automatically optimizes this complex operation for TPUs. ## Performance Results On Flash Attention workloads, Helion-generated kernels achieve 838 TFLOPS (79% tensor core MFU) on TPU v7. For different input shapes, Helion's autotuner selects pipeline schemes that best utilize available TPU VMEM and compute resources. ## Conclusion Helion makes TPU kernel authoring accessible to PyTorch developers while unlocking the hardware's full potential. With autotuning capabilities and cross-hardware portability, it fills an important gap in the machine learning ecosystem. This technology has potential to reduce costs in large-scale ML deployment. ## Sources covering this story - https://pytorch.org/blog/helion-on-tpu-towards-hardware-heterogeneous-kernel-authoring/ — --- _Summary written by Brifing from the original reporting. Copyright in the original article belongs to its publisher; read the full piece at the link above._