Skip to content

Load AOTI plugins into hosts that don't share their C++ runtime - #24818

Open
mattiasbrand wants to merge 3 commits into
pytorch:mainfrom
mattiasbrand:export-D129654474
Open

mattiasbrand wants to merge 3 commits into
pytorch:mainfrom
mattiasbrand:export-D129654474

Conversation

@mattiasbrand

Copy link
Copy Markdown

Summary:
An AOTInductor plugin is compiled separately from the program that loads it, so the two don't have to share a C++ standard library. Two places in the CUDA backend assume they do.

  • Constants pass through a std::unordered_map. Cross-method weight sharing, mutable-state sessions and the CUDA graph host-constant check read a method's constants with AOTInductorModelContainerExtractConstantsMap, which hands the plugin a pointer to the host's map to fill in. That's only safe when both sides use the same standard library. It isn't safe for a Windows plugin cross-compiled with MinGW and loaded by an MSVC host, or a plugin loaded by a libc++ host. PyTorch's AOTI interface marks that call deprecated in favour of AOTInductorModelContainerExtractConstantsMapEntries, which returns a container-owned array of {name, handle} C structs. The backend now reads constants only through that call, so no C++ container crosses the plugin boundary.
  • The plugin's C++ runtime has to already be in the process. On Linux an x86 plugin references libstdc++ (std::cerr, exception handling, operator new) without listing libstdc++.so.6 as a dependency. Loaded from a process that doesn't link libstdc++, dlopen fails with undefined symbol: _ZSt4cerr. When the host does link libstdc++, the plugin loads only if the host's dynamic symbol table happens to include every libstdc++ symbol the plugin uses. load_library now opens libstdc++.so.6 with RTLD_GLOBAL once before the first plugin, and returns AccessFailed if it can't, since such a plugin can't load anyway. The preload adds libstdc++ to the global lookup scope after the host, so symbols the host already defines still win.

Plugins that don't export the entries call still run, with reduced features:

  • each method loads its own weights from the blob, without cross-method sharing;
  • mutable-state sessions report NotSupported;
  • CUDA graphs stay off, since their constants can't be checked for host memory.

Malformed results are never dereferenced: a null array with a nonzero count, a null name or handle, or an extraction error is rejected with InvalidProgram, or keeps graphs off.

test_cuda_mutable_state gains cases for the graph check (no entries call, no constants, a host constant, a device constant, malformed results) and for mutable-state extraction failures. It also builds outside fbcode as test_cuda_mutable_state_native.

Differential Revision: D129654474

@pytorch-bot

pytorch-bot Bot commented Nov 9, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 Cancelled Job

As of commit c7a73e6 with merge base dd38e81 (image):

CANCELLED JOB - The following job was cancelled. Please retry:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 9, 2026
@meta-codesync

meta-codesync Bot commented Nov 9, 2026

Copy link
Copy Markdown
Contributor

@mattiasbrand has exported this pull request. If you are a Meta employee, you can view the originating Diff in D129654474.

@github-actions

github-actions Bot commented Nov 9, 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.

Summary:
Pull Request resolved: pytorch#24643

Make the Buck CUDA runtime link its AOTI shims the way the CMake build already does.

A compiled AOTInductor shared library does not contain the tensor runtime. When the CUDA backend loads it, the library resolves its `aoti_torch_*` shim calls from the host process by name. Nothing in the host code calls most of these functions directly, so a static link without whole-archive drops them from `common_shims_slim` and loading fails with unresolved symbols. CMake already links `aoti_common_shims_slim` with `--whole-archive`; this sets `link_whole = True` on the Buck target to match, along with `supports_python_dlopen = True` like the other force-linked AOTI and CUDA targets. An executable that loads an AOTI library also has to export these symbols, as CMake does with `-Wl,--export-dynamic`. That includes a Buck `python_binary` that uses the Python bindings, because the extension links into its native executable.

`backends/aoti/tests:test_common_shims_slim_link` looks every shim up by name in a statically linked test binary that never calls them. It fails for every shim when `link_whole` is removed. `test_common_shims_slim` cannot catch this, because it calls the shims directly and that keeps them linked.

`backends/cuda/runtime/shims/tensor_attribute.cpp` defined `aoti_torch_get_device_type` and `aoti_torch_device_type_cuda` for ETensor, which the CUDA backend no longer uses. `common_shims_slim` defines the same `extern "C"` symbols for SlimTensor, so once both are force-linked they collide. CMake already leaves this file out of the CUDA shim sources. This deletes the file and its header and removes them from the Buck target.

