From 3a401a94b9716842cb28e7fade4324e512da2f49 Mon Sep 17 00:00:00 2001 From: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:57:04 -0700 Subject: [PATCH 1/3] Refactor VisualGen Attention; Fix SageAttn LTX-2, Cosmos3 - Fix SageAttention cross attention * Separate is_cross definition in thop vs. in python: thop needs specialized is_cross treatment only if KV cache is paged or ragged, while separateQkv cross routes through regular context attention. * Fix SageAttention invalid memory at cross-attn: derive KV size from the right params (kv_len was delicated to generation. use total_kv instead) - Metadata handling: align with LLM * Assign separate attention metadata to each attention site e.g. in Wan-like models, self-attention and cross-attention will use different metadata objects. * Each transformer class implementation will expose a `create_attn_metadata` to create metadata objects for each of its attention sites * Pipeline will create, maintain, and reuse attention metadata objects throughout the denoising process. * Pipeline will register metadata as CUDA graph keys. - Cleanup VisualGen attention modules * Resolve cross-attention definition: never infer cross/self attention from qkv_format: Use an explicit is_cross flag. * qkv_format only indicates whether the module uses fused to_qkv or separate to_q, to_k, to_v Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- .../_torch/attention_backend/fmha/fallback.py | 3 +- .../_torch/attention_backend/trtllm.py | 18 +- .../_torch/visual_gen/ENGINEERING_CRITERIA.md | 5 + .../attention_backend/cute_dsl/fmha.py | 20 +- .../attention_backend/cute_dsl/vsa.py | 12 + .../attention_backend/flash_attn4.py | 13 +- .../visual_gen/attention_backend/interface.py | 24 ++ .../visual_gen/attention_backend/metadata.py | 232 +++++++++++++++ .../visual_gen/attention_backend/parallel.py | 47 +++- .../visual_gen/attention_backend/trtllm.py | 224 ++++----------- .../visual_gen/attention_backend/utils.py | 25 +- .../visual_gen/attention_backend/vanilla.py | 11 +- tensorrt_llm/_torch/visual_gen/config.py | 13 - .../models/cosmos3/pipeline_cosmos3.py | 23 +- .../models/cosmos3/transformer_cosmos3.py | 146 +++++++++- .../visual_gen/models/flux/attention.py | 11 +- .../visual_gen/models/flux/pipeline_flux.py | 9 +- .../visual_gen/models/flux/pipeline_flux2.py | 11 + .../models/flux/transformer_flux.py | 45 +++ .../models/flux/transformer_flux2.py | 43 +++ .../visual_gen/models/ltx2/pipeline_ltx2.py | 44 ++- .../models/ltx2/pipeline_ltx2_two_stages.py | 11 + .../models/ltx2/transformer_ltx2.py | 263 ++++++++++++++++-- .../_torch/visual_gen/models/modeling.py | 68 +++-- .../models/qwen_image/pipeline_qwen_image.py | 14 + .../qwen_image/pipeline_qwen_image_edit.py | 22 ++ .../qwen_image/transformer_qwen_image.py | 61 +++- .../pipeline_qwen_image_layered.py | 12 + .../transformer_qwen_image_layered.py | 4 + .../visual_gen/models/wan/pipeline_wan.py | 28 +- .../visual_gen/models/wan/pipeline_wan_i2v.py | 23 +- .../models/wan/pipeline_wan_utils.py | 32 ++- .../visual_gen/models/wan/transformer_wan.py | 100 ++++++- .../_torch/visual_gen/modules/attention.py | 63 +++-- tensorrt_llm/_torch/visual_gen/pipeline.py | 18 ++ .../_torch/visual_gen/attn_metadata_utils.py | 96 +++++++ .../multi_gpu/test_attn2d_attention.py | 86 +++++- .../test_cosmos3_transformer_parallel.py | 4 + .../test_flux2_transformer_parallel.py | 4 + .../visual_gen/multi_gpu/test_flux_tp.py | 14 +- .../visual_gen/multi_gpu/test_flux_ulysses.py | 18 +- .../multi_gpu/test_ltx2_async_ulysses.py | 8 +- .../visual_gen/multi_gpu/test_ltx2_ulysses.py | 23 +- .../multi_gpu/test_ring_attention.py | 47 +++- .../visual_gen/multi_gpu/test_tp_attention.py | 6 +- .../multi_gpu/test_ulysses_attention.py | 15 +- .../multi_gpu/test_ulysses_sage_attention.py | 4 - .../visual_gen/multi_gpu/test_wan_tp.py | 50 +++- .../test_wan_transformer_parallel.py | 19 ++ .../visual_gen/test_attention_cute_dsl_vsa.py | 12 +- .../visual_gen/test_attention_integration.py | 113 ++++++-- .../visual_gen/test_attention_metadata.py | 187 +++++++++++++ .../_torch/visual_gen/test_attention_perf.py | 50 +++- .../visual_gen/test_cosmos3_distilled.py | 9 +- .../visual_gen/test_cosmos3_transformer.py | 10 + .../visual_gen/test_fa4_key_padding_mask.py | 44 ++- .../test_flux2_image_conditioning.py | 18 +- .../_torch/visual_gen/test_flux_attention.py | 35 ++- .../_torch/visual_gen/test_flux_pipeline.py | 16 +- .../visual_gen/test_flux_transformer.py | 6 + .../_torch/visual_gen/test_ltx2_attention.py | 36 ++- .../_torch/visual_gen/test_ltx2_pipeline.py | 15 +- .../visual_gen/test_ltx2_transformer.py | 107 ++++++- .../test_qwen_image_layered_registry.py | 6 + .../visual_gen/test_qwen_image_pipeline.py | 18 +- .../test_qwen_image_pipeline_config.py | 14 +- .../visual_gen/test_qwen_image_registry.py | 6 + .../test_trtllm_attention_metadata.py | 64 ----- .../test_vanilla_key_padding_mask.py | 13 +- .../visual_gen/test_wan21_t2v_pipeline.py | 37 ++- .../_torch/visual_gen/test_wan_transformer.py | 22 ++ 71 files changed, 2411 insertions(+), 519 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/metadata.py create mode 100644 tests/unittest/_torch/visual_gen/attn_metadata_utils.py create mode 100644 tests/unittest/_torch/visual_gen/test_attention_metadata.py delete mode 100644 tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 4cfcda50bac2..88b518d5eec7 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -109,7 +109,8 @@ def forward( max_context_length=metadata.max_context_length, max_seq_len=metadata.max_seq_len, trtllm_gen_jit_warmup=metadata.trtllm_gen_jit_warmup, - is_cross=metadata.is_cross, + # Selects the KV layout, not merely unequal Q/KV lengths. + is_cross=metadata.is_cross_with_kv_cache, # --- Per-call (AttentionForwardArgs) --- out_scale=forward_args.out_scale, kv_scale_orig_quant=forward_args.kv_scale_orig_quant, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 6dfb55263658..b89c620945d6 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -86,6 +86,16 @@ class TrtllmAttentionMetadata(AttentionMetadata): # when beam search is enabled. beam_width: int = 1 + @property + def is_cross_with_kv_cache(self) -> bool: + """Cross-attention whose K/V is read from a KV-cache pool. + + thop uses ``is_cross`` flag to handle cross attention with kv cache + (``Q_CONTIGUOUS_KV`` / ``Q_PAGED_KV``, fed by ``cross_kv``), while + cross-attention with ``SEPARATE_QKV`` is handled without the flag. + """ + return self.is_cross and self.kv_cache_manager is not None + @property def effective_beam_width(self) -> int: # Only use this for the fallback kernel's beam_width argument. @@ -1798,7 +1808,13 @@ def forward( and metadata._seq_lens_cuda is not None): metadata.update_blackwell_first_sparse_mask_offset() - if metadata.is_cross: + # The ``cacheless`` (separate-QKV) cross attention does not need cross- + # attention specific handling + _cacheless_cross = (metadata.is_cross + and metadata.kv_cache_manager is None + and k is not None and v is not None) + + if metadata.is_cross and not _cacheless_cross: if k is not None and v is not None: k_flat = k.contiguous().view(k.shape[0], -1) v_flat = v.contiguous().view(v.shape[0], -1) diff --git a/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md b/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md index 240e5833ae0c..06369984ccae 100644 --- a/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md +++ b/tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md @@ -29,6 +29,11 @@ end-to-end workflow. T2V / TI2V) with LPIPS — or an equivalent metric — at an explicit threshold, comparing against a baseline image/video produced by a reference framework. +6. **Attention metadata is a forward argument.** Models take metadata as + required `forward()` parameters; the pipeline builds it per forward + from the model's `attn_metadata_spec()`. Each attention *site* + (e.g., LTX-2's Vid2Vid, Vid2Sound, etc.), with its own Q/KV lengths, + gets its own metadata object. ## 2. API diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index ada69e1acd9e..6cb65fb0af4b 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -28,7 +28,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ....attention_backend.interface import PredefinedAttentionMask +from ....attention_backend.interface import AttentionMetadata, PredefinedAttentionMask from ..interface import AttentionBackend, AttentionTensorLayout _cute_dsl_import_error: BaseException | None = None @@ -740,6 +740,7 @@ def forward( k: torch.Tensor, v: torch.Tensor, *, + attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, **kwargs, ) -> torch.Tensor: @@ -752,12 +753,17 @@ def forward( q: Query tensor [batch_size, seq_len, num_heads, head_dim] k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] + attn_metadata: Unused; the CuTe DSL kernels derive everything from + tensor shapes. Preserved for future metadata-consuming variants + of this path. attention_mask: Attention mask type (CAUSAL or FULL) Returns: Output tensor [batch_size, seq_len, num_heads, head_dim] """ - output, _ = self.forward_with_lse(q, k, v, attention_mask=attention_mask, **kwargs) + output, _ = self.forward_with_lse( + q, k, v, attn_metadata=attn_metadata, attention_mask=attention_mask, **kwargs + ) return output def forward_with_lse( @@ -765,6 +771,8 @@ def forward_with_lse( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + *, + attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -777,12 +785,20 @@ def forward_with_lse( always in float32. Used for numerically stable combination of partial attention results in Attention2D parallelism. """ + _ = attn_metadata q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask) output, lse = self._fwd(q, k, v, is_causal, **kwargs) if output.dtype != origin_dtype: output = output.to(origin_dtype) return output, lse.transpose(1, 2) + @property + def requires_metadata(self) -> bool: + """While CuTe's kernel supports varlen, its not wired to CuTeDSLAttention. + + No metadata referred.""" + return False + @classmethod def support_lse(cls) -> bool: return True diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py index 741f82097ca2..945bec112bc7 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py @@ -30,6 +30,7 @@ import torch import torch.nn.functional as F +from ....attention_backend.interface import AttentionMetadata from ..interface import AttentionBackend, AttentionTensorLayout _vsa_import_error = None @@ -300,6 +301,7 @@ def forward( k: torch.Tensor, v: torch.Tensor, *, + attn_metadata: AttentionMetadata, gate_compress: Optional[torch.Tensor] = None, gate_fine: Optional[torch.Tensor] = None, **kwargs, @@ -307,6 +309,10 @@ def forward( """ VSA forward: coarse mean-pool + fine block-sparse top-K. + VSA carries its own sparse-layout metadata through + ``set_vsa_forward_context``; ``attn_metadata`` is unused here and + preserved for future metadata-consuming variants of this path. + Args: q, k, v: [B, S, H, D] in original (un-tiled) token order. gate_compress: [B, S, H, D] G_c gate weighting the coarse branch O_c. @@ -316,6 +322,7 @@ def forward( Returns: [B, S, H, D] in the same original token order. """ + _, _ = attn_metadata, kwargs if gate_compress is None: raise ValueError( "VSAAttention requires gate_compress. " @@ -399,6 +406,11 @@ def forward( return gate_compress * o_c_full + gate_fine * o_f return gate_compress * o_c_full + o_f + @property + def requires_metadata(self) -> bool: + """Currently VSAAttention only supports non-varlen batching. No metadata referred.""" + return False + @classmethod def support_lse(cls) -> bool: return False 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 532e86e91e08..c6cad4a46e53 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -24,7 +24,7 @@ import torch -from ...attention_backend.interface import PredefinedAttentionMask +from ...attention_backend.interface import AttentionMetadata, PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout _flash_attn_fwd_import_error = None @@ -122,6 +122,7 @@ def forward( k: torch.Tensor, v: torch.Tensor, *, + attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, key_padding_mask: Optional[torch.Tensor] = None, **kwargs, @@ -135,6 +136,7 @@ def forward( q: Query tensor [batch_size, seq_len, num_heads, head_dim] k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] + attn_metadata: Unused; FA4 derives everything from tensor shapes. attention_mask: Attention mask type (CAUSAL or FULL) key_padding_mask: Optional ``[B, S_kv]`` bool tensor; True = valid, False = pad. Translated to FA4's ``seqused_k = mask.sum(dim=1)`` @@ -147,6 +149,7 @@ def forward( q, k, v, + attn_metadata=attn_metadata, attention_mask=attention_mask, key_padding_mask=key_padding_mask, **kwargs, @@ -158,6 +161,8 @@ def forward_with_lse( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + *, + attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, key_padding_mask: Optional[torch.Tensor] = None, **kwargs, @@ -171,6 +176,7 @@ def forward_with_lse( always in float32. Used for numerically stable combination of partial attention results in Attention2D parallelism. """ + _, _ = attn_metadata, kwargs q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask) seqused_k = None if key_padding_mask is not None: @@ -190,6 +196,11 @@ def forward_with_lse( output = output.to(origin_dtype) return output, lse + @property + def requires_metadata(self) -> bool: + """Currently FlashAttn4Attention only supports non-varlen batching. No metadata referred.""" + return False + @classmethod def support_lse(cls) -> bool: return True diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py b/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py index 7a866acf5879..f0c02eee7b01 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py @@ -20,9 +20,12 @@ from abc import ABC, abstractmethod from enum import Enum +from typing import Type import torch +from ...attention_backend.interface import AttentionMetadata + class AttentionTensorLayout(str, Enum): """ @@ -42,8 +45,17 @@ class AttentionBackend(ABC): Every backend must implement ``forward`` and declare a ``preferred_layout``. Backends pick the kwargs they need from the caller and ignore the rest via ``**kwargs``. + + ``attn_metadata`` is a required argument on every forward, mirroring the LLM + backend. Backends that do not consume metadata (VANILLA, FA4, CuTe DSL) accept + and ignore it; only the TRTLLM backend reads it. See + ``visual_gen/attention_backend/metadata.py`` for the site contract. """ + #: Concrete metadata type this backend consumes; callers construct from + #: it, as the LLM engine does. + Metadata: Type[AttentionMetadata] = AttentionMetadata + def __call__(self, *args, **kwargs) -> torch.Tensor: return self.forward(*args, **kwargs) @@ -53,6 +65,8 @@ def forward( q: torch.Tensor, k: torch.Tensor | None = None, v: torch.Tensor | None = None, + *, + attn_metadata: AttentionMetadata, **kwargs, ) -> torch.Tensor: ... @@ -60,11 +74,21 @@ def forward( @abstractmethod def preferred_layout(self) -> AttentionTensorLayout: ... + @property + def requires_metadata(self) -> bool: + """Queries if the configured attention module requires a prepared metadata to run. + + For size-aligned diffusion models, many attention backends can run without a metadata. + """ + return True + def forward_with_lse( self, q: torch.Tensor, k: torch.Tensor | None = None, v: torch.Tensor | None = None, + *, + attn_metadata: AttentionMetadata, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError( diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/metadata.py b/tensorrt_llm/_torch/visual_gen/attention_backend/metadata.py new file mode 100644 index 000000000000..2b9fad90c883 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/metadata.py @@ -0,0 +1,232 @@ +# 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. +"""Attention metadata helpers for VisualGen. + +VisualGen reuses the shared-core ``AttentionMetadata`` rather than defining its +own: it already models no-KV-cache operation, mixed Q/KV lengths, and several +metadata objects held side by side. + +An *attention site* is one place a model attends, and gets its own metadata +object -- e.g. WAN has ``self``, ``cross_text`` and ``cross_image``. +""" + +from typing import Sequence, Type + +import torch + +from tensorrt_llm.mapping import Mapping + +from ...attention_backend.interface import AttentionMetadata, AttentionRuntimeFeatures + +__all__ = [ + "make_diffusion_attn_metadata", + "create_diffusion_attn_metadata", + "prepare_diffusion_attn_metadata", +] + +SeqLens = int | Sequence[int] | torch.Tensor + + +@torch.compiler.disable +def _seqlens_as_s32(value: SeqLens, batch_size: int) -> torch.Tensor: + """Normalize a per-batch sequence length spec to an int32 CPU tensor.""" + if isinstance(value, torch.Tensor): + tensor = value.detach().flatten().to(device="cpu", dtype=torch.int32) + elif isinstance(value, int): + tensor = torch.full((batch_size,), value, dtype=torch.int32) + else: + tensor = torch.tensor(list(value), dtype=torch.int32) + + if tensor.shape[0] != batch_size: + raise ValueError(f"sequence lengths have batch {tensor.shape[0]} but {batch_size=}") + return tensor + + +@torch.compiler.disable +def _cu_seqlens(seq_lens: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + cu = torch.zeros(seq_lens.shape[0] + 1, dtype=torch.int32) + torch.cumsum(seq_lens, dim=0, dtype=torch.int32, out=cu[1:]) + device = getattr(attn_metadata, "seq_lens_cuda", None) + if device is not None: + cu = cu.to(device=device.device, non_blocking=True) + return cu + + +@torch.compiler.disable +def _batch_changed(attn_metadata: AttentionMetadata, batch_size: int) -> bool: + current = attn_metadata.seq_lens + return current is None or current.shape[0] != batch_size + + +@torch.compiler.disable +def _is_already_prepared( + attn_metadata: AttentionMetadata, + batch_size: int, + q_lens: torch.Tensor, + kv_lens: torch.Tensor | None, +) -> bool: + """Whether the site already carries exactly this shape.""" + if attn_metadata.num_contexts != batch_size: + return False + if _batch_changed(attn_metadata, batch_size): + return False + if not torch.equal(attn_metadata.seq_lens, q_lens): + return False + if (kv_lens is not None) != attn_metadata.is_cross: + return False + if kv_lens is not None and not torch.equal(attn_metadata.seq_lens_kv, kv_lens): + return False + return True + + +@torch.compiler.disable +def create_diffusion_attn_metadata( + metadata_cls: Type[AttentionMetadata], + *, + max_batch_size: int, + max_seq_len: int, + mapping: Mapping | None = None, +) -> AttentionMetadata: + """Allocate one no-KV-cache attention metadata site. Can't be dynamo-traced. + + Args: + metadata_cls: Concrete metadata type, from the backend's ``Metadata``. + max_batch_size: Upper bound on sequences per forward. + max_seq_len: Upper bound on the longest sequence, Q or KV. + mapping: Defaults to single-rank; diffusion parallelism lives in the + ``parallel.py`` wrappers, not the kernel. + """ + if max_batch_size <= 0: + raise ValueError(f"max_batch_size must be positive, got {max_batch_size}") + if max_seq_len <= 0: + raise ValueError(f"max_seq_len must be positive, got {max_seq_len}") + + metadata = metadata_cls( + max_num_requests=max_batch_size, + max_num_tokens=max_batch_size * max_seq_len, + max_num_sequences=max_batch_size, + kv_cache_manager=None, # Diffusion attention runs without a KV cache. + mapping=mapping if mapping is not None else Mapping(), + runtime_features=AttentionRuntimeFeatures(), + # Makes the `seq_lens` setter copy in place rather than reallocating + # `seq_lens_cuda` on every assignment. + is_cuda_graph=True, + ) + # `max_seq_len` is a property with a manual setter on TrtllmAttentionMetadata + # (required for the no-cache path) and a plain attribute elsewhere. + metadata.max_seq_len = max_seq_len + return metadata + + +@torch.compiler.disable +def prepare_diffusion_attn_metadata( + attn_metadata: AttentionMetadata, + *, + batch_size: int, + q_seq_lens: SeqLens, + kv_seq_lens: SeqLens | None = None, +) -> AttentionMetadata: + """Populate and prepare one site in place, then return it. + + Every sequence is a context request and there is no KV cache, mirroring + ``_prepare_qwen_vl_vision_attn_metadata``. Can't be dynamo-traced. + + Args: + attn_metadata: A site from :func:`create_diffusion_attn_metadata`. + batch_size: Number of sequences in this forward. + q_seq_lens: Query length per sequence. + kv_seq_lens: Key/value length per sequence; ``None`` for self-attention. + """ + q_lens = _seqlens_as_s32(q_seq_lens, batch_size) + + capacity = getattr(attn_metadata, "max_num_requests", None) + if capacity is not None and batch_size > capacity: + raise ValueError( + f"batch_size={batch_size} exceeds the site's allocated capacity " + f"({capacity}). Size the site for its actual batch in " + f"create_diffusion_attn_metadata()." + ) + + kv_lens = None if kv_seq_lens is None else _seqlens_as_s32(kv_seq_lens, batch_size) + + if _is_already_prepared(attn_metadata, batch_size, q_lens, kv_lens): + # Skipping leaves the device-side length buffers untouched between + # CUDA graph replays. + return attn_metadata + + # The in-place `seq_lens` copy needs matching shapes; drop the stale + # buffers on a batch change so the setter reallocates instead of raising. + if _batch_changed(attn_metadata, batch_size): + attn_metadata._seq_lens_cuda = None + attn_metadata._seq_lens_kv_cuda = None + + # Order matters: `num_contexts` feeds `context_lens`, which the no-cache + # branch of TrtllmAttentionMetadata.prepare() uses to derive `prompt_lens`. + attn_metadata.num_contexts = batch_size + attn_metadata.request_ids = list(range(batch_size)) + attn_metadata.seq_lens = q_lens + # Always assign: None restores the getter's identity fall-back to + # `seq_lens`, which is what makes `is_cross` False again. + attn_metadata.seq_lens_kv = kv_lens + + attn_metadata.cu_q_seqlens = _cu_seqlens(attn_metadata.seq_lens, attn_metadata) + attn_metadata.cu_kv_seqlens = ( + attn_metadata.cu_q_seqlens + if kv_lens is None + else _cu_seqlens(attn_metadata.seq_lens_kv, attn_metadata) + ) + + # Full `prepare()`, not `prepare_encoder_only()`: the latter binds KV + # lengths to the query lengths, wrong for a cross site. + attn_metadata.prepare() + return attn_metadata + + +@torch.compiler.disable +def make_diffusion_attn_metadata( + metadata_cls: Type[AttentionMetadata], + *, + batch_size: int, + q_seq_lens: SeqLens, + kv_seq_lens: SeqLens | None = None, +) -> AttentionMetadata: + """Allocate and prepare one attention site in a single call. + + The entry point a model's ``create_attn_metadata()`` uses, one call per + site. Sizes the site from the lengths it is given, so the caller only states + the shape the attention module will be called with. + + Can't be dynamo-traced; see :func:`prepare_diffusion_attn_metadata`. + + Args: + metadata_cls: Concrete metadata type, normally the model's + ``attn_backend_metadata_cls``. + batch_size: Number of sequences in this forward. + q_seq_lens: Query length, either shared by the batch or one per sequence. + kv_seq_lens: Key/value length; ``None`` for self-attention. + """ + q_lens = _seqlens_as_s32(q_seq_lens, batch_size) + max_seq_len = int(q_lens.max()) + if kv_seq_lens is not None: + max_seq_len = max(max_seq_len, int(_seqlens_as_s32(kv_seq_lens, batch_size).max())) + + return prepare_diffusion_attn_metadata( + create_diffusion_attn_metadata( + metadata_cls, max_batch_size=batch_size, max_seq_len=max_seq_len + ), + batch_size=batch_size, + q_seq_lens=q_seq_lens, + kv_seq_lens=kv_seq_lens, + ) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index fedabf4ceba8..593b6a4c3932 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -140,7 +140,8 @@ def forward( Forward pass with Ulysses sequence parallelism. q/k/v: [B, S/P, H, D] each. All other arguments are forwarded - transparently to the inner backend via ``**kwargs``. + transparently to the inner backend via ``**kwargs``, including + ``attn_metadata``, which is re-targeted to the post-all-to-all lengths. """ # Catches upstream floor-division bugs (e.g. num_heads // ulysses_size when # num_heads % ulysses_size != 0) before they corrupt the all-to-all. @@ -315,7 +316,7 @@ def forward_async( `("q", "k", "v")` to issue the small audio-Q first. Order is correctness-neutral (`_join_async` syncs all recv bufs). **attn_kwargs : forwarded to the wrapped inner attention backend - (mask, scale, etc.). + (mask, scale, ``attn_metadata``, etc.). Returns: output tensor in the caller's sharded layout `[B, S/P, H, D]`. @@ -550,6 +551,8 @@ def forward( Forward pass with Attention2D sequence parallelism. q: [B, S_q/P, H_q, D]. k/v: [B, S_kv/P, H_kv, D]. + + ``attn_metadata`` rides through ``**kwargs`` and is re-targeted to the full lengths. """ B, shard_seq_q, H_q, D = q.shape _, shard_seq_kv, H_kv, D_kv = k.shape @@ -838,6 +841,10 @@ def forward( return out.to(dtype=q.dtype) return out + @property + def requires_metadata(self) -> bool: + return self.inner.requires_metadata + @property def preferred_layout(self) -> AttentionTensorLayout: return self._preferred_layout @@ -851,6 +858,42 @@ def support_lse(cls) -> bool: return False +def get_ulysses_seq_lens( + q_seq_len: int, + kv_seq_len: int, + *, + visual_gen_mapping: Optional["VisualGenMapping"] = None, + enable_sequence_parallel: bool = True, + use_ulysses: bool = True, + ulysses_size: Optional[int] = None, +) -> tuple[int, int]: + """Computes how sharded seqlen transforms to seqlen seen by the inner attention backend. + + - Ulysses preprocessing is all2all: both local lengths scale by the Ulysses group size. + - Attention2D gathers Q along the row fiber and K/V along the column + fiber, so the two local lengths scale by different factors. + - Ring is the identity: every step feeds one local K/V chunk. + + Args: + ulysses_size: Overrides the mapping's Ulysses degree. + """ + if not enable_sequence_parallel or visual_gen_mapping is None: + return q_seq_len, kv_seq_len + + vgm = visual_gen_mapping + q_factor, kv_factor = 1, 1 + + if vgm.attn2d_row_size * vgm.attn2d_col_size > 1: + q_factor, kv_factor = vgm.attn2d_row_size, vgm.attn2d_col_size + + ulysses = vgm.ulysses_size if ulysses_size is None else ulysses_size + if use_ulysses and ulysses > 1: + q_factor *= ulysses + kv_factor *= ulysses + + return q_seq_len * q_factor, kv_seq_len * kv_factor + + def wrap_parallel_attention( attn: AttentionBackend, *, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index cd76bb2cc64f..45253faa086c 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -15,165 +15,59 @@ """ Diffusion TRTLLM Attention Backend -Wraps TrtllmAttention with simplified metadata for visual generation (diffusion) models. -Handles the specifics of no-KV-cache operation and fused QKV requirements. +Wraps TrtllmAttention for visual generation (diffusion) models, handling the +specifics of no-KV-cache operation and fused QKV requirements. """ -from typing import Optional, Union +from typing import Optional import torch -from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention_backend.interface import AttentionRuntimeFeatures, PredefinedAttentionMask +from ...attention_backend.interface import PredefinedAttentionMask from ...attention_backend.sparse.skip_softmax import SkipSoftmaxParams from ...attention_backend.trtllm import TrtllmAttention as BaseTrtllmAttention -from ...attention_backend.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata +from ...attention_backend.trtllm import TrtllmAttentionMetadata from .interface import AttentionBackend, AttentionTensorLayout -class TrtllmAttentionMetadata: - """ - Simplified metadata adapter for diffusion models using TRTLLM backend. - - Lazy initialization with auto-growing capacity: - - Metadata created only when capacity needs increase - - prepare() called only when seq_lens actually change - - Automatically reallocates when batch_size or seq_len exceeds current capacity - - Args: - device: Target device for tensors. - attention_metadata_state: Mutable model-scoped state shared by all - attention layers in one model instance. - """ - - def __init__( - self, - device: Optional[torch.device] = None, - attention_metadata_state: Optional[dict] = None, - ): - self.device = device or torch.device("cuda") - if attention_metadata_state is None: - raise ValueError( - "TRTLLM attention requires `attention_metadata_state` to be provided " - "by visual-gen config for model-scoped metadata sharing." - ) - self._metadata_state = attention_metadata_state - - # Lazily created BaseTrtllmAttentionMetadata objects. Diffusion blocks - # can launch video and audio attention back-to-back with different - # sequence lengths, so keep separate metadata buffers per shape instead - # of mutating one shared object while kernels may still be in flight. - self._metadata_cache = self._metadata_state.setdefault("metadata_cache", {}) - self._metadata: Optional[BaseTrtllmAttentionMetadata] = None - - # Track prepared state - self._cached_seq_lens: Optional[torch.Tensor] = None - self._prepared = False - - def _needs_prepare(self, batch_size: int, seq_lens: torch.Tensor) -> bool: - """Check if we need to call prepare() (current request seq_lens or shared metadata object seq_lens changed). - - Assumes uniform sequence length per batch; if per-sample lengths vary, - we may need to check seq_lens tensor instead. - - In addition, multiple visual gen attention modules share one metadata object. A - different module may have prepared it for another sequence length even - when this wrapper's local cached seq_lens are unchanged. - """ - if not self._prepared: - return True - if self._cached_seq_lens is None: - return True - if self._cached_seq_lens.shape[0] != batch_size: - return True - if not torch.equal(self._cached_seq_lens[:batch_size], seq_lens): - return True - - metadata = self._metadata - if metadata is None: - return True - if getattr(metadata, "num_contexts", None) != batch_size: - return True - - max_seq_len = seq_lens.max().item() - if getattr(metadata, "max_seq_len", None) != max_seq_len: - return True - - metadata_seq_lens = getattr(metadata, "seq_lens", None) - if metadata_seq_lens is None or metadata_seq_lens.shape[0] < batch_size: - return True - if not torch.equal(metadata_seq_lens[:batch_size].to(seq_lens.device), seq_lens): - return True - - return False - - def _create_metadata(self, batch_size: int, max_seq_len: int) -> None: - """Create new metadata with given capacity.""" - self._metadata = BaseTrtllmAttentionMetadata( - max_num_requests=batch_size, - max_num_tokens=batch_size * max_seq_len, - max_num_sequences=batch_size, - kv_cache_manager=None, # No KV cache for diffusion - mapping=Mapping(), - runtime_features=AttentionRuntimeFeatures(), +def _check_metadata( + attn_metadata: TrtllmAttentionMetadata, + batch_size: int, + q_seq_len: int, + kv_seq_len: int, +) -> None: + """Validate that the metadata describes the tensors it is used with.""" + if attn_metadata is None: + raise ValueError( + "TrtllmAttention.forward requires `attn_metadata`. Build it with " + "visual_gen.attention_backend.metadata.create_diffusion_attn_metadata() " + "and prepare it with prepare_diffusion_attn_metadata()." ) - self._prepared = False # Reset prepare state on new metadata - - def _select_cached_metadata(self, cached) -> None: - self._metadata = cached["metadata"] - self._prepared = cached["prepared"] - self._cached_seq_lens = cached["seq_lens"] - def prepare( - self, - batch_size: int, - seq_lens: Union[int, torch.Tensor], - ) -> BaseTrtllmAttentionMetadata: - """ - Prepare metadata for a forward pass. - - Lazy behavior: - - Creates metadata only when capacity needs increase - - Calls prepare() only when (batch_size, max_seq_len) actually change - """ - if isinstance(seq_lens, int): - seq_lens_tensor = torch.full((batch_size,), seq_lens, dtype=torch.int32) - else: - seq_lens_tensor = seq_lens.to(dtype=torch.int32) - max_seq_len = seq_lens_tensor.max().item() - # Keep CUDA graph-captured metadata buffers stable per batch/seq-lens shape. - cache_key = (batch_size, tuple(int(x) for x in seq_lens_tensor.tolist())) - - cached = self._metadata_cache.get(cache_key) - if cached is None: - self._create_metadata(batch_size, max_seq_len) - cached = { - "metadata": self._metadata, - "prepared": False, - "seq_lens": None, - } - self._metadata_cache[cache_key] = cached - - self._select_cached_metadata(cached) - - if self._needs_prepare(batch_size, seq_lens_tensor): - cached_seq_lens = seq_lens_tensor.clone() - self._metadata.seq_lens = cached_seq_lens - self._metadata.num_contexts = batch_size - self._metadata.max_seq_len = max_seq_len - self._metadata.request_ids = list(range(batch_size)) - self._metadata.prepare() - - # Cache per-shape state without sharing the tensor across entries. - cached["prepared"] = True - cached["seq_lens"] = cached_seq_lens - - self._select_cached_metadata(cached) + seq_lens = attn_metadata.seq_lens + if seq_lens is None: + raise ValueError( + "`attn_metadata` has no seq_lens; call prepare_diffusion_attn_metadata() " + "before the forward pass." + ) + if seq_lens.shape[0] != batch_size: + raise ValueError( + f"attn_metadata batch_size mismatch: cached {seq_lens.shape[0]} != {batch_size=}." + ) - return self._metadata + # `seq_lens` / `seq_lens_kv` are host tensors + if bool((seq_lens != q_seq_len).any()): + raise ValueError( + f"attn_metadata q length mismatch: cached {seq_lens.tolist()} != {q_seq_len=}." + ) + seq_lens_kv = attn_metadata.seq_lens_kv + if bool((seq_lens_kv != kv_seq_len).any()): + raise ValueError( + f"attn_metadata kv length mismatch: cached {seq_lens_kv.tolist()} != {kv_seq_len=}." + ) class TrtllmAttention(BaseTrtllmAttention, AttentionBackend): @@ -182,11 +76,12 @@ class TrtllmAttention(BaseTrtllmAttention, AttentionBackend): Handles: - Fused QKV requirement for TRTLLM kernel (used when no quant_attention_config is provided) - - Metadata creation and preparation - No KV cache operation - SageAttention per-block QKV quantization (when a quant_attention_config is provided. requires unfused QKV) """ + Metadata = TrtllmAttentionMetadata + def __init__( self, layer_idx: int = 0, @@ -195,10 +90,7 @@ def __init__( num_kv_heads: Optional[int] = None, quant_config: Optional[QuantConfig] = None, dtype: Optional[torch.dtype] = None, - max_batch_size: int = 16, - max_seq_len: int = 4096, quant_attention_config: Optional[QuantAttentionConfig] = None, - attention_metadata_state: Optional[dict] = None, sparse_params: Optional[SkipSoftmaxParams] = None, ): num_kv_heads = num_kv_heads or num_heads @@ -216,17 +108,12 @@ def __init__( # TRTLLM expects flat [B*S, H*D] format self._preferred_layout = AttentionTensorLayout.NHD - self.metadata = TrtllmAttentionMetadata( - attention_metadata_state=attention_metadata_state, - ) - self.quant_attention_config = quant_attention_config - # Needed to work with torch compile cause of attention metadata - # make attn metadata as input for it to work - @torch.compiler.disable - def _prepare_metadata(self, batch_size: int, seq_len: int): - return self.metadata.prepare(batch_size, seq_len) + @property + def requires_metadata(self) -> bool: + """TrtllmAttention always needs a metadata as its backends always enables varlen.""" + return True @torch.compile def _concat_qkv( @@ -250,14 +137,13 @@ def forward( q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], - batch_size: int, - seq_len: int, + *, + attn_metadata: TrtllmAttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, - seq_len_kv: Optional[int] = None, **kwargs, ) -> torch.Tensor: """ - Forward pass with automatic metadata handling. + Forward pass against caller-supplied attention metadata. Dimensions are derived from tensor shapes (NHD layout: ``[B, S, H, D]``). @@ -272,19 +158,23 @@ def forward( q: Query tensor [B, S, H, D] or fused QKV [B, S, H_qkv, D] k: Key tensor [B, S_kv, H_kv, D] or None if fused v: Value tensor [B, S_kv, H_kv, D] or None if fused - batch_size: Batch size - seq_len: Sequence length for Q + attn_metadata: Prepared metadata for this attention site; + must match the actual tensor dimensions. attention_mask: Attention mask type seq_len_kv: Sequence length for K/V (for cross-attention, defaults to seq_len) Returns: Output tensor [B, S, H*D] """ - kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len - prepared_metadata = self._prepare_metadata(batch_size, seq_len) + batch_size, seq_len, _, _ = q.shape + _, kv_seq_len, _, _ = k.shape + _check_metadata(attn_metadata, batch_size, seq_len, kv_seq_len) timestep = kwargs.pop("timestep", None) - if self.quant_attention_config is not None: + if ( + self.quant_attention_config is not None + and attention_mask == PredefinedAttentionMask.FULL + ): assert k is not None and v is not None, ( "SageAttention requires separate Q, K, V tensors" ) @@ -296,7 +186,7 @@ def forward( q=q, k=k, v=v, - metadata=prepared_metadata, + metadata=attn_metadata, attention_mask=attention_mask, timestep=timestep, sage_attn_num_elts_per_blk_q=quant_cfg.q_block_size, @@ -313,7 +203,7 @@ def forward( q=qkv, k=None, v=None, - metadata=prepared_metadata, + metadata=attn_metadata, attention_mask=attention_mask, timestep=timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index ce012e16aea3..93f2f3261b02 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -16,8 +16,7 @@ Visual Generation Attention Backend Utilities Factory functions for creating attention backends for visual generation models. -Uses diffusion-specific wrappers (TrtllmAttention, VanillaAttention) -that handle metadata preparation internally for simplified usage. +Uses diffusion-specific wrappers (TrtllmAttention, VanillaAttention). """ from typing import Optional, Type @@ -81,17 +80,14 @@ def create_attention( num_kv_heads: Optional[int] = None, quant_config: Optional[QuantConfig] = None, dtype: Optional[torch.dtype] = None, - max_batch_size: int = 16, - max_seq_len: int = 4096, attention_config: Optional[AttentionConfig] = None, - attention_metadata_state: Optional[dict] = None, **kwargs, ) -> AttentionBackend: """ Factory function to create attention backend instance for visual generation. - Creates diffusion-specific attention backends that handle metadata preparation - internally, simplifying the forward() call. + Creates diffusion-specific attention backends. The returned backend exposes a + ``Metadata`` class attribute naming the metadata type its ``forward`` expects. Args: backend: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL") @@ -101,14 +97,8 @@ def create_attention( num_kv_heads: Number of KV heads (for GQA/MQA, defaults to num_heads) quant_config: Optional quantization configuration dtype: Data type for the attention - max_batch_size: Initial batch size for metadata pre-allocation. The backend - will automatically reallocate if larger batches are encountered. - max_seq_len: Initial sequence length for metadata pre-allocation. The backend - will automatically reallocate if longer sequences are encountered. attention_config: Optional AttentionConfig used to select the attention algorithm and forward its quantization or sparsity configuration. - attention_metadata_state: Optional model-scoped metadata state from - visual-gen config. Required for TRTLLM backend. **kwargs: Additional backend-specific arguments Returns: @@ -119,13 +109,6 @@ def create_attention( # Forward the validated quantization recipe to TRTLLM or the dense CuTe DSL FMHA backend. if attention_config is not None and attention_config.quant_attention_config is not None: kwargs["quant_attention_config"] = attention_config.quant_attention_config - if backend.upper() == "TRTLLM": - if attention_metadata_state is None: - raise ValueError( - "TRTLLM backend requires `attention_metadata_state` from " - "DiffusionModelConfig; creation path must not allocate metadata implicitly." - ) - kwargs["attention_metadata_state"] = attention_metadata_state if backend.upper() == "CUTEDSL" and attention_config is not None: if ( attention_config.sparse_attention_config is not None @@ -143,7 +126,5 @@ def create_attention( num_kv_heads=num_kv_heads, quant_config=quant_config, dtype=dtype, - max_batch_size=max_batch_size, - max_seq_len=max_seq_len, **kwargs, ) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py b/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py index 043d660bf00f..e2da1afee538 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py @@ -28,7 +28,7 @@ import torch import torch.nn.functional as F -from ...attention_backend.interface import PredefinedAttentionMask +from ...attention_backend.interface import AttentionMetadata, PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout @@ -70,6 +70,7 @@ def forward( k: torch.Tensor, v: torch.Tensor, *, + attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, key_padding_mask: Optional[torch.Tensor] = None, **kwargs, @@ -83,6 +84,8 @@ def forward( q: Query tensor [batch_size, num_heads, seq_len, head_dim] k: Key tensor [batch_size, num_kv_heads, seq_len_kv, head_dim] v: Value tensor [batch_size, num_kv_heads, seq_len_kv, head_dim] + attn_metadata: Unused; SDPA derives everything from tensor shapes. + Preserved for future metadata-consuming variants of this path. attention_mask: Attention mask type (CAUSAL or FULL) key_padding_mask: Optional ``[B, S_kv]`` bool tensor; True = valid, False = pad. Expanded internally to ``[B, 1, 1, S_kv]`` and @@ -91,6 +94,7 @@ def forward( Returns: Output tensor [batch_size, num_heads, seq_len, head_dim] """ + _, _ = attn_metadata, kwargs is_causal = attention_mask == PredefinedAttentionMask.CAUSAL assert q.dim() == 4 and q.shape[3] == self.head_dim, ( @@ -123,6 +127,11 @@ def forward( q, k, v, is_causal=is_causal, scale=self.scale, enable_gqa=enable_gqa ) + @property + def requires_metadata(self) -> bool: + """Currently VanillaAttention only supports non-varlen batching. No metadata referred.""" + return False + @property def preferred_layout(self) -> AttentionTensorLayout: """Return the preferred tensor layout for this backend.""" diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index c636c618a0e2..25276731c57a 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -71,11 +71,6 @@ def discover_pipeline_components(checkpoint_path: Path) -> Dict[str, Path]: return components -def create_attention_metadata_state() -> Dict[str, Any]: - """Create model-scoped attention metadata state for TRTLLM visual-gen backend.""" - return {"metadata_cache": {}} - - def _model_config_value(value: Any, *, deep_copy: bool = True) -> Any: """Return a value for a per-component DiffusionModelConfig.""" if value is None: @@ -119,7 +114,6 @@ class DiffusionModelConfig(_VisualGenConfigBase): torch_compile: TorchCompileConfig = PydanticField(default_factory=TorchCompileConfig) cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) - attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None @@ -184,7 +178,6 @@ class DiffusionPipelineConfig(_VisualGenConfigBase): torch_compile: TorchCompileConfig = PydanticField(default_factory=TorchCompileConfig) cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) - attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None @@ -247,7 +240,6 @@ def _make_model_config( torch_compile=_model_config_value(self.torch_compile), cuda_graph=_model_config_value(self.cuda_graph), attention=_model_config_value(self.attention), - attention_metadata_state=_model_config_value(self.attention_metadata_state), parallel=_model_config_value(self.parallel), cache=_model_config_value(self.cache), enable_layerwise_nvtx_marker=_model_config_value(self.enable_layerwise_nvtx_marker), @@ -608,10 +600,6 @@ def from_pretrained( NVFP4LinearMethod.use_tunable_quantize = True - attention_metadata_state = ( - create_attention_metadata_state() if attention_cfg.backend == "TRTLLM" else None - ) - pipeline_config = cls( quant_config=quant_config, quant_config_dict=quant_config_dict, @@ -622,7 +610,6 @@ def from_pretrained( torch_compile=torch_compile_cfg, cuda_graph=cuda_graph_cfg, attention=attention_cfg, - attention_metadata_state=attention_metadata_state, parallel=parallel_cfg, cache=cache_cfg, enable_layerwise_nvtx_marker=enable_layerwise_nvtx_marker, diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index dbe25e0eec86..08a8a39bde2b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1456,6 +1456,8 @@ def forward( ) # 4. Build forward_fn for the denoise loop + # + def forward_fn( latent_input, extra_stream_latents, @@ -1471,12 +1473,16 @@ def forward_fn( """ current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None + text_mask = extra_tensors["text_mask"] result = self.transformer( hidden_states=latent_input, + attn_metadata_und=attn_metadata_sites["und"], + attn_metadata_mixed=attn_metadata_sites["mixed"], + attn_metadata_mixed_ragged=attn_metadata_sites["mixed_ragged"], timestep=timestep / self.scheduler.config.num_train_timesteps, raw_timestep=timestep, text_ids=extra_tensors["text_ids"], - text_mask=extra_tensors["text_mask"], + text_mask=text_mask, video_shape=video_shape, fps=frame_rate, noisy_frame_mask=velocity_mask, @@ -1511,6 +1517,21 @@ def post_step_fn(step_latents): self.transformer.reset_cache() + # Create attention metadata for understanding and generation towers + denoise_batch = self.denoise_batch_size(latents, guidance_scale=guidance_scale) + denoise_mask = ( + torch.cat([uncond_mask, cond_mask], dim=0) + if denoise_batch != latents.shape[0] + else cond_mask + ) + attn_metadata_sites = self.transformer.create_attn_metadata( + batch_size=denoise_batch, + text_seq_len=denoise_mask.shape[1], + text_lens=denoise_mask.sum(dim=1).tolist(), + video_shape=video_shape, + num_audio_tokens=(audio_latents.shape[2] if do_audio else 0), + ) + # 6. Denoise timer.mark_denoise_start() extra_streams = None diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index 64deb4a3f618..292e082fe2ae 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -15,19 +15,24 @@ import math from dataclasses import dataclass -from typing import Optional, Tuple, TypeVar +from typing import Dict, Optional, Tuple, TypeVar import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import TimestepEmbedding -from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionMetadata, + PredefinedAttentionMask, +) from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import Linear, WeightMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.utils import relu2 +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata +from tensorrt_llm._torch.visual_gen.attention_backend.parallel import get_ulysses_seq_lens from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode @@ -447,6 +452,7 @@ def forward_with_kv( hidden_states: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, + attn_metadata: AttentionMetadata, timestep=None, ) -> torch.Tensor: batch_size, seq_len = hidden_states.shape[:2] @@ -475,6 +481,7 @@ def forward_with_kv( q, k, v, + attn_metadata, attention_mask=PredefinedAttentionMask.CAUSAL, timestep=timestep, ) @@ -509,11 +516,6 @@ def __init__( layer_idx: int = 0, module_name: Optional[str] = None, ): - original_backend = model_config.attention.backend - if model_config.attention.backend == "TRTLLM": - # TRTLLM backend is not supported for Cosmos3CrossAttention - model_config.attention.backend = "VANILLA" - super().__init__( hidden_size=hidden_size, num_attention_heads=num_attention_heads, @@ -527,8 +529,10 @@ def __init__( layer_idx=layer_idx, module_name=module_name, enable_sequence_parallel=True, + # Generation queries attend over [text K/V; generation K/V], so KV + # is longer than Q even though the projection is fused. + is_cross=True, ) - model_config.attention.backend = original_backend # Same flavor note as Cosmos3CausalAttention: attention Q/K norms are # fp32-weight-multiply in both recipes. @@ -547,10 +551,17 @@ def forward( v_und: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, + attn_metadata: AttentionMetadata, timestep=None, real_text_lens: Optional[list[int]] = None, + attn_metadata_ragged: Optional[list[AttentionMetadata]] = None, ) -> torch.Tensor: """ + This is the Cosmos3 "mixed" attention site: the queries are generation + tokens while the keys/values are the concatenation of the text tower's + cached K/V and the generation K/V, so ``kv_seqlen != q_seqlen``. The + caller passes the matching mixed-site metadata. + Args: hidden_states: [B, S_gen, hidden_size] visual tokens k_und: [B, S_und, H_kv, D] pre-computed und keys (post-norm, post-RoPE) @@ -573,6 +584,8 @@ def forward( q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) if real_text_lens is not None and batch_size > 1: + # Ragged batch: per-sample KV lengths differ, so these batch-1 + # calls are not described by the shared site. outs = [] for b in range(batch_size): Lb = int(real_text_lens[b]) @@ -583,6 +596,7 @@ def forward( q[b : b + 1], k_all_b, v_all_b, + attn_metadata if attn_metadata_ragged is None else attn_metadata_ragged[b], attention_mask=PredefinedAttentionMask.FULL, timestep=timestep, ) @@ -596,6 +610,7 @@ def forward( q, k_all, v_all, + attn_metadata, attention_mask=PredefinedAttentionMask.FULL, timestep=timestep, ) @@ -671,6 +686,7 @@ def forward( self, hidden_states: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor], + attn_metadata: AttentionMetadata, timestep=None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ @@ -686,6 +702,7 @@ def forward( hidden_states, cos, sin, + attn_metadata, timestep=timestep, ) hidden_states = residual + attn_out @@ -738,8 +755,10 @@ def forward( k_und: torch.Tensor, v_und: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor], + attn_metadata: AttentionMetadata, timestep=None, real_text_lens: Optional[list[int]] = None, + attn_metadata_ragged: Optional[list[AttentionMetadata]] = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -751,8 +770,10 @@ def forward( v_und=v_und, freqs_cos=cos, freqs_sin=sin, + attn_metadata=attn_metadata, timestep=timestep, real_text_lens=real_text_lens, + attn_metadata_ragged=attn_metadata_ragged, ) hidden_states = residual + hidden_states @@ -901,6 +922,7 @@ def forward( text_ids: torch.Tensor, text_mask: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor], + attn_metadata: AttentionMetadata, timestep=None, ) -> list[Tuple[torch.Tensor, torch.Tensor]]: """ @@ -919,7 +941,7 @@ def forward( cached_kv: list[Tuple[torch.Tensor, torch.Tensor]] = [] for layer in self.layers: hidden = hidden * mask_3d - hidden, k, v = layer(hidden, freqs, timestep=timestep) + hidden, k, v = layer(hidden, freqs, attn_metadata, timestep=timestep) cached_kv.append((k, v)) return cached_kv @@ -1200,9 +1222,104 @@ def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None + def gen_seq_len(self, video_shape: Tuple[int, int, int], num_audio_tokens: int = 0) -> int: + """Number of generation tokens the GEN stack attends over. + + Single source of truth for the token arithmetic that ``forward`` does + inline: video patches (with H/W padded up to the patch grid), plus the + appended audio tokens. + """ + T, H, W = video_shape + Hp, Wp, _, _ = self._pad_to_patch_size(H, W) + return T * Hp * Wp + num_audio_tokens + + def create_attn_metadata( + self, + *, + batch_size: int, + text_seq_len: int, + text_lens: list[int], + video_shape: Tuple[int, int, int], + num_audio_tokens: int = 0, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + Cosmos3 is a mixed-stream model with two sites: + + * ``und`` -- the understanding tower's causal self-attention over ``S_text``. + * ``mixed`` -- generation queries over ``S_gen`` against the concatenation + of the cached text K/V and the generation K/V, i.e. ``kv_seqlen = S_text + S_gen``. + + Only generation tower supports sequence-parallel, so ``forward`` shards the generation + tokens and the cached text K/V, and the wrapper gathers them again. + """ + if not self.audio_gen: + num_audio_tokens = 0 + metadata_cls = self.attn_backend_metadata_cls + size = self.sharder.size + + # forward() pads the generation stream up to a multiple of the shard + # count before sharding, and keeps a rounded-up slice of the text K/V. + s_gen = self.gen_seq_len(video_shape, num_audio_tokens) + s_gen_local = (s_gen + (-s_gen) % size) // size + s_text_local = self._text_kv_len_sharded(max(text_lens)) // size + + q_mixed_one, _ = get_ulysses_seq_lens( + s_gen_local, + s_gen_local, + visual_gen_mapping=self.model_config.visual_gen_mapping, + ) + q_mixed, kv_mixed = get_ulysses_seq_lens( + s_gen_local, + s_text_local + s_gen_local, + visual_gen_mapping=self.model_config.visual_gen_mapping, + ) + + # The ragged loop in Cosmos3CrossAttention runs one batch-1 call per + # sample, each with its own text length, so it needs its own sites. + mixed_ragged = [ + make_diffusion_attn_metadata( + metadata_cls, + batch_size=1, + q_seq_lens=q_mixed_one, + kv_seq_lens=get_ulysses_seq_lens( + s_gen_local, + int(text_len) + s_gen_local, + visual_gen_mapping=self.model_config.visual_gen_mapping, + )[1], + ) + for text_len in text_lens + ] + + return { + "und": make_diffusion_attn_metadata( + metadata_cls, batch_size=batch_size, q_seq_lens=text_seq_len + ), + "mixed_ragged": mixed_ragged, + "mixed": make_diffusion_attn_metadata( + metadata_cls, + batch_size=batch_size, + q_seq_lens=q_mixed, + kv_seq_lens=kv_mixed, + ), + } + + def _text_kv_len_sharded(self, max_real_len: int) -> int: + """Text K/V length kept per rank when the sequence is sharded. + + Rounds the real text length up to a multiple of the shard count; at most + ``size - 1`` extra positions, which are zeroed. Mirrors the slicing in + ``forward``. + """ + size = self.sharder.size + return int(max_real_len) + (size - int(max_real_len) % size) % size + def forward( self, hidden_states: torch.Tensor, + attn_metadata_und: AttentionMetadata, + attn_metadata_mixed: AttentionMetadata, + attn_metadata_mixed_ragged: Optional[list[AttentionMetadata]] = None, timestep: Optional[torch.Tensor] = None, raw_timestep: Optional[torch.Tensor] = None, text_ids: Optional[torch.Tensor] = None, @@ -1218,6 +1335,9 @@ def forward( Args: hidden_states: [B, C, T, H, W] noisy latents + attn_metadata_und: Metadata for the text tower's causal self-attention. + attn_metadata_mixed: Metadata for the generation stack, whose keys mix + the cached text K/V with the generation K/V.. timestep: Normalized diffusion timestep in [0, 1], shape [B]. raw_timestep: Raw scheduler diffusion timestep, shape [B], used by the Cosmos3 time embedding path. @@ -1284,6 +1404,7 @@ def forward( text_ids, text_mask, freqs_und, + attn_metadata_und, timestep=timestep, ) self.cached_freqs_gen = freqs_gen @@ -1291,8 +1412,8 @@ def forward( if self.sharder.is_active: # Round max_real_len up to next multiple of sharder.size. # At most size-1 extra positions, negligible softmax dilution. - val = (self.sharder.size - max_real_len % self.sharder.size) % self.sharder.size - S_text_shard_total = int(max_real_len) + val + S_text_shard_total = self._text_kv_len_sharded(max_real_len) + val = S_text_shard_total - int(max_real_len) self.cached_kv = [] for k, v in cached_kv_full: @@ -1350,8 +1471,10 @@ def forward( k_und, v_und, freqs_gen, + attn_metadata_mixed, timestep=timestep, real_text_lens=real_text_lens, + attn_metadata_ragged=attn_metadata_mixed_ragged, ) else: hidden_gen = layer( @@ -1359,6 +1482,7 @@ def forward( k_und, v_und, freqs_gen, + attn_metadata_mixed, timestep=timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py index 31df009e0be8..bbfc40b36770 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -13,6 +13,7 @@ import torch import torch.nn.functional as F +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.modules.linear import ( Linear, TensorParallelMode, @@ -262,6 +263,7 @@ def _prepare_qkv( def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, @@ -271,6 +273,7 @@ def forward( Args: hidden_states: Image tokens [batch, img_seq, dim] + attn_metadata: Attention metadata site for this site. encoder_hidden_states: Text tokens [batch, txt_seq, dim] (for dual-stream) attention_mask: Optional attention mask (unused, for API compat) image_rotary_emb: Tuple of (cos, sin) for RoPE @@ -286,7 +289,7 @@ def forward( hidden_states, encoder_hidden_states, image_rotary_emb ) - hidden_states = self._attn_impl(query, key, value, timestep=timestep) + hidden_states = self._attn_impl(query, key, value, attn_metadata, timestep=timestep) hidden_states = hidden_states.to(query.dtype) if is_dual_stream: @@ -630,6 +633,7 @@ def _project_split_output_with_fp4_mlp( def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, @@ -637,6 +641,7 @@ def forward( """ Args: hidden_states: [batch, seq, dim] + attn_metadata: Attention metadata site for this site. attention_mask: Optional attention mask image_rotary_emb: Tuple of (freqs_cos, freqs_sin) @@ -646,7 +651,7 @@ def forward( if self._can_project_hidden_mlp_with_cute_dsl(): qkv = self.to_qkv_mlp_proj.qkv_proj(hidden_states) q, k, v = self._apply_norm_rope(qkv, image_rotary_emb) - attn_out = self._attn_impl(q, k, v, timestep=timestep) + attn_out = self._attn_impl(q, k, v, attn_metadata, timestep=timestep) attn_out = attn_out.to(q.dtype) mlp_out = self._project_hidden_mlp_with_cute_dsl(hidden_states) return self._combine_split_projection(attn_out, mlp_out) @@ -656,7 +661,7 @@ def forward( q, k, v = self._apply_norm_rope(qkv, image_rotary_emb) - attn_out = self._attn_impl(q, k, v, timestep=timestep) + attn_out = self._attn_impl(q, k, v, attn_metadata, timestep=timestep) attn_out = attn_out.to(q.dtype) # Parallel MLP path (reshape to 2D for Triton kernel, then back) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index 5881614c872b..e6df8330be13 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -357,7 +357,13 @@ def forward( [latents.shape[0]], guidance_scale, device=self.device, dtype=torch.float32 ) - # Denoising loop + # Create attention metadata + attn_metadata_sites = self.transformer.create_attn_metadata( + batch_size=latents.shape[0], + text_seq_len=prompt_embeds.shape[1], + image_seq_len=latents.shape[1], + ) + def forward_fn( latents, extra_stream_latents, @@ -369,6 +375,7 @@ def forward_fn( """Forward function for FLUX transformer.""" return self.transformer( hidden_states=latents, + attn_metadata=attn_metadata_sites["self"], encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_prompt_embeds, timestep=timestep / 1000, # FLUX expects normalized timesteps diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 85a87f05a631..c0731a6e620d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -521,7 +521,17 @@ def forward( [latents.shape[0]], guidance_scale, device=self.device, dtype=torch.float32 ) + # Create attention metadata + attn_metadata_sites = self.transformer.create_attn_metadata( + batch_size=latents.shape[0], + text_seq_len=prompt_embeds.shape[1], + # forward_fn concatenates the reference-image latents onto the noise. + image_seq_len=latents.shape[1] + + (0 if image_latents is None else image_latents.shape[1]), + ) + # Denoising loop using forward_fn callback (WAN pattern) + def forward_fn( latents, extra_stream_latents, @@ -539,6 +549,7 @@ def forward_fn( noise_pred = self.transformer( hidden_states=transformer_latents, + attn_metadata=attn_metadata_sites["self"], encoder_hidden_states=encoder_hidden_states, timestep=timestep / 1000, # FLUX.2 expects normalized timesteps img_ids=transformer_latent_ids, 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 527dce31df23..e899195dc4e4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py @@ -26,10 +26,13 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps from tqdm import tqdm +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.utils import maybe_compile +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata +from tensorrt_llm._torch.visual_gen.attention_backend.parallel import get_ulysses_seq_lens from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import FluxJointAttnMLPProj @@ -346,6 +349,7 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, + attn_metadata: AttentionMetadata, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -355,6 +359,7 @@ def forward( hidden_states: Image tokens (batch, img_seq, dim) encoder_hidden_states: Text tokens (batch, txt_seq, dim) temb: Timestep embedding (batch, dim) + attn_metadata: Attention metadata site for the joint text+image sequence. image_rotary_emb: RoPE (cos, sin) tuple joint_attention_kwargs: Additional kwargs for attention @@ -375,6 +380,7 @@ def forward( joint_attention_kwargs = joint_attention_kwargs or {} attn_output, context_attn_output = self.attn( hidden_states=norm_hidden_states, + attn_metadata=attn_metadata, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs, @@ -500,6 +506,7 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, + attn_metadata: AttentionMetadata, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -509,6 +516,7 @@ def forward( hidden_states: Image tokens (batch, img_seq, dim) encoder_hidden_states: Text tokens (batch, txt_seq, dim) temb: Timestep embedding (batch, dim) + attn_metadata: Attention metadata site for the joint text+image sequence. image_rotary_emb: RoPE (cos, sin) tuple joint_attention_kwargs: Additional kwargs for attention @@ -532,6 +540,7 @@ def forward( joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( hidden_states=norm_hidden_states, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs, ) @@ -767,9 +776,40 @@ def apply_quant_config_exclude_modules(self): if is_excluded and getattr(module, "quant_config", None) is not None: module.quant_config = no_quant_config + def create_attn_metadata( + self, + *, + batch_size: int, + text_seq_len: int, + image_seq_len: int, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + FLUX has a single attention site: both the dual-stream and single-stream + blocks attend over the concatenated text+image sequence. + """ + # forward() shards each stream before concatenating them, so a block + # attends over this rank's share of the joint sequence. + size = self.sharder.size + local_seq_len = text_seq_len // size + image_seq_len // size + q_seq_len, kv_seq_len = get_ulysses_seq_lens( + local_seq_len, + local_seq_len, + visual_gen_mapping=self.model_config.visual_gen_mapping, + ) + return { + "self": make_diffusion_attn_metadata( + self.attn_backend_metadata_cls, + batch_size=batch_size, + q_seq_lens=q_seq_len, + kv_seq_lens=None if kv_seq_len == q_seq_len else kv_seq_len, + ) + } + def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.Tensor = None, @@ -783,6 +823,9 @@ def forward( Args: hidden_states: Latent image tokens (batch, seq_len, in_channels) + attn_metadata: Attention metadata site covering the joint + text+image sequence; built by the pipeline from + :meth:`create_attn_metadata`. encoder_hidden_states: T5 text embeddings (batch, txt_seq_len, joint_attention_dim) pooled_projections: CLIP pooled text embeddings (batch, pooled_projection_dim) timestep: Normalized timestep tensor in [0, 1], shape (batch,) @@ -837,6 +880,7 @@ def forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, ) @@ -847,6 +891,7 @@ def forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, ) 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 253f8a7599dc..f6eb8f73aefc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py @@ -24,9 +24,12 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps from tqdm import tqdm +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata +from tensorrt_llm._torch.visual_gen.attention_backend.parallel import get_ulysses_seq_lens from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.flux.attention import ( Flux2ParallelSelfAttention, @@ -285,6 +288,7 @@ def forward( image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], img_mod: Tuple[Tuple[torch.Tensor, ...], ...], txt_mod: Tuple[Tuple[torch.Tensor, ...], ...], + attn_metadata: AttentionMetadata, timestep: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ @@ -292,6 +296,7 @@ def forward( hidden_states: Image features [batch, img_seq, dim] encoder_hidden_states: Text features [batch, txt_seq, dim] image_rotary_emb: Tuple of (freqs_cos, freqs_sin) + attn_metadata: Attention metadata site for the joint sequence. img_mod: Image modulation ((shift1, scale1, gate1), (shift2, scale2, gate2)) txt_mod: Text modulation ((shift1, scale1, gate1), (shift2, scale2, gate2)) @@ -316,6 +321,7 @@ def forward( # Joint attention attn_output, encoder_attn_output = self.attn( hidden_states=hidden_states, + attn_metadata=attn_metadata, encoder_hidden_states=encoder_hidden_states, image_rotary_emb=image_rotary_emb, timestep=timestep, @@ -399,6 +405,7 @@ def forward( hidden_states: torch.Tensor, image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], mod: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + attn_metadata: AttentionMetadata, timestep: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ @@ -406,6 +413,7 @@ def forward( hidden_states: [batch, seq, dim] image_rotary_emb: Tuple of (freqs_cos, freqs_sin) mod: Modulation (shift, scale, gate) + attn_metadata: Attention metadata site for the joint sequence. Returns: hidden_states [batch, seq, dim] @@ -422,6 +430,7 @@ def forward( # Parallel attention + MLP hidden_states = self.attn( hidden_states, + attn_metadata, image_rotary_emb=image_rotary_emb, timestep=timestep, ) @@ -684,9 +693,39 @@ def apply_quant_config_exclude_modules(self): if is_excluded and getattr(module, "quant_config", None) is not None: module.quant_config = no_quant_config + def create_attn_metadata( + self, + *, + batch_size: int, + text_seq_len: int, + image_seq_len: int, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + FLUX.2 has a single attention site over the concatenated text+image sequence. + """ + # forward() shards each stream before concatenating them, so a block + # attends over this rank's share of the joint sequence. + size = self.sharder.size + local_seq_len = text_seq_len // size + image_seq_len // size + q_seq_len, kv_seq_len = get_ulysses_seq_lens( + local_seq_len, + local_seq_len, + visual_gen_mapping=self.model_config.visual_gen_mapping, + ) + return { + "self": make_diffusion_attn_metadata( + self.attn_backend_metadata_cls, + batch_size=batch_size, + q_seq_lens=q_seq_len, + kv_seq_lens=None if kv_seq_len == q_seq_len else kv_seq_len, + ) + } + def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: torch.Tensor, timestep: Optional[torch.Tensor] = None, img_ids: Optional[torch.Tensor] = None, @@ -699,6 +738,8 @@ def forward( Args: hidden_states: Latent image features [batch, img_seq, in_channels] + attn_metadata: Attention metadata site covering the joint + text+image sequence; built by the pipeline from :meth:`create_attn_metadata`. encoder_hidden_states: Text features [batch, txt_seq, joint_attention_dim] timestep: Normalized diffusion timestep in [0, 1], shape [batch] img_ids: Image position IDs [img_seq, num_axes] or [batch, img_seq, num_axes] @@ -755,6 +796,7 @@ def forward( image_rotary_emb=image_rotary_emb, img_mod=img_mod, txt_mod=txt_mod, + attn_metadata=attn_metadata, timestep=timestep, ) @@ -767,6 +809,7 @@ def forward( hidden_states=hidden_states, image_rotary_emb=image_rotary_emb, mod=single_mod[0], # Single tuple of (shift, scale, gate) + attn_metadata=attn_metadata, timestep=timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 2ebee01b7d35..5bb52557c222 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -15,6 +15,7 @@ import torch.distributed as dist from transformers import Gemma3ForConditionalGeneration, GemmaTokenizerFast +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.utils import make_weak_ref from tensorrt_llm._torch.visual_gen.cache.teacache import CacheContext, register_extractor from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache @@ -308,6 +309,15 @@ def postprocess(output): # --------------------------------------------------------------------------- +def _is_attn_metadata_map(value) -> bool: + """Is *value* a per-site attention metadata mapping (see metadata.py)?""" + return ( + isinstance(value, dict) + and bool(value) + and all(isinstance(v, AttentionMetadata) for v in value.values()) + ) + + class _LTX2CUDAGraphRunner(CUDAGraphRunner): """CUDAGraphRunner extended for LTX-2's ``Modality``-based transformer. @@ -450,6 +460,8 @@ def _copy_value(dst, src): if isinstance(src, torch.Tensor) and isinstance(dst, torch.Tensor): dst.copy_(src) return dst + if _is_attn_metadata_map(src) and _is_attn_metadata_map(dst): + return dst if isinstance(src, Modality) and isinstance(dst, Modality): dst.latent.copy_(src.latent) dst.timesteps.copy_(src.timesteps) @@ -1760,7 +1772,37 @@ def forward( # Cache encoder output for two-stage Stage 2 reuse. self._cached_encoder_output = (video_embeds, audio_embeds, connector_mask) + # Multi-modal guidance runs its own cond/uncond passes, so `denoise` + # must not also batch for CFG. + effective_guidance = 1.0 if use_multi_modal_guidance else guidance_scale + + # Create attention metadata + _denoise_batch = self.denoise_batch_size(latents, guidance_scale=effective_guidance) + _attn_metadata_cond = self.transformer.create_attn_metadata( + batch_size=_denoise_batch, + video_seq_len=latents.shape[1], + audio_seq_len=audio_latents.shape[1] if has_audio else 0, + text_cache=_text_cache, + ) + _attn_metadata_uncond = ( + None + if _text_cache_uncond is None + else self.transformer.create_attn_metadata( + batch_size=_denoise_batch, + video_seq_len=latents.shape[1], + audio_seq_len=audio_latents.shape[1] if has_audio else 0, + text_cache=_text_cache_uncond, + ) + ) + + # Select the right metadata for each attention site. + def _attn_metadata_for(text_cache): + return ( + _attn_metadata_uncond if text_cache is _text_cache_uncond else _attn_metadata_cond + ) + # ---- 9. Denoising loop ------------------------------------------ + def _run_transformer( v_latents, a_latents, @@ -1822,6 +1864,7 @@ def _run_transformer( audio=audio_mod, perturbations=perturbations, text_cache=text_cache, + attn_metadata=_attn_metadata_for(text_cache), timestep=timestep_val.new_tensor(float(step_index) / num_steps), step_index=step_index, ) @@ -1993,7 +2036,6 @@ def forward_fn( # When using multi-modal guidance, we handle everything inside # forward_fn, so tell BasePipeline not to apply its own CFG. - effective_guidance = 1.0 if use_multi_modal_guidance else guidance_scale timer.mark_denoise_start() result = self.denoise( diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 181ff81af71f..fa360ccd8c1d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1716,6 +1716,15 @@ def _refinement_denoise( dtype=self.dtype, ) + # Loop-invariant, and the stage-2 topology was switched before this + # call. No guidance here, so the batch is the working latents'. + _s2_attn_metadata = self.transformer.create_attn_metadata( + batch_size=v_working.shape[0], + video_seq_len=v_working.shape[1], + audio_seq_len=a_working.shape[1] if a_working is not None else 0, + text_cache=_s2_static, + ) + # stage2_denoise measures ONLY the step loop (upsample, LoRA bind, # and text-cache/scheduler prep stay outside the bracket). if timer is not None: @@ -1726,6 +1735,7 @@ def _refinement_denoise( # drives the profiler windows itself. Without this a numeric range # would capture stage 1 only and the trace would silently omit half # the denoising. + for i, _ in self._profile_denoise_steps(range(len(sigmas) - 1)): with nvtx_range(f"refinement_step {i}"): sigma = sigmas[i] @@ -1760,6 +1770,7 @@ def _refinement_denoise( video=video_mod, audio=audio_mod, text_cache=_s2_static, + attn_metadata=_s2_attn_metadata, step_index=i, ) 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 7cba785b1bf2..f3e24e2073f2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -21,7 +21,7 @@ import os from dataclasses import dataclass, replace from enum import Enum -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional import torch import torch.distributed as torch_dist @@ -29,11 +29,14 @@ import torch.nn.functional as F from tqdm import tqdm +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod, WeightMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.utils import Fp4QuantizedTensor, gelu_tanh +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( UlyssesAttention, + get_ulysses_seq_lens, wrap_parallel_attention, ) from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention @@ -78,6 +81,31 @@ # --------------------------------------------------------------------------- +def _ltx2_enable_sp( + vgm, + *, + enable_sequence_parallel: bool, + is_cross_attn: bool, +) -> bool: + """Whether an LTX-2 attention site is wrapped for sequence parallelism. + + ``vgm`` is the model's :class:`VisualGenMapping`, or ``None``. + + Cross-attn supports Ulysses, and distributed Attention2D only WITH Ulysses + on top (the AV dispatch keys the seq-sharded K/V path off the Ulysses + wrapper); under ring CP or attn2d-without-ulysses we disable wrappers and + fall back to the plain backend + all-gather in the AV cross-attn forward + path. Read by ``LTX2Attention.__init__`` when it wraps, and by + ``BasicAVTransformerBlock.create_attn_metadata`` when it sizes the sites. + """ + if not is_cross_attn: + return enable_sequence_parallel + ulysses_size = vgm.ulysses_size if vgm is not None else 1 + cp_size = vgm.cp_size if vgm is not None else 1 + attn2d_active = vgm is not None and vgm.attn2d_row_size * vgm.attn2d_col_size > 1 + return enable_sequence_parallel and (cp_size == 1 or (attn2d_active and ulysses_size > 1)) + + class LTX2Attention(Attention): """LTX-2 attention: extends base Attention with LTX-specific RoPE, gated attention, and separate K-RoPE for audio-video cross-attention. @@ -134,20 +162,12 @@ def __init__( else: qkv_mode = QKVMode.FUSE_QKV - # Caller opts in via enable_sequence_parallel. Cross-attn supports - # Ulysses, and distributed Attention2D only WITH Ulysses on top (the AV - # dispatch keys the seq-sharded K/V path off the Ulysses wrapper); under - # ring CP or attn2d-without-ulysses we disable wrappers and fall back to - # the plain backend + all-gather in the AV cross-attn forward path. ulysses_size = vgm.ulysses_size if vgm is not None else 1 - cp_size = vgm.cp_size if vgm is not None else 1 - attn2d_active = vgm is not None and vgm.attn2d_row_size * vgm.attn2d_col_size > 1 - if self._is_cross_attn: - enable_sp = enable_sequence_parallel and ( - cp_size == 1 or (attn2d_active and ulysses_size > 1) - ) - else: - enable_sp = enable_sequence_parallel + enable_sp = _ltx2_enable_sp( + vgm, + enable_sequence_parallel=enable_sequence_parallel, + is_cross_attn=self._is_cross_attn, + ) # Map LTX RoPE type to the fused-kernel INTERLEAVE template parameter: # INTERLEAVED → pair (2i, 2i+1) pattern → kernel INTERLEAVE=true @@ -169,6 +189,7 @@ def __init__( module_name=module_name, enable_sequence_parallel=enable_sp, async_ulysses=self._use_async_ulysses, + is_cross=self._is_cross_attn, ) # Validate Ulysses head divisibility (from main). @@ -218,7 +239,6 @@ def __init__( quant_config=self.quant_config, dtype=self.dtype, attention_config=config.attention, - attention_metadata_state=config.attention_metadata_state, sparse_params=self.sparse_params, ) self._attn_stage2 = wrap_parallel_attention( @@ -328,6 +348,7 @@ def project_kv( def forward( self, x: torch.Tensor, + attn_metadata: AttentionMetadata, context: torch.Tensor | None = None, pe: tuple[torch.Tensor, torch.Tensor] | None = None, pre_projected_kv: tuple[torch.Tensor, torch.Tensor] | None = None, @@ -374,7 +395,7 @@ def forward( and pre_projected_kv is None and hasattr(self.attn, "forward_async") ): - return self.forward_async(x, freqs=pe, timestep=timestep) + return self.forward_async(x, attn_metadata, freqs=pe, timestep=timestep) # Fused gate: prod uses fused kernels (head_dim ∈ {64, 128}); mini-config # tests (head_dim=32) fall to naive ops. @@ -448,7 +469,7 @@ def forward( attn_kwargs = {} if key_padding_mask is not None: attn_kwargs["key_padding_mask"] = key_padding_mask - out = self._attn_impl(q, k, v, timestep=timestep, **attn_kwargs) + out = self._attn_impl(q, k, v, attn_metadata, timestep=timestep, **attn_kwargs) if self.to_gate_logits is not None: gate_logits = self.to_gate_logits(x) @@ -463,6 +484,7 @@ def forward( def forward_async( self, q_input: torch.Tensor, + attn_metadata: AttentionMetadata, freqs: tuple[torch.Tensor, torch.Tensor] | None = None, kv_input: torch.Tensor | None = None, kv_freqs: tuple[torch.Tensor, torch.Tensor] | None = None, @@ -553,7 +575,12 @@ def compute_v(): issue_order = ("v", "q", "k") if self_attn else ("q", "k", "v") out_4d = self.attn.forward_async( - compute_q, compute_k, compute_v, issue_order=issue_order, timestep=timestep + compute_q, + compute_k, + compute_v, + issue_order=issue_order, + attn_metadata=attn_metadata, + timestep=timestep, ) # LTX-2 gated-attention scaling in 4D before to_out (gate on the Q input). @@ -640,6 +667,7 @@ def __init__( # 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 + self._vgm = vgm self._sharder = SequenceSharder.from_vgm(vgm) self._sharder_s2 = stage2_sharder if stage2_sharder is not None else self._sharder self._active_sharder = self._sharder @@ -920,10 +948,123 @@ def _sp_all_gather(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor: # -- Forward ------------------------------------------------------------- + def create_attn_metadata( + self, + metadata_cls, + *, + batch_size: int, + video_seq: int, + audio_seq_full: int, + text_kv_video_len: int = 0, + text_kv_audio_len: int = 0, + ulysses_size: Optional[int] = None, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this block's six sites. + + Lives on the block because every predicate it needs -- the audio shard + mode, the active sharder, and the per-site sequence-parallel flags -- + is block state, and because it has to stay in step with ``forward`` + right below. ``LTXModel.create_attn_metadata`` calls this on block 0; + blocks are homogeneous. + + Each site is sized for what the kernel sees: the lengths below are the + ones the attention modules are called with, mapped through + :func:`get_ulysses_seq_lens` with the same flags the module was + built with, so the wrappers never have to rewrite the metadata. + + Args: + metadata_cls: Concrete metadata type for the active backend. + video_seq: Video tokens this rank holds (already sharded). + audio_seq_full: Audio tokens across all ranks, including the + Ulysses padding ``configure_audio_ulysses`` appended. + ulysses_size: Active Ulysses degree, when the active stack was + wrapped with an explicit group (stage 2) rather than the + mapping's. ``None`` uses the mapping's. + """ + size = self._active_sharder.size if self._active_sharder.is_active else 1 + mode = self._audio_shard_mode + # FULL shards the audio sequence too; CONDITIONAL and NONE keep it whole. + audio_seq = audio_seq_full // size if mode == AudioShardMode.FULL else audio_seq_full + + vgm = self._vgm + + def site(q_seq_lens, kv_seq_lens=None, *, enable_sequence_parallel, is_cross_attn): + q, kv = get_ulysses_seq_lens( + q_seq_lens, + q_seq_lens if kv_seq_lens is None else kv_seq_lens, + visual_gen_mapping=vgm, + enable_sequence_parallel=_ltx2_enable_sp( + vgm, + enable_sequence_parallel=enable_sequence_parallel, + is_cross_attn=is_cross_attn, + ), + ulysses_size=ulysses_size, + ) + return make_diffusion_attn_metadata( + metadata_cls, + batch_size=batch_size, + q_seq_lens=q, + kv_seq_lens=None if kv == q and kv_seq_lens is None else kv, + ) + + sites: Dict[str, AttentionMetadata] = {} + if video_seq: + # attn1: enable_sequence_parallel=True. + sites["video_self"] = site( + video_seq, enable_sequence_parallel=True, is_cross_attn=False + ) + if text_kv_video_len: + # attn2: enable_sequence_parallel=False. + sites["video_text_cross"] = site( + video_seq, + text_kv_video_len, + enable_sequence_parallel=False, + is_cross_attn=True, + ) + if audio_seq: + # audio_attn1: enable_sequence_parallel = not CONDITIONAL, i.e. FULL. + sites["audio_self"] = site( + audio_seq, + enable_sequence_parallel=(mode == AudioShardMode.FULL), + is_cross_attn=False, + ) + if text_kv_audio_len: + # audio_attn2: enable_sequence_parallel=False. + sites["audio_text_cross"] = site( + audio_seq, + text_kv_audio_len, + enable_sequence_parallel=False, + is_cross_attn=True, + ) + + if video_seq and audio_seq: + # a2v: video queries against audio K/V, which is all-gathered + # under FULL, so the KV length is the full audio length either way. + # audio_to_video_attn: enable_sequence_parallel=False. + sites["audio_to_video"] = site( + video_seq, + audio_seq_full, + enable_sequence_parallel=False, + is_cross_attn=True, + ) + # v2a: CONDITIONAL slices audio to this rank here, and the video + # K/V is all-gathered when no Ulysses wrapper does it. + v2a_q = audio_seq_full // size if mode == AudioShardMode.CONDITIONAL else audio_seq + v2a_kv = video_seq + v2a_sp = _ltx2_enable_sp(vgm, enable_sequence_parallel=True, is_cross_attn=True) + if not v2a_sp and self._active_sharder.is_active: + v2a_kv = video_seq * size + # video_to_audio_attn: enable_sequence_parallel=True. + sites["video_to_audio"] = site( + v2a_q, v2a_kv, enable_sequence_parallel=True, is_cross_attn=True + ) + return sites + def forward( self, video: TransformerArgs | None, audio: TransformerArgs | None, + attn_metadata: Dict[str, AttentionMetadata], perturbations=None, text_kv_video: tuple[torch.Tensor, torch.Tensor] | None = None, text_kv_audio: tuple[torch.Tensor, torch.Tensor] | None = None, @@ -931,9 +1072,17 @@ def forward( ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward with optional perturbation masking for STG. + LTX-2 is a mixed-stream block with six attention sites -- video and audio + self-attention, video/audio text cross-attention, and bidirectional + audio<->video cross-attention -- each with its own Q/KV lengths. Rather + than one shared metadata object thrashing between them, each site pulls + its own metadata out of ``attn_metadata`` by site name. + Args: perturbations: Optional ``BatchedPerturbationConfig`` that masks attention outputs for selected blocks/modalities. + attn_metadata: One prepared metadata object per attention site, + keyed by the names :meth:`create_attn_metadata` returns. text_kv_video: Pre-projected (K, V) for video text cross-attention. Required when the video stream runs cross-attn — built by ``LTXModel.prepare_text_cache``. @@ -992,7 +1141,10 @@ def forward( fp4_input_scale=get_nvfp4_self_attn_input_scale(self.attn1), ) v_attn_raw = self.attn1( - norm_vx, pe=video.positional_embeddings, timestep=video.timesteps + norm_vx, + attn_metadata["video_self"], + pe=video.positional_embeddings, + timestep=video.timesteps, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx @@ -1016,6 +1168,7 @@ def forward( attn2_q_input = rms_norm(vx, eps=self.norm_eps) text_v_attn_raw = self.attn2( attn2_q_input, + attn_metadata["video_text_cross"], context=video.context, pre_projected_kv=text_kv_video, timestep=video.timesteps, @@ -1047,6 +1200,7 @@ def forward( ) a_attn_raw = self.audio_attn1( norm_ax, + attn_metadata["audio_self"], pe=audio.positional_embeddings, key_padding_mask=audio.audio_padding_mask, timestep=audio.timesteps, @@ -1072,6 +1226,7 @@ def forward( audio_attn2_q_input = rms_norm(ax, eps=self.norm_eps) text_a_attn_raw = self.audio_attn2( audio_attn2_q_input, + attn_metadata["audio_text_cross"], context=audio.context, pre_projected_kv=text_kv_audio, timestep=audio.timesteps, @@ -1230,6 +1385,7 @@ def forward( a2v_attn_raw = self.audio_to_video_attn( vx_scaled_a2v, + attn_metadata["audio_to_video"], pre_projected_kv=(k_a2v, v_a2v), pe=video.cross_positional_embeddings, key_padding_mask=audio.audio_padding_mask, @@ -1264,6 +1420,7 @@ def forward( if self._async_ulysses and self.video_to_audio_attn.is_ulysses: out_local = self.video_to_audio_attn.forward_async( q_input=ax_v2a_local, + attn_metadata=attn_metadata["video_to_audio"], freqs=a_cross_pe, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, @@ -1284,6 +1441,7 @@ def forward( v_v2a = self._sp_all_gather(v_v2a) out_local = self.video_to_audio_attn( ax_v2a_local, + attn_metadata["video_to_audio"], pre_projected_kv=(k_v2a, v_v2a), pe=a_cross_pe, timestep=audio.timesteps, @@ -1296,6 +1454,7 @@ def forward( # padded audio Q stripped on exit by LTXModel.forward). v2a_attn_raw = self.video_to_audio_attn.forward_async( q_input=ax_scaled_v2a, + attn_metadata=attn_metadata["video_to_audio"], freqs=audio.cross_positional_embeddings, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, @@ -1318,6 +1477,7 @@ def forward( v2a_attn_raw = self.video_to_audio_attn( ax_scaled_v2a, + attn_metadata["video_to_audio"], pre_projected_kv=(k_v2a, v_v2a), pe=audio.cross_positional_embeddings, timestep=audio.timesteps, @@ -2216,6 +2376,64 @@ def prepare_text_cache( audio_kv=a_kv, ) + def create_attn_metadata( + self, + *, + batch_size: int, + video_seq_len: int = 0, + audio_seq_len: int = 0, + text_cache: TextCache, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + Takes plain lengths rather than the ``Modality`` bundles ``forward`` + receives, so a pipeline can build this once before its denoise loop -- + the per-step bundles carry per-step timesteps, but the lengths are fixed. + Building it outside the loop is also what keeps it out of the + CUDA-graph-captured region, where ``prepare()``'s pinned-memory staging + would raise. + + Mirrors the length arithmetic ``forward`` performs: audio is padded to a + multiple of the Ulysses size, then video (always) and audio (under + ``AudioShardMode.FULL``) are sequence-sharded. Per-site derivation lives + on the block, next to the code that consumes it. + + Sites for a modality that a given pass does not run are harmless: the + block only looks up the ones its active path needs, and the graph key + already separates passes by modality presence. + + Args: + batch_size: Sequences per forward. + video_seq_len: Unsharded video tokens, or 0 for audio-only models. + audio_seq_len: Unsharded, unpadded audio tokens, or 0 when absent. + """ + size = self._active_sharder.size if self._active_sharder.is_active else 1 + + video_seq = video_seq_len // size + audio_seq_full = audio_seq_len + self._audio_pad if audio_seq_len else 0 + + def _kv_len(per_block_kv): + # text_cache.{video,audio}_kv is a per-block list of (K, V). + if not per_block_kv: + return 0 + return per_block_kv[0][0].shape[1] + + # The stage-2 stack is wrapped with its own Ulysses group, whose degree + # is unrelated to the mapping's; the blocks size their sites from it. + ulysses_size = None + if self._active_topology == "stage2" and self._stage2_groups is not None: + ulysses_size = torch_dist.get_world_size(group=self._stage2_groups.ulysses_group) + + return self.transformer_blocks[0].create_attn_metadata( + self.attn_backend_metadata_cls, + batch_size=batch_size, + video_seq=video_seq, + audio_seq_full=audio_seq_full, + text_kv_video_len=_kv_len(text_cache.video_kv), + text_kv_audio_len=_kv_len(text_cache.audio_kv), + ulysses_size=ulysses_size, + ) + def forward( self, video: Modality | None, @@ -2223,6 +2441,7 @@ def forward( perturbations=None, *, text_cache: TextCache, + attn_metadata: Dict[str, AttentionMetadata], timestep: torch.Tensor | None = None, step_index=None, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: @@ -2232,6 +2451,10 @@ def forward( video: Video modality input (or None). audio: Audio modality input (or None). perturbations: Optional ``BatchedPerturbationConfig`` for STG. + attn_metadata: One prepared metadata object per attention site. + LTX-2 has six sites with distinct Q/KV lengths; each block looks + up the one it needs by name. Built by the pipeline from + :meth:`create_attn_metadata`. text_cache: Pre-computed step-invariant outputs from ``prepare_text_cache()``. Always required — callers must invoke ``prepare_text_cache()`` first. timestep: Normalized denoising-time coordinate in ``[0, 1]``. @@ -2332,6 +2555,7 @@ def forward( vx, ax, perturbations=perturbations, + attn_metadata=attn_metadata, step_index=step_index, ) if video_args is not None and vx is not None: @@ -2346,6 +2570,7 @@ def forward( perturbations=perturbations, text_kv_video=v_kv[i] if v_kv else None, text_kv_audio=a_kv[i] if a_kv else None, + attn_metadata=attn_metadata, step_index=step_index, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 7244a724649b..a9c66eee5445 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -19,7 +19,9 @@ import torch import torch.nn as nn +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.attention_backend.utils import get_visual_gen_attention_backend from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig @@ -27,6 +29,39 @@ from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner +def _attn_metadata_shape_key(*args, **kwargs): + """CUDA graph key from every ``attn_metadata*`` keyword, or ``None`` if absent. + + Accepts a single site or a dict of them; the tensor-shape key cannot see + metadata, so sites of differing length must not share a graph. + """ + sites: dict[str, AttentionMetadata] = {} + for name, value in kwargs.items(): + if not name.startswith("attn_metadata"): + continue + if isinstance(value, AttentionMetadata): + sites[name] = value + elif isinstance(value, dict): + sites.update( + {f"{name}.{k}": v for k, v in value.items() if isinstance(v, AttentionMetadata)} + ) + + key = [] + for name in sorted(sites): + metadata = sites[name] + seq_lens = metadata.seq_lens + seq_lens_kv = metadata.seq_lens_kv + key.append( + ( + name, + tuple(seq_lens.tolist()) if seq_lens is not None else None, + # Only a cross site's KV length is independent of the Q length. + tuple(seq_lens_kv.tolist()) if metadata.is_cross else None, + ) + ) + return tuple(key) or None + + class BaseDiffusionModel(nn.Module): """Base class for TRT-LLM VisualGen model components.""" @@ -36,12 +71,21 @@ def __init__(self, model_config: DiffusionModelConfig): self.component_name = model_config.component_name self.pretrained_config = model_config.pretrained_config + @property + def attn_backend_metadata_cls(self) -> type[AttentionMetadata]: + """Metadata type this component's attention backend expects.""" + return get_visual_gen_attention_backend(self.model_config.attention.backend).Metadata + def forward(self, *args, timestep: torch.Tensor | None = None, **kwargs): """Run the diffusion transformer. Concrete VisualGen models own their full forward signatures. This base method defines the common arguments that every forward should accept. + Attention metadata is a required argument, one object per attention + site. Callers build it from ``create_attn_metadata()``; a model must + never construct it, which is illegal under CUDA graph capture. + Args: timestep: Normalized denoising-time coordinate in ``[0, 1]``. Larger values correspond to earlier, noisier denoising steps. @@ -60,19 +104,13 @@ def forward(self, *args, timestep: torch.Tensor | None = None, **kwargs): def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: """Register CUDA graph key contributors that are not tensor shapes. - Override this hook when a model.forward input changes captured - execution without changing tensor shapes. Implementations should call - ``runner.register_extra_key_fn(name, fn)``, where ``fn`` is the - callback. - - The callback receives the wrapped model.forward ``*args`` and - ``**kwargs`` and returns either a hashable key value or ``None``. - If the callback returns ``None``, the runner omits that key part for - the current call. - - Subclasses should call ``super()`` unless they intentionally replace - the shared registrations. + Override when a forward input changes captured execution without + changing tensor shapes, calling ``runner.register_extra_key_fn(name, + fn)``; ``fn`` takes the forward args and returns a hashable key or + ``None`` to omit it. Subclasses should call ``super()``. """ + runner.register_extra_key_fn("attn_metadata_shape", _attn_metadata_shape_key) + sparse_config = self.model_config.attention.sparse_attention_config if not isinstance(sparse_config, SkipSoftmaxAttentionConfig): return @@ -81,10 +119,8 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: if disabled_until_timestep is None: return - # Skip Softmax switches graph-visible attention behavior at the - # timestep boundary while tensor shapes stay unchanged. Key the dense - # and sparse phases separately; if timestep is absent or None, the - # scheduler returns None and the runner omits this key part. + # Skip Softmax switches attention behavior at the timestep boundary + # without changing shapes, so key its dense and sparse phases apart. runner.register_extra_key_fn( "skip_softmax_phase", lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py index b7d3a69015cf..44079c053817 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py @@ -24,6 +24,9 @@ ExtractorConfig, register_extractor_from_config, ) +from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import ( + qwen_image_attn_metadata, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline @@ -625,6 +628,14 @@ def forward( separate_cfg = use_negative_prompt_cfg and not do_cfg_parallel cache_acc.refresh(len(timesteps), separate_cfg=separate_cfg) + # Create attention metadata + attn_metadata_pos = qwen_image_attn_metadata(self.transformer, latents, prompt_embeds) + attn_metadata_neg = ( + None + if neg_prompt_embeds is None + else qwen_image_attn_metadata(self.transformer, latents, neg_prompt_embeds) + ) + for i, t in self._profile_denoise_steps(timesteps): timestep = t.expand(latents.shape[0]).to(latents.dtype) if do_cfg_parallel: @@ -639,6 +650,7 @@ def forward( ) noise_pred_local = self.transformer( hidden_states=latents, + attn_metadata=(attn_metadata_pos if cfg_rank == 0 else attn_metadata_neg), timestep=timestep / 1000, encoder_hidden_states_mask=local_mask, encoder_hidden_states=local_embeds, @@ -655,6 +667,7 @@ def forward( self.transformer._cache_branch = None noise_pred = self.transformer( hidden_states=latents, + attn_metadata=attn_metadata_pos, timestep=timestep / 1000, encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, @@ -671,6 +684,7 @@ def forward( self.transformer._cache_branch = "uncond" neg_noise_pred = self.transformer( hidden_states=latents, + attn_metadata=attn_metadata_neg, timestep=timestep / 1000, encoder_hidden_states_mask=neg_prompt_embeds_mask, encoder_hidden_states=neg_prompt_embeds, diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py index 6293b28a2ddc..8831fa51e4dd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py @@ -17,6 +17,9 @@ import torch import torch.distributed as dist +from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import ( + qwen_image_attn_metadata, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline from tensorrt_llm.inputs.utils import load_image @@ -487,6 +490,7 @@ def forward( self.scheduler.set_begin_index(0) timer.mark_denoise_start() + logger.info("Denoising edit (%d steps)...", len(timesteps)) for _, t in self._profile_denoise_steps(timesteps): latent_model_input = torch.cat([latents, image_latents], dim=1) @@ -496,6 +500,11 @@ def forward( if cfg_rank == 0: local_noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, + latent_model_input, + prompt_embeds, + ), timestep=timestep / 1000, encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, @@ -505,6 +514,11 @@ def forward( else: local_noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, + latent_model_input, + neg_prompt_embeds, + ), timestep=timestep / 1000, encoder_hidden_states_mask=neg_prompt_embeds_mask, encoder_hidden_states=neg_prompt_embeds, @@ -523,6 +537,9 @@ def forward( else: noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, latent_model_input, prompt_embeds + ), timestep=timestep / 1000, encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, @@ -542,6 +559,11 @@ def forward( noise_pred = noise_pred.clone() neg_noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, + latent_model_input, + neg_prompt_embeds, + ), timestep=timestep / 1000, encoder_hidden_states_mask=neg_prompt_embeds_mask, encoder_hidden_states=neg_prompt_embeds, diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py index b8f49937beb0..69a289e960fe 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py @@ -26,14 +26,17 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.utils import gelu_tanh, maybe_compile +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( Attention2DAttention, RingAttention, UlyssesAttention, + get_ulysses_seq_lens, ) from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel @@ -505,7 +508,6 @@ def __init__( config=config, layer_idx=layer_idx, module_name=module_name, - separate_qkv_is_self_attention=True, ) self.head_dim = attention_head_dim self._supports_key_padding_mask = _supports_qwen_key_padding_mask( @@ -684,6 +686,7 @@ def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fused_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, @@ -702,7 +705,9 @@ def forward( attn_kwargs = {} if attention_mask is not None: attn_kwargs["key_padding_mask"] = attention_mask - out = self._attn_impl(joint_q, joint_k, joint_v, timestep=timestep, **attn_kwargs) + out = self._attn_impl( + joint_q, joint_k, joint_v, attn_metadata, timestep=timestep, **attn_kwargs + ) elif self._uses_sequence_parallel_attention: raise NotImplementedError( "Padded Qwen-Image prompts require a key-padding-mask-capable " @@ -832,6 +837,7 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, + attn_metadata: AttentionMetadata, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, fused_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, @@ -851,6 +857,7 @@ def forward( img_attn_output, txt_attn_output = self.attn( hidden_states=img_modulated, encoder_hidden_states=txt_modulated, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, fused_rotary_emb=fused_rotary_emb, attention_mask=attention_mask, @@ -913,6 +920,23 @@ def _build_joint_attention_mask( return torch.cat([encoder_hidden_states_mask, image_mask], dim=1) +def qwen_image_attn_metadata( + model, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, +) -> AttentionMetadata: + """The single joint text+image site a Qwen-Image forward needs. + + Shared by the Qwen-Image, Qwen-Image-Edit and Qwen-Image-Layered pipelines, + which all call the transformer from several places per denoising step. + """ + return model.create_attn_metadata( + batch_size=hidden_states.shape[0], + text_seq_len=encoder_hidden_states.shape[1], + image_seq_len=hidden_states.shape[1], + )["self"] + + class QwenImageTransformer2DModel(BaseDiffusionModel): """Qwen-Image 20B MMDiT transformer. @@ -1267,9 +1291,41 @@ def post_load_weights(self) -> None: f"weight_scale_dtype={scale_dtype}" ) from exc + def create_attn_metadata( + self, + *, + batch_size: int, + text_seq_len: int, + image_seq_len: int, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + Qwen-Image has a single attention site: the joint text+image sequence. + """ + # forward() pads each stream up to a multiple of the shard count and + # shards it before concatenating. + size = self.sharder.size + local_seq_len = (image_seq_len + (-image_seq_len) % size) // size + ( + text_seq_len + (-text_seq_len) % size + ) // size + q_seq_len, kv_seq_len = get_ulysses_seq_lens( + local_seq_len, + local_seq_len, + visual_gen_mapping=self.model_config.visual_gen_mapping, + ) + return { + "self": make_diffusion_attn_metadata( + self.attn_backend_metadata_cls, + batch_size=batch_size, + q_seq_lens=q_seq_len, + kv_seq_lens=None if kv_seq_len == q_seq_len else kv_seq_len, + ) + } + def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: Optional[torch.Tensor] = None, timestep: Optional[torch.Tensor] = None, @@ -1332,6 +1388,7 @@ def forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, fused_rotary_emb=fused_rotary_emb, attention_mask=block_attention_mask, diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py index 113a329a27f3..c9e51c41a11b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py @@ -22,6 +22,9 @@ import numpy as np import torch +from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import ( + qwen_image_attn_metadata, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline @@ -906,12 +909,16 @@ def forward( additional_t_cond = torch.zeros(batch_size, device=device, dtype=torch.long) timer.mark_denoise_start() + logger.info("Denoising layered output (%d steps)...", len(timesteps)) for _, t in self._profile_denoise_steps(timesteps): latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latents.shape[0]).to(latents.dtype) noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, latent_model_input, prompt_embeds + ), timestep=timestep / 1000, encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, @@ -924,6 +931,11 @@ def forward( if do_true_cfg: neg_noise_pred = self.transformer( hidden_states=latent_model_input, + attn_metadata=qwen_image_attn_metadata( + self.transformer, + latent_model_input, + neg_prompt_embeds, + ), timestep=timestep / 1000, encoder_hidden_states_mask=neg_prompt_embeds_mask, encoder_hidden_states=neg_prompt_embeds, diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py index 0968f3420a91..26685f2ee526 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py @@ -21,6 +21,8 @@ import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata + from ..qwen_image.transformer_qwen_image import ( QwenEmbedRope, QwenImageTransformer2DModel, @@ -197,6 +199,7 @@ def from_config_dict( def forward( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: Optional[torch.Tensor] = None, timestep: Optional[torch.Tensor] = None, @@ -245,6 +248,7 @@ def forward( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, + attn_metadata=attn_metadata, image_rotary_emb=image_rotary_emb, attention_mask=block_attention_mask, timestep=timestep, diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 8230f10823e5..b0344c817be8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -37,7 +37,10 @@ get_wan_default_params, get_wan_extra_param_specs, ) -from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import retrieve_latents +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import ( + retrieve_latents, + wan_attn_metadata_kwargs, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline @@ -554,6 +557,25 @@ def forward( last_model_used = [None] _vsa_step_counter = [0] + # Create attention metadata + denoise_batch = self.denoise_batch_size(latents, guidance_scale=guidance_scale) + attn_metadata_high = wan_attn_metadata_kwargs( + self.transformer, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + batch_size=denoise_batch, + ) + attn_metadata_low = ( + None + if self.transformer_2 is None + else wan_attn_metadata_kwargs( + self.transformer_2, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + batch_size=denoise_batch, + ) + ) + def forward_fn( latents, extra_stream_latents, @@ -587,6 +609,8 @@ def forward_fn( else: current_model = self.transformer + sites = attn_metadata_high if current_model is self.transformer else attn_metadata_low + # Build per-patch 2D timestep for Wan 2.2 TI2V-5B if self.is_wan22_5b: _, ph, pw = self.transformer.config.patch_size @@ -615,12 +639,14 @@ def forward_fn( with set_vsa_forward_context(vsa_metadata): return current_model( hidden_states=latents, + **sites, timestep=timestep / self.scheduler.config.num_train_timesteps, encoder_hidden_states=encoder_hidden_states, ) return current_model( hidden_states=latents, + **sites, timestep=timestep / self.scheduler.config.num_train_timesteps, encoder_hidden_states=encoder_hidden_states, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index f83d5be29f1b..1540054f623c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -34,7 +34,10 @@ get_wan_default_params, get_wan_extra_param_specs, ) -from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import retrieve_latents +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import ( + retrieve_latents, + wan_attn_metadata_kwargs, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline @@ -577,6 +580,21 @@ def forward( # Track which model was used in last step (for logging model transitions) last_model_used = [None] + # Build attention metadata + denoise_batch = self.denoise_batch_size(latents, guidance_scale=guidance_scale) + + def _build_sites(model): + return wan_attn_metadata_kwargs( + model, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_image=image_embeds, + batch_size=denoise_batch, + ) + + attn_metadata_high = _build_sites(self.transformer) + attn_metadata_low = None if self.transformer_2 is None else _build_sites(self.transformer_2) + def forward_fn( latents_input, extra_stream_latents, @@ -633,8 +651,11 @@ def forward_fn( repeat_factor = latents_input.shape[0] // image_embeds_to_use.shape[0] image_embeds_to_use = image_embeds_to_use.repeat(repeat_factor, 1, 1) + sites = attn_metadata_high if current_model is self.transformer else attn_metadata_low + return current_model( hidden_states=latent_model_input, + **sites, timestep=timestep_input / self.scheduler.config.num_train_timesteps, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_image=image_embeds_to_use, diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_utils.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_utils.py index b6d31e94b458..085b28127374 100755 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_utils.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_utils.py @@ -12,10 +12,40 @@ # 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. -from typing import Any, Optional +from typing import Any, Dict, Optional import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata + + +def wan_attn_metadata_kwargs( + model: Any, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, +) -> Dict[str, AttentionMetadata]: + """Name WAN's attention metadata sites for ``forward``. + + Maps the site names from ``WanTransformer3DModel.create_attn_metadata()`` + onto the ``attn_metadata_*`` keyword arguments of its ``forward``. + """ + sites = model.create_attn_metadata( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + batch_size=batch_size, + ) + kwargs = { + "attn_metadata_self": sites["self"], + "attn_metadata_cross_text": sites["cross_text"], + } + if "cross_image" in sites: + kwargs["attn_metadata_cross_image"] = sites["cross_image"] + return kwargs + def retrieve_latents( encoder_output: Any, 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 057e088138cd..d40446d69367 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -1,5 +1,5 @@ import math -from typing import Optional, Tuple +from typing import Dict, Optional, Tuple import torch import torch.nn as nn @@ -7,11 +7,14 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from tqdm import tqdm +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.models.hf_parameter_utils import get_parameter_device from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.utils import Fp4QuantizedTensor, gelu_tanh +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata +from tensorrt_llm._torch.visual_gen.attention_backend.parallel import get_ulysses_seq_lens from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.models.wan.utils_wan import ( @@ -349,6 +352,7 @@ def __init__( layer_idx=_layer_idx, module_name=f"blocks.{_layer_idx}.attn2", enable_sequence_parallel=False, + is_cross=True, ) if cross_attn_norm: @@ -516,8 +520,18 @@ def forward( temb, freqs_cos, freqs_sin, + attn_metadata_self, + attn_metadata_cross_text, + attn_metadata_cross_image=None, timestep=None, ): + """WAN block forward. + + WAN is a mixed-stream block: ``attn1`` self-attends over the video + sequence, while ``attn2`` cross-attends the same queries against the text + stream and (for I2V) the image stream, each with its own KV length. Each + site therefore gets its own metadata site rather than sharing one. + """ if temb.ndim == 4: # temb: batch_size, seq_len, 6, hidden_size shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( @@ -559,9 +573,13 @@ def forward( # so each V/Q/K GEMM + norm + RoPE overlaps with the peer push on the # side stream; both paths return 3D [B, S, H*D]. if self._use_async_ulysses: - attn1_out = self.attn1.forward_async(normed, freqs=freqs, timestep=timestep) + attn1_out = self.attn1.forward_async( + normed, attn_metadata_self, freqs=freqs, timestep=timestep + ) else: - attn1_out = self.attn1(normed, freqs=freqs, timestep=timestep, **attn1_kwargs) + attn1_out = self.attn1( + normed, attn_metadata_self, freqs=freqs, timestep=timestep, **attn1_kwargs + ) x = (x.float() + attn1_out.float() * gate_msa).to(x.dtype) @@ -598,6 +616,7 @@ def forward( q, k, v, + attn_metadata_cross_text, batch_size=batch_size, seq_len=seq_len, kv_seq_len=encoder_hidden_states_text.shape[1], @@ -609,10 +628,17 @@ def forward( key_img = self.add_k_proj(encoder_hidden_states_img) value_img = self.add_v_proj(encoder_hidden_states_img) key_img = self.norm_added_k(key_img) + if attn_metadata_cross_image is None: + raise ValueError( + "WAN I2V image cross-attention requires attn_metadata_cross_image; " + "the pipeline must include the 'cross_image' site when " + "encoder_hidden_states_image is provided." + ) attn_img_output = self.attn2._attn_impl( q, key_img, value_img, + attn_metadata_cross_image, batch_size=batch_size, seq_len=seq_len, kv_seq_len=encoder_hidden_states_img.shape[1], @@ -801,9 +827,74 @@ def unpatchify(self, x, original_shape): .reshape(N, out_channels, T, H, W) ) + #: Text context length WanBlock assumes when splitting a concatenated + #: image+text encoder state (see ``WanBlock.forward``). + TEXT_CONTEXT_LENGTH = 512 + + def video_seq_len(self, hidden_states: torch.Tensor) -> int: + """Number of video tokens after patch embedding.""" + _, _, T, H, W = hidden_states.shape + pt, ph, pw = self.config.patch_size + return (T // pt) * (H // ph) * (W // pw) + + def create_attn_metadata( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, + ) -> Dict[str, AttentionMetadata]: + """Attention metadata for this model's sites, one object per site. + + WAN is a mixed-stream model with three attention sites and three different KV lengths: + ``self`` (video self-attention), ``cross_text`` (video queries against the text stream) and, + when the checkpoint supplies CLIP image embeddings (Wan 2.1 I2V), ``cross_image``. + """ + metadata_cls = self.attn_backend_metadata_cls + vgm = self.model_config.visual_gen_mapping + if batch_size is None: + batch_size = encoder_hidden_states.shape[0] + # forward() shards the video sequence, so the blocks see this rank's share. + s_video = self.video_seq_len(hidden_states) // self.sharder.size + has_image = ( + encoder_hidden_states_image is not None + and self.condition_embedder.image_embedder is not None + ) + + # Mirrors WanBlock.forward: with an image stream the encoder state is + # [image; text] and the text tail is a fixed length. + s_text = self.TEXT_CONTEXT_LENGTH if has_image else encoder_hidden_states.shape[1] + + q_self, kv_self = get_ulysses_seq_lens(s_video, s_video, visual_gen_mapping=vgm) + sites = { + "self": make_diffusion_attn_metadata( + metadata_cls, + batch_size=batch_size, + q_seq_lens=q_self, + kv_seq_lens=None if kv_self == q_self else kv_self, + ), + "cross_text": make_diffusion_attn_metadata( + metadata_cls, batch_size=batch_size, q_seq_lens=s_video, kv_seq_lens=s_text + ), + } + if has_image: + # The image stream is projected by WanImageEmbedding, which preserves + # its sequence length, so the raw tensor's length is the KV length. + sites["cross_image"] = make_diffusion_attn_metadata( + metadata_cls, + batch_size=batch_size, + q_seq_lens=s_video, + kv_seq_lens=encoder_hidden_states_image.shape[1], + ) + return sites + def forward( self, hidden_states, + attn_metadata_self, + attn_metadata_cross_text, + attn_metadata_cross_image=None, timestep=None, encoder_hidden_states=None, encoder_hidden_states_image=None, @@ -889,6 +980,9 @@ def forward( temb_proj, freqs_cos, freqs_sin, + attn_metadata_self, + attn_metadata_cross_text, + attn_metadata_cross_image, timestep=timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index d551720cb4e6..af1f32d117db 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -8,7 +8,7 @@ from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ...utils import Fp4QuantizedTensor -from ..attention_backend.interface import AttentionTensorLayout +from ..attention_backend.interface import AttentionMetadata, AttentionTensorLayout from ..attention_backend.parallel import wrap_parallel_attention from ..attention_backend.utils import create_attention from ..config import DiffusionModelConfig @@ -21,7 +21,7 @@ class QKVMode(str, Enum): SEPARATE_QKV = "separate" -# TODO: torch compile +# FIXME: does trtllm offer fused routine for this via torch.ops.trtllm? def apply_rotary_emb( x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor ) -> torch.Tensor: @@ -56,7 +56,7 @@ def __init__( module_name: Optional[str] = None, enable_sequence_parallel: bool = True, async_ulysses: bool = False, - separate_qkv_is_self_attention: bool = False, + is_cross: bool = False, ): super().__init__() @@ -96,16 +96,32 @@ def __init__( cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config + _qa_cfg = config.attention.quant_attention_config _is_vsa = ( base_backend == "CUTEDSL" and _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. - if self.qkv_mode == QKVMode.SEPARATE_QKV and (base_backend == "TRTLLM" or _is_vsa): + # Routing depends only on `is_cross`; `qkv_mode` describes the + # projection before it and has no bearing here. + self.is_cross = is_cross + + is_sage = base_backend == "TRTLLM" and _qa_cfg is not None + + if not is_cross: + backend_name = base_backend + elif base_backend == "TRTLLM": + # Plain TRTLLM needs packed QKV, impossible at unequal lengths; + # SageAttention's separate-Q/K/V layout is the one that works. + backend_name = base_backend if is_sage else "VANILLA" + elif _is_vsa: + # VSA sparsifies a 3-D video token grid; a separate K/V stream + # has no position in it. backend_name = "VANILLA" else: + # VANILLA, FA4 and the CuTe DSL FMHA kernels take Q and K/V of + # different lengths directly. backend_name = base_backend if _is_vsa and cp_size > 1: @@ -114,6 +130,14 @@ def __init__( f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) + if _is_vsa and async_ulysses and ulysses_size > 1: + # VSA's per-token gates arrive as forward kwargs, which + # `forward_async`'s fixed argument list cannot carry. + raise ValueError( + "VSA is incompatible with async Ulysses: the gate tensors VSA " + "needs cannot be passed through the async projection path. Set " + "async_ulysses=False to use VSA." + ) self.attn_backend = backend_name self.qk_norm = qk_norm self.qk_norm_mode = qk_norm_mode @@ -141,8 +165,6 @@ def __init__( and not self.force_dynamic_quantization ) - attention_metadata_state = getattr(config, "attention_metadata_state", None) - if self.qk_norm: # "full": norm over all heads combined (e.g. WAN, dim=q_dim) # "per_head": norm over each head independently (e.g. FLUX, dim=head_dim) @@ -231,21 +253,17 @@ def __init__( quant_config=self.quant_config, dtype=self.dtype, attention_config=config.attention, - attention_metadata_state=attention_metadata_state, sparse_params=sparse_params, ) - if ( - enable_sequence_parallel - and self.qkv_mode == QKVMode.SEPARATE_QKV - and not separate_qkv_is_self_attention - and vgm is not None - ): + if enable_sequence_parallel and is_cross and vgm is not None: ring_size = vgm.ring_size if ring_size > 1: + # Ring chunks the K/V stream across ranks and accumulates over + # steps, which assumes every rank holds a slice of one sequence. raise ValueError( - "SEPARATE_QKV cross-attention does not support Ring sequence " - "parallelism; use enable_sequence_parallel=False or Ulysses/Attention2D." + "Cross-attention does not support Ring sequence parallelism; " + "use enable_sequence_parallel=False or Ulysses/Attention2D." ) self.attn = wrap_parallel_attention( @@ -515,6 +533,7 @@ def _attn_impl( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + attn_metadata: AttentionMetadata, **kwargs, ) -> torch.Tensor: """ @@ -568,7 +587,7 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor: if kwargs.get(gate_key) is not None: kwargs[gate_key] = _reshape_gate(kwargs[gate_key]) - out = self.attn.forward(q=q, k=k, v=v, **kwargs) + out = self.attn.forward(q=q, k=k, v=v, attn_metadata=attn_metadata, **kwargs) # Flatten back to [B, S, H*D] if backend_layout == AttentionTensorLayout.HND: @@ -579,6 +598,7 @@ def _reshape_gate(gate: torch.Tensor) -> torch.Tensor: def forward( self, hidden_states: torch.Tensor | Fp4QuantizedTensor, + attn_metadata: AttentionMetadata, encoder_hidden_states: Optional[torch.Tensor] = None, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, @@ -604,7 +624,7 @@ def forward( freqs_cos, freqs_sin = freqs self.apply_packed_qk_norm_rope(qkv, freqs_cos, freqs_sin) q, k, v = qkv.split([self.local_q_dim, self.local_kv_dim, self.local_kv_dim], dim=-1) - out = self._attn_impl(q, k, v, timestep=timestep, **kwargs) + out = self._attn_impl(q, k, v, attn_metadata, timestep=timestep, **kwargs) return self.to_out[0](out) # Unfused path: separate QK norm → separate RoPE → attention @@ -623,13 +643,14 @@ def forward( q = q.flatten(2) k = k.flatten(2) - out = self._attn_impl(q, k, v, timestep=timestep, **kwargs) + out = self._attn_impl(q, k, v, attn_metadata, timestep=timestep, **kwargs) out = self.to_out[0](out) return out def forward_async( self, hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -728,6 +749,8 @@ def compute_k(): def compute_v(): return self.to_v(qkv_input).view(B, S, KV, D) - out_4d = self.attn.forward_async(compute_q, compute_k, compute_v, timestep=timestep) + out_4d = self.attn.forward_async( + compute_q, compute_k, compute_v, attn_metadata=attn_metadata, timestep=timestep + ) b, t = out_4d.shape[:2] return self.to_out[0](out_4d.reshape(b, t, H * D)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 089fdd5f8034..55ec6731fd83 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -793,6 +793,24 @@ def _resolve_step_guidance_scale( current = 1.0 return current + def denoise_batch_size( + self, + latents: torch.Tensor, + *, + guidance_scale: float, + ) -> int: + """Batch size ``forward_fn`` will be called with. + + CFG parallel is supported, so each ``forward_fn`` call might not contain all CFGs for a + model. Use this routine to determine how many batches ``forward_fn`` will see from the + supplied latents tensor. + """ + vgm = self.pipeline_config.visual_gen_mapping + cfg_size = vgm.cfg_size if vgm else 1 + do_cfg_parallel = cfg_size >= 2 and guidance_scale > 1.0 + doubles_batch = guidance_scale > 1.0 and not do_cfg_parallel + return latents.shape[0] * (2 if doubles_batch else 1) + def _setup_cfg_config( self, guidance_scale, prompt_embeds, neg_prompt_embeds, extra_cfg_tensors=None ): diff --git a/tests/unittest/_torch/visual_gen/attn_metadata_utils.py b/tests/unittest/_torch/visual_gen/attn_metadata_utils.py new file mode 100644 index 000000000000..b853a9c79b8c --- /dev/null +++ b/tests/unittest/_torch/visual_gen/attn_metadata_utils.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test helpers that derive attention metadata from the tensors a test passes in.""" + +from typing import Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata +from tensorrt_llm._torch.visual_gen.attention_backend.utils import get_visual_gen_attention_backend + + +def make_attn_metadata( + backend: str, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + *, + q_seq_len: Optional[int] = None, + kv_seq_len: Optional[int] = None, +) -> AttentionMetadata: + """Metadata for one site, sized from ``hidden_states`` and an optional KV stream. + + ``q_seq_len`` / ``kv_seq_len`` override the lengths read off the tensors. + """ + if kv_seq_len is None and encoder_hidden_states is not None: + kv_seq_len = encoder_hidden_states.shape[1] + + return make_diffusion_attn_metadata( + get_visual_gen_attention_backend(backend).Metadata, + batch_size=hidden_states.shape[0], + q_seq_lens=q_seq_len if q_seq_len is not None else hidden_states.shape[1], + kv_seq_lens=kv_seq_len, + ) + + +def make_backend_attn_metadata( + backend, + q: torch.Tensor, + k: Optional[torch.Tensor] = None, +) -> AttentionMetadata: + """Metadata for ``q``/``k`` passed straight to a backend, in its own layout.""" + seq_axis = 2 if backend.preferred_layout == AttentionTensorLayout.HND else 1 + return make_diffusion_attn_metadata( + type(backend).Metadata, + batch_size=q.shape[0], + q_seq_lens=q.shape[seq_axis], + kv_seq_lens=None if k is None else k.shape[seq_axis], + ) + + +def flux_attn_metadata(model, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor): + """The single joint text+image site a FLUX / FLUX.2 transformer forward needs.""" + return model.create_attn_metadata( + batch_size=hidden_states.shape[0], + text_seq_len=encoder_hidden_states.shape[1], + image_seq_len=hidden_states.shape[1], + )["self"] + + +def cosmos3_attn_metadata_kwargs( + model, + hidden_states: torch.Tensor, + text_mask: torch.Tensor, + video_shape, + audio_latents: Optional[torch.Tensor] = None, +) -> dict: + """Cosmos3's ``und`` / ``mixed`` sites as ``forward`` keyword arguments.""" + sites = model.create_attn_metadata( + batch_size=hidden_states.shape[0], + text_seq_len=text_mask.shape[1], + text_lens=text_mask.sum(dim=1).tolist(), + video_shape=video_shape, + num_audio_tokens=(audio_latents.shape[2] if audio_latents is not None else 0), + ) + return { + "attn_metadata_und": sites["und"], + "attn_metadata_mixed": sites["mixed"], + "attn_metadata_mixed_ragged": sites["mixed_ragged"], + } + + +def ltx2_attn_metadata(model, video, audio, text_cache) -> dict: + """LTX-2's sites, unwrapping the lengths from the ``Modality`` bundles.""" + batch_size = 1 + if video is not None: + batch_size = video.latent.shape[0] + elif audio is not None: + batch_size = audio.latent.shape[0] + return model.create_attn_metadata( + batch_size=batch_size, + video_seq_len=0 if video is None else video.latent.shape[1], + audio_seq_len=0 if audio is None else audio.latent.shape[1], + text_cache=text_cache, + ) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py index dcdc4a023a45..cc781fe1d78b 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py @@ -48,7 +48,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import make_backend_attn_metadata MODULES_AVAILABLE = True except ImportError: @@ -228,11 +230,18 @@ def _logic_attn2d_forward(rank, world_size): v = torch.randn(batch, seq_per_rank, num_heads, head_dim, device=device) # No mask (None) and explicit FULL mask are both valid - output = attn(q, k, v, batch_size=batch) + output = attn(q, k, v, attn_metadata=make_backend_attn_metadata(attn, q, k), 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}: output contains non-finite values" - output_full = attn(q, k, v, batch_size=batch, attention_mask=PredefinedAttentionMask.FULL) + output_full = attn( + q, + k, + v, + attn_metadata=make_backend_attn_metadata(attn, q, k), + batch_size=batch, + attention_mask=PredefinedAttentionMask.FULL, + ) assert output_full.shape == q.shape @@ -263,7 +272,13 @@ def _logic_attn2d_vs_standard(rank, world_size): except ImportError: pytest.skip("flash_attn_combine JIT kernels not available") - attn2d_output = attn(q_shard, k_shard, v_shard, batch_size=batch) + attn2d_output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) # Reference: standard full-sequence SDPA scale = 1.0 / math.sqrt(head_dim) @@ -301,7 +316,14 @@ def _logic_attn2d_invalid_mask(rank, world_size): v = torch.randn(2, 8, 8, 64, device=device) with pytest.raises(ValueError, match="FULL"): - attn(q, k, v, batch_size=2, attention_mask=PredefinedAttentionMask.CAUSAL) + attn( + q, + k, + v, + attn_metadata=make_backend_attn_metadata(attn, q, k), + batch_size=2, + attention_mask=PredefinedAttentionMask.CAUSAL, + ) def _logic_attn2d_asymmetric_mesh_1x4(rank, world_size): @@ -329,7 +351,13 @@ def _logic_attn2d_asymmetric_mesh_1x4(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - output = attn(q_shard, k_shard, v_shard, batch_size=batch) + output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert output.shape == q_shard.shape, ( f"Rank {rank}: expected {q_shard.shape}, got {output.shape}" @@ -381,7 +409,13 @@ def _logic_attn2d_asymmetric_mesh_4x1(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - output = attn(q_shard, k_shard, v_shard, batch_size=batch) + output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert output.shape == q_shard.shape, ( f"Rank {rank}: expected {q_shard.shape}, got {output.shape}" @@ -433,7 +467,13 @@ def _logic_attn2d_gqa(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - attn2d_output = attn(q_shard, k_shard, v_shard, batch_size=batch) + attn2d_output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) scale = 1.0 / math.sqrt(head_dim) q_std = q_full.transpose(1, 2).float() @@ -479,7 +519,13 @@ def _logic_attn2d_cross_attention(rank, world_size): k_shard = k_full[:, rank * seq_per_rank_kv : (rank + 1) * seq_per_rank_kv].contiguous() v_shard = v_full[:, rank * seq_per_rank_kv : (rank + 1) * seq_per_rank_kv].contiguous() - attn2d_output = attn(q_shard, k_shard, v_shard, batch_size=batch) + attn2d_output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert attn2d_output.shape == q_shard.shape scale = 1.0 / math.sqrt(head_dim) @@ -581,7 +627,13 @@ def _logic_attn2d_fa4_vs_standard(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - output = attn(q_shard, k_shard, v_shard, batch_size=batch) + output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert output.shape == q_shard.shape, ( f"Rank {rank}: expected {q_shard.shape}, got {output.shape}" @@ -642,7 +694,13 @@ def _logic_attn2d_fa4_asymmetric_1x4(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - output = attn(q_shard, k_shard, v_shard, batch_size=batch) + output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert output.shape == q_shard.shape, ( f"Rank {rank}: expected {q_shard.shape}, got {output.shape}" @@ -703,7 +761,13 @@ def _logic_attn2d_fa4_asymmetric_4x1(rank, world_size): k_shard = k_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() v_shard = v_full[:, rank * seq_per_rank : (rank + 1) * seq_per_rank].contiguous() - output = attn(q_shard, k_shard, v_shard, batch_size=batch) + output = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) assert output.shape == q_shard.shape, ( f"Rank {rank}: expected {q_shard.shape}, got {output.shape}" diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index afc39db47aa4..3ce78c7f3e2c 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -41,7 +41,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import cosmos3_attn_metadata_kwargs from tensorrt_llm.models.modeling_utils import QuantConfig @@ -424,6 +426,7 @@ def _forward( with torch.inference_mode(): return model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -451,6 +454,7 @@ def _forward_with_audio( with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape, audio_latents), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py index e72a1cd67d7b..80f41bc58b4f 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py @@ -48,7 +48,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import flux_attn_metadata MODULES_AVAILABLE = True except ImportError: @@ -295,6 +297,7 @@ def _logic_flux2_transformer_parallel_vs_single_gpu( with torch.no_grad(): ref_out = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, guidance=guidance, @@ -303,6 +306,7 @@ def _logic_flux2_transformer_parallel_vs_single_gpu( )["sample"] dist_out = dist_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(dist_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, guidance=guidance, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py index 88bea80e3e09..5780e884cc62 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py @@ -42,7 +42,6 @@ AttentionConfig, DiffusionModelConfig, TorchCompileConfig, - create_attention_metadata_state, ) from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import ( @@ -53,7 +52,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import flux_attn_metadata from tensorrt_llm.models.modeling_utils import QuantConfig @@ -199,9 +200,6 @@ def _make_model_config(pretrained_dict, tp_size=1, ulysses_size=1, backend="VANI attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), skip_create_weights_in_init=False, ) config.mapping = vgm.to_llm_mapping() @@ -364,6 +362,7 @@ def _logic_flux1_tp_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -429,6 +428,7 @@ def _logic_flux1_tp_vs_single_gpu_with_config(rank, world_size, config_dict): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -437,6 +437,7 @@ def _logic_flux1_tp_vs_single_gpu_with_config(rank, world_size, config_dict): ) tp_output = tp_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(tp_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -485,6 +486,7 @@ def _logic_flux2_tp_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -548,6 +550,7 @@ def _logic_flux2_tp_vs_single_gpu_with_config(rank, world_size, config_dict): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -555,6 +558,7 @@ def _logic_flux2_tp_vs_single_gpu_with_config(rank, world_size, config_dict): ) tp_output = tp_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(tp_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -620,6 +624,7 @@ def _logic_flux2_tp_ulysses_vs_single_gpu(rank, world_size): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -627,6 +632,7 @@ def _logic_flux2_tp_ulysses_vs_single_gpu(rank, world_size): ) combined_output = combined_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(combined_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py index 1a9de4713baf..6383a9514809 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py @@ -26,16 +26,15 @@ import sys from pathlib import Path - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import flux_attn_metadata from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig @@ -163,9 +162,6 @@ def _make_model_config(pretrained_dict, ulysses_size=1, backend="VANILLA"): attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), skip_create_weights_in_init=False, ) config.mapping = vgm.to_llm_mapping() @@ -229,6 +225,7 @@ def _logic_flux1_ulysses_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -288,6 +285,7 @@ def _logic_flux1_ulysses_vs_single_gpu(rank, world_size): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -296,6 +294,7 @@ def _logic_flux1_ulysses_vs_single_gpu(rank, world_size): ) ulysses_output = ulysses_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ulysses_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -349,6 +348,7 @@ def _logic_flux2_ulysses_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -404,6 +404,7 @@ def _logic_flux2_ulysses_vs_single_gpu(rank, world_size): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ref_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -411,6 +412,7 @@ def _logic_flux2_ulysses_vs_single_gpu(rank, world_size): ) ulysses_output = ulysses_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(ulysses_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -477,6 +479,7 @@ def _logic_flux2_ulysses_trtllm_vs_vanilla(rank, world_size): with torch.no_grad(): vanilla_output = vanilla_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(vanilla_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, @@ -484,6 +487,7 @@ def _logic_flux2_ulysses_trtllm_vs_vanilla(rank, world_size): ) trtllm_output = trtllm_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(trtllm_model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, img_ids=img_ids, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py index fe21222a254e..0f7e2cc16b0f 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py @@ -39,10 +39,7 @@ import sys from pathlib import Path - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping # Spawn distributed workers via a helper that retries with a fresh master @@ -166,9 +163,6 @@ def _make_model_config( attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), parallel=ParallelConfig( ulysses_size=ulysses_size, async_ulysses=async_ulysses, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py index 9ace40540dd0..d308721f795f 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py @@ -27,10 +27,7 @@ import sys from pathlib import Path - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.rope import LTXRopeType @@ -155,9 +152,6 @@ def _make_model_config( attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), parallel=ParallelConfig(ulysses_size=ulysses_size), skip_create_weights_in_init=False, ) @@ -346,9 +340,6 @@ def _make_model_config_cfg( attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), parallel=ParallelConfig(cfg_size=cfg_size, ulysses_size=ulysses_size), skip_create_weights_in_init=False, ) @@ -574,7 +565,17 @@ def _logic_ltx2_full_audio_construction(rank, world_size, backend, audio_seq_len ) video_i, audio_i, *_ = _build_inputs(1, 16, (1, 4, 4), audio_seq_len, dtype, device) with torch.no_grad(): - v, a = model(video=video_i, audio=audio_i, text_cache=cache) + v, a = model( + video=video_i, + audio=audio_i, + text_cache=cache, + attn_metadata=model.create_attn_metadata( + batch_size=video_i.latent.shape[0], + video_seq_len=video_i.latent.shape[1], + audio_seq_len=audio_i.latent.shape[1], + text_cache=cache, + ), + ) assert v.shape[1] == 16 and a.shape[1] == audio_seq_len finally: tl._LTX2_AUDIO_CONDITIONAL_SHARD = orig 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 bcc4fd856a51..aea794427d4b 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 @@ -56,7 +56,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import make_backend_attn_metadata MODULES_AVAILABLE = True except ImportError: @@ -212,7 +214,7 @@ def _logic_ring_forward(rank, world_size): 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) + output = attn(q, k, v, attn_metadata=make_backend_attn_metadata(attn, q, k), 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" @@ -239,7 +241,13 @@ def _logic_ring_vs_standard(rank, world_size): 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) + ring_out = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) scale = 1.0 / math.sqrt(head_dim) q_std = q_full.transpose(1, 2).float() @@ -271,7 +279,14 @@ def _logic_ring_invalid_mask(rank, world_size): v = torch.randn(2, 8, 8, 64, device=device) with pytest.raises(NotImplementedError, match="only supports FULL attention mask"): - attn(q, k, v, batch_size=2, attention_mask=PredefinedAttentionMask.CAUSAL) + attn( + q, + k, + v, + attn_metadata=make_backend_attn_metadata(attn, q, k), + batch_size=2, + attention_mask=PredefinedAttentionMask.CAUSAL, + ) def _logic_ring_fa4_vs_standard(rank, world_size): @@ -300,7 +315,13 @@ def _logic_ring_fa4_vs_standard(rank, world_size): 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) + ring_out = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) scale = 1.0 / math.sqrt(head_dim) ref = ( @@ -346,7 +367,7 @@ def _logic_ring_ulysses_forward(rank, world_size): 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) + output = attn(q, k, v, attn_metadata=make_backend_attn_metadata(attn, q, k), 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" @@ -375,7 +396,13 @@ def _logic_ring_ulysses_vs_standard(rank, world_size): 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) + combo_out = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) scale = 1.0 / math.sqrt(head_dim) ref = ( @@ -425,7 +452,13 @@ def _logic_ring_ulysses_fa4_vs_standard(rank, world_size): 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) + combo_out = attn( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn, q_shard, k_shard), + batch_size=batch, + ) scale = 1.0 / math.sqrt(head_dim) ref = ( diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py index f090cf7522d0..d3a9a5daebe1 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py @@ -50,7 +50,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import make_attn_metadata from tensorrt_llm.mapping import Mapping from tensorrt_llm.visual_gen.args import AttentionConfig @@ -368,7 +370,7 @@ def _run_tp_with_params( head_dim, ) - tp_out = attn_tp(x) + tp_out = attn_tp(x, make_attn_metadata(attn_tp.attn_backend, x)) torch.testing.assert_close(tp_out, ref_out, rtol=1e-2, atol=1e-2) @@ -471,7 +473,7 @@ def _logic_tp_ulysses_combined( ].contiguous() expected_shard = ref_out[:, ulysses_rank * seq_per_rank : (ulysses_rank + 1) * seq_per_rank] - combined_out = attn_combined(x_shard) + combined_out = attn_combined(x_shard, make_attn_metadata(attn_combined.attn_backend, x_shard)) torch.testing.assert_close(combined_out, expected_shard, rtol=1e-2, atol=1e-2) 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 507fda5f3f4c..646abf28b24d 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 @@ -34,7 +34,9 @@ # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). sys.path.insert(0, str(Path(__file__).resolve().parent)) + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from _visual_gen_dist_utils import spawn_with_retry + from attn_metadata_utils import make_backend_attn_metadata MODULES_AVAILABLE = True except ImportError: @@ -578,8 +580,10 @@ def _logic_ulysses_fused_vs_unfused(rank, world_size): process_group=None, ) - out_unfused = attn_unfused(q, k, v) - out_fused = attn_fused(q, k, v) + out_unfused = attn_unfused( + q, k, v, attn_metadata=make_backend_attn_metadata(attn_unfused, q, k) + ) + out_fused = attn_fused(q, k, v, attn_metadata=make_backend_attn_metadata(attn_fused, q, k)) torch.testing.assert_close( out_fused, @@ -615,7 +619,12 @@ def _logic_ulysses_fused_vs_standard(rank, world_size): process_group=None, ) - fused_output = attn_fused(q_shard, k_shard, v_shard) + fused_output = attn_fused( + q_shard, + k_shard, + v_shard, + attn_metadata=make_backend_attn_metadata(attn_fused, q_shard, k_shard), + ) q_std = q_full.transpose(1, 2) k_std = k_full.transpose(1, 2) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py index 54a1f65ffc30..1790cdf3148c 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py @@ -42,7 +42,6 @@ from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention - from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state # Spawn distributed workers via a helper that retries with a fresh master # port when the c10d rendezvous TCPStore loses the bind race (EADDRINUSE). @@ -53,7 +52,6 @@ MODULES_AVAILABLE = True ATTENTION_META_DICT = threading.local() - ATTENTION_META_DICT.metadata = create_attention_metadata_state() except ImportError: MODULES_AVAILABLE = False @@ -149,7 +147,6 @@ def _logic_sage_ulysses_forward(rank, world_size, *, sage_attn_qk_int8: bool): k_block_size=blk_k, v_block_size=1, ), - attention_metadata_state=ATTENTION_META_DICT.metadata, ) attention = UlyssesAttention(inner_backend=inner, process_group=None) @@ -208,7 +205,6 @@ def _logic_sage_ulysses_vs_reference( k_block_size=sage_attn_num_elts_per_blk_k, v_block_size=1, ), - attention_metadata_state=ATTENTION_META_DICT.metadata, ) attention = UlyssesAttention(inner_backend=inner, process_group=None) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py index 7c9ce3d2ccbb..ee7532a566f8 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py @@ -42,7 +42,6 @@ AttentionConfig, DiffusionModelConfig, TorchCompileConfig, - create_attention_metadata_state, ) from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping @@ -51,6 +50,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) from _visual_gen_dist_utils import spawn_with_retry + from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import ( + wan_attn_metadata_kwargs, + ) from tensorrt_llm.models.modeling_utils import QuantConfig from .tp_shard_utils import copy_tp_parameter @@ -185,9 +187,6 @@ def _make_model_config(pretrained_dict, tp_size=1, ulysses_size=1, backend="VANI attention=AttentionConfig(backend=backend), visual_gen_mapping=vgm, cache=None, - attention_metadata_state=( - create_attention_metadata_state() if backend.upper() == "TRTLLM" else None - ), skip_create_weights_in_init=False, ) config.mapping = vgm.to_llm_mapping() @@ -272,6 +271,11 @@ def _logic_wan_t2v_tp_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) @@ -330,11 +334,21 @@ def _logic_wan_t2v_tp_vs_single_gpu_with_config(rank, world_size, config_dict): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + ref_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) tp_output = tp_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + tp_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) @@ -390,11 +404,21 @@ def _logic_wan_t2v_tp_ulysses_vs_single_gpu(rank, world_size): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + ref_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) combined_output = combined_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + combined_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) @@ -445,6 +469,12 @@ def _logic_wan_i2v_tp_forward(rank, world_size): with torch.no_grad(): output = model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image, @@ -509,12 +539,24 @@ def _logic_wan_i2v_tp_vs_single_gpu_with_config(rank, world_size, config_dict): with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + ref_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image, ) tp_output = tp_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + tp_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py index 9da4942eddaf..2615b45bbaa9 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py @@ -51,6 +51,10 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) from _visual_gen_dist_utils import spawn_with_retry + from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import ( + wan_attn_metadata_kwargs, + ) + from .tp_shard_utils import copy_tp_parameter MODULES_AVAILABLE = True @@ -334,11 +338,21 @@ def _logic_wan_transformer_parallel_vs_single_gpu( with torch.no_grad(): ref_output = ref_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + ref_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) dist_output = dist_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + dist_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) @@ -397,6 +411,11 @@ def _logic_wan_transformer_parallel_forward_sanity( with torch.no_grad(): output = model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py index 63684cd0519a..06dc7eaf1213 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py @@ -30,10 +30,7 @@ VSAMetadataBuilder, ) from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen.args import ( AttentionConfig, @@ -91,22 +88,19 @@ def _make_config( attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), skip_create_weights_in_init=False, ) - config.attention_metadata_state = ( - create_attention_metadata_state() if backend == "TRTLLM" else None - ) return config @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") def test_vsa_falls_back_to_vanilla_for_cross_attention(): - """Cross-attention (SEPARATE_QKV) falls back to VANILLA — it has no cube structure.""" + """Cross-attention falls back to VANILLA — it has no cube structure.""" device = torch.device("cuda") dtype = torch.bfloat16 cfg = _make_config( hidden_size=64, num_heads=4, head_dim=16, backend="CUTEDSL", vsa_sparsity=0.5 ) cross_attn = ( - Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg) + Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, is_cross=True, config=cfg) .to(device=device, dtype=dtype) .eval() ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index a372db9641f2..b01d17aa03bb 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -13,6 +13,9 @@ import torch.nn as nn import torch.nn.functional as F +# Import new integrated versions +from attn_metadata_utils import make_attn_metadata + from tensorrt_llm._torch.modules.rms_norm import RMSNorm # ============================================================================ @@ -27,13 +30,8 @@ ) from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention from tensorrt_llm._torch.visual_gen.attention_backend.vanilla import VanillaAttention -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - -# Import new integrated versions from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode, apply_rotary_emb from tensorrt_llm.visual_gen.args import ( AttentionConfig, @@ -158,9 +156,6 @@ def create_model_config( ), skip_create_weights_in_init=skip_create_weights_in_init, ) - config.attention_metadata_state = ( - create_attention_metadata_state() if attn_backend == "TRTLLM" else None - ) if visual_gen_mapping is not None: config.visual_gen_mapping = visual_gen_mapping return config @@ -194,6 +189,7 @@ def _make_cross_attention_with_mapping( skip_create_weights_in_init=True, ) return Attention( + is_cross=True, hidden_size=hidden_size, num_attention_heads=num_heads, head_dim=head_dim, @@ -277,11 +273,11 @@ def generate_rope_embeddings( @pytest.mark.cpu_only -class TestSeparateQkvSequenceParallelGuard: - def test_ring_with_separate_qkv_raises(self): +class TestCrossAttentionSequenceParallelGuard: + def test_ring_with_cross_attention_raises(self): vgm = VisualGenMapping(world_size=2, rank=0, ring_size=2) - with pytest.raises(ValueError, match="SEPARATE_QKV cross-attention does not support"): + with pytest.raises(ValueError, match="Cross-attention does not support Ring"): _make_cross_attention_with_mapping(vgm, enable_sequence_parallel=True) def test_ring_with_sequence_parallel_disabled_allowed(self): @@ -300,11 +296,20 @@ def test_attn2d_with_sequence_parallel_disabled_allowed(self): assert isinstance(attn.attn, VanillaAttention) -def _build_sage_routed_attention(qkv_mode: QKVMode): - """Build a TRTLLM-SAGE-configured Attention for backend-routing checks.""" - quant_cfg = QuantAttentionConfig( - qk_dtype="int8", q_block_size=1, k_block_size=16, v_block_size=1 - ) +_SAGE_DEFAULT = object() + + +def _build_sage_routed_attention( + qkv_mode: QKVMode, + *, + quant_cfg: "QuantAttentionConfig | None" = _SAGE_DEFAULT, + is_cross: bool = False, +): + """Build a TRTLLM-configured Attention for backend-routing checks.""" + if quant_cfg is _SAGE_DEFAULT: + quant_cfg = QuantAttentionConfig( + qk_dtype="int8", q_block_size=1, k_block_size=16, v_block_size=1 + ) config = create_model_config( hidden_size=512, num_heads=4, @@ -319,6 +324,7 @@ def _build_sage_routed_attention(qkv_mode: QKVMode): head_dim=128, qkv_mode=qkv_mode, config=config, + is_cross=is_cross, ) return attn, quant_cfg @@ -331,11 +337,26 @@ def test_self_attention_uses_trtllm_sage_backend(self): assert attn.attn.quant_attention_config == quant_cfg assert not attn.attn.support_fused_qkv() - def test_cross_attention_with_sage_config_falls_back_to_vanilla(self): - attn, _ = _build_sage_routed_attention(QKVMode.SEPARATE_QKV) + def test_cross_attention_uses_trtllm_sage_backend(self): + # Sage's separate-Q/K/V layout is the one TRTLLM path serving a cross site. + attn, quant_cfg = _build_sage_routed_attention(QKVMode.SEPARATE_QKV, is_cross=True) + assert attn.attn_backend == "TRTLLM" + assert isinstance(attn.attn, TrtllmAttention) + assert attn.attn.quant_attention_config == quant_cfg + + def test_cross_attention_without_sage_falls_back_to_vanilla(self): + # Plain TRTLLM expects packed QKV, which unequal Q/KV lengths rule out. + attn, _ = _build_sage_routed_attention(QKVMode.SEPARATE_QKV, quant_cfg=None, is_cross=True) assert attn.attn_backend == "VANILLA" assert isinstance(attn.attn, VanillaAttention) + def test_projection_layout_does_not_affect_backend_choice(self): + # The backend sees the same three tensors either way, so a SEPARATE_QKV + # self-attention site routes exactly like a FUSE_QKV one. + fused, _ = _build_sage_routed_attention(QKVMode.FUSE_QKV) + separate, _ = _build_sage_routed_attention(QKVMode.SEPARATE_QKV) + assert separate.attn_backend == fused.attn_backend == "TRTLLM" + # ============================================================================ # Test functions @@ -404,7 +425,11 @@ def test_self_attention_equivalence( # Forward pass with torch.no_grad(): out_naive = naive(hidden_states, freqs_cos_HSD, freqs_sin_HSD) - out_integrated = integrated(hidden_states, freqs=(freqs_cos_SHD, freqs_sin_SHD)) + out_integrated = integrated( + hidden_states, + make_attn_metadata(integrated.attn_backend, hidden_states), + freqs=(freqs_cos_SHD, freqs_sin_SHD), + ) # Compare (using looser tolerance for bf16) max_diff = (out_naive - out_integrated).abs().max().item() @@ -499,7 +524,11 @@ def test_sage_attention_self_attention(qk_dtype: str, batch_size: int, seq_len: # Forward pass with torch.no_grad(): out_naive = naive(hidden_states, freqs_cos_HSD, freqs_sin_HSD) - out_sage = integrated(hidden_states, freqs=(freqs_cos_SHD, freqs_sin_SHD)) + out_sage = integrated( + hidden_states, + make_attn_metadata(integrated.attn_backend, hidden_states), + freqs=(freqs_cos_SHD, freqs_sin_SHD), + ) # --- Assertions --- @@ -599,7 +628,11 @@ def test_cross_attention_equivalence( # Forward pass with torch.no_grad(): out_naive = naive(hidden_states, encoder_hidden_states) - out_integrated = integrated(hidden_states, encoder_hidden_states) + out_integrated = integrated( + hidden_states, + make_attn_metadata(integrated.attn_backend, hidden_states, encoder_hidden_states), + encoder_hidden_states, + ) # Compare (using looser tolerance for bf16) max_diff = (out_naive - out_integrated).abs().max().item() @@ -687,8 +720,16 @@ def test_fast_cross_attention_wan_shapes( encoder_hidden_states = torch.randn(batch, seq_len_kv, hidden_size, device=device, dtype=dtype) with torch.no_grad(): - out_ref = ref(hidden_states, encoder_hidden_states) - out_fast = fast_model(hidden_states, encoder_hidden_states) + out_ref = ref( + hidden_states, + make_attn_metadata(ref.attn_backend, hidden_states, encoder_hidden_states), + encoder_hidden_states, + ) + out_fast = fast_model( + hidden_states, + make_attn_metadata(fast_model.attn_backend, hidden_states, encoder_hidden_states), + encoder_hidden_states, + ) max_diff = (out_ref - out_fast).abs().max().item() tol = 1e-2 if quant_attention_config is None else 2e-2 @@ -765,7 +806,10 @@ def test_vsa_self_attention_equivalence_at_sparsity_zero(): out_naive = s.naive(s.hidden_states, *s.freqs_HSD) with torch.no_grad(), set_vsa_forward_context(s.metadata): out_vsa = s.integrated( - s.hidden_states, freqs=s.freqs_SHD, gate_compress=s.gate_compress_zero + s.hidden_states, + make_attn_metadata(s.integrated.attn_backend, s.hidden_states), + freqs=s.freqs_SHD, + gate_compress=s.gate_compress_zero, ) assert out_naive.shape == out_vsa.shape, ( @@ -788,7 +832,12 @@ def test_vsa_self_attention_finite(sparsity: float): s = _build_vsa_setup(sparsity=sparsity, batch_size=1, seed=0) with torch.no_grad(), set_vsa_forward_context(s.metadata): - out = s.integrated(s.hidden_states, freqs=s.freqs_SHD, gate_compress=s.gate_compress_zero) + out = s.integrated( + s.hidden_states, + make_attn_metadata(s.integrated.attn_backend, s.hidden_states), + freqs=s.freqs_SHD, + gate_compress=s.gate_compress_zero, + ) assert out.shape == s.hidden_states.shape nan_count = torch.isnan(out).sum().item() @@ -851,7 +900,11 @@ def test_trtllm_cached_prepare(): ) out_naive = naive(hidden_states, freqs_cos_HSD, freqs_sin_HSD) - out_integrated = integrated(hidden_states, freqs=(freqs_cos_SHD, freqs_sin_SHD)) + out_integrated = integrated( + hidden_states, + make_attn_metadata(integrated.attn_backend, hidden_states), + freqs=(freqs_cos_SHD, freqs_sin_SHD), + ) # Check this iteration matches naive max_diff = (out_naive - out_integrated).abs().max().item() @@ -942,7 +995,11 @@ def test_trtllm_varying_seq_len(): ) out_naive = naive(hidden_states, freqs_cos_HSD, freqs_sin_HSD) - out_integrated = integrated(hidden_states, freqs=(freqs_cos_SHD, freqs_sin_SHD)) + out_integrated = integrated( + hidden_states, + make_attn_metadata(integrated.attn_backend, hidden_states), + freqs=(freqs_cos_SHD, freqs_sin_SHD), + ) max_diff = (out_naive - out_integrated).abs().max().item() is_close = torch.allclose(out_naive, out_integrated, rtol=1e-2, atol=1e-2) diff --git a/tests/unittest/_torch/visual_gen/test_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_attention_metadata.py new file mode 100644 index 000000000000..333916f4a19a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_metadata.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for VisualGen attention metadata sites. + +VisualGen reuses the shared-core ``AttentionMetadata`` types rather than +defining its own; these tests pin the behaviour VisualGen depends on (no KV +cache, mixed Q/KV lengths, allocate-once/prepare-in-place) plus the thin +helpers in ``visual_gen/attention_backend/metadata.py``. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import ( + create_diffusion_attn_metadata, + make_diffusion_attn_metadata, + prepare_diffusion_attn_metadata, +) +from tensorrt_llm._torch.visual_gen.models.modeling import _attn_metadata_shape_key + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _site(max_batch_size=2, max_seq_len=4096, cls=TrtllmAttentionMetadata): + return create_diffusion_attn_metadata( + cls, max_batch_size=max_batch_size, max_seq_len=max_seq_len + ) + + +def test_self_attention_site_is_not_cross(): + md = _site() + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=1024) + + # `is_cross` upstream is an identity check, so a self-attention site must + # leave seq_lens_kv unset rather than pass an equal-valued tensor. + assert md.is_cross is False + assert torch.equal(md.seq_lens, torch.tensor([1024], dtype=torch.int32)) + assert torch.equal(md.kv_lens[:1], torch.tensor([1024], dtype=torch.int32)) + + +def test_mixed_site_carries_distinct_kv_length(): + md = _site() + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=1024, kv_seq_lens=1536) + + assert md.is_cross is True + assert torch.equal(md.seq_lens, torch.tensor([1024], dtype=torch.int32)) + assert torch.equal(md.seq_lens_kv, torch.tensor([1536], dtype=torch.int32)) + # prepare() must derive the KV lengths from seq_lens_kv, not seq_lens. + assert torch.equal(md.kv_lens[:1], torch.tensor([1536], dtype=torch.int32)) + + +def test_no_kv_cache_path(): + md = _site() + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=256) + + assert md.kv_cache_manager is None + assert md.kv_cache_params.use_cache is False + + +def test_reprepare_keeps_device_buffers_pointer_stable(): + """CUDA graphs capture pointers into these buffers; they must not move.""" + md = _site() + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=1024) + ptr = md.seq_lens_cuda.data_ptr() + + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=1024) # unchanged + assert md.seq_lens_cuda.data_ptr() == ptr + + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=2048) # new length + assert md.seq_lens_cuda.data_ptr() == ptr + assert torch.equal(md.seq_lens, torch.tensor([2048], dtype=torch.int32)) + + +def test_batch_size_change_is_supported(): + md = _site(max_batch_size=2) + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=512) + prepare_diffusion_attn_metadata(md, batch_size=2, q_seq_lens=512) + + assert md.num_seqs == 2 + assert md.num_contexts == 2 + + +def test_batch_size_beyond_capacity_raises(): + md = _site(max_batch_size=1) + with pytest.raises(ValueError, match="capacity"): + prepare_diffusion_attn_metadata(md, batch_size=4, q_seq_lens=512) + + +def test_seq_lens_batch_mismatch_raises(): + md = _site() + with pytest.raises(ValueError, match="batch"): + prepare_diffusion_attn_metadata(md, batch_size=2, q_seq_lens=[512]) + + +def test_base_metadata_type_works_for_backends_that_ignore_metadata(): + md = _site(max_seq_len=512, cls=AttentionMetadata) + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=512) + assert md.num_tokens == 512 + + +class TestMakeAttnMetadata: + def test_each_call_returns_an_independent_site(self): + self_site = make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=1, q_seq_lens=64 + ) + cross_site = make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=1, q_seq_lens=64, kv_seq_lens=32 + ) + assert self_site is not cross_site + assert self_site.is_cross is False + assert cross_site.is_cross is True + + def test_sizes_capacity_from_the_longer_of_q_and_kv(self): + md = make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=2, q_seq_lens=128, kv_seq_lens=4096 + ) + assert md.max_seq_len >= 4096 + assert md.max_num_requests >= 2 + + def test_accepts_per_sequence_lengths(self): + md = make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=2, q_seq_lens=[64, 128] + ) + assert torch.equal(md.seq_lens, torch.tensor([64, 128], dtype=torch.int32)) + assert md.max_seq_len >= 128 + + def test_is_legal_outside_but_not_inside_a_cuda_graph_capture(self): + """Metadata must be built before entering the captured region. + + ``prepare()`` stages lengths through pinned host memory, so building + inside a capture raises. This pins the constraint that forces + construction out of the CUDA-graph-wrapped model ``forward``. + """ + make_diffusion_attn_metadata(TrtllmAttentionMetadata, batch_size=1, q_seq_lens=128) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + make_diffusion_attn_metadata(TrtllmAttentionMetadata, batch_size=1, q_seq_lens=128) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="CUDA graph capture"): + with torch.cuda.graph(graph): + make_diffusion_attn_metadata(TrtllmAttentionMetadata, batch_size=1, q_seq_lens=128) + + +class TestCudaGraphShapeKey: + """The graph key must see metadata, which reaches models as non-tensor args.""" + + def test_skips_kv_for_self_attention(self): + md = _site() + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=64) + assert _attn_metadata_shape_key(attn_metadata=md) == (("attn_metadata", (64,), None),) + + prepare_diffusion_attn_metadata(md, batch_size=1, q_seq_lens=64, kv_seq_lens=96) + assert _attn_metadata_shape_key(attn_metadata=md) == (("attn_metadata", (64,), (96,)),) + + def test_distinguishes_lengths_and_cross_ness(self): + def key(**kw): + return _attn_metadata_shape_key( + attn_metadata=make_diffusion_attn_metadata(TrtllmAttentionMetadata, **kw) + ) + + base = key(batch_size=1, q_seq_lens=128) + assert key(batch_size=1, q_seq_lens=256) != base + assert key(batch_size=1, q_seq_lens=128, kv_seq_lens=512) != base + assert key(batch_size=1, q_seq_lens=128) == base + + def test_expands_a_dict_of_sites(self): + sites = { + "video_self": make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=1, q_seq_lens=64 + ), + "video_text_cross": make_diffusion_attn_metadata( + TrtllmAttentionMetadata, batch_size=1, q_seq_lens=64, kv_seq_lens=32 + ), + } + assert _attn_metadata_shape_key(attn_metadata=sites) == ( + ("attn_metadata.video_self", (64,), None), + ("attn_metadata.video_text_cross", (64,), (32,)), + ) + + def test_is_none_when_the_forward_carries_no_metadata(self): + assert _attn_metadata_shape_key(hidden_states=torch.zeros(1)) is None diff --git a/tests/unittest/_torch/visual_gen/test_attention_perf.py b/tests/unittest/_torch/visual_gen/test_attention_perf.py index a662b788cf0d..9555c180c6af 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_perf.py +++ b/tests/unittest/_torch/visual_gen/test_attention_perf.py @@ -35,6 +35,7 @@ import pytest import torch +from attn_metadata_utils import make_attn_metadata from tensorrt_llm._torch.visual_gen.attention_backend import ( VSAMetadataBuilder, @@ -49,10 +50,7 @@ from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( _flash_attn_fwd_import_error as _fa4_import_error, ) -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen.args import ( AttentionConfig, @@ -212,9 +210,6 @@ def create_model_config( ), skip_create_weights_in_init=False, ) - config.attention_metadata_state = ( - create_attention_metadata_state() if attn_backend == "TRTLLM" else None - ) return config @@ -361,9 +356,9 @@ def create_cross_attention_model( attn_backend=backend, quant_attention_config=quant_attention_config, ) - model = Attention(hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, config=config).to( - self.device - ) + model = Attention( + hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, is_cross=True, config=config + ).to(self.device) model.eval() return model @@ -427,7 +422,13 @@ def benchmark_cross_attn_single( # Warmup with torch.no_grad(): for _ in range(self.warmup_iterations): - _ = model(hidden_states, encoder_hidden_states=encoder_hidden_states) + _ = model( + hidden_states, + make_attn_metadata( + model.attn_backend, hidden_states, encoder_hidden_states + ), + encoder_hidden_states=encoder_hidden_states, + ) if self.device.type == "cuda": torch.cuda.synchronize() @@ -437,7 +438,13 @@ def benchmark_cross_attn_single( with torch.no_grad(): for i in range(self.benchmark_iterations): with cuda_timer(self.device) as get_time: - _ = model(hidden_states, encoder_hidden_states=encoder_hidden_states) + _ = model( + hidden_states, + make_attn_metadata( + model.attn_backend, hidden_states, encoder_hidden_states + ), + encoder_hidden_states=encoder_hidden_states, + ) times.append(get_time()) times_tensor = torch.tensor(times) @@ -545,7 +552,12 @@ def _forward(): ) kwargs = {"gate_compress": gate_compress} if gate_compress is not None else {} with ctx: - return model(hidden_states, freqs=freqs, **kwargs) + return model( + hidden_states, + make_attn_metadata(model.attn_backend, hidden_states), + freqs=freqs, + **kwargs, + ) # Warmup with nvtx_range(f"warmup_{backend}"): @@ -742,7 +754,11 @@ def test_memory_usage( # Warmup with torch.no_grad(): - _ = model(hidden_states, freqs=freqs) + _ = model( + hidden_states, + make_attn_metadata(model.attn_backend, hidden_states), + freqs=freqs, + ) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() @@ -750,7 +766,11 @@ def test_memory_usage( # Forward pass with nvtx_range(f"memory_test_{backend}"): with torch.no_grad(): - _ = model(hidden_states, freqs=freqs) + _ = model( + hidden_states, + make_attn_metadata(model.attn_backend, hidden_states), + freqs=freqs, + ) torch.cuda.synchronize() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 02280c2f4419..229bda79ffc4 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -856,11 +856,18 @@ def _forward_ready_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: """A pipeline stubbed just enough for forward() to run end to end.""" defaults = dict( sampling=_distilled_policy(), - pipeline_config=SimpleNamespace(torch_dtype=torch.float32, visual_gen_mapping=None), + pipeline_config=SimpleNamespace( + torch_dtype=torch.float32, + visual_gen_mapping=None, + # forward() allocates attention metadata sites from the configured + # backend; see visual_gen/attention_backend/metadata.py. + attention=SimpleNamespace(backend="VANILLA"), + ), transformer=SimpleNamespace( latent_channel_size=4, reset_cache=lambda: None, device=torch.device("cpu"), + create_attn_metadata=lambda **kwargs: {}, ), vae_scale_factor_temporal=4, vae_scale_factor_spatial=16, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 6041300cdf6e..743b4bb32fb6 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -26,6 +26,7 @@ import pytest import torch +from attn_metadata_utils import cosmos3_attn_metadata_kwargs from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig @@ -221,6 +222,7 @@ def test_sanity_forward(self, cosmos3_model_config): with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -239,6 +241,7 @@ def test_reset_cache(self, cosmos3_model_config): with torch.inference_mode(): out1 = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -247,6 +250,7 @@ def test_reset_cache(self, cosmos3_model_config): ) out2 = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -269,6 +273,7 @@ def test_sanity_forward_i2v_mask(self, cosmos3_model_config): with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -344,6 +349,7 @@ def test_forward_with_audio(self, audio_model_config): with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape, audio_latents), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -368,6 +374,7 @@ def test_forward_without_audio_latents_returns_none(self, audio_model_config): with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -389,6 +396,7 @@ def test_forward_with_audio_multiframe(self, audio_model_config): with torch.inference_mode(): out = model( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(model, hs, text_mask, video_shape, audio_latents), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -431,6 +439,7 @@ def test_load_weights_and_forward(self, cosmos3_transformer): with torch.inference_mode(): out = transformer( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(transformer, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, @@ -461,6 +470,7 @@ def test_load_fp8_quantization(self, quant_algo: str): with torch.inference_mode(): out = transformer( hidden_states=hs, + **cosmos3_attn_metadata_kwargs(transformer, hs, text_mask, video_shape), timestep=ts / _NUM_TRAIN_TIMESTEPS, raw_timestep=ts, text_ids=text_ids, diff --git a/tests/unittest/_torch/visual_gen/test_fa4_key_padding_mask.py b/tests/unittest/_torch/visual_gen/test_fa4_key_padding_mask.py index f0923cfb5042..42825195398b 100644 --- a/tests/unittest/_torch/visual_gen/test_fa4_key_padding_mask.py +++ b/tests/unittest/_torch/visual_gen/test_fa4_key_padding_mask.py @@ -15,6 +15,7 @@ import pytest import torch +from attn_metadata_utils import make_backend_attn_metadata try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( @@ -48,10 +49,21 @@ def _run_self_attn(B, S_real, S_pad, H, d_h, dtype=torch.bfloat16): fa4 = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H) # FA4 path with seqused_k (q_full = k_full = v_full = x_full padded; only valid prefix attends). - out_padded = fa4.forward(q=x_full, k=x_full, v=x_full, key_padding_mask=mask) + out_padded = fa4.forward( + q=x_full, + k=x_full, + v=x_full, + attn_metadata=make_backend_attn_metadata(fa4, x_full, x_full), + key_padding_mask=mask, + ) # Reference: FA4 on unpadded inputs only (Q seq = K seq = S_real). - out_ref = fa4.forward(q=x_valid, k=x_valid, v=x_valid) + out_ref = fa4.forward( + q=x_valid, + k=x_valid, + v=x_valid, + attn_metadata=make_backend_attn_metadata(fa4, x_valid, x_valid), + ) # Only the valid Q rows must match; pad Q rows are stripped downstream by # the caller and are not part of the contract. @@ -76,8 +88,16 @@ def _run_cross_attn(B, S_q, S_real_kv, S_pad_kv, H, d_h, dtype=torch.bfloat16): fa4 = FlashAttn4Attention(num_heads=H, head_dim=d_h, num_kv_heads=H) - out_padded = fa4.forward(q=q, k=k_full, v=v_full, key_padding_mask=mask) - out_ref = fa4.forward(q=q, k=k_valid, v=v_valid) + out_padded = fa4.forward( + q=q, + k=k_full, + v=v_full, + attn_metadata=make_backend_attn_metadata(fa4, q, k_full), + key_padding_mask=mask, + ) + out_ref = fa4.forward( + q=q, k=k_valid, v=v_valid, attn_metadata=make_backend_attn_metadata(fa4, q, k_valid) + ) return out_padded, out_ref @@ -124,8 +144,20 @@ def test_self_attn_pad_junk_values_dont_affect_valid_output(): x_a = torch.cat([x_valid, pad_a], dim=1) x_b = torch.cat([x_valid, pad_b], dim=1) - out_a = fa4.forward(q=x_a, k=x_a, v=x_a, key_padding_mask=mask) - out_b = fa4.forward(q=x_b, k=x_b, v=x_b, key_padding_mask=mask) + out_a = fa4.forward( + q=x_a, + k=x_a, + v=x_a, + attn_metadata=make_backend_attn_metadata(fa4, x_a, x_a), + key_padding_mask=mask, + ) + out_b = fa4.forward( + q=x_b, + k=x_b, + v=x_b, + attn_metadata=make_backend_attn_metadata(fa4, x_b, x_b), + key_padding_mask=mask, + ) # Q[:S_real] sees the same K/V[:S_real]; mask zeros K/V[S_real:] contribution. # Only the valid Q rows are part of the contract. diff --git a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py index ac4a5ee84b14..2dde5e70c5c5 100644 --- a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py +++ b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py @@ -13,6 +13,8 @@ import torch from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline @@ -136,7 +138,12 @@ def test_reference_images_run_with_cache_acceleration( reference_count: int, ) -> None: pipeline = Flux2Pipeline.__new__(Flux2Pipeline) - pipeline.pipeline_config = SimpleNamespace(cache_backend=cache_backend) + pipeline.pipeline_config = SimpleNamespace( + cache_backend=cache_backend, + # The pipeline allocates attention metadata sites from the configured + # backend; see visual_gen/attention_backend/metadata.py. + attention=SimpleNamespace(backend="VANILLA"), + ) pipeline._encode_prompt = lambda _prompt, _max_length: ( torch.zeros(1, 2, 8), torch.zeros(2, 4), @@ -160,6 +167,15 @@ class Transformer: def __init__(self) -> None: self.sequence_lengths: list[int] = [] + def create_attn_metadata(self, *, batch_size, text_seq_len, image_seq_len): + return { + "self": make_diffusion_attn_metadata( + AttentionMetadata, + batch_size=batch_size, + q_seq_lens=text_seq_len + image_seq_len, + ) + } + def parameters(self) -> Iterator[torch.Tensor]: return iter([torch.empty(0)]) diff --git a/tests/unittest/_torch/visual_gen/test_flux_attention.py b/tests/unittest/_torch/visual_gen/test_flux_attention.py index e196f83af947..4d22c8a06b2e 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_attention.py +++ b/tests/unittest/_torch/visual_gen/test_flux_attention.py @@ -20,11 +20,9 @@ import pytest import torch import torch.nn.functional as F +from attn_metadata_utils import make_attn_metadata -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import AttentionConfig @@ -99,6 +97,13 @@ def test_vanilla_backend_sanity(self): # Skip RoPE for this sanity test (pass None) output, text_output = attn( hidden_states=hidden_states, + # FLUX joint attention attends over the concatenated + # text + image sequence, so the site covers both. + attn_metadata=make_attn_metadata( + attn.attn_backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, image_rotary_emb=None, ) @@ -122,7 +127,6 @@ def test_trtllm_backend_sanity(self): torch.manual_seed(42) config = self._create_config("TRTLLM") - config.attention_metadata_state = create_attention_metadata_state() attn = ( FluxJointAttention( @@ -149,6 +153,13 @@ def test_trtllm_backend_sanity(self): # Skip RoPE for this sanity test (pass None) output, text_output = attn( hidden_states=hidden_states, + # FLUX joint attention attends over the concatenated + # text + image sequence, so the site covers both. + attn_metadata=make_attn_metadata( + attn.attn_backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, image_rotary_emb=None, ) @@ -195,7 +206,6 @@ def test_backend_equivalence(self): p.normal_(0, 0.02) config = self._create_config("TRTLLM") - config.attention_metadata_state = create_attention_metadata_state() trtllm_attn = ( FluxJointAttention( hidden_size=dim, @@ -227,6 +237,12 @@ def test_backend_equivalence(self): for backend in ["VANILLA", "TRTLLM"]: out, text_out = attns[backend]( hidden_states=hidden_states.clone(), + # Joint attention runs over text + image concatenated. + attn_metadata=make_attn_metadata( + backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states.clone(), image_rotary_emb=None, ) @@ -321,6 +337,13 @@ def test_flux2_vanilla_backend_sanity(self): # Skip RoPE for this sanity test (pass None) output, text_output = attn( hidden_states=hidden_states, + # FLUX joint attention attends over the concatenated + # text + image sequence, so the site covers both. + attn_metadata=make_attn_metadata( + attn.attn_backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, image_rotary_emb=None, ) diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py index 4e33fee0670f..933815a51fd3 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py @@ -24,6 +24,7 @@ import torch.distributed as dist import torch.multiprocessing as mp import torch.nn.functional as F +from attn_metadata_utils import flux_attn_metadata from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader @@ -95,6 +96,7 @@ def _get_flux_transformer_inputs(transformer, device="cuda", dtype=torch.bfloat1 return dict( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(transformer, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -437,7 +439,12 @@ def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_a output_bf16 = transformer_bf16(**inputs) print(f"[E2E] Running {quant_algo} transformer forward...") - inputs_fp8 = {k: v.clone() for k, v in inputs.items()} + # attn_metadata is not a tensor; clone only the tensor inputs. It is also + # rebuilt below so the site matches the FP8 transformer's backend. + inputs_fp8 = {k: (v.clone() if torch.is_tensor(v) else v) for k, v in inputs.items()} + inputs_fp8["attn_metadata"] = flux_attn_metadata( + transformer_fp8, inputs_fp8["hidden_states"], inputs_fp8["encoder_hidden_states"] + ) with torch.no_grad(): output_fp8 = transformer_fp8(**inputs_fp8) @@ -619,6 +626,13 @@ def test_attention_backend_comparison(self, flux1_checkpoint_exists): ) transformer_trtllm = pipeline_trtllm.transformer + # Rebuild the inputs for this transformer: the attention metadata strain + # is typed by the backend (VANILLA gets the base AttentionMetadata, + # TRTLLM gets TrtllmAttentionMetadata), so it cannot be shared across + # backends. `_get_flux_transformer_inputs` re-seeds, so the tensors are + # identical and the comparison below stays apples-to-apples. + inputs = _get_flux_transformer_inputs(transformer_trtllm) + print("[Attention Backend Test] Running TRTLLM transformer forward...") with torch.no_grad(): output_trtllm = transformer_trtllm(**inputs) diff --git a/tests/unittest/_torch/visual_gen/test_flux_transformer.py b/tests/unittest/_torch/visual_gen/test_flux_transformer.py index f77662e76cb5..762d629e05e5 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_flux_transformer.py @@ -17,6 +17,7 @@ import pytest import torch import torch.nn.functional as F +from attn_metadata_utils import flux_attn_metadata from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader @@ -571,6 +572,7 @@ def test_flux1_forward_sanity(self): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, @@ -625,6 +627,7 @@ def test_flux2_forward_sanity(self): with torch.no_grad(): output = model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata(model, hidden_states, encoder_hidden_states), encoder_hidden_states=encoder_hidden_states, timestep=timestep, guidance=guidance, @@ -759,6 +762,9 @@ def test_flux1_allclose_to_hf(self): trtllm_output = trtllm_model( hidden_states=hidden_states, + attn_metadata=flux_attn_metadata( + trtllm_model, hidden_states, encoder_hidden_states + ), encoder_hidden_states=encoder_hidden_states, pooled_projections=pooled_projections, timestep=timestep, diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_attention.py b/tests/unittest/_torch/visual_gen/test_ltx2_attention.py index ae1c95f1e506..a3ca35ff6ee2 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_attention.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_attention.py @@ -15,11 +15,9 @@ import pytest import torch import torch.nn.functional as F +from attn_metadata_utils import make_attn_metadata -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import AttentionConfig @@ -110,7 +108,7 @@ def test_vanilla_self_attention_sanity(self): pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=pe) + output = attn(x, make_attn_metadata(attn.attn_backend, x), context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -128,7 +126,6 @@ def test_trtllm_self_attention_sanity(self): torch.manual_seed(42) config = _create_config("TRTLLM") - config.attention_metadata_state = create_attention_metadata_state() attn = ( LTX2Attention( @@ -147,7 +144,7 @@ def test_trtllm_self_attention_sanity(self): pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=pe) + output = attn(x, make_attn_metadata(attn.attn_backend, x), context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -195,7 +192,12 @@ def test_cross_attention_sanity(self): # this via prepare_text_cache; AV cross-attn does it inline). with torch.no_grad(): k, v = attn.project_kv(ctx, pe=None) - output = attn(x, pre_projected_kv=(k, v), pe=None) + output = attn( + x, + make_attn_metadata(attn.attn_backend, x, kv_seq_len=k.shape[1]), + pre_projected_kv=(k, v), + pe=None, + ) self.assertEqual(output.shape, (batch_size, q_seq, query_dim)) @@ -234,7 +236,12 @@ def test_cross_attention_different_dims(self): with torch.no_grad(): k, v = attn.project_kv(ctx, pe=None) - output = attn(x, pre_projected_kv=(k, v), pe=None) + output = attn( + x, + make_attn_metadata(attn.attn_backend, x, kv_seq_len=k.shape[1]), + pre_projected_kv=(k, v), + pe=None, + ) self.assertEqual(output.shape, (batch_size, q_seq, query_dim)) @@ -279,7 +286,7 @@ def test_gated_self_attention_sanity(self): pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=pe) + output = attn(x, make_attn_metadata(attn.attn_backend, x), context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -321,7 +328,6 @@ def test_backend_equivalence(self): # Create TRTLLM attention and copy weights config_trtllm = _create_config("TRTLLM") - config_trtllm.attention_metadata_state = create_attention_metadata_state() trtllm_attn = ( LTX2Attention( query_dim=query_dim, @@ -340,8 +346,12 @@ def test_backend_equivalence(self): pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - out_vanilla = vanilla_attn(x.clone(), context=None, pe=pe) - out_trtllm = trtllm_attn(x.clone(), context=None, pe=pe) + out_vanilla = vanilla_attn( + x.clone(), make_attn_metadata(vanilla_attn.attn_backend, x), context=None, pe=pe + ) + out_trtllm = trtllm_attn( + x.clone(), make_attn_metadata(trtllm_attn.attn_backend, x), context=None, pe=pe + ) # Skip comparison if either has NaN/Inf (can happen with random weights) has_nan = torch.isnan(out_vanilla).any() or torch.isnan(out_trtllm).any() diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index d0bb3fe43a69..83f70045b01c 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -21,6 +21,7 @@ import pytest import torch import torch.nn.functional as F +from attn_metadata_utils import ltx2_attn_metadata from test_common.llm_data import llm_models_root from tensorrt_llm._torch.modules.linear import Linear @@ -406,7 +407,12 @@ def test_attention_backend_comparison(self, ltx2_bf16_checkpoint_exists): print("[Attention Backend Test] Running VANILLA transformer forward...") with torch.no_grad(): output_baseline = transformer_baseline( - video=video_input, audio=audio_input, text_cache=text_cache_baseline + video=video_input, + audio=audio_input, + text_cache=text_cache_baseline, + attn_metadata=ltx2_attn_metadata( + transformer_baseline, video_input, audio_input, text_cache_baseline + ), ) vout_baseline, aout_baseline = _extract_output(output_baseline) vout_baseline_cpu = vout_baseline.cpu() if vout_baseline is not None else None @@ -430,7 +436,12 @@ def test_attention_backend_comparison(self, ltx2_bf16_checkpoint_exists): _, _, text_cache_trtllm = _get_ltx2_transformer_inputs(transformer_trtllm) with torch.no_grad(): output_trtllm = transformer_trtllm( - video=video_input, audio=audio_input, text_cache=text_cache_trtllm + video=video_input, + audio=audio_input, + text_cache=text_cache_trtllm, + attn_metadata=ltx2_attn_metadata( + transformer_trtllm, video_input, audio_input, text_cache_trtllm + ), ) vout_trtllm, aout_trtllm = _extract_output(output_trtllm) vout_trtllm_cpu = vout_trtllm.cpu() if vout_trtllm is not None else None diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_transformer.py b/tests/unittest/_torch/visual_gen/test_ltx2_transformer.py index 226623d4207b..546efcfed10e 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_transformer.py @@ -11,6 +11,7 @@ import pytest import torch +from attn_metadata_utils import ltx2_attn_metadata from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm.mapping import Mapping @@ -47,6 +48,11 @@ ) +def _sites(model, video=None, audio=None, text_cache=None): + """LTX-2's per-site attention metadata for a direct LTXModel.forward call.""" + return ltx2_attn_metadata(model, video, audio, text_cache) + + def _create_model_config(backend: str = "VANILLA") -> DiffusionModelConfig: """Create a minimal DiffusionModelConfig for unit tests.""" from types import SimpleNamespace @@ -235,7 +241,12 @@ def test_video_only_forward_sanity(self): ) with torch.no_grad(): - video_out, audio_out = model(video=video_modality, audio=None, text_cache=text_cache) + video_out, audio_out = model( + video=video_modality, + audio=None, + text_cache=text_cache, + attn_metadata=_sites(model, video_modality, None, text_cache), + ) self.assertIsNotNone(video_out) self.assertIsNone(audio_out) @@ -354,7 +365,10 @@ def test_audio_video_forward_sanity(self): with torch.no_grad(): video_out, audio_out = model( - video=video_modality, audio=audio_modality, text_cache=text_cache + video=video_modality, + audio=audio_modality, + text_cache=text_cache, + attn_metadata=_sites(model, video_modality, audio_modality, text_cache), ) self.assertIsNotNone(video_out) @@ -419,7 +433,12 @@ def test_video_only_input_to_audio_video_model(self): ) with torch.no_grad(): - video_out, audio_out = model(video=video_modality, audio=None, text_cache=text_cache) + video_out, audio_out = model( + video=video_modality, + audio=None, + text_cache=text_cache, + attn_metadata=_sites(model, video_modality, None, text_cache), + ) self.assertIsNotNone(video_out) self.assertIsNone(audio_out) @@ -559,7 +578,10 @@ def _record(mod, args): try: with torch.no_grad(): video_out, audio_out = model( - video=video_modality, audio=audio_modality, text_cache=text_cache + video=video_modality, + audio=audio_modality, + text_cache=text_cache, + attn_metadata=_sites(model, video_modality, audio_modality, text_cache), ) finally: for h in handles: @@ -828,7 +850,12 @@ def _make_modalities(timestep_val=0.5): dtype=dtype, ) with torch.no_grad(): - eager_v, eager_a = model(video=video_mod, audio=audio_mod, text_cache=text_cache) + eager_v, eager_a = model( + video=video_mod, + audio=audio_mod, + text_cache=text_cache, + attn_metadata=_sites(model, video_mod, audio_mod, text_cache), + ) eager_v = eager_v.clone() eager_a = eager_a.clone() @@ -841,7 +868,12 @@ def _make_modalities(timestep_val=0.5): torch.manual_seed(100) video_mod, audio_mod = _make_modalities(0.5) with torch.no_grad(): - graph_v1, graph_a1 = model(video=video_mod, audio=audio_mod, text_cache=text_cache) + graph_v1, graph_a1 = model( + video=video_mod, + audio=audio_mod, + text_cache=text_cache, + attn_metadata=_sites(model, video_mod, audio_mod, text_cache), + ) self.assertTrue( torch.equal(eager_v, graph_v1), @@ -859,7 +891,12 @@ def _make_modalities(timestep_val=0.5): # Eager baseline for new inputs (same static — context/PE don't change) model.forward = original_forward with torch.no_grad(): - eager_v2, eager_a2 = model(video=video_mod2, audio=audio_mod2, text_cache=text_cache) + eager_v2, eager_a2 = model( + video=video_mod2, + audio=audio_mod2, + text_cache=text_cache, + attn_metadata=_sites(model, video_mod2, audio_mod2, text_cache), + ) eager_v2 = eager_v2.clone() eager_a2 = eager_a2.clone() @@ -868,7 +905,12 @@ def _make_modalities(timestep_val=0.5): torch.manual_seed(200) video_mod2, audio_mod2 = _make_modalities(0.3) with torch.no_grad(): - graph_v2, graph_a2 = model(video=video_mod2, audio=audio_mod2, text_cache=text_cache) + graph_v2, graph_a2 = model( + video=video_mod2, + audio=audio_mod2, + text_cache=text_cache, + attn_metadata=_sites(model, video_mod2, audio_mod2, text_cache), + ) self.assertTrue( torch.equal(eager_v2, graph_v2), @@ -953,11 +995,21 @@ def make_mods(ts, v_ctx, a_ctx): # -- Part 1: deterministic — same inputs + same cache → same output -- torch.manual_seed(200) v_mod, a_mod = make_mods(0.5, v_ctx_A, a_ctx_A) - v_out1, a_out1 = model(video=v_mod, audio=a_mod, text_cache=static_A) + v_out1, a_out1 = model( + video=v_mod, + audio=a_mod, + text_cache=static_A, + attn_metadata=_sites(model, v_mod, a_mod, static_A), + ) torch.manual_seed(200) v_mod, a_mod = make_mods(0.5, v_ctx_A, a_ctx_A) - v_out2, a_out2 = model(video=v_mod, audio=a_mod, text_cache=static_A) + v_out2, a_out2 = model( + video=v_mod, + audio=a_mod, + text_cache=static_A, + attn_metadata=_sites(model, v_mod, a_mod, static_A), + ) self.assertTrue( torch.equal(v_out1, v_out2), @@ -979,7 +1031,12 @@ def make_mods(ts, v_ctx, a_ctx): audio_positions=a_pos, dtype=dtype, ) - v_B, _ = model(video=v_mod_B, audio=a_mod_B, text_cache=static_B) + v_B, _ = model( + video=v_mod_B, + audio=a_mod_B, + text_cache=static_B, + attn_metadata=_sites(model, v_mod_B, a_mod_B, static_B), + ) self.assertFalse(torch.equal(v_out1, v_B), "Different context must differ") @@ -987,11 +1044,21 @@ def make_mods(ts, v_ctx, a_ctx): # -- Part 3: reuse static across steps -- torch.manual_seed(300) v_mod1, a_mod1 = make_mods(0.8, v_ctx_A, a_ctx_A) - v_step1, _ = model(video=v_mod1, audio=a_mod1, text_cache=static_A) + v_step1, _ = model( + video=v_mod1, + audio=a_mod1, + text_cache=static_A, + attn_metadata=_sites(model, v_mod1, a_mod1, static_A), + ) torch.manual_seed(400) v_mod2, a_mod2 = make_mods(0.5, v_ctx_A, a_ctx_A) - v_step2, _ = model(video=v_mod2, audio=a_mod2, text_cache=static_A) + v_step2, _ = model( + video=v_mod2, + audio=a_mod2, + text_cache=static_A, + attn_metadata=_sites(model, v_mod2, a_mod2, static_A), + ) # Different timestep/latent → different output self.assertFalse(torch.equal(v_step1, v_step2), "Steps should differ") @@ -1069,7 +1136,12 @@ def make_mods(): # Baseline: direct path (no wrappers) with torch.no_grad(): v_mod, a_mod = make_mods() - v_direct, a_direct = model(video=v_mod, audio=a_mod, text_cache=static) + v_direct, a_direct = model( + video=v_mod, + audio=a_mod, + text_cache=static, + attn_metadata=_sites(model, v_mod, a_mod, static), + ) v_direct = v_direct.clone() a_direct = a_direct.clone() @@ -1083,7 +1155,12 @@ def make_mods(): try: with torch.no_grad(): v_mod, a_mod = make_mods() - v_wrap, a_wrap = model(video=v_mod, audio=a_mod, text_cache=static) + v_wrap, a_wrap = model( + video=v_mod, + audio=a_mod, + text_cache=static, + attn_metadata=_sites(model, v_mod, a_mod, static), + ) finally: model.transformer_blocks = nn.ModuleList(original_blocks) del model._uses_cache_dit diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py b/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py index 79e61bb65067..a83b50f28331 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py @@ -20,6 +20,7 @@ import pytest import torch +from attn_metadata_utils import make_attn_metadata # Importing the models package side-effects the ``@register_pipeline`` # decorator on ``QwenImageLayeredPipeline`` being applied. @@ -275,6 +276,11 @@ def test_layered_transformer_forward_sanity(with_text_mask): with torch.inference_mode(): output = model( hidden_states=hidden_states, + attn_metadata=make_attn_metadata( + model.model_config.attention.backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, timestep=torch.tensor([0.5], device=device, dtype=dtype), diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py index 4900421d70bd..b71cb3d6ce42 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py @@ -15,6 +15,8 @@ import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.visual_gen.attention_backend.metadata import make_diffusion_attn_metadata from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImagePipeline from tensorrt_llm._torch.visual_gen.profiler import VisualGenProfiler @@ -26,10 +28,20 @@ def __init__(self): self._device_anchor = torch.nn.Parameter(torch.zeros(())) self.calls = [] + def create_attn_metadata(self, *, batch_size, text_seq_len, image_seq_len): + return { + "self": make_diffusion_attn_metadata( + AttentionMetadata, + batch_size=batch_size, + q_seq_lens=text_seq_len + image_seq_len, + ) + } + def forward( self, *, hidden_states, + attn_metadata, timestep, encoder_hidden_states_mask, encoder_hidden_states, @@ -86,7 +98,11 @@ def _pipeline_with_test_doubles(): pipe.transformer = _RecordingTransformer() pipe.scheduler = _RecordingScheduler() pipe.pipeline_config = SimpleNamespace( - cuda_graph=SimpleNamespace(enable=False), visual_gen_mapping=None + cuda_graph=SimpleNamespace(enable=False), + visual_gen_mapping=None, + # The pipeline allocates attention metadata sites from the configured + # backend; see visual_gen/attention_backend/metadata.py. + attention=SimpleNamespace(backend="VANILLA"), ) pipe._is_warmup = False pipe._profiler = VisualGenProfiler() diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py index 71ca19766751..7191602d08e8 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py @@ -11,6 +11,7 @@ import pytest import torch import torch.nn.functional as F +from attn_metadata_utils import make_attn_metadata from tensorrt_llm._torch.modules.linear import TensorParallelMode from tensorrt_llm._torch.utils import gelu_tanh @@ -345,8 +346,9 @@ def test_qwen_joint_attention_passes_padding_mask_to_backend(monkeypatch): ) captured = {} - def fake_attn_impl(q, k, v, **kwargs): + def fake_attn_impl(q, k, v, attn_metadata, **kwargs): captured["key_padding_mask"] = kwargs.get("key_padding_mask") + captured["attn_metadata"] = attn_metadata return q.new_zeros(q.shape) monkeypatch.setattr(attention, "_attn_impl", fake_attn_impl) @@ -357,11 +359,18 @@ def fake_attn_impl(q, k, v, **kwargs): attention( hidden_states=hidden_states, + attn_metadata=make_attn_metadata( + attention.attn_backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, ) assert captured["key_padding_mask"] is attention_mask + # The site reaches the backend alongside the mask. + assert captured["attn_metadata"] is not None def test_qwen_joint_attention_rejects_unsupported_masked_sequence_parallel(monkeypatch): @@ -387,6 +396,9 @@ def test_qwen_joint_attention_rejects_unsupported_masked_sequence_parallel(monke with pytest.raises(NotImplementedError, match="Padded Qwen-Image prompts"): attention( hidden_states=torch.empty(1, 4, 16), + attn_metadata=make_attn_metadata( + attention.attn_backend, torch.empty(1, 4, 16), q_seq_len=7 + ), encoder_hidden_states=torch.empty(1, 3, 16), attention_mask=torch.tensor([[True, False, True, True, True, True, True]]), ) diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_registry.py b/tests/unittest/_torch/visual_gen/test_qwen_image_registry.py index 7e23ecb93aa3..a184c40be0d5 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_registry.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_registry.py @@ -14,6 +14,7 @@ import pytest import torch +from attn_metadata_utils import make_attn_metadata from tensorrt_llm._torch.modules.linear import ( FP8QDQLinearMethod, @@ -412,6 +413,11 @@ def test_transformer_forward_sanity(with_text_mask): with torch.inference_mode(): output = model( hidden_states=hidden_states, + attn_metadata=make_attn_metadata( + model.model_config.attention.backend, + hidden_states, + q_seq_len=hidden_states.shape[1] + encoder_hidden_states.shape[1], + ), encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, timestep=torch.tensor([0.5], device=device, dtype=dtype), diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py deleted file mode 100644 index fb0b16ba40eb..000000000000 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import torch - -from tensorrt_llm._torch.visual_gen.attention_backend import trtllm as visual_trtllm - - -class _FakeBaseTrtllmAttentionMetadata: - def __init__(self, **kwargs): - self.kwargs = kwargs - self.prepare_calls = 0 - self.seq_lens = None - self.num_contexts = None - self.max_seq_len = None - self.request_ids = None - - def prepare(self): - self.prepare_calls += 1 - - -def test_trtllm_attention_metadata_caches_distinct_seq_lens(monkeypatch): - monkeypatch.setattr( - visual_trtllm, - "BaseTrtllmAttentionMetadata", - _FakeBaseTrtllmAttentionMetadata, - ) - attention_metadata_state = {} - metadata = visual_trtllm.TrtllmAttentionMetadata( - device=torch.device("cpu"), - attention_metadata_state=attention_metadata_state, - ) - - first_seq_lens = torch.tensor([64], dtype=torch.int32) - first_metadata = metadata.prepare(batch_size=1, seq_lens=first_seq_lens) - first_seq_lens.fill_(999) - - second_metadata = metadata.prepare(batch_size=1, seq_lens=torch.tensor([96], dtype=torch.int32)) - first_metadata_again = metadata.prepare( - batch_size=1, - seq_lens=torch.tensor([64], dtype=torch.int32), - ) - - assert first_metadata is first_metadata_again - assert first_metadata is not second_metadata - assert first_metadata.prepare_calls == 1 - assert second_metadata.prepare_calls == 1 - - metadata_cache = attention_metadata_state["metadata_cache"] - assert set(metadata_cache) == { - (1, (64,)), - (1, (96,)), - } - assert metadata_cache[(1, (64,))]["metadata"] is first_metadata - assert metadata_cache[(1, (96,))]["metadata"] is second_metadata - - first_cached_seq_lens = metadata_cache[(1, (64,))]["seq_lens"] - second_cached_seq_lens = metadata_cache[(1, (96,))]["seq_lens"] - assert torch.equal(first_cached_seq_lens, torch.tensor([64], dtype=torch.int32)) - assert torch.equal(second_cached_seq_lens, torch.tensor([96], dtype=torch.int32)) - assert first_cached_seq_lens is not second_cached_seq_lens - assert first_cached_seq_lens.data_ptr() != second_cached_seq_lens.data_ptr() - assert first_metadata.seq_lens is first_cached_seq_lens - assert second_metadata.seq_lens is second_cached_seq_lens diff --git a/tests/unittest/_torch/visual_gen/test_vanilla_key_padding_mask.py b/tests/unittest/_torch/visual_gen/test_vanilla_key_padding_mask.py index 2cdf9919069b..b038caa9c298 100644 --- a/tests/unittest/_torch/visual_gen/test_vanilla_key_padding_mask.py +++ b/tests/unittest/_torch/visual_gen/test_vanilla_key_padding_mask.py @@ -10,6 +10,7 @@ import pytest import torch +from attn_metadata_utils import make_backend_attn_metadata from tensorrt_llm._torch.visual_gen.attention_backend import VanillaAttention @@ -27,7 +28,9 @@ def test_padded_kv_with_mask_matches_unpadded(): k_valid = torch.randn(B, H_kv, S_kv_valid, d_h) v_valid = torch.randn(B, H_kv, S_kv_valid, d_h) - ref = attn.forward(q=q, k=k_valid, v=v_valid) + ref = attn.forward( + q=q, k=k_valid, v=v_valid, attn_metadata=make_backend_attn_metadata(attn, q, k_valid) + ) # Pad K/V with realistic magnitudes (O(1)). Real audio pad goes through # K/V Linear from a zero-latent, so it is bounded by the layer's bias norm. @@ -36,7 +39,13 @@ def test_padded_kv_with_mask_matches_unpadded(): mask = torch.zeros(B, S_kv, dtype=torch.bool) mask[:, :S_kv_valid] = True - out = attn.forward(q=q, k=k_padded, v=v_padded, key_padding_mask=mask) + out = attn.forward( + q=q, + k=k_padded, + v=v_padded, + attn_metadata=make_backend_attn_metadata(attn, q, k_padded), + key_padding_mask=mask, + ) torch.testing.assert_close( out, diff --git a/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py index 08d105a6565f..1f7aa6a83f98 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py @@ -33,6 +33,7 @@ from diffusers import DiffusionPipeline from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import wan_attn_metadata_kwargs from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import ( AttentionConfig, @@ -600,10 +601,24 @@ def test_fp8_e2e_accuracy( with torch.no_grad(): out_bf16 = wan21_t2v_bf16.transformer( - hidden_states=hs.clone(), timestep=ts, encoder_hidden_states=enc.clone() + hidden_states=hs.clone(), + **wan_attn_metadata_kwargs( + wan21_t2v_bf16.transformer, + hidden_states=hs, + encoder_hidden_states=enc, + ), + timestep=ts, + encoder_hidden_states=enc.clone(), ).float() out_quant = quant_pipeline.transformer( - hidden_states=hs.clone(), timestep=ts, encoder_hidden_states=enc.clone() + hidden_states=hs.clone(), + **wan_attn_metadata_kwargs( + quant_pipeline.transformer, + hidden_states=hs, + encoder_hidden_states=enc, + ), + timestep=ts, + encoder_hidden_states=enc.clone(), ).float() assert not torch.isnan(out_bf16).any(), "BF16 output contains NaN" @@ -634,10 +649,24 @@ def test_nvfp4_e2e_accuracy(self, wan21_t2v_bf16, wan21_t2v_nvfp4): with torch.no_grad(): out_bf16 = wan21_t2v_bf16.transformer( - hidden_states=hs.clone(), timestep=ts, encoder_hidden_states=enc.clone() + hidden_states=hs.clone(), + **wan_attn_metadata_kwargs( + wan21_t2v_bf16.transformer, + hidden_states=hs, + encoder_hidden_states=enc, + ), + timestep=ts, + encoder_hidden_states=enc.clone(), ).float() out_nvfp4 = wan21_t2v_nvfp4.transformer( - hidden_states=hs.clone(), timestep=ts, encoder_hidden_states=enc.clone() + hidden_states=hs.clone(), + **wan_attn_metadata_kwargs( + wan21_t2v_nvfp4.transformer, + hidden_states=hs, + encoder_hidden_states=enc, + ), + timestep=ts, + encoder_hidden_states=enc.clone(), ).float() assert not torch.isnan(out_nvfp4).any(), "NVFP4 output contains NaN" diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py index b3132939f3d9..bf7b728e5ae5 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py @@ -42,6 +42,7 @@ DiffusionPipelineConfig, VisualGenArgs, ) +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import wan_attn_metadata_kwargs from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel from tensorrt_llm.models.modeling_utils import QuantConfig @@ -284,6 +285,11 @@ def test_sanity_forward(self): with torch.inference_mode(): out = model( hidden_states=hs, + **wan_attn_metadata_kwargs( + model, + hidden_states=hs, + encoder_hidden_states=enc, + ), timestep=_normalize_wan_timestep(ts), encoder_hidden_states=enc, ) @@ -335,6 +341,11 @@ def test_allclose_to_hf(self): )[0].float() trt_out = trtllm( hidden_states=hs, + **wan_attn_metadata_kwargs( + trtllm, + hidden_states=hs, + encoder_hidden_states=enc, + ), timestep=_normalize_wan_timestep(ts), encoder_hidden_states=enc, ).float() @@ -393,6 +404,11 @@ def test_cosine_similarity(self, t2v_models): our_out = our_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + our_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + ), timestep=_normalize_wan_timestep(timestep), encoder_hidden_states=encoder_hidden_states, ) @@ -469,6 +485,12 @@ def test_cosine_similarity(self, i2v_models): our_out = our_model( hidden_states=hidden_states, + **wan_attn_metadata_kwargs( + our_model, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=image_embeds, + ), timestep=_normalize_wan_timestep(timestep), encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_image=image_embeds, From 26e902f29d88fee856a95984dd9294f970d6ea13 Mon Sep 17 00:00:00 2001 From: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:51:36 -0700 Subject: [PATCH 2/3] Fix CI - Make metadata optional: some pipeline can detect attention modules' requires_metadata field and skip metadata creation - Fix CI tests on thop API change, Cosmos3-distilled, and Cosmos3-edge. - Adjust code style according to auto-review. Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- .../visual_gen/attention_backend/trtllm.py | 9 ++-- .../models/cosmos3/pipeline_cosmos3.py | 13 ++++-- .../models/flux/transformer_flux2.py | 3 ++ .../_torch/visual_gen/models/modeling.py | 46 +++++++++++++------ .../visual_gen/models/wan/pipeline_fastwan.py | 7 +++ .../visual_gen/models/wan/transformer_wan.py | 13 +++--- .../_torch/visual_gen/modules/attention.py | 5 ++ .../attention/test_attention_op_sync.py | 2 + .../visual_gen/test_cosmos3_distilled.py | 6 ++- .../_torch/visual_gen/test_cosmos3_edge.py | 12 ++++- 10 files changed, 84 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 45253faa086c..2bdabdfb8c3d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -42,9 +42,8 @@ def _check_metadata( """Validate that the metadata describes the tensors it is used with.""" if attn_metadata is None: raise ValueError( - "TrtllmAttention.forward requires `attn_metadata`. Build it with " - "visual_gen.attention_backend.metadata.create_diffusion_attn_metadata() " - "and prepare it with prepare_diffusion_attn_metadata()." + "TrtllmAttention.forward requires `attn_metadata`, but this site got " + "None. Build it for this site in the model's create_attn_metadata()." ) seq_lens = attn_metadata.seq_lens @@ -166,8 +165,8 @@ def forward( Returns: Output tensor [B, S, H*D] """ - batch_size, seq_len, _, _ = q.shape - _, kv_seq_len, _, _ = k.shape + batch_size, seq_len = q.shape[0], q.shape[1] + kv_seq_len = seq_len if k is None else k.shape[1] _check_metadata(attn_metadata, batch_size, seq_len, kv_seq_len) timestep = kwargs.pop("timestep", None) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 08a8a39bde2b..d7bddd01fb0c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1519,11 +1519,14 @@ def post_step_fn(step_latents): # Create attention metadata for understanding and generation towers denoise_batch = self.denoise_batch_size(latents, guidance_scale=guidance_scale) - denoise_mask = ( - torch.cat([uncond_mask, cond_mask], dim=0) - if denoise_batch != latents.shape[0] - else cond_mask - ) + if denoise_batch != latents.shape[0]: + # Sequential CFG stacks both branches into one batch. + denoise_mask = torch.cat([uncond_mask, cond_mask], dim=0) + else: + # One CFG at a time. + vgm = self.pipeline_config.visual_gen_mapping + is_conditional = vgm.is_cfg_conditional if vgm is not None else True + denoise_mask = cond_mask if is_conditional else uncond_mask attn_metadata_sites = self.transformer.create_attn_metadata( batch_size=denoise_batch, text_seq_len=denoise_mask.shape[1], 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 f6eb8f73aefc..ce53c27a5b7c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py @@ -704,6 +704,9 @@ def create_attn_metadata( FLUX.2 has a single attention site over the concatenated text+image sequence. """ + if not self.attn_requires_metadata: + return {"self": None} + # forward() shards each stream before concatenating them, so a block # attends over this rank's share of the joint sequence. size = self.sharder.size diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index a9c66eee5445..ee41ff43093b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -14,7 +14,7 @@ # limitations under the License. """Base classes for VisualGen model components.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Optional import torch import torch.nn as nn @@ -29,22 +29,29 @@ from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner -def _attn_metadata_shape_key(*args, **kwargs): +def _collect_attn_metadata(name: str, value: Any, sites: dict[str, AttentionMetadata]) -> None: + """Flatten a metadata argument into ``sites``, keyed by its position.""" + if isinstance(value, AttentionMetadata): + sites[name] = value + elif isinstance(value, dict): + for key, item in value.items(): + _collect_attn_metadata(f"{name}.{key}", item, sites) + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + _collect_attn_metadata(f"{name}[{index}]", item, sites) + + +def _attn_metadata_shape_key(*args: Any, **kwargs: Any) -> Optional[tuple]: """CUDA graph key from every ``attn_metadata*`` keyword, or ``None`` if absent. - Accepts a single site or a dict of them; the tensor-shape key cannot see - metadata, so sites of differing length must not share a graph. + Accepts a single site, or a dict or list of them (Cosmos3 passes a per-sample + list); the tensor-shape key cannot see metadata, so sites of differing length + must not share a graph. """ sites: dict[str, AttentionMetadata] = {} for name, value in kwargs.items(): - if not name.startswith("attn_metadata"): - continue - if isinstance(value, AttentionMetadata): - sites[name] = value - elif isinstance(value, dict): - sites.update( - {f"{name}.{k}": v for k, v in value.items() if isinstance(v, AttentionMetadata)} - ) + if name.startswith("attn_metadata"): + _collect_attn_metadata(name, value, sites) key = [] for name in sorted(sites): @@ -76,7 +83,20 @@ def attn_backend_metadata_cls(self) -> type[AttentionMetadata]: """Metadata type this component's attention backend expects.""" return get_visual_gen_attention_backend(self.model_config.attention.backend).Metadata - def forward(self, *args, timestep: torch.Tensor | None = None, **kwargs): + @property + def attn_requires_metadata(self) -> bool: + """Whether any attention site in this model needs prepared metadata.""" + from tensorrt_llm._torch.visual_gen.modules.attention import Attention + + return any(m.requires_metadata for m in self.modules() if isinstance(m, Attention)) + + def create_attn_metadata(self, **kwargs) -> dict[str, AttentionMetadata]: + """Build one prepared metadata object per attention site.""" + raise NotImplementedError( + "Diffusion model subclasses must implement create_attn_metadata()." + ) + + def forward(self, *args: Any, timestep: torch.Tensor | None = None, **kwargs: Any) -> Any: """Run the diffusion transformer. Concrete VisualGen models own their full forward signatures. This base diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py index f3f31167ce04..aed6ddbf6b7c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py @@ -32,6 +32,7 @@ from .defaults import get_fastwan_default_params from .pipeline_wan import WanPipeline +from .pipeline_wan_utils import wan_attn_metadata_kwargs @register_pipeline( @@ -164,12 +165,18 @@ def _denoise( nf = latents.shape[2] nh = latents.shape[3] // ph nw = latents.shape[4] // pw + # Attention metadata: DMD runs one guidance-free pass per step at a fixed shape. + attn_metadata_sites = wan_attn_metadata_kwargs( + self.transformer, hidden_states=latents, encoder_hidden_states=prompt_embeds + ) + start = time.time() for i, t in self._profile_denoise_steps(timesteps): t_tensor = torch.full((latents.shape[0], nf * nh * nw), float(t), device=latents.device) pred_noise = self.transformer( hidden_states=latents, + **attn_metadata_sites, timestep=t_tensor / self.NUM_TRAIN_TIMESTEPS, encoder_hidden_states=prompt_embeds, ) 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 d40446d69367..4601cdbfc86d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -628,12 +628,6 @@ def forward( key_img = self.add_k_proj(encoder_hidden_states_img) value_img = self.add_v_proj(encoder_hidden_states_img) key_img = self.norm_added_k(key_img) - if attn_metadata_cross_image is None: - raise ValueError( - "WAN I2V image cross-attention requires attn_metadata_cross_image; " - "the pipeline must include the 'cross_image' site when " - "encoder_hidden_states_image is provided." - ) attn_img_output = self.attn2._attn_impl( q, key_img, @@ -866,6 +860,13 @@ def create_attn_metadata( # [image; text] and the text tail is a fixed length. s_text = self.TEXT_CONTEXT_LENGTH if has_image else encoder_hidden_states.shape[1] + if not self.attn_requires_metadata: + return { + "self": None, + "cross_text": None, + **({"cross_image": None} if has_image else {}), + } + q_self, kv_self = get_ulysses_seq_lens(s_video, s_video, visual_gen_mapping=vgm) sites = { "self": make_diffusion_attn_metadata( diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index af1f32d117db..aac2980ecc6c 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -274,6 +274,11 @@ def __init__( async_ulysses=use_ulysses and async_ulysses, ) + @property + def requires_metadata(self) -> bool: + """Whether the selected backend needs prepared metadata to run.""" + return self.attn.requires_metadata + @staticmethod def _qualified_module_name( component_name: Optional[str], diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index 66215181638c..d60102f04886 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -70,6 +70,8 @@ "host_context_lengths": ("metadata", ("prompt_lens_cpu_runtime",)), "host_past_key_value_lengths": ("metadata", ("kv_lens_runtime",)), "host_request_types": ("metadata", ("host_request_types_runtime",)), + # thop's `is_cross` selects the cache-backed KV layout, narrower than `metadata.is_cross` + "is_cross": ("metadata", ("is_cross_with_kv_cache",)), "sequence_length": ("metadata", ("kv_lens_cuda_runtime",)), "spec_decoding_target_max_draft_tokens": ( "metadata", diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index 229bda79ffc4..4995435d26eb 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -867,7 +867,11 @@ def _forward_ready_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: latent_channel_size=4, reset_cache=lambda: None, device=torch.device("cpu"), - create_attn_metadata=lambda **kwargs: {}, + create_attn_metadata=lambda **kwargs: { + "und": None, + "mixed": None, + "mixed_ragged": None, + }, ), vae_scale_factor_temporal=4, vae_scale_factor_spatial=16, diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index ffe3777774d4..e117de45d8c5 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -22,6 +22,7 @@ import numpy as np import pytest import torch +from attn_metadata_utils import cosmos3_attn_metadata_kwargs, make_attn_metadata from diffusers import UniPCMultistepScheduler from tensorrt_llm._torch.modules.mlp import MLP @@ -406,8 +407,12 @@ def test_k_norm_touches_only_gen_facing_keys(self): cos = torch.ones(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) sin = torch.zeros(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) + und_metadata = make_attn_metadata("VANILLA", hidden) + with torch.inference_mode(): - out_before, k_gen_before, v_before = attn.forward_with_kv(hidden, cos, sin) + out_before, k_gen_before, v_before = attn.forward_with_kv( + hidden, cos, sin, und_metadata + ) # The rope above is identity (cos=1, sin=0), so the cached gen K # must be exactly the Nemotron-normed raw K. _, k_raw, _ = attn.get_qkv(hidden) @@ -415,7 +420,7 @@ def test_k_norm_touches_only_gen_facing_keys(self): assert torch.equal(k_gen_before, attn.k_norm_und_for_gen(k_raw)) attn.k_norm_und_for_gen.weight.fill_(2.0) - out_after, k_gen_after, v_after = attn.forward_with_kv(hidden, cos, sin) + out_after, k_gen_after, v_after = attn.forward_with_kv(hidden, cos, sin, und_metadata) assert torch.equal(out_before, out_after) assert torch.equal(v_before, v_after) @@ -1385,6 +1390,9 @@ def test_load_weights_and_forward(self, edge_pipeline): with torch.inference_mode(): out = transformer( hidden_states=latents, + **cosmos3_attn_metadata_kwargs( + transformer, latents, text_mask, video_shape=(1, 16, 16) + ), timestep=timestep / 1000.0, raw_timestep=timestep, text_ids=text_ids, From 0f4bd1a12821e649456498e2b804909a66a9b895 Mon Sep 17 00:00:00 2001 From: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:09:57 -0700 Subject: [PATCH 3/3] Fix DMD test Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py b/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py index e60b7030cc2a..17af00a446f6 100644 --- a/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py +++ b/tests/unittest/_torch/visual_gen/test_fastwan_dmd_math.py @@ -86,7 +86,11 @@ class _Config: def parameters(self): yield self._dummy_param - def __call__(self, hidden_states, timestep, encoder_hidden_states): + def create_attn_metadata(self, **kwargs): + # Site names WanTransformer3DModel.forward expects; this fake needs no metadata. + return {"self": None, "cross_text": None} + + def __call__(self, hidden_states, timestep, encoder_hidden_states, **kwargs): self.captured_timesteps.append(timestep.detach().clone()) return torch.zeros_like(hidden_states)