diff --git a/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.cu b/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.cu new file mode 100644 index 000000000000..aba929755690 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.cu @@ -0,0 +1,107 @@ +/* + * Copyright (c) 2025-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. + */ + +#include "tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +// Vector type: int4 = 16 bytes = 8 bf16. Issues LDG.E.128 / STG.E.128. +constexpr int VEC = 8; +constexpr int BLOCK_S = 32; // rows per CTA +constexpr int THREADS_PER_BLOCK = 128; // 4 warps + +// 1 CTA handles BLOCK_S rows × 1 head × full D. +// For fixed h (per CTA), peer = h // H_local is constant — no warp divergence. +// All writes within one CTA go to the same destination slot, contiguous in +// dst space. Matches the access pattern of PyTorch inductor's permute kernel +// epilogue stores. +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + ulyssesPermuteScatterKernel(__nv_bfloat16 const* __restrict__ input, // [B, S_local, H, D] + __nv_bfloat16* __restrict__ send_buf, // [P, B, S_local, H/P, D] + __nv_bfloat16* __restrict__ recv_buf, // [P, B, S_local, H/P, D] + int const my_rank, + int const n_rows, // B * S_local + int const H, int const D, + int const H_local) // H / P +{ + int const bs_block = blockIdx.x; + int const h = blockIdx.y; + + // Scalar branch — same destination slot for all threads in this CTA. + int const peer = h / H_local; + int const h_local = h - peer * H_local; + int const slot_idx = (peer == my_rank) ? my_rank : peer; + __nv_bfloat16* __restrict__ dst_base = (peer == my_rank) ? recv_buf : send_buf; + + int const n_d_chunks = D / VEC; + int const total_tasks = BLOCK_S * n_d_chunks; + int const t = threadIdx.x; + + int const row_base = bs_block * BLOCK_S; + + int4 const* __restrict__ in_v = reinterpret_cast(input); + int4* __restrict__ dst_v = reinterpret_cast(dst_base); + + int const row_in_stride_v = (H * D) / VEC; // = H * n_d_chunks + int const head_in_off_v = h * n_d_chunks; + int const slot_off_v = slot_idx * n_rows * H_local * n_d_chunks; + int const row_dst_stride_v = H_local * n_d_chunks; + int const head_dst_off_v = h_local * n_d_chunks; + +#pragma unroll 1 + for (int idx = t; idx < total_tasks; idx += blockDim.x) + { + int const s_in_block = idx / n_d_chunks; + int const d_chunk = idx - s_in_block * n_d_chunks; + int const row = row_base + s_in_block; + if (row >= n_rows) + continue; + + int const src_idx = row * row_in_stride_v + head_in_off_v + d_chunk; + int const dst_idx = slot_off_v + row * row_dst_stride_v + head_dst_off_v + d_chunk; + + dst_v[dst_idx] = in_v[src_idx]; + } +} + +} // anonymous namespace + +void launchUlyssesPermuteScatter(void const* input, void* send_buf, void* recv_buf, int my_rank, int B, int S_local, + int H, int D, int P, cudaStream_t stream) +{ + int const n_rows = B * S_local; + int const H_local = H / P; + dim3 const grid((n_rows + BLOCK_S - 1) / BLOCK_S, H); + dim3 const block(THREADS_PER_BLOCK); + + ulyssesPermuteScatterKernel<<>>(reinterpret_cast<__nv_bfloat16 const*>(input), + reinterpret_cast<__nv_bfloat16*>(send_buf), reinterpret_cast<__nv_bfloat16*>(recv_buf), my_rank, n_rows, H, D, + H_local); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h b/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h new file mode 100644 index 000000000000..1b2738ac29b3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2025-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. + */ +#pragma once + +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Fused permute + scatter for Ulysses A2A (peer-WRITE variant). +// +// Replaces the .permute(2,0,1,3,4).contiguous() materialization that would +// otherwise happen pre-A2A. Reads `input [B, S_local, H, D]` once and +// scatters each (b,s,h,d) element to one of two destinations: +// - peer != my_rank → send_buf[peer, b, s, h-peer*H_local, d] (local) +// - peer == my_rank → recv_buf[my_rank, b, s, h-my_rank*H_local, d] (symm-mem) +// +// After this kernel runs, the caller fires (P-1) cudaMemcpyBatchAsync +// entries to push send_buf[p] → peer[p].recv_buf[my_rank], then an LSA +// barrier (both folded into ulysses_a2a_async). +// +// Layout (all contiguous bf16): +// input : [B, S_local, H, D] row-major +// send_buf : [P, B, S_local, H/P, D] row-major +// recv_buf : [P, B, S_local, H/P, D] row-major +// +// Requires: D % 8 == 0 (int4 vec load); H % P == 0; bf16 only. +void launchUlyssesPermuteScatter(void const* input, // bf16 [B, S_local, H, D] + void* send_buf, // bf16 [P, B, S_local, H/P, D] + void* recv_buf, // bf16 [P, B, S_local, H/P, D] + int my_rank, int B, int S_local, int H, int D, int P, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu new file mode 100644 index 000000000000..c5f61ff3431d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu @@ -0,0 +1,111 @@ +/* + * 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. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "ulyssesPostUnscatterKernel.h" +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +// Each block handles one (p, b, sp) tile across all H heads. +// Reads H*D bf16 contiguous from input; writes the same bytes scattered across +// H rows of the output in NHD layout [B, P*Sp, H, D]. The caller (op wrapper) +// returns this storage as a transpose-view when an HND-shape output is needed, +// so the resulting tensor is HND-shape with NHD-stride (matching what the +// sync `_forward_unfused` path produces via `q.transpose(1, 2)`). +// +// threads/block = H * (D / 8); each thread copies one uint4 (8 bf16). +template +__global__ void ulyssesPostUnscatterKernel(T const* __restrict__ q_in, T const* __restrict__ k_in, + T const* __restrict__ v_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 h = threadIdx.x / vec_per_row; + int const vec_idx = threadIdx.x - h * vec_per_row; + + int const psp = blockIdx.x; // 0 .. P*Sp-1 + int const p = psp / Sp; + int const sp = psp - p * Sp; + int const b = blockIdx.y; + int const PSp = P * Sp; + + T const* in_ptr; + T* out_ptr; + switch (blockIdx.z) + { + case 0: + in_ptr = q_in; + out_ptr = q_out; + break; + case 1: + in_ptr = k_in; + out_ptr = k_out; + break; + default: + in_ptr = v_in; + out_ptr = v_out; + break; + } + + // in[p, b, sp, h, d]: ((((p*B + b)*Sp + sp)*H + h)*D + vec_idx*VEC) + // NHD out[b, p*Sp+sp, h, d]: (((b*PSp + psp)*H + h)*D + vec_idx*VEC) + // int64_t: P*B*Sp*H*D can exceed 2^31 at large workloads. + int64_t const in_base = ((((static_cast(p) * B + b) * Sp + sp) * H + h) * D) + vec_idx * VEC; + int64_t const out_base = (((static_cast(b) * PSp + psp) * H + h) * D) + vec_idx * VEC; + + uint4 const* in_v4 = reinterpret_cast(in_ptr + 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, + 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, "ulyssesPostUnscatter: D must be a multiple of 8 (uint4 vec), got %d", D); + int const vec_per_row = D / VEC; + int const threads = H * vec_per_row; + TLLM_CHECK_WITH_INFO(threads <= 1024, + "ulyssesPostUnscatter: threads/block (H*D/8) must be <= 1024, got H=%d D=%d -> %d", H, D, threads); + + dim3 const grid(P * Sp, B, 3); + dim3 const block(threads); + + auto* q_in_typed = reinterpret_cast<__nv_bfloat16 const*>(q_in); + auto* k_in_typed = reinterpret_cast<__nv_bfloat16 const*>(k_in); + auto* v_in_typed = reinterpret_cast<__nv_bfloat16 const*>(v_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); + + ulyssesPostUnscatterKernel<__nv_bfloat16><<>>( + q_in_typed, k_in_typed, v_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 new file mode 100644 index 000000000000..f5a2932dacdf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2022-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. + */ + +#ifndef TRTLLM_ULYSSESPOSTUNSCATTERKERNEL_H +#define TRTLLM_ULYSSESPOSTUNSCATTERKERNEL_H + +#include "tensorrt_llm/common/config.h" +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Post-Ulysses A2A unscatter for Q/K/V. Pairs with ulyssesPermuteScatter: +// PermuteScatter prepares send_buf pre-A2A; PostUnscatter consumes recv_buf +// post-A2A and produces SDPA-ready tensors in NHD layout. +// +// After the head-dim → seq-dim all-to-all, each rank holds tensors of shape +// [P, B, Sp, H, D] where P = sequence-parallel world size, Sp = local seq +// len, H = heads-per-rank, D = head dim. This kernel always writes NHD-contig +// [B, P*Sp, H, D] storage. The op wrapper returns the storage as-is for NHD +// callers (TRTLLM / FA4) or as a transpose-view for HND callers (VANILLA / +// torch SDPA) — the HND-shape return is thus HND-shape with NHD-stride, +// mirroring what the sync `_forward_unfused` path produces via +// `q.transpose(1, 2)` (without `.contiguous()`). This stride pattern lets +// cudnn SDPA preserve NHD-stride through its output, so the downstream +// `_output_a2a`'s `.transpose(1, 2).contiguous()` collapses to a no-op. +// +// Equivalent eager expression this kernel replaces: +// t.permute(1, 0, 2, 3, 4).reshape(B, P*Sp, H, D).contiguous() +// +// Layout: +// - Each block reads one fully contiguous (p, b, sp, :H, :D) tile of +// H*D bf16 +// - H*(D/8) threads/block — each thread copies one uint4 (8 bf16) +// - Grid (P*Sp, B, 3): blockIdx.z selects Q / K / V +// +// Constraints: +// - dtype must be bf16 +// - D must be a multiple of 8 (uint4 vector load/store, 8 bf16 per thread) +// - threads/block = H * (D / 8) must be <= 1024 (CUDA hw limit) +void launchUlyssesPostUnscatter(void const* q_in, // [P, B, Sp, H, D] + void const* k_in, void const* v_in, + void* q_out, // [B, P*Sp, H, D] NHD-contig + void* k_out, void* v_out, int P, int B, int Sp, int H, int D, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_ULYSSESPOSTUNSCATTERKERNEL_H diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 90239c56f526..1942a00556b4 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -41,6 +41,7 @@ add_library( allgatherOp.cpp allreduceOp.cpp alltoallOp.cpp + asyncUlyssesOp.cpp attentionOp.cpp causalConv1dOp.cpp convertSpecDecodingMaskToPackedMaskOp.cpp @@ -69,6 +70,8 @@ add_library( fusedDiTQKNormRopeOp.cpp fusedDiTSplitQKNormRopeOp.cpp fusedDiTSplitNormOp.cpp + ulyssesPostUnscatterOp.cpp + ulyssesPermuteScatterOp.cpp fusedAddRMSNormQuant.cpp fusedActivationQuant.cpp fusedGatedRMSNormQuant.cpp diff --git a/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp b/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp new file mode 100644 index 000000000000..b0c01c673b1f --- /dev/null +++ b/cpp/tensorrt_llm/thop/asyncUlyssesOp.cpp @@ -0,0 +1,516 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * 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. + */ +// +// Async Ulysses A2A — PyTorch _SymmetricMemory CUDA-IPC backend. +// +// Pipeline (paired with UlyssesAttention.forward_async in +// tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py): +// +// recv, send_h = ulysses_a2a_async_prepare(input, pg) # default stream +// ev.record() +// with torch.cuda.stream(comm_stream): +// ev.wait() +// ulysses_a2a_async_push(send_h, pg) # CE push only +// ... repeat _prepare/_push for next V/Q/K ... +// with torch.cuda.stream(comm_stream): +// ulysses_a2a_async_barrier(pg) # one per deferred push +// +// Phase 1 (`_prepare`) on the caller's compute stream: +// - lazily allocate one slot of a ring of P-symmetric-memory buffers +// via empty_strided_p2p + rendezvous (PyTorch CUDA-IPC backend); +// - launch the fused permute+scatter kernel into (slot.sendBuf for +// peer chunks, slot.basePtr+my_rank for self chunk); +// - return the 5D recv view and an opaque SendHandle. +// +// Phase 2 (`_async`) on the comm stream: +// - cudaMemcpyBatchAsync (capture-safe per-peer loop fallback) pushes +// each peer's slice of sendBuf into peer.basePtr[my_rank]; +// - PT symm-mem `barrier(channel, timeout_ms)` is the cross-rank fence. +// +// No NCCL device API; no LSA barrier kernel. +// + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Opaque handle returned by `_prepare`, consumed by `_async`. Hides raw +// pointer plumbing from Python; `send_t` keeps the slot's sendBuf tensor +// view alive across the two op calls. +// +// `group_name` binds the handle to the PG that produced it: peer_recv_ptrs +// are valid only in that PG's symm-mem registration. `_async` rejects any +// PG whose group name doesn't match — two distinct PGs of the same size +// would otherwise pass the peer-pointer-count check and silently push +// into the wrong group's buffers. +struct SendHandle : torch::CustomClassHolder +{ + torch::Tensor send_t; + std::vector peer_recv_ptrs; + int64_t slot_bytes; + std::string group_name; +}; + +#if ENABLE_MULTI_DEVICE + +namespace +{ + +class AsyncUlyssesOp +{ +public: + // Slot ring depth. Minimum 3 = one slot each for V/Q/K within a single + // forward_async call. A slot is touched by 4 ops in sequence: + // (a) default-stream Phase-1 write — permute+scatter into slot.sendBuf + // (b) side-stream Phase-2 CE push — reads slot.sendBuf + // (c) side-stream Phase-2 barrier — peer writes into slot.recv + // (d) default-stream SDPA read — reads slot.recv + // Intra-layer hazard: V/Q/K must use distinct slots, otherwise (a) on the + // default stream races (b) on the side stream — they touch the same + // sendBuf and there is no stream sync between them until _join_async. + // Cross-layer hazard (Layer N+1 V reusing Layer N V's slot): safe because + // _join_async at end of Layer N waits the side stream's K barrier event, + // and SDPA on the default stream drains the recv read before Layer N+1 + // starts. So kNumSlots = 3 is the tight minimum. + static constexpr int kNumSlots = 3; + + explicit AsyncUlyssesOp(c10::intrusive_ptr pg) + : mPg(std::move(pg)) + { + } + + void initialize() + { + TLLM_CHECK_WITH_INFO(mPg, "AsyncUlyssesOp requires a torch ProcessGroup"); + TLLM_CHECK_WITH_INFO(mPg->getSize() >= 1, "ProcessGroup size must be >= 1"); + // Register the PG's group_info with PT symm-mem (one-shot per process per group). + ensureGroupRegistered(); + } + + int getPgSize() const + { + return mPg->getSize(); + } + + int getPgRank() const + { + return mPg->getRank(); + } + + // Phase 1: lazy-alloc next ring slot via PT symm-mem; return tensor + // views over send_buf (local push source) and recv_buf (peer-writable) + // plus the host-side peer-pointer array. + std::tuple, int64_t> acquireSlotPair( + at::IntArrayRef shape, c10::ScalarType dtype) + { + int64_t const elemSize = static_cast(c10::elementSize(dtype)); + TORCH_CHECK(elemSize > 0, "dtype must have positive itemsize"); + int64_t numel = 1; + for (auto d : shape) + { + TORCH_CHECK(d > 0, "shape dims must be positive"); + numel *= d; + } + int64_t const bufferBytes = numel * elemSize; + TORCH_CHECK(bufferBytes > 0, "bufferBytes must be positive"); + int const pSize = mPg->getSize(); + TORCH_CHECK(bufferBytes % pSize == 0, "bufferBytes must be divisible by world_size"); + + int const slotIdx = nextSlotIdx(); + Slot& slot = getOrAllocSlot(slotIdx, static_cast(bufferBytes)); + + auto opts = torch::dtype(dtype).device(torch::kCUDA); + auto sendT = torch::from_blob( + slot.sendBuf, shape, /*deleter=*/[](void*) {}, opts); + auto recvT = torch::from_blob( + slot.basePtr, shape, /*deleter=*/[](void*) {}, opts); + + std::vector peerRecvPtrs(pSize); + for (int p = 0; p < pSize; ++p) + { + peerRecvPtrs[p] = reinterpret_cast(slot.peerPtrs[p]); + } + + int64_t const slotBytes = bufferBytes / pSize; + return std::make_tuple(sendT, recvT, std::move(peerRecvPtrs), slotBytes); + } + + // Phase 2 (data): out-of-capture uses cudaMemcpyBatchAsync (multi-CE + // engine fan-out); under stream capture we serialize via per-peer + // cudaMemcpyAsync (cudaMemcpyBatchAsync is not graph-capture-safe). + // Self chunk is NOT pushed (already written by the upstream + // fused-permute kernel into recv_buf[my_rank]). + void runCePush(torch::Tensor send_buf, std::vector const& peer_recv_ptrs, int64_t slot_bytes) + { + int const pSize = mPg->getSize(); + int const pgRank = mPg->getRank(); + TORCH_CHECK(static_cast(peer_recv_ptrs.size()) == pSize, "peer_recv_ptrs size must equal world_size"); + + int const nPeers = pSize - 1; + if (nPeers == 0) + { + // P=1: self-only; recv_buf already populated by the permute kernel. + return; + } + + char const* sendBase = static_cast(send_buf.data_ptr()); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + cudaStreamCaptureStatus captureStatus; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(stream, &captureStatus)); + bool const underCapture = (captureStatus != cudaStreamCaptureStatusNone); + + if (!underCapture) + { + std::vector dsts; + dsts.reserve(nPeers); + std::vector srcs; + srcs.reserve(nPeers); + std::vector sizes; + sizes.reserve(nPeers); + for (int p = 0; p < pSize; ++p) + { + if (p == pgRank) + continue; + void* peerBase = reinterpret_cast(peer_recv_ptrs[p]); + dsts.push_back( + static_cast(peerBase) + static_cast(pgRank) * static_cast(slot_bytes)); + srcs.push_back(sendBase + static_cast(p) * static_cast(slot_bytes)); + sizes.push_back(static_cast(slot_bytes)); + } + cudaMemcpyAttributes attrs[1]; + std::memset(&attrs[0], 0, sizeof(attrs[0])); + attrs[0].srcAccessOrder = cudaMemcpySrcAccessOrderStream; + attrs[0].flags = 1u; + size_t attrIdxs[1] = {0}; + TLLM_CUDA_CHECK(cudaMemcpyBatchAsync( + dsts.data(), srcs.data(), sizes.data(), static_cast(nPeers), attrs, attrIdxs, 1, stream)); + } + else + { + for (int p = 0; p < pSize; ++p) + { + if (p == pgRank) + continue; + void* peerBase = reinterpret_cast(peer_recv_ptrs[p]); + void* dst + = static_cast(peerBase) + static_cast(pgRank) * static_cast(slot_bytes); + void const* src = sendBase + static_cast(p) * static_cast(slot_bytes); + TLLM_CUDA_CHECK( + cudaMemcpyAsync(dst, src, static_cast(slot_bytes), cudaMemcpyDeviceToDevice, stream)); + } + } + } + + // Phase 2 (fence): PT symm-mem barrier on the current CUDA stream. + // Any allocated slot's handle works — they all belong to the same group. + void emitBarrier() + { + TLLM_CHECK_WITH_INFO( + mCanonicalHandle, "emitBarrier: no slot allocated yet — _prepare must precede the first _async barrier."); + // 10s timeout: on hang, the kernel traps with rank+channel diagnostic instead of spinning silently + // until SLURM wall-clock kills. Generous enough to absorb first-touch IPC + first cuda_graph + // capture jitter. channel=0: V/Q/K issues all run on the same per-device side stream so + // they FIFO-serialize; channel multiplexing only matters across distinct streams. + mCanonicalHandle->barrier(/*channel=*/0, /*timeout_ms=*/10000); + } + +private: + struct Slot + { + // PT _SymmetricMemory-backed recv buffer (peer-writable). + at::Tensor symm_tensor; + c10::intrusive_ptr handle; + void* basePtr = nullptr; // aliases symm_tensor.data_ptr() + size_t size = 0; + std::vector peerPtrs; // from handle->get_buffer_ptrs() + + // Local-only push source (no symm-mem). cudaMalloc'd eagerly to + // stay cuda_graph-capture-safe. + void* sendBuf = nullptr; + size_t sendBufBytes = 0; + }; + + // One-shot per process per group: register PG's (name, rank, size, store) + // with PT symm-mem's group registry. Subsequent rendezvous() calls reuse it. + void ensureGroupRegistered() + { + static std::set sRegistered; + static std::mutex sMutex; + std::string const& name = mPg->getGroupName(); + std::lock_guard lock(sMutex); + if (sRegistered.count(name)) + { + return; + } + c10d::symmetric_memory::set_group_info(name, mPg->getRank(), mPg->getSize(), mPg->getStore()); + sRegistered.insert(name); + } + + int nextSlotIdx() + { + std::lock_guard lock(mNextMutex); + int idx = mNextIdx; + mNextIdx = (mNextIdx + 1) % kNumSlots; + return idx; + } + + // Lazy collective allocator. Cached when size is sufficient; reallocates + // (releasing the old handle) on size-up. All ranks must reach this in the + // same order (collective rendezvous). + // + // Commit-on-success: every allocation step writes to local variables + // first, and the cached `slot` is mutated only after all steps succeed. + // If `empty_strided_p2p`, `rendezvous`, `get_buffer_ptrs`, or `cudaMalloc` + // throws mid-way, the local at::Tensor / intrusive_ptr clean up via RAII + // and the previously-cached slot remains untouched (so the next call + // either retries or reuses the still-valid prior state). + Slot& getOrAllocSlot(int slotIdx, size_t requiredSize) + { + std::lock_guard lock(mSlotsMutex); + Slot& slot = mSlots[slotIdx]; + + if (slot.basePtr != nullptr && slot.size >= requiredSize) + { + return slot; + } + + // First-time / size-up allocation is NOT capture-safe: + // empty_strided_p2p + rendezvous + cudaMalloc all violate stream + // capture invariants. Caller must warm up out-of-capture so the slot + // is allocated and cached before any cuda_graph capture begins. + cudaStream_t const stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(stream, &captureStatus)); + TORCH_CHECK(captureStatus == cudaStreamCaptureStatusNone, + "async-ulysses: slot allocation (empty_strided_p2p + rendezvous + cudaMalloc) " + "is not graph-capture-safe. Warm up the model out-of-capture (run one forward " + "pass before enabling cuda_graph capture) so slots are cached."); + + int currentDev = -1; + TLLM_CUDA_CHECK(cudaGetDevice(¤tDev)); + c10::Device device(c10::DeviceType::CUDA, currentDev); + std::string const& groupName = mPg->getGroupName(); + int const pSize = mPg->getSize(); + + // Build new state in local variables — no mutation of `slot` yet. + at::Tensor newSymmTensor = c10d::symmetric_memory::empty_strided_p2p( + /*size=*/{static_cast(requiredSize)}, /*stride=*/{1}, + /*dtype=*/at::kByte, device, + /*group_name=*/std::make_optional(groupName), /*alloc_id=*/std::nullopt); + auto newHandle = c10d::symmetric_memory::rendezvous(newSymmTensor, groupName); + TLLM_CHECK_WITH_INFO(newHandle, "rendezvous returned null handle"); + + auto ptrs = newHandle->get_buffer_ptrs(); + TLLM_CHECK_WITH_INFO( + static_cast(ptrs.size()) == pSize, "get_buffer_ptrs size %zu != world_size %d", ptrs.size(), pSize); + std::vector newPeerPtrs(ptrs.begin(), ptrs.end()); + + // cudaMalloc last so any throw above is cleaned up by newSymmTensor / + // newHandle RAII without leaking GPU memory. + void* newSendBuf = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&newSendBuf, requiredSize)); + + // All allocations succeeded — commit. Free old sendBuf (the raw void* + // isn't owned by any RAII type in Slot); the at::Tensor / intrusive_ptr + // fields are released by move-assign. + if (slot.sendBuf != nullptr) + { + (void) cudaFree(slot.sendBuf); + } + slot.symm_tensor = std::move(newSymmTensor); + slot.handle = std::move(newHandle); + slot.basePtr = slot.symm_tensor.data_ptr(); + slot.size = requiredSize; + slot.peerPtrs = std::move(newPeerPtrs); + slot.sendBuf = newSendBuf; + slot.sendBufBytes = requiredSize; + + // Cache the first allocated handle for emitBarrier() (any handle from + // this PG yields the same channel-N barrier semantics). + if (!mCanonicalHandle) + { + mCanonicalHandle = slot.handle; + } + + return slot; + } + + c10::intrusive_ptr mPg; + + int mNextIdx{0}; + std::mutex mNextMutex; + + std::array mSlots{}; + std::mutex mSlotsMutex; + + // Cached on the first slot allocation. SymmetricMemory::barrier() is a + // PG-level sync (any handle from this PG triggers the same channel-N + // barrier), so emitBarrier() can use this directly instead of scanning + // mSlots for a non-null handle on every call. + c10::intrusive_ptr mCanonicalHandle; +}; + +// Process-lifetime cache of AsyncUlyssesOp instances keyed by group_name. +static std::shared_ptr getOrCreateOp(c10::intrusive_ptr const& pg) +{ + TLLM_CHECK_WITH_INFO(pg, "ProcessGroup is null"); + static std::map> sCache; + static std::mutex sMutex; + std::string const& groupName = pg->getGroupName(); + std::lock_guard lock(sMutex); + auto it = sCache.find(groupName); + if (it != sCache.end()) + { + return it->second; + } + auto op = std::make_shared(pg); + op->initialize(); + sCache[groupName] = op; + return op; +} + +// Step 1 (caller's compute stream): acquire slot ring entry + CUDA C +// permute+scatter (writes peer chunks to send_buf, self chunk directly to +// recv_buf[my_rank]). Returns the 5D recv-buf view (for downstream SDPA) +// and an opaque SendHandle that the second op consumes. +std::tuple> ulysses_a2a_async_prepare( + torch::Tensor input_4d, c10::intrusive_ptr const& pg) +{ + TORCH_CHECK(input_4d.is_cuda(), "input must be on CUDA"); + TORCH_CHECK(input_4d.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(input_4d.dim() == 4, "input must be [B, S_local, H, D]"); + TORCH_CHECK(input_4d.scalar_type() == at::ScalarType::BFloat16, "bf16 only"); + + // Bind current device + slot allocator + kernel launch to the input's + // device. `getOrAllocSlot` reads `cudaGetDevice()`, and the kernel stream + // is taken from `input_4d.get_device()`; without this guard the two can + // diverge (e.g. caller forgot a torch.cuda.set_device) → slot allocated + // on dev A, kernel launched on dev B → illegal memory access. + c10::cuda::CUDAGuard device_guard(input_4d.device()); + + int const B = static_cast(input_4d.size(0)); + int const S_local = static_cast(input_4d.size(1)); + int const H = static_cast(input_4d.size(2)); + int const D = static_cast(input_4d.size(3)); + TORCH_CHECK(D % 8 == 0, "D must be divisible by 8 (int4 vec)"); + + auto op = getOrCreateOp(pg); + + int const P = op->getPgSize(); + int const my_rank = op->getPgRank(); + TORCH_CHECK(H % P == 0, "H must be divisible by world_size"); + int const H_local = H / P; + + auto [send_t, recv_t, peer_recv_ptrs, slot_bytes] = op->acquireSlotPair( + {(int64_t) P, (int64_t) B, (int64_t) S_local, (int64_t) H_local, (int64_t) D}, input_4d.scalar_type()); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(input_4d.get_device()).stream(); + tensorrt_llm::kernels::launchUlyssesPermuteScatter( + input_4d.data_ptr(), send_t.data_ptr(), recv_t.data_ptr(), my_rank, B, S_local, H, D, P, stream); + + auto send_h = c10::make_intrusive(); + send_h->send_t = std::move(send_t); + send_h->peer_recv_ptrs = std::move(peer_recv_ptrs); + send_h->slot_bytes = slot_bytes; + send_h->group_name = pg->getGroupName(); + + return std::make_tuple(std::move(recv_t), send_h); +} + +// Step 2a (caller's comm stream): CE push only, no barrier. Issue V/Q/K +// pushes back-to-back on the side stream so they FIFO through copy-engines +// without barrier-induced stalls; defer all fences to `ulysses_a2a_async_barrier` +// at join time. Caller must event-sync from compute stream before calling. +// +// Reject cross-PG handle use: peer_recv_ptrs are valid only in the symm-mem +// group registered for the PG that produced this handle. Two PGs of the same +// size would otherwise pass the peer-count check inside runCePush and silently +// push into the wrong group's buffers. +void ulysses_a2a_async_push( + c10::intrusive_ptr const& send_h, c10::intrusive_ptr const& pg) +{ + TORCH_CHECK(send_h.get() != nullptr, "send_h is null"); + TORCH_CHECK(send_h->send_t.defined(), "send_h.send_t is undefined"); + TORCH_CHECK(send_h->group_name == pg->getGroupName(), "SendHandle was produced by ProcessGroup '", + send_h->group_name, "' but ulysses_a2a_async_push was called with ProcessGroup '", pg->getGroupName(), + "'. Handle and PG must match."); + auto op = getOrCreateOp(pg); + op->runCePush(send_h->send_t, send_h->peer_recv_ptrs, send_h->slot_bytes); +} + +// Step 2b (caller's comm stream): emit a symm-mem barrier on channel 0. +// Pairs with `ulysses_a2a_async_push`; one call per deferred push (e.g. +// V/Q/K -> 3 barriers at join). +void ulysses_a2a_async_barrier(c10::intrusive_ptr const& pg) +{ + auto op = getOrCreateOp(pg); + op->emitBarrier(); +} + +} // namespace + +#endif // ENABLE_MULTI_DEVICE + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.class_("SendHandle"); + + m.def( + "ulysses_a2a_async_prepare(Tensor input, " + "__torch__.torch.classes.c10d.ProcessGroup pg) " + "-> (Tensor, __torch__.torch.classes.trtllm.SendHandle)"); + m.def( + "ulysses_a2a_async_push(__torch__.torch.classes.trtllm.SendHandle send_h, " + "__torch__.torch.classes.c10d.ProcessGroup pg) -> ()"); + m.def("ulysses_a2a_async_barrier(__torch__.torch.classes.c10d.ProcessGroup pg) -> ()"); +} + +// Both ops take/return a custom-class handle, not tensors, so the dispatcher +// can't pick a backend from input types. Register on CompositeExplicitAutograd +// (the underlying CUDA work runs on the caller's current CUDA stream). +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ +#if ENABLE_MULTI_DEVICE + m.impl("ulysses_a2a_async_prepare", &tensorrt_llm::torch_ext::ulysses_a2a_async_prepare); + m.impl("ulysses_a2a_async_push", &tensorrt_llm::torch_ext::ulysses_a2a_async_push); + m.impl("ulysses_a2a_async_barrier", &tensorrt_llm::torch_ext::ulysses_a2a_async_barrier); +#endif +} diff --git a/cpp/tensorrt_llm/thop/ulyssesPermuteScatterOp.cpp b/cpp/tensorrt_llm/thop/ulyssesPermuteScatterOp.cpp new file mode 100644 index 000000000000..5637b49a9a73 --- /dev/null +++ b/cpp/tensorrt_llm/thop/ulyssesPermuteScatterOp.cpp @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2022-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. + */ + +#include "tensorrt_llm/kernels/ulyssesPermuteScatterKernel.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Standalone Python entry point for ulyssesPermuteScatterKernel — used by +// unit tests. The production caller is the Ulysses async A2A path in +// alltoallOp.cpp, which combines this kernel with IPC writes + an LSA +// barrier; that whole sequence is multi-rank and not amenable to a +// single-GPU pytest. By exposing just the kernel, we can validate the +// permute+scatter layout transform independently. +// +// Layout: +// input : bf16 [B, S_local, H, D] contiguous +// send_buf : bf16 [P, B, S_local, H/P, D] contiguous +// recv_buf : bf16 [P, B, S_local, H/P, D] contiguous +// For each (b, s, h, d): +// peer = h // (H/P) +// h_local = h % (H/P) +// dst = (recv_buf if peer == my_rank else send_buf) +// slot = peer // applies to both branches +// dst[slot, b, s, h_local, d] = input[b, s, h, d] +void ulysses_permute_scatter(torch::Tensor& input, // [B, S_local, H, D] + torch::Tensor& send_buf, // [P, B, S_local, H/P, D] + torch::Tensor& recv_buf, // [P, B, S_local, H/P, D] + int64_t my_rank, int64_t P) +{ + TORCH_CHECK(input.dim() == 4, "input must be 4D [B, S_local, H, D]"); + TORCH_CHECK(send_buf.dim() == 5 && recv_buf.dim() == 5, "send_buf / recv_buf must be 5D [P, B, S_local, H/P, D]"); + CHECK_INPUT(input, torch::kBFloat16); + CHECK_INPUT(send_buf, torch::kBFloat16); + CHECK_INPUT(recv_buf, torch::kBFloat16); + + // Validate P first: P=0 would crash on `H % P` below (UB; SIGFPE on x86) + // and on `my_rank < P` after `0 <= my_rank` trivially passes. + TORCH_CHECK(P > 0, "P (world_size) must be positive"); + + int64_t const B = input.size(0); + int64_t const S_local = input.size(1); + int64_t const H = input.size(2); + int64_t const D = input.size(3); + TORCH_CHECK(H % P == 0, "H must be divisible by P"); + TORCH_CHECK(D % 8 == 0, "D must be divisible by 8 (uint4 vec)"); + int64_t const H_local = H / P; + TORCH_CHECK(send_buf.size(0) == P && send_buf.size(1) == B && send_buf.size(2) == S_local + && send_buf.size(3) == H_local && send_buf.size(4) == D, + "send_buf shape mismatch"); + TORCH_CHECK(recv_buf.size(0) == P && recv_buf.size(1) == B && recv_buf.size(2) == S_local + && recv_buf.size(3) == H_local && recv_buf.size(4) == D, + "recv_buf shape mismatch"); + TORCH_CHECK(0 <= my_rank && my_rank < P, "my_rank out of range"); + + // Empty-tensor no-op: B=0 or S_local=0 produces zero grid extent in the + // kernel launcher (undefined cuLaunchKernel behavior across CUDA versions). + // Shape consistency between input/send_buf/recv_buf was already enforced + // above, so empty input implies empty buffers — nothing to write. + if (input.numel() == 0) + { + return; + } + + auto stream = at::cuda::getCurrentCUDAStream(); + tensorrt_llm::kernels::launchUlyssesPermuteScatter(input.data_ptr(), send_buf.data_ptr(), recv_buf.data_ptr(), + static_cast(my_rank), static_cast(B), static_cast(S_local), static_cast(H), + static_cast(D), static_cast(P), stream); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "ulysses_permute_scatter(Tensor(a!) input, Tensor(b!) send_buf, Tensor(c!) recv_buf, " + "int my_rank, int P) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("ulysses_permute_scatter", &ulysses_permute_scatter); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp new file mode 100644 index 000000000000..9d45b4f0b119 --- /dev/null +++ b/cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp @@ -0,0 +1,120 @@ +/* + * Copyright (c) 2022-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. + */ + +#include "tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// 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 ALWAYS writes NHD-contig storage [B, P*Sp, H, D]. The returned +// tensor shape depends on ``layout``: +// layout=0 (HND) → returns transpose-view [B, H, P*Sp, D] +// (HND-shape, NHD-stride, NON-contig — mirrors the +// `q.transpose(1, 2)` result in the sync `_forward_unfused` +// path, which lets cudnn SDPA preserve NHD-stride through +// its output and collapses the downstream +// `_output_a2a.transpose(1, 2).contiguous()` to a no-op) +// layout=1 (NHD) → returns storage as-is [B, P*Sp, H, D] (NHD contig) +// Replaces the eager chain +// t.permute(1, 0, 2, 3, 4).reshape(B, P * Sp, H, D).contiguous() // NHD +// [.transpose(1, 2)] // HND: stride view only +// for Q, K, V in one kernel launch. +std::tuple ulysses_post_unscatter_qkv( + torch::Tensor& q_in, // [P, B, Sp, H, D] + torch::Tensor& k_in, // [P, B, Sp, H, D] + torch::Tensor& v_in, // [P, B, Sp, H, D] + int64_t layout) // 0 = HND, 1 = NHD +{ + TORCH_CHECK(q_in.dim() == 5 && k_in.dim() == 5 && v_in.dim() == 5, + "ulysses_post_unscatter_qkv expects 5D tensors [P, B, Sp, H, D]"); + TORCH_CHECK(q_in.sizes() == k_in.sizes() && q_in.sizes() == v_in.sizes(), "Q/K/V must share the same shape"); + TORCH_CHECK(layout == 0 || layout == 1, "layout must be 0 (HND) or 1 (NHD), got ", layout); + + CHECK_INPUT(q_in, torch::kBFloat16); + CHECK_INPUT(k_in, torch::kBFloat16); + CHECK_INPUT(v_in, torch::kBFloat16); + + // D % 8 enforced here at op boundary (mirrors sibling ulysses_permute_scatter). + // Without this, torch::empty allocates the three output tensors before the + // kernel launcher's TLLM_CHECK_WITH_INFO fires, producing a less-actionable + // error path. Vec width is 8 elements for bf16 (16-byte vectorized stores). + TORCH_CHECK(q_in.size(-1) % 8 == 0, "D (last dim) must be divisible by 8 (bf16 vec=8)"); + + int64_t const P = q_in.size(0); + int64_t const B = q_in.size(1); + int64_t const Sp = q_in.size(2); + int64_t const H = q_in.size(3); + int64_t const D = q_in.size(4); + + bool const is_hnd = (layout == 0); + auto opts = q_in.options(); + // Always allocate NHD-contig storage [B, P*Sp, H, D]. + auto const storage_shape = std::vector{B, P * Sp, H, D}; + auto q_out = torch::empty(storage_shape, opts); + auto k_out = torch::empty(storage_shape, opts); + auto v_out = torch::empty(storage_shape, opts); + + // Empty-tensor no-op: P=0/B=0/Sp=0 produces zero grid extent in the + // kernel launcher (undefined cuLaunchKernel behavior across CUDA versions). + // The three output tensors above are already empty-shaped via P*Sp=0 or B=0, + // so returning them directly preserves the output-shape contract. + if (q_in.numel() == 0) + { + 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); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + tensorrt_llm::kernels::launchUlyssesPostUnscatter(q_in.data_ptr(), k_in.data_ptr(), v_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); + + // HND callers get a transpose-view of the NHD storage (zero-copy stride + // reinterpretation). NHD callers get the storage as-is. + 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)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("ulysses_post_unscatter_qkv", &ulysses_post_unscatter_qkv); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index b6f125ed027c..0fe823ced41b 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -172,6 +172,7 @@ Configured under `VisualGenArgs.parallel_config`. Modes can be combined: - **CFG Parallelism** (`cfg_size: 2`): Splits positive/negative guidance prompts across GPUs. - **Ulysses Parallelism** (`ulysses_size: N`): Splits the sequence dimension across GPUs for longer sequences. + - **Async Ulysses A2A pipeline** (`async_ulysses: true` in `parallel_config`): Overlaps per-rank V/Q/K projection compute with the cross-rank all-to-all on a dedicated side stream. Requires `ulysses_size > 1` and an NVLink-connected GPU domain (uses PyTorch `_SymmetricMemory` with CUDA IPC for peer pushes; not currently supported across nodes without MNNVL). Currently wired for WAN and LTX-2 self-attention. - **Parallel VAE** (`parallel_vae_size: N`): Shards the final VAE decode along a spatial axis (constraint: `parallel_vae_size ≤ world_size`; WAN/Cosmos3 only). - **Context Parallel (CP)** — Partitions the sequence into shards so that each rank computes partial attention. Requires an LSE-capable attention backend (`FA4` or `CUTEDSL`). CP can be composed with Ulysses, giving a total sequence-parallel (SP) degree = `cp_size · ulysses_size`. The CP degree depends on the implementation below: - **Attention2D** (`attn2d_size: [N, M]`): Shards the sequence axis across an `N × M` device mesh (CP degree = `N · M`; total SP degree = `N · M · ulysses_size`). diff --git a/examples/visual_gen/configs/ltx2-4gpu.yaml b/examples/visual_gen/configs/ltx2-4gpu.yaml new file mode 100644 index 000000000000..a667a147c9e5 --- /dev/null +++ b/examples/visual_gen/configs/ltx2-4gpu.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +# 4-GPU LTX-2 AudioVideo parallel config (precision-agnostic — picks up +# whatever checkpoint is passed via --model_path). Shared by offline +# examples (--extra_visual_gen_options) and trtllm-serve. +attention_config: + backend: VANILLA +parallel_config: + cfg_size: 2 + ulysses_size: 2 + async_ulysses: true +torch_compile_config: + enable: true +cuda_graph_config: + enable: true diff --git a/examples/visual_gen/configs/wan2.2-t2v-fp4-4gpu.yaml b/examples/visual_gen/configs/wan2.2-t2v-fp4-4gpu.yaml index a645bbe7794b..82a033090aed 100644 --- a/examples/visual_gen/configs/wan2.2-t2v-fp4-4gpu.yaml +++ b/examples/visual_gen/configs/wan2.2-t2v-fp4-4gpu.yaml @@ -23,6 +23,7 @@ attention_config: parallel_config: cfg_size: 2 ulysses_size: 2 + async_ulysses: true parallel_vae_size: 4 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 b31bb56230c6..30ec7bfaa358 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1013,6 +1013,20 @@ def _(workspace, cp_rank, cp_size): # This op initializes workspace in-place and returns nothing return None + @torch.library.register_fake("trtllm::ulysses_post_unscatter_qkv") + def _(q_in, k_in, v_in, layout=0): + # Storage is always NHD-contig [B, P*Sp, H, D]. HND-shape return is a + # transpose-view (HND-shape, NHD-stride, non-contig) so Inductor sees + # the same stride pattern as the real op. + P, B, Sp, H, D = q_in.shape + nhd_shape = (B, P * Sp, H, D) + + def _mk(t): + base = t.new_empty(nhd_shape) + return base.transpose(1, 2) if layout == 0 else base + + return (_mk(q_in), _mk(k_in), _mk(v_in)) + @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 241aa8b2d166..bb1a84e6c30c 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -20,7 +20,7 @@ """ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, ClassVar, Dict, Optional import torch import torch.distributed as dist @@ -42,6 +42,25 @@ _flash_attn_combine_import_error = e +def post_permute_5d_to_4d(out_5d, P): + """5D [P, B, Sp, H/P, D] → 4D [B, P*Sp, H/P, D] (block-by-rank gather). + .contiguous() copies slot data out (layout-normalize for SDPA, decoupling + from IPC slot lifetime). Inductor fuses permute+contig with downstream + SDPA input prep.""" + _P, Bt, Spt, HpP, Dt = out_5d.shape + return out_5d.permute(1, 0, 2, 3, 4).contiguous().view(Bt, _P * Spt, HpP, Dt) + + +def _ulysses_post_unscatter(q_5d, k_5d, v_5d, *, is_hnd): + """One-launch fused replacement for the post-A2A 5D -> 4D chain. + + is_hnd=True -> output [B, H, P*Sp, D] (VANILLA / torch SDPA) + is_hnd=False -> output [B, P*Sp, H, D] (TRTLLM / FA4) + """ + layout = 0 if is_hnd else 1 + return torch.ops.trtllm.ulysses_post_unscatter_qkv(q_5d, k_5d, v_5d, layout) + + class UlyssesAttention(AttentionBackend): """ Ulysses Sequence Parallelism wrapper. @@ -65,10 +84,16 @@ class UlyssesAttention(AttentionBackend): + 1 for output (2 collectives total) """ + # One side stream shared across all UlyssesAttention instances on the + # same device. Per-layer streams inflate the stream count and break + # cuda_graph capture. + _side_stream_by_device: ClassVar[Dict[int, "torch.cuda.Stream"]] = {} + def __init__( self, inner_backend: AttentionBackend, process_group: torch.distributed.ProcessGroup, + async_ulysses: bool = False, ): self.inner_backend = inner_backend self.process_group = process_group @@ -83,6 +108,24 @@ def __init__( self.num_heads = self.sharded_num_heads * self.world_size self.num_kv_heads = self.sharded_num_kv_heads * self.world_size + # Async pipeline state. Eagerly populated when async_ulysses=True; + # forward_async assumes these are set. Non-async path doesn't touch + # them. + self._pg_boxed = None + self._async_side_stream: Optional[torch.cuda.Stream] = None + # Count of deferred pushes since the last `_join_async`. `_join_async` + # drains exactly this many `ulysses_a2a_async_barrier` calls on the + # side stream so V/Q/K pushes FIFO together without intermediate + # barrier kernels. + self._pending_barriers: int = 0 + if async_ulysses: + device = torch.cuda.current_device() + if device not in UlyssesAttention._side_stream_by_device: + UlyssesAttention._side_stream_by_device[device] = torch.cuda.Stream(device=device) + self._async_side_stream = UlyssesAttention._side_stream_by_device[device] + if process_group is not None: + self._pg_boxed = process_group.boxed() + def forward( self, q: torch.Tensor, @@ -166,6 +209,120 @@ def _forward_unfused( return self._output_a2a(output, batch_size, seq_len_full) + # ------------------------------------------------------------------ + # Split-QKV async A2A pipeline. `_issue_async` and `_join_async` are + # the only stream-switch boundaries and are @torch.compiler.disable'd; + # the caller's compiled forward fuses each compute_{q,k,v} closure. + # ------------------------------------------------------------------ + + @torch.compiler.disable(recursive=False) + def _issue_async(self, perm_4d: torch.Tensor) -> torch.Tensor: + """Issue one V/Q/K async a2a (CE push only; barrier deferred to join). + Phase 1 (acquire slot + CUDA C permute+scatter) runs on the CURRENT + (default) stream. Phase 2a (cudaMemcpyBatchAsync peer push) is queued + on the comm side stream, gated by an event so it waits for Phase 1. + Phase 2b (symm-mem barrier) is NOT issued here — `_join_async` drains + all pending barriers in one shot so V/Q/K pushes FIFO through CE + without intermediate barrier kernels splitting them up. Returns the + 5D recv-buf view. + + Comm-stream FIFO serializes consecutive V/Q/K pushes in caller order; + no explicit chain event is needed between them. The default stream + is free to immediately begin the next V/Q/K compute — that's where + the V_push ∥ Q_compute ∥ K_compute overlap comes from.""" + recv, send_h = torch.ops.trtllm.ulysses_a2a_async_prepare(perm_4d, self._pg_boxed) + ev = torch.cuda.Event() + ev.record() + with torch.cuda.stream(self._async_side_stream): + ev.wait() + torch.ops.trtllm.ulysses_a2a_async_push(send_h, self._pg_boxed) + self._pending_barriers += 1 + return recv + + @torch.compiler.disable(recursive=False) + def _join_async(self) -> None: + """Drain pending symm-mem barriers (one per deferred push) on the + side stream, then have the default stream wait on the tail event. + Comm-stream FIFO preserves [push V, push Q, push K, barrier, barrier, + barrier] order; all N barriers fire on channel=0 with identical + semantics, so the default stream sees a fully-synced recv buffer.""" + with torch.cuda.stream(self._async_side_stream): + for _ in range(self._pending_barriers): + torch.ops.trtllm.ulysses_a2a_async_barrier(self._pg_boxed) + ev_done = torch.cuda.Event() + ev_done.record() + self._pending_barriers = 0 + torch.cuda.current_stream().wait_event(ev_done) + + def forward_async( + self, + compute_q: Callable[[], torch.Tensor], + compute_k: Callable[[], torch.Tensor], + compute_v: Callable[[], torch.Tensor], + **attn_kwargs, + ) -> torch.Tensor: + """Run the async ulysses attention path (V/Q/K rolling A2A). + + Args: + compute_q / compute_k / compute_v : caller-provided closures that + each return a 4D tensor `[B, S_local, H, D]`. The closure + typically does `GEMM → (RMSNorm) → (RoPE) → view(4D)`; closures + live in the caller's compiled forward so inductor fuses each + into a single Triton kernel. + **attn_kwargs : forwarded to the wrapped inner attention backend + (mask, scale, etc.). + + Returns: + output tensor in the caller's sharded layout `[B, S/P, H, D]`. + + Pipeline: V/Q/K computed in V→Q→K order on the default stream; each + compute's output is fed to `_issue_async` which queues push+barrier on + the comm side stream. Default stream proceeds to the next compute + immediately, so V's push overlaps with Q's compute, Q's push overlaps + with K's compute. `_join_async` makes default wait on the last push. + Post-attention permute / SDPA / reverse A2A run in the caller's outer + compile region for additional inductor fusion.""" + P = self.world_size + + v_4d = compute_v() + v_5d = self._issue_async(v_4d) + + q_4d = compute_q() + q_5d = self._issue_async(q_4d) + + k_4d = compute_k() + k_5d = self._issue_async(k_4d) + + self._join_async() + + # Fast path: one fused kernel replaces the eager post-A2A chain + # (6 ops for HND target: permute+reshape+contig + transpose+contig + # per Q/K/V; 3 ops for NHD target). bf16-only because the kernel is + # 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 + 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 + seq_len_full = P * Sp_q + else: + v_out = post_permute_5d_to_4d(v_5d, P) + q_out = post_permute_5d_to_4d(q_5d, P) + k_out = post_permute_5d_to_4d(k_5d, P) + + B = q_out.shape[0] + seq_len_full = q_out.shape[1] + if is_hnd: + q_out = q_out.transpose(1, 2).contiguous() + k_out = k_out.transpose(1, 2).contiguous() + v_out = v_out.transpose(1, 2).contiguous() + + attn_kwargs["seq_len"] = seq_len_full + attn_kwargs["seq_len_kv"] = seq_len_full + output = self.inner_backend.forward(q=q_out, k=k_out, v=v_out, **attn_kwargs) + return self._output_a2a(output, B, seq_len_full) + def _output_a2a( self, output: torch.Tensor, @@ -622,6 +779,7 @@ def wrap_parallel_attention( *, visual_gen_mapping: Optional["VisualGenMapping"] = None, enable_sequence_parallel: bool = True, + async_ulysses: bool = False, ) -> AttentionBackend: """Wrap a compute backend with the configured parallelism strategy. @@ -650,5 +808,9 @@ def wrap_parallel_attention( attn = RingAttention(attn, process_group=vgm.ring_group) if ulysses_size > 1: - attn = UlyssesAttention(attn, process_group=vgm.ulysses_group) + attn = UlyssesAttention( + attn, + process_group=vgm.ulysses_group, + async_ulysses=async_ulysses, + ) return attn diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.py index 6e75448d435e..646624f8bd10 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.py @@ -15,9 +15,13 @@ def to_velocity( velocity = (sample - denoised) / sigma """ + # Tensor sigma: keep on device. `.item()` would force a D2H sync that + # deadlocks under nsys profiling combined with CUDA graph replay. + # The scheduler guarantees sigma > 0 inside the denoise loop, so we + # skip the zero check on the tensor path (re-checking would re-sync). if isinstance(sigma, torch.Tensor): - sigma = sigma.to(calc_dtype).item() - if sigma == 0: + sigma = sigma.to(calc_dtype) + elif sigma == 0: raise ValueError("Sigma can't be 0.0") return ((sample.to(calc_dtype) - denoised.to(calc_dtype)) / sigma).to(sample.dtype) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index ce09797c389c..fe7f40bec00d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -29,6 +29,7 @@ from tensorrt_llm._torch.modules.linear import Linear, WeightMode from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import Fp4QuantizedTensor from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader @@ -87,6 +88,8 @@ def __init__( config: Optional["DiffusionModelConfig"] = None, layer_idx: int = 0, enable_sequence_parallel: bool = False, + use_ulysses: bool = False, + async_ulysses: bool = False, ): from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig @@ -98,10 +101,24 @@ def __init__( self.rope_type = rope_type self._is_cross_attn = context_dim is not None + # Async ulysses opt-in: V/Q/K GEMMs interleave with the all-to-all on a + # side stream. Forces SEPARATE_QKV so the 3 projections can issue + # independently. + self._use_async_ulysses = bool( + use_ulysses + and not self._is_cross_attn + and async_ulysses + and vgm is not None + and vgm.ulysses_size > 1 + ) + # Self-attention: FUSE_QKV enables the optimized backend + auto Ulysses # wrapping from the base class. - # Cross-attention: SEPARATE_QKV since K/V come from a different source. - qkv_mode = QKVMode.SEPARATE_QKV if self._is_cross_attn else QKVMode.FUSE_QKV + # Cross-attention or async ulysses: SEPARATE_QKV. + if self._is_cross_attn or self._use_async_ulysses: + qkv_mode = QKVMode.SEPARATE_QKV + else: + qkv_mode = QKVMode.FUSE_QKV # Caller opts in via enable_sequence_parallel. Cross-attn supports # Ulysses-only (SEPARATE_QKV + ring/attn2d is rejected in Attention); @@ -132,13 +149,11 @@ def __init__( config=config, layer_idx=layer_idx, enable_sequence_parallel=enable_sp, + enable_ulysses=use_ulysses, + async_ulysses=self._use_async_ulysses, ) - # Build a runtime-toggleable Ulysses ↔ plain pair. - # Self-attn: audio length isn't always divisible by ulysses_size, so - # we need a plain fallback. Cross-attn (v2a): same need — audio Q is - # padded when divisible, plain backend is used otherwise. Plain has - # to be built with the full (unsharded) head count. + # Validate Ulysses head divisibility (from main). self._has_dual_attn = False if enable_sp and ulysses_size > 1: U = ulysses_size @@ -152,6 +167,12 @@ def __init__( # Base class already built `self.attn` as the Ulysses-wrapped path # (sharded inner backend + UlyssesAttention) for both self-attn and # cross-attn paths. + + # For audio self-attention that may need a runtime Ulysses toggle + # (sequence length not always divisible by ulysses_size), create a + # plain backend as fallback. The base class already set self.attn + # to UlyssesAttention(inner_backend=sharded_backend). + if use_ulysses and not self._is_cross_attn and ulysses_size > 1: self._ulysses_attn = self.attn self._plain_attn = create_attention( backend=self.attn_backend, @@ -286,7 +307,9 @@ def forward( Caller contract: - FUSE_QKV (self-attn): pe must be set; k_pe and pre_projected_kv unused. - SEPARATE_QKV (cross-attn): cached path requires pre_projected_kv; - uncached path requires `context`. pe optional (None = norm-only). + uncached path uses ``context`` (may be None when the async-Ulysses + inner backend was swapped to a non-async one — falls back to + self-attn via kv_source=x). pe optional (None = norm-only). k_pe overrides pe for K (e.g. AV cross-attn) when provided. Args: @@ -297,46 +320,92 @@ def forward( silently ignores it. ``LTX2Attention`` constructs ``audio_attn1`` with a VANILLA backend whenever Ulysses is active under a TRTLLM backend config (see ``_init_audio_modules``). + + Routing: + 1. Async-Ulysses self-attn → ``forward_async`` (V/Q/K rolling A2A). + 2. FUSE_QKV self-attn → packed fused kernel (or naive mini-config). + 3. SEPARATE_QKV cross-attn → split fused kernel (or naive mini-config). """ - # Fallback to the naive eager rope path when fusion is disabled or - # the kernel doesn't support this head_dim. LTX-2 prod has - # fuse_qk_norm_rope=True and head_dim ∈ {64, 128}, so this branch - # only fires under mini-config unit tests (head_dim=32). - if not self.fuse_qk_norm_rope or self.head_dim not in (64, 128): - return self._forward_unfused(x, context, pe, k_pe, pre_projected_kv, key_padding_mask) + # Async-Ulysses self-attn dispatch. ``hasattr`` guard: audio_attn1 may + # have ``set_ulysses_active(False)`` swap ``self.attn`` to a plain + # backend that lacks ``forward_async`` — fall through to the sync + # uncached SEPARATE_QKV branch, which handles context=None via + # kv_source=x. + if ( + self.qkv_mode == QKVMode.SEPARATE_QKV + and self._use_async_ulysses + and context is None + and pre_projected_kv is None + and hasattr(self.attn, "forward_async") + ): + return self.forward_async(x, freqs=pe) + + # Fused gate: prod uses fused kernels (head_dim ∈ {64, 128}); mini-config + # tests (head_dim=32) fall to naive ops. + use_fused = self.fuse_qk_norm_rope and self.head_dim in (64, 128) and self.qk_norm if self.qkv_mode == QKVMode.FUSE_QKV: - # ─── self-attn → packed kernel (norm + rope on QKV in-place) ─── - qkv = self.qkv_proj(x) - cos, sin = pe - self.apply_packed_qk_norm_rope(qkv, cos, sin) - q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) + # ─── sync self-attn ─── + if use_fused and pe is not None: + # Fused packed kernel: norm + RoPE on QKV in-place. + qkv = self.qkv_proj(x) + cos, sin = pe + self.apply_packed_qk_norm_rope(qkv, cos, sin) + q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) + else: + # Naive (mini-config head_dim ∉ {64, 128}). + q, k, v = self.get_qkv(x) + if self.qk_norm: + q = self.norm_q(q) + k = self.norm_k(k) + if pe is not None: + q = apply_rotary_emb(q, pe, self.rope_type) + k = apply_rotary_emb(k, pe, self.rope_type) elif self.qkv_mode == QKVMode.SEPARATE_QKV: - # ─── cross-attn → split kernel (norm or norm+rope based on pe) ─── if pre_projected_kv is not None: - # K/V cached by caller (text cross-attn + AV cross-attn). - # The caller is responsible for any K-norm + K-rope on the - # cached tensor; we only fuse Q here. + # ─── cached cross-attn (text + AV cross-attn) ─── + # K/V cached by caller; we only norm+RoPE Q here. k, v = pre_projected_kv q = self.to_q(x) - self.apply_split_norm_or_norm_rope( - q, self.norm_q.weight, self.num_attention_heads, pe - ) + if use_fused: + self.apply_split_norm_or_norm_rope( + q, self.norm_q.weight, self.num_attention_heads, pe + ) + else: + if self.qk_norm: + q = self.norm_q(q) + if pe is not None: + q = apply_rotary_emb(q, pe, self.rope_type) else: - # Uncached cross-attn (not exercised by LTX-2 in practice; kept for fuse-dispatch consistency). + # ─── uncached cross-attn / async self-attn fallback ─── + # LTX-2 prod doesn't use uncached cross-attn (always pre-projects + # K/V). This branch also catches async self-attn when the inner + # backend lacks forward_async (audio Ulysses-inactive swap): + # context=None then, fall back to self-attn via kv_source=x. + kv_source = context if context is not None else x q = self.to_q(x) - k = self.to_k(context) - v = self.to_v(context) - self.apply_split_norm_or_norm_rope( - q, self.norm_q.weight, self.num_attention_heads, pe - ) - self.apply_split_norm_or_norm_rope( - k, - self.norm_k.weight, - self.num_key_value_heads, - k_pe if k_pe is not None else pe, - ) + k = self.to_k(kv_source) + v = self.to_v(kv_source) + if use_fused: + self.apply_split_norm_or_norm_rope( + q, self.norm_q.weight, self.num_attention_heads, pe + ) + self.apply_split_norm_or_norm_rope( + k, + self.norm_k.weight, + self.num_key_value_heads, + k_pe if k_pe is not None else pe, + ) + else: + if self.qk_norm: + q = self.norm_q(q) + k = self.norm_k(k) + if pe is not None: + q = apply_rotary_emb(q, pe, self.rope_type) + k_pe_use = k_pe if k_pe is not None else pe + if k_pe_use is not None: + k = apply_rotary_emb(k, k_pe_use, self.rope_type) attn_kwargs = {} if key_padding_mask is not None: @@ -353,58 +422,87 @@ def forward( return self.to_out[0](out) - def _forward_unfused( + def forward_async( self, x: torch.Tensor, - context: torch.Tensor | None, - pe: tuple[torch.Tensor, torch.Tensor] | None, - k_pe: tuple[torch.Tensor, torch.Tensor] | None, - pre_projected_kv: tuple[torch.Tensor, torch.Tensor] | None, - key_padding_mask: torch.Tensor | None = None, + freqs: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: - """Fallback path for unsupported configs (head_dim ∉ {64, 128} or - ``fuse_qk_norm_rope=False``). + """LTX-2 async-Ulysses self-attn driver. Structurally mirrors base + ``Attention.forward_async`` (single function, fused/unfused branches) + but uses LTX-2's ``apply_rotary_emb`` (with ``rope_type``) on the + unfused fallback and injects gated-attention scaling in 4D between + the attn output and ``to_out``. - LTX-2 prod hardcodes ``fuse_qk_norm_rope=True`` and head_dim ∈ - {64, 128}, so in production this is never entered. Exercised by the - mini-config unit tests in ``test_ltx2_transformer.py`` (head_dim=32) - and by ablation tests that explicitly disable fusion. + Precondition: caller in ``LTX2Attention.forward`` gates on + ``_use_async_ulysses`` + ``hasattr(self.attn, "forward_async")``. - Contract: caller must pass *pe* / *k_pe* in 4D layout - ([B, T, H, D] for SPLIT rope, [B, T, D] for INTERLEAVED). The fused - kernel's 2D form is not compatible with the naive ``apply_rotary_emb``. + Returns 3D ``[B, S, H*D]`` matching ``forward``'s output contract. """ - if pre_projected_kv is not None: - k, v = pre_projected_kv - q = self.to_q(x) + B, S, _ = x.shape + H = self.num_attention_heads + KV = self.num_key_value_heads + D = self.head_dim + # Mirrors LTX2Attention.forward's fused gate; qkv_mode is implicitly + # SEPARATE_QKV under async (caller-enforced). head_dim check matches + # the fused split kernel's HEAD_DIM template instantiations {64, 128}. + use_fused = ( + self.fuse_qk_norm_rope + and self.head_dim in (64, 128) + and freqs is not None + and self.qk_norm + ) + + # SEPARATE_QKV self-attn 3x fp4_quantize dedup; see Attention.forward_async. + if self._maybe_share_qkv_quantize and getattr(self.to_q, "input_scale", None) is not None: + x_2d = x.reshape(-1, x.shape[-1]) + fp4, sf = torch.ops.trtllm.tunable_fp4_quantize( + x_2d, self.to_q.input_scale, self.to_q.scaling_vector_size, False + ) + qkv_input = Fp4QuantizedTensor(fp4, sf, is_sf_swizzled=False) + else: + qkv_input = x + + def compute_q(): + q = self.to_q(qkv_input) + if q.dim() == 2: + q = q.view(B, S, -1) + if use_fused: + self.apply_split_norm_rope(q, self.norm_q.weight, H, freqs[0], freqs[1]) + return q.view(B, S, H, D) + # Unfused fallback (mini-config); LTX-2 RoPE with rope_type. if self.qk_norm: q = self.norm_q(q) - else: - q, k, v = self.get_qkv(x, context) - q, k = self.apply_qk_norm(q, k) - - if pe is not None: - q = apply_rotary_emb(q, pe, self.rope_type) - # k_pe=None with pre_projected_kv signals K already rotated. - if k_pe is not None: - k = apply_rotary_emb(k, k_pe, self.rope_type) - elif pre_projected_kv is None: - k = apply_rotary_emb(k, pe, self.rope_type) + q = q.view(B, S, H, D) + if freqs is not None: + q = apply_rotary_emb(q, freqs, self.rope_type) + return q + + def compute_k(): + k = self.to_k(qkv_input) + if k.dim() == 2: + k = k.view(B, S, -1) + if use_fused: + self.apply_split_norm_rope(k, self.norm_k.weight, KV, freqs[0], freqs[1]) + return k.view(B, S, KV, D) + if self.qk_norm: + k = self.norm_k(k) + k = k.view(B, S, KV, D) + if freqs is not None: + k = apply_rotary_emb(k, freqs, self.rope_type) + return k - attn_kwargs = {} - if key_padding_mask is not None: - attn_kwargs["key_padding_mask"] = key_padding_mask - out = self._attn_impl(q, k, v, **attn_kwargs) + def compute_v(): + return self.to_v(qkv_input).view(B, S, KV, D) + + out_4d = self.attn.forward_async(compute_q, compute_k, compute_v) + # LTX-2 gated-attention scaling in 4D before to_out. if self.to_gate_logits is not None: - gate_logits = self.to_gate_logits(x) - b, t, _ = out.shape - out = out.view(b, t, self.num_attention_heads, self.head_dim) - gates = 2.0 * torch.sigmoid(gate_logits) - out = out * gates.unsqueeze(-1) - out = out.view(b, t, self.num_attention_heads * self.head_dim) + gates = 2.0 * torch.sigmoid(self.to_gate_logits(x)) + out_4d = out_4d * gates.unsqueeze(-1) - return self.to_out[0](out) + b, t = out_4d.shape[:2] + return self.to_out[0](out_4d.reshape(b, t, H * D)) # --------------------------------------------------------------------------- @@ -471,6 +569,7 @@ def _make_mlp(cfg, model_config, idx): ) def _init_video_modules(self, cfg, rope_type, eps, model_config, idx): + _async_ulysses = model_config.parallel.async_ulysses if model_config is not None else False self.attn1 = LTX2Attention( query_dim=cfg.dim, heads=cfg.heads, @@ -482,6 +581,8 @@ def _init_video_modules(self, cfg, rope_type, eps, model_config, idx): config=model_config, layer_idx=idx, enable_sequence_parallel=True, + use_ulysses=True, + async_ulysses=_async_ulysses, ) self.attn2 = LTX2Attention( query_dim=cfg.dim, diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index 0bac474df0ff..3e89abe4c4d2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -27,7 +27,6 @@ def get_parameter_device(module): return next(module.parameters()).device - # ========================================================================= # 1. Rotary Positional Embeddings # ========================================================================= @@ -290,17 +289,30 @@ def __init__( # However, this kernel does not support TP due to the cross-head # normalization being a collective op. Thus, we must disable it if # using TP. + # When ulysses_size > 1 AND parallel.async_ulysses is set, switch + # to SEPARATE_QKV so V/Q/K projections can stream-pipeline through + # the async ulysses A2A path. tp_size = model_config.mapping.tp_size if model_config.mapping else 1 + vgm_self = model_config.visual_gen_mapping + ulysses_size_self = vgm_self.ulysses_size if vgm_self is not None else 1 + _async_a2a = model_config.parallel.async_ulysses if model_config is not None else False + self._use_async_ulysses = bool(ulysses_size_self > 1) and _async_a2a + _qkv_mode_self = QKVMode.SEPARATE_QKV if self._use_async_ulysses else QKVMode.FUSE_QKV self.attn1 = Attention( hidden_size=hidden_size, num_attention_heads=num_heads, head_dim=head_dim, - qkv_mode=QKVMode.FUSE_QKV, + qkv_mode=_qkv_mode_self, qk_norm=True, eps=eps, + # fuse_qk_norm_rope=True drives the packed kernel on sync (FUSE_QKV) + # and the split kernel on async (SEPARATE_QKV via forward_async). + # Disabled when TP>1 since the fused kernel lacks cross-rank + # all-reduce for the cross-head RMSNorm variance. fuse_qk_norm_rope=(tp_size == 1), config=model_config, layer_idx=_layer_idx, + async_ulysses=self._use_async_ulysses, ) # Cross-attention with separate Q, K, V @@ -414,15 +426,15 @@ def forward( # Prepare frequencies for Attention freqs = (freqs_cos, freqs_sin) if freqs_cos is not None and freqs_sin is not None else None - # Self-attention with RoPE - x = ( - x.float() - + self.attn1( - normed, - freqs=freqs, - ).float() - * gate_msa - ).to(x.dtype) + # Self-attention with RoPE. Async-ulysses dispatches to forward_async + # so each V/Q/K GEMM + norm + RoPE overlaps with the peer push on the + # side stream; both paths return 3D [B, S, H*D]. + if self._use_async_ulysses: + attn1_out = self.attn1.forward_async(normed, freqs=freqs) + else: + attn1_out = self.attn1(normed, freqs=freqs) + + x = (x.float() + attn1_out.float() * gate_msa).to(x.dtype) norm_x = self.norm2(x.float()).to(x.dtype) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 3465fcd63acd..37a1763ebce6 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -7,6 +7,7 @@ from tensorrt_llm.llmapi.llm_args import SkipSoftmaxAttentionConfig from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig +from ...utils import Fp4QuantizedTensor from ..attention_backend.interface import AttentionTensorLayout from ..attention_backend.parallel import wrap_parallel_attention from ..attention_backend.utils import create_attention @@ -53,6 +54,8 @@ def __init__( config: Optional[DiffusionModelConfig] = None, layer_idx: Optional[int] = None, enable_sequence_parallel: bool = True, + enable_ulysses: bool = True, + async_ulysses: bool = False, ): super().__init__() @@ -115,6 +118,20 @@ def __init__( self._init_qkv_proj() + # Structural eligibility for SEPARATE_QKV self-attn quantize dedup. + # When True, get_qkv() may pre-quantize hidden_states once and pass the + # shared Fp4QuantizedTensor to to_q/to_k/to_v (relies on Linear's + # Fp4QuantizedTensor shortcut). Numerical equality of the per-tensor + # input_scales is an invariant of modelopt's self-attn calibration + # (q/k/v share the same input distribution -> same calibrated scale). + self._maybe_share_qkv_quantize = ( + self.qkv_mode == QKVMode.SEPARATE_QKV + and self.quant_config is not None + and getattr(self.quant_config, "layer_quant_mode", None) is not None + and self.quant_config.layer_quant_mode.has_nvfp4() + and not self.force_dynamic_quantization + ) + attention_metadata_state = getattr(config, "attention_metadata_state", None) if self.qk_norm: @@ -160,9 +177,20 @@ def __init__( ] ) - # Ulysses shards heads across workers; inner backend sees sharded head count. - # Attention2D gathers sequence (not heads); see wrap_parallel_attention for nesting. - use_ulysses = ulysses_size > 1 and enable_sequence_parallel + # Ulysses auto-wrap normally skips SEPARATE_QKV (cross-attention). + # The async-ulysses path uses SEPARATE_QKV for stream-pipelined + # V/Q/K projections AND still needs the head-sharding wrap — opt in + # via async_ulysses=True. + use_ulysses = ( + ulysses_size > 1 + and enable_sequence_parallel + and enable_ulysses + and (self.qkv_mode != QKVMode.SEPARATE_QKV or async_ulysses) + ) + + # Compute head counts for the backend + # Ulysses shards heads across workers; inner backend sees sharded count + # Attention2D gathers sequence (not heads); inner backend sees full count if use_ulysses: backend_num_heads = self.local_num_attention_heads // ulysses_size backend_num_kv_heads = self.local_num_key_value_heads // ulysses_size @@ -214,6 +242,7 @@ def __init__( self.attn, visual_gen_mapping=vgm, enable_sequence_parallel=enable_sequence_parallel, + async_ulysses=use_ulysses and async_ulysses, ) def _init_qkv_proj(self) -> None: @@ -520,3 +549,104 @@ def forward( out = self._attn_impl(q, k, v) out = self.to_out[0](out) return out + + def forward_async( + self, + hidden_states: torch.Tensor, + freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + """Async-Ulysses self-attn driver. Structurally mirrors ``forward``: + each closure does ``to_{q,k,v}`` + (optional) fused norm+RoPE on the + default stream while the previous tensor's peer push runs on a side + stream, so V/Q/K projections overlap with the all-to-all. + + Fused path: ``apply_split_norm_rope`` (SEPARATE_QKV analog of + ``apply_packed_qk_norm_rope``) does in-place RMSNorm + RoPE in one + kernel launch per Q/K. Same math as the packed kernel, just split + across two launches instead of one packed launch. + + Unfused path: naive ``norm_q`` + ``apply_rotary_emb``, mirroring + ``forward``'s unfused branch. + + Precondition: caller gates on ``_use_async_ulysses`` so ``self.attn`` + is a ``UlyssesAttention`` with ``async_ulysses=True``. + + Returns 3D ``[B, S, H*D]`` matching ``forward``'s output contract. + + TODO (kernel follow-up): the fused split kernel below writes Q/K to an + intermediate tensor, then ``UlyssesAttention._issue_async`` permutes + and scatters that tensor into the symm-mem slot. Folding the + permute+scatter into a "ulyssesPermuteScatter" epilogue inside the + fused norm+RoPE kernel would let it write directly into the slot, + saving one alloc + one copy + one kernel launch per Q/K closure. + """ + # Runtime precondition guard. Without async_ulysses=True at init, + # `self.attn` is the bare backend (e.g. TrtllmAttention) which has + # no `forward_async` method — the inner call below would otherwise + # crash with a non-prescriptive AttributeError deep in the function. + if not hasattr(self.attn, "forward_async"): + raise ValueError( + "Attention.forward_async() requires the inner attention to be a " + "UlyssesAttention with async_ulysses=True. Build the Attention with " + "ParallelConfig(ulysses_size > 1, async_ulysses=True), or use " + "forward() for sync execution." + ) + + B, S, _ = hidden_states.shape + H = self.num_attention_heads + KV = self.num_key_value_heads + D = self.head_dim + # Mirrors forward()'s fused gate. qkv_mode is implicitly SEPARATE_QKV + # under async (caller-enforced), so the FUSE_QKV check in forward() + # has no async analog here. + use_fused = self.fuse_qk_norm_rope and freqs is not None and self.qk_norm + + # SEPARATE_QKV self-attn 3x fp4_quantize dedup: pre-quantize hidden_states + # once and pass the shared Fp4QuantizedTensor to to_q/to_k/to_v via Linear's + # Fp4QuantizedTensor shortcut. Saves 2 of 3 fp4_quantize launches per layer. + # Eligibility is structural (set in __init__); runtime gate checks that the + # checkpoint loaded an input_scale (some attn Linears can be excluded from + # NVFP4 per checkpoint config — e.g. LTX-2 transformer_blocks.10.attn1). + if self._maybe_share_qkv_quantize and getattr(self.to_q, "input_scale", None) is not None: + x_2d = hidden_states.reshape(-1, hidden_states.shape[-1]) + fp4, sf = torch.ops.trtllm.tunable_fp4_quantize( + x_2d, self.to_q.input_scale, self.to_q.scaling_vector_size, False + ) + qkv_input = Fp4QuantizedTensor(fp4, sf, is_sf_swizzled=False) + else: + qkv_input = hidden_states + + def compute_q(): + q = self.to_q(qkv_input) + if q.dim() == 2: + q = q.view(B, S, -1) + if use_fused: + self.apply_split_norm_rope(q, self.norm_q.weight, H, freqs[0], freqs[1]) + return q.view(B, S, H, D) + if self.qk_norm: + q = self.norm_q(q) + q = q.view(B, S, H, D) + if freqs is not None: + q = apply_rotary_emb(q, freqs[0], freqs[1]) + return q + + def compute_k(): + k = self.to_k(qkv_input) + if k.dim() == 2: + k = k.view(B, S, -1) + if use_fused: + self.apply_split_norm_rope(k, self.norm_k.weight, KV, freqs[0], freqs[1]) + return k.view(B, S, KV, D) + if self.qk_norm: + k = self.norm_k(k) + k = k.view(B, S, KV, D) + if freqs is not None: + k = apply_rotary_emb(k, freqs[0], freqs[1]) + return k + + def compute_v(): + return self.to_v(qkv_input).view(B, S, KV, D) + + out_4d = self.attn.forward_async(compute_q, compute_k, compute_v) + b, t = out_4d.shape[:2] + return self.to_out[0](out_4d.reshape(b, t, H * D)) diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8daefd5d7bb6..77ba4cbbb348 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -205,6 +205,15 @@ class ParallelConfig(StrictBaseModel): status="prototype", description=("Ulysses head-sharding degree. Heads are sharded across ulysses_size GPUs."), ) + async_ulysses: bool = Field( + False, + status="prototype", + description=( + "Enable the async Ulysses A2A pipeline: overlap per-rank V/Q/K projection compute " + "with cross-rank symm-mem all-to-all on a dedicated side stream. " + "Requires ulysses_size > 1. Defaults to False." + ), + ) ring_size: int = Field( 1, ge=1, @@ -258,6 +267,23 @@ def n_workers(self) -> int: def total_parallel_size(self) -> int: return self.cfg_size * self.seq_parallel_size + @model_validator(mode="after") + def _validate_async_ulysses(self) -> "ParallelConfig": + if self.async_ulysses: + if self.ulysses_size == 1: + raise ValueError( + "async_ulysses=True requires ulysses_size > 1; got " + f"ulysses_size={self.ulysses_size}." + ) + if self.ring_size > 1: + raise ValueError( + "async_ulysses=True is incompatible with ring_size > 1: " + "async_ulysses forces SEPARATE_QKV which bypasses the " + "RingAttention wrapper. Set ring_size=1 or async_ulysses=False " + f"(got ring_size={self.ring_size})." + ) + return self + def validate_world_size(self, world_size: int) -> None: if self.total_parallel_size > world_size: raise ValueError( diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_permute_scatter.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_permute_scatter.py new file mode 100644 index 000000000000..24d995526672 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_permute_scatter.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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 + +import pytest +import torch + + +def torch_ref(input_4d, send_ref, recv_ref, my_rank, P): + """CPU/eager reference for ulyssesPermuteScatterKernel. + + input_4d : [B, S_local, H, D] bf16 + send/recv: [P, B, S_local, H/P, D] bf16 + For each (b, s, h, d): + peer = h // (H/P) + h_local = h % (H/P) + dst = recv_ref if peer == my_rank else send_ref + dst[peer, b, s, h_local, d] = input[b, s, h, d] + """ + B, S_local, H, D = input_4d.shape + P_check, _, _, H_local, _ = send_ref.shape + assert P == P_check + assert H == H_local * P + + # 4D -> [B, S_local, P, H_local, D] via H = P * H_local + x = input_4d.view(B, S_local, P, H_local, D) + # All-peer (P, B, S_local, H_local, D) + all_dst = x.permute(2, 0, 1, 3, 4).contiguous() + + # Split: peer != my_rank -> send_ref, peer == my_rank -> recv_ref + send_ref.copy_(all_dst) + recv_ref.zero_() + recv_ref[my_rank].copy_(all_dst[my_rank]) + send_ref[my_rank].zero_() + return send_ref, recv_ref + + +@pytest.mark.parametrize( + "P,B,S_local,H,D,my_rank", + [ + # LTX-2 ws=8 shape: heads-per-rank in input is H_full=32, P=4 ulysses + (4, 2, 1024, 32, 64, 0), + (4, 2, 1024, 32, 64, 1), + (4, 2, 1024, 32, 64, 2), + (4, 2, 1024, 32, 64, 3), + # WAN-like (alternate H, D) + (4, 2, 512, 16, 128, 0), + (8, 1, 256, 16, 128, 5), + # Smaller smoke shapes + (2, 1, 128, 8, 64, 0), + (2, 1, 128, 8, 64, 1), + ], +) +@torch.inference_mode() +def test_ulysses_permute_scatter_exact_match(P, B, S_local, H, D, my_rank): + """The kernel is a pure data-movement op (no float arithmetic), so the + output must match the eager reference byte-exact (max_diff == 0).""" + torch.manual_seed(0) + H_local = H // P + input_4d = torch.randn(B, S_local, H, D, device="cuda", dtype=torch.bfloat16).contiguous() + send = torch.zeros(P, B, S_local, H_local, D, device="cuda", dtype=torch.bfloat16) + recv = torch.zeros(P, B, S_local, H_local, D, device="cuda", dtype=torch.bfloat16) + + # Reference + send_ref = torch.zeros_like(send) + recv_ref = torch.zeros_like(recv) + torch_ref(input_4d, send_ref, recv_ref, my_rank, P) + + # Kernel + torch.ops.trtllm.ulysses_permute_scatter(input_4d, send, recv, my_rank, P) + + # 1. The peer==my_rank slot must be in recv (not send), and equal to input's slice + assert torch.equal(recv[my_rank], recv_ref[my_rank]) + # 2. The peer!=my_rank slots must be in send (not recv), and equal to input's slice + for peer in range(P): + if peer == my_rank: + continue + assert torch.equal(send[peer], send_ref[peer]), ( + f"send[peer={peer}] mismatch for my_rank={my_rank}" + ) + + # 3. The kernel must leave send[my_rank] and recv[peer!=my_rank] UNTOUCHED + # (we initialised them to zero; check they're still zero). + assert torch.count_nonzero(send[my_rank]) == 0 + for peer in range(P): + if peer == my_rank: + continue + assert torch.count_nonzero(recv[peer]) == 0 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 new file mode 100644 index 000000000000..91092ce4d07f --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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 + +import pytest +import torch + + +def torch_ref(q_5d, k_5d, v_5d, is_hnd): + """Eager reference: the permute+reshape+contiguous chain the kernel replaces. + For HND the returned tensor is a transpose-view of NHD storage (HND-shape, + NHD-stride, non-contig) — matches the op's new behavior that preserves + NHD-stride into SDPA so the downstream `_output_a2a` transpose+contiguous + collapses to a no-op.""" + + def post(t): + P, B, Sp, H, D = t.shape + out = t.permute(1, 0, 2, 3, 4).reshape(B, P * Sp, H, D).contiguous() # NHD storage + if is_hnd: + return out.transpose(1, 2) # [B, H, P*Sp, D] view — NHD-stride + return out # [B, P*Sp, H, D] + + return post(q_5d), post(k_5d), post(v_5d) + + +@pytest.mark.parametrize("layout", [0, 1], ids=["HND", "NHD"]) +@pytest.mark.parametrize( + "P,B,Sp,H,D", + [ + # LTX-2 self-attn (ulysses=4): per-rank H = 32/4 = 8, D = 128 + (4, 2, 1024, 8, 128), + # LTX-2 audio attn (ulysses=4): per-rank H = 32/4 = 8, D = 64 + (4, 2, 1024, 8, 64), + # WAN-like (alternate H, D) + (4, 2, 512, 16, 128), + (8, 1, 256, 16, 128), + # H * (D/8) edge: 64 threads/block (smallest interesting tile) + (2, 1, 128, 4, 128), + # Larger H with D=64 + (4, 2, 256, 32, 64), + ], +) +@torch.inference_mode() +def test_ulysses_post_unscatter_exact_match(P, B, Sp, H, D, layout): + """The op is a pure data movement, so output must match the eager + permute+reshape+contiguous chain exactly (max_diff == 0). HND output is + a transpose-view (HND-shape, NHD-stride, non-contig); NHD output is + contig. Both are exercised.""" + is_hnd = layout == 0 + torch.manual_seed(0) + q = torch.randn(P, B, Sp, H, D, device="cuda", dtype=torch.bfloat16).contiguous() + k = torch.randn(P, B, Sp, H, D, device="cuda", dtype=torch.bfloat16).contiguous() + v = torch.randn(P, B, Sp, H, D, device="cuda", dtype=torch.bfloat16).contiguous() + + q_ref, k_ref, v_ref = torch_ref(q, k, v, is_hnd=is_hnd) + q_out, k_out, v_out = torch.ops.trtllm.ulysses_post_unscatter_qkv(q, k, v, 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: + # HND-shape, NHD-stride transpose-view of NHD-contig storage. The + # underlying storage IS contig (in NHD layout), but the HND-labeled + # tensor is non-contig — this is intentional: cudnn SDPA preserves + # this NHD-stride to its output, collapsing _output_a2a's + # transpose+contiguous to a no-op. + 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).""" + q = torch.randn(2, 1, 128, 8, 64, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError): + torch.ops.trtllm.ulysses_post_unscatter_qkv(q, q, q, 2) + + +@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).""" + q = torch.randn(2, 1, 128, 8, 60, 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_post_unscatter_rejects_oversized_block(): + """Threads/block = H * (D/8) must be <= 1024 (CUDA hw limit).""" + # H=128, D=128 -> 128 * 16 = 2048 threads, exceeds the 1024 hw cap. + 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) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py new file mode 100644 index 000000000000..301c776103c8 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-GPU tests for LTX-2 async-Ulysses parity and perf. + +Compares ``LTXModel.forward`` outputs with ``async_ulysses=True`` (V/Q/K +rolling A2A via ``Attention.forward_async``) vs ``async_ulysses=False`` +(standard sync Ulysses) at ``ulysses_size=2``. Both paths wrap the inner +backend with ``UlyssesAttention`` and produce mathematically equivalent +results — drift comes only from per-kernel accumulation order under BF16. + +Mirrors ``test_ltx2_ulysses.py`` (PR 14044) but specifically exercises the +async-Ulysses code paths added on this branch: + - ``LTX2Attention.forward`` async self-attn dispatch + - ``LTX2Attention.forward_async`` (fused split kernel or naive fallback) + - ``UlyssesAttention.forward_async`` (V/Q/K rolling side-stream pipeline) + +Uses ``attention_head_dim=64`` so ``LTX2Attention.forward_async`` takes +the fused split kernel path (``apply_split_norm_rope``) — the prod code +path. Requires LTX-2 C++ extensions to be built. + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py -v +""" + +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from types import SimpleNamespace +from typing import Callable + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +try: + from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, + ) + from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping + from tensorrt_llm._utils import get_free_port + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.visual_gen.args import AttentionConfig, ParallelConfig, TorchCompileConfig + + 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) + + +# ============================================================================= +# Distributed helpers (same pattern as test_ltx2_ulysses.py) +# ============================================================================= + + +def init_distributed_worker(rank: int, world_size: int, backend: str = "nccl", port: int = 29500): + 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=backend, rank=rank, world_size=world_size) + + +def cleanup_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_worker(rank, world_size, backend, test_fn, port, fn_args): + try: + init_distributed_worker(rank, world_size, backend, port) + test_fn(rank, world_size, *fn_args) + 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, *fn_args): + 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") + port = get_free_port() + mp.spawn( + _distributed_worker, + args=(world_size, "nccl", test_fn, port, fn_args), + nprocs=world_size, + join=True, + ) + + +# ============================================================================= +# Model config +# ============================================================================= + +# Small AudioVideo config. head_dim=64 matches the fused split kernel's +# {64, 128} template — exercises ``LTX2Attention.forward_async`` fused +# branch (apply_split_norm_rope) rather than the naive fallback. +# cross_attention_dim must equal num_heads * head_dim (= 4 * 64 = 256) so +# caption_projection's output matches block.attn2.to_k's input dim. Same +# constraint for audio_*. +_AV_CONFIG = dict( + num_attention_heads=4, + attention_head_dim=64, + in_channels=16, + out_channels=16, + num_layers=2, + cross_attention_dim=256, # = num_attention_heads * attention_head_dim + caption_channels=64, + norm_eps=1e-6, + positional_embedding_max_pos=[4, 32, 32], + timestep_scale_multiplier=1000, + use_middle_indices_grid=True, + audio_num_attention_heads=4, + audio_attention_head_dim=64, + audio_in_channels=16, + audio_out_channels=16, + audio_cross_attention_dim=256, # = audio_num_attention_heads * audio_attention_head_dim + audio_positional_embedding_max_pos=[64], + av_ca_timestep_scale_multiplier=1, +) + + +def _make_model_config( + ulysses_size: int = 1, + backend: str = "VANILLA", + async_ulysses: bool = False, +) -> "DiffusionModelConfig": + """Create DiffusionModelConfig for LTX-2 tests. + + The async_ulysses flag toggles the V/Q/K rolling A2A pipeline (this PR's + feature). ulysses_size > 1 is required for async_ulysses to fire. + """ + if ulysses_size > 1 and dist.is_initialized(): + ws = dist.get_world_size() + rk = dist.get_rank() + else: + ws = ulysses_size + rk = 0 + vgm = VisualGenMapping(world_size=ws, rank=rk, ulysses_size=ulysses_size) + + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + quant_config=QuantConfig(), + torch_compile=TorchCompileConfig(enable=False), + attention=AttentionConfig(backend=backend), + visual_gen_mapping=vgm, + cache=None, + attention_metadata_state=( + create_attention_metadata_state() if backend.upper() == "TRTLLM" else None + ), + parallel=ParallelConfig( + ulysses_size=ulysses_size, + async_ulysses=async_ulysses, + ), + skip_create_weights_in_init=False, + ) + config.mapping = vgm.to_llm_mapping() + return config + + +def _init_all_weights(model: torch.nn.Module, std: float = 0.02): + with torch.no_grad(): + for name, p in model.named_parameters(): + if "norm" in name and "weight" in name: + p.fill_(1.0) + elif p.numel() > 0: + torch.nn.init.normal_(p, mean=0.0, std=std) + + +def _make_video_positions( + batch: int, n_patches: int, n_frames: int, grid_h: int, grid_w: int, device: torch.device +) -> torch.Tensor: + positions = torch.zeros(batch, 3, n_patches, 2, device=device) + idx = 0 + for f in range(n_frames): + for h in range(grid_h): + for w in range(grid_w): + positions[:, 0, idx, :] = torch.tensor([f, f + 1], dtype=torch.float32) + positions[:, 1, idx, :] = torch.tensor([h, h + 1], dtype=torch.float32) + positions[:, 2, idx, :] = torch.tensor([w, w + 1], dtype=torch.float32) + idx += 1 + return positions + + +def _make_audio_positions(batch: int, n_patches: int, device: torch.device) -> torch.Tensor: + positions = torch.zeros(batch, 1, n_patches, 2, device=device) + for i in range(n_patches): + positions[:, 0, i, :] = torch.tensor([i, i + 1], dtype=torch.float32) + return positions + + +def _build_inputs(batch, v_patches, v_dims, a_patches, dtype, device, seed=456): + g = torch.Generator(device=device).manual_seed(seed) + v_frames, v_h, v_w = v_dims + in_channels = _AV_CONFIG["in_channels"] + audio_in_channels = _AV_CONFIG["audio_in_channels"] + caption_channels = _AV_CONFIG["caption_channels"] + text_len = 8 + + v_context = ( + torch.randn(batch, text_len, caption_channels, device=device, dtype=dtype, generator=g) + * 0.02 + ) + a_context = ( + torch.randn(batch, text_len, caption_channels, device=device, dtype=dtype, generator=g) + * 0.02 + ) + v_positions = _make_video_positions(batch, v_patches, v_frames, v_h, v_w, device) + a_positions = _make_audio_positions(batch, a_patches, device) + + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality + + video = Modality( + latent=torch.randn(batch, v_patches, in_channels, device=device, dtype=dtype, generator=g) + * 0.02, + timesteps=torch.tensor([0.5], device=device), + positions=v_positions, + context=v_context, + ) + audio = Modality( + latent=torch.randn( + batch, a_patches, audio_in_channels, device=device, dtype=dtype, generator=g + ) + * 0.02, + timesteps=torch.tensor([0.5], device=device), + positions=a_positions, + context=a_context, + ) + return video, audio, v_context, a_context, v_positions, a_positions + + +def _pack_async_state_for_sync(async_state, sync_state_keys): + """Translate a SEPARATE_QKV (async) state_dict to FUSE_QKV (sync) layout: + concatenates to_q/to_k/to_v.{weight,bias} into packed qkv_proj.{weight,bias}. + Non-self-attn weights pass through unchanged. + + Used to load identical effective weights into both async and sync LTX + models despite their different attn1 layer structure. + """ + out = {} + for k in sync_state_keys: + if k in async_state: + out[k] = async_state[k] + elif k.endswith(".qkv_proj.weight"): + prefix = k[: -len(".qkv_proj.weight")] + out[k] = torch.cat( + [ + async_state[f"{prefix}.to_q.weight"], + async_state[f"{prefix}.to_k.weight"], + async_state[f"{prefix}.to_v.weight"], + ], + dim=0, + ) + elif k.endswith(".qkv_proj.bias"): + prefix = k[: -len(".qkv_proj.bias")] + out[k] = torch.cat( + [ + async_state[f"{prefix}.to_q.bias"], + async_state[f"{prefix}.to_k.bias"], + async_state[f"{prefix}.to_v.bias"], + ], + dim=0, + ) + else: + raise KeyError(f"sync key {k!r} not present in async state and not a qkv_proj fuse") + return out + + +def _build_av_model( + world_size: int, + backend: str, + audio_seq_len: int, + async_ulysses: bool, + device: torch.device, + dtype: torch.dtype, + seed: int = 123, +): + """Build LTXModel (AudioVideo) with deterministic weights via shared seed. + + ``configure_audio_ulysses(audio_seq_len)`` gates audio_attn1's Ulysses + activity by divisibility: not divisible → ``set_ulysses_active(False)`` + swaps the audio backend to plain (no ``forward_async``), forcing async + self-attn to fall through the ``hasattr`` guard in ``LTX2Attention.forward``. + """ + from tensorrt_llm._torch.visual_gen.models.ltx2.transformer_ltx2 import LTXModel, LTXModelType + + torch.manual_seed(seed) + cfg = _make_model_config( + ulysses_size=world_size, + backend=backend, + async_ulysses=async_ulysses, + ) + model = ( + LTXModel(model_type=LTXModelType.AudioVideo, model_config=cfg, **_AV_CONFIG) + .to(device, dtype=dtype) + .eval() + ) + _init_all_weights(model) + model.configure_audio_ulysses(audio_seq_len) + return model + + +# ============================================================================= +# Test logic +# ============================================================================= + + +def _logic_async_vs_sync_parity(rank, world_size, backend, audio_seq_len): + """LTX-2 AV at ws=2: async-Ulysses output matches sync-Ulysses output. + + Both models use ``UlyssesAttention`` (forced by ``ulysses_size=2``); the + only difference is whether self-attn dispatches through + ``Attention.forward_async`` (V/Q/K closures + side-stream A2A) or + ``Attention.forward`` (precomputed Q/K/V + sync A2A). Math is equivalent; + BF16 accumulation order drift lands well under the 5e-2 tolerance. + """ + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + + batch = 1 + v_dims = (1, 4, 4) + v_patches = v_dims[0] * v_dims[1] * v_dims[2] # 16, divisible by ws=2 + + # Build async (SEPARATE_QKV self-attn → to_q/k/v) first as the canonical + # source, then mirror its weights into a sync (FUSE_QKV → qkv_proj) model + # by concatenating to_q/k/v into the packed layout. + async_model = _build_av_model(world_size, backend, audio_seq_len, True, device, dtype) + async_state = async_model.state_dict() + sync_model = _build_av_model(world_size, backend, audio_seq_len, False, device, dtype) + sync_model.load_state_dict( + _pack_async_state_for_sync(async_state, sync_model.state_dict().keys()) + ) + + # Sanity: confirm sync vs async actually take different code paths. + # transformer_blocks may be wrapped by LTX2CacheDiTPattern0BlockWrapper; + # unwrap via .inner if present. + def _block_attn1(model): + b = model.transformer_blocks[0] + b = b.inner if hasattr(b, "inner") else b + return b.attn1 + + assert _block_attn1(async_model)._use_async_ulysses is True, ( + "async model didn't enable async path" + ) + assert _block_attn1(sync_model)._use_async_ulysses is False, ( + "sync model unexpectedly enabled async path" + ) + assert _block_attn1(async_model).qkv_mode != _block_attn1(sync_model).qkv_mode, ( + "test bug: sync and async models have identical qkv_mode " + f"(both {_block_attn1(sync_model).qkv_mode})" + ) + + video, audio, v_ctx, a_ctx, v_pos, a_pos = _build_inputs( + batch, v_patches, v_dims, audio_seq_len, dtype, device + ) + + sync_cache = sync_model.prepare_text_cache( + video_context=v_ctx, + video_positions=v_pos, + audio_context=a_ctx, + audio_positions=a_pos, + dtype=dtype, + ) + async_cache = async_model.prepare_text_cache( + video_context=v_ctx, + video_positions=v_pos, + audio_context=a_ctx, + audio_positions=a_pos, + dtype=dtype, + ) + + with torch.no_grad(): + sync_v, sync_a = sync_model(video=video, audio=audio, text_cache=sync_cache) + async_v, async_a = async_model(video=video, audio=audio, text_cache=async_cache) + + assert sync_v.shape == async_v.shape, f"Rank {rank}: video shape mismatch" + assert sync_a.shape == async_a.shape, f"Rank {rank}: audio shape mismatch" + + # Diagnostic: actual BF16 drift between sync (packed kernel + sync a2a) and + # async (split kernel + side-stream a2a) — both at ws=2, same collectives. + for name, s, a in [("video", sync_v, async_v), ("audio", sync_a, async_a)]: + diff = (a.float() - s.float()).abs() + ref = s.float().abs() + print( + f"\n[LTX-2 rank={rank} backend={backend} {name}] " + f"max_abs_diff={diff.max().item():.3e} " + f"max_rel_diff={(diff / ref.clamp(min=1e-6)).max().item():.3e} " + f"sync_abs_max={ref.max().item():.3e}" + ) + + torch.testing.assert_close( + async_v, + sync_v, + rtol=1e-3, + atol=1e-3, + msg=f"Rank {rank}: LTX-2 AV async-Ulysses video differs from sync-Ulysses", + ) + torch.testing.assert_close( + async_a, + sync_a, + rtol=1e-3, + atol=1e-3, + msg=f"Rank {rank}: LTX-2 AV async-Ulysses audio differs from sync-Ulysses", + ) + + +# ============================================================================= +# Test classes +# ============================================================================= + + +class TestLTX2AsyncUlysses: + """async_ulysses=True/False parity for LTX-2 AudioVideo at ws=2.""" + + @pytest.mark.parametrize("backend", ["VANILLA", "FA4"]) + def test_av_async_vs_sync_parity(self, backend): + """ws=2, audio_seq_len % 2 == 0 — audio_attn1 uses Ulysses on both + sync and async paths.""" + run_test_in_distributed(2, _logic_async_vs_sync_parity, backend, 16) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py new file mode 100644 index 000000000000..62353b4b4741 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py @@ -0,0 +1,292 @@ +"""Multi-rank tests for the async Ulysses A2A op. + +Exercises the production ``ulysses_a2a_async_prepare`` / ``ulysses_a2a_async_push`` +/ ``ulysses_a2a_async_barrier`` chain end-to-end across multiple ranks. Two test +surfaces: + +1. ``test_slot_ring_wraparound`` — eager mode (cudaMemcpyBatchAsync path). + Loops the pair more than ``kNumSlots`` times so the slot ring wraps multiple + times, asserts byte-exact match vs ``all_to_all_4d`` on every iteration. + Catches off-by-one slot-reuse bugs. + +2. ``test_capture_smoke`` — under-capture mode (per-peer cudaMemcpyAsync path). + Warms up the slot out-of-capture, then captures a ``torch.cuda.CUDAGraph`` + containing one prepare+async pair, replays it K times with fresh inputs. + Smoke-tests the cuda_graph path used by the e2e production benchmark. + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py -v +""" + +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: + from tensorrt_llm._torch.distributed import all_to_all_4d + from tensorrt_llm._utils import get_free_port + + MODULES_AVAILABLE = True +except ImportError: + MODULES_AVAILABLE = False + + +# Loop count must comfortably exceed kNumSlots so the ring wraps at least +# twice. kNumSlots is 3 today; 8 iterations = ~2.67 full rotations. +NUM_ITERS = 8 + + +def _init_dist(rank: int, world_size: int, port: int): + 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 _slot_ring_logic(rank: int, world_size: int): + """Loop the async A2A prepare/async pair NUM_ITERS times; assert byte-exact + match against all_to_all_4d on every iteration. + + Stream pattern mirrors production (parallel.py:_issue_async): Phase 1 on + the default (compute) stream, Phase 2 on a dedicated side stream gated by + an event recorded after Phase 1. + """ + device = torch.device(f"cuda:{rank}") + torch.manual_seed(1234 + rank) + + # Production-scale shapes (LTX-2 ws=4: H=32, D=128, S_local=720 region). + B, S_local, H, D = 2, 128, world_size * 8, 128 + pg = dist.group.WORLD + pg_boxed = pg.boxed() + side_stream = torch.cuda.Stream(device=device) + + for it in range(NUM_ITERS): + # Fresh input each iteration (different bytes) so a stale slot would + # produce a stale value and the byte-exact compare would catch it. + torch.manual_seed(1000 * it + rank) + x = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device=device) + + # Phase 1 on default stream. + recv_5d, send_h = torch.ops.trtllm.ulysses_a2a_async_prepare(x, pg_boxed) + ev = torch.cuda.Event() + ev.record() + # Phase 2 on side stream. + with torch.cuda.stream(side_stream): + ev.wait() + torch.ops.trtllm.ulysses_a2a_async_push(send_h, pg_boxed) + torch.ops.trtllm.ulysses_a2a_async_barrier(pg_boxed) + ev_done = torch.cuda.Event() + with torch.cuda.stream(side_stream): + ev_done.record() + torch.cuda.current_stream().wait_event(ev_done) + torch.cuda.synchronize() + + # Async op produces [P, B, S_local, H/P, D]; permute to [B, P*S_local, H/P, D] + # to match the all_to_all_4d output layout. + async_out = ( + recv_5d.permute(1, 0, 2, 3, 4) + .contiguous() + .view(B, world_size * S_local, H // world_size, D) + ) + + ref = all_to_all_4d(x, scatter_dim=2, gather_dim=1, process_group=pg) + + assert torch.equal(async_out, ref), ( + f"rank {rank} iter {it}: async A2A ≠ all_to_all_4d reference " + f"(slot ring wrap = {it % 3})" + ) + + +def _capture_smoke_logic(rank: int, world_size: int): + """Smoke-test the under-capture branch: warm up slots out-of-capture, then + capture a torch.cuda.CUDAGraph containing one prepare+async pair and replay + K times with fresh inputs. Exercises the per-peer cudaMemcpyAsync loop + (production cuda_graph path) instead of the eager cudaMemcpyBatchAsync. + """ + device = torch.device(f"cuda:{rank}") + torch.manual_seed(1234 + rank) + + B, S_local, H, D = 2, 128, world_size * 8, 128 + P = world_size + pg = dist.group.WORLD + pg_boxed = pg.boxed() + side_stream = torch.cuda.Stream(device=device) + + # Static buffers (graph captures pointers, not values). + x_static = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device=device) + out_static = torch.empty(B, P * S_local, H // P, D, dtype=torch.bfloat16, device=device) + + # Warmup: must allocate ALL kNumSlots ring entries out-of-capture (allocation + # is not capture-safe; the cudaStreamIsCapturing guard in getOrAllocSlot + # enforces this). One warmup call only allocates 1 slot — the captured + # _prepare advances mNextIdx and would hit an unallocated slot. So warm up + # >= kNumSlots times. Mirror production's stream pattern (Phase 2 on side + # stream) so cudaMemcpyBatchAsync sees the steady-state stream context. + K_NUM_SLOTS = 3 # mirrors AsyncUlyssesOp::kNumSlots + for _ in range(K_NUM_SLOTS): + recv_w, sh_w = torch.ops.trtllm.ulysses_a2a_async_prepare(x_static, pg_boxed) + ev_w = torch.cuda.Event() + ev_w.record() + with torch.cuda.stream(side_stream): + ev_w.wait() + torch.ops.trtllm.ulysses_a2a_async_push(sh_w, pg_boxed) + torch.ops.trtllm.ulysses_a2a_async_barrier(pg_boxed) + ev_w_done = torch.cuda.Event() + with torch.cuda.stream(side_stream): + ev_w_done.record() + torch.cuda.current_stream().wait_event(ev_w_done) + torch.cuda.synchronize() + del recv_w, sh_w + + # Capture. + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + recv_5d, send_h = torch.ops.trtllm.ulysses_a2a_async_prepare(x_static, pg_boxed) + ev = torch.cuda.Event() + ev.record() + with torch.cuda.stream(side_stream): + ev.wait() + torch.ops.trtllm.ulysses_a2a_async_push(send_h, pg_boxed) + torch.ops.trtllm.ulysses_a2a_async_barrier(pg_boxed) + ev_done = torch.cuda.Event() + with torch.cuda.stream(side_stream): + ev_done.record() + torch.cuda.current_stream().wait_event(ev_done) + captured_out = recv_5d.permute(1, 0, 2, 3, 4).contiguous().view(B, P * S_local, H // P, D) + out_static.copy_(captured_out) + + # Replay K times. K > kNumSlots*2 to exercise wrap-around under capture. + K = 8 + for it in range(K): + torch.manual_seed(2000 * it + rank) + new_x = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device=device) + x_static.copy_(new_x) + g.replay() + torch.cuda.synchronize() + ref = all_to_all_4d(new_x, scatter_dim=2, gather_dim=1, process_group=pg) + assert torch.equal(out_static, ref), ( + f"rank {rank} capture replay {it}: async A2A ≠ all_to_all_4d reference" + ) + + +def _multi_pg_logic(rank: int, world_size: int): + """Two distinct PGs spanning the same ranks but with different group_names; + alternate calls between them; verify each gets its own slot ring and + byte-exact output vs all_to_all_4d on its own group. + """ + device = torch.device(f"cuda:{rank}") + torch.manual_seed(4242 + rank) + + B, S_local, H, D = 2, 64, world_size * 8, 128 + # Two PGs, same membership (all ranks), but new_group assigns distinct + # group_names — so getOrCreateOp must cache one op per PG-name and + # set_group_info must register both. + pg_a = dist.new_group(ranks=list(range(world_size))) + pg_b = dist.new_group(ranks=list(range(world_size))) + assert pg_a.group_name != pg_b.group_name, "new_group should yield distinct names" + pg_a_boxed = pg_a.boxed() + pg_b_boxed = pg_b.boxed() + side_stream = torch.cuda.Stream(device=device) + + def _issue(pg_boxed, pg_obj, x): + recv_5d, send_h = torch.ops.trtllm.ulysses_a2a_async_prepare(x, pg_boxed) + ev = torch.cuda.Event() + ev.record() + with torch.cuda.stream(side_stream): + ev.wait() + torch.ops.trtllm.ulysses_a2a_async_push(send_h, pg_boxed) + torch.ops.trtllm.ulysses_a2a_async_barrier(pg_boxed) + ev_done = torch.cuda.Event() + with torch.cuda.stream(side_stream): + ev_done.record() + torch.cuda.current_stream().wait_event(ev_done) + torch.cuda.synchronize() + async_out = ( + recv_5d.permute(1, 0, 2, 3, 4) + .contiguous() + .view(B, world_size * S_local, H // world_size, D) + ) + ref = all_to_all_4d(x, scatter_dim=2, gather_dim=1, process_group=pg_obj) + return async_out, ref + + # Alternate between the two PGs for 2*kNumSlots iterations to exercise + # both slot rings wrapping. + for it in range(NUM_ITERS): + torch.manual_seed(3000 * it + rank) + x = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device=device) + pg_boxed, pg_obj, tag = (pg_a_boxed, pg_a, "A") if it % 2 == 0 else (pg_b_boxed, pg_b, "B") + async_out, ref = _issue(pg_boxed, pg_obj, x) + assert torch.equal(async_out, ref), ( + f"rank {rank} iter {it} PG={tag}: async A2A ≠ all_to_all_4d reference" + ) + + +def _worker(rank, world_size, port): + try: + _init_dist(rank, world_size, port) + _slot_ring_logic(rank, world_size) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _worker_capture(rank, world_size, port): + try: + _init_dist(rank, world_size, port) + _capture_smoke_logic(rank, world_size) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _worker_multi_pg(rank, world_size, port): + try: + _init_dist(rank, world_size, port) + _multi_pg_logic(rank, world_size) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run(world_size: int, test_fn: Callable): + 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") + port = get_free_port() + mp.spawn(test_fn, args=(world_size, port), nprocs=world_size, join=True) + + +def test_slot_ring_wraparound(): + """Loop _prepare/_async ≥ 2*kNumSlots iterations on ws=2; assert byte-exact + match against all_to_all_4d on each iteration (eager mode). + """ + _run(2, _worker) + + +def test_capture_smoke(): + """Capture a CUDAGraph containing one _prepare/_async pair on ws=2; replay + K iterations with fresh inputs; assert byte-exact match against + all_to_all_4d (per-peer cudaMemcpyAsync path under capture). + """ + _run(2, _worker_capture) + + +def test_multi_pg(): + """Two distinct ProcessGroups on ws=2; alternate _prepare/_async calls + between them; assert byte-exact match against each group's all_to_all_4d. + + Exercises PG-name caching in ``getOrCreateOp`` and ``set_group_info`` + re-registration across multiple groups — each PG has its own group_name, + so each must yield its own cached ``AsyncUlyssesOp`` instance + slot ring. + """ + _run(2, _worker_multi_pg) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py new file mode 100644 index 000000000000..f1f38a8dcc16 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-GPU tests for WAN async-Ulysses parity and perf. + +Compares ``WanTransformer3DModel.forward`` outputs with +``async_ulysses=True`` (V/Q/K rolling A2A via ``Attention.forward_async``, +fused split RMSNorm+RoPE kernel) vs ``async_ulysses=False`` (standard sync +Ulysses with packed FUSE_QKV fused kernel) at ``ulysses_size=2``. Math is +equivalent (unify PR ties split vs packed kernel families); BF16 drift +through the 2-layer model is small. + +Exercises: + - ``WanTransformerBlock.forward`` block-level dispatch to ``forward_async`` + - ``Attention.forward_async`` fused split kernel path (head_dim=64) + - ``UlyssesAttention.forward_async`` V/Q/K rolling side-stream pipeline + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py -v +""" + +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from types import SimpleNamespace +from typing import Callable + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +try: + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping + from tensorrt_llm._utils import get_free_port + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + ParallelConfig, + TeaCacheConfig, + TorchCompileConfig, + ) + + 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) + + +# ============================================================================= +# Distributed helpers +# ============================================================================= + + +def init_distributed_worker(rank: int, world_size: int, backend: str = "nccl", port: int = 29500): + 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=backend, rank=rank, world_size=world_size) + + +def cleanup_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_worker(rank, world_size, backend, test_fn, port, fn_args): + try: + init_distributed_worker(rank, world_size, backend, port) + test_fn(rank, world_size, *fn_args) + 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, *fn_args): + 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") + port = get_free_port() + mp.spawn( + _distributed_worker, + args=(world_size, "nccl", test_fn, port, fn_args), + nprocs=world_size, + join=True, + ) + + +# ============================================================================= +# Model config +# ============================================================================= + +# Small WAN config. head_dim=64 matches the fused split kernel template; with +# num_attention_heads=4 and ulysses_size=2, each rank holds 2 heads after the +# Ulysses head-shard. Video patches yield 8 tokens, divisible by ws=2. +_WAN_CONFIG = dict( + num_attention_heads=4, + attention_head_dim=64, + num_layers=2, + in_channels=4, + out_channels=4, + patch_size=[1, 2, 2], + text_dim=64, + freq_dim=32, +) +_VIDEO_SHAPE = (1, 4, 2, 4, 4) # [B, C, T, H, W] +_TEXT_SEQ = 4 + + +def _make_model_config( + ulysses_size: int = 1, + backend: str = "FA4", + async_ulysses: bool = False, +) -> "DiffusionModelConfig": + """Create DiffusionModelConfig for WAN tests.""" + if ulysses_size > 1 and dist.is_initialized(): + ws = dist.get_world_size() + rk = dist.get_rank() + else: + ws = ulysses_size + rk = 0 + vgm = VisualGenMapping(world_size=ws, rank=rk, ulysses_size=ulysses_size) + + pretrained_config = SimpleNamespace(**_WAN_CONFIG) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + torch_compile=TorchCompileConfig(enable=False), + attention=AttentionConfig(backend=backend), + visual_gen_mapping=vgm, + cache=TeaCacheConfig(), + parallel=ParallelConfig( + ulysses_size=ulysses_size, + async_ulysses=async_ulysses, + ), + skip_create_weights_in_init=False, + ) + config.mapping = vgm.to_llm_mapping() + return config + + +def _stabilize_model_weights(model): + """Reinit weights to small values — prevents BF16 overflow through layers.""" + with torch.no_grad(): + for p in model.parameters(): + if p.ndim >= 2: + fan_in = p.shape[1] + std = 0.02 / max(1.0, fan_in**0.5) + p.data.uniform_(-std, std) + else: + p.data.uniform_(-0.01, 0.01) + + +def _build_wan_model( + world_size: int, + backend: str, + async_ulysses: bool, + device: torch.device, + dtype: torch.dtype, + seed: int = 42, +): + """Build WanTransformer3DModel with deterministic weights via shared seed.""" + from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel + + torch.manual_seed(seed) + cfg = _make_model_config(ulysses_size=world_size, backend=backend, async_ulysses=async_ulysses) + model = WanTransformer3DModel(cfg).to(device).to(dtype) + _stabilize_model_weights(model) + return model + + +def _pack_async_state_for_sync(async_state, sync_state_keys): + """Translate a SEPARATE_QKV (async) state_dict to FUSE_QKV (sync) layout: + concatenates to_q/to_k/to_v.{weight,bias} into packed qkv_proj.{weight,bias}. + Non-self-attn weights pass through unchanged. + + Used to load identical effective weights into both async and sync models + despite their different attn1 layer structure. + """ + out = {} + for k in sync_state_keys: + if k in async_state: + out[k] = async_state[k] + elif k.endswith(".qkv_proj.weight"): + prefix = k[: -len(".qkv_proj.weight")] + out[k] = torch.cat( + [ + async_state[f"{prefix}.to_q.weight"], + async_state[f"{prefix}.to_k.weight"], + async_state[f"{prefix}.to_v.weight"], + ], + dim=0, + ) + elif k.endswith(".qkv_proj.bias"): + prefix = k[: -len(".qkv_proj.bias")] + out[k] = torch.cat( + [ + async_state[f"{prefix}.to_q.bias"], + async_state[f"{prefix}.to_k.bias"], + async_state[f"{prefix}.to_v.bias"], + ], + dim=0, + ) + else: + raise KeyError(f"sync key {k!r} not present in async state and not a qkv_proj fuse") + return out + + +def _build_inputs(device, dtype, seed=100): + """Identical inputs across all ranks via shared seed.""" + B, C, T, H, W = _VIDEO_SHAPE + text_dim = _WAN_CONFIG["text_dim"] + torch.manual_seed(seed) + hidden_states = torch.randn(_VIDEO_SHAPE, device=device, dtype=dtype) * 0.1 + encoder_hidden_states = torch.randn(B, _TEXT_SEQ, text_dim, device=device, dtype=dtype) * 0.1 + timestep = torch.tensor([0.5], device=device, dtype=dtype) + return hidden_states, encoder_hidden_states, timestep + + +# ============================================================================= +# Test logic +# ============================================================================= + + +def _logic_async_vs_sync_parity(rank, world_size, backend): + """WAN at ws=2: async-Ulysses output matches sync-Ulysses output. + + Sync path: block calls ``self.attn1(...)`` → ``Attention.forward`` → + FUSE_QKV packed kernel → sync ``UlyssesAttention.forward``. + Async path: block calls ``self.attn1.forward_async(...)`` → split + kernel via closures → ``UlyssesAttention.forward_async`` rolling A2A. + Both should match within BF16 accumulation drift. + """ + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + + # Build async (SEPARATE_QKV self-attn → to_q/k/v) first as the canonical + # source, then mirror its weights into a sync (FUSE_QKV → qkv_proj) model + # by concatenating to_q/k/v into the packed layout. + async_model = _build_wan_model(world_size, backend, True, device, dtype) + async_state = async_model.state_dict() + sync_model = _build_wan_model(world_size, backend, False, device, dtype) + sync_model.load_state_dict( + _pack_async_state_for_sync(async_state, sync_model.state_dict().keys()) + ) + + # Sanity: confirm sync vs async actually take different code paths. + assert async_model.blocks[0]._use_async_ulysses is True, "async model didn't enable async path" + assert sync_model.blocks[0]._use_async_ulysses is False, ( + "sync model unexpectedly enabled async path" + ) + assert async_model.blocks[0].attn1.qkv_mode != sync_model.blocks[0].attn1.qkv_mode, ( + "test bug: sync and async models have identical qkv_mode " + f"(both {sync_model.blocks[0].attn1.qkv_mode})" + ) + + hidden_states, encoder_hidden_states, timestep = _build_inputs(device, dtype) + + with torch.no_grad(): + sync_out = sync_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + async_out = async_model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + ) + + assert sync_out.shape == async_out.shape, f"Rank {rank}: output shape mismatch" + assert not torch.isnan(async_out).any(), f"Rank {rank}: NaN in async output" + assert not torch.isinf(async_out).any(), f"Rank {rank}: Inf in async output" + + # Diagnostic: actual BF16 drift between sync (packed kernel + sync a2a) and + # async (split kernel + side-stream a2a) — both at ws=2, same collectives. + diff = (async_out.float() - sync_out.float()).abs() + ref = sync_out.float().abs() + max_abs = diff.max().item() + max_rel = (diff / ref.clamp(min=1e-6)).max().item() + print( + f"\n[WAN rank={rank} backend={backend}] " + f"max_abs_diff={max_abs:.3e} max_rel_diff={max_rel:.3e} " + f"sync_abs_max={ref.max().item():.3e}" + ) + + torch.testing.assert_close( + async_out, + sync_out, + rtol=1e-3, + atol=1e-3, + msg=f"Rank {rank}: WAN async-Ulysses output differs from sync-Ulysses", + ) + + +# ============================================================================= +# Test classes +# ============================================================================= + + +class TestWanAsyncUlysses: + """async_ulysses=True/False parity for WAN at ws=2.""" + + @pytest.mark.parametrize("backend", ["VANILLA", "FA4"]) + def test_async_vs_sync_parity(self, backend): + """ws=2: async path (block-level inline-replaced by forward_async) + matches sync path (Attention.forward FUSE_QKV packed kernel).""" + run_test_in_distributed(2, _logic_async_vs_sync_parity, backend) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"])