`executorch_pybindings(use_cuda = True)` compiles code that calls the CUDA runtime directly but did not declare the CUDA dependency, so such a target could not link. The macro now adds it when `use_cuda` is set, together with `//executorch/extension/cuda:cuda_allocator`, which `pybindings_cuda.h` includes and calls. Previously that dependency was only met indirectly through the AOTI CUDA backend.

`//executorch/extension/pybindings:cuda_lib` builds these bindings with the AOTI CUDA backend. LibTorch defines the same `aoti_torch_*` symbols with a different tensor handle ABI, so it cannot link into the same binary. `extension/pybindings/test:test_pybindings_cuda_lib` exports a CUDA-delegated program with external weights, runs it through `cuda_lib` in a separate LibTorch-free binary, and compares the output with eager PyTorch. Without `link_whole`, or without `-Wl,--export-dynamic` on that binary, it fails with `undefined symbol: aoti_torch_dtype_float32`. `pytest.ini` and `pytest-windows.ini` leave it out of the pytest runs, because it needs a GPU and the Buck-built runner.

Differential Revision: D129530257
…ytorch#24817)

Summary:
Pull Request resolved: pytorch#24817

The CUDA backend's C++ runtime builds only in fbcode, and only with the LLM extension's KV cache linked in. So a C++ application that runs an exported vision or audio model on an NVIDIA GPU can't use it from any other Buck cell, cross-compiled aarch64 targets such as Jetson included. To run the same AOTInductor-compiled kernels it has to link LibTorch instead. XNNPACK, Vulkan and QNN already define their C++ targets in both cells; this does the same for CUDA.

- **Targets in both cells.** The `aoti/slim`, `extension/cuda` and `cuda/runtime` C++ targets are now defined outside fbcode too. There they carry `compatible_with = ["ovr_config//cuda:has_cuda"]`, so CPU-only builds skip them instead of failing to compile, and SlimTensor `storage` picks up its CUDA deps through a `select`. The Python AOTI targets stay fbcode-only.
- **Public headers.** Headers that dependents include (`slim_tensor.h`, `storage.h`, `runtime_api.h`, the shim and factory headers) move from `headers` to `exported_headers`, and `runtime_api`'s CUDA dependency becomes an exported external dep. Outside fbcode `headers` are private to the target, so nothing could compile against them. fbcode already exports both, so nothing changes there.
- **LLM cache optional, as in CMake.** CMake already builds `cuda_kv_cache.cpp` and `cuda_kv_pool.cpp` only with `EXECUTORCH_BUILD_EXTENSION_LLM`. Buck now does the same outside fbcode: the cache code in `cuda_backend.cpp` and `cuda_kv_cache.h` is behind that define, and the define is exported so that dependents see the same declarations as the backend. CMake now exports it too (`PUBLIC` instead of `PRIVATE` on `aoti_cuda_backend`), since the header's declarations depend on it; `test_cuda_kv_cache` and `cuda_batching` include that header. fbcode keeps the cache. A backend built without it still refuses a program whose KV cache lives off-graph, because `requires_offgraph_kv_storage` reads that from the program's constants and sits outside the gate.
- **Build flags.** The `nvcc` architecture flags stay fbcode-only. Outside fbcode the shims undefine `CPU_CAPABILITY_AVX2` and `CPU_CAPABILITY_AVX512` for CUDA sources: c10 exports them on x86, and `Half.h` then pulls host AVX2 `vec_half` code into nvcc compiles. `cuda-lazy` maps to `libcudart`, and the shims link `curand` and `thrust` for `rand.cu` and `sort.cu`.

fbcode builds keep the same sources and flags, and the CMake build is untouched.

Differential Revision: D129654471
Summary:
An AOTInductor plugin is compiled separately from the program that loads it, so the two don't have to share a C++ standard library. Two places in the CUDA backend assume they do.

- **Constants pass through a `std::unordered_map`.** Cross-method weight sharing, mutable-state sessions and the CUDA graph host-constant check read a method's constants with `AOTInductorModelContainerExtractConstantsMap`, which hands the plugin a pointer to the host's map to fill in. That's only safe when both sides use the same standard library. It isn't safe for a Windows plugin cross-compiled with MinGW and loaded by an MSVC host, or a plugin loaded by a libc++ host. PyTorch's AOTI interface marks that call deprecated in favour of `AOTInductorModelContainerExtractConstantsMapEntries`, which returns a container-owned array of `{name, handle}` C structs. The backend now reads constants only through that call, so no C++ container crosses the plugin boundary.
- **The plugin's C++ runtime has to already be in the process.** On Linux an x86 plugin references libstdc++ (`std::cerr`, exception handling, `operator new`) without listing `libstdc++.so.6` as a dependency. Loaded from a process that doesn't link libstdc++, `dlopen` fails with `undefined symbol: _ZSt4cerr`. When the host does link libstdc++, the plugin loads only if the host's dynamic symbol table happens to include every libstdc++ symbol the plugin uses. `load_library` now opens `libstdc++.so.6` with `RTLD_GLOBAL` once before the first plugin, and returns `AccessFailed` if it can't, since such a plugin can't load anyway. The preload adds libstdc++ to the global lookup scope after the host, so symbols the host already defines still win.

Plugins that don't export the entries call still run, with reduced features:
- each method loads its own weights from the blob, without cross-method sharing;
- mutable-state sessions report `NotSupported`;
- CUDA graphs stay off, since their constants can't be checked for host memory.

Malformed results are never dereferenced: a null array with a nonzero count, a null name or handle, or an extraction error is rejected with `InvalidProgram`, or keeps graphs off.

`test_cuda_mutable_state` gains cases for the graph check (no entries call, no constants, a host constant, a device constant, malformed results) and for mutable-state extraction failures. It also builds outside fbcode as `test_cuda_mutable_state_native`.

Differential Revision: D129654474
@shoumikhin

Copy link
Copy Markdown
Contributor

@claude review effort=max

@claude

claude Bot commented Nov 10, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 58m 32s —— View job


PR Review: #24818

Request Changes: a plugin without AOTInductorModelContainerExtractConstantsMapEntries now silently loses cross-method weight sharing, mutable-state sessions and CUDA graphs, nothing shows the pinned PyTorch emits that symbol, and no test calls it.

Detailed review

Summary

This review covers all three commits in the stack (a42c487, 2b4739c, c7a73e6). Together they:

  • read AOTI constants through a flat C ABI instead of a std::unordered_map*;
  • preload libstdc++.so.6 with RTLD_GLOBAL on Linux;
  • define the CUDA C++ targets in the xplat cell;
  • add a cuda_lib pybinding with a GPU end-to-end test.

The blocking problem is that plugins without the new symbol now lose features without any warning. The other problems:

  • a test target is built without CUDA_AVAILABLE;
  • some Buck targets are missing settings that CMake or the commit messages say they have;
  • no test calls the new extraction path;
  • some code is duplicated or dead.

All 34 files were reviewed. Nothing was built or run. I couldn't read upstream PyTorch's interface.h from this runner, so what the entries API returns is inferred from how this PR uses it.

Code Quality

  • A failed libstdc++.so.6 preload is logged only once (backends/cuda/runtime/platform/platform.cpp:82-91).
    • The result is cached in a function-local static, so only the first load_library call logs dlerror().
    • Every later call returns AccessFailed with no message, and the caller (cuda_backend.cpp:547-550) logs nothing either. So a second method's init fails with nothing in the log.
    • Fix this →
  • The off-graph KV error message is now less useful (cuda_backend.cpp:599-604).
    • Every build now prints "requires an installed off-graph KV cache and a CUDA backend built with EXECUTORCH_BUILD_EXTENSION_LLM".
    • In a build that already has the LLM extension, that points at a flag that's already set. It also drops the kCacheKeyOption backend option that the old message told the caller to set. Pick the message based on the build.
    • Builds without the LLM extension now ignore a supplied cache key instead of returning NotSupported. That's harmless for programs without off-graph storage, but the description doesn't mention it.
    • Fix this →
  • The same extract-and-validate loop is written three times: cuda_backend.cpp:1604-1630, cuda_mutable_state.cpp:287-308 and cuda_delegate_handle.h:367-385.
    • Each copy has its own error messages and handles a null array its own way.
    • One helper that returns a validated name → AtenTensorHandle map (or an error) would serve all three. It would also keep the SlimTensor cast out of the exported header cuda_delegate_handle.h.
    • Fix this →
  • Dead code (Fix this →):
    • Unused member. AOTIDelegateHandle::extract_constants_map (aoti_delegate_handle.h:156) is still loaded at cuda_backend.cpp:303-304, but no production code reads it. Only the tests set it, to check that it isn't called.
    • Unreachable branch. cuda_backend.cpp:1572-1578 can never run, because line 1506 already takes the fallback when the symbol is null. Its Error log ("requires C-ABI constants extraction") also contradicts that silent fallback.
    • Redundant null checks. The new checks at cuda_backend.cpp:236-238 and 1500-1502, and at cuda_mutable_state.cpp:258-264, test values that can't be null there:
      • handle comes from new at line 556;
      • named_data_map is already dereferenced at line 487;
      • build_descriptors only runs after handle_has_symbols succeeds (line 604).

