Repository navigation
Reconcile the fork onto upstream v0.32.2 (two carried deltas) - #4
Merged
Merged
Conversation
BatchKVCache.update_and_fetch advances `offset` as an mx.array (`self.offset = self.offset + n`) and never forces it to evaluate. On decode steps that build no mask nothing else consumes it, so its lazy graph grows by one node per token for the life of the cache (a memory leak on long runs). Tying `offset` and `left_padding` into `keys` with mx.depends evaluates them with every step. Re-carried from 11648d2 onto v0.32.2. The surrounding code moved: the method now returns through `keys_and_values()` (upstream ml-explore#1778), so the mx.depends line sits just before that return. v0.32.2 still has no equivalent; upstream's attempts at this fix (ml-explore#1731, ml-explore#1780, ml-explore#1872) were closed unmerged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With left padding (batch > 1), `mx.take_along_axis(mask, topk_indices)` gathers the padded attention mask by the indexer's top-k positions and mis-selects entries, so padded rows attend to the wrong keys. Rebuild the mask from the gathered indices against each row's left padding instead, and drop it when no row is padded. Re-carried from 587b88a onto v0.32.2, where the surrounding code is unchanged and the buggy take_along_axis line is still present; upstream ml-explore#1552 proposing a fix remains open. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 57c4b0ef5c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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.
Reconciles the Foxlight fork onto upstream v0.32.2. The base branch
foxlight/v0.32.2-baseis the tag itself. It carries the fork's two remaining deltas; neither has an upstream equivalent in v0.32.2. The CVE-2026-5843 gate that the fork once carried is upstream now.update_and_fetchadvancesoffsetas anmx.arrayand never forces it to evaluate. On decode steps that build no mask nothing consumes it, so its graph grows one node per token. This is re-carried from 11648d2. Upstream Make "state" of cache return full state ml-explore/mlx-lm#1778 moved the return tokeys_and_values(), so themx.dependsline now sits just before it. Upstream attempts (Guard batch cache metadata against lazy-graph accumulation ml-explore/mlx-lm#1731, Periodically evaluate cache state during decode (#1662) ml-explore/mlx-lm#1780, Bound ArraysCache metadata evaluation during batched decoding ml-explore/mlx-lm#1872) were closed unmerged.take_along_axis(mask, topk_indices)mis-selects mask entries. The fix rebuilds the mask from the gathered indices against each row's padding. Re-carried unchanged from 587b88a. The code path runs only forL == 1decode, so dropping the mask when no row is padded is safe, because a single query's causal mask is all-true. Upstream Fix deepseek_v32 Indexer evicting attention sinks from sparse top-k ml-explore/mlx-lm#1552 is still open.Evidence (Apple Silicon test node, mlx 0.32.4)
The leak reproduces on the published
mlx-lm==0.32.2and is gone with this branch. The measurement is a decode-onlyBatchKVCacheloop that evaluates each step's keys and values and nothing else:offsetTests on this branch:
tests/test_prompt_cache.py,test_models.py,test_gated_delta.pyandtest_generate.pygive 185 passed, 1 skipped, 102 subtests passed.Follow-up outside this repo
Skulk's
fork-wheel.ymlmust stamp0.32.2.post1through setuptools-scm, because upstream ml-explore#1916 generates_version.py. That change, the wheel publish and Skulk's coordinated MLX pin move land in Skulk.