diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift new file mode 100644 index 00000000..f968bf0d --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift @@ -0,0 +1,137 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Dispatch +import Synchronization + +extension BatchScheduler { + /// Scores `ScoreRequest`s in padded batches on `decodeStep`. Each batch runs on a new session. + struct Scorer { + /// Decode step the scoring passes run on. + let decodeStep: any BatchedDecodeStep + /// Maximum number of requests per scoring batch. + let maxConcurrent: Int + /// Positions the KV cache holds. A padded chunk never runs past it. + let maxPositions: Int + /// Maximum number of positions per scoring forward pass. Each pass returns at most + /// `rows * chunkLength * vocab` logits. + let chunkLength = 128 + /// Shortest padded chunk. The last chunk of a batch pads to 16, 32, 64 or 128 positions, so + /// scoring reuses a few input shapes instead of compiling one per prompt length. + static let shortestChunk = 16 + + /// Positions in the chunk at `start` for a batch whose longest prompt has `longest` tokens: a + /// full `chunkLength`, or the remaining positions rounded up to a power of two (at least + /// `shortestChunk`), unless that would run past `maxPositions`. With `chunkLength` 128 and + /// `longest` 300, the chunks at 0, 128 and 256 take 128, 128 and 64 positions. + static func paddedChunk(start: Int, longest: Int, chunkLength: Int, maxPositions: Int) -> Int { + let remaining = min(chunkLength, longest - start) + var padded = shortestChunk + while padded < remaining { padded *= 2 } + let fitting = min(padded, chunkLength, maxPositions - start) + return max(remaining, fitting) + } + + /// Rows in a scoring pass over `requests` requests: the graph's fixed `capacity` for a static + /// graph, otherwise `requests` rounded up to a power of two, at most `maxConcurrent`. + static func paddedRows(requests: Int, capacity: Int, maxConcurrent: Int) -> Int { + guard capacity == Int.max else { return capacity } + var rows = 1 + while rows < requests { rows *= 2 } + return max(requests, min(rows, maxConcurrent)) + } + + /// Counts the leading entries of `lengths`, prompt lengths in queue order, that form the next + /// scoring batch. Takes up to `limit` entries and stops before one that would make more than a + /// third of the batch's positions padding. Returns 0 when `lengths` is empty. With `limit` 4 + /// and lengths `[7, 40, 130, 300]`, successive batches are `[7]`, `[40]` and `[130, 300]`. + static func batchCount(_ lengths: [Int], limit: Int) -> Int { + guard let first = lengths.first else { return 0 } + var count = 1 + var total = first + var longest = first + while count < min(limit, lengths.count) { + let next = lengths[count] + let padded = (count + 1) * max(longest, next) + if 2 * padded > 3 * (total + next) { break } + count += 1 + total += next + longest = max(longest, next) + } + return count + } + + /// Scores `requests` in one padded batch on a new session. Returns one `LogProbabilities` per + /// request, in request order. At most `decodeStep.maxBatchCapacity` requests. + /// + /// Short rows repeat their last token, and the last chunk pads to `paddedChunk`. For prompts of + /// lengths 3 and 5 with a chunk length of 4 and a shortest chunk of 2, the rows are + /// ``` + /// chunk 0: [a0 a1 a2 a2] [b0 b1 b2 b3] positions 0..<4 + /// chunk 1: [a2 a2] [b4 b4] positions 4..<6, only b4 is read + /// ``` + func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { + let rows = Self.paddedRows( + requests: requests.count, capacity: decodeStep.maxBatchCapacity, maxConcurrent: maxConcurrent) + let session = try decodeStep.makeSessionState(batchSize: rows) + let longest = requests.map(\.tokens.count).max() ?? 0 + var entries = Array(repeating: [LogProbabilities.Entry](), count: requests.count) + + var start = 0 + while start < longest { + let q = Self.paddedChunk( + start: start, longest: longest, chunkLength: chunkLength, maxPositions: maxPositions) + let tokensPerRow = (0.. [Int32] in + guard row < requests.count else { return Array(repeating: 0, count: q) } + let tokens = requests[row].tokens + return (start.. 0 else { return } + let rowStart = row * q * vocab + let rowEntries = LogProbabilities.compute( + logits: UnsafeBufferPointer(rebasing: buffer[rowStart.. waiting queue ---> free row? ---> prefill the prompt into that row +/// [ F ][ G ] +/// +/// batch session, N = 4 rows, one decodeBatch per step over all rows: +/// +/// row 0 A A A A A a a = token sampled this step +/// row 1 B B b every row decodes at its own cursor +/// row 2 E e E took the row that D freed when it finished +/// row 3 C C C C c +/// ``` +/// +/// Slot `i` is row `i` of the batch KV cache, and `N` is `maxConcurrent`. A slot goes through these +/// phases: +/// - Admission: a queued request takes a free slot, is prefilled into its row and samples its first +/// token. +/// - Decode: each step runs `decodeBatch` over all `N` rows, every live row at its own cursor. +/// - Retirement: the row ends on EOS, a stop sequence, `maxTokens` or cancellation, and the slot +/// takes the next queued request while the other rows keep decoding. +/// +/// A request is cancelled by `BatchRequest.cancellation` or by dropping its stream. A cancelled +/// request leaves before the next decode step, or leaves the queue without a prefill. +/// +/// `score(_:)` requests run as their own padded batch between decode steps. +package final class BatchScheduler: RequestSink, @unchecked Sendable { + private let decodeStep: any BatchedDecodeStep + private let session: GenerationSessionState + package let maxConcurrent: Int + /// Requests allowed to wait beyond the free rows, or `nil` for no limit. At least 0. + private let maxQueued: Int? + + /// True when more than one request can decode at a time (`maxConcurrent > 1`). + package var supportsContinuousBatching: Bool { maxConcurrent > 1 } + + private struct Pending { + let request: BatchRequest + let continuation: AsyncThrowingStream.Continuation + } + private struct State { + var queue: [Pending] = [] + /// Rows held by admitted requests. + var occupiedRows = 0 + /// Scoring jobs waiting for the pump. Each `score(_:)` call appends its jobs shortest prompt + /// first. + var scoring: [ScoringJob] = [] + var pumping = false + + /// Removes cancelled requests from the queue and finishes their streams. + mutating func dropCancelledRequests() { + queue.removeAll { pending in + guard pending.request.cancellation.isCancelled else { return false } + pending.continuation.finish() + return true + } + } + } + private let state = Mutex(State()) + /// Lock held by the pump during each iteration, and by an `ExclusiveEngineLease` until it is + /// released. + private let engineLock = AsyncFIFOLock() + /// Scorer that runs the scoring batches. + private let scorer: Scorer + + private struct ScoringJob { + let request: ScoreRequest + let continuation: AsyncThrowingStream.Continuation + } + + /// Creates a scheduler over `engine` with a new `maxConcurrent`-row session. + /// + /// - Parameters: + /// - maxConcurrent: Rows in the batch, from 1 to the graph's batch capacity. Above 1, + /// `engine.supportsBatching` must be true. + /// - maxQueued: Requests allowed to wait beyond the free rows, at least 0. `nil` means no limit. + /// + /// Throws when the session cannot be created. Out-of-range values fail a precondition. + package convenience init(engine: CoreAISequentialEngine, maxConcurrent: Int, maxQueued: Int? = nil) throws { + precondition( + maxConcurrent == 1 || engine.supportsBatching, + "maxConcurrent \(maxConcurrent) needs a graph with an attn_mask input and a batch dimension above 1") + self.init( + decodeStep: engine, session: try engine.makeSessionState(batchSize: maxConcurrent), + maxConcurrent: maxConcurrent, maxQueued: maxQueued, maxPositions: engine.config.maxContextLength) + } + + init( + decodeStep: any BatchedDecodeStep, session: GenerationSessionState, maxConcurrent: Int, + maxQueued: Int? = nil, maxPositions: Int = .max + ) { + precondition(maxConcurrent >= 1, "maxConcurrent must be >= 1") + precondition( + maxConcurrent <= decodeStep.maxBatchCapacity, + "maxConcurrent \(maxConcurrent) exceeds the graph's batch capacity \(decodeStep.maxBatchCapacity)") + precondition((maxQueued ?? 0) >= 0, "maxQueued must be >= 0") + self.decodeStep = decodeStep + self.session = session + self.maxConcurrent = maxConcurrent + self.maxQueued = maxQueued + self.scorer = Scorer(decodeStep: decodeStep, maxConcurrent: maxConcurrent, maxPositions: maxPositions) + } + + // MARK: - RequestSink + + /// Queues `request` and returns its event stream. + /// + /// Throws `invalidArgument` when the prompt is empty, and `InferenceRuntimeError.queueFull` when + /// admitted plus queued requests would exceed `maxConcurrent + maxQueued`. With + /// `maxConcurrent == 2` and `maxQueued == 1`, a request throws while three others are admitted or + /// queued. With `maxQueued == nil`, it throws only for an empty prompt. + package func submit(_ request: BatchRequest) throws -> AsyncThrowingStream { + try enqueue([request])[0] + } + + /// Enqueues `requests` as one unit and starts the pump if none is running. Returns one stream per + /// request, in order. Throws `invalidArgument` when a prompt is empty and `queueFull` when the + /// queue cannot take all of `requests`. A throw enqueues nothing. + private func enqueue( + _ requests: [BatchRequest] + ) throws -> [AsyncThrowingStream] { + if requests.contains(where: { $0.promptTokens.isEmpty }) { + throw InferenceRuntimeError.invalidArgument("a batched request needs a non-empty prompt") + } + var streams: [AsyncThrowingStream] = [] + var pending: [Pending] = [] + streams.reserveCapacity(requests.count) + pending.reserveCapacity(requests.count) + for request in requests { + let (stream, continuation) = AsyncThrowingStream.makeStream(of: BatchSinkEvent.self) + continuation.onTermination = { [cancellation = request.cancellation] termination in + if case .cancelled = termination { cancellation.cancel() } + } + streams.append(stream) + pending.append(Pending(request: request, continuation: continuation)) + } + let shouldStart = try state.withLock { s -> Bool in + s.dropCancelledRequests() + let freeRows = maxConcurrent - s.occupiedRows + if let maxQueued, s.queue.count + pending.count > freeRows + maxQueued { + throw InferenceRuntimeError.queueFull(depth: maxQueued) + } + s.queue.append(contentsOf: pending) + if s.pumping { return false } + s.pumping = true + return true + } + if shouldStart { Task { await self.pump() } } + return streams + } + + /// A request that holds a slot. `cursor` is the next KV write position in its row. + private struct ActiveRow { + let continuation: AsyncThrowingStream.Continuation + var adapter: BatchSessionAdapter + var lastToken: Int32 + var cursor: Int + } + + /// Scheduling loop, run by at most one task at a time. Each iteration does the following: + /// 1. Run every waiting scoring job. + /// 2. Admit queued requests into free slots, and retire cancelled rows. + /// 3. If no slot is live, loop again when requests or scoring jobs are waiting, otherwise clear + /// `pumping` and return. + /// 4. Run one `decodeBatch` over all `N` slots with `tokens` `[N][1]` and `cursors` `[N]`. A free + /// slot gets token 0 at cursor 0 and its logits are dropped. + /// 5. Advance each live cursor, then sample, emit or retire each live row. + /// + /// If `decodeBatch` throws, every resident row, queued request and queued scoring job ends with + /// that error. + private func pump() async { + var active: [ActiveRow?] = Array(repeating: nil, count: maxConcurrent) + do { + while true { + await engineLock.lock() + defer { engineLock.unlock() } + await runPendingScoring() + await admit(into: &active) + for slot in active.indices where active[slot]?.adapter.request.cancellation.isCancelled == true { + active[slot]!.continuation.finish() + retire(slot, active: &active) + } + + let liveSlots = active.indices.filter { active[$0] != nil } + if liveSlots.isEmpty { + // Requests and scoring jobs enqueued while `pumping` is set start no pump. + let more = state.withLock { s -> Bool in + if s.queue.isEmpty && s.scoring.isEmpty { + s.pumping = false + return false + } + return true + } + if more { continue } + return + } + + let tokens = active.map { [$0?.lastToken ?? 0] } + let cursors = active.map { $0?.cursor ?? 0 } + let logits = try await decodeStep.decodeBatch( + tokensPerRow: tokens, startPositions: cursors, session: session) + for slot in active.indices where active[slot] != nil { active[slot]!.cursor += 1 } + for slot in liveSlots { stepSlot(slot, logits: logits[slot], active: &active) } + } + } catch { + for row in active where row != nil { row!.continuation.finish(throwing: error) } + let (orphaned, orphanedScoring) = state.withLock { s -> ([Pending], [ScoringJob]) in + let waiting = (s.queue, s.scoring) + s.queue.removeAll() + s.scoring.removeAll() + s.occupiedRows = 0 + s.pumping = false + return waiting + } + for pending in orphaned { pending.continuation.finish(throwing: error) } + for job in orphanedScoring { job.continuation.finish(throwing: error) } + } + } + + /// Fills free slots from the queue. Each admitted request is prefilled into its row and its first + /// token is sampled immediately. A request can retire during admission. + /// + /// A prefill error fails only that request's stream and leaves its slot free. + private func admit(into active: inout [ActiveRow?]) async { + for slot in active.indices where active[slot] == nil { + let pending: Pending? = state.withLock { s in + s.dropCancelledRequests() + guard !s.queue.isEmpty else { return nil } + s.occupiedRows += 1 + return s.queue.removeFirst() + } + guard let pending else { break } + do { + let adapter = BatchSessionAdapter(request: pending.request) + let logits = try await decodeStep.prefillRow( + row: slot, promptTokens: pending.request.promptTokens, session: session) + active[slot] = ActiveRow( + continuation: pending.continuation, adapter: adapter, lastToken: 0, + cursor: pending.request.promptTokens.count) + stepSlot(slot, logits: logits, active: &active) + } catch { + state.withLock { $0.occupiedRows -= 1 } + pending.continuation.finish(throwing: error) + } + } + } + + /// Samples one token for `slot`, emits it, and retires the row when its request is done. + private func stepSlot(_ slot: Int, logits: [LogitsScalarType], active: inout [ActiveRow?]) { + guard active[slot] != nil else { return } + let token = BatchedSampler.sample( + rows: [logits], + configurations: [active[slot]!.adapter.request.sampling], + histories: [ArraySlice(active[slot]!.adapter.generated)], + steps: [active[slot]!.adapter.step])[0] + let stillLive = active[slot]!.adapter.accept(token) + active[slot]!.continuation.yield(.token(token, logits: nil)) + if stillLive { + active[slot]!.lastToken = token + } else { + active[slot]!.continuation.yield(.finished(active[slot]!.adapter.stopReason ?? .maxTokens)) + active[slot]!.continuation.finish() + retire(slot, active: &active) + } + } + + private func retire(_ slot: Int, active: inout [ActiveRow?]) { + active[slot] = nil + state.withLock { $0.occupiedRows -= 1 } + } + + // MARK: - Scoring + + /// Scores each request's targets. Returns one `LogProbabilities` per request, in request order, + /// with one entry per position: the target's log-probability and the top `topK` alternatives. + /// + /// Throws `invalidArgument` when a request has no tokens or its `targets` and `tokens` counts + /// differ. Rethrows the error of a failed scoring batch or decode step. + /// + /// Requests run between decode steps, grouped by prompt length into batches of up to + /// `maxConcurrent` requests, each batch on a new session. Decode rows pause during a scoring + /// batch and keep their KV state. A failed scoring batch fails only its own requests. + package func score(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { + if requests.contains(where: { $0.tokens.isEmpty || $0.targets.count != $0.tokens.count }) { + throw InferenceRuntimeError.invalidArgument( + "a scoring request needs a non-empty prompt and one target per token") + } + let order = requests.indices.sorted { requests[$0].tokens.count < requests[$1].tokens.count } + var jobs: [ScoringJob] = [] + var streams = [AsyncThrowingStream?](repeating: nil, count: requests.count) + for index in order { + let (stream, continuation) = AsyncThrowingStream.makeStream(of: LogProbabilities.self) + jobs.append(ScoringJob(request: requests[index], continuation: continuation)) + streams[index] = stream + } + let shouldStart = state.withLock { s -> Bool in + s.scoring.append(contentsOf: jobs) + if s.pumping { return false } + s.pumping = true + return true + } + if shouldStart { Task { await self.pump() } } + + var results: [LogProbabilities] = [] + results.reserveCapacity(streams.count) + for stream in streams { + var iterator = stream!.makeAsyncIterator() + guard let result = try await iterator.next() else { + throw InferenceRuntimeError.invalidState("scoring ended without a result") + } + results.append(result) + } + return results + } + + /// Runs every waiting scoring job, in batches of up to `maxConcurrent` sized by + /// `Scorer.batchCount(_:limit:)`. Call only while holding `engineLock`. + private func runPendingScoring() async { + while true { + let jobs = state.withLock { s -> [ScoringJob] in + let count = Scorer.batchCount(s.scoring.map(\.request.tokens.count), limit: scorer.maxConcurrent) + let taken = Array(s.scoring.prefix(count)) + s.scoring.removeFirst(taken.count) + return taken + } + if jobs.isEmpty { return } + do { + let results = try await scorer.scoreBatch(jobs.map(\.request)) + for (job, result) in zip(jobs, results) { + job.continuation.yield(result) + job.continuation.finish() + } + } catch { + for job in jobs { job.continuation.finish(throwing: error) } + } + } + } + + // MARK: - Exclusive engine access + + /// Waits behind the current scheduling step and any earlier lease, then returns a lease that + /// gives the caller sole use of the engine until `release()`. Resident rows keep their slots and + /// resume decoding after release. + /// + /// A request that runs directly on the scheduler's engine holds a lease for as long as it runs: + /// + /// ```swift + /// let lease = await scheduler.acquireExclusiveAccess() + /// defer { lease.release() } + /// let sequence = try await engine.generate( + /// with: promptTokens, samplingConfiguration: sampling, inferenceOptions: options) + /// ``` + package func acquireExclusiveAccess() async -> ExclusiveEngineLease { + await engineLock.lock() + return ExclusiveEngineLease(lock: engineLock) + } + + // MARK: - Closed batch + + /// Submits `requests` together, waits for all of them, and returns one `RowResult` per request, in + /// request order. Can be called while other requests stream. Throws `invalidArgument` when a + /// prompt is empty, `queueFull` when the queue cannot take all of `requests`, and any error that + /// ends a request's stream. + func run(_ requests: [BatchRequest]) async throws -> [RowResult] { + let streams = try enqueue(requests) + return try await withThrowingTaskGroup(of: (Int, RowResult).self) { group in + for (i, stream) in streams.enumerated() { + group.addTask { + var tokens: [Int32] = [] + var reason = StopReason.maxTokens + for try await event in stream { + switch event { + case .token(let token, _): tokens.append(token) + case .finished(let stopReason): reason = stopReason + } + } + return (i, RowResult(tokens: tokens, stopReason: reason)) + } + } + var results = Array( + repeating: RowResult(tokens: [], stopReason: .maxTokens), count: streams.count) + for try await (i, result) in group { results[i] = result } + return results + } + } +} + +/// A batched backend shared by many requests, each with its own event stream. Wrap it in a +/// `SchedulerSession`, one per request, to use it as an `InferenceEngine`. +package protocol RequestSink: Sendable { + /// Maximum number of requests decoded in one batch. At least 1. + var maxConcurrent: Int { get } + + /// Queues `request` and returns its event stream. Throws when the sink rejects the request, such + /// as `InferenceRuntimeError.queueFull` when its queue is full. + func submit(_ request: BatchRequest) throws -> AsyncThrowingStream +} + +/// A generation request for a `RequestSink`: prompt, sampling and stop conditions. +package struct BatchRequest: Sendable { + /// List of prompt tokens, cannot be empty. + let promptTokens: [Int32] + let sampling: SamplingConfiguration + let eosTokenIds: Set + let maxTokens: Int + let stopSequences: StopSequences + /// Converts a matched stop sequence's tokens to the text reported in `StopReason.stopSequence`. + /// Defaults to an empty string. + let decodeStopText: @Sendable ([Int32]) -> String + /// Cancels the request when set. Dropping the request's stream also cancels it. + let cancellation = RequestCancellation() + + init( + promptTokens: [Int32], + sampling: SamplingConfiguration, + eosTokenIds: Set, + maxTokens: Int, + stopSequences: StopSequences = StopSequences(sequences: []), + decodeStopText: @escaping @Sendable ([Int32]) -> String = { _ in "" } + ) { + self.promptTokens = promptTokens + self.sampling = sampling + self.eosTokenIds = eosTokenIds + self.maxTokens = maxTokens + self.stopSequences = stopSequences + self.decodeStopText = decodeStopText + } +} + +/// A flag that cancels a `BatchRequest`. Safe to set from any task. +final class RequestCancellation: Sendable { + private let cancelled = Atomic(false) + + var isCancelled: Bool { cancelled.load(ordering: .relaxed) } + + func cancel() { cancelled.store(true, ordering: .relaxed) } +} + +/// An event on a request's stream from a `RequestSink`. A stream yields `.token` events, then one +/// `.finished`, or ends by throwing. +package enum BatchSinkEvent: Sendable { + /// A sampled token and, when the sink provides them, its logits. `BatchScheduler` always sends + /// `nil` logits. + case token(Int32, logits: [LogitsScalarType]?) + /// The reason the request stopped. Always the last event. + case finished(StopReason) +} + +/// A lease on sole use of a `BatchScheduler`'s engine, returned by `acquireExclusiveAccess()`. +/// Calls to `release()` after the first do nothing. Deinit releases a lease that is still held. +package final class ExclusiveEngineLease: Sendable { + private let lock: AsyncFIFOLock + private let released = Mutex(false) + + init(lock: AsyncFIFOLock) { self.lock = lock } + + package func release() { + let first = released.withLock { wasReleased -> Bool in + defer { wasReleased = true } + return !wasReleased + } + if first { lock.unlock() } + } + + deinit { release() } +} + +/// An async lock granted to waiters in arrival order. `unlock()` hands the lock directly to the next +/// waiter, and a task that unlocks and locks again queues behind the existing waiters. +final class AsyncFIFOLock: Sendable { + private struct State { + var held = false + var waiters: [CheckedContinuation] = [] + } + private let state = Mutex(State()) + + func lock() async { + await withCheckedContinuation { continuation in + let granted = state.withLock { s -> Bool in + if !s.held { + s.held = true + return true + } + s.waiters.append(continuation) + return false + } + if granted { continuation.resume() } + } + } + + func unlock() { + let next = state.withLock { s -> CheckedContinuation? in + if s.waiters.isEmpty { + s.held = false + return nil + } + return s.waiters.removeFirst() + } + next?.resume() + } +} + +/// Generated tokens and stop reason for one request from `BatchScheduler.run(_:)`. +struct RowResult: Sendable { + let tokens: [Int32] + let stopReason: StopReason +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift new file mode 100644 index 00000000..d6c09590 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift @@ -0,0 +1,21 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +/// Serving lane of a request under a batched server. +package enum BatchServingLane: Equatable, Sendable { + /// Requests that share the batch scheduler through a `SchedulerSession`. + case batched + /// Requests that run one at a time on their own engine session, such as guided generation. + case fallbackBatchOne +} + +/// Lane selection for chat requests under a batched server. +package enum BatchServingRouter { + /// Returns `.fallbackBatchOne` for a request with a `response_format` JSON schema, otherwise + /// `.batched`. + package static func lane(hasResponseSchema: Bool) -> BatchServingLane { + hasResponseSchema ? .fallbackBatchOne : .batched + } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift new file mode 100644 index 00000000..dea78b50 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift @@ -0,0 +1,64 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +/// Token history and stop checks for one request in a batched decode. +/// +/// Pass each sampled token to `accept(_:)`. Once it returns false, the request is done and +/// `stopReason` holds the reason. +struct BatchSessionAdapter { + let request: BatchRequest + + private(set) var generated: [Int32] = [] + /// False once the row has retired. + private(set) var live: Bool = true + /// The reason the row retired, nil while it is live. + private(set) var stopReason: StopReason? + + /// The last `stopSequences.maxLength` tokens, matched against the stop sequences. + private var recentTokens: [Int32] = [] + + /// The row's sampler step index: the number of tokens generated so far. + var step: Int { generated.count } + + init(request: BatchRequest) { + self.request = request + } + + /// Appends `token`, then retires the row on an EOS id, a stop-sequence match or `maxTokens`, + /// checked in that order. Returns true while the row is live. After retirement, ignores tokens + /// and returns false. + @discardableResult + mutating func accept(_ token: Int32) -> Bool { + guard live else { return false } + generated.append(token) + + if request.eosTokenIds.contains(token) { + retire(.eos) + return false + } + + if request.stopSequences.maxLength > 0 { + recentTokens.append(token) + if recentTokens.count > request.stopSequences.maxLength { + recentTokens.removeFirst(recentTokens.count - request.stopSequences.maxLength) + } + if let matched = request.stopSequences.matchedSequence(recentTokens: recentTokens) { + retire(.stopSequence(request.decodeStopText(matched))) + return false + } + } + + if generated.count >= request.maxTokens { + retire(.maxTokens) + return false + } + return true + } + + private mutating func retire(_ reason: StopReason) { + live = false + stopReason = reason + } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift index ff5805e4..ac36879b 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift @@ -5,8 +5,9 @@ // MARK: - Batched Decode Step -/// Runs one forward pass over the `N` rows of a batch KV cache and returns each row's last-token -/// logits. Sampling, stop handling and admission stay with the caller. +/// Batched forward passes over the `N` rows of a batch KV cache: admitting a row, decoding all rows, +/// and reading logits at every position. Sampling, stop handling and admission order stay with the +/// caller. /// /// The caller tracks each row's cursor. Neither method advances `session.processedTokenCount`. /// @@ -45,4 +46,15 @@ protocol BatchedDecodeStep { promptTokens: [Int32], session: GenerationSessionState ) async throws -> [LogitsScalarType] + + /// Runs `tokensPerRow` (`[N][q]`, every row at `startPosition`) and returns the logits at every + /// position, flat row-major `[N, q, vocab]`. + func allPositionLogits( + tokensPerRow: [[Int32]], + startPosition: Int, + session: GenerationSessionState + ) async throws -> [LogitsScalarType] + + /// Creates a session whose KV cache holds `batchSize` rows, at most `maxBatchCapacity`. + func makeSessionState(batchSize: Int) throws -> GenerationSessionState } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 69f557c1..82edf8dc 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -98,6 +98,13 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba return 1 // Single-batch graph. } + /// True when the graph has an attention-mask input and a batch dimension above 1, static or + /// dynamic. + package var supportsBatching: Bool { + functionDescriptor.inputNames.contains(where: InputLayout.knownAttnMaskNames.contains) + && maxBatchCapacity > 1 + } + /// The prefix hit count of the most recent generation. With several sessions in use, this is /// the count of whichever session ran last. public private(set) var lastPrefixHitCount: Int = 0 @@ -410,7 +417,58 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba startPositions: [Int], session: GenerationSessionState ) async throws -> [[LogitsScalarType]] { - precondition(!tokensPerRow.isEmpty, "decodeBatch requires at least one row") + let out = try await runBatchedForward( + tokensPerRow: tokensPerRow, startPositions: startPositions, session: session) + let n = tokensPerRow.count + let q = tokensPerRow[0].count + let vocab = config.vocabSize + return Self.logitsRows( + out, count: n, rowLength: vocab, firstOffset: (q - 1) * vocab, rowStride: q * vocab) + } + + func allPositionLogits( + tokensPerRow: [[Int32]], + startPosition: Int, + session: GenerationSessionState + ) async throws -> [LogitsScalarType] { + let n = tokensPerRow.count + let out = try await runBatchedForward( + tokensPerRow: tokensPerRow, startPositions: Array(repeating: startPosition, count: n), + session: session) + return readNDArray(out, as: LogitsScalarType.self, count: n * tokensPerRow[0].count * config.vocabSize) + } + + /// Copies `count` rows of `rowLength` logits out of `array`. Row `b` starts at element + /// `firstOffset + b * rowStride` in row-major order. + private static func logitsRows( + _ array: NDArray, count: Int, rowLength: Int, firstOffset: Int, rowStride: Int + ) -> [[LogitsScalarType]] { + let view = array.view(as: LogitsScalarType.self) + let contiguous = view.withUnsafePointer { _, shape, strides in + isContiguousRowMajor(shape: shape, strides: strides) + } + guard contiguous else { + let total = firstOffset + (count - 1) * rowStride + rowLength + let flat = readNDArray(array, as: LogitsScalarType.self, count: total) + return (0.. NDArray { + precondition(!tokensPerRow.isEmpty, "a batched forward requires at least one row") let n = tokensPerRow.count precondition(startPositions.count == n, "one start position per row") let q = tokensPerRow[0].count @@ -489,11 +547,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba } } - let flat = readNDArray(out, as: LogitsScalarType.self, count: n * q * vocab) - return (0.. + + package init(sink: any RequestSink, config: ModelConfig, stopTokenIds: Set) { + self.sink = sink + self.config = config + self.stopTokenIds = stopTokenIds + } + + /// Submits `input` to the sink and returns the request's output sequence. + /// + /// The request's token limit is `inferenceOptions.maxTokens` capped at the context room left after + /// the prompt, and defaults to that room. At least 1. Throws when the sink rejects the request, + /// for example with `InferenceRuntimeError.queueFull`. + package func generate( + with input: [TokenId], + samplingConfiguration: SamplingConfiguration, + inferenceOptions: InferenceOptions + ) async throws -> GenerationSequence { + let contextRoom = Swift.max(1, config.maxContextLength - input.count) + let requested = inferenceOptions.maxTokens ?? contextRoom + let request = BatchRequest( + promptTokens: input.map { Int32($0) }, + sampling: samplingConfiguration, + eosTokenIds: stopTokenIds, + maxTokens: Swift.max(1, Swift.min(requested, contextRoom))) + return GenerationSequence(events: try sink.submit(request), cancellation: request.cancellation) + } + + /// Does nothing. + package func reset(to tokenIndex: Int) async throws {} + + /// Always false. + package var supportsLogits: Bool { false } + + /// Always false. + package var isBusy: Bool { false } + + package var maxContextLength: Int { config.maxContextLength } +} + +// MARK: - Output sequence + +extension SchedulerSession { + /// The sink's event stream for one request, as an `InferenceOutputSequence`. + /// + /// `stopReason` is the sink's `.finished` reason, or `.maxTokens` when the stream ends without + /// one. `setStopReason(_:)` overrides it and cancels the request, so the stream ends within one + /// decode step. + package struct GenerationSequence: InferenceOutputSequence { + package typealias Element = InferenceOutput + package typealias Failure = Error + + let events: AsyncThrowingStream + let cancellation: RequestCancellation + let stopReasonStore = StopReasonStore() + + package var stopReason: StopReason? { stopReasonStore.stopReason } + + package func setStopReason(_ reason: StopReason) { + stopReasonStore.set(reason) + cancellation.cancel() + } + + package func makeAsyncIterator() -> Iterator { + Iterator(upstream: events.makeAsyncIterator(), stopReasonStore: stopReasonStore) + } + } +} + +extension SchedulerSession.GenerationSequence { + package final class Iterator: AsyncIteratorProtocol { + package typealias Element = InferenceOutput + package typealias Failure = Error + + private var upstream: AsyncThrowingStream.AsyncIterator + private let stopReasonStore: StopReasonStore + + init( + upstream: AsyncThrowingStream.AsyncIterator, + stopReasonStore: StopReasonStore + ) { + self.upstream = upstream + self.stopReasonStore = stopReasonStore + } + + package func next() async throws -> InferenceOutput? { + while let event = try await upstream.next() { + switch event { + case .token(let id, let logits): + return InferenceOutput(tokenId: id, logits: logits) + case .finished(let reason): + stopReasonStore.setIfUnset(reason) + } + } + stopReasonStore.setIfUnset(.maxTokens) + return nil + } + } +} diff --git a/swift/Sources/CoreAILanguageModels/TextGeneration/LogProbabilities.swift b/swift/Sources/CoreAILanguageModels/TextGeneration/LogProbabilities.swift index 07083d5d..e4f64ae1 100644 --- a/swift/Sources/CoreAILanguageModels/TextGeneration/LogProbabilities.swift +++ b/swift/Sources/CoreAILanguageModels/TextGeneration/LogProbabilities.swift @@ -46,90 +46,113 @@ public struct LogProbabilities: Sendable { targets: [Int32], topK: Int = 0 ) -> LogProbabilities { + var scratch = Scratch() var entries: [Entry] = [] entries.reserveCapacity(min(logits.count, targets.count)) - for (logitVec, targetToken) in zip(logits, targets) { - let tokenIndex = Int(targetToken) - guard tokenIndex >= 0 && tokenIndex < logitVec.count else { - entries.append(Entry(tokenId: targetToken, value: -.infinity, alternatives: [])) - continue + logitVec.withUnsafeBufferPointer { row in + entries.append(entry(row: row, target: targetToken, topK: topK, scratch: &scratch)) } + } + return LogProbabilities(entries: entries) + } - let vocabSize = logitVec.count - let (logSumExp, floatBuffer) = computeLogSumExpVectorized(logitVec) + /// Computes one entry per target from `logits`, a row-major `[targets.count, vocabSize]` buffer + /// whose row `p` holds the logits that predict `targets[p]`. + static func compute( + logits: UnsafeBufferPointer, + vocabSize: Int, + targets: ArraySlice, + topK: Int + ) -> [Entry] { + precondition(logits.count >= targets.count * vocabSize, "logits buffer is shorter than targets") + var scratch = Scratch() + var entries: [Entry] = [] + entries.reserveCapacity(targets.count) + for (position, targetToken) in targets.enumerated() { + let start = position * vocabSize + let row = UnsafeBufferPointer(rebasing: logits[start.. 0 { - alts = findTopK(floatBuffer, k: topK, logSumExp: logSumExp, vocabSize: vocabSize) - } + private static func entry( + row: UnsafeBufferPointer, target: Int32, topK: Int, scratch: inout Scratch + ) -> Entry { + let tokenIndex = Int(target) + guard tokenIndex >= 0 && tokenIndex < row.count else { + return Entry(tokenId: target, value: -.infinity, alternatives: []) + } - entries.append(Entry(tokenId: targetToken, value: targetLogProb, alternatives: alts)) + let vocabSize = row.count + if scratch.floats.count != vocabSize { + scratch.floats = [Float](repeating: 0, count: vocabSize) + scratch.shifted = [Float](repeating: 0, count: vocabSize) + } + let logSumExp = computeLogSumExpVectorized(row, floats: &scratch.floats, shifted: &scratch.shifted) + + let targetLogProb: Double + if logSumExp.isInfinite { + // +Inf logits: tokens with +Inf get log-prob 0, others get -Inf + targetLogProb = Double(scratch.floats[tokenIndex]).isInfinite ? 0.0 : -.infinity + } else { + let raw = Double(scratch.floats[tokenIndex]) - logSumExp + targetLogProb = raw.isNaN ? 0.0 : raw } - return LogProbabilities(entries: entries) + var alts: [(tokenId: Int32, value: Double)] = [] + if topK > 0 { + alts = findTopK(scratch.floats, k: topK, logSumExp: logSumExp, vocabSize: vocabSize) + } + return Entry(tokenId: target, value: targetLogProb, alternatives: alts) } - /// Vectorized log-sum-exp using Accelerate. - /// Returns (logSumExp, floatBuffer) where floatBuffer is the Float32-converted logits. + /// Converts `logits` to Float32 into `floats` and returns their log-sum-exp, using `shifted` as + /// working space. Returns infinity when the maximum logit is infinite. private static func computeLogSumExpVectorized( - _ logits: [LogitsScalarType] - ) -> (Double, [Float]) { + _ logits: UnsafeBufferPointer, floats: inout [Float], shifted: inout [Float] + ) -> Double { let count = logits.count // Float16 → Float32 via vFloatConversion (Accelerate) - var floatBuffer = [Float](repeating: 0, count: count) - logits.withUnsafeBufferPointer { src in - src.baseAddress!.withMemoryRebound(to: UInt16.self, capacity: count) { halfPtr in - floatBuffer.withUnsafeMutableBufferPointer { dst in - var bufferSrc = vImage_Buffer( - data: UnsafeMutableRawPointer(mutating: halfPtr), - height: 1, width: vImagePixelCount(count), rowBytes: count * 2) - var bufferDst = vImage_Buffer( - data: dst.baseAddress!, height: 1, - width: vImagePixelCount(count), rowBytes: count * 4) - vImageConvert_Planar16FtoPlanarF(&bufferSrc, &bufferDst, 0) - } + logits.baseAddress!.withMemoryRebound(to: UInt16.self, capacity: count) { halfPtr in + floats.withUnsafeMutableBufferPointer { dst in + var bufferSrc = vImage_Buffer( + data: UnsafeMutableRawPointer(mutating: halfPtr), + height: 1, width: vImagePixelCount(count), rowBytes: count * 2) + var bufferDst = vImage_Buffer( + data: dst.baseAddress!, height: 1, + width: vImagePixelCount(count), rowBytes: count * 4) + vImageConvert_Planar16FtoPlanarF(&bufferSrc, &bufferDst, 0) } } var maxVal: Float = 0 - vDSP_maxv(floatBuffer, 1, &maxVal, vDSP_Length(count)) + vDSP_maxv(floats, 1, &maxVal, vDSP_Length(count)) if maxVal.isInfinite { - return (Double.infinity, floatBuffer) + return Double.infinity } - // log-sum-exp with temporary stack buffers var negMax = -maxVal let countLen = vDSP_Length(count) var countInt32 = Int32(count) - return withUnsafeTemporaryAllocation(of: Float.self, capacity: count) { shiftedBuf in - vDSP_vsadd(floatBuffer, 1, &negMax, shiftedBuf.baseAddress!, 1, countLen) - - return withUnsafeTemporaryAllocation(of: Float.self, capacity: count) { expBuf in - vvexpf(expBuf.baseAddress!, shiftedBuf.baseAddress!, &countInt32) - - var sumExp: Float = 0 - vDSP_sve(expBuf.baseAddress!, 1, &sumExp, countLen) - - let logSumExp = Double(maxVal) + Double(log(sumExp)) - return (logSumExp, floatBuffer) - } + return shifted.withUnsafeMutableBufferPointer { work in + vDSP_vsadd(floats, 1, &negMax, work.baseAddress!, 1, countLen) + vvexpf(work.baseAddress!, work.baseAddress!, &countInt32) + var sumExp: Float = 0 + vDSP_sve(work.baseAddress!, 1, &sumExp, countLen) + return Double(maxVal) + Double(log(sumExp)) } } - - /// Find top-K elements using partial sort (O(n) for small K). private static func findTopK( _ floatBuffer: [Float], k: Int, diff --git a/swift/Sources/Tools/benchmark/BenchmarkMain.swift b/swift/Sources/Tools/benchmark/BenchmarkMain.swift index 05c08e73..2f186ac4 100644 --- a/swift/Sources/Tools/benchmark/BenchmarkMain.swift +++ b/swift/Sources/Tools/benchmark/BenchmarkMain.swift @@ -64,13 +64,13 @@ struct LLMBenchmark: AsyncParsableCommand { @Option( name: .customLong("inference-engine-variant"), - help: "Engine variant. The default picks one from the model. Batch sizes above 1 need coreai-sequential.") + help: "Engine variant. The default picks one from the model: coreai-sequential when it has an attn_mask input.") var inferenceEngineVariant: String = "default" @Option( name: .customLong("batch-size"), help: - "Number of requests to run as one batch. Defaults to the graph's static batch dimension. A dynamic-batch graph reports 1, so set this explicitly. Needs --inference-engine-variant coreai-sequential." + "Number of requests to run as one batch. Defaults to the graph's static batch dimension. A dynamic-batch graph reports 1, so set this explicitly. Above 1, needs a model exported with --attention-mask batched." ) var batchSize: Int? @@ -143,8 +143,8 @@ struct LLMBenchmark: AsyncParsableCommand { if effectiveBatch > 1 { guard engine is CoreAISequentialEngine else { throw ValidationError( - "batched benchmarking with batch size \(effectiveBatch) requires " - + "--inference-engine-variant coreai-sequential") + "batched benchmarking with batch size \(effectiveBatch) needs the coreai-sequential engine " + + "and a model exported with --attention-mask batched") } let source = batchSize != nil ? "requested" : "auto-detected from model" print("🧵 Batch size: \(effectiveBatch) (\(source); benchmarking \(effectiveBatch) rows together)") diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index 3165d755..c5e5115f 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -1000,10 +1000,13 @@ struct LLMRunner: AsyncParsableCommand, Sendable { // MARK: - VLM Inference - /// Runs `rows` copies of the prompt through `lockstepBatchedGenerate` and prints each row. + /// Runs `rows` copies of `promptTokens` through `lockstepBatchedGenerate`, up to `--max-tokens` + /// new tokens per row, and prints each row's decoded text. /// - /// Row 0 is greedy, so it should match a normal greedy run of the same prompt. The other rows - /// sample at `--temperature` (0.8 if that is 0) with a fixed seed per row, so reruns repeat. + /// Row 0 is greedy. Rows 1 and up sample at `--temperature` (0.8 when it is 0 or less) with + /// seed `1234 + row`. A row stops on the tokenizer's EOS or an `additionalEosTokenIds` token. + /// Throws `ValidationError` when the engine is not a `CoreAISequentialEngine` or `rows` exceeds + /// `bundle.maxBatchSize`. private func runBatchPrototype( rows: Int, inferenceEngine: any InferenceEngine, diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 42aeba95..0685deae 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -103,45 +103,55 @@ private func handleAutoRoute(request: Request, state: ServerState) async throws private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws -> Response { - // Non-streaming paths release the permit at each return; the streaming path - // hands it to the body closure. Release-on-deinit backstops a dropped response. - let permit: QueuePermit - do { - permit = try await state.queue.acquire() - } catch let error as ServerError { - let err = ErrorResponse( - error: .init( - message: error.localizedDescription, type: "server_error", code: "queue_full")) - let data = try JSONEncoder().encode(err) - return Response( - status: .tooManyRequests, headers: [.contentType: "application/json"], - body: .init(byteBuffer: ByteBuffer(data: data))) - } let chatRequest: ChatCompletionRequest do { chatRequest = try JSONDecoder().decode(ChatCompletionRequest.self, from: body) } catch { - permit.release() let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil)) let data = try JSONEncoder().encode(err) return Response( status: .badRequest, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } + + let lane = BatchServingRouter.lane( + hasResponseSchema: chatRequest.responseFormat?.extractedSchema != nil) + let access: RequestAccess + do { + access = try await state.acquireAccess(lane: lane) + } catch let error as ServerError { + let err = ErrorResponse( + error: .init( + message: error.localizedDescription, type: "server_error", code: "queue_full")) + let data = try JSONEncoder().encode(err) + return Response( + status: .tooManyRequests, headers: [.contentType: "application/json"], + body: .init(byteBuffer: ByteBuffer(data: data))) + } + do { let shouldStream = chatRequest.stream ?? false if shouldStream { - // Ownership of the permit passes to the streaming body closure. + // Ownership of the access passes to the streaming body closure. return try await handleStreamingRequest( - chatRequest: chatRequest, state: state, sessionID: sessionID, permit: permit) + chatRequest: chatRequest, state: state, sessionID: sessionID, access: access) } else { let response = try await handleNonStreamingRequest( chatRequest: chatRequest, state: state, sessionID: sessionID) - permit.release() + access.release() return response } + } catch let error as InferenceRuntimeError where error.isQueueFull { + // Thrown by a full batch scheduler on submit, on both the streaming and non-streaming paths. + access.release() + let err = ErrorResponse( + error: .init(message: error.localizedDescription, type: "server_error", code: "queue_full")) + let data = try JSONEncoder().encode(err) + return Response( + status: .tooManyRequests, headers: [.contentType: "application/json"], + body: .init(byteBuffer: ByteBuffer(data: data))) } catch let error as ServerError { - permit.release() + access.release() let status: HTTPResponse.Status = error.isBadRequest ? .badRequest : .internalServerError let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil)) let data = try JSONEncoder().encode(err) @@ -149,7 +159,7 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, status: status, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data)) ) } catch { - permit.release() + access.release() let err = ErrorResponse(error: .init(message: "\(error)", type: "server_error", code: nil)) let data = try JSONEncoder().encode(err) return Response( @@ -223,17 +233,19 @@ func runChatCompletion(chatRequest: ChatCompletionRequest, state: ServerState, s let t0 = SuspendingClock().now let strategy: any DecodingStrategy - if let schema = chatRequest.responseFormat?.extractedSchema { + let schema = chatRequest.responseFormat?.extractedSchema + if let schema { CLILogger.log("[\(requestID)] constrained generation (json_schema)", component: "Server") strategy = ConstrainedDecodingStrategy(jsonSchema: schema, vocabSize: state.config.vocabSize) } else { strategy = VanillaDecodingStrategy() } + let lane = BatchServingRouter.lane(hasResponseSchema: schema != nil) let stream = try await strategy.decode( from: input, tokenizer: state.tokenizer, - inferenceEngine: state.engine, + inferenceEngine: state.engineForRequest(lane: lane), samplingConfiguration: samplingConfig, options: InferenceOptions(maxTokens: requestMaxTokens), stopSequences: stopSequences @@ -442,31 +454,49 @@ func prepareStreaming(chatRequest: ChatCompletionRequest, state: ServerState, se requestMaxTokens: requestMaxTokens, prefixReuseTokens: prefixReused) } -/// The streaming generation core, shared by the HTTP SSE handler and `--replay`. Drives the -/// incremental token loop with per-chunk think/tool parsing, delivering each `Delta` to `emit`, -/// and returns per-request instrumentation. The caller owns SSE framing: the leading role chunk -/// and the trailing done/`[DONE]` frames are added by the HTTP handler. -func runStreamingLoop( - prepared: PreparedStream, chatRequest: ChatCompletionRequest, state: ServerState, - emit: (ChatCompletionChunk.Delta) async throws -> Void -) async throws -> StreamingOutcome { - let genStart = SuspendingClock().now +/// Token stream and start time of a streaming decode. On the batched lane the request is already in +/// the batch scheduler. Iterate `tokens` from one task only. +struct StartedDecode: @unchecked Sendable { + let tokens: any AsyncSequence + let start: SuspendingClock.Instant +} +/// Starts decoding `prepared` on the engine for the request's lane, with constrained decoding when +/// `chatRequest` has a JSON schema. Throws `InferenceRuntimeError.queueFull` when the batch +/// scheduler is full. +func startStreamingDecode( + prepared: PreparedStream, chatRequest: ChatCompletionRequest, state: ServerState +) async throws -> StartedDecode { + let start = SuspendingClock().now let strategy: any DecodingStrategy - if let schema = chatRequest.responseFormat?.extractedSchema { + let schema = chatRequest.responseFormat?.extractedSchema + if let schema { strategy = ConstrainedDecodingStrategy(jsonSchema: schema, vocabSize: state.config.vocabSize) } else { strategy = VanillaDecodingStrategy() } - let tokenStream = try await strategy.decode( + let lane = BatchServingRouter.lane(hasResponseSchema: schema != nil) + let tokens = try await strategy.decode( from: prepared.input, tokenizer: state.tokenizer, - inferenceEngine: state.engine, + inferenceEngine: state.engineForRequest(lane: lane), samplingConfiguration: prepared.samplingConfig, options: InferenceOptions(maxTokens: prepared.requestMaxTokens), stopSequences: prepared.stopSequences ) + return StartedDecode(tokens: tokens, start: start) +} + +/// Iterates `started.tokens`, splits reasoning and tool calls out of the text unless the request is +/// raw, and passes each `Delta` to `emit`. Records the request in `state.stats` and returns its +/// timing, token counts and finish reason. Emits no role chunk, finish chunk or `[DONE]` frame. +func runStreamingLoop( + started: StartedDecode, prepared: PreparedStream, chatRequest: ChatCompletionRequest, + state: ServerState, emit: (ChatCompletionChunk.Delta) async throws -> Void +) async throws -> StreamingOutcome { + let genStart = started.start + let tokenStream = started.tokens var thinkParser = ThinkTagParser(format: state.thinkingFormat) var toolParser = state.makeToolCallParser() @@ -583,19 +613,21 @@ func runStreamingLoop( ) } +/// Starts a streaming chat completion and returns its SSE response, which releases `access` when +/// the stream ends. On a throw the caller still owns `access`. Throws `ServerError.badRequest` for +/// a bad request and `InferenceRuntimeError.queueFull` when the batch scheduler is full. private func handleStreamingRequest( - chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil, permit: QueuePermit + chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil, access: RequestAccess ) async throws -> Response { - // Preflight throws for bad requests before any SSE framing (mapped to HTTP 400 upstream). let prepared = try await prepareStreaming(chatRequest: chatRequest, state: state, sessionID: sessionID) + let started = try await startStreamingDecode(prepared: prepared, chatRequest: chatRequest, state: state) let created = Int(Date().timeIntervalSince1970) let responseBody = ResponseBody { writer in - // The closure owns the permit for the stream's lifetime; release-on-deinit - // covers a response the framework drops before entering this closure. - defer { permit.release() } + // If this closure never runs, the permit and lease release on deinit. + defer { access.release() } do { let encoder = JSONEncoder() @@ -613,7 +645,8 @@ private func handleStreamingRequest( try await writeDelta(.init(role: "assistant", content: nil)) let outcome = try await runStreamingLoop( - prepared: prepared, chatRequest: chatRequest, state: state, emit: { try await writeDelta($0) }) + started: started, prepared: prepared, chatRequest: chatRequest, state: state, + emit: { try await writeDelta($0) }) try await writeDelta(.init(role: nil, content: nil), finishReason: outcome.finishReason) try await writer.write(ByteBuffer(string: "data: [DONE]\n\n")) diff --git a/swift/Sources/Tools/llm-server/CompletionHandler.swift b/swift/Sources/Tools/llm-server/CompletionHandler.swift index 7962148f..6a981335 100644 --- a/swift/Sources/Tools/llm-server/CompletionHandler.swift +++ b/swift/Sources/Tools/llm-server/CompletionHandler.swift @@ -30,9 +30,9 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw status: .notImplemented, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - let permit: QueuePermit + let access: RequestAccess do { - permit = try await state.queue.acquire() + access = try await state.acquireScoringAccess() } catch let error as ServerError { let err = ErrorResponse( error: .init( @@ -42,7 +42,7 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw status: .tooManyRequests, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - defer { permit.release() } + defer { access.release() } let req: CompletionRequest do { req = try JSONDecoder().decode(CompletionRequest.self, from: body) @@ -81,12 +81,26 @@ private func handleLoglikelihood(req: CompletionRequest, state: ServerState) asy let t0 = SuspendingClock().now var choices: [CompletionResponse.CompletionChoice] = [] - for (idx, prompt) in req.prompts.enumerated() { - let choice = try await processOnePrompt( - prompt: prompt, index: idx, - wantsEcho: wantsEcho, topN: topN, state: state - ) - choices.append(choice) + let prepared = try req.prompts.map { try preparePrompt($0, state: state) } + let scorable = prepared.compactMap { $0 } + var scores: [LogProbabilities] = [] + if let scheduler = state.scoringScheduler { + scores = try await scheduler.score( + scorable.map { ScoreRequest(tokens: $0.tokens, targets: $0.targets, topK: topN) }) + } else { + for prompt in scorable { + scores.append(try await scoreOnEngine(prompt, topN: topN, state: state)) + } + } + + var nextScore = scores.makeIterator() + for (index, prompt) in prepared.enumerated() { + guard let prompt, let probs = nextScore.next() else { + choices.append(.init(index: index, text: "", logprobs: nil, finishReason: "stop")) + continue + } + choices.append( + makeChoice(prompt, probs: probs, index: index, wantsEcho: wantsEcho, topN: topN, state: state)) } let elapsed = SuspendingClock().now - t0 @@ -115,11 +129,16 @@ private func handleLoglikelihood(req: CompletionRequest, state: ServerState) asy ) } -private func processOnePrompt( - prompt: Prompt, index: Int, - wantsEcho: Bool, topN: Int, - state: ServerState -) async throws -> CompletionResponse.CompletionChoice { +/// Tokenized prompt to score, with the target token each position predicts. +private struct PreparedPrompt { + let tokens: [Int32] + /// `tokens.dropFirst()` followed by `tokens[0]` as padding. Same count as `tokens`. + let targets: [Int32] +} + +/// Tokenizes `prompt` into tokens and targets. Returns nil when it has fewer than two tokens. +/// Throws `ServerError.badRequest` when it has more than `state.config.maxContextLength` tokens. +private func preparePrompt(_ prompt: Prompt, state: ServerState) throws -> PreparedPrompt? { let allTokens: [Int32] switch prompt { case .tokenIds(let ids): @@ -128,31 +147,31 @@ private func processOnePrompt( allTokens = state.tokenizer.encode(text: text).map { Int32($0) } } - guard allTokens.count >= 2 else { - return .init(index: index, text: "", logprobs: nil, finishReason: "stop") - } + guard allTokens.count >= 2 else { return nil } guard allTokens.count <= state.config.maxContextLength else { throw ServerError.badRequest( "Prompt length \(allTokens.count) exceeds max context \(state.config.maxContextLength)") } + return PreparedPrompt(tokens: allTokens, targets: Array(allTokens.dropFirst()) + [allTokens[0]]) +} - // lm-eval sends echo=true, max_tokens=1 and slices logprobs[ctxlen:-1]. - // The :-1 discards the last entry. We append a padding token so :-1 - // drops it instead of discarding real data. - let paddingToken = allTokens[0] - let continuation = Array(allTokens.dropFirst()) + [paddingToken] - +/// Scores `prompt` on the `.fallbackBatchOne` engine after a reset, by feeding `tokens[0]` and +/// forcing `targets`. Reports the top `topN` alternatives per position. +private func scoreOnEngine(_ prompt: PreparedPrompt, topN: Int, state: ServerState) async throws + -> LogProbabilities +{ let options = InferenceOptions( - maxTokens: continuation.count, + maxTokens: prompt.targets.count, includeLogits: true, - forcedContinuation: continuation + forcedContinuation: prompt.targets ) - try await state.engine.reset() + let scoringEngine = state.engineForRequest(lane: .fallbackBatchOne) + try await scoringEngine.reset() - let stream = try await state.engine.generate( - with: [allTokens[0]], + let stream = try await scoringEngine.generate( + with: [prompt.tokens[0]], samplingConfiguration: .greedy, inferenceOptions: options ) @@ -163,8 +182,14 @@ private func processOnePrompt( allLogits.append(logits) } } + return LogProbabilities.compute(logits: allLogits, targets: prompt.targets, topK: topN) +} - let probs = LogProbabilities.compute(logits: allLogits, targets: continuation, topK: topN) +private func makeChoice( + _ prompt: PreparedPrompt, probs: LogProbabilities, index: Int, wantsEcho: Bool, topN: Int, + state: ServerState +) -> CompletionResponse.CompletionChoice { + let allTokens = prompt.tokens var tokens: [String] = [] var tokenLogprobs: [Double?] = [] diff --git a/swift/Sources/Tools/llm-server/LLMServerMain.swift b/swift/Sources/Tools/llm-server/LLMServerMain.swift index bdc8e87d..3a721b52 100644 --- a/swift/Sources/Tools/llm-server/LLMServerMain.swift +++ b/swift/Sources/Tools/llm-server/LLMServerMain.swift @@ -52,6 +52,13 @@ struct LLMServer: AsyncParsableCommand { @Option(name: .customLong("replay-output"), help: "Write replay results JSONL here (default: stdout)") var replayOutput: String? + @Flag( + name: .customLong("replay-concurrent"), + help: ArgumentHelp( + "Run all --replay requests concurrently instead of one at a time. Use with --concurrent-requests N.") + ) + var replayConcurrent: Bool = false + @Option(name: .customLong("variant"), help: "Engine variant: auto, coreai-pipelined, coreai-sequential") var inferenceEngineVariant: String = "default" @@ -78,6 +85,14 @@ struct LLMServer: AsyncParsableCommand { @Option(name: .customLong("max-queue-depth"), help: "Max requests queued before returning 429 (default: 16)") var maxQueueDepth: Int = 16 + @Option( + name: .customLong("concurrent-requests"), + help: ArgumentHelp( + "Max requests decoded together in one batch (default: 1). Values above 1 enable the batch " + + "scheduler when the model supports batching, and are ignored with a warning otherwise.") + ) + var concurrentRequests: Int = 1 + @Flag( name: .customLong("no-thinking"), help: "Disable thinking/reasoning (alias for --default-reasoning-effort none)" @@ -103,6 +118,9 @@ struct LLMServer: AsyncParsableCommand { guard maxQueueDepth >= 0 else { throw ValidationError("--max-queue-depth must be >= 0 (got \(maxQueueDepth))") } + guard concurrentRequests >= 1 else { + throw ValidationError("--concurrent-requests must be >= 1 (got \(concurrentRequests))") + } do { _ = try ReasoningEffort.resolveDefault( defaultReasoningEffort: defaultReasoningEffort, noThinking: noThinking) @@ -218,14 +236,39 @@ struct LLMServer: AsyncParsableCommand { maxQueueDepth: maxQueueDepth ) + // The scheduler holds up to --concurrent-requests plus --max-queue-depth requests, then rejects. + var batchScheduler: BatchScheduler? = nil + if concurrentRequests > 1 { + if let seq = engine as? CoreAISequentialEngine, seq.supportsBatching { + let assetMax = bundle.maxBatchSize + guard concurrentRequests <= assetMax else { + throw ValidationError( + "--concurrent-requests \(concurrentRequests) exceeds the asset's max batch size \(assetMax). " + + "Re-export with --dynamic-batch-size \(concurrentRequests) or higher, or lower --concurrent-requests." + ) + } + batchScheduler = try BatchScheduler( + engine: seq, maxConcurrent: concurrentRequests, maxQueued: maxQueueDepth) + print("Batched serving: \(concurrentRequests) concurrent requests (asset max batch \(assetMax))") + } else { + print( + "⚠️ --concurrent-requests \(concurrentRequests) ignored: batched serving needs the coreai-sequential engine " + + "and an asset exported with --attention-mask batched --dynamic-batch-size N. " + + "Serving one request at a time.") + } + } + let state = ServerState( engine: engine, tokenizer: tokenizer, - config: config + config: config, + batchServing: batchScheduler.map { (scheduler: $0, modelConfig: engineConfig) } ) if let replayPath { - try await ReplayRunner.run(inputPath: replayPath, outputPath: replayOutput, state: state) + try await ReplayRunner.run( + inputPath: replayPath, outputPath: replayOutput, state: state, + concurrent: replayConcurrent) return } diff --git a/swift/Sources/Tools/llm-server/ReplayRunner.swift b/swift/Sources/Tools/llm-server/ReplayRunner.swift index e40201d4..7965cd9d 100644 --- a/swift/Sources/Tools/llm-server/ReplayRunner.swift +++ b/swift/Sources/Tools/llm-server/ReplayRunner.swift @@ -4,40 +4,36 @@ // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause import CoreAILMCommon +import CoreAILanguageModels import Foundation /// Drives the server's generation core over a JSONL request file instead of HTTP. /// /// Requests run in arrival-timestamp order through the same cores the HTTP handler uses: /// `runChatCompletion` for `stream:false`, `runStreamingLoop` for `stream:true` (deltas folded -/// back into the result). The engine is single-active, so requests run sequentially, matching the -/// server. Interleaving sessions in the input reproduces the same per-session cache behavior. Each -/// result is one JSONL line with per-request instrumentation. +/// back into the result). By default requests run one at a time. With `concurrent`, all requests +/// start at once, and requests over the queue limits produce queue-full error lines. Output lines +/// keep arrival-timestamp order in both modes. Interleaving sessions in the input reproduces the +/// same per-session cache behavior. Each result is one JSONL line with per-request instrumentation. enum ReplayRunner { - static func run(inputPath: String, outputPath: String?, state: ServerState) async throws { + static func run( + inputPath: String, outputPath: String?, state: ServerState, concurrent: Bool = false + ) async throws { let text = try String(contentsOfFile: inputPath, encoding: .utf8) let requests = try ReplayIO.parseRequests(text) - var lines: [String] = [] - lines.reserveCapacity(requests.count) - for req in requests { - let sessionID = req.session ?? "default" - let result: ReplayResult - do { - if req.isStreaming { - result = try await runStreamed(req, sessionID: sessionID, state: state) - } else { - result = try await runNonStreamed(req, sessionID: sessionID, state: state) + var lines = [String](repeating: "", count: requests.count) + if concurrent { + try await withThrowingTaskGroup(of: (Int, String).self) { group in + for (i, req) in requests.enumerated() { + group.addTask { (i, await oneLine(req, state: state)) } } - } catch { - result = ReplayResult( - id: req.id, session: req.session, t: req.t, - choices: nil, systemFingerprint: nil, - promptTokens: 0, completionTokens: 0, prefixReuseTokens: 0, - ttftMs: 0, totalMs: 0, decodeTps: 0, - streamed: req.isStreaming, error: "\(error)") + for try await (i, line) in group { lines[i] = line } + } + } else { + for (i, req) in requests.enumerated() { + lines[i] = await oneLine(req, state: state) } - lines.append(try ReplayIO.encodeResult(result)) } let output = lines.joined(separator: "\n") + (lines.isEmpty ? "" : "\n") @@ -48,6 +44,33 @@ enum ReplayRunner { } } + /// Runs `req` while holding the access for its lane and returns its encoded result line. An + /// error, including queue-full, produces an error result line. + private static func oneLine(_ req: ReplayRequest, state: ServerState) async -> String { + let sessionID = req.session ?? "default" + let result: ReplayResult + do { + let lane = BatchServingRouter.lane( + hasResponseSchema: req.request.responseFormat?.extractedSchema != nil) + let access = try await state.acquireAccess(lane: lane) + defer { access.release() } + + if req.isStreaming { + result = try await runStreamed(req, sessionID: sessionID, state: state) + } else { + result = try await runNonStreamed(req, sessionID: sessionID, state: state) + } + } catch { + result = ReplayResult( + id: req.id, session: req.session, t: req.t, + choices: nil, systemFingerprint: nil, + promptTokens: 0, completionTokens: 0, prefixReuseTokens: 0, + ttftMs: 0, totalMs: 0, decodeTps: 0, + streamed: req.isStreaming, error: "\(error)") + } + return (try? ReplayIO.encodeResult(result)) ?? "{\"error\":\"encode failed\"}" + } + /// Non-streaming path: the same `runChatCompletion` core the HTTP handler uses. private static func runNonStreamed(_ req: ReplayRequest, sessionID: String, state: ServerState) async throws -> ReplayResult @@ -76,8 +99,9 @@ enum ReplayRunner { { let prepared = try await prepareStreaming(chatRequest: req.request, state: state, sessionID: sessionID) var aggregator = ReplayStreamAggregator() + let started = try await startStreamingDecode(prepared: prepared, chatRequest: req.request, state: state) let outcome = try await runStreamingLoop( - prepared: prepared, chatRequest: req.request, state: state, + started: started, prepared: prepared, chatRequest: req.request, state: state, emit: { aggregator.consume($0) }) let genSeconds = max(0, outcome.totalSeconds - outcome.ttftSeconds) let decodeTps = genSeconds > 0 ? Double(outcome.genTokenCount) / genSeconds : 0 diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index da5e8147..c789bcf3 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -127,8 +127,14 @@ final class ServerState: @unchecked Sendable { /// Stable identifier for the backend configuration, echoed on responses as /// `system_fingerprint`. Reproducible outputs (same `seed` + prompt) are only /// guaranteed while this value is unchanged; it changes when output-affecting - /// configuration (model, engine, context length, sampling defaults) changes. + /// configuration (model, engine, batch size, context length, sampling defaults) changes. The + /// batch size is part of it because batched decoding sums in a different order, so the same + /// seed can give slightly different logits. let systemFingerprint: String + /// Scheduler for the batched lane and completions scoring, and the model config for each + /// `SchedulerSession`. Nil when batched serving is off. + private let batchServing: (scheduler: BatchScheduler, modelConfig: ModelConfig)? + private var batchScheduler: BatchScheduler? { batchServing?.scheduler } private let _state = Mutex(InternalState()) private struct InternalState { @@ -140,23 +146,106 @@ final class ServerState: @unchecked Sendable { var totalToolCalls: Int = 0 } - init(engine: any InferenceEngine, tokenizer: any Tokenizer, config: ServerConfig) { + init( + engine: any InferenceEngine, + tokenizer: any Tokenizer, + config: ServerConfig, + batchServing: (scheduler: BatchScheduler, modelConfig: ModelConfig)? = nil + ) { self.engine = engine self.tokenizer = tokenizer self.config = config self.queue = RequestQueue(maxDepth: config.maxQueueDepth) self.toolCallDetection = detectToolCallFormat(using: tokenizer) self.thinkingFormat = detectThinkingFormat(using: tokenizer) - self.systemFingerprint = Self.makeSystemFingerprint(config: config, engine: engine) + self.systemFingerprint = Self.makeSystemFingerprint( + config: config, engine: engine, batchSize: batchServing?.scheduler.maxConcurrent ?? 1) + self.batchServing = batchServing + } + + /// Token ids that end a row in the scheduler: the tokenizer's EOS id and + /// `config.additionalEosTokenIds`. Stop sequences are not included. + private var schedulerStopTokenIds: Set { + var ids = Set() + if let eos = tokenizer.eosTokenId { ids.insert(Int32(eos)) } + ids.formUnion(config.additionalEosTokenIds.map { Int32($0) }) + return ids + } + + /// Returns the engine for one request on `lane`. Call once per request. + /// + /// - `.batched`: a new `SchedulerSession` over the batch scheduler, or the shared engine when + /// batched serving is off. + /// - `.fallbackBatchOne`: the shared engine. With a batch scheduler running, hold the access from + /// `acquireAccess(lane:)` while using it. + func engineForRequest(lane: BatchServingLane) -> any InferenceEngine { + switch lane { + case .batched: + guard let batchServing else { return engine } + return SchedulerSession( + sink: batchServing.scheduler, config: batchServing.modelConfig, + stopTokenIds: schedulerStopTokenIds) + case .fallbackBatchOne: + return engine + } + } + + /// Waits for and returns the access a request on `lane` holds while it uses the engine. + /// + /// - Single-slot permit from `queue`: on every lane except `.batched` with a scheduler that + /// supports continuous batching. + /// - `ExclusiveEngineLease`: on `.fallbackBatchOne` when `batchScheduler` is set. The batch + /// does not step while the lease is held. + /// + /// Throws `ServerError.queueFull` when the queue is full, and `CancellationError` when the task + /// is cancelled while waiting for the permit. + /// + /// Pass the same lane here and to `engineForRequest(lane:)`, and release the access when the + /// request is done: + /// + /// ```swift + /// let lane = BatchServingRouter.lane(hasResponseSchema: schema != nil) + /// let access = try await state.acquireAccess(lane: lane) + /// defer { access.release() } + /// let engine = state.engineForRequest(lane: lane) + /// ``` + func acquireAccess(lane: BatchServingLane) async throws -> RequestAccess { + let permit = needsServerPermit(lane: lane) ? try await queue.acquire() : nil + let lease = lane == .fallbackBatchOne ? await batchScheduler?.acquireExclusiveAccess() : nil + return RequestAccess(permit: permit, lease: lease) + } + + /// Scheduler that scores completion prompts in batches, or nil when batched serving is off. + var scoringScheduler: BatchScheduler? { batchScheduler } + + /// Waits for and returns the single-slot permit for a scoring request, with no engine lease. + /// Throws `ServerError.queueFull` when the queue is full, and `CancellationError` when the task + /// is cancelled while waiting. + func acquireScoringAccess() async throws -> RequestAccess { + RequestAccess(permit: try await queue.acquire(), lease: nil) + } + + /// True when a request on `lane` needs the single-slot permit: every lane except `.batched` + /// with a scheduler that supports continuous batching. + private func needsServerPermit(lane: BatchServingLane) -> Bool { + switch lane { + case .batched: + return !(batchScheduler?.supportsContinuousBatching ?? false) + case .fallbackBatchOne: + return true + } } /// Deterministic FNV-1a hash of a canonical config descriptor. Deterministic (not /// Swift's per-process-seeded `Hasher`) so two identical launches produce the same /// fingerprint and a config change produces a different one. - private static func makeSystemFingerprint(config: ServerConfig, engine: any InferenceEngine) -> String { + private static func makeSystemFingerprint( + config: ServerConfig, engine: any InferenceEngine, batchSize: Int + ) -> String { let descriptor = [ config.modelName, String(describing: type(of: engine)), + "batch=\(batchSize)", "ctx=\(config.maxContextLength)", "t=\(config.defaultTemperature)", "k=\(config.defaultTopK.map { "\($0)" } ?? "nil")", @@ -315,3 +404,20 @@ enum RequestID { return "coreai-\(n)" } } + +/// Single-slot permit and engine lease a request holds while it uses the engine. Either can be nil. +/// `release()` can be called more than once, and the permit and lease also release on deinit. +final class RequestAccess: Sendable { + private let permit: QueuePermit? + private let lease: ExclusiveEngineLease? + + init(permit: QueuePermit?, lease: ExclusiveEngineLease?) { + self.permit = permit + self.lease = lease + } + + func release() { + lease?.release() + permit?.release() + } +} diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift new file mode 100644 index 00000000..e91ebee3 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -0,0 +1,465 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Synchronization +import Testing + +@testable import CoreAILanguageModels + +/// Makes logits over a 64-token vocabulary with the maximum at `token`. +private func oneHot(_ token: Int32) -> [LogitsScalarType] { + var row = [LogitsScalarType](repeating: 0, count: 64) + row[Int(token)] = 1 + return row +} + +/// Makes a session state backed by a one-element mock KV cache. +private func dummySession() -> GenerationSessionState { + GenerationSessionState( + kvCache: MockStateHandler(names: ["k"], shape: [1, 1, 1, 1, 1]), + additionalStates: nil, + hasNonTruncatableStates: false) +} + +/// Makes flattened count-up logits, `row.count * 64` values per row, where position `p` peaks at +/// that row's token at `p` plus one, mod 64. +private func countUpAllPositions(_ tokensPerRow: [[Int32]]) -> [LogitsScalarType] { + tokensPerRow.flatMap { row in row.flatMap { oneHot(($0 + 1) % 64) } } +} + +/// Collects a request stream's tokens and its `.finished` reason. +private func drain( + _ stream: AsyncThrowingStream +) async throws -> (tokens: [Int32], reason: StopReason?) { + var tokens: [Int32] = [] + var reason: StopReason? + for try await event in stream { + switch event { + case .token(let token, _): tokens.append(token) + case .finished(let stopReason): reason = stopReason + } + } + return (tokens, reason) +} + +/// Mock `BatchedDecodeStep` that predicts the last input token plus one, mod 64. A request with +/// prompt `[..., k]` and `maxTokens: m` generates `k+1` through `k+m`. +/// +/// Records the forward count, each decode's `startPositions`, each scoring forward's row count, and +/// the peak number of overlapping forward calls. `forwardDelay` is a sleep added to every forward call. +private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { + struct Boom: Error {} + + let maxBatchCapacity: Int + /// `prefillRow` throws `Boom` for a prompt that starts with this token. + private let failPrefillFor: Int32? + /// When true, every `decodeBatch` throws `Boom`. + private let failDecode: Bool + private let forwardDelay: Duration? + + private struct Record { + var startPositions: [[Int]] = [] + var inFlight = 0 + var maxInFlight = 0 + var forwards = 0 + var scoringRows: [Int] = [] + var scoringPositions: [Int] = [] + } + private let record = Mutex(Record()) + + var observedStartPositions: [[Int]] { record.withLock { $0.startPositions } } + var maxInFlight: Int { record.withLock { $0.maxInFlight } } + /// Row count of each scoring forward, in call order. + var scoringRows: [Int] { record.withLock { $0.scoringRows } } + var scoringPositions: [Int] { record.withLock { $0.scoringPositions } } + /// Number of prefill, decode and scoring forward calls so far. + var forwards: Int { record.withLock { $0.forwards } } + + init( + capacity: Int, failPrefillFor: Int32? = nil, failDecode: Bool = false, + forwardDelay: Duration? = nil + ) { + self.maxBatchCapacity = capacity + self.failPrefillFor = failPrefillFor + self.failDecode = failDecode + self.forwardDelay = forwardDelay + } + + private func enter() async { + record.withLock { + $0.forwards += 1 + $0.inFlight += 1 + $0.maxInFlight = max($0.maxInFlight, $0.inFlight) + } + if let forwardDelay { try? await Task.sleep(for: forwardDelay) } + } + + private func leave() { record.withLock { $0.inFlight -= 1 } } + + func prefillRow( + row: Int, promptTokens: [Int32], session: GenerationSessionState + ) async throws -> [LogitsScalarType] { + await enter() + defer { leave() } + if let failPrefillFor, promptTokens.first == failPrefillFor { throw Boom() } + return oneHot(((promptTokens.last ?? 0) + 1) % 64) + } + + func decodeBatch( + tokensPerRow: [[Int32]], startPositions: [Int], session: GenerationSessionState + ) async throws -> [[LogitsScalarType]] { + await enter() + defer { leave() } + record.withLock { $0.startPositions.append(startPositions) } + if failDecode { throw Boom() } + return tokensPerRow.map { oneHot((($0.last ?? 0) + 1) % 64) } + } + + func allPositionLogits( + tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState + ) async throws -> [LogitsScalarType] { + await enter() + defer { leave() } + record.withLock { + $0.scoringRows.append(tokensPerRow.count) + $0.scoringPositions.append(tokensPerRow.first?.count ?? 0) + } + return countUpAllPositions(tokensPerRow) + } + + func makeSessionState(batchSize: Int) throws -> GenerationSessionState { dummySession() } +} + +/// Mock count-up `BatchedDecodeStep` whose `prefillRow` blocks until `release()`. +private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { + let maxBatchCapacity: Int + + private struct Gate { + var released = false + var releaseWaiter: CheckedContinuation? + } + private let gate = Mutex(Gate()) + + init(capacity: Int) { maxBatchCapacity = capacity } + + /// Unblocks the waiting `prefillRow` call. Later calls return without blocking. + func release() { + let waiter: CheckedContinuation? = gate.withLock { g in + g.released = true + defer { g.releaseWaiter = nil } + return g.releaseWaiter + } + waiter?.resume() + } + + func prefillRow( + row: Int, promptTokens: [Int32], session: GenerationSessionState + ) async throws -> [LogitsScalarType] { + await withCheckedContinuation { (cont: CheckedContinuation) in + gate.withLock { g in + if g.released { cont.resume() } else { g.releaseWaiter = cont } + } + } + return oneHot((promptTokens.last ?? 0) + 1) + } + + func decodeBatch( + tokensPerRow: [[Int32]], startPositions: [Int], session: GenerationSessionState + ) async throws -> [[LogitsScalarType]] { + tokensPerRow.map { oneHot(($0.last ?? 0) + 1) } + } + + func allPositionLogits( + tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState + ) async throws -> [LogitsScalarType] { + countUpAllPositions(tokensPerRow) + } + + func makeSessionState(batchSize: Int) throws -> GenerationSessionState { dummySession() } +} + +@Suite("BatchScheduler") +struct BatchSchedulerTests { + private func makeScheduler( + _ step: any BatchedDecodeStep, maxConcurrent: Int, maxQueued: Int? = nil + ) -> BatchScheduler { + BatchScheduler( + decodeStep: step, + session: dummySession(), + maxConcurrent: maxConcurrent, + maxQueued: maxQueued) + } + + private func req( + prompt: [Int32], maxTokens: Int, eos: Set = [], stops: [[Int32]] = [] + ) -> BatchRequest { + BatchRequest( + promptTokens: prompt, + sampling: .greedy, + eosTokenIds: eos, + maxTokens: maxTokens, + stopSequences: StopSequences(sequences: stops), + decodeStopText: { $0.map(String.init).joined(separator: ",") }) + } + + @Test("Rows retire independently on EOS, max tokens and stop sequence") + func independentRetirement() async throws { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 4), maxConcurrent: 4) + let results = try await scheduler.run([ + req(prompt: [9], maxTokens: 100, eos: [12]), + req(prompt: [20], maxTokens: 5), + req(prompt: [5], maxTokens: 100, stops: [[8]]), + ]) + #expect(results.map(\.tokens) == [[10, 11, 12], [21, 22, 23, 24, 25], [6, 7, 8]]) + #expect(results.map(\.stopReason) == [.eos, .maxTokens, .stopSequence("8")]) + } + + @Test("Rows decode at their own cursors, and a retired row's slot admits the next request") + func raggedCursorsAndSlotReuse() async throws { + let mock = CountUpDecodeStep(capacity: 2) + let scheduler = makeScheduler(mock, maxConcurrent: 2) + let results = try await scheduler.run([ + req(prompt: [1], maxTokens: 2), + req(prompt: [1, 2, 3], maxTokens: 3), + req(prompt: [20], maxTokens: 1), + ]) + #expect(results.map(\.tokens) == [[2, 3], [4, 5, 6], [21]]) + #expect(mock.observedStartPositions.first == [1, 3]) + } + + @Test("Forwards from run and submit never overlap") + func singlePump() async throws { + let mock = CountUpDecodeStep(capacity: 4, forwardDelay: .milliseconds(5)) + let scheduler = makeScheduler(mock, maxConcurrent: 4) + let stream = try scheduler.submit(req(prompt: [1, 2], maxTokens: 6)) + async let ran = scheduler.run([req(prompt: [1, 2], maxTokens: 6)]) + _ = try await drain(stream) + _ = try await ran + #expect(mock.maxInFlight == 1) + } + + @Test( + "A burst of submits fills the free rows plus maxQueued, then throws queueFull", + arguments: [(1, 3), (0, 2)] as [(Int, Int)]) + func backpressure(maxQueued: Int, accepted: Int) async throws { + let mock = GatedPrefillStep(capacity: 2) + let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: maxQueued) + let streams = try (0.. ScoreRequest { + ScoreRequest(tokens: tokens, targets: tokens.dropFirst() + [63], topK: 2) + } + + @Test( + "Batched scoring matches scoring each request alone, across padding and chunks", + arguments: [ + (2, [2, 2, 2]), // A static batch=2 graph pads every forward to two rows. + (Int.max, [2, 1, 1]), // A dynamic-batch graph uses one row per request. + ] as [(Int, [Int])]) + func scoringBatchInvariance(capacity: Int, expectedRows: [Int]) async throws { + let step = CountUpDecodeStep(capacity: capacity) + let scheduler = makeScheduler(step, maxConcurrent: 2) + // Three requests, at most two per batch. The 200-token one spans two 128-position chunks. + let prompts: [[Int32]] = [[1, 2, 3], [7, 8, 9, 10, 11], (0..<200).map { Int32($0 % 60) }] + let together = try await scheduler.score(prompts.map(scoreRequest)) + #expect(step.scoringRows == expectedRows) + + // Argmax targets score highest, and each entry reports topK alternatives. + let first = together[0].entries + #expect(first.map(\.tokenId) == [2, 3, 63]) + #expect(first[0].value == first[1].value && first[2].value < first[0].value) + #expect(first[0].alternatives.count == 2) + + for (prompt, batched) in zip(prompts, together) { + let alone = try await scheduler.score([scoreRequest(prompt)])[0] + #expect(batched.entries.map(\.tokenId) == alone.entries.map(\.tokenId)) + #expect(batched.entries.map(\.value) == alone.entries.map(\.value)) + } + } + + @Test("Scoring pads the last chunk and the rows to a few fixed shapes, with unchanged scores") + func scoringPadsShapes() async throws { + let step = CountUpDecodeStep(capacity: Int.max) + let scheduler = makeScheduler(step, maxConcurrent: 4) + // One batch of three short prompts: 3 rows pad to 4, 6 positions pad to 16. + let prompts: [[Int32]] = [[1, 2, 3], [7, 8, 9, 10, 11], [20, 21, 22, 23, 24, 25]] + let together = try await scheduler.score(prompts.map(scoreRequest)) + #expect(step.scoringRows == [4]) + #expect(step.scoringPositions == [16]) + for (prompt, batched) in zip(prompts, together) { + let alone = try await scheduler.score([scoreRequest(prompt)])[0] + #expect(batched.entries.count == prompt.count) + #expect(batched.entries.map(\.tokenId) == alone.entries.map(\.tokenId)) + #expect(batched.entries.map(\.value) == alone.entries.map(\.value)) + } + } + + @Test( + "The last scoring chunk pads to a power of two, at least 16, within the context", + arguments: [ + (0, 7, Int.max, 16), + (0, 17, Int.max, 32), + (0, 300, Int.max, 128), + (256, 300, Int.max, 64), + (256, 300, 300, 44), // Padding to 64 would run past a 300-position context. + (128, 256, Int.max, 128), + ] as [(Int, Int, Int, Int)]) + func scoringChunkLength(start: Int, longest: Int, maxPositions: Int, expected: Int) { + let chunk = BatchScheduler.Scorer.paddedChunk( + start: start, longest: longest, chunkLength: 128, maxPositions: maxPositions) + #expect(chunk == expected) + } + + @Test( + "Scoring rows: a static graph's capacity, else a power of two up to maxConcurrent", + arguments: [ + (1, 4, 4, 4), + (3, Int.max, 8, 4), + (3, Int.max, 3, 3), + (5, Int.max, 6, 6), + (1, Int.max, 8, 1), + ] as [(Int, Int, Int, Int)]) + func scoringRowCount(requests: Int, capacity: Int, maxConcurrent: Int, expected: Int) { + let rows = BatchScheduler.Scorer.paddedRows( + requests: requests, capacity: capacity, maxConcurrent: maxConcurrent) + #expect(rows == expected) + } + + @Test("Scoring runs between decode steps while a request streams") + func scoringBesideDecode() async throws { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + let stream = try scheduler.submit(req(prompt: [5], maxTokens: 6)) + let scores = try await scheduler.score([scoreRequest([1, 2, 3])]) + let (tokens, _) = try await drain(stream) + #expect(tokens == [6, 7, 8, 9, 10, 11]) + #expect(scores[0].entries.count == 3) + } + + @Test( + "Scoring batches stop before a third of the positions would be padding", + arguments: [ + ([7, 40, 130, 300], 1), + ([130, 300], 2), + ([64, 64, 64, 64, 64], 4), + ([], 0), + ] as [([Int], Int)]) + func scoringBatchCount(lengths: [Int], expected: Int) { + #expect(BatchScheduler.Scorer.batchCount(lengths, limit: 4) == expected) + } +} + +@Suite("SchedulerSession.GenerationSequence") +struct SchedulerSessionSequenceTests { + @Test( + "Maps tokens and takes the stop reason from .finished, else .maxTokens", + arguments: [StopReason.eos, nil]) + func mapsEvents(finished: StopReason?) async throws { + let (stream, continuation) = AsyncThrowingStream.makeStream(of: BatchSinkEvent.self) + continuation.yield(.token(10, logits: nil)) + continuation.yield(.token(20, logits: nil)) + if let finished { continuation.yield(.finished(finished)) } + continuation.finish() + + let sequence = SchedulerSession.GenerationSequence(events: stream, cancellation: RequestCancellation()) + var tokens: [Int32] = [] + for try await output in sequence { tokens.append(output.tokenId) } + #expect(tokens == [10, 20]) + #expect(sequence.stopReason == (finished ?? .maxTokens)) + } +} diff --git a/swift/Tests/LanguageModelsTests/BatchServingRouterTests.swift b/swift/Tests/LanguageModelsTests/BatchServingRouterTests.swift new file mode 100644 index 00000000..3de5f452 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/BatchServingRouterTests.swift @@ -0,0 +1,18 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Testing + +@testable import CoreAILanguageModels + +@Suite("BatchServingRouter") +struct BatchServingRouterTests { + @Test( + "Plain chat batches and a response schema takes the batch=1 lane", + arguments: [(false, BatchServingLane.batched), (true, .fallbackBatchOne)]) + func lane(hasResponseSchema: Bool, expected: BatchServingLane) { + #expect(BatchServingRouter.lane(hasResponseSchema: hasResponseSchema) == expected) + } +} diff --git a/swift/Tests/LanguageModelsTests/BatchSessionAdapterTests.swift b/swift/Tests/LanguageModelsTests/BatchSessionAdapterTests.swift new file mode 100644 index 00000000..25de2720 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/BatchSessionAdapterTests.swift @@ -0,0 +1,37 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Testing + +@testable import CoreAILanguageModels + +@Suite("BatchSessionAdapter") +struct BatchSessionAdapterTests { + private func adapter( + eos: Set = [], maxTokens: Int = 1_000, stops: [[Int32]] = [] + ) -> BatchSessionAdapter { + BatchSessionAdapter( + request: BatchRequest( + promptTokens: [], + sampling: .greedy, + eosTokenIds: eos, + maxTokens: maxTokens, + stopSequences: StopSequences(sequences: stops), + decodeStopText: { $0.map(String.init).joined(separator: ",") })) + } + + @Test("Each row retires on its own criterion and ignores later tokens") + func retirement() { + var rows = [adapter(eos: [99]), adapter(maxTokens: 2), adapter(stops: [[7, 8]])] + var live: [[Bool]] = [[], [], []] + for token: Int32 in [5, 7, 8, 99] { + for i in rows.indices { live[i].append(rows[i].accept(token)) } + } + // The stop row stays live on 7, a partial match, and retires on 8. + #expect(live == [[true, true, true, false], [true, false, false, false], [true, true, false, false]]) + #expect(rows.map(\.generated) == [[5, 7, 8, 99], [5, 7], [5, 7, 8]]) + #expect(rows.map(\.stopReason) == [.eos, .maxTokens, .stopSequence("7,8")]) + } +}