Skip to content

Add skip_h2d_for_constants to keep constants in device memory - #24144

Merged
shoumikhin merged 4 commits into
pytorch:mainfrom
shoumikhin:device-resident-constants
Oct 28, 2026
Merged

shoumikhin merged 4 commits into
pytorch:mainfrom
shoumikhin:device-resident-constants

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

What is wrong today

With skip_h2d_for_method_inputs, the inputs of a device delegate arrive in device memory. Constants (weights and read-only buffers) do not. Their data is stored on the host, so a constant that a delegate reads from outside is copied to the device on every call. The CUDA partitioner already moves constants into its delegate, so this hits constants left outside, for example with a custom partitioner.

What this change does

A new opt-in option keeps such a constant in device memory. The runtime copies it there once, when the method loads.

ExecutorchBackendConfig(
    propagate_device_config=PropagateDeviceConfig(skip_h2d_for_constants=True),
)
  • PropagateDevicePass tags a constant that only delegates on one device read with that device.
  • Memory planning gives it device memory for the whole method. The emitter stores its data as the initial value of that memory.
  • The tensor parser loads such data on the host and copies it with the registered DeviceAllocator. New DeviceAllocator::synchronize() waits for the copy, so the host buffer can be freed and the first run reads the right data. On CUDA it waits for the stream the copy used, on the right device.
  • An empty tensor with initial data no longer fails to load, also in host memory.

Limits:

  • Each method that reads the constant keeps its own device copy while it is loaded.
  • A constant that anything else also reads, or that is changed on the host first (a view, a transpose), is still copied on every call.
  • share_mutable_buffers does not work with it.

Cost: the loading code is in every runtime build. In CI, the stripped Arm bare-metal size_test grows by 560 bytes, to 117884 (limit 118000). On x86 its code grows by 727 bytes with GCC and 607 with Clang, and the stripped file size does not change.

What was tested

  • test_propagate_device_pass.py: only such constants skip the copy and get device memory; warnings and errors.
  • tensor_parser_device_test.cpp: such a constant loads from the program and from a .ptd file into device 1 of a mock CUDA device, and is waited for.
  • test_emit.py: a method whose initial data is empty loads.
  • test_cuda_allocator.cpp: synchronize() waits for the right stream, also with another GPU current (needs two GPUs; ran on H100s).
  • test_device_placement.py: such a program runs twice on an H100 and matches eager PyTorch.

Each new test fails without its change, and each commit passes on its own.

@shoumikhin shoumikhin added the release notes: exir Changes to any dialects and passes on these dialects, such as memory planning label Oct 26, 2026
Copilot AI balanced review requested due to automatic review settings October 26, 2026 02:21
@pytorch-bot

pytorch-bot Bot commented Oct 26, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 7 Unclassified Failures

As of commit 0cf1823 with merge base 0549e1e (image):

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

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 Oct 26, 2026

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from fb3f67e to f1833e1 Compare October 26, 2026 02:31
Copilot AI balanced review requested due to automatic review settings October 26, 2026 02:31

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude review

@claude

claude Bot commented Oct 26, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 4m 51s —— View job


The approach is sound, but please fix a few things before merging: skip_h2d_for_method_inputs now quietly changes where constants are stored and how much device memory they use, nothing stops a weight from being stored twice in the .pte, and the CUDA test passes even if the feature stops working.

Detailed review

PR Review: #24144

I reviewed only the last two commits (cf90e0a, f1833e1), as the description asks. Three independent reviewers went through the changes, and I checked each finding below against the code.

Summary

With skip_h2d_for_method_inputs, a constant that only delegates on one device read now lives in planned device memory and is loaded there once, when the method loads. Problems found:

  • skip_h2d_for_method_inputs now also changes how constants are stored, and there is no way to turn that off.
  • A weight shared by two methods can be stored twice in the .pte.
  • If the runner's allocator has no device info, the runtime silently treats a device buffer as host memory.
  • The CUDA end-to-end test can't tell whether the feature is actually on.

