Consolidate VLM sequential engine: shared KV cache, token box, chunked prefill (B1+B2) - #249
Conversation
|
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 |
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. |
8d9bfc2 to
f31b292
Compare
…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.
0f8dd87 to
f56e84b
Compare
Summary
First half of the VLM sequential-engine consolidation (steps B1 and B2). Removes duplicated
decode-loop code between
CoreAISequentialVLMEngineandCoreAISequentialEngine. Net -215/+168across 5 files, before tests.
Changes
B1: Replace the VLM engine's hand-rolled KV cache (
keyCache/valueCache/currentKVCapacityplusensureKVCapacity/copyCache/zeroFill) with the sharedStateHandlerFactoryandrunWithStatespath 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
GrowingNDArrayStateso both enginesemit it.
B2: Add
GenerationTokenBoxfor the_activeTokenbusy/cancel/install bookkeeping both enginesduplicated, and
runChunkedPrefill()to replace bothprocessChunkedPromptimplementations. Thetext engine passes a non-zero
heldBackand 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, theclearIfActivenewer-token guard, and a 400-task concurrency stress on the Mutex-backed state.ChunkedPrefillTests: chunk math forheldBack == 0andheldBack > 0, single-token andexact-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-sequentialengine fora 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 buildandxcodebuild testpass.Scope
Later steps are separate PRs:
SequentialIteratorscalar helpers and tidy (B3/B4), and theEngineFactorybundle entry point plus removal of the CLI fork (C1/C2).