### Quick Start: Build and Preview MagiAttention Docs

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Follow these steps to install dependencies, build the HTML documentation, and open it in your browser. Run these commands from the 'docs' directory.

```bash
# 1. Enter the docs directory
cd docs

# 2. Install dependencies (one-time)
pip install -r requirements.txt

# 3. Build the docs
make html

# 4. Open in browser
open build/html/index.html        # macOS
xdg-open build/html/index.html    # Linux
start build/html/index.html       # Windows (Git Bash)
```

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

### New User Guide Page Structure

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Example of a basic Markdown file for a new user guide page. Includes title, sections, and inline code formatting.

```markdown
# My New Page Title

## Section One

Some content here. You can use **bold**, *italic*, `inline code`, etc.

### Subsection

More content...

## Section Two

Another section...
```

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

### Install MagiAttention Extensions from Source

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Clones the MagiAttention repository and installs the extensions package from source. Ensure you are in the extensions directory.

```bash
git clone https://github.com/SandAI-org/MagiAttention.git

cd MagiAttention/extensions

pip install --no-build-isolation .
```

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

### Install and Set Up pre-commit Hooks

Source: https://github.com/sandai-org/magiattention/blob/main/CONTRIBUTING.md

Installs the pre-commit tool and sets up the Git hooks to automatically check code style before commits.

```bash
# Install pre-commit
pip install pre-commit

# Set up the hooks, which may take some time
# Luckily, this step only needs to be performed once
pre-commit install
```

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

### Install MagiAttention from Source (Blackwell)

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/install.md

Install MagiAttention from source for Blackwell architecture. It requires setting MAGI_ATTENTION_PREBUILD_FFA to 0 and MAGI_ATTENTION_FA4_BACKEND to 1.

```bash
export MAGI_ATTENTION_PREBUILD_FFA=0 # for now, native ffa does not support Blackwell
pip install --no-build-isolation .

export MAGI_ATTENTION_FA4_BACKEND=1 # always set it when using MagiAttention on Blackwell
```

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

### Install Required Packages

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Installs packages from a requirements.txt file. Note that some packages like flash_attn_3 and magi_attention may require specific installation steps.

```bash
# NOTE: some required packages might need more tailored installation
# such as flash_attn_3 and magi_attention
pip install -r requirements.txt
```

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

### Run FFA Benchmark Script Example

Source: https://github.com/sandai-org/magiattention/blob/main/exps/attn/profile_ffa/README.md

Example of running the profile_ffa.sh script to compare performance between 'main' and 'optimize_ffa' branches, redirecting output to 'output.txt'.

```shell
bash profile_ffa.sh main profile_ffa >& output.txt
```

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

### Install Required Packages

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/install.md

Use this command to install all necessary Python packages listed in the requirements.txt file.

```bash
pip install -r requirements.txt
```

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

### Install Development Requirements

Source: https://github.com/sandai-org/magiattention/blob/main/CONTRIBUTING.md

Installs the necessary Python packages for development, including those for testing and building.

```bash
pip install -r requirements.txt
pip install -r requirements_dev.txt
```

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

### Install Transformers and Accelerate Libraries

Source: https://github.com/sandai-org/magiattention/blob/main/examples/transformers/README.md

Install the necessary libraries for training with Hugging Face Transformers and Accelerate. Ensure specific versions are used for compatibility.

```shell
pip install transformers==4.51.3
pip install accelerate==1.6.0
pip install datasets==3.5.1
pip install tiktoken==0.9.0
pip install blobfile
pip install evaluate
```

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

### Install MagiAttention from Source (Ampere)

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/install.md

Install MagiAttention from source for Ampere architecture. It requires specific environment variables to be set, including disabling prebuild FFA and skipping MagiAttention communication build.

```bash
export MAGI_ATTENTION_PREBUILD_FFA=0 # for now, native ffa does not support Blackwell
export MAGI_ATTENTION_SKIP_MAGI_ATTN_COMM_BUILD=1 # for now, magi_attn_comm does not support Ampere
pip install --no-build-isolation .

export MAGI_ATTENTION_FA4_BACKEND=1 # always set it when using MagiAttention on Ampere
```

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

### Install flash_attn_cute for Ampere/Blackwell

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/install.md

Install the flash_attn_cute package to enable the FFA_FA backend on Ampere or Blackwell GPUs. Specify the target architectures (e.g., 'sm80' for Ampere, 'sm100' for Blackwell).

