Skip to content

Add bind(into:) for zero-copy state binding - #156

Merged
stikves merged 1 commit into
apple:mainfrom
stikves:sukru/state-handler-binder
Aug 8, 2026
Merged

Add bind(into:) for zero-copy state binding#156
stikves merged 1 commit into
apple:mainfrom
stikves:sukru/state-handler-binder

Conversation

@stikves

@stikves stikves commented Aug 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

State handlers are now classes that own their NDArrays at refcount 1.

bind(into:) inserts all states into MutableViews in a loop using stdlib _overrideLifetime to express disjoint
element access. Same pattern for AsyncMutableViews on the pipelined engine's FixedMTLBufferState.

Changes

  • SyncStateHandler protocol: add bind(into:), constrain to AnyObject
  • FixedNDArrayState, GrowingNDArrayState: struct → class, dictionary-backed
  • FixedMTLBufferState: struct → class, add bind(into:) for AsyncMutableViews
  • runWithStates(): sequential engine helper
  • encodeWithStates(): pipelined engine helper
  • Sequential engine: replace inline state binding with runWithStates()
  • Pipelined engine: replace 3 inline state binding blocks with encodeWithStates()

Requires Lifetimes experimental Swift feature.
Tested with 2 state and 4 state hybrid models.

@stikves
stikves force-pushed the sukru/state-handler-binder branch 3 times, most recently from b12a824 to ff3538f Compare August 6, 2026 21:16
@stikves
stikves marked this pull request as ready for review August 7, 2026 01:46
@stikves
stikves force-pushed the sukru/state-handler-binder branch 4 times, most recently from 0237172 to cb23ae4 Compare August 7, 2026 18:13
@stikves
Add bind(into:) for zero-copy state binding with loop-based insert
e2daf95 
State handlers are now classes (AnyObject) that own their NDArrays at
refcount 1. bind(into:) inserts all states into MutableViews in a loop
using stdlib _overrideLifetime to express disjoint element access.
No COW, no switch on state count, no write-back needed.

FixedMTLBufferState also gains bind(into:) for AsyncMutableViews,
eliminating the 3x duplicated switch blocks in the pipelined engine.

Engine call sites: runWithStates (sequential) and encodeWithStates
(pipelined) are now thin wrappers around bind + function.run/encode.
@stikves
stikves force-pushed the sukru/state-handler-binder branch from cb23ae4 to e2daf95 Compare August 7, 2026 18:19
@stikves stikves self-assigned this Aug 7, 2026
@stikves stikves changed the title Complete Input Handler Structures and Centralize MutableViews binding with withBoundStates Add bind(into:) for zero-copy state binding Aug 7, 2026
carinapeng
carinapeng reviewed Aug 7, 2026
View reviewed changes
Comment thread swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift
carinapeng
carinapeng reviewed Aug 7, 2026
View reviewed changes
Comment thread swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift
carinapeng
carinapeng reviewed Aug 7, 2026
View reviewed changes
Comment thread swift/Tests/LanguageModelsTests/StateHandlerTests.swift
@stikves
stikves merged commit 1677713 into apple:main Aug 8, 2026
3 checks passed
@stikves
stikves deleted the sukru/state-handler-binder branch August 8, 2026 01:58
carinapeng pushed a commit to carinapeng/coreai-models that referenced this pull request Aug 10, 2026
Carina Peng
Add static-shape inference engine
8f18e0d 
Static-shape LLM inference engine (StaticShapeEngine) built on the bind(into:)
state binding from apple#156. States are discovered by name and the KV cache is
right-sized per context bucket (each bucket graph is compiled with its own
per-ctx strides, so a max-ctx buffer sliced down corrupts KV); ctx is parsed
from the function name. Input preparation is pluggable per model family via
StaticInputProvider + StaticModelProfile; StaticInputContext composes the shared
InputContext and the provider mirrors SyncInputHandler.

Testing: LanguageModelsTests (317) pass; correct output on gemma4-E2B,
qwen2.5-1.5B, qwen3-0.6B static assets.

Known limitation: cross-bucket cache re-layout needs a channel-interleave-aware
copy; within-bucket generation is correct.

Depends on apple#156.
@carinapeng carinapeng mentioned this pull request Aug 10, 2026
Open
carinapeng pushed a commit to carinapeng/coreai-models that referenced this pull request Aug 10, 2026
Carina Peng
Add static-shape inference engine
6265330 
Static-shape LLM inference engine (StaticShapeEngine) built on the bind(into:)
state binding from apple#156. States are discovered by name and the KV cache is
right-sized per context bucket (each bucket graph is compiled with its own
per-ctx strides, so a max-ctx buffer sliced down corrupts KV); ctx is parsed
from the function name. Input preparation is pluggable per model family via
StaticInputProvider + StaticModelProfile; StaticInputContext composes the shared
InputContext and the provider mirrors SyncInputHandler.

Testing: LanguageModelsTests (317) pass; correct output on gemma4-E2B,
qwen2.5-1.5B, qwen3-0.6B static assets.

Known limitation: cross-bucket cache re-layout needs a channel-interleave-aware
copy; within-bucket generation is correct.

Depends on apple#156.
carinapeng pushed a commit to carinapeng/coreai-models that referenced this pull request Aug 10, 2026
Carina Peng
Add static-shape inference engine
50693ed 
Static-shape LLM inference engine (StaticShapeEngine). States are discovered by
name and the KV cache is right-sized per context bucket (each bucket graph is
compiled with its own per-ctx strides, so a max-ctx buffer sliced down corrupts
KV); ctx is parsed from the function name.

Input preparation extends Sukru's shared SyncInputHandler / InputContext (apple#147):
concrete handlers (position ids, causal mask, step, RoPE, PLE, sliding) conform
to SyncInputHandler; InputContext gains a per-graph descriptors map so handlers
size their own buffers. No Static input protocol/context — one input family for
dynamic and static engines. State binding uses bind(into:) (apple#156).

Testing: LanguageModelsTests pass incl. new unit tests (ctx-bucket parsing,
causal mask fill); correct output on gemma4-E2B, qwen2.5-1.5B, qwen3-0.6B static.

Depends on apple#156, apple#147.
msnabiel added a commit to msnabiel/coreai-models that referenced this pull request Aug 11, 2026
@msnabiel @claude
Merge upstream/main into fork
89f6af6 
Resolves conflicts:
- Removed ModelShapeConfig.swift (deleted upstream in apple#148)
- Updated NDArray+Helpers.swift documentation for stride-aware indexing

Pulls in upstream features:
- Parakeet speech recognition export and runtime (apple#136)
- SyncInputHandler protocol and InputContext (apple#147)
- Zero-copy state binding with bind(into:) (apple#156)
- Additional speech tests and config improvements
- Bug fixes for forced_decoder_ids parsing, topK normalization, logits token count

Preserves local iOS optimizations:
- int8 KV cache quantization (2× memory reduction)
- 16k default context with unreachable shape pruning
- CoreAI simulator availability guards

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@carinapeng carinapeng carinapeng left review comments
@tjia1818 tjia1818 tjia1818 approved these changes
@alejandro-isaza alejandro-isaza Awaiting requested review from alejandro-isaza
@kevchengcodes kevchengcodes Awaiting requested review from kevchengcodes
@blevine1 blevine1 Awaiting requested review from blevine1

Assignees

@stikves stikves

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@stikves @tjia1818 @carinapeng