qk-norm-rope

Fused RMSNorm (across all heads) + split-rotary embedding CUDA kernel for LTX-2 style audio-video DiTs, built with kernel-builder for the kernels library.

It replaces the QK-norm + RoPE sequence used by e.g. the diffusers LTX-2 attention processors:

query = attn.norm_q(query)                              # RMSNorm over heads * head_dim
query = apply_split_rotary_emb(query, (cos, sin))       # per-head first/second-half rotation

with a single pass: each row is staged in shared memory, normalized with an fp32 warp-shuffle reduction (rounded to bf16 at the norm boundary to match the reference numerics bit-for-bit within 1-2 ulps), rotated, and written once.

Two ops are provided:

rms_norm_split_rope(x, weight, cos, sin, heads, eps) — single tensor.
rms_norm_split_rope_qk(q, k, weight_q, weight_k, cos, sin, heads, eps) — self-attention fast path. Q and K share one rotary table, and the fp32 cos/sin tables are as large as the activations themselves, so reading them once for both tensors is ~1.3x faster than two single calls.

Both ops register fake implementations, so they compose with torch.compile(fullgraph=True) / compile_repeated_blocks(fullgraph=True).

Inputs are bf16 (x, weights) and fp32 (cos, sin) in the (..., heads, head_dim / 2) contiguous layout — the buffer underlying the (B, H, T, r) view that apply_split_rotary_emb receives.

Measured on NVIDIA GB10 (DGX Spark), LTX-2.3 at 768x512, 121 frames, CFG batch 2: 5.75x over the eager norm+rope sequence, and -3.3% end-to-end denoising time on top of regional compilation + cuDNN attention.

How to use

# make sure `kernels` is installed: `pip install -U kernels`
from kernels import get_kernel

kernel_module = get_kernel("sayakpaul/qk-norm-rope", version=1)
rms_norm_split_rope = kernel_module.rms_norm_split_rope

rms_norm_split_rope(...)