Mirage

https://mirage-project.readthedocs.io/en/latest/
Admin
Mirage is a tensor algebra superoptimizer that automatically discovers highly-optimized tensor...
Tokens:7,729
Snippets:41
Trust Score:8.3
Update:1 year ago
Show doc for...
Tokens:
Raw
### Install Mirage from PyPI

Source: https://mirage-project.readthedocs.io/en/latest/installation

Installs the latest stable release of the Mirage project directly from the Python Package Index (PyPI). This is the quickest method to start using Mirage.

```bash
pip install mirage-project
```

--------------------------------

### Build Mirage C++ Library from Source

Source: https://mirage-project.readthedocs.io/en/latest/installation

Compiles the Mirage C++ library from its source code. This process involves navigating to the Mirage root directory, creating a build folder, configuring with CMake, and executing the build and installation steps. A static library is generated by default.

```bash
cd $MIRAGE_ROOT
mkdir build; cd build
cmake ..
make -j
make install
```

--------------------------------

### Install Mirage from Pre-built Wheel

Source: https://mirage-project.readthedocs.io/en/latest/installation

Installs a specific pre-built binary wheel for Mirage from a GitHub release page. This command specifies the Mirage version, CUDA version, and Python version for the wheel.

```bash
pip install https://github.com/mirage-project/mirage/releases/download/v0.2.2/mirage_project-0.2.2+cu122-cp310-cp310-linux_x86_64.whl
```

--------------------------------

### Build and Install Mirage from Source

Source: https://mirage-project.readthedocs.io/en/latest/installation

Builds and installs the Mirage Python package in editable mode from the cloned source code. This command also automatically installs all necessary Python dependencies.

```bash
pip install -e . -v
```

--------------------------------

### Check Mirage Installation

Source: https://mirage-project.readthedocs.io/en/latest/installation

Verifies that the Mirage project has been successfully installed by attempting to import it within a Python interpreter. No output indicates a successful import.

```python
python -c 'import mirage'
```

--------------------------------

### Run Mirage Demo Script

Source: https://mirage-project.readthedocs.io/en/latest/installation

Executes a Python demo script for Mirage, demonstrating specific functionalities like group query attention. It requires specifying a checkpoint file for the model.

```python
python demo/demo_group_query_attention_spec_decode.py --checkpoint demo/checkpoint_group_query_attn_spec_decode.json
```

--------------------------------

### Clone Mirage Source Code

Source: https://mirage-project.readthedocs.io/en/latest/installation

Clones the Mirage project's source code repository from GitHub, including submodules, to prepare for building from source.

```bash
git clone --recursive https://www.github.com/mirage-project/mirage
cd mirage
```

--------------------------------

### Clone Mirage Git Repository

Source: https://mirage-project.readthedocs.io/en/latest/installation

Clones the Mirage project's Git repository, including all necessary submodules. This is the initial step to obtain the project files and scripts required for building or running Mirage.

```git
git clone --recursive https://www.github.com/mirage-project/mirage
```

--------------------------------

### Basic Triton Transpiler Usage Example

Source: https://mirage-project.readthedocs.io/en/latest/_sources/triton-transpiler.rst

Provides a basic example of using the Triton transpiler within the Mirage framework. It shows how to create a kernel graph, define inputs and operations (like RMS normalization and matrix multiplication), optimize the graph for the 'triton' backend, and generate/save the resulting Triton code to a file.

```python
import mirage as mi
import torch

# Create a kernel graph
graph = mi.new_kernel_graph()

# Define inputs and operations
X = graph.new_input(dims=(16, 4096), dtype=mi.float16)
W = graph.new_input(dims=(4096, 6144), dtype=mi.float16)
D = graph.rms_norm(X, normalized_shape=(4096,))
O = graph.matmul(D, W)
graph.mark_output(O)

# Optimize and generate Triton code
optimized_graph = graph.superoptimize(config="mlp", backend="triton")

# Generate and save the code
with open("triton_generated.py", "w") as f:
    f.write(mi.generate_triton_program(
        optimized_graph.cygraph,
        target_cc=10){"code"])
```

--------------------------------

