This repository provides a Ring Sliding Window Attention implementation for efficient long-sequence training with context parallelism.

For a complete implementation details, please refer to our documentation.

• Ring Sliding Window Attention: A distributed attention mechanism for sliding window attention.
• Ring Streaming LLM Attention: A distributed attention mechanism for sliding window attention with attention sink token.
• Variable Length Support: Supports variable-length sequences, similar to flash_attn_varlen_func.

pip install git+https://github.com/XunhaoLai/ring-sliding-window-attention.git

• Python >= 3.8
• torch >= 2.7.0
• flash_attn >= 2.5.8
• einops >= 0.6.0
• triton >= 3.0.0

In ring sliding window attention, the query, key, and value tensors are split into cp_size chunks. Each chunk i is then placed on its corresponding rank i within the context parallel group. For simplicity, we'll use random tensors in following examples.

• q: The Query tensor.
  • Shape (standard): [batch_size, seq_len, num_q_heads, head_dim]
  • Shape (varlen): [total_seq_len, num_q_heads, head_dim]
• k: The Key tensor.
  • Shape (standard): [batch_size, seq_len, num_kv_heads, head_dim]
  • Shape (varlen): [total_seq_len, num_kv_heads, head_dim]
• v: The Value tensor, which has the same shape as the key tensor k.
• window_size: The number of tokens that each query will attend to. This is similar to window_size[0] in flash_attn.
• cu_seqlens: A tensor of cumulative sequence lengths for the entire batch across all ranks. Its shape is [batch_size + 1]. This is only required for the varlen version.
• cp_group: The process group for context parallelism. If None, the function runs in a single-GPU mode.
• cp_local_rank: The rank of the current process within its context parallel group.
• cp_prev_global_rank: The global rank of the previous process in the ring.
• cp_next_global_rank: The global rank of the next process in the ring.
• softmax_scale: The scaling factor for the softmax function. It defaults to 1 / sqrt(head_dim).

Note on Ranks: The rank arguments (cp_local_rank, cp_prev_global_rank, and cp_next_global_rank) are optional. If you don't provide them, they're computed automatically. First, the cp_local_rank is found using torch.distributed.get_rank(group=cp_group). Then, a local-to-global rank mapping is used to determine the cp_prev_global_rank and cp_next_global_rank. You can get this mapping yourself with the following utility: ring_swa.ops.get_group_local_to_global_rank_map(cp_group).

import torch
import torch.distributed as dist
from ring_swa import ring_sliding_window_attn

# Initialize distributed training
dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank())
cp_group = dist.group.WORLD
world_size = dist.get_world_size()
rank = dist.get_rank()

# Prepare input tensors; here we use random tensors at each rank
batch_size, seq_len, num_heads, head_dim = 2, 8192, 32, 128
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
window_size = 1000

# Compute ring sliding window attention
output = ring_sliding_window_attn(
    q, k, v,
    window_size=window_size,
    cp_group=cp_group,
    cp_size=world_size,
    cp_local_rank=rank,
    cp_prev_global_rank=(rank - 1 + world_size) % world_size,
    cp_next_global_rank=(rank + 1) % world_size,
)

# Print attention output at each rank
print(f"[rank] {rank}, attention output shape: {output.shape}, norm: {output.norm()}")

dist.destroy_process_group()

This is similar to the non-varlen version, but requires an extra input, cu_seqlens. Note that cu_seqlens is based on all sequences in a batch, not just those related to the current rank.

import torch
import torch.distributed as dist
from ring_swa import ring_sliding_window_attn_varlen

# Initialize distributed training
dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank())
cp_group = dist.group.WORLD
world_size = dist.get_world_size()
rank = dist.get_rank()

# Prepare input tensors; here we use random tensors at each rank
seq_len, num_heads, head_dim = 8192, 32, 128
cu_seqlens = torch.LongTensor([0, 512, 4000, 8192]).cuda().to(torch.int32)
q = torch.randn(seq_len // world_size, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn(seq_len // world_size, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
v = torch.randn(seq_len // world_size, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
window_size = 1000

# Compute ring sliding window attention
output = ring_sliding_window_attn_varlen(
    q, k, v,
    window_size=window_size,
    cu_seqlens=cu_seqlens,
    cp_group=cp_group,
    cp_size=world_size,
    cp_local_rank=rank,
    cp_prev_global_rank=(rank - 1 + world_size) % world_size,
    cp_next_global_rank=(rank + 1) % world_size,
)

# Print attention output at each rank
print(f"[rank] {rank}, attention output shape: {output.shape}, norm: {output.norm()}")

dist.destroy_process_group()

The usage of ring_streaming_llm_attn and ring_streaming_llm_attn_varlen is similar to the standard ring sliding window attention functions. The key difference is its use of an attention sink, a mechanism that forces every query in the sequence to attend to the first token.

❗️Note: These functions currently assume that the first token for each sequence has identical content (e.g., a BOS token). Please ensure this condition is met before use.

from ring_swa import ring_streaming_llm_attn
from ring_swa import ring_streaming_llm_attn_varlen

You can run the test scripts to verify the correctness:

torchrun --nproc_per_node=8 test/test_ring_swa_nonvarlen.py
torchrun --nproc_per_node=8 test/test_ring_swa_varlen.py
torchrun --nproc_per_node=8 test/test_ring_slm_nonvarlen.py
torchrun --nproc_per_node=8 test/test_ring_slm_varlen.py

This project is licensed under the MIT License - see the LICENSE file for details.

Contributions are welcome! Please feel free to submit a Pull Request. For any questions or feedback, you can contact laixunhao@pku.edu.cn.