[TRTLLM-12580][perf] ltx2: fused RMSNorm+RoPE across all attention paths + PE pre-shard#13985

Merged

luyiyun1021 merged 13 commits into

NVIDIA:mainfrom

luyiyun1021:dev/ltx2-fused-all-paths

May 24, 2026

Commits on May 23, 2026

[None][perf] ltx2: fuse RMSNorm+RoPE across all attention paths
luyiyun1021
committed
[None][perf] ltx2: add norm-only split kernel + simplify forward dispatch
luyiyun1021
committed
[None][chore] ltx2: clean up dead branches in fused_dit kernels
luyiyun1021
committed
[None][perf] ltx2: switch SPLIT-rope cos/sin to token-major [B, T, H, D]
luyiyun1021
committed
[None][perf] ltx2: kernel accepts bf16 cos/sin, drop .float() upcast in helper
luyiyun1021
committed
[None][chore] ltx2: simplify cos_last derivation in fused norm+rope helpers
luyiyun1021
committed
[None][perf] ltx2: kernel-side cos broadcast over B, drop host .repeat(B, 1)
luyiyun1021
committed
[None][perf] ltx2: pre-shard 2D contiguous PE in TextCache, skip per-step reshape+contiguous
luyiyun1021
committed
[None][chore] ltx2: address CR review, simplify PE cache plumbing, fix kernel UB
luyiyun1021
committed
[None][perf] ltx2: restructure fused norm+rope kernels with cp.async + SMEM staging
luyiyun1021
committed
[None][chore] ltx2: post-CR cleanup of fused norm+rope plumbing
luyiyun1021
committed
[None][chore] visual_gen: add qk_norm_rope_kernel selector; route WAN to PR #13052 cross-head op
luyiyun1021
committed
[None][test] ltx2: self-attn unit tests pass valid pe to match fused contract
luyiyun1021
committed