Improve Qwen3.5 recurrent cache handling#323
Open
aleroot wants to merge 1 commit into
Open
Conversation
davidkoski
reviewed
Jun 1, 2026
davidkoski
reviewed
Jun 1, 2026
aleroot
force-pushed
the
qwen3_5perf
branch
2 times, most recently
from
d2c2851 to
c34d2d9
Compare
Store GatedDelta convolution state contiguously and advance array-cache metadata after each recurrent step, matching upstream mlx-lm behavior for Qwen3.5/Qwen3-Next style models. Keep left-padding masks active after recurrent cache state initialization, add coverage for ArraysCache metadata advancement, and align Qwen3 RoPE setup with the shared rope initializer.
Contributor
Author
|
Updated rebasing on latest main , also fixed a MTP tail-budget issue by making speculative rounds budget-aware: the iterator now only drafts when the remaining output budget can hold both the accepted draft prefix and the verifier correction/bonus. |
davidkoski
reviewed
Jun 15, 2026
Comment on lines
+116
to
+130
| fileprivate func prepareArrayMetadata(lengths: [Int]?) { | ||
| if let arrays = self as? ArraysCache { | ||
| arrays.prepare(lengths: lengths) | ||
| } else if let list = self as? CacheList { | ||
| list.prepare(lengths: lengths) | ||
| } | ||
| } | ||
|
|
||
| fileprivate func finalizeArrayMetadata() { | ||
| if let arrays = self as? ArraysCache { | ||
| arrays.finalize() | ||
| } else if let list = self as? CacheList { | ||
| list.finalize() | ||
| } | ||
| } |
Collaborator
There was a problem hiding this comment.
I wonder if this should be done through either methods on the KVCache protocol with default (empty) implementations? Or add a KVCacheLifecycle (or something) protocol. I worry that new cache types could get added (we already have quite a few) and this particular check wouldn't be updated.
davidkoski
reviewed
Jun 15, 2026
davidkoski
left a comment
davidkoski
left a comment
Collaborator
There was a problem hiding this comment.
Looks good. See what you think about my question on KVCache finalize.
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
2 participants
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.
Proposed changes
Qwen3.6 models using
model_type: qwen3_5rely heavily on hybrid linear-attention / GatedDelta layers. Keeping the convolution state contiguous avoids carrying strided slices through decode steps, which should reduce per-token cache overhead and better match upstream MLX Python behavior.Advancing array-cache metadata and preserving left-padding masks also improves stability for padded or batched generation paths.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes