Skip to content

[TRTLLM-15740][perf] VisualGen Wan: deduplicate shared RoPE SMEM staging #18036

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
yanxinzhangcs wants to merge 2 commits into
base: main
Choose a base branch
from
+96 −10
Draft

[TRTLLM-15740][perf] VisualGen Wan: deduplicate shared RoPE SMEM staging #18036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
74e6bf3
Deduplicate shared FP32 RoPE SMEM staging
yanxinzhangcs Aug 19, 2026
6507727
Merge main into perf/wan-qk-rope-shared-cos-smem
yanxinzhangcs Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 38 additions & 10 deletions cpp/tensorrt_llm/kernels/fusedDiTQKNormRopeKernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,9 @@ void launchFusedDiTQKNormRope(void* qkv, int num_tokens, int num_heads_q, int nu
// - 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.
//
// When PER_HEAD_COS is false and CosT is float, each head consumes the same HEAD_DIM-wide cos/sin row. Stage that
// row once per token instead of replicating it for every head in SMEM. Keep the existing bf16 staging unchanged.
template <int HEAD_DIM, bool INTERLEAVE, bool PER_HEAD_COS, typename CosT>
__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,
Expand All @@ -271,6 +274,7 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
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
static_assert(HEAD_DIM % CHUNK_ELEMS == 0, "HEAD_DIM must be divisible by the vector width");
// MAX_N = max(num_heads_q) * HEAD_DIM. 64 covers WAN-14B (40 heads) and
// future ≤64-head models. SMEM budget at 64h × 128 = 128 KB (bf16 cos) / 192 KB
// (fp32 cos), both within B200's 227 KB dynamic SMEM cap.
Expand Down Expand Up @@ -301,13 +305,16 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
: static_cast<int64_t>(cos_tokenIdx) * HEAD_DIM;

// SMEM layout: [Q row0][Q row1][K row0][K row1] bf16, [cos row0..1][sin row0..1] CosT, warp_sums.
// Only the shared fp32 path is deduplicated. The bf16 and per-head paths keep their existing layout.
constexpr bool kDeduplicateSharedCos = !PER_HEAD_COS && std::is_same_v<CosT, float>;
int const cosStride = kDeduplicateSharedCos ? HEAD_DIM : N;
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<CosT*>(smem_raw + 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16));
CosT* smem_sin = smem_cos + ROWS_PER_BLOCK * N;
CosT* smem_sin = smem_cos + ROWS_PER_BLOCK * cosStride;
float* warp_sums = reinterpret_cast<float*>(
smem_raw + 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * N * sizeof(CosT));
smem_raw + 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16) + 2 * ROWS_PER_BLOCK * cosStride * sizeof(CosT));

// Phase 0a: cp.async Q + K + cos + sin -> SMEM (all in one commit group).
#pragma unroll
Expand All @@ -323,7 +330,22 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
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<CosT, float>)
if constexpr (kDeduplicateSharedCos)
{
// Only the threads assigned to the first head stage the shared row.
if (elemBase < HEAD_DIM)
{
__pipeline_memcpy_async(
smem_cos + row_in_block * HEAD_DIM + elemBase, cos_emb + embBase + elemBase, 16);
__pipeline_memcpy_async(
smem_cos + row_in_block * HEAD_DIM + elemBase + 4, cos_emb + embBase + elemBase + 4, 16);
__pipeline_memcpy_async(
smem_sin + row_in_block * HEAD_DIM + elemBase, sin_emb + embBase + elemBase, 16);
__pipeline_memcpy_async(
smem_sin + row_in_block * HEAD_DIM + elemBase + 4, sin_emb + embBase + elemBase + 4, 16);
}
}
else if constexpr (std::is_same_v<CosT, float>)
{
__pipeline_memcpy_async(
smem_cos + row_in_block * N + elemBase, cos_emb + embBase + cosHeadOff + baseDim, 16);
Expand Down Expand Up @@ -360,7 +382,7 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
k_w_cache[chunk] = *reinterpret_cast<uint4 const*>(&k_weight[headIdx * HEAD_DIM + baseDim]);
}

// Phase 0c: wait + sync.
// Phase 0c: wait + sync. The barrier makes the shared row visible to threads assigned to every head.
__pipeline_wait_prior(0);
__syncthreads();

Expand Down Expand Up @@ -419,12 +441,16 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
uint4 const in_vec = *reinterpret_cast<uint4 const*>(&smem_input[row_in_block * N + elemBase]);
uint4 const w_vec = w_cache[chunk];

// CHUNK_ELEMS divides HEAD_DIM, so a vector never crosses a head boundary.
int const cosIdx
= kDeduplicateSharedCos ? (row_in_block * HEAD_DIM + elemBase % HEAD_DIM) : (row_in_block * N + elemBase);

