lablup / lablup/mlxcel

perf(qwen_vl): multimodal RoPE leaks f32 into the decode residual

Open
#1,714 1 comment 0 reactions 0 assignees View on GitHub
area:models modeltype:vlm priority:medium status:review type:performance
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.