API Design

  • exir/passes/propagate_device_pass.py:406-421: skip_h2d_for_method_inputs now does more than its name says, and can't be opted out of.
    • The flag's contract is "the caller passes method inputs on the device". Existing users of the flag now also get every qualifying constant moved into persistent planned device memory. Its lifetime ends at max_node_idx (memory_planning.py:354), and every method gets its own copy.
    • For prefill/decode-style programs with large weights that stay outside the delegate, this raises peak device memory. It is spread across all methods, and the old short-lived per-call copies could reuse their memory.
    • The constant's data also moves from the constant segment to the mutable segment, and the warning that used to flag the case is gone.
    • Suggest a separate PropagateDeviceConfig field, for example keep_constants_on_device, so memory-bound users can keep the per-call copy. If not, document the memory cost in cuda-overview.md.

Code Quality

  • exir/emit/_emitter.py:483-486 vs :505: a weight shared across methods can be stored twice in the .pte.
    • A device-resident constant has allocation_info, so it is looked up only in cached_spec_mutable_hash_values and saved to mutable_buffer.
    • The same bytes read on the host in another method are looked up only in cached_spec_hash_values and saved to constant_buffer.
    • Trigger: one weight that method A reads only from a CUDA delegate and method B reads from a host op. The program grows by the size of that weight. External .ptd constants are not affected, because they go through external_constant_hash.
  • runtime/executor/tensor_parser_exec_aten.cpp:201-209: getPlannedBufferDevice falls back to CPU without saying so.
    • When planned_buffer_devices() is empty (HierarchicalAllocator built with the spans-only constructor), it returns CPU even if the program plans that memory_id on CUDA.
    • Trigger: a custom runner that supplies a device arena through the legacy constructor. load_mutable_subsegment_into or load_data_into then writes device memory with the CPU and crashes, with no diagnostic.
    • This is a new failure mode. Before this PR, nothing wrote initial data into device-planned memory. Module is safe, because it passes devices and rejects share_memory_arenas with device buffers.
    • Suggest checking the device against MethodMeta::memory_planned_buffer_device and returning an error on a mismatch.
  • tensor_parser_exec_aten.cpp:212-241: copyToDevice frees the host staging buffer as soon as copy_host_to_device returns.
    • With a caller stream set, CudaAllocator::copy_impl issues cudaMemcpyAsync for H2D and does not synchronize (extension/cuda/cuda_allocator.cpp:341-343).
    • This is safe today only because loaders return pageable memory, which CUDA stages before the call returns. A loader or NamedDataMap that returns pinned memory, or another DeviceAllocator with truly async H2D, would read freed memory.
    • The contract of DeviceAllocator::copy_host_to_device (runtime/core/device_allocator.h:64-77) doesn't say when src may be released. Either document that H2D must finish reading src before it returns, or synchronize here.
  • exir/emit/_emitter.py:2046-2051: the new exclusion may hide cross-method mutability. (Low confidence.)
    • The not device_resident_constant check also bypasses the rule that a buffer planned elsewhere must be treated as mutable in every entry point.
    • Trigger: a buffer that one method mutates, while another method only reads it from CUDA delegates. The reading method would freeze its own copy instead of seeing the shared state.
    • This is masked today because share_mutable_buffers already conflicts with device planning. An assert, or a comment saying the case can't happen, would make it explicit.

