Sitelet https://github.com/mlx-node/mlx-node/pull/80
Skip to content

fix(qwen3.5-VL): correct image inference (interleaved M-RoPE, patch-embed, compressed-position decode + delta lifetime) - #80

Merged
Brooooooklyn merged 7 commits into
mainfrom
fix/qwen35-vision-mrope-interleaved
Jun 27, 2026
Merged

Brooooooklyn merged 7 commits into
mainfrom
fix/qwen35-vision-mrope-interleaved

Conversation

@Brooooooklyn

@Brooooooklyn Brooooooklyn commented Jun 27, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Fixes Qwen3.5-VL image inference for both the qwen3_5 (dense) and
qwen3_5_moe families. Before this branch, an image turn ran end-to-end but
produced garbage descriptions — the model could not read text and hallucinated
unrelated subjects (e.g. described a bank-reconciliation table as "dark grey
stone / white fabric"). Text inference and mlx convert were unaffected; only
the visual-feature path was wrong.

Found via ornith-1.0-35b (a Qwen3.5-VL-MoE-35B-A3B post-train). Root-caused by
numeric bisection against the reference mlx-vlm implementation.

Root causes & fixes (all verified vs mlx-vlm)

  1. Interleaved multimodal RoPE for image tokens (fb95c6d8)
    Qwen3.5-VL uses the interleaved (stride-3 per-freq) M-RoPE selector, not
    the PaddleOCR sectioned/contiguous one. Text tokens (t==h==w) are
    identical either way — which is exactly why text always worked while images
    (t≠h≠w) got the wrong rotation axis on ~17% of frequencies, every attention
    layer, destroying the 2D positional structure.

  2. Patch embedding: sum Conv3d temporal slices + load patch_embed.proj.bias (6a889bdb)
    The patch embed previously used only the first temporal slice and hard-coded
    the conv bias to None.

  3. Decode rotates at the compressed M-RoPE position (b5c4420e)
    Image prefill compresses ~754 placeholder tokens into ~29 M-RoPE positions, so
    rope_deltas = max_position + 1 − seq_len (≈ −726). Decode must rotate the
    query at physical_slot + rope_deltas while K/V still writes at the physical
    slot. Threaded a rope_position_offset: i32 through
    forward_paged / decoder_layer / paged_forward (dense + MoE share one
    forward_paged); the negative delta is carried on VisionMerge.rope_deltas
    and stored in cached_rope_deltas at VLM prefill. The physical position is
    cast u32 → i32 before the negative add to avoid underflow.

  4. rope-delta lifetime gate (f7ae00ca)
    Found by adversarial review of (3). The cross-turn delta was reset only when
    cached_prefix_len == 0, but the paged-turn planner also yields
    cached_prefix_len > 0 with continued_live_prefix == false on a non-live
    prefix-cache hit
    . Because the model instance is shared across all sessions,
    a stale negative delta from a prior image turn could then leak into an
    unrelated text-only request that merely shares a cached text prefix —
    rotating that text at physical + stale_delta → garbage.

    Factored the decision into rope_delta_for_paged_turn(cached_rope_deltas, continued_live_prefix) (keep the delta iff it is a live image continuation,
    else clear) and wired it through all six paged reset gates (dense + MoE ×
    sync / stream / engine). This is safe because image requests prefill with
    skip_lookup and never publish a hashable text stream that collides with
    their expanded-placeholder blocks, so every non-live hit restores only
    pure-text prefix blocks (delta 0); only a live continuation re-attends the
    image's compressed-position K/V. Added three model-free lifecycle regression
    tests.

  5. e2e correctness gate (ab044b88)
    qwen3_5_moe_vl_reads_document_text and qwen3_5_vl_image_chat.rs (dense) —
    #[ignore], env-gated (MLX_TEST_QWEN35MOE_VL_MODEL_PATH +
    MLX_TEST_VLM_IMAGE_PATH), asserting the model reads ≥2 DOC_KEYWORDS from
    examples/ocr.png. These are the ground-truth gates for VL image inference.

