Helion on TPU: High-Performance ML Kernel Authoring
PyTorch-based Helion simplifies writing optimized kernels for TPUs.
1Introduction
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.
2Importance 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.
3Developer 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.
4Helion 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.
5TPU 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.
6Performance 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.
7Codegen 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.
8Example: 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.
9Example: 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.
10Performance 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.
11Conclusion
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.
Original article: PyTorch Blog — Helion ile TPU Üzerinde Yüksek Performanslı ML Çekirdekleri