Give device-tagged method outputs device memory in the Python bindings by shoumikhin · Pull Request #23547 · pytorch/executorch

shoumikhin · 2026-10-07T20:30:24Z

Problem

A program exported to keep its method outputs on the GPU (skip_d2h_for_method_outputs=True with alloc_graph_output=False) leaves each output's buffer to the caller. From Python, PyMethod::allocate_output_storages() supplies that buffer, and it always allocates host memory (std::vector<uint8_t>), whatever device the program tagged the output with.

What happens next depends on the delegate:

The CUDA backend checks each CUDA-tagged tensor with cudaPointerGetAttributes and rejects host memory, so execute() fails up front with error 0x12 (InvalidArgument) on every GPU.
A delegate without that check, such as Torch-TensorRT's, writes its results into the host buffer, and the bindings return a CUDA tensor over that host pointer. Cloning it is a device-to-device copy from a host pointer, which the driver accepts only on a GPU that can read pageable host memory. Elsewhere execute() fails:

RuntimeError: CUDA error: invalid argument
The CUDA driver logged these messages, which may provide useful details:
Allocation not found
Returning 1 (CUDA_ERROR_INVALID_VALUE) from cuMemcpyDtoDAsync_v2

The second case is what fails Torch-TensorRT's executorch-runtime-test jobs since they moved to A10G runners without pageable memory access. It fails the same way on a Jetson Orin.

Part of #23101, which described this and was closed by a documentation change only. This PR fixes the Method path. The Module path (PyModule::allocate_output_tensors) has the same problem and is tracked in #23553.

Change

An unplanned tensor output tagged for a device now gets a buffer from that device's registered allocator, through DeviceMemoryBuffer. Host outputs, memory-planned outputs, non-tensor outputs and zero-byte outputs are unchanged, and so is ATen mode.

The output buffers are built in locals and moved into place only once every output has one, so an allocation that fails partway leaves the method's state as it was.

Test plan

New test test_unplanned_device_output_is_device_memory exports a CUDA-delegated program with device-resident inputs and outputs. It first asserts that the outputs are unplanned and tagged CUDA, then runs the program with CUDA inputs. It asks the CUDA driver (cuPointerGetAttribute) whether the uncloned output is device memory, and checks the values through execute(). It skips without the CUDA backend or a GPU.

I built two CUDA wheels from main at 1802bcf, one with this change and one without, and ran extension/pybindings/test/test_pybindings.py against each. The run without the change used this branch's test file. 43 tests were collected in every run. The Thor numbers with the change are from the current head; the Orin numbers are from the first version, which the review changes do not touch on this path.

Board	Pageable host memory access	Without the change	With the change
Jetson AGX Thor	yes	new test fails (`0x12`); 38 passed, 4 failed, 1 skipped	new test passes; 39 passed, 3 failed, 1 skipped
Jetson Orin Nano	no	new test fails (`0x12`); 38 passed, 4 failed, 1 skipped	new test passes; 39 passed, 3 failed, 1 skipped

The other failures are not caused by this change:

test_quantized_ops and test_method_quantized_ops fail in every run with "Missing out variants: quantized_decomposed::...". The wheel ships the quantized AOT library, but this test file never imports executorch.kernels.quantized, which is what loads it. Run alone, nothing else loads it either. Both pass once that package is imported first. CI loads it through other test modules in the same pytest session, and the Buck targets preload it.
test_device_planned_method_allocates_on_the_device measures free device memory around a method load. On a Jetson that memory is shared with the CPU and moves with whatever else runs, so the check fails on both boards with and without this change, depending on the moment. Its outputs are planned, so it never reaches this code.

A program exported to keep its outputs on the GPU (skip_d2h_for_method_outputs with alloc_graph_output=False) leaves the output buffers to the caller. The Python bindings gave every such output a host std::vector, whatever device the program tagged it with. The CUDA backend checks for this and fails execute() with InvalidArgument. A delegate without that check writes its results into the host buffer, and the bindings return a CUDA tensor over that host pointer. Cloning it is a device-to-device copy from a host pointer, which the driver accepts only on a GPU that can read pageable host memory. Elsewhere, for example on a Jetson Orin or an A10G, execute() fails with "CUDA error: invalid argument" and the driver logs "Allocation not found" from cuMemcpyDtoDAsync. Now an unplanned output whose tensor is tagged for a device gets a buffer from that device's registered allocator. Host outputs are unchanged. The new test asks the CUDA driver whether the uncloned output is device memory. It fails without this change and passes with it on a Jetson AGX Thor and a Jetson Orin Nano.

pytorch-bot · 2026-10-07T20:30:28Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 5e23d75 with merge base 1802bcf ():
💚 Looks good so far! There are no failures yet. 💚

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 reached their quota limit.

Gasoonjia · 2026-10-07T21:09:54Z

+            # CU_POINTER_ATTRIBUTE_MEMORY_TYPE is 2 and CU_MEMORYTYPE_DEVICE is 2. The driver does
+            # not know a pageable host pointer at all, and returns an error for it.
+            memory_type = ctypes.c_uint(0)
+            status = ctypes.CDLL("libcuda.so.1").cuPointerGetAttribute(


what's this test for?

It checks the bug this PR fixes, and fails without the fix.

With alloc_graph_output=False the bindings supply each output buffer, and they used to supply host memory even for a CUDA-tagged output. The tensor's device tag says cuda either way, so this line asks the CUDA driver what the pointer really is: CU_MEMORYTYPE_DEVICE with the fix, and an error for a pageable host pointer without it.

On main the test fails one line earlier, at method._method(...), with 0x12 (InvalidArgument), because the CUDA backend's own cudaPointerGetAttributes check rejects the host buffer. A delegate without that check (Torch-TensorRT's) writes into it and fails later, at the clone, on GPUs without pageable memory access. The pointer check here covers that case too, and the execute() assert below checks the values.

The comments in the test now say this more directly (70b36e1).

Gasoonjia

Thanks for the fix. The approach matches how ProgramMemory already backs device-planned buffers with DeviceMemoryBuffer. A few comments inline, plus one more here:

[Follow-up] The PyModule path has the same problem

PyModule::allocate_output_tensors (pybindings.cpp:990, the _load_for_executorch / Module path) still backs every unplanned tensor output with std::vector<uint8_t>(nbytes), whatever its device. So the same program run through the legacy API hits the same failure. Fine to leave for a follow-up, but worth tracking. Maybe reopen #23101 or file a new issue rather than closing it with this PR.

Gasoonjia · 2026-10-07T21:16:47Z