### Generated CUDA Code Example

Source: https://mirage-project.readthedocs.io/en/latest/_sources/cuda-transpiler.rst

This snippet shows an example of CUDA code generated by Mirage's transpiler. It includes runtime headers and demonstrates a function `_execute_mugraph` that calls a matrix multiplication kernel (`kn::gemm`) from the runtime library.

```cpp
#define NUM_GPUS 1
#define USE_NVSHMEM 0
#include "runtime.h"
using namespace cute;

static void _init() {
}

static void _execute_mugraph(std::vector<void const *> input_tensors, std::vector<void*> output_tensors, void* buf) {
  {
    // OP type: kn_matmul_op
    half_t *dtensor10000000 = (half_t*)input_tensors.at(0);
    half_t *dtensor10000001 = (half_t*)input_tensors.at(1);
    half_t *dtensor10000002 = (half_t*)output_tensors.at(0);
    kn::gemm<CUBLAS_COMPUTE_16F>(dtensor10000002,dtensor10000000,dtensor10000001, 6,4,9, 16,1, 8,1, 1,8, 3, 128,128,64);
  }
}
```

--------------------------------

### Build Z3 C++ Library from Source

Source: https://mirage-project.readthedocs.io/en/latest/installation

Compiles the Z3 C++ library from its source code. Requires navigating to the Z3 dependency directory, creating a build folder, configuring with CMake, and executing the build process.

```bash
cd $MIRAGE_ROOT/deps/z3
mkdir build; cd build
cmake ..
make -j
```

--------------------------------

### Export Z3 Build Directory

Source: https://mirage-project.readthedocs.io/en/latest/installation

Sets the Z3_DIR environment variable to point to the compiled Z3 build directory. This is a prerequisite for building the Mirage library.

```bash
export Z3_DIR=$MIRAGE_ROOT/deps/z3/build
```

--------------------------------

### Generated CUDA Code Example

Source: https://mirage-project.readthedocs.io/en/latest/cuda-transpiler

This snippet shows an example of CUDA code generated by the Mirage transpiler. It includes runtime headers, namespace usage, initialization, and an execution function (_execute_mugraph) demonstrating a matrix multiplication operation (kn::gemm) with specific compute capabilities and tensor dimensions.

```C++
#define NUM_GPUS 1
#define USE_NVSHMEM 0
#include "runtime.h"
using namespace cute;

static void _init() {
}

static void _execute_mugraph(std::vector<void const *> input_tensors, std::vector<void*> output_tensors, void* buf) {
  {
    // OP type: kn_matmul_op
    half_t *dtensor10000000 = (half_t*)input_tensors.at(0);
    half_t *dtensor10000001 = (half_t*)input_tensors.at(1);
    half_t *dtensor10000002 = (half_t*)output_tensors.at(0);
    kn::gemm<CUBLAS_COMPUTE_16F>(dtensor10000002,dtensor10000000,dtensor10000001, 6,4,9, 16,1, 8,1, 1,8, 3, 128,128,64);
  }
}

```

--------------------------------

### Run Mirage Docker Image

Source: https://mirage-project.readthedocs.io/en/latest/installation

Executes a Mirage Docker image, allowing users to run Mirage in a containerized environment. The command specifies the Docker image to run, defaulting to CUDA version 12.4.

```bash
/path-to-mirage/docker/run_docker.sh mlso/mirage
```

--------------------------------

### Basic Triton Transpiler Usage with Mirage

Source: https://mirage-project.readthedocs.io/en/latest/triton-transpiler

This example shows how to use Mirage's `superoptimize` function with the 'triton' backend to automatically search, profile, and select the best kernel arrangement. It then generates and saves the Triton code to a file.

```python
import mirage as mi
import torch

# Create a kernel graph
graph = mi.new_kernel_graph()

# Define inputs and operations
X = graph.new_input(dims=(16, 4096), dtype=mi.float16)
W = graph.new_input(dims=(4096, 6144), dtype=mi.float16)
D = graph.rms_norm(X, normalized_shape=(4096,))
O = graph.matmul(D, W)
graph.mark_output(O)

# Optimize and generate Triton code
optimized_graph = graph.superoptimize(config="mlp", backend="triton")

# Generate and save the code
with open("triton_generated.py", "w") as f:
    f.write(mi.generate_triton_program(
        optimized_graph.cygraph,
        target_cc=10
    )("code"))
```

