Source-linked AI summary

MAGI-1: Autoregressive Video Generation at Scale

Sand. ai, Hansi Teng, Hongyu Jia, Lei Sun, Lingzhi Li, Maolin Li, Mingqiu Tang, Shuai Han, Tianning Zhang, W. Q. Zhang, Weifeng Luo, Xiaoyang Kang, Yuchen Sun, Yue Cao, Yunpeng Huang, Yutong Lin, Yuxin Fang, Zewei Tao, Zheng Zhang, Zhongshu Wang, Zixun Liu, Dai Shi, Guoli Su, Hanwen Sun, Hong Pan, Jie Wang, Jiexin Sheng, Min Cui, Min Hu, Ming Yan, Shucheng Yin, Siran Zhang, Tingting Liu, Xianping Yin, Xiaoyu Yang, Xin Song, Xuan Hu, Yankai Zhang, Yuqiao Li

arXiv:2505.13211v1cs.CVcs.AI

TL;DR

Video generation must produce temporally coherent, photorealistic sequences from semantic inputs while remaining causally consistent and compatible with streaming. MAGI-1 autoregressively denoises fixed-length video chunks and outperforms existing large-scale video diffusion models in prompt adherence, physical plausibility, and temporal coherence.

  • Problem

    Generating video from semantic inputs requires temporal coherence and photorealism, while applications also call for causal consistency and streaming compatibility.

    Source: §1 · p. 1
  • Method

    MAGI-1 generates video autoregressively by progressively denoising fixed-length chunks under causal constraints.

    Source: §1 · p. 1, §7 · p. 37
  • Results

    MAGI-1 outperforms existing large-scale video diffusion models in prompt adherence, physical plausibility, and temporal coherence.

    Source: §7 · p. 37
  • Takeaways & Limitations

    MAGI-1 supports text-to-video, image-to-video, and video continuation under a shared training objective, with fine-grained semantic control across long-form generation.

    Source: §7 · p. 37

    MAGI-1 uses a tightly coupled decoder-style Transformer for both high-level temporal context fusion and low-level visual denoising.

    Source: §8 · p. 38

Abstract

from arXiv · show

We present MAGI-1, a world model that generates videos by autoregressively predicting a sequence of video chunks, defined as fixed-length segments of consecutive frames. Trained to denoise per-chunk noise that increases monotonically over time, MAGI-1 enables causal temporal modeling and naturally supports streaming generation. It achieves strong performance on image-to-video (I2V) tasks conditioned on text instructions, providing high temporal consistency and scalability, which are made possible by several algorithmic innovations and a dedicated infrastructure stack. MAGI-1 facilitates controllable generation via chunk-wise prompting and supports real-time, memory-efficient deployment by maintaining constant peak inference cost, regardless of video length. The largest variant of MAGI-1 comprises 24 billion parameters and supports context lengths of up to 4 million tokens, demonstrating the scalability and robustness of our approach. The code and models are available at https://github.com/SandAI-org/MAGI-1 and https://github.com/SandAI-org/MagiAttention. The product can be accessed at https://sand.ai.

1 Introduction

MAGI-1 generates video autoregressively in fixed-length chunks with temporally progressive noise, aiming to combine causal consistency with scalable, streaming-compatible generation.

1 Introduction

  • Unlike globally conditioned video diffusion models that use uniform noise and require full-sequence inference, MAGI-1 targets causal, streaming, real-time video generation.The paper frames this as addressing demand for high-fidelity, computationally efficient, causally consistent generation conditioned on inputs such as language, imagery, or short video clips.Source: §1 · p. 1
  • MAGI-1 predicts fixed-length frame chunks autoregressively, using temporally progressive noise to capture mid-range dependencies while maintaining strict left-to-right temporal consistency.Chunk-level training also supports adaptable conditional generation.Source: §1 · p. 1
  • A Transformer designed for bidirectional spatial and causal temporal denoising, plus MagiAttention and scalable training and inference infrastructure, enables streaming generation with peak resource use independent of video length.The infrastructure includes a data curation pipeline that adapts data distribution based on ongoing model evaluation.Source: §1 · pp. 1–2
  • The largest MAGI-1 variant has 24 billion parameters and supports context lengths up to 4 million tokens.The inference codebase and pretrained models, distributed attention, and a live demonstration are publicly accessible.Source: §1 · p. 2
  • MAGI-1 substantially improves over previous models on VBench-I2V and Physics-IQ Benchmark, especially for complex motion, semantic alignment, and physically plausible interactions.The evaluation also assesses prompt fidelity, temporal coherence, and subject integrity.Source: §1 · p. 2

