From a82bdd0266fa318d9a03ee4fee431212f24b0328 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 6 Aug 2026 17:55:28 +0000 Subject: [PATCH] =?UTF-8?q?perf(cuda-ep):=20make=20GatedDeltaNet=20Split?= =?UTF-8?q?=20capture-safe=20(C3)=20=E2=80=94=20collapse=2030=20decode=20g?= =?UTF-8?q?raph=20seams?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 35B-A3B QMoE (GatedDeltaNet hybrid) decode step fragments into 184 CUDA-graph segments / 183 eager seams. Steady-state capture logging (ONNX_GENAI_LOG_CAPTURE_SEGMENTS) attributed these to three independent blockers: C1 shared elementwise fixed-decode signature (~120), C2 LinearAttention trailing host synchronize (~30), C3 dynamic Split host-read + synchronize (30). This lands C3. The 30 GatedDeltaNet `Split` nodes take their sizes from a `Constant` node (value_ints=[2048,2048,4096]) that the default OptimizationLevel::None never folds into an initializer, so the kernel fell to the dynamic path: host-read the split-size input + trailing synchronize, both illegal during capture. But the executor pre-allocates each Split output at its statically-inferred shape, so the split sizes are fully determined by the output shapes. SplitKernel now derives the plan from the output shapes (no host read, no sync), caches it so capture_support() reports Supported after warmup, and copies via the same stream-ordered launch_copies — byte identical to the host-read path. Genuinely data-dependent splits (outputs not statically resolved) fall through to the old path and are structurally declined for capture, so correctness is preserved. Byte-exact (35B-A3B QMoE, GPU0, temp=0, KV_MAX_LEN=4096): generated token stream identical to baseline, token@119 == 33803 autoregressive lock holds. Decode 13.415 -> 12.132 ms/tok (-9.6%, 74.5 -> 82.4 tok/s). Segments 184 -> 154 (all 30 Split seams collapse). C2 (LinearAttention sync elision) was implemented and REJECTED: removing the trailing synchronize diverged the decode stream (token@119 became 46283) — the barrier is load-bearing, not conservative. Per "we do not trade correctness for capture," C2 is dropped pending a deeper root-cause pass. C1 (shared elementwise capture gate) remains the dominant follow-up. Test: split_runtime_sizes_with_resolved_outputs_warms_and_captures exercises the GDN case (runtime split input, constant_split_input=false) — cold Split declines capture, warms via the output-derived plan, then captures + replays byte-identically. Refs roofline + GDN-Split investigation (.squad/decisions/inbox/cohaagen-gdn-split-capture.md). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/movement.rs | 72 ++++++++ .../tests/construction_gpu.rs | 174 ++++++++++++++++++ 2 files changed, 246 insertions(+) diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/movement.rs b/crates/onnx-runtime-ep-cuda/src/kernels/movement.rs index 89033b8753..317b2b9385 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/movement.rs +++ b/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)> { + 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::() != 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 diff --git a/crates/onnx-runtime-ep-cuda/tests/construction_gpu.rs b/crates/onnx-runtime-ep-cuda/tests/construction_gpu.rs index d8029afd6b..68344d5bf6 100644 --- a/crates/onnx-runtime-ep-cuda/tests/construction_gpu.rs +++ b/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::>(); + 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::>(); + let output_strides = output_shapes + .iter() + .map(|shape| compute_contiguous_strides(shape)) + .collect::>(); + + 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::>(); + 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::>(); + + // 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::>(); + assert_eq!(replayed, eager, "replayed Split bytes must equal eager"); + + // Byte-exact vs the reference chunking. + let expected: Vec> = { + 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::>() + }) + .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,