Helion Examples

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

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

Matrix Multiplication Operations

• matmul.py: Basic matrix multiplication

• bmm.py: Batch matrix multiplication

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

• matmul_layernorm.py: Fused matrix multiplication and layer normalization

• fp8_gemm.py: Matrix multiplication using FP8 precision

Attention Operations

• attention.py: Scaled dot-product attention mechanism

• fp8_attention.py: Attention mechanism using FP8 precision

Normalization

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

Sparse and Jagged Tensors

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

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

• segment_reduction.py: Segmented reduction operation

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

Other Operations

• concatenate.py: Tensor concatenation along a dimension

• cross_entropy.py: Cross entropy loss function

• embedding.py: Embedding lookup operation

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

Exponential Function Example

Exponential Function Example

Element-wise Addition Example

Element-wise Addition Example

Tensor Concatenation Example

Tensor Concatenation Example

Sum Reduction Example

Sum Reduction Example

Embedding Lookup Example

Embedding Lookup Example

Batch Matrix Multiplication Example

Batch Matrix Multiplication Example

Helion Layer Normalization Forward Example

Helion Layer Normalization Forward Example

Cross Entropy Loss Example

Cross Entropy Loss Example

Root Mean Square Normalization Example

Root Mean Square Normalization Example

Helion Softmax Kernel Examples

Helion Softmax Kernel Examples

FP8 General Matrix Multiplication (GEMM) with Helion

FP8 General Matrix Multiplication (GEMM) with Helion

Matrix Multiplication with Layer Normalization Example

Matrix Multiplication with Layer Normalization Example

Long Dimension Sum Example

Long Dimension Sum Example

Matrix Multiplication with Split-K using Helion

Matrix Multiplication with Split-K using Helion

Attention Example

Attention Example

Helion Matmul Kernel Example

Helion Matmul Kernel Example

Jagged Dense Addition Example

Jagged Dense Addition Example

Jagged Mean Example

Jagged Mean Example

Mixture-of-Experts (MoE) Matmul with Outer-Gather-Scatter (OGS)

Mixture-of-Experts (MoE) Matmul with Outer-Gather-Scatter (OGS)

Segmented Reduction Example

Segmented Reduction Example

All-Gather Matrix Multiplication Example

All-Gather Matrix Multiplication Example

FP8 Attention Example

FP8 Attention Example