Add StaticInputHandler protocol for zero-allocation input filling by stikves · Pull Request #188 · apple/coreai-models

stikves · 2026-08-20T22:11:26Z

Adds InputBuffers and StaticInputHandler, a protocol for filling pre-allocated input buffers in-place without per-step allocation.

Design

public protocol StaticInputHandler: Sendable {
    var inputNames: [String] { get }
    func registerBuffers(into buffers: inout InputBuffers)
    func fill(_ context: InputContext, into buffers: inout InputBuffers) throws
}

Engine owns InputBuffers, creates once at init, passes inout each step
_modify accessor on InputBuffers subscript avoids COW (26× vs get/set)
Non-mutating fill — handler holds no per-step state
Backward compatible: existing SyncInputHandler unchanged

Primary benefit is architectural (no COW risk, compile-time ownership via inout) rather than raw throughput.

Introduces InputBuffers (pre-allocated, reused across steps) and the StaticInputHandler protocol (inout fill pattern). TokenStaticInputHandler fills input_ids and position_ids without per-step allocation. Wires the existing SyncInputHandler (TokenInputHandler) into the sequential engine, proving the protocol end-to-end. The engine now delegates input construction to handler.prepare() instead of inline fillNDArray calls. StaticShapeEngine (ANE) is unchanged.

Add StaticBucketInputFiller conforming to StaticInputHandler. It fills position_ids (UInt16), causal_mask (Float16), and step (Int32) into engine-owned InputBuffers in-place — no per-step allocation. The static engine now delegates input construction to the filler via the register/fill protocol. Embedding table and transformer_input (which require ANE gather invocation) remain engine-managed. Add contextBucket field to InputContext so the filler knows the current graph's context length for mask shape selection.

carinapeng

There are two blockers, suggest testing more and testing end to end before marking as review ready

For perf/CoW, I agree with you that maybe we want to switch to _modify

Rewrite StaticBucketInputFiller to pre-allocate all input NDArrays per (batchSize, contextBucket) pair at init from per-graph descriptors. Zero per-step allocation — bucket switch is an O(1) pool lookup. This matches the original RFC design (StaticBucketInputHandler with per-bucket BucketBuffers) but conforms to the StaticInputHandler protocol's inout fill pattern instead of returning a new dictionary. Add InputBuffers.preAllocate(name:descriptor:) for pooling NDArrays by shape at init, and ensureCapacity(name:descriptor:) for O(1) pool swap at runtime. These bypass resolvingDynamicDimensions entirely — static-shape models have all-fixed descriptors, no dynamic dims. Add InputBuffers.withMutableBuffer for COW-safe mutation through the dictionary's _modify accessor. Engine init now scans all extend/prompt functions to build the complete bucket descriptor table. Remove dead fillCausalMask — superseded by the handler.

Addresses review: TokenStaticInputHandler was never wired into any engine. The sequential engine uses TokenInputHandler (SyncInputHandler), the static engine uses StaticBucketInputFiller (StaticInputHandler). Drop it until the sequential engine migration is ready.

stikves · 2026-08-28T17:07:50Z

There are two blockers, suggest testing more and testing end to end before marking as review ready

For perf/CoW, I agree with you that maybe we want to switch to _modify

They should be handled now, please take a look

…ument COW Extract parseFunctionDimensions() to eliminate 5 duplicated context/query length parsing sites across StaticShapeEngine. Add resolveInputName() with known-name lists instead of fragile contains() matching. Rename asDict() to borrowedInputs() with ownership-contract docs. Add causalMaskSentinel constant shared across engines. Document COW risk on InputBuffers subscript get accessor.

carinapeng

Let's do more testing before merging :)

The model's step input is named in_step but knownStepNames only contained step, so the input was silently skipped. This caused in_step errors during graph evaluation on all static-shape models.

stikves marked this pull request as ready for review August 26, 2026 02:26

stikves force-pushed the sukru/input-handler-public branch from d25b4fa to 1a4c832 Compare August 27, 2026 05:33

stikves force-pushed the sukru/input-handler-public branch from 0d12561 to 54cff2d Compare August 27, 2026 06:21

tjia1818 requested a review from carinapeng August 27, 2026 16:22

tjia1818 approved these changes Aug 27, 2026

View reviewed changes

carinapeng reviewed Aug 27, 2026

View reviewed changes

Comment thread swift/Sources/CoreAILanguageModels/Handlers/InputHandler.swift Outdated

carinapeng reviewed Aug 27, 2026

View reviewed changes

Comment thread swift/Sources/CoreAILanguageModels/Handlers/InputHandler.swift Outdated

carinapeng reviewed Aug 27, 2026

View reviewed changes

Comment thread swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

carinapeng reviewed Aug 27, 2026

View reviewed changes

Comment thread swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift Outdated

carinapeng reviewed Aug 27, 2026

View reviewed changes

Comment thread swift/Tests/LanguageModelsTests/InputFillerTests.swift

carinapeng requested changes Aug 27, 2026

View reviewed changes

stikves force-pushed the sukru/input-handler-public branch 3 times, most recently from f0c2192 to 4c5fee0 Compare August 28, 2026 16:43

stikves force-pushed the sukru/input-handler-public branch from d39698d to 372db91 Compare August 28, 2026 17:04