feat(paged): port PagedKVCache + BlockAllocator foundation

.github/workflows/paged-kv-cache-test.yml

@@ -0,0 +1,45 @@
+name: PagedKVCache Tests
+
+# Fork-runnable Swift test job (the main Build and Test workflow is gated to
+# ml-explore/mlx-swift-lm). Runs the PagedKVCache + BlockAllocator tests on a
+# Blacksmith macOS runner so MLX has a Metal device.
+on:
+  pull_request:
+    paths:
+      - "Libraries/MLXLMCommon/PagedKVCache.swift"
+      - "Libraries/MLXLMCommon/BlockAllocator.swift"
+      - "Tests/MLXLMTests/PagedKVCacheTests.swift"
+      - ".github/workflows/paged-kv-cache-test.yml"
+  workflow_dispatch:
+
+jobs:
+  paged-kv-cache:
+    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

Libraries/MLXLMCommon/BlockAllocator.swift

@@ -0,0 +1,113 @@
+// Copyright © 2026 Eigen Labs.
+//
+// Ported from ekryski/mlx-swift-lm PR #97 (PagedKVCache + BlockAllocator
+// foundation). Data-structure layer for paged attention.
+
+import Foundation
+
+/// Shared physical block pool for `PagedKVCache`. Refcounted free list —
+/// `retain` lets multiple sequences share blocks (prefix caching).
+///
+/// **Thread safety:** all public methods serialize on an internal lock.
+public final class BlockAllocator: @unchecked Sendable {
+
+    public let numBlocks: Int
+
+    /// Refcount per block. 0 = free; >0 = referenced by N sequences.
+    private var refcounts: [Int]
+
+    /// LIFO free list — recently freed blocks reuse first for cache locality.
+    private var freeStack: [Int]
+
+    private let lock = NSLock()
+
+    public init(numBlocks: Int) {
+        precondition(numBlocks > 0, "numBlocks must be > 0")
+        self.numBlocks = numBlocks
+        self.refcounts = Array(repeating: 0, count: numBlocks)
+        // Reverse so id 0 pops first — predictable in tests.
+        self.freeStack = (0 ..< numBlocks).reversed()
+    }
+
+    /// Allocate `count` fresh blocks. Throws if pool exhausted.
+    public func allocate(_ count: Int) throws -> [Int] {
+        lock.lock()
+        defer { lock.unlock() }
+        guard count <= freeStack.count else {
+            throw AllocatorError.exhausted(requested: count, available: freeStack.count)
+        }
+        var allocated: [Int] = []
+        allocated.reserveCapacity(count)
+        for _ in 0 ..< count {
+            let id = freeStack.removeLast()
+            refcounts[id] += 1
+            allocated.append(id)
+        }
+        return allocated
+    }
+
+    /// Decrement refcount; return to free list when it hits 0.
+    public func free(_ id: Int) {
+        lock.lock()
+        defer { lock.unlock() }
+        precondition(id >= 0 && id < numBlocks, "block id out of range")
+        precondition(refcounts[id] > 0, "double-free of block \(id)")
+        refcounts[id] -= 1
+        if refcounts[id] == 0 {
+            freeStack.append(id)
+        }
+    }
+
+    public func free(_ ids: [Int]) {
+        lock.lock()
+        defer { lock.unlock() }
+        for id in ids {
+            precondition(id >= 0 && id < numBlocks, "block id out of range")
+            precondition(refcounts[id] > 0, "double-free of block \(id)")
+            refcounts[id] -= 1
+            if refcounts[id] == 0 {
+                freeStack.append(id)
+            }
+        }
+    }
+
+    /// Increment refcount on already-allocated blocks (prefix sharing).
+    public func retain(_ ids: [Int]) {
+        lock.lock()
+        defer { lock.unlock() }
+        for id in ids {
+            precondition(id >= 0 && id < numBlocks, "block id out of range")
+            precondition(refcounts[id] >= 1, "retain on free block \(id)")
+            refcounts[id] += 1
+        }
+    }
+
+    public var freeCount: Int {
+        lock.lock()
+        defer { lock.unlock() }
+        return freeStack.count
+    }
+
+    public var allocatedCount: Int {
+        lock.lock()
+        defer { lock.unlock() }
+        return numBlocks - freeStack.count
+    }
+
+    public func refcount(of id: Int) -> Int {
+        lock.lock()
+        defer { lock.unlock() }
+        return refcounts[id]
+    }
+}
+
+public enum AllocatorError: Error, CustomStringConvertible {
+    case exhausted(requested: Int, available: Int)
+
+    public var description: String {
+        switch self {
+        case let .exhausted(req, avail):
+            return "BlockAllocator exhausted: requested \(req), \(avail) free"
+        }
+    }
+}