--------------------------------

### Python Superoptimization Example

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/multi-latent-attention.rst

This Python code snippet demonstrates the superoptimization of a multi-latent attention graph using the 'mirage' library. It defines input tensors, performs attention calculations, and optimizes the resulting graph.

```Python
import torch
import mirage as mi
graph = mi.new_kernel_graph()
Q = graph.new_input(dims=(2, 256, 64), dtype=mi.float16)
K = graph.new_input(dims=(2, 64, 4096), dtype=mi.float16)
V = graph.new_input(dims=(2, 4096, 64), dtype=mi.float16)
A = graph.matmul(Q, K)
E = graph.exp(A)
S = graph.reduction(E, 2)
D = graph.div(E, S)
O = graph.matmul(D, V)
graph.mark_output(O)
optimized_graph = mi.superoptimize(graph)
```

--------------------------------

### Triton Profiler Implementation Steps

Source: https://mirage-project.readthedocs.io/en/latest/_sources/triton-transpiler.rst

Demonstrates the key steps involved in the Triton profiler, including making temporary files, measuring execution time, filtering invalid kernels, and selecting the best performing variant. It shows how to iterate through graphs, perform warmup and profiling runs, and calculate performance metrics.

```python
for idx, g in enumerate(all_graphs):
    if not g.valid_kernels():
        print("muGraph {}: skipping since its shared memory usage exceed limit".format(idx))
        continue
    # Warmup runs
    for _ in range(16):
        g(inputs=input_tensors)
    torch.cuda.synchronize()
    starter.record()
    for _ in range(1000):
        g(inputs=input_tensors)
    ender.record()
    torch.cuda.synchronize()
    perf = starter.elapsed_time(ender) / 1000
```

--------------------------------

### Create and Optimize Mirage Kernel Graph

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/rms-norm-linear

Demonstrates creating a computation graph in Mirage for RMSNorm followed by MatMul. It shows how to define inputs, fuse operations, and optimize the graph for performance by leveraging Mirage's uGraph structure.

```Python
import mirage as mi
graph = mi.new_kernel_graph()
X = graph.new_input(dims=(16, 4096), dtype=mi.float16)
G = graph.new_input(dims=(1, 4096), dtype=mi.float16)
W = graph.new_input(dims=(406, 4096), dtype=mi.float16)
O = graph.matmul(graph.rms_norm(X, G), W)
graph.mark_output(O)
optimized_graph = mi.superoptimize(graph)
```

--------------------------------

### Profiling Kernel Variants with Triton

Source: https://mirage-project.readthedocs.io/en/latest/triton-transpiler

This snippet demonstrates the core profiling steps for evaluating multiple kernel variants. It includes warmup runs, timing execution, and calculating performance metrics. Dependencies include PyTorch and timing utilities.

```python
for idx, g in enumerate(all_graphs):
    if not g.valid_kernels():
        print("muGraph {}: skipping since its shared memory usage exceed limit".format(idx))
        continue
    # Warmup runs
    for _ in range(16):
        g(inputs=input_tensors)
    torch.cuda.synchronize()
    starter.record()
    for _ in range(1000):
        g(inputs=input_tensors)
    ender.record()
    torch.cuda.synchronize()
    perf = starter.elapsed_time(ender) / 1000
```

--------------------------------

### Superoptimize LoRA Adapters in Mirage

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/lora

Demonstrates the process of defining and superoptimizing LoRA adapters using the Mirage framework. It initializes a kernel graph, sets up input tensors for LoRA computation (X, W, A, B), constructs the computation graph using matrix multiplications and additions, and then applies Mirage's superoptimization to generate an efficient kernel.

