From 78f1d537aa080e2e187f93a934b2d9c14e8216c2 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Fri, 31 Jul 2026 10:04:49 -0700 Subject: [PATCH 1/2] Add state handlers for hybrid model support (>2 states) Both engines previously hardcoded exactly 2 states (key_cache + value_cache). Models with sliding window caches, recurrent states, or other persistent states beyond the KV pair would fail at load. StateHandler protocol + implementations: - GrowingNDArrayState: dynamic KV cache with capacity doubling - FixedNDArrayState: static states (sliding caches, conv/recurrent) - FixedMTLBufferState: MTLBuffer states for pipelined engine - StateHandlerFactory: shape-based classification with metadata override Engine changes: - Sequential: uses StateHandlerFactory, supports 2-4 states - Pipelined: additionalStates?.insertAll(into:) at all encode sites - Both: hasNonTruncatableStates guard forces full reset for hybrid models - LanguageConfig: optional "states" field for explicit classification --- .../Bundle/LanguageConfig.swift | 10 +- .../Handlers/StateHandler+MTLBuffer.swift | 67 +++++ .../Handlers/StateHandler+NDArray.swift | 195 +++++++++++++++ .../Handlers/StateHandler.swift | 52 ++++ .../Handlers/StateHandlerFactory.swift | 177 +++++++++++++ .../CoreAIPipelinedEngine.swift | 76 +++++- .../CoreAISequentialEngine.swift | 233 ++++++++---------- .../StateHandlerTests.swift | 90 +++++++ 8 files changed, 760 insertions(+), 140 deletions(-) create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.swift create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift create mode 100644 swift/Tests/LanguageModelsTests/StateHandlerTests.swift diff --git a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift b/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift index 09cc8996..784fa90f 100644 --- a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift +++ b/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift @@ -25,13 +25,18 @@ public struct LanguageConfig: Codable, Sendable, Equatable { /// Vision-specific configuration. Nil for text-only language models. public let vision: VisionConfig? + /// Explicit state classification. Nil = use shape-based heuristic. + /// Keys are state names from the model descriptor, values are StateKind. + public let states: [String: StateKind]? + public init( tokenizer: String, vocabSize: Int, maxContextLength: Int, embeddedTokenizer: Bool = true, functionMap: FunctionMap? = nil, - vision: VisionConfig? = nil + vision: VisionConfig? = nil, + states: [String: StateKind]? = nil ) { self.tokenizer = tokenizer self.vocabSize = vocabSize @@ -39,6 +44,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable { self.embeddedTokenizer = embeddedTokenizer self.functionMap = functionMap self.vision = vision + self.states = states } enum CodingKeys: String, CodingKey { @@ -48,6 +54,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable { case embeddedTokenizer = "embedded_tokenizer" case functionMap = "function_map" case vision + case states } public init(from decoder: Swift.Decoder) throws { @@ -58,6 +65,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable { self.embeddedTokenizer = try c.decodeIfPresent(Bool.self, forKey: .embeddedTokenizer) ?? true self.functionMap = try c.decodeIfPresent(FunctionMap.self, forKey: .functionMap) self.vision = try c.decodeIfPresent(VisionConfig.self, forKey: .vision) + self.states = try c.decodeIfPresent([String: StateKind].self, forKey: .states) } // MARK: - Additional Stop Tokens diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.swift new file mode 100644 index 00000000..c0c94c13 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.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 + +import CoreAI +import CoreAIShared +import Metal + +/// Fixed-size MTLBuffer state for non-truncatable persistent states. +/// Allocated once at init, zero-initialized, never grows. +public struct FixedMTLBufferState { + public let stateNames: [String] + public var stateCount: Int { bindings.count } + + private var bindings: + [(name: String, buffer: MTLBuffer, scalarType: NDArray.ScalarType, shape: [Int], strides: [Int])] + + public init( + states: [(name: String, descriptor: NDArrayDescriptor)], + device: MTLDevice + ) throws { + var bindings: [(String, MTLBuffer, NDArray.ScalarType, [Int], [Int])] = [] + for (name, desc) in states { + guard !desc.shape.contains(where: { $0 < 0 }) else { + throw InferenceRuntimeError.invalidOutputType( + "FixedMTLBufferState '\(name)' has dynamic shape \(desc.shape)") + } + let resolved = desc.resolvingDynamicDimensions(desc.shape) + let strides = resolved.preferredStrides + let byteCount = resolved.minimumByteCount + guard let buffer = device.makeBuffer(length: max(byteCount, 64), options: .storageModeShared) else { + throw InferenceRuntimeError.bufferAllocationFailed("\(name) (\(byteCount) bytes)") + } + memset(buffer.contents(), 0, buffer.length) + bindings.append((name, buffer, desc.scalarType, desc.shape, strides)) + } + self.bindings = bindings + self.stateNames = states.map(\.name) + } + + /// Access a state binding by index for AsyncMutableValue construction. + /// The engine builds AsyncMutableValue views from these in the same scope as encode(). + /// Note: MTLBuffer is a reference type — the returned buffer is shared, not copied. + public subscript(stateIndex index: Int) -> ( + name: String, buffer: MTLBuffer, scalarType: NDArray.ScalarType, shape: [Int], strides: [Int] + ) { + get { bindings[index] } + } + + /// Insert all states into AsyncMutableViews for pipelined encoding. + public func insertAll(into views: inout InferenceFunction.AsyncMutableViews) { + for (name, buffer, scalarType, shape, strides) in bindings { + var value = unsafe InferenceFunction.AsyncMutableValue( + unsafeBuffer: buffer, byteOffset: 0, + scalarType: scalarType, shape: shape, strides: strides) + views.insert(&value, for: name) + } + } + + /// Zero all state buffers. Caller must ensure no in-flight GPU work references these. + public mutating func reset() { + for (_, buffer, _, _, _) in bindings { + memset(buffer.contents(), 0, buffer.length) + } + } +} diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift new file mode 100644 index 00000000..752c0f0a --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift @@ -0,0 +1,195 @@ +// 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 CoreAI +import CoreAIShared +import Darwin + +// MARK: - Fixed NDArray State + +/// Fixed-size state for non-truncatable persistent states. +/// Allocated at full size on init, zero-initialized. No capacity management needed. +public struct FixedNDArrayState: SyncStateHandler { + public let stateNames: [String] + public let supportsTruncation: Bool = false + public let currentCapacity: Int = .max + public var stateCount: Int { arrays.count } + + private var arrays: [(name: String, array: NDArray)] + + public init(states: [(name: String, descriptor: NDArrayDescriptor)]) { + var arrays: [(String, NDArray)] = [] + for (name, desc) in states { + var array = NDArray(descriptor: desc) + zeroFillNDArray(&array) + arrays.append((name, array)) + } + self.arrays = arrays + self.stateNames = states.map(\.name) + } + + public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool { + false + } + + public subscript(stateIndex index: Int) -> (name: String, array: NDArray) { + get { arrays[index] } + set { arrays[index] = newValue } + } + + public mutating func reset() { + for i in arrays.indices { + zeroFillNDArray(&arrays[i].array) + } + } + + public mutating func truncate(to tokenCount: Int) { + preconditionFailure("truncate(to:) called on non-truncatable FixedNDArrayState") + } +} + +// MARK: - Growing NDArray State + +/// Dynamically-growing KV cache state. Starts small and doubles capacity +/// when more context is needed. +public struct GrowingNDArrayState: SyncStateHandler { + public let stateNames: [String] + public let supportsTruncation: Bool = true + public private(set) var currentCapacity: Int + public var stateCount: Int { arrays.count } + + private var arrays: [(name: String, array: NDArray)] + private let descriptors: [NDArrayDescriptor] + private let maxCapacity: Int + private let sequenceDimIndex: Int + + public init( + states: [(name: String, descriptor: NDArrayDescriptor)], + initialCapacity: Int, + maxCapacity: Int + ) { + self.maxCapacity = maxCapacity + self.descriptors = states.map(\.descriptor) + self.stateNames = states.map(\.name) + + let firstDesc = states[0].descriptor + self.sequenceDimIndex = firstDesc.shape.firstIndex(where: { $0 < 0 }) ?? max(0, firstDesc.shape.count - 2) + + let capacity = min(initialCapacity, maxCapacity) + self.currentCapacity = capacity + + var arrays: [(String, NDArray)] = [] + for (name, desc) in states { + let resolved = desc.resolvingDynamicDimensions( + desc.shape.map { $0 < 0 ? capacity : $0 }) + arrays.append((name, NDArray(descriptor: resolved))) + } + self.arrays = arrays + } + + public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool { + guard contextLength > currentCapacity else { return false } + guard contextLength <= maxCapacity else { + throw InferenceRuntimeError.invalidState( + "Context length \(contextLength) exceeds maximum \(maxCapacity)") + } + + var newCapacity = max(currentCapacity, 1) + while newCapacity < contextLength { + newCapacity = min(newCapacity * 2, maxCapacity) + } + + for i in arrays.indices { + let desc = descriptors[i] + let newShape = desc.shape.map { $0 < 0 ? newCapacity : $0 } + let resolvedDesc = desc.resolvingDynamicDimensions(newShape) + var newArray = NDArray(descriptor: resolvedDesc) + // Force backing allocation before copy + _ = newArray.mutableRawView() + + copyCache(from: arrays[i].array, to: &newArray, sequenceDim: sequenceDimIndex) + arrays[i].array = newArray + } + + currentCapacity = newCapacity + return true + } + + public subscript(stateIndex index: Int) -> (name: String, array: NDArray) { + get { arrays[index] } + set { arrays[index] = newValue } + } + + public mutating func reset() { + for i in arrays.indices { + zeroFillNDArray(&arrays[i].array) + } + } + + public mutating func truncate(to tokenCount: Int) { + // KV cache truncation is a no-op on the backing storage. + // The causal mask hides positions beyond processedTokenCount. + } + + // MARK: - Private + + private func copyCache(from source: NDArray, to destination: inout NDArray, sequenceDim: Int) { + let srcShape = source.shape + let dstShape = destination.shape + guard let headDim = srcShape.last else { return } + + let numBlocks = srcShape[...size) + } + case .float32: + var view = array.mutableView(as: Float.self) + view.withUnsafeMutablePointer { ptr, _, _ in + memset(ptr, 0, count * MemoryLayout.size) + } + default: + preconditionFailure("Unsupported scalar type for state: \(array.scalarType)") + } +} diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift new file mode 100644 index 00000000..f1f2b715 --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift @@ -0,0 +1,52 @@ +// 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 CoreAI + +/// Persistent model state that the engine carries across inference steps. +/// +/// Each handler manages one or more named state tensors (e.g., key_cache + value_cache, +/// or a single convolution_state). The engine owns an array of handlers and delegates +/// all state lifecycle to them — allocation, growth, and reset. +/// +/// ## MutableViews Lifetime Constraint +/// +/// `InferenceFunction.MutableViews.insert` creates a lifetime dependency on each `inout` +/// variable. This means state binding CANNOT be abstracted into a method that takes +/// `inout MutableViews` — the inserts must happen in the same scope as `function.run()`. +/// Handlers therefore expose their arrays directly via subscript for the engine to bind. +public protocol SyncStateHandler { + /// Names of the states managed by this handler. + var stateNames: [String] { get } + + /// Number of state arrays managed. + var stateCount: Int { get } + + /// Current capacity in the sequence/context dimension. + /// For fixed-size states this equals max capacity. + var currentCapacity: Int { get } + + /// Whether this state supports in-place truncation (cursor rewind). + /// KV cache: true (causal mask hides positions beyond the cursor). + /// Recurrent/conv: false (no independent token axis). + var supportsTruncation: Bool { get } + + /// Ensure the state can accommodate `contextLength` tokens. + /// Returns true if reallocation occurred. + mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool + + /// Access a state array by index for binding into MutableViews. + /// The engine calls this to get name + array pairs for insert. + subscript(stateIndex index: Int) -> (name: String, array: NDArray) { get set } + + /// Full reset — zero all backing storage, rewind to position 0. + mutating func reset() + + /// Truncate to a given token position. + /// Only valid when `supportsTruncation == true`. For KV cache, this is a no-op + /// on the backing storage (causal mask handles visibility); the engine just + /// updates its processedTokenCount. + mutating func truncate(to tokenCount: Int) +} diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift new file mode 100644 index 00000000..faf7055f --- /dev/null +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift @@ -0,0 +1,177 @@ +// 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 CoreAI +import CoreAIShared + +/// Classification of a model state's lifecycle behavior. +public enum StateKind: String, Codable, Sendable { + /// KV cache — grows dynamically with context, supports truncation (causal mask). + case kvCache = "kv_cache" + /// Sliding window cache — fixed size, supports truncation (causal mask). + case slidingCache = "sliding_cache" + /// Fixed state (conv, recurrent) — fixed size, does NOT support truncation. + case fixed +} + +/// Result of state handler creation. +struct SyncStateHandlerSet { + /// Growing states (KV caches with dynamic sequence dimension). + var kvCache: any SyncStateHandler + /// Fixed states (sliding caches + recurrent/conv). Nil for transformer-only models. + var additionalStates: FixedNDArrayState? + /// Whether any state is non-truncatable (triggers full-reset-only mode). + var hasNonTruncatableStates: Bool +} + +/// Creates state handlers from a model's function descriptor. +/// +/// Classification priority: +/// 1. Explicit metadata (`"states"` field in metadata.json) — preferred +/// 2. Shape-based heuristic — dynamic dim → kvCache, static + "cache" in name → slidingCache, else → fixed +/// 3. Legacy (2 states) — both kvCache, no warning +enum StateHandlerFactory { + + /// Classify states using metadata or heuristic fallback. + static func classifyStates( + descriptor: InferenceFunctionDescriptor, + stateKinds: [String: StateKind]? = nil, + verbose: Bool = false + ) -> [(name: String, kind: StateKind)] { + let names = descriptor.stateNames + + if let kinds = stateKinds { + // Explicit metadata — validate and use + return names.map { name in + let kind = kinds[name] ?? inferKind(name: name, descriptor: descriptor) + return (name, kind) + } + } + + // Heuristic fallback + if names.count == 2 { + // Legacy: 2 states = both KV cache (no warning) + return names.map { ($0, .kvCache) } + } + + // N states: classify by shape + name + let classified = names.map { name -> (String, StateKind) in + (name, inferKind(name: name, descriptor: descriptor)) + } + + if verbose { + CLILogger.log("State classification (heuristic):", component: "StateHandlerFactory") + for (name, kind) in classified { + guard case .ndArray(let desc) = descriptor.stateDescriptor(of: name) else { continue } + let shapeStr = desc.shape.map { $0 < 0 ? "?" : "\($0)" }.joined(separator: "×") + let growth = desc.shape.contains(where: { $0 < 0 }) ? "GROWING" : "FIXED" + CLILogger.log(" \(name): \(growth) \(kind.rawValue) (\(shapeStr))", + component: "StateHandlerFactory") + } + CLILogger.log(" Add \"states\" to metadata.json for explicit control.", + component: "StateHandlerFactory") + } + if !verbose && names.count > 2 { + CLILogger.log( + "StateHandlerFactory: \(names.count) states classified by heuristic. " + + "Add \"states\" to metadata.json for explicit control.", + component: "StateHandlerFactory") + } + + return classified + } + + /// Infer state kind from shape and name. + private static func inferKind(name: String, descriptor: InferenceFunctionDescriptor) -> StateKind { + guard case .ndArray(let desc) = descriptor.stateDescriptor(of: name) else { + return .fixed + } + let hasDynamicDim = desc.shape.contains(where: { $0 < 0 }) + if hasDynamicDim { + return .kvCache + } + let lower = name.lowercased() + if lower.contains("cache") || lower.contains("kv") { + return .slidingCache + } + return .fixed + } + + /// Create sync state handlers from classified states. + static func createSyncHandlers( + descriptor: InferenceFunctionDescriptor, + maxContextLength: Int, + stateKinds: [String: StateKind]? = nil, + options: EngineOptions = EngineOptions(), + verbose: Bool = false + ) throws -> SyncStateHandlerSet { + guard !descriptor.stateNames.isEmpty else { + throw InferenceRuntimeError.invalidOutputType( + "Expected states but found none") + } + + let classified = classifyStates( + descriptor: descriptor, stateKinds: stateKinds, verbose: verbose) + + // Separate into growing (kvCache) and fixed (slidingCache + fixed) + var growingPairs: [(name: String, descriptor: NDArrayDescriptor)] = [] + var fixedPairs: [(name: String, descriptor: NDArrayDescriptor)] = [] + var hasNonTruncatable = false + + for (name, kind) in classified { + guard case .ndArray(let desc) = descriptor.stateDescriptor(of: name) else { + throw InferenceRuntimeError.invalidOutputType( + "Cannot get state descriptor for '\(name)'") + } + + switch kind { + case .kvCache: + growingPairs.append((name, desc)) + case .slidingCache: + fixedPairs.append((name, desc)) + case .fixed: + fixedPairs.append((name, desc)) + hasNonTruncatable = true + } + } + + // Build growing handler (KV caches) + let kvCache: any SyncStateHandler + if !growingPairs.isEmpty { + if options.kvCacheStrategy == .fixedSize { + let resolved = growingPairs.map { (name, desc) -> (name: String, descriptor: NDArrayDescriptor) in + let resolvedDesc = desc.resolvingDynamicDimensions( + desc.shape.map { $0 < 0 ? maxContextLength : $0 }) + return (name, resolvedDesc) + } + kvCache = FixedNDArrayState(states: resolved) + } else { + let initial = min(256, maxContextLength) + kvCache = GrowingNDArrayState( + states: growingPairs, + initialCapacity: initial, + maxCapacity: maxContextLength + ) + } + } else { + // All states are fixed (e.g., all sliding caches at fixed-size mode) + // Use the fixed pairs as KV cache too + kvCache = FixedNDArrayState(states: fixedPairs) + fixedPairs = [] + } + + // Build fixed handler (sliding caches + recurrent/conv) + var additionalStates: FixedNDArrayState? = nil + if !fixedPairs.isEmpty { + additionalStates = FixedNDArrayState(states: fixedPairs) + } + + return SyncStateHandlerSet( + kvCache: kvCache, + additionalStates: additionalStates, + hasNonTruncatableStates: hasNonTruncatable + ) + } +} diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 2c287844..75471e16 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -178,6 +178,16 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { self.history.clear() resolvedNewTokens = input[...] commonPrefix = 0 + } else if self.engine.hasNonTruncatableStates { + // Hybrid model: recurrent state can't be partially rewound. + // Full reset and replay the entire prompt. + if commonPrefix < self.engine.processedTokenCount { + await self.engine.computeStream.currentWorkCompleted() + self.engine.reset() + self.history.clear() + resolvedNewTokens = input[...] + commonPrefix = 0 + } } else if commonPrefix < self.engine.processedTokenCount { // Pure extension — partial rewind (buffer phase preserved) await self.engine.computeStream.currentWorkCompleted() @@ -281,6 +291,11 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { // Partial reset: wait for generation to finish naturally, then rewind counter. // Do NOT cancel — cancelling corrupts the pipeline's double-buffer state. // The KV cache is valid up to processedTokenCount after natural completion. + if engine.hasNonTruncatableStates { + throw InferenceRuntimeError.invalidState( + "Partial reset is not supported for hybrid models with recurrent state. " + + "Use reset(to: 0) and replay the prefix.") + } drain() await engine.computeStream.currentWorkCompleted() guard tryAcquireEngine() else { return } @@ -431,6 +446,11 @@ private struct EngineImpl: ~Copyable { // KV cache — reuses CoreAIKVCache protocol from KVCache+CoreAI.swift var kvCache: any CoreAIKVCache + // Linear attention state bindings for hybrid models (nil for pure transformer models). + // States 0/1 are KV cache; additional states handled by handler. + var additionalStates: FixedMTLBufferState? + var hasNonTruncatableStates: Bool + // Logits — reuses GrowingLogitsBuffer from TensorStorage+CoreAI.swift var logits: GrowingLogitsBuffer @@ -474,16 +494,36 @@ private struct EngineImpl: ~Copyable { throw InferenceRuntimeError.invalidOutputType( "Expected at least 1 output, got \(descriptor.outputNames.count)") } - guard descriptor.stateNames.count == 2 else { + guard descriptor.stateNames.count >= 2 && descriptor.stateNames.count <= 4 else { throw InferenceRuntimeError.invalidOutputType( - "Expected 2 states (KV cache), got \(descriptor.stateNames.count): \(descriptor.stateNames)") + "Expected 2–4 states, got \(descriptor.stateNames.count): \(descriptor.stateNames)" + ) } + // Classify states using the shared factory logic + let classified = StateHandlerFactory.classifyStates( + descriptor: descriptor, stateKinds: nil, verbose: descriptor.stateNames.count > 2) + + // Find the growing KV pair (first two states with .kvCache kind) + let growingNames = classified.filter { $0.kind == .kvCache }.map(\.name) + guard growingNames.count >= 2 else { + throw InferenceRuntimeError.invalidOutputType( + "Expected at least 2 growing KV cache states, found \(growingNames.count) " + + "in: \(classified.map { "\($0.name)=\($0.kind.rawValue)" })") + } + let keyCacheName = growingNames[0] + let valueCacheName = growingNames[1] + + // Fixed states: everything that isn't the primary growing KV pair + let fixedNames = classified + .filter { $0.kind == .slidingCache || $0.kind == .fixed } + .map(\.name) + // Additional growing states beyond the primary pair + let extraGrowingNames = Array(growingNames.dropFirst(2)) + // Extract names let inputIdsName = descriptor.inputNames[0] let positionIdsName = descriptor.inputNames[1] - let keyCacheName = descriptor.stateNames[0] - let valueCacheName = descriptor.stateNames[1] let logitsOutputName = descriptor.outputNames[0] // Extract state descriptors for KV cache shape/type @@ -568,6 +608,27 @@ private struct EngineImpl: ~Copyable { let resolvedSize = options.resolvedKVCacheSize(maxContextLength: config.maxContextLength) CLILogger.log("Created \(options.kvCacheStrategy) KV cache with size \(resolvedSize, default: "nil")") + // Allocate fixed-size buffers for additional persistent states (sliding caches, hybrid states). + var additionalStatesLocal: FixedMTLBufferState? = nil + let allFixedNames = fixedNames + extraGrowingNames // extra growing get resolved to max size + if !allFixedNames.isEmpty { + var extraStates: [(name: String, descriptor: NDArrayDescriptor)] = [] + for name in allFixedNames { + guard case .ndArray(let desc) = descriptor.stateDescriptor(of: name) else { + throw InferenceRuntimeError.invalidOutputType( + "Cannot get descriptor for persistent state '\(name)'") + } + // Resolve dynamic dims to max for any extra growing states + let resolved = desc.shape.contains(where: { $0 < 0 }) + ? desc.resolvingDynamicDimensions(desc.shape.map { $0 < 0 ? config.maxContextLength : $0 }) + : desc + extraStates.append((name, resolved)) + } + additionalStatesLocal = try FixedMTLBufferState(states: extraStates, device: device) + CLILogger.log( + "Pipelined additional states: \(allFixedNames.joined(separator: ", "))") + } + // Create growing logits buffer (reuses TensorStorage+CoreAI.swift) let logitsRef = try GrowingLogitsBuffer( device: device, @@ -613,6 +674,8 @@ private struct EngineImpl: ~Copyable { self.decodeOutputBuffers = decodeOutBuffers self.decodeLogitsBuffers = decodeLogBufs self.kvCache = kvCacheLocal + self.additionalStates = additionalStatesLocal + self.hasNonTruncatableStates = classified.contains(where: { $0.kind == .fixed }) self.logits = logitsRef self.cachedSampler = nil self.cachedSamplerTemperature = nil @@ -764,6 +827,7 @@ private struct EngineImpl: ~Copyable { var asyncStates = InferenceFunction.AsyncMutableViews() asyncStates.insert(&keyState, for: keyCacheName) asyncStates.insert(&valState, for: valueCacheName) + additionalStates?.insertAll(into: &asyncStates) // Build Output as AsyncMutableValue (logits) // Decode uses per-step rotating buffer; prefill uses the shared growing buffer. @@ -1075,6 +1139,7 @@ private struct EngineImpl: ~Copyable { var asyncStates = InferenceFunction.AsyncMutableViews() asyncStates.insert(&keyState, for: keyCacheName) asyncStates.insert(&valState, for: valueCacheName) + additionalStates?.insertAll(into: &asyncStates) let logitsShape = [1, queryLength, vocabSize] let logitsStrides = try resolvedStrides(descriptor: logitsBaseDesc, shape: logitsShape) @@ -1102,6 +1167,8 @@ private struct EngineImpl: ~Copyable { step = 0 cachedSampler = nil cachedSamplerTemperature = nil + // Zero SSM states so the next conversation starts from a clean slate. + additionalStates?.reset() span.end() } @@ -1179,6 +1246,7 @@ private struct EngineImpl: ~Copyable { var asyncStates = InferenceFunction.AsyncMutableViews() asyncStates.insert(&keyState, for: keyCacheName) asyncStates.insert(&valState, for: valueCacheName) + additionalStates?.insertAll(into: &asyncStates) let lShape = [1, shape, vocabSize] let lStrides = try resolvedStrides(descriptor: logitsBaseDesc, shape: lShape) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 1a768dc3..5c7e03fb 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -26,7 +26,7 @@ enum PrefillStrategy { /// Expects a `.aimodel` with: /// - **2 inputs**: `input_ids` (Int32), `position_ids` (Int32) /// - **1 output**: `logits` (LogitsScalarType) -/// - **2 states**: `keyCache`, `valueCache` — persistent across steps, updated in-place +/// - **2–4 states**: KV cache pair + optional persistent states (hybrid models), updated in-place /// /// KV cache NDArrays start small (256 tokens) and grow dynamically with 2× expansion. /// Passed as `states` on every forward pass; the model graph updates them in-place. @@ -44,18 +44,19 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable // I/O names from descriptor private let inputIdsName: String private let positionIdsName: String - private let keyCacheName: String - private let valueCacheName: String private let logitsName: String + // State management — handlers own allocation, growth, and reset + private var kvCache: any SyncStateHandler + private var additionalStates: FixedNDArrayState? + private var hasNonTruncatableStates: Bool + // Descriptors for dynamic shape resolution private let inputIdsDescriptor: NDArrayDescriptor private let positionIdsDescriptor: NDArrayDescriptor private let logitsDescriptor: NDArrayDescriptor // Persistent state — reused across steps - private var keyCache: NDArray - private var valueCache: NDArray private var logitsArray: NDArray // Pre-allocated input_ids reused across decode steps. Only reallocated when // batch size changes (i.e., once when transitioning from prefill to decode). @@ -64,9 +65,6 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable private var inputIdsArray: NDArray private var cachedInputBatchSize: Int private var cachedLogitsBatchSize: Int - private var currentKVCapacity: Int - private let keyCacheDescriptor: NDArrayDescriptor - private let valueCacheDescriptor: NDArrayDescriptor // Track processed tokens for incremental inference public private(set) var processedTokenCount: Int = 0 @@ -109,7 +107,8 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable } self.functionDescriptor = descriptor - // Validate model architecture: 2 inputs, 1+ output, 2 states + // Validate model architecture: 2 inputs, 1+ output, at least KV cache pair. + // Hybrid models may declare additional persistent fixed-shape states. guard descriptor.inputNames.count == 2 else { throw InferenceRuntimeError.invalidInputType( "Expected 2 inputs, got \(descriptor.inputNames.count): \(descriptor.inputNames)") @@ -118,17 +117,15 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable throw InferenceRuntimeError.invalidOutputType( "Expected at least 1 output, got \(descriptor.outputNames.count): \(descriptor.outputNames)") } - guard descriptor.stateNames.count == 2 else { + guard descriptor.stateNames.count >= 2 && descriptor.stateNames.count <= 4 else { throw InferenceRuntimeError.invalidOutputType( - "Expected 2 states (KV cache), got \(descriptor.stateNames.count): " + "Expected 2–4 states (KV cache + optional persistent states), got \(descriptor.stateNames.count): " + "states=\(descriptor.stateNames), outputs=\(descriptor.outputNames)") } // Extract names self.inputIdsName = descriptor.inputNames[0] self.positionIdsName = descriptor.inputNames[1] - self.keyCacheName = descriptor.stateNames[0] - self.valueCacheName = descriptor.stateNames[1] self.logitsName = descriptor.outputNames[0] // Extract and validate input descriptors @@ -152,37 +149,23 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable } self.logitsDescriptor = logitsDesc - // Extract KV cache state descriptors - guard case .ndArray(let keyCacheDesc) = descriptor.stateDescriptor(of: keyCacheName), - case .ndArray(let valueCacheDesc) = descriptor.stateDescriptor(of: valueCacheName) - else { - throw InferenceRuntimeError.invalidOutputType("Cannot get KV cache state descriptors") - } - - // Store unresolved descriptors for dynamic reallocation - self.keyCacheDescriptor = keyCacheDesc - self.valueCacheDescriptor = valueCacheDesc - - let isDynamic = keyCacheDesc.shape.contains(where: { $0 < 0 }) - - // Allocate KV cache at initial size (grow on demand unless fixedSize requested) - let initialCapacity: Int - if options.kvCacheStrategy == .fixedSize || !isDynamic { - initialCapacity = config.maxContextLength - } else { - initialCapacity = min(256, config.maxContextLength) - } - self.currentKVCapacity = initialCapacity - let resolvedKeyDesc = keyCacheDesc.resolvingDynamicDimensions( - keyCacheDesc.shape.map { $0 < 0 ? initialCapacity : $0 }) - let resolvedValueDesc = valueCacheDesc.resolvingDynamicDimensions( - valueCacheDesc.shape.map { $0 < 0 ? initialCapacity : $0 }) - self.keyCache = NDArray(descriptor: resolvedKeyDesc) - self.valueCache = NDArray(descriptor: resolvedValueDesc) + // Create state handlers from descriptor + let stateHandlers = try StateHandlerFactory.createSyncHandlers( + descriptor: descriptor, + maxContextLength: config.maxContextLength, + options: options + ) + self.kvCache = stateHandlers.kvCache + self.additionalStates = stateHandlers.additionalStates + self.hasNonTruncatableStates = stateHandlers.hasNonTruncatableStates CLILogger.log( - "KV cache: dynamic=\(isDynamic), initial=\(initialCapacity), key=\(keyCacheDesc.shape) → \(resolvedKeyDesc.shape)" + "KV cache: capacity=\(kvCache.currentCapacity), states=\(kvCache.stateNames)" ) + if let additional = additionalStates { + CLILogger.log( + "Additional persistent states: \(additional.stateNames.joined(separator: ", "))") + } // Allocate initial logits (1 token — will be reallocated per batch) let initLogitsDesc = logitsDesc.resolvingDynamicDimensions([1, 1, config.vocabSize]) @@ -241,7 +224,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable throw InferenceRuntimeError.invalidState("Cannot process empty token batch") } - try ensureKVCapacity(forContextLength: processedTokenCount + batchSize) + _ = try kvCache.ensureCapacity(forContextLength: processedTokenCount + batchSize) let batchSignpost = InstrumentsProfiler.beginCustomInterval( name: "CoreAIClean Batch", @@ -273,21 +256,61 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable cachedLogitsBatchSize = batchSize } - // Build states (KV cache — persistent, inout) - var states = InferenceFunction.MutableViews() - states.insert(&keyCache, for: keyCacheName) - states.insert(&valueCache, for: valueCacheName) - // Build output backings (logits — written in-place) var outputViews = InferenceFunction.MutableViews() outputViews.insert(&logitsArray, for: logitsName) - // Execute - _ = try await function.run( - inputs: [inputIdsName: inputIdsArray, positionIdsName: positionIds], - states: consume states, - outputViews: consume outputViews - ) + // Build states and execute. MutableViews lifetime requires explicit local + // bindings — see StateHandler.swift for why this can't be abstracted. + var kv0 = kvCache[stateIndex: 0] + var kv1 = kvCache[stateIndex: 1] + if var handler = additionalStates { + switch handler.stateCount { + case 1: + var s0 = handler[stateIndex: 0] + var states = InferenceFunction.MutableViews() + states.insert(&kv0.array, for: kv0.name) + states.insert(&kv1.array, for: kv1.name) + states.insert(&s0.array, for: s0.name) + _ = try await function.run( + inputs: [inputIdsName: inputIdsArray, positionIdsName: positionIds], + states: consume states, + outputViews: consume outputViews + ) + handler[stateIndex: 0] = s0 + case 2: + var s0 = handler[stateIndex: 0] + var s1 = handler[stateIndex: 1] + var states = InferenceFunction.MutableViews() + states.insert(&kv0.array, for: kv0.name) + states.insert(&kv1.array, for: kv1.name) + states.insert(&s0.array, for: s0.name) + states.insert(&s1.array, for: s1.name) + _ = try await function.run( + inputs: [inputIdsName: inputIdsArray, positionIdsName: positionIds], + states: consume states, + outputViews: consume outputViews + ) + handler[stateIndex: 0] = s0 + handler[stateIndex: 1] = s1 + default: + preconditionFailure( + "Unsupported additional state count \(handler.stateCount). " + + "Add a new case to handle \(handler.stateCount) states.") + } + additionalStates = handler + } else { + var states = InferenceFunction.MutableViews() + states.insert(&kv0.array, for: kv0.name) + states.insert(&kv1.array, for: kv1.name) + _ = try await function.run( + inputs: [inputIdsName: inputIdsArray, positionIdsName: positionIds], + states: consume states, + outputViews: consume outputViews + ) + } + kvCache[stateIndex: 0] = kv0 + kvCache[stateIndex: 1] = kv1 // Read logits from NDArray let totalLogits = batchSize * config.vocabSize @@ -356,17 +379,30 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable } // Implicit prefix caching: resolve input against history. + // For hybrid models with recurrent states, we must full-reset on any + // rewind because recurrent state summarizes the whole prefix and cannot + // be truncated by moving a KV cursor. Future: checkpoint/restore. if history.count > 0 { let (commonPrefix, _) = history.resolve(input: input) - if commonPrefix < input.count && commonPrefix < history.count { + if hasNonTruncatableStates { + // Hybrid model: recurrent state can't be partially rewound. + // Full reset and replay the entire prompt. + if commonPrefix < history.count || processedTokenCount >= input.count { + internalReset(to: 0) + } + lastPrefixHitCount = 0 + } else if commonPrefix < input.count && commonPrefix < history.count { // Divergence: input differs from history. Full reset needed. internalReset(to: 0) + lastPrefixHitCount = commonPrefix } else if processedTokenCount >= input.count { // Pure extension: all input tokens match history. Rewind for seeding. let resetTo = Swift.max(0, commonPrefix - 1) internalReset(to: resetTo) + lastPrefixHitCount = commonPrefix + } else { + lastPrefixHitCount = commonPrefix } - lastPrefixHitCount = commonPrefix } let token = GenerationToken() @@ -405,6 +441,11 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable precondition( tokenIndex >= 0 && tokenIndex <= processedTokenCount, "reset(to: \(tokenIndex)) out of range [0, \(processedTokenCount)]") + if tokenIndex != 0 && hasNonTruncatableStates { + throw InferenceRuntimeError.invalidState( + "Partial reset is not supported for hybrid models with recurrent state. " + + "Use reset(to: 0) and replay the prefix.") + } _activeToken.withLock { $0?.cancel() $0 = nil @@ -419,8 +460,8 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable if tokenIndex == 0 { processedTokenCount = 0 history.clear() - zeroFill(&keyCache) - zeroFill(&valueCache) + kvCache.reset() + additionalStates?.reset() } else { processedTokenCount = tokenIndex history.truncate(to: tokenIndex) @@ -434,84 +475,6 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable cleanupSpan.end() } - // MARK: - KV Cache (dynamic growth) - - private func ensureKVCapacity(forContextLength needed: Int) throws { - guard needed > currentKVCapacity else { return } - guard needed <= config.maxContextLength else { - throw InferenceRuntimeError.invalidState( - "Context length \(needed) exceeds maximum \(config.maxContextLength)") - } - - var newCapacity = currentKVCapacity - while newCapacity < needed { newCapacity *= 2 } - newCapacity = min(newCapacity, config.maxContextLength) - - let resolvedKeyDesc = keyCacheDescriptor.resolvingDynamicDimensions( - keyCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 }) - let resolvedValueDesc = valueCacheDescriptor.resolvingDynamicDimensions( - valueCacheDescriptor.shape.map { $0 < 0 ? newCapacity : $0 }) - - var newKeyCache = NDArray(descriptor: resolvedKeyDesc) - var newValueCache = NDArray(descriptor: resolvedValueDesc) - _ = newKeyCache.mutableRawView() - _ = newValueCache.mutableRawView() - - try Self.copyCache(from: keyCache, to: &newKeyCache) - try Self.copyCache(from: valueCache, to: &newValueCache) - - CLILogger.log("KV cache grew: \(currentKVCapacity) → \(newCapacity)") - keyCache = newKeyCache - valueCache = newValueCache - currentKVCapacity = newCapacity - } - - private static func copyCache(from source: NDArray, to destination: inout NDArray) throws { - let srcShape = source.shape - let dstShape = destination.shape - guard let headDim = srcShape.last else { - throw InferenceRuntimeError.invalidState("KV cache has empty shape — cannot copy") - } - let seqDim = KVCacheFactory.detectSequenceDim(shape: srcShape) - - // Number of independent blocks before the sequence dimension (L * B * H or B * H) - let numBlocks = srcShape[.. LogitsScalarType` closure is invoked per element (no inlining), - // which made zeroing the KV cache (~14.7M elements for a 32K-context - // Qwen3) take ~6 seconds per `reset()`. Direct loop keeps this in - // the few-ms range even unoptimized; under -O it lowers to memset. - view.withUnsafeMutablePointer { ptr, _, _ in - for i in 0.. Date: Fri, 31 Jul 2026 10:30:35 -0700 Subject: [PATCH 2/2] Fix Swift format --- .../Handlers/StateHandler+NDArray.swift | 2 +- .../Handlers/StateHandlerFactory.swift | 11 ++++++----- .../InferenceEngines/CoreAIPipelinedEngine.swift | 6 ++++-- .../InferenceEngines/CoreAISequentialEngine.swift | 1 - .../Tests/LanguageModelsTests/StateHandlerTests.swift | 4 ++-- 5 files changed, 13 insertions(+), 11 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift index 752c0f0a..beafc846 100644 --- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift @@ -176,7 +176,7 @@ public struct GrowingNDArrayState: SyncStateHandler { // MARK: - Shared Utilities /// Zero-initialize an NDArray, dispatching on scalar type. -fileprivate func zeroFillNDArray(_ array: inout NDArray) { +private func zeroFillNDArray(_ array: inout NDArray) { let count = array.shape.reduce(1, *) switch array.scalarType { case .float16, .bfloat16: diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift index faf7055f..10a3d7bd 100644 --- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift +++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift @@ -33,7 +33,6 @@ struct SyncStateHandlerSet { /// 2. Shape-based heuristic — dynamic dim → kvCache, static + "cache" in name → slidingCache, else → fixed /// 3. Legacy (2 states) — both kvCache, no warning enum StateHandlerFactory { - /// Classify states using metadata or heuristic fallback. static func classifyStates( descriptor: InferenceFunctionDescriptor, @@ -67,11 +66,13 @@ enum StateHandlerFactory { guard case .ndArray(let desc) = descriptor.stateDescriptor(of: name) else { continue } let shapeStr = desc.shape.map { $0 < 0 ? "?" : "\($0)" }.joined(separator: "×") let growth = desc.shape.contains(where: { $0 < 0 }) ? "GROWING" : "FIXED" - CLILogger.log(" \(name): \(growth) \(kind.rawValue) (\(shapeStr))", - component: "StateHandlerFactory") + CLILogger.log( + " \(name): \(growth) \(kind.rawValue) (\(shapeStr))", + component: "StateHandlerFactory") } - CLILogger.log(" Add \"states\" to metadata.json for explicit control.", - component: "StateHandlerFactory") + CLILogger.log( + " Add \"states\" to metadata.json for explicit control.", + component: "StateHandlerFactory") } if !verbose && names.count > 2 { CLILogger.log( diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index 75471e16..79b90a5a 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -515,7 +515,8 @@ private struct EngineImpl: ~Copyable { let valueCacheName = growingNames[1] // Fixed states: everything that isn't the primary growing KV pair - let fixedNames = classified + let fixedNames = + classified .filter { $0.kind == .slidingCache || $0.kind == .fixed } .map(\.name) // Additional growing states beyond the primary pair @@ -619,7 +620,8 @@ private struct EngineImpl: ~Copyable { "Cannot get descriptor for persistent state '\(name)'") } // Resolve dynamic dims to max for any extra growing states - let resolved = desc.shape.contains(where: { $0 < 0 }) + let resolved = + desc.shape.contains(where: { $0 < 0 }) ? desc.resolvingDynamicDimensions(desc.shape.map { $0 < 0 ? config.maxContextLength : $0 }) : desc extraStates.append((name, resolved)) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift index 5c7e03fb..af8d2e1a 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift @@ -474,7 +474,6 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable CLILogger.log("CoreAI clean engine cleanup complete") cleanupSpan.end() } - } extension CoreAISequentialEngine { diff --git a/swift/Tests/LanguageModelsTests/StateHandlerTests.swift b/swift/Tests/LanguageModelsTests/StateHandlerTests.swift index 842a0d45..5e53472e 100644 --- a/swift/Tests/LanguageModelsTests/StateHandlerTests.swift +++ b/swift/Tests/LanguageModelsTests/StateHandlerTests.swift @@ -60,8 +60,8 @@ struct StateKindTests { @Test("StateKind decodes from JSON") func decodable() throws { let json = """ - {"key": "kv_cache", "sliding": "sliding_cache", "fix": "fixed"} - """ + {"key": "kv_cache", "sliding": "sliding_cache", "fix": "fixed"} + """ struct Wrapper: Decodable { let key: StateKind let sliding: StateKind