From fc37994bd93bf3054a84f9f397b2dff6b4e6d439 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 23 Sep 2026 18:38:59 -0700 Subject: [PATCH] Document both shapes for device resident outputs The CUDA page tells you to export device resident inputs and outputs with unplanned graph outputs. A Python caller cannot follow that. An unplanned output means the program reserves no buffer, so the caller has to supply one, and the Python bindings supply host memory. A CUDA delegated method then fails to run. The setting the page warns against is the one that works from Python. With non CPU memory planning on, a planned output buffer is allocated on the device, and skipping the device to host copy means nothing brings it back, so the result stays on the device and the caller receives a pointer to the runtime's own buffer. That is also what the exporters in this repository already do. The input half of the recipe was correct and is unchanged. The output half now says there are two shapes: planned when the caller cannot supply device memory, which is every Python caller today, and unplanned when it can, which today means a C++ caller handing the runtime a device tensor. Test plan: read against the example exporters in the repository, which use unplanned inputs and planned outputs with both copy skips. --- docs/source/backends/cuda/cuda-overview.md | 38 +++++++++++++++------- 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index a4076251ea9..ef57aca7ae5 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -140,13 +140,13 @@ exec_program = et_program.to_executorch( skip_d2h_for_method_outputs=True, ), enable_non_cpu_memory_planning=True, - # Required alongside the skips. Memory planning reserves a buffer for graph inputs and - # outputs by default, and the runtime fills a planned input by copying the caller's - # memory into that buffer, which reintroduces the copy you just asked to skip. On device - # memory that copy is a host memcpy into a device pointer, which is undefined. - memory_planning_pass=MemoryPlanningPass( - alloc_graph_input=False, alloc_graph_output=False - ), + # Required alongside the skips. Memory planning reserves a buffer for graph inputs by + # default, and the runtime fills a planned input by copying the caller's memory into that + # buffer, which reintroduces the copy you just asked to skip. On device memory that copy is + # a host memcpy into a device pointer, which is undefined. + # Outputs stay planned, so the runtime allocates them on the device and hands back a + # pointer to its own buffer. See the note below for when to unplan them instead. + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), ) ) ``` @@ -169,11 +169,25 @@ out works, and the snippet above sets it explicitly only to make the requirement it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens during device-aware memory planning. -**Both flags also need unplanned graph inputs and outputs**, via -`MemoryPlanningPass(alloc_graph_input=False, alloc_graph_output=False)` as shown above. Without -it the program still reserves its own buffer and the runtime copies into it, so the copy comes -back at run time. On device memory that copy is a host memcpy into a device pointer, which is -undefined, so the program crashes rather than returning a wrong answer. +**Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the +program reserves its own input buffer and the runtime copies the caller's memory into it, so the +copy comes back at run time. On device memory that copy is a host memcpy into a device pointer, +which is undefined, so the program crashes rather than returning a wrong answer. + +**Outputs have two valid shapes, and which one you want depends on the language.** + +Leave outputs planned, which is the default, when the caller cannot supply device memory itself. +The planned buffer is allocated on the device because `enable_non_cpu_memory_planning` is on, and +`skip_d2h_for_method_outputs` means nothing copies it back, so the result stays on the device and +the caller receives a pointer to the runtime's own buffer. That buffer belongs to the method and is +reused by the next execution, so read or copy the result before running again. This is the shape +every Python caller wants, and it is what the exporters in this repository use, for example +`examples/models/gemma4_31b/export.py`. + +Set `alloc_graph_output=False` only when the caller supplies the output memory itself, which today +means a C++ caller handing the runtime a device tensor through `Module::set_output`. The Python +bindings allocate host memory for an unplanned output, so a device resident method exported this +way cannot run from Python. **Per-method selection is on the outer config, not on the flags.** Pass a dict of `PropagateDeviceConfig` keyed by method name: