Skip to content

CoreAISequentialEngine: idempotent generation via external session state - #258

Merged
stikves merged 6 commits into
apple:mainfrom
stikves:sukru/idempotent-engine
Sep 22, 2026
Merged

stikves merged 6 commits into
apple:mainfrom
stikves:sukru/idempotent-engine

Conversation

@stikves

@stikves stikves commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

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

func generate(
  with input: [TokenId],
  sessionState: GenerationSessionState,
  samplingConfiguration: SamplingConfiguration,
  inferenceOptions: InferenceOptions
) async throws -> OutputSequence

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
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.

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. All new and existing tests pass.

Limitations
Sessions run one at a time. lastPrefixHitCount and the shared GenerationTokenBox stay on the engine, so a concurrent scheduler would need to move or disable those too.

(I would later suggest moving them).

@stikves
stikves force-pushed the sukru/idempotent-engine branch 3 times, most recently from 125fbf8 to 71f29e3 Compare September 18, 2026 01:44
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
stikves force-pushed the sukru/idempotent-engine branch from 71f29e3 to b7e127b Compare September 18, 2026 01:44
@stikves stikves self-assigned this Sep 18, 2026
@stikves
stikves marked this pull request as ready for review September 18, 2026 02:27
stikves and others added 4 commits September 17, 2026 20:28
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).
@stikves
stikves merged commit c2ef849 into apple:main Sep 22, 2026
3 checks passed
@stikves
stikves deleted the sukru/idempotent-engine branch September 22, 2026 19:23
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.

2 participants