From 7e1c19b332028efd1335eb23bcabbcb8e1751426 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 01:37:01 -0700 Subject: [PATCH 01/10] Add a continuous-batching scheduler BatchScheduler serves several requests on one batched CoreAISequentialEngine session. Each request takes a free row, is prefilled into it with prefillRow, and then advances one token per batched decode step at its own cursor. A row retires on EOS, a stop sequence or its token limit, and the next waiting request takes its slot. - SchedulerSession presents one scheduled request as an InferenceEngine, so decoding strategies and stop handling work unchanged. - IsolatedSessionEngine runs a request on its own session of the same engine, for requests that need per-step logits. - BatchServingRouter picks between the two. - submit throws queueFull once the scheduler holds maxConcurrent + maxQueued requests. Free rows absorb a burst before the queue limit applies. Empty prompts are rejected. - acquireExclusiveAccess returns a lease that pauses the scheduler at a step boundary, so a request on its own session never runs on the engine at the same time as the batch. - CoreAISequentialEngine.supportsBatching reports whether a graph can serve several requests at once. --- .../InferenceEngines/BatchScheduler.swift | 389 ++++++++++++++++++ .../InferenceEngines/BatchServingRouter.swift | 26 ++ .../BatchSessionAdapter.swift | 65 +++ .../CoreAISequentialEngine.swift | 7 + .../InferenceEngines/InferenceEngine.swift | 10 + .../IsolatedSessionEngine.swift | 67 +++ .../InferenceEngines/SchedulerSession.swift | 126 ++++++ .../BatchSchedulerTests.swift | 366 ++++++++++++++++ .../BatchServingRouterTests.swift | 18 + .../BatchSessionAdapterTests.swift | 37 ++ 10 files changed, 1111 insertions(+) create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift create mode 100644 swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift create mode 100644 swift/Tests/LanguageModelsTests/BatchServingRouterTests.swift create mode 100644 swift/Tests/LanguageModelsTests/BatchSessionAdapterTests.swift diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift new file mode 100644 index 00000000..80e3304c --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -0,0 +1,389 @@ +// 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 + +/// One request submitted to a `RequestSink`. +package struct BatchRequest: Sendable { + /// Must not be empty; `submit(_:)` throws `invalidArgument` for an empty prompt. + let promptTokens: [Int32] + let sampling: SamplingConfiguration + let eosTokenIds: Set + let maxTokens: Int + let stopSequences: StopSequences + /// Turns a matched stop sequence's tokens into the text reported in `StopReason.stopSequence`, + /// so the scheduler needs no tokenizer. The default returns an empty string. + let decodeStopText: @Sendable ([Int32]) -> String + + 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 + } +} + +/// The tokens one request generated in `BatchScheduler.run(_:)` and why it stopped. +struct RowResult: Sendable { + let tokens: [Int32] + let stopReason: StopReason +} + +/// One event on a request's stream from a `RequestSink`. A stream yields `.token` events and then +/// one `.finished`, or ends by throwing. +package enum BatchSinkEvent: Sendable { + /// A sampled token. `BatchScheduler` always sends nil `logits` because it samples internally. + case token(Int32, logits: [LogitsScalarType]?) + /// Why the request stopped. Nothing follows it. + case finished(StopReason) +} + +/// A shared batched backend that gives each submitted request its own event stream. +/// +/// The server does not call `submit(_:)` directly. It wraps the sink in one `SchedulerSession` per +/// request, which submits on `generate` and presents the stream as an ordinary `InferenceEngine`. +/// +/// ```swift +/// let engine = SchedulerSession(sink: scheduler, config: modelConfig, stopTokenIds: eosIds) +/// let sequence = try await engine.generate( +/// with: promptTokens, samplingConfiguration: sampling, inferenceOptions: options) +/// ``` +package protocol RequestSink: Sendable { + /// The most requests that decode in the same batch. + var maxConcurrent: Int { get } + + /// Queues `request` and returns its event stream. When the sink rejects the request, for example + /// with `InferenceRuntimeError.queueFull`, this throws before any stream exists. + func submit(_ request: BatchRequest) throws -> AsyncThrowingStream +} + +/// Runs up to `maxConcurrent` requests through one shared batch-N session and admits new requests +/// while others are still decoding. +/// +/// Slot `i` is row `i` of the batch KV cache. A slot goes through these phases: +/// - Admission: a queued request takes a free slot and is prefilled into its row with +/// `BatchedDecodeStep.prefillRow`. +/// - Decode: each step runs `decodeBatch` over all `N` rows, every live row at its own cursor, so a +/// short request admitted late runs beside an older long one. +/// - Retirement: `BatchSessionAdapter` ends the row on EOS, a stop sequence or `maxTokens`. The next +/// pump iteration refills the slot from the queue while the other rows keep decoding. +/// +/// The class is `@unchecked Sendable` because the queue and row count sit behind a `Mutex`, and at +/// most one pump task runs at a time. That task alone drives the session. +package final class BatchScheduler: RequestSink, @unchecked Sendable { + private let decodeStep: any BatchedDecodeStep + private let session: GenerationSessionState + package let maxConcurrent: Int + /// How many requests may wait beyond the free rows before `submit(_:)` throws. `nil` means + /// unbounded. + private let maxQueued: Int? + + /// Whether the server can skip its single-slot permit for the batched lane. This scheduler queues + /// and bounds requests itself, so that holds whenever it has more than one row. + 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 + var pumping = false + } + private let state = Mutex(State()) + /// Held by the pump for each step and by an `ExclusiveEngineLease` for its whole lifetime, so + /// the engine never runs work from two tasks at once. + private let engineLock = AsyncFIFOLock() + + /// Creates a scheduler with a new `maxConcurrent`-row session on `engine`. + /// `engine.supportsBatching` must be true, and `maxConcurrent` must not exceed the graph's batch + /// capacity or a precondition fails. + package convenience init(engine: CoreAISequentialEngine, maxConcurrent: Int, maxQueued: Int? = nil) throws { + self.init( + decodeStep: engine, session: try engine.makeSessionState(batchSize: maxConcurrent), + maxConcurrent: maxConcurrent, maxQueued: maxQueued) + } + + init( + decodeStep: any BatchedDecodeStep, session: GenerationSessionState, maxConcurrent: Int, + maxQueued: Int? = nil + ) { + 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 + } + + // MARK: - RequestSink + + /// Queues `request` and returns its event stream. + /// + /// Throws `InferenceRuntimeError.queueFull`, which the server maps to HTTP 429, when + /// `queue.count + 1 > freeRows + maxQueued`. With `maxConcurrent == 2`, `maxQueued == 1` and both + /// rows busy, one request may wait and the next one throws. With `maxQueued == nil` it never throws. + /// + /// `queue.count` includes requests the pump has not admitted yet, so free rows absorb a burst + /// before `maxQueued` applies. In the same configuration with both rows free, a burst of four + /// requests accepts three and the fourth throws. + package func submit(_ request: BatchRequest) throws -> AsyncThrowingStream { + try enqueue([request])[0] + } + + /// Enqueues `requests` atomically and starts the pump if none is running. Returns one stream per + /// request, in order. `submit(_:)` and `run(_:)` both go through here, so only the pump ever + /// drives `session`. + 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) + streams.append(stream) + pending.append(Pending(request: request, continuation: continuation)) + } + let shouldStart = try state.withLock { s -> Bool in + 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 + } + + /// The single pump. Each iteration does the following: + /// 1. Admit queued requests into free slots. + /// 2. If no slot is live and the queue is empty, clear `pumping` and return. + /// 3. 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. + /// 4. Advance each live cursor, then sample, emit or retire each live row. + /// + /// The step always spans `maxConcurrent` rows, so a lone request still pays for a full batch. + /// If `decodeBatch` throws, every resident row and every queued request 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 admit(into: &active) + + let liveSlots = active.indices.filter { active[$0] != nil } + if liveSlots.isEmpty { + // Recheck the queue under the lock. A request enqueued after `admit` saw + // `pumping == true` and did not start a pump, so this one must keep going. + let more = state.withLock { s -> Bool in + if s.queue.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) } + // Drain the queue under the same lock that clears `pumping`. A queued request left + // behind would wait forever on a pump that has stopped. + let orphaned = state.withLock { s -> [Pending] in + let waiting = s.queue + s.queue.removeAll() + s.occupiedRows = 0 + s.pumping = false + return waiting + } + for pending in orphaned { pending.continuation.finish(throwing: error) } + } + } + + /// Fills free slots from the queue. Each admitted request is prefilled into its row and its first + /// token is sampled at once, so a request can also retire here. + /// + /// 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 + 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 if the adapter says it 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() + active[slot] = nil + state.withLock { $0.occupiedRows -= 1 } + } + } + + // MARK: - Exclusive engine access + + /// Waits until the pump finishes its current step, then gives the caller sole use of the engine + /// until the lease is released. Resident rows stay in their slots and resume afterwards. + /// + /// A request that runs on its own session of the same engine, such as one that needs per-step + /// logits, holds a lease for as long as it runs: + /// + /// ```swift + /// let lease = await scheduler.acquireExclusiveAccess() + /// defer { lease.release() } + /// // Generate on an IsolatedSessionEngine over the same CoreAISequentialEngine. + /// ``` + package func acquireExclusiveAccess() async -> ExclusiveEngineLease { + await engineLock.lock() + return ExclusiveEngineLease(lock: engineLock) + } + + // MARK: - Closed batch + + /// Submits `requests` together and waits for all of them. Returns one `RowResult` per request, in + /// order. It shares the pump with `submit(_:)`, so it is safe to call while other requests 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 + } + } +} + +/// Sole use of a `BatchScheduler`'s engine, from `acquireExclusiveAccess()`. Releasing it more than +/// once is harmless, and it releases on deinit, so a dropped lease cannot stall the batch. +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 that grants waiters in arrival order. `unlock()` hands the lock straight to the +/// next waiter, so a task that unlocks and locks again goes to the back of the line. +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() + } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift new file mode 100644 index 00000000..de3a6b39 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift @@ -0,0 +1,26 @@ +// 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 + +/// Which serving lane a request takes under a batched server. +package enum BatchServingLane: Equatable, Sendable { + /// Served by the batch scheduler through a `SchedulerSession`. + case batched + /// Served by a single-request engine that returns per-step logits. + case fallbackBatchOne +} + +/// Chooses the serving lane for a request when batched serving is enabled. +/// +/// The batched lane samples inside the scheduler and returns no per-step logits, so a request with a +/// JSON schema takes the fallback lane, where constrained decoding can mask the logits at each step. +/// Loglikelihood scoring also needs logits, but the completions handler sends it to the fallback lane +/// directly without asking the router. +package enum BatchServingRouter { + /// Returns `.fallbackBatchOne` when the chat request has 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..39717150 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift @@ -0,0 +1,65 @@ +// 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 + +/// Per-row token bookkeeping and stop checks for one request in a batched decode. +/// +/// `BatchScheduler` keeps one adapter per occupied slot and feeds it each sampled token. When +/// `accept(_:)` returns false, the scheduler emits `.finished` and frees the slot for the next queued +/// request while the other rows keep decoding. +struct BatchSessionAdapter { + let request: BatchRequest + + private(set) var generated: [Int32] = [] + /// False once the row has retired. + private(set) var live: Bool = true + /// Set when the row retires. + private(set) var stopReason: StopReason? + + /// The last `stopSequences.maxLength` tokens, matched against the stop sequences. + private var recentTokens: [Int32] = [] + + /// The per-row step index for the sampler, equal to `generated.count`. + 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 reaching + /// `maxTokens`, checked in that order. Returns whether the row is still live. Once the row has + /// retired, it ignores further 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/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 69f557c1..0ccf9c76 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. } + /// Whether this graph can run several requests in one forward. It needs 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 diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 9fc24ddd..213f7427 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -272,8 +272,16 @@ public enum InferenceRuntimeError: Error, LocalizedError { case unsupportedEngine(String) case unsupportedEngineVariant(String) case bufferAllocationFailed(String) + /// The scheduler rejected the request because its wait queue is at its limit of `depth`. + case queueFull(depth: Int) case genericError(String) + /// Whether this is `.queueFull`, which the server reports as HTTP 429. + package var isQueueFull: Bool { + if case .queueFull = self { return true } + return false + } + public var errorDescription: String? { switch self { case .functionNotFound(let name): @@ -305,6 +313,8 @@ public enum InferenceRuntimeError: Error, LocalizedError { return "Unknown variant \(variant)" case .bufferAllocationFailed(let details): return "GPU buffer allocation failed: \(details)" + case .queueFull(let depth): + return "Request queue full (depth: \(depth))" case .genericError(let message): return message } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift new file mode 100644 index 00000000..a96e72ae --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift @@ -0,0 +1,67 @@ +// 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 + +/// An `InferenceEngine` that runs one fallback-lane request (guided generation or loglikelihood +/// scoring) on a `CoreAISequentialEngine` through its own `GenerationSessionState`. +/// +/// Create one per request. The KV cache, token history and prefix-hit count belong to this session +/// rather than the engine's default session, so other sessions keep their state. It does not make +/// concurrent use of the engine safe: next to a `BatchScheduler`, hold an `ExclusiveEngineLease` +/// while the request runs. +/// +/// The class is `@unchecked Sendable` because a single request owns it, so no two tasks drive it at +/// once. +package final class IsolatedSessionEngine: InferenceEngine, @unchecked Sendable { + package typealias ConfigType = ModelConfig + + private let base: CoreAISequentialEngine + /// Created on first use so `init` cannot throw. `reset(to:)` replaces it with a fresh one. + private var sessionState: GenerationSessionState? + + package init(base: CoreAISequentialEngine) { + self.base = base + } + + private func ensureSession() throws -> GenerationSessionState { + if let sessionState { return sessionState } + let fresh = try base.makeSessionState() + sessionState = fresh + return fresh + } + + package var config: ModelConfig { base.config } + + package func generate( + with input: [TokenId], + samplingConfiguration: SamplingConfiguration, + inferenceOptions: InferenceOptions + ) async throws -> CoreAISequentialEngine.GenerationSequence { + try await base.generate( + with: input, + sessionState: try ensureSession(), + samplingConfiguration: samplingConfiguration, + inferenceOptions: inferenceOptions) + } + + /// Replaces this request's session with a fresh one, whatever `tokenIndex` is. Guided generation + /// and loglikelihood scoring both reset to 0 before generating, so this lane never needs a + /// partial rewind. + package func reset(to tokenIndex: Int) async throws { + sessionState = try base.makeSessionState() + } + + package var supportsLogits: Bool { base.supportsLogits } + + /// This session's busy state. The base engine's flag would reflect other sessions too. + package var isBusy: Bool { sessionState?.tokenBox.isBusy ?? false } + + package var lastPrefixHitCount: Int { sessionState?.lastPrefixHitCount ?? 0 } + + package func cancel() async throws { + sessionState?.tokenBox.cancelActive() + } + + package var maxContextLength: Int { base.config.maxContextLength } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift new file mode 100644 index 00000000..e5c072d6 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift @@ -0,0 +1,126 @@ +// 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 CoreAIShared + +/// An `InferenceEngine` that represents one request inside a `RequestSink` such as `BatchScheduler`. +/// +/// Create one per request. Decoding strategies and the server see an ordinary engine while several +/// requests share one batched forward. `generate` submits a `BatchRequest` and returns the sink's +/// per-request event stream as a `GenerationSequence`. +/// +/// Differences from a single-request engine: +/// - `reset(to:)` does nothing. The sink owns the shared KV cache, so there is no per-request +/// history to rewind. +/// - `supportsLogits` is false. The sink samples internally, so guided generation and +/// loglikelihood scoring must use the fallback lane (see `BatchServingRouter`). +/// - The sink stops a row only on `stopTokenIds` or `maxTokens`. Multi-token stop sequences and +/// detokenization stay in the per-request `DecodingStrategy`. +/// +/// ```swift +/// let engine = SchedulerSession(sink: scheduler, config: modelConfig, stopTokenIds: eosIds) +/// let sequence = try await engine.generate( +/// with: promptTokens, samplingConfiguration: sampling, inferenceOptions: options) +/// for try await output in sequence { /* output.tokenId */ } +/// ``` +package final class SchedulerSession: InferenceEngine, @unchecked Sendable { + package typealias ConfigType = ModelConfig + + package let config: ModelConfig + private let sink: any RequestSink + /// Token ids that end the row inside the sink, usually EOS plus any additional EOS ids. + private let stopTokenIds: Set + + package init(sink: any RequestSink, config: ModelConfig, stopTokenIds: Set) { + self.sink = sink + self.config = config + self.stopTokenIds = stopTokenIds + } + + package func generate( + with input: [TokenId], + samplingConfiguration: SamplingConfiguration, + inferenceOptions: InferenceOptions + ) async throws -> GenerationSequence { + // Cap `maxTokens` at the room left in the context window after the prompt. Decoding past + // the window would make the shared batched forward throw, and a failed forward fails every + // resident row, so each row must retire with `.maxTokens` before that point. + let contextRoom = Swift.max(1, config.maxContextLength - input.count) + let requested = inferenceOptions.maxTokens ?? contextRoom + let stream = try sink.submit( + BatchRequest( + promptTokens: input.map { Int32($0) }, + sampling: samplingConfiguration, + eosTokenIds: stopTokenIds, + maxTokens: Swift.max(1, Swift.min(requested, contextRoom)))) + return GenerationSequence(events: stream) + } + + /// Does nothing. The sink reclaims this request's state, and resetting the shared engine would + /// destroy the other rows. + package func reset(to tokenIndex: Int) async throws {} + + /// Always false because the sink samples internally and returns no per-step logits. + package var supportsLogits: Bool { false } + + /// Always false. Each session is one request, and the sink runs several at once. + package var isBusy: Bool { false } + + package var maxContextLength: Int { config.maxContextLength } +} + +// MARK: - Output sequence + +extension SchedulerSession { + /// The sink's per-request event stream, presented as an engine output sequence. + package struct GenerationSequence: InferenceOutputSequence { + package typealias Element = InferenceOutput + package typealias Failure = Error + + let events: AsyncThrowingStream + let stopReasonStore = StopReasonStore() + + package var stopReason: StopReason? { stopReasonStore.stopReason } + + package func setStopReason(_ reason: StopReason) { stopReasonStore.set(reason) } + + 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): + // Record the sink's reason unless the consumer already set one. The + // `.maxTokens` default below only covers a stream that ends without `.finished`. + stopReasonStore.setIfUnset(reason) + } + } + stopReasonStore.setIfUnset(.maxTokens) + return nil + } + } +} diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift new file mode 100644 index 00000000..ff0b6e37 --- /dev/null +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -0,0 +1,366 @@ +// 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 + +/// Logits over a 64-token vocabulary that greedy sampling resolves to `token`. +private func oneHot(_ token: Int32) -> [LogitsScalarType] { + var row = [LogitsScalarType](repeating: 0, count: 64) + row[Int(token)] = 1 + return row +} + +/// 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) +} + +/// A `BatchedDecodeStep` that predicts the input token plus one, so a request with prompt `[..., k]` +/// and `maxTokens: m` generates `k+1` through `k+m` whatever its slot or admission time. +/// +/// It records each decode's `startPositions` and the peak number of overlapping forward calls. +/// `forwardDelay` keeps each call in flight longer so that an overlap is reliably observed. +private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { + struct Boom: Error {} + + let maxBatchCapacity: Int + /// `prefillRow` throws for prompts that start with this token. + private let failPrefillFor: Int32? + /// Every `decodeBatch` throws. + private let failDecode: Bool + private let forwardDelay: Duration? + + private struct Record { + var startPositions: [[Int]] = [] + var inFlight = 0 + var maxInFlight = 0 + var forwards = 0 + } + private let record = Mutex(Record()) + + var observedStartPositions: [[Int]] { record.withLock { $0.startPositions } } + var maxInFlight: Int { record.withLock { $0.maxInFlight } } + /// Prefill and decode 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) + } + + 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) } + } +} + +/// A count-up `BatchedDecodeStep` whose `prefillRow` blocks until `release()`, so a test can keep a +/// row occupying its slot while it fills the queue. +private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { + let maxBatchCapacity: Int + + private struct Gate { + var entered = false + var released = false + var enteredWaiter: CheckedContinuation? + var releaseWaiter: CheckedContinuation? + } + private let gate = Mutex(Gate()) + + init(capacity: Int) { maxBatchCapacity = capacity } + + /// Returns once a `prefillRow` call has started, which means its row holds a slot. + func waitEntered() async { + await withCheckedContinuation { cont in + gate.withLock { g in + if g.entered { cont.resume() } else { g.enteredWaiter = cont } + } + } + } + + /// Unblocks the waiting `prefillRow`. Later calls no longer block. + 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] { + let enteredWaiter: CheckedContinuation? = gate.withLock { g in + g.entered = true + defer { g.enteredWaiter = nil } + return g.enteredWaiter + } + enteredWaiter?.resume() + 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) } + } +} + +@Suite("BatchScheduler") +struct BatchSchedulerTests { + private func makeScheduler( + _ step: any BatchedDecodeStep, maxConcurrent: Int, maxQueued: Int? = nil + ) -> BatchScheduler { + BatchScheduler( + decodeStep: step, + session: GenerationSessionState( + kvCache: MockStateHandler(names: ["k"], shape: [1, 1, 1, 1, 1]), + additionalStates: nil, + hasNonTruncatableStates: false), + 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("Submit streams the tokens, then .finished") + func streamingSubmit() async throws { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + let (tokens, reason) = try await drain( + scheduler.submit(req(prompt: [5], maxTokens: 100, eos: [8]))) + // The EOS token itself is streamed. + #expect(tokens == [6, 7, 8]) + #expect(reason == .eos) + } + + @Test("A retired row's slot admits the next queued request") + func slotReuse() async throws { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + let results = try await scheduler.run([ + req(prompt: [0], maxTokens: 2), + req(prompt: [10], maxTokens: 3), + req(prompt: [20], maxTokens: 1), + req(prompt: [30], maxTokens: 4), + ]) + #expect(results.map(\.tokens) == [[1, 2], [11, 12, 13], [21], [31, 32, 33, 34]]) + #expect(results.allSatisfy { $0.stopReason == .maxTokens }) + } + + @Test("Rows decode at their own cursors in one forward") + func raggedCursors() async throws { + let mock = CountUpDecodeStep(capacity: 2) + let scheduler = makeScheduler(mock, maxConcurrent: 2) + let results = try await scheduler.run([ + req(prompt: [1], maxTokens: 3), + req(prompt: [1, 2, 3], maxTokens: 3), + ]) + #expect(results.map(\.tokens) == [[2, 3, 4], [4, 5, 6]]) + // Each cursor starts at the row's prompt length. + #expect(mock.observedStartPositions == [[1, 3], [2, 4]]) + } + + @Test("Run and submit share one pump, so forwards 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("Submit throws queueFull past free rows plus maxQueued", arguments: [false, true]) + func backpressure(waitForAdmission: Bool) async throws { + // Queued plus admitted requests may total maxConcurrent + maxQueued, here 3. The limit is + // the same whether the first request is still queued (a burst) or already holds a row. + let mock = GatedPrefillStep(capacity: 2) + let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: 1) + var streams = [try scheduler.submit(req(prompt: [0], maxTokens: 1))] + if waitForAdmission { await mock.waitEntered() } + for _ in 0..<2 { streams.append(try scheduler.submit(req(prompt: [0], maxTokens: 1))) } + #expect(throws: InferenceRuntimeError.self) { + _ = try scheduler.submit(req(prompt: [0], maxTokens: 1)) + } + + mock.release() + for stream in streams { _ = try await drain(stream) } + } + + @Test("maxQueued 0 accepts only as many requests as there are free rows") + func noWaiting() async throws { + let mock = GatedPrefillStep(capacity: 2) + let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: 0) + let streams = try (0..<2).map { _ in try scheduler.submit(req(prompt: [0], maxTokens: 1)) } + #expect(throws: InferenceRuntimeError.self) { + _ = try scheduler.submit(req(prompt: [0], maxTokens: 1)) + } + + mock.release() + for stream in streams { _ = try await drain(stream) } + } + + @Test("An exclusive lease pauses the pump until it is released") + func exclusiveLeasePausesPump() async throws { + let step = CountUpDecodeStep(capacity: 2) + let scheduler = makeScheduler(step, maxConcurrent: 2) + let lease = await scheduler.acquireExclusiveAccess() + let stream = try scheduler.submit(req(prompt: [5], maxTokens: 3)) + + try await Task.sleep(for: .milliseconds(50)) + #expect(step.forwards == 0) + + lease.release() + let (tokens, reason) = try await drain(stream) + #expect(tokens == [6, 7, 8]) + #expect(reason == .maxTokens) + } + + @Test("AsyncFIFOLock grants waiters in arrival order") + func fifoLockOrder() async { + let lock = AsyncFIFOLock() + let order = Mutex<[Int]>([]) + await lock.lock() + await withTaskGroup(of: Void.self) { group in + for i in 0..<3 { + group.addTask { + await lock.lock() + order.withLock { $0.append(i) } + lock.unlock() + } + // Let each waiter enqueue before starting the next one. + try? await Task.sleep(for: .milliseconds(20)) + } + lock.unlock() + } + #expect(order.withLock { $0 } == [0, 1, 2]) + } + + @Test("An empty prompt is rejected at submit") + func emptyPromptRejected() { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + #expect(throws: InferenceRuntimeError.self) { + _ = try scheduler.submit(req(prompt: [], maxTokens: 1)) + } + } + + @Test("A prefill failure fails only its own request") + func prefillFailureIsolated() async throws { + let scheduler = makeScheduler( + CountUpDecodeStep(capacity: 2, failPrefillFor: 50), maxConcurrent: 2) + let failing = try scheduler.submit(req(prompt: [50], maxTokens: 4)) + let healthy = try scheduler.submit(req(prompt: [5], maxTokens: 3)) + + await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(failing) } + let (tokens, reason) = try await drain(healthy) + #expect(tokens == [6, 7, 8]) + #expect(reason == .maxTokens) + } + + @Test("A failed shared decode also fails the queued requests") + func decodeFailureDrainsQueue() async throws { + let scheduler = makeScheduler( + CountUpDecodeStep(capacity: 1, failDecode: true), maxConcurrent: 1) + let streams = try [Int32(1), 2, 3].map { + try scheduler.submit(req(prompt: [$0], maxTokens: 4)) + } + // With one slot, requests 2 and 3 wait in the queue while request 1's decode throws. If + // the queue were not drained, iterating their streams would hang. + for stream in streams { + await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(stream) } + } + } +} + +@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) + 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")]) + } +} From b34cca67ad6e386f235d569bf8c2cbcf9cf7293f Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 01:37:01 -0700 Subject: [PATCH 02/10] Serve concurrent requests through the batch scheduler in llm-server --concurrency N serves chat requests through the batch scheduler when the model supports batching, meaning it has an attn_mask input and a batch dimension above 1. Otherwise the server warns and serves one request at a time. Requests that need per-step logits, chat with a response schema and completions scoring, run on an isolated session while holding an exclusive lease, so the batch pauses for them instead of sharing the engine. A full scheduler returns HTTP 429. Streaming requests start decoding before the SSE response is sent, so they get the 429 too. --replay-concurrent sends all replay requests at once, so replay runs can exercise the scheduler. Without --concurrency, requests take the same path as before. --- .../Tools/llm-server/ChatHandler.swift | 120 ++++++++++++------ .../Tools/llm-server/CompletionHandler.swift | 13 +- .../Tools/llm-server/LLMServerMain.swift | 50 +++++++- .../Tools/llm-server/ReplayRunner.swift | 71 +++++++---- .../Tools/llm-server/ServerState.swift | 98 +++++++++++++- 5 files changed, 282 insertions(+), 70 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 42aeba95..2457db02 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -103,45 +103,59 @@ 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))) } + + // Requests with a JSON schema take the fallback lane. The streaming path hands the access to its + // body closure, and the non-streaming path releases it at each return. + 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 { + // The batch scheduler rejects a submit once it holds maxConcurrent + maxQueued requests, so + // free rows absorb a burst before the queue limit applies. Streaming requests submit before + // their SSE response starts, so they get the same 429. + 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 +163,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 +237,20 @@ 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() } + // Guided generation needs per-step logits, which only the fallback lane provides. + 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 +459,51 @@ 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 +/// A streaming decode that has started. On the batched lane, starting it submits the request to +/// the batch scheduler, so a full scheduler throws before the handler sends any SSE headers. +/// +/// `@unchecked Sendable` because it is created once and then handed to a single consumer, the +/// response body or the replay runner, which iterates it on one task. +struct StartedDecode: @unchecked Sendable { + let tokens: any AsyncSequence + let start: SuspendingClock.Instant +} +/// Picks the decoding strategy for `chatRequest` and starts it on the engine for its lane. +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) +} + +/// 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( + 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() @@ -584,18 +621,20 @@ func runStreamingLoop( } 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). + // Preflight throws for bad requests before any SSE framing (mapped to HTTP 400 upstream), and + // starting the decode throws queueFull (mapped to HTTP 429) when the batch scheduler is full. 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() } + // The closure owns the access for the stream's lifetime. Its deinit releases it if the + // framework drops the response before running this closure. + defer { access.release() } do { let encoder = JSONEncoder() @@ -613,7 +652,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..0f2483de 100644 --- a/swift/Sources/Tools/llm-server/CompletionHandler.swift +++ b/swift/Sources/Tools/llm-server/CompletionHandler.swift @@ -30,9 +30,10 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw status: .notImplemented, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - let permit: QueuePermit + // Scoring needs per-step logits, so it takes the fallback lane. + let access: RequestAccess do { - permit = try await state.queue.acquire() + access = try await state.acquireAccess(lane: .fallbackBatchOne) } catch let error as ServerError { let err = ErrorResponse( error: .init( @@ -42,7 +43,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) @@ -149,9 +150,11 @@ private func processOnePrompt( forcedContinuation: continuation ) - try await state.engine.reset() + // Scoring needs per-token logits and a forced continuation, which only the fallback lane supports. + let scoringEngine = state.engineForRequest(lane: .fallbackBatchOne) + try await scoringEngine.reset() - let stream = try await state.engine.generate( + let stream = try await scoringEngine.generate( with: [allTokens[0]], samplingConfiguration: .greedy, inferenceOptions: options diff --git a/swift/Sources/Tools/llm-server/LLMServerMain.swift b/swift/Sources/Tools/llm-server/LLMServerMain.swift index bdc8e87d..39174866 100644 --- a/swift/Sources/Tools/llm-server/LLMServerMain.swift +++ b/swift/Sources/Tools/llm-server/LLMServerMain.swift @@ -52,6 +52,14 @@ 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, so the batch scheduler can " + + "fill several rows. Use with --concurrency N.") + ) + var replayConcurrent: Bool = false + @Option(name: .customLong("variant"), help: "Engine variant: auto, coreai-pipelined, coreai-sequential") var inferenceEngineVariant: String = "default" @@ -78,6 +86,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("concurrency"), + 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 concurrency: Int = 1 + @Flag( name: .customLong("no-thinking"), help: "Disable thinking/reasoning (alias for --default-reasoning-effort none)" @@ -103,6 +119,9 @@ struct LLMServer: AsyncParsableCommand { guard maxQueueDepth >= 0 else { throw ValidationError("--max-queue-depth must be >= 0 (got \(maxQueueDepth))") } + guard concurrency >= 1 else { + throw ValidationError("--concurrency must be >= 1 (got \(concurrency))") + } do { _ = try ReasoningEffort.resolveDefault( defaultReasoningEffort: defaultReasoningEffort, noThinking: noThinking) @@ -218,14 +237,41 @@ struct LLMServer: AsyncParsableCommand { maxQueueDepth: maxQueueDepth ) + // Batched serving needs --concurrency above 1 and a CoreAISequentialEngine whose graph has an + // attention mask input and a batch dimension above 1. The scheduler holds up to + // concurrency + --max-queue-depth requests, decoding or waiting, and rejects the next one. + var batchScheduler: BatchScheduler? = nil + if concurrency > 1 { + if let seq = engine as? CoreAISequentialEngine, seq.supportsBatching { + let assetMax = bundle.maxBatchSize + guard concurrency <= assetMax else { + throw ValidationError( + "--concurrency \(concurrency) exceeds the asset's max batch size \(assetMax). " + + "Re-export with --dynamic-batch-size \(concurrency) or higher, or lower --concurrency.") + } + batchScheduler = try BatchScheduler( + engine: seq, maxConcurrent: concurrency, maxQueued: maxQueueDepth) + print("Batched serving: concurrency \(concurrency) (asset max batch \(assetMax))") + } else { + print( + "⚠️ --concurrency \(concurrency) 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, + batchScheduler: batchScheduler, + batchModelConfig: batchScheduler != nil ? engineConfig : nil ) 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..458115b0 100644 --- a/swift/Sources/Tools/llm-server/ReplayRunner.swift +++ b/swift/Sources/Tools/llm-server/ReplayRunner.swift @@ -4,40 +4,37 @@ // 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 +/// back into the result). By default requests run one at a time, matching the single-request server. +/// With `concurrent`, they all start at once so the batch scheduler can fill several rows. +/// 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 { + // Lines keep input order. Requests past the queue limits come back as queue-full error lines. + 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 +45,35 @@ enum ReplayRunner { } } + /// Runs one replay request and encodes its result line. Errors become an error result line. + private static func oneLine(_ req: ReplayRequest, state: ServerState) async -> String { + let sessionID = req.session ?? "default" + let result: ReplayResult + do { + // runChatCompletion and runStreamingLoop do not take the request access, so take it here + // as the HTTP handler does. Without it, concurrent replay would drive the shared engine + // from several tasks at once. + 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 +102,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..44f6fcf2 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -129,6 +129,11 @@ final class ServerState: @unchecked Sendable { /// guaranteed while this value is unchanged; it changes when output-affecting /// configuration (model, engine, context length, sampling defaults) changes. let systemFingerprint: String + /// Scheduler for the batched lane, or nil when the server handles one request at a time. + private let batchScheduler: BatchScheduler? + /// Config for each `SchedulerSession`. Set it whenever `batchScheduler` is set. If it is nil + /// while a scheduler is present, the batched lane runs on the shared engine without the permit. + private let batchModelConfig: ModelConfig? private let _state = Mutex(InternalState()) private struct InternalState { @@ -140,7 +145,13 @@ 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, + batchScheduler: BatchScheduler? = nil, + batchModelConfig: ModelConfig? = nil + ) { self.engine = engine self.tokenizer = tokenizer self.config = config @@ -148,6 +159,74 @@ final class ServerState: @unchecked Sendable { self.toolCallDetection = detectToolCallFormat(using: tokenizer) self.thinkingFormat = detectThinkingFormat(using: tokenizer) self.systemFingerprint = Self.makeSystemFingerprint(config: config, engine: engine) + self.batchScheduler = batchScheduler + self.batchModelConfig = batchModelConfig + } + + /// Token ids that retire a row in the scheduler: EOS and the additional EOS ids. Multi-token + /// stop sequences are still matched by the decoding strategy. + 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 one request on `lane` should run on. + /// + /// - `.batched`: a new `SchedulerSession` over `batchScheduler`, or the shared base engine when + /// no scheduler is configured. + /// - `.fallbackBatchOne`: with a scheduler running, a new `IsolatedSessionEngine` with its own + /// session, so the batched rows keep their KV state while the batch is paused for it (see + /// `acquireAccess(lane:)`). Without a scheduler, the shared base engine, + /// which keeps its prefix cache. + func engineForRequest(lane: BatchServingLane) -> any InferenceEngine { + switch lane { + case .batched: + guard let scheduler = batchScheduler, let modelConfig = batchModelConfig else { + return engine + } + return SchedulerSession( + sink: scheduler, config: modelConfig, stopTokenIds: schedulerStopTokenIds) + case .fallbackBatchOne: + if batchScheduler != nil, let seq = engine as? CoreAISequentialEngine { + return IsolatedSessionEngine(base: seq) + } + return engine + } + } + + /// Waits for what a request on `lane` must hold while it uses the engine, and returns it. + /// + /// - The single-slot permit from `queue`, except on the batched lane when a continuous-batching + /// scheduler admits the request. Throws `ServerError.queueFull` when the queue is full. + /// - On the fallback lane with a scheduler running, an `ExclusiveEngineLease` as well, so the + /// batch pauses at a step boundary instead of sharing the engine. + /// + /// 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) + } + + /// Whether a request on `lane` needs the single-slot permit. The batched lane does not when a + /// continuous-batching scheduler admits it, because the scheduler enforces its own queue limit. + 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 @@ -315,3 +394,20 @@ enum RequestID { return "coreai-\(n)" } } + +/// What one request holds while it uses the engine. Releasing it more than once is harmless, and +/// the permit and lease each 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() + } +} From bd06a9875592c8937a5cb0503cf03464f8ef039d Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 10:10:46 -0700 Subject: [PATCH 03/10] Score completions prompts in batches With a batch scheduler running, /v1/completions scoring no longer pauses the batch with an exclusive lease. BatchScheduler.score(_:) queues the prompts, and between decode steps the pump scores them as one padded batch on a separate session. - Prompts are queued shortest first, and a batch stops before a third of its positions would be padding. - The batch advances 128 positions per forward. Each chunk's logits are reduced to log-probabilities, rows in parallel, before the next chunk runs. - CoreAISequentialEngine.allPositionLogits returns every position's logits from one batched forward. decodeBatch now slices the same forward. All prompts in one completions request score together. Scoring requests still take the server permit, so one runs at a time. Without a scheduler, scoring runs on the shared engine as before. --- .../InferenceEngines/BatchScheduler.swift | 176 +++++++++++++++++- .../InferenceEngines/BatchServingRouter.swift | 4 +- .../InferenceEngines/BatchedDecodeStep.swift | 11 ++ .../CoreAISequentialEngine.swift | 39 +++- .../IsolatedSessionEngine.swift | 7 +- .../InferenceEngines/SchedulerSession.swift | 4 +- .../Tools/llm-server/CompletionHandler.swift | 77 +++++--- .../Tools/llm-server/ServerState.swift | 10 + .../BatchSchedulerTests.swift | 111 ++++++++++- 9 files changed, 391 insertions(+), 48 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift index 80e3304c..a2b15a1d 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -3,6 +3,7 @@ // 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 /// One request submitted to a `RequestSink`. @@ -35,6 +36,24 @@ package struct BatchRequest: Sendable { } /// The tokens one request generated in `BatchScheduler.run(_:)` and why it stopped. +/// One scoring request for `BatchScheduler.score(_:)`. +/// +/// `targets[p]` is scored against the logits at position `p` of `tokens`, so both have the same +/// length. For a prompt `[t0, t1, t2]` the usual targets are `[t1, t2, x]`, where the last entry has +/// no real next token and the caller ignores it. +package struct ScoreRequest: Sendable { + let tokens: [Int32] + let targets: [Int32] + /// How many alternatives to report per position. + let topK: Int + + package init(tokens: [Int32], targets: [Int32], topK: Int) { + self.tokens = tokens + self.targets = targets + self.topK = topK + } +} + struct RowResult: Sendable { let tokens: [Int32] let stopReason: StopReason @@ -101,12 +120,22 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { var queue: [Pending] = [] /// Rows held by admitted requests. var occupiedRows = 0 + /// Scoring requests waiting for the pump, in arrival order. + var scoring: [ScoringJob] = [] var pumping = false } private let state = Mutex(State()) /// Held by the pump for each step and by an `ExclusiveEngineLease` for its whole lifetime, so /// the engine never runs work from two tasks at once. private let engineLock = AsyncFIFOLock() + /// Positions per scoring forward. Each chunk's logits are reduced to log-probabilities and + /// dropped, so memory stays at `rows * scoringChunkLength * vocab` logits however long the prompt. + private let scoringChunkLength = 128 + + private struct ScoringJob { + let request: ScoreRequest + let continuation: AsyncThrowingStream.Continuation + } /// Creates a scheduler with a new `maxConcurrent`-row session on `engine`. /// `engine.supportsBatching` must be true, and `maxConcurrent` must not exceed the graph's batch @@ -202,6 +231,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { while true { await engineLock.lock() defer { engineLock.unlock() } + await runPendingScoring() await admit(into: &active) let liveSlots = active.indices.filter { active[$0] != nil } @@ -209,7 +239,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { // Recheck the queue under the lock. A request enqueued after `admit` saw // `pumping == true` and did not start a pump, so this one must keep going. let more = state.withLock { s -> Bool in - if s.queue.isEmpty { + if s.queue.isEmpty && s.scoring.isEmpty { s.pumping = false return false } @@ -230,14 +260,16 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { for row in active where row != nil { row!.continuation.finish(throwing: error) } // Drain the queue under the same lock that clears `pumping`. A queued request left // behind would wait forever on a pump that has stopped. - let orphaned = state.withLock { s -> [Pending] in - let waiting = s.queue + 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) } } } @@ -288,6 +320,144 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } + // MARK: - Scoring + + /// Returns the log-probability of each request's targets, in request order. + /// + /// The requests join the pump's scoring queue. Between decode steps the pump takes up to + /// `maxConcurrent` waiting requests and runs them as one batch on a separate session: + /// - Rows shorter than the longest prompt are padded at the end. The causal mask keeps padding + /// from affecting real positions, and its logits are ignored. + /// - The batch advances `scoringChunkLength` positions per forward, and each chunk's logits are + /// reduced to log-probabilities before the next chunk runs. + /// + /// Running decode rows wait for the scoring batch but keep their KV state. A failed scoring + /// forward fails only the requests in that batch. + 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") + } + // Queue the requests shortest first so each batch holds similar lengths and pads less. The + // sort is stable, so the same requests always form the same batches. + 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 request, `maxConcurrent` at a time. Called by the pump while it + /// holds the engine lock. + private func runPendingScoring() async { + while true { + let jobs = state.withLock { s -> [ScoringJob] in + let count = Self.scoringBatchCount(s.scoring.map(\.request.tokens.count), limit: maxConcurrent) + let taken = Array(s.scoring.prefix(count)) + s.scoring.removeFirst(taken.count) + return taken + } + if jobs.isEmpty { return } + do { + let results = try await 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) } + } + } + } + + /// How many of the waiting scoring requests, with prompt lengths `lengths` in queue order, go into + /// the next batch. It takes up to `limit` requests and stops before one that would make more than + /// a third of the batch's positions padding, where batching stops paying for itself. For lengths + /// `[7, 40, 130, 300]` the batches are `[7]`, `[40]` and `[130, 300]`. + static func scoringBatchCount(_ 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 up to `maxConcurrent` requests in one padded batch. + /// + /// For prompts of lengths 3 and 5 with a chunk length of 4, the rows are + /// ``` + /// chunk 0: [a0 a1 a2 a2] [b0 b1 b2 b3] positions 0..<4 + /// chunk 1: [a2] [b4] position 4, row a is padding only + /// ``` + private func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { + // A static batch=N graph needs exactly N rows, so fill it with padding rows. + let capacity = decodeStep.maxBatchCapacity + let rows = capacity == Int.max ? requests.count : capacity + 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 = min(scoringChunkLength, longest - start) + 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 rowLogits = (0.. [LogitsScalarType] + + /// Runs `tokensPerRow` (`[N][q]`, every row at `startPosition`) and returns each row's logits at + /// every position, flat `[q * vocab]`. Scoring uses this to read all positions in one pass. + func allPositionLogits( + tokensPerRow: [[Int32]], + startPosition: Int, + session: GenerationSessionState + ) async throws -> [[LogitsScalarType]] + + /// A new session whose KV cache holds `batchSize` rows. + func makeSessionState(batchSize: Int) throws -> GenerationSessionState } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 0ccf9c76..140cb03f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -417,7 +417,38 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba startPositions: [Int], session: GenerationSessionState ) async throws -> [[LogitsScalarType]] { - precondition(!tokensPerRow.isEmpty, "decodeBatch requires at least one row") + let flat = try await runBatchedForward( + tokensPerRow: tokensPerRow, startPositions: startPositions, session: session) + let n = tokensPerRow.count + let q = tokensPerRow[0].count + let vocab = config.vocabSize + return (0.. [[LogitsScalarType]] { + let n = tokensPerRow.count + let flat = try await runBatchedForward( + tokensPerRow: tokensPerRow, startPositions: Array(repeating: startPosition, count: n), + session: session) + let rowLength = flat.count / n + return (0.. [LogitsScalarType] { + 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 @@ -496,11 +527,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba } } - let flat = readNDArray(out, as: LogitsScalarType.self, count: n * q * vocab) - return (0.. CompletionResponse.CompletionChoice { +/// One prompt ready to score: its tokens and the target each position predicts. +private struct PreparedPrompt { + let tokens: [Int32] + /// `tokens.dropFirst()` plus a padding token. lm-eval sends echo=true, max_tokens=1 and slices + /// logprobs[ctxlen:-1], so the extra last entry is the one it drops. + let targets: [Int32] +} + +/// Tokenizes `prompt`, or returns nil when it has fewer than two tokens and nothing to score. +private func preparePrompt(_ prompt: Prompt, state: ServerState) throws -> PreparedPrompt? { let allTokens: [Int32] switch prompt { case .tokenIds(let ids): @@ -129,33 +149,30 @@ 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 one prompt on the shared engine, which runs it in one prefill with every position's logits. +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 ) - // Scoring needs per-token logits and a forced continuation, which only the fallback lane supports. let scoringEngine = state.engineForRequest(lane: .fallbackBatchOne) try await scoringEngine.reset() let stream = try await scoringEngine.generate( - with: [allTokens[0]], + with: [prompt.tokens[0]], samplingConfiguration: .greedy, inferenceOptions: options ) @@ -166,8 +183,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/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 44f6fcf2..320e0994 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -218,6 +218,16 @@ final class ServerState: @unchecked Sendable { return RequestAccess(permit: permit, lease: lease) } + /// The scheduler that scores completions prompts in batches, or nil when the server scores them + /// one at a time on the shared engine. + var scoringScheduler: BatchScheduler? { batchScheduler } + + /// Waits for the single-slot permit a scoring request holds. Scoring takes no lease: with a + /// scheduler, its prompts run on the scheduler's own task between decode steps. + func acquireScoringAccess() async throws -> RequestAccess { + RequestAccess(permit: try await queue.acquire(), lease: nil) + } + /// Whether a request on `lane` needs the single-slot permit. The batched lane does not when a /// continuous-batching scheduler admits it, because the scheduler enforces its own queue limit. private func needsServerPermit(lane: BatchServingLane) -> Bool { diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift index ff0b6e37..f5f57912 100644 --- a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -15,6 +15,19 @@ private func oneHot(_ token: Int32) -> [LogitsScalarType] { return row } +/// A session with no real KV cache, for mock decode steps. +private func dummySession() -> GenerationSessionState { + GenerationSessionState( + kvCache: MockStateHandler(names: ["k"], shape: [1, 1, 1, 1, 1]), + additionalStates: nil, + hasNonTruncatableStates: false) +} + +/// Count-up logits at every position: position `p` of a row predicts that row's token at `p` plus one. +private func countUpAllPositions(_ tokensPerRow: [[Int32]]) -> [[LogitsScalarType]] { + tokensPerRow.map { row in row.flatMap { oneHot(($0 + 1) % 64) } } +} + /// Collects a request stream's tokens and its `.finished` reason. private func drain( _ stream: AsyncThrowingStream @@ -50,11 +63,14 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { var inFlight = 0 var maxInFlight = 0 var forwards = 0 + var scoringRows: [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 so far. + var scoringRows: [Int] { record.withLock { $0.scoringRows } } /// Prefill and decode calls so far. var forwards: Int { record.withLock { $0.forwards } } @@ -97,6 +113,17 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { if failDecode { throw Boom() } return tokensPerRow.map { oneHot(($0.last ?? 0) + 1) } } + + func allPositionLogits( + tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState + ) async throws -> [[LogitsScalarType]] { + await enter() + defer { leave() } + record.withLock { $0.scoringRows.append(tokensPerRow.count) } + return countUpAllPositions(tokensPerRow) + } + + func makeSessionState(batchSize: Int) throws -> GenerationSessionState { dummySession() } } /// A count-up `BatchedDecodeStep` whose `prefillRow` blocks until `release()`, so a test can keep a @@ -155,6 +182,14 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { ) 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") @@ -164,10 +199,7 @@ struct BatchSchedulerTests { ) -> BatchScheduler { BatchScheduler( decodeStep: step, - session: GenerationSessionState( - kvCache: MockStateHandler(names: ["k"], shape: [1, 1, 1, 1, 1]), - additionalStates: nil, - hasNonTruncatableStates: false), + session: dummySession(), maxConcurrent: maxConcurrent, maxQueued: maxQueued) } @@ -309,6 +341,77 @@ struct BatchSchedulerTests { #expect(order.withLock { $0 } == [0, 1, 2]) } + /// Count-up targets for `tokens`: each position's argmax, except the last, which has no real + /// next token. + private func scoreRequest(_ tokens: [Int32]) -> ScoreRequest { + ScoreRequest(tokens: tokens, targets: tokens.dropFirst() + [63], topK: 2) + } + + @Test("Scoring reports a log-probability per target, highest where the target is the argmax") + func scoringValues() async throws { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + let result = try await scheduler.score([scoreRequest([3, 4, 5])])[0] + #expect(result.entries.map(\.tokenId) == [4, 5, 63]) + #expect(result.entries[0].value == result.entries[1].value) + #expect(result.entries[2].value < result.entries[0].value) + #expect(result.entries[0].alternatives.count == 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 over two rows. The third is longer than a 128-position chunk. + 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) + + 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)) + #expect(batched.entries.count == prompt.count) + } + } + + @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, reason) = try await drain(stream) + #expect(tokens == [6, 7, 8, 9, 10, 11]) + #expect(reason == .maxTokens) + #expect(scores[0].entries.count == 3) + } + + @Test( + "Scoring batches stop before a third of the positions would be padding", + arguments: [ + ([7, 40, 130, 300], 4, 1), + ([130, 300], 4, 2), + ([64, 64, 64, 64, 64], 4, 4), + ([10, 200], 4, 1), + ([5], 4, 1), + ([], 4, 0), + ] as [([Int], Int, Int)]) + func scoringBatchCount(lengths: [Int], limit: Int, expected: Int) { + #expect(BatchScheduler.scoringBatchCount(lengths, limit: limit) == expected) + } + + @Test("A scoring request needs one target per token") + func scoringValidation() async { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + await #expect(throws: InferenceRuntimeError.self) { + _ = try await scheduler.score([ScoreRequest(tokens: [1, 2], targets: [2], topK: 1)]) + } + } + @Test("An empty prompt is rejected at submit") func emptyPromptRejected() { let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) From af1fa6d00e6f2bb9fc7fd9a1ad3c42862352f7d6 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 13:01:52 -0700 Subject: [PATCH 04/10] Clarify comments on batched serving Doc comments on the batch scheduler, llm-server, engine selection and llm-runner --batch-prototype now state what each type and function does, its inputs, outputs and throw conditions, and how to call it. Two small code changes come with it: - BatchScheduler(engine:maxConcurrent:maxQueued:) requires engine.supportsBatching when maxConcurrent is above 1. - ServerState holds the batch scheduler and its model config as one optional, so one cannot be set without the other. --- .../InferenceEngines/BatchScheduler.swift | 205 +++++++++--------- .../InferenceEngines/BatchServingRouter.swift | 17 +- .../BatchSessionAdapter.swift | 17 +- .../InferenceEngines/BatchedDecodeStep.swift | 9 +- .../CoreAISequentialEngine.swift | 6 +- .../InferenceEngines/EngineFactory.swift | 8 +- .../InferenceEngines/InferenceEngine.swift | 4 +- .../IsolatedSessionEngine.swift | 21 +- .../InferenceEngines/SchedulerSession.swift | 39 ++-- .../Tools/llm-runner/LLMRunnerMain.swift | 9 +- .../Tools/llm-server/ChatHandler.swift | 33 ++- .../Tools/llm-server/CompletionHandler.swift | 13 +- .../Tools/llm-server/LLMServerMain.swift | 10 +- .../Tools/llm-server/ReplayRunner.swift | 15 +- .../Tools/llm-server/ServerState.swift | 68 +++--- .../BatchSchedulerTests.swift | 48 ++-- 16 files changed, 243 insertions(+), 279 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift index a2b15a1d..4b3f8822 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -6,16 +6,16 @@ import Dispatch import Synchronization -/// One request submitted to a `RequestSink`. +/// A generation request for a `RequestSink`: prompt, sampling and stop conditions. package struct BatchRequest: Sendable { - /// Must not be empty; `submit(_:)` throws `invalidArgument` for an empty prompt. + /// List of prompt tokens, cannot be empty. let promptTokens: [Int32] let sampling: SamplingConfiguration let eosTokenIds: Set let maxTokens: Int let stopSequences: StopSequences - /// Turns a matched stop sequence's tokens into the text reported in `StopReason.stopSequence`, - /// so the scheduler needs no tokenizer. The default returns an empty string. + /// Converts a matched stop sequence's tokens to the text reported in `StopReason.stopSequence`. + /// Defaults to an empty string. let decodeStopText: @Sendable ([Int32]) -> String init( @@ -35,16 +35,18 @@ package struct BatchRequest: Sendable { } } -/// The tokens one request generated in `BatchScheduler.run(_:)` and why it stopped. -/// One scoring request for `BatchScheduler.score(_:)`. +/// A scoring request for `BatchScheduler.score(_:)`: prompt tokens and the target token to score at +/// each position. /// -/// `targets[p]` is scored against the logits at position `p` of `tokens`, so both have the same -/// length. For a prompt `[t0, t1, t2]` the usual targets are `[t1, t2, x]`, where the last entry has -/// no real next token and the caller ignores it. +/// For a prompt `[t0, t1, t2]` the usual targets are `[t1, t2, x]`. `x` can be any token, and the +/// caller ignores the last position's entry in the result. package struct ScoreRequest: Sendable { + /// Prompt tokens, cannot be empty. let tokens: [Int32] + /// Token scored at each position of `tokens`, same length as `tokens`. `targets[p]` is scored + /// against the logits at position `p`. let targets: [Int32] - /// How many alternatives to report per position. + /// Number of top alternatives reported per position. let topK: Int package init(tokens: [Int32], targets: [Int32], topK: Int) { @@ -54,62 +56,51 @@ package struct ScoreRequest: Sendable { } } +/// Generated tokens and stop reason for one request from `BatchScheduler.run(_:)`. struct RowResult: Sendable { let tokens: [Int32] let stopReason: StopReason } -/// One event on a request's stream from a `RequestSink`. A stream yields `.token` events and then -/// one `.finished`, or ends by throwing. +/// 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. `BatchScheduler` always sends nil `logits` because it samples internally. + /// A sampled token and, when the sink provides them, its logits. `BatchScheduler` always sends + /// `nil` logits. case token(Int32, logits: [LogitsScalarType]?) - /// Why the request stopped. Nothing follows it. + /// The reason the request stopped. Always the last event. case finished(StopReason) } -/// A shared batched backend that gives each submitted request its own event stream. -/// -/// The server does not call `submit(_:)` directly. It wraps the sink in one `SchedulerSession` per -/// request, which submits on `generate` and presents the stream as an ordinary `InferenceEngine`. -/// -/// ```swift -/// let engine = SchedulerSession(sink: scheduler, config: modelConfig, stopTokenIds: eosIds) -/// let sequence = try await engine.generate( -/// with: promptTokens, samplingConfiguration: sampling, inferenceOptions: options) -/// ``` +/// 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 { - /// The most requests that decode in the same batch. + /// Maximum number of requests decoded in one batch. At least 1. var maxConcurrent: Int { get } - /// Queues `request` and returns its event stream. When the sink rejects the request, for example - /// with `InferenceRuntimeError.queueFull`, this throws before any stream exists. + /// 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 } -/// Runs up to `maxConcurrent` requests through one shared batch-N session and admits new requests -/// while others are still decoding. +/// A `RequestSink` that decodes up to `maxConcurrent` requests in one shared batch session and +/// admits new requests while others are still decoding. Safe to call from any task. /// -/// Slot `i` is row `i` of the batch KV cache. A slot goes through these phases: -/// - Admission: a queued request takes a free slot and is prefilled into its row with -/// `BatchedDecodeStep.prefillRow`. -/// - Decode: each step runs `decodeBatch` over all `N` rows, every live row at its own cursor, so a -/// short request admitted late runs beside an older long one. -/// - Retirement: `BatchSessionAdapter` ends the row on EOS, a stop sequence or `maxTokens`. The next -/// pump iteration refills the slot from the queue while the other rows keep decoding. -/// -/// The class is `@unchecked Sendable` because the queue and row count sit behind a `Mutex`, and at -/// most one pump task runs at a time. That task alone drives the session. +/// 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 or `maxTokens`, and the slot takes the next +/// queued request while the other rows keep decoding. package final class BatchScheduler: RequestSink, @unchecked Sendable { private let decodeStep: any BatchedDecodeStep private let session: GenerationSessionState package let maxConcurrent: Int - /// How many requests may wait beyond the free rows before `submit(_:)` throws. `nil` means - /// unbounded. + /// Requests allowed to wait beyond the free rows, or `nil` for no limit. At least 0. private let maxQueued: Int? - /// Whether the server can skip its single-slot permit for the batched lane. This scheduler queues - /// and bounds requests itself, so that holds whenever it has more than one row. + /// True when more than one request can decode at a time (`maxConcurrent > 1`). package var supportsContinuousBatching: Bool { maxConcurrent > 1 } private struct Pending { @@ -120,16 +111,17 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { var queue: [Pending] = [] /// Rows held by admitted requests. var occupiedRows = 0 - /// Scoring requests waiting for the pump, in arrival order. + /// Scoring jobs waiting for the pump. Each `score(_:)` call appends its jobs shortest prompt + /// first. var scoring: [ScoringJob] = [] var pumping = false } private let state = Mutex(State()) - /// Held by the pump for each step and by an `ExclusiveEngineLease` for its whole lifetime, so - /// the engine never runs work from two tasks at once. + /// Lock held by the pump during each iteration, and by an `ExclusiveEngineLease` until it is + /// released. private let engineLock = AsyncFIFOLock() - /// Positions per scoring forward. Each chunk's logits are reduced to log-probabilities and - /// dropped, so memory stays at `rows * scoringChunkLength * vocab` logits however long the prompt. + /// Maximum number of positions per scoring forward pass. Each pass returns at most + /// `rows * scoringChunkLength * vocab` logits. private let scoringChunkLength = 128 private struct ScoringJob { @@ -137,10 +129,18 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { let continuation: AsyncThrowingStream.Continuation } - /// Creates a scheduler with a new `maxConcurrent`-row session on `engine`. - /// `engine.supportsBatching` must be true, and `maxConcurrent` must not exceed the graph's batch - /// capacity or a precondition fails. + /// 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) @@ -165,20 +165,17 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { /// Queues `request` and returns its event stream. /// - /// Throws `InferenceRuntimeError.queueFull`, which the server maps to HTTP 429, when - /// `queue.count + 1 > freeRows + maxQueued`. With `maxConcurrent == 2`, `maxQueued == 1` and both - /// rows busy, one request may wait and the next one throws. With `maxQueued == nil` it never throws. - /// - /// `queue.count` includes requests the pump has not admitted yet, so free rows absorb a burst - /// before `maxQueued` applies. In the same configuration with both rows free, a burst of four - /// requests accepts three and the fourth throws. + /// 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` atomically and starts the pump if none is running. Returns one stream per - /// request, in order. `submit(_:)` and `run(_:)` both go through here, so only the pump ever - /// drives `session`. + /// 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] { @@ -216,15 +213,17 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { var cursor: Int } - /// The single pump. Each iteration does the following: - /// 1. Admit queued requests into free slots. - /// 2. If no slot is live and the queue is empty, clear `pumping` and return. - /// 3. Run one `decodeBatch` over all `N` slots with `tokens` `[N][1]` and `cursors` `[N]`. A free + /// 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. + /// 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. - /// 4. Advance each live cursor, then sample, emit or retire each live row. + /// 5. Advance each live cursor, then sample, emit or retire each live row. /// - /// The step always spans `maxConcurrent` rows, so a lone request still pays for a full batch. - /// If `decodeBatch` throws, every resident row and every queued request ends with that error. + /// 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 { @@ -236,8 +235,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { let liveSlots = active.indices.filter { active[$0] != nil } if liveSlots.isEmpty { - // Recheck the queue under the lock. A request enqueued after `admit` saw - // `pumping == true` and did not start a pump, so this one must keep going. + // 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 @@ -258,8 +256,6 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } catch { for row in active where row != nil { row!.continuation.finish(throwing: error) } - // Drain the queue under the same lock that clears `pumping`. A queued request left - // behind would wait forever on a pump that has stopped. let (orphaned, orphanedScoring) = state.withLock { s -> ([Pending], [ScoringJob]) in let waiting = (s.queue, s.scoring) s.queue.removeAll() @@ -274,7 +270,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } /// Fills free slots from the queue. Each admitted request is prefilled into its row and its first - /// token is sampled at once, so a request can also retire here. + /// 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 { @@ -300,7 +296,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } - /// Samples one token for `slot`, emits it, and retires the row if the adapter says it is done. + /// 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( @@ -322,24 +318,20 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { // MARK: - Scoring - /// Returns the log-probability of each request's targets, in request order. + /// 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. /// - /// The requests join the pump's scoring queue. Between decode steps the pump takes up to - /// `maxConcurrent` waiting requests and runs them as one batch on a separate session: - /// - Rows shorter than the longest prompt are padded at the end. The causal mask keeps padding - /// from affecting real positions, and its logits are ignored. - /// - The batch advances `scoringChunkLength` positions per forward, and each chunk's logits are - /// reduced to log-probabilities before the next chunk runs. + /// 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. /// - /// Running decode rows wait for the scoring batch but keep their KV state. A failed scoring - /// forward fails only the requests in that batch. + /// 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") } - // Queue the requests shortest first so each batch holds similar lengths and pads less. The - // sort is stable, so the same requests always form the same batches. let order = requests.indices.sorted { requests[$0].tokens.count < requests[$1].tokens.count } var jobs: [ScoringJob] = [] var streams = [AsyncThrowingStream?](repeating: nil, count: requests.count) @@ -368,8 +360,8 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { return results } - /// Runs every waiting scoring request, `maxConcurrent` at a time. Called by the pump while it - /// holds the engine lock. + /// Runs every waiting scoring job, in batches of up to `maxConcurrent` sized by + /// `scoringBatchCount(_:limit:)`. Call only while holding `engineLock`. private func runPendingScoring() async { while true { let jobs = state.withLock { s -> [ScoringJob] in @@ -391,10 +383,10 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } - /// How many of the waiting scoring requests, with prompt lengths `lengths` in queue order, go into - /// the next batch. It takes up to `limit` requests and stops before one that would make more than - /// a third of the batch's positions padding, where batching stops paying for itself. For lengths - /// `[7, 40, 130, 300]` the batches are `[7]`, `[40]` and `[130, 300]`. + /// 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 scoringBatchCount(_ lengths: [Int], limit: Int) -> Int { guard let first = lengths.first else { return 0 } var count = 1 @@ -411,15 +403,17 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { return count } - /// Scores up to `maxConcurrent` requests in one padded batch. + /// Scores `requests` in one padded batch on a new session. At most `decodeStep.maxBatchCapacity` + /// requests. /// - /// For prompts of lengths 3 and 5 with a chunk length of 4, the rows are + /// Short rows repeat their last token. For prompts of lengths 3 and 5 with a chunk length of 4, + /// the rows are /// ``` /// chunk 0: [a0 a1 a2 a2] [b0 b1 b2 b3] positions 0..<4 /// chunk 1: [a2] [b4] position 4, row a is padding only /// ``` private func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { - // A static batch=N graph needs exactly N rows, so fill it with padding rows. + // Static graphs take exactly `capacity` rows; dynamic graphs (`Int.max`) take one per request. let capacity = decodeStep.maxBatchCapacity let rows = capacity == Int.max ? requests.count : capacity let session = try decodeStep.makeSessionState(batchSize: rows) @@ -437,7 +431,6 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { let logits = try await decodeStep.allPositionLogits( tokensPerRow: tokensPerRow, startPosition: start, session: session) let vocab = logits[0].count / q - // Reducing logits to log-probabilities is CPU work, so rows run in parallel. let chunkEntries = Mutex([[LogProbabilities.Entry]](repeating: [], count: requests.count)) let chunkStart = start DispatchQueue.concurrentPerform(iterations: requests.count) { row in @@ -460,11 +453,11 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { // MARK: - Exclusive engine access - /// Waits until the pump finishes its current step, then gives the caller sole use of the engine - /// until the lease is released. Resident rows stay in their slots and resume afterwards. + /// 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 on its own session of the same engine, such as one that needs per-step - /// logits, holds a lease for as long as it runs: + /// A request running on its own session of the same engine holds a lease for as long as it runs: /// /// ```swift /// let lease = await scheduler.acquireExclusiveAccess() @@ -478,8 +471,10 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { // MARK: - Closed batch - /// Submits `requests` together and waits for all of them. Returns one `RowResult` per request, in - /// order. It shares the pump with `submit(_:)`, so it is safe to call while other requests stream. + /// 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 @@ -504,8 +499,8 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } -/// Sole use of a `BatchScheduler`'s engine, from `acquireExclusiveAccess()`. Releasing it more than -/// once is harmless, and it releases on deinit, so a dropped lease cannot stall the batch. +/// 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) @@ -523,8 +518,8 @@ package final class ExclusiveEngineLease: Sendable { deinit { release() } } -/// An async lock that grants waiters in arrival order. `unlock()` hands the lock straight to the -/// next waiter, so a task that unlocks and locks again goes to the back of the line. +/// 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 diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift index 8918dde5..d6c09590 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchServingRouter.swift @@ -3,23 +3,18 @@ // 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 -/// Which serving lane a request takes under a batched server. +/// Serving lane of a request under a batched server. package enum BatchServingLane: Equatable, Sendable { - /// Served by the batch scheduler through a `SchedulerSession`. + /// Requests that share the batch scheduler through a `SchedulerSession`. case batched - /// Served by a single-request engine that returns per-step logits. + /// Requests that run one at a time on their own engine session, such as guided generation. case fallbackBatchOne } -/// Chooses the serving lane for a request when batched serving is enabled. -/// -/// The batched lane samples inside the scheduler and returns no per-step logits, so a request with a -/// JSON schema takes the fallback lane, where constrained decoding can mask the logits at each step. -/// Completions scoring does not use the router: with a scheduler it runs through -/// `BatchScheduler.score(_:)`, otherwise on the shared engine. +/// Lane selection for chat requests under a batched server. package enum BatchServingRouter { - /// Returns `.fallbackBatchOne` when the chat request has a `response_format` JSON schema, - /// otherwise `.batched`. + /// 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 index 39717150..dea78b50 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchSessionAdapter.swift @@ -3,33 +3,32 @@ // 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 -/// Per-row token bookkeeping and stop checks for one request in a batched decode. +/// Token history and stop checks for one request in a batched decode. /// -/// `BatchScheduler` keeps one adapter per occupied slot and feeds it each sampled token. When -/// `accept(_:)` returns false, the scheduler emits `.finished` and frees the slot for the next queued -/// request while the other rows keep decoding. +/// 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 - /// Set when the row retires. + /// 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 per-row step index for the sampler, equal to `generated.count`. + /// 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 reaching - /// `maxTokens`, checked in that order. Returns whether the row is still live. Once the row has - /// retired, it ignores further tokens and returns false. + /// 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 } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift index aaed7703..c75ca603 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`. /// @@ -47,13 +48,13 @@ protocol BatchedDecodeStep { ) async throws -> [LogitsScalarType] /// Runs `tokensPerRow` (`[N][q]`, every row at `startPosition`) and returns each row's logits at - /// every position, flat `[q * vocab]`. Scoring uses this to read all positions in one pass. + /// every position as `[N][q * vocab]`. func allPositionLogits( tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState ) async throws -> [[LogitsScalarType]] - /// A new session whose KV cache holds `batchSize` rows. + /// 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 140cb03f..0123ffd7 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -98,8 +98,8 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba return 1 // Single-batch graph. } - /// Whether this graph can run several requests in one forward. It needs an attention-mask input - /// and a batch dimension above 1, static or dynamic. + /// 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 @@ -442,7 +442,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba } /// Runs one batched forward step and returns the logits for every row and position, flat - /// `[N, q, vocab]`. `decodeBatch` and `allPositionLogits` slice this buffer. + /// `[N, q, vocab]`. private func runBatchedForward( tokensPerRow: [[Int32]], startPositions: [Int], diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/EngineFactory.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/EngineFactory.swift index 4a85a2f8..08fdb610 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/EngineFactory.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/EngineFactory.swift @@ -208,10 +208,12 @@ public struct EngineFactory: Sendable { // MARK: - Variant Resolution /// Resolves the engine variant based on override or auto-detection. - // Internal for EngineSupportTests. /// - /// A dynamic graph with an `attn_mask` input resolves to the sequential engine, which is the - /// only engine that fills the mask. + /// `takesAttentionMask` is true when the graph has an `attn_mask` input. Auto-detection picks + /// `.sequential` for a dynamic graph with one. Throws `unsupportedEngineVariant` for an unknown + /// or incompatible override, including `coreai-pipelined` on a dynamic graph with an `attn_mask` + /// input. + // Internal for EngineSupportTests. static func resolveVariant( override: String?, detectedStructure: ModelStructure, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 213f7427..e0425494 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -272,11 +272,11 @@ public enum InferenceRuntimeError: Error, LocalizedError { case unsupportedEngine(String) case unsupportedEngineVariant(String) case bufferAllocationFailed(String) - /// The scheduler rejected the request because its wait queue is at its limit of `depth`. + /// The scheduler's wait queue is full. `depth` is the queue limit. case queueFull(depth: Int) case genericError(String) - /// Whether this is `.queueFull`, which the server reports as HTTP 429. + /// True for `.queueFull`. package var isQueueFull: Bool { if case .queueFull = self { return true } return false diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift index cc5cfe4e..8e61449f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift @@ -3,21 +3,17 @@ // 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 -/// An `InferenceEngine` that runs one fallback-lane request, such as guided generation, on a -/// `CoreAISequentialEngine` through its own `GenerationSessionState`. +/// An `InferenceEngine` that runs one request on a `CoreAISequentialEngine` with its own +/// `GenerationSessionState`. Used for `.fallbackBatchOne` lane requests, such as guided generation. /// -/// Create one per request. The KV cache, token history and prefix-hit count belong to this session -/// rather than the engine's default session, so other sessions keep their state. It does not make -/// concurrent use of the engine safe: next to a `BatchScheduler`, hold an `ExclusiveEngineLease` -/// while the request runs. -/// -/// The class is `@unchecked Sendable` because a single request owns it, so no two tasks drive it at -/// once. +/// Create one per request and drive it from one task at a time. The KV cache, token history and +/// `lastPrefixHitCount` belong to this engine's session, separate from `base`'s. When `base` also +/// backs a `BatchScheduler`, hold the scheduler's `ExclusiveEngineLease` while the request runs. package final class IsolatedSessionEngine: InferenceEngine, @unchecked Sendable { package typealias ConfigType = ModelConfig private let base: CoreAISequentialEngine - /// Created on first use so `init` cannot throw. `reset(to:)` replaces it with a fresh one. + /// Session for the request, created on first use. `reset(to:)` replaces it. private var sessionState: GenerationSessionState? package init(base: CoreAISequentialEngine) { @@ -45,15 +41,14 @@ package final class IsolatedSessionEngine: InferenceEngine, @unchecked Sendable inferenceOptions: inferenceOptions) } - /// Replaces this request's session with a fresh one, whatever `tokenIndex` is. Guided generation - /// resets to 0 before generating, so this lane never needs a partial rewind. + /// Replaces this request's session with a fresh one. Ignores `tokenIndex`. package func reset(to tokenIndex: Int) async throws { sessionState = try base.makeSessionState() } package var supportsLogits: Bool { base.supportsLogits } - /// This session's busy state. The base engine's flag would reflect other sessions too. + /// True while a generation on this request's session is in flight. package var isBusy: Bool { sessionState?.tokenBox.isBusy ?? false } package var lastPrefixHitCount: Int { sessionState?.lastPrefixHitCount ?? 0 } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift index ba0256b5..a9140f7c 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift @@ -5,19 +5,10 @@ import CoreAIShared -/// An `InferenceEngine` that represents one request inside a `RequestSink` such as `BatchScheduler`. +/// An `InferenceEngine` that runs one request on a `RequestSink`, such as `BatchScheduler`. /// -/// Create one per request. Decoding strategies and the server see an ordinary engine while several -/// requests share one batched forward. `generate` submits a `BatchRequest` and returns the sink's -/// per-request event stream as a `GenerationSequence`. -/// -/// Differences from a single-request engine: -/// - `reset(to:)` does nothing. The sink owns the shared KV cache, so there is no per-request -/// history to rewind. -/// - `supportsLogits` is false. The sink samples internally, so guided generation must use the -/// fallback lane (see `BatchServingRouter`). Scoring goes through `BatchScheduler.score(_:)`. -/// - The sink stops a row only on `stopTokenIds` or `maxTokens`. Multi-token stop sequences and -/// detokenization stay in the per-request `DecodingStrategy`. +/// Create one per request. The sink ends the request only on `stopTokenIds` or `maxTokens`. Match +/// multi-token stop sequences in the caller's `DecodingStrategy`. /// /// ```swift /// let engine = SchedulerSession(sink: scheduler, config: modelConfig, stopTokenIds: eosIds) @@ -30,7 +21,7 @@ package final class SchedulerSession: InferenceEngine, @unchecked Sendable { package let config: ModelConfig private let sink: any RequestSink - /// Token ids that end the row inside the sink, usually EOS plus any additional EOS ids. + /// Token ids that end the request inside the sink. private let stopTokenIds: Set package init(sink: any RequestSink, config: ModelConfig, stopTokenIds: Set) { @@ -39,14 +30,16 @@ package final class SchedulerSession: InferenceEngine, @unchecked Sendable { 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 { - // Cap `maxTokens` at the room left in the context window after the prompt. Decoding past - // the window would make the shared batched forward throw, and a failed forward fails every - // resident row, so each row must retire with `.maxTokens` before that point. let contextRoom = Swift.max(1, config.maxContextLength - input.count) let requested = inferenceOptions.maxTokens ?? contextRoom let stream = try sink.submit( @@ -58,14 +51,13 @@ package final class SchedulerSession: InferenceEngine, @unchecked Sendable { return GenerationSequence(events: stream) } - /// Does nothing. The sink reclaims this request's state, and resetting the shared engine would - /// destroy the other rows. + /// Does nothing. package func reset(to tokenIndex: Int) async throws {} - /// Always false because the sink samples internally and returns no per-step logits. + /// Always false. package var supportsLogits: Bool { false } - /// Always false. Each session is one request, and the sink runs several at once. + /// Always false. package var isBusy: Bool { false } package var maxContextLength: Int { config.maxContextLength } @@ -74,7 +66,10 @@ package final class SchedulerSession: InferenceEngine, @unchecked Sendable { // MARK: - Output sequence extension SchedulerSession { - /// The sink's per-request event stream, presented as an engine output sequence. + /// 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. A reason set with `setStopReason(_:)` takes precedence. package struct GenerationSequence: InferenceOutputSequence { package typealias Element = InferenceOutput package typealias Failure = Error @@ -114,8 +109,6 @@ extension SchedulerSession.GenerationSequence { case .token(let id, let logits): return InferenceOutput(tokenId: id, logits: logits) case .finished(let reason): - // Record the sink's reason unless the consumer already set one. The - // `.maxTokens` default below only covers a stream that ends without `.finished`. stopReasonStore.setIfUnset(reason) } } diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index 1d680846..4798fd03 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -980,10 +980,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 2457db02..0685deae 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -114,8 +114,6 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, body: .init(byteBuffer: ByteBuffer(data: data))) } - // Requests with a JSON schema take the fallback lane. The streaming path hands the access to its - // body closure, and the non-streaming path releases it at each return. let lane = BatchServingRouter.lane( hasResponseSchema: chatRequest.responseFormat?.extractedSchema != nil) let access: RequestAccess @@ -144,9 +142,7 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, return response } } catch let error as InferenceRuntimeError where error.isQueueFull { - // The batch scheduler rejects a submit once it holds maxConcurrent + maxQueued requests, so - // free rows absorb a burst before the queue limit applies. Streaming requests submit before - // their SSE response starts, so they get the same 429. + // 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")) @@ -245,7 +241,6 @@ func runChatCompletion(chatRequest: ChatCompletionRequest, state: ServerState, s strategy = VanillaDecodingStrategy() } - // Guided generation needs per-step logits, which only the fallback lane provides. let lane = BatchServingRouter.lane(hasResponseSchema: schema != nil) let stream = try await strategy.decode( from: input, @@ -459,17 +454,16 @@ func prepareStreaming(chatRequest: ChatCompletionRequest, state: ServerState, se requestMaxTokens: requestMaxTokens, prefixReuseTokens: prefixReused) } -/// A streaming decode that has started. On the batched lane, starting it submits the request to -/// the batch scheduler, so a full scheduler throws before the handler sends any SSE headers. -/// -/// `@unchecked Sendable` because it is created once and then handed to a single consumer, the -/// response body or the replay runner, which iterates it on one task. +/// 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 } -/// Picks the decoding strategy for `chatRequest` and starts it on the engine for its lane. +/// 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 { @@ -494,10 +488,9 @@ func startStreamingDecode( return StartedDecode(tokens: tokens, start: start) } -/// 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. +/// 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 @@ -620,20 +613,20 @@ 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, access: RequestAccess ) async throws -> Response { - // Preflight throws for bad requests before any SSE framing (mapped to HTTP 400 upstream), and - // starting the decode throws queueFull (mapped to HTTP 429) when the batch scheduler is full. 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 access for the stream's lifetime. Its deinit releases it if the - // framework drops the response before running this closure. + // If this closure never runs, the permit and lease release on deinit. defer { access.release() } do { let encoder = JSONEncoder() diff --git a/swift/Sources/Tools/llm-server/CompletionHandler.swift b/swift/Sources/Tools/llm-server/CompletionHandler.swift index 325079e6..6a981335 100644 --- a/swift/Sources/Tools/llm-server/CompletionHandler.swift +++ b/swift/Sources/Tools/llm-server/CompletionHandler.swift @@ -30,8 +30,6 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw status: .notImplemented, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))) } - // One scoring request runs at a time. With a batch scheduler its prompts score together - // between decode steps, so it holds the permit but no lease. let access: RequestAccess do { access = try await state.acquireScoringAccess() @@ -131,15 +129,15 @@ private func handleLoglikelihood(req: CompletionRequest, state: ServerState) asy ) } -/// One prompt ready to score: its tokens and the target each position predicts. +/// Tokenized prompt to score, with the target token each position predicts. private struct PreparedPrompt { let tokens: [Int32] - /// `tokens.dropFirst()` plus a padding token. lm-eval sends echo=true, max_tokens=1 and slices - /// logprobs[ctxlen:-1], so the extra last entry is the one it drops. + /// `tokens.dropFirst()` followed by `tokens[0]` as padding. Same count as `tokens`. let targets: [Int32] } -/// Tokenizes `prompt`, or returns nil when it has fewer than two tokens and nothing to score. +/// 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 { @@ -158,7 +156,8 @@ private func preparePrompt(_ prompt: Prompt, state: ServerState) throws -> Prepa return PreparedPrompt(tokens: allTokens, targets: Array(allTokens.dropFirst()) + [allTokens[0]]) } -/// Scores one prompt on the shared engine, which runs it in one prefill with every position's logits. +/// 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 { diff --git a/swift/Sources/Tools/llm-server/LLMServerMain.swift b/swift/Sources/Tools/llm-server/LLMServerMain.swift index 39174866..4a9db7a5 100644 --- a/swift/Sources/Tools/llm-server/LLMServerMain.swift +++ b/swift/Sources/Tools/llm-server/LLMServerMain.swift @@ -55,8 +55,7 @@ struct LLMServer: AsyncParsableCommand { @Flag( name: .customLong("replay-concurrent"), help: ArgumentHelp( - "Run all --replay requests concurrently instead of one at a time, so the batch scheduler can " - + "fill several rows. Use with --concurrency N.") + "Run all --replay requests concurrently instead of one at a time. Use with --concurrency N.") ) var replayConcurrent: Bool = false @@ -237,9 +236,7 @@ struct LLMServer: AsyncParsableCommand { maxQueueDepth: maxQueueDepth ) - // Batched serving needs --concurrency above 1 and a CoreAISequentialEngine whose graph has an - // attention mask input and a batch dimension above 1. The scheduler holds up to - // concurrency + --max-queue-depth requests, decoding or waiting, and rejects the next one. + // The scheduler holds up to --concurrency plus --max-queue-depth requests, then rejects. var batchScheduler: BatchScheduler? = nil if concurrency > 1 { if let seq = engine as? CoreAISequentialEngine, seq.supportsBatching { @@ -264,8 +261,7 @@ struct LLMServer: AsyncParsableCommand { engine: engine, tokenizer: tokenizer, config: config, - batchScheduler: batchScheduler, - batchModelConfig: batchScheduler != nil ? engineConfig : nil + batchServing: batchScheduler.map { (scheduler: $0, modelConfig: engineConfig) } ) if let replayPath { diff --git a/swift/Sources/Tools/llm-server/ReplayRunner.swift b/swift/Sources/Tools/llm-server/ReplayRunner.swift index 458115b0..7965cd9d 100644 --- a/swift/Sources/Tools/llm-server/ReplayRunner.swift +++ b/swift/Sources/Tools/llm-server/ReplayRunner.swift @@ -11,10 +11,10 @@ import Foundation /// /// 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). By default requests run one at a time, matching the single-request server. -/// With `concurrent`, they all start at once so the batch scheduler can fill several rows. -/// 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, concurrent: Bool = false @@ -24,7 +24,6 @@ enum ReplayRunner { var lines = [String](repeating: "", count: requests.count) if concurrent { - // Lines keep input order. Requests past the queue limits come back as queue-full error lines. try await withThrowingTaskGroup(of: (Int, String).self) { group in for (i, req) in requests.enumerated() { group.addTask { (i, await oneLine(req, state: state)) } @@ -45,14 +44,12 @@ enum ReplayRunner { } } - /// Runs one replay request and encodes its result line. Errors become an error result line. + /// 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 { - // runChatCompletion and runStreamingLoop do not take the request access, so take it here - // as the HTTP handler does. Without it, concurrent replay would drive the shared engine - // from several tasks at once. let lane = BatchServingRouter.lane( hasResponseSchema: req.request.responseFormat?.extractedSchema != nil) let access = try await state.acquireAccess(lane: lane) diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 320e0994..623b5dcb 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -129,11 +129,10 @@ final class ServerState: @unchecked Sendable { /// guaranteed while this value is unchanged; it changes when output-affecting /// configuration (model, engine, context length, sampling defaults) changes. let systemFingerprint: String - /// Scheduler for the batched lane, or nil when the server handles one request at a time. - private let batchScheduler: BatchScheduler? - /// Config for each `SchedulerSession`. Set it whenever `batchScheduler` is set. If it is nil - /// while a scheduler is present, the batched lane runs on the shared engine without the permit. - private let batchModelConfig: ModelConfig? + /// 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 { @@ -149,8 +148,7 @@ final class ServerState: @unchecked Sendable { engine: any InferenceEngine, tokenizer: any Tokenizer, config: ServerConfig, - batchScheduler: BatchScheduler? = nil, - batchModelConfig: ModelConfig? = nil + batchServing: (scheduler: BatchScheduler, modelConfig: ModelConfig)? = nil ) { self.engine = engine self.tokenizer = tokenizer @@ -159,12 +157,11 @@ final class ServerState: @unchecked Sendable { self.toolCallDetection = detectToolCallFormat(using: tokenizer) self.thinkingFormat = detectThinkingFormat(using: tokenizer) self.systemFingerprint = Self.makeSystemFingerprint(config: config, engine: engine) - self.batchScheduler = batchScheduler - self.batchModelConfig = batchModelConfig + self.batchServing = batchServing } - /// Token ids that retire a row in the scheduler: EOS and the additional EOS ids. Multi-token - /// stop sequences are still matched by the decoding strategy. + /// 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)) } @@ -172,22 +169,20 @@ final class ServerState: @unchecked Sendable { return ids } - /// Returns the engine one request on `lane` should run on. + /// Returns the engine for one request on `lane`. Call once per request. /// - /// - `.batched`: a new `SchedulerSession` over `batchScheduler`, or the shared base engine when - /// no scheduler is configured. - /// - `.fallbackBatchOne`: with a scheduler running, a new `IsolatedSessionEngine` with its own - /// session, so the batched rows keep their KV state while the batch is paused for it (see - /// `acquireAccess(lane:)`). Without a scheduler, the shared base engine, - /// which keeps its prefix cache. + /// - `.batched`: a new `SchedulerSession` over the batch scheduler, or the shared engine when + /// batched serving is off. + /// - `.fallbackBatchOne`: a new `IsolatedSessionEngine` with its own KV session when + /// `batchScheduler` is set and the shared engine is a `CoreAISequentialEngine`. Otherwise the + /// shared engine. func engineForRequest(lane: BatchServingLane) -> any InferenceEngine { switch lane { case .batched: - guard let scheduler = batchScheduler, let modelConfig = batchModelConfig else { - return engine - } + guard let batchServing else { return engine } return SchedulerSession( - sink: scheduler, config: modelConfig, stopTokenIds: schedulerStopTokenIds) + sink: batchServing.scheduler, config: batchServing.modelConfig, + stopTokenIds: schedulerStopTokenIds) case .fallbackBatchOne: if batchScheduler != nil, let seq = engine as? CoreAISequentialEngine { return IsolatedSessionEngine(base: seq) @@ -196,12 +191,15 @@ final class ServerState: @unchecked Sendable { } } - /// Waits for what a request on `lane` must hold while it uses the engine, and returns it. + /// 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. /// - /// - The single-slot permit from `queue`, except on the batched lane when a continuous-batching - /// scheduler admits the request. Throws `ServerError.queueFull` when the queue is full. - /// - On the fallback lane with a scheduler running, an `ExclusiveEngineLease` as well, so the - /// batch pauses at a step boundary instead of sharing the engine. + /// 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: @@ -218,18 +216,18 @@ final class ServerState: @unchecked Sendable { return RequestAccess(permit: permit, lease: lease) } - /// The scheduler that scores completions prompts in batches, or nil when the server scores them - /// one at a time on the shared engine. + /// Scheduler that scores completion prompts in batches, or nil when batched serving is off. var scoringScheduler: BatchScheduler? { batchScheduler } - /// Waits for the single-slot permit a scoring request holds. Scoring takes no lease: with a - /// scheduler, its prompts run on the scheduler's own task between decode steps. + /// 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) } - /// Whether a request on `lane` needs the single-slot permit. The batched lane does not when a - /// continuous-batching scheduler admits it, because the scheduler enforces its own queue limit. + /// 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: @@ -405,8 +403,8 @@ enum RequestID { } } -/// What one request holds while it uses the engine. Releasing it more than once is harmless, and -/// the permit and lease each release on deinit. +/// 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? diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift index f5f57912..67c9703c 100644 --- a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -8,14 +8,14 @@ import Testing @testable import CoreAILanguageModels -/// Logits over a 64-token vocabulary that greedy sampling resolves to `token`. +/// 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 } -/// A session with no real KV cache, for mock decode steps. +/// 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]), @@ -23,7 +23,8 @@ private func dummySession() -> GenerationSessionState { hasNonTruncatableStates: false) } -/// Count-up logits at every position: position `p` of a row predicts that row's token at `p` plus one. +/// 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.map { row in row.flatMap { oneHot(($0 + 1) % 64) } } } @@ -43,18 +44,18 @@ private func drain( return (tokens, reason) } -/// A `BatchedDecodeStep` that predicts the input token plus one, so a request with prompt `[..., k]` -/// and `maxTokens: m` generates `k+1` through `k+m` whatever its slot or admission time. +/// Mock `BatchedDecodeStep` that predicts the last input token plus one. A request with prompt +/// `[..., k]` and `maxTokens: m` generates `k+1` through `k+m`. /// -/// It records each decode's `startPositions` and the peak number of overlapping forward calls. -/// `forwardDelay` keeps each call in flight longer so that an overlap is reliably observed. +/// 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 for prompts that start with this token. + /// `prefillRow` throws `Boom` for a prompt that starts with this token. private let failPrefillFor: Int32? - /// Every `decodeBatch` throws. + /// When true, every `decodeBatch` throws `Boom`. private let failDecode: Bool private let forwardDelay: Duration? @@ -69,9 +70,9 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { var observedStartPositions: [[Int]] { record.withLock { $0.startPositions } } var maxInFlight: Int { record.withLock { $0.maxInFlight } } - /// Row count of each scoring forward so far. + /// Row count of each scoring forward, in call order. var scoringRows: [Int] { record.withLock { $0.scoringRows } } - /// Prefill and decode calls so far. + /// Number of prefill, decode and scoring forward calls so far. var forwards: Int { record.withLock { $0.forwards } } init( @@ -126,8 +127,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { func makeSessionState(batchSize: Int) throws -> GenerationSessionState { dummySession() } } -/// A count-up `BatchedDecodeStep` whose `prefillRow` blocks until `release()`, so a test can keep a -/// row occupying its slot while it fills the queue. +/// Mock count-up `BatchedDecodeStep` whose `prefillRow` blocks until `release()`. private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { let maxBatchCapacity: Int @@ -141,7 +141,7 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { init(capacity: Int) { maxBatchCapacity = capacity } - /// Returns once a `prefillRow` call has started, which means its row holds a slot. + /// Waits until the first `prefillRow` call has started. func waitEntered() async { await withCheckedContinuation { cont in gate.withLock { g in @@ -150,7 +150,7 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { } } - /// Unblocks the waiting `prefillRow`. Later calls no longer block. + /// Unblocks the waiting `prefillRow` call. Later calls return without blocking. func release() { let waiter: CheckedContinuation? = gate.withLock { g in g.released = true @@ -260,11 +260,11 @@ struct BatchSchedulerTests { req(prompt: [1, 2, 3], maxTokens: 3), ]) #expect(results.map(\.tokens) == [[2, 3, 4], [4, 5, 6]]) - // Each cursor starts at the row's prompt length. + // Each row's first decode starts at its prompt length. #expect(mock.observedStartPositions == [[1, 3], [2, 4]]) } - @Test("Run and submit share one pump, so forwards never overlap") + @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) @@ -277,8 +277,7 @@ struct BatchSchedulerTests { @Test("Submit throws queueFull past free rows plus maxQueued", arguments: [false, true]) func backpressure(waitForAdmission: Bool) async throws { - // Queued plus admitted requests may total maxConcurrent + maxQueued, here 3. The limit is - // the same whether the first request is still queued (a burst) or already holds a row. + // Accepts 3 requests, maxConcurrent plus maxQueued, whether or not the first holds a row yet. let mock = GatedPrefillStep(capacity: 2) let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: 1) var streams = [try scheduler.submit(req(prompt: [0], maxTokens: 1))] @@ -333,7 +332,7 @@ struct BatchSchedulerTests { order.withLock { $0.append(i) } lock.unlock() } - // Let each waiter enqueue before starting the next one. + // Waiter i enqueues on the lock before waiter i + 1 starts. try? await Task.sleep(for: .milliseconds(20)) } lock.unlock() @@ -341,8 +340,8 @@ struct BatchSchedulerTests { #expect(order.withLock { $0 } == [0, 1, 2]) } - /// Count-up targets for `tokens`: each position's argmax, except the last, which has no real - /// next token. + /// Makes a top-2 score request for `tokens` whose target at each position is the next token, + /// and 63 at the last position. private func scoreRequest(_ tokens: [Int32]) -> ScoreRequest { ScoreRequest(tokens: tokens, targets: tokens.dropFirst() + [63], topK: 2) } @@ -366,7 +365,7 @@ struct BatchSchedulerTests { func scoringBatchInvariance(capacity: Int, expectedRows: [Int]) async throws { let step = CountUpDecodeStep(capacity: capacity) let scheduler = makeScheduler(step, maxConcurrent: 2) - // Three requests over two rows. The third is longer than a 128-position chunk. + // 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) @@ -440,8 +439,7 @@ struct BatchSchedulerTests { let streams = try [Int32(1), 2, 3].map { try scheduler.submit(req(prompt: [$0], maxTokens: 4)) } - // With one slot, requests 2 and 3 wait in the queue while request 1's decode throws. If - // the queue were not drained, iterating their streams would hang. + // With one slot, requests 2 and 3 are still queued when request 1's decode throws. for stream in streams { await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(stream) } } From b005497a46f3cf2ae08caa41e6e2715cd74ac2d8 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 14:13:33 -0700 Subject: [PATCH 05/10] Simplify batched serving and cut logits copies - Fallback requests (chat with a JSON schema) run on the shared engine under the exclusive lease. IsolatedSessionEngine is removed. - runBatchedForward returns the logits NDArray. decodeBatch copies only each row's last position, and allPositionLogits returns one flat [N, q, vocab] buffer. - LogProbabilities computes from a contiguous [positions, vocab] buffer and reuses its Float32 scratch buffers across positions. Scoring passes buffer slices instead of per-position arrays. - BatchScheduler.swift lists BatchScheduler first, then the types it uses. The class doc has a diagram of how requests fill the batch rows. - llm-benchmark help and errors no longer ask for --inference-engine-variant coreai-sequential, which the default now picks for attn_mask graphs. - BatchSchedulerTests is trimmed from 19 tests to 12 with the same coverage. --- .../InferenceEngines/BatchScheduler.swift | 199 ++++++++-------- .../InferenceEngines/BatchedDecodeStep.swift | 6 +- .../CoreAISequentialEngine.swift | 46 ++-- .../InferenceEngines/InferenceEngine.swift | 6 +- .../IsolatedSessionEngine.swift | 61 ----- .../TextGeneration/LogProbabilities.swift | 131 ++++++----- .../Tools/benchmark/BenchmarkMain.swift | 8 +- .../Tools/llm-server/ServerState.swift | 8 +- .../BatchSchedulerTests.swift | 217 ++++++------------ 9 files changed, 297 insertions(+), 385 deletions(-) delete mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift index 4b3f8822..9ce73478 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -6,86 +6,21 @@ import Dispatch import Synchronization -/// 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 - - 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 scoring request for `BatchScheduler.score(_:)`: prompt tokens and the target token to score at -/// each position. -/// -/// For a prompt `[t0, t1, t2]` the usual targets are `[t1, t2, x]`. `x` can be any token, and the -/// caller ignores the last position's entry in the result. -package struct ScoreRequest: Sendable { - /// Prompt tokens, cannot be empty. - let tokens: [Int32] - /// Token scored at each position of `tokens`, same length as `tokens`. `targets[p]` is scored - /// against the logits at position `p`. - let targets: [Int32] - /// Number of top alternatives reported per position. - let topK: Int - - package init(tokens: [Int32], targets: [Int32], topK: Int) { - self.tokens = tokens - self.targets = targets - self.topK = topK - } -} - -/// Generated tokens and stop reason for one request from `BatchScheduler.run(_:)`. -struct RowResult: Sendable { - let tokens: [Int32] - let stopReason: StopReason -} - -/// 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 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 `RequestSink` that decodes up to `maxConcurrent` requests in one shared batch session and /// admits new requests while others are still decoding. Safe to call from any task. /// +/// ``` +/// submit(_:) ---> 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 @@ -93,6 +28,8 @@ package protocol RequestSink: Sendable { /// - 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 or `maxTokens`, and the slot takes the next /// queued request while the other rows keep decoding. +/// +/// `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 @@ -430,20 +367,22 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } let logits = try await decodeStep.allPositionLogits( tokensPerRow: tokensPerRow, startPosition: start, session: session) - let vocab = logits[0].count / q + let vocab = logits.count / (rows * q) let chunkEntries = Mutex([[LogProbabilities.Entry]](repeating: [], count: requests.count)) let chunkStart = start - DispatchQueue.concurrentPerform(iterations: requests.count) { row in - let request = requests[row] - let real = min(q, request.tokens.count - chunkStart) - guard real > 0 else { return } - let rowLogits = (0.. 0 else { return } + let rowStart = row * q * vocab + let rowEntries = LogProbabilities.compute( + logits: UnsafeBufferPointer(rebasing: buffer[rowStart.. ExclusiveEngineLease { await engineLock.lock() @@ -499,6 +439,56 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } +/// 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 + + 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 + } +} + +/// 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 { @@ -552,3 +542,30 @@ final class AsyncFIFOLock: Sendable { next?.resume() } } + +/// A scoring request for `BatchScheduler.score(_:)`: prompt tokens and the target token to score at +/// each position. +/// +/// For a prompt `[t0, t1, t2]` the usual targets are `[t1, t2, x]`. `x` can be any token, and the +/// caller ignores the last position's entry in the result. +package struct ScoreRequest: Sendable { + /// Prompt tokens, cannot be empty. + let tokens: [Int32] + /// Token scored at each position of `tokens`, same length as `tokens`. `targets[p]` is scored + /// against the logits at position `p`. + let targets: [Int32] + /// Number of top alternatives reported per position. + let topK: Int + + package init(tokens: [Int32], targets: [Int32], topK: Int) { + self.tokens = tokens + self.targets = targets + self.topK = topK + } +} + +/// 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/BatchedDecodeStep.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift index c75ca603..ac36879b 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift @@ -47,13 +47,13 @@ protocol BatchedDecodeStep { session: GenerationSessionState ) async throws -> [LogitsScalarType] - /// Runs `tokensPerRow` (`[N][q]`, every row at `startPosition`) and returns each row's logits at - /// every position as `[N][q * vocab]`. + /// 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]] + ) 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 0123ffd7..82edf8dc 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -417,37 +417,57 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba startPositions: [Int], session: GenerationSessionState ) async throws -> [[LogitsScalarType]] { - let flat = try await runBatchedForward( + 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 (0.. [[LogitsScalarType]] { + ) async throws -> [LogitsScalarType] { let n = tokensPerRow.count - let flat = try await runBatchedForward( + let out = try await runBatchedForward( tokensPerRow: tokensPerRow, startPositions: Array(repeating: startPosition, count: n), session: session) - let rowLength = flat.count / n - return (0.. [[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.. [LogitsScalarType] { + ) async throws -> 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") @@ -527,7 +547,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, Ba } } - return readNDArray(out, as: LogitsScalarType.self, count: n * q * vocab) + return out } /// See ``BatchedDecodeStep`` for the contract. diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index e0425494..7270866a 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -278,8 +278,10 @@ public enum InferenceRuntimeError: Error, LocalizedError { /// True for `.queueFull`. package var isQueueFull: Bool { - if case .queueFull = self { return true } - return false + switch self { + case .queueFull: true + default: false + } } public var errorDescription: String? { diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift deleted file mode 100644 index 8e61449f..00000000 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/IsolatedSessionEngine.swift +++ /dev/null @@ -1,61 +0,0 @@ -// 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 - -/// An `InferenceEngine` that runs one request on a `CoreAISequentialEngine` with its own -/// `GenerationSessionState`. Used for `.fallbackBatchOne` lane requests, such as guided generation. -/// -/// Create one per request and drive it from one task at a time. The KV cache, token history and -/// `lastPrefixHitCount` belong to this engine's session, separate from `base`'s. When `base` also -/// backs a `BatchScheduler`, hold the scheduler's `ExclusiveEngineLease` while the request runs. -package final class IsolatedSessionEngine: InferenceEngine, @unchecked Sendable { - package typealias ConfigType = ModelConfig - - private let base: CoreAISequentialEngine - /// Session for the request, created on first use. `reset(to:)` replaces it. - private var sessionState: GenerationSessionState? - - package init(base: CoreAISequentialEngine) { - self.base = base - } - - private func ensureSession() throws -> GenerationSessionState { - if let sessionState { return sessionState } - let fresh = try base.makeSessionState() - sessionState = fresh - return fresh - } - - package var config: ModelConfig { base.config } - - package func generate( - with input: [TokenId], - samplingConfiguration: SamplingConfiguration, - inferenceOptions: InferenceOptions - ) async throws -> CoreAISequentialEngine.GenerationSequence { - try await base.generate( - with: input, - sessionState: try ensureSession(), - samplingConfiguration: samplingConfiguration, - inferenceOptions: inferenceOptions) - } - - /// Replaces this request's session with a fresh one. Ignores `tokenIndex`. - package func reset(to tokenIndex: Int) async throws { - sessionState = try base.makeSessionState() - } - - package var supportsLogits: Bool { base.supportsLogits } - - /// True while a generation on this request's session is in flight. - package var isBusy: Bool { sessionState?.tokenBox.isBusy ?? false } - - package var lastPrefixHitCount: Int { sessionState?.lastPrefixHitCount ?? 0 } - - package func cancel() async throws { - sessionState?.tokenBox.cancelActive() - } - - package var maxContextLength: Int { base.config.maxContextLength } -} 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-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 623b5dcb..26eb8498 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -173,9 +173,8 @@ final class ServerState: @unchecked Sendable { /// /// - `.batched`: a new `SchedulerSession` over the batch scheduler, or the shared engine when /// batched serving is off. - /// - `.fallbackBatchOne`: a new `IsolatedSessionEngine` with its own KV session when - /// `batchScheduler` is set and the shared engine is a `CoreAISequentialEngine`. Otherwise the - /// shared engine. + /// - `.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: @@ -184,9 +183,6 @@ final class ServerState: @unchecked Sendable { sink: batchServing.scheduler, config: batchServing.modelConfig, stopTokenIds: schedulerStopTokenIds) case .fallbackBatchOne: - if batchScheduler != nil, let seq = engine as? CoreAISequentialEngine { - return IsolatedSessionEngine(base: seq) - } return engine } } diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift index 67c9703c..61be9d62 100644 --- a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -25,8 +25,8 @@ private func dummySession() -> GenerationSessionState { /// 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.map { row in row.flatMap { oneHot(($0 + 1) % 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. @@ -117,7 +117,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { func allPositionLogits( tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState - ) async throws -> [[LogitsScalarType]] { + ) async throws -> [LogitsScalarType] { await enter() defer { leave() } record.withLock { $0.scoringRows.append(tokensPerRow.count) } @@ -132,24 +132,13 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { let maxBatchCapacity: Int private struct Gate { - var entered = false var released = false - var enteredWaiter: CheckedContinuation? var releaseWaiter: CheckedContinuation? } private let gate = Mutex(Gate()) init(capacity: Int) { maxBatchCapacity = capacity } - /// Waits until the first `prefillRow` call has started. - func waitEntered() async { - await withCheckedContinuation { cont in - gate.withLock { g in - if g.entered { cont.resume() } else { g.enteredWaiter = cont } - } - } - } - /// Unblocks the waiting `prefillRow` call. Later calls return without blocking. func release() { let waiter: CheckedContinuation? = gate.withLock { g in @@ -163,12 +152,6 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { func prefillRow( row: Int, promptTokens: [Int32], session: GenerationSessionState ) async throws -> [LogitsScalarType] { - let enteredWaiter: CheckedContinuation? = gate.withLock { g in - g.entered = true - defer { g.enteredWaiter = nil } - return g.enteredWaiter - } - enteredWaiter?.resume() await withCheckedContinuation { (cont: CheckedContinuation) in gate.withLock { g in if g.released { cont.resume() } else { g.releaseWaiter = cont } @@ -185,7 +168,7 @@ private final class GatedPrefillStep: BatchedDecodeStep, @unchecked Sendable { func allPositionLogits( tokensPerRow: [[Int32]], startPosition: Int, session: GenerationSessionState - ) async throws -> [[LogitsScalarType]] { + ) async throws -> [LogitsScalarType] { countUpAllPositions(tokensPerRow) } @@ -228,40 +211,17 @@ struct BatchSchedulerTests { #expect(results.map(\.stopReason) == [.eos, .maxTokens, .stopSequence("8")]) } - @Test("Submit streams the tokens, then .finished") - func streamingSubmit() async throws { - let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) - let (tokens, reason) = try await drain( - scheduler.submit(req(prompt: [5], maxTokens: 100, eos: [8]))) - // The EOS token itself is streamed. - #expect(tokens == [6, 7, 8]) - #expect(reason == .eos) - } - - @Test("A retired row's slot admits the next queued request") - func slotReuse() async throws { - let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) - let results = try await scheduler.run([ - req(prompt: [0], maxTokens: 2), - req(prompt: [10], maxTokens: 3), - req(prompt: [20], maxTokens: 1), - req(prompt: [30], maxTokens: 4), - ]) - #expect(results.map(\.tokens) == [[1, 2], [11, 12, 13], [21], [31, 32, 33, 34]]) - #expect(results.allSatisfy { $0.stopReason == .maxTokens }) - } - - @Test("Rows decode at their own cursors in one forward") - func raggedCursors() async throws { + @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: 3), + req(prompt: [1], maxTokens: 2), req(prompt: [1, 2, 3], maxTokens: 3), + req(prompt: [20], maxTokens: 1), ]) - #expect(results.map(\.tokens) == [[2, 3, 4], [4, 5, 6]]) - // Each row's first decode starts at its prompt length. - #expect(mock.observedStartPositions == [[1, 3], [2, 4]]) + #expect(results.map(\.tokens) == [[2, 3], [4, 5, 6], [21]]) + #expect(mock.observedStartPositions.first == [1, 3]) } @Test("Forwards from run and submit never overlap") @@ -275,27 +235,13 @@ struct BatchSchedulerTests { #expect(mock.maxInFlight == 1) } - @Test("Submit throws queueFull past free rows plus maxQueued", arguments: [false, true]) - func backpressure(waitForAdmission: Bool) async throws { - // Accepts 3 requests, maxConcurrent plus maxQueued, whether or not the first holds a row yet. - let mock = GatedPrefillStep(capacity: 2) - let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: 1) - var streams = [try scheduler.submit(req(prompt: [0], maxTokens: 1))] - if waitForAdmission { await mock.waitEntered() } - for _ in 0..<2 { streams.append(try scheduler.submit(req(prompt: [0], maxTokens: 1))) } - #expect(throws: InferenceRuntimeError.self) { - _ = try scheduler.submit(req(prompt: [0], maxTokens: 1)) - } - - mock.release() - for stream in streams { _ = try await drain(stream) } - } - - @Test("maxQueued 0 accepts only as many requests as there are free rows") - func noWaiting() async throws { + @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: 0) - let streams = try (0..<2).map { _ in try scheduler.submit(req(prompt: [0], maxTokens: 1)) } + let scheduler = makeScheduler(mock, maxConcurrent: 2, maxQueued: maxQueued) + let streams = try (0..([]) - await lock.lock() - await withTaskGroup(of: Void.self) { group in - for i in 0..<3 { - group.addTask { - await lock.lock() - order.withLock { $0.append(i) } - lock.unlock() - } - // Waiter i enqueues on the lock before waiter i + 1 starts. - try? await Task.sleep(for: .milliseconds(20)) - } - lock.unlock() + @Test("Empty prompts and mismatched scoring targets are rejected") + func invalidRequests() async { + let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) + #expect(throws: InferenceRuntimeError.self) { + _ = try scheduler.submit(req(prompt: [], maxTokens: 1)) + } + await #expect(throws: InferenceRuntimeError.self) { + _ = try await scheduler.score([ScoreRequest(tokens: [1, 2], targets: [2], topK: 1)]) } - #expect(order.withLock { $0 } == [0, 1, 2]) } + @Test("A prefill failure fails only its own request") + func prefillFailureIsolated() async throws { + let scheduler = makeScheduler( + CountUpDecodeStep(capacity: 2, failPrefillFor: 50), maxConcurrent: 2) + let failing = try scheduler.submit(req(prompt: [50], maxTokens: 4)) + let healthy = try scheduler.submit(req(prompt: [5], maxTokens: 3)) + + await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(failing) } + let (tokens, reason) = try await drain(healthy) + #expect(tokens == [6, 7, 8]) + #expect(reason == .maxTokens) + } + + @Test("A failed shared decode also fails the queued requests") + func decodeFailureDrainsQueue() async throws { + let scheduler = makeScheduler( + CountUpDecodeStep(capacity: 1, failDecode: true), maxConcurrent: 1) + let streams = try [Int32(1), 2, 3].map { + try scheduler.submit(req(prompt: [$0], maxTokens: 4)) + } + // With one slot, requests 2 and 3 are still queued when request 1's decode throws. + for stream in streams { + await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(stream) } + } + } + + // MARK: - Scoring + /// Makes a top-2 score request for `tokens` whose target at each position is the next token, /// and 63 at the last position. private func scoreRequest(_ tokens: [Int32]) -> ScoreRequest { ScoreRequest(tokens: tokens, targets: tokens.dropFirst() + [63], topK: 2) } - @Test("Scoring reports a log-probability per target, highest where the target is the argmax") - func scoringValues() async throws { - let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) - let result = try await scheduler.score([scoreRequest([3, 4, 5])])[0] - #expect(result.entries.map(\.tokenId) == [4, 5, 63]) - #expect(result.entries[0].value == result.entries[1].value) - #expect(result.entries[2].value < result.entries[0].value) - #expect(result.entries[0].alternatives.count == 2) - } - @Test( "Batched scoring matches scoring each request alone, across padding and chunks", arguments: [ @@ -370,11 +325,16 @@ struct BatchSchedulerTests { 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)) - #expect(batched.entries.count == prompt.count) } } @@ -383,66 +343,21 @@ struct BatchSchedulerTests { 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, reason) = try await drain(stream) + let (tokens, _) = try await drain(stream) #expect(tokens == [6, 7, 8, 9, 10, 11]) - #expect(reason == .maxTokens) #expect(scores[0].entries.count == 3) } @Test( "Scoring batches stop before a third of the positions would be padding", arguments: [ - ([7, 40, 130, 300], 4, 1), - ([130, 300], 4, 2), - ([64, 64, 64, 64, 64], 4, 4), - ([10, 200], 4, 1), - ([5], 4, 1), - ([], 4, 0), - ] as [([Int], Int, Int)]) - func scoringBatchCount(lengths: [Int], limit: Int, expected: Int) { - #expect(BatchScheduler.scoringBatchCount(lengths, limit: limit) == expected) - } - - @Test("A scoring request needs one target per token") - func scoringValidation() async { - let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) - await #expect(throws: InferenceRuntimeError.self) { - _ = try await scheduler.score([ScoreRequest(tokens: [1, 2], targets: [2], topK: 1)]) - } - } - - @Test("An empty prompt is rejected at submit") - func emptyPromptRejected() { - let scheduler = makeScheduler(CountUpDecodeStep(capacity: 2), maxConcurrent: 2) - #expect(throws: InferenceRuntimeError.self) { - _ = try scheduler.submit(req(prompt: [], maxTokens: 1)) - } - } - - @Test("A prefill failure fails only its own request") - func prefillFailureIsolated() async throws { - let scheduler = makeScheduler( - CountUpDecodeStep(capacity: 2, failPrefillFor: 50), maxConcurrent: 2) - let failing = try scheduler.submit(req(prompt: [50], maxTokens: 4)) - let healthy = try scheduler.submit(req(prompt: [5], maxTokens: 3)) - - await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(failing) } - let (tokens, reason) = try await drain(healthy) - #expect(tokens == [6, 7, 8]) - #expect(reason == .maxTokens) - } - - @Test("A failed shared decode also fails the queued requests") - func decodeFailureDrainsQueue() async throws { - let scheduler = makeScheduler( - CountUpDecodeStep(capacity: 1, failDecode: true), maxConcurrent: 1) - let streams = try [Int32(1), 2, 3].map { - try scheduler.submit(req(prompt: [$0], maxTokens: 4)) - } - // With one slot, requests 2 and 3 are still queued when request 1's decode throws. - for stream in streams { - await #expect(throws: CountUpDecodeStep.Boom.self) { _ = try await drain(stream) } - } + ([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.scoringBatchCount(lengths, limit: 4) == expected) } } From d161a4a49090069bb4000f82cc6c0a987c2ac7c9 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 14:38:26 -0700 Subject: [PATCH 06/10] Move batched scoring into BatchScheduler.Scorer BatchScheduler+Scorer.swift holds Scorer, which runs padded scoring batches and picks how many queued requests go into each batch, and ScoreRequest. BatchScheduler keeps the scoring queue and calls into its scorer. --- .../BatchScheduler+Scorer.swift | 111 ++++++++++++++++++ .../InferenceEngines/BatchScheduler.swift | 104 +--------------- .../BatchSchedulerTests.swift | 2 +- 3 files changed, 118 insertions(+), 99 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift new file mode 100644 index 00000000..30fd9e94 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift @@ -0,0 +1,111 @@ +// 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 + /// Maximum number of positions per scoring forward pass. Each pass returns at most + /// `rows * chunkLength * vocab` logits. + let chunkLength = 128 + + /// 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. For prompts of lengths 3 and 5 with a chunk length of 4, + /// the rows are + /// ``` + /// chunk 0: [a0 a1 a2 a2] [b0 b1 b2 b3] positions 0..<4 + /// chunk 1: [a2] [b4] position 4, row a is padding only + /// ``` + func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { + // Static graphs take exactly `capacity` rows; dynamic graphs (`Int.max`) take one per request. + let capacity = decodeStep.maxBatchCapacity + let rows = capacity == Int.max ? requests.count : capacity + 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 = min(chunkLength, longest - start) + 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.. [ScoringJob] in - let count = Self.scoringBatchCount(s.scoring.map(\.request.tokens.count), limit: maxConcurrent) + 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 scoreBatch(jobs.map(\.request)) + let results = try await scorer.scoreBatch(jobs.map(\.request)) for (job, result) in zip(jobs, results) { job.continuation.yield(result) job.continuation.finish() @@ -320,76 +319,6 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } } - /// 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 scoringBatchCount(_ 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. At most `decodeStep.maxBatchCapacity` - /// requests. - /// - /// Short rows repeat their last token. For prompts of lengths 3 and 5 with a chunk length of 4, - /// the rows are - /// ``` - /// chunk 0: [a0 a1 a2 a2] [b0 b1 b2 b3] positions 0..<4 - /// chunk 1: [a2] [b4] position 4, row a is padding only - /// ``` - private func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { - // Static graphs take exactly `capacity` rows; dynamic graphs (`Int.max`) take one per request. - let capacity = decodeStep.maxBatchCapacity - let rows = capacity == Int.max ? requests.count : capacity - 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 = min(scoringChunkLength, longest - start) - 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.. Date: Thu, 22 Oct 2026 17:31:40 -0700 Subject: [PATCH 07/10] Rename llm-server --concurrency to --concurrent-requests --- .../Tools/llm-server/LLMServerMain.swift | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/swift/Sources/Tools/llm-server/LLMServerMain.swift b/swift/Sources/Tools/llm-server/LLMServerMain.swift index 4a9db7a5..3a721b52 100644 --- a/swift/Sources/Tools/llm-server/LLMServerMain.swift +++ b/swift/Sources/Tools/llm-server/LLMServerMain.swift @@ -55,7 +55,7 @@ struct LLMServer: AsyncParsableCommand { @Flag( name: .customLong("replay-concurrent"), help: ArgumentHelp( - "Run all --replay requests concurrently instead of one at a time. Use with --concurrency N.") + "Run all --replay requests concurrently instead of one at a time. Use with --concurrent-requests N.") ) var replayConcurrent: Bool = false @@ -86,12 +86,12 @@ struct LLMServer: AsyncParsableCommand { var maxQueueDepth: Int = 16 @Option( - name: .customLong("concurrency"), + 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 concurrency: Int = 1 + var concurrentRequests: Int = 1 @Flag( name: .customLong("no-thinking"), @@ -118,8 +118,8 @@ struct LLMServer: AsyncParsableCommand { guard maxQueueDepth >= 0 else { throw ValidationError("--max-queue-depth must be >= 0 (got \(maxQueueDepth))") } - guard concurrency >= 1 else { - throw ValidationError("--concurrency must be >= 1 (got \(concurrency))") + guard concurrentRequests >= 1 else { + throw ValidationError("--concurrent-requests must be >= 1 (got \(concurrentRequests))") } do { _ = try ReasoningEffort.resolveDefault( @@ -236,22 +236,23 @@ struct LLMServer: AsyncParsableCommand { maxQueueDepth: maxQueueDepth ) - // The scheduler holds up to --concurrency plus --max-queue-depth requests, then rejects. + // The scheduler holds up to --concurrent-requests plus --max-queue-depth requests, then rejects. var batchScheduler: BatchScheduler? = nil - if concurrency > 1 { + if concurrentRequests > 1 { if let seq = engine as? CoreAISequentialEngine, seq.supportsBatching { let assetMax = bundle.maxBatchSize - guard concurrency <= assetMax else { + guard concurrentRequests <= assetMax else { throw ValidationError( - "--concurrency \(concurrency) exceeds the asset's max batch size \(assetMax). " - + "Re-export with --dynamic-batch-size \(concurrency) or higher, or lower --concurrency.") + "--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: concurrency, maxQueued: maxQueueDepth) - print("Batched serving: concurrency \(concurrency) (asset max batch \(assetMax))") + engine: seq, maxConcurrent: concurrentRequests, maxQueued: maxQueueDepth) + print("Batched serving: \(concurrentRequests) concurrent requests (asset max batch \(assetMax))") } else { print( - "⚠️ --concurrency \(concurrency) ignored: batched serving needs the coreai-sequential engine " + "⚠️ --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.") } From b4c8d2d43ea0e080e2231b6e8fcf230dcb9410fb Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 22 Oct 2026 18:41:41 -0700 Subject: [PATCH 08/10] Free a batch row when its request is cancelled or its stream is dropped A dropped stream, or setStopReason on the scheduler's output sequence (as the decoding strategy does on a stop-sequence match), now cancels the request. The scheduler retires a cancelled row before the next decode step and skips cancelled requests in the queue. --- .../InferenceEngines/BatchScheduler.swift | 46 +++++++++++++++++-- .../InferenceEngines/SchedulerSession.swift | 22 +++++---- .../BatchSchedulerTests.swift | 40 ++++++++++++++-- 3 files changed, 89 insertions(+), 19 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift index c60f5c92..57efd875 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -25,8 +25,11 @@ import Synchronization /// - 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 or `maxTokens`, and the slot takes the next -/// queued request while the other rows keep decoding. +/// - 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 { @@ -51,6 +54,15 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { /// 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 @@ -124,10 +136,14 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { 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) @@ -151,7 +167,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { /// 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. + /// 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 @@ -168,6 +184,10 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { 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 { @@ -212,6 +232,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { 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() @@ -247,11 +268,15 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { } else { active[slot]!.continuation.yield(.finished(active[slot]!.adapter.stopReason ?? .maxTokens)) active[slot]!.continuation.finish() - active[slot] = nil - state.withLock { $0.occupiedRows -= 1 } + 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, @@ -390,6 +415,8 @@ package struct BatchRequest: Sendable { /// 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], @@ -408,6 +435,15 @@ package struct BatchRequest: Sendable { } } +/// 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 { diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift index a9140f7c..5dc61985 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/SchedulerSession.swift @@ -42,13 +42,12 @@ package final class SchedulerSession: InferenceEngine, @unchecked Sendable { ) async throws -> GenerationSequence { let contextRoom = Swift.max(1, config.maxContextLength - input.count) let requested = inferenceOptions.maxTokens ?? contextRoom - let stream = try sink.submit( - BatchRequest( - promptTokens: input.map { Int32($0) }, - sampling: samplingConfiguration, - eosTokenIds: stopTokenIds, - maxTokens: Swift.max(1, Swift.min(requested, contextRoom)))) - return GenerationSequence(events: stream) + 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. @@ -69,17 +68,22 @@ 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. A reason set with `setStopReason(_:)` takes precedence. + /// 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) } + package func setStopReason(_ reason: StopReason) { + stopReasonStore.set(reason) + cancellation.cancel() + } package func makeAsyncIterator() -> Iterator { Iterator(upstream: events.makeAsyncIterator(), stopReasonStore: stopReasonStore) diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift index c0a8e8d5..dd116c64 100644 --- a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -44,8 +44,8 @@ private func drain( return (tokens, reason) } -/// Mock `BatchedDecodeStep` that predicts the last input token plus one. A request with prompt -/// `[..., k]` and `maxTokens: m` generates `k+1` through `k+m`. +/// 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. @@ -102,7 +102,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { await enter() defer { leave() } if let failPrefillFor, promptTokens.first == failPrefillFor { throw Boom() } - return oneHot((promptTokens.last ?? 0) + 1) + return oneHot(((promptTokens.last ?? 0) + 1) % 64) } func decodeBatch( @@ -112,7 +112,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { defer { leave() } record.withLock { $0.startPositions.append(startPositions) } if failDecode { throw Boom() } - return tokensPerRow.map { oneHot(($0.last ?? 0) + 1) } + return tokensPerRow.map { oneHot((($0.last ?? 0) + 1) % 64) } } func allPositionLogits( @@ -303,6 +303,36 @@ struct BatchSchedulerTests { } } + @Test("Dropped streams free their row, and a dropped queued request is never decoded") + func droppedStreamsFreeRows() async throws { + let mock = CountUpDecodeStep(capacity: 1, forwardDelay: .milliseconds(1)) + let scheduler = makeScheduler(mock, maxConcurrent: 1) + do { + var reading = try scheduler.submit(req(prompt: [1], maxTokens: 1000)).makeAsyncIterator() + _ = try await reading.next() + _ = try scheduler.submit(req(prompt: [2], maxTokens: 1000)) + } + let (tokens, _) = try await drain(try scheduler.submit(req(prompt: [5], maxTokens: 2))) + #expect(tokens == [6, 7]) + #expect(mock.forwards < 100) + } + + @Test("setStopReason on the output sequence ends the request within one decode step") + func setStopReasonCancels() async throws { + let mock = CountUpDecodeStep(capacity: 1, forwardDelay: .milliseconds(1)) + let scheduler = makeScheduler(mock, maxConcurrent: 1) + let request = req(prompt: [1], maxTokens: 1000) + let sequence = SchedulerSession.GenerationSequence( + events: try scheduler.submit(request), cancellation: request.cancellation) + var forwardsAtStop = 0 + for try await _ in sequence where forwardsAtStop == 0 { + forwardsAtStop = mock.forwards + sequence.setStopReason(.stopSequence("4")) + } + #expect(mock.forwards - forwardsAtStop <= 1) + #expect(sequence.stopReason == .stopSequence("4")) + } + // MARK: - Scoring /// Makes a top-2 score request for `tokens` whose target at each position is the next token, @@ -373,7 +403,7 @@ struct SchedulerSessionSequenceTests { if let finished { continuation.yield(.finished(finished)) } continuation.finish() - let sequence = SchedulerSession.GenerationSequence(events: stream) + 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]) From acb4c05be3cd089499abbf5061f802eccc46bc50 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Fri, 23 Oct 2026 14:31:05 -0700 Subject: [PATCH 09/10] Pad scoring passes to a few fixed shapes Scoring ran each pass at [requests, positions], so on a dynamic-shape graph every new prompt length or batch size was a new input shape, and a new shape can cost a compile or an allocation. The last chunk now pads to 16, 32, 64 or 128 positions, never past the context length, and a dynamic graph's rows round up to a power of two, at most maxConcurrent. Padded positions repeat the row's last token and padded rows are never read, so scores do not change. --- .../BatchScheduler+Scorer.swift | 40 +++++++++++--- .../InferenceEngines/BatchScheduler.swift | 6 +- .../BatchSchedulerTests.swift | 55 ++++++++++++++++++- 3 files changed, 90 insertions(+), 11 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift index 30fd9e94..f968bf0d 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler+Scorer.swift @@ -13,9 +13,35 @@ extension BatchScheduler { 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 @@ -40,23 +66,23 @@ extension BatchScheduler { /// 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. For prompts of lengths 3 and 5 with a chunk length of 4, - /// the rows are + /// 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] [b4] position 4, row a is padding only + /// chunk 1: [a2 a2] [b4 b4] positions 4..<6, only b4 is read /// ``` func scoreBatch(_ requests: [ScoreRequest]) async throws -> [LogProbabilities] { - // Static graphs take exactly `capacity` rows; dynamic graphs (`Int.max`) take one per request. - let capacity = decodeStep.maxBatchCapacity - let rows = capacity == Int.max ? requests.count : capacity + 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 = min(chunkLength, longest - start) + 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 diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift index 57efd875..2fd6c4cb 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchScheduler.swift @@ -90,12 +90,12 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { "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) + maxConcurrent: maxConcurrent, maxQueued: maxQueued, maxPositions: engine.config.maxContextLength) } init( decodeStep: any BatchedDecodeStep, session: GenerationSessionState, maxConcurrent: Int, - maxQueued: Int? = nil + maxQueued: Int? = nil, maxPositions: Int = .max ) { precondition(maxConcurrent >= 1, "maxConcurrent must be >= 1") precondition( @@ -106,7 +106,7 @@ package final class BatchScheduler: RequestSink, @unchecked Sendable { self.session = session self.maxConcurrent = maxConcurrent self.maxQueued = maxQueued - self.scorer = Scorer(decodeStep: decodeStep, maxConcurrent: maxConcurrent) + self.scorer = Scorer(decodeStep: decodeStep, maxConcurrent: maxConcurrent, maxPositions: maxPositions) } // MARK: - RequestSink diff --git a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift index dd116c64..e91ebee3 100644 --- a/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift +++ b/swift/Tests/LanguageModelsTests/BatchSchedulerTests.swift @@ -65,6 +65,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { var maxInFlight = 0 var forwards = 0 var scoringRows: [Int] = [] + var scoringPositions: [Int] = [] } private let record = Mutex(Record()) @@ -72,6 +73,7 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { 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 } } @@ -120,7 +122,10 @@ private final class CountUpDecodeStep: BatchedDecodeStep, @unchecked Sendable { ) async throws -> [LogitsScalarType] { await enter() defer { leave() } - record.withLock { $0.scoringRows.append(tokensPerRow.count) } + record.withLock { + $0.scoringRows.append(tokensPerRow.count) + $0.scoringPositions.append(tokensPerRow.first?.count ?? 0) + } return countUpAllPositions(tokensPerRow) } @@ -368,6 +373,54 @@ struct BatchSchedulerTests { } } + @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) From 4f0f89d5167ee6edcc8ede3ba263da5cae9f2c52 Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Fri, 23 Oct 2026 14:31:05 -0700 Subject: [PATCH 10/10] Include the batch size in llm-server's system_fingerprint Batched decoding sums in a different order than one request at a time, so the same seed can give slightly different logits. The fingerprint descriptor now carries batch=, so a batched server and a single-request server report different fingerprints. --- swift/Sources/Tools/llm-server/ServerState.swift | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 26eb8498..c789bcf3 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -127,7 +127,9 @@ 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. @@ -156,7 +158,8 @@ final class ServerState: @unchecked Sendable { 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 } @@ -236,10 +239,13 @@ final class ServerState: @unchecked Sendable { /// 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")",