Testing

  • backends/cuda/tests/test_device_placement.py:149-177: the test only compares outputs.
    • If PropagateDevicePass stopped tagging table and fell back to the per-call H2D copy, the outputs would match and the test would still pass.
    • Set error_on_per_call_copies=True (added earlier in this stack), or assert that no warning is logged, as the exir test does.
    • Nothing exercises the .ptd (external_constants) path on real CUDA memory; only the mock-allocator C++ test does.
  • No tests for these cases in exir/tests/test_propagate_device_pass.py:
    • An nn.Parameter or a lifted tensor constant. Only a register_buffer constant is covered.
    • A multi-method program sharing one weight, which would catch the duplicate-storage issue above.
    • The emitter case that the device_resident_constant exclusion changes.
  • runtime/executor/test/CMakeLists.txt:96-99: method_meta_test_pte_files still DEPENDS only on ModuleAddWithDevice.pte. If one of the three new outputs is deleted, the files are not regenerated and tensor_parser_device_test fails. List all four outputs. The target name is also misleading now that another test depends on it.

Performance

  • tensor_parser_exec_aten.cpp:307-314, 340-349: a full host copy of every device-resident constant at method load.
    • get_data and load_mutable_subsegment allocate a host buffer the size of the tensor before the H2D copy, with no chunking.
    • For large weights this adds peak host memory and churn at load time. Worth mentioning in the docs next to the per-method device copy.

Recommendation

Request Changes: the main blockers are the missing opt-out (or documented cost) for keeping constants on the device, the duplicate storage of a shared weight, and making the CUDA test actually check the feature.

@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from f1833e1 to fc3ce9c Compare October 26, 2026 05:43
Copilot AI balanced review requested due to automatic review settings October 26, 2026 05:43

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin shoumikhin changed the title Keep constants in device memory with skip_h2d_for_method_inputs Add skip_h2d_for_constants to keep constants in device memory Oct 26, 2026
@shoumikhin

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Changed in the new push:

  • Keeping constants on the device is now opt-in, with the new PropagateDeviceConfig(skip_h2d_for_constants=True), so skip_h2d_for_method_inputs alone behaves as before. The CUDA docs list the cost: device memory for as long as the method is loaded, one copy per method, the load through host memory, and the duplicate storage below.
  • The CUDA test sets error_on_per_call_copies=True, so it fails if the constant is copied. The Python test also covers a parameter and the case without the option. The CMake target now depends on all four generated files.
  • copyToDevice now says why freeing the host buffer right away is safe.

Not changed, and why:

  • A weight kept on the device in one method and read on the host in another is stored twice when its data is in the program. Sharing it there would need the runtime to fill planned memory from the constant segment, which the program format cannot express. The docs say so.
  • An allocator without device metadata describes host buffers, which is also what planning without device memory gives. An error there would break mutable buffers that carry a device tag in that mode. Module and the Python bindings pass the metadata.
  • With share_mutable_buffers, memory planning rejects a buffer that already has planned memory, so a constant kept on the device cannot be shared with a method that mutates it.
  • The .ptd path on real CUDA uses the same copy as the in-program path, which the CUDA test runs. The parse path itself is covered by the mock test.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude review

@claude

claude Bot commented Oct 26, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 3m 52s —— View job


Before this can merge, please delete the exir/tests/test_propagate_device_pass.py.orig file that was committed by mistake. The other findings are minor: one is about which readers make a constant fall back to a per-call copy, one asks for a per-method form of the new option, and the rest are gaps in test coverage.

Detailed review

PR Review: #24144

Summary

This PR adds the opt-in option skip_h2d_for_constants. A constant that only delegates on one device read is then kept in that device's memory and loaded there once, when the method loads. I reviewed only the last two commits (4089b91, fc3ce9c). Three independent reviewers went through the diff, and I checked each finding below against the code.

Infrastructure

  • exir/tests/test_propagate_device_pass.py.orig: a 2472-line backup file from a patch run was committed in fc3ce9c. All three reviewers flagged it.
    • It is an older version of the test module that nothing uses, and pytest does not collect it.
    • It accounts for about half of the PR's added lines (+5008), and it will go stale as the real test file changes.
    • Delete it from the commit. Fix this →

API Design

  • exir/passes/propagate_device_config.py:87: skip_h2d_for_constants accepts only a bool. The flags next to it, skip_h2d_for_method_inputs (:66) and skip_d2h_for_method_outputs (:79), accept Union[bool, Dict[str, bool]].
    • The cost that makes this option opt-in applies to each method separately: each method that reads the constant keeps its own copy in device memory.
    • Example: a program with prefill and decode that share a large weight, where only decode should keep it on the device. That cannot be expressed today.
    • Making it per-method later would mean widening a public field's type, so it is easier to match the other flags now.

Code Quality

  • exir/passes/propagate_device_pass.py:427-431: an op that only reads a constant's shape stops it from staying on the device. (Medium confidence; I have not run a repro.)
    • The all(...) check requires every user of the constant to be a delegate on the target device, so a host op that only reads metadata counts as a reader.
    • The input path in the same file, _place_input_on_device (:524-530), leaves out _METADATA_READERS when it collects the other readers.
    • Trigger: a constant read by a cuda:0 delegate and by a host aten.sym_size or size node. It falls back to a copy on every call, or raises under error_on_per_call_copies, even though no host op reads its data.
    • Fix: apply the same user.target not in _METADATA_READERS filter here.

Testing

  • exir/tests/test_propagate_device_pass.py: three cases are not covered.
    • Two delegates on the same device read one constant. The constant should stay on the device. In this case the loop at propagate_device_pass.py:423 reaches a constant it has already tagged, and no test runs that path. Using PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:0"]) with a model that has two readers would cover it.
    • skip_h2d_for_constants=True set without skip_h2d_for_method_inputs. Every new test (_skip_h2d_for_constants_config), the CUDA test, and export_program_with_device_info.py set both. Two things are untested in that setting: method inputs should still get H2D copies, and a constant that is still copied should raise under error_on_per_call_copies. The second relies on the or at :449.
    • The zero-byte early return in copyToDevice (runtime/executor/tensor_parser_exec_aten.cpp) has no test.

Recommendation

Request Changes: the stray .orig file has to be removed. The other findings are minor and do not block the merge.

@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from fc3ce9c to cd8691c Compare October 26, 2026 05:57
Copilot AI balanced review requested due to automatic review settings October 26, 2026 05:57

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Thanks. Changed in the new push: the stray .orig file is gone, and new tests cover a constant read by two delegates on one device, and skip_h2d_for_constants without skip_h2d_for_method_inputs, where the input is still copied and a copied constant raises with error_on_per_call_copies.

Not changed, and why:

  • Per method: ExecutorchBackendConfig(propagate_device_config={"decode": PropagateDeviceConfig(skip_h2d_for_constants=True)}) already turns it on for one method only.
  • Shape-only readers of a constant: a constant has a fixed shape, so export folds its sizes into numbers and emits no shape ops on it.
  • The zero-byte early return has no test, since it would need an exported model with an empty constant just for that line.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude review

@claude

claude Bot commented Oct 26, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 4m 17s —— View job


I found no blockers in the last two commits; what is left is one point about the DeviceAllocator contract, one header cleanup, and three gaps in test coverage.

Detailed review

PR Review: #24144

Summary

This PR adds the opt-in option PropagateDeviceConfig(skip_h2d_for_constants=True). With it, a constant that only delegates on one device read stays in that device's planned memory and is loaded there once, when the method loads. I reviewed only 4089b91 and cd8691c. Three independent reviewers went through the diff, and I re-read the code behind each finding below. None of the remaining issues blocks the merge.

API Design

  • runtime/executor/tensor_parser_exec_aten.cpp:239-243: the rule that makes freeing data safe is written only at this call site.
    • The comment relies on CUDA finishing reading pageable memory before cudaMemcpyAsync returns.
    • DeviceAllocator::copy_host_to_device (runtime/core/device_allocator.h:64-77) says nothing about when src may be released.
    • Trigger: an out-of-tree DeviceAllocator that queues an async DMA, or a DataLoader/NamedDataMap that returns pinned memory under CUDA. Either one would read the FreeableBuffer after it is freed or unmapped.
    • Fix: add one line to the copy_host_to_device doc, for example "src may be released as soon as this returns". That makes it a requirement on every allocator, not an assumption made in one place.