Infrastructure

  • test_cuda_mutable_state_native is built without -DCUDA_AVAILABLE=1 (backends/cuda/runtime/targets.bzl:180-191).
    • Every fbcode CUDA test in the same file sets it (lines 211, 232, 253, 272, 290), and so does CMake's test_cuda_mutable_state (backends/cuda/CMakeLists.txt:490).
    • cuda_backend sets the define only for its own sources. The test file therefore compiles the stub DeviceTraits<CUDA> from slim/core/storage.h:204-233, which aborts, while the backend sources in the same binary compile the real one.
    • That breaks the one-definition rule. The linker keeps only one copy of each inline member, so in unoptimized builds the backend's CUDA storage calls can end up in the stub (ET_CHECK_MSG(false, "Build with CUDA_AVAILABLE=1 …")) on GPU runners.
    • Fix this →
  • The rules for using the PUBLIC cuda_lib target appear only in a commit message.
    • a42c487 says two things:
      • any executable that loads a plugin needs -Wl,--export-dynamic;
      • LibTorch's aoti_torch_* definitions mean cuda_lib can't link into the same binary.
    • CMake passes the export flag on to every consumer automatically (aoti_common PUBLIC, backends/aoti/CMakeLists.txt:53-55).
    • In Buck, only run_cuda_program sets it (extension/pybindings/test/BUCK:103-104). Nothing on cuda_lib (extension/pybindings/BUCK:74-82) or common_shims_slim enforces or mentions either rule. The next binary built like run_cuda_program will fail with undefined symbol: aoti_torch_*.
    • Consider adding exported_linker_flags to //executorch/backends/aoti:common_shims_slim, or a dynamic list limited to aoti_torch_*. Also put a comment on cuda_lib about LibTorch.
  • The xplat CUDA targets lack compatible_with: extension/cuda/targets.bzl, backends/aoti/slim/cuda/targets.bzl and backends/aoti/slim/c10/cuda/targets.bzl.
    • Their BUCK files now call define_common_targets() in both cells, but none of these targets sets compatible_with = ["ovr_config//cuda:has_cuda"].
    • The 2b4739c message says they have it, and the cuda/runtime targets do get it, through cuda_kwargs.
    • The comment this PR removes from extension/cuda/BUCK warned that in xplat caller_stream "cannot compile ("cuda_runtime.h not found")".
    • Unless the xplat target behind cuda-lazy is itself limited to has_cuda (I couldn't check), CPU-only wildcard builds will try to compile caller_stream.cpp, current_stream.cpp, device_guard.cpp, cuda_allocator.cpp and guard.cpp.
    • Fix this →

Testing

  • No test runs the new extraction path in load_constants_with_cache (cuda_backend.cpp:1504-1630).
    • Neither the fallback for a missing symbol nor the InvalidProgram rejections at lines 1614-1627 is tested.
    • The fake-handle tests cover only describe_host_constants_for_graph and mutable state.
  • The end-to-end test never calls the new symbol.
    • test_pybindings_cuda_lib.py exports one method with no backend options. So weight_sharing_across_methods_ stays false (cuda_backend.cpp:1367), init takes load_constants_legacy, and CUDA graphs and mutable state stay off.
    • The test would still pass if the pinned PyTorch didn't emit AOTInductorModelContainerExtractConstantsMapEntries.
    • Add a two-method case with weight sharing on, or at least check that a freshly exported AOTI .so exports the symbol, so a PyTorch pin without it fails CI.
    • Fix this →

API Design

  • The new C ABI is described in a single line (backends/aoti/aoti_delegate_handle.h:117): "Extracts borrowed entries through a flat C ABI; the container owns the array." The code relies on two things that line doesn't say:

    • entries[i].name is the original FQN.
      • cuda_backend.cpp:1637 and cuda_mutable_state.cpp:313 look entries up by FQN. But the same AOTInductorConstantMapEntry struct carries internal names into UpdateUserManagedConstantBufferPairs (cuda_backend.cpp:1696).
      • If the extraction returned internal names, every lookup at line 1637 would miss and hit the continue at line 1641. Sharing would then switch off without any error.
    • How long the array stays valid: until the next extract, a constants update, or deletion of the container. The comment also doesn't say whether folded constants are included.

    Document both, and cite the upstream PyTorch declaration this mirrors. My confidence here is medium, since I couldn't read upstream. Fix this →

Backward Compatibility

  • Plugins without AOTInductorModelContainerExtractConstantsMapEntries lose three features, with no warning. The description calls these "reduced features", but at runtime nothing reports them. No code reads the old map API any more. So a plugin compiled by a PyTorch that doesn't emit the new symbol now behaves as follows (this includes the plugin inside any .pte exported before the symbol existed):

    • Weights are no longer shared, and nothing is logged.
      • cuda_backend.cpp:1504-1510 falls back to load_constants_legacy, so every method uploads its own copy of the weights.
      • Weight sharing is turned on by qwen35_moe_engine.cpp:125, gemma4_31b_engine.cpp:143 and muse_glimmer_engine.cpp:319, and by cuda_executor.h:69, where it defaults to true.
      • For those multi-method models this puts a second copy of the weights in device memory. That is likely to run out of memory, with nothing pointing at the cause.
    • Mutable-state sessions fail with NotSupported, and nothing is logged (cuda_mutable_state.cpp:103-107, 227-229).
    • CUDA graphs are off. The only report is an Info log saying "its constants could not be listed" (cuda_delegate_handle.h:364-366, cuda_backend.cpp:643-648).

    This also applies on libstdc++ hosts loading libstdc++ plugins, where the old map API was safe.

    Neither torch_pin.py nor .ci/docker/ci_commit_pins/pytorch.txt changes, and nothing in the repo shows which PyTorch version emits the symbol. Please:

    • confirm that the pinned nightly emits it, and bump the pin if it doesn't;
    • log a Warning at each fallback, naming the missing symbol and the first PyTorch version that has it.

    If older artifacts matter, the map path could also stay as the fallback when __GLIBCXX__ is defined.

    Fix this →

Specific Comments

  • backends/cuda/runtime/cuda_delegate_handle.h:360: "Compiled plugins return flat entries so C++ containers stay host-owned" describes the extraction ABI, not this function. It belongs on the typedef.
  • backends/aoti/aoti_delegate_handle.h:157-158: extract_constants_map_entries{nullptr} is the only member with a default initializer. Give all the function pointers one, or none.
  • backends/cuda/runtime/cuda_kv_cache.h:196-199: attach_offgraph_kv_cache is declared outside the #if at line 201, although its comment says it's defined only with the LLM extension. Moving it inside turns a link error into a compile error.
  • backends/cuda/runtime/platform/platform.cpp:80-81:
    • The comment doesn't say why the runtime is loaded RTLD_GLOBAL rather than left to the plugin.
    • Line 85 should name libstdc++.so.6.
    • The description talks about x86 plugins, but the guard is __linux__, so it also runs on aarch64 (Jetson). Say whether those plugins need it too.
  • backends/cuda/runtime/test/test_cuda_mutable_state.cpp:224, 334, 348: the test names call the same API both "SafeExtraction" and "EntryExtraction".
  • backends/aoti/tests/test_common_shims_slim_link.cpp:17-46: the hardcoded list of 28 shims will drift from common_shims_slim.h. Either say it's a sample or point at the header.
  • backends/cuda/runtime/targets.bzl:44: "use scalar Half conversions, not host SIMD" leaves out the real reason given in 2b4739c. c10 exports CPU_CAPABILITY_AVX2/AVX512 on x86, and Half.h then pulls host AVX2 code into nvcc compiles.
  • backends/cuda/CMakeLists.txt:422-424 and backends/cuda/runtime/targets.bzl:133-135: these export EXECUTORCH_BUILD_EXTENSION_LLM, which is a build-option name, to every dependent. Any #if on it that a dependent adds later will flip too. A macro specific to this backend would avoid that.
  • Description and commit messages:
    • c7a73e6 says an extraction error is rejected with InvalidProgram. But both call sites return the extractor's own error, and InvalidEntryExtractionPreventsSessions expects Error::Internal (test_cuda_mutable_state.cpp:392-393).
    • 2b4739c says "the CMake build is untouched" right after it describes the PRIVATE → PUBLIC change.
    • 2b4739c also says "fbcode builds keep the same … flags". Yet fbcode dependents of cuda_backend now inherit -DEXECUTORCH_BUILD_EXTENSION_LLM through exported_preprocessor_flags.

Recommendation

Request Changes.

  • Blocking (Backward Compatibility): confirm the pinned PyTorch emits AOTInductorModelContainerExtractConstantsMapEntries, and log a warning at each of the three fallbacks. If the pin already emits it, only the warnings remain.
  • Should land with it: the one-line CUDA_AVAILABLE fix for test_cuda_mutable_state_native, and a test that calls the entries path.
  • Not blocking: everything else.

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. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants