Skip to content

Commit f469e0a

Browse files
stikvesstikves
committed
Add bind(into:) for zero-copy state binding with loop-based insert
State handlers are now classes (AnyObject) that own their NDArrays at refcount 1. bind(into:) inserts all states into MutableViews in a loop using _overrideLifetime to express disjoint element access (pattern from Ben Levine). No COW, no switch on state count, no write-back. Engine call sites: runWithStates (sequential) and encodeWithStates (pipelined) replace 3x duplicated 60-line switch blocks.
1 parent 4236629 commit f469e0a

8 files changed

Lines changed: 301 additions & 290 deletions

File tree

‎Package.swift‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ let package = Package(
5757
swiftSettings: [
5858
.define("CXGRAMMAR_IMPORT"),
5959
.enableUpcomingFeature("MemberImportVisibility"),
60+
.enableExperimentalFeature("Lifetimes"),
6061
],
6162
linkerSettings: [
6263
.linkedLibrary("c++")
@@ -225,6 +226,9 @@ let package = Package(
225226
resources: [
226227
.copy("Resources/MinimalTokenizer")
227228
],
229+
swiftSettings: [
230+
.enableExperimentalFeature("Lifetimes"),
231+
],
228232
linkerSettings: [
229233
.linkedLibrary("c++")
230234
]

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler+Encode.swift‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import CoreAI
7+
import Metal
8+
9+
/// Encode an inference step with KV cache states, optional additional MTLBuffer
10+
/// states, and logits output — all in one scope to satisfy ~Escapable constraints.
11+
func encodeWithStates(
12+
function: InferenceFunction,
13+
inputs: [String: InferenceFunction.AsyncValue],
14+
keyState: inout InferenceFunction.AsyncMutableValue,
15+
keyCacheName: String,
16+
valState: inout InferenceFunction.AsyncMutableValue,
17+
valueCacheName: String,
18+
additionalStates: FixedMTLBufferState?,
19+
logitsBuffer: MTLBuffer,
20+
logitsName: String,
21+
logitsShape: [Int],
22+
logitsStrides: [Int],
23+
computeStream: ComputeStream
24+
) throws {
25+
switch additionalStates?.stateCount ?? 0 {
26+
case 1:
27+
let extra0 = additionalStates![stateIndex: 0]
28+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
29+
unsafeBuffer: extra0.buffer, byteOffset: 0,
30+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
31+
var asyncStates = InferenceFunction.AsyncMutableViews()
32+
asyncStates.insert(&keyState, for: keyCacheName)
33+
asyncStates.insert(&valState, for: valueCacheName)
34+
asyncStates.insert(&extraState0, for: extra0.name)
35+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
36+
unsafeBuffer: logitsBuffer, byteOffset: 0,
37+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
38+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
39+
asyncOutputs.insert(&logitsOutput, for: logitsName)
40+
let _ = try function.encode(
41+
inputs: inputs, states: consume asyncStates,
42+
outputViews: consume asyncOutputs, to: computeStream)
43+
case 2:
44+
let extra0 = additionalStates![stateIndex: 0]
45+
let extra1 = additionalStates![stateIndex: 1]
46+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
47+
unsafeBuffer: extra0.buffer, byteOffset: 0,
48+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
49+
var extraState1 = unsafe InferenceFunction.AsyncMutableValue(
50+
unsafeBuffer: extra1.buffer, byteOffset: 0,
51+
scalarType: extra1.scalarType, shape: extra1.shape, strides: extra1.strides)
52+
var asyncStates = InferenceFunction.AsyncMutableViews()
53+
asyncStates.insert(&keyState, for: keyCacheName)
54+
asyncStates.insert(&valState, for: valueCacheName)
55+
asyncStates.insert(&extraState0, for: extra0.name)
56+
asyncStates.insert(&extraState1, for: extra1.name)
57+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
58+
unsafeBuffer: logitsBuffer, byteOffset: 0,
59+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
60+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
61+
asyncOutputs.insert(&logitsOutput, for: logitsName)
62+
let _ = try function.encode(
63+
inputs: inputs, states: consume asyncStates,
64+
outputViews: consume asyncOutputs, to: computeStream)
65+
default:
66+
if (additionalStates?.stateCount ?? 0) > 0 {
67+
preconditionFailure(
68+
"encodeWithStates: unsupported additional state count "
69+
+ "\(additionalStates!.stateCount)")
70+
}
71+
var asyncStates = InferenceFunction.AsyncMutableViews()
72+
asyncStates.insert(&keyState, for: keyCacheName)
73+
asyncStates.insert(&valState, for: valueCacheName)
74+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
75+
unsafeBuffer: logitsBuffer, byteOffset: 0,
76+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
77+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
78+
asyncOutputs.insert(&logitsOutput, for: logitsName)
79+
let _ = try function.encode(
80+
inputs: inputs, states: consume asyncStates,
81+
outputViews: consume asyncOutputs, to: computeStream)
82+
}
83+
}

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift‎

Lines changed: 41 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -11,56 +11,63 @@ import Darwin
1111

1212
/// Fixed-size state for non-truncatable persistent states.
1313
/// Allocated at full size on init, zero-initialized. No capacity management needed.
14-
public struct FixedNDArrayState: SyncStateHandler {
14+
public final class FixedNDArrayState: SyncStateHandler {
1515
public let stateNames: [String]
1616
public let supportsTruncation: Bool = false
1717
public let currentCapacity: Int = .max
1818
public var stateCount: Int { arrays.count }
1919

20-
private var arrays: [(name: String, array: NDArray)]
20+
private var arrays: [String: NDArray]
2121

2222
public init(states: [(name: String, descriptor: NDArrayDescriptor)]) {
23-
var arrays: [(String, NDArray)] = []
23+
var arrays: [String: NDArray] = [:]
2424
for (name, desc) in states {
2525
var array = NDArray(descriptor: desc)
2626
zeroFillNDArray(&array)
27-
arrays.append((name, array))
27+
arrays[name] = array
2828
}
2929
self.arrays = arrays
3030
self.stateNames = states.map(\.name)
3131
}
3232

33-
public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
33+
public func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
3434
false
3535
}
3636

3737
public subscript(stateIndex index: Int) -> (name: String, array: NDArray) {
38-
get { arrays[index] }
39-
set { arrays[index] = newValue }
38+
get { (stateNames[index], arrays[stateNames[index]]!) }
39+
set { arrays[stateNames[index]] = newValue.array }
4040
}
4141

42-
public mutating func reset() {
43-
for i in arrays.indices {
44-
zeroFillNDArray(&arrays[i].array)
42+
@_lifetime(views: borrow self)
43+
public func bind(into views: inout InferenceFunction.MutableViews) {
44+
for name in stateNames {
45+
let view = _overrideLifetime(arrays[name]!.mutableRawView(), borrowing: Void())
46+
views.insert(view, for: name)
4547
}
4648
}
4749

48-
public mutating func truncate(to tokenCount: Int) {
50+
public func reset() {
51+
for name in stateNames {
52+
zeroFillNDArray(&arrays[name]!)
53+
}
54+
}
55+
56+
public func truncate(to tokenCount: Int) {
4957
preconditionFailure("truncate(to:) called on non-truncatable FixedNDArrayState")
5058
}
5159
}
5260

5361
// MARK: - Growing NDArray State
5462

55-
/// Dynamically-growing KV cache state. Starts small and doubles capacity
56-
/// when more context is needed.
57-
public struct GrowingNDArrayState: SyncStateHandler {
63+
/// Dynamically-growing KV cache state. Starts small and doubles capacity.
64+
public final class GrowingNDArrayState: SyncStateHandler {
5865
public let stateNames: [String]
5966
public let supportsTruncation: Bool = true
6067
public private(set) var currentCapacity: Int
6168
public var stateCount: Int { arrays.count }
6269

63-
private var arrays: [(name: String, array: NDArray)]
70+
private var arrays: [String: NDArray]
6471
private let descriptors: [NDArrayDescriptor]
6572
private let maxCapacity: Int
6673
private let sequenceDimIndex: Int
@@ -80,16 +87,16 @@ public struct GrowingNDArrayState: SyncStateHandler {
8087
let capacity = min(initialCapacity, maxCapacity)
8188
self.currentCapacity = capacity
8289

83-
var arrays: [(String, NDArray)] = []
90+
var arrays: [String: NDArray] = [:]
8491
for (name, desc) in states {
8592
let resolved = desc.resolvingDynamicDimensions(
8693
desc.shape.map { $0 < 0 ? capacity : $0 })
87-
arrays.append((name, NDArray(descriptor: resolved)))
94+
arrays[name] = NDArray(descriptor: resolved)
8895
}
8996
self.arrays = arrays
9097
}
9198

92-
public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
99+
public func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
93100
guard contextLength > currentCapacity else { return false }
94101
guard contextLength <= maxCapacity else {
95102
throw InferenceRuntimeError.invalidState(
@@ -101,38 +108,41 @@ public struct GrowingNDArrayState: SyncStateHandler {
101108
newCapacity = min(newCapacity * 2, maxCapacity)
102109
}
103110

104-
for i in arrays.indices {
111+
for (i, name) in stateNames.enumerated() {
105112
let desc = descriptors[i]
106113
let newShape = desc.shape.map { $0 < 0 ? newCapacity : $0 }
107114
let resolvedDesc = desc.resolvingDynamicDimensions(newShape)
108115
var newArray = NDArray(descriptor: resolvedDesc)
109-
// Force backing allocation before copy
110116
_ = newArray.mutableRawView()
111-
112-
copyCache(from: arrays[i].array, to: &newArray, sequenceDim: sequenceDimIndex)
113-
arrays[i].array = newArray
117+
copyCache(from: arrays[name]!, to: &newArray, sequenceDim: sequenceDimIndex)
118+
arrays[name] = newArray
114119
}
115120

116121
currentCapacity = newCapacity
117122
return true
118123
}
119124

120125
public subscript(stateIndex index: Int) -> (name: String, array: NDArray) {
121-
get { arrays[index] }
122-
set { arrays[index] = newValue }
126+
get { (stateNames[index], arrays[stateNames[index]]!) }
127+
set { arrays[stateNames[index]] = newValue.array }
123128
}
124129

125-
public mutating func reset() {
126-
for i in arrays.indices {
127-
zeroFillNDArray(&arrays[i].array)
130+
@_lifetime(views: borrow self)
131+
public func bind(into views: inout InferenceFunction.MutableViews) {
132+
for name in stateNames {
133+
let view = _overrideLifetime(arrays[name]!.mutableRawView(), borrowing: Void())
134+
views.insert(view, for: name)
128135
}
129136
}
130137

131-
public mutating func truncate(to tokenCount: Int) {
132-
// KV cache truncation is a no-op on the backing storage.
133-
// The causal mask hides positions beyond processedTokenCount.
138+
public func reset() {
139+
for name in stateNames {
140+
zeroFillNDArray(&arrays[name]!)
141+
}
134142
}
135143

144+
public func truncate(to tokenCount: Int) {}
145+
136146
// MARK: - Private
137147

138148
private func copyCache(from source: NDArray, to destination: inout NDArray, sequenceDim: Int) {
@@ -175,7 +185,6 @@ public struct GrowingNDArrayState: SyncStateHandler {
175185

176186
// MARK: - Shared Utilities
177187

178-
/// Zero-initialize an NDArray, dispatching on scalar type.
179188
func zeroFillNDArray(_ array: inout NDArray) {
180189
let count = array.shape.reduce(1, *)
181190
switch array.scalarType {

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler+Run.swift‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import CoreAI
7+
8+
/// Run an inference step with combined primary + secondary states and output.
9+
/// Zero-copy: bind(into:) uses reference-backed storage and _overrideLifetime.
10+
func runWithStates(
11+
function: InferenceFunction,
12+
inputs: [String: NDArray],
13+
primary: any SyncStateHandler,
14+
secondary: FixedNDArrayState?,
15+
outputArray: inout NDArray,
16+
outputName: String
17+
) async throws {
18+
var states = InferenceFunction.MutableViews()
19+
primary.bind(into: &states)
20+
secondary?.bind(into: &states)
21+
22+
var outputViews = InferenceFunction.MutableViews()
23+
outputViews.insert(&outputArray, for: outputName)
24+
25+
_ = try await function.run(
26+
inputs: inputs,
27+
states: _unsafeEscapeMutableViews(consume states),
28+
outputViews: consume outputViews)
29+
}

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift‎

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,10 @@ import CoreAI
1111
/// or a single convolution_state). The engine owns an array of handlers and delegates
1212
/// all state lifecycle to them — allocation, growth, and reset.
1313
///
14-
/// ## MutableViews Lifetime Constraint
15-
///
16-
/// `InferenceFunction.MutableViews.insert` creates a lifetime dependency on each `inout`
17-
/// variable. This means state binding CANNOT be abstracted into a method that takes
18-
/// `inout MutableViews` — the inserts must happen in the same scope as `function.run()`.
19-
/// Handlers therefore expose their arrays directly via subscript for the engine to bind.
20-
public protocol SyncStateHandler {
14+
/// Handlers use reference-backed storage (NDArrayRef) so that mutableRawView()
15+
/// never triggers COW. The `bind(into:)` method inserts all states into
16+
/// MutableViews in a loop using `_overrideLifetime` for disjoint element access.
17+
public protocol SyncStateHandler: AnyObject {
2118
/// Names of the states managed by this handler.
2219
var stateNames: [String] { get }
2320

@@ -35,18 +32,46 @@ public protocol SyncStateHandler {
3532

3633
/// Ensure the state can accommodate `contextLength` tokens.
3734
/// Returns true if reallocation occurred.
38-
mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool
35+
func ensureCapacity(forContextLength contextLength: Int) throws -> Bool
3936

40-
/// Access a state array by index for binding into MutableViews.
41-
/// The engine calls this to get name + array pairs for insert.
37+
/// Access a state array by name+index (value copy — use bind(into:) on hot paths).
4238
subscript(stateIndex index: Int) -> (name: String, array: NDArray) { get set }
4339

40+
/// Insert all managed states into `views`. Zero-copy: uses reference-backed
41+
/// storage internally, so mutableRawView() never triggers COW.
42+
@_lifetime(views: borrow self)
43+
func bind(into views: inout InferenceFunction.MutableViews)
44+
4445
/// Full reset — zero all backing storage, rewind to position 0.
45-
mutating func reset()
46+
func reset()
4647

4748
/// Truncate to a given token position.
48-
/// Only valid when `supportsTruncation == true`. For KV cache, this is a no-op
49-
/// on the backing storage (causal mask handles visibility); the engine just
50-
/// updates its processedTokenCount.
51-
mutating func truncate(to tokenCount: Int)
49+
func truncate(to tokenCount: Int)
50+
}
51+
52+
// MARK: - Lifetime Helpers
53+
54+
/// Override the lifetime of a ~Escapable value to borrow from a different source.
55+
/// Used inside bind(into:) to detach per-element raw views from the array they
56+
/// came from, allowing multiple inserts in a loop without exclusivity violations.
57+
@unsafe
58+
@_unsafeNonescapableResult
59+
@_alwaysEmitIntoClient
60+
@_transparent
61+
@_lifetime(borrow source)
62+
func _overrideLifetime<T: ~Copyable & ~Escapable, U: ~Copyable & ~Escapable>(
63+
_ dependent: consuming T, borrowing source: borrowing U
64+
) -> T {
65+
dependent
66+
}
67+
68+
/// Detach lifetime dependencies from MutableViews so it can cross scope
69+
/// boundaries (closures, await). Caller must ensure inserted arrays remain valid.
70+
@inline(__always)
71+
@_unsafeNonescapableResult
72+
@_lifetime(immortal)
73+
func _unsafeEscapeMutableViews(
74+
_ views: consuming InferenceFunction.MutableViews
75+
) -> InferenceFunction.MutableViews {
76+
views
5277
}

0 commit comments

Comments
 (0)