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

[yanxinzhangcs]

@coderabbitai summary

Publication gate: obtain approval on the linked performance request and replace every pending exact-candidate result before marking this PR ready for review.

Description

Closes #18035.

The full-dimension fused DiT QK-norm/RoPE kernel currently expands a shared FP32 cosine/sine row across every head in dynamic shared memory. Wan A14B uses one head_dim=128 row shared by 40 heads, so the expansion moves and reserves data that every head reads identically.

This change:

• stages one shared FP32 cosine and sine row per token/CTA row;
• indexes that row by dimension for every head;
• reduces the shared-FP32 dynamic shared-memory allocation from the full Q/K width to head_dim;
• leaves BF16 cosine/sine and per-head layouts unchanged;
• keeps the existing generic full-dimension entry point without an SM100-only launch-bound specialization; and
• adds exact shared-vs-expanded-layout tests across two RoPE styles and three head shapes.

No public API, mathematical algorithm, or architecture-specific dispatch is added. The existing generic dispatch remains in place for every layout and architecture.

This is a focused optimization of the current two-row full-dimension kernel introduced in #13985 and routed to Wan by the unified dispatch in #14529. The earlier #13052 cross-head path was separate and was removed by #14529.

Attribution

The initial optimization idea, prototype, and B200 research measurements were contributed by @chang-l. This PR reshapes that work into a shared-FP32-only upstream change with explicit fallbacks and layout-parity coverage.

Performance and quality evidence

Exact upstream candidate

Candidate commit: 6507727246e91e80855d398694e8bd4a585610e1.

Gate	Result
B200 SM100 build and source-tree identity	PENDING — update before review
New shared-vs-expanded parity cases	PENDING — expected 6/6
Existing fused DiT QK-norm/RoPE regression tests	PENDING — expected 11/11 in the focused gate
Dispatch/resource check (40x128, shared FP32, interleaved)	PENDING — must show generic <128,true,false,float>, 43,072 B dynamic SMEM, STACK=LOCAL=0, no LDL/STL, and the pinned two-call internal rsqrt topology
Expanded-layout control	PENDING — must show generic <128,true,true,float>, 122,944 B dynamic SMEM, STACK=LOCAL=0, no LDL/STL, and the same canonical call/return skeleton; a direct return or one adjacent preserved-link register bridge is allowed
Candidate kernel benchmark	PENDING
Candidate Wan end-to-end denoise benchmark	PENDING
Candidate LPIPS against the existing golden	PENDING

Do not substitute the research data below for these candidate-specific gates. Build-only diagnostics on the identical source diff passed the source/tree gate, full SM100 build, 6/6 new tests, 11/11 existing focused tests, and the profiled dispatch test. They reported candidate REG93/STACK0/LOCAL0 at 43,072 B dynamic SMEM and expanded control REG95/STACK0/LOCAL0 at 122,944 B, with no LDL/STL in either focused SASS body. The first diagnostic validator incorrectly rejected any compiler-generated internal CALL; its successor then incorrectly required a direct return-register match and rejected the control specialization's valid one-hop return bridge. These runs are diagnostic evidence, not passing receipts. The final fetchable commit must pass the corrected dual-role gate before review.

Earlier research prototype (directional only)

Measured by @chang-l on B200; this was not the final upstream diff.

Measurement	Baseline	Research prototype	Change
Standalone fused QK-norm/RoPE kernel	975.0 us	486.98 us	2.00x faster
In-model fused QK-norm/RoPE kernel	1303.7 us	681 us	1.91x faster
Wan 480x832x81-frame denoise	107.118 s	106.629 s	-0.46%
Dynamic shared memory, A14B dispatch	122,944 B	43,072 B	-65.0%

The research parity check reported a bitwise-identical full QKV output.
The same research sweep measured LB2 at 504 us and LB3 at 486.98 us, while LB4/LB5 regressed to 542/611 us. The final candidate deliberately carries no launch bound; these sweep points are directional context, not its performance result.

A superseded research prototype included the launch-bound specialization and
also changed shared-BF16 staging. It is not the PR patch. The proposed diff is
FP32-only, keeps the BF16 layout unchanged, and is the only source candidate
covered by the pending gates.

Test Coverage

• test_full_dim_shared_fp32_cos_matches_expanded compares shared FP32 embeddings against materialized per-head embeddings with rtol=0, atol=0 for:
  • interleaved and rotate-half RoPE;
  • (num_heads, head_dim) = (24, 64), (24, 128), (40, 128); and
  • an unchanged V slice.
• File-scoped repository pre-commit passed locally, including clang-format, Ruff, codespell, whitespace, conflict, and test-list hooks.
• Python syntax compilation passed for the modified test.
• This development host has no CUDA toolkit or GPU; the exact B200 build/runtime and VisualGen end-to-end gates are pending above.

PR Checklist

• PR description explains the redundant staging and shared-FP32-only scope.
• Change follows the TRT-LLM coding guidelines and has focused unit coverage.
• No public API or dependency change.
• No CODEOWNERS or architecture-diagram update is required.
• Linked performance request is approved.
• Exact upstream candidate passes B200 build, parity, dispatch, performance, and LPIPS gates.
• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, comment /bot help.

[yanxinzhangcs]

Updated the branch to current main in signed merge commit 6507727246e91e80855d398694e8bd4a585610e1; the reported merge conflict is resolved and the PR-owned diff remains limited to the kernel plus its focused parity test. On the merged tree, Python syntax compilation, git diff --check, pre-commit, and static shared-memory/address-boundary checks pass. I also corrected the two publication-gate checkboxes in the PR body: performance-request approval and exact-candidate B200/perf/LPIPS evidence are still pending, so this PR remains Draft. CI has not started because this account is not in the Blossom-CI authorization allowlist.

[yanxinzhangcs]

CI note: this account is not in the repository’s Blossom-CI authorization allowlist, so the earlier command did not start a pipeline. An authorized owner must run /bot run on current head 6507727246 after the performance request is approved.

[yanxinzhangcs]

CI handoff: yanxinzhangcs is not present in the repository’s Blossom-CI authorization allowlist, so my /bot run comment cannot start the pipeline. @byshiue or @kaiyux, once the performance request is approved, please run /bot run on this updated head (6507727246) so the B200 build/test gates can be collected.