DOCSEARCHbyAlgolia
Submit Docs

Helion

repository·main·Indexed 21 days ago

https://github.com/pytorch/helion

A high-level kernel programming framework and Python-embedded DSL for writing ML kernels. Helion allows developers to write GPU kernels that can be lowered to multiple backends, such as Triton or CuTe, and features an autotuning system to optimize tiling, loop, and memory indexing strategies.

Tokens
37.4K
Snippets
101
Records
157
Agent score
75%

What's inside helion

  1. What is Helion?

    main

    Helion is a Python-embedded domain-specific language (DSL) for authoring machine learning kernels. It compiles down to [Triton], providing a higher level of abstraction than raw Triton.

    Helion automates several complex kernel engineering tasks through autotuning, including:

    • Tensor Indexing: Automatically calculates strides and indices, supporting various methods like pointers, block pointers, and TensorDescriptors.
    • Masking: Most masking is implicit and optimized away when unnecessary.
    • Grid Sizes and PID Calculations: Automatically determines grid sizes and maps Program IDs (PIDs) to data tiles.
    • Kernel Arguments Management: Automates handling of tensor sizes, strides, and lifting global variables/closures into kernel arguments.
    • Automated Optimizations: Includes PID swizzling, loop reordering, persistent kernel strategies, and warp specialization.
    • Looping Reductions: Automatically converts large reductions into looped implementations.

  2. Explore Helion operation examples by category

    main

    Helion provides a wide range of example implementations for various tensor operations. You can use these as templates for your own high-performance kernels:

    Matrix Multiplication

    • matmul.py: Basic matrix multiplication
    • bmm.py: Batch matrix multiplication
    • fp8_gemm.py: FP8 precision matrix multiplication
    • int4_gemm.py: INT4 quantized matrix multiplication
    • nvfp4_gemm.py: NVFP4 (E2M1) quantized matrix multiplication
    • grouped_gemm.py: Grouped matrix multiplication

    Attention Mechanisms

    • attention.py: Scaled dot-product attention
    • fp8_attention.py: FP8 precision attention
    • blackwell_attention.py: Optimized for Blackwell architecture
    • flex_attention.py: Flex attention with score modification and block masking

    Sparse and Jagged Tensors

    • jagged_dense_add.py: Addition between jagged and dense tensors
    • jagged_softmax.py: Softmax for jagged tensors
    • moe_matmul_ogs.py: Mixture-of-Experts using Outer-Gather-Scatter

    Distributed Operations

    • distributed/all_gather_matmul.py: All-gather followed by matmul
    • distributed/all_reduce.py: One-shot all-reduce
    • distributed/matmul_reduce_scatter.py: Fused matmul with reduce-scatter

  3. System requirements for Helion

    main

    Helion targets Linux-based systems and requires the following environment:

    Operating System

    • Linux-based OS (other Unix-like systems are not officially supported).

    Python Environment

    • Python 3.10–3.14 (using uv is recommended).

    Core Dependencies

    • PyTorch 2.9 or later
    • Triton 3.5 or later

    Note: Older versions may lack support for features like TMA on Hopper/Blackwell GPUs and may exhibit lower performance.

    GPU Requirements

    • NVIDIA GPUs: Compute Capability 8.0+
    • AMD GPUs: ROCm 6.2+
    • Also supports Intel XPU, CPU, and other architectures via third-party forks.

  4. Use the helion.runtime module for kernel execution and configuration

    main
    The helion.runtime module is the central entry point for managing kernel execution and configuration within Helion. It provides the core execution infrastructure and integrates with low-level utilities for hardware and allocator management.

  5. What is Ahead-of-Time (AOT) Heuristic Tuning?

    main

    Standard Helion autotuning runs at the first call to a kernel for the exact arguments seen. This causes high latency during 'cold-starts' and doesn't generalize to different input shapes.

    AOT autotuning solves this by performing an offline sweep over a representative set of shapes. It tunes each shape, then distills the results into a small decision-tree heuristic. At runtime, this heuristic selects the optimal configuration for a given shape in microseconds without running the autotuner.

    Use AOT when:

    • You are deploying a service/library handling variable input shapes (e.g., LLM serving with variable token counts).
    • You require near-zero-cost configuration selection at runtime.
    • You can define a representative shape sweep offline.

  6. Precedence of Settings vs Environment Variables

    main
    When configuring Helion, if both an environment variable and a kernel decorator argument are provided for the same setting, the kernel decorator argument takes precedence and the environment variable is ignored.

  7. Distinguish between Config and Settings

    main

    Helion uses two distinct parameter types for kernel creation:

    Config (GPU Execution Parameters)

    These control how kernels execute on the hardware. They are performance-focused, hardware-dependent, and are the primary targets for the autotuner.

    • Examples: block_sizes, num_warps, indexing, loop_orders, num_stages.

    Settings (Compilation Control)

    These control how kernels are compiled and the development environment. They are development-focused, not autotuned, and often set via environment variables.

    • Examples: autotune_effort (e.g., "none" to skip autotuning), print_output_code (to debug generated code).
    @helion.kernel(
    # Settings: Control compilation behavior
    autotune_effort="none",      # Skip autotuning for development
    print_output_code=True,       # Debug: show generated code
    # Config: Control GPU execution (when not using default)
    # config=helion.Config(block_sizes=[64, 32], num_warps=8)
)
def debug_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    pass

  8. Understand the Linear Attention module architecture

    main

    The linear attention implementation is organized into two primary modules:

    • linear_attention_engine.py: This is the core module. It contains all 15 @helion.kernel() functions (including the fused recurrent step), the ChunkedLinearAttnFn autograd wrapper, the LinearAttentionEngine class, and the primary public entry points: chunked_linear_attn() and recurrent_step().
    • linear_attention_utils.py: Contains supporting logic, including the pure-PyTorch chunked reference implementation, the naive recurrent reference, WY decomposition helpers, and input generators.

  9. Understand the difference between Settings and Config

    main

    In Helion, it is important to distinguish between Settings and Config to manage your development workflow versus performance optimization:

    • Settings: Control the compilation process and the development environment. They are not autotuned and remain constant across all kernel configurations. Use settings for debugging, logging, and environment setup (e.g., print_output_code, autotune_effort).
    • Config: Control execution performance. These are automatically optimized during the autotuning process (e.g., block_sizes, num_warps).

    Use Settings when you want to change how Helion builds kernels, and use Config when you want to change how the kernel runs on hardware.

  10. Manage kernel execution and compilation settings via helion.runtime

    main

    The helion.runtime module coordinates several specialized sub-modules to manage the lifecycle of a kernel. To perform specific tasks, use the following related modules:

    • Configuration management: Use the helion.runtime.config module.
    • Kernel execution: Use the helion.runtime.kernel module.
    • Compilation settings: Use the helion.runtime.settings module.

  11. Understand tuning knob modifications for TileIR

    main

    The TileIR backend modifies several standard Helion tuning knobs. These changes are exclusive to the TileIR backend and do not affect other backends.

    New Knobs

    • num_ctas: Number of CTAs in one CGA.
    • occupancy: Hardware utilization/occupancy.

    Modified Knobs

    • num_warps: Constrained to 4 (acts as a placeholder).
    • num_stages: Changed from IntegerFragment to EnumFragment (range: 1-10). This is analogous to latency in cuTile.
    • indexing: The block_ptr type is unsupported. Use pointer or tensor_descriptor instead.

    Unsupported Knobs

    The following knobs are removed from the autotuning search space for TileIR to improve efficiency:

    • static_ranges
    • range_unroll_factors
    • range_multi_buffers
    • range_flattens
    • range_warp_specialize
    • load_eviction_policies