```bash
ARCHS="sm80,sm100" # if you only want to use in Blackwell, you can set it to "sm100" to speed up the installation; and if you only want to use in Ampere, you can set it to "sm80" accordingly

bash scripts/install_flash_attn_cute.sh $ARCHS
```

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

### Install Megatron-Core for Hybrid DCP

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Install the necessary Megatron-Core version from the dev branch for Hybrid DCP benchmarks. Uninstall existing versions if conflicts are expected.

```bash
pip install Megatron==0.5.1
pip install -e git+https://github.com/NVIDIA/Megatron-LM.git@dev#egg=megatron-core
```

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

### Install clang-format for C++ Code Formatting

Source: https://github.com/sandai-org/magiattention/blob/main/CONTRIBUTING.md

Installs clang-format, which is used for formatting C++ source code within the project.

```bash
# Install clang-format with llvm for csrc code format
bash ./scripts/install_clang_format.sh
```

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

### Install MagiAttention from Source (Hopper)

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/install.md

Install MagiAttention from source using pip. This command is suitable for Hopper architecture and assumes no special environment variables are needed for the build.

```bash
pip install --no-build-isolation .
```

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

### Distributed MagiAttention Workflow

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/quickstart.md

Illustrates a full distributed workflow using MagiAttention, including environment setup, distributed attention calculation, and handling gradient reduction for attention sink. This example is intended to be run with `torchrun`.

```python
# run this python script with the command like:
# torchrun --standalone --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 ${SCRIPT_PATH}
import torch
import torch.nn as nn
import torch.distributed as dist

import magi_attention
from magi_attention.api import (
    magi_attn_flex_key, dispatch, calc_attn, undispatch, roll, # interface functions
    compute_pad_size, # helper functions
)
from magi_attention.common import AttnRanges
from magi_attention.common.enum import AttnMaskType
from magi_attention.utils import setup_dist_env, clearup_dist_env

# --- Set up distributed environment ---

rank, local_rank, world_size, num_nodes, num_local_ranks, world_group, device, seed = setup_dist_env()

# --- Define attention config ---

total_seqlen = 32 * 1024   # 32k tokens, if we dispatch it to 8 GPUs, then each GPU holds 4k tokens
seqlen_sink = 4            # 4 sink tokens
num_heads_q = 48           # number of attention (query) heads
num_heads_kv = 8           # number of key/value heads (GQA)
head_dim = 128             # dimension of each attention head
chunk_size = 512           # chunk size to chunk the input tensor x along the seqlen dim for dispatch to control the granularity of computation load-balance.
dtype = torch.bfloat16     # attention activation / computation dtype (while the reduction dtype for partial attention outputs is always fp32 for magi_attention right now)
has_sink = True            # whether to apply attention sink

# --- Initialize token embedding tensor ---

embed_dim = 4096
x = torch.randn(total_seqlen, embed_dim, device=device, dtype=dtype, requires_grad=True)


```

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

### New Blog Post with YAML Frontmatter

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Example of a blog post Markdown file including YAML frontmatter for metadata and a BibTeX citation block.

```markdown
---
blogpost: true
date: Mar 31, 2026
author: Your Name
location: China
category: MagiAttention
tags: Tag1, Tag2, Tag3
language: English
---

# My Blog Post Title

## Introduction

Your content here...

## Citation

If you find MagiAttention useful in your research, please cite:

```bibtex
@misc{magiattention2025,
  title={MagiAttention: ...},
  author={...},
  year={2025},
  howpublished={\url{https://github.com/SandAI-org/MagiAttention/}},
}
```

## References

```{bibliography}
:filter: docname in docnames
```
```

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

### Example .po File Structure

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Illustrates the structure of a .po file used for translations. 'msgid' contains the original English string, and 'msgstr' is where the Chinese translation should be entered.

```po
#: ../../source/user_guide/install.md:1
msgid "Installation"
msgstr "<Chinese translation>"

#: ../../source/user_guide/install.md:14
msgid "Activate an NGC-PyTorch Container"
msgstr "<Chinese translation>"
```

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

### Test Method with Dual Backend Parameter

Source: https://github.com/sandai-org/magiattention/blob/main/tests/test_common/README.md

Example of a test method designed to run against both Python and C++ backends by accepting the 'backend' fixture parameter.

```python
class TestMyFeature:
    def test_something(self, backend):
        from magi_attention.common import AttnRange
        r = AttnRange(0, 10)
        assert r.seqlen == 10
```

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

### Run Normal Attention Benchmark

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Execute the normal attention benchmark script. Note the special installation instructions for transformer-engine's flash-attention-3.

```bash
cd exps/attn

# NOTE: since transformer-engine has its own customized way to install flash-attention-3
# which has conflict with the original repo
# you need to get into the package directory and manually copy the `flash_attn_interface.py` file into `flash_attn_3/` subdirectory
bash run_benchmark.sh
```

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

### Run Unit Tests with PyTest

Source: https://github.com/sandai-org/magiattention/blob/main/CONTRIBUTING.md

Executes the project's unit tests using the PyTest framework. Ensure PyTest is installed first.

```bash
pytest tests/
```

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

### Run Multiple Test Cases with Environment Variable

Source: https://github.com/sandai-org/magiattention/blob/main/tests/README.md

Filter test cases to include multiple values for a dimension. This example runs tests for world_size 2 and 4.

```bash
MAGI_ATTENTION_TEST_WORLD_SIZE=2,4 pytest tests/test_pipeline.py
```

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

### Build and Preview Documentation

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Commands to clean, build, and preview the generated HTML documentation locally.

```bash
make clean && make html
open build/html/index.html
```

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

### Initialize Distributed Environment and Build Model

Source: https://github.com/sandai-org/magiattention/blob/main/examples/torch_native/README.md

Sets up the distributed environment, builds the device mesh, and initializes the Llama model for training.

```python
# ---   initialize distributed env   --- #
init_env(backend="nccl")

# ---   build device mesh  --- #
device_mesh = build_mesh()

# ---   set seed   --- #
torch.manual_seed(SEED)

# ---   build llama model  --- #
model = build_llama_model()   # build model from modeling llama.

# --   apply parallisim(fsdp + magi_attention)   --- #
parallize_model(model, device_mesh)

# ---   build optimizer and lr_scheduler   --- #
optimizer, lr_scheduler = build_optimizer(model, train_config["optimizer_config"])

# ---   main training loop   --- #
train(model, optimizer, lr_scheduler, device_mesh, train_config["train_iters"])
```

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

### Local HTTP Server for Preview

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Serves the built documentation over HTTP to simulate the production environment. This is necessary for the language switcher to function correctly, as it relies on path manipulation.

```bash
# Create a mirror of the production path layout
mkdir -p build/MagiAttention/docs/main
cp -r build/html/. build/MagiAttention/docs/main/

# Serve over HTTP from the build/ root
python3 -m http.server 8080 --directory build/
```

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

### Run Device All-to-All-V PoC Test

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Navigate to the device all-to-all-v experiments directory and execute the main run script.

```bash
cd exps/device_a2av

bash run.sh
```

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

### Get Most Recent Key

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/magi_api.md

Utility function to retrieve the most recent key.

```APIDOC
## Get Most Recent Key

### Description
Utility function to get the most recent key.

### Interface
`get_most_recent_key`
```

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

### Create Deploy Branch

Source: https://github.com/sandai-org/magiattention/blob/main/docs/deploy_docs.md

Create a new branch for deploying the documentation version from your source branch.

```bash
git checkout -b deploy_v1.2.0
```

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

### Get Position Ids

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/magi_api.md

Utility function to retrieve position IDs.

```APIDOC
## Get Position Ids

### Description
Utility function to get position IDs.

### Interface
`get_position_ids`
```

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

### Find PyTorch Dependency

Source: https://github.com/sandai-org/magiattention/blob/main/magi_attention/csrc/extensions/CMakeLists.txt

Finds the PyTorch library required for the extension. If Torch is not found, set CMAKE_PREFIX_PATH to your torch installation.

```cmake
# Find PyTorch
# Tip: If Torch is not found, set CMAKE_PREFIX_PATH to your torch installation.
# e.g., run: python -c "import torch; print(torch.utils.cmake_prefix_path)"
find_package(Torch REQUIRED)
```

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

### Workflow for Adding a New Page

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Steps to follow when a new documentation page is added, including updating translation files and rebuilding both language versions of the site.

```bash
make update-po                    # generates a new .po file for the new page
# ... translate the new .po file ...
DOCS_LANGUAGE=en sphinx-build -b html source/ build/html/
DOCS_LANGUAGE=zh_CN sphinx-build -b html source/ build/html/zh_CN/
```

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

### Build Ninja Manifest for Compilation

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/blog/jit_compile.md

Illustrates the preparation of a Ninja build manifest (`build.ninja`) with aggressive compilation flags for CUDA kernels, including optimization and architecture-specific options.

```text
# Generating build.ninja: It prepares Ninja build manifests containing aggressive CUTLASS and performance flags (-O3, -use_fast_math, -DCUTLASS_ENABLE_GDC_FOR_SM90, etc.).
```

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

### Run Single Test Case with Environment Variable

Source: https://github.com/sandai-org/magiattention/blob/main/tests/README.md

Use environment variables to filter test cases. This example runs tests only for world_size=2.

```bash
MAGI_ATTENTION_TEST_WORLD_SIZE=2 pytest tests/test_pipeline.py
```

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

### Flexible Dispatch Key

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/magi_api.md

Interface to get a key for flexible dispatch, supporting diverse mask types including sliding window masks.

```APIDOC
## Flexible Dispatch Key

### Description
Generates a key for flexible dispatch, supporting a wider range of mask types beyond varlen masks, such as sliding window masks.

### Interface
`magi_attn_flex_key`
```

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

### Build Multi-language Documentation

Source: https://github.com/sandai-org/magiattention/blob/main/docs/locale/README.md

Execute this command after editing translation files to build documentation in multiple languages.

```bash
make html-multilang
```

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

### Basic Usage of fa2_varlen_kvpacked_func_with_sink

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Demonstrates how to use the fa2_varlen_kvpacked_func_with_sink function for attention with variable sequence lengths and sink tokens. Ensure correct tensor shapes, data types, and device placement.

```python
import torch
from magi_attn_extensions import fa2_varlen_kvpacked_func_with_sink

sq, sk, s_sink = 2048, 2048, 2
nhq, nhk, hd = 8, 4, 128
dtype = torch.bfloat16
device = torch.cuda.current_device()
causal = True
sink_layout = "sh" # options: {"sh", "ssh"}

q = torch.randn((sq, nhq, hd), dtype=dtype, device=device, requires_grad=True)
kv = torch.randn((sk, 2, nhk, hd), dtype=dtype, device=device, requires_grad=True)
do = torch.randn_like(q)
match sink_layout:
    case "sh":
        sink = torch.randn((s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case "ssh":
        sink = torch.randn((sq, s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case _:
        raise ValueError(f"Invalid sink layout: {sink_layout}")

cu_seqlens_q = torch.tensor([0, sq // 2, sq], dtype=torch.int32, device=device)
cu_seqlens_k = torch.tensor([0, sk // 2, sk], dtype=torch.int32, device=device)
max_seqlen_q = sq // 2
max_seqlen_k = sk // 2


out = fa2_varlen_kvpacked_func_with_sink(
    q=q,
    kv=kv,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=max_seqlen_q,
    max_seqlen_k=max_seqlen_k,
    sink=sink,
    sink_layout=sink_layout,
    causal=causal,
    return_attn_probs=False,
)

out.backward(do)

dq, dkv, dsink = q.grad, kv.grad, sink.grad

```

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

### Instantiate and Use MagiTrainer in run_magi_clm.py

Source: https://github.com/sandai-org/magiattention/blob/main/examples/transformers/README.md

Import and use the custom `MagiTrainer` when initializing the trainer in your main training script.

```python
+ from magi_trainer import MagiTrainer

+ trainer = MagiTrainer(...)
- trainer = Trainer(...)

...

trainer.train()
```

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

### Registering a New User Guide Page in TOC

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

How to add a new Markdown file to the table of contents in `source/user_guide/toc.md`. This involves adding the filename to the toctree directive.

```markdown
# User Guide

```{toctree}
:caption: User Guide
:maxdepth: 2

install.md
quickstart.md
magi_api.md
env_variables.md
my_new_page.md          # <-- add your new file here
```
```

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

### Configure and Run Distributed Training with Torchrun

Source: https://github.com/sandai-org/magiattention/blob/main/examples/torch_native/README.md

Sets up environment variables for distributed training and defines the command to launch the main training script using torchrun. This is the standard way to initiate distributed PyTorch training jobs.

```shell
export GPUS_PER_NODE=${GPUS_PER_NODE:-8} # just set this to 1 to enable the non-dist training
export NNODES=${WORLD_SIZE:-1}
export NODE_RANK=${RANK:-0}
export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
export MASTER_PORT=${MASTER_PORT:-16989}

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

TORCHRUN_CMD="torchrun $DISTRIBUTED_ARGS main.py"
$TORCHRUN_CMD
```

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

### Scale Loss in training_step for Backward Pass

Source: https://github.com/sandai-org/magiattention/blob/main/examples/transformers/README.md

Modify `training_step` to scale the loss by `CP_SIZE` before the backward pass. This corrects for averaging by DP/FSDP in distributed training setups.

```python
def training_step():
    ...
    cp_size = int(os.environ.get("CP_SIZE", 1))
    backward_loss = loss * cp_size
    self.accelerator.backward(backward_loss, **kwargs)
-   self.accelerator.backward(loss, **kwargs)

    return loss.detach()
```

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

### Basic Usage for fa2_varlen_qkvpacked_func_with_sink

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Demonstrates the basic usage of fa2_varlen_qkvpacked_func_with_sink for attention with variable sequence lengths, packed QKV inputs, and a sink token. This function is suitable for batched inputs with varying sequence lengths.

```python
import torch
from magi_attn_extensions import fa2_varlen_qkvpacked_func_with_sink

s, s_sink = 2048, 2
n h, hd = 8, 128
dtype = torch.bfloat16
device = torch.cuda.current_device()
causal = True
sink_layout = "sh" # options: {"sh", "ssh"}

qkv = torch.randn((s, 3, nh, hd), dtype=dtype, device=device, requires_grad=True)
do = torch.randn((s, nh, hd), dtype=dtype, device=device, requires_grad=True)
match sink_layout:
    case "sh":
        sink = torch.randn((s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case "ssh":
        sink = torch.randn((s, s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case _:
        raise ValueError(f"Invalid sink layout: {sink_layout}")

cu_seqlens = torch.tensor([0, s // 2, s], dtype=torch.int32, device=device)
max_seqlen = s // 2

out = fa2_varlen_qkvpacked_func_with_sink(
    qkv=qkv,
    cu_seqlens=cu_seqlens,
    max_seqlen=max_seqlen,
    sink=sink,
    sink_layout=sink_layout,
    causal=causal,
    return_attn_probs=False,
)

out.backward(do)

dqkv, dsink = qkv.grad, sink.grad
```

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

### Basic Usage of fa4_varlen_func_with_sink

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Demonstrates the basic usage of the `fa4_varlen_func_with_sink` function for variable-length inputs. This includes setting up sequence lengths and maximum sequence lengths.

```python
import torch
from magi_attn_extensions import fa4_varlen_func_with_sink

sq, sk, s_sink = 2048, 2048, 2
nhq, nhk, hd = 8, 4, 128
dtype = torch.bfloat16
device = torch.cuda.current_device()
causal = True
sink_layout = "sh"  # options: {"sh", "ssh"}

q = torch.randn((sq, nhq, hd), dtype=dtype, device=device, requires_grad=True)
k = torch.randn((sk, nhk, hd), dtype=dtype, device=device, requires_grad=True)
v = torch.randn((sk, nhk, hd), dtype=dtype, device=device, requires_grad=True)
do = torch.randn_like(q)
match sink_layout:
    case "sh":
        sink = torch.randn((s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case "ssh":
        sink = torch.randn((sq, s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case _:
        raise ValueError(f"Invalid sink layout: {sink_layout}")

cu_seqlens_q = torch.tensor([0, sq // 2, sq], dtype=torch.int32, device=device)
cu_seqlens_k = torch.tensor([0, sk // 2, sk], dtype=torch.int32, device=device)
max_seqlen_q = sq // 2
max_seqlen_k = sk // 2

out, lse = fa4_varlen_func_with_sink(
    q=q,
    k=k,
    v=v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=max_seqlen_q,
    max_seqlen_k=max_seqlen_k,
    sink=sink,
    sink_layout=sink_layout,
    causal=causal,
    return_attn_probs=True,
)

out.backward(do)

dq, dk, dv, dsink = q.grad, k.grad, v.grad, sink.grad
```

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

### FlashAttention 2: Basic Usage with Sink

Source: https://github.com/sandai-org/magiattention/blob/main/extensions/README.md

Demonstrates the basic usage of `fa2_func_with_sink` for standard attention with sink tokens in FlashAttention 2. Note that `return_attn_probs` is set to `False` in this example.

```python
import torch
from magi_attn_extensions import fa2_func_with_sink

b = 2
sq, sk, s_sink = 2048, 2048, 2
nhq, nhk, hd = 8, 4, 128
dtype = torch.bfloat16
device = torch.cuda.current_device()
causal = True
sink_layout = "sh" # options: {"sh", "ssh"}

q = torch.randn((b, sq, nhq, hd), dtype=dtype, device=device, requires_grad=True)
k = torch.randn((b, sk, nhk, hd), dtype=dtype, device=device, requires_grad=True)
v = torch.randn((b, sk, nhk, hd), dtype=dtype, device=device, requires_grad=True)
do = torch.randn_like(q)
match sink_layout:
    case "sh":
        sink = torch.randn((s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case "ssh":
        sink = torch.randn((b, sq, s_sink, nhq), dtype=torch.float32, device=device, requires_grad=True)
    case _:
        raise ValueError(f"Invalid sink layout: {sink_layout}")

out = fa2_func_with_sink(
    q=q,
    k=k,
    v=v,
    sink=sink,
    sink_layout=sink_layout,
    causal=causal,
    return_attn_probs=False,
)
out.backward(do)

dq, dk, dv, dsink = q.grad, k.grad, v.grad, sink.grad
```

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

### Run Distributed Attention Benchmark with Profiling

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Execute the benchmark with profiling enabled. Specify an output name for the profile results or use the default. Ensure BENCH_MODE.enable_profile is set in the config file.

```bash
cd exps/dist_attn

# 1. Enable profiling with default output name (cp_benchmark)
bash run_benchmark.sh --profile

# 2. Enable profiling and specify an output name
bash run_benchmark.sh --profile output_name

# 3. Equivalent syntax using '='
bash run_benchmark.sh --profile=output_name

# 4. Disable profiling by default
bash run_benchmark.sh
```

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

### Varlen Dispatch Key

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/magi_api.md

Interface to get a key for varlen dispatch, suitable for masks defined by cu_seqlens like varlen full or varlen causal masks. Inspired by FlashAttention's API.

```APIDOC
## Varlen Dispatch Key

### Description
Generates a key for varlen dispatch, compatible with masks defined by `cu_seqlens` (e.g., varlen full, varlen causal).

### Interface
`magi_attn_varlen_key`
```

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

### Run All Tests

Source: https://github.com/sandai-org/magiattention/blob/main/tests/test_common/README.md

Execute all unit tests for both Python and C++ backends.

```bash
pytest tests/test_common/ -v
```

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

### Configure Small Q Block Sizes for Sparse Attention

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Example Python code snippet to configure small Q block sizes for uniform block sparse attention. This is suitable for NSA in GQA scenarios.

```python
q_block_sizes = [64, 32, 16, 8]
k_block_sizes = [64, 64, 64, 64]
```

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

### Configure Small K Block Sizes for Sparse Attention

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Example Python code snippet to configure small K block sizes for uniform block sparse attention. This is useful for token-level sparse attention.

```python
q_block_sizes = [64, 64, 64, 64, 64]
k_block_sizes = [64, 32, 16, 8, 1]
```

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

### Sphinx Build Command for English

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Builds the HTML documentation for the English version. The output is placed at the root of the build directory, mirroring production URL structures.

```bash
DOCS_LANGUAGE=en sphinx-build -b html source/ build/html/
```

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

### Copy JIT Compiled Kernels for AOT Distribution

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/blog/jit_compile.md

This code demonstrates how compiled JIT kernels are copied to a designated AOT (Ahead-Of-Time) directory during the project installation process. It ensures that commonly used configurations are available as pre-compiled binaries.

```python
compute_dtype, output_dtype = compute_output_dtype_tuple
spec, uri = get_ffa_jit_spec( ... )  # Employs the identical spec builder!
spec.build()
src_dir = (jit_env.MAGI_ATTENTION_JIT_DIR / uri).resolve()
dst_dir = (jit_env.MAGI_ATTENTION_AOT_DIR / uri).resolve()
shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True)
```

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

### Implement _prepare_magi_attention Helper

Source: https://github.com/sandai-org/magiattention/blob/main/examples/transformers/README.md

This helper function dispatches input data and prepares the runtime key for MagiAttention, handling tensor reshaping and configuration.

```python
# dispatch data and prepare key
def _prepare_magi_attention(
    self,
    inputs: torch.Tensor,
    cu_seqlens_q: torch.Tensor,
    cu_seqlens_k: torch.Tensor,
    num_heads_q: int,
    num_heads_kv: int,
    head_dim: int,
    pad_size: int,
):
    # ---   magi_attn_varlen_dispatch   --- #

    dist_attn_config = DistAttnConfig()
    cp_group = self._build_cp_group()

    inputs = squash_batch_dim(inputs)

    dist_attn_runtime_key = magi_attn_varlen_key(
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        num_heads_q=num_heads_q,
        num_heads_kv=num_heads_kv,
        head_dim=head_dim,
        chunk_size=512,
        pad_size=pad_size,
        cp_group_or_mesh=cp_group,
        causal=True,
        dist_attn_config=dist_attn_config,
    )

    x_padded = dispatch(inputs, key=dist_attn_runtime_key)
    x_padded = x_padded.unsqueeze(0)

    return x_padded, dist_attn_runtime_key
```

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

### Configure Large Q and K Blocks for Sparse Attention

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Example Python code snippet to configure large Q and K block sizes for uniform block sparse attention. This is used for traditional block sparse patterns.

```python
q_block_sizes = [64, 128]
k_block_sizes = [64, 128]
```

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

### Initialize MagiAttention Meta Configs

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/quickstart.md

Sets up customized attention mask configurations (block-causal) using AttnRanges and computes padding size for distributed embeddings.

```python
# --- Initialize MagiAttention meta configs for customized attention mask ---
# the following customized attention mask is known as `block-causal` mask where `block_size` = 4096 (4k),
# which looks like (`*` for unmasked, `0` for masked):
#     - - - - - - - - -> (k)
#   | * * 0 0 0 0 0 0
#   | * * 0 0 0 0 0 0
#   | * * * * 0 0 0 0
#   | * * * * 0 0 0 0
#   | * * * * * * 0 0
#   | * * * * * * 0 0
#   | * * * * * * * *
#   | * * * * * * * *
#   V
#  (q)
q_ranges = AttnRanges.from_ranges(
    [
        [0, 4096], # 0~4k
        [4096, 8192], # 4k~8k
        [8192, 12288], # 8k~12k
        [12288, 16384], # 12k~16k
        [16384, 20480], # 16k~20k
        [20480, 24576], # 20k~24k
        [24576, 28672], # 24k~28k
        [28672, 32768], # 28k~32k
    ]
)
k_ranges = AttnRanges.from_ranges(
    [
        [0, 4096], # 0~4k
        [0, 8192], # 0~8k
        [0, 12288], # 0~12k
        [0, 16384], # 0~16k
        [0, 20480], # 0~20k
        [0, 24576], # 0~24k
        [0, 28672], # 0~28k
        [0, 32768], # 0~32k
    ]
)
attn_mask_type = [AttnMaskType.FULL] * len(q_ranges)
total_seqlen_q = total_seqlen_k = total_seqlen
pad_size = compute_pad_size( # pad embeds along seqlen dim for better performance
    total_seqlen_q=total_seqlen_q,
    cp_size=world_size, # assuming we only have 1-dim context parallelism (cp)
    chunk_size=chunk_size,
)
```

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

### Calculate KV Interval Boundaries for Mask Types

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/blog/blackwell_ffa_fa4.md

This C++ code calculates the start and end boundaries for Key-Value (KV) intervals based on the query token index and mask type. It is used within the `magi_to_hstu` CUDA kernel for flexible masking patterns.

```cpp
   int k_interval_start = k_start + ((mask_type & 2) ? (q_idx - q_start) : 0);
   int k_interval_end   = k_end   - ((mask_type & 1) ? (q_end - q_idx - 1) : 0);
```

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

### Force Rebuild Documentation

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Use this command to clear the build cache and perform a full rebuild if you encounter issues with stale content.

```bash
make clean && make html
```

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

### Single-GPU Flex-Flash-Attention

Source: https://github.com/sandai-org/magiattention/blob/main/docs/source/user_guide/quickstart.md

Demonstrates the minimal workflow for calling `flex_flash_attn_func` on a single GPU, including configuration, tensor initialization, and forward/backward passes.

```python
import torch
from magi_attention.api import flex_flash_attn_func

# --- Define attention config ---

total_seqlen = 2048    # 2k tokens
seqlen_sink = 4        # 4 sink tokens
num_heads_q = 8        # number of attention (query) heads
num_heads_kv = 2       # number of key/value heads (GQA)
head_dim = 128         # dimension of each attention head
dtype = torch.bfloat16 # attention activation / computation dtype (while the reduction dtype is always fp32 for ffa right now)
device = "cuda"
has_sink = True        # whether to apply attention sink

# --- Initialize q,k,v,do tensors ---

q = torch.randn(total_seqlen, num_heads_q, head_dim, dtype=dtype, device=device, requires_grad=True)
k = torch.randn(total_seqlen, num_heads_kv, head_dim, dtype=dtype, device=device, requires_grad=True)
v = torch.randn(total_seqlen, num_heads_kv, head_dim, dtype=dtype, device=device, requires_grad=True)
do = torch.randn_like(q)

# --- Initialize optional sink tensor ---

sink = torch.randn(seqlen_sink, num_heads_q, dtype=torch.float32, device=device, requires_grad=True) if has_sink else None

# --- Initialize FFA meta args for customized attention mask ---

# the following customized attention mask looks like (`*` for unmasked, `0` for masked):
#     - - - - - - - - -> (k)
#   | * * * * 0 0 0 0
#   | * * * * 0 0 0 0
#   | * * * * 0 0 0 0
#   | * * * * 0 0 0 0
#   | * * * * * 0 0 0
#   | * * * * * * 0 0
#   | * * * * * * * 0
#   | * * * * * * * *
#   V
#  (q)
q_ranges_tensor = torch.tensor([[0, 1024], [1024, 2048]], dtype=torch.int32, device=device)
k_ranges_tensor = torch.tensor([[0, 1024], [0, 2048]], dtype=torch.int32, device=device)
attn_type_map_tensor = torch.tensor([0, 1], dtype=torch.int32, device=device) # full mask for 1st slice, causal mask for 2nd

# --- Attention computation ---

out, meta = flex_flash_attn_func(
    q=q,
    k=k,
    v=v,
    q_ranges=q_ranges_tensor,
    k_ranges=k_ranges_tensor,
    attn_type_map=attn_type_map_tensor,
    sink=sink, # Defaults to None to not apply attention sink
    softmax_scale=None, # Defaults to 1/sqrt(head_dim)
    softcap=0, # Defaults to 0
)
lse = meta.lse

out.backward(do)

dq, dk, dv = q.grad, k.grad, v.grad
dsink = sink.grad if has_sink else None

```

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

### Sphinx Build Command for Chinese

Source: https://github.com/sandai-org/magiattention/blob/main/docs/README.md

Builds the HTML documentation for the Chinese version. The output is placed under a language-specific subdirectory (e.g., build/html/zh_CN/).

```bash
DOCS_LANGUAGE=zh_CN sphinx-build -b html source/ build/html/zh_CN/
```

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

### Draw Benchmark Results from CSV

Source: https://github.com/sandai-org/magiattention/blob/main/exps/README.md

Generate visualizations from existing benchmark results stored in CSV files. Navigate to the 'exps/dist_attn' directory before running.

```bash
cd exps/dist_attn

python draw_benchmark.py
```

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

### Add New Version to versions.json

Source: https://github.com/sandai-org/magiattention/blob/main/docs/deploy_docs.md

Register a new documentation version by adding an entry to the versions.json file.

```json
{
    "name": "v1.2.0",
    "version": "v1.2.0",
    "url": "https://sandai-org.github.io/MagiAttention/docs/v1.2.0/"
}
```

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

### Main Training Loop with MagiAttention Preparation

Source: https://github.com/sandai-org/magiattention/blob/main/examples/torch_native/README.md

The main training loop iterates through training steps, preparing data and conditionally preparing MagiAttention inputs if context parallelization is enabled.

```python
def train(model, optimizer, lr_scheduler, device_mesh, train_iter):
    """main training loop"""
    model.train()

    for iter in range(train_iter):
        input, label, cu_seqlens_q, cu_seqlens_k, pad_size = prepare_data(
            device_mesh, train_iter
        )

        dist_attn_runtime_key = None

        if (
            parallel_config["context_parallel_size"] > 1
            and parallel_config["context_parallel_backend"] == "magi_attention"
        ):
            # dispatched input and prepare magi_attn key.
            input, dist_attn_runtime_key = prepare_magi_attention(
                input=input,
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                num_heads_q=model.config.num_attention_heads,
                num_heads_kv=model.config.num_key_value_heads,
                head_dim=model.config.head_dim,
                pad_size=pad_size,
                cp_group=device_mesh.get_group("cp"),
            )

        output = model(input, dist_attn_runtime_key)
        loss = loss_func(output, label, device_mesh, dist_attn_runtime_key)
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
```