Note on the addmm commit (12e89b3a + 8e6d69d9)

12e89b3a replaced the fused mlx_array_addmm primitive with an explicit
matmul + add, originally attributed to a bug in mlx::core::addmm. That
premise was wrong
— PyPI MLX's addmm applies the C term correctly
(maxdiff 0) and the FFI wrapper passes its arguments correctly. The real cause
was a corrupt local metallib that miscompiled the fused GEMM kernels (the same
bad build also miscompiled the NAX gemm). The explicit form is kept as
robustness against this project's documented non-deterministic metallib
corruption — it is correctness-equivalent and vision/bias-only so the perf cost
is negligible — and the nn::linear C-application tests double as a build
canary. 8e6d69d9 corrects the misleading comment so it no longer claims an mlx
source bug.

Validation

  • e2e READS-DOC on the mxfp8-converted ornith-1.0-35b checkpoint: the model
    reads the document, matching all 7 keywords
    (reconciliation, bank, council, trunch, october, 2019, balance).
  • nn::linear addmm/bias tests, paged_forward rope-offset + delta-lifetime
    tests, and the broader rope/m-rope/delta unit sweep (70 tests) pass.
  • cargo clippy -p mlx-core --all-targets -D warnings clean; cargo fmt clean.
  • Adversarial review of the branch: approve, no material findings.

The full cargo test -p mlx-core --lib (debug) shows pre-existing Metal-f32
failures (banded-attn / attention-vjp) plus a few cross-family numerical tests
that fail only against the corrupted local debug metallib; all of them pass
in release against a correctly-built metallib (and reproduce on main), so the
branch introduces no new failures.

Test plan

# unit (release, correct metallib)
cargo test --release -p mlx-core --lib -- nn::linear paged_forward rope delta

# e2e image-correctness (env-gated, sequential — large models)
MLX_TEST_QWEN35MOE_VL_MODEL_PATH=<converted-qwen3.5-vl-moe> \
MLX_TEST_VLM_IMAGE_PATH=$PWD/examples/ocr.png \
  cargo test --release -p mlx-core --test qwen3_5_moe_vl_image_chat -- \
  --ignored --nocapture qwen3_5_moe_vl_reads_document_text

🤖 Generated with Claude Code


Note

High Risk
Changes core attention RoPE, vision weights, and shared paged inference state across dense/MoE VL; incorrect delta lifetime or offset math would corrupt multi-turn or cached-prefix text, though extensive unit and env-gated e2e tests mitigate this.

Overview
Fixes Qwen3.5-VL image inference (dense and MoE) by aligning vision, RoPE, and paged decode with mlx-vlm.

Interleaved M-RoPE — Adds apply_multimodal_rotary_pos_emb_interleaved and switches Qwen3.5 attention from the PaddleOCR sectioned apply so image tokens get stride-3 per-frequency axis selection; text-only paths stay unchanged via invariance tests.

Vision tower — addmm is implemented as explicit matmul + scaled add (avoids local metallib fused-GEMM corruption that dropped biases). Patch embed loads optional conv bias and collapses Conv3d weights by summing temporal slices instead of taking slice 0.

Compressed M-RoPE decode — VisionMerge carries rope_deltas; get_rope_index uses the global max over t/h/w axes. Paged prefill/decode/MTP thread cached_rope_deltas through paged_rope_offset and rope_position_offset so rotation uses compressed positions while KV stays at physical slots. rope_delta_for_paged_turn keeps the delta only on continued_live_prefix so stale image deltas do not leak into unrelated text prefix-cache hits.

Tests — Unit tests for interleaved RoPE, rope offset/delta lifecycle, patch embed, linear bias; env-gated e2e reads_document_text gates for dense and MoE VL.

Reviewed by Cursor Bugbot for commit 8e6d69d. Bugbot is set up for automated code reviews on this repo. Configure here.