Libraries/MLXLMCommon/PagedKVCache.swift

@@ -0,0 +1,178 @@
+// Copyright © 2026 Eigen Labs.
+//
+// Ported from ekryski/mlx-swift-lm PR #97 (PagedKVCache + BlockAllocator
+// foundation) and adapted to this repo's `KVCache` surface.
+
+import Foundation
+import MLX
+
+/// Per-layer paged KV cache. One instance per transformer layer.
+///
+/// Stores K and V in fixed-size blocks of `blockSize` tokens. Sequences
+/// address blocks via a `blockTable` populated externally from a shared
+/// `BlockAllocator`. The forward path here gathers blocks into a contiguous
+/// `[1, kvHeads, T, headDim]` view and hands it to MLX SDPA — a Metal paged
+/// kernel will replace `gather()` for the production decode path.
+///
+/// Block layout matches vLLM's MetalPagedKVCache exactly so a future kernel
+/// port can read this storage directly:
+///   keyBlocks:   [numBlocks, blockSize, numKVHeads, headDim]
+///   valueBlocks: [numBlocks, blockSize, numKVHeads, headDim]
+public class PagedKVCache: BaseKVCache {
+
+    public let numBlocks: Int
+    public let blockSize: Int
+    public let numKVHeads: Int
+    public let headDim: Int
+
+    public internal(set) var keyBlocks: MLXArray
+    public internal(set) var valueBlocks: MLXArray
+
+    /// Block ids assigned to this cache, in logical token order.
+    /// One request per cache instance — multi-tenancy is handled upstream.
+    public internal(set) var blockTable: [Int]
+
+    public override var maxSize: Int? { numBlocks * blockSize }
+
+    public init(
+        numBlocks: Int,
+        blockSize: Int = 16,
+        numKVHeads: Int,
+        headDim: Int,
+        dtype: DType = .bfloat16
+    ) {
+        precondition(numBlocks > 0, "numBlocks must be > 0")
+        precondition(blockSize > 0, "blockSize must be > 0")
+        precondition(numKVHeads > 0, "numKVHeads must be > 0")
+        precondition(headDim > 0, "headDim must be > 0")
+
+        self.numBlocks = numBlocks
+        self.blockSize = blockSize
+        self.numKVHeads = numKVHeads
+        self.headDim = headDim
+
+        let shape = [numBlocks, blockSize, numKVHeads, headDim]
+        self.keyBlocks = MLXArray.zeros(shape, dtype: dtype)
+        self.valueBlocks = MLXArray.zeros(shape, dtype: dtype)
+        self.blockTable = []
+
+        super.init()
+    }
+
+    /// Append block ids from the allocator to this cache's table.
+    public func appendBlocks(_ ids: [Int]) {
+        precondition(ids.allSatisfy { $0 >= 0 && $0 < numBlocks }, "block id out of range")
+        blockTable.append(contentsOf: ids)
+    }
+
+    /// Drop the last `n` block table entries. Block reclamation is the
+    /// allocator's responsibility — this only severs the link.
+    public func dropLastBlocks(_ n: Int) {
+        let drop = Swift.min(n, blockTable.count)
+        blockTable.removeLast(drop)
+    }
+
+    /// Scatter `[1, kvHeads, T, headDim]` K/V into block storage starting at
+    /// the current `offset`. Caller must have already appended enough blocks
+    /// to cover `offset + T` positions.
+    public func scatter(keys: MLXArray, values: MLXArray) {
+        precondition(keys.dim(0) == 1, "scatter requires B=1")
+        precondition(keys.dim(1) == numKVHeads, "kv_heads mismatch on scatter")
+        precondition(keys.dim(3) == headDim, "head_dim mismatch on scatter")
+
+        let numNew = keys.dim(2)
+        let startTok = offset
+
+        for i in 0 ..< numNew {
+            let absTok = startTok + i
+            let blockIdx = absTok / blockSize
+            let slotIdx = absTok % blockSize
+            precondition(
+                blockIdx < blockTable.count,
+                "scatter past end of block table — caller must allocate blocks first")
+            let physicalBlock = blockTable[blockIdx]
+
+            let kSlice = keys[0..., 0..., i ..< (i + 1), 0...]
+            let vSlice = values[0..., 0..., i ..< (i + 1), 0...]
+
+            // Block layout is [block_size, kv_heads, head_dim] per block;
+            // squeeze input from [1, kv_heads, 1, head_dim] to match.
+            keyBlocks[physicalBlock, slotIdx, 0..., 0...] = kSlice.squeezed(axis: 0).squeezed(axis: 1)
+            valueBlocks[physicalBlock, slotIdx, 0..., 0...] =
+                vSlice.squeezed(axis: 0).squeezed(axis: 1)
+        }
+
+        offset += numNew
+    }
+
+    /// Reconstruct contiguous `[1, kvHeads, offset, headDim]` K/V from the
+    /// blocks listed in `blockTable`. Output matches what `KVCacheSimple`
+    /// would return for the same input sequence.
+    public func gather() -> (MLXArray, MLXArray) {
+        let tokenCount = offset
+        let numUsedBlocks = (tokenCount + blockSize - 1) / blockSize
+        precondition(numUsedBlocks <= blockTable.count, "block table too short for offset")
+
+        let physicalIds = MLXArray(blockTable.prefix(numUsedBlocks).map { Int32($0) })
+
+        let kGathered = keyBlocks[physicalIds, 0..., 0..., 0...]
+        let vGathered = valueBlocks[physicalIds, 0..., 0..., 0...]
+
+        let totalSlots = numUsedBlocks * blockSize
+        let kFlat = kGathered.reshaped([1, totalSlots, numKVHeads, headDim])
+        let vFlat = vGathered.reshaped([1, totalSlots, numKVHeads, headDim])
+
+        let kTransposed = kFlat.transposed(0, 2, 1, 3)
+        let vTransposed = vFlat.transposed(0, 2, 1, 3)
+
+        // Trim padding from the partially-filled last block.
+        let kTrimmed = kTransposed[0..., 0..., ..<tokenCount, 0...]
+        let vTrimmed = vTransposed[0..., 0..., ..<tokenCount, 0...]
+
+        return (kTrimmed, vTrimmed)
+    }
+
+    public override func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) {
+        scatter(keys: keys, values: values)
+        return gather()
+    }
+
+    public override var state: [MLXArray] {
+        get { [keyBlocks, valueBlocks] }
+        set {
+            guard newValue.count == 2 else {
+                fatalError("PagedKVCache state requires 2 arrays (key_blocks, value_blocks)")
+            }
+            keyBlocks = newValue[0]
+            valueBlocks = newValue[1]
+        }
+    }
+
+    public override func innerState() -> [MLXArray] {
+        [keyBlocks, valueBlocks]
+    }
+
+    public override var isTrimmable: Bool { true }
+
+    @discardableResult
+    public override func trim(_ n: Int) -> Int {
+        let trimCount = Swift.min(n, offset)
+        offset -= trimCount
+        return trimCount
+    }
+
+    public override func copy() -> any KVCache {
+        let new = PagedKVCache(
+            numBlocks: numBlocks,
+            blockSize: blockSize,
+            numKVHeads: numKVHeads,
+            headDim: headDim,
+            dtype: keyBlocks.dtype
+        )
+        new.keyBlocks = keyBlocks[.ellipsis]
+        new.valueBlocks = valueBlocks[.ellipsis]
+        new.blockTable = blockTable
+        new.offset = offset
+        return new
+    }
+}