Helion on TPU: Towards Hardware Heterogeneous Kernel Authoring

Helion, PyTorch's high-level DSL for writing performance-portable ML kernels, partnered with Google to build a TPU backend that compiles Helion kernels to Pallas, enabling PyTorch-friendly TPU kernel authoring. On a flash attention workload, the Helion-generated kernel achieves 838 TFLOPs (~79% MFU of one tensor core) on TPU v7, with autotuning over different code-generation strategies to select the optimal pipelining schema.

Featured projects TL;DR Helion https://helionlang.com/ is PyTorch’s high-level DSL for writing performance-portable ML kernels. Partnering with Google, we have built a TPU backend that compiles Helion kernels to Pallas https://docs.jax.dev/en/latest/pallas/index.html , providing a PyTorch-friendly way to author performant TPU kernels. On a flash attention workload, the Helion-generated kernel achieves 838 TFLOPs ~79% MFU of one tensor core on TPU v7. On different input shapes, Helion autotunes over different code-generation strategies to select the optimal pipelining schema, making the most use of TPU’s available VMEM and compute. Introduction TPUs are increasingly important as an ML compute platform to complement GPUs. Google’s latest TPU v7 Ironwood delivers comparable performance to NVIDIA B200 with a potentially lower total cost of ownership TCO , making TPUs an appealing option for large-scale training and inference workloads. However, authoring TPU kernels traditionally requires expertise in Pallas, a low-level DSL that comes with a steep learning curve and code complexity. Helion https://helionlang.com/ bridges this gap. As PyTorch’s portable DSL for ML kernels, Helion lets users write familiar PyTorch-style code and compiles it to optimized TPU code. Paired with performance wins brought by its autotuner, Helion is evolving towards an attractive option for authoring TPU kernels. Specifically, Helion TPU targets three main use cases: Performance-critical use cases where autotuning is required to explore the configuration space Non-Pallas experts hoping to onboard TPU kernel authoring quickly Cross-hardware users who prefer to maintain the same set of kernels across TPU and GPU This article starts with a brief overview of TPU’s hardware features and programming models as compared to GPUs, and then demonstrates how Helion generates performant Pallas code with ideal pipelining characteristics for different input shapes. TPU Primer TPUs are highly specialized accelerators designed and optimized specifically for machine learning workloads. The architecture https://jax-ml.github.io/scaling-book/tpus/ and programming model of TPUs differ significantly from GPUs. The most prominent difference is that a TPU is a sequential machine featuring wide vector registers and compute units. This contrasts with GPUs, which achieve performance via both massively parallel execution CUDA cores and specialized tensor units tensor cores . | TPU Pallas | GPU CUDA | | |---|---|---| Threading | Sequential Few large workers | Parallel SIMT + tensor core Many small workers | Memory Hierarchy | Explicit memory spaces persistent vs scratchpad memory , Async mem copies required for pipelining. | Implicit caches, HW-managed | As a result, TPUs feature a memory hierarchy that kernel authors must deeply understand, so that the kernels they write can orchestrate when and how data is loaded from the off-chip HBM to the fast on-chip VMEM. A performant Pallas kernel would overlap these HBM< VMEM memory transfers with floating point computation happening in the matrix MXU and vector compute units. Despite the architectural differences, current-generation TPUs and GPUs are highly comparable in raw performance. TPU7x and NVIDIA B200 have very similar BF16 compute TFLOPS and HBM bandwidth — the two most important hardware metrics for modern ML workloads. Helion’s Pallas Codegen To extract maximum performance out of a TPU, Helion’s Pallas codegen aims to maximize software pipelining , ensuring that memory transfers and computation overlap as much as possible. This section illustrates Helion’s three-fold strategy for generating pipelined kernels: - Outer loop: pallas-provided pipelined device invocation pallas call / emit pipeline - Inner loop: autotuned between: - pallas-provided pipelined device-side loop emit pipeline - Pre-fetching all values into VMEM, if possible unroll - pallas-provided pipelined device-side loop - Auto-tuned pipeline buffer sizes Example: add As a simple example, consider the following helion kernel for adding two tensors. php @helion.kernel def add x: torch.Tensor, y: torch.Tensor - torch.Tensor: out = torch.empty like x for tile in hl.tile out.size : out tile = x tile + y tile return out The Helion compiler translates this into two functions: a host-side launcher that tiles the input and invokes the device function in a pipelined fashion, and a device function that operates on VMEM-resident tiles: python def helion add x, y, out : out : = x : + y : def add x: torch.Tensor, y: torch.Tensor : BLOCK SIZE 0 = <autotuner-selected value out = torch.empty ... out = launcher wraps around pallas call helion add, x.shape 0 + BLOCK SIZE 0 - 1 // BLOCK SIZE 0, , grid size x, y, out, block spec info= BLOCK SIZE 0, ... , ... return out Within the generated code: - The hl.tile loop in the Helion source becomes a grid on the host side. The launcher wrapping pallas call invokes helion add once per tile, with each invocation automatically pipelined — while one tile is being computed, the next tile’s data is being loaded from HBM into VMEM. - The device function helion add is simple: it receives VMEM references not HBM pointers , so the kernel body is a simple addition. BLOCK SIZE 0 the tile/buffer size is selected by the autotuner, which explores different sizes to find the best overlap between memory transfers and compute for the target hardware. This results in a pipelined execution as illustrated below. Example: Flash Attention Attention is one of the key operations in modern language models. Production implementations follow the “Flash Attention” pattern – a memory-efficient technique that computes attention in tiles to avoid materializing the full S×S attention matrix. The structure of a flash attention kernel in Helion is illustrated below: B, H, S, D = 8, 32, 8192, 256 batch, head, sequence length, head dimension @helion.kernel def attention q: torch.Tensor, k: torch.Tensor, v: torch.Tensor - torch.Tensor: out = torch.empty ... for tile b, tile q in hl.tile B H, S : this q = q tile b, tile q, : acc = ... for tile kv in hl.tile S : this k = k tile b, tile kv, : this v = v tile b, tile kv, : <qk matmul, online softmax, v matmul, update acc out tile b, tile q = acc return out Compared to the “add” example discussed previously, the flash attention kernel https://github.com/pytorch/helion/blob/main/examples/attention.py contains an additional inner loop which performs tiled accesses across the entire K and V sequences. How we pipeline the memory and compute within the inner loop is key to the performance of this kernel. In Helion, the compiler autotunes over two different strategies for translating this kernel to Pallas. This is keyed on the pallas loop type autotuner config. With the default pallas loop type == emit pipeline option, Helion relies on Pallas’ device-side emit pipeline API to pipeline an inner loop body function, similarly to how the host-side logic uses pallas call to pipeline the device function invocation: python def helion attention q VMEM, k HBM, v HBM, out VMEM : acc = ... this q = q VMEM :, :, : def inner pipeline body k VMEM, v VMEM : this k = k VMEM :, :, : this v = v VMEM :, :, : <matmul, online softmax, matmul, update acc pallas.tpu.emit pipeline inner pipeline body, k HBM, v HBM, block spec info= BLOCK SIZE KV ,... , ... out VMEM = acc def attention q: torch.Tensor, k: torch.Tensor, v:torch.Tensor : out = torch.empty ... out = launcher wraps around pallas call helion attention, q, k, v, out, block spec info= BLOCK SIZE Q ,... , ... return out This generated kernel follows a nested pipeline structure: - Outer pipeline The host uses pallas call to invoke helion attention . The HBM reference of q is tiled, and each invocation of helion attention receives a VMEM tile of q. For k and v, helion attention receives HBM references directly. - Inner Pipeline Within helion attention , the device uses emit pipeline to invoke inner pipeline body , which receives VMEM tiles of k and v. This results in a pipelined execution as illustrated in this image: One obvious point of inefficiency in this pipeline is that there are bubbles in the compute units – for every new Q tile, while we fetch the 0th KV tile, there is no work available for the compute units. This comes down to the fact that we are re-loading the KV tiles from HBM to VMEM for every new Q tile. Helion offers an alternative pallas loop type == unroll config which avoids this bubbling. With unroll , we translate the inner for loop into a simple Python for loop: python def helion attention q VMEM, k VMEM FULL, v VMEM FULL, out VMEM : acc = ... this q = q VMEM :, :, : for offset in range 0, k VMEM FULL.size 1 , BLOCK SIZE KV : this k = k VMEM FULL :, pallas.dslice offset, BLOCK SIZE KV , : this v = v VMEM FULL :, pallas.dslice offset, BLOCK SIZE KV , : <matmul, online softmax, matmul, update acc out VMEM = acc def attention q: torch.Tensor, k: torch.Tensor, v:torch.Tensor : out = torch.empty ... out = launcher wraps around pallas call helion attention, q, k, v, out, block spec info= BLOCK SIZE Q, None, None , ... return out The name “unroll” reflects the fact that Pallas device functions are traced by JAX’s JIT — the Python for loop is effectively unrolled at trace time into a flat sequence of operations. In this version of the generated kernel: K and V are pre-fetched in full: The host passes None as the block spec for K and V, instructing pallas call to load them entirely into VMEM. The full VMEM references persist across all device function invocations. The inner loop slices locally: Each iteration uses pallas.dslice to select the relevant KV tile from the already-resident VMEM buffer. No HBM traffic occurs during the inner loop. This results in a different pipelining scheme as illustrated below: In this workflow, there are no longer bubbles in the compute pipeline. The trade-off is that this requires more VMEM usage, as the entire K and V sequences need to be present. This means that although more performant, this translation isn’t always possible. The VMEM usage is linear with respect to the input sequence lengths as opposed to tile size , which is prohibitive with longer sequences. The performance difference between these strategies is significant — the table below shows results on workloads with B=8, H=32, D=256: | S = 8k | S = 32k | | |---|---|---| emit pipeline TFLOPs | 653 | 695 | unroll TFLOPs | 892 | OOM | The benefit of Helion lies in its ability to autotune and select the best autotuner config. So that with smaller sequences, it makes use of the VMEM available and generates pipelined code with no compute bubbles. For longer sequences, it falls back to emit pipeline which scales to arbitrary context lengths. The following graph plots the performance of this attention kernel compared to various other Pallas attention implementations, on varying sequence lengths: The autotuner’s ability codegen different loop and pipelining strategies depending on the input length is what gives Helion its edge even when compared to highly optimized implementations such as Tokamax. Broader Kernel Benchmarks We benchmark Helion across a variety of kernels, tracked on our dashboard https://helionlang.com/dashboard/ . The table below compares Helion against TorchTPU eager and torch.compile using XLA across a range of different kernels. Helion shows a geometric average speed-up of 1.55x compared to eager, and 1.12x compared to compiled. | kernel | shape | torch tpu eager ms | torch.compile tpu ms | Helion ms | Helion vs torch tpu eager | Helion vs torch.compile | |---|---|---|---|---|---|---| | attention | 8,32,8192,256 | 87.77 | 88.28 | 19.72 | 4.45× | 4.48× | | softmax | 65536,2560 | 0.712 | 0.743 | 0.477 | 1.49× | 1.56× | | batch softmax | 64,2048,4096 | 1.888 | 1.373 | 0.982 | 1.92× | 1.40× | | softmax two pass | 8192,8192 | 0.386 | 0.417 | 0.334 | 1.16× | 1.25× | | bmm | 64,2048,2048,2048 | 3.211 | 1.860 | 1.527 | 2.10× | 1.22× | | rms norm-bwd | 8192,8192 | 1.792 | 0.789 | 0.661 | 2.71x | 1.19x | | epilogue subtiling | 4096,4096,4096 | 0.850 | 0.462 | 0.417 | 2.04x | 1.11x | | matmul layernorm | 4096,4096,4096 | 0.535 | 0.523 | 0.489 | 1.10× | 1.07× | | welford | 524288,512 | 1.330 | 1.357 | 1.316 | 1.01x | 1.03x | | swiglu | 16,16384,4096 | 3.510 | 2.244 | 2.295 | 1.53× | 0.98× | | matmul | 8192,8192,8192 | 1.552 | 1.527 | 1.597 | 0.97× | 0.96× | | geglu | 16,8192,8192 | 3.779 | 2.240 | 2.424 | 1.56× | 0.92× | | cross entropy | 128,2048 | 0.363 | 0.264 | 0.320 | 1.13× | 0.82× | | broadcast matmul | 64,2048,2048,2048 | 1.817 | 1.440 | 1.806 | 1.01× | 0.80× | | layer norm | 16384,16384 | 1.253 | 0.779 | 1.126 | 1.11× | 0.69× | | rms norm | 8192,8192 | 1.194 | 0.419 | 0.617 | 1.94× | 0.68× | Helion shows the largest gains on kernels that employ fusion or optimization patterns that are difficult for XLA to discover automatically — flash attention is a prominent example. For the more standard operations like matmul and layer norm , XLA’s compiler already produces high-quality code, and Helion performs comparably. What’s Next Helion on TPU is under active development. Here’s a non-exhaustive list of things we are working on: - Expand kernel coverage: Get more Helion examples working on TPU - Further performance improvements - Better support for jagged and sparse operations - Support for distributed TPU computing Getting Started Helion is open source and available on GitHub. Its TPU backend has a dependency on TorchTPU https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/ , which is expected to be released publicly later this year. When it does, we encourage you to try-out Helion on TPU and share your feedback. Resources: Acknowledgements This project was made possible through the invaluable collaboration and technical insights of our peers. A special thank you to Joe Pamer, Robert Hundt, Claudio Basile and Adam Paszke at Google, as well as Jana van Greunen, Gregory Chanan, Peng Wu, and Zongwei Zhou at Meta, for their feedback and support in bringing this to fruition.