`export_llm` Core ML builds fail at execute: "must be an MLState" even though `take_over_mutable_buffer=False`

[john-rocky]

Exporting any KV-cache LLM through export_llm with the Core ML backend produces a
.pte that cannot execute. The delegate reports a state input that the runtime
never binds:

[backend_delegate.mm:435] [Core ML] Model execution failed
    The input feature for layers_9_attention_recurrent_state must be an MLState, but it was not.
[coreml_backend_delegate.mm:388] CoreMLBackend: Failed to run the model.
[method.cpp:1530] CALL_DELEGATE execute failed at instruction 0: 0x32

partitioner_lib.py already anticipates this and disables the takeover:

# ExecuTorch does not build CoreML delegate runtime to handle state
# when using OSS scripts, so we define take_over_mutable_buffer = False,
# even when target is iOS18
take_over_mutable_buffer = False

and the flag really is off in the path export_llama_lib.py uses:

>>> from executorch.extension.llm.export.partitioner_lib import get_coreml_partitioner
>>> get_coreml_partitioner(ios=18, coreml_quantize="c4w",
...                        coreml_compute_units="all").take_over_mutable_buffer
False

So the guard is in place and the state is created anyway.

Minimal reproduction (executorch 1.4.0, coremltools 9.0, macOS arm64) — the
trigger is in-place mutation of a slice of a buffer, which is how Qwen3.5's
gated delta-net updates its recurrent state and how LFM2.5 updates its conv state:

class Cache(torch.nn.Module):
    def __init__(self, n=16, d=8):
        super().__init__()
        self.register_buffer("cache", torch.zeros(1, n, d))

    def forward(self, x, pos):
        self.cache[:1].mul_(0.9)
        self.cache[:1].copy_(self.cache[:1] + x.unsqueeze(1))
        return self.cache.sum(dim=1) + x

specs = CoreMLBackend.generate_compile_specs(
    compute_precision=ct.precision.FLOAT16, compute_unit=ct.ComputeUnit.ALL,
    minimum_deployment_target=ct.target.iOS18)
part = CoreMLPartitioner(compile_specs=specs, take_over_mutable_buffer=False)
et = to_edge_transform_and_lower(torch.export.export(Cache().eval(), (x, pos)),
                                 partitioner=[part]).to_executorch()

• minimum_deployment_target=iOS18: lowers, then execute fails with
  CALL_DELEGATE execute failed at instruction 0: 0x32.
• minimum_deployment_target=iOS17: lowering itself fails with
  State model is supported only >= iOS18, which says plainly that a state model
  was built even though the takeover is off.

Mutating the whole buffer instead (cache.index_put_(...)) executes fine, so the
slice is what does it. Full script: convert/repro_coreml_state.py in the report.

The same happens end to end through export_llm on Qwen3.5-0.8B
(layers_9_attention_recurrent_state) and LFM2.5-350M (layers_7_conv_conv_state).
The XNNPACK builds of both run — 172.8 tok/s decode for LFM2.5-350M on the same
machine.

A fix I tried and would not send. Dropping delegation_tag from the nodes
listed in graph_signature.buffers_to_mutate when the takeover is off leaves the
mutation outside Core ML, but the node sits inside the proposed partition and
removing it alone produces Invalid partition, found dependency cycles. Excluding
the mutation properly means excluding its dependent chain, or teaching the
delegate runtime to bind Core ML state — which the comment in partitioner_lib.py
implies is the actual missing piece.

What I have not ruled out. examples/models/llama/non_cpu_backends.md says to
run Core ML LLM builds through a llama runner built with Core ML flags, and I have
only tried the Python runtime. But the comment above says the OSS runtime does not
handle Core ML state at all, which suggests the runner would hit the same wall.
Is the OSS runner expected to execute these files, and if so, what binds the state?

Related: the tied-embedding conflict that also blocks this path is filed
separately as #21856 — it reproduces with no mutable buffer, so the two are
independent.

Edited twice. An earlier version of this section listed two more items. The
dynamic-shape one was wrong: _validate_args runs as the first statement of
_export_llama and raises immediately with the fix in the message. The
compute_units: cpu_and_ne one turned out not to be independent — a model with no
mutable buffer lowers and runs fine under cpu_and_ne, so that failure is a
symptom of the state problem above, not a separate defect. Apologies for the
noise.

[john-rocky]

Some more digging, in case it saves someone time. The short version is that
take_over_mutable_buffer=False cannot be honoured today, because the decision is
not ExecuTorch's to make.

Where the state comes from. coreml_preprocess.py calls ct.convert(model=edge_program, ...)
without passing states, and coremltools then does this (frontend/torch/converter.py):

if states is None or len(states) == 0:
    # For torch.export, we default to create states from torch mutable buffers
    self.states = []
    for name, tensor in self.graph.buffers.items():
        ...

Every registered buffer in the delegated program becomes a Core ML state, and
states=[] takes the same branch as states=None, so there is no way to opt out
from the caller's side. The partitioner flag controls tagging; it has no influence
over this.

Two fixes I tried and would not send, both at the wrong layer:

1. Dropping delegation_tag from the nodes that write the buffer. The nodes sit
   inside the proposed partition, so removing them alone yields
   Invalid partition, found dependency cycles.
2. Reporting mutating ops (schema.is_mutable) as unsupported so
   CapabilityBasedPartitioner routes around them. This partitions cleanly, but
   the buffer is still read inside the subgraph, so coremltools still turns it into
   a state — the failure just renames itself from cache_state to cache.

So the fix looks like one of three, none of them local to the partitioner:

• coremltools grows a way to say "do not create states from buffers" (a sentinel
  distinct from the empty list);
• the Core ML delegate runtime learns to bind MLState, which is what the comment
  in partitioner_lib.py implies is missing;
• or the backend rewrites buffers into plain inputs of the delegated sub-program
  before calling ct.convert, so there is nothing for coremltools to promote.

I am happy to prototype whichever direction you prefer, but which one is right
seems like a call for the people who own these two pieces rather than something to
guess at from outside.