Skip to content

Consolidate VLM sequential engine: shared KV cache, token box, chunked prefill (B1+B2) - #249

Merged
stikves merged 2 commits into
apple:mainfrom
stikves:sukru/vlm-consolidation-b1b2
Sep 14, 2026
Merged

stikves merged 2 commits into
apple:mainfrom
stikves:sukru/vlm-consolidation-b1b2

Conversation

@stikves

@stikves stikves commented Sep 13, 2026

Copy link
Copy Markdown
Contributor

Summary

First half of the VLM sequential-engine consolidation (steps B1 and B2). Removes duplicated
decode-loop code between CoreAISequentialVLMEngine and CoreAISequentialEngine. Net -215/+168
across 5 files, before tests.

Changes

B1: Replace the VLM engine's hand-rolled KV cache (keyCache/valueCache/currentKVCapacity plus
ensureKVCapacity/copyCache/zeroFill) with the shared StateHandlerFactory and runWithStates
path already used by the text engine. Growth semantics are unchanged (initial 256, 2x growth).
Restores the per-realloc "KV cache grew" log in the shared GrowingNDArrayState so both engines
emit it.

B2: Add GenerationTokenBox for the _activeToken busy/cancel/install bookkeeping both engines
duplicated, and runChunkedPrefill() to replace both processChunkedPrompt implementations. The
text engine passes a non-zero heldBack and routes non-held-back chunks through the prefill graph;
the VLM engine passes heldBack = 0.

Tests

Adds 16 unit tests (swift-testing, no model assets required):

  • GenerationTokenBoxTests: install/busy transitions, cancel and clear semantics, the
    clearIfActive newer-token guard, and a 400-task concurrency stress on the Mutex-backed state.
  • ChunkedPrefillTests: chunk math for heldBack == 0 and heldBack > 0, single-token and
    exact-multiple edges, and a coverage invariant.

Validation

Bit-exact against the pre-refactor baseline at this tip: text greedy 128/128, chunked text 64/64,
VLM greedy 115/115, VLM text-only chunked 64/64. Re-confirmed on the coreai-sequential engine for
a text model (Qwen3-0.6B, dynamic KV) and a VLM (qwen3-vl): coherent output, no crashes, KV growth
256 to 512 with correct copy-on-grow.

xcodebuild build and xcodebuild test pass.

Scope

Later steps are separate PRs: SequentialIterator scalar helpers and tidy (B3/B4), and the
EngineFactory bundle entry point plus removal of the CLI fork (C1/C2).

@carinapeng

Copy link
Copy Markdown
Contributor

The direction looks good, main design gap as you iterate is that seems like the VLM adopts the shared factory but not the full shared contract

The text engine reads back this flag hasNonTruncatableStates after the factory builds the states which is the careful approach, but the VLM engine throws that away, this is risky

@stikves

stikves commented Sep 14, 2026

Copy link
Copy Markdown
Contributor Author

The direction looks good, main design gap as you iterate is that seems like the VLM adopts the shared factory but not the full shared contract

The text engine reads back this flag hasNonTruncatableStates after the factory builds the states which is the careful approach, but the VLM engine throws that away, this is risky

This is the first step in series of changes. I don't want to drop one large PR, which would be much harder to read.

@stikves
stikves force-pushed the sukru/vlm-consolidation-b1b2 branch 2 times, most recently from 8d9bfc2 to f31b292 Compare September 14, 2026 21:09
@stikves
stikves marked this pull request as ready for review September 14, 2026 21:10
@stikves stikves self-assigned this Sep 14, 2026
…d prefill (B1+B2)

B1: Replace the VLM engine's hand-rolled KV cache (keyCache/valueCache/currentKVCapacity
plus ensureKVCapacity/copyCache/zeroFill) with the shared StateHandlerFactory and
runWithStates path used by the text engine. Growth semantics unchanged (initial 256, 2x).
Restore the per-realloc "KV cache grew" log in GrowingNDArrayState. Read back
hasNonTruncatableStates and guard reset(to:) against partial reset for recurrent-state
models, matching the text engine.

B2: Add GenerationTokenBox for the active-token busy/cancel/install bookkeeping both engines
duplicated, and runChunkedPrefill() to replace both processChunkedPrompt implementations.
The text engine passes a non-zero heldBack and routes non-held-back chunks through the
prefill graph; the VLM engine passes heldBack = 0.

Adds GenerationTokenBox and chunked-prefill unit tests.
@stikves
stikves force-pushed the sukru/vlm-consolidation-b1b2 branch from 0f8dd87 to f56e84b Compare September 14, 2026 21:26
@stikves
stikves merged commit 7359dbc into apple:main Sep 14, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants