Skip to content

QNN: construct partitioners from options and allocate LLM I/O in rpcmem - #24569

Draft
psiddh wants to merge 4 commits into
pytorch:mainfrom
psiddh:stack2/pr04-qnn-options-allocator
Draft

psiddh wants to merge 4 commits into
pytorch:mainfrom
psiddh:stack2/pr04-qnn-options-allocator

Conversation

@psiddh

@psiddh psiddh commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

QNN: construct partitioners from options and allocate LLM I/O in rpcmem

QnnPartitioner.from_options() routes an exporter's option dictionary to the partitioner or compiler-spec builder and rejects unknown names. The rpcmem allocator registers through the generic hook from PR 3, keeping LLM runner code independent of QNN. It forwards alignment, allocation, binding and free; a null allocation permits heap fallback.

Review option routing, llm/io_buffer_allocator.cpp, build/llama link wiring, then tests. Delegate output ordering remains separate in PR 5 for Qualcomm review.

Files

backends/qualcomm/CMakeLists.txt
backends/qualcomm/llm/BUCK
backends/qualcomm/llm/io_buffer_allocator.cpp
backends/qualcomm/llm/targets.bzl
backends/qualcomm/partition/qnn_partitioner.py
backends/qualcomm/tests/CMakeLists.txt
backends/qualcomm/tests/test_io_buffer_allocator.cpp
backends/qualcomm/tests/test_partitioner_options.py
backends/qualcomm/tests/test_passes.py
examples/models/llama/CMakeLists.txt
tools/cmake/executorch-config.cmake

Test Plan

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

['flock', '/data/users/sidart/qnnflow/gh_prs/model-export.lock', '/data/users/sidart/qnnflow/gh_prs/venv-rebase312/bin/python', '-c', 'import torch,pytest,sys;torch.ops.load_library("/data/users/sidart/qnnflow/gh_prs/build-rebase-native-exir-test/libexir_test.so");raise SystemExit(pytest.main(sys.argv[1:]))', '-q', '-rs', 'backends/qualcomm/tests/test_partitioner_options.py', 'exir/tests/test_partitioner_transform_passes.py', '--junitxml=/data/users/sidart/qnnflow/gh_prs/logs/split-pr04-python-1.xml']
13 passed in 5.16s
['flock', '/data/users/sidart/qnnflow/gh_prs/model-export.lock', '/data/users/sidart/qnnflow/gh_prs/venv-rebase312/bin/python', '-c', 'import torch,pytest,sys;torch.ops.load_library("/data/users/sidart/qnnflow/gh_prs/build-rebase-native-exir-test/libexir_test.so");raise SystemExit(pytest.main(sys.argv[1:]))', '-q', '-rs', 'backends/qualcomm/tests/test_passes.py', '-k', 'partitioner_from_options or partitioner_carries or partitioner_runs', '--junitxml=/data/users/sidart/qnnflow/gh_prs/logs/split-pr04-python-2.xml']
3 passed, 53 deselected in 3.78s

Real QNN backend builds; 2 allocator host tests passed; 0 skips/failures.

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

pytorch-bot Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure, 1 Unrelated Failure, 1 Unclassified Failure

As of commit 1a30de6 with merge base 765e973 (image):

NEW FAILURE - The following job has failed:

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

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

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

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 winskuo-quic left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @psiddh,
Thanks for the PR.
My understanding of this PR is that it is designed for shared buffer(zero-copy) where KV cache is IO instead of mutable buffer, so can I assume the mutable buffer KV Cache won't be supported in this case?

Another question is:
If I want to fallback last index_copy node to cpu, which is the node that updates new_kv_cache to past_kv_cache, will the current flow be able to handle this situation?

Run each partitioner’s edge transforms after caller-supplied passes so backends can use the standard lowering entry point. QNN contributes its existing transforms through the new hook.
Reserve names for constant tensors already present in the exported program before lifting new constants, so methods with folded tables retain distinct constants.
Choose the registered IOManager and its logits, method names and step sizes, rejecting unknown contracts and invalid metadata. Load and reset the selected methods; route prompt tails to decode with a side-effect-free capacity check. Preserve module and I/O-buffer lifetimes and expose a backend-neutral allocator hook.
Route backend options through QnnPartitioner.from_options and register an allocator that supplies QNN shared-memory buffers through the generic IOManager hook. The LLM runner remains backend-neutral.
@psiddh

psiddh commented Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Thanks. The version you reviewed was out of date. This PR is now just QnnPartitioner.from_options plus an allocator that puts the runner's I/O buffers in rpcmem.

  1. Right, it covers caches the runner owns as I/O. A mutable-buffer cache is allocated by the program, so this hook doesn't touch it. That part of the API is still under internal discussion.
  2. Yes. from_options forwards skip_node_op_set / skip_node_id_set to the partitioner, so you can keep that last index_copy on CPU.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants