[TRTLLM-11457][feat] Async Ulysses pipeline (Enabled for LTX-2 + WAN)

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 <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+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<int4 const*>(input);
+    int4* __restrict__ dst_v = reinterpret_cast<int4*>(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<<<grid, block, 0, stream>>>(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

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 <cuda_runtime.h>
+
+#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

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 <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+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 <typename T>
+__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<int64_t>(p) * B + b) * Sp + sp) * H + h) * D) + vec_idx * VEC;
+    int64_t const out_base = (((static_cast<int64_t>(b) * PSp + psp) * H + h) * D) + vec_idx * VEC;
+
+    uint4 const* in_v4 = reinterpret_cast<uint4 const*>(in_ptr + in_base);
+    uint4* out_v4 = reinterpret_cast<uint4*>(out_ptr + out_base);
+    *out_v4 = *in_v4;
+}
+
+} // namespace
+
+void launchUlyssesPostUnscatter(void const* q_in, void const* k_in, void const* v_in, void* q_out, void* k_out,
+    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><<<grid, block, 0, stream>>>(
+        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

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 <cuda_runtime.h>
+
+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

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