feat(pipeline): drive every_step components via ComponentSession trait (native multi-component inc1)

.squad/decisions/inbox/mary-native-pipeline-plan.md

@@ -0,0 +1,101 @@
+### 2026-07-30: Native multi-component pipeline decode — refactor plan (issue #384)
+**By:** Mary
+**What:** Scoped plan to make `PipelineDecodeLoopBackend` drive **native** component
+sessions (nxrt custom-EP), unblocking native decode of pipelined multi-component
+models (Qwen3.6-35B-A3B = embedding + decoder + vision, and any multimodal pipeline).
+This is the GAP 3 work referenced in `mary-35b-a3b-blocker.md`.
+
+#### 1. Where the pipeline loop hardcodes the concrete ORT `Session` / ORT `Value`
+
+Grepped `&'a Session`, `.run(`, `Session,`, `Value` across the pipeline module.
+The decode loop that must become backend-neutral is **`flat_autoregressive.rs` +
+`paged_decode.rs`** (only these two construct/use `PipelineDecodeLoopBackend`;
+`nested_autoregressive.rs` and `iterative.rs` drive their own loops and are out of
+scope for the AR-decoder path).
+
+Concrete-ORT couplings in `paged_decode.rs::PipelineDecodeLoopBackend`:
+
+| Field / call | Coupling |
+|---|---|
+| `decoder: &'a Session` | concrete ORT decoder session (GAP: Inc2) |
+| `step_components: Vec<(StepComponentBinding, &'a Session)>` | concrete ORT **every_step** sessions (GAP: **Inc1**) |
+| `pool: &'a mut PipelineTensors` = `HashMap<String, ort::Value>` | pool holds ORT `Value` (the value-type seam) |
+| `static_cross_kv: Vec<(String, Arc<Value>)>` | ORT `Value` cross-attn KV (Inc3) |
+| `run_step_components`: builds `Vec<(String, Value)>`, calls `session.run(&refs)`, inserts `Value` | ORT value + ORT session run (Inc1) |
+| `decoder_extras`: `clone_value`, `Value::alias_with_shape` | ORT `Value` decoder binding (Inc2) |
+| `next_logits`: `run_decode_step_with_extra(self.decoder, ...)`, `mirror_present_kv_to_pages`, `extract_next_token_logits_with_io` | ORT decoder step + ORT KV mirror (Inc2/Inc3) |
+
+`flat_autoregressive.rs` pairs each binding with `self.models.session(name) -> &Session`
+(lines ~152–169) and constructs the backend (line 187). `pipeline/mod.rs` imports the
+concrete `Session, Value, DataType` from `onnx_genai_ort` and defines
+`PipelineTensors = HashMap<String, Value>` (mod.rs:31–51).
+
+#### 2. Trait boundary — THE VALUE-TYPE SEAM (verdict)
+
+**Does the ORT `Session` implement `ComponentSession`?** No — but an adapter already
+exists: `onnx_genai_ort::OrtComponentSession` wraps an **owned** `Session` and
+implements `ComponentSession` (ort/src/component.rs:130). `NativeComponentSession`
+implements the same trait (engine/src/native_component.rs:164). Both are object-safe
+(`Box<dyn ComponentSession>`).
+
+**Do `ComponentSession::run` inputs/outputs use a backend-neutral `Value` or ORT
+`Value`?** **VERDICT: the seam is a backend-neutral, host-resident `ComponentTensor`
+(raw little-endian element bytes + dtype + static shape), NOT ORT `Value` and NOT an
+nxrt tensor.** `ComponentSession::run(&mut self, &[(&str, &ComponentTensor)]) ->
+Vec<(String, ComponentTensor)>` (metadata/src/component.rs:283). Each backend adapter
+translates at its own boundary: ORT via `Value::to_raw_bytes`/`from_raw_bytes`
+(ort/src/component.rs `to_value`/`from_value`); native via `Tensor::from_raw`/`as_bytes`
+(native_component.rs `to_native_tensor`/`from_native_tensor`).
+
+Consequence: the pipeline **pool holds ORT `Value`**, but the trait speaks
+`ComponentTensor`. So routing step components through `dyn ComponentSession` requires a
+**pool-`Value` ⇄ `ComponentTensor` conversion at the loop boundary** — this host
+round-trip **is the crux of the work**. `DataType ⇄ ComponentDataType` `From` impls
+already exist in `onnx_genai_ort`. The conversion is a host copy (`numel * dtype.size`);
+for the decoder (Inc2) this is the KV-cache-sized cost and must eventually be avoided by
+keeping tensors backend-native across the seam, but for small every_step embedding
+outputs it is negligible (same order as the existing per-step `clone_value` the decoder
+already pays).
+
+#### 3. Increment breakdown
+
+- **Inc1 (THIS task):** Route the `every_step` (step) components through
+  `dyn ComponentSession`. Change `step_components` to
+  `Vec<(StepComponentBinding, Box<dyn ComponentSession + 'a>)>`. Rewrite
+  `run_step_components` to convert pool `Value → ComponentTensor`, call the trait, and
+  convert `ComponentTensor → Value` back into the pool. ORT default path wraps
+  `&Session` in a **borrowing** ORT adapter (`OrtComponentSessionRef`), behaviour
+  unchanged. Native path loads the component via `NativeComponentSession`. **Prove one
+  NATIVE every_step component (the gemma4-vlm `embedding`) runs inside the loop with
+  ORT-vs-native token parity** while the decoder stays ORT (hybrid). Deliverable =
+  solve the value seam for step components + wire embedding + parity test.
+- **Inc2:** The decoder itself (`decoder: &Session`, `run_decode_step_with_extra`,
+  `decoder_extras`, logits extraction) becomes trait-driven. This forces the KV-cache
+  ownership question (see risks) and the decoder-sized value-seam copy — the real perf
+  seam. Reconcile with `NativeDecodeSession` (the single-graph native decode path).
+- **Inc3:** Cross-component value handoff — `static_cross_kv` (encoder cross-attn KV),
+  device placement/transfer between components, and the vision `prompt_only` stage; the
+  full 35B-A3B embedding+decoder+vision chain end-to-end on native.
+
+#### 4. Risks
+
+- **Value-seam copies:** every seam crossing is a host copy. Fine for embeddings (Inc1),
+  costly for the decoder KV (Inc2) — Inc2 should keep tensors backend-native across the
+  seam or add a zero-copy fast path behind the trait rather than always round-tripping
+  through host bytes.
+- **KV-cache ownership (Inc2):** the single-graph native path is `NativeDecodeSession`,
+  which owns its own persistent KV state. A pipeline-driven decoder needs the loop
+  (`DecodeState`, paged mirror) to own/advance KV. Reconciling these two ownership
+  models is the central Inc2 design decision.
+- **Device placement:** native EP tensors may be device-resident; the `ComponentTensor`
+  seam is host-only (`as_raw_bytes`/`to_raw_bytes` reject device tensors). Cross-device
+  step→decoder handoff (Inc3) needs explicit host staging or a device-aware seam.
+- **Feature gating (#436/#441 class):** native paths are `#[cfg(feature =
+  "native-backend")]`; imports must be cfg-correct so both cuda and non-cuda,
+  native and non-native builds compile without unused-import warnings.
+
+**Why:** `PipelineDecodeLoopBackend` owning ORT `Value`/`Session` is the last blocker to
+native pipelined decode. Establishing the value-type seam verdict (neutral
+`ComponentTensor`, not ORT `Value`) up front prevents Inc2/Inc3 from re-litigating the
+boundary, and the Inc1 every_step slice proves the seam end-to-end with token parity on a
+tiny CPU fixture before the heavier decoder/KV work.

crates/onnx-genai-engine/Cargo.toml

@@ -20,6 +20,10 @@ required-features = ["cuda", "native-backend"]
 name = "weight_offload_native_cuda_e2e"
 required-features = ["cuda", "native-backend"]
 
+[[test]]
+name = "native_step_component_parity"
+required-features = ["native-backend"]
+
 [features]
 default = []
 native-backend = [

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

@@ -153,18 +153,23 @@ impl PipelineEngine {
             .models
             .session(&ar.decoder)
             .with_context(|| format!("pipeline decoder '{}' was not loaded", ar.decoder))?;
-        // Pair every `every_step` binding with its loaded session. This is the
-        // generic replacement for the old one-output `inputs_embeds` fusion.
+        // Pair every `every_step` binding with a backend-neutral component
+        // session. By default this borrows the already-loaded ORT session
+        // (behaviour unchanged); components named in
+        // `ONNX_GENAI_PIPELINE_NATIVE_STEP_COMPONENTS` are instead loaded and
+        // driven through the native nxrt backend, proving the value-type seam
+        // (native multi-component inc1) while the decoder stays on ORT.
+        let native_step_components = native_step_component_set();
         let step_components = step_bindings
             .into_iter()
             .map(|binding| {
-                let session = self.models.session(&binding.component).with_context(|| {
-                    format!(
-                        "pipeline every_step component '{}' was not loaded",
-                        binding.component
-                    )
-                })?;
-                Ok((binding, session))
+                let component: Box<dyn onnx_genai_metadata::ComponentSession> =
+                    build_step_component_session(
+                        &self.models,
+                        &binding.component,
+                        &native_step_components,
+                    )?;
+                Ok((binding, component))
             })
             .collect::<anyhow::Result<Vec<_>>>()?;
         let tokenizer = self
@@ -234,6 +239,12 @@ impl PipelineEngine {
                 mirror.mirrored_tokens
             }
         });
+        // The backend now owns its every_step component sessions behind
+        // `Box<dyn ComponentSession>`, so it carries drop glue and its borrows of
+        // `tensors` / `self` would otherwise live to end of scope. Everything
+        // needed downstream has been copied out above, so release it explicitly
+        // before the paged-sequence retirement and the `tensors` move below.
+        drop(backend);
         let result = match result {
             Ok(result) => result,
             Err(error) => {

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

@@ -365,7 +365,76 @@ fn build_native_pipeline_components(
     Ok(components)
 }
 
-/// Components whose graphs contain only deterministic operators.
+/// Every_step components the operator explicitly requested be run on the native
+/// nxrt backend, from `ONNX_GENAI_PIPELINE_NATIVE_STEP_COMPONENTS` (a
+/// comma-separated list of component names). Empty/unset means all every_step
+/// components run on the default ORT backend, so the ORT decode path is
+/// unchanged. This is the injection seam for the native multi-component inc1
+/// hybrid: it drives named every_step components natively inside an
+/// otherwise-ORT pipeline decode loop, proving the value-type seam.
+fn native_step_component_set() -> BTreeSet<String> {
+    std::env::var("ONNX_GENAI_PIPELINE_NATIVE_STEP_COMPONENTS")
+        .ok()
+        .map(|list| {
+            list.split(',')
+                .map(str::trim)
+                .filter(|name| !name.is_empty())
+                .map(str::to_string)
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
+/// Build the backend-neutral [`ComponentSession`](onnx_genai_metadata::ComponentSession)
+/// for one every_step component.
+///
+/// By default this borrows the already-loaded ORT session
+/// ([`OrtComponentSessionRef`]) so the ORT decode path is behaviour-identical.
+/// A component named in `native_components` is instead loaded and driven through
+/// the native nxrt backend, so the same decode loop drives both backends through
+/// the trait with no forked code path.
+fn build_step_component_session<'a>(
+    models: &'a PipelineModels,
+    component: &str,
+    native_components: &BTreeSet<String>,
+) -> anyhow::Result<Box<dyn onnx_genai_metadata::ComponentSession + 'a>> {
+    if native_components.contains(component) {
+        #[cfg(feature = "native-backend")]
+        {
+            let path = models
+                .directory
+                .model_paths
+                .get(component)
+                .with_context(|| {
+                    format!("native every_step component '{component}' has no model path")
+                })?;
+            // The every_step slice (inc1) stages tensors through the host
+            // `ComponentTensor` seam, so the native component runs on CPU; device
+            // placement of the pipeline decoder is the later (inc2/inc3) work.
+            let native = crate::native_component::NativeComponentSession::load(
+                path,
+                crate::native_decode::NativeDecodeDevice::Cpu,
+            )
+            .with_context(|| format!("failed to load native every_step component '{component}'"))?;
+            return Ok(Box::new(native));
+        }
+        #[cfg(not(feature = "native-backend"))]
+        {
+            anyhow::bail!(
+                "every_step component '{component}' was requested on the native backend via \
+                 ONNX_GENAI_PIPELINE_NATIVE_STEP_COMPONENTS, but this build was compiled without \
+                 the 'native-backend' feature. Rebuild with `--features native-backend`."
+            );
+        }
+    }
+    let session = models
+        .session(component)
+        .with_context(|| format!("pipeline every_step component '{component}' was not loaded"))?;
+    Ok(Box::new(onnx_genai_ort::OrtComponentSessionRef::new(
+        session,
+    )))
+}
+
 ///
 /// Read once at load, because a component's declared phase says when it runs,
 /// never that it is pure — a graph with `RandomNormal` in it would otherwise be
@@ -1303,6 +1372,46 @@ fn coerce_value_to_dtype(value: &Value, target: DataType) -> anyhow::Result<Valu
     }
 }
 
+/// Convert a pool ORT [`Value`] into a backend-neutral [`ComponentTensor`] for
+/// the [`ComponentSession`](onnx_genai_metadata::ComponentSession) seam.
+///
+/// This is the pipeline side of the value-type seam: the decode-loop pool holds
+/// ORT `Value`s, but every_step components run through the backend-neutral trait
+/// whose boundary is a host-resident `ComponentTensor` (raw little-endian element
+/// bytes). The copy is `numel * dtype.size_of()` bytes; for the small every_step
+/// embedding outputs this is negligible relative to the decoder step.
+fn value_to_component_tensor(
+    value: &Value,
+) -> anyhow::Result<onnx_genai_metadata::ComponentTensor> {
+    let dtype = onnx_genai_metadata::ComponentDataType::from(value.dtype());
+    let bytes = value.to_raw_bytes()?;
+    onnx_genai_metadata::ComponentTensor::from_raw(dtype, value.shape().to_vec(), bytes)
+        .map_err(Into::into)
+}
+
+/// Convert a [`ComponentTensor`] produced by a component back into a pool ORT
+/// [`Value`]. Inverse of [`value_to_component_tensor`].
+fn component_tensor_to_value(
+    tensor: &onnx_genai_metadata::ComponentTensor,
+) -> anyhow::Result<Value> {
+    let dtype = DataType::from(tensor.dtype());
+    Value::from_raw_bytes(tensor.as_bytes().to_vec(), tensor.shape(), dtype).map_err(Into::into)
+}
+
+/// Build the running-token seed as a neutral `int64` [`ComponentTensor`] of shape
+/// `[1, seq]`, matching the ORT `Value::from_slice_i64` the loop previously fed.
+fn token_seed_component_tensor(
+    ids: &[i64],
+) -> anyhow::Result<onnx_genai_metadata::ComponentTensor> {
+    let bytes: Vec<u8> = ids.iter().flat_map(|v| v.to_le_bytes()).collect();
+    onnx_genai_metadata::ComponentTensor::from_raw(
+        onnx_genai_metadata::ComponentDataType::Int64,
+        vec![1, ids.len() as i64],
+        bytes,
+    )
+    .map_err(Into::into)
+}
+
 #[derive(Debug, Clone)]
 struct IterativePlan {
     /// The component re-invoked once per step.

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

@@ -7,6 +7,7 @@
 //! autoregressive driver constructs these adapters across the module boundary.
 
 use super::*;
+use onnx_genai_metadata::ComponentSession;
 
 /// Whether this error is the KV page pool being full, rather than a fault.
 ///
@@ -83,9 +84,12 @@ pub(crate) struct PipelineDecodeLoopBackend<'a> {
     /// Shared tensor pool: external inputs + prompt-phase outputs + the
     /// per-step outputs of the `every_step` components (refreshed each step).
     pub(crate) pool: &'a mut PipelineTensors,
-    /// Declared `every_step` components (with their loaded sessions), executed in
-    /// topological order on every step before the decoder runs.
-    pub(crate) step_components: Vec<(StepComponentBinding, &'a Session)>,
+    /// Declared `every_step` components (paired with their backend-neutral
+    /// [`ComponentSession`]), executed in topological order on every step before
+    /// the decoder runs. Boxed behind the trait so the same loop drives an ORT
+    /// session (via `OrtComponentSessionRef`) or a native nxrt component with no
+    /// forked code path — this is the value-type seam for the every_step slice.
+    pub(crate) step_components: Vec<(StepComponentBinding, Box<dyn ComponentSession + 'a>)>,
     /// `(source_endpoint, decoder_input_port)` routing recomputed each step.
     pub(crate) decoder_in_edges: Vec<(String, String)>,
     /// Static encoder-produced cross-attention KV bound to the decoder every
@@ -124,42 +128,44 @@ impl PipelineDecodeLoopBackend<'_> {
             return Ok(());
         }
         let ids: Vec<i64> = seed.iter().map(|&t| i64::from(t)).collect();
-        let seq = ids.len() as i64;
-        for (binding, session) in &self.step_components {
-            let mut inputs: Vec<(String, Value)> =
+        // Disjoint borrows: the step sessions run `&mut` while the pool is read
+        // for inputs and written for outputs.
+        let Self {
+            step_components,
+            pool,
+            ..
+        } = self;
+        for (binding, session) in step_components.iter_mut() {
+            let mut inputs: Vec<(String, onnx_genai_metadata::ComponentTensor)> =
                 Vec::with_capacity(binding.routed_inputs.len() + 1);
             for routed in &binding.routed_inputs {
-                let value = self
-                    .pool
+                let value = pool
                     .get(&routed.endpoint)
                     .or_else(|| {
                         routed
                             .routed_from
                             .as_deref()
-                            .and_then(|from| self.pool.get(from))
+                            .and_then(|from| pool.get(from))
                     })
                     .with_context(|| routed.missing_message.clone())?;
-                inputs.push((
-                    routed.port.clone(),
-                    coerce_value_to_dtype(value, routed.dtype)?,
-                ));
+                let coerced = coerce_value_to_dtype(value, routed.dtype)?;
+                inputs.push((routed.port.clone(), value_to_component_tensor(&coerced)?));
             }
             if let Some(port) = &binding.token_input {
-                inputs.push((port.clone(), Value::from_slice_i64(&ids, &[1, seq])?));
+                inputs.push((port.clone(), token_seed_component_tensor(&ids)?));
             }
             let refs = inputs
                 .iter()
-                .map(|(name, value)| (name.as_str(), value))
+                .map(|(name, tensor)| (name.as_str(), tensor))
                 .collect::<Vec<_>>();
             let outputs = session.run(&refs).map_err(|e| {
-                anyhow::anyhow!(
-                    "ORT every_step component '{}' failed: {e}",
-                    binding.component
-                )
+                anyhow::anyhow!("every_step component '{}' failed: {e}", binding.component)
             })?;
-            for (name, value) in session.output_names().iter().zip(outputs) {
-                self.pool
-                    .insert(format!("{}.{}", binding.component, name), value);
+            for (name, tensor) in outputs {
+                pool.insert(
+                    format!("{}.{}", binding.component, name),
+                    component_tensor_to_value(&tensor)?,
+                );
             }
         }
         Ok(())