diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 48147d5d3234..833ca696496b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -370,8 +370,7 @@ /examples/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /scripts/visualgen_eval @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/_torch/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs -/tensorrt_llm/_torch/visual_gen/attention_backend @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs -/tensorrt_llm/_torch/visual_gen/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs +/tensorrt_llm/_torch/visual_gen/attention @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs /cpp/tensorrt_llm/thop/*lysses* @NVIDIA/trt-llm-torch-visual-gen-devs /cpp/tensorrt_llm/kernels/ulysses* @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/media @NVIDIA/trt-llm-torch-visual-gen-devs diff --git a/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp b/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp index 48cf2a1d2ec9..0dd076f62578 100644 --- a/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp +++ b/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp @@ -18,7 +18,7 @@ // Async Ulysses A2A — PyTorch _SymmetricMemory CUDA-IPC backend. // // Pipeline (paired with UlyssesAttention.forward_async in -// tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py): +// tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py): // // recv, send_h = ulysses_a2a_async_prepare(input, pg) # default stream // ev.record() diff --git a/docs/source/blogs/tech_blog/blog25_Scaling_Video_Generation_Across_NVL72_Rack_with_TensorRT-LLM.md b/docs/source/blogs/tech_blog/blog25_Scaling_Video_Generation_Across_NVL72_Rack_with_TensorRT-LLM.md index 7fdafdfdb2c1..1641234131ef 100644 --- a/docs/source/blogs/tech_blog/blog25_Scaling_Video_Generation_Across_NVL72_Rack_with_TensorRT-LLM.md +++ b/docs/source/blogs/tech_blog/blog25_Scaling_Video_Generation_Across_NVL72_Rack_with_TensorRT-LLM.md @@ -166,7 +166,7 @@ Rank r holds Q_r (fixed) and K/V block for step i running output merged after each step --> final O_r ``` -**Online-softmax merge.** Each step yields a partial output and a log-sum-exp (LSE) statistic; [`RingAttention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py) merges them with a streaming softmax recurrence in **fp32**, block by block, until the result matches one global attention pass. +**Online-softmax merge.** Each step yields a partial output and a log-sum-exp (LSE) statistic; [`RingAttention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py) merges them with a streaming softmax recurrence in **fp32**, block by block, until the result matches one global attention pass. **Overlapping communication and compute.** The K/V exchange for the next block is posted with non-blocking point-to-point (`batch_isend_irecv`) *before* the current block's FA4 kernel runs, so the neighbor transfer overlaps with attention compute. Even ranks send-then-recv and odd ranks recv-then-send to avoid deadlock; the final ring step skips the exchange entirely. @@ -191,7 +191,7 @@ So **Q expands by `col_size`** (a row holds `col_size` ranks → $S/\text{row si Members of a **row (Q) group** end up with identical Q but **disjoint** K/V slices; members of a **column (K/V) group** share identical K/V but hold different Q. -**Local Attention.** The inner backend ([`Attention2DAttention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py)) runs attention with LSE returning a partial output $[B, S/\text{row size}, H, D]$ and a log-sum-exp (LSE) statistic $[B, H, S/\text{row size}]$. The LSE records the softmax normalizer for that partial pass. +**Local Attention.** The inner backend ([`Attention2DAttention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py)) runs attention with LSE returning a partial output $[B, S/\text{row size}, H, D]$ and a log-sum-exp (LSE) statistic $[B, H, S/\text{row size}]$. The LSE records the softmax normalizer for that partial pass. **Merge phase.** The `col_size` ranks that share the same Q each attended over a **disjoint slice of K/V** (above: ranks `0`, `1`, `2` all hold queries `{0,1,2}` but keys `{0,3}`, `{1,4}`, `{2,5}` respectively). Their FA4 outputs are partial results for the **same query positions** over **complementary key spans**, so they must be combined **within that row group** to recover the full softmax. For `col_size = 3`, the partial results are spread across 3 GPUs. These are combined into the final output as follows: @@ -462,7 +462,7 @@ Future directions: - VisualGen runtime: [`tensorrt_llm/_torch/visual_gen/`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm/_torch/visual_gen) - Mesh: [`mapping.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/mapping.py) - - Parallelism wrappers: [`attention_backend/parallel.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py) + - Parallelism wrappers: [`attention/backends/parallel.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py) - Pipeline base: [`pipeline.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/pipeline.py) - Config models: [`config.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/config.py) - Parallel VAE: [`modules/vae/parallel_vae_interface.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/visual_gen/modules/vae/parallel_vae_interface.py) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index bc2da9b9c713..9d4f33a061e8 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -63,7 +63,7 @@ After a recipe meets the quality target, benchmark its end-to-end throughput wit | `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | | `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. Keep `0` if `v_dtype` defines its own scaling format (e.g., `mxfp8` or `nvfp4`) | -Routing (`tensorrt_llm/_torch/visual_gen/attention_backend/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. +Routing (`tensorrt_llm/_torch/visual_gen/attention/backends/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. ## QK16PV8 Attention Kernels in the CUTEDSL Backend diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index dbe283471a84..c532a78f0474 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1651,9 +1651,9 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/_torch/pyexecutor/model_engine.py", "tensorrt_llm/_torch/pyexecutor/py_executor.py", "tensorrt_llm/_torch/weight_sharing/", - "tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py", + "tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py", "tensorrt_llm/_torch/visual_gen/modules/vae/", - "tensorrt_llm/_torch/visual_gen/modules/attention.py", + "tensorrt_llm/_torch/visual_gen/attention/attention.py", "tensorrt_llm/_torch/visual_gen/executor.py", "tensorrt_llm/_torch/visual_gen/mapping.py", "tensorrt_llm/_torch/visual_gen/pipeline.py", diff --git a/tensorrt_llm/_torch/visual_gen/attention/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/__init__.py new file mode 100644 index 000000000000..d418485c1572 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention/__init__.py @@ -0,0 +1,22 @@ +# 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. +"""VisualGen attention. + +Holds the DiT attention module (``attention.py``) and the kernel backends it +dispatches to (the ``backends`` sub-package). + +Nothing is re-exported from this file on purpose: importing the package root +stays cheap, and callers import the module or backend they need directly. +""" diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/attention/attention.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/modules/attention.py rename to tensorrt_llm/_torch/visual_gen/attention/attention.py index cd91df8ae11f..2f1d97c0f6e4 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/attention/attention.py @@ -24,11 +24,11 @@ from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ...utils import Fp4QuantizedTensor -from ..attention_backend.interface import AttentionTensorLayout -from ..attention_backend.parallel import wrap_parallel_attention -from ..attention_backend.utils import create_attention from ..config import DiffusionModelConfig from ..modules.rms_norm import RMSNormTPAware +from .backends.interface import AttentionTensorLayout +from .backends.parallel import wrap_parallel_attention +from .backends.utils import create_attention class QKVMode(str, Enum): diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/backends/__init__.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/__init__.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py b/tensorrt_llm/_torch/visual_gen/attention/backends/cudnn.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/cudnn.py index 238978c49033..139393c24cb5 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/cudnn.py @@ -43,7 +43,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention.backends.interface import PredefinedAttentionMask +from ....attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/backends/cute_dsl/__init__.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/cute_dsl/__init__.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention/backends/cute_dsl/fmha.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/cute_dsl/fmha.py index edebb9da7078..7039ff7b1872 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/cute_dsl/fmha.py @@ -17,7 +17,7 @@ JIT-compiles dense or SkipSoftmax FMHA and caches the compiled artifact for each kernel configuration. Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA -sparse backend uses `VSACuTeDSLAttention` in `attention_backend.sparse.vsa` instead. +sparse backend uses `VSACuTeDSLAttention` in `attention.backends.sparse.vsa` instead. """ import math @@ -30,7 +30,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ....attention.backends.interface import PredefinedAttentionMask +from .....attention.backends.interface import PredefinedAttentionMask from ..interface import AttentionBackend, AttentionTensorLayout _cute_dsl_import_error: BaseException | None = None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py index 1ec1b2db77be..de2be0c1e112 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py @@ -25,7 +25,7 @@ import torch -from ...attention.backends.interface import PredefinedAttentionMask +from ....attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout # 2CTA regresses long sequences; must precede the import below. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py b/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py index caf3ddb71f65..168f36d64f64 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/flashinfer.py @@ -10,7 +10,7 @@ from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention.backends.interface import PredefinedAttentionMask +from ....attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout _WORKSPACE_BYTES = 128 * 1024 * 1024 diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/interface.py b/tensorrt_llm/_torch/visual_gen/attention/backends/interface.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/interface.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/interface.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py index 3a8ff1f72bc2..bea2871bd072 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/parallel.py @@ -29,11 +29,11 @@ import torch.nn.functional as F if TYPE_CHECKING: - from ..mapping import VisualGenMapping + from ...mapping import VisualGenMapping from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d -from ...attention.backends.interface import PredefinedAttentionMask +from ....attention.backends.interface import PredefinedAttentionMask from .flash_attn4 import _install_cutlass_dsl_compatibility from .interface import AttentionBackend, AttentionTensorLayout diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/__init__.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/__init__.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/__init__.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/__init__.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py similarity index 99% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py index 5872ccbbcb6b..f8070fb9fecd 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py @@ -256,7 +256,7 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: """Whether the sparse kernel applies to this particular call. Everything false here is delegated (see ``_delegate``). Deciding it from the - tensors, per call, is deliberate, and it is why ``modules/attention.py`` + tensors, per call, is deliberate, and it is why ``attention/attention.py`` has no ``SEPARATE_QKV`` rule for SOL: ``qkv_mode`` describes how Q/K/V are *projected*, not whether K/V come from another sequence, so a construction-time rule keyed on it mistakes self-attention for diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/kernels.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/kernels.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/params.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/params.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/predictor.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/predictor.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/__init__.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/__init__.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py similarity index 98% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py index 6968424dac95..386eaef42ad9 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py @@ -24,9 +24,9 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from .....attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha -from .....attention.backends.interface import PredefinedAttentionMask -from .....attention.backends.sparse.params import SparseBackendForwardArgs +from ......attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from ......attention.backends.interface import PredefinedAttentionMask +from ......attention.backends.sparse.params import SparseBackendForwardArgs from ...cute_dsl import CuTeDSLAttention from ...trtllm import TrtllmAttention from .metadata import VSA_BLOCK_SIZE, VSAMetadata diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/kernels.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/kernels.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/metadata.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/metadata.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/predictor.py similarity index 98% rename from tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/predictor.py index 6fa244039981..a0d340813931 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/predictor.py @@ -21,8 +21,8 @@ import torch -from .....attention.backends.interface import PredefinedAttentionMask -from .....attention.backends.sparse.params import BlockSparseForwardInputs +from ......attention.backends.interface import PredefinedAttentionMask +from ......attention.backends.sparse.params import BlockSparseForwardInputs from .kernels import blend_coarse_fine, sort_last_dim, tile_and_pool_cubes from .metadata import VSA_BLOCK_SIZE, VSAMetadata, get_vsa_forward_context diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py similarity index 97% rename from tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index 6e5371487c84..3502ca6f9a4a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -27,16 +27,16 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention.backends.interface import ( +from ....attention.backends.interface import ( AttentionForwardArgs, AttentionRuntimeFeatures, PredefinedAttentionMask, ) -from ...attention.backends.sparse.params import SparseBackendForwardArgs, SparseParams -from ...attention.backends.sparse.timestep_phase import graph_phase_for_timestep -from ...attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention -from ...attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata -from ..cuda_graph_runner import resolved_extra_key +from ....attention.backends.sparse.params import SparseBackendForwardArgs, SparseParams +from ....attention.backends.sparse.timestep_phase import graph_phase_for_timestep +from ....attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention +from ....attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata +from ...cuda_graph_runner import resolved_extra_key from .interface import AttentionBackend, AttentionTensorLayout diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention/backends/utils.py similarity index 100% rename from tensorrt_llm/_torch/visual_gen/attention_backend/utils.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/utils.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py b/tensorrt_llm/_torch/visual_gen/attention/backends/vanilla.py similarity index 96% rename from tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py rename to tensorrt_llm/_torch/visual_gen/attention/backends/vanilla.py index 8fc94ec126be..b6b109c55bcf 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/vanilla.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +28,7 @@ import torch import torch.nn.functional as F -from ...attention.backends.interface import PredefinedAttentionMask +from ....attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 926ca6e5f957..a466873d027e 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -44,7 +44,7 @@ ``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, ``kv_splits``, and an optional exact KV sink range. -TRT-LLM's dispatch path (``attention_backend/sparse/sol/backend.py``, +TRT-LLM's dispatch path (``attention/backends/sparse/sol/backend.py``, ``SOLCuTeDSLAttention``) consumes exactly two names from this module: ``_run_sol_attn_bthd`` and ``sol_attn_supported``. The dense-prefix decision lives there too, keyed off the normalized timestep forward kwarg. 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 750bd51fb648..ca049e9e1013 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -30,6 +30,7 @@ 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.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.cosmos3.step_precision import ( StepPrecisionController, @@ -37,7 +38,6 @@ parse_diffusion_step_policy, ) from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py index 0d3b3054fd0d..e7cc1e141959 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -22,12 +22,12 @@ from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.modules.swiglu import swiglu from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode, apply_rotary_emb from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import ( FluxJointAttnMLPProj, FluxJointQKVMLPProj, ) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode, apply_rotary_emb from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.quantization.mode import QuantAlgo diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py index db8ed3e04472..5bc1ff35d927 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py @@ -7,7 +7,7 @@ - FluxPosEmbed: Multi-axis rotary position embeddings (FLUX.1: 3-axis, FLUX.2: 4-axis) - get_1d_rotary_pos_embed: 1D rotary position embedding computation -RoPE application uses the shared apply_rotary_emb from modules/attention.py. +RoPE application uses the shared apply_rotary_emb from attention/attention.py. """ from typing import List, Tuple diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index 32c183229748..9013b8c1ba6b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -25,9 +25,9 @@ from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode 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 from tensorrt_llm._torch.visual_gen.quantization import DynamicLinearWeightLoader from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py index 6a393da02d17..0e3ea948fe3b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py +++ b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py @@ -33,9 +33,9 @@ from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode 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 from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig 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 2ff493af9900..7be3b9f16506 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -7,7 +7,7 @@ # - Linear projections → tensorrt_llm._torch.modules.linear.Linear # - RMSNorm (QK norm) → tensorrt_llm._torch.modules.rms_norm.RMSNorm # - FeedForward (MLP) → tensorrt_llm._torch.modules.mlp.MLP -# - Attention backend → tensorrt_llm._torch.visual_gen.attention_backend +# - Attention backend → tensorrt_llm._torch.visual_gen.attention.backends # # Architecture-specific components (RoPE, AdaLN, timestep/text embeddings, # modality dataclass, transformer args) are ported from LTX-2 and live @@ -32,13 +32,13 @@ 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.parallel import ( +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( UlyssesAttention, wrap_parallel_attention, ) -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import create_attention from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py index fd6a5bdc3c4a..3bc5c3e4008b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py @@ -35,9 +35,9 @@ from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode 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 from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig 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 c0733a570078..2f109a0447fa 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 @@ -30,14 +30,14 @@ 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.parallel import ( +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( Attention2DAttention, RingAttention, UlyssesAttention, ) 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 from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig 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 5f78ed1d5be0..ae61c2da8496 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -26,7 +26,7 @@ from diffusers.video_processor import VideoProcessor from transformers import AutoTokenizer, UMT5EncoderModel -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) 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 1c8ae1c2e6da..ca1e9ac97f92 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -27,6 +27,7 @@ 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.attention import Attention, QKVMode 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 ( @@ -36,7 +37,6 @@ apply_fused_layernorm_affine_quant, get_nvfp4_input_scale, ) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.modules.rms_norm import RMSNormTPAware from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm._torch.visual_gen.utils import SequenceSharder diff --git a/tensorrt_llm/_torch/visual_gen/modules/__init__.py b/tensorrt_llm/_torch/visual_gen/modules/__init__.py index 93636ddb53f4..12b4fb6f08c4 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/modules/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,10 +17,3 @@ This module provides modular neural network components for visual generation models. """ - -from .attention import Attention, QKVMode - -__all__ = [ - "Attention", - "QKVMode", -] diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 24e26316b96c..782eeed16bf2 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -24,7 +24,7 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel if TYPE_CHECKING: - from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.params import SolParams class BaseSparseAttentionConfig(StrictBaseModel): @@ -257,7 +257,7 @@ class SolAttentionConfig(BaseSparseAttentionConfig): dense attention and are counted; GQA/MQA raises at construction. Context parallelism (cp_size > 1) is rejected for both backends in - visual_gen/modules/attention.py. + visual_gen/attention/attention.py. """ algorithm: Literal["sol_attn"] = "sol_attn" @@ -314,7 +314,7 @@ def _validate_dense_layers(cls, layers: Optional[list[int]]) -> Optional[list[in def to_sparse_params(self, **kwargs: Any) -> "SolParams": """Lower the public recipe into the SOL parameters shared by both backends.""" del kwargs - from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.params import SolParams return SolParams( tau=self.tau, 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 443a398e2bd6..c0c60c5bc9dd 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 @@ -34,10 +34,12 @@ import torch.nn.functional as F from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask -from tensorrt_llm._torch.visual_gen.attention_backend import Attention2DAttention -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd -from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention.backends import Attention2DAttention +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import FlashAttn4Attention +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, +) +from tensorrt_llm._torch.visual_gen.attention.backends.interface import AttentionTensorLayout # FA4 kernel availability (separate from the combine kernel used by Attention2DAttention) _flash_attn4_available = _fa4_fwd is not None 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 a678403fa5b2..0c9d4fd8b9ae 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 @@ -44,10 +44,10 @@ # attn2d tests fail loudly when the kernels are missing on their Blackwell runner — # a build/dependency problem, not a reason to skip. try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, ) - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( _flash_attn_combine as _fa_combine, ) @@ -680,7 +680,7 @@ def _logic_cosmos3_ulysses_vs_single_gpu(rank, world_size): text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) ref_out = _forward(ref_model, device, text_seed) - from tensorrt_llm._torch.visual_gen.attention_backend import parallel as parallel_attention + from tensorrt_llm._torch.visual_gen.attention.backends import parallel as parallel_attention original_prefix_helper = parallel_attention._run_attention_with_replicated_kv 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 14bd1bfc184c..efda6abe5170 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 @@ -42,7 +42,7 @@ from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, ) 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 14204a027207..007b17bf5fc5 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 @@ -439,7 +439,7 @@ def _logic_ltx2_dual_topology(rank, world_size, backend, audio_seq_len): expect = blk.attn1._attn_stage2 if is_stage2 else blk.attn1._attn_default assert blk.attn1.attn is expect, f"Rank {rank}: attn stack not switched" # is_ulysses must track the ACTIVE stack (the pair can differ in type). - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import UlyssesAttention + from tensorrt_llm._torch.visual_gen.attention.backends.parallel import UlyssesAttention for name in ("attn1", "video_to_audio_attn", "audio_attn1"): mod = getattr(blk, name, None) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py index 87c709731272..01ec1550b863 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py @@ -25,7 +25,7 @@ import torch.multiprocessing as mp from utils.util import skip_num_gpus_less_than, skip_single_gpu -from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( +from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( Attention2DAttention, RingAttention, UlyssesAttention, @@ -36,10 +36,10 @@ from tensorrt_llm.visual_gen.args import ParallelConfig try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, ) - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( _flash_attn_combine as _fa_combine, ) 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 ceeb15d1897d..36d1f3d3873d 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 @@ -42,10 +42,12 @@ from torch.distributed.device_mesh import init_device_mesh from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask -from tensorrt_llm._torch.visual_gen.attention_backend import RingAttention, UlyssesAttention -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd -from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention.backends import RingAttention, UlyssesAttention +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import FlashAttn4Attention +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, +) +from tensorrt_llm._torch.visual_gen.attention.backends.interface import AttentionTensorLayout _flash_attn4_available = _fa4_fwd is not None 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 8878cd788963..d93542ae37d8 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 @@ -39,9 +39,9 @@ import torch.nn.functional as F from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.mapping import Mapping from tensorrt_llm.visual_gen.args import AttentionConfig 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 49660abdcb41..ff4fc879619f 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 @@ -44,8 +44,8 @@ def _disable_mpi() -> Iterator[None]: with _disable_mpi(): from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d - from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention, VanillaAttention - from tensorrt_llm._torch.visual_gen.attention_backend.interface import ( + from tensorrt_llm._torch.visual_gen.attention.backends import UlyssesAttention, VanillaAttention + from tensorrt_llm._torch.visual_gen.attention.backends.interface import ( AttentionBackend, AttentionTensorLayout, ) @@ -116,7 +116,7 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = def test_forward_async_redistributes_vsa_gates(monkeypatch: pytest.MonkeyPatch) -> None: - import tensorrt_llm._torch.visual_gen.attention_backend.parallel as parallel_backend + import tensorrt_llm._torch.visual_gen.attention.backends.parallel as parallel_backend class _CaptureBackend: preferred_layout = AttentionTensorLayout.NHD 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 7559c8182162..d4c82123dc3b 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 @@ -36,8 +36,8 @@ import torch.multiprocessing as mp import torch.nn.functional as F -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.attention.backends import UlyssesAttention +from tensorrt_llm._torch.visual_gen.attention.backends.trtllm import TrtllmAttention from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state from tensorrt_llm.visual_gen.args import QuantAttentionConfig diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py index 4584cca7bc86..276912736dce 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py @@ -318,7 +318,7 @@ def test_async_vs_sync_parity(self, backend): def test_forward_async_uses_tp_local_heads_for_qkv_gates_and_output() -> None: - from tensorrt_llm._torch.visual_gen.modules.attention import Attention + from tensorrt_llm._torch.visual_gen.attention.attention import Attention class _CaptureAsyncAttention(torch.nn.Module): def forward_async( diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py index c5e1aed1e148..f268297e1365 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py @@ -42,10 +42,10 @@ ) try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, ) - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( _flash_attn_combine as _fa_combine, ) 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 45db554d0014..0f7664fed5cf 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 @@ -529,8 +529,8 @@ def test_wan_vsa_gates_follow_ulysses_aligned_tp_q_shard( monkeypatch: pytest.MonkeyPatch, tp_rank: int, expected_head_range: tuple[int, int] ) -> None: """VSA gates must select the same TP-local heads as the Q projection.""" + from tensorrt_llm._torch.visual_gen.attention import attention as attention_module from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanBlock - from tensorrt_llm._torch.visual_gen.modules import attention as attention_module from tensorrt_llm.mapping import Mapping from tensorrt_llm.visual_gen.args import VideoSparseAttentionConfig 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 a4db10e4a50e..4d311e762865 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 @@ -45,10 +45,10 @@ from .tp_shard_utils import copy_tp_parameter try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, ) - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( _flash_attn_combine as _fa_combine, ) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py index 187b3ae6ebe0..65b82714148d 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py @@ -35,8 +35,8 @@ import torch.multiprocessing as mp from utils.llm_data import get_checkpoint -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import _cute_dsl_import_error +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py index ddd1cb0ad219..a0937ed14966 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py @@ -15,14 +15,14 @@ from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import fmha as cute_dsl_fmha -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import ( +from tensorrt_llm._torch.visual_gen.attention.attention import Attention +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import fmha as cute_dsl_fmha +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl.fmha import ( CuTeDSLAttention, _resolve_skip_softmax_threshold_scale_factor, ) 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 from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig, VisualGenArgs from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig @@ -517,7 +517,7 @@ def _fake_cute_dsl_fmha_fwd(q, k, v, o, **kwargs): o.zero_() monkeypatch.setattr( - "tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha.cute_dsl_fmha_fwd", + "tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl.fmha.cute_dsl_fmha_fwd", _fake_cute_dsl_fmha_fwd, ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py index 35a8034bc8c5..31158399a46b 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py @@ -41,22 +41,24 @@ from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention as CoreTrtllmAttention -from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import backend as sol_backend -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import predictor as sol_predictor -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend import ( +from tensorrt_llm._torch.visual_gen.attention import attention as attention_module +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.attention.backends import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol import backend as sol_backend +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol import predictor as sol_predictor +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.backend import ( SOLCuTeDSLAttention, SOLTrtllmAttention, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.params import SolParams +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.predictor import ( SolPredictorOutputs, ) -from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import ( +from tensorrt_llm._torch.visual_gen.attention.backends.trtllm import ( TrtllmAttention, TrtllmAttentionMetadata, ) -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import create_attention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, @@ -68,8 +70,6 @@ resolved_extra_keys_scope, ) from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel -from tensorrt_llm._torch.visual_gen.modules import attention as attention_module -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen import SolAttentionConfig from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py index df6735fdc696..c4f25d813db5 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py @@ -31,15 +31,15 @@ import torch from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import predictor as sol_predictor -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol import predictor as sol_predictor +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.backend import ( SOLCuTeDSLAttention, SOLTrtllmAttention, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.predictor import ( SolPredictorOutputs, ) -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import create_attention from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_keys_scope from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell import sol_attn_backend diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py index f5ce73cc777b..83326c3259f6 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py @@ -10,11 +10,11 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.kernels import ( block_pool, block_thresholds, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.predictor import ( BLOCK_SIZE, SolPredictorOutputs, predict, diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py index 330b23a2fd05..a56fffdc64cb 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol.kernels import ( _LOG2_E, _block_pool_torch, _block_thresholds_torch, diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py index ba7195246c58..24aa3b32f774 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py @@ -28,34 +28,34 @@ from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import CuTeDSLAttention -from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as vsa_backend -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import kernels as vsa_kernels -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import predictor as vsa_predictor -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( +from tensorrt_llm._torch.visual_gen.attention import attention as attention_module +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention.backends.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import backend as vsa_backend +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import kernels as vsa_kernels +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import predictor as vsa_predictor +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.backend import ( VSACuTeDSLAttention, VSATrtllmAttention, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.kernels import tile_and_pool_cubes -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.metadata import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.kernels import tile_and_pool_cubes +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.metadata import ( VSA_BLOCK_SIZE, VSAMetadataBuilder, set_vsa_forward_context, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.predictor import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.predictor import ( VSAForwardInputs, VSAPredictor, ) -from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention -from tensorrt_llm._torch.visual_gen.attention_backend.vanilla import VanillaAttention +from tensorrt_llm._torch.visual_gen.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import create_attention +from tensorrt_llm._torch.visual_gen.attention.backends.vanilla import VanillaAttention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, ) -from tensorrt_llm._torch.visual_gen.modules import attention as attention_module -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py index 657beb19b670..0637ea33cd89 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py @@ -26,7 +26,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.kernels import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.kernels import ( blend_coarse_fine, sort_last_dim, tile_and_pool_cubes, diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py index 0fe6b494a212..90e2af853903 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py @@ -27,16 +27,16 @@ from utils.util import isSM100Family from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode, apply_rotary_emb +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import VSACuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.backend import VSACuTeDSLAttention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, ) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode, apply_rotary_emb from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig _REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") diff --git a/tests/unittest/_torch/visual_gen/test_attention_cudnn.py b/tests/unittest/_torch/visual_gen/test_attention_cudnn.py index 56554d46577f..88c4a945179c 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cudnn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cudnn.py @@ -19,13 +19,13 @@ from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.utils import unswizzle_sf -from tensorrt_llm._torch.visual_gen.attention_backend import create_attention -from tensorrt_llm._torch.visual_gen.attention_backend.cudnn import ( +from tensorrt_llm._torch.visual_gen.attention.backends import create_attention +from tensorrt_llm._torch.visual_gen.attention.backends.cudnn import ( CuDNNAttention, _quantize_mxfp8_qk, _quantize_mxfp8_v, ) -from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention.backends.interface import AttentionTensorLayout from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig # Quantization dtype -> quant_attention_config accepted by AttentionConfig. diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py index dcdedfcc2ef5..83523cde2ae9 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py @@ -18,7 +18,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import ( +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl.fmha import ( _COMPILE_CACHE, _quantize_blockscaled_one, _quantize_fp8_v, diff --git a/tests/unittest/_torch/visual_gen/test_attention_fa4.py b/tests/unittest/_torch/visual_gen/test_attention_fa4.py index ecd00827b059..db6a0c0ea37d 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_fa4.py +++ b/tests/unittest/_torch/visual_gen/test_attention_fa4.py @@ -24,7 +24,7 @@ import torch try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( FlashAttn4Attention, _flash_attn_fwd, ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py b/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py index a6dc0d2f4736..f1c08e944093 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py +++ b/tests/unittest/_torch/visual_gen/test_attention_flashinfer.py @@ -11,8 +11,8 @@ import torch.nn.functional as F from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask -from tensorrt_llm._torch.visual_gen.attention_backend.flashinfer import FlashInferAttention -from tensorrt_llm._torch.visual_gen.attention_backend.utils import get_visual_gen_attention_backend +from tensorrt_llm._torch.visual_gen.attention.backends.flashinfer import FlashInferAttention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import get_visual_gen_attention_backend from tensorrt_llm.visual_gen.args import QuantAttentionConfig diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index 64fcb0d14f14..dd47b6871277 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -14,23 +14,25 @@ import torch.nn.functional as F from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.visual_gen.attention_backend.cudnn import CuDNNAttention -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd -from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + +# Import new integrated versions +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode, apply_rotary_emb +from tensorrt_llm._torch.visual_gen.attention.backends.cudnn import CuDNNAttention +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import _cute_dsl_import_error +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, +) +from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( Attention2DAttention, RingAttention, UlyssesAttention, ) -from tensorrt_llm._torch.visual_gen.attention_backend.vanilla import VanillaAttention +from tensorrt_llm._torch.visual_gen.attention.backends.vanilla import VanillaAttention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, ) 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._utils import get_sm_version from tensorrt_llm.visual_gen.args import ( AttentionConfig, diff --git a/tests/unittest/_torch/visual_gen/test_attention_perf.py b/tests/unittest/_torch/visual_gen/test_attention_perf.py index 855366914932..8fe38d2c95b7 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_perf.py +++ b/tests/unittest/_torch/visual_gen/test_attention_perf.py @@ -36,15 +36,19 @@ import pytest import torch +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode + # ============================================================================ # Flash Attention 4 availability # ============================================================================ -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import _cute_dsl_import_error +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( + _flash_attn_fwd as _fa4_fwd, +) +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( _flash_attn_fwd_import_error as _fa4_import_error, ) -from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) @@ -52,7 +56,6 @@ DiffusionModelConfig, create_attention_metadata_state, ) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen.args import ( AttentionConfig, QuantAttentionConfig, @@ -1079,7 +1082,7 @@ def test_vsa_kernel_vs_fa4( block_size: int, sparsity: float, ): - from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.backend import ( VSA_KERNEL_MAX_CUBES, ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py b/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py index 2372019b748b..251476a8b957 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py +++ b/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py @@ -27,8 +27,8 @@ import pytest import torch +from tensorrt_llm._torch.visual_gen.attention.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo diff --git a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py index 427787dbe073..25c38a467f92 100644 --- a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py +++ b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py @@ -11,8 +11,8 @@ cute = pytest.importorskip("cutlass.cute") -from tensorrt_llm._torch.visual_gen.attention_backend import flash_attn4, parallel # noqa: E402 -from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( # noqa: E402 +from tensorrt_llm._torch.visual_gen.attention.backends import flash_attn4, parallel # noqa: E402 +from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import ( # noqa: E402 _install_cutlass_dsl_compatibility, _install_flash_attn_tile_scheduler_compatibility, ) @@ -85,7 +85,7 @@ def test_fa4_work_tile_info_survives_task_scheduling_imported_first() -> None: import cutlass.experimental.task_scheduling # noqa: F401 from cutlass.cutlass_dsl import Boolean - import tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 # noqa: F401 + import tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 # noqa: F401 from flash_attn.cute.tile_scheduler import WorkTileInfo tile_idx = (cutlass.Int32(1), cutlass.Int32(2), cutlass.Int32(3), cutlass.Int32(0)) @@ -120,7 +120,7 @@ def test_fa4_kernel_traces_after_task_scheduling_imported_first() -> None: import cutlass.experimental.task_scheduling # noqa: F401 import torch - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import _flash_attn_fwd torch.manual_seed(0) # S_kv stays above FA4's short-KV short-circuit so num_splits=8 reaches the diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py index 64ba148c4cd6..7e33e3feae94 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py @@ -36,7 +36,7 @@ from tensorrt_llm.visual_gen.args import AttentionConfig try: - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd + from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import _flash_attn_fwd FA4_AVAILABLE = _flash_attn_fwd is not None except ImportError: 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 8a7e35c9167b..792586431d80 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 @@ -17,8 +17,9 @@ # Importing the models package applies the Qwen-Image registration side effect. from tensorrt_llm._torch.visual_gen import models # noqa: F401 -from tensorrt_llm._torch.visual_gen import modules as visual_gen_modules -from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( +from tensorrt_llm._torch.visual_gen.attention import attention as visual_gen_attention +from tensorrt_llm._torch.visual_gen.attention.attention import QKVMode, apply_rotary_emb +from tensorrt_llm._torch.visual_gen.attention.backends.parallel import ( Attention2DAttention, RingAttention, UlyssesAttention, @@ -44,7 +45,6 @@ qwen_complex_freqs_to_cos_sin, qwen_joint_freqs_to_cos_sin, ) -from tensorrt_llm._torch.visual_gen.modules.attention import QKVMode, apply_rotary_emb from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.mode import QuantAlgo @@ -259,7 +259,7 @@ def capture_parallel_attention(attn, **kwargs): return attn monkeypatch.setattr( - visual_gen_modules.attention, + visual_gen_attention, "wrap_parallel_attention", capture_parallel_attention, ) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index 37796447dd8a..056cfb5d39d8 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -16,7 +16,7 @@ SkipSoftmaxParams, SkipSoftmaxScheduler, ) -from tensorrt_llm._torch.visual_gen.attention_backend import trtllm as visual_trtllm +from tensorrt_llm._torch.visual_gen.attention.backends import trtllm as visual_trtllm from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_keys_scope 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..c025e6da5c92 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 @@ -11,7 +11,7 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.attention_backend import VanillaAttention +from tensorrt_llm._torch.visual_gen.attention.backends import VanillaAttention pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py index 330e5a355305..f67b74f7727b 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -26,7 +26,7 @@ import torch.nn.functional as F from utils.llm_data import get_checkpoint -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error +from tensorrt_llm._torch.visual_gen.attention.backends.cute_dsl import _cute_dsl_import_error from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import ( AttentionConfig, @@ -120,7 +120,7 @@ def _count_cute_fine_stage_launches() -> Iterator[SimpleNamespace]: """ from unittest.mock import patch - from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.backend import ( VSACuTeDSLAttention, ) @@ -147,7 +147,7 @@ def _assert_vsa_matches_dense( """Compare CuTe-DSL VSA against SDPA-fallback VSA (same gated formulation, different fine kernel).""" from unittest.mock import patch - from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as _vsa_module + from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa import backend as _vsa_module common_kwargs = dict( prompt=PROMPT,