Brooooooklyn and others added 7 commits June 26, 2026 16:10
The existing VL image-chat e2e tests (`*_t0_capture`) only assert paged==flat
byte-identity, which passes even when the shared vision path produces garbage
features (the model describes the ocr.png financial table as "stone/fabric").

Add a correctness gate per file — `qwen3_5_moe_vl_reads_document_text` (MoE,
primary) and `qwen3_5_vl_reads_document_text` (dense, shared vision path). Each
sends ONE image+prompt turn at T=0 ("Transcribe the text in this document...")
with max_new_tokens=512 past the small thinking budget, then asserts the output
contains >=2 of the document keywords (reconciliation/bank/council/trunch/
october/2019/balance), printing the full model output on failure.

This is the TDD ground-truth gate for the qwen3.5 VL image fix; it is
`#[ignore]`-gated on MLX_TEST_QWEN35{,MOE}_VL_MODEL_PATH + MLX_TEST_VLM_IMAGE_PATH
(image defaults to examples/ocr.png) and never runs under a plain `cargo test`.
Reuses the existing cfg/user_msg/resolve_image_path helpers; no new infra.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
qwen3_5 / qwen3_5_moe image inference rotated Q/K with PaddleOCR-VL's
sectioned (contiguous-chunk) multimodal-RoPE selector instead of
qwen3_5's interleaved (stride-3) selector. For image tokens (where
temporal != height != width) this assigns ~17% of frequencies to the
wrong spatial axis, producing garbage visual features. Text tokens were
unaffected because temporal == height == width makes the three cos/sin
axis rows bit-identical, so any selector yields the same angles — which
is exactly why text worked and images did not.

Add `apply_multimodal_rotary_pos_emb_interleaved`, which builds the
per-frequency axis selector matching mlx-vlm's
`_interleaved_position_selector` (height: idx 1,4,7,… up to
section[1]*3; width: idx 2,5,8,… up to section[2]*3; temporal
otherwise), mirrors it across the doubled cos/sin, and gathers the
selected t/h/w axis per frequency via `take_along_axis` before the same
rotate_half + partial-rotary tail. Switch qwen3_5 attention's two M-RoPE
call sites (forward + forward_paged) to it; this covers qwen3_5_moe too
(shared `Qwen3_5Attention`). PaddleOCR-VL keeps the sectioned path
(unchanged; its production forward uses the C++ sectioned kernel).

Tests (no model): interleaved selection correctness against the
hand-computed selector, and a text-invariance hard gate asserting the
interleaved apply is bit-identical to the sectioned apply when
t==h==w (proving the text path cannot regress).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The qwen3_5 shared vision patch-embed loader collapsed the 5D Conv3d
weight [out, kD=2, kH, kW, in] by taking only temporal slice 0 and never
loaded the Conv bias. The image processor duplicates the static frame
across the temporal axis (mlx-vlm qwen3_vl Conv3d(bias=True), kD=2), so
the effective 2D kernel is the SUM of the temporal slices plus the bias.

Sum over the temporal axis (robust to kD != 2) and plumb the optional
patch_embed.proj.bias through set_patch_embed / PatchEmbedding::new into
the Conv2d bias (was hardcoded None). get_parameters now round-trips the
bias. Shared path (qwen3_5 / qwen3_5_moe / Qwen3.6-35B-A3B), not forked.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
`MxArray::addmm` called `mlx_array_addmm`, but the vendored `mlx::core::addmm`
in this build returns only `alpha*(A@B)` and silently drops `beta*C`. That was
invisible for the bias-free LM linears (Q/K/V/O/MLP and the bias-free MoE router
gate all pass a zero C), but it corrupted every biased linear — most visibly the
Qwen3.5-VL vision tower, whose qkv/proj/fc1/fc2 and merger projections all carry
a bias. Dropping those biases produced semantically wrong image features (the
model described the ocr.png financial table as "stone/fabric").

