PyTorch’s Helion, a domain-specific language for performance-portable ML kernels, now compiles to optimized TPU code via Pallas using a two-level pipelining strategy. Outer-loop pipelining uses pallas_call to overlap HBM-to-VMEM transfers with compute, while inner-loop pipelining autotunes between device-side emit_pipeline (memory-efficient for variable sequence lengths) and unroll (prefetches all data into VMEM to eliminate compute bubbles). The autotuner picks the optimal strategy per input shape, letting a flash attention kernel hit 892 TFLOPs on shorter sequences and achieve a geometric speedup of 1.55x over eager execution and 1.12x over torch.compile across varied kernel workloads.