Code Quality

  • runtime/executor/program.h:301-312: MutableSubsegment and locate_mutable_subsegment() are declared in the friend-visible section.
    • Only load_mutable_subsegment_into() and load_mutable_subsegment() use them.
    • Because they sit next to those two, Method, TensorParser, BackendDelegate, and the other friends can call them too.
    • Move them to the second private: block (:314).
    • Their doc comment also repeats the one on load_mutable_subsegment.

Testing

  • runtime/executor/test/tensor_parser_device_test.cpp: none of the new error returns in copyToDevice is tested.
    • No allocator registered for the device type (NotFound).
    • A loader or data map that returns fewer than nbytes bytes (InvalidProgram).
    • A load error that is passed through.
    • The mock CUDA fixture can already reach the first two: leave the allocator unregistered, or use a NamedDataMap that returns a short buffer.
  • exir/tests/test_propagate_device_pass.py: two kinds of constant have no test.
    • A constant read by delegates on two different devices, for example PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]). The all(...) check at propagate_device_pass.py:423-431 should fail for both delegates, so the constant should fall back to the per-call copy and raise under error_on_per_call_copies. No test pins this down.
    • A lifted tensor constant (a plain tensor attribute that is neither a parameter nor a buffer). The new tests cover register_buffer and nn.Parameter. A lifted constant reaches the emitter through inputs_to_lifted_tensor_constants, and the new allocation_info / mutable-segment path has not been run for it.

Recommendation

Approve: none of the findings blocks the merge. The copy_host_to_device doc line is the one I'd most like to see added before more device allocators come in.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Thanks. Answers, with no code change:

  • copy_host_to_device: CudaAllocator copies asynchronously on a caller stream, so promising in the interface that src can be freed at return would not hold for pinned memory there. The comment states the condition this call relies on: loader buffers are pageable.
  • MutableSubsegment and locate_mutable_subsegment() sit next to the two loaders that use them, in the section where the friend classes can already call every private loader.
  • The new error returns in copyToDevice pass on the registry's or the loader's error. I did not add tests that fake those failures.
  • A constant read by delegates on two devices is the "two devices" case of test_skip_h2d_for_constants_copies_constants_read_elsewhere, which expects one copy per device.
  • A plain tensor attribute (a lifted constant) takes the same path. I checked it end to end on an H100 with error_on_per_call_copies=True: it stays in CUDA memory, is stored as initial state, and matches eager PyTorch over two calls.

it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens
during device-aware memory planning.

**With `skip_h2d_for_method_inputs`, an input that a device delegate reads stays on that device.**

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.

can you keep this .md file more succinct and readable by humans

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Shortened, see the new version. On top of the shorter text from #24088, this PR now adds only one short paragraph on skip_h2d_for_constants, what it does and what it costs, and a pointer to it from the list of copies that stay.

@mergennachin

Copy link
Copy Markdown
Contributor

it's hard to review stacked PR. not clear what is in the previous and what is new

Can you use ghstack?

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude review

About the last review: I kept load_segment_into(). Calling copy_from_host(loader_->load(...)) from Program moves the host buffer and its cleanup into code that every build carries, and that measured 160 bytes more in the minimal runtime. Its doc now says why it exists. The docs point is fixed in the last commit.

@claude

claude Bot commented Oct 27, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 4m 55s —— View job


Nothing in this PR's own six commits (a8592ed..1f367f7) blocks the merge; the only thing left is one stale sentence in the CUDA docs that still talks about "two flags".

Detailed review

PR Review: #24144

Summary