2 MAGI-1

This section describes MAGI-1’s chunk-wise autoregressive video generation, including its model, training objectives, and distillation. It also covers generation controls and inference strategies for temporal coherence, longer videos, prompt refinement, and streaming.

2 MAGI-1

  • MAGI-1 denoises 24-frame chunks conditioned on preceding chunks, starting the next chunk before the current one is fully clean to enable concurrent, streaming generation.Up to four chunks can be inferred simultaneously; block-causal attention enforces temporal causality across chunks.Source: §2 · pp. 2–3, Figure 1 · p. 2

2.1 Transformer-based Variational Auto-Encoder

  • The transformer-based VAE embeds video with a 3D convolution, stacks 24 transformer blocks, and downsamples encoded features by 8× spatially and 4× temporally relative to raw video input.At inference, sliding spatial windows of 256 × 256 pixels with 192-pixel stride support arbitrary resolutions; no temporal overlap is applied.Source: §2.1 · pp. 3–4
  • Despite having the largest model size, MAGI-1’s transformer VAE achieves the fastest average decoding time, significantly outperforms most baselines in encoding speed, and ranks second in PSNR.All models were evaluated on one NVIDIA H800 GPU using 169 videos of 25 frames at 256×256 resolution.Source: Table 1 · p. 3, §2.1 · p. 4

2.2.1 Training objective

  • MAGI-1 trains with flow matching on independently noised chunks whose timesteps increase over time, predicting each chunk’s velocity conditioned on all preceding chunks.This monotonic noise ordering makes earlier chunks cleaner than later chunks.Source: §2.2.1 · pp. 4–5

2.2.2 Model Architecture

  • MAGI-1 combines block-causal attention and parallel spatial-temporal and cross-attention with QK normalization, grouped-query attention, sandwich normalization, and softcapped timestep modulation.Parallel attention reuses query projections to reduce tensor-parallel communication from two rounds to one per block; softcap constrains modulation to [−1, 1].Source: §2.2.2 · p. 6

2.2.3 Training Recipes

  • Text-to-video, image-to-video, and video continuation share one training objective by varying clean-chunk proportions, while chunk-wise prompts enable fine-grained text control.Image-to-video is treated as continuation with only the first frame of the first chunk clean; no task-specific fine-tuning is required.Source: §2.2.3 · p. 7

2.3 Distillation Using Shortcut Model

  • Shortcut distillation trains a model conditioned on timestep and step size, using bootstrapped targets so one step matches two consecutive half-steps and supports 64, 32, 16, or 8 inference steps.The smallest step size, 1/64, corresponds to the standard 64-evaluation flow-matching setting.Source: §2.3 · pp. 9–10

2.4.1 Diffusion Guidance

  • Classifier-free guidance mediates the fidelity-diversity trade-off in generative modeling and is particularly effective for text-to-video generation.Source: §2.4.1 · p. 10
  • MAGI-1 separates unconditional, temporal-context, and prompt-conditioned guidance, with wprev = 1.5 improving chunk alignment and mitigating flicker.The paper sets wtext = 7.5; increasing wprev beyond the useful range can cause saturation or static frames.Source: §2.4.1 · pp. 11–12, Figure 7 · p. 11

2.4.2 Inference Timestep Sampler

  • The inference timestep sampler applies a tunable power transformation, with w = 1/3 and k = 2 yielding the best visual quality in the reported experiments.Figure 8 visualizes timestep sampling for the non-distilled model.Source: §2.4.2 · p. 12, Figure 8 · p. 12

