feat(pipeline): native CUDA device-KV decoder via inputs_embeds (native multi-component inc3a)

.squad/decisions/inbox/mary-pipeline-inc3-cuda-design.md

@@ -0,0 +1,103 @@
+# Inc3 design — CUDA-target native decoder with `inputs_embeds` / routed step inputs
+
+Author: Mary (native-pipeline). Refs #384. Stacks on #479 (Inc2b native device-KV
+decoder). Scope: **text decoder on the CUDA EP only** — paged cross-request reuse
+and vision cross-KV stay deferred (Inc4+).
+
+## Why this is the critical-path unblock
+
+Inc2b proved the native device-KV decoder end-to-end **on CPU**. The whole point of
+the native path is to run the 35B-A3B decoder **on-GPU** with a device-resident KV
+cache. Today the native decoder *refuses* a CUDA target the moment the decoder is
+driven by `inputs_embeds` (a fused VLM decoder) or any routed host step input — so
+the GPU native-decode path for a multi-component / VLM model is blocked at load.
+
+## The refusal — exact sites and root cause
+
+Two guards, both hit for a CUDA + `inputs_embeds` decoder:
+
+1. **Load-time (the load-bearing one):**
+   `native_decode/load.rs:434-446` — when `device == Cuda` and
+   (`sequence_source != TokenIds` **or** any `step_input.source` is
+   `InputsEmbeds | Routed`), it `bail!`s before constructing `DecodeCudaState`.
+2. **Runtime:** `native_decode/mod.rs:275-280` — `decode_with_step_inputs` bails if
+   `self.cuda.is_some()` and `step_inputs` is non-empty.
+
+**Root cause (verified by reading `cuda.rs`): an unimplemented on-device binding
+path, NOT a correctness/architecture barrier.** `DecodeCudaState`
+(`cuda.rs:731-1046`) hardwires the *sequence* input as an `Int64 [1,1]` `input_ids`
+device binding (`input_ids_binding`, `cuda.rs:944-951`) and writes it each step via
+`write_decode_inputs(token_id, position)` (`cuda.rs:1118-1125`). There is:
+  - no allocation of a **float `inputs_embeds [1, 1, hidden]`** device binding, and
+  - no path to **upload a routed host tensor** into a device binding each step.
+
+`DecodeCudaIo` (`cuda.rs:174-179`) only carries `input_ids / attention_mask /
+position_ids / logits`. Everything else (KV device-residency, mask growth,
+CUDA-graph capture, logits readback) is **already generic** over the bindings vector
+and does not care whether the sequence input is a token id or an embedding — only the
+*sequence-input binding + its per-step write* are token-specific.
+
+### What must change to accept `inputs_embeds` on CUDA
+- `DecodeCudaIo`: carry the sequence source (token vs `inputs_embeds`) + the embed
+  port name, dtype and hidden width from `ModelIoSpec`.
+- `DecodeCudaState::new`: when the source is `InputsEmbeds`, allocate the sequence
+  binding as the graph's `inputs_embeds` dtype/shape `[1, 1, hidden]` **instead of**
+  the `Int64 [1,1]` token binding. KV / mask / position / logits bindings are
+  unchanged.
+- A per-step **`write_decode_embeds(&host_embed_bytes)`** that copies the one-token
+  embedding into the (fixed-address, stable) device binding — exactly parallel to
+  `write_decode_inputs` writing a token id. Because the binding *address* is stable
+  and only the *contents* change per step, this stays **CUDA-graph-capture eligible**
+  (same property as the token-id write).
+- `decode_cuda` / `run_cuda_eager_rows`: resolve the sequence binding by source; when
+  `InputsEmbeds`, pull the embedding tensor from `step_inputs` (validated exactly like
+  the CPU path in `cpu.rs:205-213`) and upload it, rather than building a token tensor.
+- Lift the two refusals **for the `inputs_embeds` case**; keep bailing for
+  *arbitrary* `Routed` ports (defer generic routing to a later slice).
+
+### The device-residency guarantee (the whole point)
+Per step, exactly **one token's embedding** crosses host→device: `[1, 1, hidden]`
+(e.g. `hidden * 4` bytes for f32, `* 2` for f16) — negligible. The **KV cache never
+leaves the device**: it lives in the persistent `kv_binding_range` device bindings
+and grows in place across steps (`extend_mask` + `set_logical_len`), identical to the
+existing token-id CUDA path. This is the same guarantee Inc2b gave on CPU, now on the
+CUDA EP (device 4).
+
+## Honest split
+
+- **Inc3a (this PR — the smallest PROVEN slice):** implement the CUDA `inputs_embeds`
+  sequence-input device binding + per-step upload; lift the refusals for the
+  `inputs_embeds` case; keep KV device-resident. **Prove native-CUDA decode token IDs
+  == CPU-native decode token IDs** on a tiny CUDA-capable `inputs_embeds` fixture
+  (`inputs_embeds` + `attention_mask` + `position_ids` + one KV pair, closed-form
+  deterministic head) driven directly through `decode_with_step_inputs` on both
+  devices. This isolates the new binding path from pipeline wiring.
+- **Inc3b (next):** wire `NativePipelineDecoder` to load the decoder on the CUDA
+  device (via the existing `ONNX_GENAI_PIPELINE_NATIVE_DECODER` flag + a device
+  selector) and prove **native-CUDA-decoder-in-pipeline** token parity vs the ORT
+  pipeline on GPU. Then generic `Routed` ports.
+- **Deferred (Inc4+):** native present-KV exposure for paged cross-request reuse,
+  vision cross-KV, non-`inputs_embeds` arbitrary routed device bindings.
+
+## Risks
+- **KV op support on CUDA persistent bindings.** The fixture's KV must grow correctly
+  under the CUDA persistent-device-KV contract. Mitigation: model the fixture KV on a
+  pattern the CUDA path already handles; if the toy `Concat` KV does not bind cleanly
+  on CUDA, the fixture uses the graph shape the CUDA persistent-KV path expects, or the
+  slice is reported as blocked with the exact failure rather than shipped half-wired.
+- **dtype at the seam.** The every_step embedding is produced on host as the pool
+  dtype (f32/f16); the device binding dtype must match the decoder's declared
+  `inputs_embeds` dtype. Reuse the Inc1 `coerce_value_to_dtype` seam; upload bytes
+  only after coercion.
+- **Graph capture.** The per-step embedding write keeps a stable device address, so
+  capture eligibility is unchanged; verified by the parity run (capture on vs eager).
+- **CUDA + routed generic ports** remain refused — explicitly out of Inc3a.
+
+## Verification plan
+- New parity test (feature-gated `cuda,native-backend`, skips when no GPU):
+  native-CUDA `decode_with_step_inputs` tokens **==** CPU-native tokens on the tiny
+  fixture, ≥2 decode steps (real growing KV).
+- ORT-path goldens unchanged; `cargo test -p onnx-genai-engine` default +
+  `native-backend`; `fmt --check`; clippy default / native-backend / cuda /
+  cuda,native-backend (cfg-correct: all CUDA-embeds code behind `cfg(feature =
+  "cuda")`, all native behind `cfg(feature = "native-backend")`).

