Add state handlers for hybrid model support (>2 states)

swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift

@@ -25,20 +25,26 @@ 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
         self.maxContextLength = maxContextLength
         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

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)
+        }
+    }
+}

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[..<sequenceDim].reduce(1, *)
+        let oldSeqLen = srcShape[sequenceDim]
+        let copyElements = oldSeqLen * headDim
+        let srcBlockStride = srcShape[sequenceDim...].reduce(1, *)
+        let dstBlockStride = dstShape[sequenceDim...].reduce(1, *)
+
+        switch source.scalarType {
+        case .float16, .bfloat16:
+            source.view(as: Float16.self).withUnsafePointer { srcPtr, _, _ in
+                var dstView = destination.mutableView(as: Float16.self)
+                dstView.withUnsafeMutablePointer { dstPtr, _, _ in
+                    for block in 0..<numBlocks {
+                        dstPtr.advanced(by: block * dstBlockStride).update(
+                            from: srcPtr.advanced(by: block * srcBlockStride), count: copyElements)
+                    }
+                }
+            }
+        case .float32:
+            source.view(as: Float.self).withUnsafePointer { srcPtr, _, _ in
+                var dstView = destination.mutableView(as: Float.self)
+                dstView.withUnsafeMutablePointer { dstPtr, _, _ in
+                    for block in 0..<numBlocks {
+                        dstPtr.advanced(by: block * dstBlockStride).update(
+                            from: srcPtr.advanced(by: block * srcBlockStride), count: copyElements)
+                    }
+                }
+            }
+        default:
+            preconditionFailure("Unsupported scalar type for state copy: \(source.scalarType)")
+        }
+    }
+}
+
+// MARK: - Shared Utilities
+
+/// Zero-initialize an NDArray, dispatching on scalar type.
+private func zeroFillNDArray(_ array: inout NDArray) {
+    let count = array.shape.reduce(1, *)
+    switch array.scalarType {
+    case .float16, .bfloat16:
+        var view = array.mutableView(as: Float16.self)
+        view.withUnsafeMutablePointer { ptr, _, _ in
+            memset(ptr, 0, count * MemoryLayout<Float16>.size)
+        }
+    case .float32:
+        var view = array.mutableView(as: Float.self)
+        view.withUnsafeMutablePointer { ptr, _, _ in
+            memset(ptr, 0, count * MemoryLayout<Float>.size)
+        }
+    default:
+        preconditionFailure("Unsupported scalar type for state: \(array.scalarType)")
+    }
+}

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)
+}