2.4.3 Fine-Grained Control of Guidance Strength

  • For non-distilled models, reducing guidance after t = 0.3 mitigates artifacts in videos longer than 5 seconds; distilled models attenuate previous-chunk guidance in the final stage.The non-distilled schedule uses wprev = 1.0 and wtext = 0.0 when t > 0.3; the distilled model uses wprev = 0.7 in its final denoising stage.Source: §2.4.3 · pp. 12–13, Figure 9 · p. 13

2.4.4 KV Cache

  • KV caching reuses denoised chunk features, and restricting the KV range to preceding content supports long-video generation with computation scaling linearly with duration.The model is theoretically compatible with KV-compression methods, but their exploration is left for future work.Source: §2.4.4 · pp. 13–14

2.5 Prompt-Enhancement Strategy

  • Prompt enhancement uses an MLLM to describe the input image and predict scene evolution, then distills enhanced prompts into a smaller model for lower-latency deployment.The distilled model (~7B) is trained on approximately 2 million examples and achieves comparable video-generation quality to the larger model in human evaluation.Source: §2.5 · p. 14

2.6 Model Capability Study

  • Chunk-wise pipelining enables real-time streaming, while stage-specific KV ranges enable shot transitions that preserve either object identity or scene layout.Limiting each chunk to the preceding 8 seconds keeps computation linear with video duration; changing KV range at high- versus low-noise stages controls which visual properties persist.Source: §2.4.4 · p. 13, §2.6 · pp. 14, 16, Figure 12 · p. 17
  • MAGI-1 natively continues video from its full prefix, preserving motion continuity where image-conditioned continuation loses historical motion context.The examples cover pen rotation and post-occlusion behavior. Chunk-wise prompts also support complex action sequences and narrative structures in near-30-second videos.Source: §2.6 · pp. 14, 16, Figure 10 · p. 15, Figure 11 · p. 16

3 DATA

MAGI-1’s data pipeline turns tens of petabytes of videos and images into curated, captioned training data, then adjusts data quality, duration, and distribution across training stages.

3 DATA

  • The pipeline processes tens of petabytes of raw videos and images through shot segmentation, quality filtering, de-duplication, MLLM filtering, and captioning.PySceneDetect segments videos into single-shot clips; the overview notes that shot cutting applies only to video data.Source: §3 · p. 17, Figure 13 · p. 17

3.1 Filter Actors

  • For video quality, the pipeline uses DOVER’s technical score, brightness checks, and optical-flow measures to select clips with moderate motion and stable camera movement.Motion is measured overall and separately for foreground and background; clips outside thresholds for the three statistics are excluded.Source: §3.1 · p. 18
  • Additional filter actors remove clips with excessive text, subtitles, logos, corner-face narrators, slide movements, borders, or complex transitions.These actors use text size and placement, Florence-2 logo detection, face location and confidence, optical-flow divergence, edge and Hough detection, and CLIP similarity between keyframes.Source: §3.1 · pp. 18–19

3.2 De-duplication

3.3 MLLM as Advanced Filter

  • After initial filtering and de-duplication, an MLLM performs another filtering pass for complex low-quality cases and is integrated with captioning to reduce cost.The extra pass follows a substantial reduction in dataset size from the preceding steps.Source: §3.3 · p. 19

3.4 Caption

  • Video captions use 4–12 keyframes and a two-stage prompt: structured questions about predefined attributes followed by a descriptive caption; images receive direct caption prompts.The number of frames varies with clip duration to balance descriptive accuracy and computational efficiency.Source: §3.4 · pp. 19–20
  • Second-by-second autoregressive captions describe the first second in detail and subsequent seconds as changes from the previous second, enabling distinct text prompts for video parts.This captioning design supports the model’s conditioning of different video parts on different prompts.Source: §3.4 · p. 20

3.5 Data Adjustment in Training

  • Across three training stages, resolution rises from 256p to 480p to 720p while data volume falls and filtering becomes stricter; the final stage extends video duration to 16 seconds.The earlier stages cap duration at 8 seconds; the final-stage extension is intended to capture richer temporal dynamics. Data proportions also adapt during training to strengthen underperforming subsets identified by evaluation.Source: §3.5 · p. 20

