Do RoPE with mul instead of matmul (#49265)
* Compute RoPE frequencies with a broadcast multiply instead of a matmul
RoPE frequencies were computed as `inv_freq_expanded @ position_ids_expanded`, a matmul with an
inner dimension of 1. Depending on `torch.set_float32_matmul_precision`, `torch.compile` settings
and the kernel that gets picked, that matmul can run in TF32, which rounds both the positions
(integers above 2048 stop being exact) and the inverse frequencies to a 10-bit mantissa and
corrupts long-context positions. The `maybe_autocast(enabled=False)` block around it only
protected against autocast, not TF32.
Use `position_ids[..., None].float() * inv_freq` instead. It is an elementwise op, so it is never
computed in TF32 or downcast by autocast, and the autocast block in these forwards is removed. In
full float32 the result is bit-identical to the old matmul.
- M-RoPE rotaries broadcast `position_ids` to the number of M-RoPE axes explicitly; the matmul
against the expanded `inv_freq` used to do this implicitly, and generation can pass one row.
- Parakeet and Nemotron ASR relative positional encodings (same matmul) keep their batch
dimension with an explicit `expand`.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Comment cleanup, add explicit kwargs
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