```python
import mirage as mi
graph = mi.new_kernel_graph()
X = graph.new_input(dims=(16, 256), dtype=mi.float16)
W = graph.new_input(dims=(256, 4096), dtype=mi.float16)
A = graph.new_input(dims=(256, 16), dtype=mi.float16)
B = graph.new_input(dims=(16, 4096), dtype=mi.float16)
D = graph.matmul(X, A)
E = graph.matmul(D, B)
C = graph.matmul(X, W)
O = graph.add(C, E)
opt_graph = graph.superoptimize()
```

--------------------------------

### Execute Optimized Mirage Kernel

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/rms-norm-linear

Shows how to execute a pre-optimized Mirage kernel graph. The optimized graph is called with input tensors, triggering Mirage to transpile the uGraph into CUDA code, compile it, and launch the compiled kernel for execution.

```Python
import torch
input_tensors = [
    torch.randn(16, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(1, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(406, 4096, dtype=torch.float16, device='cuda:0')
]
optimized_graph(input_tensors)
```

--------------------------------

### Superoptimize Gated MLP with Mirage Python

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/gated-mlp

Demonstrates how to construct a kernel graph in Mirage for a Gated MLP computation, superoptimize it, and then execute the optimized graph. This process is crucial for efficient execution of complex neural network layers in large language models.

```python
import torch
import mirage as mi

# Construct a kernel graph for the gated MLP computation
graph = mi.new_kernel_graph()
X = graph.new_input(dims=(8, 4096), dtype=mi.float16)
W1 = graph.new_input(dims=(4096, 4096), dtype=mi.float16)
W2 = graph.new_input(dims=(4096, 4096), dtype=mi.float16)

D1 = graph.matmul(X, W1)
D2 = graph.matmul(X, W2)

# Apply SiLU activation and element-wise multiplication
O = graph.mul(graph.silu(D1, D2))
graph.mark_output(O)

# Superoptimize the constructed graph
optimized_graph = mi.superoptimize(graph)

# Prepare input tensors for execution
input_tensors = [
    torch.randn(8, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(4096, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(4096, 4096, dtype=torch.float16, device='cuda:0')
]

# Execute the superoptimized graph
optimized_graph(input_tensors)
```

--------------------------------

### Accumulator Storage in Register File vs. Shared Memory

Source: https://mirage-project.readthedocs.io/en/latest/cuda-transpiler

Illustrates how to manage accumulator storage for matrix multiplication, comparing the use of register files versus shared memory. It shows the typical workflow using `get_mma_rc`, `run`, and `write_back_mma_rc` methods for each scenario.

```cpp
auto t = get_mma_rc();
for (...) {
  run(t, ...);
}
write_back_mma_rc(t, ...);
```

```cpp
for (...) {
  auto t = get_mma_rc();
  run(t, ...);
  write_back_mma_rc(t, ...);
}
```

--------------------------------

### Implement QK Normalization and Optimize Attention Kernel in Python

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/attention-with-kv-norm.rst

This Python code demonstrates the application of Query-Key (QK) normalization before attention computation, as used in the Chameleon model. It utilizes the Mirage library and PyTorch to build and optimize an attention kernel, normalizing Q and K tensors before matrix multiplication.

```Python
import torch
import mirage as mi
graph = mi.new_kernel_graph()
Q = graph.new_input(dims=(2, 256, 64), dtype=mi.float16)
K = graph.new_input(dims=(2, 64, 4096), dtype=mi.float16)
V = graph.new_input(dims=(2, 4096, 64), dtype=mi.float16)
Q = graph.rms_norm(Q)
K = graph.rms_norm(K)
A = graph.matmul(Q, K)
E = graph.exp(A)
S = graph.reduction(E, 2)
D = graph.div(E, S)
O = graph.matmul(D, V)
graph.mark_output(O)
optimized_graph = mi.superoptimize(graph)
```

--------------------------------

### Optimizing thr_layout for make_tiled_mma

Source: https://mirage-project.readthedocs.io/en/latest/cuda-transpiler

Explains how to select the `thr_layout` parameter for CuTe's `make_tiled_mma` function to optimize data copying costs when distributing MMA (Matrix Multiply-Accumulate) operations across thread groups. It involves enumerating layouts and calculating costs based on MMA atom calls and data volume.

