Place pipeline components on the requested device and charge their pool

crates/onnx-genai-engine/src/engine/memory_plan.rs

@@ -80,6 +80,17 @@ pub(crate) enum Holder {
     /// Seeding that tier with the KV sub-budget instead was what stopped every
     /// other holder from joining it.
     FixedDeviceReservation,
+    /// The standing device pool held by a pipeline component session (a vision
+    /// encoder, say) that is loaded lazily alongside the decoder.
+    ///
+    /// These sessions build their own execution provider, so without a lease
+    /// their device memory is spent behind the ledger's back: the decoder's
+    /// governor keeps sizing its ceiling from a `measured_free` that silently
+    /// includes bytes another session on the same device already holds. Two
+    /// co-resident tenants each admitting against the whole card is how a
+    /// 3264-patch vision prefill drove a 30B decoder into `CUDA_ERROR_OUT_OF_MEMORY`
+    /// with the card genuinely full.
+    PipelineComponentPool,
 }
 
 impl Holder {
@@ -97,7 +108,8 @@ impl Holder {
     /// the variants are unconditional even when only one build configuration
     /// constructs them.
     #[allow(dead_code)]
-    pub(crate) const ALL: [Holder; 10] = [
+    pub(crate) const ALL: [Holder; 11] = [
+        Holder::PipelineComponentPool,
         Holder::KvPool,
         Holder::PipelineKvPool,
         Holder::DraftKvPool,
@@ -126,6 +138,7 @@ impl Holder {
             Holder::NativeKvCache => 8,
             Holder::Activations => 6,
             Holder::FixedDeviceReservation => 7,
+            Holder::PipelineComponentPool => 11,
         })
     }
 
@@ -143,6 +156,9 @@ impl Holder {
             Holder::NativeKvCache => MemoryRole::KvCache,
             Holder::Activations => MemoryRole::Activation,
             Holder::FixedDeviceReservation => MemoryRole::Weights,
+            // A standing pool the component keeps for as long as the session
+            // lives, not a per-step scratch buffer.
+            Holder::PipelineComponentPool => MemoryRole::Workspace { step_scoped: false },
         }
     }
 
@@ -159,6 +175,7 @@ impl Holder {
             Holder::NativeKvCache => "native decode KV tensors",
             Holder::Activations => "activations",
             Holder::FixedDeviceReservation => "fixed device reservation",
+            Holder::PipelineComponentPool => "pipeline component device pool",
         }
     }
 
@@ -184,7 +201,8 @@ impl Holder {
             // for a device session.
             | Holder::NativeKvCache
             | Holder::Activations
-            | Holder::FixedDeviceReservation => Tier::Device,
+            | Holder::FixedDeviceReservation
+            | Holder::PipelineComponentPool => Tier::Device,
         }
     }
 }

crates/onnx-genai-engine/src/native_component.rs

@@ -119,14 +119,25 @@ impl NativeComponentSession {
     }
 
     /// Load an ONNX model on the requested native device as a neutral component.
-    pub fn load(path: &Path, device: NativeDecodeDevice) -> anyhow::Result<Self> {
+    ///
+    /// `governor`, when present, is the ledger the rest of the engine admits
+    /// against on this device. A component session builds its own execution
+    /// provider, which sizes any standing device pool for itself; adopting the
+    /// governor turns that pool into a claim the decoder's admission control can
+    /// see. Without it two co-resident sessions each measure the same free VRAM
+    /// and each conclude they may use ~90% of it.
+    pub fn load(
+        path: &Path,
+        device: NativeDecodeDevice,
+        governor: Option<&dyn onnx_runtime_memory_governor::MemoryGovernor>,
+    ) -> anyhow::Result<Self> {
         let requested_cuda = matches!(&device, NativeDecodeDevice::Cuda { .. });
         let preference = match device {
             NativeDecodeDevice::Cpu => DevicePreference::Cpu,
             NativeDecodeDevice::Cuda { index } => DevicePreference::Gpu { index },
             NativeDecodeDevice::Plugin { .. } => DevicePreference::Cpu,
         };
-        let session = InferenceSession::builder()
+        let mut session = InferenceSession::builder()
             .model(path)
             .device(preference)
             .build()
@@ -136,13 +147,40 @@ impl NativeComponentSession {
                     path.display()
                 )
             })?;
+        // A component graph is run once per request and never records a device
+        // graph, so its intermediates can be freed as they die. The vision
+        // encoder is why this matters: holding every one of its 2545 node
+        // outputs at once cost ~23 GB at 448px and put full-resolution images
+        // out of reach entirely.
+        session.set_release_dead_values(true);
         if requested_cuda && let Some(report) = session.execution_provider_fallback_report() {
             tracing::warn!(
                 model = %path.display(),
                 fallback = %report,
                 "native CUDA pipeline component fell back to CPU"
             );
         }
+        if let Some(governor) = governor {
+            let holder = crate::engine::memory_plan::Holder::PipelineComponentPool;
+            // Not fatal: a provider that holds no standing pool 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.
+            match session.adopt_memory_governor(governor, holder.tier(), holder.id()) {
+                Ok(bytes) => tracing::debug!(
+                    model = %path.display(),
+                    governed_bytes = bytes,
+                    holder = holder.name(),
+                    "pipeline component device pool adopted by the engine memory governor"
+                ),
+                Err(error) => tracing::warn!(
+                    model = %path.display(),
+                    %error,
+                    "pipeline component device pool could not be charged to the engine memory \
+                     governor; admission on this device stays optimistic"
+                ),
+            }
+        }
         Self::new(session).map_err(anyhow::Error::from)
     }
 }

crates/onnx-genai-engine/src/pipeline/mod.rs

@@ -411,16 +411,19 @@ fn build_native_pipeline_components(
     use crate::native_component::NativeComponentSession;
     use onnx_genai_metadata::ComponentSession;
 
+    // Deliberately CPU-neutral: this path exists only to name the components in
+    // an unsupported-plan error, so it must not claim device memory to do it.
     let device = crate::engine::resolve_native_decode_device(
         config.native_device.clone(),
         &SessionOptions::default(),
     )?;
     let mut components: std::collections::BTreeMap<String, Box<dyn ComponentSession>> =
         std::collections::BTreeMap::new();
     for (name, path) in &directory.model_paths {
-        let session = NativeComponentSession::load(path, device.clone()).with_context(|| {
-            format!("failed to construct pipeline component '{name}' on the native backend")
-        })?;
+        let session =
+            NativeComponentSession::load(path, device.clone(), None).with_context(|| {
+                format!("failed to construct pipeline component '{name}' on the native backend")
+            })?;
         components.insert(name.clone(), Box::new(session));
     }
     Ok(components)
@@ -519,11 +522,17 @@ fn build_step_component_session<'a>(
             // CPU pipeline keeps it on CPU. Inputs/outputs still cross the host
             // `ComponentTensor` seam, so only one small `inputs_embeds` row
             // round-trips host<->device per step; the decoder's KV never does.
-            let native =
-                crate::native_component::NativeComponentSession::load(path, native_device.clone())
-                    .with_context(|| {
-                        format!("failed to load native every_step component '{component}'")
-                    })?;
+            let native = crate::native_component::NativeComponentSession::load(
+                path,
+                native_device.clone(),
+                // The governor does not exist yet at this point in construction:
+                // every_step components are built before the resource governor,
+                // which is itself sized from the models on disk. The lazily
+                // loaded prompt components (routing.rs) are the ones that can and
+                // do adopt it.
+                None,
+            )
+            .with_context(|| format!("failed to load native every_step component '{component}'"))?;
             return Ok(Box::new(native));
         }
         #[cfg(not(feature = "native-backend"))]
@@ -924,6 +933,23 @@ impl PipelineEngine {
             EngineDecodeBackend::Auto => resolve_auto_pipeline_backend(&directory)?,
         };
         let plan = PipelinePlan::from_spec(&directory.spec, schedulers)?;
