[shoumikhin]

This builds on #24088, which must merge first. Review only the last 6 commits: shoumikhin/executorch@device-pass-constants...device-resident-constants

What is wrong today

With skip_h2d_for_method_inputs, the inputs of a device delegate arrive in device memory. But constants (weights and read-only buffers) are stored on the host, in the program or its data file, so the method copies them to the device on every call.

What this change does

New option PropagateDeviceConfig(skip_h2d_for_constants=True) keeps a constant that only delegates on one device read in that device's memory. The runtime loads it there once, at load.

• PropagateDevicePass tags it with the device instead of adding a copy.
• Memory planning gives it device memory for the whole method, also when graph inputs are not planned.
• The emitter stores its data as the initial value of that memory.
• A data loader cannot write device memory, so the tensor parser hands such data to the device allocator. New DeviceAllocator::copy_from_host() and load_segment_into() load it on the host, copy it, and wait. Their default bodies stay in the header, so builds without an allocator skip them.
• The wait lets a first run on another stream read the right data. New DeviceAllocator::synchronize() does nothing by default, and on CUDA waits for the stream the copy used.

Limits:

• It is opt-in: the constant holds device memory while the method is loaded, and each method reading it has its own copy.
• If another method reads it on the host, a .pte holding its data stores it twice.
• A constant that anything else also reads (a host op, or a delegate on another device) is still copied on every call, with a warning, or an error with error_on_per_call_copies.

What was tested

• New exir/tests/test_propagate_device_pass.py tests check that only such constants skip the copy and get device memory, even without skip_h2d_for_method_inputs. They fail without this change.
• New tensor_parser_device_test.cpp tests load such a constant from the program and from a .ptd file onto a mock CUDA device, checking its data and the wait. They fail without the runtime change.
• test_cuda_allocator.cpp checks that synchronize() waits for the caller and default streams. program_test.cpp checks that program data loaded into device memory goes through the allocator.
• A CUDA test, with error_on_per_call_copies set, runs such a program twice and matches eager PyTorch on an H100.
• The stripped minimal runtime (test/size_test.cpp, unoptimized x86-64) grows by 1,632 bytes.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 7 Unclassified Failures

As of commit 7f587a9 with merge base 5fa57d1 ():

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:

• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (facebook, dinov2-small-imagenet1k-1-layer, non-quantized) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (mistralai, Voxtral-Mini-3B-2507, non-quantized) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (mistralai, Voxtral-Mini-3B-2507, quantized-int4-weight-only) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (mistralai, Voxtral-Mini-4B-Realtime-2602, quantized-int4-tile-packed) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (nvidia, parakeet-tdt, non-quantized) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / test-model-cuda-windows-e2e (nvidia, parakeet-tdt, quantized-int4-weight-only) / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch
• Test CUDA Windows Export and E2E / unittest-cuda-runtime-windows / windows-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  executorch

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

[Copilot]

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

[Copilot]

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

[shoumikhin]

@claude review

[claude]

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.

[Copilot]

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

[shoumikhin]

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]

@claude review

[claude]

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.

[Copilot]

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

[shoumikhin]

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]

@claude review

[claude]

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]

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.

[mergennachin]

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

[shoumikhin]

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]

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

Can you use ghstack?

[Copilot]

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

[shoumikhin]

Thanks. In the latest push:

• CudaAllocator::synchronize() now waits only on the stream the copies used (the caller stream, or else the default stream), as its doc says. The test covers both streams.
• The warning and the decision now share one "other reader" check. The hint names PropagateDeviceConfig, and a test checks it is not suggested when it would not help.
• The mock records the synchronized device, and the test checks it.

HIP aliases for the test are left for a separate change, since I cannot run it on ROCm.

@claude review

[claude]

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

---

Nothing in the last two commits blocks the merge; what's left is a load-time sync that still runs once per tensor, a warning that can name the graph output node as a "reader", and two small test gaps.

Detailed review

PR Review: #24144

Summary

I reviewed only c5eec90 (shared "other reader" check and warning) and fc0d846 (DeviceAllocator::synchronize(); the CUDA override now waits on the copy stream only). Three independent reviewers went through both commits, and I re-read the code behind each finding below. I left out points settled in earlier rounds (tests for error paths, HIP aliases, the copy_host_to_device contract). I also dropped one finding that all three reviewers raised: the new hint condition does not repeat _is_placeholder(arg). That has no effect today, because spec.const is set only on placeholders (spec_prop_pass.py). All findings are minor.

Code Quality

• exir/passes/propagate_device_pass.py:383-388, 482-485: the warning can name the graph output node as the reader.
  • _constant_reader_off_device returns any user whose _delegate_device is not the target device. That includes the output node.
  • Trigger: with skip_h2d_for_constants=True, a buffer that a cuda:0 delegate reads and that the method also returns directly.
  • Result: the message says "because 'output' also reads it. Have the partitioner take it into the delegate…". Neither part helps the user: the real blocker is that the method returns the constant.
  • Fix: special-case reader.op == "output", for example "because the method returns it".
• extension/cuda/cuda_allocator.cpp:579-584: with a caller stream set, synchronize(index) accepts any index.
  • copy_impl (:293-301) rejects any index other than -1 or 0 when a caller stream is set. synchronize(1) instead returns Ok after waiting on a stream that may belong to another device.
  • The tensor parser cannot reach this today, because the copy would already have failed. But DeviceAllocator::synchronize documents per-index behaviour (runtime/core/device_allocator.h).
  • Fix: apply the same index == -1 || index == 0 check, so the two entry points agree.

