ring-sliding-window-attention

KO | EN

gitlite — search

#javascript #hacktoberfest #golang #machine-learning #deep-learning #linux #docker #nodejs #awesome #vue #awesome-list #ios #openai #open-source

ring-sliding-window-attention

(★ 11)

Ring sliding window attention implementation with flash attention

ring-sliding-window-attention Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 align="center">Ring Sliding Window Attention</h1>

</div>

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](docs/implementation_details.md).

## Features

- **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`.

</div>

## Installation

### Install with pip

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

### Dependencies

- Python \>= 3.8
  - torch \>= 2.7.0
  - flash\_attn \>= 2.5.8
  - einops \>= 0.6.0
  - triton \>= 3.0.0

## Usage

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.

### Function Arguments

* `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)`.

### Non-Varlen Example

```python
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()
```

### Varlen Example

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.

```python
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()
```

### Streaming LLM Attention

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.

```python
from ring_swa import ring_streaming_llm_attn
from ring_swa import ring_streaming_llm_attn_varlen
```

## Testing

You can run the test scripts to verify the correctness:

```bash
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
```

## License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

## Contributing

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

ring-sliding-window-attention

ring-sliding-window-attention Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

ring-sliding-window-attention

ring-sliding-window-attention Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README