Helion Introduces Compile-Time Heuristics for Kernel Optimization
Eliminating autotuning bottlenecks with smarter kernel configuration.
1Introduction to Helion
Helion is PyTorch's domain-specific language designed for performance-portable machine learning kernels. Traditionally, it relies on autotuning to optimize kernel performance, which can be time-consuming. For instance, tuning a single paged-attention kernel across 72 shapes previously took up to 25 hours in full effort mode.
2Compile-Time Heuristics
To address the inefficiencies of autotuning, Helion now employs compile-time heuristics that analyze a kernel's intermediate representation (IR) and dataflow. These heuristics classify kernels into categories like matmul, reduction, or pointwise, and automatically select optimal configurations such as block sizes, warps, and stages.
3Performance Gains
The new heuristics have shown significant performance improvements. For example, in linear attention kernels, the heuristic configuration outperformed hand-written Triton kernels by 1.18x in forward passes and 1.12x when including backward passes on H100 GPUs. This represents a 2.13x speedup over Helion's previous default settings.
4Reduction Kernels
In vLLM reduction kernels, the heuristic configuration achieved a 1.27x speedup over vLLM's hand-written CUDA/C++ operators and a 2.2x speedup over Helion's previous default. This demonstrates the heuristic's ability to deliver near-autotuned performance without the associated tuning time.
5Compile-Time Analysis
The heuristic approach involves analyzing the kernel's IR to determine structural facts such as tiled axes, tunable block sizes, and live tiles. This analysis allows the heuristic to emit a configuration directly at compile time, bypassing the need for extensive autotuning.
6Heuristic Categories
Kernels are classified into three categories: matmul, reduction, and pointwise. Each category has a specific heuristic that sets values for config knobs like block sizes, num_warps, and num_stages. This classification ensures that the heuristic can handle a wide range of custom kernels written by users.
7Resource Estimation
The heuristic estimates resource usage, such as registers, shared memory, and tensor memory, by running a liveness analysis on the compiler IR. This helps in predicting Triton's behavior regarding register usage, shared memory staging, and value materialization.
8Occupancy and Register Pressure
The heuristic considers occupancy (how many CTAs fit on an SM) and register pressure (peak live registers per thread) to inform config knob selection. For example, if a CTA needs six-tenths of tensor memory, only one CTA per SM is allowed to avoid exceeding the budget.
9Matmul Heuristic Details
The matmul heuristic adjusts block sizes to ensure sufficient CTAs per SM, sets num_stages to hide load latency, and chooses num_warps based on instruction type and register pressure. It also performs a feasibility check to confirm the config can launch successfully.
10Reduction Heuristic Details
The reduction heuristic sets num_warps based on reduction width, residency, and register pressure. It adjusts block sizes to fit more CTAs per SM and considers persistence to save memory trips. The heuristic iterates until both warp count and block sizes stabilize.
11Seeding Autotuner
When autotuning is enabled, the heuristic generates 5-20 related or alternative configs to seed the autotuner. This speeds up the autotuning process by providing promising starting points, though the gains diminish as more configs are tried.
12Conclusion
Helion's compile-time heuristics offer a significant improvement in kernel optimization by reducing or eliminating the need for autotuning. This approach delivers near-autotuned performance with minimal compile-time overhead, making it a valuable tool for machine learning practitioners.
Original article: PyTorch Blog — Helion Introduces Compile-Time Heuristics for Kernel Optimization