4 Infrastructure

This section describes MAGI-1’s purpose-built infrastructure for efficient distributed training and real-time or cost-efficient inference that addresses video-specific demands. It covers streaming-latency optimization, memory-efficient inference, and hybrid parallelism for deployment on RTX 4090 GPUs.

4.1 Training Infrastructure

  • Training infrastructure combines data, context, and tensor parallelism with online batching to handle variable-length video sequences and improve GPU utilization.The system also introduces MagiAttention for flexible attention masks and ultra-long sequences, while identifying testability and infrastructure–algorithm coupling as broader framework challenges.Source: §4.1 · p. 22

4.1.1 Distributed Packing and Padding

  • 99% capacity utilization is achieved by online Packing and Padding with negligible differences between data-parallel groups.The method extends First-Fit Decreasing heuristics to pack a large candidate pool into bins while satisfying DP, TP, and CP divisibility constraints.Source: §4.1.1 · p. 23

4.1.2 MagiAttention: Towards Linear Scalability for Ultra-Long and Heterogeneous Mask Training.

  • MagiAttention combines balanced workload dispatch, zero-redundant communication primitives, and adaptive multi-stage overlap to target linear scalability across context-parallel sizes.Its group-cast and group-reduce primitives eliminate redundant communication, though the all-to-all-v prototype adds preprocessing overhead.Source: §4.1.2 · pp. 23, 25, 26, Figure 14 · p. 24
  • Flex-Flash-Attention supports diverse distributed attention masks, including irregular block-causal patterns, while achieving comparable MFU to FA3.It decomposes masks into AttnSlices and uses a tailored kernel for distributed mask partitioning.Source: §4.1.2 · pp. 24–25, Figure 15 · p. 25

4.1.3 Rethinking System Design for Robust Distributed Training Frameworks with DTensor

  • The proposed DTensor and Parallel Plan extensions decouple model design from parallelization and provide a non-distributed oracle for checking distributed correctness.Higher-precision tensors, deterministic algorithms, and consistent seeds target alignment within relative error 10^-8.Source: §4.1.3 · p. 27

4.2 Inference Infrastructure

  • The inference infrastructure targets real-time streaming on H100/H800 GPUs and cost-efficient deployment on RTX 4090 GPUs.Multi-chunk parallel inference and KV cache support these deployment scenarios.Source: §4.2 · p. 28

4.2.1 Real-Time Streaming Video Generation

  • The largest 24B MAGI-1 model supports real-time streaming on 3 nodes with 24 H100 GPUs through a heterogeneous serving pipeline.T5 and MAGI-1 run on high-performance GPUs while the VAE uses cost-efficient hardware, enabling concurrent inference and decoding; the system addresses the under-1-second TPOC target and CPU-bound first-chunk latency.Source: §4.2.1 · p. 28
  • W8A8 SmoothQuant quantizes weights and activations to FP8 except in the first and last layers, delivering a 30% speedup without compromising generation quality.Ulysses-based multi-node inference overlaps computation and communication, leaving less than 3% of communication time unoverlapped.Source: §4.2.1 · p. 28
  • TPOC is within 1 second for 480p video on 24 H100/H800 GPUs, while TTFC is reduced to 2.3 seconds.The TPOC condition uses 16 denoising steps and KV range 5. CUDA Graphs reduce latency by 30.4%, and tile-based parallel VAE decoding with torch.compile lowers decoding latency from 1 second to around 70 milliseconds.Source: §4.2.1 · p. 29

