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

Author: stikves

Package.resolved

@@ -0,0 +1,107 @@
+// Asynchronous token input handler for the Pipelined engine (MTLBuffer-based).
+//
+// 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
+
+/// Input handler for the pipelined engine. Owns the token and position MTLBuffers,
+/// handles buffer rotation, and the prefill/decode source split.
+///
+/// `decodeOutputBuffers` is shared with the engine (GPU sampler writes next token
+/// there; this handler reads the previous step's token during decode).
+struct PipelinedTokenInputHandler {
+    let inputIdsName: String
+    let positionIdsName: String
+    let inputIdsDescriptor: NDArrayDescriptor
+    let positionIdsDescriptor: NDArrayDescriptor
+
+    let inputTokensBuffer: MTLBuffer
+    let cachePositionBuffers: [MTLBuffer]
+    let decodeOutputBuffers: [MTLBuffer]
+    let pipelineDepth: Int
+
+    /// Build async input values for one encode step.
+    ///
+    /// - Prefill (`!tokens.isEmpty`): writes tokens to `inputTokensBuffer` at their
+    ///   natural position (disjoint from prior chunks still in-flight on GPU).
+    /// - Decode (`tokens.isEmpty`): reads from previous step's `decodeOutputBuffer`.
+    func prepare(
+        tokens: some Collection<Int32>,
+        processedTokenCount: Int,
+        step: Int
+    ) throws -> [String: InferenceFunction.AsyncValue] {
+        let queryLength = tokens.isEmpty ? 1 : tokens.count
+
+        // Write tokens (prefill only — decode reads from GPU sampler output)
+        if !tokens.isEmpty {
+            let dst = inputTokensBuffer.contents()
+                .assumingMemoryBound(to: Int32.self)
+                .advanced(by: processedTokenCount)
+            let tokenArray = Array(tokens)
+            tokenArray.withUnsafeBufferPointer { src in
+                memcpy(dst, src.baseAddress!, tokenArray.count * MemoryLayout<Int32>.size)
+            }
+        }
+
+        // Token input
+        let tokenShape = [1, queryLength]
+        let tokenStrides = try resolvedStrides(descriptor: inputIdsDescriptor, shape: tokenShape)
+        let tokenValue: InferenceFunction.AsyncValue
+        if tokens.isEmpty {
+            tokenValue = unsafe InferenceFunction.AsyncValue(
+                unsafeBuffer: decodeOutputBuffers[(step + pipelineDepth - 1) % pipelineDepth],
+                byteOffset: 0, scalarType: .int32, shape: tokenShape, strides: tokenStrides)
+        } else {
+            tokenValue = unsafe InferenceFunction.AsyncValue(
+                unsafeBuffer: inputTokensBuffer,
+                byteOffset: processedTokenCount * MemoryLayout<Int32>.size,
+                scalarType: .int32, shape: tokenShape, strides: tokenStrides)
+        }
+
+        // Position input (rotating buffer)
+        let posLength = processedTokenCount + queryLength
+        let posShape = [1, posLength]
+        let posStrides = try resolvedStrides(descriptor: positionIdsDescriptor, shape: posShape)
+        let posValue = unsafe InferenceFunction.AsyncValue(
+            unsafeBuffer: cachePositionBuffers[step % pipelineDepth],
+            byteOffset: 0, scalarType: .int32, shape: posShape, strides: posStrides)
+
+        return [
+            inputIdsName: tokenValue,
+            positionIdsName: posValue,
+        ]
+    }
+
+    /// Prepare inputs for warmup (always prefill mode, writes dummy tokens).
+    func prepareWarmup(
+        shape: Int,
+        processedTokenCount: Int,
+        step: Int
+    ) throws -> [String: InferenceFunction.AsyncValue] {
+        let ptr = inputTokensBuffer.contents().assumingMemoryBound(to: Int32.self)
+        for i in 0..<shape { ptr[i] = 1 }
+
+        let tokenShape = [1, shape]
+        let tokenStrides = try resolvedStrides(descriptor: inputIdsDescriptor, shape: tokenShape)
+        let tokenValue = unsafe InferenceFunction.AsyncValue(
+            unsafeBuffer: inputTokensBuffer, byteOffset: 0,
+            scalarType: .int32, shape: tokenShape, strides: tokenStrides)
+
+        let posLength = processedTokenCount + shape
+        let posShape = [1, posLength]
+        let posStrides = try resolvedStrides(descriptor: positionIdsDescriptor, shape: posShape)
+        let posValue = unsafe InferenceFunction.AsyncValue(
+            unsafeBuffer: cachePositionBuffers[step % pipelineDepth],
+            byteOffset: 0, scalarType: .int32, shape: posShape, strides: posStrides)
+
+        return [
+            inputIdsName: tokenValue,
+            positionIdsName: posValue,
+        ]
+    }
+}

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