```APIDOC
CuTe MMA Primitives:

  MMAAtom:
    Represents a basic matrix multiplication operation performed by a group of threads.
    Example: SM80_16x8x16_F16F16F16F16_TN (group size: 32 threads).

  TiledMMA:
    Manages distribution of MMA jobs across multiple thread groups.

  make_tiled_mma(..., thr_layout, ...):
    - `thr_layout`: Parameter to define how MMA atoms are assigned to thread groups.
    - Goal: Minimize data copying cost.

  Layout Optimization Strategy:
    1. Enumerate the number of groups along the m and n axes (k-axis groups usually set to 1).
    2. For each layout, calculate:
       - Number of MMA atom calls.
       - Total volume of data copied.
    3. Choose the layout with the minimum calculated cost.

  Example Scenario:
    - MMAAtom: SM80_16x8x16_F16F16F16F16_TN (group size = 32 threads).
    - Total threads: 128 (4 groups).
    - MMA job size: 64x32x16 (m x n x k).
    - Four methods exist to assign MMA atoms to thread groups, visualized by group assignments to matrix tiles.
```

--------------------------------

### Superoptimize Attention with QK Normalization using Mirage

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/attention-with-kv-norm

This snippet demonstrates the process of superoptimizing an attention kernel that incorporates Query-Key (QK) normalization. It utilizes the Mirage library to build and optimize a computational graph, applying RMS normalization to query and key inputs before calculating attention scores and outputs. The code requires PyTorch and Mirage.

```python
import torch
import mirage as mi
graph = mi.new_kernel_graph()
Q = graph.new_input(dims=(2, 256, 64), dtype=mi.float16)
K = graph.new_input(dims=(2, 64, 4096), dtype=mi.float16)
V = graph.new_input(dims=(2, 4096, 64), dtype=mi.float16)
Q = graph.rms_norm(Q)
K = graph.rms_norm(K)
A = graph.matmul(Q, K)
E = graph.exp(A)
S = graph.reduction(E, 2)
D = graph.div(E, S)
O = graph.matmul(D, V)
graph.mark_output(O)
optimized_graph = mi.superoptimize(graph)
```

--------------------------------

### Execute Optimized Mirage Kernel with PyTorch

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/rms-norm-linear.rst

This snippet demonstrates how to execute the Mirage-optimized kernel using PyTorch tensors. The optimized graph object can be called directly, which triggers Mirage to transpile, compile, and launch the CUDA kernel.

```Python
import torch
input_tensors = [
    torch.randn(16, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(1, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(406, 4096, dtype=torch.float16, device='cuda:0')
]
optimized_graph(input_tensors)
```

--------------------------------

### Launch Mirage-Generated GQA Kernel

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/group-query-attention

This snippet shows how to execute the `optimized_graph` generated by Mirage. It prepares input tensors using PyTorch, ensuring they match the expected format and device (CUDA). The `optimized_graph` is then called with these tensors, triggering Mirage to transpile, compile, and launch the optimized CUDA kernel for the GQA operation.

```Python
import torch
input_tensors = [
    torch.randn(64, 1, 128, dtype=torch.float16, device='cuda:0'),
    torch.randn(64, 128, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(64, 4096, 128, dtype=torch.float16, device='cuda:0')
]
# Launch the Mirage-generated kernel to perform attention
output = optimized_graph(input_tensors)
```

--------------------------------

### Mirage: Run Superoptimized GQA Kernel

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/group-query-attention.rst

Executes the superoptimized GQA kernel generated by Mirage. It prepares input tensors using PyTorch, ensuring they match the expected format and device. The `optimized_graph` object is then called like a function, triggering Mirage to transpile, compile, and launch the optimized CUDA kernel.

```Python
import torch
input_tensors = [
    torch.randn(64, 1, 128, dtype=torch.float16, device='cuda:0'),
    torch.randn(64, 128, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(64, 4096, 128, dtype=torch.float16, device='cuda:0')
]
# Launch the Mirage-generated kernel to perform attention
output = optimized_graph(input_tensors)
```

--------------------------------

### CuTe MMA Primitives and make_tiled_mma