Compute the result explicitly (matmul, optional alpha scale, then add beta*C) so
the C term is actually applied. Numeric bisection vs mlx-vlm on converted bf16
ornith: vision-tower image_embeds per-token cosine median 0.9996 after the fix
(block0 rel-err 35x -> 0.0078). Add nn::linear unit tests proving addmm applies
a [4]/[1,4]/[2,4] C and Linear::forward applies its bias.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
An image run compresses its placeholder tokens into fewer M-RoPE positions,
so image prefill rotates Q/K at the compressed positions and records
rope_deltas = max_position + 1 - seq_len (negative). Decode and warm
continuation must keep rotating at that compressed position
(physical_slot + rope_deltas) while K/V is still written at the physical
slot. The scalar-offset RoPE path rotated at the raw physical token count,
ignoring rope_deltas, so post-image decode ran ~725 positions off and
produced repetition garbage.

Carry rope_deltas on VisionMerge, store it as cached_rope_deltas at every
VLM prefill, and thread a rope_position_offset (physical position +
cached_rope_deltas, cast u32->i32 before the negative add) through
forward_paged, the decoder-layer paged forward, and the paged decode/prefill
drivers for both qwen3_5 and qwen3_5_moe. get_rope_index takes the global
max over the (t,h,w) axes for the delta. Text turns store no delta, so the
offset equals the physical position and non-VLM paths stay byte-identical.

Matches paddleocr_vl (cache_offset + rope_deltas) and mlx-vlm
(base_offset + rope_delta).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…che hits

The cross-turn M-RoPE delta an image prefill bakes into the shared
`cached_rope_deltas` was reset only when `cached_prefix_len == 0`. But the
paged turn planner also produces `cached_prefix_len > 0` with
`continued_live_prefix == false` on a non-live prefix-cache hit: a later
text request that merely shares a cached text prefix with an earlier image
request (the model is shared across all sessions). There the old gate did
not fire, so a stale negative delta survived and rotated unrelated text at
`physical_slot + stale_delta` instead of the raw physical position.

Only a live continuation (`continued_live_prefix`) re-attends the image's
physically-resident compressed-position K/V, so only it needs the delta:
image requests prefill with `skip_lookup` and never publish a text stream
that collides with their expanded-placeholder blocks, so every non-live hit
restores pure-text prefix blocks (delta 0).

Factor the decision into `rope_delta_for_paged_turn(current,
continued_live_prefix)` and wire all six paged gates (dense + MoE sync,
stream, engine paths) through it. Add model-free lifecycle regression tests
covering the live-continuation, cold-start, and stale-delta-on-text-hit
cases.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… mlx bug

The earlier addmm rewrite attributed the dropped `beta*C` to a bug in
`mlx::core::addmm`. That premise was wrong: PyPI MLX's addmm applies `C`
correctly and the FFI wrapper passes its args correctly. The real cause was a
corrupt local metallib that miscompiled the fused GEMM kernels (the same bad
build also miscompiled the NAX gemm). The explicit matmul+add form is kept as
robustness against this project's documented non-deterministic metallib
corruption, and the nn::linear C-application tests double as a build canary —
but the code comment no longer claims an mlx source bug.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@coderabbitai

coderabbitai Bot commented Jun 27, 2026

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: d6962a69-9c42-4c42-8cc5-58035b22e444

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/qwen35-vision-mrope-interleaved

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@Brooooooklyn Brooooooklyn added the model-e2e Run the heavy Model E2E workflow (per-family real-checkpoint tests) on this PR label Jun 27, 2026
@Brooooooklyn
Brooooooklyn merged commit 855b91c into main Jun 27, 2026
16 checks passed
@Brooooooklyn
Brooooooklyn deleted the fix/qwen35-vision-mrope-interleaved branch June 27, 2026 02:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

model-e2e Run the heavy Model E2E workflow (per-family real-checkpoint tests) on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant