Helion Examples

This directory contains examples demonstrating how to use Helion for high-performance tensor operations. The examples are organized into the following categories:

Pretuned Kernels (run as-is, no autotuning)

The pretuned_kernels/ directory holds runnable kernels that ship with checked-in AOT heuristic files (currently tuned for NVIDIA B200 / sm100). Helion picks the checked-in config at startup, so these kernels run immediately without any online autotuning — useful as copy/paste recipes for common patterns or as a quick way to try Helion on a supported GPU. Each kernel module has a main() that benchmarks against the PyTorch eager baseline.

• vector_add — element-wise addition.

• softmax — softmax with a long-context shape sweep.

• layer_norm — layer normalization across realistic hidden sizes.

• rms_norm — RMS normalization with NPOT and LLM-shaped inputs.

• cross_entropy — cross-entropy across LLM vocabulary sizes.

To pretune one of these kernels for a different GPU (or to ship a heuristic for your own kernel), see the AOT Heuristic Tuning section of the deployment guide.

Basic Operations

• add.py: Element-wise addition with broadcasting support

• exp.py: Element-wise exponential function

• sum.py: Sum reduction along the last dimension

• long_sum.py: Efficient sum reduction along a long dimension

• softmax.py: Different implementations of the softmax function

• batch_softmax.py: Batched (3D) softmax with arithmetic broadcasting

• concatenate.py: Tensor concatenation along a dimension

• low_mem_dropout.py: Memory-efficient dropout implementation

Matrix Multiplication Operations

• matmul.py: Basic matrix multiplication

• bmm.py: Batch matrix multiplication

• broadcast_matmul.py: Batch matrix multiplication with broadcasting (weight without batch dimension)

• matmul_split_k.py: Matrix multiplication using split-K algorithm for better parallelism

• split_k_barrier.py: Split-K matmul with barrier synchronization for deterministic results

• matmul_layernorm.py: Fused matrix multiplication and layer normalization

• fp8_gemm.py: Matrix multiplication using FP8 precision

• bf16xint16_gemm.py: BF16 x INT16 matrix multiplication

• int4_gemm.py: INT4 quantized matrix multiplication

• nvfp4_gemm.py: NVFP4 (E2M1) quantized matrix multiplication

• grouped_gemm.py: Grouped matrix multiplication

• gather_gemv.py: Gather-based matrix-vector multiplication

Attention Operations

• attention.py: Scaled dot-product attention mechanism

• fp8_attention.py: Attention mechanism using FP8 precision

• blackwell_attention.py: Attention optimized for Blackwell architecture

• flex_attention.py: Flex attention with score modification and block masking

Normalization

• rms_norm.py: Root Mean Square (RMS) normalization

• layer_norm.py: Layer normalization

Activation Functions

• geglu.py: Gated Linear Unit (GEGLU) activation

• swiglu.py: SwiGLU activation function

Loss Functions

• cross_entropy.py: Cross entropy loss function

• grpo_loss.py: Group Relative Policy Optimization (GRPO) loss function

• jsd.py: Jensen-Shannon Divergence

• fused_linear_jsd.py: Fused linear layer with JSD loss

• kl_div.py: Kullback-Leibler divergence

Sparse and Jagged Tensors

• jagged_dense_add.py: Addition between a jagged tensor and a dense tensor

• jagged_dense_bmm.py: Batch matrix multiplication with jagged tensors

• jagged_mean.py: Computing the mean of each row in a jagged tensor

• jagged_sum.py: Sum reduction for jagged tensors

• jagged_softmax.py: Softmax for jagged tensors

• jagged_layer_norm.py: Layer normalization for jagged tensors

• jagged_hstu_attn.py: HSTU attention for jagged tensors

• segment_reduction.py: Segmented reduction operation

• moe_matmul_ogs.py: Mixture-of-Experts matrix multiplication using Outer-Gather-Scatter

Sequence Models

• mamba2_chunk_scan.py: Mamba2 chunk scan operation

• mamba2_chunk_state.py: Mamba2 chunk state operation

Statistics

• welford.py: Welford’s online algorithm for computing variance

Neural Network Components

• embedding.py: Embedding lookup operation

• squeeze_and_excitation_net.py: Squeeze-and-Excitation network

• gdn_fwd_h.py: Generalized Divisive Normalization (GDN) forward pass

Advanced Usage

• aot_example.py: Ahead-of-time (AOT) autotuning workflow with batch-aware heuristics

• acfs/softmax_acf.py: Using Advanced Controls Files (ACFs) with kernel configurations and autotuning

Distributed Operations

• distributed/all_gather_matmul.py: All-gather operation followed by matrix multiplication

• distributed/all_reduce.py: All-reduce operation (one-shot)

• distributed/matmul_reduce_scatter.py: Fused matmul with reduce-scatter

• distributed/one_shot_allreduce_bias_rmsnorm.py: Fused all-reduce, bias add, and RMS normalization