Skip to content

Commit a2dc47c

Browse files
sukru tikvesstikves
sukru tikves
authored and
stikves
committed
Add SyncInputHandler protocol and InputContext for engine extensibility
Introduces a composable input preparation layer for inference engines: - SyncInputHandler protocol: prepares named NDArray inputs each step - InputContext: carries tokens, position, batch size, sliding window - InputContext.dynamic() for sequential/pipelined engines - InputContext.static() for static-shape (ANE) engines - TokenInputHandler: standard token ID input with batch size caching - CompositeInputHandler: wraps a base handler with extra inputs (RoPE, PLE) - InputCoverage.verify(): fail-fast check at engine init - PipelinedTokenInputHandler: MTLBuffer-based for pipelined engine Also hardens NDArray helpers with stride-aware fill/read: - fillNDArray(count:using:) now checks contiguity and falls back to stride-aware indexing for 4D+ tensors with GPU alignment padding - readNDArray checks contiguity similarly - TokenInputHandler.prepare() validates non-empty batch
1 parent 57a0a63 commit a2dc47c

6 files changed

Lines changed: 402 additions & 27 deletions

File tree

‎swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Pipelined.swift‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
// Asynchronous token input handler for the Pipelined engine (MTLBuffer-based).
2+
//
3+
// Copyright 2026 Apple Inc.
4+
//
5+
// Use of this source code is governed by a BSD-3-clause license that can
6+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
7+
8+
import CoreAI
9+
import CoreAIShared
10+
import Metal
11+
12+
/// Input handler for the pipelined engine. Owns the token and position MTLBuffers,
13+
/// handles buffer rotation, and the prefill/decode source split.
14+
///
15+
/// `decodeOutputBuffers` is shared with the engine (GPU sampler writes next token
16+
/// there; this handler reads the previous step's token during decode).
17+
struct PipelinedTokenInputHandler {
18+
let inputIdsName: String
19+
let positionIdsName: String
20+
let inputIdsDescriptor: NDArrayDescriptor
21+
let positionIdsDescriptor: NDArrayDescriptor
22+
23+
let inputTokensBuffer: MTLBuffer
24+
let cachePositionBuffers: [MTLBuffer]
25+
let decodeOutputBuffers: [MTLBuffer]
26+
let pipelineDepth: Int
27+
28+
/// Build async input values for one encode step.
29+
///
30+
/// - Prefill (`!tokens.isEmpty`): writes tokens to `inputTokensBuffer` at their
31+
/// natural position (disjoint from prior chunks still in-flight on GPU).
32+
/// - Decode (`tokens.isEmpty`): reads from previous step's `decodeOutputBuffer`.
33+
func prepare(
34+
tokens: some Collection<Int32>,
35+
processedTokenCount: Int,
36+
step: Int
37+
) throws -> [String: InferenceFunction.AsyncValue] {
38+
let queryLength = tokens.isEmpty ? 1 : tokens.count
39+
40+
// Write tokens (prefill only — decode reads from GPU sampler output)
41+
if !tokens.isEmpty {
42+
let dst = inputTokensBuffer.contents()
43+
.assumingMemoryBound(to: Int32.self)
44+
.advanced(by: processedTokenCount)
45+
let tokenArray = Array(tokens)
46+
tokenArray.withUnsafeBufferPointer { src in
47+
memcpy(dst, src.baseAddress!, tokenArray.count * MemoryLayout<Int32>.size)
48+
}
49+
}
50+
51+
// Token input
52+
let tokenShape = [1, queryLength]
53+
let tokenStrides = try resolvedStrides(descriptor: inputIdsDescriptor, shape: tokenShape)
54+
let tokenValue: InferenceFunction.AsyncValue
55+
if tokens.isEmpty {
56+
tokenValue = unsafe InferenceFunction.AsyncValue(
57+
unsafeBuffer: decodeOutputBuffers[(step + pipelineDepth - 1) % pipelineDepth],
58+
byteOffset: 0, scalarType: .int32, shape: tokenShape, strides: tokenStrides)
59+
} else {
60+
tokenValue = unsafe InferenceFunction.AsyncValue(
61+
unsafeBuffer: inputTokensBuffer,
62+
byteOffset: processedTokenCount * MemoryLayout<Int32>.size,
63+
scalarType: .int32, shape: tokenShape, strides: tokenStrides)
64+
}
65+
66+
// Position input (rotating buffer)
67+
let posLength = processedTokenCount + queryLength
68+
let posShape = [1, posLength]
69+
let posStrides = try resolvedStrides(descriptor: positionIdsDescriptor, shape: posShape)
70+
let posValue = unsafe InferenceFunction.AsyncValue(
71+
unsafeBuffer: cachePositionBuffers[step % pipelineDepth],
72+
byteOffset: 0, scalarType: .int32, shape: posShape, strides: posStrides)
73+
74+
return [
75+
inputIdsName: tokenValue,
76+
positionIdsName: posValue,
77+
]
78+
}
79+
80+
/// Prepare inputs for warmup (always prefill mode, writes dummy tokens).
81+
func prepareWarmup(
82+
shape: Int,
83+
processedTokenCount: Int,
84+
step: Int
85+
) throws -> [String: InferenceFunction.AsyncValue] {
86+
let ptr = inputTokensBuffer.contents().assumingMemoryBound(to: Int32.self)
87+
for i in 0..<shape { ptr[i] = 1 }
88+
89+
let tokenShape = [1, shape]
90+
let tokenStrides = try resolvedStrides(descriptor: inputIdsDescriptor, shape: tokenShape)
91+
let tokenValue = unsafe InferenceFunction.AsyncValue(
92+
unsafeBuffer: inputTokensBuffer, byteOffset: 0,
93+
scalarType: .int32, shape: tokenShape, strides: tokenStrides)
94+
95+
let posLength = processedTokenCount + shape
96+
let posShape = [1, posLength]
97+
let posStrides = try resolvedStrides(descriptor: positionIdsDescriptor, shape: posShape)
98+
let posValue = unsafe InferenceFunction.AsyncValue(
99+
unsafeBuffer: cachePositionBuffers[step % pipelineDepth],
100+
byteOffset: 0, scalarType: .int32, shape: posShape, strides: posStrides)
101+
102+
return [
103+
inputIdsName: tokenValue,
104+
positionIdsName: posValue,
105+
]
106+
}
107+
}

