Repository navigation
Conversation
… mRoPE VLMs sparse_prefill computes selected_indices + position_offset for every token it keeps, then consumes those positions through one mechanism: a _PositionMappedRoPE installed on attn.rope, behind hasattr(first_attn, "rope"). mlx_vlm's mRoPE attention has no attn.rope. It holds a rotary_emb and rotates from an explicit position_ids argument. So no wrapper was installed, no positions were passed, and the language model derived them from the cache offset instead — a contiguous run. selected_positions was computed and then discarded, and every selected token after the first was written at the wrong position. On a cold-cache request, a 3,122-token selection out of 15,635 conversation tokens was prefilled at positions 0..3,121 instead of across the prompt it was selected from. Decode then continued from a cache offset short of the true timeline by M - N. Greedy output changed. Pass the positions sparse_prefill already computed through the adapter's explicit-position seam, laid out by the same per-model rules the offset-derived paths use. Give the decode adjustment to the request's mRoPE delta, which expresses the identical "cache offset + delta" contract on the carrier those models decode from, and give it the prefill's lifetime: recorded only when the forward was really driven from supplied positions, cleared in cleanup_rope, and reset on every path that abandons a sparse prefill. Models with a wrappable attn.rope are unchanged. A restored prefix only moves where the wrong run starts; the cold-cache case has no cache at all.
Formatting only, and only on lines this PR added. The file's pre-existing Black findings are left alone so the diff carries no unrelated churn. The parsed AST is unchanged.
…requeue reset Cut the new tests to three: selected positions with and without a cached prefix, the adapter passing explicit position_ids through, and the rope wrapper path staying unchanged. Drop the rope_deltas reset in _requeue_or_fail_prefill and its test. A request whose sparse prefill succeeded skips external and chunked prefill, so it never reaches the requeue path; the test only passed with hand-set state. Drop the reset on the abort path too, since the request is gone. Shorten the comments.
This was referenced Oct 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
sparse_prefill()computesselected_positions = selected_indices + position_offset, but only uses them through the_PositionMappedRoPEwrapper, which it installs when the attention module has a.rope. mlx_vlm's mRoPE attention (Qwen3.5 here) has arotary_emband takesposition_idsinstead. So no wrapper is installed, no positions are passed, and the model falls back to positions from the cache offset. The selection is written at dense positionsoffset..offset+N-1. Every selected token after the first gap in the selection ends up at the wrong position.This hits text-only requests to a served mRoPE VLM with a SpecPrefill draft loaded. Image requests never reach sparse prefill because
_try_specprefill_scoringin the scheduler returns early whenvlm_inputs_embedsis set.Change
VLMModelAdapter.position_ids_for_absolute()lays out a caller-computed position row with the same per-model rules the offset-derived paths use (moved into_shape_position_ids()). It returns None when the model is not mRoPE.VLMModelAdapter.__call__passes an explicitposition_idsthrough instead of rebuilding it from the cache offset.sparse_prefill()passesselected_positionsper chunk when there is no.ropeto wrap.request.rope_deltas(minus the kickoff token). Decode already reads that throughregister_rope_deltaas cache offset + delta. It is reset if sparse prefill fails and falls back to dense prefill.Models with a
.ropekeep the wrapper path as before.Measured
On a served Qwen3.5 VLM, cold cache, 3,122 tokens selected out of 15,635 conversation tokens:
With a restored 12,288-token prefix, chunk 0 went from 12,288..14,335 (dense) to 12,288..15,039. Across six requests in one session, 30 of 30 multi-token chunks were dense before. After, 28 of 29 spanned more positions than tokens; the other one is the trailing window chunk selection always keeps (#2442). Greedy output for the same prompt and selection changed.
Only Qwen3.5 was measured. Other mRoPE families in mlx-vlm (qwen2_vl, qwen2_5_vl, qwen3_vl, glm4v and others) go through the same gate, so I expect the same bug and the same fix there, but I have not run them.
Tests
Three tests in
tests/test_vlm_specprefill.py: positions reaching the forward with offset 0 and 12,288, the adapter passing explicitposition_idsthrough, and a.ropemodel still using the wrapper. All three fail on main; the position test shows the dense run: