feat(paged): port PagedKVCache + BlockAllocator foundation - #36
Closed
devin-ai-integration[bot] wants to merge 2 commits into
Closed
feat(paged): port PagedKVCache + BlockAllocator foundation#36devin-ai-integration[bot] wants to merge 2 commits into
devin-ai-integration[bot] wants to merge 2 commits into
Conversation
Port the paged-attention data-structure layer from ekryski/mlx-swift-lm PR #97, adapted to this repo's KVCache surface (KVCacheSimple reference, no peek/memoryBytes/storageKind): - BlockAllocator: refcounted free-list block pool with retain() for prefix sharing. - PagedKVCache: per-layer block-based KV storage (block layout matches vLLM's MetalPagedKVCache) with a scatter -> gather -> SDPA forward path. - Tests: allocator semantics + forward-equivalence vs KVCacheSimple. No model is wired to this cache and there is no Metal paged kernel yet; gather()+SDPA is a foundation path, not a production decode path.
Author
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Comment on lines
+17
to
+45
| name: PagedKVCache Swift tests | ||
| runs-on: blacksmith-12vcpu-macos-26 | ||
| steps: | ||
| - uses: actions/checkout@v4 | ||
|
|
||
| - name: Build tests | ||
| run: swift build --build-tests | ||
|
|
||
| - name: Extract and colocate mlx.metallib | ||
| run: | | ||
| python3 -m venv /tmp/mlxvenv | ||
| /tmp/mlxvenv/bin/pip install 'mlx==0.31.1' | ||
| pyver="$(/tmp/mlxvenv/bin/python -c 'import sys;print(f"{sys.version_info.major}.{sys.version_info.minor}")')" | ||
| metallib="/tmp/mlxvenv/lib/python${pyver}/site-packages/mlx/lib/mlx.metallib" | ||
| test -f "$metallib" || { echo "::error::mlx.metallib not found at $metallib"; exit 1; } | ||
| # MLX's C++ loader looks for mlx.metallib next to the running binary. | ||
| # Under `swift test` that is the xctest runner inside the bundle, so | ||
| # place a copy there as well as in .build/debug. | ||
| cp "$metallib" .build/debug/mlx.metallib | ||
| for bundle in .build/debug/*PackageTests.xctest; do | ||
| macos="$bundle/Contents/MacOS" | ||
| if [ -d "$macos" ]; then | ||
| cp "$metallib" "$macos/mlx.metallib" | ||
| echo "placed metallib in $macos" | ||
| fi | ||
| done | ||
|
|
||
| - name: Run PagedKVCache tests | ||
| run: swift test --filter PagedKVCache |
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
Ports the PagedAttention foundation layer from ekryski/mlx-swift-lm#97 into
MLXLMCommon, adapted to this repo'sKVCachesurface. This is the data-structure layer only — there is no Metal paged-attention kernel (still unsolved in Swift; tracked as ekryski issue ml-explore#127) and no model is wired to this cache.What's added
Libraries/MLXLMCommon/BlockAllocator.swift— refcounted, lock-guarded physical block pool (LIFO free list).retain()bumps refcounts so multiple sequences can share blocks (prefix sharing / copy-on-write hook). ThrowsAllocatorError.exhaustedon pool exhaustion.Libraries/MLXLMCommon/PagedKVCache.swift— per-layer block KV storage subclassingBaseKVCache. Block layout matches vLLM'sMetalPagedKVCacheexactly so a future Metal kernel can read it directly:scatter()(write new tokens into blocks atoffset) →gather()(reconstruct contiguous[1, kvHeads, T, headDim], trimming the partial last block) → SDPA.update()composes the two and is element-wise equivalent toKVCacheSimple.update().Tests/MLXLMTests/PagedKVCacheTests.swift— allocator semantics (allocate/free/retain/exhaustion) + round-trip identity + forward-equivalence vsKVCacheSimpleover a varying-length chunk sequence (7,1,1,4,1,1) that crosses block boundaries.Adaptation notes (vs ekryski source)
KVCacheSimple(this repo) rather thanStandardKVCache.peek()/memoryBytes/storageKindoverrides — those members don't exist on this repo'sKVCache/BaseKVCache.copy()implementation (required by this repo'sKVCacheprotocol;BaseKVCache's defaultfatalErrors).Scope / non-goals
retain()hook).Validation
Could not build/test in this environment (Linux, no Swift toolchain; MLX needs Metal/Apple Silicon). Needs
swift test --filter PagedKVCacheon a Mac. Note: theBuild and Testworkflow jobs are gatedif: github.repository == 'ml-explore/mlx-swift-lm', so CI does not run on this fork.Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changesLink to Devin session: https://app.devin.ai/sessions/95935953320b49d1a083d0385740ae2c