MagiAttention Documentation

MagiAttention is a distributed attention (Context Parallelism) solution designed for ultra-long sequences and heterogeneous masking patterns. It achieves linear scalability through several key architectural components:

• Flex-Flash-Attention (FFA): A kernel that supports distributable and flexible mask representations.
• Dispatch Solver: Manages load-balanced computation.
• Group Collective Primitives: Enables zero-redundant communication.
• Adaptive Multi-stage Overlap Strategy: Coordinates components to optimize performance.

It is optimized for large-scale training scenarios, such as video generation (e.g., Magi-1).

MagiAttention is a distributed attention solution designed for linear scalability in training with ultra-long contexts (up to ~4M tokens) and heterogeneous, irregular attention masks. It consists of five core components:

1. Flex-Flash-Attention (FFA): An optimized kernel based on Flash-Attention 3 that supports flexible mask patterns and distributable mask representations.
2. Dispatch Solver: Shards ultra-long data and dispatches it to ensure load-balanced computation across ranks.
3. Zero-Redundant Communication Primitives: Uses GroupCast and GroupReduce to eliminate unnecessary data movement.
4. Overlap Solver: An adaptive strategy that partitions multi-stage computation and communication to maximize overlap.
5. Scheduled Timelines: Manages forward and backward pass timelines for optimal execution.

The GroupCast kernel is designed to logically chunk an input buffer along the sequence length dimension into several input_splits. Each split contains its size and a list of destination peers (dst_indices).

Workflow:

1. Producer (Sender SM): Loads an input_split from global memory to shared memory via TMA. It assigns a warp to send the data to each destination peer via NVLink or RDMA.
2. Consumer (Receiver SM): Waits for its receive buffer to be filled by a unique sender. It then assigns a warp to load the data from shared memory into the corresponding output_split in the output buffer via TMA (using the src_index list).

To avoid the redundant communication found in standard Ring-style P2P implementations (where data is often broadcasted even if only specific ranks need it), MagiAttention uses Group Collective Primitives.

• GroupCast: Used in the forward pass to model low-demand KV requests. It builds a transfer table for KV send/receive buffers and uses AlltoAll-v (or a native CUDA implementation) to send data exclusively to target ranks.
• GroupReduce: Used in the backward pass to aggregate partial dKV gradients. It collects and reduces partial outputs back to the source rank using AlltoAll-v or native kernels.

These primitives achieve zero-redundant communication by ensuring data only moves between the ranks that actually require it.

To calculate the final throughput metrics from elapsed time:

• Kernel Throughput (wd ∈ {fwd, bwd}): $$\text{TFLOPs/s}^{(wd)} = \frac{\text{FLOPs}^{(wd)}}{\text{ElapsedTime}^{(wd)}}$$

• Distributed Throughput (wd ∈ {fwd, bwd}): $$\text{TFLOPs/s/GPU}^{(wd)} = \frac{\text{FLOPs}^{(wd)}}{\text{ElapsedTime}^{(wd)} \times cp_size}$$

Note on ElapsedTime: For distributed runs, ElapsedTime is defined as the maximum elapsed time across all ranks: $\max_{rank \in [0, cp_size)} \text{ElapsedTime}_{rank}^{(wd)}$.

For diverse mask types beyond varlen full/causal (such as sliding window masks), use magi_attn_flex_key.

To apply multiple flexible masks in one training pass, use make_flex_key_for_new_mask_after_dispatch. This creates a new key for a new mask based on existing mask arguments and the current dispatch key, reusing the same dispatch solution with updated meta arguments.

from magi_attention.api.magi_attn_interface import magi_attn_flex_key, make_flex_key_for_new_mask_after_dispatch

# 1. Create the initial flexible key
key = magi_attn_flex_key(...)

# 2. Create a new key for a different mask in the same pass
new_key = make_flex_key_for_new_mask_after_dispatch(key, ...)

MagiAttention is designed to integrate with mainstream deep learning training frameworks. Supported integrations include:

• Megatron-LM
• PyTorch FSDP
• HuggingFace Transformers

The FFA_FA4 backend implements several optimizations to handle flexible masking efficiently on Blackwell:

1. create_block_mask Kernel: A high-performance kernel that parses the HSTU Function directly to categorize blocks into Full (no masking), Partial (masking required), or Empty (skipped). It includes q2k (forward) and k2q (backward) implementations. An optimization in the forward pass treats blocks that are out-of-bounds only in the q direction as Full to improve throughput.
2. CSR Compression: To prevent memory scaling issues with long sequences, sparsity metadata is stored in a Compressed Sparse Row (CSR) format. This uses full_block_offset and mask_block_offset to locate valid n-block indices, storing only non-empty blocks instead of fixed-size tensors.