SandAI-org/MagiAttention

A Distributed Attention Towards Linear Scalability for Ultra-Long Context, Heterogeneous Data Training

What it solves

MagiAttention addresses the challenge of scaling attention mechanisms for ultra-long contexts and heterogeneous mask patterns in distributed training. It aims to provide linear scalability and high performance, reducing the computational and communication overhead typically associated with context-parallel (CP) training setups.

How it works

It implements a distributed attention mechanism through several key innovations:

  • Flex-Flash-Attention (FFA): A generalized attention mask formulation (AttnSlice) and a tailored kernel that allows for compact expression of diverse mask types and efficient distributed partitioning.
  • Computation Load Balancing: Uses a fine-grained chunk-level sharding strategy and a dispatch solver to ensure workloads are evenly distributed across CP ranks.
  • Zero-Redundant Communication: Replaces traditional Ring-style P2P communication with novel GroupCast and GroupReduce primitives to eliminate redundant communication volume during forward and backward passes.
  • Adaptive Multi-Stage Overlap: Schedules computation and communication to hide latency and maximize GPU utilization.

Who it’s for

It is designed for researchers and developers training large-scale models with ultra-long contexts, such as autoregressive video generation models (e.g., Magi-1), and those integrating with frameworks like Megatron-LM, PyTorch FSDP, and HuggingFace Transformers.

Highlights

  • Linear Scalability: Delivers near-linear scalability across distributed training setups on H100 and B200 GPUs.
  • Broad Hardware Support: Performance comparable to Flash-Attention 3 on Hopper GPUs, with early support for Blackwell GPUs via Flash-Attention 4.
  • Framework Integration: Easy integration with Megatron-LM, PyTorch FSDP, and HuggingFace Transformers.
  • Advanced Masking: Native support for diverse and overlapping attention mask patterns.
  • Attention Sink: Includes extensions for distributed learnable attention sink mechanisms.

Related

  • Project
    beilusaiying/always-accompany

    A multi-purpose AI agent framework focused on context and attention mechanisms, providing layered memory and dynamic retrieval for long-term companionship, coding, and work.

  • Project
    EvolvingLMMs-Lab/lmms-engine

    A unified, high-performance training engine for multimodal models that integrates distributed training, kernel fusion, and advanced optimizers to enable scalable pretraining and fine-tuning.

  • Project
    Dao-AILab/flash-attention

    A fast and memory-efficient implementation of the attention mechanism for Transformers, reducing memory IO to accelerate LLM training and inference.

  • Dispatch
    Ulysses Sequence Parallelism for Million-Token Context Training

    Hugging Face has integrated Ulysses Sequence Parallelism into Accelerate, Transformers, and TRL, enabling the training of LLMs with million-token contexts by distributing attention computation across multiple GPUs.

  • Project
    fla-org/flash-linear-attention

    A library of hardware-efficient, training-ready implementations for linear attention, state space models, and hybrid LLM architectures across multiple hardware platforms.