This PR adds the opt-in PropagateDeviceConfig(skip_h2d_for_constants=True). With it, a constant that only delegates on one device read is planned in that device's memory and loaded there once, at method load, through new DeviceAllocator hooks.

I reviewed 964a611..1f367f7. Three independent reviewers went through the diff, and I re-read the code behind each finding.

I accept the author's reason for keeping load_segment_into(), so I am not raising it again. I also left out points settled in earlier rounds.

I checked and dropped these findings:

  • The ATen variant of program_no_prim_ops now depends on the non-ATen //executorch/runtime/core:device_allocator. etdump_aten already does the same (devtools/etdump/targets.bzl:135), so this follows an existing pattern.
  • The new registry lookup in getPlannedBufferAllocator could add size to the minimal runtime. It is already part of the 1,632 bytes measured in the description.
  • The emitter's SEGMENT + allocation_info branch could look in the wrong cache. I found no pass that puts SEGMENT info on a non-mutable constant, so nothing can reach this branch today.
  • export_program_with_device_info.py renames the .ptd through _tensor_data. This was settled earlier: it follows test/models/export_program.py.
  • The test uses std::optional without #include <optional>. It gets the header through caller_stream.h, and the missing include is only a style nit.

API Design

  • docs/source/backends/cuda/cuda-overview.md:164: "The two flags are independent." is now stale.
    • Three lines below, the doc says "Every skip flag…", and the third flag, skip_h2d_for_constants, is introduced at line 228.
    • Someone reading this section will think there are only two flags, or will wonder whether the constant flag depends on the other two. The paragraph at line 228 then has to say that it doesn't.
    • Fix: name the two flags here ("skip_h2d_for_method_inputs and skip_d2h_for_method_outputs are independent"). Optionally, move the last sentence of the skip_h2d_for_constants paragraph ("This flag works on its own…") up next to it, so all the facts about which flag depends on which are in one place. Fix this →

Recommendation

Approve: the one finding is a wording fix in the docs and does not block the merge.

@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from 1f367f7 to 7f587a9 Compare October 27, 2026 00:38
Copilot AI balanced review requested due to automatic review settings October 27, 2026 00:38

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude review

@claude

claude Bot commented Oct 27, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 6m 21s —— View job


I found nothing to fix in this PR's own six commits (a8592ed..7f587a9), and the stale "two flags" sentence in the CUDA docs from the last round is fixed.

Detailed review

PR Review: #24144

Summary

This PR adds the opt-in option PropagateDeviceConfig(skip_h2d_for_constants=True). With it, a constant that only delegates on one device read is planned in that device's memory, and the runtime loads it there once, at method load, through new DeviceAllocator hooks. No issues were found.

I reviewed 964a611..7f587a9. Three independent reviewers went through every changed file, and I re-read the code behind each finding.

I left out points that earlier rounds already settled with the author:

  • one sync per load-time copy
  • load_segment_into() kept for binary size
  • shape-only readers
  • share_mutable_buffers together with device planning
  • the index check in synchronize() when a caller stream is set
  • the ATen dependency on device_allocator
  • the _tensor_data rename
  • the <optional> include
  • the placeholder condition on the hint
  • the negative stream test, and the HIP aliases
  • tests for the error paths

I checked and dropped these new findings:

  • A buffer mutated inside a delegate could be tagged device_resident_constant. It can't. SpecPropPass.update_placeholder_tensor_specs (exir/passes/spec_prop_pass.py:127-136) sets spec.const only for parameters, lifted constants, and buffers that are not mutated. The new branch at propagate_device_pass.py:456-461 requires arg_spec.const.
  • SynchronizeWaitsForTheCopyStream depends on timing. If synchronize() did not wait, the 200 ms host callback would still be running when done is checked, so the test fails reliably. The only cost is about 0.4 s of run time.
  • _constant_reader_off_device runs even when its result is unused. The warning text uses it when the option is off too, and it only runs for constant inputs of delegates.