+      if (!device.is_cpu()) {
+        auto buffer = DeviceMemoryBuffer::create(
+            output_size, device.type(), device.index());
+        THROW_IF_ERROR(


[Bug] Partial failure leaves allocate_output_storages() in a broken state

allocate_output_storages() only returns early when output_storages_.size() == num_outputs, and it never clears output_storages_ before appending. If a THROW_IF_ERROR fires partway through the loop, output_storages_ keeps the entries pushed so far. The next execute() appends another full set, so output_storages_.size() ends up greater than num_outputs. From then on the early-return check never matches, and the loop in execute() indexes output_storage_spans[i] / device_output_storages_[i] out of bounds.

This was latent before, but this PR adds a throw site that can actually happen: DeviceMemoryBuffer::create returns NotFound when no allocator is registered, or fails on device OOM. A Python caller that catches the RuntimeError and retries hits UB. Could you build both vectors in locals and move them into the members only after the loop succeeds, or at least clear() both at the top?

Good catch, fixed in f4c1d82. Both vectors are built in locals and moved into the members only after every output has its buffer, so a throw partway leaves the members untouched and the next call starts over.

Gasoonjia · 2026-10-07T21:16:47Z

+      // read pageable host memory, and fails everywhere else.
+      const auto device = method_->get_output(i).toTensor().device();
+      if (!device.is_cpu()) {
+        auto buffer = DeviceMemoryBuffer::create(


[Nit] Zero-byte outputs

For a device-tagged output with nbytes() == 0, this now calls DeviceMemoryBuffer::create(0, ...), and whether that succeeds depends on the allocator. Before, a 0-size host vector gave an empty span, and setup_output_storage skipped it. Consider keeping the output_size == 0 case on the existing empty-storage path.

Done in 09eeffa. A zero-byte device output keeps the empty host buffer, which setup_output_storage skips, the same as before this PR.

Gasoonjia · 2026-10-07T21:16:47Z

+            status = ctypes.CDLL("libcuda.so.1").cuPointerGetAttribute(
+                ctypes.byref(memory_type), 2, ctypes.c_uint64(output.data_ptr())
+            )
+            self.assertEqual((status, memory_type.value), (0, 2))


[Nit] Assert the test's premise

The neighbouring device tests assert their premise first ("Without this the test would quietly become..."). This test would be more robust if it first checked that forward's output in exported_program.executorch_program is actually unplanned and tagged CUDA. If the planner ever started planning these outputs, this cuPointerGetAttribute check would still pass, because a planned device arena is also device memory, and the test would silently stop covering the new code.

Done in 5e23d75. Before the pointer check, the test asserts that every output of forward has no allocation_info and is tagged DeviceType.CUDA. I checked that the assert can fail: exporting the same program with alloc_graph_output=True gives a planned output, which it rejects.

The test's comments said what memory the bindings must supply, but not that the pointer check is the assertion, or why it asks the CUDA driver rather than reading the tensor's device, which says cuda either way.

allocate_output_storages() returned early only once it held one buffer per output, and appended to the member vectors as it went. A throw partway through left the buffers made so far, the next call appended a full set after them, and from then on the vectors were longer than the outputs, so execute() indexed past the end of its spans. The device allocation added before can fail, when no allocator is registered for the device or it runs out of memory. Build both vectors in locals and move them into the members only after every output has its buffer.

An unplanned device output of zero bytes needs no memory, and whether a device allocator accepts a zero size is up to the allocator. Before the device path it got an empty host buffer, which setup_output_storage skips. Keep it there.

If the planner started placing these outputs, the pointer check would still pass, since a planned device arena is device memory too, and the test would stop covering the unplanned path. Check the premise first, as the neighbouring device tests do.

Gasoonjia · 2026-10-07T21:38:08Z

Another thing in my mind is cuda allocator registration. Now for cuda backend we register it whenever people import cuda backend. Did we register the allocator in the torch-tensorrt as well? Or trt need to depend on cuda backend (hope not lol)?

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

shoumikhin · 2026-10-07T21:57:44Z

Thanks for the review. All four inline points are addressed. Rebuilt and rerun on a Jetson AGX Thor: the new test passes, and the only other failures are the two quantized tests, which need executorch.kernels.quantized imported, and test_device_planned_method_allocates_on_the_device. That test fails on the Thor with or without this change, depending on how much shared memory is free at the time, and never reaches this code, since its outputs are planned.

On the PyModule path: agreed, allocate_output_tensors has the same host-vector problem. I filed #23553 to track it and changed this PR from "Fixes #23101" to "Part of #23101", so the issue stays open until both paths are fixed. That path needs a little more than swapping the buffer: make_tensor_ptr takes ownership of host bytes, so the device buffer has to be owned by the tensor the same way, or the next call writes through a result the caller still holds.

On allocator registration: Torch-TensorRT registers it itself and does not depend on the CUDA backend. Its delegate has RegisterCudaDeviceAllocator.cpp, which registers CudaAllocator::instance() at static init, unguarded on purpose: if the CUDA backend is linked into the same binary, the registry's "Allocator already registered" check aborts instead of silently picking one. Two gaps, both outside this PR:

Only the Bazel build compiles that file. The CMake build and the runtime wheel do not, so there the allocator comes only from the CUDA backend, when it is loaded.
A process that loads both the CUDA backend and the TensorRT delegate registers twice and aborts.

Moving the registration into extension_cuda, where CudaAllocator now lives, would give every CUDA delegate one registration and fix both. I can open that as a separate change.

Summary: ## Ulterior Motive Restore the "Test WebGPU Native (Dawn)" workflow on `main`. After "[Vulkan] Allow assert-only SymBool chains in the partitioner (#23545)" fixed the RoPE export, the job advanced to a second failure in the SwiGLU dynamic-shape tests. ## Rationale **What**: Register `et_vk.swiglu.default` in the WebGPU runtime, backed by the existing `silu_mul_fused` shader. **Why**: The WebGPU backend reuses the Vulkan partitioner. "[Vulkan] Fuse SwiGLU elementwise operations (#22979)" made that partitioner fuse `sigmoid`/`mul`/`mul` into `et_vk.swiglu`, which the WebGPU runtime did not implement, so loading any program containing a SwiGLU failed with `DelegateInvalidCompatibility` (`0x30`). `DynamicShape.SwiGluReusedGraph` and `DynamicShape.SwiGluCommutativeAndOwnership` failed at `load_forward()`. ## Details - Programs containing SwiGLU - **Before**: Load fails with "unsupported op: et_vk.swiglu.default". - **After**: Lowered to one `silu_mul_fused` dispatch, the same kernel the runtime's load-time SwiGLU fusion already emits. - `add_silu_mul_fused_dispatch` - **Before**: Registered its resize hook on a single tensor (the shared q4 linear input). - **After**: Takes a list of trigger tensors. The load-time fusion passes its shared input; the new op passes both `gate` and `up`, matching how two-input ops such as `mul` register resize hooks. The op requires fp32 operands with identical `gate`, `up`, and output shapes, which is what the flat `silu_mul_fused` kernel supports. Test Plan: ## Observed Results ### CI ("Test WebGPU Native (Dawn)", Linux, Dawn + SwiftShader) | Run | Result | |---|---| | `main` at "Give device-tagged method outputs device memory in the Python bindings (#23547)" (run 37704744609) | Failed: `DynamicShape.SwiGluReusedGraph`, `DynamicShape.SwiGluCommutativeAndOwnership` (`load_forward()` returned `0x30`) | | This branch via `workflow_dispatch` (run 37718579736) | Passed | Per-binary results for run 37718579736: | Test binary | Passed | Skipped | Failed | |---|---|---|---| | webgpu_native_test | 29 | 3 | 0 | | webgpu_update_cache_test | 3 | 0 | 0 | | webgpu_dispatch_order_test | 9 | 0 | 0 | | webgpu_index_test | 4 | 0 | 0 | | webgpu_dynamic_shape_test | 40 | 4 | 0 | | webgpu_scratch_buffer_test | 8 | 0 | 0 | | webgpu_dispatch_2d_test | 13 | 0 | 0 | | webgpu_compute_dispatch_test | 20 | 0 | 0 | | webgpu_execution_options_test | 12 | 0 | 0 | | webgpu_output_suppression_test | 3 | 0 | 0 | | webgpu_op_test_util_test | 8 | 0 | 0 | | webgpu_op_test (op-test framework) | 391 | 0 | 0 | The 4 skipped dynamic-shape tests are the same ones skipped on `main`. ### Local (macOS, Apple Silicon, Dawn on Metal) | Validation | Result | |---|---| | SwiGLU fixtures, parent commit | All SwiGLU programs fail to load (reproduces CI) | | SwiGLU fixtures, this commit | All load and match goldens except `dyn_swiglu_qkv_overlap` output 1 (max err 0.025 vs 0.01) | | `QuantizedLinearBk64QkvReusedGraphAndFallbacks`, `QuantizedLinearSweep` | Fail on Metal on both parent and this commit; pass in CI. `dyn_swiglu_qkv_overlap` uses the same BK64 linear path, and passes in CI | | Performance | Not run: the op emits the same single `silu_mul_fused` dispatch as the existing load-time fusion | ## Reproduction Guide ### Run the CI workflow on this branch ```bash gh workflow run test-webgpu-native.yml --repo pytorch/executorch --ref webgpu-swiglu-op gh run list --repo pytorch/executorch --workflow test-webgpu-native.yml --branch webgpu-swiglu-op --limit 1 ``` ### Run locally on Linux (mirrors the CI job) ```bash git clone --recursive https://github.com/pytorch/executorch.git cd executorch git fetch origin webgpu-swiglu-op && git checkout FETCH_HEAD ./install_executorch.sh source .ci/scripts/setup-webgpu-linux-deps.sh bash backends/webgpu/scripts/test_webgpu_native_ci.sh ``` ### Run only the SwiGLU tests after the script has built and exported fixtures ```bash cmake-out-webgpu-dawn-ci/backends/webgpu/webgpu_dynamic_shape_test /tmp/dynamic_shape \ --gtest_filter='*SwiGlu*' ``` Reviewers: This PR was authored with Claude Code.

## Problem `test_unplanned_device_output_is_device_memory`, added in #23547, needs the CUDA backend and a visible GPU. The CPU `unittest` jobs have neither, so it skips there, and no CI job runs it. ## Change The `unittest-cuda` job already runs one pybindings device test by name and treats a skip as a failure. This runs the new test in the same command, so it gets the same check. ## Test plan On a Jetson AGX Thor, with a wheel that includes #23547, I ran the command from this change. The new test passes.

Summary: ## Ulterior Motive Restore the "Test WebGPU Native (Dawn)" workflow on `main`. After "[Vulkan] Allow assert-only SymBool chains in the partitioner (#23545)" fixed the RoPE export, the job advanced to a second failure in the SwiGLU dynamic-shape tests. ## Rationale **What**: Register `et_vk.swiglu.default` in the WebGPU runtime, backed by the existing `silu_mul_fused` shader. **Why**: The WebGPU backend reuses the Vulkan partitioner. "[Vulkan] Fuse SwiGLU elementwise operations (#22979)" made that partitioner fuse `sigmoid`/`mul`/`mul` into `et_vk.swiglu`, which the WebGPU runtime did not implement, so loading any program containing a SwiGLU failed with `DelegateInvalidCompatibility` (`0x30`). `DynamicShape.SwiGluReusedGraph` and `DynamicShape.SwiGluCommutativeAndOwnership` failed at `load_forward()`. ## Details - Programs containing SwiGLU - **Before**: Load fails with "unsupported op: et_vk.swiglu.default". - **After**: Lowered to one `silu_mul_fused` dispatch, the same kernel the runtime's load-time SwiGLU fusion already emits. - `add_silu_mul_fused_dispatch` - **Before**: Registered its resize hook on a single tensor (the shared q4 linear input). - **After**: Takes a list of trigger tensors. The load-time fusion passes its shared input; the new op passes both `gate` and `up`, matching how two-input ops such as `mul` register resize hooks. The op requires fp32 operands with identical `gate`, `up`, and output shapes, which is what the flat `silu_mul_fused` kernel supports. Test Plan: ## Observed Results ### CI ("Test WebGPU Native (Dawn)", Linux, Dawn + SwiftShader) | Run | Result | |---|---| | `main` at "Give device-tagged method outputs device memory in the Python bindings (#23547)" (run 37704744609) | Failed: `DynamicShape.SwiGluReusedGraph`, `DynamicShape.SwiGluCommutativeAndOwnership` (`load_forward()` returned `0x30`) | | This branch via `workflow_dispatch` (run 37718579736) | Passed | Per-binary results for run 37718579736: | Test binary | Passed | Skipped | Failed | |---|---|---|---| | webgpu_native_test | 29 | 3 | 0 | | webgpu_update_cache_test | 3 | 0 | 0 | | webgpu_dispatch_order_test | 9 | 0 | 0 | | webgpu_index_test | 4 | 0 | 0 | | webgpu_dynamic_shape_test | 40 | 4 | 0 | | webgpu_scratch_buffer_test | 8 | 0 | 0 | | webgpu_dispatch_2d_test | 13 | 0 | 0 | | webgpu_compute_dispatch_test | 20 | 0 | 0 | | webgpu_execution_options_test | 12 | 0 | 0 | | webgpu_output_suppression_test | 3 | 0 | 0 | | webgpu_op_test_util_test | 8 | 0 | 0 | | webgpu_op_test (op-test framework) | 391 | 0 | 0 | The 4 skipped dynamic-shape tests are the same ones skipped on `main`. ### Local (macOS, Apple Silicon, Dawn on Metal) | Validation | Result | |---|---| | SwiGLU fixtures, parent commit | All SwiGLU programs fail to load (reproduces CI) | | SwiGLU fixtures, this commit | All load and match goldens except `dyn_swiglu_qkv_overlap` output 1 (max err 0.025 vs 0.01) | | `QuantizedLinearBk64QkvReusedGraphAndFallbacks`, `QuantizedLinearSweep` | Fail on Metal on both parent and this commit; pass in CI. `dyn_swiglu_qkv_overlap` uses the same BK64 linear path, and passes in CI | | Performance | Not run: the op emits the same single `silu_mul_fused` dispatch as the existing load-time fusion | ## Reproduction Guide ### Run the CI workflow on this branch ```bash gh workflow run test-webgpu-native.yml --repo pytorch/executorch --ref webgpu-swiglu-op gh run list --repo pytorch/executorch --workflow test-webgpu-native.yml --branch webgpu-swiglu-op --limit 1 ``` ### Run locally on Linux (mirrors the CI job) ```bash git clone --recursive https://github.com/pytorch/executorch.git cd executorch git fetch origin webgpu-swiglu-op && git checkout FETCH_HEAD ./install_executorch.sh source .ci/scripts/setup-webgpu-linux-deps.sh bash backends/webgpu/scripts/test_webgpu_native_ci.sh ``` ### Run only the SwiGLU tests after the script has built and exported fixtures ```bash cmake-out-webgpu-dawn-ci/backends/webgpu/webgpu_dynamic_shape_test /tmp/dynamic_shape \ --gtest_filter='*SwiGlu*' ``` Reviewers: This PR was authored with Claude Code.

Summary: ## Ulterior Motive Restore the "Test WebGPU Native (Dawn)" workflow on `main`. After "[Vulkan] Allow assert-only SymBool chains in the partitioner (#23545)" fixed the RoPE export, the job advanced to a second failure in the SwiGLU dynamic-shape tests. ## Rationale **What**: Register `et_vk.swiglu.default` in the WebGPU runtime, backed by the existing `silu_mul_fused` shader, and delete the runtime SwiGLU graph fusion it replaces. **Why**: The WebGPU backend reuses the Vulkan partitioner. "[Vulkan] Fuse SwiGLU elementwise operations (#22979)" made Vulkan preprocessing fuse `sigmoid`/`mul`/`mul` into `et_vk.swiglu`, which the WebGPU runtime did not implement, so loading any program containing a SwiGLU failed with `DelegateInvalidCompatibility` (`0x30`). With fusion now done at export time, the runtime pattern matcher can no longer match: everything it accepts (fp32, q4 gate/up producers, gate used only by the pattern, equal shapes) is a strict subset of what the export-time pattern fuses. Instrumenting it on the current fixtures confirmed zero matches across all 89 graph builds in `webgpu_dynamic_shape_test`. ## Details - Programs containing SwiGLU - **Before**: Load fails with "unsupported op: et_vk.swiglu.default". - **After**: Lowered to one `silu_mul_fused` dispatch, with the resize hook registered on both `gate` and `up`, matching how two-input ops such as `mul` register resize hooks. - Runtime SwiGLU graph fusion (`passes/SwiGLU.{h,cpp}`) - **Before**: Matched q4 gate/up + `sigmoid` + 2x `mul` at load time, emitted `silu_mul_fused`, and placed `gate` in a scratch buffer. - **After**: Deleted. The dispatch code moved into `ops/mul/SwiGLU.cpp`. - QKV-BK64 vs. SwiGLU precedence - **Before**: The runtime SwiGLU pass claimed the gate/up linears, and `retain_unclaimed_qkv_fusions` dropped any QKV-BK64 group overlapping them. - **After**: Removed. Since #22979 the elementwise fusion no longer owns the linears, so both fusions apply: q/gate/up share one `linear_q4gsw_bk64_qkv` dispatch followed by `silu_mul_fused`. QKV candidates are grouped by input tensor and cannot overlap each other, so `retain_unclaimed_qkv_fusions` had no other effect. - Profiling-only tests (`SwiGluFusionProfile`, `SwiGluQkvOverlapProfile`) - **Before**: Encoded the runtime matcher's rules: no fusion when `gate` has an extra consumer, is a graph output, or when `up` comes from a separate input, and no QKV-BK64 dispatch in the overlap case. These tests only build with `EXECUTORCH_BUILD_WEBGPU_PROFILING`, so CI does not run them and they had gone stale since #22979. - **After**: Expect fusion in those three cases (`gate` stays materialized, so fusing is valid) and one `linear_q4gsw_bk64_qkv` dispatch in the overlap case. Cases where `sigmoid` or `silu` has another user still expect no fusion. Test Plan: ## Observed Results ### CI ("Test WebGPU Native (Dawn)", Linux, Dawn + SwiftShader) | Run | Result | |---|---| | `main` at "Give device-tagged method outputs device memory in the Python bindings (#23547)" (run 37704744609) | Failed: `DynamicShape.SwiGluReusedGraph`, `DynamicShape.SwiGluCommutativeAndOwnership` (`load_forward()` returned `0x30`) | | Op only, before removing the runtime fusion (run 37718579736) | Passed (all 11 native test binaries; `webgpu_dynamic_shape_test` 40 passed / 4 skipped, same skips as `main`; op-test framework 391 passed) | | This commit, rebased on "Vulkan: narrow the int4 linear output tile on Mali" | Pending | ### Local (macOS, Apple Silicon, Dawn on Metal, `EXECUTORCH_BUILD_WEBGPU_PROFILING=ON`, `WEBGPU_TIMESTAMP_QUERY=1`) | Validation | Result | |---|---| | Runtime SwiGLU matcher hit count on current fixtures (temporary instrumentation, before deletion) | 0 matches in 89 graph builds | | `SwiGluReusedGraph`, `SwiGluFusionProfile` | Pass | | `SwiGluQkvOverlapProfile` dispatch counts (1 `linear_q4gsw_bk64_qkv`, 1 `silu_mul_fused`, 0 `mul`) | Pass | | `SwiGluCommutativeAndOwnership`, `SwiGluQkvOverlapProfile` numerics | Fail only on `dyn_swiglu_qkv_overlap` output 1 (max err 0.025 vs 0.01) | | `QuantizedLinearBk64QkvReusedGraphAndFallbacks`, `QkvBk64LiveRoutesProfile`, `QuantizedLinearBk64QkvProfileSoleWriter` | Fail on Metal with and without this change; BK64 tests pass in CI. `dyn_swiglu_qkv_overlap` uses the same BK64 path and passed in CI run 37718579736 | | Performance | Not run. New exports lose only the `gate` scratch-buffer reuse; programs exported before #22979 lose the load-time fusion and run `sigmoid`/`mul`/`mul` as three dispatches | ## Reproduction Guide ### Run the CI workflow on this branch ```bash gh workflow run test-webgpu-native.yml --repo pytorch/executorch --ref webgpu-swiglu-op gh run list --repo pytorch/executorch --workflow test-webgpu-native.yml --branch webgpu-swiglu-op --limit 1 ``` ### Run locally on Linux (mirrors the CI job) ```bash git clone --recursive https://github.com/pytorch/executorch.git cd executorch git fetch origin webgpu-swiglu-op && git checkout FETCH_HEAD ./install_executorch.sh source .ci/scripts/setup-webgpu-linux-deps.sh bash backends/webgpu/scripts/test_webgpu_native_ci.sh ``` ### Run the profiling-only SwiGLU tests After the script above has exported fixtures to `/tmp/dynamic_shape`: ```bash cmake -DPYTHON_EXECUTABLE=python \ -DEXECUTORCH_BUILD_WEBGPU=ON -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \ -DEXECUTORCH_BUILD_WEBGPU_PROFILING=ON -DEXECUTORCH_BUILD_TESTS=ON \ -DDawn_DIR="${Dawn_DIR}" \ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON -DCMAKE_BUILD_TYPE=Release \ -B cmake-out-webgpu-prof . cmake --build cmake-out-webgpu-prof --target webgpu_dynamic_shape_test WEBGPU_TIMESTAMP_QUERY=1 cmake-out-webgpu-prof/backends/webgpu/webgpu_dynamic_shape_test \ /tmp/dynamic_shape --gtest_filter='*SwiGlu*' ``` Reviewers: This PR was authored with Claude Code.

…ction The ExecuTorch pages of the user guide were hard to follow and partly out of date: - The saving section carried its own install command, which disagreed with the guide's. - The install section opened with a command that gives no published wheel its delegate: only a nightly extra brings it today. No command on the page installed the delegate for Torch-TensorRT 2.15. The section also listed Windows, which 2.15 has no delegate for. - It said each delegate of a coalesced program runs on its own CUDA stream by default. Both now run on the calling thread's default stream. - The guide said engines on the shared activation scratch never replay a CUDA graph. In builds newer than 2.15 they do. - The guide said Python cannot run a program whose outputs are not planned. It can now, with ExecuTorch builds that have pytorch/executorch#23547, which both 2.15 and nightly install. - Features that 2.15 lacks were not marked, and two were missing: engines in a separate .ptd file, and optimization profiles. Weight streaming was not described at all. - The zero-copy KV cache section was a wall of text, and the list of save() options for this format named only six of them. - The comparison table said a .pt2 runs on Windows. save() refuses that format outside Linux. - The save() docstring said a target_device compile spec overrides the default target device. It does not move the engines: an engine runs on the GPU it was built on. - The saving section showed how to write an ETRecord with generate_etrecord=True. For a program with TensorRT engines that fails today with a SerializeError, on nightly and on 2.15. - Nothing said that TensorRT warns, on a green context stream, that the engine was built without green context constraints and may crash. - The guide's note said ordinary Torch-TensorRT builds keep their CUDA 12 support. The 2.15 release candidate and current nightly builds are built for CUDA 13 only. The guide is rewritten in shorter sentences. It keeps one install section, which the other pages now link to. It starts with the commands that work today. The 2.15 command pins the release and names the 0.2.1 delegate, which comes with the next 2.15 release candidate, from the test channel. It also lists the nightly channel for the ExecuTorch build that delegate pins, which that channel keeps only for a limited time, and the section says so. The nightly command follows, and then the plain form, which is enough once a release's extra brings the delegate. Each feature that 2.15 lacks is marked "Not in 2.15", which stays true after the next release branches, and a new section shows how to set a runtime option from C++ and from Python. The saving section now covers only the saving step: save() versus export(), every option of the executorch format, several methods in one file, and a condensed zero-copy KV cache section. The API page no longer says the module is Linux only, and the save() docstring points to the deployment guide, where the CUDA backend recipe moved. That docstring also stops saying that target_device moves the engines, and says that generate_etrecord does not work yet for programs with TensorRT engines. The export() docstring says the same. Windows is marked "not in 2.15", because the 2.15 release candidate publishes no Windows delegate wheel. The test that checks every printed install command requires --pre on any command that installs the extra from the nightly channel. The 2.15 command lists that channel only for the ExecuTorch build, and must not take --pre, or pip takes nightly builds of everything else. The test now skips the --pre requirement, and only that, for a user docs command that also names the delegate with == or ~=. Every other check, including the supported channel check, still runs on it. New cases pin the 2.15 command, the same command with an unsupported channel, and a nightly command without --pre. Tested on Linux x86-64 with an H100 and nightly Torch-TensorRT. Ran every Python block of the guide in order: they export the plain, coalesced, GPU-resident and separate engine file programs, run the first three, and set the runtime options from Python. A separate run loaded the separate engine file program, which fails without its .ptd and works with it. Built and ran the C++ snippets against the installed wheels, all but the optimization profile one, and ran one save() example and one export() example from the saving section. Checked these 2.15 differences against the Torch-TensorRT 2.15.0 release candidate: the re-export default, the Python runtime options, the separate engine file, and a CPU input to a GPU-resident program. On both versions, ran a GPU-resident program whose last part runs on the CUDA backend, with planned and with unplanned outputs. Windows was not tested. The pages build with Sphinx with no new warnings, and the tests that read these files pass. The new cases fail without the test change: the 2.15 command is rejected for lacking --pre.

…unner READMEs The three READMEs next to the ExecuTorch delegate had grown long and hard to read. The delegate README was over 900 lines, with broken line wraps and long measurement notes. The runtime wheel README repeated its Python example three times. It and the reference runner README each carried their own install command, and those disagreed with the guide. Each README now says what its directory is, links to the deployment guide for installing, and keeps to its own topic: - The delegate README: how to build the delegate from source, the rules a C++ caller has to keep (threads, buffers, streams), load errors, known issues, and a reference for each runtime option. Implementation notes that already live in the code comments are gone. So is a known issue about Python outputs in host memory, which pytorch/executorch#23547 fixed. The capture rules and the notes on how the shared scratch pool grows stay, shorter, because the delegate's error message and header point there. The build recipe now sets EXECUTORCH_SOURCE_DIR, which the delegate needs to find the ExecuTorch checkout when the CUDA backend is off. The recipe for building the archive alone now passes the prebuilt extension_cuda library from libtorchtrt.tar.gz, and says how to build that library from the checkout instead. - Features that 2.15 lacks are marked "Not in 2.15", in all three READMEs. - The runtime wheel README: version matching, use from Python and C++, runtime library paths, and how to build the wheel. It no longer says that building the wheel needs NVIDIA's package index: TensorRT also installs from PyPI alone, more slowly. - The reference runner README: what each example shows, one build recipe for all of them, and how to run each one. It also says that the prebuilt runner has no CUDA backend, so it cannot run a program split with the CUDA backend, and that the JetPack and TensorRT-RTX builds do not ship it. Tested on Linux x86-64 with an H100 and nightly Torch-TensorRT: checked the wheel's dependency metadata, runtime paths and exported symbols against the README, built and ran the C++ link recipe, ran the device-resident export and the Python loader from the reference runner README, and ran the prebuilt runner. Built the delegate archive alone with a prebuilt extension_cuda library. On Linux aarch64 with CUDA 13.2, built the archive alone without the prebuilt library, and the delegate inside a project that adds ExecuTorch without its CUDA backend. Both work. The other source build recipes were not run. The tests that read these READMEs (the ExecuTorch pin checks and the pin updater) pass.

Each ExecuTorch export example told readers to install Torch-TensorRT with its own command, an editable install from a source checkout, which is not how most readers get it and disagreed with the guide. Two docstrings were also out of date. One said that Python cannot run a program whose outputs are not planned, which it can with ExecuTorch builds that have pytorch/executorch#23547. The other said a CPU input to a program that expects GPU inputs still works, with a copy, but ExecuTorch Python builds newer than the one 2.15 installs refuse it. A code comment in export_device_resident.py said both memory planning flags are needed, when only the input one is. The examples now link to the guide's install section, the multi-profile example says that it needs a build newer than 2.15 and links to the nightly install, and the docstrings and the comment say what happens today, including when a CPU input is copied and when it is refused. The device-resident example names the ExecuTorch change and the first build that has it. Tested on Linux x86-64 with an H100 and nightly Torch-TensorRT: ran export_device_resident.py and load_model_device_resident.py, which pass, and checked that the Python runtime refuses a CPU input for that program, and that a planned input fails while planned outputs work. The test files that read these examples pass.

…ction The ExecuTorch pages of the user guide were hard to follow and partly out of date: - The saving section carried its own install command, which disagreed with the guide's. - The install section opened with a command that gives no published wheel its delegate: only a nightly extra brings it today. No command on the page installed the delegate for Torch-TensorRT 2.15. The section also listed Windows, which 2.15 has no delegate for. - It said each delegate of a coalesced program runs on its own CUDA stream by default. Both now run on the calling thread's default stream. - The guide said engines on the shared activation scratch never replay a CUDA graph. In builds newer than 2.15 they do. - The guide said Python cannot run a program whose outputs are not planned. It can now, with ExecuTorch builds that have pytorch/executorch#23547, which both 2.15 and nightly install. - Features that 2.15 lacks were not marked, and two were missing: engines in a separate .ptd file, and optimization profiles. Weight streaming was not described at all. - The zero-copy KV cache section was a wall of text, and the list of save() options for this format named only six of them. - The comparison table said a .pt2 runs on Windows. save() refuses that format outside Linux. - The save() docstring said a target_device compile spec overrides the default target device. It does not move the engines: an engine runs on the GPU it was built on. - The saving section showed how to write an ETRecord with generate_etrecord=True. For a program with TensorRT engines that fails today with a SerializeError, on nightly and on 2.15. - Nothing said that TensorRT warns, on a green context stream, that the engine was built without green context constraints and may crash. - The guide's note said ordinary Torch-TensorRT builds keep their CUDA 12 support. The 2.15 release candidate and current nightly builds are built for CUDA 13 only. The guide is rewritten in shorter sentences. It keeps one install section, which the other pages now link to. It starts with the commands that work today. The 2.15 command pins the release and names its 0.2 delegate, from the test channel, so it keeps working when a newer 0.2 delegate ships with the next release candidate. It also lists the nightly channel for the ExecuTorch build that delegate pins, which that channel keeps only for a limited time, and the section says so. The nightly command follows, and then the plain form, which is enough once a release's extra brings the delegate. Each feature that 2.15 lacks is marked "Not in 2.15", which stays true after the next release branches, and a new section shows how to set a runtime option from C++ and from Python. The saving section now covers only the saving step: save() versus export(), every option of the executorch format, several methods in one file, and a condensed zero-copy KV cache section. The API page no longer says the module is Linux only, and the save() docstring points to the deployment guide, where the CUDA backend recipe moved. That docstring also stops saying that target_device moves the engines, and says that generate_etrecord does not work yet for programs with TensorRT engines. The export() docstring says the same. Windows is marked "not in 2.15", because the 2.15 release candidate publishes no Windows delegate wheel. The test that checks every printed install command requires --pre on any command that installs the extra from the nightly channel. The 2.15 command lists that channel only for the ExecuTorch build, and must not take --pre, or pip takes nightly builds of everything else. The test now skips the --pre requirement, and only that, for a user docs command that also names the delegate with == or ~=. Every other check, including the supported channel check, still runs on it. New cases pin the 2.15 command, the same command with an unsupported channel, and a nightly command without --pre. Tested on Linux x86-64 with an H100 and nightly Torch-TensorRT. Ran every Python block of the guide in order: they export the plain, coalesced, GPU-resident and separate engine file programs, run the first three, and set the runtime options from Python. A separate run loaded the separate engine file program, which fails without its .ptd and works with it. Built and ran the C++ snippets against the installed wheels, all but the optimization profile one, and ran one save() example and one export() example from the saving section. Checked these 2.15 differences against the Torch-TensorRT 2.15.0 release candidate: the re-export default, the Python runtime options, the separate engine file, and a CPU input to a GPU-resident program. On both versions, ran a GPU-resident program whose last part runs on the CUDA backend, with planned and with unplanned outputs. Windows was not tested. The pages build with Sphinx with no new warnings, and the tests that read these files pass. The new cases fail without the test change: the 2.15 command is rejected for lacking --pre. On Linux aarch64, the 2.15 command as written resolves in a fresh environment with pip --dry-run: torch-tensorrt 2.15.0, torch-tensorrt-executorch-runtime 0.2.0 and torch 2.15.0 from the test channel, ExecuTorch 1.6.0.dev20261009 from the nightly channel, TensorRT 11.3 from PyPI.