Testing

• exir/tests/test_propagate_device_pass.py:2019-2057: the hint cases never check which reader is named.
  • "host reader" and "two devices" only assert the shared kept_elsewhere text and assertNotIn("_h2d_copy", ...).
  • The look-through-_h2d_copy logic exists so that "two devices" names the cuda:1 delegate. A regression that named some other node would still pass.
  • Fix: assert the expected node name, for example executorch_call_delegate_1, or the host op's name.
• backends/cuda/runtime/test/test_cuda_allocator.cpp:189-217: the test would also pass with cudaDeviceSynchronize().
  • It shows that synchronize() waits for the copy stream. It does not show the point of this push: that synchronize() does not wait for unrelated streams.
  • Fix: queue a long host callback on a second non-blocking stream that is not the caller stream, and assert it is still pending after synchronize() returns.

Performance

• runtime/executor/tensor_parser_exec_aten.cpp:244: one stream sync per tensor at load.
  • copyToDevice syncs after every copy, at both call sites (:314 and :348). With a caller stream set, N device-resident constants mean N back-to-back cudaStreamSynchronize calls, so the copies cannot overlap with each other or with host-side loading of the next tensor.
  • This is much cheaper than the earlier cudaDeviceSynchronize, and the per-copy sync is what currently makes freeing data right away safe. So this is a trade-off, not a bug.
  • If load time matters for large models: keep the host buffers alive and sync once per device at the end of Method::init.

Recommendation

Approve: none of these blocks the merge.

[shoumikhin]

Thanks. Answers to the rest, without another push:

• Naming output already says the method returns the constant, which is the reason it stays a copy.
• synchronize() with a caller stream and a nonzero index cannot happen after a successful copy, since the copy rejects that index first.
• One stream sync per load-time copy is the trade-off that lets the parser free each host buffer right away. Syncing once in Method::init would need the buffers kept alive.
• A negative check that other streams are not waited on is a fair test addition for later.

[Copilot]

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

[shoumikhin]

@claude review

[claude]

Claude finished @shoumikhin's task in 5m 15s —— View job

---

Nothing in this PR's own five commits (a8592ed..a094831) blocks the merge. Two minor points are left: DeviceAllocator gains a one-line virtual, load_segment_into(), that only Program calls, and the CUDA docs still say "both flags" now that there are three.

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 loaded there once, at method load, through new DeviceAllocator hooks. The newest commit, a094831, moves the device loading behind DeviceAllocator virtuals so that builds without device memory don't carry it.

I reviewed 964a611..a094831, with the most attention on a094831. Three independent reviewers went through the diff, and I re-read the code behind each finding below.

I left out points settled in earlier rounds:

• one sync per load-time copy
• the per-method dict
• tests for the error paths
• the HIP aliases
• the placeholder condition on the hint

I also checked and dropped these findings:

• The CUDA test doesn't keep the constant's memory from being reused. It does: EachAddOnCuda gives each add its own delegate.
• The .pte file name is backwards. It isn't: …Program.pte with ….ptd follows test/models/export_program.py:398.
• Nothing guards the header-only rule in device_allocator.h. The size thresholds in pull.yml would catch the regression.
• A Buck dependency is missing for data_loader.h. It isn't: :core exports that header.

API Design

• runtime/core/device_allocator.h:177-186: load_segment_into() is a public virtual with one caller and a one-line body.
  • Its only caller is Program::load_mutable_subsegment_into (runtime/executor/program.cpp:750-753).
  • Its body is copy_from_host(loader.load(...), ...), which program.cpp could call directly: device_allocator->copy_from_host(loader_->load(total_offset, size, info), buffer, size, device_index).
  • Because copy_from_host is already virtual, the default still lives only where a concrete allocator's vtable is emitted. Inlining the call should keep the 1.6 KB result.
  • The cost of keeping it:
    • Every device backend (CUDA today, others later) gets a second override point that knows about DataLoader and SegmentInfo.
    • device_allocator.h now includes data_loader.h, which ties the device-memory interface to the program-loading layer.
    • CLAUDE.md asks for no trivial single-use helpers.
  • Suggest dropping load_segment_into() and the data_loader.h include. If it is meant as a hook for a future zero-copy loader-to-device path, say so in its doc. Fix this →
• docs/source/backends/cuda/cuda-overview.md:167-175, 228-233: the requirements still describe two flags.
  • Line 167 says "Both flags need enable_non_cpu_memory_planning=True" and "asking for either skip raises a ValueError". PropagateDevicePass.__init__ (exir/passes/propagate_device_pass.py:336-346) now raises the same error for skip_h2d_for_constants.
  • The new paragraph doesn't say that skip_h2d_for_constants works on its own: it needs neither skip_h2d_for_method_inputs nor alloc_graph_input=False. test_skip_h2d_for_constants_on_its_own covers that case.
  • Because the option sits right after the inputs-only requirements, readers will likely assume it needs them too.
  • Change "Both flags" to "The skip flags", and add one short sentence that the constant option is independent of the input and output flags. That keeps the doc short, as @mergennachin asked. Fix this →

Recommendation

Approve: neither finding blocks the merge.

[Copilot]

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

[shoumikhin]

@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 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.

[Copilot]

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

[shoumikhin]

@claude review

[claude]

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.