A transformer can look fast in its first few inference calls and still become unusable during a long interaction. If the app submits the full history on every step, the model repeats attention work for information it has already processed. Latency grows with context length, and an experience that begins responsive can gradually lose its rhythm.
Core AI supports stateful inference for this situation. Model states are values that the inference function reads and updates in place. A key-value cache stored as state lets the next token or observation reuse work from the preceding steps instead of rebuilding the complete history.
Confirm that history is the cost
Do not add a cache because a model happens to be called a transformer. Record the real loop in the Core AI instrument and inspect inference intervals over time. A steadily increasing interval is a strong signal that the input history or repeated attention work is growing.
Then verify the model's contract. A stateful version has different inputs, states, and reset requirements from a stateless version. The cache is not an invisible app optimization. It becomes part of the model interface and must have a clear owner.
Author mutable cache buffers in Python
The state must originate in the model authoring code. In the Core AI session example, key and value tensors are registered as buffers in the PyTorch module. The forward pass reads the valid cache prefix, uses it while calculating the new step, then writes the new keys and values back.
class StatefulTransformer(nn.Module):
def __init__(self, layers, max_context, hidden_size):
super().__init__()
self.register_buffer(
"key_cache",
torch.zeros(layers, max_context, hidden_size),
)
self.register_buffer(
"value_cache",
torch.zeros(layers, max_context, hidden_size),
)
def forward(self, features, position_ids):
# Read the prior cache, process the new features, then update both buffers.
...The important decision is the maximum context size. A fixed maximum makes allocation predictable, but the product must define what happens when the interaction reaches it. Resetting, truncating, or summarizing context are different application behaviours, not implementation details to leave to a tensor shape.
Export states as part of the function signature
When converting the exported program, pass the cache names as state_names. This tells Core AI that these values are mutable state arguments rather than ordinary model inputs or returned outputs.
program = coreai_torch.TorchConverter().add_exported_program(
exported,
input_names=["features", "position_ids"],
state_names=["keyCache", "valueCache"],
output_names=["logits"],
).to_coreai()Keep the conversion names aligned with the Swift code. An InferenceFunctionDescriptor can inspect the resulting state names and descriptors before an app allocates the buffers. This is a good place to catch a cache shape mismatch rather than discover it during a user interaction.
Give each conversation or loop its own state
In Swift, store the cache arrays with the object that owns one inference sequence. On each call, insert mutable views for the named states and supply them to run. Core AI updates those views in place.
var stateViews = InferenceFunction.MutableViews()
stateViews.insert(&keyCache, for: "keyCache")
stateViews.insert(&valueCache, for: "valueCache")
var outputs = try await function.run(
inputs: ["features": nextFeatures],
states: stateViews
)Do not share one mutable cache across independent chats, games, or background jobs. The cache contains the history for exactly one sequence. Start a fresh cache when that sequence starts, and reset it whenever the app intentionally discards the old context.
Measure the updated loop, not only the first token
After conversion, record the same long interaction again. The intended result is stable or much more slowly growing inference latency as context accumulates. Also watch memory usage. A cache can trade repeated computation for a predictable memory allocation, so its size needs to fit the target device and the number of simultaneous sessions.
If the model remains slow, the next issue may be model size, layout conversions, output allocation, or a different operation. The cache is a targeted optimization for repeated history work, not a substitute for profiling.
Key takeaway
Stateful Core AI inference turns a growing history into an owned resource. Model the cache in Python, export it as named state, keep one cache per sequence, and prove the gain with a trace that covers the whole interaction.