…e cross_head op

Goal: collapse the Python-side fused QK-norm + RoPE dispatch in visual_gen
to a single op call with no per-model branching, and remove the orphaned
fused_dit_cross_head_qk_norm_rope op once the full-dim template covers
its envelope.

Kernel changes
- fusedDiTQKNormRopeKernel.cu / .h:
  • Full-dim template: lift num_heads cap 32 -> 64 (MAX_N = 64 * HEAD_DIM)
    so WAN-14B (40 heads * 128) fits. SMEM budget stays well within B200's
    227 KB dynamic SMEM cap.
  • Per-head template: add cos_seq_per_batch runtime parameter for
    kernel-side cos broadcast over B (matching what the full-dim template
    already did). Python dispatcher no longer needs host-side .repeat(B, 1)
    for FLUX dual-stream CFG paths.
  • Delete fusedDiTCrossHeadQKNormRopeKernel + launchFusedDiTCrossHeadQKNormRope
    (kernel + launcher + .h declaration).
- fusedDiTSplitNormKernel.cu + fusedDiTSplitQKNormRopeKernel.cu:
  • Lift num_heads cap 32 -> 64 (these caps were SMEM-budget driven, not
    block-size driven as the misleading comments suggested -- comments
    corrected).

C++ op surface
- fusedDiTQKNormRopeOp.cpp:
  • Accept cos_emb / sin_emb with rank in [2, 4]; the op flattens internally
    based on whether shape[-2] == num_heads, so callers can pass raw cos
    tensors without reshaping.
  • Drop the per-head launcher's cos_seq_per_batch == 0 reject (kernel
    now supports broadcast).
  • Remove the fused_dit_cross_head_qk_norm_rope op entirely (schema +
    impl + registration).
- fusedDiTSplitQKNormRopeOp.cpp: same rank-flexible cos handling.

Python -- modules/attention.py
- Remove the qk_norm_rope_kernel: str = ... selector parameter and its
  validation assert; the C++ op now auto-dispatches by tensor shape/dtype.
- apply_packed_qk_norm_rope / apply_split_norm_rope / apply_split_norm
  collapse to ~6-line passthroughs -- no cos shape inference, no
  host-side broadcast tile, no kernel-name branching.

Python -- model files
- models/wan/transformer_wan.py: drop the
  qk_norm_rope_kernel="fused_dit_cross_head_qk_norm_rope" arg from
  attn1's Attention(...). All WAN variants (12 / 24 / 40 heads * 128)
  now route through the full-dim template via the default op.
- models/ltx2/transformer_ltx2.py: remove the head_dim in (64, 128)
  hardcoded check from both LTX2Attention.forward's guard and
  project_kv. The kernel templates already enforce the head_dim
  envelope; if a caller constructs LTX2Attention with an unsupported
  head_dim the kernel throws a clear TLLM_THROW.

Test / utility hygiene
- _torch/compilation/utils.py: drop fused_dit_cross_head_qk_norm_rope
  from the torch.compile op-allowlist mapping.
- tests/integration/defs/examples/test_visual_gen.py: delete the
  _disable_wan_fused_qk_norm_rope_if_unavailable runtime fallback helper
  and its call site (cross-head op no longer exists, so the
  graceful-fallback wrapping is moot).
- tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py:
  delete the cross-head test block.
- benchmarks/bench_fused_dit_cross_head_qk_norm_rope.py: delete entire file.

Validation
- Microbenchmark (B200, 30-40 iters, min): the full-dim template is the
  fastest fused QK-norm + RoPE kernel for every WAN num_heads value we
  measured (12 / 24 / 32 / 40 / 48 / 56 / 64 heads x 128). Direct vs the
  deleted cross-head op on representative WAN shapes:
    WAN-1.3B 12h, 32K tok, B=1:      0.166 ms vs 0.674 ms -> 4.06x faster
    WAN-TI2V-5B 24h, 32K tok, B=1:   0.233 ms vs 0.872 ms -> 3.75x faster
    WAN-14B family 40h, 32K tok, B=1: 0.492 ms vs 1.093 ms -> 2.22x faster
  CFG=2 (B=2) shapes: the cross-head op fails outright (no batch broadcast
  path); the full-dim path handles it natively via cos_seq_per_batch.
- Bit-similarity full-dim vs cross-head on 5 WAN shapes (12 / 24 / 40 / 48
  / 64 heads * 128): max_abs_diff <= 1.6e-2 (1 bf16 ULP at output ~ 1),
  mean ~ 1.24e-8.
- tests/unittest/_torch/visual_gen/test_ltx2_attention.py: 6/6 PASS.
- tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py:
  55 passed, 24 skipped (per-head + full-dim suites; cross-head suite
  deleted with the op).
- LTX-2 nvfp4 single-stage e2e smoke earlier on this branch (40 steps,
  768x1280 x 121 frames, 1 GPU): 12.5 s, 12.9 MB mp4. No regression.

Follow-ups (deferred to keep diff focused):
- PR NVIDIA#13978 (async Ulysses, LTX-2 + WAN): move per-model compute_q/k/v
  closure into a base Attention.forward_async() now that the closure's
  norm+rope step is a single apply_split_norm_rope call.
- Merge fused_dit_split_norm + fused_dit_split_norm_rope into one op.
- Wire WAN cross-attn (attn2) through fused_dit_split_norm instead of
  eager apply_qk_norm.

Signed-off-by: Yiyun Lu <[email protected]>

…ape equality

Both fused_dit_qk_norm_rope and fused_dit_split_norm_rope previously checked only rank parity and post-flatten 2D-shape equality. That permits silently-bad layouts (e.g. cos=[B,S,D] vs sin=[S,B,D]) to reshape to the same [B*S, D] buffer but pair the wrong sin row with each cos row, rotating the embedding incorrectly. Reject mismatched raw shapes before flatten so this fails fast with a clear error.

Signed-off-by: Yiyun Lu <[email protected]>

…tests

The fused split-norm kernel only supports head_dim in {64, 128}. Earlier in this PR I removed the Python-side 'head_dim in (64, 128)' guards from LTX2Attention.forward and project_kv on the assumption that 'production always uses 64 or 128, so the guard is dead code'. That assumption was wrong for the test suite: tests/unittest/_torch/visual_gen/test_ltx2_transformer.py builds mini LTX-2 models with head_dim=32 (TestLTX2AudioVideoModel, TestLTX2VideoOnlyModel, TestLTX2CUDAGraphCapture, TestLTX2TextCache, TestLTX2CacheDiTWrapperPassthrough — 6 failing tests) to keep unit-test runtime down. Without the Python-side guard those tests hit the kernel's clear TLLM_THROW. Restore the head_dim guard in both call sites so the eager path is reachable for head_dim != 64/128, and update _forward_unfused docstring to mention the mini-config path explicitly.

Signed-off-by: Yiyun Lu <[email protected]>