:orphan:

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

.. toctree::
   :maxdepth: 2
   :caption: Contents
   :hidden:

   add
   all_gather_matmul
   attention
   bmm
   concatenate
   cross_entropy
   embedding
   exp
   fp8_attention
   fp8_gemm
   jagged_dense_add
   jagged_mean
   long_sum
   matmul
   matmul_layernorm
   matmul_split_k
   moe_matmul_ogs
   rms_norm
   segment_reduction
   softmax
   sum



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement an element-wise exponential function using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_exp_thumb.png
    :alt:

  :ref:`sphx_glr_examples_exp.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Exponential Function Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement an element-wise addition kernel using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_add_thumb.png
    :alt:

  :ref:`sphx_glr_examples_add.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Element-wise Addition Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a tensor concatenation operation using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_concatenate_thumb.png
    :alt:

  :ref:`sphx_glr_examples_concatenate.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Tensor Concatenation Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a sum reduction operation along the last dimension using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_sum_thumb.png
    :alt:

  :ref:`sphx_glr_examples_sum.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Sum Reduction Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement an embedding lookup operation using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_embedding_thumb.png
    :alt:

  :ref:`sphx_glr_examples_embedding.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Embedding Lookup Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a batch matrix multiplication kernel using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_bmm_thumb.png
    :alt:

  :ref:`sphx_glr_examples_bmm.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Batch Matrix Multiplication Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Helion Layer Normalization Forward Example">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_layer_norm_thumb.png
    :alt:

  :ref:`sphx_glr_examples_layer_norm.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Helion Layer Normalization Forward Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a cross entropy loss function using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_cross_entropy_thumb.png
    :alt:

  :ref:`sphx_glr_examples_cross_entropy.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Cross Entropy Loss Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a Root Mean Square (RMS) normalization operation using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_rms_norm_thumb.png
    :alt:

  :ref:`sphx_glr_examples_rms_norm.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Root Mean Square Normalization Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Helion Softmax Kernel Examples">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_softmax_thumb.png
    :alt:

  :ref:`sphx_glr_examples_softmax.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Helion Softmax Kernel Examples</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="FP8 General Matrix Multiplication (GEMM) with Helion">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_fp8_gemm_thumb.png
    :alt:

  :ref:`sphx_glr_examples_fp8_gemm.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">FP8 General Matrix Multiplication (GEMM) with Helion</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a fused matrix multiplication and layer normalization operation using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_matmul_layernorm_thumb.png
    :alt:

  :ref:`sphx_glr_examples_matmul_layernorm.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Matrix Multiplication with Layer Normalization Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement efficient sum reduction along a long dimension using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_long_sum_thumb.png
    :alt:

  :ref:`sphx_glr_examples_long_sum.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Long Dimension Sum Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Matrix Multiplication with Split-K using Helion">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_matmul_split_k_thumb.png
    :alt:

  :ref:`sphx_glr_examples_matmul_split_k.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Matrix Multiplication with Split-K using Helion</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This code implements a custom attention kernel using Helion and PyTorch for efficient computation of scaled dot-product attention, with support for both static and dynamic input shapes.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_attention_thumb.png
    :alt:

  :ref:`sphx_glr_examples_attention.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Attention Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Helion Matmul Kernel Example">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_matmul_thumb.png
    :alt:

  :ref:`sphx_glr_examples_matmul.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Helion Matmul Kernel Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement an addition operation between a jagged tensor and a dense tensor using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_jagged_dense_add_thumb.png
    :alt:

  :ref:`sphx_glr_examples_jagged_dense_add.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Jagged Dense Addition Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute the mean of each row in a jagged tensor with variable features per row using Helion.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_jagged_mean_thumb.png
    :alt:

  :ref:`sphx_glr_examples_jagged_mean.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Jagged Mean Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Mixture-of-Experts (MoE) Matmul with Outer-Gather-Scatter (OGS)">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_moe_matmul_ogs_thumb.png
    :alt:

  :ref:`sphx_glr_examples_moe_matmul_ogs.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Mixture-of-Experts (MoE) Matmul with Outer-Gather-Scatter (OGS)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to implement a segmented reduction operation using Helion, comparing it with Triton and PyTorch implementations. Code based on https://github.com/pytorch/helion/issues/237">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_segment_reduction_thumb.png
    :alt:

  :ref:`sphx_glr_examples_segment_reduction.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Segmented Reduction Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="All-Gather Matrix Multiplication Example">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_all_gather_matmul_thumb.png
    :alt:

  :ref:`sphx_glr_examples_all_gather_matmul.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">All-Gather Matrix Multiplication Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="FP8 Attention Example">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_fp8_attention_thumb.png
    :alt:

  :ref:`sphx_glr_examples_fp8_attention.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">FP8 Attention Example</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>


.. toctree::
   :hidden:

   /examples/exp
   /examples/add
   /examples/concatenate
   /examples/sum
   /examples/embedding
   /examples/bmm
   /examples/layer_norm
   /examples/cross_entropy
   /examples/rms_norm
   /examples/softmax
   /examples/fp8_gemm
   /examples/matmul_layernorm
   /examples/long_sum
   /examples/matmul_split_k
   /examples/attention
   /examples/matmul
   /examples/jagged_dense_add
   /examples/jagged_mean
   /examples/moe_matmul_ogs
   /examples/segment_reduction
   /examples/all_gather_matmul
   /examples/fp8_attention


.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-gallery

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download all examples in Python source code: examples_python.zip </examples/examples_python.zip>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download all examples in Jupyter notebooks: examples_jupyter.zip </examples/examples_jupyter.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_