Add opt-in Wired Memory

Libraries/MLXLMCommon/BatchGenerator.swift

@@ -4,6 +4,18 @@
 import Foundation
 import MLX
 
+public enum BatchGeneratorError: Error, CustomStringConvertible, Equatable {
+    case unsupportedCacheTopology(layer: Int, path: String, cacheType: String, reason: String)
+
+    public var description: String {
+        switch self {
+        case .unsupportedCacheTopology(let layer, let path, let cacheType, let reason):
+            return "Unsupported cache topology at layer \(layer), \(path): "
+                + "\(cacheType). \(reason)"
+        }
+    }
+}
+
 /// Continuous-batching engine.
 ///
 ///   1. `insert(prompts:)` queues new requests and returns their UIDs.
@@ -28,6 +40,7 @@ public final class BatchGenerator: @unchecked Sendable {
     private var unprocessed: [QueuedRequest] = []
     private var promptBatch: PromptProcessingBatch
     private var generationBatch: GenerationBatch?
+    private let cacheFactories: [BatchedCacheFactory]
 
     public private(set) var promptTokensProcessed: Int = 0
     public private(set) var generatedTokens: Int = 0
@@ -38,15 +51,19 @@ public final class BatchGenerator: @unchecked Sendable {
         defaultMaxTokens: Int = 128,
         prefillStepSize: Int = 2048,
         prefillBatchSize: Int = 8,
-        completionBatchSize: Int = 32
-    ) {
+        completionBatchSize: Int = 32,
+        cacheParameters: GenerateParameters? = nil
+    ) throws {
         self.model = model
         self.prefillStepSize = prefillStepSize
         self.prefillBatchSize = prefillBatchSize
         self.completionBatchSize = max(completionBatchSize, prefillBatchSize)
         self.defaultMaxTokens = defaultMaxTokens
         self.defaultEosTokens = eosTokens
         self.defaultSampler = greedySampler
+        self.cacheFactories = try Self.makeBatchedCacheFactories(
+            for: model.newCache(parameters: cacheParameters)
+        )
 
         if eosTokens.isEmpty {
             self.defaultStateMachine = SequenceStateMachine()
@@ -119,6 +136,68 @@ public final class BatchGenerator: @unchecked Sendable {
         return next()
     }
 
+    /// Drain all currently queued and active work.
+    ///
+    /// If `wiredMemoryTicket` is provided, the ticket is held for the full
+    /// drain loop rather than started and ended for each `next()` step. This
+    /// is best suited to bounded batches. Long-running servers that keep
+    /// inserting work should scope wired-memory tickets around bounded driver
+    /// windows instead.
+    ///
+    /// `BatchGenerator` remains a single-driver type while draining: do not
+    /// call `insert`, `cancel`, `next`, or `close` concurrently with this
+    /// method. Serialize server-side admission and driving through one owner,
+    /// such as an actor.
+    ///
+    /// If `onResponse` throws, or if the surrounding task is cancelled, the
+    /// error propagates and the generator is left partially drained at the
+    /// point where draining stopped. Responses delivered before the throw or
+    /// cancellation are not replayed or rolled back.
+    public func drain(
+        wiredMemoryTicket: WiredMemoryTicket? = nil,
+        onResponse: @escaping (GenerationBatchResponse) async throws -> Void
+    ) async throws {
+        if let wiredMemoryTicket {
+            try await Self.drainResponses(
+                hasWork: { self.hasWork },
+                next: { self.next() },
+                withScope: { body in
+                    try await wiredMemoryTicket.withWiredLimit(body)
+                },
+                onResponse: onResponse
+            )
+        } else {
+            try await Self.drainResponses(
+                hasWork: { self.hasWork },
+                next: { self.next() },
+                withScope: { body in
+                    try await body()
+                },
+                onResponse: onResponse
+            )
+        }
+    }
+
+    internal static func drainResponses(
+        hasWork: @escaping () -> Bool,
+        next: @escaping () -> [GenerationBatchResponse],
+        withScope: (@escaping () async throws -> Void) async throws -> Void,
+        checkCancellation: @escaping () throws -> Void = {
+            try Task.checkCancellation()
+        },
+        onResponse: @escaping (GenerationBatchResponse) async throws -> Void
+    ) async throws {
+        try await withScope {
+            while hasWork() {
+                try checkCancellation()
+                for response in next() {
+                    try checkCancellation()
+                    try await onResponse(response)
+                }
+            }
+        }
+    }
+
     public var hasWork: Bool {
         !unprocessed.isEmpty
             || (generationBatch?.isEmpty == false)
@@ -204,30 +283,92 @@ public final class BatchGenerator: @unchecked Sendable {
         }
     }
 
-    /// Allocate one batched cache per layer. The model's
-    /// `newCache(parameters:)` describes the per-layer cache topology
-    /// (`KVCacheSimple` / `RotatingKVCache` for full attention, `MambaCache`
-    /// or other `ArraysCache` subclasses for SSM-style layers); we build a
-    /// batched analog of each.
+    /// Allocate one batched cache per layer using the topology validated at init time.
     private func makeBatchedCache(batchSize B: Int) -> [any BatchedCache] {
-        let probe = model.newCache(parameters: nil)
         let zeroLeftPadding = Array(repeating: 0, count: B)
-        return probe.map { layer -> any BatchedCache in
-            if layer is MambaCache {
-                return MambaCache(leftPadding: zeroLeftPadding)
+        return cacheFactories.map { $0(zeroLeftPadding) }
+    }
+
+    private static func makeBatchedCacheFactories(
+        for probe: [any KVCache]
+    ) throws -> [BatchedCacheFactory] {
+        try probe.enumerated().map { layer, cache in
+            try makeBatchedCacheFactory(for: cache, layer: layer, path: "layer")
+        }
+    }
+
+    private static func makeBatchedCacheFactory(
+        for cache: any KVCache,
+        layer: Int,
+        path: String
+    ) throws -> BatchedCacheFactory {
+        let cacheType = String(describing: Swift.type(of: cache))
+
+        func unsupported(_ reason: String) -> BatchGeneratorError {
+            .unsupportedCacheTopology(
+                layer: layer,
+                path: path,
+                cacheType: cacheType,
+                reason: reason
+            )
+        }
+
+        if cache is QuantizedKVCache {
+            throw unsupported("Quantized KV caches are not supported by continuous batching.")
+        }
+
+        if cache is ChunkedKVCache {
+            throw unsupported("Chunked KV caches are not supported by continuous batching.")
+        }
+
+        if let cacheList = cache as? CacheList {
+            let childFactories = try cacheList.children.enumerated().map { childIndex, child in
+                try makeBatchedCacheFactory(
+                    for: child,
+                    layer: layer,
+                    path: "\(path).children[\(childIndex)]"
+                )
             }
-            if let arrays = layer as? ArraysCache {
-                return ArraysCache(size: arrays.slotCount, leftPadding: zeroLeftPadding)
+            return { leftPadding in
+                BatchedCacheList(caches: childFactories.map { $0(leftPadding) })
             }
-            if let rotating = layer as? RotatingKVCache, let maxSize = rotating.maxSize {
-                let keep = Int(rotating.metaState.first ?? "0") ?? 0
-                precondition(keep == 0, "RotatingKVCache with keep tokens is not supported")
-                return BatchRotatingKVCache(maxSize: maxSize, leftPadding: zeroLeftPadding)
+        }
+
+        // Exact-type matches avoid misclassifying subclasses such as
+        // MambaCache : ArraysCache and ChunkedKVCache : KVCacheSimple.
+        if Swift.type(of: cache) == MambaCache.self {
+            return { leftPadding in MambaCache(leftPadding: leftPadding) }
+        }
+
+        if Swift.type(of: cache) == ArraysCache.self, let arrays = cache as? ArraysCache {
+            let slotCount = arrays.slotCount
+            return { leftPadding in ArraysCache(size: slotCount, leftPadding: leftPadding) }
+        }
+
+        if let rotating = cache as? RotatingKVCache {
+            guard let maxSize = rotating.maxSize else {
+                throw unsupported("RotatingKVCache must have a non-nil maxSize.")
             }
-            return BatchKVCache(leftPadding: zeroLeftPadding)
+
+            let keep = Int(rotating.metaState.first ?? "0") ?? 0
+            guard keep == 0 else {
+                throw unsupported("RotatingKVCache with keep tokens is not supported.")
+            }
+
+            return { leftPadding in
+                BatchRotatingKVCache(maxSize: maxSize, leftPadding: leftPadding)
+            }
+        }
+
+        if Swift.type(of: cache) == KVCacheSimple.self {
+            return { leftPadding in BatchKVCache(leftPadding: leftPadding) }
         }
+
+        throw unsupported("No batched cache implementation exists for this cache type.")
     }
 
+    private typealias BatchedCacheFactory = (_ leftPadding: [Int]) -> any BatchedCache
+
     private struct QueuedRequest: Sendable {
         let uid: Int
         let tokens: [Int]

Libraries/MLXLMCommon/BatchKVCache.swift

@@ -28,6 +28,65 @@ public protocol BatchedCache: KVCache {
     func advanceBatched(_ n: Int)
 }
 
+/// Batched wrapper for composite per-layer caches. Some hybrid models keep
+/// multiple cache objects per logical layer, so batching has to preserve that
+/// nested topology instead of treating the composite as full attention.
+public final class BatchedCacheList: CacheList, BatchedCache {
+
+    private let batchedCaches: [any BatchedCache]
+
+    internal init(caches: [any BatchedCache]) {
+        self.batchedCaches = caches
+        super.init(caches: caches.map { $0 as any KVCache })
+    }
+
+    public func filterBatched(batchIndices: MLXArray) {
+        for cache in batchedCaches {
+            cache.filterBatched(batchIndices: batchIndices)
+        }
+    }
+
+    public func extendBatched(_ other: any BatchedCache) {
+        guard let other = other as? BatchedCacheList else {
+            preconditionFailure("BatchedCacheList.extendBatched requires another BatchedCacheList")
+        }
+        precondition(
+            batchedCaches.count == other.batchedCaches.count,
+            "Cannot extend BatchedCacheList with different child count"
+        )
+
+        for (a, b) in zip(batchedCaches, other.batchedCaches) {
+            a.extendBatched(b)
+        }
+    }
+
+    public func prepareBatched(leftPadding: [Int]?, lengths: [Int]?, rightPadding: [Int]?) {
+        for cache in batchedCaches {
+            cache.prepareBatched(
+                leftPadding: leftPadding,
+                lengths: lengths,
+                rightPadding: rightPadding
+            )
+        }
+    }
+
+    public func finalizeBatched() {
+        for cache in batchedCaches {
+            cache.finalizeBatched()
+        }
+    }
+
+    public func extractBatched(_ idx: Int) -> any KVCache {
+        CacheList(caches: batchedCaches.map { $0.extractBatched(idx) })
+    }
+
+    public func advanceBatched(_ n: Int) {
+        for cache in batchedCaches {
+            cache.advanceBatched(n)
+        }
+    }
+}
+
 /// Continuous-batching KV cache.
 ///
 /// Storage is right-justified along axis=2: for each row `b`, real keys

Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md

@@ -8,7 +8,7 @@ Use it when you need explicit control over request admission and per-row
 streaming responses:
 
 ```swift
-let generator = BatchGenerator(
+let generator = try BatchGenerator(
     model: model,
     eosTokens: [[eosToken]],
     defaultMaxTokens: 128,
@@ -30,6 +30,44 @@ context at a time. Each call to `next()` may prefill queued prompts, run one
 decode step for active rows, and return one response per active row. A response
 with a non-`nil` `finishReason` is the final response for that UID.
 
+Use `drain(wiredMemoryTicket:onResponse:)` when you want the generator to hold
+a wired-memory ticket across the whole drain loop:
+
+```swift
+let policy = WiredBudgetPolicy(baseBytes: baseBytes)
+let ticket = policy.ticket(size: estimatedBatchBytes)
+
+try await generator.drain(wiredMemoryTicket: ticket) { response in
+    print(response.uid, response.token, response.finishReason as Any)
+}
+```
+
+The ticket stays active until the generator has no queued or active rows, or
+until the response handler throws or the task is cancelled. In those cases, the
+error propagates, the ticket is released, and the generator keeps any remaining
+state. Responses delivered before the throw or cancellation are not replayed or
+rolled back. Do not call `insert`, `cancel`, `next`, or `close` concurrently
+with `drain`; serialize server-side admission and driving through one owner,
+such as an actor.
+
+For long-running services that keep inserting work, wrap shorter driver windows
+directly in `ticket.withWiredLimit` so one never-ending drain does not block
+other wired-memory admission waiters:
+
+```swift
+while shouldContinueServing {
+    try await waitForMoreWork()
+
+    try await ticket.withWiredLimit {
+        while generator.hasWork {
+            for response in generator.next() {
+                try await send(response)
+            }
+        }
+    }
+}
+```
+
 Call `cancel(uid:)` to remove a queued or active row. The method returns `true`
 when it found the UID and filtered that row out of the generator state.
 
@@ -50,6 +88,9 @@ request.
 ## Cache Model
 
 Continuous batching uses `BatchKVCache` for attention layers and `ArraysCache`
-or `MambaCache` for array-backed state-space layers. `BatchKVCache` keeps rows
-right-aligned so requests with different prompt lengths can share one cache and
-attention mask while still preserving per-row positions for RoPE.
+or `MambaCache` for array-backed state-space layers. Composite `CacheList`
+layers are preserved when every child cache has a supported batched analog.
+Unsupported cache topologies fail during `BatchGenerator` initialization.
+`BatchKVCache` keeps rows right-aligned so requests with different prompt
+lengths can share one cache and attention mask while still preserving per-row
+positions for RoPE.