Abstract:

This video introduces Helium, a new Python-embedded Domain-Specific Language (DSL) for authoring high-performance GPU kernels. Presented by its creators from the PyTorch team, Helium is positioned at a higher level of abstraction than Triton but lower than PyTorch, giving developers more explicit control over kernel fusion than torch.compile while maintaining ease of use.

Helium's core innovation is its sophisticated, integrated autotuner, which explores a vast search space of thousands of Triton kernel variants from a single high-level Helium kernel. This process automates complex, hardware-specific optimization choices like indexing modes, block sizes, loop structures, and reduction strategies, significantly reducing developer effort and improving performance portability across different hardware, including NVIDIA B200 and AMD GPUs. The language uses familiar PyTorch syntax and leverages the torch.inductor compiler backend, allowing users to write intuitive code that operates on "tiles" of data. Performance benchmarks demonstrate that Helium is highly competitive, often outperforming handwritten Triton and other compiler solutions on tasks like RMS Normalization and fused attention.

An Introduction to Helium: A High-Level DSL for GPU Kernel Authoring

• 0:00:45 What is Helium? Helium is a new DSL for writing GPU kernels that is easier and more portable than low-level languages like Triton. It provides more control than torch.compile by guaranteeing that a single Helium function compiles to a single, fused GPU kernel.
• 0:01:52 The Power of Autotuning: Helium's key feature is a powerful autotuner that automatically searches thousands of possible low-level implementations for a given kernel. This saves developers from tedious manual tuning and ensures optimal performance across different hardware generations (e.g., H100, B200, AMD).
• 0:03:52 Anatomy of a Helium Kernel: A kernel is written as a Python function with two main parts:
  • Host Code: Regular Python and PyTorch for setup and memory allocation.
  • Device Code: An inner region, defined by an hl.tile loop, which is compiled into the final Triton kernel. This region uses standard PyTorch operations on data tiles.
• 0:06:08 Automating Complex Decisions: The autotuner handles many complex choices automatically, including indexing methods, block sizes, loop flattening, reduction strategies, and hardware-specific Triton flags. This allows the developer to focus on the high-level algorithm.
• 0:19:52 Autotuning in Practice: The autotuning process takes about 20 minutes to find an optimal configuration. This configuration can then be saved and hardcoded into the application, eliminating tuning overhead during production runs.
• 0:22:37 Impressive Performance Gains: Benchmarks show Helium achieving significant speedups. On B200 GPUs, it achieves a geometric mean speedup of 3.2x over eager PyTorch and is nearly twice as fast as expertly handwritten Triton kernels, thanks to the autotuner's comprehensive search.
• 0:50:14 Writing an Attention Kernel: A practical example demonstrates how to implement a fused attention kernel. The core principle is to use the outer hl.tile loop to parallelize over dimensions like batch and heads, and inner loops for sequential operations like the Softmax reduction.
• 1:01:17 Composability and Debugging: Helium kernels work seamlessly within torch.compile. The language provides powerful debugging tools, including an option to print the generated Triton code, disable autotuning for quick correctness checks, or run the kernel in an interpreted eager mode.
• 0:26:38 Beta Release: The team announces the beta release of Helium and encourages community feedback and contributions.