PyTorch adds compile-time heuristics to Helion for faster kernel defaults

PyTorch team introduced compile-time heuristics to Helion DSL to provide fast default configurations, addressing the bottleneck of autotuning which previously took 10 to 25 hours for a single paged-attention kernel.
Heuristics classify kernels into matmul, reduction, and pointwise categories, using static analysis of device IR and dataflow to determine block sizes, warps, and stages with negligible compile-time overhead.
On H100, the approach achieved 1.18x forward speedup over Flash Linear Attention (FLA) and 1.27x over vLLM's handwritten CUDA operators for reduction kernels.
Matmul performance reached 1.88x compared to torch.compile (max-autotune), though a regression occurred for jagged HSTU attention where heuristics (3.27x) underperformed the previous default (3.72x).

Implementation Details

Liveness analysis estimates register, shared memory, and tensor memory usage to calculate occupancy and register pressure.
Matmul heuristics adjust block sizes based on CTA counts and maximize num_stages within shared memory limits, while num_warps are determined by Triton's MMA instruction selection.
Reduction heuristics optimize warp counts based on reduction width and apply persistent reduction for kernels like RMSNorm that require multiple reads.

Currently enabled only for NVIDIA H100 and B200 using the Triton backend.
Heuristics rely on predicting other compilers' memory layout decisions, which can introduce errors, particularly on B200 where matmul results reside in tensor memory.
Seeding the autotuner with heuristic results speeds up initial exploration by 35x to 41x but does not change the final configuration quality.