perf(qwen_vl): text decode falls to half of mlx-lm as context grows
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
Qwen VL text decode is not slow against mlx-lm by a constant factor. The ratio degrades monotonically with prompt length, on every Qwen VL checkpoint measured. These are TEXT-mode runs with no image, so this is the language-model decode path of these checkpoints and not the vision tower.
Environment: mlxcel 0.7.0-beta.1 at `03cfcc61`, MacBook Pro M5 Max 128GB, macOS 26.6.2, MLX pin `9a795735`, built with `cargo build --release --features metal,accelerate`. Baseline is mlx-lm 0.31.3 in the repo's `.venv-mlxlm` (Python 3.12, mlx 0.32.2). Both sides decode 128 tokens greedily from a synthetic prompt padded to the stated length. Cells are mlxcel decode tok/s against mlx-lm decode tok/s, with the ratio.
| Checkpoint | prompt 64 | prompt 128 | prompt 512 | prompt 2048 |
|---|---|---|---|---|
| `qwen2.5-vl-3b-4bit` | 160.5 / 225.5 = 71% | 156.2 / 224.4 = 70% | 137.8 / 218.0 = 63% | 95.8 / 206.1 = 46% |
| `qwen3-vl-4b-4bit` | 153.8 / 172.7 = 89% | | | 88.4 / 168.1 = 53% |
| `qwen2-vl-2b-4bit` | 267.7 / 376.2 = 71% | | | 157.3 / 341.8 = 46% |
The shape is what matters. Over a 32x increase in context, mlxcel loses 40 to 43% of its decode rate (160.5 to 95.8, 153.8 to 88.4, 267.7 to 157.3) while mlx-lm loses 9 to 11% (225.5 to 206.1, 172.7 to 168.1, 376.2 to 341.8). Something in the mlxcel decode path scales with context length far more steeply than the reference does.
Controls, from the same-day sweep at 512 prompt tokens (`benchmarks/metal_m5max_2026-09-06.csv` against `benchmarks/pylm_m5max_2026-09-06.csv`): the non-VL Qwen decoders are at parity, `qwen3-4b-4bit` at 186.47 against 185.95 (100%) and `qwen2.5-7b-4bit` at 124.00 against 123.67 (100%). More pointed still: mlx-lm decodes `qwen3-vl-4b-4bit` at 186.78 and `qwen3-4b-4bit` at 185.95, so on the reference side a Qwen3-VL text decoder costs the same as the plain Qwen3 decoder of the same size. mlxcel runs that same pair at 135.89 and 186.47. The gap belongs to our Qwen VL text-decode path, not to the architecture.
Why this matters for the published numbers: `docs/benchmark_results/model_tests_m5max.md:90-109` already records this cluster as context sensitivity rather than a regression and calls it "a standing optimization target for the M-RoPE attention path", but it quotes a single ratio per model taken at 512 prompt tokens (`qwen2.5-vl-3b-4bit` 63%, `qwen2-vl-2b-4bit` 64%, `qwen3-vl-4b-4bit` 73%, `qwen3-vl-8b-4bit` 79%). One point at 512 describes the 512-token case only, and understates the gap at the context lengths these models are actually used at.
This is not a regression. It reproduces at the current commit and no earlier measurement shows these models at parity across context lengths. It is an optimization target, and the per-length curve above is the measurement to move.
## Current Behavior
Three per-decode-step costs sit in the Qwen VL decoder that neither the plain Qwen decoders nor mlx-vlm pay. Which of the three a given checkpoint pays explains both the level and the slope in the table above.
**1. `repeat_kv` materializes the whole live cache n_rep times, per layer, per decode step.** Every Qwen VL attention path expands the fetched K and V to full head count before SDPA: `src/models/qwen2_vl.rs:362-371`, `src/models/qwen3_vl.rs:374-383` and `:448-457`, `src/models/qwen3_vl_moe.rs:390-399` and `:464-473`. `repeat_kv` (`src/lib/mlxcel-core/src/utils.rs:954-975`) is a reshape, then `broadcast_to`, then a reshape, and the trailing reshape of a broadcast array is a copy, so this writes an n_rep-sized duplicate of the entire live cache. It is also unnecessary. The doc comment directly above it (`src/lib/mlxcel-core/src/utils.rs:946-953`) already states the rule: "fused SDPA broadcasts KV heads internally, so only models that materialize attention scores themselves need an explicit repeat". These models do not materialize scores. They call `attention_from_ptr` (`src/lib/mlxcel-core/src/layers.rs:4547`), which reaches `attention_dispatch` (`:4252`) and, at `softcap == 0`, `ffi::fast_scaled_dot_product_attention` (`:4277`).
This is the term that scales with context, and it is the one thing all three checkpoints share. Sizing it for `qwen2.5-vl-3b-4bit` (36 layers, 16 Q heads, 2 KV heads so n_rep 8, head_dim 128): live K plus V is 1024 bytes per token per layer in fp16, so the repeat writes 16 MiB per layer and 576 MiB per decode token at 2048 tokens, against 18 MiB per decode token at 64. Set that 558 MiB difference against the measured per-token difference (6.23 ms at 64, 10.44 ms at 2048) and it implies roughly 133 GiB/s of copy traffic, doubled once SDPA reads the expanded buffer back. That is the right order of magnitude for the observed falloff on this hardware, which makes it the leading suspect rather than a proof.
**2. An all-permitting causal mask rebuilt on every decode step.** `src/models/qwen2_vl.rs:729-733`, `src/models/qwen3_vl.rs:963-967` and `src/models/qwen3_vl_moe.rs:1124-1128` call `create_causal_mask(seq_len, caches[0].live_len())` unconditionally, including at `seq_len == 1`. `create_causal_mask` (`src/lib/mlxcel-core/src/utils.rs:118-158`) builds a `[seq_len, seq_len + offset]` f32 array from index comparisons. At `seq_len == 1` the single query row sits at logical position `live_len` and every key column `k <= live_len` is permitted, so all `live_len + 1` columns are permitted: the mask is uniformly zero, constrains nothing, and only forces the masked arm of the fused SDPA. mlx-lm returns `None` for this case (`create_attention_mask`, `mlx_lm/models/base.py:51-52` in `.venv-mlxlm`), and `src/models/qwen3_5.rs:1394-1403` already carries the `seq_len > 1` guard in-tree.
**3. MRoPE cos/sin recomputed once per layer.** `src/models/qwen2_vl.rs:355`, `src/models/qwen3_vl.rs:367` and `src/models/qwen3_vl_moe.rs:383` call `self.mrope.forward(position_ids)` inside each layer's attention. `MRoPE::forward` (`src/models/qwen2_vl.rs:132-174`) is a `from_slice_f32`, an `astype`, two reshapes, two broadcasts, a matmul, a transpose, `apply_mrope`'s slice/squeeze/slice_update loop over `mrope_section`, a concatenate, a cos and a sin. On a 36-layer model that is 36 identical builds per decode token from identical inputs. mlx-vlm builds it once per forward and threads `(cos, sin)` into every layer (`mlx_vlm/models/qwen2_5_vl/language.py:195` and `:198`).
**Why that decomposition fits the data.** `qwen3_vl.rs` and `qwen3_vl_moe.rs` have a text-only fast path: `forward_hidden` routes to `forward_text_only_hidden` (`src/models/qwen3_vl.rs:894-895`, gate at `:717-732`) whenever there is no MRoPE state, no visual mask and no deepstack embeds, which is exactly a text run. That path uses `fast_rope` instead of MRoPE and passes `mask` straight through without building one (`src/models/qwen3_vl.rs:992-1005`), so it pays neither cost 2 nor cost 3. `qwen2_vl.rs` has no such path at all (zero occurrences of `forward_text_only`) and pays both. That is the level: `qwen3-vl-4b-4bit` starts at 89% of mlx-lm at 64 tokens while `qwen2.5-vl-3b-4bit` and `qwen2-vl-2b-4bit` start at 71%. The slope is the same 40 to 43% for all three, because `repeat_kv` is on every one of these paths including the text-only one.
## Proposed Solution
Three independent changes, in the order they should land, each measurable on its own.
1. **Drop `repeat_kv` from the Qwen VL decoder attention and pass GQA-shaped K and V straight to `attention_from_ptr`.** Sites: `src/models/qwen2_vl.rs:362-371`, `src/models/qwen3_vl.rs:374-383` and `:448-457`, `src/models/qwen3_vl_moe.rs:390-399` and `:464-473`. Where the current code takes `mlxcel_core::copy(&k)` in the `n_rep == 1` arm, use the fetched array directly rather than reintroducing a copy. Then regenerate the `// Used by:` list on `repeat_kv` (`src/lib/mlxcel-core/src/utils.rs:946-953`) with the `grep` that comment itself specifies, per `docs/code-guidelines.md`.
2. **Guard the auto causal mask on `seq_len > 1`**, mirroring `src/models/qwen3_5.rs:1394-1403`, at `src/models/qwen2_vl.rs:729-733`, `src/models/qwen3_vl.rs:963-967` and `src/models/qwen3_vl_moe.rs:1124-1128`. An explicitly supplied `mask` still wins as it does today; only the auto path changes.
3. **Hoist the MRoPE cos/sin build out of the per-layer attention.** Compute it once in `forward_hidden` beside `position_ids` and thread the `(cos, sin)` pair through the decoder layer into `Attention::forward`, which is what mlx-vlm does. This touches `src/models/qwen2_vl.rs`, `src/models/qwen3_vl.rs` and `src/models/qwen3_vl_moe.rs`, and is the only one of the three that changes function signatures.
Rejected: writing a fused GQA decode kernel for these models. MLX's fused SDPA already broadcasts KV heads internally, so the fix is to stop defeating it, not to add a kernel.
## Scope
**In scope:** `src/models/qwen2_vl.rs`, `src/models/qwen3_vl.rs`, `src/models/qwen3_vl_moe.rs`, the `// Used by:` comment at `src/lib/mlxcel-core/src/utils.rs:946-953`, and the doc update listed under Acceptance Criteria.
**Out of scope:** `src/models/qwen3_5.rs`, which already guards the mask and never calls `repeat_kv`; `repeat_kv` itself, which stays for the models that genuinely materialize attention scores; the vision towers, which no measurement here touches; the other `repeat_kv` callers (`deepseek_v2`, `minicpm3`, `nemotron_nas`, `recurrent_gemma`, `glm4v`, `glm4v_moe`, `ernie4_5_moe_vl`, `hunyuan_vl`, `paddleocr_vl`), which may carry the same defect but have no measurement in this issue and should be filed separately if so; prefill, which is already at or above parity in the same sweeps.
## Implementation Notes
- **Reuse**: `src/models/qwen3_5.rs:1394-1403` is the in-tree pattern for the mask guard. `src/models/qwen3.rs` is the shape to converge on for the attention body, since it hands the cache output to `attention_from_ptr` with no repeat and is the checkpoint measuring 100% above.
- **Both reachable SDPA arms already accept GQA shapes**, so change 1 needs no new arm. `metal4_attention` has a GQA regression test (`metal4_attention_gqa_shape`, `src/lib/mlxcel-core/src/layers.rs:7133`), and the chunked arm cannot be reached at decode because `materializing_sdpa_query_chunk` returns `None` when `q_len < 2` (`src/lib/mlxcel-core/src/layers.rs:4376`).
- **The `--max-kv-size` trim invariant must survive change 2.** The comment at `src/models/qwen2_vl.rs:718-727` explains why the mask is sized from `live_len()` and not from the monotonic `offset` (issue #421). Guarding on `seq_len > 1` leaves that sizing untouched on the prefill path where it matters.
- **Edge cases**: `n_rep == 1` already skipped the repeat and only paid a copy; multimodal prefill takes the MRoPE path in all three files and must keep byte-identical positions; chunked prefill has `cache_offset > 0` with `seq_len > 1` and still needs the mask; the server's per-sequence path (`forward_with_sequence_id`) must resolve `position_ids` per sequence before change 3 hoists the cos/sin, so the hoisted pair belongs to the right row; `qwen3_vl` deepstack and visual-mask runs bypass the text-only fast path and must still work.
- **Error handling**: no new failure paths. A shape mismatch from change 1 surfaces as an MLX broadcast error at the SDPA call, which is a visible test failure rather than a silent wrong answer.
## Acceptance Criteria
- [ ] `repeat_kv` no longer appears in `src/models/qwen2_vl.rs`, `src/models/qwen3_vl.rs` or `src/models/qwen3_vl_moe.rs`, and the `// Used by:` list at `src/lib/mlxcel-core/src/utils.rs:946-953` matches the regenerated `grep` output.
- [ ] `create_causal_mask` is not reached at `seq_len == 1` in any of the three files.
- [ ] MRoPE cos/sin is built once per forward rather than once per layer in all three files.
- [ ] A profile at 64 and at 2048 prompt tokens on `qwen2.5-vl-3b-4bit` attributes the extra per-token cost to these specific code paths rather than asserting it, and is recorded under `docs/benchmark_results/`.
- [ ] The per-length curve is re-measured at 64, 128, 512 and 2048 for `qwen2.5-vl-3b-4bit`, `qwen3-vl-4b-4bit` and `qwen2-vl-2b-4bit`, and either the 64 to 2048 falloff moves toward mlx-lm's 9 to 11% or the residual is explained in writing as structural for this family.
- [ ] Output is unchanged: a teacher-forced logit trace per `docs/benchmarks.md` shows no disagreement at decided positions against the pre-change build, at forward width 1 behind a 2048-token context, for at least `qwen2.5-vl-3b-4bit` and `qwen3-vl-4b-4bit`.
- [ ] Image input still works end to end on `qwen2.5-vl-3b-4bit` and `qwen3-vl-4b-4bit`, since the vision path shares these attention bodies.
- [ ] `docs/benchmark_results/model_tests_m5max.md:90-109`, which currently calls this a standing optimization target with no curve, is updated with the outcome and the per-length numbers.
- [ ] `cargo test --release`, `cargo clippy --all-targets -- -D warnings` and `cargo fmt --check` pass.
## Verification
```bash
cargo build --release --features metal,accelerate
cargo test --release qwen2_vl
cargo test --release qwen3_vl
cargo clippy --all-targets -- -D warnings && cargo fmt --check
# mlxcel side, one run per prompt length, on an idle machine.
for N in 64 128 512 2048; do
./scripts/bench_decode.sh --cooldown 0 --big-cooldown 0 --prompt-tokens "$N" models/qwen2.5-vl-3b-4bit
done
# Output equivalence at decode width behind a realistic context, per docs/benchmarks.md.
cargo run --release --features metal,accelerate --example logit_trace -- \
models/qwen2.5-vl-3b-4bit trace_input.txt 1 4 8 2048 > after.tsv
python3 scripts/compare_logit_traces.py before.tsv after.tsv
# The vision path shares these attention bodies, so check it still produces the same description.
./target/release/mlxcel generate -m models/qwen2.5-vl-3b-4bit \
--image tests/fixtures/test_image.png -p "Describe this image." -n 50
```
A pass looks like: the 64 to 2048 decode falloff on `qwen2.5-vl-3b-4bit` sitting closer to mlx-lm's 9 to 11% than to the current 40%, zero disagreement at decided positions in the trace comparison, and the image run producing the same description as the pre-change build.
For the mlx-lm baseline, note that `scripts/bench_mlxlm.py` has no `--prompt-tokens` flag: the target length is the module constant `PROMPT_TOKENS` at `scripts/bench_mlxlm.py:56`, so set it per run. That environment also needs `torch`, `torchvision`, `timm` and `numba` installed beyond the `mlx-lm mlx-vlm` pair the harness comment mentions, or several checkpoints fail to load at all.
## Technical Considerations
Related but distinct: #1685 is also an M5 Max decode gap against mlx-lm, but its mechanism is a per-layer host materialization in the SSM mixers and its shape is a fixed additive per-token cost, not a context-scaling one. No overlap in files or fix.
Contributor guide
Research direction
Start with the attention paths in src/models/qwen2_vl.rs, src/models/qwen3_vl.rs, and src/models/qwen3_vl_moe.rs, comparing them with src/models/qwen3.rs and the mask guard in src/models/qwen3_5.rs. Run the existing metal4_attention_gqa_shape test and the 64- and 2048-token benchmarks before changing the three proposed paths. Done means the acceptance criteria pass, including the regenerated repeat_kv usage list and recorded profiles under docs/benchmark_results/.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100