Source: https://mirage-project.readthedocs.io/en/latest/_sources/cuda-transpiler.rst

Explains CuTe primitives like MMAAtom and TiledMMA, used for matrix multiplication. It details how to decide the 'thr_layout' parameter for make_tiled_mma by enumerating thread group distributions and minimizing data copying costs.

```APIDOC
CuTe MMA Primitives:

  MMAAtom:
    Represents a basic matrix multiplication operation performed by a group of threads.
    Example: SM80_16x8x16_F16F16F16F16_TN (group size: 32 threads).

  TiledMMA:
    Used for distributing larger MMA jobs across multiple thread groups.

  make_tiled_mma Function:
    - Parameter: thr_layout (thread layout)
    - Purpose: Distributes MMAAtom operations to thread groups.
    - Decision Strategy for thr_layout:
      1. Enumerate the number of groups along the m and n axes (k axis usually 1).
      2. Calculate the number of MMA atom calls and total data copied for each layout.
      3. Choose the layout with the minimum cost (data copying).

  Example Scenario:
    - MMA Job Size: 64x32x16 (m x n x k)
    - Group Size: 32 threads (e.g., SM80_16x8x16_F16F16F16F16_TN)
    - Total Threads: 128 (4 groups)
    - Layouts: Discusses 4 methods to assign MMAAtoms to thread groups, illustrated by a figure.
```

--------------------------------

### Superoptimize Gated MLP with Mirage Python

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/gated-mlp.rst

Constructs a kernel graph in Mirage representing a gated MLP computation, superoptimizes it, and executes the optimized graph. This process leverages PyTorch for tensor operations and Mirage for graph optimization. The optimized graph can be called directly to launch efficient kernels.

```Python
import torch
import mirage as mi

graph = mi.new_kernel_graph()
X = graph.new_input(dims=(8, 4096), dtype=mi.float16)
W1 = graph.new_input(dims=(4096, 4096), dtype=mi.float16)
W2 = graph.new_input(dims=(4096, 4096), dtype=mi.float16)

D1 = graph.matmul(X, W1)
D2 = graph.matmul(X, W2)

O = graph.mul(graph.silu(D1, D2))
graph.mark_output(O)

optimized_graph = mi.superoptimize(graph)

input_tensors = [
    torch.randn(8, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(4096, 4096, dtype=torch.float16, device='cuda:0'),
    torch.randn(4096, 4096, dtype=torch.float16, device='cuda:0')
]

optimized_graph(input_tensors)
```

--------------------------------

### Superoptimize LoRA Adapters in Mirage

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/lora.rst

This Python code snippet demonstrates the process of creating a kernel graph in Mirage, defining input tensors with specified dimensions and data types (float16), performing matrix multiplication and addition operations, and finally applying superoptimization to the graph. This process aims to generate highly efficient kernels for LoRA computations.

```Python
import mirage as mi
graph = mi.new_kernel_graph()
X = graph.new_input(dims=(16, 256), dtype=mi.float16)
W = graph.new_input(dims=(256, 4096), dtype=mi.float16)
A = graph.new_input(dims=(256, 16), dtype=mi.float16)
B = graph.new_input(dims=(16, 4096), dtype=mi.float16)
D = graph.matmul(X, A)
E = graph.matmul(D, B)
C = graph.matmul(X, W)
O = graph.add(C, E)
opt_graph = graph.superoptimize()
```

--------------------------------

### Build Mirage Kernel Graph for RMSNorm and MatMul

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/rms-norm-linear.rst

This snippet shows how to define a computation graph using Mirage to fuse RMSNorm and a MatMul operation. It involves creating input tensors and chaining operations like rms_norm and matmul before optimizing the graph.

```Python
import mirage as mi
graph = mi.new_kernel_graph()
X = graph.new_input(dims=(16, 4096), dtype=mi.float16)
G = graph.new_input(dims=(1, 4096), dtype=mi.float16)
W = graph.new_input(dims=(406, 4096), dtype=mi.float16)
O = graph.matmul(graph.rms_norm(X, G), W)
graph.mark_output(O)
optimized_graph = mi.superoptimize(graph)
```

