fix(specprefill): keep manual_rope output in the input dtype - #4461

Problem

manual_rope and manual_rope_with_freqs build cos and sin in float32 and return the product with the bf16/fp16 x without casting back. mx.fast.rope, which the normal path uses, returns the input dtype. After a SpecPrefill sparse prefill the KV cache holds float32 keys, and from the first full-attention layer on, every decode forward runs in float32.

On Qwen3.5 this matters because the verify-width quantized kernels (sg8, mma, vk in qwen35_verify_qmm) only accept bf16/fp16. I counted which route each MTP verify projection took with SpecPrefill on (batched engine, Qwen3.5 oQ4e, 16K prompt). Before the change about 95% of the calls had float32 activations and used the stock qmm. The other 5% matches the three linear-attention layers ahead of the first full-attention layer (3 of 64). After it, all calls are bf16 and 48.7% of the 232,596 calls use a fast kernel (two-row verifies have none, since vk_eligible needs at least three rows).

Change

Cast the result of both functions back to x.dtype.

Reproduction

The dtype needs no model:

import mlx.core as mx
from omlx.patches.specprefill import manual_rope

x = mx.random.normal((1, 4, 6, 64)).astype(mx.bfloat16)
print(manual_rope(x, mx.array([0, 3, 4, 9, 10, 11]), dims=64).dtype)
# main: mlx.core.float32   this change: mlx.core.bfloat16

End to end: serve a Qwen3.5-family MTP checkpoint on the batched engine (a text-only copy, with vision_config and the image/video/vision token ids removed from config.json, is served by BatchedEngine), enable mtp_enabled and specprefill_enabled with a draft model and specprefill_threshold 8192 in the model settings, send a chat request of about 16K prompt tokens with temperature 0 and max_tokens 400, and compare generation_tokens_per_second in usage with and without this change.

Measurements

Upstream main f3ad34d2 with and without this change, batched engine, Mac Studio M4 Max (40-core GPU) 64 GB. 16K prompt, 400 output tokens, temperature 0, one request at a time, MTP on, SpecPrefill on (threshold 8192, keep 20%, Qwen3.5-0.8B 4-bit draft). Sparse prefill bypasses the prefix cache, so every run is a cold prefill. Decode and TTFT are the median of three runs after one warmup; per-cycle figures are the range over those runs.

Qwen3.5 oQ4e text-only checkpoint with MTP head, 16,043-token prompt:

Two other checkpoints, served text-only so the batched engine runs them (vision config removed from a copy, weights unchanged), with 16,053- and 16,031-token prompts:

manual_rope returned float32 on every call before and never after. With SpecPrefill off neither function is called, and decode doesn't change beyond run-to-run noise: 43.0 and 44.6 tok/s for the 27B, 136.0 and 133.1 tok/s for the 35B-A3B. For the Qwen3.5 oQ4e checkpoint SpecPrefill off gives 38.1 tok/s. Acceptance is a little lower after the change, with about twice as many drafted tokens per cycle.

For reference, the same checkpoint on the VLM engine, which doesn't call these functions, decodes at 47.6 tok/s with SpecPrefill on. This isn't an equal-work comparison: on main the VLM path doesn't map the selected tokens back to their original positions (#3811), so its decode runs over a differently positioned context.

Scope

_PositionMappedRoPE is the only caller of the two functions and is installed only for attention modules with a .rope attribute. That covers mlx-lm RoPE models, including the custom-frequency ones that use manual_rope_with_freqs. mRoPE attention (rotary_emb) doesn't reach them. I measured three Qwen3.5-family checkpoints. None of them uses manual_rope_with_freqs, which I only checked numerically. Other families that reach these functions weren't measured.

Tests

Eleven new tests in tests/test_specprefill.py: both functions keep bf16, fp16 and float32 inputs in their dtype, _PositionMappedRoPE keeps bf16 and fp16, and contiguous bf16/fp16 calls agree with mx.fast.rope (including a partial-rotary case). Nine fail on main; the other two use float32 inputs. pytest tests -k "specprefill or rope" passes (282 passed, 17 skipped). I didn't run the whole default suite locally.

Numerical agreement

• Exact: across the call paths of both functions (base and scale, custom frequencies, pre_scale, partial rotary), in bf16 and fp16, the new output equals the old float32 output cast to the input dtype.
• Within tolerance: against mx.fast.rope on contiguous positions the largest absolute difference was 0.0625 for bf16 and 0.0078 for fp16 in the cases I ran. The new tests use a tolerance.

Not verified

• Generation quality. Greedy decoding (MTP off, 160 tokens, one 16K prompt) diverges from the old output after about 180 characters; both stay on topic. I didn't run a quality benchmark.
• Memory end to end. By arithmetic the float32 keys cost about 105 MB more than bf16 for this geometry (16 full-attention layers, 4 KV heads, head dim 256, 3,211 kept tokens).

Refs #4454