4.2.2 Cost-effective Inference on RTX 4090

  • The 4.5B model runs on one RTX 4090 and the 24B model on 8×RTX 4090 GPUs.The RTX 4090’s 24G memory made memory insufficiency the primary serving bottleneck addressed by the inference architecture.Source: §4.2.2 · p. 29
  • Context Shuffle Overlap scatters each chunk evenly across GPUs to enable finer-grained computation–communication overlap and improve MFU under RTX 4090 PCIe bandwidth limits.The 24B setup combines pipeline parallelism for model weights with context parallelism for activations; task interleaving mitigates pipeline idle time.Source: §4.2.2 · pp. 29–30
  • Peak memory is 19.07GB for the 4.5B model on one RTX 4090 and 19.29 GB for the 24B model on 8×RTX 4090 GPUs; maximum MFU reaches 66%.KV-offload stores the KV cache in CPU memory and reloads it to the GPU as needed; the 24B deployment also uses hybrid parallelism.Source: §4.2.2 · pp. 29–30

5 Evaluation

MAGI-1 is evaluated through complementary perceptual and physics-based assessments, combining expert comparisons, automated I2V benchmarking, and tests of physical prediction. The results show strong benchmark performance while also identifying limitations in visual quality and complex physical effects.

5.1.1 In-house Human Evaluation Benchmark

  • The in-house I2V benchmark uses 100 diverse image-prompt pairs, four evaluation dimensions, and double-blind expert comparisons under standardized output conditions.Its dimensions are Overall, Motion Quality, Instruction Following, and Visual Quality; samples come from user-submitted inputs, FLUX images, public photographs, and cinematographic materials.Source: §5.1.1 · pp. 31–32
  • In overall performance, MAGI-1 outperforms open-source Wan-2.1, performs slightly behind commercial Kling1.6 (HD), and achieves clearly better results than Hailuo (i2v-01) and HunyuanVideo; it excels in instruction following and motion quality.Visual quality remains an area for improvement compared with top models.Source: §5.1.1 · p. 32

5.1.2 VBench

  • VBench evaluation focuses on I2V and compares 4-second, 24-FPS, 16:9 outputs from MAGI-1’s 1× and 2× decoder configurations.The 2× configuration changes the VAE decoder to provide 2× upsampling while leaving the core autoregressive denoising model identical.Source: §5.1.2 · p. 33
  • 89.28 gives MAGI-1 (2× decoder) the top overall VBench-I2V score, with advantages in dynamic Degree while maintaining strong aesthetic quality and motion smoothness.The 1× and 2× decoder models share the same autoregressive denoising model; the 2× version uses an enhanced VAE decoder.Source: §5.1.2 · p. 33, Table 9 · p. 34

5.2 Physical Evaluation

  • Physics-IQ tests prediction of the final 5 seconds of 8-second physical-event videos, using either the first 3 seconds of video or a single frame as input.The authors attribute MAGI-1’s strong physical results to autoregressive modeling’s emphasis on causality and its video-focused combination of autoregressive and denoising-based modeling.Source: §5.2 · p. 34
  • MAGI-1 captures primary dynamics but struggles with precise collisions, material-specific reactions, and post-deformation behavior, while sometimes producing coherent, physically plausible alternatives.For example, it misses match ignition and balloon popping but generates a coherent rod rotation, contact, and bending sequence.Source: §5.2 · pp. 34–35
  • 56.02 for video-conditioned MAGI-1 and 30.23 with image conditioning lead all models on Physics-IQ; the video-conditioned score exceeds VideoPoet by approximately 27 points.The video-conditioned setting uses the first 3 seconds of a video to predict the remaining 5 seconds; the image-conditioned setting uses only a frame at the 3-second mark.Source: §5.2 · p. 34, Table 10 · p. 35
  • Increasing MAGI-1’s accessible historical context generally improves physical prediction, with the largest gain at KV range = 2, where short-term history is often sufficient.The study varied the accessible history by adjusting the model’s KV range during inference.Source: §5.2 · p. 35

6 Related Works

Video generation has advanced through proprietary systems, open-source models, and emerging causal approaches, but challenges in scalability, causality, and streaming remain. MAGI-1 is presented as a chunk-wise autoregressive diffusion model designed to address these gaps.

Proprietary Systems

  • Proprietary systems advance generation length, resolution, fidelity, and control, but most remain closed-source and opaque, limiting reproducibility and extensibility.The cited systems include Sora, Kling, Gen-3, DreamMachine, Pika 1.5, MovieGen (Polyak et al., 2024), and Veo 2.Source: §6 · pp. 35–36

Open-Source Ecosystem

  • Open-source work progressed from latent diffusion to temporal and transformer-based video models, yet systems largely retain bidirectional denoising and global conditioning that limit causal, real-time use.Recent open implementations introduced modular advances in chunking, compression, and streaming.Source: §6 · p. 37

Autoregressive and Causal Modeling

  • Causal approaches use per-token noise schedules, timestep vectorization, or distillation for streaming, but remain limited in scale, chunk-wise abstraction, and unified video-task support.The passage attributes these approaches to Diffusion Forcing (Chen et al., 2024a), FVDM (Liu et al., 2024), and CausVid (Yin et al., 2024), respectively.Source: §6 · p. 37

MAGI-1: Scalable Autoregressive Diffusion

  • MAGI-1 is presented as the first large-scale, from-scratch chunk-wise autoregressive diffusion model to unify text-to-video, image-to-video, and video continuation under strict causal constraints.It supports real-time streaming and long-horizon synthesis through chunk-wise denoising, shortcut distillation, and KV-cached inference.Source: §6 · p. 37

7 Conclusion

MAGI-1 is a scalable chunk-wise autoregressive diffusion framework for high-fidelity video synthesis, designed for causal, streaming generation with fixed computational overhead. It unifies multiple video tasks and reports stronger performance than existing large-scale video diffusion models on two benchmarks.

  • Progressively denoising fixed-length chunks under strict causal constraints enables real-time, streaming-compatible generation with fixed computational overhead regardless of video length.Source: §7 · p. 37
  • Block-causal and parallel attention modules, distributed attention, and an efficient training strategy support MAGI-1’s Transformer architecture and ultra-long contexts.Source: §7 · p. 37
  • A shared training objective supports text-to-video, image-to-video, and video continuation without task-specific modifications, while chunk-wise text conditioning enables fine-grained semantic control.Source: §7 · p. 37
  • Shortcut distillation reduces inference diffusion steps while maintaining temporal consistency and sample quality.Source: §7 · p. 37
  • MAGI-1 outperforms existing large-scale video diffusion models on VBench-I2V and Physics-IQ in prompt adherence, physical plausibility, and temporal coherence.Source: §7 · p. 37

8 Limitation and Future Work · A Inference Infra · A.1 W8A8 Quantization

MAGI-1’s monolithic architecture creates latency, optimization, and controllability limitations, motivating a decoupled design. The W8A8 quantization strategy improves performance without compromising generation quality.

8 Limitation and Future Work

  • A single decoder-style Transformer handles both high-level temporal context fusion and low-level visual denoising, coupling heterogeneous objectives.It integrates static conditioning signals with progressively noisier visual inputs while reconstructing fine-grained visual details.Source: §8 · p. 38
  • Repeatedly invoking the same large model at every denoising step wastes compute, particularly when only minor refinements are needed in streaming settings.This creates an inference-latency bottleneck where low-latency frame delivery is critical.Source: §8 · p. 38
  • Jointly optimizing global semantic planning and pixel-level restoration exacerbates objective interference and can produce suboptimal scaling behavior.Source: §8 · p. 38
  • The monolithic architecture constrains auxiliary controls such as confidence-based guidance modulation and dynamic temporal constraints.Entangled latent pathways and overlapping functional scopes make these signals difficult to insert.Source: §8 · p. 38
  • The authors propose exploring decoupled semantic reasoning and visual synthesis, as video understanding becomes a bottleneck for coherent long-form generation.They view a modular architecture as a step toward closing the loop between video understanding and generation.Source: §8 · p. 38

A.1 W8A8 Quantization

  • A8W8 SmoothQuant uses calibration-derived per-channel scaling factors to mitigate outlier effects in channel-wise activations.It transforms Y as (X · diag(s)−1) · (diag(s)W).Source: §A.1 · p. 46
  • 30% performance improvement was achieved without compromising generation quality with the selected quantization strategy.Calibration covered T2V and I2V scenarios, with I2V making up approximately 30% of samples and step sizes uniformly sampled from [12, 32]. FP8 was used because INT8 introduced noticeable visual artifacts; α = 0.45 was selected from a search over α ∈(0.4, 0.6), and all weights except those in the first and last layers were quantized.Source: §A.1 · p. 46

A.2 Multi-Node Parallel Inference · A.3 Context Shuffle Overlap

The inference stack overlaps Ulysses communication with computation to reduce communication overhead, while Context Shuffle Overlap (CSO) repartitions context to enable chunk-level overlap and balanced all-to-all operations.

A.2 Multi-Node Parallel Inference

  • Ulysses overlaps v-communication with k-computation and k-communication with q-computation.Ulysses performs four all-to-all steps for q, k, v, and o tensors.Source: §A.2 · p. 46
  • Ulysses also overlaps q-communication with KV cache updates and o-communication with cross-attention computation.Source: §A.2 · p. 46
  • Communication overhead fell below 3% of total execution time with the overlapping multi-node inference strategy.The framework ran Ulysses-based parallel inference across 3 nodes (24 GPUs), where communication and attention compute were the primary bottlenecks.Source: §A.2 · p. 46

A.3 Context Shuffle Overlap

  • CSO was proposed for PCIe-connected RTX 4090 GPUs, and assigns each rank a partial view of every chunk rather than distributing whole chunks sequentially as Ulysses does.This context partitioning enables computation and communication to overlap at the chunk level, as shown in Figure 19.Source: §A.3 · p. 47, Figure 19 · p. 47
  • In the five-chunk CSO schedule, k- and v-communication across all chunks overlaps with q-computation, while q-communication for chunk 1 overlaps with KV cache updates.Source: §A.3 · p. 47
  • The five-chunk schedule then overlaps q-communication for later chunks and o-communication for earlier chunks with o-computation of successive chunks.Specifically, q-communication of chunks 2–5 and o-communication of chunks 1–3 overlap with o-computation of chunks 1–4.Source: §A.3 · p. 47
  • CSO finishes by overlapping o-communication of chunk 4 with chunk 5’s o-computation, then overlapping chunk 5’s o-communication with cross-attention computation.Source: §A.3 · p. 47
  • CSO splits communication into balanced all-to-all operations, which perform better than unbalanced all-to-all communication and support efficient subsequent merging.Source: §A.3 · p. 47

B Training Infrastructure · B.1 MagiAttention Materials

MagiAttention supports distributed attention across irregular masks by slicing and sharding work, reducing redundant communication, and overlapping computation with communication. Its materials also describe algorithms for balancing chunk workloads and selecting overlap stages.

B.1 MagiAttention Materials

  • AttnSlice decomposes irregular masks into slices that can be rearranged across context-parallel ranks for distributed attention; the illustration does not account for load balance.Source: Figure 20 · p. 48
  • FFA slice-level parallelism overlaps QRange rows and KRange columns, requiring atomic reductions in both forward and backward kernels.Source: Figure 21 · p. 48
  • Ring-Attention uses sequential sharding for full masks and zigzag sharding for causal masks; varlen causal zigzag sharding can degrade performance through fragmentation and padding.Source: Figure 22 · p. 48
  • Ring P2P incurs 25% redundant communication for a causal mask and over 33% for a varlen block-causal mask with a last global block.Source: Figure 23 · p. 49
  • Group-cast and group-reduce use transfer tables and all-to-all-v with Range Gather and Range Scatter-Reduce kernels to handle KV transfers and dKV reductions without redundancy.Source: Figure 24 · p. 49
  • Forward four-stage and backward three-stage schedules overlap computation with KV prefetching, hiding all communication except the final-stage backward dKV reduction.Source: Figure 25 · p. 50
  • The greedy dispatch algorithm sorts chunks by descending area and assigns each to the least-loaded non-full bucket to minimize maximum bucket workload under equal chunk capacity.Source: §B.1 · p. 50
  • Offline-profiled kernel costs guide bounded communication-package partitioning and stage-count search; ranks select forward and backward counts, then synchronize them by all-reduce.Source: §B.1 · p. 51

B.2 MagiAttention Experiments · B.2.1 Benchmarking MagiAttention kernel-level performance and flexibility

The MagiAttention experiments benchmark FFA kernel throughput and flexibility on Hopper GPUs across standard, variable-length, and irregular attention masks. The supplied passages describe the benchmark setup and figure scopes but do not provide measured throughput values or establish a winning kernel.

B.2.1 Benchmarking MagiAttention kernel-level performance and flexibility

  • FFA is benchmarked against other leading attention kernels on Hopper GPUs for full, causal, variable-length, and sliding-window causal masks.Figures 27–31 cover these mask scenarios; the supplied captions do not report comparative throughput values.Source: §B.2.1 · p. 52, Figure 27 · p. 53, Figure 28 · p. 54
  • The benchmark also covers irregular mask patterns, including sliding-window causal masks.Figure 31 notes that E marks a configuration where the corresponding distributed attention implementation raises a CUDA out-of-memory error.Source: §B.2.1 · p. 52, Figure 31 · p. 57
  • Throughput is measured in TFLOPs/s for both forward and backward passes, with sequence lengths from 4k to 128k and seqlenq = seqlenk.Other configurations are fixed to common training settings to isolate the effects of sequence length and mask pattern.Source: §B.2.1 · p. 52, Table 11 · p. 53
  • Variable-length packed inputs follow the sequence-length distribution of the open-sourced dataset (Xu et al., 2024).Samples from that distribution are packed and padded to the required sequence length; Figure 26 identifies the distribution used for kernel- and module-level experiments.Source: §B.2.1 · p. 52, Figure 26 · p. 52
  • Forward FLOPs are calculated from two matrix multiplications and the mask area, while backward FLOPs are 2.5 times the forward FLOPs.The calculation follows Flash-Attention (Dao, 2023), accounting for recomputation in the backward pass.Source: §B.2.1 · p. 52
  • For varlen full-mask and causal-mask scenarios, E marks configurations where the corresponding distributed attention implementation raises a CUDA out-of-memory error.These scenarios are shown in Figures 29 and 30.Source: Figure 29 · p. 55, Figure 30 · p. 56

B.2.2 Benchmarking MagiAttention module-level scalability

MagiAttention scales linearly with context length and context-parallel size in forward and backward propagation for full-mask and variable-length full-mask configurations, while baselines face scaling limits or performance degradation.

B.2.2 Benchmarking MagiAttention module-level scalability

  • The benchmark measures attention-module throughput in TFLOPs/s/GPU against Ring-Attention and Ulysses, but covers only full-mask and varlen full-mask scenarios because irregular masks are unsupported by baselines.Variable sequence lengths follow the distribution used in the kernel-level experiments in B.2.1.Source: §B.2.2 · p. 58
  • Experiments scale sequence length from 64k to 3072k, CP size from 1 to 48, and node count from 1 to 48, with TP size fixed at 8 and sequence parallelism enabled.The productive GPU cluster’s specific GPU count and type are withheld for business and confidentiality reasons.Source: §B.2.2 · p. 58
  • Each rank is assigned seqlen = 64k, num_headsq = 8, and num_headsk = 1 for attention propagation; remaining activations use seqlen = 8k, num_headsq = 64, and num_headsk = 8 with SP enabled.The setup simulates a common training configuration.Source: §B.2.2 · p. 58
  • MagiAttention exhibits linear scalability as context length and CP size increase for both full-mask and varlen full-mask configurations in forward and backward passes; baselines face scaling limits or degrade.Baseline degradation worsens with varlen masks at ultra-long contexts.Source: §B.2.2 · p. 58, Figure 32 · p. 59, Figure 33 · p. 60

B.3 Other Materials

Section B.3 provides examples of DTensor parallel placements and propagation used in modern training workflows, identifying the placements the authors complemented.

B.3 Other Materials

  • Table 12 shows examples of DTensor parallel placements and propagation common in modern training workflows, with bold text marking the placements the authors complemented.Source: Table 12 · p. 61, §B.3 · p. 61
Loading 2505.13211v1…