--------------------------------

### Matmul Kernel Accumulator Storage Logic

Source: https://mirage-project.readthedocs.io/en/latest/_sources/cuda-transpiler.rst

Demonstrates the usage of tb::Matmul kernel methods for managing accumulators, either in the register file or written back to shared memory.

```APIDOC
tb::Matmul Kernel Methods for Accumulator Management:

- get_mma_rC()
  - Description: Retrieves a fragment (tensor on register) that stores the accumulator.
  - Usage: Called to obtain the accumulator storage.

- run(accumulator_fragment, ...)
  - Description: Performs the matmul operation using the provided accumulator fragment.
  - Usage: Called within a loop to process data.

- write_back_mma_rC(accumulator_fragment, ...)
  - Description: Writes the accumulator fragment back to shared memory.
  - Usage: Called after processing to persist the accumulator.

Example: Storing accumulator in register file:
```cpp
auto t = get_mma_rc();
for (...) {
  run(t, ...);
}
write_back_mma_rc(t, ...);
```

Example: Storing accumulator in shared memory:
```cpp
for (...) {
  auto t = get_mma_rc();
  run(t, ...);
  write_back_mma_rc(t, ...);
}
```

Heuristic for decision:
- If per-thread register burden < 192, store accumulator in register file.
- Otherwise, store accumulator in shared memory.
```

--------------------------------

### Mirage: Define GQA Computation Graph

Source: https://mirage-project.readthedocs.io/en/latest/_sources/tutorials/group-query-attention.rst

Defines the computation graph for Group-Query Attention (GQA) using the Mirage library. It sets up input tensors for Q, K, and V with specified dimensions and data types, then constructs the attention calculation through matrix multiplications and reductions. The `superoptimize` function is called to find an optimized kernel configuration.

```Python
import mirage as mi
graph = mi.new_kernel_graph()
Q = graph.new_input(dims=(2, 256, 64), dtype=mi.float16)
K = graph.new_input(dims=(2, 64, 4096), dtype=mi.float16)
V = graph.new_input(dims=(2, 4096, 64), dtype=mi.float16)
A = graph.matmul(Q, K)
E = graph.exp(A)
S = graph.reduction(E, 2)
D = graph.div(E, S)
O = graph.matmul(D, V)
optimized_graph = graph.superoptimize(config="attention")
```

--------------------------------

### Chunked Copy Feasibility and Innermost Dimension

Source: https://mirage-project.readthedocs.io/en/latest/_sources/cuda-transpiler.rst

Discusses the conditions for performing chunked copying (e.g., in uint128_t) for G->S or S->G transfers to improve performance. It defines 'chunk congruent' layout and the 'real innermost dimension' for feasibility checks.

```APIDOC
Chunked Copying:
  - Goal: Improve performance for G->S or S->G transfers by copying in chunks (e.g., uint128_t).
  - Requirement: Layout must be 'chunk congruent'.
    - Definition: Every 16 Bytes (8 halfs) in the DTensor are contiguous in the STensor.

  Feasibility Check:
    - Slow Method: Iterate every 16 bytes in DTensor and check contiguity in STensor.
    - Efficient Method:
      1. Find the 'real innermost dimension': Defined as the non-1 dimension with a stride of 1 (must be unique).
      2. Chunked copy is performed 'along' this dimension.
```

--------------------------------

### Triton Transpiler Debug Mode Configuration

Source: https://mirage-project.readthedocs.io/en/latest/triton-transpiler

This snippet illustrates how to enable debug mode for the Triton transpiler. Setting `debug_mode=True` provides detailed output for each generated code variant and saves intermediate files, aiding in debugging transpiled Triton code.

```python
elif backend == "triton":
            return profile_and_select_best_graph(all_graphs, target_cc=torch.cuda.get_device_properties(0).major * 10 + torch.cuda.get_device_properties(0).minor, warmup_iters=16, profile_iters=1000, debug_mode=False)
```

--------------------------------

### Define GQA Computation Graph with Mirage

Source: https://mirage-project.readthedocs.io/en/latest/tutorials/group-query-attention