+        // Resolve the native device from the session's execution providers, the
+        // same way `engine::load` does for the standalone decoder.
+        //
+        // `config.native_device` is only ever `Some` when an operator named a
+        // device explicitly, so reading it alone made `--ep cuda` resolve to
+        // `Cpu` for every native pipeline session. The decoder survived that on
+        // its own load path, which does consult the providers; the lazily loaded
+        // prompt components (a vision tower) did not, and ran the whole encoder
+        // on CPU with no warning -- a >20 minute prefill where the GPU took 13s.
+        #[cfg(feature = "native-backend")]
+        let resolved_native_device = match backend {
+            PipelineBackend::Native => Some(crate::engine::resolve_native_decode_device(
+                config.native_device.clone(),
+                &session_options,
+            )?),
+            PipelineBackend::Ort => config.native_device.clone(),
+        };
         let graph_memory = combine_graph_memory(
             directory
                 .model_paths
@@ -983,7 +1009,7 @@ impl PipelineEngine {
         // resident weights on any GPU.
         #[cfg(all(feature = "cuda", feature = "native-backend"))]
         let pipeline_cuda_index = if backend == PipelineBackend::Native {
-            native_decoder_device(config.native_device.as_ref()).cuda_index()
+            native_decoder_device(resolved_native_device.as_ref()).cuda_index()
         } else {
             None
         };
@@ -1008,7 +1034,7 @@ impl PipelineEngine {
         #[cfg(all(feature = "cuda", feature = "native-backend"))]
         let native_cuda_plan = backend == PipelineBackend::Native
             && matches!(
-                native_decoder_device(config.native_device.as_ref()),
+                native_decoder_device(resolved_native_device.as_ref()),
                 crate::native_decode_device::NativeDecodeDevice::Cuda { .. }
             );
         #[cfg(not(all(feature = "cuda", feature = "native-backend")))]
@@ -1091,7 +1117,7 @@ impl PipelineEngine {
         );
         #[cfg(all(feature = "cuda", feature = "native-backend"))]
         let authority_domain = if backend == PipelineBackend::Native {
-            match native_decoder_device(config.native_device.as_ref()) {
+            match native_decoder_device(resolved_native_device.as_ref()) {
                 crate::native_decode_device::NativeDecodeDevice::Cuda { index } => {
                     crate::memory_authority::DeviceCompatibilityDomain::Cuda(index.unwrap_or(0))
                 }
@@ -1113,7 +1139,7 @@ impl PipelineEngine {
         )?;
         #[cfg(all(feature = "cuda", feature = "native-backend"))]
         let native_cuda_authority = if backend == PipelineBackend::Native {
-            match native_decoder_device(config.native_device.as_ref()) {
+            match native_decoder_device(resolved_native_device.as_ref()) {
                 crate::native_decode_device::NativeDecodeDevice::Cuda { .. } => {
                     Some(resource_governor.device_authority())
                 }
@@ -1323,6 +1349,9 @@ impl PipelineEngine {
             native_retained_decoder: None,
             prefill_chunk_size,
             paged,
+            #[cfg(feature = "native-backend")]
+            native_device: resolved_native_device,
+            #[cfg(not(feature = "native-backend"))]
             native_device: config.native_device.clone(),
             native_prompt_sessions: RefCell::new(BTreeMap::new()),
             #[cfg(feature = "cuda")]

crates/onnx-genai-engine/src/pipeline/routing.rs

@@ -316,10 +316,26 @@ impl PipelineEngine {
                 format!("native pipeline prompt component '{component}' has no model path")
             })?;
         let device = super::native_decoder_device(self.native_device.as_ref());
-        let session = crate::native_component::NativeComponentSession::load(path, device)
-            .with_context(|| {
-                format!("failed to load native pipeline prompt component '{component}'")
-            })?;
+        if matches!(device, crate::native_decode_device::NativeDecodeDevice::Cpu) {
+            // The silent case is the expensive one: a vision tower on CPU still
+            // produces correct embeddings, just two orders of magnitude slower,
+            // so nothing fails and nothing is logged. Say it once, at load.
+            tracing::warn!(
+                component,
+                model = %path.display(),
+                "native pipeline prompt component is running on CPU; if this package has a CUDA \
+                 execution provider available, prompt-phase encoders (a vision tower) will be far \
+                 slower than the decoder that follows them"
+            );
+        }
+        let session = crate::native_component::NativeComponentSession::load(
+            path,
+            device,
+            Some(self.resource_governor.memory()),
+        )
+        .with_context(|| {
+            format!("failed to load native pipeline prompt component '{component}'")
+        })?;
         self.native_prompt_sessions
             .borrow_mut()
             .insert(component.to_string(), Box::new(session));

crates/onnx-genai-server/src/driver.rs

@@ -210,7 +210,11 @@ impl DriverFailure {
             )
         });
         Self {
-            message: error.to_string(),
+            // Anyhow's Display shows only the outermost context, which for a
+            // decode failure is the generic "forward pass failed" wrapper. The
+            // alternate form keeps the whole chain, and this message is the
+            // only thing the client ever sees.
+            message: format!("{error:#}"),
             kind: if memory_overload {
                 DriverFailureKind::MemoryOverload
             } else {
@@ -1123,7 +1127,7 @@ fn run_static_batch_until_idle(
         route_continuous_admissions(manager.poll_admissions(), &mut routes);
         if let Err(err) = manager.step() {
             let mut failure = DriverFailure::from_engine_error(&err);
-            failure.message = format!("continuous batch generation failed: {err}");
+            failure.message = format!("continuous batch generation failed: {err:#}");
             for (_, mut route) in routes.drain() {
                 if let Some(sender) = route.admission.take() {
                     let _ = sender.send(Err(failure.clone()));