diff --git a/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.cu index d5690c81438a..ebf88704e844 100644 --- a/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.cu @@ -20,7 +20,9 @@ #include "tensorrt_llm/common/mathUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" #include +#include #include +#include TRTLLM_NAMESPACE_BEGIN @@ -430,6 +432,283 @@ __global__ void fusedDiTCrossHeadQKNormRopeKernel(__nv_bfloat16* qkv, // [num_to } } +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Fused full-dim RMSNorm + RoPE on packed QKV; Q and K share a single cos/sin pair. +// Strategy: +// - 2 rows per CTA (256 threads = 2 rows x 128 threads x 4 warps). +// - Phase 0a: cp.async Q + K + cos + sin (HBM -> SMEM) in one commit group. +// - Phase 0b: sync load q_weight + k_weight -> regs (overlaps the cp.async transfers). +// - Phase 1: sum^2_Q and sum^2_K together from SMEM, per-row reduce with packed (Q, K) warp slots. +// - Phase 2: applies norm + RoPE to Q and K via shared cos/sin SMEM stage; writes HBM in place. +template +__global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const num_heads_q, int const num_heads_k, + int const num_heads_v, float const eps, __nv_bfloat16 const* q_weight, __nv_bfloat16 const* k_weight, + CosT const* cos_emb, CosT const* sin_emb, int const num_tokens, int const cos_seq_per_batch) +{ + constexpr int BLOCK_SIZE = 256; + constexpr int ROWS_PER_BLOCK = 2; + constexpr int THREADS_PER_ROW = BLOCK_SIZE / ROWS_PER_BLOCK; // 128 + constexpr int WARPS_PER_ROW = THREADS_PER_ROW / 32; // 4 + constexpr int CHUNK_ELEMS = 8; // uint4 = 8 bf16 + constexpr int MAX_N = 32 * HEAD_DIM; + constexpr int MAX_CHUNKS = (MAX_N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;"); +#endif + + int const tid = threadIdx.x; + int const row_in_block = tid / THREADS_PER_ROW; + int const lane_in_row = tid % THREADS_PER_ROW; + int const row_warp = lane_in_row >> 5; + int const row_lane = lane_in_row & 31; + + int const tokenIdx = blockIdx.x * ROWS_PER_BLOCK + row_in_block; + if (tokenIdx >= num_tokens) + return; + + int const N = num_heads_q * HEAD_DIM; // num_heads_q == num_heads_k (enforced by launcher) + int const chunks_per_row = (N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + int const num_heads_total = num_heads_q + num_heads_k + num_heads_v; + int64_t const tokenBaseQ = static_cast(tokenIdx) * num_heads_total * HEAD_DIM; + int64_t const tokenBaseK = tokenBaseQ + N; + int const cos_tokenIdx = (cos_seq_per_batch > 0) ? (tokenIdx % cos_seq_per_batch) : tokenIdx; + int64_t const embBase = PER_HEAD_COS ? static_cast(cos_tokenIdx) * num_heads_q * HEAD_DIM + : static_cast(cos_tokenIdx) * HEAD_DIM; + + // SMEM layout: [Q row0][Q row1][K row0][K row1] bf16, [cos row0..1][sin row0..1] CosT, warp_sums. + extern __shared__ __align__(16) unsigned char smem_raw[]; + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw); + __nv_bfloat16* smem_k = smem_q + ROWS_PER_BLOCK * N; + CosT* smem_cos = reinterpret_cast(smem_raw + 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16)); + CosT* smem_sin = smem_cos + ROWS_PER_BLOCK * N; + float* warp_sums = reinterpret_cast( + smem_raw + 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * N * sizeof(CosT)); + + // Phase 0a: cp.async Q + K + cos + sin -> SMEM (all in one commit group). +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + __pipeline_memcpy_async(smem_q + row_in_block * N + elemBase, qkv + tokenBaseQ + elemBase, 16); + __pipeline_memcpy_async(smem_k + row_in_block * N + elemBase, qkv + tokenBaseK + elemBase, 16); + int const headIdx = elemBase / HEAD_DIM; + int const baseDim = elemBase - headIdx * HEAD_DIM; + int const cosHeadOff = PER_HEAD_COS ? headIdx * HEAD_DIM : 0; + if constexpr (std::is_same_v) + { + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase, cos_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase + 4, cos_emb + embBase + cosHeadOff + baseDim + 4, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase, sin_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase + 4, sin_emb + embBase + cosHeadOff + baseDim + 4, 16); + } + else + { + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase, cos_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase, sin_emb + embBase + cosHeadOff + baseDim, 16); + } + } + __pipeline_commit(); + + // Phase 0b: sync load q_weight + k_weight -> regs (overlaps cp.async transfers). + uint4 q_w_cache[MAX_CHUNKS], k_w_cache[MAX_CHUNKS]; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + int const headIdx = elemBase / HEAD_DIM; + int const baseDim = elemBase - headIdx * HEAD_DIM; + q_w_cache[chunk] = *reinterpret_cast(&q_weight[headIdx * HEAD_DIM + baseDim]); + k_w_cache[chunk] = *reinterpret_cast(&k_weight[headIdx * HEAD_DIM + baseDim]); + } + + // Phase 0c: wait + sync. + __pipeline_wait_prior(0); + __syncthreads(); + + // Phase 1: compute sum²_Q and sum²_K together from SMEM. + float q_sum2 = 0.0f, k_sum2 = 0.0f; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + uint4 const qv = *reinterpret_cast(&smem_q[row_in_block * N + elemBase]); + uint4 const kv = *reinterpret_cast(&smem_k[row_in_block * N + elemBase]); + uint const* qu = reinterpret_cast(&qv); + uint const* ku = reinterpret_cast(&kv); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 qv2 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&qu[i])); + float2 kv2 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&ku[i])); + q_sum2 += qv2.x * qv2.x + qv2.y * qv2.y; + k_sum2 += kv2.x * kv2.x + kv2.y * kv2.y; + } + } + + // Per-row reduce both at once: pack (q_sum, k_sum) per warp slot. + q_sum2 = tensorrt_llm::common::warpReduceSum(q_sum2); + k_sum2 = tensorrt_llm::common::warpReduceSum(k_sum2); + // warp_sums layout: [row_in_block][2 * warp + (0=Q, 1=K)] + if (row_lane == 0) + { + warp_sums[row_in_block * (2 * WARPS_PER_ROW) + 2 * row_warp + 0] = q_sum2; + warp_sums[row_in_block * (2 * WARPS_PER_ROW) + 2 * row_warp + 1] = k_sum2; + } + __syncthreads(); + float q_total = 0.0f, k_total = 0.0f; +#pragma unroll + for (int w = 0; w < WARPS_PER_ROW; w++) + { + q_total += warp_sums[row_in_block * (2 * WARPS_PER_ROW) + 2 * w + 0]; + k_total += warp_sums[row_in_block * (2 * WARPS_PER_ROW) + 2 * w + 1]; + } + float const q_rms_rcp = rsqrtf(q_total / static_cast(N) + eps); + float const k_rms_rcp = rsqrtf(k_total / static_cast(N) + eps); + + // Phase 2: apply norm + RoPE to Q and K, writing to HBM. + // Cos/sin loaded from SMEM (same stage as Q+K), converted to fp32 at use. + auto apply_chunk + = [&](int chunk, __nv_bfloat16 const* smem_input, uint4 const* w_cache, int64_t tokenBaseOut, float rms_rcp) + { + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + return; + uint4 const in_vec = *reinterpret_cast(&smem_input[row_in_block * N + elemBase]); + uint4 const w_vec = w_cache[chunk]; + + float cos_vals[CHUNK_ELEMS]; + float sin_vals[CHUNK_ELEMS]; + if constexpr (std::is_same_v) + { + float4 const* cs = reinterpret_cast(&smem_cos[row_in_block * N + elemBase]); + float4 const* ss = reinterpret_cast(&smem_sin[row_in_block * N + elemBase]); + float4 c0 = cs[0], c1 = cs[1]; + float4 s0 = ss[0], s1 = ss[1]; + cos_vals[0] = c0.x; + cos_vals[1] = c0.y; + cos_vals[2] = c0.z; + cos_vals[3] = c0.w; + cos_vals[4] = c1.x; + cos_vals[5] = c1.y; + cos_vals[6] = c1.z; + cos_vals[7] = c1.w; + sin_vals[0] = s0.x; + sin_vals[1] = s0.y; + sin_vals[2] = s0.z; + sin_vals[3] = s0.w; + sin_vals[4] = s1.x; + sin_vals[5] = s1.y; + sin_vals[6] = s1.z; + sin_vals[7] = s1.w; + } + else + { + uint4 const cp = *reinterpret_cast(&smem_cos[row_in_block * N + elemBase]); + uint4 const sp = *reinterpret_cast(&smem_sin[row_in_block * N + elemBase]); + uint const* cu = reinterpret_cast(&cp); + uint const* su = reinterpret_cast(&sp); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 cv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&cu[i])); + float2 sv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&su[i])); + cos_vals[2 * i] = cv.x; + cos_vals[2 * i + 1] = cv.y; + sin_vals[2 * i] = sv.x; + sin_vals[2 * i + 1] = sv.y; + } + } + + float elements[CHUNK_ELEMS]; + float w_vals[CHUNK_ELEMS]; + uint const* x_uints = reinterpret_cast(&in_vec); + uint const* w_uints = reinterpret_cast(&w_vec); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 xv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&x_uints[i])); + float2 wv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&w_uints[i])); + elements[2 * i] = xv.x; + elements[2 * i + 1] = xv.y; + w_vals[2 * i] = wv.x; + w_vals[2 * i + 1] = wv.y; + } + +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i++) + elements[i] *= rms_rcp * w_vals[i]; + + if constexpr (INTERLEAVE) + { +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i += 2) + { + float const x = elements[i], y = elements[i + 1]; + elements[i] = x * cos_vals[i] + (-y) * sin_vals[i]; + elements[i + 1] = y * cos_vals[i + 1] + x * sin_vals[i + 1]; + } + } + else + { + constexpr int xor_mask = HEAD_DIM / 16; + bool const negate = ((row_lane & xor_mask) == 0); + unsigned const activeMask = __activemask(); +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i++) + { + float p = __shfl_xor_sync(activeMask, elements[i], xor_mask); + if (negate) + { + p = -p; + } + elements[i] = elements[i] * cos_vals[i] + p * sin_vals[i]; + } + } + + uint4 out_vec; + uint* o_uints = reinterpret_cast(&out_vec); +#pragma unroll + for (int i = 0; i < 4; i++) + { + __nv_bfloat162 vals = __float22bfloat162_rn(make_float2(elements[2 * i], elements[2 * i + 1])); + reinterpret_cast<__nv_bfloat162&>(o_uints[i]) = vals; + } + *reinterpret_cast(&qkv[tokenBaseOut + elemBase]) = out_vec; + }; + +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + apply_chunk(chunk, smem_q, q_w_cache, tokenBaseQ, q_rms_rcp); + apply_chunk(chunk, smem_k, k_w_cache, tokenBaseK, k_rms_rcp); + } + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + //////////////////////////////////////////////////////////////////////////////////////////////////// void launchFusedDiTCrossHeadQKNormRope(void* qkv, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, @@ -478,6 +757,84 @@ void launchFusedDiTCrossHeadQKNormRope(void* qkv, int num_tokens, int num_heads_ #undef LAUNCH_CROSS_HEAD_KERNEL } +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Full-dim launch: norm range = num_heads_q * head_dim (LTX-2 mode). +// Requires num_heads_q == num_heads_k (block_size identical for grid.y=0/1). + +void launchFusedDiTQKNormRopeFullDim(void* qkv, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, + int head_dim, float eps, void const* q_weight, void const* k_weight, void const* cos_emb, void const* sin_emb, + bool interleave, bool per_head_cos, bool cos_is_bf16, int cos_seq_per_batch, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(num_heads_q == num_heads_k, + "fusedDiTQKNormRopeFullDim: requires num_heads_q == num_heads_k (got %d, %d)", num_heads_q, num_heads_k); + TLLM_CHECK_WITH_INFO(num_heads_q <= 32, "fusedDiTQKNormRopeFullDim: num_heads must be <= 32, got %d", num_heads_q); + + int const N = num_heads_q * head_dim; + constexpr int ROWS_PER_BLOCK = 2; + + cudaLaunchAttribute attrs[1] = {}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK); + cfg.blockDim = dim3(256); + // SMEM = Q+K stage (bf16) + cos+sin stage (CosT) + warp_sums. + size_t const cos_elem_size = cos_is_bf16 ? sizeof(__nv_bfloat16) : sizeof(float); + cfg.dynamicSmemBytes = 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * N * cos_elem_size + + ROWS_PER_BLOCK * 2 * 4 /*WARPS_PER_ROW*/ * sizeof(float); + cfg.stream = stream; + cfg.attrs = attrs; + cfg.numAttrs = 1; + // Default per-CTA dynamic SMEM cap is 48 KB; raise it per kernel specialization. +#define LAUNCH(HEAD_DIM, INTERLEAVE, PER_HEAD, COS_T) \ + do \ + { \ + auto* kptr = fusedDiTQKNormFullDimRopeKernel; \ + cudaFuncSetAttribute( \ + reinterpret_cast(kptr), cudaFuncAttributeMaxDynamicSharedMemorySize, cfg.dynamicSmemBytes); \ + cudaLaunchKernelEx(&cfg, kptr, reinterpret_cast<__nv_bfloat16*>(qkv), num_heads_q, num_heads_k, num_heads_v, \ + eps, reinterpret_cast<__nv_bfloat16 const*>(q_weight), reinterpret_cast<__nv_bfloat16 const*>(k_weight), \ + reinterpret_cast(cos_emb), reinterpret_cast(sin_emb), num_tokens, \ + cos_seq_per_batch); \ + } while (0) +#define DISPATCH(INTERLEAVE, PER_HEAD, COS_T) \ + do \ + { \ + switch (head_dim) \ + { \ + case 64: LAUNCH(64, INTERLEAVE, PER_HEAD, COS_T); break; \ + case 128: LAUNCH(128, INTERLEAVE, PER_HEAD, COS_T); break; \ + default: TLLM_THROW("Unsupported head_dim for fusedDiTQKNormRopeFullDim: %d", head_dim); \ + } \ + } while (0) +#define DISPATCH_DTYPE(INTERLEAVE, PER_HEAD) \ + do \ + { \ + if (cos_is_bf16) \ + DISPATCH(INTERLEAVE, PER_HEAD, __nv_bfloat16); \ + else \ + DISPATCH(INTERLEAVE, PER_HEAD, float); \ + } while (0) + if (interleave) + { + if (per_head_cos) + DISPATCH_DTYPE(true, true); + else + DISPATCH_DTYPE(true, false); + } + else + { + if (per_head_cos) + DISPATCH_DTYPE(false, true); + else + DISPATCH_DTYPE(false, false); + } +#undef DISPATCH_DTYPE +#undef DISPATCH +#undef LAUNCH +} + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.h index 0c249b86a352..d937c512cc7d 100644 --- a/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.h @@ -74,6 +74,24 @@ void launchFusedDiTCrossHeadQKNormRope(void* qkv, // [num_tokens, (Hq+Hk+Hv)*hea bool interleave, // true = interleaved pairs, false = rotate_half cudaStream_t stream); +// Full-dim variant for LTX-2: RMSNorm range = num_heads_per_side * head_dim. +// Requires num_heads_q == num_heads_k. No dual-stream support. +// per_head_cos=false: cos/sin shape [num_tokens, head_dim] (head broadcast). +// per_head_cos=true: cos/sin shape [num_tokens, num_heads*head_dim] +// (LTX-2 INTERLEAVED 3D RoPE — different freqs per head). +void launchFusedDiTQKNormRopeFullDim(void* qkv, // [num_tokens, (Hq+Hk+Hv)*head_dim], in-place + int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, + int head_dim, // Must be 64 or 128 + float eps, + void const* q_weight, // [num_heads_q * head_dim] + void const* k_weight, // [num_heads_k * head_dim] + void const* cos_emb, // float32 or bfloat16 (selected by cos_is_bf16) + void const* sin_emb, // same dtype as cos_emb + bool interleave, bool per_head_cos, + bool cos_is_bf16, // true → cos/sin are bf16 + int cos_seq_per_batch, // 0 = flat cos [num_tokens, …]; >0 = cos broadcast over B + cudaStream_t stream); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.cu b/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.cu new file mode 100644 index 000000000000..3c2fe0c87010 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.cu @@ -0,0 +1,211 @@ +/* + * 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 "fusedDiTSplitNormKernel.h" +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/mathUtils.h" +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Fused full-dim RMSNorm only (no RoPE) on a SINGLE Q or K tensor. +// Strategy: +// - 2 rows per CTA (256 threads = 2 rows × 128 threads × 4 warps). +// - cp.async X HBM → SMEM (Phase 0a) overlaps with sync weight load → regs (Phase 0b). +// - Phase 1: sum^2 from SMEM, per-row reduce. +// - Phase 2: re-read X from SMEM, multiply by cached weight regs, write HBM. +template +__global__ void fusedDiTSplitNormFullDimKernel(__nv_bfloat16* __restrict__ tensor, int const num_tokens, + int const num_heads, float const eps, __nv_bfloat16 const* __restrict__ weight) +{ + constexpr int BLOCK_SIZE = 256; + constexpr int ROWS_PER_BLOCK = 2; + constexpr int THREADS_PER_ROW = BLOCK_SIZE / ROWS_PER_BLOCK; // 128 + constexpr int WARPS_PER_ROW = THREADS_PER_ROW / 32; // 4 + constexpr int CHUNK_ELEMS = 8; // uint4 = 8 bf16 + constexpr int MAX_N = 32 * HEAD_DIM; + constexpr int MAX_CHUNKS_PER_ROW = (MAX_N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;"); +#endif + + int const tid = threadIdx.x; + int const row_in_block = tid / THREADS_PER_ROW; // 0 or 1 + int const lane_in_row = tid % THREADS_PER_ROW; // 0..127 + int const row_warp = lane_in_row >> 5; // 0..3 + int const row_lane = lane_in_row & 31; + + int const tokenIdx = blockIdx.x * ROWS_PER_BLOCK + row_in_block; + if (tokenIdx >= num_tokens) + return; + + int const N = num_heads * HEAD_DIM; + int const chunks_per_row = (N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + int64_t const tokenBase = static_cast(tokenIdx) * N; + + extern __shared__ __align__(16) unsigned char smem_raw[]; + __nv_bfloat16* smem_input = reinterpret_cast<__nv_bfloat16*>(smem_raw); + float* warp_sums = reinterpret_cast(smem_raw + ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16)); + + // Phase 0a: issue cp.async X HBM -> SMEM (one commit group). +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS_PER_ROW; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + __pipeline_memcpy_async(smem_input + row_in_block * N + elemBase, tensor + tokenBase + elemBase, 16); + } + __pipeline_commit(); + + // Phase 0b: SYNC load weight into registers (overlaps with cp.async X HBM transfer). + // weight_cache[chunk] holds 8 bf16 weight elements per chunk. + uint4 weight_cache[MAX_CHUNKS_PER_ROW]; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS_PER_ROW; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + int const headIdx = elemBase / HEAD_DIM; + int const baseDim = elemBase - headIdx * HEAD_DIM; + weight_cache[chunk] = *reinterpret_cast(&weight[headIdx * HEAD_DIM + baseDim]); + } + + // Phase 0c: wait for X cp.async + sync block. + __pipeline_wait_prior(0); + __syncthreads(); + + // Phase 1: sum^2 from SMEM. + float sum2 = 0.0f; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS_PER_ROW; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + uint4 const v = *reinterpret_cast(&smem_input[row_in_block * N + elemBase]); + uint const* uints = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&uints[i])); + sum2 += vals.x * vals.x + vals.y * vals.y; + } + } + + // Per-row warp reduce + cross-warp reduce. + sum2 = tensorrt_llm::common::warpReduceSum(sum2); + if (row_lane == 0) + warp_sums[row_in_block * WARPS_PER_ROW + row_warp] = sum2; + __syncthreads(); + float total = 0.0f; +#pragma unroll + for (int w = 0; w < WARPS_PER_ROW; w++) + total += warp_sums[row_in_block * WARPS_PER_ROW + w]; + float const rms_rcp = rsqrtf(total / static_cast(N) + eps); + + // Phase 2: re-read X from SMEM, multiply by cached weight regs, write to HBM. +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS_PER_ROW; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + + uint4 const in_vec = *reinterpret_cast(&smem_input[row_in_block * N + elemBase]); + uint4 const w_vec = weight_cache[chunk]; + + uint const* x_uints = reinterpret_cast(&in_vec); + uint const* w_uints = reinterpret_cast(&w_vec); + uint4 out_vec; + uint* o_uints = reinterpret_cast(&out_vec); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 x_vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&x_uints[i])); + float2 w_vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&w_uints[i])); + float2 y_vals; + y_vals.x = x_vals.x * rms_rcp * w_vals.x; + y_vals.y = x_vals.y * rms_rcp * w_vals.y; + __nv_bfloat162 bf = __float22bfloat162_rn(y_vals); + reinterpret_cast<__nv_bfloat162&>(o_uints[i]) = bf; + } + *reinterpret_cast(&tensor[tokenBase + elemBase]) = out_vec; + } + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void launchFusedDiTSplitNormFullDim( + void* tensor, int num_tokens, int num_heads, int head_dim, float eps, void const* weight, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(num_heads <= 32, + "fusedDiTSplitNormFullDim: num_heads (%d) must be <= 32 (block_size = num_heads*32 <= 1024)", num_heads); + TLLM_CHECK_WITH_INFO(num_heads >= 1, "fusedDiTSplitNormFullDim: num_heads must be >= 1, got %d", num_heads); + + int const N = num_heads * head_dim; + constexpr int ROWS_PER_BLOCK = 2; + + cudaLaunchAttribute attrs[1] = {}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes + = ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + ROWS_PER_BLOCK * 4 /*warps_per_row*/ * sizeof(float); + cfg.stream = stream; + cfg.attrs = attrs; + cfg.numAttrs = 1; +#define LAUNCH(HEAD_DIM) \ + cudaLaunchKernelEx(&cfg, fusedDiTSplitNormFullDimKernel, reinterpret_cast<__nv_bfloat16*>(tensor), \ + num_tokens, num_heads, eps, reinterpret_cast<__nv_bfloat16 const*>(weight)) + switch (head_dim) + { + case 64: LAUNCH(64); break; + case 128: LAUNCH(128); break; + default: TLLM_THROW("Unsupported head_dim for fusedDiTSplitNormFullDim: %d (only 64, 128)", head_dim); + } +#undef LAUNCH +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.h b/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.h new file mode 100644 index 000000000000..51f88df345c5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/fusedDiTSplitNormKernel.h @@ -0,0 +1,56 @@ +/* + * 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. + */ + +#ifndef TRTLLM_FUSEDDITSPLITNORMKERNEL_H +#define TRTLLM_FUSEDDITSPLITNORMKERNEL_H + +#include "tensorrt_llm/common/config.h" +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Fused full-dim RMSNorm (no RoPE) for a SINGLE Q or K tensor. +// +// Mirror of fusedDiTSplitNormFullDimRopeKernel but without the RoPE step -- +// used by LTX-2 cross-attn paths where Q/K need norm but no RoPE (e.g. text +// cross-attention where positional info is already baked into the text +// encoder output). +// +// Layout: input is a contiguous 2D tensor [num_tokens, num_heads * head_dim] +// (e.g. output of self.to_q / self.to_k). Block=256 with chunked reduce; cross +// -warp shared-memory sum^2 reduction over the full inner dim +// (num_heads * head_dim). +// +// Constraints: +// - num_heads <= 32 +// - head_dim ∈ {64, 128} +// - num_heads * head_dim <= 4096 (chunk count cap) + +void launchFusedDiTSplitNormFullDim(void* tensor, // [num_tokens, num_heads * head_dim], bf16, contiguous, in-place + int num_tokens, int num_heads, + int head_dim, // 64 or 128 + float eps, + void const* weight, // bf16, [num_heads * head_dim] (full-dim norm weight) + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_FUSEDDITSPLITNORMKERNEL_H diff --git a/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.cu new file mode 100644 index 000000000000..11b7d4f849b4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.cu @@ -0,0 +1,362 @@ +/* + * 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 "fusedDiTSplitQKNormRopeKernel.h" +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/mathUtils.h" +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Fused full-dim RMSNorm + RoPE on a SINGLE Q or K tensor (LTX-2 SEPARATE_QKV cross-attn). +// Strategy: +// - 2 rows per CTA (256 threads = 2 rows x 128 threads x 4 warps). +// - Phase 0a: cp.async X + cos + sin (HBM -> SMEM) in one commit group. +// - Phase 0b: sync load weight -> regs (overlaps the cp.async transfers). +// - Phase 1: sum^2 reads X from SMEM (no HBM re-read). +// - Phase 2: reads X + cos + sin from SMEM, multiplies by cached weight regs, writes HBM. +template +__global__ void fusedDiTSplitNormFullDimRopeKernel(__nv_bfloat16* __restrict__ tensor, int const num_tokens, + int const num_heads, float const eps, __nv_bfloat16 const* __restrict__ weight, CosT const* __restrict__ cos_emb, + CosT const* __restrict__ sin_emb, int const cos_seq_per_batch) +{ + constexpr int BLOCK_SIZE = 256; + constexpr int ROWS_PER_BLOCK = 2; + constexpr int THREADS_PER_ROW = BLOCK_SIZE / ROWS_PER_BLOCK; // 128 + constexpr int WARPS_PER_ROW = THREADS_PER_ROW / 32; // 4 + constexpr int CHUNK_ELEMS = 8; // uint4 = 8 bf16 + constexpr int MAX_N = 32 * HEAD_DIM; + constexpr int MAX_CHUNKS = (MAX_N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;"); +#endif + + int const tid = threadIdx.x; + int const row_in_block = tid / THREADS_PER_ROW; // 0 or 1 + int const lane_in_row = tid % THREADS_PER_ROW; // 0..127 + int const row_warp = lane_in_row >> 5; // 0..3 + int const row_lane = lane_in_row & 31; + + int const tokenIdx = blockIdx.x * ROWS_PER_BLOCK + row_in_block; + if (tokenIdx >= num_tokens) + return; + + int const N = num_heads * HEAD_DIM; + int const chunks_per_row = (N + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); + int64_t const tokenBase = static_cast(tokenIdx) * N; + int const cos_tokenIdx = (cos_seq_per_batch > 0) ? (tokenIdx % cos_seq_per_batch) : tokenIdx; + int64_t const embBase = PER_HEAD_COS ? static_cast(cos_tokenIdx) * num_heads * HEAD_DIM + : static_cast(cos_tokenIdx) * HEAD_DIM; + + // SMEM layout: [X bf16 stage][cos CosT stage][sin CosT stage][warp_sums fp32]. + extern __shared__ __align__(16) unsigned char smem_raw[]; + __nv_bfloat16* smem_input = reinterpret_cast<__nv_bfloat16*>(smem_raw); + CosT* smem_cos = reinterpret_cast(smem_raw + ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16)); + CosT* smem_sin = smem_cos + ROWS_PER_BLOCK * N; + float* warp_sums = reinterpret_cast( + smem_raw + ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * N * sizeof(CosT)); + + // Phase 0a: cp.async X + cos + sin HBM -> SMEM. +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + __pipeline_memcpy_async(smem_input + row_in_block * N + elemBase, tensor + tokenBase + elemBase, 16); + int const headIdx = elemBase / HEAD_DIM; + int const baseDim = elemBase - headIdx * HEAD_DIM; + int const cosHeadOff = PER_HEAD_COS ? headIdx * HEAD_DIM : 0; + if constexpr (std::is_same_v) + { + // fp32 cos: 8 floats per chunk = 32 bytes = 2x 16-byte cp.async per array. + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase, cos_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase + 4, cos_emb + embBase + cosHeadOff + baseDim + 4, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase, sin_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase + 4, sin_emb + embBase + cosHeadOff + baseDim + 4, 16); + } + else + { + // bf16 cos: 8 bf16 = 16 bytes = 1x cp.async per array. + __pipeline_memcpy_async( + smem_cos + row_in_block * N + elemBase, cos_emb + embBase + cosHeadOff + baseDim, 16); + __pipeline_memcpy_async( + smem_sin + row_in_block * N + elemBase, sin_emb + embBase + cosHeadOff + baseDim, 16); + } + } + __pipeline_commit(); + + // Phase 0b: sync load weight -> regs (overlaps with cp.async transfers above). + uint4 weight_cache[MAX_CHUNKS]; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + int const headIdx = elemBase / HEAD_DIM; + int const baseDim = elemBase - headIdx * HEAD_DIM; + weight_cache[chunk] = *reinterpret_cast(&weight[headIdx * HEAD_DIM + baseDim]); + } + + // Phase 0c: wait + sync. + __pipeline_wait_prior(0); + __syncthreads(); + + // Phase 1: sum^2 from SMEM. + float sum2 = 0.0f; +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + uint4 const v = *reinterpret_cast(&smem_input[row_in_block * N + elemBase]); + uint const* uints = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&uints[i])); + sum2 += vals.x * vals.x + vals.y * vals.y; + } + } + + // Per-row warp reduce + cross-warp reduce. + sum2 = tensorrt_llm::common::warpReduceSum(sum2); + if (row_lane == 0) + warp_sums[row_in_block * WARPS_PER_ROW + row_warp] = sum2; + __syncthreads(); + float total = 0.0f; +#pragma unroll + for (int w = 0; w < WARPS_PER_ROW; w++) + total += warp_sums[row_in_block * WARPS_PER_ROW + w]; + float const rms_rcp = rsqrtf(total / static_cast(N) + eps); + + // Phase 2: read X + cos + sin from SMEM, apply cached weight, RoPE, write to HBM. +#pragma unroll + for (int chunk = 0; chunk < MAX_CHUNKS; chunk++) + { + if (chunk >= chunks_per_row) + continue; + int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; + if (elemBase >= N) + continue; + + uint4 const in_vec = *reinterpret_cast(&smem_input[row_in_block * N + elemBase]); + uint4 const w_vec = weight_cache[chunk]; + + float cos_vals[CHUNK_ELEMS]; + float sin_vals[CHUNK_ELEMS]; + if constexpr (std::is_same_v) + { + float4 const* cs = reinterpret_cast(&smem_cos[row_in_block * N + elemBase]); + float4 const* ss = reinterpret_cast(&smem_sin[row_in_block * N + elemBase]); + float4 c0 = cs[0], c1 = cs[1]; + float4 s0 = ss[0], s1 = ss[1]; + cos_vals[0] = c0.x; + cos_vals[1] = c0.y; + cos_vals[2] = c0.z; + cos_vals[3] = c0.w; + cos_vals[4] = c1.x; + cos_vals[5] = c1.y; + cos_vals[6] = c1.z; + cos_vals[7] = c1.w; + sin_vals[0] = s0.x; + sin_vals[1] = s0.y; + sin_vals[2] = s0.z; + sin_vals[3] = s0.w; + sin_vals[4] = s1.x; + sin_vals[5] = s1.y; + sin_vals[6] = s1.z; + sin_vals[7] = s1.w; + } + else + { + uint4 const cp = *reinterpret_cast(&smem_cos[row_in_block * N + elemBase]); + uint4 const sp = *reinterpret_cast(&smem_sin[row_in_block * N + elemBase]); + uint const* cu = reinterpret_cast(&cp); + uint const* su = reinterpret_cast(&sp); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 cv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&cu[i])); + float2 sv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&su[i])); + cos_vals[2 * i] = cv.x; + cos_vals[2 * i + 1] = cv.y; + sin_vals[2 * i] = sv.x; + sin_vals[2 * i + 1] = sv.y; + } + } + + float elements[CHUNK_ELEMS]; + float w_vals[CHUNK_ELEMS]; + uint const* x_uints = reinterpret_cast(&in_vec); + uint const* w_uints = reinterpret_cast(&w_vec); +#pragma unroll + for (int i = 0; i < 4; i++) + { + float2 xv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&x_uints[i])); + float2 wv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&w_uints[i])); + elements[2 * i] = xv.x; + elements[2 * i + 1] = xv.y; + w_vals[2 * i] = wv.x; + w_vals[2 * i + 1] = wv.y; + } + +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i++) + elements[i] *= rms_rcp * w_vals[i]; + + if constexpr (INTERLEAVE) + { +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i += 2) + { + float const x = elements[i], y = elements[i + 1]; + elements[i] = x * cos_vals[i] + (-y) * sin_vals[i]; + elements[i + 1] = y * cos_vals[i + 1] + x * sin_vals[i + 1]; + } + } + else + { + constexpr int xor_mask = HEAD_DIM / 16; + bool const negate = ((row_lane & xor_mask) == 0); + unsigned const activeMask = __activemask(); +#pragma unroll + for (int i = 0; i < CHUNK_ELEMS; i++) + { + float p = __shfl_xor_sync(activeMask, elements[i], xor_mask); + if (negate) + { + p = -p; + } + elements[i] = elements[i] * cos_vals[i] + p * sin_vals[i]; + } + } + + uint4 out_vec; + uint* o_uints = reinterpret_cast(&out_vec); +#pragma unroll + for (int i = 0; i < 4; i++) + { + __nv_bfloat162 vals = __float22bfloat162_rn(make_float2(elements[2 * i], elements[2 * i + 1])); + reinterpret_cast<__nv_bfloat162&>(o_uints[i]) = vals; + } + *reinterpret_cast(&tensor[tokenBase + elemBase]) = out_vec; + } + +#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void launchFusedDiTSplitNormFullDimRope(void* tensor, int num_tokens, int num_heads, int head_dim, float eps, + void const* weight, void const* cos_emb, void const* sin_emb, bool interleave, bool per_head_cos, bool cos_is_bf16, + int cos_seq_per_batch, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(num_heads <= 32, + "fusedDiTSplitNormFullDimRope: num_heads (%d) must be <= 32 (block_size = num_heads*32 <= 1024)", num_heads); + TLLM_CHECK_WITH_INFO(num_heads >= 1, "fusedDiTSplitNormFullDimRope: num_heads must be >= 1, got %d", num_heads); + + int const N = num_heads * head_dim; + constexpr int ROWS_PER_BLOCK = 2; + + cudaLaunchAttribute attrs[1] = {}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK); + cfg.blockDim = dim3(256); + // SMEM = X stage (bf16) + cos stage (CosT) + sin stage (CosT) + warp_sums. + size_t const cos_elem_size = cos_is_bf16 ? sizeof(__nv_bfloat16) : sizeof(float); + cfg.dynamicSmemBytes = ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * N * cos_elem_size + + ROWS_PER_BLOCK * 4 /*WARPS_PER_ROW*/ * sizeof(float); + cfg.stream = stream; + cfg.attrs = attrs; + cfg.numAttrs = 1; + // Default per-CTA dynamic SMEM cap is 48 KB; raise it per kernel specialization. +#define LAUNCH(HEAD_DIM, INTERLEAVE, PER_HEAD, COS_T) \ + do \ + { \ + auto* kptr = fusedDiTSplitNormFullDimRopeKernel; \ + cudaFuncSetAttribute( \ + reinterpret_cast(kptr), cudaFuncAttributeMaxDynamicSharedMemorySize, cfg.dynamicSmemBytes); \ + cudaLaunchKernelEx(&cfg, kptr, reinterpret_cast<__nv_bfloat16*>(tensor), num_tokens, num_heads, eps, \ + reinterpret_cast<__nv_bfloat16 const*>(weight), reinterpret_cast(cos_emb), \ + reinterpret_cast(sin_emb), cos_seq_per_batch); \ + } while (0) +#define DISPATCH(INTERLEAVE, PER_HEAD, COS_T) \ + do \ + { \ + switch (head_dim) \ + { \ + case 64: LAUNCH(64, INTERLEAVE, PER_HEAD, COS_T); break; \ + case 128: LAUNCH(128, INTERLEAVE, PER_HEAD, COS_T); break; \ + default: TLLM_THROW("Unsupported head_dim for fusedDiTSplitNormFullDimRope: %d (only 64, 128)", head_dim); \ + } \ + } while (0) +#define DISPATCH_DTYPE(INTERLEAVE, PER_HEAD) \ + do \ + { \ + if (cos_is_bf16) \ + DISPATCH(INTERLEAVE, PER_HEAD, __nv_bfloat16); \ + else \ + DISPATCH(INTERLEAVE, PER_HEAD, float); \ + } while (0) + if (interleave) + { + if (per_head_cos) + DISPATCH_DTYPE(true, true); + else + DISPATCH_DTYPE(true, false); + } + else + { + if (per_head_cos) + DISPATCH_DTYPE(false, true); + else + DISPATCH_DTYPE(false, false); + } +#undef DISPATCH_DTYPE +#undef DISPATCH +#undef LAUNCH +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.h new file mode 100644 index 000000000000..4523ea7dcaff --- /dev/null +++ b/cpp/tensorrt_llm/kernels/fusedDiTSplitQKNormRopeKernel.h @@ -0,0 +1,61 @@ +/* + * 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. + */ + +#ifndef TRTLLM_FUSEDDITSPLITQKNORMROPEKERNEL_H +#define TRTLLM_FUSEDDITSPLITQKNORMROPEKERNEL_H + +#include "tensorrt_llm/common/config.h" +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// Fused full-dim RMSNorm + RoPE for a SINGLE Q or K tensor (DiT models, e.g. LTX-2). +// +// Designed for SEPARATE_QKV layout: input is a contiguous 2D tensor +// [num_tokens, num_heads * head_dim] (e.g. output of self.to_q / self.to_k). +// Block=256 with chunked reduce; cross-warp shared-memory sum² reduction over +// the full inner dim (num_heads * head_dim). +// +// For FUSE_QKV layout (packed QKV buffer), use fusedDiTQKNormFullDimRopeKernel +// (full-dim variant in fusedDiTQKNormRopeKernel.h) -- it processes Q+K together +// in one launch. +// +// Constraints: +// - num_heads ≤ 32 +// - head_dim ∈ {64, 128} +// - num_heads * head_dim ≤ 4096 (chunk count cap) + +void launchFusedDiTSplitNormFullDimRope(void* tensor, // [num_tokens, num_heads * head_dim], bf16, contiguous, in-place + int num_tokens, int num_heads, + int head_dim, // 64 or 128 + float eps, + void const* weight, // bf16, [num_heads * head_dim] (full-dim norm weight) + void const* cos_emb, // float32 or bfloat16 (selected by cos_is_bf16) + void const* sin_emb, // same dtype as cos_emb + bool interleave, // true = pair (2i, 2i+1); false = rotate_half + bool per_head_cos, // false: cos shape [N, head_dim]; true: [N, num_heads*head_dim] + bool cos_is_bf16, // true → cos/sin are bf16 (no fp32 cast needed upstream) + int cos_seq_per_batch, // 0 = flat cos [num_tokens, …]; >0 = cos broadcast over B (cos rows = cos_seq_per_batch) + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_FUSEDDITSPLITQKNORMROPEKERNEL_H diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index a31e92ec095e..6849415992b7 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -67,6 +67,8 @@ add_library( dsv3FusedAGemmOp.cpp fusedQKNormRopeOp.cpp fusedDiTQKNormRopeOp.cpp + fusedDiTSplitQKNormRopeOp.cpp + fusedDiTSplitNormOp.cpp fusedAddRMSNormQuant.cpp fusedActivationQuant.cpp fusedGatedRMSNormQuant.cpp diff --git a/cpp/tensorrt_llm/thop/fusedDiTQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedDiTQKNormRopeOp.cpp index 56dadf5f4290..fb046f455530 100644 --- a/cpp/tensorrt_llm/thop/fusedDiTQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedDiTQKNormRopeOp.cpp @@ -46,31 +46,81 @@ void fused_dit_qk_norm_rope(torch::Tensor& qkv, // [num_tokens, (Hq+Hk+Hv)*head_ TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, total_heads*head_dim]"); TORCH_CHECK(q_weight.dim() == 1, "q_weight must be 1D"); TORCH_CHECK(k_weight.dim() == 1, "k_weight must be 1D"); - TORCH_CHECK(cos_emb.dim() == 2, "cos_emb must be 2D: [num_tokens, head_dim]"); - TORCH_CHECK(sin_emb.dim() == 2, "sin_emb must be 2D: [num_tokens, head_dim]"); + TORCH_CHECK(cos_emb.dim() == 2, "cos_emb must be 2D: [num_tokens, head_dim] or [num_tokens, num_heads*head_dim]"); + TORCH_CHECK(sin_emb.dim() == 2, "sin_emb must be 2D: [num_tokens, head_dim] or [num_tokens, num_heads*head_dim]"); CHECK_INPUT(qkv, torch::kBFloat16); CHECK_INPUT(q_weight, torch::kBFloat16); CHECK_INPUT(k_weight, torch::kBFloat16); - CHECK_INPUT(cos_emb, torch::kFloat32); - CHECK_INPUT(sin_emb, torch::kFloat32); + // Cos/sin may be fp32 (per-head FLUX path) or bf16 (B-2 full-dim LTX-2 path). + // Per-head path requires fp32 (kernel has no bf16 branch); enforced below. + auto const cos_dtype = cos_emb.scalar_type(); + TORCH_CHECK(cos_dtype == torch::kFloat32 || cos_dtype == torch::kBFloat16, + "cos_emb dtype must be float32 or bfloat16, got ", cos_dtype); + TORCH_CHECK(sin_emb.scalar_type() == cos_dtype, "sin_emb dtype must match cos_emb"); + bool const cos_is_bf16 = (cos_dtype == torch::kBFloat16); + if (cos_is_bf16) + { + CHECK_INPUT(cos_emb, torch::kBFloat16); + CHECK_INPUT(sin_emb, torch::kBFloat16); + } + else + { + CHECK_INPUT(cos_emb, torch::kFloat32); + CHECK_INPUT(sin_emb, torch::kFloat32); + } int64_t num_tokens = qkv.size(0); int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; TORCH_CHECK(qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total_heads * head_dim"); - TORCH_CHECK(cos_emb.size(0) == num_tokens && cos_emb.size(1) == head_dim, "cos_emb must be [num_tokens, head_dim]"); - TORCH_CHECK(sin_emb.size(0) == num_tokens && sin_emb.size(1) == head_dim, "sin_emb must be [num_tokens, head_dim]"); - - // Only per-head norm supported - TORCH_CHECK(q_weight.size(0) == head_dim, - "fused_dit_qk_norm_rope only supports per-head norm (q_weight must be [head_dim]). " - "Full-dim norm (q_weight [num_heads * head_dim]) is not yet supported. Got q_weight size: ", - q_weight.size(0), ", head_dim: ", head_dim); - TORCH_CHECK(k_weight.size(0) == head_dim, - "fused_dit_qk_norm_rope only supports per-head norm (k_weight must be [head_dim]). Got k_weight size: ", - k_weight.size(0)); + // Auto-detect broadcast: cos rows == num_tokens (flat) or num_tokens / B (broadcast over B). + int64_t const cos_rows = cos_emb.size(0); + int cos_seq_per_batch = 0; + if (cos_rows != num_tokens) + { + TORCH_CHECK(cos_rows > 0 && num_tokens % cos_rows == 0, "cos_emb.size(0) (", cos_rows, + ") must equal num_tokens (", num_tokens, ") or evenly divide it (broadcast); got non-divisor count"); + cos_seq_per_batch = static_cast(cos_rows); + } + bool const per_head_cos = (cos_emb.size(1) == num_heads_q * head_dim); + TORCH_CHECK(per_head_cos || cos_emb.size(1) == head_dim, "cos_emb last dim must be head_dim (", head_dim, + ") or num_heads_q*head_dim (", num_heads_q * head_dim, "); got ", cos_emb.size(1)); + TORCH_CHECK(sin_emb.size(0) == cos_rows && sin_emb.size(1) == cos_emb.size(1), "sin_emb shape must match cos_emb"); + + // Auto-dispatch by weight shape: + // weight.size(0) == head_dim → per-head norm (FLUX/Cosmos3, original kernel) + // weight.size(0) == num_heads_per_side*head_dim → full-dim norm (LTX-2) + bool const is_full_dim_q = (q_weight.size(0) == num_heads_q * head_dim); + bool const is_full_dim_k = (k_weight.size(0) == num_heads_k * head_dim); + bool const is_per_head_q = (q_weight.size(0) == head_dim); + bool const is_per_head_k = (k_weight.size(0) == head_dim); + TORCH_CHECK(is_full_dim_q == is_full_dim_k && is_per_head_q == is_per_head_k, + "q_weight and k_weight must use the same norm mode (both per-head or both full-dim)."); + TORCH_CHECK(is_per_head_q || is_full_dim_q, + "q_weight size must be [head_dim] (per-head) or [num_heads*head_dim] (full-dim); got ", q_weight.size(0), + " head_dim=", head_dim, " num_heads_q=", num_heads_q); + + if (is_full_dim_q) + { + TORCH_CHECK(!q_add_weight.has_value() && !k_add_weight.has_value(), + "Full-dim norm does not support dual-stream add_weights"); + TORCH_CHECK(num_txt_tokens <= 0, "Full-dim norm does not support dual-stream (num_txt_tokens must be -1)"); + auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + tensorrt_llm::kernels::launchFusedDiTQKNormRopeFullDim(qkv.data_ptr(), static_cast(num_tokens), + static_cast(num_heads_q), static_cast(num_heads_k), static_cast(num_heads_v), + static_cast(head_dim), static_cast(eps), q_weight.data_ptr(), k_weight.data_ptr(), + cos_emb.data_ptr(), sin_emb.data_ptr(), interleave, per_head_cos, cos_is_bf16, cos_seq_per_batch, stream); + return; + } - // Validate optional add_weights (dual-stream) + // Per-head path (original FLUX/Cosmos3 kernel) — only fp32 cos supported here, no broadcast. + TORCH_CHECK(cos_seq_per_batch == 0, + "Per-head fused_dit_qk_norm_rope (FLUX/Cosmos) does not support cos broadcast over B; " + "got cos_emb rows = ", + cos_rows, ", num_tokens = ", num_tokens); + TORCH_CHECK(!cos_is_bf16, + "Per-head fused_dit_qk_norm_rope (FLUX/Cosmos) requires fp32 cos/sin; bf16 cos is only supported " + "by the full-dim path (LTX-2)"); void const* q_add_ptr = nullptr; void const* k_add_ptr = nullptr; if (q_add_weight.has_value()) diff --git a/cpp/tensorrt_llm/thop/fusedDiTSplitNormOp.cpp b/cpp/tensorrt_llm/thop/fusedDiTSplitNormOp.cpp new file mode 100644 index 000000000000..36030ebbb024 --- /dev/null +++ b/cpp/tensorrt_llm/thop/fusedDiTSplitNormOp.cpp @@ -0,0 +1,66 @@ +/* + * 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/kernels/fusedDiTSplitNormKernel.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Fused full-dim RMSNorm only (no RoPE) for a single Q or K tensor. +// Mirror of fused_dit_split_norm_rope but without the RoPE step -- used by +// LTX-2 paths that need norm-only (e.g. text cross-attn Q where positional +// info is already baked into the text encoder output). +// +// Input must be a contiguous 2D tensor [num_tokens, num_heads * head_dim]. +void fused_dit_split_norm(torch::Tensor& tensor, int64_t num_heads, int64_t head_dim, double eps, torch::Tensor& weight) +{ + TORCH_CHECK(tensor.dim() == 2, "tensor must be 2D: [num_tokens, num_heads*head_dim]"); + TORCH_CHECK(weight.dim() == 1, "weight must be 1D"); + + CHECK_INPUT(tensor, torch::kBFloat16); + CHECK_INPUT(weight, torch::kBFloat16); + + int64_t const num_tokens = tensor.size(0); + TORCH_CHECK( + tensor.size(1) == num_heads * head_dim, "tensor inner dim must be num_heads*head_dim; got ", tensor.size(1)); + TORCH_CHECK(weight.size(0) == num_heads * head_dim, "weight must be [num_heads*head_dim] (full-dim norm), got ", + weight.size(0), " expected ", num_heads * head_dim); + + auto stream = at::cuda::getCurrentCUDAStream(tensor.get_device()); + + tensorrt_llm::kernels::launchFusedDiTSplitNormFullDim(tensor.data_ptr(), static_cast(num_tokens), + static_cast(num_heads), static_cast(head_dim), static_cast(eps), weight.data_ptr(), stream); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("fused_dit_split_norm(Tensor(a!) tensor, int num_heads, int head_dim, float eps, Tensor weight) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("fused_dit_split_norm", &fused_dit_split_norm); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fusedDiTSplitQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedDiTSplitQKNormRopeOp.cpp new file mode 100644 index 000000000000..e942ea10580b --- /dev/null +++ b/cpp/tensorrt_llm/thop/fusedDiTSplitQKNormRopeOp.cpp @@ -0,0 +1,102 @@ +/* + * 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/kernels/fusedDiTSplitQKNormRopeKernel.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Fused full-dim RMSNorm + RoPE for a single Q or K tensor (DiT SEPARATE_QKV +// layout, e.g. LTX-2 cross-attn). Input must be a contiguous 2D tensor +// [num_tokens, num_heads * head_dim]. For FUSE_QKV (packed buffer) use +// fused_dit_qk_norm_rope instead. +void fused_dit_split_norm_rope(torch::Tensor& tensor, int64_t num_heads, int64_t head_dim, double eps, + torch::Tensor& weight, torch::Tensor& cos_emb, torch::Tensor& sin_emb, bool interleave) +{ + TORCH_CHECK(tensor.dim() == 2, "tensor must be 2D: [num_tokens, num_heads*head_dim]"); + TORCH_CHECK(weight.dim() == 1, "weight must be 1D"); + TORCH_CHECK(cos_emb.dim() == 2, "cos_emb must be 2D"); + TORCH_CHECK(sin_emb.dim() == 2, "sin_emb must be 2D"); + + CHECK_INPUT(tensor, torch::kBFloat16); + CHECK_INPUT(weight, torch::kBFloat16); + // Cos/sin may be fp32 or bf16 (kernel upcasts bf16 to fp32 in registers, lossless). + auto const cos_dtype = cos_emb.scalar_type(); + TORCH_CHECK(cos_dtype == torch::kFloat32 || cos_dtype == torch::kBFloat16, + "cos_emb dtype must be float32 or bfloat16, got ", cos_dtype); + TORCH_CHECK(sin_emb.scalar_type() == cos_dtype, "sin_emb dtype must match cos_emb (", sin_emb.scalar_type(), " vs ", + cos_dtype, ")"); + bool const cos_is_bf16 = (cos_dtype == torch::kBFloat16); + if (cos_is_bf16) + { + CHECK_INPUT(cos_emb, torch::kBFloat16); + CHECK_INPUT(sin_emb, torch::kBFloat16); + } + else + { + CHECK_INPUT(cos_emb, torch::kFloat32); + CHECK_INPUT(sin_emb, torch::kFloat32); + } + + int64_t const num_tokens = tensor.size(0); + TORCH_CHECK( + tensor.size(1) == num_heads * head_dim, "tensor inner dim must be num_heads*head_dim; got ", tensor.size(1)); + // Auto-detect broadcast: cos may carry one row per token (num_tokens) or one row + // per token-in-batch (num_tokens / B); in the latter case the kernel broadcasts + // cos across B via cos_tokenIdx = tokenIdx % cos_seq_per_batch. + int64_t const cos_rows = cos_emb.size(0); + int cos_seq_per_batch = 0; + if (cos_rows != num_tokens) + { + TORCH_CHECK(cos_rows > 0 && num_tokens % cos_rows == 0, "cos_emb.size(0) (", cos_rows, + ") must equal num_tokens (", num_tokens, ") or evenly divide it (broadcast); got non-divisor count"); + cos_seq_per_batch = static_cast(cos_rows); + } + bool const per_head_cos = (cos_emb.size(1) == num_heads * head_dim); + TORCH_CHECK(per_head_cos || cos_emb.size(1) == head_dim, "cos_emb last dim must be head_dim (", head_dim, + ") or num_heads*head_dim (", num_heads * head_dim, "); got ", cos_emb.size(1)); + TORCH_CHECK(sin_emb.size(0) == cos_rows && sin_emb.size(1) == cos_emb.size(1), "sin_emb shape must match cos_emb"); + TORCH_CHECK(weight.size(0) == num_heads * head_dim, "weight must be [num_heads*head_dim] (full-dim norm), got ", + weight.size(0), " expected ", num_heads * head_dim); + + auto stream = at::cuda::getCurrentCUDAStream(tensor.get_device()); + + tensorrt_llm::kernels::launchFusedDiTSplitNormFullDimRope(tensor.data_ptr(), static_cast(num_tokens), + static_cast(num_heads), static_cast(head_dim), static_cast(eps), weight.data_ptr(), + cos_emb.data_ptr(), sin_emb.data_ptr(), interleave, per_head_cos, cos_is_bf16, cos_seq_per_batch, stream); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "fused_dit_split_norm_rope(Tensor(a!) tensor, int num_heads, int head_dim, float eps, " + "Tensor weight, Tensor cos_emb, Tensor sin_emb, bool interleave) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("fused_dit_split_norm_rope", &fused_dit_split_norm_rope); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index 900175931fcb..83c912abda47 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -85,6 +85,12 @@ def inplace_info(): torch.ops.trtllm.fused_dit_cross_head_qk_norm_rope.default: { 1: "qkv" }, + torch.ops.trtllm.fused_dit_split_norm_rope.default: { + 1: "tensor" + }, + torch.ops.trtllm.fused_dit_split_norm.default: { + 1: "tensor" + }, torch.ops.trtllm.flashinfer_apply_rope_with_cos_sin_cache_inplace.default: { 1: "query", diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py index d890b28754cd..6652b0b0af95 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -192,7 +192,7 @@ def _prepare_qkv_fused( q_add = self.norm_added_q.weight if hasattr(self, "norm_added_q") else None k_add = self.norm_added_k.weight if hasattr(self, "norm_added_k") else None - self.apply_qk_norm_rope( + self.apply_packed_qk_norm_rope( qkv, freqs_cos, freqs_sin, @@ -362,7 +362,7 @@ def _apply_norm_rope_fused( # torch.split produces non-contiguous views; fused kernel requires contiguous qkv = qkv.contiguous() - self.apply_qk_norm_rope(qkv, freqs_cos, freqs_sin, num_txt_tokens=-1) + self.apply_packed_qk_norm_rope(qkv, freqs_cos, freqs_sin, num_txt_tokens=-1) q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) return q, k, v diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py index ae453f194a0a..5b9f301fbc77 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py @@ -52,11 +52,20 @@ def _apply_split_rotary_emb( ) -> torch.Tensor: needs_reshape = False if input_tensor.ndim != 4 and cos_freqs.ndim == 4: - _, h, t, _ = cos_freqs.shape + # cos/sin are token-major [B, T, H, D]; reshape input to match. + _, t, h, _ = cos_freqs.shape b = input_tensor.shape[0] - input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) + input_tensor = input_tensor.reshape(b, t, h, -1) needs_reshape = True + # cos/sin are stored block-duplicated to head_dim (see _split_freqs_cis) + # so the fused kernel can read directly. The SPLIT formula uses only the + # first half (head_dim/2) since both halves are identical. + if cos_freqs.shape[-1] == input_tensor.shape[-1]: + half = cos_freqs.shape[-1] // 2 + cos_freqs = cos_freqs[..., :half] + sin_freqs = sin_freqs[..., :half] + split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) first_half_input = split_input[..., :1, :] second_half_input = split_input[..., 1:, :] @@ -70,7 +79,7 @@ def _apply_split_rotary_emb( output = rearrange(output, "... d r -> ... (d r)") if needs_reshape: - output = output.swapaxes(1, 2).reshape(b, t, -1) + output = output.reshape(b, t, -1) return output @@ -163,8 +172,18 @@ def _split_freqs_cis( sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) b, t = cos_freq.shape[0], cos_freq.shape[1] - cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1).swapaxes(1, 2) - sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1).swapaxes(1, 2) + # Token-major layout [B, T, H, D]: matches the layout the fused norm+RoPE + # kernel consumes after reshape(-1, H*D), so no permute/contiguous is needed + # in the helper before kernel launch. + cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) + # Block-duplicate per-head cos/sin from head_dim/2 to head_dim along the last + # dim so the fused norm+RoPE kernel (rotate-half / INTERLEAVE=false branch) + # can consume the layout directly. The eager _apply_split_rotary_emb slices + # back to head_dim/2 when needed; both halves are identical so the operation + # is bit-identical to the original SPLIT-format cos/sin. + cos_freq = torch.cat([cos_freq, cos_freq], dim=-1) + sin_freq = torch.cat([sin_freq, sin_freq], dim=-1) return cos_freq, sin_freq diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/text_cache.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/text_cache.py index 54b6231dc736..fedf3e7626a9 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/text_cache.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/text_cache.py @@ -20,15 +20,26 @@ class TextCache: """Pre-computed text-derived tensors that are constant across denoise steps. + The ``*_pe`` fields hold sharded-local positional embeddings in the form + the consumer wants: + + - ``fuse_qk_norm_rope=True`` (LTX-2 default): 2D ``[T_local, H*D]`` + contiguous, fed directly to the fused norm+rope kernel. + - ``fuse_qk_norm_rope=False``: 4D ``[B, T_local, H, D]`` sharded but + otherwise unchanged, for the naive ``apply_rotary_emb`` path. + + Form is decided at cache-build time (``LTXModel.prepare_text_cache``); no + per-step reshape, ``.contiguous()``, or shard slicing. + Attributes: video_context: Projected text embedding for video cross-attention. video_mask: Attention mask for video text cross-attention. - video_pe: RoPE (cos, sin) for video. + video_pe: Sharded-local RoPE (cos, sin) for video self-attn. + video_cross_pe: Sharded-local RoPE for video AV cross-attn (audio-video model only). audio_context: Projected text embedding for audio cross-attention. audio_mask: Attention mask for audio text cross-attention. - audio_pe: RoPE (cos, sin) for audio. - video_cross_pe: Cross-modal RoPE for video (audio-video model only). - audio_cross_pe: Cross-modal RoPE for audio (audio-video model only). + audio_pe: Sharded-local RoPE (cos, sin) for audio self-attn. + audio_cross_pe: Sharded-local RoPE for audio AV cross-attn (audio-video model only). video_kv: Per-layer pre-projected text K/V for video cross-attention. audio_kv: Per-layer pre-projected text K/V for audio cross-attention. """ 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 9fb1b2e80349..eb65911818bc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -103,6 +103,10 @@ def __init__( # 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 + # Map LTX RoPE type to the fused-kernel INTERLEAVE template parameter: + # INTERLEAVED → pair (2i, 2i+1) pattern → kernel INTERLEAVE=true + # SPLIT → rotate-half pattern → kernel INTERLEAVE=false + # (cos/sin are stored block-duplicated for SPLIT; see _split_freqs_cis.) super().__init__( hidden_size=query_dim, num_attention_heads=heads, @@ -112,6 +116,8 @@ def __init__( qk_norm_mode="full", eps=norm_eps, bias=True, + interleave=(rope_type == LTXRopeType.INTERLEAVED), + fuse_qk_norm_rope=True, config=config, layer_idx=layer_idx, ) @@ -204,17 +210,30 @@ def _init_qkv_proj(self): def project_kv( self, context: torch.Tensor, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Project and normalize K/V from context. - - Used by the project-before-gather pattern in AV cross-attention: - project K/V on sharded data, then all-gather the smaller projected - tensors instead of all-gathering the full context first. + """Project K/V from context, optionally apply RMSNorm + RoPE on K. + + Used by the project-before-gather pattern in AV cross-attention. + When *pe* is given, RoPE is applied on the LOCAL K shard (Ulysses) + before all-gather. RoPE is per-token element-wise so it commutes with + seq-dim concat — bit-identical to the post-gather rope while saving + the cos/sin all-gather collective and reducing K-rope compute by U×. + The forward() consumer should pass ``k_pe=None`` to signal that K is + already rotated. """ k = self.to_k(context) v = self.to_v(context) - if self.qk_norm: - k = self.norm_k(k) + + # All cross-attn K-norm paths (with or without RoPE) go through the + # split-fuse kernels. fallback only kicks in for unsupported head_dim. + if self.qk_norm and self.head_dim in (64, 128): + self.apply_split_norm_or_norm_rope(k, self.norm_k.weight, self.num_key_value_heads, pe) + else: + if self.qk_norm: + k = self.norm_k(k) + if pe is not None: + k = apply_rotary_emb(k, pe, self.rope_type) return k, v def forward( @@ -227,13 +246,82 @@ def forward( ) -> torch.Tensor: """Forward pass. - Args: - x: Query input [B, T, D]. - context: Key/value input [B, S, C]. None → self-attention. - pe: (cos, sin) RoPE embeddings for Q (and K when k_pe is None). - k_pe: Separate (cos, sin) RoPE embeddings for K (for AV cross-attn). - pre_projected_kv: Pre-projected (k, v) tuple from project_kv(). - When provided, skips K/V projection and K-norm (already done). + 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). + k_pe overrides pe for K (e.g. AV cross-attn) when provided. + """ + # 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 + # never fires in production. + 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) + + 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) + + 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. + 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 + ) + else: + # Uncached cross-attn (not exercised by LTX-2 in practice; kept for fuse-dispatch consistency). + 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, + ) + + out = self._attn_impl(q, k, v) + + 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) + + return self.to_out[0](out) + + def _forward_unfused( + 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, + ) -> torch.Tensor: + """Fallback path for unsupported configs (e.g. head_dim ∉ {64, 128} or + fuse_qk_norm_rope=False). + + LTX-2 prod uses fused (head_dim ∈ {64, 128} and fuse_qk_norm_rope=True + by default), so in practice this is never entered. Kept for safety in + case the class is reused for other models or fusion is explicitly off. + + 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``. """ if pre_projected_kv is not None: k, v = pre_projected_kv @@ -246,7 +334,11 @@ def forward( if pe is not None: q = apply_rotary_emb(q, pe, self.rope_type) - k = apply_rotary_emb(k, pe if k_pe is None else k_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) out = self._attn_impl(q, k, v) @@ -463,18 +555,6 @@ def _sp_all_gather(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor: """All-gather *x* along *dim* across sequence-parallel ranks.""" return self._sharder.gather(x, dim=dim) - def _sp_gather_pe(self, pe): - """All-gather RoPE (cos, sin) tuple along its sequence dim. - - Split RoPE is ``[B, H, S, D]`` (dim 2); interleaved RoPE is - ``[B, S, D]`` (dim 1). Inferred from ``cos.ndim``. - """ - if pe is None: - return None - cos, sin = pe - seq_dim = 2 if cos.ndim == 4 else 1 - return (self._sharder.gather(cos, dim=seq_dim), self._sharder.gather(sin, dim=seq_dim)) - # -- Forward ------------------------------------------------------------- def forward( @@ -603,22 +683,23 @@ def forward( ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v # Project-before-gather: K/V projections run on sharded data - # so they benefit from Ulysses scaling. Only the smaller - # projected tensors are all-gathered. - k_a2v, v_a2v = self.audio_to_video_attn.project_kv(ax_scaled) + # so they benefit from Ulysses scaling. RoPE is applied to K + # inside project_kv on the sharded shard (RoPE commutes with + # seq-dim concat), so the cos/sin all-gather is unneeded and K + # rope work is U× cheaper. + k_a2v, v_a2v = self.audio_to_video_attn.project_kv( + ax_scaled, pe=audio.cross_positional_embeddings + ) if self._audio_is_sharded: k_a2v = self._sp_all_gather(k_a2v) v_a2v = self._sp_all_gather(v_a2v) - k_pe_a2v = self._sp_gather_pe(audio.cross_positional_embeddings) - else: - k_pe_a2v = audio.cross_positional_embeddings a2v_out = ( self.audio_to_video_attn( vx_scaled, pre_projected_kv=(k_a2v, v_a2v), pe=video.cross_positional_embeddings, - k_pe=k_pe_a2v, + k_pe=None, # K already rotated in project_kv ) * gate_out_a2v ) @@ -634,21 +715,21 @@ def forward( ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a - # Project-before-gather (video → audio direction). - k_v2a, v_v2a = self.video_to_audio_attn.project_kv(vx_scaled) + # Project-before-gather (video → audio direction). RoPE applied + # to K in project_kv on local shard; see audio→video branch above. + k_v2a, v_v2a = self.video_to_audio_attn.project_kv( + vx_scaled, pe=video.cross_positional_embeddings + ) if self._sharder.is_active: k_v2a = self._sp_all_gather(k_v2a) v_v2a = self._sp_all_gather(v_v2a) - k_pe_v2a = self._sp_gather_pe(video.cross_positional_embeddings) - else: - k_pe_v2a = video.cross_positional_embeddings v2a_out = ( self.video_to_audio_attn( ax_scaled, pre_projected_kv=(k_v2a, v_v2a), pe=audio.cross_positional_embeddings, - k_pe=k_pe_v2a, + k_pe=None, # K already rotated in project_kv ) * gate_out_v2a ) @@ -1142,41 +1223,64 @@ def _init_transformer_blocks( # -- Sequence sharding / gathering ---------------------------------------- def _shard_transformer_args(self, args: TransformerArgs) -> TransformerArgs: - """Shard sequence-dependent fields of *args* across sequence-parallel ranks. + """Shard step-dependent fields of *args* across sequence-parallel ranks. - Fields whose dim-1 doesn't match ``args.x``'s sequence length are passed - through unchanged (broadcast-compatible scalars, etc.). + PE (``positional_embeddings`` / ``cross_positional_embeddings``) is + already sharded-local in ``TextCache`` (one-time in + ``prepare_text_cache``) so we leave it untouched. Only step-varying + fields (``x``, timesteps, etc.) need slicing each step. """ seq_len = args.x.shape[1] sh = self._sharder - pe_seq_dim = ( - 2 - if args.positional_embeddings is not None and args.positional_embeddings[0].ndim == 4 - else 1 - ) - cross_pe_seq_dim = ( - 2 - if args.cross_positional_embeddings is not None - and args.cross_positional_embeddings[0].ndim == 4 - else 1 - ) return replace( args, x=sh.shard(args.x, dim=1), timesteps=sh.shard(args.timesteps, dim=1, expected_seq_len=seq_len), embedded_timestep=sh.shard(args.embedded_timestep, dim=1, expected_seq_len=seq_len), - positional_embeddings=sh.shard_rope( - args.positional_embeddings, seq_len=seq_len, seq_dim=pe_seq_dim - ), - cross_positional_embeddings=sh.shard_rope( - args.cross_positional_embeddings, seq_len=seq_len, seq_dim=cross_pe_seq_dim - ), cross_scale_shift_timestep=sh.shard( args.cross_scale_shift_timestep, dim=1, expected_seq_len=seq_len ), cross_gate_timestep=sh.shard(args.cross_gate_timestep, dim=1, expected_seq_len=seq_len), ) + def _make_pe_local( + self, + pe: tuple[torch.Tensor, torch.Tensor] | None, + *, + is_audio: bool, + fuse: bool, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Sharded-local PE for the attention consumer. + + Slices the source 4D PE along seq dim by Ulysses rank (one-time, in + ``prepare_text_cache``), then either reshapes to 2D ``[T_local, H*D]`` + for the fused kernel or keeps 4D for the eager apply_rotary_emb path. + LTX-2 SPLIT rope produces 4D PE; INTERLEAVED is not used in prod. + + ``_audio_is_sharded`` (set in ``configure_audio_ulysses``) already + encodes whether audio_seq_len is divisible by ulysses_size, so we + gate sharding on that flag alone — no second divisibility check. + """ + if pe is None: + return None + cos, sin = pe + sh = self._sharder + if sh.is_active and (not is_audio or self._audio_is_sharded): + chunk = cos.shape[1] // sh.size + s = sh.rank * chunk + e = s + chunk + cos = cos[:, s:e] + sin = sin[:, s:e] + cos = cos.contiguous() + sin = sin.contiguous() + if fuse: + # [B, T_local, H, D] -> [B*T_local, H*D]. PE source from + # precompute_freqs_cis has B=1 so this collapses to [T_local, H*D]; + # the fused kernel broadcasts cos over B internally. + cos = cos.reshape(cos.shape[0] * cos.shape[1], -1) + sin = sin.reshape(sin.shape[0] * sin.shape[1], -1) + return (cos, sin) + def _gather_sequence(self, x: torch.Tensor) -> torch.Tensor: """All-gather hidden states along the sequence dim.""" return self._sharder.gather(x, dim=1) @@ -1290,6 +1394,24 @@ def prepare_text_cache( ) a_kv = [block.audio_attn2.project_kv(a_ctx) for block in self.transformer_blocks] + # Build sharded-local PE in the form the attention consumer expects. + # fuse_qk_norm_rope=True (LTX-2 default) -> 2D [T_local, H*D] contiguous, + # ready for the fused kernel; False -> 4D [B, T_local, H, D] for the + # naive apply_rotary_emb path. Done one-time here, so the inner loop + # has no reshape/contiguous/shard work on PE. + # Inspect any LTX2Attention to learn whether fusion is on (per-modality + # attentions are constructed with the same flag in this codepath). + fuse_video = self.transformer_blocks[0].attn1.fuse_qk_norm_rope + fuse_audio = ( + self.transformer_blocks[0].audio_attn1.fuse_qk_norm_rope + if hasattr(self.transformer_blocks[0], "audio_attn1") + else True + ) + v_pe = self._make_pe_local(v_pe, is_audio=False, fuse=fuse_video) + v_cross_pe = self._make_pe_local(v_cross_pe, is_audio=False, fuse=fuse_video) + a_pe = self._make_pe_local(a_pe, is_audio=True, fuse=fuse_audio) + a_cross_pe = self._make_pe_local(a_cross_pe, is_audio=True, fuse=fuse_audio) + return TextCache( video_context=v_ctx, video_mask=v_mask, 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 09a5ac02579e..b7af3f6d8464 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -282,8 +282,10 @@ def __init__( ) # Self-attention with fused QKV. - # fuse_qk_norm_rope=True: use fused cross-head QK Norm + RoPE CUDA kernel - # to eliminate extra global memory round-trip between separate norm and RoPE. + # WAN-14B is 40 heads × 128, which exceeds the default fused op's + # num_heads<=32 / head_dim in {64,128} envelope, so route to + # PR #13052's fused_dit_cross_head_qk_norm_rope op (validated for + # WAN-1.3B and WAN-14B sizes). self.attn1 = Attention( hidden_size=hidden_size, num_attention_heads=num_heads, @@ -292,6 +294,7 @@ def __init__( qk_norm=True, eps=eps, fuse_qk_norm_rope=True, + qk_norm_rope_kernel="fused_dit_cross_head_qk_norm_rope", config=model_config, layer_idx=_layer_idx, ) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index abb6f3254155..dc01f8926403 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -47,6 +47,7 @@ def __init__( bias: bool = True, interleave: bool = True, fuse_qk_norm_rope: Optional[bool] = None, + qk_norm_rope_kernel: str = "fused_dit_qk_norm_rope", config: Optional[DiffusionModelConfig] = None, layer_idx: Optional[int] = None, ): @@ -70,6 +71,23 @@ def __init__( # Supported for both per_head (FLUX) and full/cross-head (WAN) norm modes. # Defaults to False; models that want the fused kernel must pass True explicitly. self.fuse_qk_norm_rope = fuse_qk_norm_rope if fuse_qk_norm_rope is not None else False + # Which trtllm op backs the fused path. Value is the op name itself + # (1:1 with torch.ops.trtllm registrations): + # "fused_dit_qk_norm_rope" (default) + # — Per-head autodispatch for FLUX/Cosmos and full-dim path for + # LTX-2. Bounded to num_heads<=32 and head_dim in {64,128}. + # "fused_dit_cross_head_qk_norm_rope" + # — PR #13052 cross-head kernel for WAN sizes outside the default + # op's range (head_dim=256, or num_heads>32). fp32 head-broadcast + # cos only. + assert qk_norm_rope_kernel in ( + "fused_dit_qk_norm_rope", + "fused_dit_cross_head_qk_norm_rope", + ), ( + f"qk_norm_rope_kernel must be 'fused_dit_qk_norm_rope' or " + f"'fused_dit_cross_head_qk_norm_rope', got {qk_norm_rope_kernel}" + ) + self.qk_norm_rope_kernel = qk_norm_rope_kernel self.interleave = interleave # Select compute backend (orthogonal to parallelism) @@ -250,7 +268,7 @@ def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, k = self.norm_k(k) return q, k - def apply_qk_norm_rope( + def apply_packed_qk_norm_rope( self, qkv: torch.Tensor, freqs_cos: torch.Tensor, @@ -259,23 +277,28 @@ def apply_qk_norm_rope( q_add_weight: Optional[torch.Tensor] = None, k_add_weight: Optional[torch.Tensor] = None, ) -> None: - """Apply fused QK Norm + RoPE in-place on packed QKV tensor. + """Apply fused QK Norm + RoPE in-place on packed QKV tensor (FUSE_QKV self-attn). - Dispatches to per-head kernel (FLUX) or cross-head kernel (WAN) - based on qk_norm_mode. + cos/sin can be either shape (per-token total elements): + - [..., head_dim] : shared across heads (FLUX/Cosmos style) + - [..., num_heads*head_dim] : per-head freqs (LTX-2 3D RoPE style) + Op auto-detects via cos_emb.size(1) and dispatches the kernel template. """ - cos_2d = freqs_cos.reshape(-1, self.head_dim).float().contiguous() - sin_2d = freqs_sin.reshape(-1, self.head_dim).float().contiguous() - B, S, D = qkv.shape - assert cos_2d.shape == (S, self.head_dim), ( - f"cos_emb shape mismatch: expected [{S}, {self.head_dim}], got {list(cos_2d.shape)}" - ) - qkv_2d = qkv.view(B * S, D) - cos_tiled = cos_2d.repeat(B, 1) if B > 1 else cos_2d - sin_tiled = sin_2d.repeat(B, 1) if B > 1 else sin_2d - if self.qk_norm_mode == "full": + if self.qk_norm_rope_kernel == "fused_dit_cross_head_qk_norm_rope": + # PR #13052 cross-head op path (WAN sizes that exceed the default + # op's num_heads<=32 / head_dim in {64,128} envelope). Op requires + # head-broadcast fp32 cos: [num_tokens, head_dim]. + cos_2d = freqs_cos.reshape(-1, self.head_dim).float().contiguous() + sin_2d = freqs_sin.reshape(-1, self.head_dim).float().contiguous() + if cos_2d.shape[0] == S and B > 1: + cos_tiled = cos_2d.repeat(B, 1) + sin_tiled = sin_2d.repeat(B, 1) + else: + cos_tiled = cos_2d + sin_tiled = sin_2d + qkv_2d = qkv.view(B * S, D) torch.ops.trtllm.fused_dit_cross_head_qk_norm_rope( qkv_2d, self.num_attention_heads, @@ -289,30 +312,124 @@ def apply_qk_norm_rope( sin_tiled, self.interleave, ) + return + + # cos last-dim is fixed by qk_norm_mode: + # "full" → num_heads * head_dim (LTX-2 / WAN per-head cos) + # "per_head" → head_dim (FLUX / Cosmos shared-across-heads cos) + cos_last = self.q_dim if self.qk_norm_mode == "full" else self.head_dim + # cos/sin are token-major [B, T, H, D] (or shared per-token [B, T, D]); + # reshape(-1, cos_last) yields the [B*T, cos_last] layout the kernel reads. + # Full-dim LTX-2 / WAN path accepts bf16 cos (kernel upcasts in registers, lossless); + # FLUX per-head path requires fp32 (and cos is already fp32 upstream there). + cos_2d = freqs_cos.reshape(-1, cos_last).contiguous() + sin_2d = freqs_sin.reshape(-1, cos_last).contiguous() + # LTX-2 / WAN full-dim path: kernel broadcasts cos over B internally + # (cos_tokenIdx = tokenIdx % cos_seq_per_batch in the fused kernel), so we + # pass cos as-is regardless of B. FLUX / Cosmos per-head path: kernel does + # not support broadcast, so host still has to tile when B > 1. + if self.qk_norm_mode == "full" or cos_2d.shape[0] != S or B == 1: + cos_tiled = cos_2d + sin_tiled = sin_2d else: - # Dual-stream batch correction: when B>1 and dual-stream is active, - # the kernel uses modulo (tokenIdx % tokens_per_batch) to find the - # local position within each batch element for the text/image boundary. - # 0 = no dual-stream (single-stream or batch=1). - tokens_per_batch = S if num_txt_tokens > 0 else 0 + cos_tiled = cos_2d.repeat(B, 1) + sin_tiled = sin_2d.repeat(B, 1) + qkv_2d = qkv.view(B * S, D) - torch.ops.trtllm.fused_dit_qk_norm_rope( - qkv_2d, - self.num_attention_heads, - self.num_key_value_heads, - self.num_key_value_heads, - self.head_dim, - self.eps, - self.norm_q.weight, - self.norm_k.weight, - q_add_weight, - k_add_weight, - cos_tiled, - sin_tiled, - num_txt_tokens, - self.interleave, - tokens_per_batch, - ) + # Dual-stream batch correction: when B>1 and dual-stream is active, + # the kernel uses modulo (tokenIdx % tokens_per_batch) to find the + # local position within each batch element for the text/image boundary. + # 0 = no dual-stream (single-stream or batch=1). + tokens_per_batch = S if num_txt_tokens > 0 else 0 + + torch.ops.trtllm.fused_dit_qk_norm_rope( + qkv_2d, + self.num_attention_heads, + self.num_key_value_heads, + self.num_key_value_heads, + self.head_dim, + self.eps, + self.norm_q.weight, + self.norm_k.weight, + q_add_weight, + k_add_weight, + cos_tiled, + sin_tiled, + num_txt_tokens, + self.interleave, + tokens_per_batch, + ) + + def apply_split_norm_rope( + self, + tensor: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> None: + """In-place fused RMSNorm + RoPE on a single Q or K tensor [B, T, H*D] (SEPARATE_QKV cross-attn). + + Calls trtllm.fused_dit_split_norm_rope. Full-dim per-head cos in + [B, T, H, D] reshape(-1, H*D) -> [B*T, H*D] is what the kernel reads; + the kernel broadcasts cos over B internally + (cos_tokenIdx = tokenIdx % cos_seq_per_batch), so we pass cos as-is. + bf16 cos is also accepted (kernel upcasts to fp32 in registers). + """ + B, T, _ = tensor.shape + cos_last = num_heads * self.head_dim + cos_2d = cos.reshape(-1, cos_last).contiguous() + sin_2d = sin.reshape(-1, cos_last).contiguous() + tensor_2d = tensor.view(B * T, -1) + torch.ops.trtllm.fused_dit_split_norm_rope( + tensor_2d, + num_heads, + self.head_dim, + self.eps, + weight, + cos_2d, + sin_2d, + self.interleave, + ) + + def apply_split_norm( + self, + tensor: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + ) -> None: + """In-place fused full-dim RMSNorm only (no RoPE) on a single Q or K tensor [B, T, H*D]. + + Calls trtllm.fused_dit_split_norm. Used by paths that need norm but + no RoPE -- e.g. LTX-2 text cross-attn (Q-norm with pe=None). + """ + B, T, _ = tensor.shape + tensor_2d = tensor.view(B * T, -1) + torch.ops.trtllm.fused_dit_split_norm( + tensor_2d, + num_heads, + self.head_dim, + self.eps, + weight, + ) + + def apply_split_norm_or_norm_rope( + self, + tensor: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + pe: tuple[torch.Tensor, torch.Tensor] | None, + ) -> None: + """Dispatcher: in-place norm-only when pe is None, else norm + RoPE. + + Used to dispatch all SEPARATE_QKV cross-attn norm paths through the + split-fuse kernels regardless of whether the path needs RoPE. + """ + if pe is None: + self.apply_split_norm(tensor, weight, num_heads) + else: + cos, sin = pe + self.apply_split_norm_rope(tensor, weight, num_heads, cos, sin) def _attn_impl( self, @@ -389,7 +506,7 @@ def forward( ): qkv = self.qkv_proj(hidden_states) freqs_cos, freqs_sin = freqs - self.apply_qk_norm_rope(qkv, freqs_cos, freqs_sin) + self.apply_packed_qk_norm_rope(qkv, freqs_cos, freqs_sin) q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) out = self._attn_impl(q, k, v) return self.to_out[0](out) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py index 2310c3d68b9a..b5da528c3581 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py @@ -732,3 +732,245 @@ def test_cross_head_batched(batch_size): True, ) torch.testing.assert_close(qkv, ref, rtol=5e-2, atol=1e-1) + + +# ============================================================================ +# Full-dim norm tests (LTX-2 / WAN packed FUSE_QKV) +# weight is [num_heads*head_dim] (full-dim per-token norm), no dual-stream, +# cos/sin can be [num_tokens, num_heads*head_dim] (per-head) or +# broadcast [num_tokens/B, num_heads*head_dim] (kernel modulo over B). +# ============================================================================ + + +@torch.inference_mode() +def torch_ref_full_dim( + qkv, + num_heads, + head_dim, + eps, + q_weight, + k_weight, + cos_emb, + sin_emb, + interleave, +): + """Reference: full-dim RMSNorm (weight shape [num_heads*head_dim]) on Q and K + of a packed qkv tensor [N, 3*num_heads*head_dim], then rope; V untouched. + + cos/sin shape: [num_tokens, num_heads*head_dim] (per-head, full-dim layout). + """ + num_tokens = qkv.shape[0] + hidden = num_heads * head_dim + q = qkv[:, :hidden].float() + k = qkv[:, hidden : 2 * hidden].float() + v = qkv[:, 2 * hidden :] + + # full-dim norm: stat over all num_heads*head_dim elements per token. + # Keep everything in fp32 through norm + rope to match kernel internal + # precision (single bf16 cast at the end). + var_q = q.pow(2).mean(-1, keepdim=True) + q = q * torch.rsqrt(var_q + eps) * q_weight.float() + var_k = k.pow(2).mean(-1, keepdim=True) + k = k * torch.rsqrt(var_k + eps) * k_weight.float() + + cos_3d = cos_emb.float().view(num_tokens, num_heads, head_dim) + sin_3d = sin_emb.float().view(num_tokens, num_heads, head_dim) + q_4d = q.view(num_tokens, num_heads, head_dim) + k_4d = k.view(num_tokens, num_heads, head_dim) + if interleave: + # pair (2i, 2i+1) — INTERLEAVED rope + def _rope_interleaved(x_4d): + rot = torch.empty_like(x_4d) + rot[..., 0::2] = -x_4d[..., 1::2] + rot[..., 1::2] = x_4d[..., 0::2] + return x_4d * cos_3d + rot * sin_3d + + q_4d = _rope_interleaved(q_4d) + k_4d = _rope_interleaved(k_4d) + else: + # rotate-half: pair (i, i+D/2) within head — LTX-2 SPLIT + half = head_dim // 2 + + def _rope_rotate_half(x_4d): + x1 = x_4d[..., :half] + x2 = x_4d[..., half:] + rot = torch.cat([-x2, x1], dim=-1) + return x_4d * cos_3d + rot * sin_3d + + q_4d = _rope_rotate_half(q_4d) + k_4d = _rope_rotate_half(k_4d) + + q_out = q_4d.reshape(num_tokens, -1).to(qkv.dtype) + k_out = k_4d.reshape(num_tokens, -1).to(qkv.dtype) + return torch.cat([q_out, k_out, v], dim=1) + + +def _make_per_head_cos_full_dim(num_tokens, num_heads, head_dim, device, dtype=torch.float32): + """Per-head cos: [num_tokens, num_heads*head_dim] with rotate-half block-duplicate pattern. + + cos/sin generated as (num_heads, num_tokens, head_dim/2) and block-duplicated + along the last dim (cos = cat([cos_half, cos_half], -1)) so the kernel's + rotate-half read pattern works on the flat 2D layout. + """ + half = head_dim // 2 + freqs = torch.randn(num_tokens, num_heads, half, device=device, dtype=torch.float32) + cos_h = freqs.cos() + sin_h = freqs.sin() + cos = torch.cat([cos_h, cos_h], dim=-1) # [N, H, D] + sin = torch.cat([sin_h, sin_h], dim=-1) + cos_2d = cos.reshape(num_tokens, num_heads * head_dim).contiguous().to(dtype) + sin_2d = sin.reshape(num_tokens, num_heads * head_dim).contiguous().to(dtype) + return cos_2d, sin_2d + + +@pytest.mark.parametrize("cos_dtype", [torch.float32, torch.bfloat16], ids=["fp32cos", "bf16cos"]) +@pytest.mark.parametrize( + "label,B,T,num_heads,head_dim", + [ + # LTX-2 video self-attn (FUSE_QKV packed, full-dim norm, rotate-half) + ("ltx2_video_self_attn", 2, 12288, 32, 128), + # LTX-2 audio self-attn + ("ltx2_audio_self_attn", 2, 504, 32, 64), + # B=1 sanity (no broadcast) + ("ltx2_video_b1", 1, 12288, 32, 128), + ], +) +def test_full_dim_norm_packed_rotate_half(label, B, T, num_heads, head_dim, cos_dtype): + """LTX-2 packed FUSE_QKV full-dim norm + rotate-half rope (no broadcast).""" + device = "cuda" + torch.random.manual_seed(0) + + hidden = num_heads * head_dim + num_tokens = B * T + qkv = torch.randn(num_tokens, 3 * hidden, dtype=torch.bfloat16, device=device) * 0.5 + qkv_copy = qkv.clone() + q_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + k_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + # Tile cos/sin to num_tokens (no kernel-side broadcast in this test) + cos_2d_T, sin_2d_T = _make_per_head_cos_full_dim( + T, num_heads, head_dim, device, dtype=cos_dtype + ) + if B > 1: + cos_2d = cos_2d_T.repeat(B, 1) + sin_2d = sin_2d_T.repeat(B, 1) + else: + cos_2d, sin_2d = cos_2d_T, sin_2d_T + + _call_fused_kernel( + qkv, + num_heads, + num_heads, + num_heads, + head_dim, + 1e-6, + q_weight, + k_weight, + None, + None, + cos_2d, + sin_2d, + -1, + interleave=False, + ) + ref = torch_ref_full_dim( + qkv_copy, + num_heads, + head_dim, + 1e-6, + q_weight, + k_weight, + cos_2d, + sin_2d, + interleave=False, + ) + torch.testing.assert_close(qkv, ref, rtol=2e-2, atol=5e-3) + + +@pytest.mark.parametrize("cos_dtype", [torch.float32, torch.bfloat16], ids=["fp32cos", "bf16cos"]) +@pytest.mark.parametrize( + "label,B,T,num_heads,head_dim", + [ + # C7 broadcast: cos has T rows, qkv has B*T tokens; kernel does tokenIdx % T + ("ltx2_video_bcast", 2, 12288, 32, 128), + ("ltx2_audio_bcast", 2, 504, 32, 64), + ], +) +def test_full_dim_norm_packed_rotate_half_broadcast(label, B, T, num_heads, head_dim, cos_dtype): + """C7: kernel-side cos broadcast over B (cos.size(0) == T, not B*T).""" + device = "cuda" + torch.random.manual_seed(0) + + hidden = num_heads * head_dim + num_tokens = B * T + qkv = torch.randn(num_tokens, 3 * hidden, dtype=torch.bfloat16, device=device) * 0.5 + qkv_copy = qkv.clone() + q_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + k_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos_2d_T, sin_2d_T = _make_per_head_cos_full_dim( + T, num_heads, head_dim, device, dtype=cos_dtype + ) + + # Reference: tile cos to B*T tokens (kernel does this in-place via modulo) + cos_2d_full = cos_2d_T.repeat(B, 1) + sin_2d_full = sin_2d_T.repeat(B, 1) + + # Call kernel with the unbroadcast T-row cos (kernel detects via cos.size(0) != num_tokens) + _call_fused_kernel( + qkv, + num_heads, + num_heads, + num_heads, + head_dim, + 1e-6, + q_weight, + k_weight, + None, + None, + cos_2d_T, + sin_2d_T, + -1, + interleave=False, + ) + ref = torch_ref_full_dim( + qkv_copy, + num_heads, + head_dim, + 1e-6, + q_weight, + k_weight, + cos_2d_full, + sin_2d_full, + interleave=False, + ) + torch.testing.assert_close(qkv, ref, rtol=2e-2, atol=5e-3) + + +def test_full_dim_norm_packed_v_unchanged(): + """Full-dim path leaves V slice untouched (sanity).""" + device = "cuda" + torch.random.manual_seed(0) + num_heads, head_dim, num_tokens = 32, 128, 256 + hidden = num_heads * head_dim + qkv = torch.randn(num_tokens, 3 * hidden, dtype=torch.bfloat16, device=device) * 0.5 + v_original = qkv[:, 2 * hidden :].clone() + q_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + k_weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos_2d, sin_2d = _make_per_head_cos_full_dim(num_tokens, num_heads, head_dim, device) + + _call_fused_kernel( + qkv, + num_heads, + num_heads, + num_heads, + head_dim, + 1e-6, + q_weight, + k_weight, + None, + None, + cos_2d, + sin_2d, + -1, + interleave=False, + ) + torch.testing.assert_close(qkv[:, 2 * hidden :], v_original, rtol=0, atol=0) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_norm.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_norm.py new file mode 100644 index 000000000000..0be790485a49 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_norm.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Tests for the LTX-2 split full-dim RMSNorm-only kernel (no RoPE). +# Mirror of test_fused_dit_split_qk_norm_rope.py but for the norm-only path. + +import pytest +import torch + +import tensorrt_llm # noqa: F401 -- triggers libth_common.so load (registers trtllm ops) + +# ============================================================================ +# Reference implementation (PyTorch fp32) +# ============================================================================ + + +@torch.inference_mode() +def torch_ref(x_2d, weight, eps): + """Reference: full-dim RMSNorm only on a single Q-or-K tensor. + + x_2d: [T, num_heads * head_dim], bf16 + weight: [num_heads * head_dim] bf16 (full-dim) + """ + out = x_2d.float() + var = out.pow(2).mean(-1, keepdim=True) + out = out * torch.rsqrt(var + eps) * weight.float() + return out.to(x_2d.dtype) + + +# ============================================================================ +# Helper +# ============================================================================ + + +def _call_norm_op(tensor, weight, num_heads, head_dim, eps): + torch.ops.trtllm.fused_dit_split_norm(tensor, num_heads, head_dim, eps, weight) + + +# ============================================================================ +# Full-dim norm tests +# ============================================================================ + + +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("num_heads", [1, 8, 32]) +@pytest.mark.parametrize("num_tokens", [1, 64, 1024]) +def test_full_dim_norm_only(head_dim, num_heads, num_tokens): + """Full-dim RMSNorm only (no RoPE) on contiguous 2D tensor.""" + device = "cuda" + torch.random.manual_seed(42) + + hidden = num_heads * head_dim + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + eps = 1e-6 + + _call_norm_op(x, weight, num_heads, head_dim, eps) + ref = torch_ref(x_copy, weight, eps) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_ltx2_video_shape(): + """LTX-2 video self/cross-attn shape: H=32, D=128, T=12288.""" + device = "cuda" + torch.random.manual_seed(0) + + num_tokens = 12288 + num_heads = 32 + head_dim = 128 + hidden = num_heads * head_dim + + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + + _call_norm_op(x, weight, num_heads, head_dim, 1e-6) + ref = torch_ref(x_copy, weight, 1e-6) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_ltx2_audio_shape(): + """LTX-2 audio self/cross-attn shape: H=32, D=64, T=504.""" + device = "cuda" + torch.random.manual_seed(0) + + num_tokens = 504 + num_heads = 32 + head_dim = 64 + hidden = num_heads * head_dim + + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + + _call_norm_op(x, weight, num_heads, head_dim, 1e-6) + ref = torch_ref(x_copy, weight, 1e-6) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_rejects_non_contiguous(): + """fused_dit_split_norm requires contiguous SEPARATE_QKV input.""" + device = "cuda" + num_heads, head_dim, num_tokens = 32, 128, 64 + qkv = torch.randn(num_tokens, 3 * num_heads * head_dim, dtype=torch.bfloat16, device=device) + q_view = qkv[:, : num_heads * head_dim] + assert not q_view.is_contiguous() + + weight = torch.randn(num_heads * head_dim, dtype=torch.bfloat16, device=device) * 5.0 + + with pytest.raises(RuntimeError, match=r"contiguous"): + _call_norm_op(q_view, weight, num_heads, head_dim, 1e-6) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_qk_norm_rope.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_qk_norm_rope.py new file mode 100644 index 000000000000..f84c9bf9833b --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_split_qk_norm_rope.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Tests for the LTX-2 split Q/K fused full-dim RMSNorm + RoPE kernel. + +import pytest +import torch + +import tensorrt_llm # noqa: F401 — triggers libth_common.so load (registers trtllm ops) + +# ============================================================================ +# Reference implementation (PyTorch fp32) +# ============================================================================ + + +@torch.inference_mode() +def torch_ref(x_2d, weight, cos, sin, num_heads, head_dim, eps, interleave): + """Reference: full-dim RMSNorm then RoPE on a single Q-or-K tensor. + + x_2d: [T, num_heads * head_dim], bf16 + weight: [num_heads * head_dim] bf16 (full-dim) + cos, sin: [T, head_dim], float32 + """ + T = x_2d.shape[0] + # Match kernel: do reduce/scale/RoPE in fp32, cast to bf16 only at the end. + out = x_2d.float() + var = out.pow(2).mean(-1, keepdim=True) + out = out * torch.rsqrt(var + eps) * weight.float() + + # RoPE: cos/sin are [T, head_dim], broadcast over all heads. + out_4d = out.view(T, num_heads, head_dim) + cos_3d = cos.unsqueeze(1) + sin_3d = sin.unsqueeze(1) + if interleave: + # pair (2i, 2i+1) — LTX-2 INTERLEAVED + rot = torch.empty_like(out_4d) + rot[..., 0::2] = -out_4d[..., 1::2] + rot[..., 1::2] = out_4d[..., 0::2] + out_4d = out_4d * cos_3d + rot * sin_3d + else: + # rotate_half: pair (i, i+D/2) + half = head_dim // 2 + x1 = out_4d[..., :half] + x2 = out_4d[..., half:] + rot = torch.cat([-x2, x1], dim=-1) + out_4d = out_4d * cos_3d + rot * sin_3d + + return out_4d.reshape(T, -1).to(x_2d.dtype) + + +# ============================================================================ +# Helper +# ============================================================================ + + +def _call_split_op(tensor, weight, cos, sin, num_heads, head_dim, eps, interleave=True): + torch.ops.trtllm.fused_dit_split_norm_rope( + tensor, num_heads, head_dim, eps, weight, cos, sin, interleave + ) + + +def _generate_cos_sin(num_tokens, head_dim, device, dtype=torch.float32): + """Generate paired cos/sin (matches original DiT test format: freqs.cos/sin in [-1,1]). + + dtype: cos/sin output dtype — fp32 (default) or bf16 (kernel upcasts bf16 to fp32 in registers). + """ + half_dim = head_dim // 2 + freqs = torch.randn(num_tokens, half_dim, device=device, dtype=torch.float32) + freqs = freqs.repeat_interleave(2, dim=-1) + return freqs.cos().to(dtype), freqs.sin().to(dtype) + + +# ============================================================================ +# Full-dim norm tests (the only mode the kernel supports) +# ============================================================================ + + +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("num_heads", [1, 8, 32]) +@pytest.mark.parametrize("num_tokens", [1, 64, 1024]) +@pytest.mark.parametrize("interleave", [True, False]) +@pytest.mark.parametrize("cos_dtype", [torch.float32, torch.bfloat16], ids=["fp32cos", "bf16cos"]) +def test_full_dim_norm_self_attn(head_dim, num_heads, num_tokens, interleave, cos_dtype): + """Full-dim RMSNorm + RoPE on contiguous 2D tensor (LTX-2 self-attn shape).""" + device = "cuda" + torch.random.manual_seed(42) + + hidden = num_heads * head_dim + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos, sin = _generate_cos_sin(num_tokens, head_dim, device, dtype=cos_dtype) + eps = 1e-6 + + _call_split_op(x, weight, cos, sin, num_heads, head_dim, eps, interleave=interleave) + ref = torch_ref(x_copy, weight, cos, sin, num_heads, head_dim, eps, interleave=interleave) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_ltx2_video_shape(): + """LTX-2 video self-attn shape: H=32, D=128, S=12288.""" + device = "cuda" + torch.random.manual_seed(0) + + num_tokens = 12288 + num_heads = 32 + head_dim = 128 + hidden = num_heads * head_dim + + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos, sin = _generate_cos_sin(num_tokens, head_dim, device) + + _call_split_op(x, weight, cos, sin, num_heads, head_dim, 1e-6, interleave=True) + ref = torch_ref(x_copy, weight, cos, sin, num_heads, head_dim, 1e-6, interleave=True) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_ltx2_audio_shape(): + """LTX-2 audio self-attn shape: H=32, D=64, S=504.""" + device = "cuda" + torch.random.manual_seed(0) + + num_tokens = 504 + num_heads = 32 + head_dim = 64 + hidden = num_heads * head_dim + + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=device) + x_copy = x.clone() + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos, sin = _generate_cos_sin(num_tokens, head_dim, device) + + _call_split_op(x, weight, cos, sin, num_heads, head_dim, 1e-6, interleave=True) + ref = torch_ref(x_copy, weight, cos, sin, num_heads, head_dim, 1e-6, interleave=True) + torch.testing.assert_close(x, ref, rtol=1e-2, atol=5e-3) + + +def test_full_dim_norm_rejects_non_contiguous(): + """Split kernel only supports contiguous SEPARATE_QKV input. + For FUSE_QKV packed-view, callers must use fused_dit_qk_norm_rope instead.""" + device = "cuda" + num_heads, head_dim, num_tokens = 32, 128, 64 + qkv = torch.randn(num_tokens, 3 * num_heads * head_dim, dtype=torch.bfloat16, device=device) + q_view = qkv[:, : num_heads * head_dim] + assert not q_view.is_contiguous() + + weight = torch.randn(num_heads * head_dim, dtype=torch.bfloat16, device=device) * 5.0 + cos, sin = _generate_cos_sin(num_tokens, head_dim, device) + + with pytest.raises(RuntimeError, match=r"contiguous"): + _call_split_op(q_view, weight, cos, sin, num_heads, head_dim, 1e-6, interleave=True) + + +# ============================================================================ +# LTX-2 production scenarios: rotate-half (SPLIT) + per-head cos +# Tolerance rtol=2e-2 atol=5e-3 (consistent with existing fuse-kernel checks). +# ============================================================================ + + +def _torch_ref_per_head_cos(x_2d, weight, cos_2d, sin_2d, num_heads, head_dim, eps, interleave): + """Reference for per-head cos: cos_2d shape [T, num_heads*head_dim].""" + T = x_2d.shape[0] + out = x_2d.float() + var = out.pow(2).mean(-1, keepdim=True) + out = out * torch.rsqrt(var + eps) * weight.float() + out_4d = out.view(T, num_heads, head_dim) + cos_3d = cos_2d.float().view(T, num_heads, head_dim) + sin_3d = sin_2d.float().view(T, num_heads, head_dim) + if interleave: + rot = torch.empty_like(out_4d) + rot[..., 0::2] = -out_4d[..., 1::2] + rot[..., 1::2] = out_4d[..., 0::2] + out_4d = out_4d * cos_3d + rot * sin_3d + else: + half = head_dim // 2 + x1 = out_4d[..., :half] + x2 = out_4d[..., half:] + rot = torch.cat([-x2, x1], dim=-1) + out_4d = out_4d * cos_3d + rot * sin_3d + return out_4d.reshape(T, -1).to(x_2d.dtype) + + +def _make_per_head_cos(B, T, num_heads, head_dim, device, dtype=torch.float32): + """LTX-2-style per-head cos: per (B, head, T) freqs, block-duplicated to head_dim. + Returns cos_2d, sin_2d of shape (B*T, num_heads*head_dim).""" + half = head_dim // 2 + freqs = torch.randn(B, num_heads, T, half, dtype=torch.float32, device=device) + cos_h = freqs.cos() + sin_h = freqs.sin() + cos = torch.cat([cos_h, cos_h], dim=-1) # (B, H, T, D) + sin = torch.cat([sin_h, sin_h], dim=-1) + cos_2d = ( + cos.permute(0, 2, 1, 3) + .contiguous() + .reshape(B * T, num_heads * head_dim) + .contiguous() + .to(dtype) + ) + sin_2d = ( + sin.permute(0, 2, 1, 3) + .contiguous() + .reshape(B * T, num_heads * head_dim) + .contiguous() + .to(dtype) + ) + return cos_2d, sin_2d + + +@pytest.mark.parametrize("cos_dtype", [torch.float32, torch.bfloat16], ids=["fp32cos", "bf16cos"]) +@pytest.mark.parametrize( + "label,B,T,num_heads,head_dim", + [ + # video self-attn: 121 frames @ 768x1024 → 12288 tokens, 32 heads × 128 dim + ("ltx2_video_self_attn", 2, 12288, 32, 128), + # audio self-attn: 504 tokens, 32 heads × 64 dim + ("ltx2_audio_self_attn", 2, 504, 32, 64), + # text→video cross-attn (Q on video shape, K on text but using video cos for Q) + ("ltx2_text_cross_video", 2, 12288, 32, 128), + # AV a2v Q-only fused: video Q with video cos + ("ltx2_av_a2v_q", 2, 12288, 32, 128), + ], +) +def test_ltx2_split_rotate_half_per_head_cos(label, B, T, num_heads, head_dim, cos_dtype): + """LTX-2 production scenario: rotate-half (SPLIT) RoPE + per-head cos + full-dim norm.""" + device = "cuda" + torch.random.manual_seed(0) + + hidden = num_heads * head_dim + x = torch.randn(B * T, hidden, dtype=torch.bfloat16, device=device) * 0.5 + x_copy = x.clone() + weight = torch.randn(hidden, dtype=torch.bfloat16, device=device) * 5.0 + cos_2d, sin_2d = _make_per_head_cos(B, T, num_heads, head_dim, device, dtype=cos_dtype) + eps = 1e-6 + + torch.ops.trtllm.fused_dit_split_norm_rope( + x, + num_heads, + head_dim, + eps, + weight, + cos_2d, + sin_2d, + False, # interleave=False → rotate-half (SPLIT) + ) + ref = _torch_ref_per_head_cos( + x_copy, weight, cos_2d, sin_2d, num_heads, head_dim, eps, interleave=False + ) + torch.testing.assert_close(x, ref, rtol=2e-2, atol=5e-3) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_attention.py b/tests/unittest/_torch/visual_gen/test_ltx2_attention.py index 4e1b58c58b4d..e420b535d332 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_attention.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_attention.py @@ -52,6 +52,27 @@ def _init_weights(module: torch.nn.Module, std: float = 0.02): torch.nn.init.normal_(p, mean=0.0, std=std) +def _make_pe( + batch_size: int, + seq_len: int, + heads: int, + head_dim: int, + dtype: torch.dtype, + device: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build an identity-rotation (cos=1, sin=0) RoPE tuple for self-attn tests. + + LTX-2 self-attn forward (fuse_qk_norm_rope=True, head_dim ∈ {64, 128}) requires + ``pe`` to be a non-None ``(cos, sin)`` tuple in token-major [B, T, H, D] layout — + the same shape ``_split_freqs_cis`` produces in production. cos=1, sin=0 makes + the RoPE step an identity, so shape-only sanity checks remain meaningful while + still exercising the fused norm+RoPE kernel. + """ + cos = torch.ones(batch_size, seq_len, heads, head_dim, device=device, dtype=dtype) + sin = torch.zeros(batch_size, seq_len, heads, head_dim, device=device, dtype=dtype) + return cos, sin + + class TestLTX2SelfAttention(unittest.TestCase): """Test LTX2Attention self-attention with different backends.""" @@ -86,9 +107,10 @@ def test_vanilla_self_attention_sanity(self): ) x = torch.randn(batch_size, seq_len, query_dim, device=self.DEVICE, dtype=dtype) * 0.02 + pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=None) + output = attn(x, context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -122,9 +144,10 @@ def test_trtllm_self_attention_sanity(self): ) x = torch.randn(batch_size, seq_len, query_dim, device=self.DEVICE, dtype=dtype) * 0.02 + pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=None) + output = attn(x, context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -248,9 +271,10 @@ def test_gated_self_attention_sanity(self): self.assertIsNotNone(attn.to_gate_logits, "Gated attention should create to_gate_logits") x = torch.randn(batch_size, seq_len, query_dim, device=self.DEVICE, dtype=dtype) * 0.02 + pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - output = attn(x, context=None, pe=None) + output = attn(x, context=None, pe=pe) self.assertEqual(output.shape, (batch_size, seq_len, query_dim)) @@ -308,10 +332,11 @@ def test_backend_equivalence(self): trtllm_attn.load_state_dict(vanilla_attn.state_dict()) x = torch.randn(batch_size, seq_len, query_dim, device=self.DEVICE, dtype=dtype) * 0.02 + pe = _make_pe(batch_size, seq_len, heads, head_dim, dtype, self.DEVICE) with torch.no_grad(): - out_vanilla = vanilla_attn(x.clone(), context=None, pe=None) - out_trtllm = trtllm_attn(x.clone(), context=None, pe=None) + out_vanilla = vanilla_attn(x.clone(), context=None, pe=pe) + out_trtllm = trtllm_attn(x.clone(), context=None, pe=pe) # Skip comparison if either has NaN/Inf (can happen with random weights) has_nan = torch.isnan(out_vanilla).any() or torch.isnan(out_trtllm).any()