[TRTLLM-15864][chore] Consolidate the VisualGen attention domain under _torch/visual_gen/attention/

.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

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()

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)  

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
 

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",

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.
+"""

tensorrt_llm/_torch/visual_gen/modules/attention.py → 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):

tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py → 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
 
 

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py → 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

tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py → 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.

tensorrt_llm/_torch/visual_gen/attention_backend/flashinfer.py → 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

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py → 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
 

tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py → 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

tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py → 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

tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py → 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
 

tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py → 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
 
 

tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py → 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
 
 

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.

tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py

@@ -30,14 +30,14 @@
 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,
     install_step_precision,
     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

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
 

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

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
 

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

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