diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu index 2d0714a890ec..f5eb084b3afd 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 +__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(p) * B + b) * Sp + sp) * 3 + qkv_idx) * H + h) * D) + vec_idx * VEC; + int64_t const out_base = (((static_cast(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(qkv_in + in_base); + uint4* out_v4 = reinterpret_cast(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> + <<>>(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..dafc1988e8dc 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..77cc6fec4c0b 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 +#include #include 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::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::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::max() - b && a + b <= std::numeric_limits::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 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 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 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(P), static_cast(B), static_cast(Sp), + static_cast(H), static_cast(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/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 5ef301de8624..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 | No | No | Yes | Yes | Yes | 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 3e80c120b2ef..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 | No | No | Yes | Yes | Yes | 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/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..5e9edf8a10cf --- /dev/null +++ b/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 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/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 fedabf4ceba8..545d2b9fafeb 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 @@ -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. @@ -78,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 @@ -156,8 +184,40 @@ 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: + # 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, + AttentionTensorLayout.NHD, + ): + return False + if q.shape != k.shape or q.shape != v.shape: + 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: + return False + return kwargs.get("gate_compress") is None and kwargs.get("gate_fine") is None def _forward_fused( self, @@ -198,6 +258,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, @@ -344,7 +426,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 = q_5d.dtype == torch.bfloat16 + 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/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 63755e48bd6e..81ff413a1f17 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,13 @@ 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 - unittest/_torch/visual_gen/multi_gpu/test_tp_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/_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 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 new file mode 100644 index 000000000000..24f03b33502b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py @@ -0,0 +1,266 @@ +# 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 coverage for Qwen-Image-Edit Ulysses attention.""" + +import os +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 + +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. +_OLD_TLLM_DISABLE_MPI = os.environ.get("TLLM_DISABLE_MPI") +os.environ["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 + 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) + + +@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() + + +@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, + 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: + test_fn(rank, world_size, **kwargs) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +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, + args=(world_size, "nccl", test_fn, port, kwargs), + nprocs=world_size, + join=True, + ) + ) + + +def _make_config(rank: int, world_size: int, ulysses_size: int) -> DiffusionModelConfig: + mapping = VisualGenMapping( + world_size=world_size, + rank=rank, + tp_size=1, + cp_size=1, + pp_size=1, + cfg_size=1, + sp_size=ulysses_size, + ulysses_size=ulysses_size, + device_mesh=DeviceMeshTopologyImpl.create().initialize(world_size), + ) + return DiffusionModelConfig( + mapping=mapping.to_llm_mapping(), + visual_gen_mapping=mapping, + attention=AttentionConfig(backend="VANILLA"), + dtype="bfloat16", + ) + + +def _make_attention(config: DiffusionModelConfig) -> QwenJointAttention: + return QwenJointAttention( + dim=16, + num_attention_heads=2, + attention_head_dim=8, + config=config, + ).cuda() + + +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, + ) + + 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)