EMLX.Fast.rope*/mlx::core::fast::rope miscomputes rotations for multi-head (H>1) {B,T,H,D} tensors

[polvalente]

Summary

EMLX.Fast.rope/6, rope_with_positions/6, and rope_with_freqs/6 (decode /
T=1 fast paths) return correct results only for head 0 on a multi-head
{B, T, H, D} tensor (H>1). Every other head is rotated by the wrong angle,
with error growing with the position offset. Reproduces on both :cpu and
:gpu (Metal), for all three fast::rope overloads (scalar offset, array
offset, freqs tensor).

Severity: high. Silently wrong numbers, not a crash — every real
multi-head attention decode step (H>1 is universal) calling these entry
points on Bumblebee's standard heads-not-yet-transposed layout is affected.

Root cause

mlx::core::fast::rope (mlx/backend/metal/rope.cpp, mirrored in the CPU
backend) has a head_seq_transpose stride-detection special case meant to
recognize the {B, T, H, D} (heads-not-transposed) convention, but it only
triggers for a specific non-row-contiguous stride signature. A plain,
freshly-allocated {B, T, H, D} tensor is row-contiguous, so the guard never
fires, and the kernel falls into the row_contiguous branch that iterates
its own "T" axis over what is actually our H axis — applying rotation angle
position + head_index instead of position.

Minimal repro

a = Nx.iota({1, 1, 2, 64}, type: :f32) |> Nx.divide(100)
      |> Nx.backend_transfer({EMLX.Backend, device: :gpu})

joint = EMLX.fast_rope(EMLX.Backend.from_nx(a), 64, false, 10_000.0, 1.0, 6)
        |> EMLX.Backend.to_nx()

head1_alone = a[[.., .., 1..1, ..]]
alone = EMLX.fast_rope(EMLX.Backend.from_nx(head1_alone), 64, false, 10_000.0, 1.0, 6)
        |> EMLX.Backend.to_nx()

# joint's head-1 slice != alone, even though both represent "head 1 at the
# same position". `alone` matches the textbook RoPE formula; `joint` doesn't.

Why existing tests didn't catch it

The existing equivalence tests compare the compiled-graph opcode against the
eager EMLX.Fast NIF — both call the same buggy mlx::core::fast::rope
primitive, so they trivially agree with each other while both disagree with
the textbook formula. It only surfaced when compared against an independent,
hand-written cos/sin/rotate primitive (a separate prefill/T>1 lowering path
that doesn't call fast::rope at all).

Practical impact

validate_qwen3.exs doesn't hit this because Qwen3's RoPE base (1e6) routes
decode through the hand-written fast_rope_positions NIF rather than
fast::rope directly. Any config that does route decode through fast::rope
directly (this is the default path for EMLX.Fast.rope*) is affected.

Suggested fix

One of:

1. Fix mlx::core::fast::rope's head_seq_transpose detection to also
   trigger for the row-contiguous {B, T, H, D} case (upstream MLX fix).
2. Transpose Q/K to {B, H, T, D} before calling fast::rope and transpose
   back after (extra copy, but correct) — EMLX-side workaround.
3. Route all multi-head EMLX.Fast.rope* calls through the existing
   hand-written cos/sin/rotate composition (fast_rope_positions-style)
   instead of mlx::core::fast::rope, eliminating the dependency entirely —
   EMLX-side workaround.

[polvalente]

This is conditioned on merging #120