PyTorch Helion Gains TPU Backend, Autotunes Pallas Kernels for Speed

PyTorch and Google built a TPU backend for Helion, PyTorch's high-level DSL for writing performance-portable ML kernels. The backend compiles Helion kernels to Pallas, Google's low-level TPU DSL, giving PyTorch users a way to author TPU kernels without learning Pallas directly. On a flash attention workload, the Helion-generated kernel reaches 838 TFLOPs, about 79% of one tensor core's MFU, on TPU v7. The compiler autotunes over code-generation strategies to pick the best pipelining scheme for different input shapes, using TPU's VMEM and compute efficiently.

TPU Architecture and Helion's Codegen

TPUs differ from GPUs in being sequential machines with wide vector registers and compute units, unlike GPUs which rely on massively parallel execution and tensor cores. TPU memory is explicit, with persistent and scratchpad spaces, requiring async copies for pipelining, whereas GPUs use hardware-managed caches. Kernel authors must orchestrate data movement from HBM to on-chip VMEM, overlapping transfers with computation on the MXU and vector units. Current TPU7x and NVIDIA B200 have similar BF16 compute TFLOPS and HBM bandwidth.

Helion's Pallas codegen maximizes software pipelining through a three-fold strategy: an outer loop using Pallas' pipelined device invocation, an inner loop autotuned between a pipelined device-side loop and pre-fetching all values into VMEM, and autotuned pipeline buffer sizes. For a simple add kernel, the compiler generates a host-side launcher that tiles the input and invokes the device function in a pipelined fashion, while the device function operates on VMEM-resident tiles. The tile size is selected by the autotuner to find the best overlap between memory transfers and compute.

Flash Attention and Autotuning Strategies

Flash attention kernels contain an inner loop over K and V sequences, and how that loop is pipelined determines performance. Helion autotunes over two strategies keyed on the pallas_loop_type config. The default emit_pipeline option uses Pallas' device-side emit_pipeline API to pipeline the inner loop, creating a nested pipeline: the host invokes the device function with a VMEM tile of Q, while K and V are fetched from HBM in tiles. This causes bubbles in compute when a new Q tile starts, as KV tiles are reloaded from HBM.

The alternative unroll config translates the inner loop into a Python for loop, pre-fetching K and V entirely into VMEM. The host passes None as the block spec for K and V, loading them fully, and the inner loop slices locally with no HBM traffic. This eliminates compute bubbles but requires more VMEM, scaling linearly with sequence length. On workloads with batch 8, heads 32, and head dimension 256, emit_pipeline achieves 653 TFLOPs at sequence length 8k and 695 at 32k, while unroll reaches 892 TFLOPs at 8k but runs out of memory at 32k. The autotuner selects unroll for smaller sequences and emit_pipeline for longer ones, giving Helion an edge over highly optimized implementations like Tokamax.

Benchmark Results Across Kernels

Helion benchmarks against TorchTPU eager and torch.compile with XLA across several kernels. The geometric average speed-up is 1.55x over eager and 1.12x over compiled. For attention with shape [8,32,8192,256], Helion runs in 19.72 ms versus 87.77 ms for eager and 88.28 ms for compiled, a 4.45x and 4.48x improvement. Other kernels show gains: softmax at [65536,2560] is 1.49x over eager and 1.56x over compiled, batch_softmax at [64,2048,4096] is 1.92x and 1.40x, bmm at [64,2048,2048,2048] is 2.10x and 1.22x, and rms_norm backward at [8192,8192] is 2.71x and 1.19x. These results come from the autotuner's ability to generate different loop and pipelining strategies based on input length, making Helion a viable option for TPU kernel authoring without Pallas expertise.