[feat] Add Ring Attention support

fastvideo/attention/AGENTS.md

@@ -10,7 +10,9 @@ Backend registry + selector wrapping FlashAttn / SageAttn / SageAttn3 / SDPA / V
 attention/
 ├── __init__.py            # Exports DistributedAttention, LocalAttention, get_attn_backend
 ├── layer.py               # DistributedAttention, DistributedAttention_VSA, LocalAttention
+├── ring_attention.py      # RingAttention: owns every Ring/USP-specific decision (see below)
 ├── selector.py            # get_attn_backend (cached) + _component_attention_backend_scope
+├── ring/                  # Vendored Ring FlashAttention kernel (ring_flash_attn_func)
 ├── backends/
 │   ├── abstract.py        #   AttentionBackend / AttentionMetadata / AttentionMetadataBuilder
 │   ├── flash_attn.py      #   FA2/FA3
@@ -97,6 +99,35 @@ config. Thread the request alongside that tuple, one model family at a time —
 Wan, LTX-2 and Kandinsky5 first, since those carry per-role requests — and the
 scope goes away when the last family lands. Do not build on it.
 
+## Ring Attention / USP Ownership Boundary
+
+Ring Attention (and its Ring+Ulysses/USP hybrid) is a second delegate on
+`DistributedAttention`, with the same shape as `self.attn_impl`:
+
+- `DistributedAttention` reads only `get_ring_size()`, and only to fail loudly
+  if a subclass overrides `forward()` wholesale (e.g. `DistributedAttention_VSA`)
+  without wiring in the Ring dispatch. Everything else about Ring — reading
+  the process-wide topology, the Ulysses-within-ring all-to-all, Ring-local
+  RoPE slicing, construction-time validation, and the direct call into the
+  vendored Ring FlashAttention kernel (bypassing `self.attn_impl` /
+  `AttentionBackend` entirely, since Ring picks *which* backend runs per
+  shard rather than being one) — lives in `RingAttention`
+  (`ring_attention.py`). `DistributedAttention.forward()` dispatches to
+  `self._ring_attention.forward(...)` when Ring is enabled, exactly the way
+  it dispatches to `self.attn_impl.forward(...)` otherwise.
+- The Ring x Ulysses topology itself (deriving ring/ulysses sizes from
+  `sp_size`, the pure-Ring / pure-Ulysses degenerate cases, building the 2D
+  mesh of subgroups) is policy owned by `fastvideo/distributed/usp_topology.py`,
+  not `parallel_state.py`. `parallel_state.py` calls
+  `build_usp_topology()` once from `initialize_model_parallel()` and stores
+  the resulting `USPTopology`, then hands out thin accessors
+  (`get_ring_group()`, `get_ulysses_group()`, `get_ring_rank()`, ...) the same
+  way it does for `_TP`/`_SP`/`_DP` — it stays a process-group registry, not a
+  sequence-parallel policy module.
+
+A new Ring/USP-specific decision belongs in `RingAttention` or
+`usp_topology.py`, not back in `layer.py` or `parallel_state.py`.
+
 ## Adding a Backend
 
 1. Subclass `AttentionBackend` in `backends/<name>.py`.

fastvideo/attention/layer.py

@@ -6,10 +6,11 @@
 import torch
 import torch.nn as nn
 
+from fastvideo.attention.ring_attention import RingAttention
 from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
 from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather,
                                                     sequence_model_parallel_all_to_all_4D)
-from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, get_sp_world_size)
+from fastvideo.distributed.parallel_state import get_ring_size, get_sp_parallel_rank, get_sp_world_size
 from fastvideo.forward_context import ForwardContext, get_forward_context
 from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.utils import get_compute_dtype