@@ -0,0 +1,99 @@
+// Synchronous token input handler for Sequential and StaticShape engines.
+//
+// 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
+
+/// Standard input handler for text LLMs: `input_ids` (Int32) + `position_ids` (Int32).
+///
+/// Pre-allocates the `input_ids` NDArray and reuses it when batch size is unchanged.
+public struct TokenInputHandler: SyncInputHandler {
+    public let inputNames: [String]
+
+    private let inputIdsName: String
+    private let positionIdsName: String
+    private let inputIdsDescriptor: NDArrayDescriptor
+    private let positionIdsDescriptor: NDArrayDescriptor
+
+    private var inputIdsArray: NDArray
+    private var cachedBatchSize: Int
+
+    public init(
+        inputIdsName: String,
+        positionIdsName: String,
+        inputIdsDescriptor: NDArrayDescriptor,
+        positionIdsDescriptor: NDArrayDescriptor
+    ) {
+        self.inputIdsName = inputIdsName
+        self.positionIdsName = positionIdsName
+        self.inputIdsDescriptor = inputIdsDescriptor
+        self.positionIdsDescriptor = positionIdsDescriptor
+        self.inputNames = [inputIdsName, positionIdsName]
+
+        let initDesc = inputIdsDescriptor.resolvingDynamicDimensions([1, 1])
+        self.inputIdsArray = NDArray(descriptor: initDesc)
+        self.cachedBatchSize = 1
+    }
+
+    public mutating func prepare(_ context: InputContext) async throws -> [String: NDArray] {
+        let tokens = context.tokens
+        let batchSize = tokens.count
+        precondition(batchSize > 0, "TokenInputHandler: empty token batch")
+
+        if cachedBatchSize != batchSize {
+            let resolved = inputIdsDescriptor.resolvingDynamicDimensions([1, batchSize])
+            inputIdsArray = NDArray(descriptor: resolved)
+            cachedBatchSize = batchSize
+        }
+        fillNDArray(&inputIdsArray, as: Int32.self, with: tokens)
+
+        let totalPositions = context.processedTokenCount + batchSize
+        let resolvedPosDesc = positionIdsDescriptor.resolvingDynamicDimensions([1, totalPositions])
+        var positionIds = NDArray(descriptor: resolvedPosDesc)
+        fillNDArray(&positionIds, as: Int32.self, count: totalPositions) { Int32($0) }
+
+        return [
+            inputIdsName: inputIdsArray,
+            positionIdsName: positionIds,
+        ]
+    }
+}
+
+/// Wraps a base input handler and appends model-specific extra inputs.
+///
+/// Use for any input that needs per-step computation beyond standard token/position IDs.
+public struct CompositeInputHandler<Base: SyncInputHandler>: SyncInputHandler {
+    public var inputNames: [String] {
+        base.inputNames + extras.map(\.name)
+    }
+
+    private var base: Base
+    private let extras: [ExtraInput]
+
+    public struct ExtraInput: Sendable {
+        public let name: String
+        public let prepare: @Sendable (InputContext) throws -> NDArray
+
+        public init(name: String, prepare: @Sendable @escaping (InputContext) throws -> NDArray) {
+            self.name = name
+            self.prepare = prepare
+        }
+    }
+
+    public init(base: Base, extras: [ExtraInput]) {
+        self.base = base
+        self.extras = extras
+    }
+
+    public mutating func prepare(_ context: InputContext) async throws -> [String: NDArray] {
+        var inputs = try await base.prepare(context)
+        for extra in extras {
+            inputs[extra.name] = try extra.prepare(context)
+        }
+        return inputs
+    }
+}