Repository navigation
Merge upstream main (#603 model cache extraction) - #68
Merged
Merged
Conversation
TokenIterator.next() cleared the buffer cache when tokenCount % 256 == 0, but checked after incrementing tokenCount, so the first clear came at token 256. mlx-lm checks before incrementing and also clears on the first token, right after prefill. A generation shorter than 256 tokens never cleared, so repeated short requests grew the cache without bound. Co-authored-by: Joshua Langberg <negativetime@me.com>
…n-flight load (ml-explore#603) * Move the model cache into its own file - Moves the constraint kind, the tokenizer bias, and the cache actor out of the model file, changing only the access the cross-file use requires: `private actor ModelCache` becomes `actor ModelCache`, because a `private` top-level type is visible only inside its own file. No other content changes. The model cache grew large enough inside MLXLanguageModel.swift to warrant its own file, separate from the model's own responsibilities. * Catch only the fork-unavailable error when a constraint clone fails - Narrows the freshly-compiled constraint's clone attempt from a blanket `try?` to a catch on `GrammarError.forkFailed`, matching the cached-template clone site above it, which already narrowed this way. Every other clone failure now reaches the caller instead of being silently treated as the expected one. A blanket `try?` treated every clone failure as the one expected case, xgrammar lacking `Fork()`. `GrammarConstraint.clone()` has exactly one throw site today, so this narrowing changes no observable behavior yet; it only stops a future second failure mode from being swallowed the same way. * Cancel every registered load from evictAll() - Cancels every registered load task, which remove(modelID:) already did for one model. - Corrects the eviction documentation, which claimed an in-flight load completes. - Adds a test that evictAll() cancels an in-flight load, and the same probe against evict(). - Bounds both new tests with a time limit, because each parks its loader and a regression would hang. evictAll() dropped every load registration without cancelling, so nothing in the process could reach a load that was still running. A caller awaiting preload() or respond() now fails, and callers coalesced onto one load fail together. Cancellation is cooperative. A load already inside loadWeights runs to completion and allocates, and that limit applies to evict() too. * Rename the grammar tokenizer members away from an xgrammar abbreviation - Renames the grammar tokenizer members and their storage away from an abbreviation of the library name: `xgTokenizers` becomes `grammarTokenizers`, `makeXGTokenizer` becomes `makeGrammarTokenizer`, and `hasCachedXGTokenizer` becomes `hasCachedGrammarTokenizer`, at every call site. - Renames the same local variables at every call site: `xgTokenizer` and `xgTok` become `grammarTokenizer`, and the `vocab` local inside `makeGrammarTokenizer` becomes `vocabulary` to avoid shadowing the `vocab:` parameter label once its neighbor is renamed. The cache's tokenizer members were named after xgrammar's own abbreviation rather than the type they store, which the style guide's initialism rule flags. * Name the download progress throughput property in full - Renames the observable throughput property to throughputBytesPerSecond, dropping the abbreviation. - Drops public from the two reporting entry points, leaving the observable state and the shared instance public. Every caller of the two entry points is in this module, and the tests reach them through @testable import, so neither needs to be public. No release contains this type: the newest tag is 3.31.4 from June, and the whole MLXFoundationModels library arrived in July, so no consumer can depend on either name. * Track download progress for each model - Keys the progress state on the model ID, with one published value per model and accessors that add every entry. - Reports the end of a load from a defer, so a failed or cancelled load clears its own entry. - Sends every report through one stream and one main-actor consumer, which applies them in the order the loader sent them. - Gives every load its own identifier, so a report from a cancelled load cannot change the load that replaced it. - Removes isActive and modelName, and adds activeDownloads and download(forModelID:). One fraction and one identifier cannot describe two loads, and the cache loads models concurrently, so two downloads overwrote each other's state. reportCompleted() ran only after a successful load, so a cancelled load left the state active for the life of the process. Each report also created its own main-actor task, and those tasks run in no fixed order, so a progress report could apply after the end and leave a finished download on screen. Eviction cancels a load and forgets it at once, so a replacement can start while the cancelled load still transfers bytes, and both shared one entry. No release contains this type, so none of these changes can break a consumer.
solderzzc
added a commit
to SharpAI/SwiftLM
that referenced
this pull request
Sep 29, 2026
Picks up SharpAI/mlx-swift-lm#72: rejected MTP drafts are rolled back layer by layer, including sliding-window layers after the ring wraps (#184). The range also has SharpAI/mlx-swift-lm#68 (upstream #603 sync, not used by SwiftLM). README: replace the "avoid --mtp" warning with the measured result. Rollback is exact; the remaining differences from plain decoding are batch-shape bf16 numerics and the text-only vs vision model path. Fixes #184 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
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.
Proposed changes
Merges
ml-explore/mlx-swift-lmmain (ee673d6) into our main with a real merge commit, so later upstream merges start from a shared base.Upstream brings two commits:
ModelCache.swift.evictAll()now cancels in-flight loads instead of abandoning them. Download progress is tracked per load.The merge had no conflicts.
MLXLanguageModel.swiftis the only file both sides changed, and ourcase .chunk(let text, _):(forGeneration.chunk(String, tokenId:)) is kept. SwiftLM does not useMLXFoundationModels.Testing
xcodebuild test(serial): 1,702 tests, 11 failures. All 11 also fail on our main (7cc37a0), so none come from this merge. Most are numeric tolerances. One is aWiredMemory"Ticket not active" assertion that looks like the bug upstream mlx-swift#471 fixes.swift build -c releaseand--build-testssucceed with this commit as the submodule.MLXFoundationModelsis behindcanImport(FoundationModels, _version: 2), which needs the macOS 27 SDK. So the Extract the model cache and stop whole-cache eviction abandoning an in-flight load ml-explore/mlx-swift-lm#603 code and its new tests (ModelCacheEvictionTests,MLXDownloadProgressTests) did not compile or run here.ModelCache.swiftandMLXDownloadProgress.swiftwere type-checked separately against our modules, with the gate removed.Found during review (not fixed here)
case .chunk(let text):, for exampleGuidedGenerationBenchmarkTests.swift:427,Gemma4EseriesLoadIntegrationTests.swift:60,StaleWeightIndexIntegrationTests.swift:52andMTPIteratorEndToEndDiagnosticTests.swift. That does not compile against our two-valueGeneration.chunk. CI never builds that project. Separate fix..startedcan overwrite a newer load's progress entry for the same model, which hides the newer load's download progress. To be reported upstream.Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changesAI usage
accurately describes the code changes.
🤖 Generated with Claude Code