Repository navigation
Aligned weight loading and ring KV cache for sliding-window layers (1.2.1) - #7
Merged
Merged
Conversation
…w layers (1.2.1) MLX writes safetensors without the 16-byte alignment the kernels' vector loads need, so every tensor was copied into its own Metal buffer while the mapped file stayed referenced by the command buffer: the GPU held two copies of the weights, and a 9 GB Gemma 4 12B ran out of GPU memory on a 16 GB Mac. Shards are now copied once, in file order, into a page-aligned buffer and the mapping is released. Loading is faster (12B: 156 s -> 25 s from disk, 11 s from the page cache). Sliding-window layers keep a ring of window + 512 K/V rows (position p at row p mod rows) instead of the whole context: the 12B at 8K context needs 0.5 GB of cache instead of 2.7 GB and runs at 32K. Prefix reuse is refused when the previous run went on for more than a batch past the divergence point, since those rows are gone. A GPU out-of-memory failure is reported as such, with advice; the load-time budget includes the scratch buffers. README and the landing page: load-time and no-copy wording, Gemma rows in the performance tables. Existing models: greedy output token-identical; Qwen3 0.6B and Mistral 7B decode within 1% (3 alternating pairs); Gemma 3 1B and Gemma 4 E2B unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
🚀 Deploying Preview to Cloudflare 🚀Preview Deployments by commit
|
The model tracks its own KV write high-water mark, so prefix reuse accounts for writes the engine's bookkeeping dropped (a cancelled prefill, rejected speculative drafts, a failed in-flight step); the engine resets it when a run starts from position 0. The ring size is derived from the one maxBatch constant and checked at load. GPU out-of-memory is recognised by the MTLCommandBufferError code, with the text match as fallback. The load-time scratch estimate is Transformer.scratchBytes(config:), and a test checks it covers what init allocates on every installed checkpoint. flash_prefill computes the ring slot per row. Prefill and decode A/B (3 alternating pairs): Mistral 7B and Gemma 3 1B within 1% of the previous build; prefill unchanged everywhere. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
A single arena for all of a shard's tensors measured 6% slower at 1K context on Qwen3 0.6B (3 alternating pairs, repeatedly); separate page-aligned buffers, the layout the engine always had, are within 1%. The mapping is still released after the copy, so the GPU holds one copy. Co-Authored-By: Claude Fable 5.1 <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.
Summary
Fixes "GPU error: Insufficient Memory" when running Gemma 4 12B (8.96 GB) in the app on a 16 GB Mac, and makes long contexts cheap on every Gemma.
safetensorsare never 16-byte aligned (8-byte header padding; Gemma 4's 2-bytelayer_scalartensors shift everything after them), soSafeTensorFilecopied every tensor into its own buffer while the mapped file stayed referenced by the command buffer — the GPU working set held the weights twice. Shards are now copied once, in file order, into a page-aligned arena and the mapping is released. Load times: Gemma 4 12B 156 s → 25 s from disk / 11–16 s from the page cache; Mistral 7B 3.3 → 2.9 s.MTLCommandBufferErrorcode) with advice to lower the context length; the load-time budget usesTransformer.scratchBytes(config:), which a test keeps above whatinitallocates.Reviewed with
/code-review(6 findings, all addressed in the second commit).Test plan
swift test: 135 tests — decode and split decode over a ring cache, a 1,200-token Gemma 3 1B prompt prefilled vs decoded incrementally with the ring active, the prefix-reuse and write-tracking rules, the scratch estimate on every installed checkpoint🤖 Generated with Claude Code