[Swift runtime] CoreAISequentialEngine rejects hybrid models with four persistent states

[massif-01]

Summary

CoreAISequentialEngine currently rejects a valid hybrid-decoder model that
has four persistent states (KV cache plus convolution and recurrent state).
The failure occurs before inference because the engine requires exactly two
states.

Affected Apple source commit:
04a3fd6cfe9bfae9cf05b1f246cf915d930d1c0a.

Minimal reproduction

I prepared a 16 KB .aimodel with no LLM weights. Its function contract is:

inputs: input_ids, position_ids
output: logits
states: keyCache, valueCache, convState, recState

Reproducer, fixture, iPhone harness, patch, and evidence:
https://github.com/massif-01/coreai-hybrid-state-runtime

Run the unmodified Apple runtime with the included iPhone harness. It fails at
engine initialization with:

Expected 2 states (KV cache), got 4: states=["keyCache", "valueCache", "convState", "recState"], outputs=["logits"]

Expected behavior

The sequential engine should bind and retain the fixed-shape convolution and
recurrent states in addition to the first two KV states, or document that this
otherwise valid Core AI state contract is intentionally unsupported.

Root cause

The engine checks descriptor.stateNames.count == 2. Simply relaxing the guard
is not sufficient: extra states must be allocated, zeroed, bound on every call,
retained after calls, and reset correctly.

There is also a correctness issue with KV-only prefix rewind. Recurrent state is
a summary of the full prefix and cannot be rewound by changing a KV token
cursor. A safe implementation must replay the prompt or maintain recurrent
checkpoints.

Proposed behavior

The attached patch:

• preserves the existing two-state path;
• supports up to two extra fixed-shape FP16/FP32 states;
• uses explicit lexical state bindings required by
  InferenceFunction.MutableViews.insert lifetime semantics;
• carries all states across decode calls;
• zeroes all states on full reset;
• rejects nonzero partial reset for recurrent models;
• fully replays a later prompt instead of performing an invalid KV-only rewind.

Device A/B evidence

Same fixture and harness on iPhone 15 Pro, iOS 27.0 (24A5380h), Xcode 27 beta
(27A5218g):

• unmodified runtime: fails with the four-state guard above;
• patched runtime: two decode steps pass;
• first logits: [1, 2, 3, 4, 1, 1, 3, 4];
• second logits: [5, 10, 15, 20, 4, 5, 15, 20];
• recurrent partial reset is rejected and full reset succeeds.

The accumulated second-step values prove that all four state writes persisted.
The same runtime path also completed real Qwen3.5-2B sequential structured
generation on device, but the tiny fixture is the self-contained regression.

Scope

This report concerns CoreAISequentialEngine, including constrained/structured
generation paths that use it. A community project has separately explored extra
states in CoreAIPipelinedEngine; that does not remove the sequential failure.

Would Apple confirm whether four-state hybrid decoders are intended to be
supported by the sequential runtime and consider adopting this behavior or an
equivalent lifetime-safe state-binding API?

[massif-01]

A corresponding Feedback Assistant report has now been filed as FB23893830.

The submission includes:

• the minimal 16 KB, no-weights .aimodel reproducer;
• the proposed runtime patch and iPhone A/B evidence; and
• an iPhone 15 Pro sysdiagnose captured after reproducing the baseline failure on iOS 27.0 Seed 3 (24A5380h).

The public reproducer and implementation remain available at https://github.com/massif-01/coreai-hybrid-state-runtime.

[stikves]

Thanks for the report.

The check for only 2 states is deliberate. We currently do not have support for linear attention or similar hybrid state models.

Keeping this open for potential future changes.

[massif-01]

Thanks for the report.

The check for only 2 states is deliberate. We currently do not have support for linear attention or similar hybrid state models.

Keeping this open for potential future changes.

Thank you for your reply. Additionally, I have another issue ; please take a look at it as well.
#116

[stikves]

Sure, thanks for the model suggestion. Let us take a look.

[stikves]

This is now implemented in #132