Recommendation

Approve: no findings remain. This is still a stacked PR on #24088, so it needs the planned rebase onto main after that PR lands.

@pytorch-bot
pytorch-bot Bot deployed to upload-benchmark-results October 27, 2026 05:17 Active
@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from 7f587a9 to be1bb26 Compare October 27, 2026 18:11
Copilot AI balanced review requested due to automatic review settings October 27, 2026 18:11

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

The next commit copies the initial data of a tensor planned in device
memory there when a method loads, through the device allocator. On CUDA
that copy can still be running when copy_host_to_device() returns: it is
asynchronous on a caller stream, and from pageable memory even a plain
cudaMemcpy may return before the data lands. Two things can then go
wrong. The host buffer the data came from may be freed while the copy
still reads it. And the first run reads the tensor on its own stream; if
that stream is not ordered after the copy, for example when the method
loads with one caller stream and runs with another, or runs on a
non-blocking stream such as a green context stream, it can read the
tensor before its data is there.

DeviceAllocator gets synchronize(), which waits for the copies that the
allocator made on the calling thread, with the same per-thread settings.
The default does nothing, for allocators whose copies finish before they
return. CudaAllocator switches to the device and waits for the stream
its copies use: the caller stream when one is set, otherwise the default
stream of the device. It switches even with a caller stream, because a
handle such as nullptr or cudaStreamPerThread names a stream of the
current device. Without a caller stream the copy runs on the legacy
default stream, so waiting for it can also wait for work that other
threads queued on blocking streams of that device.

test_cuda_allocator checks that synchronize() waits for the caller
stream and for the default stream, and, with two GPUs, for the default
stream of device 0 while device 1 is current and the caller stream is
nullptr.
A memory-planned tensor can have initial data, such as a mutable buffer
with a starting value. The runtime loads that data into the tensor's
planned memory when it loads the method, with the data loader or the named
data map. Both write with the CPU, so they cannot fill a planned buffer
that lives in device memory, such as CUDA memory.

The tensor parser now checks the device of the planned buffer, which the
memory manager's planned memory records. For a device buffer it loads the
data into host memory and copies it with the device allocator registered
for that device, once, at method load. It then waits for the copy with
the allocator's synchronize(), so the host buffer can be freed and the
first run reads the right data. A host buffer is loaded as before, and
the copies made on every call are unchanged.

An empty tensor now loads nothing. When all the initial data of a program
is empty, the program has no mutable segment at all, so loading such a
method failed with "No mutable data segments in program", also for a host
buffer. The emitter test now loads such a program.

Program gets a private load_mutable_subsegment(), which loads a part of a
mutable segment into a new buffer and checks the same bounds as
load_mutable_subsegment_into(). Both share the code that finds the part.
The program test loads the same weights both ways and expects the same
bytes, and expects both to fail for bad indices and sizes.
With skip_h2d_for_method_inputs, a method input that a device delegate
reads stays in device memory. A constant, such as a weight or a buffer the
method does not change, was still copied to the device on every call,
because its data is stored in the program, on the host.

The new opt-in PropagateDeviceConfig option skip_h2d_for_constants keeps
a constant that only delegates on one device read in device memory.
PropagateDevicePass tags it with that device instead of adding a copy.
Memory planning gives it planned device memory for the whole method,
also when graph inputs are not planned, because the method reads it on
every call. The emitter stores its data as the initial state of that
memory, as it does for a mutable buffer with an initial value, and the
runtime loads it there once, when it loads the method. Each method that
reads the constant keeps its own copy, so this is opt-in. A constant that
anything else reads, such as a host op or a delegate on another device,
is still copied, with the same warning, and error_on_per_call_copies
turns that copy into an error under this option too. A constant that is
changed on the host before a delegate reads it, such as by a view, a
transpose or a clone, is not read by the delegate directly, so it is
still copied; the docs say so. They also say that share_mutable_buffers,
which needs every planned tensor in host memory, does not work with this
option.

The CUDA guide now says that every skip flag needs
enable_non_cpu_memory_planning, and that skip_h2d_for_constants works
on its own, without skip_h2d_for_method_inputs or unplanned inputs.

The new Python test checks that a buffer or parameter read only by
delegates on one device has no copy and no warning, has planned device
memory for the whole method with its data as the initial state, with and
without planned graph inputs, and that a constant also read on the host
or on a second device, or any constant without the option, is still
copied. It fails without this change.

export_program_with_device_info.py now also exports such a program, on
cuda:1, once with the constant in the program and once in a separate .ptd
file. Two new C++ tests parse the constant of each into planned memory on
a mock CUDA device and check that its data was copied there, to device 1,
and that the copy was waited for.
tensor_parser_device_test now also builds and runs with CMake. A new CUDA
test runs a program whose constant is kept in device memory twice, with
different inputs, and compares with eager PyTorch, with
error_on_per_call_copies set so it fails if the constant is copied.
The warning for a constant copied to the device on every call only said
to move it into the delegate. Without skip_h2d_for_constants, it now
also suggests the option when only delegates on that device read the
constant, unless the constant is empty and costs nothing to copy. With
the option, it names the other reader that keeps the constant off the
device.
@shoumikhin
shoumikhin force-pushed the device-resident-constants branch from be1bb26 to 0cf1823 Compare October 28, 2026 00:55
Copilot AI balanced review requested due to automatic review settings October 28, 2026 00:55

Copilot AI 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.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@pytorch-bot
pytorch-bot Bot deployed to upload-benchmark-results October 28, 2026 07:19 Active
@shoumikhin
shoumikhin merged commit 92721a8 into pytorch:main Oct 28, 2026
565 of 573 checks passed
@shoumikhin

Copy link
Copy Markdown
Contributor Author

@pytorchbot cherry-pick --onto release/1.6 -c fixnewfeature

@pytorchbot

Copy link
Copy Markdown
Collaborator

Cherry picking #24144

Command git -C /home/runner/work/executorch/executorch cherry-pick -x -X theirs 92721a88b4097e557c7d21f085227bd26c2cab32 returned non-zero exit code 1

CONFLICT (modify/delete): backends/cuda/tests/test_device_placement.py deleted in HEAD and modified in 92721a88b4 (Add skip_h2d_for_constants to keep constants in device memory (#24144)).  Version 92721a88b4 (Add skip_h2d_for_constants to keep constants in device memory (#24144)) of backends/cuda/tests/test_device_placement.py left in tree.
Auto-merging docs/source/backends/cuda/cuda-overview.md
Auto-merging exir/emit/_emitter.py
Auto-merging exir/passes/propagate_device_config.py
Auto-merging exir/passes/propagate_device_pass.py
Auto-merging exir/program/_program.py
Auto-merging exir/tests/test_propagate_device_pass.py
error: could not apply 92721a88b4... Add skip_h2d_for_constants to keep constants in device memory (#24144)
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".
hint: Disable this message with "git config set advice.mergeConflict false"
Details for Dev Infra team Raised by workflow job

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@pytorchbot cherry-pick --onto release/1.6 -c fixnewfeature

@pytorchbot

Copy link
Copy Markdown
Collaborator

Cherry picking #24144

The cherry pick PR is at #24274 and it is recommended to link a fixnewfeature cherry pick PR with an issue.

Details for Dev Infra team Raised by workflow job

This branch was successfully deployed

1 active deployment
upload-benchmark-results — 0cf1823d Deployed Oct 28, 2026 by pytorch-bot[bot] via upload-benchmark-results #21385
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/cuda ciflow/trunk CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. release notes: exir Changes to any dialects and passes on these dialects, such as memory planning

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants