[feat] Add Ring Attention support - #1694
Conversation
mergify
Bot
added
type: feat
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
1 similar comment
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
1 similar comment
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
klhhhhh
changed the title
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
Run ring-atten and usp algorithm successfully.4-GPU Ring Attention Inference TestThis PR was manually validated with four GPUs using the following sequence-parallel configuration:
Because Inference example
from fastvideo import VideoGenerator
# from fastvideo.api.sampling_param import SamplingParam
OUTPUT_PATH = "video_samples"
def main():
# FastVideo will automatically use the optimal default arguments for the
# model.
# If a local path is provided, FastVideo will make a best-effort attempt
# to identify the optimal arguments.
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
# FastVideo will automatically handle distributed setup.
num_gpus=4,
sp_size=4,
ring_size=2,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set to False when CPU RAM is limited or an obscure
# "CUDA error: Invalid argument" is encountered.
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
# sampling_param = SamplingParam.from_pretrained(
# "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# )
# sampling_param.num_frames = 45
# sampling_param.image_path = (
# "https://huggingface.co/datasets/huggingface/"
# "documentation-images/resolve/main/diffusers/astronaut.jpg"
# )
# Generate a video using the same API regardless of the GPU count.
prompt = (
"A curious raccoon peers through a vibrant field of yellow sunflowers, "
"its eyes wide with interest. The playful yet serene atmosphere is "
"complemented by soft natural light filtering through the petals. "
"Mid-shot, warm and cheerful tones."
)
generator.generate_video(
prompt,
output_path=OUTPUT_PATH,
save_video=True,
)
# Generate another video without reloading the model.
prompt2 = (
"A majestic lion strides across the golden savanna, its powerful frame "
"glistening under the warm afternoon sun. The tall grass ripples gently "
"in the breeze, enhancing the lion's commanding presence. The tone is "
"vibrant, embodying the raw energy of the wild. Low angle, steady "
"tracking shot, cinematic."
)
generator.generate_video(
prompt2,
output_path=OUTPUT_PATH,
save_video=True,
)
if __name__ == "__main__":
main()Running scriptThe following script was used to run the four-GPU inference test: #!/usr/bin/env bash
set -euo pipefail
cd /FastVideo
export CUDA_VISIBLE_DEVICES=0,1,2,3
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
export NCCL_DEBUG=INFO
export PYTHONUNBUFFERED=1
export NCCL_CUMEM_ENABLE=0
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
echo "===== GPU information ====="
nvidia-smi \
--query-gpu=index,name,memory.total \
--format=csv
echo "===== FastVideo source ====="
python -c 'import fastvideo; print(fastvideo.__file__)'
echo "===== FlashAttention ====="
python -c 'import flash_attn; print("flash_attn:", flash_attn.__version__)'
echo "===== Start 4-GPU USP/Ring Attention inference ====="
echo "num_gpus=4, sp_size=4, ring_size=2"
mkdir -p video_samples logs
python examples/inference/basic/basic.py \
2>&1 | tee logs/ring_4gpu_pipeline.log
echo "===== Finished ====="
echo "Videos: /FastVideo/video_samples"
echo "Log: /FastVideo/logs/ring_4gpu_pipeline.log"Run commandchmod +x run_ring_4gpu.sh
./run_ring_4gpu.shOutputsGenerated videos: Execution log: |
|
And the running log is here. |
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
Pre-commit checks failedHi @klhhhhh, the pre-commit checks have failed. To fix them locally: # Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install
# Run all checks and auto-fix what's possible
pre-commit run --all-filesCommon fixes:
After fixing, commit and push the changes. The checks will re-run automatically. For future commits, |
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
…nd skip all-to-all when ulysses_size == 1
3 participants
Motivation
Video diffusion models process long spatiotemporal token sequences, making attention increasingly memory- and communication-intensive as video resolution and duration grow. Existing Ulysses sequence parallelism distributes sequence tokens across GPUs by exchanging sequence and attention-head dimensions, but its parallel degree is constrained by the number of attention heads.
This PR adds Ring Attention as an additional sequence-parallel strategy in FastVideo. Ring Attention keeps queries local while circulating sharded keys and values across a Ring process group, allowing every query shard to attend to the complete global sequence without materializing the full sequence on each GPU.
The PR supports both:
Pure Ring Attention, where the entire sequence-parallel group forms one Ring group (
ring_size == sp_size).Unified Sequence Parallelism (USP), which combines Ulysses and Ring Attention in a two-dimensional process mesh:
USP enables more flexible sequence-parallel configurations and avoids requiring the full sequence-parallel degree to be provided exclusively by either Ulysses or Ring Attention.
The implementation integrates Ring Attention with FastVideo's distributed process-group management, attention layer, and local RoPE handling. It also adds tests covering configuration validation, RoPE slicing, Ring subgroup and Ulysses subgroup construction, the single-GPU Ring fallback, blockwise log-sum-exp merging, and two-GPU pure-Ring numerical parity against full FlashAttention.
The current numerical parity test validates the pure-Ring kernel. The hybrid USP path is currently covered by process-group layout tests and multi-GPU integration runs.