From 5904c40a8c611a6baae531642264d2f20fb61eaa Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 2 May 2026 16:26:34 -0700 Subject: [PATCH 1/5] Fix BatchKVCache fall through --- Libraries/MLXLMCommon/BatchGenerator.swift | 113 +++++++++++-- Libraries/MLXLMCommon/BatchKVCache.swift | 67 ++++++++ .../Documentation.docc/continuous-batching.md | 11 +- .../MLXLMTests/ContinuousBatchingTests.swift | 159 +++++++++++++++++- 4 files changed, 322 insertions(+), 28 deletions(-) diff --git a/Libraries/MLXLMCommon/BatchGenerator.swift b/Libraries/MLXLMCommon/BatchGenerator.swift index 49798e52b..4039511a8 100644 --- a/Libraries/MLXLMCommon/BatchGenerator.swift +++ b/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,8 +51,9 @@ 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 @@ -47,6 +61,9 @@ public final class BatchGenerator: @unchecked Sendable { 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() @@ -204,30 +221,90 @@ 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)]" + ) + } + return { leftPadding in + BatchedCacheList(caches: childFactories.map { $0(leftPadding) }) + } + } + + 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.") } - if let arrays = layer as? ArraysCache { - return ArraysCache(size: arrays.slotCount, leftPadding: zeroLeftPadding) + + let keep = Int(rotating.metaState.first ?? "0") ?? 0 + guard keep == 0 else { + throw unsupported("RotatingKVCache with keep tokens is not supported.") } - 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) + + return { leftPadding in + BatchRotatingKVCache(maxSize: maxSize, leftPadding: leftPadding) } - return BatchKVCache(leftPadding: zeroLeftPadding) } + + 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] diff --git a/Libraries/MLXLMCommon/BatchKVCache.swift b/Libraries/MLXLMCommon/BatchKVCache.swift index d120c6dd5..3d58a9f02 100644 --- a/Libraries/MLXLMCommon/BatchKVCache.swift +++ b/Libraries/MLXLMCommon/BatchKVCache.swift @@ -28,6 +28,73 @@ 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 { + + internal init(caches: [any BatchedCache]) { + super.init(caches: caches.map { $0 as any KVCache }) + } + + private var batchedChildren: [any BatchedCache] { + children.map { child in + guard let batched = child as? any BatchedCache else { + preconditionFailure("BatchedCacheList contains a non-batched child cache") + } + return batched + } + } + + public func filterBatched(batchIndices: MLXArray) { + for cache in batchedChildren { + cache.filterBatched(batchIndices: batchIndices) + } + } + + public func extendBatched(_ other: any BatchedCache) { + guard let other = other as? BatchedCacheList else { + preconditionFailure("BatchedCacheList.extendBatched requires another BatchedCacheList") + } + let lhs = batchedChildren + let rhs = other.batchedChildren + precondition( + lhs.count == rhs.count, + "Cannot extend BatchedCacheList with different child count" + ) + + for (a, b) in zip(lhs, rhs) { + a.extendBatched(b) + } + } + + public func prepareBatched(leftPadding: [Int]?, lengths: [Int]?, rightPadding: [Int]?) { + for cache in batchedChildren { + cache.prepareBatched( + leftPadding: leftPadding, + lengths: lengths, + rightPadding: rightPadding + ) + } + } + + public func finalizeBatched() { + for cache in batchedChildren { + cache.finalizeBatched() + } + } + + public func extractBatched(_ idx: Int) -> any KVCache { + CacheList(caches: batchedChildren.map { $0.extractBatched(idx) }) + } + + public func advanceBatched(_ n: Int) { + for cache in batchedChildren { + cache.advanceBatched(n) + } + } +} + /// Continuous-batching KV cache. /// /// Storage is right-justified along axis=2: for each row `b`, real keys diff --git a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md b/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md index 66767e13a..e8a84b602 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md +++ b/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, @@ -50,6 +50,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. diff --git a/Tests/MLXLMTests/ContinuousBatchingTests.swift b/Tests/MLXLMTests/ContinuousBatchingTests.swift index d37d892be..a8cccac7e 100644 --- a/Tests/MLXLMTests/ContinuousBatchingTests.swift +++ b/Tests/MLXLMTests/ContinuousBatchingTests.swift @@ -145,8 +145,94 @@ final class ContinuousBatchingTests: XCTestCase { } } - func testBatchGeneratorAdmitsQueuedRowsAndReportsFinishReasons() { - let generator = BatchGenerator( + func testBatchGeneratorAcceptsSupportedCacheTopologies() throws { + _ = try BatchGenerator( + model: CacheTopologyLanguageModel { _ in [KVCacheSimple()] } + ) + + _ = try BatchGenerator( + model: CacheTopologyLanguageModel { _ in [ArraysCache(size: 3)] } + ) + + _ = try BatchGenerator( + model: CacheTopologyLanguageModel { _ in [MambaCache()] } + ) + + _ = try BatchGenerator( + model: CacheTopologyLanguageModel { _ in [RotatingKVCache(maxSize: 8, keep: 0)] } + ) + + _ = try BatchGenerator( + model: CacheTopologyLanguageModel { _ in [CacheList(MambaCache(), KVCacheSimple())] } + ) + } + + func testBatchGeneratorRejectsUnsupportedCacheTopologies() { + assertBatchGeneratorRejectsCache( + QuantizedKVCache(), + expectedType: "QuantizedKVCache", + expectedPath: "layer" + ) + assertBatchGeneratorRejectsCache( + ChunkedKVCache(), + expectedType: "ChunkedKVCache", + expectedPath: "layer" + ) + assertBatchGeneratorRejectsCache( + RotatingKVCache(maxSize: 8, keep: 4), + expectedType: "RotatingKVCache", + expectedPath: "layer" + ) + assertBatchGeneratorRejectsCache( + CacheList(MambaCache(), QuantizedKVCache()), + expectedType: "QuantizedKVCache", + expectedPath: "layer.children[1]" + ) + } + + func testBatchGeneratorPassesCacheParametersToModel() throws { + let model = CacheTopologyLanguageModel { _ in [KVCacheSimple()] } + + _ = try BatchGenerator( + model: model, + cacheParameters: GenerateParameters(maxKVSize: 17) + ) + + XCTAssertEqual(model.receivedParameters?.maxKVSize, 17) + } + + func testBatchGeneratorRejectsUnsupportedCacheParameters() { + XCTAssertThrowsError( + try BatchGenerator( + model: IncrementingLanguageModel(), + cacheParameters: GenerateParameters(maxKVSize: 17) + ) + ) { error in + guard + case BatchGeneratorError.unsupportedCacheTopology( + _, + let + path, + let + cacheType, + let + reason + ) = error + else { + XCTFail( + "Expected BatchGeneratorError.unsupportedCacheTopology, got \(error)" + ) + return + } + + XCTAssertEqual(path, "layer") + XCTAssertEqual(cacheType, "RotatingKVCache") + XCTAssertTrue(reason.contains("keep tokens")) + } + } + + func testBatchGeneratorAdmitsQueuedRowsAndReportsFinishReasons() throws { + let generator = try BatchGenerator( model: IncrementingLanguageModel(), eosTokens: [[5]], defaultMaxTokens: 4, @@ -181,8 +267,8 @@ final class ContinuousBatchingTests: XCTestCase { XCTAssertFalse(generator.hasWork) } - func testBatchGeneratorCancelRemovesQueuedRequest() { - let generator = BatchGenerator( + func testBatchGeneratorCancelRemovesQueuedRequest() throws { + let generator = try BatchGenerator( model: IncrementingLanguageModel(), defaultMaxTokens: 3, prefillBatchSize: 1, @@ -206,8 +292,8 @@ final class ContinuousBatchingTests: XCTestCase { XCTAssertEqual(seenUIDs, [uids[0]]) } - func testBatchGeneratorCancelRemovesActiveRequest() { - let generator = BatchGenerator( + func testBatchGeneratorCancelRemovesActiveRequest() throws { + let generator = try BatchGenerator( model: IncrementingLanguageModel(), defaultMaxTokens: 4, prefillBatchSize: 2, @@ -235,6 +321,45 @@ final class ContinuousBatchingTests: XCTestCase { } } +private func assertBatchGeneratorRejectsCache( + _ cache: any KVCache, + expectedType: String, + expectedPath: String, + file: StaticString = #filePath, + line: UInt = #line +) { + let model = CacheTopologyLanguageModel { _ in [cache] } + + XCTAssertThrowsError( + try BatchGenerator(model: model), + file: file, + line: line + ) { error in + guard + case BatchGeneratorError.unsupportedCacheTopology( + _, + let + path, + let + cacheType, + let + reason + ) = error + else { + XCTFail( + "Expected BatchGeneratorError.unsupportedCacheTopology, got \(error)", + file: file, + line: line + ) + return + } + + XCTAssertEqual(path, expectedPath, file: file, line: line) + XCTAssertEqual(cacheType, expectedType, file: file, line: line) + XCTAssertFalse(reason.isEmpty, file: file, line: line) + } +} + private func makeCache(keys: [Float], values: [Float]) -> KVCacheSimple { let cache = KVCacheSimple() _ = cache.update( @@ -296,3 +421,25 @@ private final class IncrementingLanguageModel: Module, LanguageModel, KVCacheDim return MLXArray(logits).reshaped([batchSize, sequenceLength, vocabularySize]) } } + +private final class CacheTopologyLanguageModel: Module, LanguageModel { + private let cacheFactory: (GenerateParameters?) -> [any KVCache] + private(set) var receivedParameters: GenerateParameters? + + init(_ cacheFactory: @escaping (GenerateParameters?) -> [any KVCache]) { + self.cacheFactory = cacheFactory + } + + func prepare(_ input: LMInput, cache: [any KVCache], windowSize: Int?) throws -> PrepareResult { + .tokens(input.text) + } + + func callAsFunction(_ inputs: MLXArray, cache: [any KVCache]?) -> MLXArray { + fatalError("CacheTopologyLanguageModel is only used for cache topology tests") + } + + func newCache(parameters: GenerateParameters?) -> [any KVCache] { + receivedParameters = parameters + return cacheFactory(parameters) + } +} From 60c8705c074565bcd6109e3608b339f4425c2e7d Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 2 May 2026 16:32:44 -0700 Subject: [PATCH 2/5] Simplify BatchKVCache --- Libraries/MLXLMCommon/BatchGenerator.swift | 2 ++ Libraries/MLXLMCommon/BatchKVCache.swift | 28 +++++++------------ .../MLXLMTests/ContinuousBatchingTests.swift | 28 ++++++++++--------- 3 files changed, 27 insertions(+), 31 deletions(-) diff --git a/Libraries/MLXLMCommon/BatchGenerator.swift b/Libraries/MLXLMCommon/BatchGenerator.swift index 4039511a8..526511a1a 100644 --- a/Libraries/MLXLMCommon/BatchGenerator.swift +++ b/Libraries/MLXLMCommon/BatchGenerator.swift @@ -272,6 +272,8 @@ public final class BatchGenerator: @unchecked Sendable { } } + // 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) } } diff --git a/Libraries/MLXLMCommon/BatchKVCache.swift b/Libraries/MLXLMCommon/BatchKVCache.swift index 3d58a9f02..ce0f1d3de 100644 --- a/Libraries/MLXLMCommon/BatchKVCache.swift +++ b/Libraries/MLXLMCommon/BatchKVCache.swift @@ -33,21 +33,15 @@ public protocol BatchedCache: KVCache { /// 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 }) } - private var batchedChildren: [any BatchedCache] { - children.map { child in - guard let batched = child as? any BatchedCache else { - preconditionFailure("BatchedCacheList contains a non-batched child cache") - } - return batched - } - } - public func filterBatched(batchIndices: MLXArray) { - for cache in batchedChildren { + for cache in batchedCaches { cache.filterBatched(batchIndices: batchIndices) } } @@ -56,20 +50,18 @@ public final class BatchedCacheList: CacheList, BatchedCache { guard let other = other as? BatchedCacheList else { preconditionFailure("BatchedCacheList.extendBatched requires another BatchedCacheList") } - let lhs = batchedChildren - let rhs = other.batchedChildren precondition( - lhs.count == rhs.count, + batchedCaches.count == other.batchedCaches.count, "Cannot extend BatchedCacheList with different child count" ) - for (a, b) in zip(lhs, rhs) { + for (a, b) in zip(batchedCaches, other.batchedCaches) { a.extendBatched(b) } } public func prepareBatched(leftPadding: [Int]?, lengths: [Int]?, rightPadding: [Int]?) { - for cache in batchedChildren { + for cache in batchedCaches { cache.prepareBatched( leftPadding: leftPadding, lengths: lengths, @@ -79,17 +71,17 @@ public final class BatchedCacheList: CacheList, BatchedCache { } public func finalizeBatched() { - for cache in batchedChildren { + for cache in batchedCaches { cache.finalizeBatched() } } public func extractBatched(_ idx: Int) -> any KVCache { - CacheList(caches: batchedChildren.map { $0.extractBatched(idx) }) + CacheList(caches: batchedCaches.map { $0.extractBatched(idx) }) } public func advanceBatched(_ n: Int) { - for cache in batchedChildren { + for cache in batchedCaches { cache.advanceBatched(n) } } diff --git a/Tests/MLXLMTests/ContinuousBatchingTests.swift b/Tests/MLXLMTests/ContinuousBatchingTests.swift index a8cccac7e..431782b8d 100644 --- a/Tests/MLXLMTests/ContinuousBatchingTests.swift +++ b/Tests/MLXLMTests/ContinuousBatchingTests.swift @@ -186,7 +186,7 @@ final class ContinuousBatchingTests: XCTestCase { assertBatchGeneratorRejectsCache( CacheList(MambaCache(), QuantizedKVCache()), expectedType: "QuantizedKVCache", - expectedPath: "layer.children[1]" + expectedPathContains: "children" ) } @@ -324,7 +324,8 @@ final class ContinuousBatchingTests: XCTestCase { private func assertBatchGeneratorRejectsCache( _ cache: any KVCache, expectedType: String, - expectedPath: String, + expectedPath: String? = nil, + expectedPathContains: String? = nil, file: StaticString = #filePath, line: UInt = #line ) { @@ -335,16 +336,12 @@ private func assertBatchGeneratorRejectsCache( file: file, line: line ) { error in - guard - case BatchGeneratorError.unsupportedCacheTopology( - _, - let - path, - let - cacheType, - let - reason - ) = error + guard case let BatchGeneratorError.unsupportedCacheTopology( + _, + path, + cacheType, + reason + ) = error else { XCTFail( "Expected BatchGeneratorError.unsupportedCacheTopology, got \(error)", @@ -354,7 +351,12 @@ private func assertBatchGeneratorRejectsCache( return } - XCTAssertEqual(path, expectedPath, file: file, line: line) + if let expectedPath { + XCTAssertEqual(path, expectedPath, file: file, line: line) + } + if let expectedPathContains { + XCTAssertTrue(path.contains(expectedPathContains), file: file, line: line) + } XCTAssertEqual(cacheType, expectedType, file: file, line: line) XCTAssertFalse(reason.isEmpty, file: file, line: line) } From 76e1fa05c5b728db0379dfbee987839ed8a7f18b Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 2 May 2026 19:01:34 -0700 Subject: [PATCH 3/5] Add WiredMemoryTicket --- Libraries/MLXLMCommon/BatchGenerator.swift | 61 +++++++++++++ .../Documentation.docc/continuous-batching.md | 37 ++++++++ .../MLXLMTests/ContinuousBatchingTests.swift | 85 ++++++++++++++++++- 3 files changed, 182 insertions(+), 1 deletion(-) diff --git a/Libraries/MLXLMCommon/BatchGenerator.swift b/Libraries/MLXLMCommon/BatchGenerator.swift index 526511a1a..587e9059e 100644 --- a/Libraries/MLXLMCommon/BatchGenerator.swift +++ b/Libraries/MLXLMCommon/BatchGenerator.swift @@ -136,6 +136,67 @@ 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. + 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) diff --git a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md b/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md index e8a84b602..5fb6ef069 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md +++ b/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md @@ -30,6 +30,43 @@ 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. 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. diff --git a/Tests/MLXLMTests/ContinuousBatchingTests.swift b/Tests/MLXLMTests/ContinuousBatchingTests.swift index 431782b8d..7752f9f43 100644 --- a/Tests/MLXLMTests/ContinuousBatchingTests.swift +++ b/Tests/MLXLMTests/ContinuousBatchingTests.swift @@ -1,6 +1,6 @@ import Foundation import MLX -import MLXLMCommon +@testable import MLXLMCommon import MLXNN import XCTest @@ -145,6 +145,77 @@ final class ContinuousBatchingTests: XCTestCase { } } + func testDrainResponsesHelperScopesWrapperAroundFullLoop() async throws { + var batches = [ + [makeResponse(uid: 1, token: 10)], + [ + makeResponse(uid: 2, token: 20), + makeResponse(uid: 3, token: 30), + ], + ] + var events: [String] = [] + + try await BatchGenerator.drainResponses( + hasWork: { !batches.isEmpty }, + next: { + events.append("next") + return batches.removeFirst() + }, + withScope: { body in + events.append("enter") + defer { events.append("exit") } + try await body() + }, + onResponse: { response in + events.append("response:\(response.uid)") + } + ) + + XCTAssertEqual( + events, + [ + "enter", + "next", + "response:1", + "next", + "response:2", + "response:3", + "exit", + ] + ) + } + + func testDrainResponsesHelperChecksCancellationBeforeNextStep() async { + var cancellationChecks = 0 + var didCallNext = false + + do { + try await BatchGenerator.drainResponses( + hasWork: { true }, + next: { + didCallNext = true + return [] + }, + withScope: { body in + try await body() + }, + checkCancellation: { + cancellationChecks += 1 + throw CancellationError() + }, + onResponse: { _ in + XCTFail("Cancellation should stop before response handling.") + } + ) + XCTFail("Expected cancellation to stop the drain loop.") + } catch is CancellationError { + XCTAssertEqual(cancellationChecks, 1) + XCTAssertFalse(didCallNext) + } catch { + XCTFail("Expected CancellationError, got \(error).") + } + } + func testBatchGeneratorAcceptsSupportedCacheTopologies() throws { _ = try BatchGenerator( model: CacheTopologyLanguageModel { _ in [KVCacheSimple()] } @@ -362,6 +433,18 @@ private func assertBatchGeneratorRejectsCache( } } +private func makeResponse(uid: Int, token: Int) -> GenerationBatchResponse { + GenerationBatchResponse( + uid: uid, + token: token, + finishReason: nil, + matchedSequence: nil, + currentState: nil, + allTokens: nil, + promptCache: nil + ) +} + private func makeCache(keys: [Float], values: [Float]) -> KVCacheSimple { let cache = KVCacheSimple() _ = cache.update( From f6298308969ccc6e06b67a83f6795853eb57a12b Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 2 May 2026 19:09:46 -0700 Subject: [PATCH 4/5] Update docs --- Libraries/MLXLMCommon/BatchGenerator.swift | 3 ++- .../MLXLMCommon/Documentation.docc/continuous-batching.md | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/Libraries/MLXLMCommon/BatchGenerator.swift b/Libraries/MLXLMCommon/BatchGenerator.swift index 587e9059e..042a331f1 100644 --- a/Libraries/MLXLMCommon/BatchGenerator.swift +++ b/Libraries/MLXLMCommon/BatchGenerator.swift @@ -151,7 +151,8 @@ public final class BatchGenerator: @unchecked Sendable { /// /// 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. + /// 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 diff --git a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md b/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md index 5fb6ef069..f99f27d91 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md +++ b/Libraries/MLXLMCommon/Documentation.docc/continuous-batching.md @@ -45,9 +45,10 @@ try await generator.drain(wiredMemoryTicket: ticket) { response in 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. Do not call `insert`, `cancel`, `next`, or `close` concurrently with -`drain`; serialize server-side admission and driving through one owner, such as -an actor. +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 From f2c2b5b2a2953a907c099d470c3a4c8035cd1a98 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 2 May 2026 19:10:10 -0700 Subject: [PATCH 5/5] swift lint --- .../MLXLMTests/ContinuousBatchingTests.swift | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/Tests/MLXLMTests/ContinuousBatchingTests.swift b/Tests/MLXLMTests/ContinuousBatchingTests.swift index 7752f9f43..c07db7ce5 100644 --- a/Tests/MLXLMTests/ContinuousBatchingTests.swift +++ b/Tests/MLXLMTests/ContinuousBatchingTests.swift @@ -1,9 +1,10 @@ import Foundation import MLX -@testable import MLXLMCommon import MLXNN import XCTest +@testable import MLXLMCommon + final class ContinuousBatchingTests: XCTestCase { func testBatchKVCacheMergeExtendFilterAndExtract() { @@ -407,12 +408,16 @@ private func assertBatchGeneratorRejectsCache( file: file, line: line ) { error in - guard case let BatchGeneratorError.unsupportedCacheTopology( - _, - path, - cacheType, - reason - ) = error + guard + case BatchGeneratorError.unsupportedCacheTopology( + _, + let + path, + let + cacheType, + let + reason + ) = error else { XCTFail( "Expected BatchGeneratorError.unsupportedCacheTopology, got \(error)",