This snippet demonstrates how to define a computation graph for Group-Query Attention (GQA) using the Mirage library. It initializes a kernel graph, defines input tensors for Q, K, and V with specified dimensions and data types, and constructs the attention computation. Finally, it calls `superoptimize` to find an optimized version of the graph.

```Python
import mirage as mi
graph = mi.new_kernel_graph()
Q = graph.new_input(dims=(2, 256, 64), dtype=mi.float16)
K = graph.new_input(dims=(2, 64, 4096), dtype=mi.float16)
V = graph.new_input(dims=(2, 4096, 64), dtype=mi.float16)
A = graph.matmul(Q, K)
E = graph.exp(A)
S = graph.reduction(E, 2)
D = graph.div(E, S)
O = graph.matmul(D, V)
optimized_graph = graph.superoptimize(config="attention")
```

--------------------------------

### CUDA ldmatrix Instruction Variants and Constraints

Source: https://mirage-project.readthedocs.io/en/latest/_sources/cuda-transpiler.rst

Details the variants of the CUDA ldmatrix instruction, specifically focusing on the .num = .x4 variant. It highlights the requirement for all threads in a warp to execute ldmatrix identically and the implications of the .aligned qualifier for conditional execution.

```APIDOC
CUDA ldmatrix Instruction:
  Variants:
    - .num = .x1
    - .num = .x2
    - .num = .x4 (focus of discussion)

  Functionality:
    Each thread provides an address pointing to 8 elements in shared memory (consecutive in one row or column). After execution, each thread holds 8 elements in registers for MMA instructions.

  Constraints:
    - .aligned qualifier: All threads in a warp must execute the same ldmatrix instruction.
    - Conditional Execution: ldmatrix should only be used if all threads in the warp evaluate the condition identically; otherwise, behavior is undefined.

  Memory Handling:
    - Out-of-bound addresses must be mapped to valid memory locations containing zero.
    - Red areas (out-of-bound elements affecting results) must be zero-filled.
    - Innermost dimension is padded to a multiple of 8 for safe addressing.
```

--------------------------------

### Handling Non-Divisible Matrix Sizes with ldmatrix

Source: https://mirage-project.readthedocs.io/en/latest/cuda-transpiler

Discusses challenges and solutions for performing threadblock-level matrix multiplication in CUDA when matrix dimensions are not divisible by the MMA (Matrix Multiply-Accumulate) tile size. It focuses on the `ldmatrix` instruction, explaining how out-of-bound memory accesses can occur and how to mitigate them by providing valid addresses or zero-filled data.

```APIDOC
Warp-Level Matrix Instructions (ldmatrix):

  Variants:
    .num = .x1
    .num = .x2
    .num = .x4

  Usage Example (.num = .x4):
    Each thread provides an address pointing to 8 elements in shared memory.
    After the instruction, each thread holds 8 elements in registers, fed to MMA.

  Handling Non-Divisible Sizes:
    - Examine if the start address for 8 elements is out of bounds.
    - If out of bounds, feed a special, valid address guaranteed to contain zero.
    - Ensure elements in the 'red' area (affecting final result) are zero.

  Limitations:
    - `ldmatrix` requires all threads in a warp to execute the same instruction.
    - Using `copy_if` with conditional execution is unsafe due to undefined behavior if conditions differ across threads in a warp.
    - Padding dimensions to multiples of 16k (where k is an integer) might be required, potentially wasting memory and performance.
```

--------------------------------

### Triton Transpiler Debug Mode Configuration

Source: https://mirage-project.readthedocs.io/en/latest/_sources/triton-transpiler.rst

Illustrates how to enable debug mode for the Triton transpiler. By setting `debug_mode=True` in the `profile_and_select_best_graph` function, users can obtain more detailed output for each generated code and inspect generated files, aiding in the discovery of potential issues in transpiled Triton code.

```python
elif backend == "triton":
    return profile_and_select_best_graph(all_graphs, target_cc=torch.cuda.get_device_properties(0).major * 10 + torch.cuda.get_device_properties(0).minor, warmup_iters=16, profile_iters=1000, debug_mode=True)
```
Tokens:
Raw