From 2adad54b26cda8db4f3a4c632cc795003a81a7e4 Mon Sep 17 00:00:00 2001 From: Venmugil Elango <498703+venmugil@users.noreply.github.com> Date: Thu, 16 Apr 2026 15:30:09 -0700 Subject: [PATCH 01/13] Add multi-node support for VisualGen diffusion workers Detect torchrun/SLURM external launchers via environment variables (RANK/WORLD_SIZE and SLURM_PROCID/SLURM_NTASKS). Rank 0 acts as coordinator; ranks 1..N-1 run as pure workers and exit after completion. Adds local_rank parameter to run_diffusion_worker for correct per-node GPU device assignment in multi-node runs. Signed-off-by: Venmugil Elango <498703+venmugil@users.noreply.github.com> --- tensorrt_llm/visual_gen/visual_gen.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 85f0ed88ae6d..23c8097ca5d2 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -144,7 +144,6 @@ def __init__( args: VisualGenArgs, ): self.args = args - self.n_workers = args.parallel.n_workers # --- Detect external launcher (torchrun / srun) --- ext = _detect_external_launch() From eae2ee622a56f63d94ef19d64a28c9f0ea35ee8e Mon Sep 17 00:00:00 2001 From: Venmugil Elango <498703+venmugil@users.noreply.github.com> Date: Fri, 17 Apr 2026 15:32:59 -0700 Subject: [PATCH 02/13] Address PR comments on multi-node VisualGen support - Raise RuntimeError when MASTER_ADDR is unset in torchrun path (mirrors SLURM check) - Validate world_size == n_workers on all ranks before worker launch - Replace master_port+1/+2 ZMQ port derivation with find_free_port() on rank 0 - Make request_queue_addr/response_queue_addr Optional[str]; pass None for non-zero ranks - Store n_workers on DiffusionRemoteClient; use self.n_workers consistently - Pass log_level to rank-0 external worker thread for consistency Signed-off-by: Venmugil Elango <498703+venmugil@users.noreply.github.com> --- tensorrt_llm/visual_gen/visual_gen.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 23c8097ca5d2..85f0ed88ae6d 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -144,6 +144,7 @@ def __init__( args: VisualGenArgs, ): self.args = args + self.n_workers = args.parallel.n_workers # --- Detect external launcher (torchrun / srun) --- ext = _detect_external_launch() From 7fae1d1bfa0bbdb976260a46cf2af4114c95b726 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 21 Apr 2026 13:12:21 -0700 Subject: [PATCH 03/13] add ring attention Signed-off-by: Shreyas Misra --- .../visual_gen/attention_backend/__init__.py | 3 +- .../attention_backend/flash_attn4.py | 4 + .../visual_gen/attention_backend/parallel.py | 144 ++++++++++++++++++ tensorrt_llm/_torch/visual_gen/config.py | 4 +- .../_torch/visual_gen/modules/attention.py | 17 ++- 5 files changed, 163 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index df07ed5f3276..1391a65a4259 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/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", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index 8981faa0de11..7060602208f7 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -180,3 +180,7 @@ def preferred_layout(self) -> AttentionTensorLayout: @classmethod def support_fused_qkv(cls) -> bool: return False + + @classmethod + def supports_lse(cls) -> bool: + return True diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index 00e0d5cd8415..ee85fed1e93d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -23,6 +23,8 @@ 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 @@ -439,3 +441,145 @@ 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 True + + +class RingComm: + def __init__(self, pg: Optional[dist.ProcessGroup]): + self.pg = pg + self.world_size = dist.get_world_size(group=pg) + self.rank = dist.get_rank(group=pg) + self.send_rank = dist.get_global_rank(pg, (self.rank + 1) % self.world_size) + self.recv_rank = dist.get_global_rank(pg, (self.rank - 1) % self.world_size) + self._reqs = [] + + def send_recv(self, send: torch.Tensor, recv: torch.Tensor) -> None: + ops = [] + if self.rank % 2 == 0: + ops.append(dist.P2POp(dist.isend, send, self.send_rank, group=self.pg)) + ops.append(dist.P2POp(dist.irecv, recv, self.recv_rank, group=self.pg)) + else: + ops.append(dist.P2POp(dist.irecv, recv, self.recv_rank, group=self.pg)) + ops.append(dist.P2POp(dist.isend, send, self.send_rank, group=self.pg)) + self._reqs = dist.batch_isend_irecv(ops) + + def wait(self) -> None: + for r in self._reqs: + r.wait() + self._reqs.clear() + + +def _update_out_and_lse( + 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)) + + +class RingAttention(AttentionBackend): + """Ring sequence parallelism around an LSE-capable attention backend.""" + + def __init__( + self, + inner_backend: AttentionBackend, + process_group: Optional[dist.ProcessGroup] = None, + ): + if not type(inner_backend).support_lse(): + raise ValueError( + f"RingAttention requires an LSE-capable inner backend (FA4); " + f"got {type(inner_backend).__name__}" + ) + self.inner = inner_backend + self.pg = process_group + try: + self.world_size = dist.get_world_size(group=process_group) + except (RuntimeError, ValueError): + self.world_size = 1 + 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 + + self._comm = RingComm(process_group) if self.world_size > 1 else None + self._buf_key = None + self._kv_bufs = None + self._out_buf = None + self._lse_buf = None + + 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 = torch.empty(2, 2, B, S, H_kv, D, device=k.device, dtype=k.dtype) + self._out_buf = torch.empty(B, S, H, D, device=q.device, dtype=q.dtype) + self._lse_buf = torch.empty(B, S, H, device=q.device, dtype=torch.float32) + self._buf_key = key + + 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 or single-rank groups. + if self.world_size == 1 or 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.CAUSAL: + raise NotImplementedError("Causal ring attention not implemented yet.") + + 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 + assert self._comm is not None + + 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._comm.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: + _update_out_and_lse(out, lse, block_out, block_lse) + if step < self.world_size - 1: + self._comm.wait() + 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 True diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index b339e197c302..2029fa07f912 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -186,10 +186,10 @@ class ParallelConfig(StrictBaseModel): parallel_vae_split_dim: Literal["width", "height"] = "width" # DiT Parallelism - dit_dp_size: int = PydanticField(1, ge=1) + dit_dp_size: int = PydanticField(1, ge=1) # Not yet supported dit_tp_size: int = PydanticField(1, ge=1) # Not yet supported dit_ulysses_size: int = PydanticField(1, ge=1) # Supported - dit_ring_size: int = PydanticField(1, ge=1) # Not yet supported + dit_ring_size: int = PydanticField(1, ge=1) # Supported dit_attn2d_row_size: int = PydanticField(1, ge=1) # Supported dit_attn2d_col_size: int = PydanticField(1, ge=1) # Supported dit_cfg_size: int = PydanticField(1, ge=1) # Supported diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 815bef713cee..abb6f3254155 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -74,6 +74,7 @@ def __init__( # Select compute backend (orthogonal to parallelism) vgm = config.visual_gen_mapping + ring_size = vgm.ring_size if vgm else 1 ulysses_size = vgm.ulysses_size if vgm else 1 base_backend = config.attention.backend @@ -161,13 +162,17 @@ def __init__( row_process_group=vgm.attn2d_row_group, col_process_group=vgm.attn2d_col_group, ) - elif use_ulysses: - from ..attention_backend.parallel import UlyssesAttention + else: + # Wrap with parallelism strategies (orthogonal to backend choice) + if (ring_size > 1 or ulysses_size > 1) and self.qkv_mode != QKVMode.SEPARATE_QKV: + if ring_size > 1: + from ..attention_backend.parallel import RingAttention - self.attn = UlyssesAttention( - inner_backend=self.attn, - process_group=vgm.ulysses_group, - ) + self.attn = RingAttention(self.attn, process_group=vgm.ring_group) + if ulysses_size > 1: + from ..attention_backend.parallel import UlyssesAttention + + self.attn = UlyssesAttention(self.attn, process_group=vgm.ulysses_group) def _init_qkv_proj(self) -> None: if self.qkv_mode == QKVMode.FUSE_QKV: From db554761fe2bb1f972e4855dbcb0ca920d2da85c Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 22 Apr 2026 12:04:23 -0700 Subject: [PATCH 04/13] fix wan transformer to consider ring size Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/visual_gen/mapping.py | 64 ++++++++++++++++++- .../visual_gen/models/wan/transformer_wan.py | 12 +++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/mapping.py b/tensorrt_llm/_torch/visual_gen/mapping.py index 7bd58e0dbb17..cbbd6b5c7b09 100644 --- a/tensorrt_llm/_torch/visual_gen/mapping.py +++ b/tensorrt_llm/_torch/visual_gen/mapping.py @@ -35,7 +35,7 @@ class VisualGenMapping(DeviceMeshTopologyImpl): Parallelism Hierarchy: total_workers = cfg × sp sp (sequence parallelism) = cp × ulysses [mutually exclusive today; TODO to combine] - cp (context parallelism) = ring [ring attention, not yet implemented] + cp (context parallelism) = ring [ring attention, 1D cp group] | attn2d [Attention2D 2D mesh, row_size × col_size] cfg: Splits positive/negative CFG prompts across GPUs (independent streams). @@ -57,6 +57,10 @@ class VisualGenMapping(DeviceMeshTopologyImpl): innermost / most contiguous). """ + # Flattened (cp, ulysses) mesh, cached after build_mesh(). Shared + # across instances via the class object. + seq_mesh = None + def __init__( self, world_size: int, @@ -161,6 +165,22 @@ def build_mesh(self): f"shape={shape}, mesh={cls.device_mesh}" ) + # Combined sequence-parallel mesh (cp × ulysses) for token sharding (e.g. WAN). + # ``_flatten`` is collective; every rank must call with the same dim names/order. + # Requires ``cp`` and ``ulysses`` adjacent (default ``cfg-tp-cp-ulysses``). + if self.ring_size * self.ulysses_size > 1: + cp_idx = self._dim_names.index("cp") + uly_idx = self._dim_names.index("ulysses") + if abs(cp_idx - uly_idx) != 1: + raise ValueError( + "seq_group construction requires cp and ulysses dims to be adjacent " + f"in the mesh, got order={self._order}" + ) + # Linearisation: seq_rank = cp_rank * ulysses_size + ulysses_rank (cp outer). + VisualGenMapping.seq_mesh = cls.device_mesh["cp", "ulysses"]._flatten( + mesh_dim_name="seq" + ) + if self.attn2d_row_size * self.attn2d_col_size > 1: self._build_attn2d_groups() @@ -288,6 +308,26 @@ def cp_rank(self) -> int: def ulysses_rank(self) -> int: return self._local_rank("ulysses") + # Combined sequence-parallel dimension: cp × ulysses (cp is ring when + # Attention2D is off). This is the dimension along which the *input token + # sequence* is sharded. Ring and Ulysses then re-shard internally, but from + # the transformer's point of view there is one group of size + # ``ring_size * ulysses_size``. + @property + def seq_size(self) -> int: + return self.ring_size * self.ulysses_size + + @property + def ring_rank(self) -> int: + """Rank along the CP (ring) mesh axis; ``cp_rank`` when ring uses the ``cp`` dim.""" + return self.cp_rank + + @property + def seq_rank(self) -> int: + # Must match the linearisation used by ``_flatten("cp", "ulysses")``: + # cp is outer (slower-varying), ulysses is inner (faster-varying). + return self.ring_rank * self.ulysses_size + self.ulysses_rank + @property def attn2d_mesh_rank(self) -> int: """Rank within the Attention2D CP group (same as cp_rank).""" @@ -324,6 +364,13 @@ def tp_group_pg(self) -> Optional[ProcessGroup]: def cfg_group(self) -> Optional[ProcessGroup]: return self._group("cfg") + @property + def ring_group(self) -> Optional[ProcessGroup]: + """Process group for RingAttention (1D ``cp`` mesh when ``ring_size > 1``).""" + if self.ring_size <= 1: + return None + return self.cp_group + @property def attn2d_row_group(self) -> Optional[ProcessGroup]: return self._attn2d_row_group @@ -348,6 +395,21 @@ def vae_group(self) -> Optional[ProcessGroup]: @property def vae_adj_groups(self) -> list[Optional[ProcessGroup]]: return self._vae_adj_groups + + def seq_group(self) -> Optional[ProcessGroup]: + """Process group spanning (cp × ulysses) for combined sequence-axis sharding.""" + cls = DeviceMeshTopologyImpl + if cls.device_mesh is None: + if self.world_size == 1: + return SingleProcessGroup.get_group() + return None + if self.ring_size * self.ulysses_size == 1: + # Degenerate: single rank along both dims. Fall back to the + # ulysses group (equivalent at size-1) to keep call sites simple. + return self._group("ulysses") + if VisualGenMapping.seq_mesh is None: + return None + return VisualGenMapping.seq_mesh.get_group() # ------------------------------------------------------------------ # Bridge to LLM Mapping (for Linear layers) diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index eca63b945807..f69c92328cc8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -483,7 +483,6 @@ def __init__( f"num_attention_heads ({num_heads}) must be divisible by " f"ulysses_size ({ulysses_size})" ) - if use_attn2d: self.use_seq_parallel = True self.seq_parallel_size = attn2d_mesh_size @@ -494,6 +493,13 @@ def __init__( self.seq_parallel_size = ulysses_size self.seq_parallel_pg = vgm.ulysses_group self.seq_parallel_rank = vgm.ulysses_rank + elif vgm is not None and vgm.ring_size > 1: + # Ring-only: shard tokens once over (cp × ulysses) == ring group so + # attention’s ring pass matches sequence length per rank. + self.use_seq_parallel = True + self.seq_parallel_size = vgm.seq_size + self.seq_parallel_pg = vgm.seq_group + self.seq_parallel_rank = vgm.seq_rank else: self.use_seq_parallel = False self.seq_parallel_size = 1 @@ -650,9 +656,9 @@ def forward( **kwargs, ): """ - Forward pass with optional parallelism (Ulysses head-sharding or Attention2D context parallelism). + Forward pass with optional parallelism (Ulysses, Ring, or Attention2D). - With parallelism enabled (seq_parallel_size > 1): + With sequence-axis sharding enabled (``seq_parallel_size > 1``): 1. Shard input sequence across ranks: [B, S] -> [B, S/P] 2. Each block's attention handles communication internally 3. Gather output sequence: [B, S/P] -> [B, S] From 31d9174a6a97204a46547a1014e1ad941f70dd73 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 30 Apr 2026 14:24:48 -0700 Subject: [PATCH 05/13] expose arg to example Signed-off-by: Shreyas Misra --- examples/visual_gen/visual_gen_wan_t2v.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index f1170895d41b..5569131b2bef 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -239,6 +239,7 @@ 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 parallel size") parser.add_argument( "--parallel_vae_size", type=int, @@ -326,11 +327,18 @@ def main(): attn2d_size = args.attn2d_row_size * args.attn2d_col_size if attn2d_size > 1 and args.ulysses_size > 1: raise ValueError( - "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented." + "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented.") + + if args.ulysses_size > 1: + 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: - 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 +378,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, From 17a099079dfffae8c16aede24ee9e19e40093d6e Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 4 May 2026 13:35:06 -0700 Subject: [PATCH 06/13] redesign seq sharding; make it work with attn2d Signed-off-by: Shreyas Misra --- .../visual_gen/attention_backend/interface.py | 1 + .../visual_gen/attention_backend/parallel.py | 159 +++++++------ tensorrt_llm/_torch/visual_gen/mapping.py | 60 ++--- .../models/flux/transformer_flux.py | 70 +----- .../models/flux/transformer_flux2.py | 75 +------ .../models/ltx2/transformer_ltx2.py | 171 ++++++-------- .../visual_gen/models/wan/transformer_wan.py | 82 ++----- tensorrt_llm/_torch/visual_gen/utils.py | 212 ++++++++++++++++++ 8 files changed, 439 insertions(+), 391 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py b/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py index 80a508c6f1f9..7a866acf5879 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py +++ b/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 diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index ee85fed1e93d..01872832087d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -20,8 +20,6 @@ """ -from typing import Optional - import torch import torch.distributed as dist import torch.nn.functional as F @@ -67,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 @@ -77,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 @@ -98,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) @@ -111,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 @@ -134,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] @@ -175,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 @@ -444,44 +449,7 @@ def support_fused_qkv(cls) -> bool: @classmethod def support_lse(cls) -> bool: - return True - - -class RingComm: - def __init__(self, pg: Optional[dist.ProcessGroup]): - self.pg = pg - self.world_size = dist.get_world_size(group=pg) - self.rank = dist.get_rank(group=pg) - self.send_rank = dist.get_global_rank(pg, (self.rank + 1) % self.world_size) - self.recv_rank = dist.get_global_rank(pg, (self.rank - 1) % self.world_size) - self._reqs = [] - - def send_recv(self, send: torch.Tensor, recv: torch.Tensor) -> None: - ops = [] - if self.rank % 2 == 0: - ops.append(dist.P2POp(dist.isend, send, self.send_rank, group=self.pg)) - ops.append(dist.P2POp(dist.irecv, recv, self.recv_rank, group=self.pg)) - else: - ops.append(dist.P2POp(dist.irecv, recv, self.recv_rank, group=self.pg)) - ops.append(dist.P2POp(dist.isend, send, self.send_rank, group=self.pg)) - self._reqs = dist.batch_isend_irecv(ops) - - def wait(self) -> None: - for r in self._reqs: - r.wait() - self._reqs.clear() - - -def _update_out_and_lse( - 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)) + return False class RingAttention(AttentionBackend): @@ -490,30 +458,74 @@ class RingAttention(AttentionBackend): def __init__( self, inner_backend: AttentionBackend, - process_group: Optional[dist.ProcessGroup] = None, + process_group: dist.ProcessGroup, ): + # Invariant: only instantiated when ring_size > 1 (see attention.py), + # so distributed must be initialized and the group must be non-trivial. 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 - try: - self.world_size = dist.get_world_size(group=process_group) - except (RuntimeError, ValueError): - self.world_size = 1 + 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 - self._comm = RingComm(process_group) if self.world_size > 1 else None + # 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] @@ -525,6 +537,18 @@ def _ensure_buffers(self, q: torch.Tensor, k: torch.Tensor) -> None: self._lse_buf = torch.empty(B, S, H, device=q.device, 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, @@ -534,8 +558,8 @@ def forward( attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, **kwargs, ) -> torch.Tensor: - # Bypass ring for cross-attention or single-rank groups. - if self.world_size == 1 or k.shape[1] != q.shape[1]: + # 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.CAUSAL: raise NotImplementedError("Causal ring attention not implemented yet.") @@ -546,14 +570,13 @@ def forward( kv_bufs = self._kv_bufs out = self._out_buf lse = self._lse_buf - assert self._comm is not None 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._comm.send_recv(kv_bufs[cur], kv_bufs[nxt]) + 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], @@ -567,9 +590,9 @@ def forward( out.copy_(block_out) lse.copy_(block_lse) else: - _update_out_and_lse(out, lse, block_out, block_lse) + self._update_out_and_lse(out, lse, block_out, block_lse) if step < self.world_size - 1: - self._comm.wait() + self._ring_wait() return out @property @@ -582,4 +605,4 @@ def support_fused_qkv(cls) -> bool: @classmethod def support_lse(cls) -> bool: - return True + return False diff --git a/tensorrt_llm/_torch/visual_gen/mapping.py b/tensorrt_llm/_torch/visual_gen/mapping.py index cbbd6b5c7b09..368ce0c5ce82 100644 --- a/tensorrt_llm/_torch/visual_gen/mapping.py +++ b/tensorrt_llm/_torch/visual_gen/mapping.py @@ -34,9 +34,10 @@ class VisualGenMapping(DeviceMeshTopologyImpl): Parallelism Hierarchy: total_workers = cfg × sp - sp (sequence parallelism) = cp × ulysses [mutually exclusive today; TODO to combine] + sp (sequence parallelism) = cp × ulysses cp (context parallelism) = ring [ring attention, 1D cp group] - | attn2d [Attention2D 2D mesh, row_size × col_size] + | attn2d [Attention2D 2D mesh, row_size × col_size. + Cannot be combined with Ulysses right now.] cfg: Splits positive/negative CFG prompts across GPUs (independent streams). tp: Tensor parallelism all-reduce within tp groups. @@ -82,19 +83,23 @@ def __init__( "Ring and Attention2D are mutually exclusive: both shard the sequence " f"dimension. Got ring_size={ring_size}, attn2d={attn2d_row_size}x{attn2d_col_size}." ) - cp_size = attn2d_size if attn2d_size > 1 else ring_size - if cp_size > 1 and ulysses_size > 1: - raise NotImplementedError( - "Combining CP and Ulysses is not yet supported. " - "They are orthogonal (CP shards sequence; Ulysses shards heads) " - "but the combined wrapper is not implemented." - ) - if attn2d_size > 1 and tp_size > 1: - raise NotImplementedError( - "Combining Attention2D and TP is not yet supported. " - "The row/col group construction in _build_attn2d_groups does not account " - "for TP ranks." - ) + if attn2d_size > 1: + cp_size = attn2d_size + if ulysses_size > 1: + raise NotImplementedError( + "Combining Attention2D and Ulysses is not yet supported. " + "They are orthogonal (Attention2D shards sequence; Ulysses shards heads) " + "but the combined wrapper is not implemented." + ) + if tp_size > 1: + raise NotImplementedError( + "Combining Attention2D and TP is not yet supported. " + "The row/col group construction in _build_attn2d_groups does not account " + "for TP ranks." + ) + else: + cp_size = ring_size + product = cfg_size * tp_size * cp_size * ulysses_size if product != world_size: raise ValueError( @@ -168,7 +173,7 @@ def build_mesh(self): # Combined sequence-parallel mesh (cp × ulysses) for token sharding (e.g. WAN). # ``_flatten`` is collective; every rank must call with the same dim names/order. # Requires ``cp`` and ``ulysses`` adjacent (default ``cfg-tp-cp-ulysses``). - if self.ring_size * self.ulysses_size > 1: + if self.cp_size * self.ulysses_size > 1: cp_idx = self._dim_names.index("cp") uly_idx = self._dim_names.index("ulysses") if abs(cp_idx - uly_idx) != 1: @@ -312,26 +317,25 @@ def ulysses_rank(self) -> int: # Attention2D is off). This is the dimension along which the *input token # sequence* is sharded. Ring and Ulysses then re-shard internally, but from # the transformer's point of view there is one group of size - # ``ring_size * ulysses_size``. + # ``cp_size * ulysses_size``. @property def seq_size(self) -> int: - return self.ring_size * self.ulysses_size - - @property - def ring_rank(self) -> int: - """Rank along the CP (ring) mesh axis; ``cp_rank`` when ring uses the ``cp`` dim.""" - return self.cp_rank + return self.cp_size * self.ulysses_size @property def seq_rank(self) -> int: # Must match the linearisation used by ``_flatten("cp", "ulysses")``: # cp is outer (slower-varying), ulysses is inner (faster-varying). - return self.ring_rank * self.ulysses_size + self.ulysses_rank + return self.cp_rank * self.ulysses_size + self.ulysses_rank + + @property + def ring_rank(self) -> int: + return self.cp_rank @property def attn2d_mesh_rank(self) -> int: """Rank within the Attention2D CP group (same as cp_rank).""" - return self._local_rank("cp") + return self.cp_rank @property def is_cfg_conditional(self) -> bool: @@ -367,8 +371,6 @@ def cfg_group(self) -> Optional[ProcessGroup]: @property def ring_group(self) -> Optional[ProcessGroup]: """Process group for RingAttention (1D ``cp`` mesh when ``ring_size > 1``).""" - if self.ring_size <= 1: - return None return self.cp_group @property @@ -382,7 +384,7 @@ def attn2d_col_group(self) -> Optional[ProcessGroup]: @property def attn2d_mesh_group(self) -> Optional[ProcessGroup]: """Full CP group for Attention2D (same as cp_group).""" - return self._group("cp") + return self.cp_group @property def vae_ranks(self) -> list[int]: @@ -403,7 +405,7 @@ def seq_group(self) -> Optional[ProcessGroup]: if self.world_size == 1: return SingleProcessGroup.get_group() return None - if self.ring_size * self.ulysses_size == 1: + if self.cp_size * self.ulysses_size == 1: # Degenerate: single rank along both dims. Fall back to the # ulysses group (equivalent at size-1) to keep call sites simple. return self._group("ulysses") diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py index 02cf57b77930..94c3dcf2ffd5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py @@ -34,6 +34,7 @@ from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention from tensorrt_llm._torch.visual_gen.models.flux.pos_embed_flux import FluxPosEmbed from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig # HF checkpoint key → our module attribute name @@ -565,34 +566,7 @@ def __init__(self, model_config: DiffusionModelConfig): vgm = model_config.visual_gen_mapping num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 24) - attn2d_row_size = vgm.attn2d_row_size if vgm else 1 - attn2d_col_size = vgm.attn2d_col_size if vgm else 1 - attn2d_mesh_size = attn2d_row_size * attn2d_col_size - ulysses_size = vgm.ulysses_size if vgm else 1 - use_attn2d = attn2d_mesh_size > 1 - use_ulysses = ulysses_size > 1 - - if use_ulysses and num_heads % ulysses_size != 0: - raise ValueError( - f"num_attention_heads ({num_heads}) must be divisible by " - f"ulysses_size ({ulysses_size})" - ) - - if use_attn2d: - self.use_seq_parallel = True - self.seq_parallel_size = attn2d_mesh_size - self.seq_parallel_pg = vgm.attn2d_mesh_group - self.seq_parallel_rank = vgm.attn2d_mesh_rank - elif use_ulysses: - self.use_seq_parallel = True - self.seq_parallel_size = ulysses_size - self.seq_parallel_pg = vgm.ulysses_group - self.seq_parallel_rank = vgm.ulysses_rank - else: - self.use_seq_parallel = False - self.seq_parallel_size = 1 - self.seq_parallel_pg = None - self.seq_parallel_rank = 0 + self.sharder = SequenceSharder.from_vgm(vgm, num_attention_heads=num_heads) # Extract pretrained config from model_config pretrained_config = model_config.pretrained_config @@ -828,33 +802,11 @@ def forward( if img_ids.ndim == 3: img_ids = img_ids[0] - # Shard sequences and position IDs before RoPE - if self.use_seq_parallel: - img_seq_len = img_ids.shape[0] - txt_seq_len = txt_ids.shape[0] - - if img_seq_len % self.seq_parallel_size != 0: - raise ValueError( - f"Image seq len ({img_seq_len}) not divisible by " - f"seq_parallel_size ({self.seq_parallel_size})" - ) - if txt_seq_len % self.seq_parallel_size != 0: - raise ValueError( - f"Text seq len ({txt_seq_len}) not divisible by " - f"seq_parallel_size ({self.seq_parallel_size})" - ) - - img_chunk = img_seq_len // self.seq_parallel_size - txt_chunk = txt_seq_len // self.seq_parallel_size - r = self.seq_parallel_rank - - # Shard position IDs (before RoPE computation) - img_ids = img_ids[r * img_chunk : (r + 1) * img_chunk] - txt_ids = txt_ids[r * txt_chunk : (r + 1) * txt_chunk] - - # Shard hidden states - hidden_states = hidden_states[:, r * img_chunk : (r + 1) * img_chunk, :] - encoder_hidden_states = encoder_hidden_states[:, r * txt_chunk : (r + 1) * txt_chunk, :] + # Shard sequences and position IDs before RoPE (no-op when sharder is inactive). + img_ids = self.sharder.shard(img_ids, dim=0) + txt_ids = self.sharder.shard(txt_ids, dim=0) + hidden_states = self.sharder.shard(hidden_states, dim=1) + encoder_hidden_states = self.sharder.shard(encoder_hidden_states, dim=1) # Compute RoPE embeddings (from potentially sharded IDs) ids = torch.cat((txt_ids, img_ids), dim=0) @@ -880,12 +832,8 @@ def forward( joint_attention_kwargs=joint_attention_kwargs, ) - # Gather output sequence from all ranks - if self.use_seq_parallel: - hidden_states = hidden_states.contiguous() - gathered = [torch.zeros_like(hidden_states) for _ in range(self.seq_parallel_size)] - torch.distributed.all_gather(gathered, hidden_states, group=self.seq_parallel_pg) - hidden_states = torch.cat(gathered, dim=1) + # All-gather hidden states across ranks (no-op when sharder is inactive). + hidden_states = self.sharder.gather(hidden_states, dim=1) # Output projection hidden_states = self.norm_out(hidden_states, temb) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py index 772a0532ff41..d276089a23ba 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py @@ -38,6 +38,7 @@ _remap_checkpoint_keys, ) from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig # HF FLUX.2 uses Flux2FeedForward with linear_in/linear_out attribute names. @@ -431,34 +432,7 @@ def __init__(self, model_config: DiffusionModelConfig): vgm = model_config.visual_gen_mapping num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 48) - attn2d_row_size = vgm.attn2d_row_size if vgm else 1 - attn2d_col_size = vgm.attn2d_col_size if vgm else 1 - attn2d_mesh_size = attn2d_row_size * attn2d_col_size - ulysses_size = vgm.ulysses_size if vgm else 1 - use_attn2d = attn2d_mesh_size > 1 - use_ulysses = ulysses_size > 1 - - if use_ulysses and num_heads % ulysses_size != 0: - raise ValueError( - f"num_attention_heads ({num_heads}) must be divisible by " - f"ulysses_size ({ulysses_size})" - ) - - if use_attn2d: - self.use_seq_parallel = True - self.seq_parallel_size = attn2d_mesh_size - self.seq_parallel_pg = vgm.attn2d_mesh_group - self.seq_parallel_rank = vgm.attn2d_mesh_rank - elif use_ulysses: - self.use_seq_parallel = True - self.seq_parallel_size = ulysses_size - self.seq_parallel_pg = vgm.ulysses_group - self.seq_parallel_rank = vgm.ulysses_rank - else: - self.use_seq_parallel = False - self.seq_parallel_size = 1 - self.seq_parallel_pg = None - self.seq_parallel_rank = 0 + self.sharder = SequenceSharder.from_vgm(vgm, num_attention_heads=num_heads) # Extract pretrained config from model_config (following WAN/FLUX.1 pattern) pretrained_config = model_config.pretrained_config @@ -731,36 +705,13 @@ def forward( if img_ids.ndim == 3: img_ids = img_ids[0] - # Shard sequences and position IDs before RoPE - if self.use_seq_parallel: - img_seq_len = img_ids.shape[0] - _txt_seq_len = txt_ids.shape[0] - - if img_seq_len % self.seq_parallel_size != 0: - raise ValueError( - f"Image seq len ({img_seq_len}) not divisible by " - f"seq_parallel_size ({self.seq_parallel_size})" - ) - if _txt_seq_len % self.seq_parallel_size != 0: - raise ValueError( - f"Text seq len ({_txt_seq_len}) not divisible by " - f"seq_parallel_size ({self.seq_parallel_size})" - ) - - img_chunk = img_seq_len // self.seq_parallel_size - txt_chunk = _txt_seq_len // self.seq_parallel_size - r = self.seq_parallel_rank - - # Shard position IDs (before RoPE computation) - img_ids = img_ids[r * img_chunk : (r + 1) * img_chunk] - txt_ids = txt_ids[r * txt_chunk : (r + 1) * txt_chunk] - - # Shard hidden states - hidden_states = hidden_states[:, r * img_chunk : (r + 1) * img_chunk, :] - encoder_hidden_states = encoder_hidden_states[:, r * txt_chunk : (r + 1) * txt_chunk, :] - - # Update txt_seq_len to local (sharded) length for single-stream split - txt_seq_len = txt_chunk + # Shard sequences and position IDs before RoPE (no-op when sharder is inactive). + img_ids = self.sharder.shard(img_ids, dim=0) + txt_ids = self.sharder.shard(txt_ids, dim=0) + hidden_states = self.sharder.shard(hidden_states, dim=1) + encoder_hidden_states = self.sharder.shard(encoder_hidden_states, dim=1) + # Update txt_seq_len to local (sharded) length for single-stream split. + txt_seq_len = encoder_hidden_states.shape[1] # Compute RoPE embeddings (4-axis, from potentially sharded IDs) ids = torch.cat([txt_ids, img_ids], dim=0) @@ -795,12 +746,8 @@ def forward( # Extract image features (discard text) hidden_states = hidden_states[:, txt_seq_len:, :] - # Gather output sequence from all ranks - if self.use_seq_parallel: - hidden_states = hidden_states.contiguous() - gathered = [torch.zeros_like(hidden_states) for _ in range(self.seq_parallel_size)] - torch.distributed.all_gather(gathered, hidden_states, group=self.seq_parallel_pg) - hidden_states = torch.cat(gathered, dim=1) + # All-gather hidden states across ranks (no-op when sharder is inactive). + hidden_states = self.sharder.gather(hidden_states, dim=1) # Output projection hidden_states = self.norm_out(hidden_states, temb) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index ae5b808ff2d6..774ec82cc872 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -23,7 +23,6 @@ from typing import TYPE_CHECKING, Any, Optional import torch -import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm @@ -33,6 +32,7 @@ from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -295,19 +295,12 @@ def __init__( self.idx = idx self.norm_eps = norm_eps - self._use_seq_parallel = False - self._seq_parallel_size = 1 - self._seq_parallel_pg = None - self._audio_is_sharded = False + # Per-block sharder: mirrors the root sharder's topology so the block + # can run cross-attention all-gathers independently. Head-divisibility + # is checked once at the root model — skip num_heads here. vgm = config.visual_gen_mapping if config is not None else None - if vgm is not None and vgm.ulysses_size > 1: - self._use_seq_parallel = True - self._seq_parallel_size = vgm.ulysses_size - self._seq_parallel_pg = vgm.ulysses_group - elif vgm is not None and vgm.attn2d_row_size * vgm.attn2d_col_size > 1: - self._use_seq_parallel = True - self._seq_parallel_size = vgm.attn2d_row_size * vgm.attn2d_col_size - self._seq_parallel_pg = vgm.attn2d_mesh_group + self._sharder = SequenceSharder.from_vgm(vgm) + self._audio_is_sharded = False if video is not None: self._init_video_modules(video, rope_type, norm_eps, config, idx) @@ -468,20 +461,19 @@ def _get_av_ca_ada_values( def _sp_all_gather(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor: """All-gather *x* along *dim* across sequence-parallel ranks.""" - x = x.contiguous() - gathered = [torch.empty_like(x) for _ in range(self._seq_parallel_size)] - dist.all_gather(gathered, x, group=self._seq_parallel_pg) - return torch.cat(gathered, dim=dim) + return self._sharder.gather(x, dim=dim) def _sp_gather_pe(self, pe): - """All-gather RoPE (cos, sin) tuple along the sequence dim.""" + """All-gather RoPE (cos, sin) tuple along its sequence dim. + + Split RoPE is ``[B, H, S, D]`` (dim 2); interleaved RoPE is + ``[B, S, D]`` (dim 1). Inferred from ``cos.ndim``. + """ if pe is None: return None cos, sin = pe - # Split RoPE: [B, H, S, D] — sequence at dim 2 - # Interleaved RoPE: [B, S, D] — sequence at dim 1 seq_dim = 2 if cos.ndim == 4 else 1 - return (self._sp_all_gather(cos, dim=seq_dim), self._sp_all_gather(sin, dim=seq_dim)) + return (self._sharder.gather(cos, dim=seq_dim), self._sharder.gather(sin, dim=seq_dim)) # -- Forward ------------------------------------------------------------- @@ -644,7 +636,7 @@ def forward( # Project-before-gather (video → audio direction). k_v2a, v_v2a = self.video_to_audio_attn.project_kv(vx_scaled) - if self._use_seq_parallel: + if self._sharder.is_active: k_v2a = self._sp_all_gather(k_v2a) v_v2a = self._sp_all_gather(v_v2a) k_pe_v2a = self._sp_gather_pe(video.cross_positional_embeddings) @@ -835,38 +827,23 @@ def __init__( self._init_preprocessors(cross_pe_max_pos) vgm = model_config.visual_gen_mapping - attn2d_row_size = vgm.attn2d_row_size if vgm else 1 - attn2d_col_size = vgm.attn2d_col_size if vgm else 1 - attn2d_size = attn2d_row_size * attn2d_col_size - ulysses_size = vgm.ulysses_size if vgm else 1 - - use_attn2d = attn2d_size > 1 - use_ulysses = ulysses_size > 1 - self.use_seq_parallel = use_attn2d or use_ulysses - self.seq_parallel_size = attn2d_size if use_attn2d else ulysses_size - if use_attn2d: - self.seq_parallel_pg = vgm.attn2d_mesh_group - self.seq_parallel_rank = vgm.attn2d_mesh_rank - elif use_ulysses: - self.seq_parallel_pg = vgm.ulysses_group if vgm else None - self.seq_parallel_rank = vgm.ulysses_rank if vgm else 0 - else: - self.seq_parallel_pg = None - self.seq_parallel_rank = 0 - - # Ulysses shards heads; Attention2D shards sequence — no head-count constraint for the latter. - if use_ulysses and model_type.is_video_enabled(): - if num_attention_heads % ulysses_size != 0: - raise ValueError( - f"num_attention_heads ({num_attention_heads}) " - f"must be divisible by dit_ulysses_size ({ulysses_size})" - ) - if use_ulysses and model_type.is_audio_enabled(): - if audio_num_attention_heads % ulysses_size != 0: - raise ValueError( - f"audio_num_attention_heads ({audio_num_attention_heads}) " - f"must be divisible by dit_ulysses_size ({ulysses_size})" - ) + # Validate video-head divisibility through the factory; audio heads + # are checked separately because the factory only accepts one count. + self._sharder = SequenceSharder.from_vgm( + vgm, + num_attention_heads=num_attention_heads if model_type.is_video_enabled() else None, + ) + if ( + self._sharder.is_active + and vgm is not None + and vgm.ulysses_size > 1 + and model_type.is_audio_enabled() + and audio_num_attention_heads % vgm.ulysses_size != 0 + ): + raise ValueError( + f"audio_num_attention_heads ({audio_num_attention_heads}) " + f"must be divisible by dit_ulysses_size ({vgm.ulysses_size})" + ) self._audio_is_sharded = False self._cache_dit_video_args: Optional[TransformerArgs] = None @@ -1163,46 +1140,31 @@ def _init_transformer_blocks( # -- Sequence sharding / gathering ---------------------------------------- def _shard_transformer_args(self, args: TransformerArgs) -> TransformerArgs: - """Shard sequence-dependent fields of *args* across sequence-parallel ranks.""" - seq_len = args.x.shape[1] - chunk = seq_len // self.seq_parallel_size - s = self.seq_parallel_rank * chunk - e = s + chunk - - def _shard(t): - if t is None or t.ndim < 2 or t.shape[1] != seq_len: - return t - return t[:, s:e] - - def _shard_pe(pe): - if pe is None: - return None - cos, sin = pe - if cos.ndim == 4 and cos.shape[2] == seq_len: - # Split RoPE: [B, H, S, D] — sequence dim at index 2 - return (cos[:, :, s:e], sin[:, :, s:e]) - elif cos.ndim == 3 and cos.shape[1] == seq_len: - # Interleaved RoPE: [B, S, D] — sequence dim at index 1 - return (cos[:, s:e], sin[:, s:e]) - return pe + """Shard sequence-dependent fields of *args* across sequence-parallel ranks. + Fields whose dim-1 doesn't match ``args.x``'s sequence length are passed + through unchanged (broadcast-compatible scalars, etc.). + """ + seq_len = args.x.shape[1] + sh = self._sharder return replace( args, - x=args.x[:, s:e], - timesteps=_shard(args.timesteps), - embedded_timestep=_shard(args.embedded_timestep), - positional_embeddings=_shard_pe(args.positional_embeddings), - cross_positional_embeddings=_shard_pe(args.cross_positional_embeddings), - cross_scale_shift_timestep=_shard(args.cross_scale_shift_timestep), - cross_gate_timestep=_shard(args.cross_gate_timestep), + x=sh.shard(args.x, dim=1), + timesteps=sh.shard(args.timesteps, dim=1, expected_seq_len=seq_len), + embedded_timestep=sh.shard(args.embedded_timestep, dim=1, expected_seq_len=seq_len), + positional_embeddings=sh.shard_rope(args.positional_embeddings, seq_len=seq_len), + cross_positional_embeddings=sh.shard_rope( + args.cross_positional_embeddings, seq_len=seq_len + ), + cross_scale_shift_timestep=sh.shard( + args.cross_scale_shift_timestep, dim=1, expected_seq_len=seq_len + ), + cross_gate_timestep=sh.shard(args.cross_gate_timestep, dim=1, expected_seq_len=seq_len), ) def _gather_sequence(self, x: torch.Tensor) -> torch.Tensor: """All-gather hidden states along the sequence dim.""" - x = x.contiguous() - gathered = [torch.empty_like(x) for _ in range(self.seq_parallel_size)] - dist.all_gather(gathered, x, group=self.seq_parallel_pg) - return torch.cat(gathered, dim=1) + return self._sharder.gather(x, dim=1) def configure_audio_ulysses(self, audio_seq_len: int) -> None: """Configure whether audio uses Ulysses based on sequence length. @@ -1211,11 +1173,11 @@ def configure_audio_ulysses(self, audio_seq_len: int) -> None: known. The decision is cached — ``forward()`` uses it without re-checking. """ - if not self.use_seq_parallel: + if not self._sharder.is_active: self._audio_is_sharded = False return - self._audio_is_sharded = audio_seq_len % self.seq_parallel_size == 0 + self._audio_is_sharded = audio_seq_len % self._sharder.size == 0 for block in self.transformer_blocks: target = block.inner if isinstance(block, LTX2CacheDiTPattern0BlockWrapper) else block target._audio_is_sharded = self._audio_is_sharded @@ -1231,21 +1193,26 @@ def set_ulysses_enabled(self, enabled: bool) -> None: multi-rank operation; audio sharding will be reconfigured by the next :meth:`configure_audio_ulysses` call. """ - if self.seq_parallel_size <= 1: + if self._sharder.size <= 1: return - self.use_seq_parallel = enabled - if not enabled: + if enabled: + self._sharder.enable() + else: + self._sharder.disable() self._audio_is_sharded = False for block in self.transformer_blocks: - block._use_seq_parallel = enabled - if not enabled: - block._audio_is_sharded = False - if hasattr(block, "attn1"): - block.attn1.set_ulysses_active(enabled) - if hasattr(block, "audio_attn1") and not enabled: - block.audio_attn1.set_ulysses_active(False) + target = block.inner if isinstance(block, LTX2CacheDiTPattern0BlockWrapper) else block + if enabled: + target._sharder.enable() + else: + target._sharder.disable() + target._audio_is_sharded = False + if hasattr(target, "attn1"): + target.attn1.set_ulysses_active(enabled) + if hasattr(target, "audio_attn1") and not enabled: + target.audio_attn1.set_ulysses_active(False) # -- Output processing --------------------------------------------------- @@ -1361,10 +1328,10 @@ def forward( else None ) - # Shard sequences for parallelism (Ulysses head-sharding or Attention2D context parallelism). + # Shard sequences for parallelism (Ulysses head-sharding, ring CP, or Attention2D). # Video is always sharded. Audio sharding is decided once by # configure_audio_ulysses() and cached in self._audio_is_sharded. - if self.use_seq_parallel: + if self._sharder.is_active: if video_args is not None: video_args = self._shard_transformer_args(video_args) if self._audio_is_sharded and audio_args is not None: @@ -1406,7 +1373,7 @@ def forward( # Only gather embedded_timestep if it was actually sharded (dim-1 # matches x); scalar timestep embeddings [B, 1, D] are # broadcast-compatible and must not be gathered. - if self.use_seq_parallel: + if self._sharder.is_active: if video_args is not None: gathered_vx = self._gather_sequence(video_args.x) v_et = video_args.embedded_timestep diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index f69c92328cc8..efdd725c0bf1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -15,6 +15,7 @@ from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig @@ -471,40 +472,7 @@ def __init__( raise ValueError(f"WAN does not support tensor parallelism. Got tp_size={vgm.tp_size}") num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 12) - attn2d_row_size = vgm.attn2d_row_size if vgm else 1 - attn2d_col_size = vgm.attn2d_col_size if vgm else 1 - attn2d_mesh_size = attn2d_row_size * attn2d_col_size - ulysses_size = vgm.ulysses_size if vgm else 1 - use_attn2d = attn2d_mesh_size > 1 - use_ulysses = ulysses_size > 1 - - if use_ulysses and num_heads % ulysses_size != 0: - raise ValueError( - f"num_attention_heads ({num_heads}) must be divisible by " - f"ulysses_size ({ulysses_size})" - ) - if use_attn2d: - self.use_seq_parallel = True - self.seq_parallel_size = attn2d_mesh_size - self.seq_parallel_pg = vgm.attn2d_mesh_group - self.seq_parallel_rank = vgm.attn2d_mesh_rank - elif use_ulysses: - self.use_seq_parallel = True - self.seq_parallel_size = ulysses_size - self.seq_parallel_pg = vgm.ulysses_group - self.seq_parallel_rank = vgm.ulysses_rank - elif vgm is not None and vgm.ring_size > 1: - # Ring-only: shard tokens once over (cp × ulysses) == ring group so - # attention’s ring pass matches sequence length per rank. - self.use_seq_parallel = True - self.seq_parallel_size = vgm.seq_size - self.seq_parallel_pg = vgm.seq_group - self.seq_parallel_rank = vgm.seq_rank - else: - self.use_seq_parallel = False - self.seq_parallel_size = 1 - self.seq_parallel_pg = None - self.seq_parallel_rank = 0 + self.sharder = SequenceSharder.from_vgm(vgm, num_attention_heads=num_heads) config = model_config.pretrained_config @@ -656,12 +624,12 @@ def forward( **kwargs, ): """ - Forward pass with optional parallelism (Ulysses, Ring, or Attention2D). + Forward pass with optional sequence parallelism (Ring, Ulysses, or Attention2D). - With sequence-axis sharding enabled (``seq_parallel_size > 1``): - 1. Shard input sequence across ranks: [B, S] -> [B, S/P] - 2. Each block's attention handles communication internally - 3. Gather output sequence: [B, S/P] -> [B, S] + When sharder is active: + 1. Shard input sequence (and matching RoPE) across ranks: [B, S] -> [B, S/P] + 2. Each block's attention handles cross-rank communication internally + 3. All-gather output sequence: [B, S/P] -> [B, S] When TeaCache is enabled, TeaCacheHook intercepts and replaces this call. """ @@ -675,27 +643,12 @@ def forward( # Patchify and flatten: [B, C, T, H, W] -> [B, S, hidden_size] x = self.patch_embedding(hidden_states).flatten(2).transpose(1, 2) - # Shard sequence across ranks: [B, S] -> [B, S/P] - chunk_size = None - if self.use_seq_parallel: - seq_len = x.shape[1] - if seq_len % self.seq_parallel_size != 0: - raise ValueError( - f"Sequence length ({seq_len}) is not divisible by " - f"seq_parallel_size ({self.seq_parallel_size}). " - f"Adjust video dimensions or parallelism settings." - ) - - chunk_size = seq_len // self.seq_parallel_size - chunk_start = self.seq_parallel_rank * chunk_size - chunk_end = chunk_start + chunk_size - x = x[:, chunk_start:chunk_end, :] - - # Shard RoPE frequencies to match sequence sharding - # RoPE freqs shape: [B, S, ...], so shard along dim 1 (sequence dimension) - if freqs_cos is not None and freqs_sin is not None: - freqs_cos = freqs_cos[:, chunk_start:chunk_end] - freqs_sin = freqs_sin[:, chunk_start:chunk_end] + # Shard sequence + matching RoPE across ranks (no-op when sharder is inactive). + seq_len = x.shape[1] + x = self.sharder.shard(x, dim=1) + rope = self.sharder.shard_rope((freqs_cos, freqs_sin), seq_len=seq_len) + if rope is not None: + freqs_cos, freqs_sin = rope # Time and text/image embeddings # Timestep shape: [batch_size] or [batch_size, seq_len] @@ -748,13 +701,8 @@ def forward( freqs_sin, ) - # Gather sequence from all ranks: [B, S/P] -> [B, S] - if self.use_seq_parallel: - # Ensure tensor is contiguous before all_gather - x = x.contiguous() - x_list = [torch.zeros_like(x) for _ in range(self.seq_parallel_size)] - torch.distributed.all_gather(x_list, x, group=self.seq_parallel_pg) - x = torch.cat(x_list, dim=1) + # All-gather sequence from all ranks: [B, S/P] -> [B, S] (no-op when inactive). + x = self.sharder.gather(x, dim=1) # Output projection and unpatchify if temb.ndim == 3: diff --git a/tensorrt_llm/_torch/visual_gen/utils.py b/tensorrt_llm/_torch/visual_gen/utils.py index 0c96e93ff143..cff6c3757659 100644 --- a/tensorrt_llm/_torch/visual_gen/utils.py +++ b/tensorrt_llm/_torch/visual_gen/utils.py @@ -1,6 +1,14 @@ """Utility functions for visual generation pipelines.""" +from __future__ import annotations + +from typing import Optional, Tuple + import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping @torch.compile @@ -33,3 +41,207 @@ def postprocess_video_tensor(video: torch.Tensor) -> torch.Tensor: def as_tuple(x): return x if isinstance(x, tuple) else (x, x) + + +class SequenceSharder: + """Block-shard / all-gather a tensor along its sequence dimension. + + A single ``SequenceSharder`` collapses the per-model + ``if attn2d / elif ulysses / elif ring / else`` dispatch and the hand-rolled + shard / gather of hidden states + RoPE into one model-agnostic helper. + + Built from a :class:`VisualGenMapping` via :meth:`from_vgm`, the sharder + uses ``vgm.seq_size / seq_rank / seq_group`` so the same call sites work + uniformly for ring, attn2d, ulysses, and ring + ulysses. + + Models call ``shard(...)`` / ``gather(...)`` / ``shard_rope(...)`` directly; + when the sharder is inactive (``size == 1`` or runtime-disabled) every + method is a no-op pass-through so the call sites do not need an + ``if is_active`` guard. + + The sharder is intentionally model-agnostic: dimensions are passed + explicitly at every call site and no model-specific shape conventions + leak in (the sole exception is :meth:`shard_rope`, which infers the + seq axis from a ``seq_len`` argument). + """ + + def __init__(self, size: int, rank: int, group: Optional[ProcessGroup]): + self._size = size + self._rank = rank + self._group = group + self._enabled = size > 1 + + # ------------------------------------------------------------------ + # Factory + # ------------------------------------------------------------------ + @classmethod + def from_vgm( + cls, + vgm: Optional[VisualGenMapping], + *, + num_attention_heads: Optional[int] = None, + num_kv_heads: Optional[int] = None, + ) -> "SequenceSharder": + """Build a sharder from a :class:`VisualGenMapping`. + + Uses ``(cp_size * ulysses_size, seq_rank, seq_group)`` so the same + sharder works for ring, attn2d, ulysses, and ring + ulysses. + + Validates head divisibility only when Ulysses is part of the seq + group — ring and attn2d shard the sequence axis and have no + head-count constraint. + """ + if vgm is None: + return cls(size=1, rank=0, group=None) + + size = vgm.seq_size + rank = vgm.seq_rank + group = vgm.seq_group + + if size > 1 and vgm.ulysses_size > 1: + for label, count in ( + ("num_attention_heads", num_attention_heads), + ("num_kv_heads", num_kv_heads), + ): + if count is not None and count % vgm.ulysses_size != 0: + raise ValueError( + f"{label}={count} must be divisible by ulysses_size={vgm.ulysses_size}" + ) + + return cls(size=size, rank=rank, group=group) + + # ------------------------------------------------------------------ + # State + # ------------------------------------------------------------------ + @property + def is_active(self) -> bool: + return self._enabled and self._size > 1 + + @property + def size(self) -> int: + return self._size + + @property + def rank(self) -> int: + return self._rank + + @property + def group(self) -> Optional[ProcessGroup]: + return self._group + + def disable(self) -> None: + """Run as if ``size == 1``. + + Used by LTX2's stage-2 single-rank execution path where the + non-primary workers have already exited. + """ + self._enabled = False + + def enable(self) -> None: + """Re-enable sharding after :meth:`disable` (no-op if size == 1).""" + self._enabled = self._size > 1 + + # ------------------------------------------------------------------ + # Shard + # ------------------------------------------------------------------ + def shard( + self, + tensor: Optional[torch.Tensor], + dim: int = 1, + *, + expected_seq_len: Optional[int] = None, + pad_to_multiple: bool = False, + ) -> Optional[torch.Tensor]: + """Contiguous block-shard ``tensor`` along ``dim``. + + Returns ``tensor`` unchanged when: + * the sharder is inactive (``size == 1`` or runtime-disabled), + * ``tensor is None``, + * ``expected_seq_len`` is given and ``tensor.shape[dim]`` doesn't + match — used by LTX2 to skip dataclass fields whose seq axis + doesn't line up with the field being sharded. + + When ``pad_to_multiple`` is ``True``, the sequence dim is right-padded + with zeros to a multiple of ``size`` before sharding. The matching + :meth:`gather` call must then pass ``unpad_to`` to slice the padding + back off. + """ + if tensor is None or not self.is_active: + return tensor + + seq_len = tensor.shape[dim] + if expected_seq_len is not None and seq_len != expected_seq_len: + return tensor + + if pad_to_multiple and seq_len % self._size != 0: + pad = self._size - (seq_len % self._size) + pad_shape = list(tensor.shape) + pad_shape[dim] = pad + tensor = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim) + seq_len = tensor.shape[dim] + + if seq_len % self._size != 0: + raise ValueError( + f"Sequence length ({seq_len}) along dim {dim} is not " + f"divisible by SequenceSharder.size ({self._size}). " + f"Pass pad_to_multiple=True or adjust input dimensions." + ) + + chunk = seq_len // self._size + start = self._rank * chunk + idx = [slice(None)] * tensor.ndim + idx[dim] = slice(start, start + chunk) + return tensor[tuple(idx)] + + def shard_rope( + self, + rope: Optional[Tuple[torch.Tensor, torch.Tensor]], + seq_len: int, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Shard a ``(cos, sin)`` RoPE pair along whichever axis equals ``seq_len``. + + Handles both ``[B, S, D]`` (dim 1) and ``[B, H, S, D]`` (dim 2) + RoPE layouts that occur across visual_gen models. Returns the + pair unchanged when the seq axis cannot be inferred unambiguously + (e.g. ``seq_len == head_dim``); callers can fall back to explicit + ``shard(..., dim=...)`` calls in that case. + """ + if rope is None or not self.is_active: + return rope + + cos, sin = rope + candidates = [d for d, n in enumerate(cos.shape) if n == seq_len] + if len(candidates) != 1: + return rope + d = candidates[0] + return (self.shard(cos, dim=d), self.shard(sin, dim=d)) + + # ------------------------------------------------------------------ + # Gather + # ------------------------------------------------------------------ + def gather( + self, + tensor: torch.Tensor, + dim: int = 1, + *, + unpad_to: Optional[int] = None, + ) -> torch.Tensor: + """All-gather ``tensor`` along ``dim``. + + No-op when sharder is inactive. ``unpad_to`` slices the gathered + tensor's ``dim`` back to the given length; pair with + ``shard(..., pad_to_multiple=True)`` to round-trip through padding. + """ + if not self.is_active: + return tensor + + tensor = tensor.contiguous() + parts = [torch.empty_like(tensor) for _ in range(self._size)] + dist.all_gather(parts, tensor, group=self._group) + out = torch.cat(parts, dim=dim) + + if unpad_to is not None: + idx = [slice(None)] * out.ndim + idx[dim] = slice(0, unpad_to) + out = out[tuple(idx)] + return out From 641523bbe7c79819afdbb2bbe03891985cd1d1ae Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 6 May 2026 14:20:37 -0700 Subject: [PATCH 07/13] make dim ordering internal Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/visual_gen/config.py | 8 --- tensorrt_llm/_torch/visual_gen/mapping.py | 28 ++++----- .../visual_gen/models/wan/transformer_wan.py | 3 +- .../_torch/visual_gen/pipeline_loader.py | 1 - .../multi_gpu/test_visual_gen_mapping.py | 59 ++----------------- 5 files changed, 17 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 2029fa07f912..d39c2127afdb 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -23,7 +23,6 @@ from pydantic import BaseModel, ConfigDict, model_validator from pydantic import Field as PydanticField -from tensorrt_llm._torch.visual_gen.mapping import DEFAULT_DIM_ORDER from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.logger import logger @@ -194,13 +193,6 @@ class ParallelConfig(StrictBaseModel): dit_attn2d_col_size: int = PydanticField(1, ge=1) # Supported dit_cfg_size: int = PydanticField(1, ge=1) # Supported dit_fsdp_size: int = PydanticField(1, ge=1) - dit_dim_order: str = PydanticField( - DEFAULT_DIM_ORDER, - description=( - "Outermost-to-innermost ordering of parallelism axes for the " - "DeviceMesh. Innermost = most contiguous ranks." - ), - ) # Refiner Parallelism (Optional) refiner_dit_dp_size: int = 1 diff --git a/tensorrt_llm/_torch/visual_gen/mapping.py b/tensorrt_llm/_torch/visual_gen/mapping.py index 368ce0c5ce82..4360ec890f55 100644 --- a/tensorrt_llm/_torch/visual_gen/mapping.py +++ b/tensorrt_llm/_torch/visual_gen/mapping.py @@ -21,7 +21,8 @@ from tensorrt_llm.mapping import Mapping _VALID_DIM_NAMES = frozenset({"cfg", "tp", "cp", "ulysses"}) -DEFAULT_DIM_ORDER = "cfg-tp-cp-ulysses" +# Outermost-to-innermost for init_device_mesh +_DEVICE_MESH_DIM_ORDER = "cfg-tp-cp-ulysses" class VisualGenMapping(DeviceMeshTopologyImpl): @@ -47,15 +48,14 @@ class VisualGenMapping(DeviceMeshTopologyImpl): attn2d: 2D mesh; Q all-gathered within row group, K/V within col group. ulysses: Shards heads via all-to-all (head-sharding, not sequence-sharding). - Ordering rationale (default ``"cfg-tp-cp-ulysses"``): + Fixed mesh axis ordering (implementation detail, see + ``_DEVICE_MESH_DIM_ORDER``): - Ulysses innermost: all-to-all is latency-sensitive, contiguous ranks - CP next: KV streaming (ring) or sequence shard communication (Attention2D) - TP next: all-reduce for Linear - CFG outermost: independent until final all-gather - The *order* string maps directly to ``init_device_mesh``'s - ``mesh_shape`` tuple (first = outermost / slowest-varying, last = - innermost / most contiguous). + Callers should use rank and process-group properties only, not mesh layout. """ # Flattened (cp, ulysses) mesh, cached after build_mesh(). Shared @@ -73,7 +73,6 @@ def __init__( attn2d_row_size: int = 1, attn2d_col_size: int = 1, parallel_vae_size: int = 1, - order: str = DEFAULT_DIM_ORDER, ): # cp_size unifies ring and Attention2D under one context-parallelism mesh dimension. # Ring and Attention2D are mutually exclusive: both shard the sequence axis. @@ -113,12 +112,8 @@ def __init__( f"parallel_vae_size ({parallel_vae_size}) cannot exceed world_size ({world_size})" ) - dims = order.split("-") - if set(dims) != _VALID_DIM_NAMES or len(dims) != len(_VALID_DIM_NAMES): - raise ValueError( - f"order must be a '-'-separated permutation of " - f"{sorted(_VALID_DIM_NAMES)}, got '{order}'" - ) + dims = _DEVICE_MESH_DIM_ORDER.split("-") + assert set(dims) == _VALID_DIM_NAMES and len(dims) == len(_VALID_DIM_NAMES) self.world_size = world_size self._rank = rank @@ -135,7 +130,6 @@ def __init__( self._vae_adj_groups: list[Optional[ProcessGroup]] = [] self._attn2d_row_group: Optional[ProcessGroup] = None self._attn2d_col_group: Optional[ProcessGroup] = None - self._order = order self._dim_names = tuple(dims) # cp_size covers both ring (1D) and Attention2D (2D, row_size * col_size). # For Attention2D, _build_attn2d_groups() creates row/col sub-groups; @@ -172,14 +166,14 @@ def build_mesh(self): # Combined sequence-parallel mesh (cp × ulysses) for token sharding (e.g. WAN). # ``_flatten`` is collective; every rank must call with the same dim names/order. - # Requires ``cp`` and ``ulysses`` adjacent (default ``cfg-tp-cp-ulysses``). + # Requires ``cp`` and ``ulysses`` adjacent in ``_DEVICE_MESH_DIM_ORDER``. if self.cp_size * self.ulysses_size > 1: cp_idx = self._dim_names.index("cp") uly_idx = self._dim_names.index("ulysses") if abs(cp_idx - uly_idx) != 1: - raise ValueError( - "seq_group construction requires cp and ulysses dims to be adjacent " - f"in the mesh, got order={self._order}" + raise RuntimeError( + "seq_group requires cp and ulysses adjacent; " + f"fix _DEVICE_MESH_DIM_ORDER (got {self._dim_names!r})" ) # Linearisation: seq_rank = cp_rank * ulysses_size + ulysses_rank (cp outer). VisualGenMapping.seq_mesh = cls.device_mesh["cp", "ulysses"]._flatten( diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index efdd725c0bf1..7e02b883dd86 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -671,8 +671,7 @@ def forward( # batch_size, seq_len, 6, hidden_size temb_proj = temb_proj.unflatten(2, (6, self.config.hidden_size)) # Shard per-patch temb_proj to match the local sequence chunk - if chunk_size is not None: - temb_proj = temb_proj[:, chunk_start:chunk_end] + temb_proj = self.sharder.shard(temb_proj, dim=1, expected_seq_len=seq_len) else: # batch_size, 6, hidden_size temb_proj = temb_proj.unflatten(1, (6, self.config.hidden_size)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index e626b4fdb24d..0ff5be1dd203 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -116,7 +116,6 @@ def _setup_visual_gen_mapping(self, config: DiffusionModelConfig) -> None: attn2d_row_size=self.args.parallel.dit_attn2d_row_size, attn2d_col_size=self.args.parallel.dit_attn2d_col_size, parallel_vae_size=self.args.parallel.parallel_vae_size, - order=self.args.parallel.dit_dim_order, ) else: # Single-GPU fallback. no args = no parallelism. diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 31203f999ed9..6bad931bc85e 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -3,7 +3,6 @@ Single-GPU tests run without dist. Multi-GPU tests use mp.spawn with NCCL. """ -import itertools import os os.environ["TLLM_DISABLE_MPI"] = "1" @@ -14,7 +13,7 @@ import torch.multiprocessing as mp try: - from tensorrt_llm._torch.visual_gen.mapping import _VALID_DIM_NAMES, VisualGenMapping + from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping from tensorrt_llm._utils import get_free_port MODULES_AVAILABLE = True @@ -122,24 +121,10 @@ def test_product_mismatch_raises(self): def test_parallel_vae_size_cannot_exceed_world_size(self): with pytest.raises(ValueError, match="cannot exceed world_size"): VisualGenMapping(world_size=1, rank=0, parallel_vae_size=2) - - def test_invalid_order_raises(self): - with pytest.raises(ValueError, match="permutation"): - VisualGenMapping(world_size=1, rank=0, order="cfg-tp-ulysses") - - def test_duplicate_dim_raises(self): - with pytest.raises(ValueError, match="permutation"): - VisualGenMapping(world_size=1, rank=0, order="cfg-cfg-tp-ulysses") - - def test_custom_order_stored(self): - vgm = VisualGenMapping(world_size=1, rank=0, order="ulysses-cp-tp-cfg") - assert vgm._dim_names == ("ulysses", "cp", "tp", "cfg") - - def test_all_valid_orders(self): - for perm in itertools.permutations(sorted(_VALID_DIM_NAMES)): - order = "-".join(perm) - vgm = VisualGenMapping(world_size=1, rank=0, order=order) - assert vgm._dim_names == perm + + def test_internal_dim_order(self): + vgm = VisualGenMapping(world_size=1, rank=0) + assert vgm._dim_names == ("cfg", "tp", "cp", "ulysses") def test_ulysses_and_attn2d_raises(self): """Combining Attention2D and Ulysses raises NotImplementedError.""" @@ -263,37 +248,6 @@ def _logic_default_order_cfg2_ulysses2(rank, world_size): assert m.tp_size == 1 -def _logic_custom_order_ulysses_outermost(rank, world_size): - """Custom order ulysses-cp-tp-cfg with cfg=2, ulysses=2 on 4 GPUs. - - Expected rank layout (outermost=ulysses, innermost=cfg): - Rank 0: ulysses=0, cfg=0 - Rank 1: ulysses=0, cfg=1 - Rank 2: ulysses=1, cfg=0 - Rank 3: ulysses=1, cfg=1 - """ - from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl - - DeviceMeshTopologyImpl.device_mesh = None - - vgm = VisualGenMapping( - world_size=world_size, - rank=rank, - cfg_size=2, - ulysses_size=2, - order="ulysses-cp-tp-cfg", - ) - - assert vgm.ulysses_rank == rank // 2 - assert vgm.cfg_rank == rank % 2 - assert vgm.is_cfg_conditional == (rank % 2 == 0) - - cfg_pg_size = dist.get_world_size(vgm.cfg_group) - ulysses_pg_size = dist.get_world_size(vgm.ulysses_group) - assert cfg_pg_size == 2 - assert ulysses_pg_size == 2 - - def _logic_allreduce_over_tp_group(rank, world_size): """Verify TP group works for collective ops (tp=2, ulysses=2 on 4 GPUs). @@ -441,9 +395,6 @@ class TestMultiGPU: def test_default_order_cfg2_ulysses2(self): _run_multi_gpu(4, _logic_default_order_cfg2_ulysses2) - def test_custom_order_ulysses_outermost(self): - _run_multi_gpu(4, _logic_custom_order_ulysses_outermost) - def test_allreduce_over_tp_group(self): _run_multi_gpu(4, _logic_allreduce_over_tp_group) From 9d821cc087011b0e4be499d43e665312e6712a3c Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 6 May 2026 17:03:26 -0700 Subject: [PATCH 08/13] tests Signed-off-by: Shreyas Misra --- .../multi_gpu/test_ring_attention.py | 524 ++++++++++++++++++ .../visual_gen/multi_gpu/test_wan_ring.py | 322 +++++++++++ 2 files changed, 846 insertions(+) create mode 100644 tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py create mode 100644 tests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.py diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py new file mode 100644 index 000000000000..c5443b8d3b72 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-GPU tests for RingAttention and Ring + Ulysses composition. + +Uses the same spawn pattern as test_attn2d_attention.py. Process groups for +ring (``cp`` dim) and Ulysses (``ulysses`` dim) match ``VisualGenMapping``: +mesh layout ``cfg-tp-cp-ulysses`` with shape ``(1, 1, ring_size, ulysses_size)``. + +Ring-only: ``ulysses_size=1``, ``ring_size=world_size``. + +Composition matches ``attention.py``: ``UlyssesAttention(RingAttention(inner))``. + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py -v +""" + +import math +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from typing import Callable, Optional, Tuple + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import init_device_mesh + +try: + from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask + from tensorrt_llm._torch.visual_gen.attention_backend import RingAttention, UlyssesAttention + from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention + from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, + ) + from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout + from tensorrt_llm._utils import get_free_port + + MODULES_AVAILABLE = True +except ImportError: + MODULES_AVAILABLE = False + _fa4_fwd = None + +_flash_attn4_available = _fa4_fwd is not None + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# ============================================================================= +# Test-only inner backend: VanillaAttention with LSE output (RingAttention) +# ============================================================================= + + +class _LSEVanillaAttention(nn.Module): + """Inner backend with ``forward_with_lse`` for RingAttention tests.""" + + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.scale = 1.0 / math.sqrt(head_dim) + self._preferred_layout = AttentionTensorLayout.NHD + + @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 True + + def forward(self, q, k, v, batch_size=None, seq_len=None, **kwargs): + q_t = q.transpose(1, 2).float() + k_t = k.transpose(1, 2).float() + v_t = v.transpose(1, 2).float() + out = F.scaled_dot_product_attention(q_t, k_t, v_t, scale=self.scale) + return out.to(q.dtype).transpose(1, 2).contiguous() + + def forward_with_lse(self, q, k, v, batch_size=None, seq_len=None, **kwargs): + q_t = q.transpose(1, 2).float() + k_t = k.transpose(1, 2).float() + v_t = v.transpose(1, 2).float() + scores = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale + lse = torch.logsumexp(scores, dim=-1) + attn = torch.softmax(scores, dim=-1) + out = torch.matmul(attn, v_t) + return out.to(q.dtype).transpose(1, 2).contiguous(), lse + + +# ============================================================================= +# Distributed helpers +# ============================================================================= + + +def _init_worker(rank: int, world_size: int, backend: str, port: int): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + if backend == "nccl" and torch.cuda.is_available(): + torch.cuda.set_device(rank % torch.cuda.device_count()) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + else: + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + +def _cleanup(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_worker(rank, world_size, backend, test_fn, port): + try: + _init_worker(rank, world_size, backend, port) + test_fn(rank, world_size) + except Exception as e: + print(f"Rank {rank} failed: {e}") + raise + finally: + _cleanup() + + +def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + if use_cuda and torch.cuda.device_count() < world_size: + pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") + + backend = "nccl" if use_cuda else "gloo" + port = get_free_port() + mp.spawn( + _distributed_worker, + args=(world_size, backend, test_fn, port), + nprocs=world_size, + join=True, + ) + + +def _cp_ulysses_mesh_groups( + ring_size: int, ulysses_size: int +) -> Tuple[dist.ProcessGroup, Optional[dist.ProcessGroup]]: + """DeviceMesh (1,1,ring_size,ulysses_size) matching VisualGenMapping axis order. + + Returns ``(ring_group, ulysses_group)``. ``ulysses_group`` is ``None`` when + ``ulysses_size == 1``. + """ + device_type = "cuda" if torch.cuda.is_available() else "cpu" + mesh = init_device_mesh( + device_type, + (1, 1, ring_size, ulysses_size), + mesh_dim_names=("cfg", "tp", "cp", "ulysses"), + ) + ring_pg = mesh["cp"].get_group() + uly_pg = mesh["ulysses"].get_group() if ulysses_size > 1 else None + return ring_pg, uly_pg + + +def _seq_shard_bounds(rank: int, world_size: int, seq_full: int) -> Tuple[int, int]: + """Contiguous global-token shard for rank (same linear order as cp-major, uly-minor).""" + assert seq_full % world_size == 0 + chunk = seq_full // world_size + return rank * chunk, (rank + 1) * chunk + + +# ============================================================================= +# Ring only (1D cp group) +# ============================================================================= + + +def _logic_ring_forward(rank, world_size): + """Ring CP size ``world_size``: output shape [B, S/P, H, D], finite values.""" + ring_size, ulysses_size = world_size, 1 + assert dist.get_world_size() == ring_size * ulysses_size + + batch, seq_per_rank, num_heads, head_dim = 2, 8, 8, 64 + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + + ring_pg, _ = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + inner = _LSEVanillaAttention(num_heads=num_heads, head_dim=head_dim) + attn = RingAttention(inner, ring_pg) + + q = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + k = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + v = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + + output = attn(q, k, v, batch_size=batch) + assert output.shape == q.shape, f"Rank {rank}: expected {q.shape}, got {output.shape}" + assert torch.isfinite(output).all(), f"Rank {rank}: non-finite output" + + +def _logic_ring_vs_standard(rank, world_size): + """Ring output on each rank matches the corresponding SDPA sequence shard.""" + ring_size, ulysses_size = world_size, 1 + batch, num_heads, head_dim = 2, 8, 64 + seq_per_rank = 8 + seq_full = seq_per_rank * world_size + + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + ring_pg, _ = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + + torch.manual_seed(42) + q_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + k_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + v_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + + lo, hi = _seq_shard_bounds(rank, world_size, seq_full) + q_shard = q_full[:, lo:hi].contiguous() + k_shard = k_full[:, lo:hi].contiguous() + v_shard = v_full[:, lo:hi].contiguous() + + inner = _LSEVanillaAttention(num_heads=num_heads, head_dim=head_dim) + attn = RingAttention(inner, ring_pg) + ring_out = attn(q_shard, k_shard, v_shard, batch_size=batch) + + scale = 1.0 / math.sqrt(head_dim) + q_std = q_full.transpose(1, 2).float() + k_std = k_full.transpose(1, 2).float() + v_std = v_full.transpose(1, 2).float() + std_output = F.scaled_dot_product_attention(q_std, k_std, v_std, scale=scale) + std_output = std_output.transpose(1, 2).to(ring_out.dtype) + expected = std_output[:, lo:hi] + + torch.testing.assert_close( + ring_out, + expected, + rtol=1e-3, + atol=1e-3, + msg=f"Rank {rank}: RingAttention differs from SDPA reference", + ) + + +def _logic_ring_invalid_mask(rank, world_size): + """RingAttention rejects non-FULL masks for self-attention.""" + ring_size, ulysses_size = world_size, 1 + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + ring_pg, _ = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + inner = _LSEVanillaAttention(num_heads=8, head_dim=64) + attn = RingAttention(inner, ring_pg) + + q = torch.randn(2, 8, 8, 64, device=device) + k = torch.randn(2, 8, 8, 64, device=device) + v = torch.randn(2, 8, 8, 64, device=device) + + with pytest.raises(NotImplementedError, match="Causal"): + attn(q, k, v, batch_size=2, attention_mask=PredefinedAttentionMask.CAUSAL) + + +def _logic_ring_fa4_vs_standard(rank, world_size): + """Ring with FlashAttn4 inner matches SDPA reference shards.""" + if not _flash_attn4_available: + pytest.skip("FlashAttn4 JIT kernels not available") + + ring_size, ulysses_size = world_size, 1 + batch, num_heads, head_dim = 1, 8, 128 + seq_per_rank = 16 + seq_full = seq_per_rank * world_size + + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + ring_pg, _ = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + + inner = FlashAttn4Attention(num_heads=num_heads, head_dim=head_dim) + attn = RingAttention(inner, ring_pg) + + torch.manual_seed(42) + q_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + k_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + v_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + + lo, hi = _seq_shard_bounds(rank, world_size, seq_full) + q_shard = q_full[:, lo:hi].contiguous() + k_shard = k_full[:, lo:hi].contiguous() + v_shard = v_full[:, lo:hi].contiguous() + + ring_out = attn(q_shard, k_shard, v_shard, batch_size=batch) + + scale = 1.0 / math.sqrt(head_dim) + ref = ( + F.scaled_dot_product_attention( + q_full.transpose(1, 2).float(), + k_full.transpose(1, 2).float(), + v_full.transpose(1, 2).float(), + scale=scale, + ) + .transpose(1, 2) + .to(ring_out.dtype) + ) + expected = ref[:, lo:hi] + torch.testing.assert_close( + ring_out, + expected, + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: Ring FA4 differs from SDPA reference", + ) + + +# ============================================================================= +# Ring + Ulysses: UlyssesAttention(RingAttention(inner)) +# ============================================================================= + + +def _logic_ring_ulysses_forward(rank, world_size): + """ring_size=2, ulysses_size=2 on 4 GPUs: finite output, correct shape.""" + ring_size, ulysses_size = 2, 2 + assert world_size == ring_size * ulysses_size + + batch, seq_per_rank, num_heads, head_dim = 2, 8, 8, 64 + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + + ring_pg, uly_pg = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + assert uly_pg is not None + + inner = _LSEVanillaAttention(num_heads=num_heads // ulysses_size, head_dim=head_dim) + attn = UlyssesAttention(RingAttention(inner, ring_pg), uly_pg) + + q = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + k = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + v = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) + + output = attn(q, k, v, batch_size=batch) + assert output.shape == q.shape, f"Rank {rank}: expected {q.shape}, got {output.shape}" + assert torch.isfinite(output).all(), f"Rank {rank}: non-finite output" + + +def _logic_ring_ulysses_vs_standard(rank, world_size): + """Combined parallelism matches SDPA on full sequence; compare local shard.""" + ring_size, ulysses_size = 2, 2 + assert world_size == ring_size * ulysses_size + + batch, num_heads, head_dim = 2, 8, 64 + seq_per_rank = 8 + seq_full = seq_per_rank * world_size + + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + ring_pg, uly_pg = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + + torch.manual_seed(42) + q_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + k_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + v_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device) + + lo, hi = _seq_shard_bounds(rank, world_size, seq_full) + q_shard = q_full[:, lo:hi].contiguous() + k_shard = k_full[:, lo:hi].contiguous() + v_shard = v_full[:, lo:hi].contiguous() + + inner = _LSEVanillaAttention(num_heads=num_heads // ulysses_size, head_dim=head_dim) + attn = UlyssesAttention(RingAttention(inner, ring_pg), uly_pg) + combo_out = attn(q_shard, k_shard, v_shard, batch_size=batch) + + scale = 1.0 / math.sqrt(head_dim) + ref = ( + F.scaled_dot_product_attention( + q_full.transpose(1, 2).float(), + k_full.transpose(1, 2).float(), + v_full.transpose(1, 2).float(), + scale=scale, + ) + .transpose(1, 2) + .to(combo_out.dtype) + ) + expected = ref[:, lo:hi] + + torch.testing.assert_close( + combo_out, + expected, + rtol=1e-3, + atol=1e-3, + msg=f"Rank {rank}: Ring+Ulysses differs from SDPA reference", + ) + + +def _logic_ring_ulysses_fa4_vs_standard(rank, world_size): + """Ring + Ulysses with FA4 inner matches SDPA shards.""" + if not _flash_attn4_available: + pytest.skip("FlashAttn4 JIT kernels not available") + + ring_size, ulysses_size = 2, 2 + batch, num_heads, head_dim = 1, 8, 128 + seq_per_rank = 16 + seq_full = seq_per_rank * world_size + + device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") + ring_pg, uly_pg = _cp_ulysses_mesh_groups(ring_size, ulysses_size) + + inner = FlashAttn4Attention(num_heads=num_heads // ulysses_size, head_dim=head_dim) + attn = UlyssesAttention(RingAttention(inner, ring_pg), uly_pg) + + torch.manual_seed(42) + q_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + k_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + v_full = torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16) + + lo, hi = _seq_shard_bounds(rank, world_size, seq_full) + q_shard = q_full[:, lo:hi].contiguous() + k_shard = k_full[:, lo:hi].contiguous() + v_shard = v_full[:, lo:hi].contiguous() + + combo_out = attn(q_shard, k_shard, v_shard, batch_size=batch) + + scale = 1.0 / math.sqrt(head_dim) + ref = ( + F.scaled_dot_product_attention( + q_full.transpose(1, 2).float(), + k_full.transpose(1, 2).float(), + v_full.transpose(1, 2).float(), + scale=scale, + ) + .transpose(1, 2) + .to(combo_out.dtype) + ) + expected = ref[:, lo:hi] + + torch.testing.assert_close( + combo_out, + expected, + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: Ring+Ulysses FA4 differs from SDPA reference", + ) + + +# ============================================================================= +# Init guards (RingAttention) +# ============================================================================= + + +def _logic_ring_init_guard_no_lse(rank, world_size): + class _NoLSEBackend(nn.Module): + num_heads = 8 + head_dim = 64 + _preferred_layout = AttentionTensorLayout.NHD + + @property + def preferred_layout(self): + return self._preferred_layout + + @classmethod + def support_lse(cls): + return False + + @classmethod + def support_fused_qkv(cls): + return False + + pg = dist.new_group(list(range(world_size)), use_local_synchronization=True) + with pytest.raises(ValueError, match="LSE"): + RingAttention(_NoLSEBackend(), pg) + + +# ============================================================================= +# Test classes +# ============================================================================= + + +class TestRingAttention: + """Ring CP only (``ulysses_size=1``), mesh ``cp = world_size``.""" + + def test_ring_forward(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_forward, use_cuda=True) + + def test_ring_vs_standard_attention(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_vs_standard, use_cuda=True) + + def test_ring_invalid_mask_raises(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_invalid_mask, use_cuda=True) + + +class TestRingAttentionFlashAttn4: + """Ring with FA4 inner (production-style).""" + + def test_ring_fa4_vs_standard(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_fa4_vs_standard, use_cuda=True) + + +class TestRingUlyssesAttention: + """``UlyssesAttention(RingAttention(inner))`` on 4 GPUs (2×2 cp × ulysses).""" + + def test_ring_ulysses_forward(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_ulysses_forward, use_cuda=True) + + def test_ring_ulysses_vs_standard_attention(self): + run_test_in_distributed( + world_size=4, test_fn=_logic_ring_ulysses_vs_standard, use_cuda=True + ) + + +class TestRingUlyssesAttentionFlashAttn4: + """Ring + Ulysses with FA4 inner.""" + + def test_ring_ulysses_fa4_vs_standard(self): + run_test_in_distributed( + world_size=4, test_fn=_logic_ring_ulysses_fa4_vs_standard, use_cuda=True + ) + + +class TestRingAttentionInitGuards: + def test_inner_without_lse_raises(self): + run_test_in_distributed(world_size=4, test_fn=_logic_ring_init_guard_no_lse, use_cuda=False) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.py new file mode 100644 index 000000000000..f1d66c29b19b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-GPU tests for WAN ring context parallelism. + +Tests that WanTransformer3DModel produces correct outputs when using +ring attention (1D CP group, KV streamed around the ring), optionally +combined with Ulysses head parallelism (``UlyssesAttention(RingAttention(FA4))``). + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_wan_ring.py -v +""" + +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from types import SimpleNamespace +from typing import Callable + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +try: + from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TeaCacheConfig, + TorchCompileConfig, + ) + from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping + from tensorrt_llm._utils import get_free_port + from tensorrt_llm.models.modeling_utils import QuantConfig + + MODULES_AVAILABLE = True +except ImportError: + MODULES_AVAILABLE = False + +try: + from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, + ) + + _flash_attn4_available = _fa4_fwd is not None +except (ImportError, OSError): + _flash_attn4_available = False + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + """Clean up TLLM_DISABLE_MPI env var after tests complete.""" + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# ============================================================================= +# Distributed helpers (same pattern as test_wan_attn2d.py) +# ============================================================================= + + +def init_distributed_worker(rank: int, world_size: int, backend: str = "nccl", port: int = 29500): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + torch.cuda.set_device(rank % torch.cuda.device_count()) + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + + +def cleanup_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_worker(rank, world_size, backend, test_fn, port): + """Worker function run in each spawned process. Module-level for pickling.""" + try: + init_distributed_worker(rank, world_size, backend, port) + test_fn(rank, world_size) + except Exception as e: + print(f"Rank {rank} failed with error: {e}") + raise + finally: + cleanup_distributed() + + +def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + if use_cuda and torch.cuda.device_count() < world_size: + pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") + backend = "nccl" if use_cuda else "gloo" + port = get_free_port() + mp.spawn( + _distributed_worker, args=(world_size, backend, test_fn, port), nprocs=world_size, join=True + ) + + +# ============================================================================= +# Model config helpers +# ============================================================================= + +# Small WAN config for testing. +# hidden_size = num_attention_heads * attention_head_dim = 4 * 64 = 256. +# Ring shards sequence; Ulysses shards heads — num_attention_heads must divide ulysses_size. +_WAN_TEST_CONFIG = dict( + num_attention_heads=4, + attention_head_dim=64, + num_layers=2, + in_channels=4, + out_channels=4, + patch_size=[1, 2, 2], + text_dim=64, + freq_dim=32, +) + +# Video input: [B=1, C=4, T=2, H=4, W=4] +# After patchify with patch_size=[1,2,2]: seq_len = T * (H/2) * (W/2) = 2*2*2 = 8 +# SequenceSharder uses seq_size = ring_size * ulysses_size: +# ring_size=4, ulysses=1 -> chunk 8/4 = 2 tokens/rank +# ring_size=2, ulysses=2 -> chunk 8/4 = 2 tokens/rank +_VIDEO_SHAPE = (1, 4, 2, 4, 4) +_TEXT_SEQ = 4 + + +def _make_model_config(pretrained_dict, ring_size=1, ulysses_size=1, backend="FA4"): + """Create DiffusionModelConfig for testing. + + With ``ring_size=ulysses_size=1``, single-GPU config. + Otherwise builds ``VisualGenMapping`` from the active process group's + ``world_size`` / ``rank`` (must equal ``cfg*tp*ring_size*ulysses_size``). + """ + pretrained_config = SimpleNamespace(**pretrained_dict) + use_dist = (ring_size > 1 or ulysses_size > 1) and dist.is_initialized() + if use_dist: + ws = dist.get_world_size() + rk = dist.get_rank() + else: + ws = 1 + rk = 0 + vgm = VisualGenMapping( + world_size=ws, + rank=rk, + ring_size=ring_size, + ulysses_size=ulysses_size, + ) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + torch_compile=TorchCompileConfig(enable_torch_compile=False), + attention=AttentionConfig(backend=backend), + visual_gen_mapping=vgm, + teacache=TeaCacheConfig(), + skip_create_weights_in_init=False, + ) + config.mapping = vgm.to_llm_mapping() + return config + + +def _stabilize_model_weights(model): + """Reinitialize weights to prevent BF16 overflow through multiple transformer blocks.""" + with torch.no_grad(): + for p in model.parameters(): + if p.ndim >= 2: + fan_in = p.shape[1] + std = 0.02 / max(1.0, fan_in**0.5) + p.data.uniform_(-std, std) + else: + p.data.uniform_(-0.01, 0.01) + + +# ============================================================================= +# Test logic (module-level for mp.spawn pickling) +# ============================================================================= + + +def _logic_wan_ring_vs_single_gpu(rank, world_size): + """WAN ring CP (size 4) output matches single-GPU FA4 reference. + + Both models use FA4 as the inner attention backend; ring wraps it with + KV streaming and online LSE merge. The WAN forward all-gathers the + sequence at the end, so each rank can compare to the single-GPU reference. + """ + from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel + + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + + B, C, T, H, W = _VIDEO_SHAPE + text_dim = _WAN_TEST_CONFIG["text_dim"] + + torch.manual_seed(42) + ref_config = _make_model_config(_WAN_TEST_CONFIG) + ref_model = WanTransformer3DModel(ref_config).to(device).to(dtype) + _stabilize_model_weights(ref_model) + ref_state = ref_model.state_dict() + + torch.manual_seed(42) + ring_config = _make_model_config(_WAN_TEST_CONFIG, ring_size=4) + try: + ring_model = WanTransformer3DModel(ring_config).to(device).to(dtype) + except (ImportError, ValueError, NotImplementedError) as e: + pytest.skip(f"Ring attention / FA4 not available: {e}") + ring_model.load_state_dict(ref_state) + + torch.manual_seed(100) + hidden_states = torch.randn(_VIDEO_SHAPE, device=device, dtype=dtype) * 0.1 + encoder_hidden_states = torch.randn(B, _TEXT_SEQ, text_dim, device=device, dtype=dtype) * 0.1 + timestep = torch.tensor([0.5], device=device, dtype=dtype) + + with torch.no_grad(): + ref_output = ref_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + ring_output = ring_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + + torch.testing.assert_close( + ring_output, + ref_output, + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: WAN ring CP output differs from single-GPU FA4 reference", + ) + + +def _logic_wan_ring_ulysses_vs_single_gpu(rank, world_size): + """WAN ring+Ulysses (2×2) matches single-GPU FA4 reference on each rank.""" + from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel + + assert world_size == 4 + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + + B, C, T, H, W = _VIDEO_SHAPE + text_dim = _WAN_TEST_CONFIG["text_dim"] + + torch.manual_seed(42) + ref_config = _make_model_config(_WAN_TEST_CONFIG) + ref_model = WanTransformer3DModel(ref_config).to(device).to(dtype) + _stabilize_model_weights(ref_model) + ref_state = ref_model.state_dict() + + torch.manual_seed(42) + combo_config = _make_model_config(_WAN_TEST_CONFIG, ring_size=2, ulysses_size=2) + try: + combo_model = WanTransformer3DModel(combo_config).to(device).to(dtype) + except (ImportError, ValueError, NotImplementedError) as e: + pytest.skip(f"Ring+Ulysses / FA4 not available: {e}") + combo_model.load_state_dict(ref_state) + + torch.manual_seed(100) + hidden_states = torch.randn(_VIDEO_SHAPE, device=device, dtype=dtype) * 0.1 + encoder_hidden_states = torch.randn(B, _TEXT_SEQ, text_dim, device=device, dtype=dtype) * 0.1 + timestep = torch.tensor([0.5], device=device, dtype=dtype) + + with torch.no_grad(): + ref_output = ref_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + combo_output = combo_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + + torch.testing.assert_close( + combo_output, + ref_output, + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: WAN ring+Ulysses output differs from single-GPU FA4 reference", + ) + + +# ============================================================================= +# Test classes +# ============================================================================= + + +class TestWanRing: + """Ring context parallelism tests for WAN transformer (ring_size=4, 4 GPUs).""" + + def test_wan_ring_vs_single_gpu(self): + """WAN ring CP (4 ranks) output matches single-GPU FA4 reference.""" + if not _flash_attn4_available: + pytest.skip("FlashAttn4 JIT kernels not available") + run_test_in_distributed(world_size=4, test_fn=_logic_wan_ring_vs_single_gpu) + + +class TestWanRingUlysses: + """WAN with ring CP (2) and Ulysses (2) on 4 GPUs.""" + + def test_wan_ring_ulysses_vs_single_gpu(self): + """Distributed output matches single-GPU FA4 reference.""" + if not _flash_attn4_available: + pytest.skip("FlashAttn4 JIT kernels not available") + run_test_in_distributed(world_size=4, test_fn=_logic_wan_ring_ulysses_vs_single_gpu) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 57fb19ebedd3b734fddf9468997b027ff33b22e2 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 6 May 2026 17:30:17 -0700 Subject: [PATCH 09/13] seq sharder test, minor fix Signed-off-by: Shreyas Misra --- .../attention_backend/flash_attn4.py | 4 - .../visual_gen/attention_backend/parallel.py | 8 +- .../unittest/_torch/visual_gen/test_utils.py | 188 ++++++++++++++++++ 3 files changed, 191 insertions(+), 9 deletions(-) create mode 100644 tests/unittest/_torch/visual_gen/test_utils.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index 7060602208f7..8981faa0de11 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -180,7 +180,3 @@ def preferred_layout(self) -> AttentionTensorLayout: @classmethod def support_fused_qkv(cls) -> bool: return False - - @classmethod - def supports_lse(cls) -> bool: - return True diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index 01872832087d..ddaab41bd6b0 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -460,8 +460,6 @@ def __init__( inner_backend: AttentionBackend, process_group: dist.ProcessGroup, ): - # Invariant: only instantiated when ring_size > 1 (see attention.py), - # so distributed must be initialized and the group must be non-trivial. if not type(inner_backend).support_lse(): raise ValueError( f"RingAttention requires an LSE-capable inner backend (FA4); " @@ -532,9 +530,9 @@ def _ensure_buffers(self, q: torch.Tensor, k: torch.Tensor) -> None: key = (B, S, H, H_kv, D, q.device, q.dtype, k.dtype) if key == self._buf_key: return - self._kv_bufs = torch.empty(2, 2, B, S, H_kv, D, device=k.device, dtype=k.dtype) - self._out_buf = torch.empty(B, S, H, D, device=q.device, dtype=q.dtype) - self._lse_buf = torch.empty(B, S, H, device=q.device, dtype=torch.float32) + self._kv_bufs = k.new_empty(2, 2, B, S, H_kv, D) + self._out_buf = q.new_empty(B, S, H, D) + self._lse_buf = q.new_empty(B, S, H, dtype=torch.float32) self._buf_key = key def _update_out_and_lse( diff --git a/tests/unittest/_torch/visual_gen/test_utils.py b/tests/unittest/_torch/visual_gen/test_utils.py new file mode 100644 index 000000000000..594b6c530063 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_utils.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for ``tensorrt_llm._torch.visual_gen.utils`` (SequenceSharder, etc.).""" + +import os +from types import SimpleNamespace +from typing import Callable + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +try: + from tensorrt_llm._torch.visual_gen.utils import SequenceSharder + from tensorrt_llm._utils import get_free_port + + MODULES_AVAILABLE = True +except ImportError: + MODULES_AVAILABLE = False + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + """Avoid leaking TLLM_DISABLE_MPI if other modules set it.""" + yield + + +# ============================================================================= +# SequenceSharder — single-process (no distributed) +# ============================================================================= + + +class TestSequenceSharderInactive: + def test_shard_gather_identity_when_size_one(self): + s = SequenceSharder(size=1, rank=0, group=None) + assert not s.is_active + x = torch.randn(2, 8, 4) + assert s.shard(x, dim=1) is x + assert s.gather(x, dim=1) is x + assert s.gather(x, dim=1, unpad_to=3) is x + + def test_shard_none_returns_none(self): + s = SequenceSharder(size=4, rank=1, group=None) + assert s.shard(None, dim=1) is None + + def test_shard_rope_passthrough_when_inactive(self): + s = SequenceSharder(size=1, rank=0, group=None) + cos = torch.randn(1, 4, 8) + rope = (cos, cos) + assert s.shard_rope(rope, seq_len=4) is rope + + def test_disable_enable_no_collectives(self): + s = SequenceSharder(size=4, rank=0, group=None) + assert s.is_active + s.disable() + assert not s.is_active + x = torch.randn(1, 8, 2) + assert s.shard(x, dim=1) is x + s.enable() + assert s.is_active + + +class TestSequenceSharderShardSlices: + """Active sharder: block slice math without ``gather``.""" + + def test_shard_dim1_four_ranks(self): + s = SequenceSharder(size=4, rank=2, group=None) + x = torch.arange(24).view(2, 12, 1) + part = s.shard(x, dim=1) + assert part.shape == (2, 3, 1) + assert torch.equal(part.squeeze(-1), torch.tensor([[6, 7, 8], [18, 19, 20]])) + + def test_expected_seq_len_skip_shard(self): + s = SequenceSharder(size=4, rank=0, group=None) + x = torch.randn(1, 7, 2) + out = s.shard(x, dim=1, expected_seq_len=8) + assert out.shape == (1, 7, 2) + + def test_not_divisible_raises(self): + s = SequenceSharder(size=4, rank=0, group=None) + x = torch.randn(1, 10, 2) + with pytest.raises(ValueError, match="divisible"): + s.shard(x, dim=1) + + +class TestSequenceSharderShardRope: + def test_shard_rope_bsd_layout(self): + s = SequenceSharder(size=2, rank=1, group=None) + B, S, D = 1, 8, 16 + cos = torch.arange(B * S * D).view(B, S, D).float() + sin = cos + 1000 + out = s.shard_rope((cos, sin), seq_len=S) + assert out is not None + oc, osin = out + assert oc.shape == (1, 4, 16) + assert torch.equal(oc, cos[:, 4:8].contiguous()) + + def test_shard_rope_ambiguous_returns_unchanged(self): + """Two axes equal seq_len → no unique dim → passthrough.""" + s = SequenceSharder(size=2, rank=0, group=None) + S = 8 + cos = torch.zeros(2, S, S) + sin = torch.ones(2, S, S) + rope = (cos, sin) + assert s.shard_rope(rope, seq_len=S) is rope + + +class TestSequenceSharderFromVgm: + def test_from_vgm_none(self): + s = SequenceSharder.from_vgm(None) + assert s.size == 1 and s.rank == 0 and not s.is_active + + def test_from_vgm_head_divisibility(self): + vgm = SimpleNamespace( + seq_size=4, + seq_rank=0, + seq_group=None, + ulysses_size=2, + ) + SequenceSharder.from_vgm(vgm, num_attention_heads=8, num_kv_heads=4) + with pytest.raises(ValueError, match="num_attention_heads"): + SequenceSharder.from_vgm(vgm, num_attention_heads=7) + + +# ============================================================================= +# SequenceSharder — distributed shard / gather round-trip (gloo, CPU) +# ============================================================================= + + +def _dist_worker_shard_gather_combined(rank: int, world_size: int, port: int): + """Single spawn: basic shard/gather round-trip and pad/gather/unpad.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + group = dist.group.WORLD + sharder = SequenceSharder(size=world_size, rank=rank, group=group) + + x = torch.arange(world_size * 5, dtype=torch.float32).view(1, world_size * 5, 1) + local = sharder.shard(x, dim=1) + assert local.shape == (1, 5, 1) + exp = torch.arange(rank * 5, (rank + 1) * 5, dtype=torch.float32).view(1, 5, 1) + assert torch.allclose(local, exp) + full = sharder.gather(local, dim=1) + assert full.shape == x.shape + assert torch.allclose(full, x) + + original_len = 10 + xp = torch.ones(1, original_len, 2) + local_p = sharder.shard(xp, dim=1, pad_to_multiple=True) + restored = sharder.gather(local_p, dim=1, unpad_to=original_len) + assert restored.shape == xp.shape + assert torch.allclose(restored, xp) + finally: + dist.destroy_process_group() + + +def _spawn_entry_combined(rank: int, world_size: int, port: int): + _dist_worker_shard_gather_combined(rank, world_size, port) + + +def _run_dist(world_size: int, entry: Callable[[int, int, int], None]): + if not MODULES_AVAILABLE: + pytest.skip("SequenceSharder import failed") + port = get_free_port() + mp.spawn(entry, args=(world_size, port), nprocs=world_size, join=True) + + +class TestSequenceSharderDistributed: + def test_shard_gather_pad_unpad_four_ranks(self): + _run_dist(4, _spawn_entry_combined) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 6e3c4d1123517f4493fc904bb00c2e92b4e5c6e8 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 6 May 2026 17:35:11 -0700 Subject: [PATCH 10/13] pre-commit and coderabbit suggestion Signed-off-by: Shreyas Misra --- examples/visual_gen/visual_gen_wan_t2v.py | 5 +++-- .../_torch/visual_gen/attention_backend/parallel.py | 6 ++++-- tensorrt_llm/_torch/visual_gen/utils.py | 7 +++++++ 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index 5569131b2bef..b30d9d2d9116 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -327,8 +327,9 @@ def main(): attn2d_size = args.attn2d_row_size * args.attn2d_col_size if attn2d_size > 1 and args.ulysses_size > 1: raise ValueError( - "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented.") - + "Combining --ulysses_size with --attn2d_row_size/--attn2d_col_size is not yet implemented." + ) + if args.ulysses_size > 1: num_heads = 40 logger.info( diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index ddaab41bd6b0..29c28ef05057 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -559,8 +559,10 @@ def forward( # 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.CAUSAL: - raise NotImplementedError("Causal ring attention not implemented yet.") + 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"} diff --git a/tensorrt_llm/_torch/visual_gen/utils.py b/tensorrt_llm/_torch/visual_gen/utils.py index cff6c3757659..e4572ed1c1d3 100644 --- a/tensorrt_llm/_torch/visual_gen/utils.py +++ b/tensorrt_llm/_torch/visual_gen/utils.py @@ -98,6 +98,13 @@ def from_vgm( rank = vgm.seq_rank group = vgm.seq_group + if size > 1 and group is None: + raise ValueError( + "SequenceSharder.from_vgm requires vgm.seq_group to be set when " + f"vgm.seq_size ({size}) > 1; otherwise gather() would call " + "dist.all_gather(..., group=None) and use the default process group." + ) + if size > 1 and vgm.ulysses_size > 1: for label, count in ( ("num_attention_heads", num_attention_heads), From 7f6e5ef0ced49e79f1250c72b9b01c97564adb5e Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 14 May 2026 11:00:21 -0700 Subject: [PATCH 11/13] update docs Signed-off-by: Shreyas Misra --- docs/source/models/visual-generation.md | 18 ++++++++++-------- examples/visual_gen/visual_gen_wan_i2v.py | 15 +++++++++++++-- examples/visual_gen/visual_gen_wan_t2v.py | 11 ++++++++++- 3 files changed, 33 insertions(+), 11 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 4230e60a3948..9486912b76f8 100644 --- a/docs/source/models/visual-generation.md +++ b/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 | +| 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 | -| **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 | +| **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 - + - **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 diff --git a/examples/visual_gen/visual_gen_wan_i2v.py b/examples/visual_gen/visual_gen_wan_i2v.py index ee43c8906e5f..4a1a63beea37 100644 --- a/examples/visual_gen/visual_gen_wan_i2v.py +++ b/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={ diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index b30d9d2d9116..db21b6187bb4 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -239,7 +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 parallel size") + 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, @@ -329,6 +334,10 @@ 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: num_heads = 40 From 4691a81770a7008947d4cd3ad5703535cb9c8d3c Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 14 May 2026 11:05:59 -0700 Subject: [PATCH 12/13] fix tests Signed-off-by: Shreyas Misra --- .../multi_gpu/test_ring_attention.py | 2 +- .../multi_gpu/test_ulysses_attention.py | 44 ------------------- 2 files changed, 1 insertion(+), 45 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py index c5443b8d3b72..8880965850b7 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py @@ -262,7 +262,7 @@ def _logic_ring_invalid_mask(rank, world_size): k = torch.randn(2, 8, 8, 64, device=device) v = torch.randn(2, 8, 8, 64, device=device) - with pytest.raises(NotImplementedError, match="Causal"): + with pytest.raises(NotImplementedError, match="only supports FULL attention mask"): attn(q, k, v, batch_size=2, attention_mask=PredefinedAttentionMask.CAUSAL) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py index 9576022e4e64..0e8f4e830c87 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py @@ -620,50 +620,6 @@ def test_ulysses_attention_with_mask(self): """Test UlyssesAttention with attention mask.""" run_test_in_distributed(world_size=2, test_fn=_logic_ulysses_with_mask, use_cuda=True) - def test_ulysses_vs_standard_attention_single_gpu(self): - """Compare UlyssesAttention with standard attention on single GPU.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - - if not torch.cuda.is_available(): - pytest.skip("Test requires CUDA") - - batch = 2 - seq = 16 - num_heads = 8 - head_dim = 64 - device = torch.device("cuda:0") - - inner = VanillaAttention(num_heads=num_heads, head_dim=head_dim) - ulysses_attn = UlyssesAttention( - inner_backend=inner, - process_group=None, - ) - - torch.manual_seed(42) - q = torch.randn(batch, seq, num_heads, head_dim, device=device) - k = torch.randn(batch, seq, num_heads, head_dim, device=device) - v = torch.randn(batch, seq, num_heads, head_dim, device=device) - - ulysses_output = ulysses_attn(q, k, v) - - q_std = q.transpose(1, 2) # [B, H, S, D] - k_std = k.transpose(1, 2) - v_std = v.transpose(1, 2) - - std_output = F.scaled_dot_product_attention( - q_std, k_std, v_std, scale=1.0 / math.sqrt(head_dim), dropout_p=0.0 - ) - std_output = std_output.transpose(1, 2).contiguous() # [B, S, H, D] - - torch.testing.assert_close( - ulysses_output, - std_output, - rtol=1e-4, - atol=1e-4, - msg="Ulysses attention output differs from standard attention", - ) - def test_ulysses_vs_standard_attention_multi_gpu(self): """Compare UlyssesAttention across GPUs with standard attention on full sequence.""" run_test_in_distributed( From e5d0bc2494158b5954f1a5f648437e806aa16531 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 15 May 2026 15:05:49 -0700 Subject: [PATCH 13/13] address comments Signed-off-by: Shreyas Misra --- docs/source/models/visual-generation.md | 4 +- .../visual_gen/attention_backend/parallel.py | 6 +- tensorrt_llm/_torch/visual_gen/config.py | 10 +-- tensorrt_llm/_torch/visual_gen/mapping.py | 13 +++- .../models/ltx2/transformer_ltx2.py | 29 +++++++- .../visual_gen/models/wan/transformer_wan.py | 2 +- .../_torch/visual_gen/pipeline_loader.py | 2 + tensorrt_llm/_torch/visual_gen/utils.py | 24 ++++--- .../multi_gpu/test_visual_gen_mapping.py | 69 ++++++++++++++++++- 9 files changed, 135 insertions(+), 24 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 9486912b76f8..943016b9f97d 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -38,7 +38,7 @@ 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 | 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 | @@ -118,7 +118,7 @@ TeaCache caches transformer outputs when timestep embeddings change slowly betwe - **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 - +- **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 diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index 29c28ef05057..efa282350808 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -531,7 +531,9 @@ def _ensure_buffers(self, q: torch.Tensor, k: torch.Tensor) -> None: if key == self._buf_key: return self._kv_bufs = k.new_empty(2, 2, B, S, H_kv, D) - self._out_buf = q.new_empty(B, S, H, 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 @@ -593,6 +595,8 @@ def forward( 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 diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index d39c2127afdb..7ddb2f26f488 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -215,10 +215,12 @@ def seq_parallel_size(self) -> int: """ attn2d_size = self.dit_attn2d_row_size * self.dit_attn2d_col_size if attn2d_size > 1: - return attn2d_size - if self.dit_ring_size > 1: - return self.dit_ring_size - return self.dit_ulysses_size + cp_size = attn2d_size + elif self.dit_ring_size > 1: + cp_size = self.dit_ring_size + else: + cp_size = 1 + return cp_size * self.dit_ulysses_size @property def n_workers(self) -> int: diff --git a/tensorrt_llm/_torch/visual_gen/mapping.py b/tensorrt_llm/_torch/visual_gen/mapping.py index 4360ec890f55..2cf291227ffc 100644 --- a/tensorrt_llm/_torch/visual_gen/mapping.py +++ b/tensorrt_llm/_torch/visual_gen/mapping.py @@ -150,7 +150,18 @@ def __init__( # ------------------------------------------------------------------ def build_mesh(self): cls = DeviceMeshTopologyImpl + expected_shape = tuple(self._dim_sizes[d] for d in self._dim_names) if cls.device_mesh is not None: + cached_dim_names = tuple(cls.device_mesh.mesh_dim_names) + cached_shape = tuple(int(x) for x in cls.device_mesh.mesh.shape) + if cached_dim_names != self._dim_names or cached_shape != expected_shape: + raise RuntimeError( + "VisualGenMapping.build_mesh reusing incompatible cached device_mesh: " + f"cached dims={cached_dim_names}, shape={cached_shape}; " + f"requested dims={self._dim_names}, shape={expected_shape}. " + "Create a mesh cache keyed by (dim_names, dim_sizes) or reset " + "DeviceMeshTopologyImpl.device_mesh before constructing a different topology." + ) return shape = tuple(self._dim_sizes[d] for d in self._dim_names) @@ -391,7 +402,7 @@ def vae_group(self) -> Optional[ProcessGroup]: @property def vae_adj_groups(self) -> list[Optional[ProcessGroup]]: return self._vae_adj_groups - + def seq_group(self) -> Optional[ProcessGroup]: """Process group spanning (cp × ulysses) for combined sequence-axis sharding.""" cls = DeviceMeshTopologyImpl diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 774ec82cc872..981cb3c35557 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -833,6 +833,8 @@ def __init__( vgm, num_attention_heads=num_attention_heads if model_type.is_video_enabled() else None, ) + self._cp_size = vgm.cp_size if vgm is not None else 1 + self._ulysses_size = vgm.ulysses_size if vgm is not None else 1 if ( self._sharder.is_active and vgm is not None @@ -1147,14 +1149,27 @@ def _shard_transformer_args(self, args: TransformerArgs) -> TransformerArgs: """ seq_len = args.x.shape[1] sh = self._sharder + pe_seq_dim = ( + 2 + if args.positional_embeddings is not None and args.positional_embeddings[0].ndim == 4 + else 1 + ) + cross_pe_seq_dim = ( + 2 + if args.cross_positional_embeddings is not None + and args.cross_positional_embeddings[0].ndim == 4 + else 1 + ) return replace( args, x=sh.shard(args.x, dim=1), timesteps=sh.shard(args.timesteps, dim=1, expected_seq_len=seq_len), embedded_timestep=sh.shard(args.embedded_timestep, dim=1, expected_seq_len=seq_len), - positional_embeddings=sh.shard_rope(args.positional_embeddings, seq_len=seq_len), + positional_embeddings=sh.shard_rope( + args.positional_embeddings, seq_len=seq_len, seq_dim=pe_seq_dim + ), cross_positional_embeddings=sh.shard_rope( - args.cross_positional_embeddings, seq_len=seq_len + args.cross_positional_embeddings, seq_len=seq_len, seq_dim=cross_pe_seq_dim ), cross_scale_shift_timestep=sh.shard( args.cross_scale_shift_timestep, dim=1, expected_seq_len=seq_len @@ -1177,7 +1192,15 @@ def configure_audio_ulysses(self, audio_seq_len: int) -> None: self._audio_is_sharded = False return - self._audio_is_sharded = audio_seq_len % self._sharder.size == 0 + is_divisible = (audio_seq_len % self._sharder.size) == 0 + if not is_divisible and self._cp_size > 1 and self._ulysses_size == 1: + raise ValueError( + f"audio_seq_len ({audio_seq_len}) must be divisible by seq_size " + f"({self._sharder.size}) when CP is active without Ulysses " + "(Ring/Attention2D has no plain-attention fallback)." + ) + + self._audio_is_sharded = is_divisible for block in self.transformer_blocks: target = block.inner if isinstance(block, LTX2CacheDiTPattern0BlockWrapper) else block target._audio_is_sharded = self._audio_is_sharded diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index 7e02b883dd86..09a5ac02579e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -646,7 +646,7 @@ def forward( # Shard sequence + matching RoPE across ranks (no-op when sharder is inactive). seq_len = x.shape[1] x = self.sharder.shard(x, dim=1) - rope = self.sharder.shard_rope((freqs_cos, freqs_sin), seq_len=seq_len) + rope = self.sharder.shard_rope((freqs_cos, freqs_sin), seq_len=seq_len, seq_dim=1) if rope is not None: freqs_cos, freqs_sin = rope diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index 0ff5be1dd203..34923fc51e71 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -106,6 +106,8 @@ def _setup_visual_gen_mapping(self, config: DiffusionModelConfig) -> None: if self.args is not None: ws = dist.get_world_size() if dist.is_initialized() else 1 rk = dist.get_rank() if dist.is_initialized() else 0 + # NOTE: Might need to instantiate multiple VisualGenMapping in the future to + # handle different parallelism strategies for different models/pipelines. vgm = VisualGenMapping( ws, rk, diff --git a/tensorrt_llm/_torch/visual_gen/utils.py b/tensorrt_llm/_torch/visual_gen/utils.py index e4572ed1c1d3..6837b83b60ee 100644 --- a/tensorrt_llm/_torch/visual_gen/utils.py +++ b/tensorrt_llm/_torch/visual_gen/utils.py @@ -96,7 +96,7 @@ def from_vgm( size = vgm.seq_size rank = vgm.seq_rank - group = vgm.seq_group + group = vgm.seq_group() if size > 1 and group is None: raise ValueError( @@ -204,23 +204,25 @@ def shard_rope( self, rope: Optional[Tuple[torch.Tensor, torch.Tensor]], seq_len: int, + *, + seq_dim: int, ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Shard a ``(cos, sin)`` RoPE pair along whichever axis equals ``seq_len``. + """Shard a ``(cos, sin)`` RoPE pair along its sequence axis. - Handles both ``[B, S, D]`` (dim 1) and ``[B, H, S, D]`` (dim 2) - RoPE layouts that occur across visual_gen models. Returns the - pair unchanged when the seq axis cannot be inferred unambiguously - (e.g. ``seq_len == head_dim``); callers can fall back to explicit - ``shard(..., dim=...)`` calls in that case. + Callers must pass ``seq_dim`` explicitly based on the known RoPE + layout at the call site. """ if rope is None or not self.is_active: return rope cos, sin = rope - candidates = [d for d, n in enumerate(cos.shape) if n == seq_len] - if len(candidates) != 1: - return rope - d = candidates[0] + d = seq_dim if seq_dim >= 0 else cos.ndim + seq_dim + if d < 0 or d >= cos.ndim: + raise ValueError(f"seq_dim ({seq_dim}) is out of range for RoPE ndim ({cos.ndim}).") + if cos.shape[d] != seq_len: + raise ValueError( + f"RoPE seq_dim ({d}) has size {cos.shape[d]}, expected seq_len ({seq_len})." + ) return (self.shard(cos, dim=d), self.shard(sin, dim=d)) # ------------------------------------------------------------------ diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 6bad931bc85e..85768f36e55a 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -121,7 +121,7 @@ def test_product_mismatch_raises(self): def test_parallel_vae_size_cannot_exceed_world_size(self): with pytest.raises(ValueError, match="cannot exceed world_size"): VisualGenMapping(world_size=1, rank=0, parallel_vae_size=2) - + def test_internal_dim_order(self): vgm = VisualGenMapping(world_size=1, rank=0) assert vgm._dim_names == ("cfg", "tp", "cp", "ulysses") @@ -390,6 +390,70 @@ def _logic_vae_group_full_world_cfg2_ulysses2(rank, world_size): ) +def _logic_cfg2_ring2_ulysses2(rank, world_size): + """Validate combined CFG+Ring+Ulysses topology on 8 GPUs. + + Layout (cfg-tp-cp-ulysses) with cfg=2, tp=1, cp(ring)=2, ulysses=2: + rank = cfg * 4 + cp * 2 + ulysses + """ + from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl + + DeviceMeshTopologyImpl.device_mesh = None + VisualGenMapping.seq_mesh = None + + vgm = VisualGenMapping( + world_size=world_size, + rank=rank, + cfg_size=2, + ring_size=2, + ulysses_size=2, + ) + + assert world_size == 8 + expected_cfg_rank = rank // 4 + expected_cp_rank = (rank // 2) % 2 + expected_ulysses_rank = rank % 2 + expected_seq_rank = expected_cp_rank * 2 + expected_ulysses_rank + + assert vgm.cfg_rank == expected_cfg_rank + assert vgm.cp_rank == expected_cp_rank + assert vgm.ring_rank == expected_cp_rank + assert vgm.ulysses_rank == expected_ulysses_rank + assert vgm.seq_rank == expected_seq_rank + assert vgm.seq_size == 4 + + assert vgm.cfg_group is not None + assert vgm.cp_group is not None + assert vgm.ring_group is not None + assert vgm.ulysses_group is not None + assert vgm.seq_group() is not None + + assert dist.get_world_size(vgm.cfg_group) == 2 + assert dist.get_world_size(vgm.cp_group) == 2 + assert dist.get_world_size(vgm.ring_group) == 2 + assert dist.get_world_size(vgm.ulysses_group) == 2 + assert dist.get_world_size(vgm.seq_group()) == 4 + + device = torch.device(f"cuda:{rank}") + one = torch.ones(1, device=device) + + x = one.clone() + dist.all_reduce(x, group=vgm.cfg_group) + assert x.item() == 2.0, f"Rank {rank}: cfg all_reduce expected 2, got {x.item()}" + + x = one.clone() + dist.all_reduce(x, group=vgm.cp_group) + assert x.item() == 2.0, f"Rank {rank}: cp all_reduce expected 2, got {x.item()}" + + x = one.clone() + dist.all_reduce(x, group=vgm.ulysses_group) + assert x.item() == 2.0, f"Rank {rank}: ulysses all_reduce expected 2, got {x.item()}" + + x = one.clone() + dist.all_reduce(x, group=vgm.seq_group()) + assert x.item() == 4.0, f"Rank {rank}: seq all_reduce expected 4, got {x.item()}" + + @pytest.mark.skipif(not MODULES_AVAILABLE, reason="Modules not available") class TestMultiGPU: def test_default_order_cfg2_ulysses2(self): @@ -407,3 +471,6 @@ def test_vae_group_and_adj_groups(self): def test_vae_group_full_world_cfg2_ulysses2(self): _run_multi_gpu(4, _logic_vae_group_full_world_cfg2_ulysses2) + + def test_cfg2_ring2_ulysses2(self): + _run_multi_gpu(8, _logic_cfg2_ring2_ulysses2)