perf(qwen_vl): multimodal RoPE leaks f32 into the decode residual
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
`apply_multimodal_rotary_pos_emb` in `src/models/qwen2_vl.rs:225`, `src/models/qwen3_vl.rs:226` and `src/models/qwen3_vl_moe.rs:244` multiplied f16/bf16 queries and keys by f32 `cos`/`sin` tables and returned the promoted result. The tables are f32 on purpose (widened `inv_freq` and position ids), so the arithmetic is correct; the defect is that the promoted value escaped the helper into the residual stream and forced every later matmul to promote its own half-precision weight.
Fourth instance of the class fixed in #1709 (bridge activation helpers), #1710 (Phi/StableLM attention guard) and #1711 (attention scale promoting the query). The invariant is that a helper returns the dtype it was given: promotion inside a compiled graph is free, only promotion that escapes a function boundary costs bandwidth. Upstream mlx-vlm computes in float32 and then casts back under a `cast_output` flag (`q_embed = q_embed.astype(q.dtype)`) in https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/rope_utils.py; the port took the arithmetic and dropped the restore.
Only the image path is affected: the text-only fast path (`can_use_text_only_fast_path` / `forward_text_only_hidden`) uses ordinary 1-D rope and never enters this function, so text throughput looked healthy while image decode was slow. Once an image sets `position_ids` the MRoPE path is sticky for the rest of the sequence.
## Fix
Cast `q_embed` and `k_embed` back to `array_dtype(q)` / `array_dtype(k)` before returning, in all three files. Committed as 2f4fbabbc on `bench/0.7.0-refresh` (PR #1617).
## Measurement
M1 Ultra, decode tok/s with an image versus the same prompt without one:
| Checkpoint | Before | After | Ratio |
|---|---|---|---|
| qwen2.5-vl-3b-hf | 19.67 | 72.21 | 0.25x to 0.93x |
| qwen2.5-vl-3b-4bit | 123.62 | 145.30 | 0.74x to 0.87x |
| qwen2-vl-2b-4bit | 183.34 | 217.73 | 0.72x to 0.85x |
| qwen3-vl-4b-instruct-4bit | 112.91 | 116.90 | 0.87x to 0.90x |
The bf16 checkpoint gains most: it is the only non-quantized one in the family, so the promotion reached the full weight set rather than just the scales.
## Verification
Three-colour fixtures (blue, green, purple) against mlx-vlm as the reference. `qwen2.5-vl-3b-hf` answers Blue/Green/Purple, matching mlx-vlm exactly. `qwen2-vl-2b-4bit` answers Blue/Blue/Blue, exactly what mlx-vlm answers for that checkpoint; the pre-fix v0.6.0 binary answered Blue/White/Blue, further from the reference. The test is agreement with the reference implementation, not whether the answer is right.
## Do not generalize this
`src/vision/encoders/gemma3n.rs:58` has a similar-looking f32 promotion that must NOT be removed. It is load-bearing: mlxcel converts that tower's bf16 weights to f16 and its residuals reach about 213055, past the f16 max of 65504. A static audit cannot separate the two cases; only reading the surrounding code or measuring can.
Contributor guide
Research direction
Read the helpers at src/models/qwen2_vl.rs:225, src/models/qwen3_vl.rs:226, and src/models/qwen3_vl_moe.rs:244, then compare their returned dtypes with the mlx-vlm rope implementation. Verify the three-colour fixtures for the listed checkpoints against mlx-vlm, while preserving the separate gemma3n.rs:58 behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100