perf(cuda-ep): make GatedDeltaNet Split capture-safe (C3) — collapse 30 decode graph seams

crates/onnx-runtime-ep-cuda/src/kernels/movement.rs

@@ -1442,6 +1442,61 @@ impl SplitKernel {
         }
         Ok(())
     }
+
+    /// Derive the split plan from the already-resolved output shapes.
+    ///
+    /// The executor pre-allocates each Split output at its statically-inferred
+    /// shape, so the split sizes are fully determined by `outputs[i].shape[axis]`
+    /// — no host read of the split-size input and no trailing synchronization are
+    /// required. Returns `None` (falling back to the host-read path) whenever the
+    /// outputs are not concretely and consistently sized, e.g. a genuinely
+    /// data-dependent split whose output shapes shape-inference could not resolve;
+    /// such regions are structurally declined for capture anyway.
+    ///
+    /// This is byte-identical to the host-read path: `launch_copies` performs the
+    /// same per-output stream-ordered copies with the same offsets.
+    fn plan_from_outputs(
+        &self,
+        inputs: &[TensorView],
+        outputs: &[TensorMut],
+    ) -> Option<(usize, Vec<usize>)> {
+        if outputs.is_empty() {
+            return None;
+        }
+        let in_shape = &inputs[0].shape;
+        let rank = in_shape.len();
+        if rank == 0 {
+            return None;
+        }
+        let axis = if self.axis < 0 {
+            self.axis + rank as i64
+        } else {
+            self.axis
+        };
+        if !(0..rank as i64).contains(&axis) {
+            return None;
+        }
+        let axis = axis as usize;
+        let mut sizes = Vec::with_capacity(outputs.len());
+        for output in outputs {
+            if output.shape.len() != rank {
+                return None;
+            }
+            // Every non-split dimension must match the input; only `axis` differs.
+            for (dim, (&out_extent, &in_extent)) in
+                output.shape.iter().zip(in_shape.iter()).enumerate()
+            {
+                if dim != axis && out_extent != in_extent {
+                    return None;
+                }
+            }
+            sizes.push(output.shape[axis]);
+        }
+        if sizes.iter().sum::<usize>() != in_shape[axis] {
+            return None;
+        }
+        Some((axis, sizes))
+    }
 }
 impl Kernel for SplitKernel {
     fn set_constant_inputs(&mut self, constant_inputs: &[bool]) {
@@ -1455,6 +1510,23 @@ impl Kernel for SplitKernel {
         }
         require_dense("Split", inputs, outputs)?;
         let rank = inputs[0].shape.len();
+        // Capture-safe fast path: when the outputs are concretely sized (the
+        // common case, including this model's GatedDeltaNet Split whose sizes come
+        // from a `Constant` node rather than a graph initializer), the split sizes
+        // are fully determined by the output shapes. Copying from them avoids the
+        // host read of the split-size input and the trailing synchronization that
+        // otherwise de-captures the surrounding decode block. Caching the derived
+        // plan lets `capture_support` report `Supported` once the shape is warmed.
+        if let Some((axis, sizes)) = self.plan_from_outputs(inputs, outputs) {
+            *self.static_plan.lock().map_err(|_| {
+                EpError::KernelFailed("cuda_ep Split: static plan lock was poisoned".into())
+            })? = Some(StaticSplitPlan {
+                axis,
+                axis_extent: inputs[0].shape[axis],
+                sizes: sizes.clone(),
+            });
+            return self.launch_copies(inputs, outputs, axis, &sizes);
+        }
         // Capturable fast path: the static, single-data-input form was planned
         // at build time, so no host read of split sizes and no trailing
         // synchronization are needed. Copies are ordered on the stream, which is

crates/onnx-runtime-ep-cuda/tests/construction_gpu.rs

@@ -798,6 +798,180 @@ fn split_constant_input_warms_and_captures() {
     }
 }
 
+#[cfg_attr(
+    not(feature = "gpu-tests"),
+    ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner"
+)]
+#[test]
+fn split_runtime_sizes_with_resolved_outputs_warms_and_captures() {
+    // The GatedDeltaNet decode Split (C3 of the GDN capture fix): a *runtime*
+    // split-sizes input that is NOT flagged constant (its sizes come from a
+    // `Constant` node the default `OptimizationLevel::None` never folds into an
+    // initializer). The old path host-read the sizes and synchronized every
+    // step, de-capturing the surrounding decode block. Because the executor
+    // pre-allocates each output at its statically-inferred shape, the sizes are
+    // fully determined by the output shapes — so after warmup the kernel derives
+    // a static plan and becomes capture-safe with byte-identical results.
+    let ep = require_cuda();
+    let runtime = ep.runtime();
+    let device = ep.device_id();
+    // Qwen3.5 GatedDeltaNet split: [1,1,H,320] -> [1,1,H,128] + [1,1,H,128] +
+    // [1,1,H,64] (q/k/v-style fan-out), axis=-1.
+    let heads = 4usize;
+    let data_shape = [1, 1, heads, 320];
+    let split_shape = [3];
+    let output_shapes = [
+        vec![1, 1, heads, 128],
+        vec![1, 1, heads, 128],
+        vec![1, 1, heads, 64],
+    ];
+    let data = (0..heads * 320)
+        .map(|value| value as f32)
+        .collect::<Vec<_>>();
+    let data_bytes = raw(&data);
+    let split_bytes = raw(&[128_i64, 128, 64]);
+    // NOTE: no `set_constant_inputs` — constant_split_input stays false, which is
+    // exactly the GDN case the C3 fix targets.
+    let kernel = build_split_kernel(
+        &ep,
+        &data_shape,
+        &output_shapes,
+        &[("axis", Attribute::Int(-1))],
+        Some(&split_shape),
+    );
+    // Before warmup the kernel has no static plan, so it must decline capture
+    // (the dynamic host-read/synchronize path).
+    assert!(
+        matches!(kernel.capture_support(), CaptureSupport::Unsupported { .. }),
+        "a cold runtime-split Split must decline capture until warmed"
+    );
+
+    let data_buffer = ep.allocate(data_bytes.len(), 256).unwrap();
+    let split_buffer = ep.allocate(split_bytes.len(), 256).unwrap();
+    unsafe {
+        runtime
+            .htod(&data_bytes, cuptr(data_buffer.as_ptr()))
+            .unwrap();
+        runtime
+            .htod(&split_bytes, cuptr(split_buffer.as_ptr()))
+            .unwrap();
+    }
+    let data_strides = compute_contiguous_strides(&data_shape);
+    let split_strides = compute_contiguous_strides(&split_shape);
+    let inputs = [
+        TensorView::new(
+            DevicePtr(data_buffer.as_ptr()),
+            DataType::Float32,
+            &data_shape,
+            &data_strides,
+            device,
+        ),
+        TensorView::new(
+            DevicePtr(split_buffer.as_ptr()),
+            DataType::Int64,
+            &split_shape,
+            &split_strides,
+            device,
+        ),
+    ];
+    let mut output_buffers = output_shapes
+        .iter()
+        .map(|shape| {
+            ep.allocate(DataType::Float32.storage_bytes(shape.iter().product()), 256)
+                .unwrap()
+        })
+        .collect::<Vec<_>>();
+    let output_strides = output_shapes
+        .iter()
+        .map(|shape| compute_contiguous_strides(shape))
+        .collect::<Vec<_>>();
+
+    macro_rules! execute {
+        () => {{
+            let mut outputs = output_buffers
+                .iter_mut()
+                .zip(&output_shapes)
+                .zip(&output_strides)
+                .map(|((buffer, shape), strides)| {
+                    TensorMut::new(
+                        DevicePtrMut(buffer.as_mut_ptr()),
+                        DataType::Float32,
+                        shape,
+                        strides,
+                        device,
+                    )
+                })
+                .collect::<Vec<_>>();
+            kernel.execute(&inputs, &mut outputs).unwrap();
+        }};
+    }
+
+    // Warm eagerly: the output-derived plan is cached, flipping capture_support
+    // to Supported WITHOUT ever host-reading the split-size input.
+    execute!();
+    assert_eq!(
+        kernel.capture_support(),
+        CaptureSupport::Supported,
+        "after warmup the output-derived static plan must admit capture"
+    );
+
+    let eager = output_buffers
+        .iter()
+        .zip(&output_shapes)
+        .map(|(buffer, shape)| {
+            let mut bytes = vec![0; DataType::Float32.storage_bytes(shape.iter().product())];
+            unsafe {
+                runtime.dtoh(&mut bytes, cuptr(buffer.as_ptr())).unwrap();
+            };
+            bytes
+        })
+        .collect::<Vec<_>>();
+
+    // Capture + replay must succeed (no host-read/sync in the captured region)
+    // and reproduce the eager bytes exactly.
+    runtime.begin_graph_capture(&[kernel.as_ref()]).unwrap();
+    execute!();
+    runtime.end_graph_capture().unwrap();
+    assert!(runtime.has_graph_executable().unwrap());
+    runtime.replay_graph().unwrap();
+    let replayed = output_buffers
+        .iter()
+        .zip(&output_shapes)
+        .map(|(buffer, shape)| {
+            let mut bytes = vec![0; DataType::Float32.storage_bytes(shape.iter().product())];
+            unsafe {
+                runtime.dtoh(&mut bytes, cuptr(buffer.as_ptr())).unwrap();
+            };
+            bytes
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(replayed, eager, "replayed Split bytes must equal eager");
+
+    // Byte-exact vs the reference chunking.
+    let expected: Vec<Vec<f32>> = {
+        let bounds = [(0usize, 128usize), (128, 256), (256, 320)];
+        bounds
+            .iter()
+            .map(|&(lo, hi)| {
+                data.chunks_exact(320)
+                    .flat_map(|head| &head[lo..hi])
+                    .copied()
+                    .collect::<Vec<_>>()
+            })
+            .collect()
+    };
+    for (idx, want) in expected.iter().enumerate() {
+        assert_eq!(eager[idx], raw(want), "output {idx} byte mismatch");
+    }
+    assert!(runtime.reset_graph().unwrap());
+
+    ep.deallocate(data_buffer).unwrap();
+    ep.deallocate(split_buffer).unwrap();
+    for buffer in output_buffers {
+        ep.deallocate(buffer).unwrap();
+    }
+}
+
 fn build_movement_kernel(
     ep: &CudaExecutionProvider,
     op: &str,