crates/onnx-genai-engine/Cargo.toml

@@ -28,6 +28,10 @@ required-features = ["native-backend"]
 name = "native_pipeline_decoder_parity"
 required-features = ["native-backend"]
 
+[[test]]
+name = "native_cuda_pipeline_decoder_parity"
+required-features = ["cuda", "native-backend"]
+
 [features]
 default = []
 native-backend = [

crates/onnx-genai-engine/src/native_decode/backend.rs

@@ -24,7 +24,7 @@ impl DecodeBackend for NativeDecodeSession {
                 return self.decode_cuda_greedy(token_ids[0], past_len).map(Some);
             }
             let token = self
-                .decode_cuda(token_ids, past_len)?
+                .decode_cuda(token_ids, past_len, &[])?
                 .pop()
                 .map(|logits| sample_greedy(&logits))
                 .context("native CUDA decoder produced no logits")?;

crates/onnx-genai-engine/src/native_decode/cuda.rs

@@ -173,11 +173,26 @@ pub(crate) struct DecodeCudaState {
 
 pub(crate) struct DecodeCudaIo<'a> {
     pub(crate) input_ids: &'a str,
+    /// Present only when the decoder's sequence source is `inputs_embeds` (a
+    /// fused VLM decoder) rather than raw token ids. Carries the embed port
+    /// name plus its element dtype and hidden width so the sequence device
+    /// binding is allocated as a float `[1, 1, hidden]` embedding instead of an
+    /// `Int64 [1, 1]` token id. `None` keeps the token-id path byte-identical.
+    pub(crate) inputs_embeds: Option<CudaEmbedsBinding<'a>>,
     pub(crate) attention_mask: &'a str,
     pub(crate) position_ids: Option<&'a str>,
     pub(crate) logits: &'a str,
 }
 
+/// Metadata for a CUDA `inputs_embeds` sequence input (Inc3a). The native CUDA
+/// decoder receives one token's embedding per step as a routed host tensor and
+/// binds it on-device, keeping the KV cache device-resident.
+pub(crate) struct CudaEmbedsBinding<'a> {
+    pub(crate) name: &'a str,
+    pub(crate) dtype: DataType,
+    pub(crate) hidden: usize,
+}
+
 pub(crate) fn trace_capture_declines(trace: &TraceContext, report: &CaptureDeclineReport) {
     for decline in &report.entries {
         if let Some(node_id) = decline.node_id {
@@ -217,11 +232,6 @@ impl NativeDecodeSession {
         position_input: Option<&str>,
         error_context: &str,
     ) -> anyhow::Result<Vec<Vec<f32>>> {
-        let state = self
-            .cuda
-            .as_mut()
-            .context("CUDA decode state is not initialized")?;
-        state.invalidate_graph(&mut self.session)?;
         let ids = token_ids
             .iter()
             .map(|&id| i64::from(id))
@@ -238,6 +248,28 @@ impl NativeDecodeSession {
                 Tensor::from_i64(&[1, token_ids.len()], &positions)?,
             ));
         }
+        self.run_cuda_eager_rows_owned(owned, total_len, error_context)
+    }
+
+    /// Shared eager (uncaptured) device forward body: bind the caller-provided
+    /// non-KV inputs (`owned` — a token or `inputs_embeds` sequence tensor plus
+    /// optional position ids) against the persistent device KV/mask bindings,
+    /// run, collect host `[K, vocab]` logits, and advance the KV logical length.
+    ///
+    /// Both the token path (via [`Self::run_cuda_eager_rows`]) and the
+    /// `inputs_embeds` path (Inc3a) share this body; only the construction of
+    /// `owned` differs, so the KV device-residency guarantee is identical.
+    fn run_cuda_eager_rows_owned(
+        &mut self,
+        owned: Vec<(String, Tensor)>,
+        total_len: usize,
+        error_context: &str,
+    ) -> anyhow::Result<Vec<Vec<f32>>> {
+        let state = self
+            .cuda
+            .as_mut()
+            .context("CUDA decode state is not initialized")?;
+        state.invalidate_graph(&mut self.session)?;
         let bindings = owned
             .iter()
             .map(|(name, tensor)| (name.as_str(), tensor))
@@ -285,10 +317,34 @@ impl NativeDecodeSession {
         &mut self,
         token_ids: &[TokenId],
         past_len: usize,
+        step_inputs: &[(String, Tensor)],
     ) -> anyhow::Result<Vec<Vec<f32>>> {
         let total_len = past_len
             .checked_add(token_ids.len())
             .context("native decode context length overflow")?;
+        // A fused VLM decoder is driven by `inputs_embeds` (Inc3a): there is no
+        // token-id sequence input, so its per-step embedding arrives as a routed
+        // host tensor. Route it through the eager device path — the KV/mask
+        // bindings stay device-resident exactly as in the token path; only the
+        // sequence tensor bound each step differs. Generic (non-embeds) routed
+        // ports on CUDA are still unsupported and are refused below.
+        if let Some(embeds_input) = self
+            .step_input_name(NativeStepInputSource::InputsEmbeds)
+            .map(str::to_owned)
+        {
+            return self.decode_cuda_inputs_embeds(
+                &embeds_input,
+                token_ids,
+                past_len,
+                total_len,
+                step_inputs,
+            );
+        }
+        if let Some((name, _)) = step_inputs.first() {
+            bail!(
+                "native CUDA target decode does not yet accept generic routed host step input '{name}'; only inputs_embeds and generated roles are supported on CUDA (use the CPU native device for arbitrary routed ports)"
+            );
+        }
         let token_input = self
             .step_input_name(NativeStepInputSource::TokenIds)
             .context("native CUDA decoder has no token input binding")?
@@ -356,6 +412,67 @@ impl NativeDecodeSession {
         )
     }
 
+    /// CUDA `inputs_embeds` decode step (Inc3a). The routed one-token embedding
+    /// (`[1, K, hidden]`, produced host-side by the every_step embedding
+    /// component) is bound as the sequence input for an eager device forward,
+    /// while the attention mask and KV cache stay resident in the persistent
+    /// device bindings — so only the small embedding crosses host→device each
+    /// step and the KV never round-trips. Generated position ids (when the
+    /// decoder declares them) are supplied alongside, mirroring the CPU path.
+    fn decode_cuda_inputs_embeds(
+        &mut self,
+        embeds_input: &str,
+        token_ids: &[TokenId],
+        past_len: usize,
+        total_len: usize,
+        step_inputs: &[(String, Tensor)],
+    ) -> anyhow::Result<Vec<Vec<f32>>> {
+        let position_input = self
+            .step_input_name(NativeStepInputSource::PositionIds)
+            .map(str::to_owned);
+        let embeds = step_inputs
+            .iter()
+            .find(|(name, _)| name == embeds_input)
+            .map(|(_, tensor)| tensor.clone())
+            .with_context(|| {
+                format!(
+                    "declared inputs_embeds input '{embeds_input}' was not supplied to the native CUDA decode step; route the current embedding component output to this exact decoder port"
+                )
+            })?;
+        for (name, _) in step_inputs {
+            if name != embeds_input {
+                bail!(
+                    "native CUDA target decode does not yet accept generic routed host step input '{name}'; only inputs_embeds and generated roles are supported on CUDA (use the CPU native device for arbitrary routed ports)"
+                );
+            }
+        }
+        let state = self
+            .cuda
+            .as_mut()
+            .context("CUDA decode state is not initialized")?;
+        if total_len > state.capacity.max_len {
+            bail!("{}", state.capacity_exceeded_error(total_len));
+        }
+        let grew = state.ensure_capacity(&mut self.session, total_len)?;
+        // The embedding is a fresh device upload every step, so this path is not
+        // CUDA-graph captured; expose the growing logical mask width (matching the
+        // eager token forward) rather than freezing to physical capacity.
+        state.extend_mask(if grew { 0 } else { past_len }, total_len, total_len)?;
+
+        let mut owned = Vec::with_capacity(2);
+        owned.push((embeds_input.to_owned(), embeds));
+        if let Some(position_ids_name) = position_input {
+            let positions = (past_len..total_len)
+                .map(|position| i64::try_from(position).context("position id exceeds i64 range"))
+                .collect::<anyhow::Result<Vec<_>>>()?;
+            owned.push((
+                position_ids_name,
+                Tensor::from_i64(&[1, token_ids.len()], &positions)?,
+            ));
+        }
+        self.run_cuda_eager_rows_owned(owned, total_len, "decoder")
+    }
+
     /// Speculative **verify** primitive (option (b): the safe eager M=K path).
     ///
     /// Runs the `draft` candidate tokens (K = `draft.len()`) through the target
@@ -847,10 +964,15 @@ impl DecodeCudaState {
         // data-dependent and must not be collapsed. See RULES.md §2 — this is a
         // purely structural signal, never a model-name gate.
         let mut unit_symbols: HashSet<SymbolId> = HashSet::new();
+        let sequence_input_name = io
+            .inputs_embeds
+            .as_ref()
+            .map(|embeds| embeds.name)
+            .unwrap_or(io.input_ids);
         if let Some(meta) = session
             .inputs()
             .iter()
-            .find(|meta| meta.name == io.input_ids)
+            .find(|meta| meta.name == sequence_input_name)
         {
             Self::collect_unit_symbols(&meta.shape, false, &mut unit_symbols);
         }
@@ -942,13 +1064,29 @@ impl DecodeCudaState {
         };
 
         let input_ids_binding = bindings.len();
-        bindings.push(session.allocate_device_binding(
-            io.input_ids,
-            None::<String>,
-            DataType::Int64,
-            vec![1, 1],
-            vec![1, 1],
-        )?);
+        if let Some(embeds) = &io.inputs_embeds {
+            // Fused VLM decoder (Inc3a): the sequence input is a float
+            // `[1, 1, hidden]` embedding, not an `Int64 [1, 1]` token id. This
+            // persistent binding keeps the sequence-binding index valid; the
+            // eager inputs_embeds decode path binds the per-step embedding as an
+            // owned input rather than through this slot, so it is never the
+            // captured token-write target.
+            bindings.push(session.allocate_device_binding(
+                embeds.name,
+                None::<String>,
+                embeds.dtype,
+                vec![1, 1, embeds.hidden],
+                vec![1, 1, embeds.hidden],
+            )?);
+        } else {
+            bindings.push(session.allocate_device_binding(
+                io.input_ids,
+                None::<String>,
+                DataType::Int64,
+                vec![1, 1],
+                vec![1, 1],
+            )?);
+        }
         let position_ids_binding = if let Some(position_ids) = io.position_ids {
             let index = bindings.len();
             bindings.push(session.allocate_device_binding(