Repository navigation
Place pipeline components on the requested device and charge their pool - #1499
Merged
Merged
Conversation
A native CUDA pipeline run that failed anywhere below the driver surfaced
as "native CUDA decoder forward pass failed" and nothing else. anyhow's
Display prints only the outermost context, so every layer of diagnosis
the lower crates had carefully attached was discarded at the boundary
where a human would read it. Switching the two driver formats to the
alternate `{err:#}` selector prints the whole chain.
With the chain visible, four real dtype gaps became legible instead of
being guessed at:
* Reduce* on CPU accepted Int64 only for Sum. A graph reducing Int32,
or taking Max/Min/Prod over integers, hit the dtype guard rather
than the integer path that was already sitting there.
* SkipLayerNormalization on CPU required its optional `mean` and
`inv_std_var` outputs to carry the input dtype. ORT's schema types
them `float`, but exporters routinely emit them in the activation
dtype instead, and the write already narrows to whichever was
declared -- so both should be accepted. The bf16 skip-RMSNorm path
documents the same allowance.
* SkipLayerNormalization on CUDA made the mirror-image assumption,
deriving the stat pointers' dtype from the input rather than from
the declared output. Threading the resolved stat dtype through to
the kernel keeps a float32 stat pair on a bf16 input from being
written as bf16.
* Range on CPU had no Int32 branch.
Each gap gets tests at the boundary that was previously unreachable.
Note for reviewers: `onnx-runtime-ep-cuda --lib` has five failures on
this machine (two in graph::tests, three in matmul_nbits fp16 tests).
They reproduce identically on an unmodified main, so they are not from
this change.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e
The executor already computed `last_use` for every value, but the only consumer was `inplace_dead_inputs`, whose conditions are narrow enough (output must match an input in both dtype and shape) that in practice it covers elementwise chains and nothing else. Any operator that changes shape -- MatMul, Reshape, Attention -- allocated a buffer that then sat resident until the run ended. A graph is therefore charged for its peak node count rather than its peak live set. The vision encoder made the difference concrete. Its 2545 node outputs were all held simultaneously, costing roughly 23 GB at 448px and making full-resolution images unservable: a 960x686 image OOMed an 80 GB card. With dead values freed the same run peaks at 55.8 GB and completes. Four details worth calling out, because each is easy to get wrong: `DeviceBuffer` is a bare handle with no `Drop`. Dropping it from the value map looks like a release and is actually a leak -- the first version of this change did exactly that and reported 7075 releases while device memory did not move at all. The buffer has to go back through `ep.deallocate`. `buffer_shapes` has to be cleared alongside it. `ensure_buffer` treats a surviving shape entry as proof the allocation is still there and skips sizing it, so leaving one behind hands the next run a value it believes is backed and is not. Every other path that takes a buffer out of the map does the same thing for the same reason. A memoized loop-invariant `If` skips its branch on later runs and serves its outputs straight from the buffers an earlier run left resident, and that memo survives an eager run. Freeing one of those outputs turns the next skip into a missing-buffer error on the *second* request rather than the first. `try_move_host_output` already declines for exactly this reason; this pass is the second buffer-stealing path and now carries the same guard. The release is per-session and off by default because it is not safe under device graph capture. Capture bakes device pointers into the recorded graph, so freeing a buffer lets a later allocation land at a different address and replay then reads the wrong memory; the decoder fails immediately with a capture validation violation. Restricting this to the eager path is not sufficient either, since one session runs eager before it captures. Components that never capture -- the pipeline components, which run once per request -- opt in; the decoder does not. Aliased values (`pinned`) are never released, even after every alias has died. That is deliberate: `pinned` is a monotone record of "was ever an alias source", and releasing on the exact liveness test instead would be wrong for the decode memo, which restores a step's view table on the next run and expects the source buffer to still be resident. An earlier draft did try to reclaim view sources once their views died; it turned out to be unreachable (every view source is pinned) and, when measured, worth 36 MiB out of ~12 GB. Not worth a cross-run aliasing analysis. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e
Two independent gaps in how pipeline components were brought up. The device was never resolved. A component was loaded without consulting the session options that the rest of the engine had already used to pick a device, so `--device cuda:3` placed the decoder on GPU 3 and the vision encoder wherever the default landed. `resolved_native_device` reads the same session options, and routing now warns when a component that asked for CUDA has fallen back to CPU instead of failing silently and running two orders of magnitude slower. Everything else built in the same constructor reads the resolved device too. Leaving `pipeline_cuda_index`, the native CUDA memory plan, the authority domain and the CUDA authority on the raw `config.native_device` would have been worse than not resolving at all: components would land on the GPU while the governor meant to charge them was constructed as if no GPU were in play, sizing the device tier from the provisional constant rather than the card. The device pool was never accounted either. A component session builds its own execution provider, which sizes a standing pool for itself. Nothing told the memory governor about it, so the decoder's admission control and the component each measured the same free VRAM and each concluded it could take ~90% of it. Adopting the governor under a new `Holder::PipelineComponentPool` turns that pool into a claim the rest of the engine can see. Adoption failure is logged, not fatal: a provider holding no standing pool legitimately reports zero, and a governor that refuses the claim leaves the component exactly as unaccounted as it was before. That is worth saying out loud and not worth failing a load over. `build_native_pipeline_components` deliberately keeps resolving against default session options. Its only job is to name the components in an unsupported-plan error, so it must not claim device memory to do it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e
justinchuby
force-pushed
the
fix/pipeline-component-device-accounting
branch
from
August 19, 2026 21:31
c7d7da2 to
4cd82e2
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1499 +/- ##
=========================================
+ Coverage 0 80.25% +80.25%
=========================================
Files 0 377 +377
Lines 0 166504 +166504
Branches 0 166504 +166504
=========================================
+ Hits 0 133629 +133629
- Misses 0 28036 +28036
- Partials 0 4839 +4839
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two independent gaps in how pipeline components were brought up.
The device was never resolved. A component was loaded without consulting the session options the rest of the engine had already used to pick a device, so
--device cuda:3put the decoder on GPU 3 and the vision encoder wherever the default landed.resolved_native_devicenow reads the same session options, and routing warns when a component that asked for CUDA fell back to CPU -- previously that was silent, and the only symptom was running two orders of magnitude slower.The device pool was never accounted. A component session builds its own execution provider, which sizes a standing pool for itself. Nothing told the memory governor, so the decoder's admission control and the component each measured the same free VRAM and each concluded it could take ~90% of it. Adopting the governor under a new
Holder::PipelineComponentPoolmakes that pool a claim the rest of the engine can see.Adoption failure is logged rather than fatal: a provider holding no standing pool legitimately reports zero, and a governor that refuses the claim leaves the component exactly as unaccounted as it was before. Worth saying out loud; not worth failing a load over.
Builds on #1498 (which added
set_release_dead_valueson the sameloadpath).Validation
onnx-genai-engine --features native-backend,cuda: 575 + others, 0 failed. End-to-end on a 30B INT4 pipeline with--device cuda:3, confirmed vianvidia-smithat both decoder and vision encoder land on GPU 3.CI runners are still down, so this was validated locally.