float cos_vals[CHUNK_ELEMS];
float sin_vals[CHUNK_ELEMS];
if constexpr (std::is_same_v<CosT, float>)
{
float4 const* cs = reinterpret_cast<float4 const*>(&smem_cos[row_in_block * N + elemBase]);
float4 const* ss = reinterpret_cast<float4 const*>(&smem_sin[row_in_block * N + elemBase]);
float4 const* cs = reinterpret_cast<float4 const*>(&smem_cos[cosIdx]);
float4 const* ss = reinterpret_cast<float4 const*>(&smem_sin[cosIdx]);
float4 c0 = cs[0], c1 = cs[1];
float4 s0 = ss[0], s1 = ss[1];
cos_vals[0] = c0.x;
Expand All @@ -446,8 +472,8 @@ __global__ void fusedDiTQKNormFullDimRopeKernel(__nv_bfloat16* qkv, int const nu
}
else
{
uint4 const cp = *reinterpret_cast<uint4 const*>(&smem_cos[row_in_block * N + elemBase]);
uint4 const sp = *reinterpret_cast<uint4 const*>(&smem_sin[row_in_block * N + elemBase]);
uint4 const cp = *reinterpret_cast<uint4 const*>(&smem_cos[cosIdx]);
uint4 const sp = *reinterpret_cast<uint4 const*>(&smem_sin[cosIdx]);
uint const* cu = reinterpret_cast<uint const*>(&cp);
uint const* su = reinterpret_cast<uint const*>(&sp);
#pragma unroll
Expand Down Expand Up @@ -557,8 +583,10 @@ void launchFusedDiTQKNormRopeFullDim(void* qkv, int num_tokens, int num_heads_q,
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);
bool const deduplicateSharedCos = !per_head_cos && !cos_is_bf16;
int const cosStride = deduplicateSharedCos ? head_dim : N;
cfg.dynamicSmemBytes = 2 * ROWS_PER_BLOCK * N * sizeof(__nv_bfloat16)
+ 2 * ROWS_PER_BLOCK * cosStride * cos_elem_size + ROWS_PER_BLOCK * 2 * 4 /*WARPS_PER_ROW*/ * sizeof(float);
cfg.stream = stream;
cfg.attrs = attrs;
cfg.numAttrs = 1;
Expand Down
58 changes: 58 additions & 0 deletions tests/unittest/_torch/visual_gen/kernels/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -678,6 +678,64 @@ def test_full_dim_norm_packed_rotate_half_broadcast(label, B, T, num_heads, head
torch.testing.assert_close(qkv, ref, rtol=2e-2, atol=5e-3)


@pytest.mark.parametrize("interleave", [True, False], ids=["interleaved", "rotate_half"])
@pytest.mark.parametrize(
"num_heads,head_dim",
[(24, 64), (24, 128), (40, 128)],
ids=["head64", "head128", "wan_a14b"],
)
def test_full_dim_shared_fp32_cos_matches_expanded(interleave, num_heads, head_dim):
"""Shared cos/sin must be bitwise-equivalent to the per-head layout."""
device = "cuda"
torch.random.manual_seed(0)
batch_size, tokens_per_batch = 2, 32
num_tokens = batch_size * tokens_per_batch
hidden = num_heads * head_dim
qkv_shared = torch.randn(num_tokens, 3 * hidden, dtype=torch.bfloat16, device=device) * 0.5
qkv_expanded = qkv_shared.clone()
v_original = qkv_shared[:, 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_shared, sin_shared = _generate_cos_sin(tokens_per_batch, head_dim, device)
cos_expanded = (
cos_shared.unsqueeze(1)
.expand(-1, num_heads, -1)
.reshape(tokens_per_batch, hidden)
.contiguous()
)
sin_expanded = (
sin_shared.unsqueeze(1)
.expand(-1, num_heads, -1)
.reshape(tokens_per_batch, hidden)
.contiguous()
)

for qkv, cos, sin in (
(qkv_shared, cos_shared, sin_shared),
(qkv_expanded, cos_expanded, sin_expanded),
):
_call_fused_kernel(
qkv,
num_heads,
num_heads,
num_heads,
head_dim,
1e-6,
q_weight,
k_weight,
None,
None,
cos,
sin,
-1,
interleave=interleave,
tokens_per_batch=tokens_per_batch,
)

torch.testing.assert_close(qkv_shared, qkv_expanded, rtol=0, atol=0)
torch.testing.assert_close(qkv_shared[:, 2 * hidden :], v_original, rtol=0, atol=0)


def test_full_dim_norm_packed_v_unchanged():
"""Full-dim path leaves V slice untouched (sanity)."""
device = "cuda"
Expand Down