‎swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
// Synchronous token input handler for Sequential and StaticShape engines.
2+
//
3+
// Copyright 2026 Apple Inc.
4+
//
5+
// Use of this source code is governed by a BSD-3-clause license that can
6+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
7+
8+
import CoreAI
9+
import CoreAIShared
10+
11+
/// Standard input handler for text LLMs: `input_ids` (Int32) + `position_ids` (Int32).
12+
///
13+
/// Pre-allocates the `input_ids` NDArray and reuses it when batch size is unchanged.
14+
public struct TokenInputHandler: SyncInputHandler {
15+
public let inputNames: [String]
16+
17+
private let inputIdsName: String
18+
private let positionIdsName: String
19+
private let inputIdsDescriptor: NDArrayDescriptor
20+
private let positionIdsDescriptor: NDArrayDescriptor
21+
22+
private var inputIdsArray: NDArray
23+
private var cachedBatchSize: Int
24+
25+
public init(
26+
inputIdsName: String,
27+
positionIdsName: String,
28+
inputIdsDescriptor: NDArrayDescriptor,
29+
positionIdsDescriptor: NDArrayDescriptor
30+
) {
31+
self.inputIdsName = inputIdsName
32+
self.positionIdsName = positionIdsName
33+
self.inputIdsDescriptor = inputIdsDescriptor
34+
self.positionIdsDescriptor = positionIdsDescriptor
35+
self.inputNames = [inputIdsName, positionIdsName]
36+
37+
let initDesc = inputIdsDescriptor.resolvingDynamicDimensions([1, 1])
38+
self.inputIdsArray = NDArray(descriptor: initDesc)
39+
self.cachedBatchSize = 1
40+
}
41+
42+
public mutating func prepare(_ context: InputContext) async throws -> [String: NDArray] {
43+
let tokens = context.tokens
44+
let batchSize = tokens.count
45+
precondition(batchSize > 0, "TokenInputHandler: empty token batch")
46+
47+
if cachedBatchSize != batchSize {
48+
let resolved = inputIdsDescriptor.resolvingDynamicDimensions([1, batchSize])
49+
inputIdsArray = NDArray(descriptor: resolved)
50+
cachedBatchSize = batchSize
51+
}
52+
fillNDArray(&inputIdsArray, as: Int32.self, with: tokens)
53+
54+
let totalPositions = context.processedTokenCount + batchSize
55+
let resolvedPosDesc = positionIdsDescriptor.resolvingDynamicDimensions([1, totalPositions])
56+
var positionIds = NDArray(descriptor: resolvedPosDesc)
57+
fillNDArray(&positionIds, as: Int32.self, count: totalPositions) { Int32($0) }
58+
59+
return [
60+
inputIdsName: inputIdsArray,
61+
positionIdsName: positionIds,
62+
]
63+
}
64+
}
65+
66+
/// Wraps a base input handler and appends model-specific extra inputs.
67+
///
68+
/// Use for any input that needs per-step computation beyond standard token/position IDs.
69+
public struct CompositeInputHandler<Base: SyncInputHandler>: SyncInputHandler {
70+
public var inputNames: [String] {
71+
base.inputNames + extras.map(\.name)
72+
}
73+
74+
private var base: Base
75+
private let extras: [ExtraInput]
76+
77+
public struct ExtraInput: Sendable {
78+
public let name: String
79+
public let prepare: @Sendable (InputContext) throws -> NDArray
80+
81+
public init(name: String, prepare: @Sendable @escaping (InputContext) throws -> NDArray) {
82+
self.name = name
83+
self.prepare = prepare
84+
}
85+
}
86+
87+
public init(base: Base, extras: [ExtraInput]) {
88+
self.base = base
89+
self.extras = extras
90+
}
91+
92+
public mutating func prepare(_ context: InputContext) async throws -> [String: NDArray] {
93+
var inputs = try await base.prepare(context)
94+
for extra in extras {
95+
inputs[extra.name] = try extra.prepare(context)
96+
}
97+
return inputs
98+
}
99+
}

