CoreAISequentialEngine: idempotent generation via external session state - #258
Merged
Merged
Conversation
stikves
force-pushed
the
sukru/idempotent-engine
branch
3 times, most recently
from
September 18, 2026 01:44
125fbf8 to
71f29e3
Compare
Summary Move per-request generation state off the engine into a caller-owned object, and add a protocol that drives generation over that object. Behavior does not change: the existing generate() API and its callers work as before. This is the interface a future scheduler needs to run several sessions on one engine. Background The engine kept each conversation's state on the instance: the KV cache, optional persistent (hybrid) states, the processed-token cursor, and the token history used for prefix reuse. The iterator mutates that state as a side effect, so one engine serves one conversation at a time and prefix reuse only applies to the request that ran last. That blocks batched or multi-session serving at the engine boundary. Changes GenerationSessionState holds the relocated state (KV cache, additional states, recurrent flag, cursor, history). It is a class so the iterator's cursor and history writes stay visible to the shim across calls; a struct would copy them and lose prefix reuse. It is not Sendable: it holds non-Sendable NDArray state and stays in the engine's one isolation domain, so region isolation covers the async generate() call and no @unchecked is needed. IdempotentEngine refines InferenceEngine, the same shape as ConstrainedGenerationCapable: makeSessionState() plus generate(with:sessionState:). It is detected with `engine is any IdempotentEngine`. Only CoreAISequentialEngine conforms. The engine now routes every forward-path, reset, and prefix-resolution site through the passed session. The public generate() is a shim over one internal session, which keeps today's prefix reuse, the GenerationTokenBox single-active-generation contract, and reset(to:) semantics as they were. The prefix-routing decision (full reset on divergence, full reset for hybrid models, rewind on pure extension, and the matching prefix-hit counts) moves into a pure prefixResetPlan() function so it can be unit-tested. Tests Builds clean with no new warnings. Seven model-free tests cover what this change touches: session independence, the prefixResetPlan routing table, the shim's preservation of the single-active-generation contract, and capability routing (the other three engines do not conform). Tests do not re-cover TokenHistory prefix resolution or the GenerationTokenBox primitive, which this change leaves unmodified. An on-device parity suite (idempotent path equals the shim; snapshot, restore, and continue equal an uninterrupted run) is gated behind a model fixture. Limitations Sessions run one at a time. lastPrefixHitCount and the shared GenerationTokenBox stay on the engine, so a concurrent scheduler would need to move those too. Run the on-device parity suite before merge to confirm KV bitwise equivalence.
stikves
force-pushed
the
sukru/idempotent-engine
branch
from
September 18, 2026 01:44
71f29e3 to
b7e127b
Compare
stikves
requested review from
alejandro-isaza,
carinapeng,
kevchengcodes,
srjoglekar246 and
tjia1818
September 18, 2026 02:27
stikves
marked this pull request as ready for review
September 18, 2026 02:27
The suite had accumulated cases that exercise unchanged primitives rather than the idempotent refactor: reference-type alias semantics, the eight TokenHistory.resolve cases (resolve() is untouched by this change), and the GenerationTokenBox unit tests. Remove them. What stays covers the refactor: session independence (sessionsDoNotAlias), the prefixResetPlan routing table (the pure decision extracted from generate(with:sessionState:)), the supersede contract, capability routing via `is any IdempotentEngine`, and the gated on-device parity scaffold. Also promote the EngineOptions comment to a doc comment. Model-free suite: 7 tests in 4 suites pass; build clean.
A fresh session (empty history) left lastPrefixHitCount holding the previous session's value. Default it to 0 before the history check. Remove sessionsDoNotAlias (checked two distinct objects aren't the same instance) and supersedeCancelsPrevious (exercised pre-existing GenerationTokenBox/GenerationToken, unrelated to this change).
tjia1818
approved these changes
Sep 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Move per-request generation state off the engine into a caller-owned object, and add a protocol that drives generation over that object. The existing
generate()API keeps its current signature. This is the interface a future scheduler needs to run several sessions on one engine.The new interface
augments the older one with explicit
sessionState.Background
The engine kept each conversation's state on the instance: the KV cache, optional persistent (hybrid) states, the processed-token cursor, and the token history used for prefix reuse. The iterator mutates that state as a side effect, so one engine serves one conversation at a time and prefix reuse only applies to the request that ran last. That blocks batched or multi-session serving at the engine boundary.
Changes
GenerationSessionStateholds the relocated state (KV cache, additional states, recurrent flag, cursor, history).It is a
classso the iterator's cursor and history writes stay visible to the shim across calls; a struct would copy them and lose prefix reuse.It is not
Sendable: it holds non-SendableNDArraystate and stays in the engine's one isolation domain, so region isolation covers theasync generate()call.IdempotentEnginerefinesInferenceEngine, the same shape asConstrainedGenerationCapable:makeSessionState()plusgenerate(with:sessionState:). It is detected withengine is any IdempotentEngine. OnlyCoreAISequentialEngineconforms.The engine now routes every forward-path, reset, and prefix-resolution site through the passed session. The public
generate()is a shim over one internal session, which keeps today's prefix reuse, theGenerationTokenBoxsingle-active-generation contract, andreset(to:)semantics as they were.The prefix-routing decision (full reset on divergence, full reset for hybrid models, rewind on pure extension, and the matching prefix-hit counts) moves into a pure
prefixResetPlan()function so it can be unit-tested.Tests
Builds clean with no new warnings. All new and existing tests pass.
Limitations
Sessions run one at a time.
lastPrefixHitCountand the sharedGenerationTokenBoxstay on the engine, so a concurrent scheduler would need to move or disable those too.(I would later suggest moving them).