From 8846a96f736cc60ac5407b0e1a0cb462a1a10852 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Fri, 31 Jul 2026 22:11:41 +0000
Subject: [PATCH 01/18] Add Qwen Image Edit Ulysses support

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 docs/source/models/supported-models.md        |   2 +-
 docs/source/models/visual-generation.md       |   2 +-
 ...qwen-image-edit-2511-fp8-2gpu-ulysses.yaml |  25 +++
 .../qwen_image/pipeline_qwen_image_edit.py    |  23 ++-
 .../multi_gpu/test_qwen_image_edit_ulysses.py | 149 ++++++++++++++++++
 .../test_qwen_image_pipeline_config.py        |  11 ++
 6 files changed, 209 insertions(+), 3 deletions(-)
 create mode 100644 examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
 create mode 100644 tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index 5ef301de8624..363f172e95f2 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -232,7 +232,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image-Layered** [^vg2] | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | No | No | Yes | Yes | Yes | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | No | No | No | No |
 | **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 
 [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md
index 3e80c120b2ef..29349994b969 100644
--- a/docs/source/models/visual-generation.md
+++ b/docs/source/models/visual-generation.md
@@ -66,7 +66,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 | **LTX-2** | Yes | Yes | Yes [^4] | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | No | No | Yes | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | No | No | No | No | No |
 | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |
 | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
 | **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
diff --git a/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml b/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
new file mode 100644
index 000000000000..018bc31f8b3a
--- /dev/null
+++ b/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
@@ -0,0 +1,25 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+# 2-GPU Qwen-Image-Edit-2511 with FP8 blockwise dynamic quantization and Ulysses.
+quant_config:
+  quant_algo: FP8_BLOCK_SCALES
+  dynamic: true
+attention_config:
+  backend: VANILLA
+parallel_config:
+  cfg_size: 1
+  ulysses_size: 2
+cuda_graph_config:
+  enable: false
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
index 7cf64dabb121..9d9f76a52e34 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
@@ -11,7 +11,7 @@
 import math
 import time
 from io import BytesIO
-from typing import Any
+from typing import Any, Optional
 
 import numpy as np
 import PIL.Image
@@ -279,6 +279,21 @@ def _encode_edit_prompt(
             prompt_embeds_mask = None
         return prompt_embeds, prompt_embeds_mask
 
+    @staticmethod
+    def _validate_ulysses_prompt_masks(
+        ulysses_size: int,
+        *prompt_masks: Optional[torch.Tensor],
+    ) -> None:
+        if ulysses_size <= 1:
+            return
+        if any(mask is not None for mask in prompt_masks):
+            raise ValueError(
+                "Qwen-Image-Edit Ulysses parallelism requires unmasked prompt "
+                "conditioning. The supported single-prompt edit path produces "
+                "all-valid prompt masks; use ulysses_size=1 for padded or "
+                "batched prompt conditioning."
+            )
+
     def _encode_vae_image(
         self,
         image: torch.Tensor,
@@ -432,6 +447,7 @@ def forward(
         cfg_size = vgm.cfg_size if vgm else 1
         cfg_rank = vgm.cfg_rank if vgm else 0
         cfg_pg = vgm.cfg_group if vgm else None
+        ulysses_size = vgm.ulysses_size if vgm else 1
         do_cfg_parallel = do_true_cfg and cfg_size > 1
         if do_cfg_parallel:
             if cfg_size != 2:
@@ -464,6 +480,11 @@ def forward(
                 max_sequence_length,
                 num_images_per_prompt=1,
             )
+        self._validate_ulysses_prompt_masks(
+            ulysses_size,
+            prompt_embeds_mask,
+            neg_prompt_embeds_mask,
+        )
 
         num_channels_latents = self.transformer.in_channels // 4
         latents, image_latents = self._prepare_edit_latents(
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
new file mode 100644
index 000000000000..0925ccb1047c
--- /dev/null
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Multi-GPU tests for Qwen-Image-Edit Ulysses sequence parallelism."""
+
+import os
+
+os.environ["TLLM_DISABLE_MPI"] = "1"
+
+from typing import Callable
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+try:
+    import sys
+    from pathlib import Path
+
+    from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl
+    from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
+    from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping
+    from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import (
+        QwenImageEditPlusPipeline,
+    )
+    from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (
+        QwenJointAttention,
+    )
+
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+    from _visual_gen_dist_utils import spawn_with_retry
+
+    MODULES_AVAILABLE = True
+except ImportError:
+    MODULES_AVAILABLE = False
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _cleanup_mpi_env():
+    yield
+    os.environ.pop("TLLM_DISABLE_MPI", None)
+
+
+def init_distributed_worker(rank: int, world_size: int, port: int) -> None:
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+
+
+def cleanup_distributed() -> None:
+    if dist.is_initialized():
+        dist.destroy_process_group()
+    DeviceMeshTopologyImpl.device_mesh = None
+    VisualGenMapping.seq_mesh = None
+
+
+def _distributed_worker(rank: int, world_size: int, test_fn: Callable, port: int) -> None:
+    try:
+        init_distributed_worker(rank, world_size, port)
+        test_fn(rank, world_size)
+    except Exception as e:
+        print(f"Rank {rank} failed with error: {e}")
+        raise
+    finally:
+        cleanup_distributed()
+
+
+def run_test_in_distributed(world_size: int, test_fn: Callable) -> None:
+    if not MODULES_AVAILABLE:
+        pytest.skip("Required modules not available")
+    if torch.cuda.device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available")
+    spawn_with_retry(
+        lambda port: mp.spawn(
+            _distributed_worker,
+            args=(world_size, test_fn, port),
+            nprocs=world_size,
+            join=True,
+        )
+    )
+
+
+def _stabilize_attention_weights(module: torch.nn.Module) -> None:
+    """Use small deterministic weights for stable BF16 synthetic forward."""
+    with torch.no_grad():
+        for parameter in module.parameters():
+            if parameter.ndim >= 2:
+                fan_in = parameter.shape[1]
+                std = 0.02 / max(1.0, fan_in**0.5)
+                parameter.data.uniform_(-std, std)
+            else:
+                parameter.data.uniform_(-0.01, 0.01)
+
+
+def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
+    torch.manual_seed(1234)
+    vgm = VisualGenMapping(
+        world_size=world_size,
+        rank=rank,
+        cfg_size=1,
+        tp_size=1,
+        ring_size=1,
+        ulysses_size=world_size,
+    )
+    config = DiffusionModelConfig(
+        mapping=vgm.to_llm_mapping(),
+        visual_gen_mapping=vgm,
+        torch_dtype=torch.bfloat16,
+    )
+
+    attn = QwenJointAttention(
+        dim=16,
+        num_attention_heads=4,
+        attention_head_dim=4,
+        dtype=torch.bfloat16,
+        config=config,
+    ).cuda()
+    _stabilize_attention_weights(attn)
+    assert attn.attn.__class__.__name__ == "UlyssesAttention"
+
+    hidden_states = torch.randn(1, 4, 16, device="cuda", dtype=torch.bfloat16) * 0.1
+    encoder_hidden_states = torch.randn(1, 4, 16, device="cuda", dtype=torch.bfloat16) * 0.1
+    image_out, text_out = attn(
+        hidden_states,
+        encoder_hidden_states,
+        image_rotary_emb=None,
+        attention_mask=None,
+        timestep=None,
+    )
+
+    assert image_out.shape == hidden_states.shape
+    assert text_out.shape == encoder_hidden_states.shape
+    assert torch.isfinite(image_out).all()
+    assert torch.isfinite(text_out).all()
+
+    QwenImageEditPlusPipeline._validate_ulysses_prompt_masks(world_size, None, None)
+    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
+        QwenImageEditPlusPipeline._validate_ulysses_prompt_masks(
+            world_size,
+            torch.ones(1, 4, dtype=torch.bool, device="cuda"),
+        )
+
+
+def test_qwen_image_edit_ulysses_attention_2gpu():
+    run_test_in_distributed(2, _test_qwen_image_edit_ulysses_attention)
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 035a00b6e7f7..6f86863ab9cf 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
@@ -815,3 +815,14 @@ def test_qwen_image_edit_rejects_multiple_images_per_prompt() -> None:
 
     with pytest.raises(ValueError, match="num_images_per_prompt=1 only"):
         pipeline.infer(req)
+
+
+def test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning() -> None:
+    from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImageEditPlusPipeline
+
+    pipeline = QwenImageEditPlusPipeline.__new__(QwenImageEditPlusPipeline)
+
+    pipeline._validate_ulysses_prompt_masks(2, None, None)
+    pipeline._validate_ulysses_prompt_masks(1, object())
+    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
+        pipeline._validate_ulysses_prompt_masks(2, object())

From af07a6ab6348e0ea82ce59c6db2802f761a67dc5 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 4 Aug 2026 03:35:44 +0000
Subject: [PATCH 02/18] Address Qwen Image Ulysses review comments

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../models/qwen_image/pipeline_qwen_image.py  | 21 +++++++++
 .../multi_gpu/test_qwen_image_edit_ulysses.py | 43 ++++++++++---------
 .../test_qwen_image_pipeline_config.py        | 11 +++++
 3 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
index b7d3a69015cf..8bf1758db4da 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
@@ -350,6 +350,20 @@ def _encode_prompt(
             return prompt_embeds, None
         return prompt_embeds, prompt_embeds_mask
 
+    @staticmethod
+    def _validate_ulysses_prompt_masks(
+        ulysses_size: int,
+        *prompt_masks: Optional[torch.Tensor],
+    ) -> None:
+        if ulysses_size <= 1:
+            return
+        if any(mask is not None for mask in prompt_masks):
+            raise ValueError(
+                "Qwen-Image Ulysses parallelism requires unmasked prompt "
+                "conditioning. Use prompts that produce all-valid prompt "
+                "masks, or set ulysses_size=1 for padded prompt conditioning."
+            )
+
     # ------------------------------------------------------------------
     # Latent prep / packing (identical shape to FLUX).
     # ------------------------------------------------------------------
@@ -558,6 +572,8 @@ def forward(
         do_cfg_parallel, cfg_size, cfg_rank, cfg_pg = self._cfg_parallel_state(
             use_negative_prompt_cfg
         )
+        vgm = self.pipeline_config.visual_gen_mapping
+        ulysses_size = vgm.ulysses_size if vgm else 1
 
         device = self.device
         generator = torch.Generator(device=device).manual_seed(seed)
@@ -571,6 +587,11 @@ def forward(
             neg_prompt_embeds, neg_prompt_embeds_mask = self._encode_prompt(
                 negative_prompt, device, max_sequence_length
             )
+        self._validate_ulysses_prompt_masks(
+            ulysses_size,
+            prompt_embeds_mask,
+            neg_prompt_embeds_mask,
+        )
 
         # Latents.
         num_channels_latents = self.transformer.in_channels // 4  # 16
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index 0925ccb1047c..8d739e10b9a8 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -4,32 +4,31 @@
 """Multi-GPU tests for Qwen-Image-Edit Ulysses sequence parallelism."""
 
 import os
-
-os.environ["TLLM_DISABLE_MPI"] = "1"
-
-from typing import Callable
+import sys
+from collections.abc import Callable, Generator
+from pathlib import Path
 
 import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-try:
-    import sys
-    from pathlib import Path
+_ORIGINAL_TLLM_DISABLE_MPI = os.environ.get("TLLM_DISABLE_MPI")
+os.environ["TLLM_DISABLE_MPI"] = "1"
 
-    from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl
-    from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
-    from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping
-    from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import (
+try:
+    from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl  # noqa: E402
+    from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig  # noqa: E402
+    from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping  # noqa: E402
+    from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import (  # noqa: E402
         QwenImageEditPlusPipeline,
     )
-    from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (
+    from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (  # noqa: E402
         QwenJointAttention,
     )
 
     sys.path.insert(0, str(Path(__file__).resolve().parent))
-    from _visual_gen_dist_utils import spawn_with_retry
+    from _visual_gen_dist_utils import spawn_with_retry  # noqa: E402
 
     MODULES_AVAILABLE = True
 except ImportError:
@@ -37,9 +36,12 @@
 
 
 @pytest.fixture(autouse=True, scope="module")
-def _cleanup_mpi_env():
+def _cleanup_mpi_env() -> Generator[None, None, None]:
     yield
-    os.environ.pop("TLLM_DISABLE_MPI", None)
+    if _ORIGINAL_TLLM_DISABLE_MPI is None:
+        os.environ.pop("TLLM_DISABLE_MPI", None)
+    else:
+        os.environ["TLLM_DISABLE_MPI"] = _ORIGINAL_TLLM_DISABLE_MPI
 
 
 def init_distributed_worker(rank: int, world_size: int, port: int) -> None:
@@ -58,18 +60,17 @@ def cleanup_distributed() -> None:
     VisualGenMapping.seq_mesh = None
 
 
-def _distributed_worker(rank: int, world_size: int, test_fn: Callable, port: int) -> None:
+def _distributed_worker(
+    rank: int, world_size: int, test_fn: Callable[[int, int], None], port: int
+) -> None:
     try:
         init_distributed_worker(rank, world_size, port)
         test_fn(rank, world_size)
-    except Exception as e:
-        print(f"Rank {rank} failed with error: {e}")
-        raise
     finally:
         cleanup_distributed()
 
 
-def run_test_in_distributed(world_size: int, test_fn: Callable) -> None:
+def run_test_in_distributed(world_size: int, test_fn: Callable[[int, int], None]) -> None:
     if not MODULES_AVAILABLE:
         pytest.skip("Required modules not available")
     if torch.cuda.device_count() < world_size:
@@ -145,5 +146,5 @@ def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
         )
 
 
-def test_qwen_image_edit_ulysses_attention_2gpu():
+def test_qwen_image_edit_ulysses_attention_2gpu() -> None:
     run_test_in_distributed(2, _test_qwen_image_edit_ulysses_attention)
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 6f86863ab9cf..01d35c5aab8a 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
@@ -817,6 +817,17 @@ def test_qwen_image_edit_rejects_multiple_images_per_prompt() -> None:
         pipeline.infer(req)
 
 
+def test_qwen_image_ulysses_rejects_masked_prompt_conditioning() -> None:
+    from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImagePipeline
+
+    pipeline = QwenImagePipeline.__new__(QwenImagePipeline)
+
+    pipeline._validate_ulysses_prompt_masks(2, None, None)
+    pipeline._validate_ulysses_prompt_masks(1, object())
+    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
+        pipeline._validate_ulysses_prompt_masks(2, object())
+
+
 def test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning() -> None:
     from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImageEditPlusPipeline
 

From ca42121c124ba43cf3ce1a5398100b1c7a8f4d75 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 4 Aug 2026 20:52:20 +0000
Subject: [PATCH 03/18] Register Qwen Image Ulysses tests

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_b200.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index 908331efffb3..5d05e08faf6c 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -271,6 +271,10 @@ l0_b200:
   # ------------- Visual Gen tests ---------------
   - unittest/_torch/visual_gen/test_media_decode.py
   - unittest/_torch/visual_gen/test_runtime_lora.py
+  - unittest/_torch/visual_gen/test_visual_gen_args.py
+  - unittest/_torch/visual_gen/test_visual_gen_params.py
+  - unittest/_torch/visual_gen/test_visual_gen_utils.py
+  - unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
   - unittest/_torch/visual_gen/test_warmup.py
   - unittest/_torch/visual_gen/test_cache_dit.py
   - unittest/_torch/visual_gen/test_quant_ops.py

From 036c70bcee104290ebbb7928c1bd753e5c4cdb13 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Thu, 6 Aug 2026 01:28:54 +0000
Subject: [PATCH 04/18] Remove Qwen Image Ulysses prompt mask guard

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../models/qwen_image/pipeline_qwen_image.py  | 22 ------------------
 .../qwen_image/pipeline_qwen_image_edit.py    | 23 +------------------
 .../multi_gpu/test_qwen_image_edit_ulysses.py | 10 --------
 .../test_qwen_image_pipeline_config.py        | 22 ------------------
 4 files changed, 1 insertion(+), 76 deletions(-)

diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
index 8bf1758db4da..1defe920b55f 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
@@ -350,20 +350,6 @@ def _encode_prompt(
             return prompt_embeds, None
         return prompt_embeds, prompt_embeds_mask
 
-    @staticmethod
-    def _validate_ulysses_prompt_masks(
-        ulysses_size: int,
-        *prompt_masks: Optional[torch.Tensor],
-    ) -> None:
-        if ulysses_size <= 1:
-            return
-        if any(mask is not None for mask in prompt_masks):
-            raise ValueError(
-                "Qwen-Image Ulysses parallelism requires unmasked prompt "
-                "conditioning. Use prompts that produce all-valid prompt "
-                "masks, or set ulysses_size=1 for padded prompt conditioning."
-            )
-
     # ------------------------------------------------------------------
     # Latent prep / packing (identical shape to FLUX).
     # ------------------------------------------------------------------
@@ -572,9 +558,6 @@ def forward(
         do_cfg_parallel, cfg_size, cfg_rank, cfg_pg = self._cfg_parallel_state(
             use_negative_prompt_cfg
         )
-        vgm = self.pipeline_config.visual_gen_mapping
-        ulysses_size = vgm.ulysses_size if vgm else 1
-
         device = self.device
         generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -587,11 +570,6 @@ def forward(
             neg_prompt_embeds, neg_prompt_embeds_mask = self._encode_prompt(
                 negative_prompt, device, max_sequence_length
             )
-        self._validate_ulysses_prompt_masks(
-            ulysses_size,
-            prompt_embeds_mask,
-            neg_prompt_embeds_mask,
-        )
 
         # Latents.
         num_channels_latents = self.transformer.in_channels // 4  # 16
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
index 9d9f76a52e34..7cf64dabb121 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py
@@ -11,7 +11,7 @@
 import math
 import time
 from io import BytesIO
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -279,21 +279,6 @@ def _encode_edit_prompt(
             prompt_embeds_mask = None
         return prompt_embeds, prompt_embeds_mask
 
-    @staticmethod
-    def _validate_ulysses_prompt_masks(
-        ulysses_size: int,
-        *prompt_masks: Optional[torch.Tensor],
-    ) -> None:
-        if ulysses_size <= 1:
-            return
-        if any(mask is not None for mask in prompt_masks):
-            raise ValueError(
-                "Qwen-Image-Edit Ulysses parallelism requires unmasked prompt "
-                "conditioning. The supported single-prompt edit path produces "
-                "all-valid prompt masks; use ulysses_size=1 for padded or "
-                "batched prompt conditioning."
-            )
-
     def _encode_vae_image(
         self,
         image: torch.Tensor,
@@ -447,7 +432,6 @@ def forward(
         cfg_size = vgm.cfg_size if vgm else 1
         cfg_rank = vgm.cfg_rank if vgm else 0
         cfg_pg = vgm.cfg_group if vgm else None
-        ulysses_size = vgm.ulysses_size if vgm else 1
         do_cfg_parallel = do_true_cfg and cfg_size > 1
         if do_cfg_parallel:
             if cfg_size != 2:
@@ -480,11 +464,6 @@ def forward(
                 max_sequence_length,
                 num_images_per_prompt=1,
             )
-        self._validate_ulysses_prompt_masks(
-            ulysses_size,
-            prompt_embeds_mask,
-            neg_prompt_embeds_mask,
-        )
 
         num_channels_latents = self.transformer.in_channels // 4
         latents, image_latents = self._prepare_edit_latents(
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index 8d739e10b9a8..b6fa4a51a157 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -20,9 +20,6 @@
     from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl  # noqa: E402
     from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig  # noqa: E402
     from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping  # noqa: E402
-    from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import (  # noqa: E402
-        QwenImageEditPlusPipeline,
-    )
     from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (  # noqa: E402
         QwenJointAttention,
     )
@@ -138,13 +135,6 @@ def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
     assert torch.isfinite(image_out).all()
     assert torch.isfinite(text_out).all()
 
-    QwenImageEditPlusPipeline._validate_ulysses_prompt_masks(world_size, None, None)
-    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
-        QwenImageEditPlusPipeline._validate_ulysses_prompt_masks(
-            world_size,
-            torch.ones(1, 4, dtype=torch.bool, device="cuda"),
-        )
-
 
 def test_qwen_image_edit_ulysses_attention_2gpu() -> None:
     run_test_in_distributed(2, _test_qwen_image_edit_ulysses_attention)
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 01d35c5aab8a..035a00b6e7f7 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
@@ -815,25 +815,3 @@ def test_qwen_image_edit_rejects_multiple_images_per_prompt() -> None:
 
     with pytest.raises(ValueError, match="num_images_per_prompt=1 only"):
         pipeline.infer(req)
-
-
-def test_qwen_image_ulysses_rejects_masked_prompt_conditioning() -> None:
-    from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImagePipeline
-
-    pipeline = QwenImagePipeline.__new__(QwenImagePipeline)
-
-    pipeline._validate_ulysses_prompt_masks(2, None, None)
-    pipeline._validate_ulysses_prompt_masks(1, object())
-    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
-        pipeline._validate_ulysses_prompt_masks(2, object())
-
-
-def test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning() -> None:
-    from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImageEditPlusPipeline
-
-    pipeline = QwenImageEditPlusPipeline.__new__(QwenImageEditPlusPipeline)
-
-    pipeline._validate_ulysses_prompt_masks(2, None, None)
-    pipeline._validate_ulysses_prompt_masks(1, object())
-    with pytest.raises(ValueError, match="requires unmasked prompt conditioning"):
-        pipeline._validate_ulysses_prompt_masks(2, object())

From be8e741dd7e548698b49f436b39daf8d89089860 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Sat, 8 Aug 2026 03:56:42 +0000
Subject: [PATCH 05/18] Address Qwen Image Edit Ulysses review comments

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 ...qwen-image-edit-2511-fp8-2gpu-ulysses.yaml |   2 +
 .../test_lists/test-db/l0_dgx_b200.yml        |   2 +
 .../multi_gpu/test_qwen_image_edit_ulysses.py | 294 +++++++++++++-----
 3 files changed, 212 insertions(+), 86 deletions(-)

diff --git a/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml b/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
index 018bc31f8b3a..5e9edf8a10cf 100644
--- a/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
+++ b/examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml
@@ -17,6 +17,8 @@ quant_config:
   quant_algo: FP8_BLOCK_SCALES
   dynamic: true
 attention_config:
+  # Ulysses pads/shards Qwen-Image-Edit text and image streams, so the
+  # backend must support key padding masks; currently only VANILLA does.
   backend: VANILLA
 parallel_config:
   cfg_size: 1
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 63755e48bd6e..4ad81c835253 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -216,6 +216,8 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py
+  - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py
+  - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py
   - unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index b6fa4a51a157..25e96fbeb3da 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -1,7 +1,19 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-"""Multi-GPU tests for Qwen-Image-Edit Ulysses sequence parallelism."""
+"""Multi-GPU coverage for Qwen-Image-Edit Ulysses attention."""
 
 import os
 import sys
@@ -13,127 +25,237 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-_ORIGINAL_TLLM_DISABLE_MPI = os.environ.get("TLLM_DISABLE_MPI")
-os.environ["TLLM_DISABLE_MPI"] = "1"
+from tests.unittest._torch.visual_gen.conftest import spawn_with_retry
+
+# The unit test creates its own torch.distributed NCCL process group. Disable the
+# TRT-LLM MPI bootstrap path so importing tensorrt_llm does not initialize MPI.
+os.environ.setdefault("TLLM_DISABLE_MPI", "1")
+
+REPO_ROOT = Path(__file__).resolve().parents[5]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
 
 try:
-    from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl  # noqa: E402
-    from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig  # noqa: E402
-    from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping  # noqa: E402
-    from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (  # noqa: E402
+    from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl
+    from tensorrt_llm._torch.visual_gen.attention_backend.parallel import UlyssesAttention
+    from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionModelConfig
+    from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping
+    from tensorrt_llm._torch.visual_gen.models.qwen_image.transformer_qwen_image import (
         QwenJointAttention,
     )
+except ImportError as e:  # pragma: no cover - import guard for direct collection
+    pytest.skip(f"TensorRT-LLM modules unavailable: {e}", allow_module_level=True)
 
-    sys.path.insert(0, str(Path(__file__).resolve().parent))
-    from _visual_gen_dist_utils import spawn_with_retry  # noqa: E402
-
-    MODULES_AVAILABLE = True
-except ImportError:
-    MODULES_AVAILABLE = False
-
-
-@pytest.fixture(autouse=True, scope="module")
-def _cleanup_mpi_env() -> Generator[None, None, None]:
-    yield
-    if _ORIGINAL_TLLM_DISABLE_MPI is None:
-        os.environ.pop("TLLM_DISABLE_MPI", None)
-    else:
-        os.environ["TLLM_DISABLE_MPI"] = _ORIGINAL_TLLM_DISABLE_MPI
 
+def _get_free_port() -> int:
+    import socket
 
-def init_distributed_worker(rank: int, world_size: int, port: int) -> None:
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = str(port)
-    os.environ["RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    torch.cuda.set_device(rank % torch.cuda.device_count())
-    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("", 0))
+        return sock.getsockname()[1]
 
 
-def cleanup_distributed() -> None:
-    if dist.is_initialized():
+@pytest.fixture(autouse=True)
+def _cleanup_distributed_env() -> Generator[None, None, None]:
+    old_env = {
+        key: os.environ.get(key) for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")
+    }
+    yield
+    for key, value in old_env.items():
+        if value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = value
+    if dist.is_available() and dist.is_initialized():
         dist.destroy_process_group()
-    DeviceMeshTopologyImpl.device_mesh = None
-    VisualGenMapping.seq_mesh = None
 
 
 def _distributed_worker(
-    rank: int, world_size: int, test_fn: Callable[[int, int], None], port: int
+    rank: int,
+    world_size: int,
+    backend: str,
+    test_fn: Callable,
+    port: int,
+    kwargs: dict,
 ) -> None:
+    os.environ.update(
+        {
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(port),
+            "RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+        }
+    )
+    torch.cuda.set_device(rank)
+    dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
     try:
-        init_distributed_worker(rank, world_size, port)
-        test_fn(rank, world_size)
+        test_fn(rank, world_size, **kwargs)
     finally:
-        cleanup_distributed()
+        if dist.is_initialized():
+            dist.destroy_process_group()
 
 
-def run_test_in_distributed(world_size: int, test_fn: Callable[[int, int], None]) -> None:
-    if not MODULES_AVAILABLE:
-        pytest.skip("Required modules not available")
-    if torch.cuda.device_count() < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available")
+def run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs) -> None:
+    port = _get_free_port()
     spawn_with_retry(
         lambda port: mp.spawn(
             _distributed_worker,
-            args=(world_size, test_fn, port),
+            args=(world_size, "nccl", test_fn, port, kwargs),
             nprocs=world_size,
             join=True,
-        )
+        ),
+        port,
     )
 
 
-def _stabilize_attention_weights(module: torch.nn.Module) -> None:
-    """Use small deterministic weights for stable BF16 synthetic forward."""
-    with torch.no_grad():
-        for parameter in module.parameters():
-            if parameter.ndim >= 2:
-                fan_in = parameter.shape[1]
-                std = 0.02 / max(1.0, fan_in**0.5)
-                parameter.data.uniform_(-std, std)
-            else:
-                parameter.data.uniform_(-0.01, 0.01)
-
-
-def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
-    torch.manual_seed(1234)
-    vgm = VisualGenMapping(
+def _make_config(rank: int, world_size: int, ulysses_size: int) -> DiffusionModelConfig:
+    mapping = VisualGenMapping(
         world_size=world_size,
         rank=rank,
-        cfg_size=1,
         tp_size=1,
-        ring_size=1,
-        ulysses_size=world_size,
+        cp_size=1,
+        pp_size=1,
+        cfg_size=1,
+        sp_size=ulysses_size,
+        ulysses_size=ulysses_size,
+        device_mesh=DeviceMeshTopologyImpl.create().initialize(world_size),
     )
-    config = DiffusionModelConfig(
-        mapping=vgm.to_llm_mapping(),
-        visual_gen_mapping=vgm,
-        torch_dtype=torch.bfloat16,
+    return DiffusionModelConfig(
+        mapping=mapping.to_llm_mapping(),
+        visual_gen_mapping=mapping,
+        attention=AttentionConfig(backend="VANILLA"),
+        dtype="bfloat16",
     )
 
-    attn = QwenJointAttention(
+
+def _make_attention(config: DiffusionModelConfig) -> QwenJointAttention:
+    return QwenJointAttention(
         dim=16,
-        num_attention_heads=4,
-        attention_head_dim=4,
-        dtype=torch.bfloat16,
+        num_attention_heads=2,
+        attention_head_dim=8,
         config=config,
     ).cuda()
-    _stabilize_attention_weights(attn)
-    assert attn.attn.__class__.__name__ == "UlyssesAttention"
-
-    hidden_states = torch.randn(1, 4, 16, device="cuda", dtype=torch.bfloat16) * 0.1
-    encoder_hidden_states = torch.randn(1, 4, 16, device="cuda", dtype=torch.bfloat16) * 0.1
-    image_out, text_out = attn(
-        hidden_states,
-        encoder_hidden_states,
-        image_rotary_emb=None,
-        attention_mask=None,
-        timestep=None,
+
+
+def _rank_ordered_joint_mask(
+    text_mask: torch.Tensor,
+    image_seq_len: int,
+    world_size: int,
+) -> torch.Tensor:
+    image_mask = torch.ones(
+        (text_mask.shape[0], image_seq_len),
+        device=text_mask.device,
+        dtype=torch.bool,
+    )
+    text_chunks = text_mask.chunk(world_size, dim=1)
+    image_chunks = image_mask.chunk(world_size, dim=1)
+    return torch.cat(
+        [chunk for pair in zip(text_chunks, image_chunks) for chunk in pair],
+        dim=1,
+    )
+
+
+def _assert_ulysses_matches_reference(
+    rank: int,
+    world_size: int,
+    ulysses_attn: QwenJointAttention,
+    reference_attn: QwenJointAttention,
+    full_hidden_states: torch.Tensor,
+    full_encoder_hidden_states: torch.Tensor,
+    text_mask: torch.Tensor | None,
+) -> None:
+    local_hidden_states = full_hidden_states.chunk(world_size, dim=1)[rank].contiguous()
+    local_encoder_hidden_states = full_encoder_hidden_states.chunk(world_size, dim=1)[
+        rank
+    ].contiguous()
+
+    reference_attention_mask = None
+    ulysses_attention_mask = None
+    if text_mask is not None:
+        image_mask = torch.ones(
+            (text_mask.shape[0], full_hidden_states.shape[1]),
+            device=text_mask.device,
+            dtype=torch.bool,
+        )
+        reference_attention_mask = torch.cat([text_mask, image_mask], dim=1)
+        ulysses_attention_mask = _rank_ordered_joint_mask(
+            text_mask,
+            full_hidden_states.shape[1],
+            world_size,
+        )
+
+    with torch.no_grad():
+        image_out, text_out = ulysses_attn(
+            hidden_states=local_hidden_states,
+            encoder_hidden_states=local_encoder_hidden_states,
+            image_rotary_emb=None,
+            attention_mask=ulysses_attention_mask,
+            timestep=None,
+        )
+        expected_image_out, expected_text_out = reference_attn(
+            hidden_states=full_hidden_states,
+            encoder_hidden_states=full_encoder_hidden_states,
+            image_rotary_emb=None,
+            attention_mask=reference_attention_mask,
+            timestep=None,
+        )
+
+    expected_image_out = expected_image_out.chunk(world_size, dim=1)[rank].contiguous()
+    expected_text_out = expected_text_out.chunk(world_size, dim=1)[rank].contiguous()
+
+    torch.testing.assert_close(image_out, expected_image_out, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(text_out, expected_text_out, rtol=2e-2, atol=2e-2)
+
+
+def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
+    torch.manual_seed(2026)
+    ulysses_attn = _make_attention(_make_config(rank, world_size, world_size))
+    reference_attn = _make_attention(_make_config(rank=0, world_size=1, ulysses_size=1))
+
+    with torch.no_grad():
+        for name, parameter in ulysses_attn.named_parameters():
+            if name.endswith("bias"):
+                parameter.zero_()
+            elif "norm" in name and name.endswith("weight"):
+                parameter.fill_(1)
+            else:
+                parameter.normal_(mean=0.0, std=0.02)
+    reference_attn.load_state_dict(ulysses_attn.state_dict())
+
+    assert isinstance(ulysses_attn.attn, UlyssesAttention)
+
+    full_hidden_states = torch.randn(
+        1,
+        4,
+        16,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    full_encoder_hidden_states = torch.randn(
+        1,
+        4,
+        16,
+        device="cuda",
+        dtype=torch.bfloat16,
     )
+    masked_prompt = torch.tensor(
+        [[True, True, True, False]],
+        device="cuda",
+        dtype=torch.bool,
+    )
+
+    for text_mask in (None, masked_prompt):
+        _assert_ulysses_matches_reference(
+            rank,
+            world_size,
+            ulysses_attn,
+            reference_attn,
+            full_hidden_states,
+            full_encoder_hidden_states,
+            text_mask,
+        )
 
-    assert image_out.shape == hidden_states.shape
-    assert text_out.shape == encoder_hidden_states.shape
-    assert torch.isfinite(image_out).all()
-    assert torch.isfinite(text_out).all()
+    dist.barrier()
 
 
 def test_qwen_image_edit_ulysses_attention_2gpu() -> None:

From 7bf0038586e7cc9a59fb7e04179b1049403b6828 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Mon, 10 Aug 2026 18:51:36 +0000
Subject: [PATCH 06/18] Fix Qwen Image multi-GPU unit test collection

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../test_qwen_image_attention_parallel.py         |  6 +++---
 .../multi_gpu/test_qwen_image_edit_ulysses.py     | 15 +++------------
 2 files changed, 6 insertions(+), 15 deletions(-)

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 850f5b6089b8..2424932f0fa8 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
@@ -170,13 +170,13 @@ def _test_qwen_image_attention_parallel_topology(
 @pytest.mark.parametrize(
     "world_size,parallel,backend,topology",
     [
-        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=pytest.mark.gpu2, id="tp2"),
+        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=[pytest.mark.gpu2], id="tp2"),
         pytest.param(
             4,
             {"ring_size": 2, "ulysses_size": 2},
             "FA4",
             "ring",
-            marks=pytest.mark.gpu4,
+            marks=[pytest.mark.gpu4],
             id="ring2_ulysses2",
         ),
         pytest.param(
@@ -184,7 +184,7 @@ def _test_qwen_image_attention_parallel_topology(
             {"attn2d_size": (2, 1), "ulysses_size": 2},
             "FA4",
             "attn2d",
-            marks=pytest.mark.gpu4,
+            marks=[pytest.mark.gpu4],
             id="attn2d_2x1_ulysses2",
         ),
     ],
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index 25e96fbeb3da..8b1c3f1c23e3 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -25,7 +25,8 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from tests.unittest._torch.visual_gen.conftest import spawn_with_retry
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from _visual_gen_dist_utils import spawn_with_retry
 
 # The unit test creates its own torch.distributed NCCL process group. Disable the
 # TRT-LLM MPI bootstrap path so importing tensorrt_llm does not initialize MPI.
@@ -47,14 +48,6 @@
     pytest.skip(f"TensorRT-LLM modules unavailable: {e}", allow_module_level=True)
 
 
-def _get_free_port() -> int:
-    import socket
-
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind(("", 0))
-        return sock.getsockname()[1]
-
-
 @pytest.fixture(autouse=True)
 def _cleanup_distributed_env() -> Generator[None, None, None]:
     old_env = {
@@ -96,15 +89,13 @@ def _distributed_worker(
 
 
 def run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs) -> None:
-    port = _get_free_port()
     spawn_with_retry(
         lambda port: mp.spawn(
             _distributed_worker,
             args=(world_size, "nccl", test_fn, port, kwargs),
             nprocs=world_size,
             join=True,
-        ),
-        port,
+        )
     )
 
 

From ee5152aad209a9658c62f2906a66e2c59edc1f66 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Fri, 14 Aug 2026 21:49:50 +0000
Subject: [PATCH 07/18] Optimize Qwen Image Ulysses post-unscatter

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../kernels/ulyssesPostUnscatterKernel.cu     | 56 +++++++++++++++
 .../kernels/ulyssesPostUnscatterKernel.h      |  6 ++
 .../thop/ulyssesPostUnscatterOp.cpp           | 43 +++++++++++
 .../_torch/custom_ops/cpp_custom_ops.py       | 11 +++
 .../visual_gen/attention_backend/parallel.py  | 71 ++++++++++++++++++-
 5 files changed, 185 insertions(+), 2 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
index 2d0714a890ec..d259f5026a9a 100644
--- a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
+++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
@@ -95,6 +95,39 @@ __global__ void ulyssesPostUnscatterKernel(T const* __restrict__ q_in, T const*
     *out_v4 = *in_v4;
 }
 
+template <typename T>
+__global__ void ulyssesPackedQkvPostUnscatterKernel(T const* __restrict__ qkv_in, T* __restrict__ q_out,
+    T* __restrict__ k_out, T* __restrict__ v_out, int const P, int const B, int const Sp, int const H, int const D,
+    int const vec_per_row)
+{
+    constexpr int VEC = 8;
+
+    int const PSp = P * Sp;
+    int const bx = blockIdx.x;
+    int const qkv_idx = bx / PSp;
+    int const psp = bx - qkv_idx * PSp;
+
+    int const h = threadIdx.x / vec_per_row;
+    int const vec_idx = threadIdx.x - h * vec_per_row;
+
+    int const p = psp / Sp;
+    int const sp = psp - p * Sp;
+    int const b = blockIdx.y;
+
+    // qkv_in[p, b, sp, qkv, h, d]:
+    // (((((p*B + b)*Sp + sp)*3 + qkv)*H + h)*D + vec_idx*VEC)
+    // out[b, p*Sp+sp, h, d]:
+    // (((b*PSp + psp)*H + h)*D + vec_idx*VEC)
+    int64_t const in_base = (((((static_cast<int64_t>(p) * B + b) * Sp + sp) * 3 + qkv_idx) * H + h) * D)
+        + vec_idx * VEC;
+    int64_t const out_base = (((static_cast<int64_t>(b) * PSp + psp) * H + h) * D) + vec_idx * VEC;
+
+    T* out_ptr = qkv_idx == 0 ? q_out : (qkv_idx == 1 ? k_out : v_out);
+    uint4 const* in_v4 = reinterpret_cast<uint4 const*>(qkv_in + in_base);
+    uint4* out_v4 = reinterpret_cast<uint4*>(out_ptr + out_base);
+    *out_v4 = *in_v4;
+}
+
 } // namespace
 
 void launchUlyssesPostUnscatter(void const* q_in, void const* k_in, void const* v_in, void* q_out, void* k_out,
@@ -123,6 +156,29 @@ void launchUlyssesPostUnscatter(void const* q_in, void const* k_in, void const*
         q_out_typed, k_out_typed, v_out_typed, P, B, D, vec_per_row, Sp_q, H_q, Sp_k, H_k, Sp_v, H_v);
 }
 
+void launchUlyssesPackedQkvPostUnscatter(
+    void const* qkv_in, void* q_out, void* k_out, void* v_out, int P, int B, int Sp, int H, int D, cudaStream_t stream)
+{
+    constexpr int VEC = 8;
+    TLLM_CHECK_WITH_INFO(D % VEC == 0, "ulyssesPackedQkvPostUnscatter: D must be a multiple of 8, got %d", D);
+    int const vec_per_row = D / VEC;
+    int const threads = H * vec_per_row;
+    TLLM_CHECK_WITH_INFO(threads <= 1024,
+        "ulyssesPackedQkvPostUnscatter: threads/block (H*D/8) must be <= 1024, got H=%d D=%d -> %d", H, D,
+        threads);
+
+    dim3 const grid(3 * P * Sp, B, 1);
+    dim3 const block(threads);
+
+    auto* qkv_in_typed = reinterpret_cast<__nv_bfloat16 const*>(qkv_in);
+    auto* q_out_typed = reinterpret_cast<__nv_bfloat16*>(q_out);
+    auto* k_out_typed = reinterpret_cast<__nv_bfloat16*>(k_out);
+    auto* v_out_typed = reinterpret_cast<__nv_bfloat16*>(v_out);
+
+    ulyssesPackedQkvPostUnscatterKernel<__nv_bfloat16><<<grid, block, 0, stream>>>(
+        qkv_in_typed, q_out_typed, k_out_typed, v_out_typed, P, B, Sp, H, D, vec_per_row);
+}
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
index d7ff250c37dd..bc66eb5be280 100644
--- a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
+++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
@@ -60,6 +60,12 @@ void launchUlyssesPostUnscatter(void const* q_in, // [P, B, Sp_q, H_q, D]
     void* k_out, void* v_out, int P, int B, int D, int Sp_q, int H_q, int Sp_k, int H_k, int Sp_v, int H_v,
     cudaStream_t stream);
 
+// Packed-QKV variant for self-attention. Consumes the raw packed all-to-all
+// receive buffer [P, B, Sp, 3, H, D] and writes per-tensor NHD-contig outputs
+// [B, P*Sp, H, D].
+void launchUlyssesPackedQkvPostUnscatter(void const* qkv_in, void* q_out, void* k_out, void* v_out, int P, int B,
+    int Sp, int H, int D, cudaStream_t stream);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END
diff --git a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
index 21837f833493..b61ae32fe867 100644
--- a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
@@ -96,17 +96,60 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_post_unscatter_q
     return std::make_tuple(q_out, k_out, v_out);
 }
 
+// Self-attention-only variant that consumes the raw packed 5D all-to-all receive buffer
+// [P, B, Sp, 3, H, D]. This removes the Python-side packed-QKV permute/contiguous/unbind
+// chain before VANILLA/HND SDPA.
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_packed_qkv_post_unscatter(
+    torch::Tensor& qkv_in, int64_t layout)
+{
+    TORCH_CHECK(qkv_in.dim() == 6, "ulysses_packed_qkv_post_unscatter expects [P, B, Sp, 3, H, D]");
+    TORCH_CHECK(qkv_in.size(3) == 3, "ulysses_packed_qkv_post_unscatter expects qkv dim size 3, got ",
+        qkv_in.size(3));
+    TORCH_CHECK(layout == 0 || layout == 1, "layout must be 0 (HND) or 1 (NHD), got ", layout);
+
+    CHECK_INPUT(qkv_in, torch::kBFloat16);
+
+    int64_t const P = qkv_in.size(0);
+    int64_t const B = qkv_in.size(1);
+    int64_t const Sp = qkv_in.size(2);
+    int64_t const H = qkv_in.size(4);
+    int64_t const D = qkv_in.size(5);
+    TORCH_CHECK(D % 8 == 0, "D (last dim) must be divisible by 8 (bf16 vec=8)");
+
+    bool const is_hnd = (layout == 0);
+    auto opts = qkv_in.options();
+    auto q_out = torch::empty({B, P * Sp, H, D}, opts);
+    auto k_out = torch::empty({B, P * Sp, H, D}, opts);
+    auto v_out = torch::empty({B, P * Sp, H, D}, opts);
+
+    if (qkv_in.numel() != 0)
+    {
+        auto stream = at::cuda::getCurrentCUDAStream();
+        tensorrt_llm::kernels::launchUlyssesPackedQkvPostUnscatter(qkv_in.data_ptr(), q_out.data_ptr(),
+            k_out.data_ptr(), v_out.data_ptr(), static_cast<int>(P), static_cast<int>(B), static_cast<int>(Sp),
+            static_cast<int>(H), static_cast<int>(D), stream);
+    }
+
+    if (is_hnd)
+    {
+        return std::make_tuple(q_out.transpose(1, 2), k_out.transpose(1, 2), v_out.transpose(1, 2));
+    }
+    return std::make_tuple(q_out, k_out, v_out);
+}
+
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     // layout: 0 = HND [B, H, P*Sp, D], 1 = NHD [B, P*Sp, H, D]. Default 0 keeps
     // backward compatibility with the original HND-only callers.
     m.def(
         "ulysses_post_unscatter_qkv(Tensor q_in, Tensor k_in, Tensor v_in, int layout=0) -> (Tensor, Tensor, Tensor)");
+    m.def("ulysses_packed_qkv_post_unscatter(Tensor qkv_in, int layout=0) -> (Tensor, Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("ulysses_post_unscatter_qkv", &ulysses_post_unscatter_qkv);
+    m.impl("ulysses_packed_qkv_post_unscatter", &ulysses_packed_qkv_post_unscatter);
 }
 
 } // namespace torch_ext
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index f5086beecbf5..4da5bce7468a 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1258,6 +1258,17 @@ def _mk(t):
 
         return (_mk(q_in), _mk(k_in), _mk(v_in))
 
+    @torch.library.register_fake("trtllm::ulysses_packed_qkv_post_unscatter")
+    def _(qkv_in, layout=0):
+        P, B, Sp, qkv_count, H, D = qkv_in.shape
+        assert qkv_count == 3
+
+        def _mk():
+            base = qkv_in.new_empty((B, P * Sp, H, D))
+            return base.transpose(1, 2) if layout == 0 else base
+
+        return (_mk(), _mk(), _mk())
+
     @torch.library.register_fake("trtllm::helix_post_process")
     def _(gathered_o, gathered_stats, scale):
         return gathered_o.new_empty(*gathered_o.shape[1:])
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index fedabf4ceba8..5183246896b5 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -61,6 +61,29 @@ def _ulysses_post_unscatter(q_5d, k_5d, v_5d, *, is_hnd):
     return torch.ops.trtllm.ulysses_post_unscatter_qkv(q_5d, k_5d, v_5d, layout)
 
 
+def _all_to_all_5d_raw(
+    input: torch.Tensor,
+    process_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> torch.Tensor:
+    """Packed QKV A2A without the final Python unscatter.
+
+    Input is [B, S/P, 3, H, D]. Output is the raw receive buffer
+    [P, B, S/P, 3, H/P, D], ready for the native packed post-unscatter op.
+    """
+    world_size = torch.distributed.get_world_size(group=process_group)
+    if world_size == 1:
+        return input.unsqueeze(0)
+
+    batch, seq, qkv_count, heads, head_dim = input.shape
+    sharded_heads = heads // world_size
+    inp = input.reshape(batch, seq, qkv_count, world_size, sharded_heads, head_dim)
+    inp = inp.permute(3, 0, 1, 2, 4, 5).contiguous()
+
+    out_flat = torch.empty_like(inp.flatten())
+    torch.distributed.all_to_all_single(out_flat, inp.flatten(), group=process_group)
+    return out_flat.view_as(inp)
+
+
 class UlyssesAttention(AttentionBackend):
     """
     Ulysses Sequence Parallelism wrapper.
@@ -156,8 +179,30 @@ def forward(
             )
 
         if self.inner_backend.support_fused_qkv():
-            return self._forward_fused(q, k, v, **kwargs)
-        return self._forward_unfused(q, k, v, **kwargs)
+            out = self._forward_fused(q, k, v, **kwargs)
+        elif self._supports_packed_self_attention(q, k, v, kwargs):
+            out = self._forward_packed_self_attention(q, k, v, **kwargs)
+        else:
+            out = self._forward_unfused(q, k, v, **kwargs)
+        return out
+
+    def _supports_packed_self_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        kwargs: Dict,
+    ) -> bool:
+        if self.inner_backend.preferred_layout not in (
+            AttentionTensorLayout.HND,
+            AttentionTensorLayout.NHD,
+        ):
+            return False
+        if q.shape != k.shape or q.shape != v.shape:
+            return False
+        if q.dtype != torch.bfloat16 or q.shape[-1] % 8 != 0:
+            return False
+        return kwargs.get("gate_compress") is None and kwargs.get("gate_fine") is None
 
     def _forward_fused(
         self,
@@ -198,6 +243,28 @@ def _forward_fused(
 
         return self._output_a2a(output, batch_size, seq_len)
 
+    def _forward_packed_self_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        batch_size = q.shape[0]
+        qkv = torch.stack([q, k, v], dim=2)
+        qkv_5d = _all_to_all_5d_raw(qkv, self.process_group)
+        is_hnd = self.inner_backend.preferred_layout == AttentionTensorLayout.HND
+        q, k, v = torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv_5d, 0 if is_hnd else 1)
+
+        seq_len = q.shape[2] if is_hnd else q.shape[1]
+        kwargs["batch_size"] = batch_size
+        kwargs["seq_len"] = seq_len
+        kwargs["seq_len_kv"] = seq_len
+
+        output = self.inner_backend.forward(q=q, k=k, v=v, **kwargs)
+
+        return self._output_a2a(output, batch_size, seq_len)
+
     def _forward_unfused(
         self,
         q: torch.Tensor,

From d9f81295cd47708ec6d87fe83474b8dc28139a8b Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Wed, 19 Aug 2026 16:52:11 +0000
Subject: [PATCH 08/18] update post_scatter op score

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../_torch/visual_gen/attention_backend/parallel.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index 5183246896b5..93b394e9ec8c 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -20,6 +20,7 @@
 
 """
 
+import os
 from typing import TYPE_CHECKING, Callable, ClassVar, Dict, Optional
 
 import torch
@@ -51,6 +52,10 @@ def post_permute_5d_to_4d(out_5d, P):
     return out_5d.permute(1, 0, 2, 3, 4).contiguous().view(Bt, _P * Spt, HpP, Dt)
 
 
+def _disable_ulysses_post_unscatter_op() -> bool:
+    return os.environ.get("TRTLLM_DISABLE_ULYSSES_POST_UNSCATTER_OP", "0") == "1"
+
+
 def _ulysses_post_unscatter(q_5d, k_5d, v_5d, *, is_hnd):
     """One-launch fused replacement for the post-A2A 5D -> 4D chain.
 
@@ -193,6 +198,8 @@ def _supports_packed_self_attention(
         v: torch.Tensor,
         kwargs: Dict,
     ) -> bool:
+        if self.world_size <= 1 or _disable_ulysses_post_unscatter_op():
+            return False
         if self.inner_backend.preferred_layout not in (
             AttentionTensorLayout.HND,
             AttentionTensorLayout.NHD,
@@ -411,7 +418,11 @@ def forward_async(
         # only instantiated for __nv_bfloat16.
         _, B_q, Sp_q, HpP_q, D_q = q_5d.shape
         is_hnd = self.inner_backend.preferred_layout == AttentionTensorLayout.HND
-        use_fused_post_unscatter = q_5d.dtype == torch.bfloat16
+        use_fused_post_unscatter = (
+            self.world_size > 1
+            and q_5d.dtype == torch.bfloat16
+            and not _disable_ulysses_post_unscatter_op()
+        )
         if use_fused_post_unscatter:
             q_out, k_out, v_out = _ulysses_post_unscatter(q_5d, k_5d, v_5d, is_hnd=is_hnd)
             B = B_q

From 73ecf986b88d61dfa421cfb88c84e073f520518c Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Wed, 19 Aug 2026 17:20:53 +0000
Subject: [PATCH 09/18] Address Qwen Image Ulysses review findings

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../thop/ulyssesPostUnscatterOp.cpp           | 44 +++++++++++++
 .../visual_gen/attention_backend/parallel.py  | 15 ++---
 .../models/qwen_image/pipeline_qwen_image.py  |  1 +
 .../test_lists/test-db/l0_b200.yml            |  4 --
 .../test_lists/test-db/l0_dgx_b200.yml        |  1 -
 .../test_ulysses_post_unscatter.py            | 64 +++++++++++++++++++
 .../test_qwen_image_attention_parallel.py     |  6 +-
 .../multi_gpu/test_qwen_image_edit_ulysses.py | 13 ++++
 8 files changed, 129 insertions(+), 19 deletions(-)

diff --git a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
index b61ae32fe867..fa11e6c67d21 100644
--- a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
@@ -18,6 +18,7 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/CUDAContext.h>
+#include <limits>
 #include <torch/extension.h>
 
 TRTLLM_NAMESPACE_BEGIN
@@ -25,6 +26,29 @@ TRTLLM_NAMESPACE_BEGIN
 namespace torch_ext
 {
 
+namespace
+{
+
+void checkInt32Dim(char const* name, int64_t value)
+{
+    TORCH_CHECK(value <= std::numeric_limits<int>::max(), name, " must fit in int32, got ", value);
+}
+
+void checkInt32Product(char const* name, int64_t lhs, int64_t rhs)
+{
+    TORCH_CHECK(lhs == 0 || rhs <= std::numeric_limits<int>::max() / lhs, name, " must fit in int32, got ", lhs, "*",
+        rhs);
+}
+
+void checkInt32SumProduct(char const* name, int64_t lhs, int64_t a, int64_t b, int64_t c)
+{
+    TORCH_CHECK(a <= std::numeric_limits<int64_t>::max() - b && a + b <= std::numeric_limits<int64_t>::max() - c,
+        name, " sum overflows int64");
+    checkInt32Product(name, lhs, a + b + c);
+}
+
+} // namespace
+
 // Post-Ulysses A2A unscatter: take Q/K/V tensors of shape [P, B, Sp, H, D]
 // (output of the head-dim -> seq-dim all-to-all) and produce SDPA-ready Q/K/V.
 // The kernel writes NHD-contig storage [B, P*Sp, H, D]; the return depends on ``layout``:
@@ -61,6 +85,19 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_post_unscatter_q
     int64_t const Sp_q = q_in.size(2), H_q = q_in.size(3);
     int64_t const Sp_k = k_in.size(2), H_k = k_in.size(3);
     int64_t const Sp_v = v_in.size(2), H_v = v_in.size(3);
+    checkInt32Dim("P", P);
+    checkInt32Dim("B", B);
+    checkInt32Dim("D", D);
+    checkInt32Dim("Sp_q", Sp_q);
+    checkInt32Dim("H_q", H_q);
+    checkInt32Dim("Sp_k", Sp_k);
+    checkInt32Dim("H_k", H_k);
+    checkInt32Dim("Sp_v", Sp_v);
+    checkInt32Dim("H_v", H_v);
+    checkInt32Product("P*Sp_q", P, Sp_q);
+    checkInt32Product("P*Sp_k", P, Sp_k);
+    checkInt32Product("P*Sp_v", P, Sp_v);
+    checkInt32SumProduct("P*(Sp_q+Sp_k+Sp_v)", P, Sp_q, Sp_k, Sp_v);
 
     bool const is_hnd = (layout == 0);
     auto opts = q_in.options();
@@ -115,6 +152,13 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_packed_qkv_post_
     int64_t const H = qkv_in.size(4);
     int64_t const D = qkv_in.size(5);
     TORCH_CHECK(D % 8 == 0, "D (last dim) must be divisible by 8 (bf16 vec=8)");
+    checkInt32Dim("P", P);
+    checkInt32Dim("B", B);
+    checkInt32Dim("Sp", Sp);
+    checkInt32Dim("H", H);
+    checkInt32Dim("D", D);
+    checkInt32Product("P*Sp", P, Sp);
+    checkInt32Product("3*P*Sp", 3, P * Sp);
 
     bool const is_hnd = (layout == 0);
     auto opts = qkv_in.options();
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index 93b394e9ec8c..7da73583f798 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -20,7 +20,6 @@
 
 """
 
-import os
 from typing import TYPE_CHECKING, Callable, ClassVar, Dict, Optional
 
 import torch
@@ -52,10 +51,6 @@ def post_permute_5d_to_4d(out_5d, P):
     return out_5d.permute(1, 0, 2, 3, 4).contiguous().view(Bt, _P * Spt, HpP, Dt)
 
 
-def _disable_ulysses_post_unscatter_op() -> bool:
-    return os.environ.get("TRTLLM_DISABLE_ULYSSES_POST_UNSCATTER_OP", "0") == "1"
-
-
 def _ulysses_post_unscatter(q_5d, k_5d, v_5d, *, is_hnd):
     """One-launch fused replacement for the post-A2A 5D -> 4D chain.
 
@@ -198,7 +193,9 @@ def _supports_packed_self_attention(
         v: torch.Tensor,
         kwargs: Dict,
     ) -> bool:
-        if self.world_size <= 1 or _disable_ulysses_post_unscatter_op():
+        # self.world_size is the Ulysses process-group size here, not the
+        # global distributed world size.
+        if self.world_size <= 1:
             return False
         if self.inner_backend.preferred_layout not in (
             AttentionTensorLayout.HND,
@@ -418,11 +415,7 @@ def forward_async(
         # only instantiated for __nv_bfloat16.
         _, B_q, Sp_q, HpP_q, D_q = q_5d.shape
         is_hnd = self.inner_backend.preferred_layout == AttentionTensorLayout.HND
-        use_fused_post_unscatter = (
-            self.world_size > 1
-            and q_5d.dtype == torch.bfloat16
-            and not _disable_ulysses_post_unscatter_op()
-        )
+        use_fused_post_unscatter = self.world_size > 1 and q_5d.dtype == torch.bfloat16
         if use_fused_post_unscatter:
             q_out, k_out, v_out = _ulysses_post_unscatter(q_5d, k_5d, v_5d, is_hnd=is_hnd)
             B = B_q
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
index 1defe920b55f..b7d3a69015cf 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
@@ -558,6 +558,7 @@ def forward(
         do_cfg_parallel, cfg_size, cfg_rank, cfg_pg = self._cfg_parallel_state(
             use_negative_prompt_cfg
         )
+
         device = self.device
         generator = torch.Generator(device=device).manual_seed(seed)
 
diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index 5d05e08faf6c..908331efffb3 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -271,10 +271,6 @@ l0_b200:
   # ------------- Visual Gen tests ---------------
   - unittest/_torch/visual_gen/test_media_decode.py
   - unittest/_torch/visual_gen/test_runtime_lora.py
-  - unittest/_torch/visual_gen/test_visual_gen_args.py
-  - unittest/_torch/visual_gen/test_visual_gen_params.py
-  - unittest/_torch/visual_gen/test_visual_gen_utils.py
-  - unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
   - unittest/_torch/visual_gen/test_warmup.py
   - unittest/_torch/visual_gen/test_cache_dit.py
   - unittest/_torch/visual_gen/test_quant_ops.py
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 4ad81c835253..52b246423e6a 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -216,7 +216,6 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py
-  - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py
   - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py
   - unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py
diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.py
index c96b25b6cd9a..5daee53afcb0 100644
--- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.py
+++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.py
@@ -28,6 +28,11 @@ def post(t):
     return post(q_5d), post(k_5d), post(v_5d)
 
 
+def torch_ref_packed(qkv_6d, is_hnd):
+    q, k, v = qkv_6d.unbind(dim=3)
+    return torch_ref(q.contiguous(), k.contiguous(), v.contiguous(), is_hnd)
+
+
 @pytest.mark.parametrize("layout", [0, 1], ids=["HND", "NHD"])
 @pytest.mark.parametrize(
     "P,B,Sp,H,D",
@@ -118,6 +123,37 @@ def test_ulysses_post_unscatter_cross_attn_varshape(P, B, D, Sp_q, H_q, Sp_kv, H
         assert max_diff == 0, f"{name}: max_diff={max_diff} (expected exact match)"
 
 
+@pytest.mark.parametrize("layout", [0, 1], ids=["HND", "NHD"])
+@pytest.mark.parametrize(
+    "P,B,Sp,H,D",
+    [
+        (2, 1, 128, 8, 128),
+        (4, 2, 256, 16, 64),
+        (8, 1, 128, 4, 128),
+    ],
+)
+@torch.inference_mode()
+def test_ulysses_packed_qkv_post_unscatter_exact_match(P, B, Sp, H, D, layout):
+    """Packed self-attention op must exactly match unbind + eager unscatter."""
+    is_hnd = layout == 0
+    torch.manual_seed(0)
+    qkv = torch.randn(P, B, Sp, 3, H, D, device="cuda", dtype=torch.bfloat16).contiguous()
+
+    q_ref, k_ref, v_ref = torch_ref_packed(qkv, is_hnd=is_hnd)
+    q_out, k_out, v_out = torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv, layout)
+
+    expected_shape = (B, H, P * Sp, D) if is_hnd else (B, P * Sp, H, D)
+    assert q_out.shape == expected_shape
+    if is_hnd:
+        assert not q_out.is_contiguous() and not k_out.is_contiguous() and not v_out.is_contiguous()
+    else:
+        assert q_out.is_contiguous() and k_out.is_contiguous() and v_out.is_contiguous()
+    assert q_out.dtype == torch.bfloat16
+    for name, ref, got in [("Q", q_ref, q_out), ("K", k_ref, k_out), ("V", v_ref, v_out)]:
+        max_diff = (ref - got).abs().max().item()
+        assert max_diff == 0, f"{name}: max_diff={max_diff} (expected exact match)"
+
+
 @torch.inference_mode()
 def test_ulysses_post_unscatter_rejects_invalid_layout():
     """layout must be 0 (HND) or 1 (NHD)."""
@@ -126,6 +162,20 @@ def test_ulysses_post_unscatter_rejects_invalid_layout():
         torch.ops.trtllm.ulysses_post_unscatter_qkv(q, q, q, 2)
 
 
+@torch.inference_mode()
+def test_ulysses_packed_qkv_post_unscatter_rejects_invalid_layout():
+    qkv = torch.randn(2, 1, 128, 3, 8, 64, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv, 2)
+
+
+@torch.inference_mode()
+def test_ulysses_packed_qkv_post_unscatter_rejects_invalid_qkv_dim():
+    qkv = torch.randn(2, 1, 128, 2, 8, 64, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv)
+
+
 @torch.inference_mode()
 def test_ulysses_post_unscatter_rejects_d_not_multiple_of_8():
     """D must be a multiple of 8 (uint4 vec load constraint)."""
@@ -134,6 +184,13 @@ def test_ulysses_post_unscatter_rejects_d_not_multiple_of_8():
         torch.ops.trtllm.ulysses_post_unscatter_qkv(q, q, q)
 
 
+@torch.inference_mode()
+def test_ulysses_packed_qkv_post_unscatter_rejects_d_not_multiple_of_8():
+    qkv = torch.randn(2, 1, 128, 3, 8, 60, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv)
+
+
 @torch.inference_mode()
 def test_ulysses_post_unscatter_rejects_oversized_block():
     """Threads/block = H * (D/8) must be <= 1024 (CUDA hw limit)."""
@@ -141,3 +198,10 @@ def test_ulysses_post_unscatter_rejects_oversized_block():
     q = torch.randn(2, 1, 64, 128, 128, device="cuda", dtype=torch.bfloat16)
     with pytest.raises(RuntimeError):
         torch.ops.trtllm.ulysses_post_unscatter_qkv(q, q, q)
+
+
+@torch.inference_mode()
+def test_ulysses_packed_qkv_post_unscatter_rejects_oversized_block():
+    qkv = torch.randn(2, 1, 64, 3, 128, 128, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError):
+        torch.ops.trtllm.ulysses_packed_qkv_post_unscatter(qkv)
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 2424932f0fa8..850f5b6089b8 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
@@ -170,13 +170,13 @@ def _test_qwen_image_attention_parallel_topology(
 @pytest.mark.parametrize(
     "world_size,parallel,backend,topology",
     [
-        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=[pytest.mark.gpu2], id="tp2"),
+        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=pytest.mark.gpu2, id="tp2"),
         pytest.param(
             4,
             {"ring_size": 2, "ulysses_size": 2},
             "FA4",
             "ring",
-            marks=[pytest.mark.gpu4],
+            marks=pytest.mark.gpu4,
             id="ring2_ulysses2",
         ),
         pytest.param(
@@ -184,7 +184,7 @@ def _test_qwen_image_attention_parallel_topology(
             {"attn2d_size": (2, 1), "ulysses_size": 2},
             "FA4",
             "attn2d",
-            marks=[pytest.mark.gpu4],
+            marks=pytest.mark.gpu4,
             id="attn2d_2x1_ulysses2",
         ),
     ],
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index 8b1c3f1c23e3..0855da88edf3 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -30,6 +30,7 @@
 
 # The unit test creates its own torch.distributed NCCL process group. Disable the
 # TRT-LLM MPI bootstrap path so importing tensorrt_llm does not initialize MPI.
+_OLD_TLLM_DISABLE_MPI = os.environ.get("TLLM_DISABLE_MPI")
 os.environ.setdefault("TLLM_DISABLE_MPI", "1")
 
 REPO_ROOT = Path(__file__).resolve().parents[5]
@@ -63,6 +64,15 @@ def _cleanup_distributed_env() -> Generator[None, None, None]:
         dist.destroy_process_group()
 
 
+@pytest.fixture(autouse=True, scope="module")
+def _cleanup_mpi_env() -> Generator[None, None, None]:
+    yield
+    if _OLD_TLLM_DISABLE_MPI is None:
+        os.environ.pop("TLLM_DISABLE_MPI", None)
+    else:
+        os.environ["TLLM_DISABLE_MPI"] = _OLD_TLLM_DISABLE_MPI
+
+
 def _distributed_worker(
     rank: int,
     world_size: int,
@@ -89,6 +99,8 @@ def _distributed_worker(
 
 
 def run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs) -> None:
+    if torch.cuda.device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available")
     spawn_with_retry(
         lambda port: mp.spawn(
             _distributed_worker,
@@ -249,5 +261,6 @@ def _test_qwen_image_edit_ulysses_attention(rank: int, world_size: int) -> None:
     dist.barrier()
 
 
+@pytest.mark.gpu2
 def test_qwen_image_edit_ulysses_attention_2gpu() -> None:
     run_test_in_distributed(2, _test_qwen_image_edit_ulysses_attention)

From b87a2e8a09ed5700418311c4a4e5299a9b0fbec5 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Fri, 21 Aug 2026 03:33:25 +0000
Subject: [PATCH 10/18] Apply clang-format to Ulysses post-unscatter op

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../kernels/ulyssesPostUnscatterKernel.cu     |  11 +-
 .../kernels/ulyssesPostUnscatterKernel.h      |   4 +-
 .../thop/ulyssesPostUnscatterOp.cpp           |  11 +-
 .../_torch/custom_ops/torch_custom_ops.py     |  83 ++++++
 tensorrt_llm/_torch/modules/linear.py         | 135 +++++++---
 tensorrt_llm/_torch/utils.py                  |  26 ++
 .../attention_backend/flash_attn4.py          |   2 +-
 .../visual_gen/attention_backend/parallel.py  |   2 +-
 .../qwen_image/transformer_qwen_image.py      | 241 +++++++++++++++++-
 9 files changed, 453 insertions(+), 62 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
index d259f5026a9a..f5eb084b3afd 100644
--- a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
+++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu
@@ -118,8 +118,8 @@ __global__ void ulyssesPackedQkvPostUnscatterKernel(T const* __restrict__ qkv_in
     // (((((p*B + b)*Sp + sp)*3 + qkv)*H + h)*D + vec_idx*VEC)
     // out[b, p*Sp+sp, h, d]:
     // (((b*PSp + psp)*H + h)*D + vec_idx*VEC)
-    int64_t const in_base = (((((static_cast<int64_t>(p) * B + b) * Sp + sp) * 3 + qkv_idx) * H + h) * D)
-        + vec_idx * VEC;
+    int64_t const in_base
+        = (((((static_cast<int64_t>(p) * B + b) * Sp + sp) * 3 + qkv_idx) * H + h) * D) + vec_idx * VEC;
     int64_t const out_base = (((static_cast<int64_t>(b) * PSp + psp) * H + h) * D) + vec_idx * VEC;
 
     T* out_ptr = qkv_idx == 0 ? q_out : (qkv_idx == 1 ? k_out : v_out);
@@ -164,8 +164,7 @@ void launchUlyssesPackedQkvPostUnscatter(
     int const vec_per_row = D / VEC;
     int const threads = H * vec_per_row;
     TLLM_CHECK_WITH_INFO(threads <= 1024,
-        "ulyssesPackedQkvPostUnscatter: threads/block (H*D/8) must be <= 1024, got H=%d D=%d -> %d", H, D,
-        threads);
+        "ulyssesPackedQkvPostUnscatter: threads/block (H*D/8) must be <= 1024, got H=%d D=%d -> %d", H, D, threads);
 
     dim3 const grid(3 * P * Sp, B, 1);
     dim3 const block(threads);
@@ -175,8 +174,8 @@ void launchUlyssesPackedQkvPostUnscatter(
     auto* k_out_typed = reinterpret_cast<__nv_bfloat16*>(k_out);
     auto* v_out_typed = reinterpret_cast<__nv_bfloat16*>(v_out);
 
-    ulyssesPackedQkvPostUnscatterKernel<__nv_bfloat16><<<grid, block, 0, stream>>>(
-        qkv_in_typed, q_out_typed, k_out_typed, v_out_typed, P, B, Sp, H, D, vec_per_row);
+    ulyssesPackedQkvPostUnscatterKernel<__nv_bfloat16>
+        <<<grid, block, 0, stream>>>(qkv_in_typed, q_out_typed, k_out_typed, v_out_typed, P, B, Sp, H, D, vec_per_row);
 }
 
 } // namespace kernels
diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
index bc66eb5be280..dafc1988e8dc 100644
--- a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
+++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h
@@ -63,8 +63,8 @@ void launchUlyssesPostUnscatter(void const* q_in, // [P, B, Sp_q, H_q, D]
 // Packed-QKV variant for self-attention. Consumes the raw packed all-to-all
 // receive buffer [P, B, Sp, 3, H, D] and writes per-tensor NHD-contig outputs
 // [B, P*Sp, H, D].
-void launchUlyssesPackedQkvPostUnscatter(void const* qkv_in, void* q_out, void* k_out, void* v_out, int P, int B,
-    int Sp, int H, int D, cudaStream_t stream);
+void launchUlyssesPackedQkvPostUnscatter(
+    void const* qkv_in, void* q_out, void* k_out, void* v_out, int P, int B, int Sp, int H, int D, cudaStream_t stream);
 
 } // namespace kernels
 
diff --git a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
index fa11e6c67d21..77cc6fec4c0b 100644
--- a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp
@@ -36,14 +36,14 @@ void checkInt32Dim(char const* name, int64_t value)
 
 void checkInt32Product(char const* name, int64_t lhs, int64_t rhs)
 {
-    TORCH_CHECK(lhs == 0 || rhs <= std::numeric_limits<int>::max() / lhs, name, " must fit in int32, got ", lhs, "*",
-        rhs);
+    TORCH_CHECK(
+        lhs == 0 || rhs <= std::numeric_limits<int>::max() / lhs, name, " must fit in int32, got ", lhs, "*", rhs);
 }
 
 void checkInt32SumProduct(char const* name, int64_t lhs, int64_t a, int64_t b, int64_t c)
 {
-    TORCH_CHECK(a <= std::numeric_limits<int64_t>::max() - b && a + b <= std::numeric_limits<int64_t>::max() - c,
-        name, " sum overflows int64");
+    TORCH_CHECK(a <= std::numeric_limits<int64_t>::max() - b && a + b <= std::numeric_limits<int64_t>::max() - c, name,
+        " sum overflows int64");
     checkInt32Product(name, lhs, a + b + c);
 }
 
@@ -140,8 +140,7 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_packed_qkv_post_
     torch::Tensor& qkv_in, int64_t layout)
 {
     TORCH_CHECK(qkv_in.dim() == 6, "ulysses_packed_qkv_post_unscatter expects [P, B, Sp, 3, H, D]");
-    TORCH_CHECK(qkv_in.size(3) == 3, "ulysses_packed_qkv_post_unscatter expects qkv dim size 3, got ",
-        qkv_in.size(3));
+    TORCH_CHECK(qkv_in.size(3) == 3, "ulysses_packed_qkv_post_unscatter expects qkv dim size 3, got ", qkv_in.size(3));
     TORCH_CHECK(layout == 0 || layout == 1, "layout must be 0 (HND) or 1 (NHD), got ", layout);
 
     CHECK_INPUT(qkv_in, torch::kBFloat16);
diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index 2ed9beafc8b3..4c2de2232703 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -2100,6 +2100,89 @@ def _(
     return input.new_empty((input.size(0), weight.size(0)), dtype=output_dtype)
 
 
+class Fp8PrequantizedGemmRunner(TunableRunner):
+    """Runs DeepGEMM FP8 GEMM with a pre-quantized activation."""
+
+    tuning_config = TuningConfig(
+        dynamic_tensor_specs=(DynamicTensorSpec(
+            0, 0, deep_gemm_gen_tuning_buckets), ),
+        exclude_from_cache=True,
+    )
+
+    def __init__(self, output_dtype: torch.dtype, disable_ue8m0_cast: bool):
+        self.output_dtype = output_dtype
+        self.disable_ue8m0_cast = disable_ue8m0_cast
+
+    def unique_id(self):
+        return (
+            self.output_dtype,
+            self.disable_ue8m0_cast,
+        )
+
+    def get_valid_tactics(
+        self,
+        inputs: List[torch.Tensor],
+        profile: OptimizationProfile,
+    ) -> List[int]:
+        return [0]
+
+    def forward(
+        self,
+        inputs: List[torch.Tensor],
+        tactic: int = -1,
+    ) -> torch.Tensor:
+        act_fp8, act_sf, weight, weight_scale = inputs
+        output = torch.empty(
+            (act_fp8.size(0), weight.size(0)),
+            device=act_fp8.device,
+            dtype=self.output_dtype,
+        )
+
+        deep_gemm.fp8_gemm_nt(
+            (act_fp8, act_sf),
+            (weight, weight_scale),
+            output,
+            disable_ue8m0_cast=self.disable_ue8m0_cast,
+        )
+        return output
+
+
+@torch.library.custom_op("trtllm::fp8_prequantized_gemm", mutates_args=())
+def fp8_prequantized_gemm(
+    act_fp8: torch.Tensor,
+    act_sf: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+    disable_ue8m0_cast: bool = False,
+) -> torch.Tensor:
+    tuner = AutoTuner.get()
+    gemm_runner = Fp8PrequantizedGemmRunner(output_dtype, disable_ue8m0_cast)
+    _, best_tactic = tuner.choose_one(
+        "trtllm::fp8_prequantized_gemm",
+        [gemm_runner],
+        Fp8PrequantizedGemmRunner.tuning_config,
+        [act_fp8, act_sf, weight, weight_scale],
+    )
+    return gemm_runner(
+        inputs=[act_fp8, act_sf, weight, weight_scale],
+        tactic=best_tactic,
+    )
+
+
+@fp8_prequantized_gemm.register_fake
+def _(
+    act_fp8: torch.Tensor,
+    act_sf: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+    disable_ue8m0_cast: bool = False,
+) -> torch.Tensor:
+    return act_fp8.new_empty((act_fp8.size(0), weight.size(0)),
+                             dtype=output_dtype)
+
+
 # The runner is used to trigger deepgemm jit during autotune.
 class Fp8BlockScalingGemmRunner(TunableRunner):
     tuning_config = TuningConfig(
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index dcbac9dac253..5bdeb4c19f6b 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -33,7 +33,8 @@
 
 from ..._utils import get_sm_version, is_sm_100f
 from ...models.modeling_utils import QuantConfig
-from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
+from ..utils import (Fp4QuantizedTensor, Fp8BlockScalesQuantizedTensor,
+                     get_model_extra_attrs, is_nvfp4_marlin_supported_sm,
                      replace_parameter_and_save_metadata, unswizzle_sf)
 from .low_m_gemm import _MAX_M as _LOW_M_GEMM_MAX_M
 from .low_m_gemm import LOW_M_GEMM_ACTIVE, apply_low_m_gemm
@@ -1214,6 +1215,21 @@ def create_weights(self, module: Linear, in_features: int,
         else:
             module.register_parameter("bias", None)
 
+    @staticmethod
+    def quantize_deep_gemm_input(
+            input: torch.Tensor) -> Fp8BlockScalesQuantizedTensor:
+        original_shape = input.shape
+        if input.dim() > 2:
+            input = input.reshape(-1, input.shape[-1])
+        assert input.dtype == torch.bfloat16
+        act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(
+            input)
+        return Fp8BlockScalesQuantizedTensor(
+            act_input_fp8,
+            act_input_sf,
+            original_shape,
+        )
+
     def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor]):
         # fp8_block_scaling_gemm does not support writing into an NCCL window
@@ -1221,36 +1237,48 @@ def apply(self, module: Linear, input: torch.Tensor,
         # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden)
         # GEMM ops require 2D matrices
         original_shape = input.shape
-        if input.dim() > 2:
-            input = input.reshape(-1, input.shape[-1])
-
-        if input.dtype == torch.float8_e4m3fn:
-            input = input.to(torch.bfloat16) * module.input_scale
-        assert input.dtype == torch.bfloat16
-
-        if is_sm_100f():
-            if module.use_cute_dsl_blockscaling_mm or module.disable_deep_gemm:
+        if isinstance(input, Fp8BlockScalesQuantizedTensor):
+            output = torch.ops.trtllm.fp8_prequantized_gemm(
+                input.fp8_tensor,
+                input.scaling_factor,
+                module.weight,
+                module.weight_scale,
+                output_dtype=torch.bfloat16,
+                disable_ue8m0_cast=True,
+            )
+            original_shape = input.original_shape
+        else:
+            if input.dim() > 2:
+                input = input.reshape(-1, input.shape[-1])
+            if input.dtype == torch.float8_e4m3fn:
+                input = input.to(torch.bfloat16) * module.input_scale
+            assert input.dtype == torch.bfloat16
+
+            if is_sm_100f():
+                if module.use_cute_dsl_blockscaling_mm or module.disable_deep_gemm:
+                    act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
+                        input)
+                    output = torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell(
+                        act_input_fp8, module.weight, act_input_sf,
+                        module.weight_scale)
+                else:
+                    output = torch.ops.trtllm.fp8_swap_ab_gemm(
+                        input,
+                        module.weight,
+                        module.weight_scale,
+                        disable_ue8m0_cast=True,
+                    )
+            elif get_sm_version() == 120:
+                act_input_fp8, act_input_sf = per_token_quant_and_transform(input)
+                output = torch.ops.trtllm.fp8_block_scaling_gemm(
+                    act_input_fp8, module.weight, act_input_sf,
+                    module.weight_scale)
+            else:
                 act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
                     input)
-                output = torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell(
+                output = torch.ops.trtllm.fp8_block_scaling_gemm(
                     act_input_fp8, module.weight, act_input_sf,
                     module.weight_scale)
-            else:
-                output = torch.ops.trtllm.fp8_swap_ab_gemm(
-                    input,
-                    module.weight,
-                    module.weight_scale,
-                    disable_ue8m0_cast=True,
-                )
-        elif get_sm_version() == 120:
-            act_input_fp8, act_input_sf = per_token_quant_and_transform(input)
-            output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
-        else:
-            act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
-                input)
-            output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
 
         # Reshape output back to original shape (with out_features as last dim)
         if len(original_shape) > 2:
@@ -1423,6 +1451,8 @@ class NVFP4LinearMethod(LinearMethodBase):
 
     supports_nccl_symmetric_memory_window_output: ClassVar[bool] = True
     quantizes_nvfp4_activations: ClassVar[bool] = True
+    _FP8_MAX: ClassVar[float] = 448.0
+    _E2M1_MAX: ClassVar[float] = 6.0
 
     # Temporary workaround which will be resolved by TRTLLM-11958
     # When True, use tunable_fp4_quantize (AutoTuner selects TRTLLM vs
@@ -1488,6 +1518,32 @@ def create_weights(self, module: Linear, in_features: int,
         else:
             module.register_parameter("bias", None)
 
+    @staticmethod
+    def quantize_dynamic_input(input: torch.Tensor,
+                               scaling_vector_size: int) -> Fp4QuantizedTensor:
+        amax_input = torch.amax(torch.abs(input)).float()
+        dynamic_alpha_scale = amax_input / (NVFP4LinearMethod._FP8_MAX *
+                                            NVFP4LinearMethod._E2M1_MAX)
+        input_scale = 1.0 / dynamic_alpha_scale
+        original_shape = input.shape
+        input_2d = input.reshape(-1, input.shape[-1])
+
+        if NVFP4LinearMethod.use_tunable_quantize:
+            act_fp4, act_sf = torch.ops.trtllm.tunable_fp4_quantize(
+                input_2d, input_scale, scaling_vector_size, False)
+        else:
+            act_fp4, act_sf = torch.ops.trtllm.fp4_quantize(
+                input_2d, input_scale, scaling_vector_size, False)
+
+        if len(original_shape) > 2:
+            act_fp4 = act_fp4.reshape(*original_shape[:-1], act_fp4.shape[-1])
+        return Fp4QuantizedTensor(
+            act_fp4,
+            act_sf,
+            is_sf_swizzled=False,
+            dynamic_alpha_scale=dynamic_alpha_scale,
+        )
+
     def _input_prepare(self, module: Linear, input: torch.Tensor):
         """Quantize input tensor to FP4 format.
 
@@ -1500,12 +1556,19 @@ def _input_prepare(self, module: Linear, input: torch.Tensor):
         """
         if isinstance(input, Fp4QuantizedTensor):
             # Input is already quantized - this should not happen if pre_quant_scale exists
-            if module.pre_quant_scale is not None or module.force_dynamic_quantization:
+            if module.pre_quant_scale is not None:
                 raise RuntimeError(
                     "Received pre-quantized FP4 input for a layer that must quantize activations locally "
-                    "(pre_quant_scale is set or dynamic quantization is forced). "
-                    "This indicates FP4 output was not disabled in the previous layer."
+                    "(pre_quant_scale is set). This indicates FP4 output was "
+                    "not disabled in the previous layer."
                 )
+            if module.input_scale is None or module.force_dynamic_quantization:
+                if input.dynamic_alpha_scale is None:
+                    raise RuntimeError(
+                        "Received pre-quantized FP4 input for dynamic NVFP4 "
+                        "without dynamic_alpha_scale metadata.")
+                return (input.fp4_tensor, input.scaling_factor,
+                        input.dynamic_alpha_scale * module.weight_scale_2)
             return input.fp4_tensor, input.scaling_factor, module.alpha
         elif isinstance(input, tuple):
             # Input is a tuple of (fp4_tensor, scaling_factor)
@@ -1525,11 +1588,11 @@ def _input_prepare(self, module: Linear, input: torch.Tensor):
             # Dynamic vs static quantization
             if module.input_scale is None or module.force_dynamic_quantization:
                 # Dynamic mode: compute input_scale and alpha from current input
-                FP8_MAX, E2M1_MAX = 448.0, 6.0
                 amax_input = torch.amax(torch.abs(input)).float()
-                input_scale = FP8_MAX * E2M1_MAX / amax_input
-                alpha = (amax_input /
-                         (FP8_MAX * E2M1_MAX)) * module.weight_scale_2
+                dynamic_alpha_scale = amax_input / (
+                    NVFP4LinearMethod._FP8_MAX * NVFP4LinearMethod._E2M1_MAX)
+                input_scale = 1.0 / dynamic_alpha_scale
+                alpha = dynamic_alpha_scale * module.weight_scale_2
             else:
                 # Static mode: use pre-computed values
                 input_scale = module.input_scale
@@ -1555,6 +1618,8 @@ def apply(self, module: Linear, input: torch.Tensor,
                 input.fp4_tensor.reshape(-1, input.fp4_tensor.shape[-1]),
                 input.scaling_factor,
                 input.is_sf_swizzled,
+                unquantized_hidden_states=input.unquantized_hidden_states,
+                dynamic_alpha_scale=input.dynamic_alpha_scale,
             )
         elif not isinstance(input,
                             (tuple, Fp4QuantizedTensor)) and input.dim() > 2:
@@ -1568,6 +1633,8 @@ def apply(self, module: Linear, input: torch.Tensor,
                                                     input.fp4_tensor.shape[-1]),
                 scaling_factor=input.scaling_factor,
                 is_sf_swizzled=input.is_sf_swizzled,
+                unquantized_hidden_states=input.unquantized_hidden_states,
+                dynamic_alpha_scale=input.dynamic_alpha_scale,
             )
 
         act_fp4, act_sf, alpha = self._input_prepare(module, input)
diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py
index 0a7db9162ee6..189978804fd9 100644
--- a/tensorrt_llm/_torch/utils.py
+++ b/tensorrt_llm/_torch/utils.py
@@ -3,6 +3,7 @@
 
 import contextlib
 import functools
+import math
 import os
 import threading
 from collections.abc import Callable
@@ -212,12 +213,37 @@ class Fp4QuantizedTensor:
     # needing the un-quantized form (e.g. DSv3.2's DSA indexer at
     # sparse/dsa.py:pre_indexer_proj) can use it without dequantizing FP4.
     unquantized_hidden_states: Optional[torch.Tensor] = None
+    # Optional runtime activation scale for dynamic NVFP4. This is
+    # amax(input)/(FP8_MAX*E2M1_MAX), shared by projections that consume the
+    # same pre-quantized activation and combined with each layer's weight scale
+    # to produce that layer's GEMM alpha.
+    dynamic_alpha_scale: Optional[torch.Tensor] = None
 
     @property
     def shape(self):
         return self.fp4_tensor.shape
 
 
+@dataclass
+class Fp8BlockScalesQuantizedTensor:
+    """FP8_BLOCK_SCALES activation and its per-1x128 scaling factors."""
+
+    fp8_tensor: torch.Tensor
+    scaling_factor: torch.Tensor
+    original_shape: torch.Size
+
+    @property
+    def shape(self) -> torch.Size:
+        return self.original_shape
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.fp8_tensor.dtype
+
+    def numel(self) -> int:
+        return math.prod(self.original_shape)
+
+
 @dataclass
 class MxFp8QuantizedTensor:
     """MXFP8 activation and its per-1x32 UE8M0 scaling factors.
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
index b65eb1e9a645..37dd3997c5e3 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
@@ -30,7 +30,7 @@
 _flash_attn_fwd_import_error = None
 try:
     from flash_attn.cute.interface import _flash_attn_fwd
-except (ImportError, OSError) as e:
+except (AttributeError, ImportError, OSError) as e:
     _flash_attn_fwd = None
     _flash_attn_fwd_import_error = e
 
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index 7da73583f798..019031b808cd 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -37,7 +37,7 @@
 _flash_attn_combine_import_error = None
 try:
     from flash_attn.cute.interface import flash_attn_combine as _flash_attn_combine
-except (ImportError, OSError) as e:
+except (AttributeError, ImportError, OSError) as e:
     _flash_attn_combine = None
     _flash_attn_combine_import_error = e
 
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
index b8f49937beb0..fb9d9e7b89a1 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
@@ -26,10 +26,20 @@
 import torch.nn.functional as F
 from torch import nn
 
-from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
+from tensorrt_llm._torch.modules.linear import (
+    FP8BlockScalesLinearMethod,
+    Linear,
+    NVFP4LinearMethod,
+    TensorParallelMode,
+)
 from tensorrt_llm._torch.modules.mlp import MLP
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-from tensorrt_llm._torch.utils import gelu_tanh, maybe_compile
+from tensorrt_llm._torch.utils import (
+    Fp4QuantizedTensor,
+    Fp8BlockScalesQuantizedTensor,
+    gelu_tanh,
+    maybe_compile,
+)
 from tensorrt_llm._torch.visual_gen.attention_backend.parallel import (
     Attention2DAttention,
     RingAttention,
@@ -40,6 +50,7 @@
 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._utils import is_sm_100f
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 _WEIGHT_KEY_REMAPS = [
@@ -62,6 +73,19 @@
     "inv_kv_scales",
 )
 
+_ULYSSES_DYNAMIC_QUANT_BF16_PATTERNS = (
+    "txt_in",
+    "transformer_blocks.*.txt_mod.*",
+    "transformer_blocks.*.txt_mlp.*",
+    "transformer_blocks.*.attn.add_q_proj",
+    "transformer_blocks.*.attn.add_k_proj",
+    "transformer_blocks.*.attn.add_v_proj",
+    "transformer_blocks.*.attn.to_add_out",
+    "transformer_blocks.*.attn.to_q",
+    "transformer_blocks.*.attn.to_k",
+    "transformer_blocks.*.attn.to_v",
+    "transformer_blocks.*.img_mlp.*",
+)
 
 def _remap_checkpoint_keys(weights: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
     remapped = {}
@@ -512,6 +536,12 @@ def __init__(
             self.attn_backend, self.attn
         )
         self._uses_sequence_parallel_attention = _is_qwen_sequence_parallel_attention(self.attn)
+        self._packed_image_fp8_qkv: Optional[
+            Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]
+        ] = None
+        self._packed_text_fp8_qkv: Optional[
+            Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]
+        ] = None
 
         tp_mode = TensorParallelMode.COLUMN if self.tp_size > 1 else None
 
@@ -579,6 +609,174 @@ def __init__(
     def _apply_rms_norm(x: torch.Tensor, norm: RMSNorm) -> torch.Tensor:
         return F.rms_norm(x, (x.shape[-1],), norm.weight, norm.variance_epsilon)
 
+    @staticmethod
+    def _can_share_dynamic_nvfp4_qkv(input: torch.Tensor, *projections: Linear) -> bool:
+        if isinstance(input, Fp4QuantizedTensor):
+            return True
+        if not isinstance(input, torch.Tensor):
+            return False
+        if input.dtype != torch.bfloat16:
+            return False
+        for projection in projections:
+            if not projection.has_nvfp4:
+                return False
+            if not projection.force_dynamic_quantization:
+                return False
+            if projection.pre_quant_scale is not None:
+                return False
+            if getattr(projection, "scaling_vector_size", None) is None:
+                return False
+            if getattr(projection, "weight_scale_2", None) is None:
+                return False
+        return True
+
+    @staticmethod
+    def _shared_dynamic_nvfp4_input(input: torch.Tensor, reference_projection: Linear):
+        if isinstance(input, Fp4QuantizedTensor):
+            return input
+        return NVFP4LinearMethod.quantize_dynamic_input(
+            input, reference_projection.scaling_vector_size
+        )
+
+    @staticmethod
+    def _can_share_fp8_block_scales_qkv(input: torch.Tensor, *projections: Linear) -> bool:
+        if isinstance(input, Fp8BlockScalesQuantizedTensor):
+            return True
+        if not isinstance(input, torch.Tensor):
+            return False
+        if input.dtype != torch.bfloat16 or not is_sm_100f():
+            return False
+        for projection in projections:
+            if not projection.has_fp8_block_scales:
+                return False
+            if projection.use_cute_dsl_blockscaling_mm:
+                return False
+            if projection.disable_deep_gemm:
+                return False
+        return True
+
+    @staticmethod
+    def _shared_fp8_block_scales_input(input: torch.Tensor):
+        if isinstance(input, Fp8BlockScalesQuantizedTensor):
+            return input
+        return FP8BlockScalesLinearMethod.quantize_deep_gemm_input(input)
+
+    @staticmethod
+    def _build_packed_fp8_qkv(
+        *projections: Linear,
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]]:
+        if len(projections) != 3:
+            return None
+        for projection in projections:
+            if not projection.has_fp8_block_scales:
+                return None
+            if projection.use_cute_dsl_blockscaling_mm:
+                return None
+            if projection.disable_deep_gemm:
+                return None
+            if projection.weight_scale.dim() != 2:
+                return None
+            if projection.weight.shape[0] != projection.out_features:
+                return None
+
+        first = projections[0]
+        for projection in projections[1:]:
+            if projection.weight.shape[1] != first.weight.shape[1]:
+                return None
+            if projection.weight_scale.shape[1] != first.weight_scale.shape[1]:
+                return None
+            if (projection.bias is None) != (first.bias is None):
+                return None
+
+        weight = torch.cat([projection.weight for projection in projections], dim=0)
+        scale_m = sum(projection.weight_scale.shape[0] for projection in projections)
+        scale_k = first.weight_scale.shape[1]
+        weight_scale_physical = torch.cat(
+            [projection.weight_scale.transpose(0, 1) for projection in projections],
+            dim=1,
+        ).contiguous()
+        weight_scale = torch.as_strided(
+            weight_scale_physical,
+            (scale_m, scale_k),
+            (1, scale_m),
+        )
+        bias = (
+            None
+            if first.bias is None
+            else torch.cat([projection.bias for projection in projections], dim=0)
+        )
+        return weight, weight_scale, bias
+
+    def cache_packed_projection_weights(self) -> None:
+        self._packed_image_fp8_qkv = self._build_packed_fp8_qkv(self.to_q, self.to_k, self.to_v)
+        self._packed_text_fp8_qkv = self._build_packed_fp8_qkv(
+            self.add_q_proj, self.add_k_proj, self.add_v_proj
+        )
+
+    def _packed_fp8_block_scales_qkv(
+        self,
+        input: torch.Tensor,
+        packed_qkv: Optional[Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]],
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        if packed_qkv is None:
+            return None
+        if isinstance(input, Fp8BlockScalesQuantizedTensor):
+            pass
+        elif not isinstance(input, torch.Tensor):
+            return None
+        elif input.dtype != torch.bfloat16 or not is_sm_100f():
+            return None
+        if not isinstance(input, Fp8BlockScalesQuantizedTensor):
+            input = self._shared_fp8_block_scales_input(input)
+
+        weight, weight_scale, bias = packed_qkv
+        original_shape = input.original_shape
+        qkv = torch.ops.trtllm.fp8_prequantized_gemm(
+            input.fp8_tensor,
+            input.scaling_factor,
+            weight,
+            weight_scale,
+            output_dtype=torch.bfloat16,
+            disable_ue8m0_cast=True,
+        )
+        if bias is not None:
+            qkv = qkv + bias
+        if len(original_shape) > 2:
+            qkv = qkv.reshape(*original_shape[:-1], qkv.shape[-1])
+        return qkv.split([self.local_q_dim, self.local_kv_dim, self.local_kv_dim], dim=-1)
+
+    def _get_image_qkv(self, hidden_states: torch.Tensor):
+        packed_qkv = self._packed_fp8_block_scales_qkv(hidden_states, self._packed_image_fp8_qkv)
+        if packed_qkv is not None:
+            return packed_qkv
+        if self._can_share_dynamic_nvfp4_qkv(hidden_states, self.to_q, self.to_k, self.to_v):
+            hidden_states = self._shared_dynamic_nvfp4_input(hidden_states, self.to_q)
+        elif self._can_share_fp8_block_scales_qkv(hidden_states, self.to_q, self.to_k, self.to_v):
+            hidden_states = self._shared_fp8_block_scales_input(hidden_states)
+        return self.get_qkv(hidden_states)
+
+    def _get_text_qkv(self, encoder_hidden_states: torch.Tensor):
+        packed_qkv = self._packed_fp8_block_scales_qkv(
+            encoder_hidden_states, self._packed_text_fp8_qkv
+        )
+        if packed_qkv is not None:
+            return packed_qkv
+        if self._can_share_dynamic_nvfp4_qkv(
+            encoder_hidden_states, self.add_q_proj, self.add_k_proj, self.add_v_proj
+        ):
+            encoder_hidden_states = self._shared_dynamic_nvfp4_input(
+                encoder_hidden_states, self.add_q_proj
+            )
+        elif self._can_share_fp8_block_scales_qkv(
+            encoder_hidden_states, self.add_q_proj, self.add_k_proj, self.add_v_proj
+        ):
+            encoder_hidden_states = self._shared_fp8_block_scales_input(encoder_hidden_states)
+        return (
+            self.add_q_proj(encoder_hidden_states),
+            self.add_k_proj(encoder_hidden_states),
+            self.add_v_proj(encoder_hidden_states),
+        )
+
     def _use_fused_qk_norm_rope(
         self,
         hidden_states: torch.Tensor,
@@ -600,10 +798,8 @@ def _prepare_qkv_fused(
         image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
         fused_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        img_q, img_k, img_v = self.get_qkv(hidden_states)
-        txt_q = self.add_q_proj(encoder_hidden_states)
-        txt_k = self.add_k_proj(encoder_hidden_states)
-        txt_v = self.add_v_proj(encoder_hidden_states)
+        img_q, img_k, img_v = self._get_image_qkv(hidden_states)
+        txt_q, txt_k, txt_v = self._get_text_qkv(encoder_hidden_states)
 
         txt_qkv = torch.cat([txt_q, txt_k, txt_v], dim=-1)
         img_qkv = torch.cat([img_q, img_k, img_v], dim=-1)
@@ -630,11 +826,9 @@ def _prepare_qkv_unfused(
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]],
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # Image QKV.
-        img_q, img_k, img_v = self.get_qkv(hidden_states)
+        img_q, img_k, img_v = self._get_image_qkv(hidden_states)
         # Text QKV.
-        txt_q = self.add_q_proj(encoder_hidden_states)
-        txt_k = self.add_k_proj(encoder_hidden_states)
-        txt_v = self.add_v_proj(encoder_hidden_states)
+        txt_q, txt_k, txt_v = self._get_text_qkv(encoder_hidden_states)
 
         # Reshape to (B, S, H, D).
         img_q = img_q.unflatten(-1, (self.local_num_attention_heads, -1))
@@ -1151,15 +1345,26 @@ def to_inference_dtype(self) -> "QwenImageTransformer2DModel":
 
     def apply_quant_config_exclude_modules(self) -> None:
         quant_config = self.model_config.quant_config
-        if quant_config is None or quant_config.exclude_modules is None:
+        if quant_config is None:
+            return
+
+        exclude_modules = list(quant_config.exclude_modules or [])
+        if self._keep_ulysses_dynamic_quant_modules_bf16(quant_config):
+            exclude_modules.extend(_ULYSSES_DYNAMIC_QUANT_BF16_PATTERNS)
+
+        if not exclude_modules:
             return
 
         kv_cache_quant_algo = quant_config.kv_cache_quant_algo if quant_config else None
         no_quant_config = QuantConfig(kv_cache_quant_algo=kv_cache_quant_algo)
+        exclusion_config = QuantConfig(
+            kv_cache_quant_algo=kv_cache_quant_algo,
+            exclude_modules=exclude_modules,
+        )
 
         for name, module in self.named_modules():
             if isinstance(module, Linear):
-                is_excluded = quant_config.is_module_excluded_from_quantization(name)
+                is_excluded = exclusion_config.is_module_excluded_from_quantization(name)
                 if is_excluded and getattr(module, "quant_config", None) is not None:
                     module.quant_config = no_quant_config
                     if getattr(module, "_weights_created", False):
@@ -1169,6 +1374,15 @@ def apply_quant_config_exclude_modules(self) -> None:
                         module._buffers.clear()
                         module.create_weights()
 
+    def _keep_ulysses_dynamic_quant_modules_bf16(self, quant_config: QuantConfig) -> bool:
+        if not self.model_config.dynamic_weight_quant:
+            return False
+        vgm = self.model_config.visual_gen_mapping
+        if getattr(vgm, "ulysses_size", 1) <= 1:
+            return False
+        quant_mode = quant_config.layer_quant_mode
+        return quant_mode.has_nvfp4() or quant_mode.has_fp8_block_scales()
+
     def _non_serialized_quant_parameter_names(self) -> set[str]:
         """Return shared Linear parameters absent from ModelOpt checkpoints."""
         non_serialized = set()
@@ -1266,6 +1480,9 @@ def post_load_weights(self) -> None:
                         f"weight_scale_shape={scale_shape}, "
                         f"weight_scale_dtype={scale_dtype}"
                     ) from exc
+        for module in self.modules():
+            if isinstance(module, QwenJointAttention):
+                module.cache_packed_projection_weights()
 
     def forward(
         self,

From c941f708a22877c5789061ab126dc132795dba4d Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Mon, 24 Aug 2026 22:37:16 +0000
Subject: [PATCH 11/18] Address Qwen Image Ulysses review comments

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 docs/source/models/supported-models.md        |  2 +-
 .../visual_gen/attention_backend/parallel.py  | 19 ++++++++++++++-----
 .../multi_gpu/test_qwen_image_edit_ulysses.py |  2 +-
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index 363f172e95f2..d3220ad32786 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -232,7 +232,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image-Layered** [^vg2] | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No |
 | **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 
 [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index 019031b808cd..99ff7953a92e 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -101,10 +101,15 @@ class UlyssesAttention(AttentionBackend):
         Step 3: All-to-All → [B, S/P, H, D] (restore sequence sharding)
         Output: [B, S/P, H, D] (sequence sharded)
 
-    Two modes (auto-selected via ``inner_backend.support_fused_qkv()``):
-    - Unfused: 3 separate all-to-all for Q/K/V + 1 for output (4 collectives)
-    - Fused: stacks Q/K/V into [B, S/P, 3, H, D], 1 fused 5D all-to-all
-      + 1 for output (2 collectives total)
+    Three modes:
+    - Unfused: 3 separate all-to-all for Q/K/V + 1 for output (4 collectives).
+    - Inner-backend fused QKV: selected via ``inner_backend.support_fused_qkv()``.
+      Stacks Q/K/V into [B, S/P, 3, H, D], uses 1 fused 5D all-to-all, then
+      calls the inner backend with packed QKV.
+    - Packed self-attention: selected after fused QKV is unavailable when Q/K/V
+      have the same BF16 CUDA shape and the post-unscatter kernel constraints
+      are satisfied. It uses 1 fused 5D all-to-all, then a native post-unscatter
+      layout conversion before calling the inner backend with separate Q/K/V.
     """
 
     # One side stream shared across all UlyssesAttention instances on the
@@ -204,7 +209,11 @@ def _supports_packed_self_attention(
             return False
         if q.shape != k.shape or q.shape != v.shape:
             return False
-        if q.dtype != torch.bfloat16 or q.shape[-1] % 8 != 0:
+        if not q.is_cuda or q.dtype != torch.bfloat16 or q.shape[-1] % 8 != 0:
+            return False
+        # ulysses_packed_qkv_post_unscatter launches one block of
+        # H_local * (D / 8) threads.
+        if (q.shape[2] // self.world_size) * (q.shape[-1] // 8) > 1024:
             return False
         return kwargs.get("gate_compress") is None and kwargs.get("gate_fine") is None
 
diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
index 0855da88edf3..24f03b33502b 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
@@ -31,7 +31,7 @@
 # The unit test creates its own torch.distributed NCCL process group. Disable the
 # TRT-LLM MPI bootstrap path so importing tensorrt_llm does not initialize MPI.
 _OLD_TLLM_DISABLE_MPI = os.environ.get("TLLM_DISABLE_MPI")
-os.environ.setdefault("TLLM_DISABLE_MPI", "1")
+os.environ["TLLM_DISABLE_MPI"] = "1"
 
 REPO_ROOT = Path(__file__).resolve().parents[5]
 if str(REPO_ROOT) not in sys.path:

From 515b5be7ff4b1ae0a446037a0df12d3180fe7e58 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 25 Aug 2026 00:56:10 +0000
Subject: [PATCH 12/18] Register missing VisualGen multi-GPU tests

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 tests/integration/test_lists/test-db/l0_dgx_b200.yml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 52b246423e6a..de7e709beeef 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -212,10 +212,12 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py
   - unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py
   - unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py
+  - unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py
   - unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py
+  - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py
   - unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py
   - unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py
   - unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py
@@ -228,6 +230,7 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py
   - unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py
   - unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
+  - unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py
 - condition:
     ranges:
       system_gpu_count:

From d6a14653af1ff0f5376230fa194862627ebbae5b Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 25 Aug 2026 01:27:18 +0000
Subject: [PATCH 13/18] Prune quantized QKV projection optimization

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../_torch/custom_ops/torch_custom_ops.py     |  83 ------
 tensorrt_llm/_torch/modules/linear.py         | 136 +++-------
 tensorrt_llm/_torch/utils.py                  |  26 --
 .../qwen_image/transformer_qwen_image.py      | 241 +-----------------
 4 files changed, 47 insertions(+), 439 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index 4c2de2232703..2ed9beafc8b3 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -2100,89 +2100,6 @@ def _(
     return input.new_empty((input.size(0), weight.size(0)), dtype=output_dtype)
 
 
-class Fp8PrequantizedGemmRunner(TunableRunner):
-    """Runs DeepGEMM FP8 GEMM with a pre-quantized activation."""
-
-    tuning_config = TuningConfig(
-        dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
-        exclude_from_cache=True,
-    )
-
-    def __init__(self, output_dtype: torch.dtype, disable_ue8m0_cast: bool):
-        self.output_dtype = output_dtype
-        self.disable_ue8m0_cast = disable_ue8m0_cast
-
-    def unique_id(self):
-        return (
-            self.output_dtype,
-            self.disable_ue8m0_cast,
-        )
-
-    def get_valid_tactics(
-        self,
-        inputs: List[torch.Tensor],
-        profile: OptimizationProfile,
-    ) -> List[int]:
-        return [0]
-
-    def forward(
-        self,
-        inputs: List[torch.Tensor],
-        tactic: int = -1,
-    ) -> torch.Tensor:
-        act_fp8, act_sf, weight, weight_scale = inputs
-        output = torch.empty(
-            (act_fp8.size(0), weight.size(0)),
-            device=act_fp8.device,
-            dtype=self.output_dtype,
-        )
-
-        deep_gemm.fp8_gemm_nt(
-            (act_fp8, act_sf),
-            (weight, weight_scale),
-            output,
-            disable_ue8m0_cast=self.disable_ue8m0_cast,
-        )
-        return output
-
-
-@torch.library.custom_op("trtllm::fp8_prequantized_gemm", mutates_args=())
-def fp8_prequantized_gemm(
-    act_fp8: torch.Tensor,
-    act_sf: torch.Tensor,
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    output_dtype: torch.dtype = torch.bfloat16,
-    disable_ue8m0_cast: bool = False,
-) -> torch.Tensor:
-    tuner = AutoTuner.get()
-    gemm_runner = Fp8PrequantizedGemmRunner(output_dtype, disable_ue8m0_cast)
-    _, best_tactic = tuner.choose_one(
-        "trtllm::fp8_prequantized_gemm",
-        [gemm_runner],
-        Fp8PrequantizedGemmRunner.tuning_config,
-        [act_fp8, act_sf, weight, weight_scale],
-    )
-    return gemm_runner(
-        inputs=[act_fp8, act_sf, weight, weight_scale],
-        tactic=best_tactic,
-    )
-
-
-@fp8_prequantized_gemm.register_fake
-def _(
-    act_fp8: torch.Tensor,
-    act_sf: torch.Tensor,
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    output_dtype: torch.dtype = torch.bfloat16,
-    disable_ue8m0_cast: bool = False,
-) -> torch.Tensor:
-    return act_fp8.new_empty((act_fp8.size(0), weight.size(0)),
-                             dtype=output_dtype)
-
-
 # The runner is used to trigger deepgemm jit during autotune.
 class Fp8BlockScalingGemmRunner(TunableRunner):
     tuning_config = TuningConfig(
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 5bdeb4c19f6b..93ea47b61643 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -33,8 +33,8 @@
 
 from ..._utils import get_sm_version, is_sm_100f
 from ...models.modeling_utils import QuantConfig
-from ..utils import (Fp4QuantizedTensor, Fp8BlockScalesQuantizedTensor,
-                     get_model_extra_attrs, is_nvfp4_marlin_supported_sm,
+from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
+                     is_nvfp4_marlin_supported_sm,
                      replace_parameter_and_save_metadata, unswizzle_sf)
 from .low_m_gemm import _MAX_M as _LOW_M_GEMM_MAX_M
 from .low_m_gemm import LOW_M_GEMM_ACTIVE, apply_low_m_gemm
@@ -1215,21 +1215,6 @@ def create_weights(self, module: Linear, in_features: int,
         else:
             module.register_parameter("bias", None)
 
-    @staticmethod
-    def quantize_deep_gemm_input(
-            input: torch.Tensor) -> Fp8BlockScalesQuantizedTensor:
-        original_shape = input.shape
-        if input.dim() > 2:
-            input = input.reshape(-1, input.shape[-1])
-        assert input.dtype == torch.bfloat16
-        act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(
-            input)
-        return Fp8BlockScalesQuantizedTensor(
-            act_input_fp8,
-            act_input_sf,
-            original_shape,
-        )
-
     def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor]):
         # fp8_block_scaling_gemm does not support writing into an NCCL window
@@ -1237,48 +1222,36 @@ def apply(self, module: Linear, input: torch.Tensor,
         # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden)
         # GEMM ops require 2D matrices
         original_shape = input.shape
-        if isinstance(input, Fp8BlockScalesQuantizedTensor):
-            output = torch.ops.trtllm.fp8_prequantized_gemm(
-                input.fp8_tensor,
-                input.scaling_factor,
-                module.weight,
-                module.weight_scale,
-                output_dtype=torch.bfloat16,
-                disable_ue8m0_cast=True,
-            )
-            original_shape = input.original_shape
-        else:
-            if input.dim() > 2:
-                input = input.reshape(-1, input.shape[-1])
-            if input.dtype == torch.float8_e4m3fn:
-                input = input.to(torch.bfloat16) * module.input_scale
-            assert input.dtype == torch.bfloat16
-
-            if is_sm_100f():
-                if module.use_cute_dsl_blockscaling_mm or module.disable_deep_gemm:
-                    act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
-                        input)
-                    output = torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell(
-                        act_input_fp8, module.weight, act_input_sf,
-                        module.weight_scale)
-                else:
-                    output = torch.ops.trtllm.fp8_swap_ab_gemm(
-                        input,
-                        module.weight,
-                        module.weight_scale,
-                        disable_ue8m0_cast=True,
-                    )
-            elif get_sm_version() == 120:
-                act_input_fp8, act_input_sf = per_token_quant_and_transform(input)
-                output = torch.ops.trtllm.fp8_block_scaling_gemm(
-                    act_input_fp8, module.weight, act_input_sf,
-                    module.weight_scale)
-            else:
+        if input.dim() > 2:
+            input = input.reshape(-1, input.shape[-1])
+
+        if input.dtype == torch.float8_e4m3fn:
+            input = input.to(torch.bfloat16) * module.input_scale
+        assert input.dtype == torch.bfloat16
+
+        if is_sm_100f():
+            if module.use_cute_dsl_blockscaling_mm or module.disable_deep_gemm:
                 act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
                     input)
-                output = torch.ops.trtllm.fp8_block_scaling_gemm(
+                output = torch.ops.trtllm.cute_dsl_fp8_gemm_blackwell(
                     act_input_fp8, module.weight, act_input_sf,
                     module.weight_scale)
+            else:
+                output = torch.ops.trtllm.fp8_swap_ab_gemm(
+                    input,
+                    module.weight,
+                    module.weight_scale,
+                    disable_ue8m0_cast=True,
+                )
+        elif get_sm_version() == 120:
+            act_input_fp8, act_input_sf = per_token_quant_and_transform(input)
+            output = torch.ops.trtllm.fp8_block_scaling_gemm(
+                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
+        else:
+            act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
+                input)
+            output = torch.ops.trtllm.fp8_block_scaling_gemm(
+                act_input_fp8, module.weight, act_input_sf, module.weight_scale)
 
         # Reshape output back to original shape (with out_features as last dim)
         if len(original_shape) > 2:
@@ -1451,8 +1424,6 @@ class NVFP4LinearMethod(LinearMethodBase):
 
     supports_nccl_symmetric_memory_window_output: ClassVar[bool] = True
     quantizes_nvfp4_activations: ClassVar[bool] = True
-    _FP8_MAX: ClassVar[float] = 448.0
-    _E2M1_MAX: ClassVar[float] = 6.0
 
     # Temporary workaround which will be resolved by TRTLLM-11958
     # When True, use tunable_fp4_quantize (AutoTuner selects TRTLLM vs
@@ -1518,32 +1489,6 @@ def create_weights(self, module: Linear, in_features: int,
         else:
             module.register_parameter("bias", None)
 
-    @staticmethod
-    def quantize_dynamic_input(input: torch.Tensor,
-                               scaling_vector_size: int) -> Fp4QuantizedTensor:
-        amax_input = torch.amax(torch.abs(input)).float()
-        dynamic_alpha_scale = amax_input / (NVFP4LinearMethod._FP8_MAX *
-                                            NVFP4LinearMethod._E2M1_MAX)
-        input_scale = 1.0 / dynamic_alpha_scale
-        original_shape = input.shape
-        input_2d = input.reshape(-1, input.shape[-1])
-
-        if NVFP4LinearMethod.use_tunable_quantize:
-            act_fp4, act_sf = torch.ops.trtllm.tunable_fp4_quantize(
-                input_2d, input_scale, scaling_vector_size, False)
-        else:
-            act_fp4, act_sf = torch.ops.trtllm.fp4_quantize(
-                input_2d, input_scale, scaling_vector_size, False)
-
-        if len(original_shape) > 2:
-            act_fp4 = act_fp4.reshape(*original_shape[:-1], act_fp4.shape[-1])
-        return Fp4QuantizedTensor(
-            act_fp4,
-            act_sf,
-            is_sf_swizzled=False,
-            dynamic_alpha_scale=dynamic_alpha_scale,
-        )
-
     def _input_prepare(self, module: Linear, input: torch.Tensor):
         """Quantize input tensor to FP4 format.
 
@@ -1556,19 +1501,12 @@ def _input_prepare(self, module: Linear, input: torch.Tensor):
         """
         if isinstance(input, Fp4QuantizedTensor):
             # Input is already quantized - this should not happen if pre_quant_scale exists
-            if module.pre_quant_scale is not None:
+            if module.pre_quant_scale is not None or module.force_dynamic_quantization:
                 raise RuntimeError(
                     "Received pre-quantized FP4 input for a layer that must quantize activations locally "
-                    "(pre_quant_scale is set). This indicates FP4 output was "
-                    "not disabled in the previous layer."
+                    "(pre_quant_scale is set or dynamic quantization is forced). "
+                    "This indicates FP4 output was not disabled in the previous layer."
                 )
-            if module.input_scale is None or module.force_dynamic_quantization:
-                if input.dynamic_alpha_scale is None:
-                    raise RuntimeError(
-                        "Received pre-quantized FP4 input for dynamic NVFP4 "
-                        "without dynamic_alpha_scale metadata.")
-                return (input.fp4_tensor, input.scaling_factor,
-                        input.dynamic_alpha_scale * module.weight_scale_2)
             return input.fp4_tensor, input.scaling_factor, module.alpha
         elif isinstance(input, tuple):
             # Input is a tuple of (fp4_tensor, scaling_factor)
@@ -1588,11 +1526,11 @@ def _input_prepare(self, module: Linear, input: torch.Tensor):
             # Dynamic vs static quantization
             if module.input_scale is None or module.force_dynamic_quantization:
                 # Dynamic mode: compute input_scale and alpha from current input
+                FP8_MAX, E2M1_MAX = 448.0, 6.0
                 amax_input = torch.amax(torch.abs(input)).float()
-                dynamic_alpha_scale = amax_input / (
-                    NVFP4LinearMethod._FP8_MAX * NVFP4LinearMethod._E2M1_MAX)
-                input_scale = 1.0 / dynamic_alpha_scale
-                alpha = dynamic_alpha_scale * module.weight_scale_2
+                input_scale = FP8_MAX * E2M1_MAX / amax_input
+                alpha = (amax_input /
+                         (FP8_MAX * E2M1_MAX)) * module.weight_scale_2
             else:
                 # Static mode: use pre-computed values
                 input_scale = module.input_scale
@@ -1618,8 +1556,6 @@ def apply(self, module: Linear, input: torch.Tensor,
                 input.fp4_tensor.reshape(-1, input.fp4_tensor.shape[-1]),
                 input.scaling_factor,
                 input.is_sf_swizzled,
-                unquantized_hidden_states=input.unquantized_hidden_states,
-                dynamic_alpha_scale=input.dynamic_alpha_scale,
             )
         elif not isinstance(input,
                             (tuple, Fp4QuantizedTensor)) and input.dim() > 2:
@@ -1633,8 +1569,6 @@ def apply(self, module: Linear, input: torch.Tensor,
                                                     input.fp4_tensor.shape[-1]),
                 scaling_factor=input.scaling_factor,
                 is_sf_swizzled=input.is_sf_swizzled,
-                unquantized_hidden_states=input.unquantized_hidden_states,
-                dynamic_alpha_scale=input.dynamic_alpha_scale,
             )
 
         act_fp4, act_sf, alpha = self._input_prepare(module, input)
diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py
index 189978804fd9..0a7db9162ee6 100644
--- a/tensorrt_llm/_torch/utils.py
+++ b/tensorrt_llm/_torch/utils.py
@@ -3,7 +3,6 @@
 
 import contextlib
 import functools
-import math
 import os
 import threading
 from collections.abc import Callable
@@ -213,37 +212,12 @@ class Fp4QuantizedTensor:
     # needing the un-quantized form (e.g. DSv3.2's DSA indexer at
     # sparse/dsa.py:pre_indexer_proj) can use it without dequantizing FP4.
     unquantized_hidden_states: Optional[torch.Tensor] = None
-    # Optional runtime activation scale for dynamic NVFP4. This is
-    # amax(input)/(FP8_MAX*E2M1_MAX), shared by projections that consume the
-    # same pre-quantized activation and combined with each layer's weight scale
-    # to produce that layer's GEMM alpha.
-    dynamic_alpha_scale: Optional[torch.Tensor] = None
 
     @property
     def shape(self):
         return self.fp4_tensor.shape
 
 
-@dataclass
-class Fp8BlockScalesQuantizedTensor:
-    """FP8_BLOCK_SCALES activation and its per-1x128 scaling factors."""
-
-    fp8_tensor: torch.Tensor
-    scaling_factor: torch.Tensor
-    original_shape: torch.Size
-
-    @property
-    def shape(self) -> torch.Size:
-        return self.original_shape
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.fp8_tensor.dtype
-
-    def numel(self) -> int:
-        return math.prod(self.original_shape)
-
-
 @dataclass
 class MxFp8QuantizedTensor:
     """MXFP8 activation and its per-1x32 UE8M0 scaling factors.
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 fb9d9e7b89a1..b8f49937beb0 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
@@ -26,20 +26,10 @@
 import torch.nn.functional as F
 from torch import nn
 
-from tensorrt_llm._torch.modules.linear import (
-    FP8BlockScalesLinearMethod,
-    Linear,
-    NVFP4LinearMethod,
-    TensorParallelMode,
-)
+from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode
 from tensorrt_llm._torch.modules.mlp import MLP
 from tensorrt_llm._torch.modules.rms_norm import RMSNorm
-from tensorrt_llm._torch.utils import (
-    Fp4QuantizedTensor,
-    Fp8BlockScalesQuantizedTensor,
-    gelu_tanh,
-    maybe_compile,
-)
+from tensorrt_llm._torch.utils import gelu_tanh, maybe_compile
 from tensorrt_llm._torch.visual_gen.attention_backend.parallel import (
     Attention2DAttention,
     RingAttention,
@@ -50,7 +40,6 @@
 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._utils import is_sm_100f
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 _WEIGHT_KEY_REMAPS = [
@@ -73,19 +62,6 @@
     "inv_kv_scales",
 )
 
-_ULYSSES_DYNAMIC_QUANT_BF16_PATTERNS = (
-    "txt_in",
-    "transformer_blocks.*.txt_mod.*",
-    "transformer_blocks.*.txt_mlp.*",
-    "transformer_blocks.*.attn.add_q_proj",
-    "transformer_blocks.*.attn.add_k_proj",
-    "transformer_blocks.*.attn.add_v_proj",
-    "transformer_blocks.*.attn.to_add_out",
-    "transformer_blocks.*.attn.to_q",
-    "transformer_blocks.*.attn.to_k",
-    "transformer_blocks.*.attn.to_v",
-    "transformer_blocks.*.img_mlp.*",
-)
 
 def _remap_checkpoint_keys(weights: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
     remapped = {}
@@ -536,12 +512,6 @@ def __init__(
             self.attn_backend, self.attn
         )
         self._uses_sequence_parallel_attention = _is_qwen_sequence_parallel_attention(self.attn)
-        self._packed_image_fp8_qkv: Optional[
-            Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]
-        ] = None
-        self._packed_text_fp8_qkv: Optional[
-            Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]
-        ] = None
 
         tp_mode = TensorParallelMode.COLUMN if self.tp_size > 1 else None
 
@@ -609,174 +579,6 @@ def __init__(
     def _apply_rms_norm(x: torch.Tensor, norm: RMSNorm) -> torch.Tensor:
         return F.rms_norm(x, (x.shape[-1],), norm.weight, norm.variance_epsilon)
 
-    @staticmethod
-    def _can_share_dynamic_nvfp4_qkv(input: torch.Tensor, *projections: Linear) -> bool:
-        if isinstance(input, Fp4QuantizedTensor):
-            return True
-        if not isinstance(input, torch.Tensor):
-            return False
-        if input.dtype != torch.bfloat16:
-            return False
-        for projection in projections:
-            if not projection.has_nvfp4:
-                return False
-            if not projection.force_dynamic_quantization:
-                return False
-            if projection.pre_quant_scale is not None:
-                return False
-            if getattr(projection, "scaling_vector_size", None) is None:
-                return False
-            if getattr(projection, "weight_scale_2", None) is None:
-                return False
-        return True
-
-    @staticmethod
-    def _shared_dynamic_nvfp4_input(input: torch.Tensor, reference_projection: Linear):
-        if isinstance(input, Fp4QuantizedTensor):
-            return input
-        return NVFP4LinearMethod.quantize_dynamic_input(
-            input, reference_projection.scaling_vector_size
-        )
-
-    @staticmethod
-    def _can_share_fp8_block_scales_qkv(input: torch.Tensor, *projections: Linear) -> bool:
-        if isinstance(input, Fp8BlockScalesQuantizedTensor):
-            return True
-        if not isinstance(input, torch.Tensor):
-            return False
-        if input.dtype != torch.bfloat16 or not is_sm_100f():
-            return False
-        for projection in projections:
-            if not projection.has_fp8_block_scales:
-                return False
-            if projection.use_cute_dsl_blockscaling_mm:
-                return False
-            if projection.disable_deep_gemm:
-                return False
-        return True
-
-    @staticmethod
-    def _shared_fp8_block_scales_input(input: torch.Tensor):
-        if isinstance(input, Fp8BlockScalesQuantizedTensor):
-            return input
-        return FP8BlockScalesLinearMethod.quantize_deep_gemm_input(input)
-
-    @staticmethod
-    def _build_packed_fp8_qkv(
-        *projections: Linear,
-    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]]:
-        if len(projections) != 3:
-            return None
-        for projection in projections:
-            if not projection.has_fp8_block_scales:
-                return None
-            if projection.use_cute_dsl_blockscaling_mm:
-                return None
-            if projection.disable_deep_gemm:
-                return None
-            if projection.weight_scale.dim() != 2:
-                return None
-            if projection.weight.shape[0] != projection.out_features:
-                return None
-
-        first = projections[0]
-        for projection in projections[1:]:
-            if projection.weight.shape[1] != first.weight.shape[1]:
-                return None
-            if projection.weight_scale.shape[1] != first.weight_scale.shape[1]:
-                return None
-            if (projection.bias is None) != (first.bias is None):
-                return None
-
-        weight = torch.cat([projection.weight for projection in projections], dim=0)
-        scale_m = sum(projection.weight_scale.shape[0] for projection in projections)
-        scale_k = first.weight_scale.shape[1]
-        weight_scale_physical = torch.cat(
-            [projection.weight_scale.transpose(0, 1) for projection in projections],
-            dim=1,
-        ).contiguous()
-        weight_scale = torch.as_strided(
-            weight_scale_physical,
-            (scale_m, scale_k),
-            (1, scale_m),
-        )
-        bias = (
-            None
-            if first.bias is None
-            else torch.cat([projection.bias for projection in projections], dim=0)
-        )
-        return weight, weight_scale, bias
-
-    def cache_packed_projection_weights(self) -> None:
-        self._packed_image_fp8_qkv = self._build_packed_fp8_qkv(self.to_q, self.to_k, self.to_v)
-        self._packed_text_fp8_qkv = self._build_packed_fp8_qkv(
-            self.add_q_proj, self.add_k_proj, self.add_v_proj
-        )
-
-    def _packed_fp8_block_scales_qkv(
-        self,
-        input: torch.Tensor,
-        packed_qkv: Optional[Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]],
-    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
-        if packed_qkv is None:
-            return None
-        if isinstance(input, Fp8BlockScalesQuantizedTensor):
-            pass
-        elif not isinstance(input, torch.Tensor):
-            return None
-        elif input.dtype != torch.bfloat16 or not is_sm_100f():
-            return None
-        if not isinstance(input, Fp8BlockScalesQuantizedTensor):
-            input = self._shared_fp8_block_scales_input(input)
-
-        weight, weight_scale, bias = packed_qkv
-        original_shape = input.original_shape
-        qkv = torch.ops.trtllm.fp8_prequantized_gemm(
-            input.fp8_tensor,
-            input.scaling_factor,
-            weight,
-            weight_scale,
-            output_dtype=torch.bfloat16,
-            disable_ue8m0_cast=True,
-        )
-        if bias is not None:
-            qkv = qkv + bias
-        if len(original_shape) > 2:
-            qkv = qkv.reshape(*original_shape[:-1], qkv.shape[-1])
-        return qkv.split([self.local_q_dim, self.local_kv_dim, self.local_kv_dim], dim=-1)
-
-    def _get_image_qkv(self, hidden_states: torch.Tensor):
-        packed_qkv = self._packed_fp8_block_scales_qkv(hidden_states, self._packed_image_fp8_qkv)
-        if packed_qkv is not None:
-            return packed_qkv
-        if self._can_share_dynamic_nvfp4_qkv(hidden_states, self.to_q, self.to_k, self.to_v):
-            hidden_states = self._shared_dynamic_nvfp4_input(hidden_states, self.to_q)
-        elif self._can_share_fp8_block_scales_qkv(hidden_states, self.to_q, self.to_k, self.to_v):
-            hidden_states = self._shared_fp8_block_scales_input(hidden_states)
-        return self.get_qkv(hidden_states)
-
-    def _get_text_qkv(self, encoder_hidden_states: torch.Tensor):
-        packed_qkv = self._packed_fp8_block_scales_qkv(
-            encoder_hidden_states, self._packed_text_fp8_qkv
-        )
-        if packed_qkv is not None:
-            return packed_qkv
-        if self._can_share_dynamic_nvfp4_qkv(
-            encoder_hidden_states, self.add_q_proj, self.add_k_proj, self.add_v_proj
-        ):
-            encoder_hidden_states = self._shared_dynamic_nvfp4_input(
-                encoder_hidden_states, self.add_q_proj
-            )
-        elif self._can_share_fp8_block_scales_qkv(
-            encoder_hidden_states, self.add_q_proj, self.add_k_proj, self.add_v_proj
-        ):
-            encoder_hidden_states = self._shared_fp8_block_scales_input(encoder_hidden_states)
-        return (
-            self.add_q_proj(encoder_hidden_states),
-            self.add_k_proj(encoder_hidden_states),
-            self.add_v_proj(encoder_hidden_states),
-        )
-
     def _use_fused_qk_norm_rope(
         self,
         hidden_states: torch.Tensor,
@@ -798,8 +600,10 @@ def _prepare_qkv_fused(
         image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
         fused_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        img_q, img_k, img_v = self._get_image_qkv(hidden_states)
-        txt_q, txt_k, txt_v = self._get_text_qkv(encoder_hidden_states)
+        img_q, img_k, img_v = self.get_qkv(hidden_states)
+        txt_q = self.add_q_proj(encoder_hidden_states)
+        txt_k = self.add_k_proj(encoder_hidden_states)
+        txt_v = self.add_v_proj(encoder_hidden_states)
 
         txt_qkv = torch.cat([txt_q, txt_k, txt_v], dim=-1)
         img_qkv = torch.cat([img_q, img_k, img_v], dim=-1)
@@ -826,9 +630,11 @@ def _prepare_qkv_unfused(
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]],
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # Image QKV.
-        img_q, img_k, img_v = self._get_image_qkv(hidden_states)
+        img_q, img_k, img_v = self.get_qkv(hidden_states)
         # Text QKV.
-        txt_q, txt_k, txt_v = self._get_text_qkv(encoder_hidden_states)
+        txt_q = self.add_q_proj(encoder_hidden_states)
+        txt_k = self.add_k_proj(encoder_hidden_states)
+        txt_v = self.add_v_proj(encoder_hidden_states)
 
         # Reshape to (B, S, H, D).
         img_q = img_q.unflatten(-1, (self.local_num_attention_heads, -1))
@@ -1345,26 +1151,15 @@ def to_inference_dtype(self) -> "QwenImageTransformer2DModel":
 
     def apply_quant_config_exclude_modules(self) -> None:
         quant_config = self.model_config.quant_config
-        if quant_config is None:
-            return
-
-        exclude_modules = list(quant_config.exclude_modules or [])
-        if self._keep_ulysses_dynamic_quant_modules_bf16(quant_config):
-            exclude_modules.extend(_ULYSSES_DYNAMIC_QUANT_BF16_PATTERNS)
-
-        if not exclude_modules:
+        if quant_config is None or quant_config.exclude_modules is None:
             return
 
         kv_cache_quant_algo = quant_config.kv_cache_quant_algo if quant_config else None
         no_quant_config = QuantConfig(kv_cache_quant_algo=kv_cache_quant_algo)
-        exclusion_config = QuantConfig(
-            kv_cache_quant_algo=kv_cache_quant_algo,
-            exclude_modules=exclude_modules,
-        )
 
         for name, module in self.named_modules():
             if isinstance(module, Linear):
-                is_excluded = exclusion_config.is_module_excluded_from_quantization(name)
+                is_excluded = quant_config.is_module_excluded_from_quantization(name)
                 if is_excluded and getattr(module, "quant_config", None) is not None:
                     module.quant_config = no_quant_config
                     if getattr(module, "_weights_created", False):
@@ -1374,15 +1169,6 @@ def apply_quant_config_exclude_modules(self) -> None:
                         module._buffers.clear()
                         module.create_weights()
 
-    def _keep_ulysses_dynamic_quant_modules_bf16(self, quant_config: QuantConfig) -> bool:
-        if not self.model_config.dynamic_weight_quant:
-            return False
-        vgm = self.model_config.visual_gen_mapping
-        if getattr(vgm, "ulysses_size", 1) <= 1:
-            return False
-        quant_mode = quant_config.layer_quant_mode
-        return quant_mode.has_nvfp4() or quant_mode.has_fp8_block_scales()
-
     def _non_serialized_quant_parameter_names(self) -> set[str]:
         """Return shared Linear parameters absent from ModelOpt checkpoints."""
         non_serialized = set()
@@ -1480,9 +1266,6 @@ def post_load_weights(self) -> None:
                         f"weight_scale_shape={scale_shape}, "
                         f"weight_scale_dtype={scale_dtype}"
                     ) from exc
-        for module in self.modules():
-            if isinstance(module, QwenJointAttention):
-                module.cache_packed_projection_weights()
 
     def forward(
         self,

From b8d9d3d8e4f3e3e8b0c84fe2ff87061d37afefa8 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Thu, 27 Aug 2026 14:43:12 +0000
Subject: [PATCH 14/18] Guard packed Ulysses attention head divisibility

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
index 99ff7953a92e..545d2b9fafeb 100644
--- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
+++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
@@ -211,6 +211,8 @@ def _supports_packed_self_attention(
             return False
         if not q.is_cuda or q.dtype != torch.bfloat16 or q.shape[-1] % 8 != 0:
             return False
+        if q.shape[2] % self.world_size != 0:
+            return False
         # ulysses_packed_qkv_post_unscatter launches one block of
         # H_local * (D / 8) threads.
         if (q.shape[2] // self.world_size) * (q.shape[-1] // 8) > 1024:

From 034a1e6c61a4f98854051ea1b08a35cfdcb5b4f7 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Mon, 31 Aug 2026 23:53:38 +0000
Subject: [PATCH 15/18] Fix VisualGen spawn helper import order

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 .../multi_gpu/_visual_gen_dist_utils.py       | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/_visual_gen_dist_utils.py b/tests/unittest/_torch/visual_gen/multi_gpu/_visual_gen_dist_utils.py
index 5dfedb664606..fe5a0bb65380 100644
--- a/tests/unittest/_torch/visual_gen/multi_gpu/_visual_gen_dist_utils.py
+++ b/tests/unittest/_torch/visual_gen/multi_gpu/_visual_gen_dist_utils.py
@@ -26,17 +26,25 @@
 single allocation is not enough on busy nodes -- we must re-allocate and retry.
 """
 
-import torch.multiprocessing as mp
+import sys
+from pathlib import Path
 
-# The CI-aware allocator lives in tests/integration/defs/common.py. Adding that
-# directory to sys.path lets us reuse it (it tracks allocated ports per-process
-# so sequential tests don't collide, and honors CONTAINER_PORT_START/NUM).
-__extra_import_path__ = ["~/tests/integration"]
-from defs.common import get_free_port_in_ci
+import torch.multiprocessing as mp
 
 _ADDR_IN_USE_MARKERS = ("EADDRINUSE", "address already in use")
 
 
+def _get_free_port_in_ci() -> int:
+    # Import lazily so mp.spawn child imports do not load tensorrt_llm before
+    # the test process has established its package environment.
+    integration_dir = Path(__file__).resolve().parents[4] / "integration"
+    if str(integration_dir) not in sys.path:
+        sys.path.insert(0, str(integration_dir))
+    from defs.common import get_free_port_in_ci
+
+    return get_free_port_in_ci()
+
+
 def _is_addr_in_use(exc: BaseException) -> bool:
     msg = str(exc)
     return any(marker in msg for marker in _ADDR_IN_USE_MARKERS)
@@ -53,7 +61,7 @@ def spawn_with_retry(spawn_fn, max_retries: int = 10):
     """
     last_exc: BaseException | None = None
     for _ in range(max_retries):
-        port = get_free_port_in_ci()
+        port = _get_free_port_in_ci()
         try:
             spawn_fn(port)
             return

From 8b3dbc6f8550f1e7763e87a71ccb0004615574d0 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 1 Sep 2026 00:17:36 +0000
Subject: [PATCH 16/18] Fix Qwen Image Edit docs support matrix

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 docs/source/models/supported-models.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index d3220ad32786..363f172e95f2 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -232,7 +232,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image-Layered** [^vg2] | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | No | No | No | No |
 | **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 
 [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.

From 1d8a3b9911f12f0e26f985352ccb8c0792fb0c97 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Tue, 1 Sep 2026 00:31:45 +0000
Subject: [PATCH 17/18] Fix pre-commit import cleanup

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/modules/linear.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 93ea47b61643..dcbac9dac253 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -34,7 +34,6 @@
 from ..._utils import get_sm_version, is_sm_100f
 from ...models.modeling_utils import QuantConfig
 from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
-                     is_nvfp4_marlin_supported_sm,
                      replace_parameter_and_save_metadata, unswizzle_sf)
 from .low_m_gemm import _MAX_M as _LOW_M_GEMM_MAX_M
 from .low_m_gemm import LOW_M_GEMM_ACTIVE, apply_low_m_gemm

From b11b79a0a0ec909fa92ef81fec694dc9c77c7d43 Mon Sep 17 00:00:00 2001
From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Date: Fri, 4 Sep 2026 05:53:49 +0000
Subject: [PATCH 18/18] Address VisualGen CI and docs review comments

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
---
 docs/source/models/supported-models.md                      | 2 +-
 docs/source/models/visual-generation.md                     | 2 +-
 tests/integration/test_lists/test-db/l0_dgx_b200.yml        | 1 -
 .../multi_gpu/test_qwen_image_attention_parallel.py         | 6 +++---
 4 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index 363f172e95f2..d3220ad32786 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -232,7 +232,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image-Layered** [^vg2] | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No |
 | **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 
 [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md
index 29349994b969..a903387e44f9 100644
--- a/docs/source/models/visual-generation.md
+++ b/docs/source/models/visual-generation.md
@@ -66,7 +66,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 | **LTX-2** | Yes | Yes | Yes [^4] | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | No | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No |
 | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |
 | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
 | **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No |
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index de7e709beeef..81ff413a1f17 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -230,7 +230,6 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py
   - unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py
   - unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
-  - unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py
 - condition:
     ranges:
       system_gpu_count:
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 850f5b6089b8..2424932f0fa8 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
@@ -170,13 +170,13 @@ def _test_qwen_image_attention_parallel_topology(
 @pytest.mark.parametrize(
     "world_size,parallel,backend,topology",
     [
-        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=pytest.mark.gpu2, id="tp2"),
+        pytest.param(2, {"tp_size": 2}, "VANILLA", "tp", marks=[pytest.mark.gpu2], id="tp2"),
         pytest.param(
             4,
             {"ring_size": 2, "ulysses_size": 2},
             "FA4",
             "ring",
-            marks=pytest.mark.gpu4,
+            marks=[pytest.mark.gpu4],
             id="ring2_ulysses2",
         ),
         pytest.param(
@@ -184,7 +184,7 @@ def _test_qwen_image_attention_parallel_topology(
             {"attn2d_size": (2, 1), "ulysses_size": 2},
             "FA4",
             "attn2d",
-            marks=pytest.mark.gpu4,
+            marks=[pytest.mark.gpu4],
             id="attn2d_2x1_ulysses2",
         ),
     ],