[TRTLLM-12342][feat] Ring Attention, Unified Context Parallel for VisualGen

docs/source/models/visual-generation.md

@@ -37,13 +37,13 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 
 ### Feature Matrix
 
-| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve |
-|---|---|---|---|---|---|---|---|---|---|
-| **FLUX.1** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes |
-| **FLUX.2** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes |
-| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| **Wan 2.2** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
-| **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes |
+| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention |
+|---|---|---|---|---|---|---|---|---|---|--|--|
+| **FLUX.1** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes | Yes | Yes |
+| **FLUX.2** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes | Yes | Yes |
+| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
+| **Wan 2.2** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
+| **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes |
 
 [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
 
@@ -113,12 +113,14 @@ TeaCache caches transformer outputs when timestep embeddings change slowly betwe
 
 ### Multi-GPU Parallelism
 
-Three parallelism modes can be combined:
+5 parallelism modes can be combined:
 
 - **CFG Parallelism** (`--cfg_size 2`): Splits positive/negative guidance prompts across GPUs.
 - **Ulysses Parallelism** (`--ulysses_size N`): Splits the sequence dimension across GPUs for longer sequences.
 - **Parallel VAE** (`--parallel_vae_size N`): Shards the final VAE decode along a spatial axis across GPUs, useful to reduce VAE latency and improve GPU utilization (Constraint: `parallel_vae_size ≤ world_size`). Currently only supported for WAN models.
-
+- **Attention Parallel**: There are 2 methods supported to run attention parallel. Both of these methods require the attention backend to support LSE (only FA4 currently) - 
+    - **Attention2D Parallelism** (`--attn2d_row_size N`, `--attn2d_col_size M`): Shards the sequence axis across a 2D `N x M` device mesh, all-gathering Q along rows and K/V along columns so each rank computes a sub-block of the attention matrix (total CP degree = `N * M`; not currently combinable with Ulysses).
+    - **Ring Attention Parallelism** (`--ring_size N`): Shards the sequence axis across a 1D ring of `N` ranks and streams K/V blocks around the ring so each rank computes its attention output without materializing the full K/V (mutually exclusive with Attention2D).
 ## Developer Guide
 
 ### Architecture Overview

examples/visual_gen/visual_gen_wan_i2v.py

@@ -240,6 +240,12 @@ def parse_args():
         "Can be set independently of --attn2d_row_size; asymmetric meshes (e.g. 1x4 or 4x1) are valid. "
         "Cannot be combined with --ulysses_size (not yet implemented).",
     )
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="Ring Attention parallel size. Cannot be combined with --attn2d_row_size / --attn2d_col_size.",
+    )
     parser.add_argument(
         "--parallel_vae_size",
         type=int,
@@ -351,9 +357,13 @@ def main():
         raise ValueError(
             "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented."
         )
+    if args.ring_size > 1 and attn2d_size > 1:
+        raise ValueError(
+            "Combining --ring_size with --attn2d_row_size/--attn2d_col_size is not yet implemented."
+        )
 
-    if args.ulysses_size > 1:
-        parallel_str = f"Ulysses(size={args.ulysses_size})"
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        parallel_str = f"Ulysses(size={args.ulysses_size}), Ring(size={args.ring_size})"
     elif attn2d_size > 1:
         parallel_str = (
             f"Attention2D(row={args.attn2d_row_size}, col={args.attn2d_col_size}, "
@@ -370,6 +380,7 @@ def main():
             "dit_ulysses_size": args.ulysses_size,
             "dit_attn2d_row_size": args.attn2d_row_size,
             "dit_attn2d_col_size": args.attn2d_col_size,
+            "dit_ring_size": args.ring_size,
             "parallel_vae_size": args.parallel_vae_size,
         },
         torch_compile={

examples/visual_gen/visual_gen_wan_t2v.py

@@ -239,6 +239,12 @@ def parse_args():
         "Can be set independently of --attn2d_row_size; asymmetric meshes (e.g. 1x4 or 4x1) are valid. "
         "Cannot be combined with --ulysses_size (not yet implemented).",
     )
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="Ring Attention parallel size. Cannot be combined with --attn2d_row_size / --attn2d_col_size.",
+    )
     parser.add_argument(
         "--parallel_vae_size",
         type=int,
@@ -328,9 +334,21 @@ def main():
         raise ValueError(
             "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented."
         )
+    if args.ring_size > 1 and attn2d_size > 1:
+        raise ValueError(
+            "Combining --ring_size with --attn2d_row_size/--attn2d_col_size is not yet implemented."
+        )
 
     if args.ulysses_size > 1:
-        parallel_str = f"Ulysses(size={args.ulysses_size})"
+        num_heads = 40
+        logger.info(
+            f"Using Ulysses sequence parallelism: "
+            f"{num_heads} heads / {args.ulysses_size} ranks = "
+            f"{num_heads // args.ulysses_size} heads per GPU"
+        )
+
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        parallel_str = f"Ulysses(size={args.ulysses_size}), Ring(size={args.ring_size})"
     elif attn2d_size > 1:
         parallel_str = (
             f"Attention2D(row={args.attn2d_row_size}, col={args.attn2d_col_size}, "
@@ -370,6 +388,7 @@ def main():
             "dit_attn2d_row_size": args.attn2d_row_size,
             "dit_attn2d_col_size": args.attn2d_col_size,
             "parallel_vae_size": args.parallel_vae_size,
+            "dit_ring_size": args.ring_size,
         },
         torch_compile={
             "enable_torch_compile": not args.disable_torch_compile,

tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py

@@ -22,7 +22,7 @@
 
 from .flash_attn4 import FlashAttn4Attention
 from .interface import AttentionBackend, AttentionTensorLayout
-from .parallel import Attention2DAttention, UlyssesAttention
+from .parallel import Attention2DAttention, RingAttention, UlyssesAttention
 from .trtllm import TrtllmAttention, TrtllmAttentionMetadata
 from .utils import create_attention, get_visual_gen_attention_backend
 from .vanilla import VanillaAttention
@@ -38,4 +38,5 @@
     "TrtllmAttentionMetadata",
     "UlyssesAttention",
     "VanillaAttention",
+    "RingAttention",
 ]

tensorrt_llm/_torch/visual_gen/attention_backend/interface.py

@@ -78,4 +78,5 @@ def support_fused_qkv(cls) -> bool:
 
     @classmethod
     def support_lse(cls) -> bool:
+        """Whether the backend supports returning the softmax log-sum-exp (LSE) of the attention weights."""
         return False

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py

@@ -20,9 +20,9 @@
 
 """
 
-from typing import Optional
-
 import torch
+import torch.distributed as dist
+import torch.nn.functional as F
 
 from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d
 
@@ -65,7 +65,7 @@ class UlyssesAttention(AttentionBackend):
     def __init__(
         self,
         inner_backend: AttentionBackend,
-        process_group: Optional[torch.distributed.ProcessGroup] = None,
+        process_group: torch.distributed.ProcessGroup,
     ):
         self.inner_backend = inner_backend
         self.process_group = process_group
@@ -75,10 +75,7 @@ def __init__(
         self.sharded_num_heads = inner_backend.num_heads
         self.sharded_num_kv_heads = getattr(inner_backend, "num_kv_heads", self.sharded_num_heads)
 
-        try:
-            self.world_size = torch.distributed.get_world_size(group=process_group)
-        except (RuntimeError, ValueError):
-            self.world_size = 1
+        self.world_size = torch.distributed.get_world_size(group=process_group)
 
         self.num_heads = self.sharded_num_heads * self.world_size
         self.num_kv_heads = self.sharded_num_kv_heads * self.world_size
@@ -96,6 +93,19 @@ def forward(
         q/k/v: [B, S/P, H, D] each.  All other arguments are forwarded
         transparently to the inner backend via ``**kwargs``.
         """
+        # Catches upstream floor-division bugs (e.g. num_heads // ulysses_size when
+        # num_heads % ulysses_size != 0) before they corrupt the all-to-all.
+        if q.shape[2] % self.world_size != 0:
+            raise ValueError(
+                f"UlyssesAttention: q num_heads ({q.shape[2]}) must be divisible "
+                f"by world_size ({self.world_size})."
+            )
+        if k.shape[2] % self.world_size != 0:
+            raise ValueError(
+                f"UlyssesAttention: k num_kv_heads ({k.shape[2]}) must be divisible "
+                f"by world_size ({self.world_size})."
+            )
+
         if self.inner_backend.support_fused_qkv():
             return self._forward_fused(q, k, v, **kwargs)
         return self._forward_unfused(q, k, v, **kwargs)
@@ -109,8 +119,7 @@ def _forward_fused(
     ) -> torch.Tensor:
         batch_size = q.shape[0]
         qkv = torch.stack([q, k, v], dim=2)
-        if self.world_size > 1:
-            qkv = all_to_all_5d(qkv, scatter_dim=3, gather_dim=1, process_group=self.process_group)
+        qkv = all_to_all_5d(qkv, scatter_dim=3, gather_dim=1, process_group=self.process_group)
 
         B, seq_len, _, Hp, D = qkv.shape
 
@@ -132,10 +141,9 @@ def _forward_unfused(
         **kwargs,
     ) -> torch.Tensor:
         batch_size = q.shape[0]
-        if self.world_size > 1:
-            q = all_to_all_4d(q, scatter_dim=2, gather_dim=1, process_group=self.process_group)
-            k = all_to_all_4d(k, scatter_dim=2, gather_dim=1, process_group=self.process_group)
-            v = all_to_all_4d(v, scatter_dim=2, gather_dim=1, process_group=self.process_group)
+        q = all_to_all_4d(q, scatter_dim=2, gather_dim=1, process_group=self.process_group)
+        k = all_to_all_4d(k, scatter_dim=2, gather_dim=1, process_group=self.process_group)
+        v = all_to_all_4d(v, scatter_dim=2, gather_dim=1, process_group=self.process_group)
 
         seq_len_full = q.shape[1]
         kv_seq_len_full = k.shape[1]
@@ -173,10 +181,9 @@ def _output_a2a(
                 )
             output = output.contiguous()
 
-        if self.world_size > 1:
-            output = all_to_all_4d(
-                output, scatter_dim=1, gather_dim=2, process_group=self.process_group
-            )
+        output = all_to_all_4d(
+            output, scatter_dim=1, gather_dim=2, process_group=self.process_group
+        )
 
         return output
 
@@ -439,3 +446,167 @@ def support_fused_qkv(cls) -> bool:
         # If a future backend supports both LSE and fused QKV with a faster kernel,
         # add fused QKV support.
         return False
+
+    @classmethod
+    def support_lse(cls) -> bool:
+        return False
+
+
+class RingAttention(AttentionBackend):
+    """Ring sequence parallelism around an LSE-capable attention backend."""
+
+    def __init__(
+        self,
+        inner_backend: AttentionBackend,
+        process_group: dist.ProcessGroup,
+    ):
+        if not type(inner_backend).support_lse():
+            raise ValueError(
+                f"RingAttention requires an LSE-capable inner backend (FA4); "
+                f"got {type(inner_backend).__name__}"
+            )
+
+        # Required attributes for buffer allocation in _ensure_buffers.
+        for attr in ("head_dim", "num_heads"):
+            if not hasattr(inner_backend, attr):
+                raise RuntimeError(
+                    f"{type(inner_backend).__name__} is missing required attribute "
+                    f"'{attr}'. RingAttention needs the inner backend to expose "
+                    "'head_dim' and 'num_heads' as instance attributes."
+                )
+
+        # Ring's _ensure_buffers / _update_out_and_lse assume NHD ([B, S, H, D]).
+        # No transpose is applied around the inner forward, so an HND backend would
+        # silently produce wrong results.
+        if inner_backend.preferred_layout != AttentionTensorLayout.NHD:
+            raise NotImplementedError(
+                f"RingAttention requires an NHD inner backend; "
+                f"{type(inner_backend).__name__} prefers {inner_backend.preferred_layout}."
+            )
+
+        self.inner = inner_backend
+        self.pg = process_group
+        self.world_size = dist.get_world_size(group=process_group)
+        self.num_heads = inner_backend.num_heads
+        self.num_kv_heads = getattr(inner_backend, "num_kv_heads", self.num_heads)
+        self.head_dim = inner_backend.head_dim
+        self._preferred_layout = AttentionTensorLayout.NHD
+
+        # P2P ring topology cached at construction time (avoids per-step lookups).
+        ring_rank = dist.get_rank(group=process_group)
+        self._send_rank = dist.get_global_rank(process_group, (ring_rank + 1) % self.world_size)
+        self._recv_rank = dist.get_global_rank(process_group, (ring_rank - 1) % self.world_size)
+        self._send_first = ring_rank % 2 == 0
+        self._p2p_reqs: list = []
+
+        self._buf_key = None
+        self._kv_bufs = None
+        self._out_buf = None
+        self._lse_buf = None
+
+    def _ring_send_recv(self, send: torch.Tensor, recv: torch.Tensor) -> None:
+        """Post a non-blocking neighbor exchange. Even ranks send-then-recv to
+        avoid deadlock against odd ranks doing recv-then-send."""
+        if self._send_first:
+            ops = [
+                dist.P2POp(dist.isend, send, self._send_rank, group=self.pg),
+                dist.P2POp(dist.irecv, recv, self._recv_rank, group=self.pg),
+            ]
+        else:
+            ops = [
+                dist.P2POp(dist.irecv, recv, self._recv_rank, group=self.pg),
+                dist.P2POp(dist.isend, send, self._send_rank, group=self.pg),
+            ]
+        self._p2p_reqs = dist.batch_isend_irecv(ops)
+
+    def _ring_wait(self) -> None:
+        for r in self._p2p_reqs:
+            r.wait()
+        self._p2p_reqs.clear()
+
+    def _ensure_buffers(self, q: torch.Tensor, k: torch.Tensor) -> None:
+        B, S, H, D = q.shape
+        H_kv = k.shape[2]
+        key = (B, S, H, H_kv, D, q.device, q.dtype, k.dtype)
+        if key == self._buf_key:
+            return
+        self._kv_bufs = k.new_empty(2, 2, B, S, H_kv, D)
+        # Accumulate ring blocks in fp32 to avoid repeated bf16<->fp32 rounding
+        # across online-softmax merges
+        self._out_buf = q.new_empty(B, S, H, D, dtype=torch.float32)
+        self._lse_buf = q.new_empty(B, S, H, dtype=torch.float32)
+        self._buf_key = key
+
+    def _update_out_and_lse(
+        self,
+        out: torch.Tensor,
+        lse: torch.Tensor,
+        block_out: torch.Tensor,
+        block_lse: torch.Tensor,
+    ) -> None:
+        """Online-softmax merge of (out, lse) with (block_out, block_lse). In-place on out/lse."""
+        c = torch.sigmoid(block_lse.unsqueeze(-1) - lse.unsqueeze(-1))
+        out.sub_(c * (out - block_out))
+        lse.sub_(F.logsigmoid(lse - block_lse))
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
+        **kwargs,
+    ) -> torch.Tensor:
+        # Bypass ring for cross-attention (Q/KV seq lengths differ).
+        if k.shape[1] != q.shape[1]:
+            return self.inner.forward(q=q, k=k, v=v, attention_mask=attention_mask, **kwargs)
+        if attention_mask != PredefinedAttentionMask.FULL:
+            raise NotImplementedError(
+                f"RingAttention only supports FULL attention mask, got {attention_mask}."
+            )
+
+        inner_kw = {kk: vv for kk, vv in kwargs.items() if kk != "attention_mask"}
+
+        self._ensure_buffers(q, k)
+        kv_bufs = self._kv_bufs
+        out = self._out_buf
+        lse = self._lse_buf
+
+        kv_bufs[0, 0].copy_(k)
+        kv_bufs[0, 1].copy_(v)
+        for step in range(self.world_size):
+            cur, nxt = step % 2, 1 - step % 2
+            if step < self.world_size - 1:
+                self._ring_send_recv(kv_bufs[cur], kv_bufs[nxt])
+            block_out, block_lse_bh = self.inner.forward_with_lse(
+                q=q,
+                k=kv_bufs[cur, 0],
+                v=kv_bufs[cur, 1],
+                attention_mask=PredefinedAttentionMask.FULL,
+                **inner_kw,
+            )
+            # Inner backend returns LSE as [B, H, S]; merge uses [B, S, H] with out [B, S, H, D].
+            block_lse = block_lse_bh.transpose(1, 2).contiguous()
+            if step == 0:
+                out.copy_(block_out)
+                lse.copy_(block_lse)
+            else:
+                self._update_out_and_lse(out, lse, block_out, block_lse)
+            if step < self.world_size - 1:
+                self._ring_wait()
+        if out.dtype != q.dtype:
+            return out.to(dtype=q.dtype)
+        return out
+
+    @property
+    def preferred_layout(self) -> AttentionTensorLayout:
+        return self._preferred_layout
+
+    @classmethod
+    def support_fused_qkv(cls) -> bool:
+        return False
+
+    @classmethod
+    def support_lse(cls) -> bool:
+        return False