[psiddh]

LLM runner: select I/O contracts and accept custom buffer allocators

Add a backend-neutral IOManager registry and buffer-allocator hook. The generic runner selects a registered manager from the program metadata and rejects unknown contracts or invalid metadata. This PR adds no concrete cache adapter; the hf_static implementation follows in PR 7.

The runner loads and resets selected methods, preserves module/buffer lifetimes, and uses a side-effect-free capacity preflight to route a prompt tail through decode. Execution errors are not retried. Builtin registration is an inert definition until PR 7 registers the concrete manager.

Review the interface/registry, runner load/reset and lifetime handling, prompt capacity preflight, then tests. QNN allocation follows in PR 4. Eviction, state adoption and fork policy remain IOManager follow-ups.

Files

extension/llm/runner/constants.h
extension/llm/runner/io_manager/io_manager.h
extension/llm/runner/llm_runner_helper.cpp
extension/llm/runner/test/CMakeLists.txt
extension/llm/runner/test/targets.bzl
extension/llm/runner/test/test_io_manager.cpp
extension/llm/runner/test/test_text_llm_runner.cpp
extension/llm/runner/test/test_text_prefiller.cpp
extension/llm/runner/text_decoder_runner.cpp
extension/llm/runner/text_decoder_runner.h
extension/llm/runner/text_llm_runner.cpp
extension/llm/runner/text_llm_runner.h
extension/llm/runner/text_prefiller.cpp
extension/llm/runner/text_prefiller.h

Test Plan

Commit 395a6e50450f4f0d065b89fe09074ba98db5056a on 892e50425f5a3e90ff8d3c8802fdfb31f8599518. Validated locally with Python 3.12, pinned Torch/TorchAO, the local transformers exporter branch and QAIRT SDK 2.48.

146 runner tests passed; 2 existing skips; 0 failures.

cmake --build /data/users/sidart/qnnflow/gh_prs/build-split-native-pr03 -j4 --target test_runner
ctest --test-dir /data/users/sidart/qnnflow/gh_prs/build-split-native-pr03 -R "^test_runner$" --output-on-failure -V

The two existing runner skips are TextDecoderRunnerTest.LoadUsesMethodName and StepWithAllModels.

No device/physical rpcmem, two-target FCB, real-model export or full QNN device-suite result is claimed. Buck could not start because the local launcher requires a missing .buck2. The base already contains #23563 and #23183; neither is a separate prerequisite commit here.

Related PRs (review in order)

1. Partitioner: run backend-contributed edge transform passes #24566 Partitioner: run backend-contributed edge transform passes
2. Export: avoid reused lifted-constant names #24567 Export: avoid reused lifted-constant names
3. LLM runner: select I/O contracts and accept custom buffer allocators #24568 LLM runner: select I/O contracts and accept custom buffer allocators
4. QNN: construct partitioners from options and allocate LLM I/O in rpcmem #24569 QNN: construct partitioners from options and allocate LLM I/O in rpcmem
5. QNN: bind delegate outputs in exported program order #24570 QNN: bind delegate outputs in exported program order
6. Export hf_static float-cache programs and register their Python runner #24571 Export hf_static float-cache programs and register their Python runner
7. LLM runner: validate and drive hf_static float-cache programs #24572 LLM runner: validate and drive hf_static float-cache programs
8. Document the version-1 hf_static I/O contract #24573 Document the version-1 hf_static I/O contract

Each PR includes the commits of the ones above it; review only its own top commit.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24568

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ You can merge normally! (1 Unrelated Failure), 1 Unclassified Failure

As of commit b294f72 with merge base 765e973 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• pull / test-arm-backend-docgen / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[winskuo-quic]

Hi @psiddh,
Since this is a draft PR, I am focusing on the high-level design.
My understanding of this PR is to introduce a KvCacheIOManager which accepts KVcache as inputs/outputs.

I would like to know what PTE (lowering flow) you are currently using for this runner?
I tried lower a XNNPACK model and it seems like the IO contract doesn't really match KvCacheIOManager.
The official documentation of HuggingFace states that there are actually a lot more inputs and outputs. Reference: https://github.com/huggingface/transformers/blob/e20e0787c9863f4b78781e603adfe6648415a854/docs/source/en/exporters.md?plain=1#L784

Difference 1: The Cache input/output actually consists of 3 things instead of 2: k_cache, v_cache, cumulative_length.
Difference 2: The HF flow will output 2 sets of Cache: regular output Cache and mutated Cache.

I would like to know whether these differences are taken into consideration or not. Personally, I think some of these HF I/Os are unnecessary and could cause memory spikes (or bloat the memory footprint) at runtime, which is problematic on resource-constrained edge devices.

It would be appreciated if you could share some of your thoughts on this. Thanks

[winskuo-quic]

I think the num_inputs consists of cumulative_length too when I tried the HF lowering flow with XNNPACK. If cumulative length is also one of the inputs maybe we should /3 instead?

[psiddh]

Thanks for trying it with XNNPACK. The version you reviewed was out of date, and I've updated it. KvCacheIOManager is gone. This PR now only adds the generic part: the runner picks an IOManager based on the contract the .pte declares, and backends can plug in a buffer allocator. The cache-as-I/O manager is in a later PR and targets programs exported with that contract, not the HF wrapper flow XNNPACK uses. The API is still under internal discussion.