@@ -99,6 +100,28 @@ def __init__(self,
         self.num_kv_heads = num_kv_heads
         self.backend = backend_name_to_enum(attn_backend.get_name())
         self.dtype = dtype
+        self.causal = causal
+
+        # Ring Attention (and its USP hybrid with Ulysses) is a process-wide
+        # parallelism setting (like SP itself), so RingAttention itself reads
+        # it from the distributed runtime rather than it being threaded
+        # through every model's constructor. The one thing only this class
+        # can decide is the guard below: a subclass that overrides forward()
+        # wholesale (e.g. LTXDistributedAttention, DistributedAttention_VSA)
+        # would otherwise never reach the Ring dispatch in forward() below
+        # and silently keep running plain Ulysses SP -- fail loudly here
+        # instead of letting that divergence pass unnoticed.
+        if get_ring_size() > 1 and type(self).forward is not DistributedAttention.forward:
+            raise NotImplementedError(f"Ring Attention is not implemented for {type(self).__name__}, which overrides "
+                                      "DistributedAttention.forward() directly instead of using the base dispatch. "
+                                      "Disable Ring Attention (ring_size=1) for this model, or wire Ring Attention "
+                                      "dispatch into its forward().")
+        self._ring_attention = RingAttention.create_if_enabled(num_heads=num_heads,
+                                                               num_kv_heads=num_kv_heads,
+                                                               softmax_scale=self.softmax_scale,
+                                                               causal=causal,
+                                                               backend=self.backend)
+        self.use_ring_attention = self._ring_attention is not None
         # Preserve the historical compiler-disabled default. The regional
         # inference loader may enable this one instance after validating the
         # transformer's resolved backend; no process-global default changes.
@@ -120,7 +143,7 @@ def forward(
         freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass for distributed attention.
-        
+
         Args:
             q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
             k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
@@ -129,14 +152,50 @@ def forward(
             replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
             replicated_k (Optional[torch.Tensor]): Replicated key tensor
             replicated_v (Optional[torch.Tensor]): Replicated value tensor
-            
+
         Returns:
             Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
                 - o (torch.Tensor): Output tensor after attention for the main sequence
                 - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
         """
         # Check input shapes
         assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
+
+        if self._ring_attention is not None:
+            return self._ring_attention.forward(
+                q,
+                k,
+                v,
+                training=self.training,
+                original_seq_len=original_seq_len,
+                replicated_q=replicated_q,
+                replicated_k=replicated_k,
+                replicated_v=replicated_v,
+                freqs_cis=freqs_cis,
+            )
+
+        return self._forward_ulysses_attention(
+            q=q,
+            k=k,
+            v=v,
+            original_seq_len=original_seq_len,
+            replicated_q=replicated_q,
+            replicated_k=replicated_k,
+            replicated_v=replicated_v,
+            freqs_cis=freqs_cis,
+        )
+
+    def _forward_ulysses_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        original_seq_len: int | None,
+        replicated_q: torch.Tensor | None,
+        replicated_k: torch.Tensor | None,
+        replicated_v: torch.Tensor | None,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         batch_size, _, num_heads, _ = q.shape
         local_rank = get_sp_parallel_rank()
         world_size = get_sp_world_size()

fastvideo/attention/ring/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# This package vendors Ring Attention kernels for FastVideo's Ring Attention
+# integration. See each module's header for its upstream source and license.
+#
+# Only ``ring_flash_attn`` (Ring Attention, optionally combined with Ulysses
+# as the USP hybrid, over plain FlashAttention) is vendored here — it is what
+# ``fastvideo.attention.ring_attention.RingAttention`` uses. Upstream yunchang also has zig-zag,
+# striped, variable-length, PyTorch, NPU, and FlashInfer variants; FastVideo
+# does not vendor them (they would depend on the ``yunchang`` package or
+# other optional backends FastVideo does not want as hard runtime
+# dependencies), and already has its own backend-dispatch system
+# (``fastvideo.attention.selector``) for backends other than FlashAttention.
+
+from .ring_flash_attn import ring_flash_attn_func
+
+__all__ = [
+    "ring_flash_attn_func",
+]

fastvideo/attention/ring/capabilities.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Adapted from:
+# https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
+#
+# FastVideo keeps only yunchang's optional-kernel capability detection here,
+# trimmed to the single backend the vendored Ring Attention kernels dispatch
+# to (plain FlashAttention). Sequence-parallel process groups are managed by
+# fastvideo.distributed.parallel_state.
+
+try:
+    import flash_attn  # noqa: F401
+    from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward  # noqa: F401
+
+    HAS_FLASH_ATTN = True
+except ImportError:
+    HAS_FLASH_ATTN = False

fastvideo/attention/ring/kernels/__init__.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Adapted from:
+# https://github.com/feifeibear/long-context-attention/blob/main/yunchang/kernels/__init__.py
+#
+# FastVideo keeps only the FA dispatch path used by the vendored pure-Ring
+# Attention implementation (fastvideo.attention.layer). Upstream yunchang
+# dispatches to many more backends (FA3, FlashInfer, aiter, SageAttention,
+# torch SDPA); FastVideo already has its own backend-dispatch system
+# (fastvideo.attention.selector / backends/), so those paths were dropped
+# here rather than carried as a second, unreachable dispatcher.
+
+from __future__ import annotations
+
+from enum import Enum
+
+from .attention import flash_attn_backward, flash_attn_forward
+
+
+class AttnType(Enum):
+    FA = "fa"
+
+
+def select_flash_attn_impl(impl_type: AttnType, stage: str = "fwd-bwd"):
+    if impl_type != AttnType.FA:
+        raise ValueError(f"Unknown flash attention implementation: {impl_type}")
+    if stage == "fwd-only":
+        return flash_attn_forward
+    elif stage == "bwd-only":
+        return flash_attn_backward
+    else:
+        raise ValueError(f"Unknown stage: {stage}")
+
+
+__all__ = [
+    "AttnType",
+    "select_flash_attn_impl",
+]

fastvideo/attention/ring/kernels/attention.py

@@ -0,0 +1,117 @@
+import re
+from typing import Any
+
+import torch
+
+from ..capabilities import HAS_FLASH_ATTN
+
+if HAS_FLASH_ATTN:
+    import flash_attn
+    from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward
+
+    def _parse_version(version: str) -> tuple[int, ...]:
+        # A plain string compare (e.g. "2.10.0" <= "2.6.3") is lexicographic
+        # and misclassifies once a two-digit minor version ships, so parse
+        # to an int tuple instead. Non-numeric suffixes (e.g. "2.6.3.post1")
+        # are tolerated by defaulting an unparsable segment to 0.
+        parts = []
+        for part in version.split(".")[:3]:
+            match = re.match(r"\d+", part)
+            parts.append(int(match.group()) if match else 0)
+        return tuple(parts)
+
+    _FLASH_ATTN_VERSION = _parse_version(flash_attn.__version__)
+
+
+def flash_attn_forward(q: torch.Tensor,
+                       k: torch.Tensor,
+                       v: torch.Tensor,
+                       dropout_p: float = 0.0,
+                       softmax_scale: float | None = None,
+                       causal: bool = False,
+                       window_size: tuple[int, int] = (-1, -1),
+                       softcap: float | None = None,
+                       alibi_slopes: torch.Tensor | None = None,
+                       return_softmax: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    assert HAS_FLASH_ATTN, "FlashAttention is not available"
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1]**(-0.5)
+    if _FLASH_ATTN_VERSION <= (2, 6, 3):
+        block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax,
+        )
+    else:
+        block_out, block_lse, _, _ = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax,
+        )
+    return block_out, block_lse
+
+
+def flash_attn_backward(dout: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor,
+                        softmax_lse: torch.Tensor, block_dq_buffer: torch.Tensor, block_dk_buffer: torch.Tensor,
+                        block_dv_buffer: torch.Tensor, dropout_p: float, softmax_scale: float | None, bwd_causal: bool,
+                        window_size: tuple[int, int], softcap: float | None, alibi_slopes: torch.Tensor | None,
+                        deterministic: bool, rng_state: Any) -> None:
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1]**(-0.5)
+    assert HAS_FLASH_ATTN
+    if _FLASH_ATTN_VERSION <= (2, 6, 3):
+        _flash_attn_backward(
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            block_dq_buffer,
+            block_dk_buffer,
+            block_dv_buffer,
+            dropout_p,
+            softmax_scale,
+            bwd_causal,
+            window_size,
+            softcap,
+            alibi_slopes,
+            deterministic,
+            rng_state,
+        )
+    else:
+        _flash_attn_backward(
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            block_dq_buffer,
+            block_dk_buffer,
+            block_dv_buffer,
+            dropout_p,
+            softmax_scale,
+            bwd_causal,
+            window_size[0],  # Pass window_size_left
+            window_size[1],  # Pass window_size_right
+            softcap,
+            alibi_slopes,
+            deterministic,
+            rng_state,
+        )