‎swift/Sources/CoreAILanguageModels/Handlers/InputHandler.swift‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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 CoreAIShared
8+
9+
/// Context for each inference step, used by both dynamic (GPU) and static (ANE) engines.
10+
public struct InputContext: Sendable {
11+
/// Tokens to process in this step.
12+
public let tokens: ArraySlice<Int32>
13+
/// Number of tokens already processed before this step.
14+
public let processedTokenCount: Int
15+
/// Batch-aligned start position.
16+
public let alignedStep: Int
17+
/// Total batch slot count (may exceed tokens.count for static-shape padding).
18+
public let batchSize: Int
19+
/// Sliding window size (nil for models without sliding attention).
20+
public let slidingWindow: Int?
21+
22+
/// For dynamic-shape engines (Sequential, Pipelined).
23+
/// alignedStep = processedTokenCount, batchSize = tokens.count.
24+
public static func dynamic(
25+
tokens: ArraySlice<Int32>,
26+
processedTokenCount: Int
27+
) -> InputContext {
28+
InputContext(
29+
tokens: tokens,
30+
processedTokenCount: processedTokenCount,
31+
alignedStep: processedTokenCount,
32+
batchSize: tokens.count,
33+
slidingWindow: nil)
34+
}
35+
36+
/// For static-shape engines. Batch is fixed-size and aligned.
37+
public static func `static`(
38+
tokens: ArraySlice<Int32>,
39+
alignedStep: Int,
40+
batchSize: Int,
41+
slidingWindow: Int?
42+
) -> InputContext {
43+
InputContext(
44+
tokens: tokens,
45+
processedTokenCount: alignedStep,
46+
alignedStep: alignedStep,
47+
batchSize: batchSize,
48+
slidingWindow: slidingWindow)
49+
}
50+
}
51+
52+
/// Prepares model inputs for each inference step.
53+
///
54+
/// The engine calls `prepare(...)` each step and passes the result to `function.run()`.
55+
/// Standard models use `TokenInputHandler`; models with extra inputs (RoPE,
56+
/// sliding step, PLE) wrap it with `CompositeInputHandler`.
57+
public protocol SyncInputHandler {
58+
/// Input names this handler produces.
59+
var inputNames: [String] { get }
60+
61+
/// Prepare inputs for the current step.
62+
mutating func prepare(_ context: InputContext) async throws -> [String: NDArray]
63+
}
64+
65+
// MARK: - Load-time Coverage Check
66+
67+
public enum InputCoverage {
68+
/// Verify that a set of handlers covers all required inputs declared by the model descriptor.
69+
/// Call at engine init to fail fast on missing handlers rather than producing NaN at runtime.
70+
///
71+
/// - Parameters:
72+
/// - handlers: All input handlers the engine will use.
73+
/// - descriptor: The model function's descriptor declaring required inputs.
74+
/// - ignoring: Input names to exclude from the check (e.g. "embedding_table" passed directly).
75+
/// - Throws: If any declared input is not produced by any handler.
76+
public static func verify(
77+
handlers: [any SyncInputHandler],
78+
descriptor: InferenceFunctionDescriptor,
79+
ignoring: Set<String> = []
80+
) throws {
81+
let produced = handlers.reduce(into: Set<String>()) { $0.formUnion($1.inputNames) }
82+
let declared = Set(descriptor.inputNames).subtracting(ignoring)
83+
let uncovered = declared.subtracting(produced)
84+
guard uncovered.isEmpty else {
85+
throw InferenceRuntimeError.invalidState(
86+
"No input handler produces required input(s): \(uncovered.sorted()). "
87+
+ "Produced: \(produced.sorted())")
88+
}
89+
}
90+
}

0 commit comments

Comments
 (0)