RoPE helpers make float32 copies of q / k on every call #15220

Is your feature request related to a problem? Please describe.

apply_rotary_emb (use_real=True) and several model-local RoPE helpers (LTX, LTX-2, QwenImage, LongCat-AudioDiT, JoyImage, Allegro, and their # Copied from copies) compute

out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)

cos / sin are float32, so type promotion already does this math in float32. The .float() calls only materialize full-size float32 copies of q / k on every attention call. On an RTX 5090, apply_rotary_emb takes about 8% of a Flux.1 / Chroma transformer forward at 1024x1024 (36.5 ms of 439 ms, 114 calls).

Describe the solution you'd like.

Drop the .float() calls and cast cos / sin to at least float32 instead, so bf16 / fp16 freqs keep float32 math. The output is bit-identical (bf16 / fp16 to float32 is exact, and the products and the sum are the same operations). Measured: full pipelines 1.3-3.7% faster on RTX 5090 and Radeon AI PRO R9700, op alone 1.15-2.0x, gradients bit-identical, torch.compile unchanged. Implemented in #15215.

Describe alternatives you've considered.

A fused Triton / CUDA RoPE kernel would be faster but adds a kernel dependency; this change keeps the same PyTorch ops and only removes the copies.

Additional context.

The same idea also applies to RoPE helpers written in other shapes (even / odd halves in Mochi and the LTX-2 diffusion decoder, the LTX-2 split RoPE, Stable Audio 3 partial RoPE, the 2x2 rotation-table form in Kandinsky / HiDream / PRX). Those are not in #15215.