Helion, a High-Performance TPU Kernel DSL for PyTorch, Released

Built in collaboration with Google, Helion is a high-level DSL for PyTorch for writing ML kernels with performance portability. It addresses the high entry barrier of Pallas, the existing low-level DSL, allowing users to write familiar PyTorch-style code that compiles into optimized TPU code.

Key features and achievements include the following:

• High Performance Achieved: On Flash Attention workloads, it achieved 838 TFLOPs on TPU v7 (approximately 79% MFU on a single tensor core).
• Autotuning: It explores various code generation strategies depending on input shape to select the optimal pipelining approach, maximizing TPU's VMEM and compute resources.
• Hardware Portability: It allows the same set of kernels to be maintained across TPU and GPU, simplifying workflows for cross-hardware users.
• Use Cases: It is suited for cases requiring performance optimization, non-experts in Pallas, and users who want to maintain identical kernels across multiple hardware platforms.