[None][feat] Qwen Image Ulysses perf enhancement

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,28 @@ 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

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

cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp

@@ -18,13 +18,37 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/CUDAContext.h>
+#include <limits>
 #include <torch/extension.h>
 
 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();
@@ -96,17 +133,66 @@ 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)");
+    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();
+    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

docs/source/models/supported-models.md

@@ -227,7 +227,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.

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 |

examples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yaml

@@ -0,0 +1,27 @@
+# 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:
+  # 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
+  ulysses_size: 2
+cuda_graph_config:
+  enable: false

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -1275,6 +1275,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:])

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