diff --git a/.squad/decisions/inbox/copilot-qmoe-workspace-stage0.md b/.squad/decisions/inbox/copilot-qmoe-workspace-stage0.md new file mode 100644 index 0000000000..ed433d5084 --- /dev/null +++ b/.squad/decisions/inbox/copilot-qmoe-workspace-stage0.md @@ -0,0 +1,4 @@ +### 2026-08-09: QMoE workspace is prepared before admission +**By:** Copilot +**What:** Kernels declare typed workspace requirements from concrete tensor metadata. Native CUDA prefill resolves QMoE shapes and reserves one reusable session-persistent workspace peak before the admission callback; execution receives only that prepared buffer. +**Why:** QMoE prompt-sized scratch previously allocated inside execute, after HTTP/SSE admission. Sharing the checked layout helper between planning and execution prevents formula drift and turns omissions or shape changes into explicit invariant failures. diff --git a/crates/onnx-genai-engine/src/engine/runtime.rs b/crates/onnx-genai-engine/src/engine/runtime.rs index 62f687f8d8..914fc961f7 100644 --- a/crates/onnx-genai-engine/src/engine/runtime.rs +++ b/crates/onnx-genai-engine/src/engine/runtime.rs @@ -16,6 +16,26 @@ fn generation_budget_cap(cap: ScheduledBudgetCap) -> GenerationBudgetCap { } } +#[cfg(feature = "native-backend")] +fn native_workspace_query_rows( + prompt_rows: usize, + plan: Option<&NativeSpeculationPlan>, + effective_max_new_tokens: usize, + max_context: Option, +) -> usize { + let remaining_context = max_context + .map(|limit| limit.saturating_sub(prompt_rows)) + .unwrap_or(effective_max_new_tokens); + let verify_rows = plan.map_or(0, |plan| { + crate::native_speculative::verification_width( + plan.width, + effective_max_new_tokens, + remaining_context, + ) + }); + prompt_rows.max(verify_rows) +} + fn metadata_eos_token_ids(metadata: &InferenceMetadata) -> Vec { metadata .tokens @@ -135,6 +155,7 @@ impl Engine { } options.max_context = self.max_context_for_request(&options); let chain = build_processor_chain(&options, Some(&self.tokenizer))?; + let speculation_plan = native_speculation_plan(&options, &chain); let scheduler_session_id = self.next_native_session_id(); let scheduled = self.admit_generate_request_with_scheduler( scheduler_session_id, @@ -144,6 +165,21 @@ impl Engine { )?; let budget_cap = scheduled.budget_cap.map(generation_budget_cap); options.max_new_tokens = scheduled.max_tokens; + let workspace_query_rows = native_workspace_query_rows( + prompt_tokens.len(), + speculation_plan.as_ref(), + options.max_new_tokens, + options.max_context, + ); + if let Err(error) = self + .native_session + .as_mut() + .context("native decoder session is unavailable")? + .prepare_generation_workspace_for_query_rows(&prompt_tokens, workspace_query_rows) + { + self.scheduler.complete(scheduler_session_id); + return Err(error); + } if let Some(callback) = admission_callback.as_mut() { callback(); } @@ -151,7 +187,7 @@ impl Engine { // Speculation ON (implemented greedy prompt-lookup) → the native // speculative driver. Every other request stays on the untouched plain // M=1 fast path below, preserving the 762 tok/s non-regression guarantee. - let result = if let Some(plan) = native_speculation_plan(&options, &chain) { + let result = if let Some(plan) = speculation_plan { let mut stats = SpeculativeStats::default(); let result = (|| { let native_session = self @@ -2035,6 +2071,15 @@ impl Engine { )?; let budget_cap = scheduled.budget_cap.map(generation_budget_cap); options.max_new_tokens = scheduled.max_tokens; + if let Err(error) = self + .native_session + .as_mut() + .context("native decoder session is unavailable")? + .prepare_generation_workspace_preserving_state(&prompt_tokens) + { + self.scheduler.complete(session_id); + return Err(error); + } if let Some(callback) = admission_callback.as_mut() { callback(); } @@ -2226,6 +2271,10 @@ impl Engine { #[cfg(test)] mod tests { use super::*; + #[cfg(feature = "native-backend")] + use crate::ProcessorChain; + #[cfg(feature = "native-backend")] + use std::path::PathBuf; #[test] fn ort_and_native_generation_use_scheduler_admission() { @@ -2234,6 +2283,56 @@ mod tests { assert!(generate_uses_scheduler(EngineDecodeBackend::Auto)); } + #[cfg(feature = "native-backend")] + #[test] + fn speculative_workspace_rows_follow_the_exact_runtime_plan() { + let chain = ProcessorChain::new(); + let mut options = GenerateOptions { + greedy: true, + temperature: 0.0, + speculative_mode: Some(SpeculativeMode::PromptLookup { + ngram: 1, + max_tokens: 4, + }), + ..GenerateOptions::default() + }; + let plan = native_speculation_plan(&options, &chain).unwrap(); + assert_eq!(native_workspace_query_rows(1, Some(&plan), 8, None), 4); + assert_eq!(native_workspace_query_rows(8, Some(&plan), 8, None), 8); + + options.num_speculative_tokens = Some(1024); + let widened = native_speculation_plan(&options, &chain).unwrap(); + assert_eq!(widened.width, 1024); + assert_eq!(native_workspace_query_rows(1, Some(&widened), 1, None), 1); + assert_eq!( + native_workspace_query_rows(1, Some(&widened), 8, Some(3)), + 2 + ); + let runtime_width = crate::native_speculative::verification_width(widened.width, 5, 3); + assert_eq!( + native_workspace_query_rows(1, Some(&widened), 5, Some(4)), + 1usize.max(runtime_width) + ); + + options.num_speculative_tokens = None; + options.speculative_mode = Some(SpeculativeMode::SharedKv( + crate::config::SharedKvProposerConfig { + assistant_model: PathBuf::from("assistant.onnx"), + target_hidden_output: "hidden".to_string(), + input_embedding_weights: PathBuf::from("embedding.bin"), + backbone_hidden_size: 1, + vocab_size: 1, + num_speculative_tokens: 3, + shared_kv: Vec::new(), + }, + )); + let shared_kv = native_speculation_plan(&options, &chain).unwrap(); + assert_eq!(shared_kv.width, 4); + assert_eq!(native_workspace_query_rows(1, Some(&shared_kv), 8, None), 4); + + assert_eq!(native_workspace_query_rows(3, None, 8, None), 3); + } + #[cfg(feature = "native-backend")] #[test] fn native_generate_rejects_over_kv_byte_budget_before_backend_run() -> anyhow::Result<()> { diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index 21661988f5..13f88b7803 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -479,6 +479,99 @@ pub(crate) fn trace_capture_declines(trace: &TraceContext, report: &CaptureDecli } impl NativeDecodeSession { + pub(crate) fn prepare_cuda_prefill_workspace( + &mut self, + token_ids: &[TokenId], + ) -> anyhow::Result { + self.prepare_cuda_prefill_workspace_with_step_inputs(token_ids, 0, &[]) + } + + pub(crate) fn prepare_cuda_prefill_workspace_with_step_inputs( + &mut self, + token_ids: &[TokenId], + past_len: usize, + step_inputs: &[(String, Tensor)], + ) -> anyhow::Result { + if self.has_eager_step_inputs() { + let workspace_nodes = self.session.workspace_node_locations(); + if workspace_nodes.is_empty() { + return Ok(onnx_runtime_session::WorkspaceRequirement::NONE); + } + } + let position_input = self + .step_input_name(NativeStepInputSource::PositionIds) + .map(str::to_owned); + let input_len = token_ids.len(); + let total_len = past_len + .checked_add(input_len) + .context("native CUDA workspace preparation length overflow")?; + let supplied = step_inputs + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + let mut owned = + if let Some(name) = self.step_input_name(NativeStepInputSource::InputsEmbeds) { + if let Some(tensor) = supplied.get(name) { + vec![(name.to_owned(), (*tensor).try_clone()?)] + } else { + let meta = self + .session + .inputs() + .iter() + .find(|meta| meta.name == name) + .with_context(|| { + format!("missing native CUDA inputs_embeds metadata for '{name}'") + })?; + let hidden = match meta.shape.last() { + Some(Dim::Static(hidden)) => *hidden, + _ => bail!( + "prepare-only QMoE workspace planning cannot conservatively bind \ + inputs_embeds '{name}': expected a static hidden width, got {:?}", + meta.shape + ), + }; + vec![( + name.to_owned(), + Tensor::zeros(meta.dtype, vec![1, input_len, hidden])?, + )] + } + } else { + let token_input = self + .step_input_name(NativeStepInputSource::TokenIds) + .context("native CUDA decoder has no token or inputs_embeds input binding")? + .to_owned(); + let ids = token_ids + .iter() + .map(|&id| i64::from(id)) + .collect::>(); + vec![(token_input, Tensor::from_i64(&[1, input_len], &ids)?)] + }; + for (name, tensor) in step_inputs { + if !owned.iter().any(|(bound, _)| bound == name) { + owned.push((name.clone(), tensor.try_clone()?)); + } + } + if let Some(position_input) = position_input { + owned.push(( + position_input, + self.build_step_positions(past_len, total_len)?, + )); + } + let state = self + .cuda + .as_mut() + .context("CUDA decode state is not initialized")?; + state.ensure_capacity(&mut self.session, total_len)?; + state.extend_mask(past_len, total_len, total_len)?; + let inputs = owned + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + self.session + .prepare_with_device_bindings(&inputs, &mut state.bindings[..state.base_binding_count]) + .context("prepare native CUDA prefill workspace") + } + /// Run one eager (uncaptured) `[1, K]` device forward pass and return host /// `[K, vocab]` logits. /// diff --git a/crates/onnx-genai-engine/src/native_decode/mod.rs b/crates/onnx-genai-engine/src/native_decode/mod.rs index 27409de891..8e958f6f09 100644 --- a/crates/onnx-genai-engine/src/native_decode/mod.rs +++ b/crates/onnx-genai-engine/src/native_decode/mod.rs @@ -854,6 +854,63 @@ impl NativeDecodeSession { ) } + pub(crate) fn prepare_generation_workspace_for_query_rows( + &mut self, + prompt_tokens: &[TokenId], + query_rows: usize, + ) -> anyhow::Result { + self.prepare_generation_workspace_inner(prompt_tokens, query_rows, true) + } + + pub(crate) fn prepare_generation_workspace_preserving_state( + &mut self, + prompt_tokens: &[TokenId], + ) -> anyhow::Result { + self.prepare_generation_workspace_inner(prompt_tokens, prompt_tokens.len(), false) + } + + pub(crate) fn prepare_generation_workspace_with_step_inputs( + &mut self, + tokens: &[TokenId], + past_len: usize, + step_inputs: &[(String, Tensor)], + ) -> anyhow::Result { + if tokens.is_empty() { + bail!("native workspace preparation requires at least one input token"); + } + if self.cuda.is_some() { + self.prepare_cuda_prefill_workspace_with_step_inputs(tokens, past_len, step_inputs) + } else { + Ok(onnx_runtime_session::WorkspaceRequirement::NONE) + } + } + + fn prepare_generation_workspace_inner( + &mut self, + prompt_tokens: &[TokenId], + query_rows: usize, + reset: bool, + ) -> anyhow::Result { + if prompt_tokens.is_empty() { + bail!("native workspace preparation requires at least one prompt token"); + } + if reset { + self.reset()?; + } + if self.cuda.is_some() { + if query_rows <= prompt_tokens.len() { + self.prepare_cuda_prefill_workspace(prompt_tokens) + } else { + let mut planning_tokens = Vec::with_capacity(query_rows); + planning_tokens.extend_from_slice(prompt_tokens); + planning_tokens.resize(query_rows, prompt_tokens[0]); + self.prepare_cuda_prefill_workspace(&planning_tokens) + } + } else { + Ok(onnx_runtime_session::WorkspaceRequirement::NONE) + } + } + /// Generate incrementally: reuse KV state up to `resume_from` and only /// prefill `prompt_tokens[resume_from..]`. The caller must ensure that /// `prompt_tokens[..resume_from]` matches the tokens already in the KV cache. diff --git a/crates/onnx-genai-engine/src/native_speculative.rs b/crates/onnx-genai-engine/src/native_speculative.rs index d2b5560a8f..3e55832d78 100644 --- a/crates/onnx-genai-engine/src/native_speculative.rs +++ b/crates/onnx-genai-engine/src/native_speculative.rs @@ -42,6 +42,14 @@ use crate::speculative::{ use anyhow::Context; use onnx_genai_ort::Tokenizer; +pub(crate) fn verification_width( + draft_width: usize, + remaining_tokens: usize, + remaining_context: usize, +) -> usize { + draft_width.min(remaining_tokens).min(remaining_context) +} + /// Outer speculative token loop bound to a single [`NativeDecodeSession`]. /// /// Peer to the plain [`NativeDecodeSession::generate_with_callback`] loop; it @@ -173,11 +181,8 @@ impl<'a> NativeSpeculativeDriver<'a> { .max_context .map(|limit| limit.saturating_sub(context_len)) .unwrap_or(remaining_tokens); - let width = self - .draft_width - .min(remaining_tokens) - .min(remaining_context) - .max(1); + let width = verification_width(self.draft_width, remaining_tokens, remaining_context); + debug_assert!(width > 0); let context_tokens: Vec = prompt_tokens .iter() diff --git a/crates/onnx-genai-engine/src/pipeline/decoder_component.rs b/crates/onnx-genai-engine/src/pipeline/decoder_component.rs index cf790f7b52..f02e98fec9 100644 --- a/crates/onnx-genai-engine/src/pipeline/decoder_component.rs +++ b/crates/onnx-genai-engine/src/pipeline/decoder_component.rs @@ -33,6 +33,17 @@ use onnx_genai_kv::{PagedKvCache, SequenceId}; /// and (when paging) [`mirror_last_present_kv`](Self::mirror_last_present_kv) /// without ever handling a concrete tensor type. pub(crate) trait PipelineDecoderComponent { + /// Prepare any governed decoder workspace using the exact routed values for + /// this step. Backends without a workspace contract remain a no-op. + fn prepare_step( + &mut self, + _input_tokens: &[TokenId], + _past_len: usize, + _extras: &[(String, Value)], + ) -> anyhow::Result<()> { + Ok(()) + } + /// Run one decoder step over `input_tokens` at absolute `past_len`, binding /// the routed `extras` (every_step outputs, `inputs_embeds`, routed /// positions, static cross-attention KV), advancing the internal KV. The @@ -211,6 +222,19 @@ pub(crate) struct NativePipelineDecoder { #[cfg(feature = "native-backend")] impl NativePipelineDecoder { + fn native_step_inputs( + extras: &[(String, Value)], + ) -> anyhow::Result> { + let mut step_inputs = Vec::with_capacity(extras.len()); + for (port, value) in extras { + let component = value_to_component_tensor(value)?; + let tensor = + crate::native_component::component_tensor_to_native_tensor(port, &component)?; + step_inputs.push((port.clone(), tensor)); + } + Ok(step_inputs) + } + /// Load the decoder ONNX model as a native decode session on the given /// device, keeping its KV cache resident across the generation. The /// pipeline-declared `io` spec is threaded so an `inputs_embeds` decoder @@ -250,6 +274,21 @@ impl NativePipelineDecoder { #[cfg(feature = "native-backend")] impl PipelineDecoderComponent for NativePipelineDecoder { + fn prepare_step( + &mut self, + input_tokens: &[TokenId], + past_len: usize, + extras: &[(String, Value)], + ) -> anyhow::Result<()> { + let step_inputs = Self::native_step_inputs(extras)?; + self.session.prepare_generation_workspace_with_step_inputs( + input_tokens, + past_len, + &step_inputs, + )?; + Ok(()) + } + fn step( &mut self, input_tokens: &[TokenId], @@ -261,13 +300,7 @@ impl PipelineDecoderComponent for NativePipelineDecoder { // value -> ComponentTensor -> native Tensor reuses the Inc1 value-type // seam; this is one token's embedding per decode step (small upload), // while the KV cache stays resident inside the native session. - let mut step_inputs = Vec::with_capacity(extras.len()); - for (port, value) in extras { - let component = value_to_component_tensor(value)?; - let tensor = - crate::native_component::component_tensor_to_native_tensor(port, &component)?; - step_inputs.push((port.clone(), tensor)); - } + let step_inputs = Self::native_step_inputs(extras)?; let rows = self .session .decode_with_step_inputs(input_tokens, past_len, &step_inputs)?; diff --git a/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs b/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs index 396bc7a6a8..a239f49053 100644 --- a/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs +++ b/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs @@ -36,6 +36,7 @@ impl PipelineEngine { pub(crate) fn run_autoregressive( &mut self, pipeline_request: PipelineGenerateRequest, + admission: Option<&mut dyn FnMut()>, callback: Option<&mut GenerateTokenCallback<'_>>, ) -> anyhow::Result<(GenerateResult, PipelineTensors)> { // Guard first: a non-autoregressive pipeline (single-pass / iterative @@ -279,6 +280,7 @@ impl PipelineEngine { }; let mut backend = PipelineDecodeLoopBackend { decoder: decoder_component, + admission, paged: paged_mirror, pool: &mut tensors, step_components, diff --git a/crates/onnx-genai-engine/src/pipeline/mod.rs b/crates/onnx-genai-engine/src/pipeline/mod.rs index 8ca5df526a..06766c6354 100644 --- a/crates/onnx-genai-engine/src/pipeline/mod.rs +++ b/crates/onnx-genai-engine/src/pipeline/mod.rs @@ -1055,16 +1055,31 @@ impl PipelineEngine { &mut self, pipeline_request: PipelineGenerateRequest, callback: Option<&mut GenerateTokenCallback<'_>>, + ) -> anyhow::Result { + self.generate_with_callbacks(pipeline_request, None, callback) + } + + /// Generate text with a pre-token admission callback. For a native routed + /// decoder, admission fires only after exact step inputs have prepared and + /// reserved governed workspace. + pub fn generate_with_callbacks( + &mut self, + pipeline_request: PipelineGenerateRequest, + mut admission: Option<&mut dyn FnMut()>, + callback: Option<&mut GenerateTokenCallback<'_>>, ) -> anyhow::Result { // A nested-AR (multi-decoder TTS) pipeline drives its own outer/inner // loops; `generate` returns the flattened per-frame code tokens (use // `synthesize` to also run the post-decode vocoder into a waveform). if matches!(self.plan, PipelinePlan::NestedAutoregressive(_)) { + if let Some(admitted) = admission.as_mut() { + admitted(); + } return self .run_nested_autoregressive(pipeline_request) .map(|(result, _pool)| result); } - self.run_autoregressive(pipeline_request, callback) + self.run_autoregressive(pipeline_request, admission, callback) .map(|(result, _pool)| result) } @@ -1092,7 +1107,7 @@ impl PipelineEngine { if matches!(self.plan, PipelinePlan::NestedAutoregressive(_)) { return self.synthesize_nested(pipeline_request); } - let (generation, mut tensors) = self.run_autoregressive(pipeline_request, None)?; + let (generation, mut tensors) = self.run_autoregressive(pipeline_request, None, None)?; let ar = self.plan.autoregressive_plan()?.clone(); // Publish the AR decoder's generated code sequence into the shared pool diff --git a/crates/onnx-genai-engine/src/pipeline/paged_decode.rs b/crates/onnx-genai-engine/src/pipeline/paged_decode.rs index 7d6009818b..108927f594 100644 --- a/crates/onnx-genai-engine/src/pipeline/paged_decode.rs +++ b/crates/onnx-genai-engine/src/pipeline/paged_decode.rs @@ -78,13 +78,16 @@ pub(crate) struct PagedMirror<'a> { pub(crate) windowed: bool, } -pub(crate) struct PipelineDecodeLoopBackend<'a> { +pub(crate) struct PipelineDecodeLoopBackend<'a, 'admission> { /// The decoder driven each step, owning its KV state behind a stateful, /// backend-neutral seam. Boxed behind [`PipelineDecoderComponent`] so the /// same loop drives the ORT decoder (via [`OrtPipelineDecoder`]) or, in a /// follow-up, a native decoder keeping KV device-resident — one code path, /// no forked native copy (mirrors the Inc1 every_step seam, but stateful). pub(crate) decoder: Box, + /// One-shot admission notification fired only after the first decoder step's + /// exact routed inputs have successfully prepared governed workspace. + pub(crate) admission: Option<&'admission mut dyn FnMut()>, /// 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, @@ -122,7 +125,7 @@ pub(crate) struct PipelineDecodeLoopBackend<'a> { pub(crate) kv_len: usize, } -impl PipelineDecodeLoopBackend<'_> { +impl PipelineDecodeLoopBackend<'_, '_> { /// Run every declared `every_step` component over `seed` (the full prompt on /// prefill, the single running token on decode), publishing all of their /// outputs into the shared pool. Topological order ensures a component sees @@ -200,7 +203,7 @@ impl PipelineDecodeLoopBackend<'_> { } } -impl DecodeLoopBackend for PipelineDecodeLoopBackend<'_> { +impl DecodeLoopBackend for PipelineDecodeLoopBackend<'_, '_> { fn context_len(&self) -> usize { self.context_tokens.len() } @@ -235,6 +238,11 @@ impl DecodeLoopBackend for PipelineDecodeLoopBackend<'_> { // into the decoder for this step. self.run_step_components(&input_tokens)?; let extras = self.decoder_extras()?; + self.decoder + .prepare_step(&input_tokens, past_len, &extras)?; + if let Some(admitted) = self.admission.take() { + admitted(); + } // Advance the decoder one step; it retains this step's outputs internally // so the loop never handles a concrete tensor type (see // `PipelineDecoderComponent`). diff --git a/crates/onnx-genai-engine/tests/glm_tiny_qmoe_native_cuda_e2e.rs b/crates/onnx-genai-engine/tests/glm_tiny_qmoe_native_cuda_e2e.rs index 33d27b4bed..098fa98e07 100644 --- a/crates/onnx-genai-engine/tests/glm_tiny_qmoe_native_cuda_e2e.rs +++ b/crates/onnx-genai-engine/tests/glm_tiny_qmoe_native_cuda_e2e.rs @@ -24,6 +24,7 @@ use std::path::{Path, PathBuf}; use onnx_genai_engine::{ Engine, EngineConfig, EngineDecodeBackend, GeneratePrompt, GenerateRequest, NativeDecodeDevice, + SpeculativeMode, }; const ANCHOR_IDS: &[u32] = &[62, 164, 59, 205, 48, 166, 27, 9, 221, 190, 123, 108]; @@ -152,3 +153,71 @@ fn glm_tiny_qmoe_native_cuda_eager_matches_cpu_and_declines_capture() -> anyhow: assert_eq!(stats.graph.replays, 0); Ok(()) } + +#[test] +fn glm_tiny_qmoe_native_cuda_session_admits_after_workspace_preparation() -> anyhow::Result<()> { + let Some(dir) = fixture_dir() else { + return Ok(()); + }; + if let Err(error) = onnx_runtime_ep_cuda::CudaExecutionProvider::new(0) { + eprintln!("skipping GLM-5.2 native CUDA admission regression: {error}"); + return Ok(()); + } + let mut engine = native_engine(&dir, NativeDecodeDevice::Cuda { index: Some(0) })?; + for turn in 0..2 { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let admission_events = events.clone(); + let mut admitted = move || admission_events.lock().unwrap().push("admission"); + let token_events = events.clone(); + let mut token = move |_| { + token_events.lock().unwrap().push("token"); + Ok(()) + }; + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![123])); + request.options.max_new_tokens = 2; + request.options.temperature = 0.0; + request.options.greedy = true; + request.options.stop_on_eos = false; + let result = + engine.generate_with_callbacks(request, Some(&mut admitted), Some(&mut token))?; + assert_eq!(result.token_ids, ANCHOR_IDS[..2]); + let events = events.lock().unwrap(); + assert_eq!(events.first(), Some(&"admission"), "turn {turn}"); + assert!(events[1..].contains(&"token"), "turn {turn}: {events:?}"); + } + Ok(()) +} + +#[test] +fn glm_tiny_qmoe_native_cuda_speculative_verify_uses_width_workspace() -> anyhow::Result<()> { + let Some(dir) = fixture_dir() else { + return Ok(()); + }; + if let Err(error) = onnx_runtime_ep_cuda::CudaExecutionProvider::new(0) { + eprintln!("skipping GLM-5.2 native CUDA speculative regression: {error}"); + return Ok(()); + } + let device = NativeDecodeDevice::Cuda { index: Some(0) }; + let mut baseline = native_engine(&dir, device.clone())?; + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![123])); + request.options.max_new_tokens = ANCHOR_IDS.len() * 2; + request.options.temperature = 0.0; + request.options.greedy = true; + request.options.stop_on_eos = false; + let expected = baseline.generate(request.clone())?; + + let mut engine = native_engine(&dir, device)?; + request.options.speculative_mode = Some(SpeculativeMode::PromptLookup { + ngram: 1, + max_tokens: 4, + }); + + let result = engine.generate(request)?; + let stats = engine.last_speculative_stats(); + assert_eq!(result.token_ids, expected.token_ids); + assert!( + stats.verification_steps > 0 && stats.proposed_tokens > 1, + "prompt-length-one regression must execute a multi-row verify: {stats:?}" + ); + Ok(()) +} diff --git a/crates/onnx-genai-engine/tests/native_cuda_routed_pipeline_decoder_parity.rs b/crates/onnx-genai-engine/tests/native_cuda_routed_pipeline_decoder_parity.rs index 6f2f6b51d7..e730ba0db0 100644 --- a/crates/onnx-genai-engine/tests/native_cuda_routed_pipeline_decoder_parity.rs +++ b/crates/onnx-genai-engine/tests/native_cuda_routed_pipeline_decoder_parity.rs @@ -19,14 +19,16 @@ //! by a zero matrix, so it flows through a CUDA op but contributes nothing). use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; -use onnx_genai_engine::pipeline::PipelineGenerateRequest; -use onnx_genai_engine::{Engine, EngineConfig, GenerateOptions, GeneratePrompt, GenerateRequest}; +use onnx_genai_engine::pipeline::{PipelineEngine, PipelineGenerateRequest}; +use onnx_genai_engine::{ + Engine, EngineConfig, EngineDecodeBackend, GenerateOptions, GeneratePrompt, GenerateRequest, + NativeDecodeDevice, +}; use onnx_genai_ort::Value; -const NATIVE_DECODER_ENV: &str = "ONNX_GENAI_PIPELINE_NATIVE_DECODER"; -const NATIVE_DECODER_DEVICE_ENV: &str = "ONNX_GENAI_PIPELINE_NATIVE_DECODER_DEVICE"; - fn fixture_dir() -> PathBuf { Path::new(env!("CARGO_MANIFEST_DIR")).join("../../tests/fixtures/tiny-gemma4-vlm-cuda-routed") } @@ -42,31 +44,35 @@ fn cuda_device_index() -> Option { } fn generate_tokens(device: &str) -> anyhow::Result> { - unsafe { - std::env::set_var(NATIVE_DECODER_ENV, "decoder"); - std::env::set_var(NATIVE_DECODER_DEVICE_ENV, device); - } - - let result = (|| { - let mut engine = Engine::from_pipeline_dir(&fixture_dir(), EngineConfig::default())?; - let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![3, 7])); - request.options = GenerateOptions { - max_new_tokens: 4, - temperature: 0.0, - stop_on_eos: false, - ..GenerateOptions::default() - }; - let pipeline_request = PipelineGenerateRequest::new(request) - .with_input("vision_encoder.pixel_values", tiny_pixels()?); - let result = engine.generate_with_pipeline_request(pipeline_request)?; - Ok::<_, anyhow::Error>(result.token_ids) - })(); - - unsafe { - std::env::remove_var(NATIVE_DECODER_ENV); - std::env::remove_var(NATIVE_DECODER_DEVICE_ENV); - } - result + let native_device = if device == "cpu" { + NativeDecodeDevice::Cpu + } else { + NativeDecodeDevice::Cuda { + index: device + .strip_prefix("cuda:") + .and_then(|value| value.parse().ok()), + } + }; + let mut engine = Engine::from_pipeline_dir( + &fixture_dir(), + EngineConfig { + decode_backend: EngineDecodeBackend::Native, + native_device: Some(native_device), + ..EngineConfig::default() + }, + )?; + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![3, 7])); + request.options = GenerateOptions { + max_new_tokens: 4, + temperature: 0.0, + stop_on_eos: false, + ..GenerateOptions::default() + }; + let pipeline_request = PipelineGenerateRequest::new(request) + .with_input("vision_encoder.pixel_values", tiny_pixels()?); + Ok(engine + .generate_with_pipeline_request(pipeline_request)? + .token_ids) } #[test] @@ -75,6 +81,14 @@ fn native_cuda_routed_pipeline_decoder_matches_cpu_token_ids() -> anyhow::Result eprintln!("skipping: no CUDA GPU visible (set CUDA_VISIBLE_DEVICES)"); return Ok(()); }; + let graph = onnx_runtime_loader::load_model(fixture_dir().join("decoder.onnx.textproto"))?; + assert!( + graph + .nodes + .values() + .any(|node| { node.domain == "pkg.nxrt" && node.op_type == "BlockQuantizedMoE" }), + "routed pipeline regression must contain a workspace-bearing QMoE" + ); // Native decoder on CPU binding the routed port host-side — the reference. let cpu_tokens = generate_tokens("cpu")?; @@ -93,3 +107,124 @@ fn native_cuda_routed_pipeline_decoder_matches_cpu_token_ids() -> anyhow::Result ); Ok(()) } + +#[test] +fn native_cuda_routed_qmoe_pipeline_admits_after_workspace() -> anyhow::Result<()> { + let Some(index) = cuda_device_index() else { + eprintln!("skipping: no CUDA GPU visible (set CUDA_VISIBLE_DEVICES)"); + return Ok(()); + }; + let result = (|| { + let mut engine = PipelineEngine::from_dir_with_config( + &fixture_dir(), + EngineConfig { + decode_backend: EngineDecodeBackend::Native, + native_device: Some(NativeDecodeDevice::Cuda { index: Some(index) }), + allow_runtime_override: true, + ..EngineConfig::default() + }, + )?; + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![3, 7])); + request.options = GenerateOptions { + max_new_tokens: 2, + temperature: 0.0, + stop_on_eos: false, + ..GenerateOptions::default() + }; + let pipeline_request = PipelineGenerateRequest::new(request) + .with_input("vision_encoder.pixel_values", tiny_pixels()?); + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let admitted_events = Arc::clone(&events); + let mut admitted = move || admitted_events.lock().unwrap().push("admission"); + let token_events = Arc::clone(&events); + let mut token = move |_| { + token_events.lock().unwrap().push("token"); + Ok(()) + }; + let generated = engine.generate_with_callbacks( + pipeline_request, + Some(&mut admitted), + Some(&mut token), + )?; + assert_eq!(generated.token_ids, vec![0, 5]); + let events = events.lock().unwrap(); + assert_eq!(events.first(), Some(&"admission")); + assert!(events[1..].contains(&"token"), "{events:?}"); + Ok::<_, anyhow::Error>(()) + })(); + result +} + +#[test] +fn native_cuda_routed_qmoe_exhaustion_precedes_admission_and_recovers() -> anyhow::Result<()> { + let Some(index) = cuda_device_index() else { + eprintln!("skipping: no CUDA GPU visible (set CUDA_VISIBLE_DEVICES)"); + return Ok(()); + }; + let result = (|| { + let mut engine = PipelineEngine::from_dir_with_config( + &fixture_dir(), + EngineConfig { + decode_backend: EngineDecodeBackend::Native, + native_device: Some(NativeDecodeDevice::Cuda { index: Some(index) }), + allow_runtime_override: true, + ..EngineConfig::default() + }, + )?; + let baseline = engine.resource_snapshot().vram.used; + engine.set_vram_limit(onnx_genai_engine::ResourceLimit::Bytes(baseline))?; + assert_eq!(engine.resource_snapshot().vram.limit, baseline); + + let request = || { + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![3, 7])); + request.options = GenerateOptions { + max_new_tokens: 1, + temperature: 0.0, + stop_on_eos: false, + ..GenerateOptions::default() + }; + PipelineGenerateRequest::new(request) + .with_input("vision_encoder.pixel_values", tiny_pixels().unwrap()) + }; + let admitted = Arc::new(AtomicBool::new(false)); + let admitted_flag = Arc::clone(&admitted); + let mut on_admitted = move || admitted_flag.store(true, Ordering::Relaxed); + let token = Arc::new(AtomicBool::new(false)); + let token_flag = Arc::clone(&token); + let mut on_token = move |_| { + token_flag.store(true, Ordering::Relaxed); + Ok(()) + }; + let error = engine + .generate_with_callbacks(request(), Some(&mut on_admitted), Some(&mut on_token)) + .expect_err("workspace exhaustion must reject the routed QMoE pipeline"); + assert!( + error.chain().any(|cause| { + matches!( + cause.downcast_ref::(), + Some(onnx_runtime_memory_governor::MemoryError::TierExhausted { + role: onnx_runtime_memory_governor::MemoryRole::Workspace { + step_scoped: false + }, + .. + }) + ) + }), + "{error:#}" + ); + assert!(!admitted.load(Ordering::Relaxed)); + assert!(!token.load(Ordering::Relaxed)); + assert_eq!(engine.resource_snapshot().vram.used, baseline); + + engine.set_vram_limit(onnx_genai_engine::ResourceLimit::Bytes(u64::MAX))?; + let admitted = Arc::new(AtomicBool::new(false)); + let admitted_flag = Arc::clone(&admitted); + let mut on_admitted = move || admitted_flag.store(true, Ordering::Relaxed); + let generated = engine.generate_with_callbacks(request(), Some(&mut on_admitted), None)?; + assert_eq!(generated.token_ids, vec![0]); + assert!(admitted.load(Ordering::Relaxed)); + assert_eq!(engine.resource_snapshot().vram.used, baseline); + Ok::<_, anyhow::Error>(()) + })(); + result +} diff --git a/crates/onnx-genai-server/Cargo.toml b/crates/onnx-genai-server/Cargo.toml index 284bdd5404..f371ee7d50 100644 --- a/crates/onnx-genai-server/Cargo.toml +++ b/crates/onnx-genai-server/Cargo.toml @@ -45,6 +45,7 @@ onnx-genai-metadata = { workspace = true } # generated bindings allow exactly one). Matches onnx-genai-engine / onnx-genai. onnx-genai-ort = { path = "../onnx-genai-ort", version = "=0.1.0-dev.5", default-features = false } onnx-runtime-ep-cuda = { workspace = true, optional = true } +onnx-runtime-memory-governor = { workspace = true } onnx-genai-preprocess = { workspace = true } tokio = { workspace = true } axum = { workspace = true, features = ["multipart"] } @@ -61,7 +62,6 @@ reqwest = { version = "0.12", default-features = false, features = ["rustls-tls" toml = { version = "0.8", default-features = false, features = ["parse"] } [dev-dependencies] -onnx-runtime-memory-governor = { workspace = true } onnx-runtime-virtual-memory = { workspace = true } tower = { workspace = true } image = { workspace = true, features = ["png"] } diff --git a/crates/onnx-genai-server/src/driver.rs b/crates/onnx-genai-server/src/driver.rs index 00a7b8631b..28682651bb 100644 --- a/crates/onnx-genai-server/src/driver.rs +++ b/crates/onnx-genai-server/src/driver.rs @@ -143,6 +143,9 @@ impl DriverFailure { matches!( source.downcast_ref::(), Some(SchedulerAdmissionError::ByteBudget { .. }) + ) || matches!( + source.downcast_ref::(), + Some(onnx_runtime_memory_governor::MemoryError::TierExhausted { .. }) ) }); Self { @@ -1353,20 +1356,32 @@ fn run_pipeline_generation( return; } }; - let _ = admission.send(Ok(())); + let mut admission = Some(admission); let mut callback = |token: GenerateToken| -> anyhow::Result<()> { metrics.token(); events .try_send(DriverEvent::Token(token)) .context("stream receiver closed") }; - match engine.generate_with_callback(pipeline_request, Some(&mut callback)) { + let result = { + let mut admitted = || { + if let Some(sender) = admission.take() { + let _ = sender.send(Ok(())); + } + }; + engine.generate_with_callbacks(pipeline_request, Some(&mut admitted), Some(&mut callback)) + }; + match result { Ok(result) => { metrics.result(result.token_ids.len(), result.prefix_cache_hit_len); let _ = events.try_send(DriverEvent::Finished(result)); } Err(err) => { - let _ = events.try_send(DriverEvent::Error(DriverFailure::from_engine_error(&err))); + let failure = DriverFailure::from_engine_error(&err); + if let Some(sender) = admission.take() { + let _ = sender.send(Err(failure.clone())); + } + let _ = events.try_send(DriverEvent::Error(failure)); } } } @@ -1558,6 +1573,20 @@ mod admission_tests { DriverFailure::from_engine_error(&memory_error).kind, DriverFailureKind::MemoryOverload ); + let workspace_error: anyhow::Error = + onnx_runtime_memory_governor::MemoryError::TierExhausted { + tier: "device", + requested: 4096, + used: 8192, + limit: 8192, + available: 0, + role: onnx_runtime_memory_governor::MemoryRole::Workspace { step_scoped: false }, + } + .into(); + assert_eq!( + DriverFailure::from_engine_error(&workspace_error).kind, + DriverFailureKind::MemoryOverload + ); let batch_error: anyhow::Error = SchedulerAdmissionError::BatchFull { running: 32, diff --git a/crates/onnx-runtime-ep-api/src/kernel.rs b/crates/onnx-runtime-ep-api/src/kernel.rs index 695e50ecc2..8ebe441fa3 100644 --- a/crates/onnx-runtime-ep-api/src/kernel.rs +++ b/crates/onnx-runtime-ep-api/src/kernel.rs @@ -3,9 +3,10 @@ use std::borrow::Cow; use onnx_runtime_ir::TensorLayout; +use onnx_runtime_memory_governor::MemoryRole; use crate::error::Result; -use crate::tensor::{TensorMut, TensorView}; +use crate::tensor::{DevicePtrMut, TensorMut, TensorView}; use crate::weight::WeightHandle; /// A cost estimate for running a kernel, consumed by the placement cost model @@ -446,6 +447,70 @@ impl<'a> KernelInput<'a> { } } +/// Concrete tensor facts available during prepare-only graph traversal. +#[derive(Clone, Copy, Debug)] +pub struct TensorMetadata<'a> { + pub dtype: onnx_runtime_ir::DataType, + pub shape: &'a [usize], + pub present: bool, +} + +impl<'a> TensorMetadata<'a> { + pub const fn new(dtype: onnx_runtime_ir::DataType, shape: &'a [usize], present: bool) -> Self { + Self { + dtype, + shape, + present, + } + } +} + +/// How long prepared kernel workspace remains live. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WorkspaceLifetime { + StepScoped, + SessionPersistent, +} + +/// A kernel's exact owned-scratch request for one concrete input geometry. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WorkspaceRequirement { + pub bytes: u64, + pub alignment: usize, + pub lifetime: WorkspaceLifetime, + pub role: MemoryRole, +} + +impl WorkspaceRequirement { + pub const NONE: Self = Self { + bytes: 0, + alignment: 1, + lifetime: WorkspaceLifetime::StepScoped, + role: MemoryRole::Workspace { step_scoped: true }, + }; +} + +/// Executor-owned prepared workspace handed to a kernel for one dispatch. +#[derive(Clone, Copy, Debug)] +pub struct WorkspaceView { + ptr: DevicePtrMut, + bytes: usize, +} + +impl WorkspaceView { + pub const fn new(ptr: DevicePtrMut, bytes: usize) -> Self { + Self { ptr, bytes } + } + + pub const fn ptr(self) -> DevicePtrMut { + self.ptr + } + + pub const fn bytes(self) -> usize { + self.bytes + } +} + /// A kernel ready to execute a specific op with specific shapes (§4.2). pub trait Kernel: Send { /// Tell the kernel which positional inputs are immutable graph constants. @@ -460,6 +525,23 @@ pub trait Kernel: Send { /// Execute over device-resident inputs/outputs. fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()>; + /// Report exact owned scratch needed for these concrete inputs. + fn workspace_requirement(&self, inputs: &[TensorMetadata<'_>]) -> Result { + let _ = inputs; + Ok(WorkspaceRequirement::NONE) + } + + /// Execute using workspace prepared before request admission. + fn execute_with_workspace( + &self, + inputs: &[TensorView], + outputs: &mut [TensorMut], + workspace: Option, + ) -> Result<()> { + let _ = workspace; + self.execute(inputs, outputs) + } + /// Execute through the general weight-delivery seam. /// /// The default adapter accepts only resident tensor inputs and forwards to diff --git a/crates/onnx-runtime-ep-api/src/lib.rs b/crates/onnx-runtime-ep-api/src/lib.rs index ab5b20aa01..340b1b32c9 100644 --- a/crates/onnx-runtime-ep-api/src/lib.rs +++ b/crates/onnx-runtime-ep-api/src/lib.rs @@ -46,8 +46,9 @@ pub use error::{EpError, Result}; pub use kernel::{ ARG_BYTES, ARG_DEVICE, ARG_FLOPS, ARG_KERNEL_VARIANT, ARG_KERNEL_VARIANT_REASON, CAT_KERNEL_WORKER, CaptureSupport, Cost, Kernel, KernelInput, KernelMatch, - KernelVariantSelection, ViewOutput, kernel_variant_tracing_enabled, kernel_worker_span, - record_kernel_metrics, record_kernel_variant_selection, record_kernel_variant_stage_selection, + KernelVariantSelection, TensorMetadata, ViewOutput, WorkspaceLifetime, WorkspaceRequirement, + WorkspaceView, kernel_variant_tracing_enabled, kernel_worker_span, record_kernel_metrics, + record_kernel_variant_selection, record_kernel_variant_stage_selection, }; pub use onnx_runtime_optimizer::OptimizationPass as OptimizerPass; pub use provider::{ @@ -127,5 +128,8 @@ mod error { #[error(transparent)] WeightHandle(#[from] crate::weight::WeightHandleError), + + #[error("{0}")] + Memory(#[from] onnx_runtime_memory_governor::MemoryError), } } diff --git a/crates/onnx-runtime-ep-api/src/provider.rs b/crates/onnx-runtime-ep-api/src/provider.rs index 15d63e7476..340e31b11e 100644 --- a/crates/onnx-runtime-ep-api/src/provider.rs +++ b/crates/onnx-runtime-ep-api/src/provider.rs @@ -692,6 +692,20 @@ pub trait ExecutionProvider: Send + Sync { None } + /// Reserve governed bytes for executor-owned kernel workspace. + /// + /// Providers whose allocator already charges committed bytes may return + /// `None`; providers backed by an eager allocator retain the returned lease + /// alongside the allocation. The default preserves compatibility for + /// providers without a device-memory governor. + fn reserve_workspace( + &self, + _bytes: u64, + _role: onnx_runtime_memory_governor::MemoryRole, + ) -> Result> { + Ok(None) + } + /// Place any long-lived device memory this provider holds under `governor`. /// /// Some providers keep a standing pool for as long as a model is loaded -- diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/block_quantized_moe.rs b/crates/onnx-runtime-ep-cuda/src/kernels/block_quantized_moe.rs index 9cf2c7d192..904fb69a55 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/block_quantized_moe.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/block_quantized_moe.rs @@ -18,13 +18,16 @@ use std::borrow::Cow; use std::ffi::c_void; -use std::sync::Arc; -use std::sync::{Mutex, OnceLock}; +use std::sync::{Arc, OnceLock}; use cudarc::driver::sys::CUdeviceptr; use cudarc::driver::{LaunchConfig, PushKernelArg}; -use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView}; +use onnx_runtime_ep_api::{ + EpError, Kernel, KernelFactory, Result, TensorMetadata, TensorMut, TensorView, + WorkspaceLifetime, WorkspaceRequirement, WorkspaceView, +}; use onnx_runtime_ir::{DataType, Node, Shape}; +use onnx_runtime_memory_governor::MemoryRole; use crate::error::driver_err; use crate::kernels::block_quantized_matmul::{BlockFormat, decoder_prelude}; @@ -458,7 +461,6 @@ impl KernelFactory for BlockQuantizedMoEFactory { runtime: self.runtime.clone(), attributes, format, - scratch: Mutex::new(ScratchPool::default()), })) } } @@ -538,7 +540,132 @@ pub struct BlockQuantizedMoEKernel { runtime: Arc, attributes: MoeAttributes, format: BlockFormat, - scratch: Mutex, +} + +const WORKSPACE_ALIGNMENT: usize = 256; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct QmoeWorkspaceLayout { + offsets: [usize; 6], + bytes: usize, +} + +fn align_up(value: usize, alignment: usize) -> Result { + let mask = alignment - 1; + value + .checked_add(mask) + .map(|value| value & !mask) + .ok_or_else(|| error("BlockQuantizedMoE workspace alignment overflow")) +} + +fn qmoe_workspace_layout( + input_shape: &[usize], + router_shape: &[usize], + fc1_shape: &[usize], + top_k: usize, + swiglu_fusion: usize, + has_fc3: bool, +) -> Result { + if !matches!(input_shape.len(), 2 | 3) { + return Err(error(format!( + "input must be 2-D [rows, hidden] or 3-D [batch, sequence, hidden], got {input_shape:?}" + ))); + } + if router_shape.len() != 2 { + return Err(error(format!( + "router_logits must be rank 2, got {router_shape:?}" + ))); + } + if fc1_shape.len() != 4 { + return Err(error(format!( + "fc1_experts_weights must be rank 4, got {fc1_shape:?}" + ))); + } + let hidden = *input_shape + .last() + .ok_or_else(|| error("input rank unexpectedly empty"))?; + let rows = checked_product( + &input_shape[..input_shape.len() - 1], + "flattened input row count", + )?; + if router_shape[0] != rows { + return Err(error(format!( + "router_logits rows {} must equal flattened input rows {rows}", + router_shape[0] + ))); + } + if top_k == 0 || top_k > router_shape[1] { + return Err(error(format!( + "requires 0 < k <= num_experts, got k={top_k} and num_experts={}", + router_shape[1] + ))); + } + if rows == 0 || hidden == 0 { + return Ok(QmoeWorkspaceLayout { + offsets: [0; 6], + bytes: 0, + }); + } + let fc1_out = fc1_shape[1]; + let inter = if swiglu_fusion == 0 { + fc1_out + } else { + if !fc1_out.is_multiple_of(2) { + return Err(error(format!( + "fused SwiGLU fc1_out must be even, got {fc1_out}" + ))); + } + fc1_out / 2 + }; + if inter == 0 { + return Err(error("inferred inter dimension must be non-zero")); + } + let routes = checked_product(&[rows, top_k], "route count")?; + let sizes = [ + checked_bytes(routes, std::mem::size_of::(), "route indices")?, + checked_bytes(routes, std::mem::size_of::(), "route weights")?, + checked_bytes( + checked_product(&[routes, fc1_out], "FC1 scratch element count")?, + 4, + "FC1 scratch", + )?, + if has_fc3 { + checked_bytes( + checked_product(&[routes, inter], "FC3 scratch element count")?, + 4, + "FC3 scratch", + )? + } else { + 0 + }, + checked_bytes( + checked_product(&[routes, inter], "intermediate element count")?, + 4, + "intermediate scratch", + )?, + checked_bytes( + checked_product(&[routes, hidden], "route output element count")?, + 4, + "route output scratch", + )?, + ]; + let mut offsets = [0; 6]; + let mut cursor = 0usize; + for (index, size) in sizes.into_iter().enumerate() { + if size == 0 { + offsets[index] = 0; + continue; + } + cursor = align_up(cursor, WORKSPACE_ALIGNMENT)?; + offsets[index] = cursor; + cursor = cursor + .checked_add(size) + .ok_or_else(|| error("BlockQuantizedMoE workspace byte count overflow"))?; + } + Ok(QmoeWorkspaceLayout { + offsets, + bytes: cursor, + }) } /// A validated per-projection packed expert-weight tensor. `packed` is the @@ -615,6 +742,43 @@ impl<'a> PackedExperts<'a> { impl Kernel for BlockQuantizedMoEKernel { fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> { + let _ = (inputs, outputs); + Err(error( + "BlockQuantizedMoE requires workspace prepared before execution", + )) + } + + fn workspace_requirement(&self, inputs: &[TensorMetadata<'_>]) -> Result { + if inputs.len() < 5 { + return Err(error(format!( + "expected at least 5 input metadata entries, got {}", + inputs.len() + ))); + } + let has_fc3 = inputs.get(6).is_some_and(|input| input.present); + let layout = qmoe_workspace_layout( + inputs[0].shape, + inputs[1].shape, + inputs[2].shape, + self.attributes.k, + self.attributes.swiglu_fusion, + has_fc3, + )?; + Ok(WorkspaceRequirement { + bytes: u64::try_from(layout.bytes) + .map_err(|_| error("BlockQuantizedMoE workspace does not fit u64"))?, + alignment: WORKSPACE_ALIGNMENT, + lifetime: WorkspaceLifetime::SessionPersistent, + role: MemoryRole::Workspace { step_scoped: false }, + }) + } + + fn execute_with_workspace( + &self, + inputs: &[TensorView], + outputs: &mut [TensorMut], + workspace: Option, + ) -> Result<()> { if !(5..=9).contains(&inputs.len()) || outputs.len() != 1 { return Err(error(format!( "expected 5 to 9 inputs and exactly 1 output, got {} inputs and {} outputs", @@ -774,30 +938,33 @@ impl Kernel for BlockQuantizedMoEKernel { return Ok(()); } + let layout = qmoe_workspace_layout( + input_shape, + inputs[1].shape, + inputs[2].shape, + self.attributes.k, + self.attributes.swiglu_fusion, + fc3.is_some(), + )?; let routes = checked_product(&[rows, self.attributes.k], "route count")?; - let route_index_bytes = checked_bytes(routes, std::mem::size_of::(), "route indices")?; - let route_weight_bytes = - checked_bytes(routes, std::mem::size_of::(), "route weights")?; - let fc1_elements = checked_product(&[routes, fc1_out], "FC1 scratch element count")?; - let fc1_bytes = checked_bytes(fc1_elements, 4, "FC1 scratch")?; - let inter_elements = checked_product(&[routes, inter], "intermediate element count")?; - let inter_bytes = checked_bytes(inter_elements, 4, "intermediate scratch")?; - let route_output_elements = - checked_product(&[routes, hidden], "route output element count")?; - let route_output_bytes = checked_bytes(route_output_elements, 4, "route output scratch")?; - - let mut scratch = self - .scratch - .lock() - .map_err(|_| error("BlockQuantizedMoE scratch pool mutex poisoned"))?; - let route_indices = scratch.ensure(&self.runtime, 0, route_index_bytes)?; - let route_weights_ptr = scratch.ensure(&self.runtime, 1, route_weight_bytes)?; - let fc1_output = scratch.ensure(&self.runtime, 2, fc1_bytes)?; - let fc3_output = fc3 - .map(|_| scratch.ensure(&self.runtime, 3, inter_bytes)) - .transpose()?; - let activated = scratch.ensure(&self.runtime, 4, inter_bytes)?; - let route_output = scratch.ensure(&self.runtime, 5, route_output_bytes)?; + let workspace = workspace.ok_or_else(|| { + error("BlockQuantizedMoE execute reached compute without prepared workspace") + })?; + if workspace.bytes() < layout.bytes { + return Err(error(format!( + "BlockQuantizedMoE prepared-workspace invariant violated: execution requires {} bytes but only {} were prepared", + layout.bytes, + workspace.bytes() + ))); + } + let base = cuptr(workspace.ptr().0.cast_const()); + let ptr = |index: usize| base + layout.offsets[index] as u64; + let route_indices = ptr(0); + let route_weights_ptr = ptr(1); + let fc1_output = ptr(2); + let fc3_output = fc3.map(|_| ptr(3)); + let activated = ptr(4); + let route_output = ptr(5); self.launch_route( &inputs[1], @@ -834,7 +1001,6 @@ impl Kernel for BlockQuantizedMoEKernel { rows, hidden, )?; - drop(scratch); self.runtime.synchronize() } @@ -1042,66 +1208,6 @@ impl BlockQuantizedMoEKernel { } } -const SCRATCH_SLOTS: usize = 6; - -#[derive(Clone, Copy, Debug, Default)] -struct ScratchSlot { - ptr: CUdeviceptr, - capacity: usize, -} - -#[derive(Debug)] -struct ScratchPool { - slots: [ScratchSlot; SCRATCH_SLOTS], -} - -impl Default for ScratchPool { - fn default() -> Self { - Self { - slots: [ScratchSlot::default(); SCRATCH_SLOTS], - } - } -} - -impl ScratchPool { - fn ensure(&mut self, runtime: &CudaRuntime, index: usize, bytes: usize) -> Result { - let slot = &mut self.slots[index]; - let bytes = bytes.max(1); - if slot.ptr != 0 && slot.capacity >= bytes { - return Ok(slot.ptr); - } - let fresh = runtime.alloc_raw(bytes)?; - if slot.ptr != 0 { - // SAFETY: the previous pointer came from this runtime and is replaced - // only after the new allocation succeeds. - unsafe { - let _ = runtime.free_raw(slot.ptr); - } - } - slot.ptr = fresh; - slot.capacity = bytes; - Ok(fresh) - } -} - -impl Drop for BlockQuantizedMoEKernel { - fn drop(&mut self) { - let scratch = self - .scratch - .get_mut() - .expect("cuda_ep BlockQuantizedMoE scratch pool poisoned"); - for slot in scratch.slots.iter_mut().rev() { - if slot.ptr != 0 { - // SAFETY: every non-zero pointer came from this runtime and is - // freed exactly once when the kernel is dropped. - let _ = unsafe { self.runtime.free_raw(slot.ptr) }; - slot.ptr = 0; - slot.capacity = 0; - } - } - } -} - fn tensor_ptr(tensor: &TensorView) -> CUdeviceptr { cuptr(tensor.data_ptr::() as *const c_void) } @@ -1220,3 +1326,55 @@ fn as_i32(name: &str, value: usize) -> Result { fn error(message: impl Into) -> EpError { EpError::KernelFailed(format!("cuda_ep {DOMAIN}::{OP}: {}", message.into())) } + +#[cfg(test)] +mod workspace_tests { + use super::*; + + #[test] + fn workspace_layout_varies_with_prompt_top_k_and_expert_width() { + for (rows, top_k, hidden, fc1_out) in + [(1, 1, 128, 256), (17, 2, 512, 1024), (4096, 8, 2048, 6144)] + { + let layout = qmoe_workspace_layout( + &[1, rows, hidden], + &[rows, 64], + &[64, fc1_out, 1, 1], + top_k, + 1, + false, + ) + .unwrap(); + assert!(layout.bytes > 0); + assert!( + layout + .offsets + .iter() + .filter(|&&offset| offset != 0) + .all(|offset| offset % WORKSPACE_ALIGNMENT == 0) + ); + } + } + + #[test] + fn workspace_layout_zero_rows_requires_no_bytes() { + let layout = + qmoe_workspace_layout(&[1, 0, 128], &[0, 8], &[8, 256, 1, 1], 2, 1, false).unwrap(); + assert_eq!(layout.bytes, 0); + assert_eq!(layout.offsets, [0; 6]); + } + + #[test] + fn workspace_layout_rejects_overflow() { + let error = qmoe_workspace_layout( + &[usize::MAX, 2], + &[usize::MAX, 8], + &[8, 256, 1, 1], + 2, + 1, + false, + ) + .unwrap_err(); + assert!(error.to_string().contains("exceeds usize limits")); + } +} diff --git a/crates/onnx-runtime-ep-cuda/src/provider.rs b/crates/onnx-runtime-ep-cuda/src/provider.rs index cf95851465..08210b73e4 100644 --- a/crates/onnx-runtime-ep-cuda/src/provider.rs +++ b/crates/onnx-runtime-ep-cuda/src/provider.rs @@ -56,6 +56,7 @@ pub struct CudaExecutionProvider { /// Defaults to CudaDeviceAllocator, which is the cuMemAlloc call this /// EP used to make directly. memory: Arc, + governor: Option>, /// Installed by `adopt_memory_governor` when VMM is enabled, and read in /// place of `memory` from then on. /// @@ -150,6 +151,7 @@ impl CudaExecutionProvider { memory: Arc::new(crate::device_allocator::CudaDeviceAllocator::new( runtime.cuda_context(), )), + governor: governor.clone(), // Built here rather than at governor adoption, because on the // native path the session allocates every tensor it will use while // loading -- which is before any governor reaches this provider. @@ -721,16 +723,7 @@ impl ExecutionProvider for CudaExecutionProvider { let ptr = self .memory() .allocate_committed(size, alignment, committed_ranges) - .map_err(|error| { - // Keep what the allocator said. Reporting every failure as "out of - // memory" would describe an alignment rejection or a dead context - // as exhausted VRAM and send the reader looking in the wrong place. - EpError::KernelFailed(format!( - "cuda_ep: could not allocate {size} bytes aligned to {alignment} on CUDA device \ - {}: {error}", - self.device.index - )) - })?; + .map_err(EpError::Memory)?; self.ep_allocations.fetch_add(1, Ordering::Relaxed); // SAFETY: `ptr` is a fresh, unique, non-null device allocation of // >= `size` bytes owned by this EP and freed exactly once in @@ -1062,6 +1055,28 @@ impl ExecutionProvider for CudaExecutionProvider { )) } + fn reserve_workspace( + &self, + bytes: u64, + role: onnx_runtime_memory_governor::MemoryRole, + ) -> Result> { + if self.memory().commits_on_demand() { + return Ok(None); + } + self.governor + .as_deref() + .map(|governor| { + governor.reserve( + onnx_runtime_memory_governor::Tier::Device, + bytes, + role, + onnx_runtime_memory_governor::HolderId::new(64), + ) + }) + .transpose() + .map_err(Into::into) + } + /// True when the VMM arena is in use: it maps 2 MiB granules as spans are /// handed out and leases each one before mapping it, so committed memory /// tracks real use rather than the largest request anyone might make. diff --git a/crates/onnx-runtime-ep-cuda/tests/block_quantized_moe_gpu.rs b/crates/onnx-runtime-ep-cuda/tests/block_quantized_moe_gpu.rs index 19df4dc1b5..4d90d74172 100644 --- a/crates/onnx-runtime-ep-cuda/tests/block_quantized_moe_gpu.rs +++ b/crates/onnx-runtime-ep-cuda/tests/block_quantized_moe_gpu.rs @@ -23,7 +23,8 @@ //! these as ignored unless `gpu-tests` is enabled. use onnx_runtime_ep_api::{ - DeviceBuffer, DevicePtr, DevicePtrMut, ExecutionProvider, KernelMatch, TensorMut, TensorView, + DeviceBuffer, DevicePtr, DevicePtrMut, ExecutionProvider, KernelMatch, TensorMetadata, + TensorMut, TensorView, WorkspaceView, }; use onnx_runtime_ep_cpu::CpuExecutionProvider; use onnx_runtime_ep_cuda::CudaExecutionProvider; @@ -374,7 +375,22 @@ fn run_gpu( &output_strides, ep.device_id(), ); - kernel.execute(&views, &mut [output_view])?; + let metadata = views + .iter() + .map(|view| TensorMetadata::new(view.dtype, view.shape, !view.is_absent())) + .collect::>(); + let requirement = kernel.workspace_requirement(&metadata)?; + let workspace_bytes = usize::try_from(requirement.bytes) + .map_err(|_| onnx_runtime_ep_api::EpError::KernelFailed("workspace too large".into()))?; + let mut workspace = ep.allocate(workspace_bytes, requirement.alignment)?; + kernel.execute_with_workspace( + &views, + &mut [output_view], + Some(WorkspaceView::new( + DevicePtrMut(workspace.as_mut_ptr()), + workspace_bytes, + )), + )?; let mut output = vec![0u8; output_len * 4]; // SAFETY: the destination exactly covers the f32 output allocation. unsafe { runtime.dtoh(&mut output, cuptr(output_buffer.as_ptr()))? }; @@ -383,6 +399,7 @@ fn run_gpu( ep.deallocate(buffer)?; } ep.deallocate(output_buffer)?; + ep.deallocate(workspace)?; Ok(output .chunks_exact(4) .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap())) diff --git a/crates/onnx-runtime-session/src/executor/bindings.rs b/crates/onnx-runtime-session/src/executor/bindings.rs index 65151f0073..5fa18705bf 100644 --- a/crates/onnx-runtime-session/src/executor/bindings.rs +++ b/crates/onnx-runtime-session/src/executor/bindings.rs @@ -1,6 +1,233 @@ use super::*; impl Executor { + pub(crate) fn workspace_node_locations(&self) -> Vec { + fn collect(graph: &Graph, scope: &str, out: &mut Vec) { + for (node_id, node) in graph.nodes.iter() { + if node.domain == onnx_runtime_ir::RUNTIME_DOMAIN + && node.op_type == "BlockQuantizedMoE" + { + out.push(format!( + "{scope}node#{} '{}::{}'", + node_id.0, node.domain, node.op_type + )); + } + } + for ((node_id, attribute), child) in &graph.subgraphs { + collect( + child, + &format!("{scope}node#{}/{attribute}/", node_id.0), + out, + ); + } + } + let mut locations = Vec::new(); + collect(&self.graph, "", &mut locations); + locations + } + + /// Resolve concrete metadata and reserve kernel workspace without executing + /// any graph node. + pub(crate) fn prepare_with_device_bindings( + &mut self, + inputs: &[(&str, &Tensor)], + bindings: &mut [DeviceIoBinding], + ) -> Result { + self.workspace_preparation_required = true; + let external = self.prepare_external_bindings(bindings)?; + let symbols = self.bind_symbols(inputs, &external)?; + let resolved = self.resolve_soft(&symbols); + let mut peak = WorkspaceRequirement::NONE; + + for pi in 0..self.plan.len() { + let node_id = self.plan[pi].node_id; + let node = self.graph.node(node_id); + if node.domain != onnx_runtime_ir::RUNTIME_DOMAIN || node.op_type != "BlockQuantizedMoE" + { + continue; + } + let input_shapes = self.plan[pi] + .inputs + .iter() + .enumerate() + .map(|(index, input)| match input { + None => Ok(Vec::new()), + Some(value) => resolved.get(value).cloned().ok_or_else(|| { + let value_name = self + .graph + .value(*value) + .name + .clone() + .unwrap_or_else(|| format!("value#{}", value.0)); + SessionError::Internal(format!( + "prepare-only workspace planning cannot resolve input {index} \ + '{value_name}' for node {} ('{}::{}'); its shape is runtime-dependent \ + and no exact graph-metadata bound is available", + node_id.0, node.domain, node.op_type + )) + }), + }) + .collect::>>()?; + let constant_inputs = self.plan[pi] + .inputs + .iter() + .map(|input| { + input.is_some_and(|value| self.graph.initializers.contains_key(&value)) + }) + .collect::>(); + let opset = effective_opset(&self.graph, node); + let (kernel, key) = self.cache.get_or_create( + node_id, + node, + &input_shapes, + &self.plan[pi].input_dtypes, + &constant_inputs, + opset, + self.ep.as_ref(), + )?; + self.kernel_bindings[pi] = Some(key); + let metadata = input_shapes + .iter() + .zip(&self.plan[pi].input_dtypes) + .zip(&self.plan[pi].inputs) + .map(|((shape, dtype), input)| TensorMetadata::new(*dtype, shape, input.is_some())) + .collect::>(); + let requirement = kernel.workspace_requirement(&metadata)?; + match requirement.role { + onnx_runtime_memory_governor::MemoryRole::Workspace { step_scoped } + if step_scoped + == matches!(requirement.lifetime, WorkspaceLifetime::StepScoped) => {} + _ => { + return Err(SessionError::Internal(format!( + "node {} ('{}::{}') returned an inconsistent workspace role/lifetime", + node_id.0, node.domain, node.op_type + ))); + } + } + if requirement.bytes > peak.bytes { + peak = requirement; + } else if requirement.bytes == peak.bytes { + peak.alignment = peak.alignment.max(requirement.alignment); + } + } + let child_graphs = self.graph.subgraphs.values().collect::>(); + for child in child_graphs { + self.collect_nested_workspace_requirement(child, &symbols, "control-flow/", &mut peak)?; + } + + if peak.bytes == 0 { + return Ok(peak); + } + let bytes = usize::try_from(peak.bytes).map_err(|_| { + SessionError::Internal(format!( + "workspace requirement {} does not fit usize", + peak.bytes + )) + })?; + let slot = match peak.lifetime { + WorkspaceLifetime::SessionPersistent => &mut self.persistent_workspace, + WorkspaceLifetime::StepScoped => &mut self.step_workspace, + }; + if slot + .as_ref() + .is_some_and(|prepared| prepared.bytes >= bytes && prepared.alignment >= peak.alignment) + { + return Ok(peak); + } + if let Some(old) = slot.take() { + self.ep.deallocate(old.buffer)?; + } + let lease = self.ep.reserve_workspace(peak.bytes, peak.role)?; + let fresh = self.ep.allocate(bytes, peak.alignment)?; + *slot = Some(PreparedWorkspace { + buffer: fresh, + _lease: lease, + bytes, + alignment: peak.alignment, + }); + Ok(peak) + } + + fn collect_nested_workspace_requirement( + &self, + graph: &Graph, + symbols: &HashMap, + scope: &str, + peak: &mut WorkspaceRequirement, + ) -> Result<()> { + for (node_id, node) in graph.nodes.iter() { + if node.domain != onnx_runtime_ir::RUNTIME_DOMAIN || node.op_type != "BlockQuantizedMoE" + { + continue; + } + let input_shapes = node + .inputs + .iter() + .enumerate() + .map(|(index, input)| { + match input { + None => Ok(Vec::new()), + Some(value) => substitute(&graph.value(*value).shape, symbols).ok_or_else(|| { + let value_name = graph + .value(*value) + .name + .clone() + .unwrap_or_else(|| format!("value#{}", value.0)); + SessionError::Internal(format!( + "prepare-only workspace planning cannot resolve nested input {index} \ + '{value_name}' for {scope}node#{} ('{}::{}'); its formal/captured \ + shape is runtime-dependent and no exact graph-metadata bound is \ + available", + node_id.0, node.domain, node.op_type + )) + }), + } + }) + .collect::>>()?; + let mut kernel = + self.ep + .get_kernel(node, &input_shapes, effective_opset(&self.graph, node))?; + let constant_inputs = node + .inputs + .iter() + .map(|input| input.is_some_and(|value| graph.initializers.contains_key(&value))) + .collect::>(); + kernel.set_constant_inputs(&constant_inputs); + let metadata = input_shapes + .iter() + .zip(&node.inputs) + .map(|(shape, input)| { + let dtype = input + .map(|value| graph.value(value).dtype) + .unwrap_or(DataType::Undefined); + TensorMetadata::new(dtype, shape, input.is_some()) + }) + .collect::>(); + let requirement = kernel.workspace_requirement(&metadata)?; + if requirement.bytes > peak.bytes { + *peak = requirement; + } else if requirement.bytes == peak.bytes { + peak.alignment = peak.alignment.max(requirement.alignment); + } + } + for ((node_id, attribute), child) in &graph.subgraphs { + self.collect_nested_workspace_requirement( + child, + symbols, + &format!("{scope}node#{}/{attribute}/", node_id.0), + peak, + )?; + } + Ok(()) + } + + pub(super) fn release_step_workspace(&mut self) -> Result<()> { + if let Some(workspace) = self.step_workspace.take() { + self.ep.deallocate(workspace.buffer)?; + } + Ok(()) + } + /// Bind the graph's symbols to concrete sizes from the actual bound-input /// shapes, validating rank and static dims and detecting symbol conflicts. pub(super) fn bind_symbols( @@ -105,7 +332,9 @@ impl Executor { } pub(crate) fn run_outputs(&mut self, inputs: &[(&str, &Tensor)]) -> Result> { - self.run_scoped(inputs, &HashMap::new(), &ExternalBindings::default())? + let result = self.run_scoped(inputs, &HashMap::new(), &ExternalBindings::default()); + self.release_step_workspace()?; + result? .into_iter() .map(|output| { output.ok_or_else(|| { @@ -123,7 +352,9 @@ impl Executor { bindings: &mut [DeviceIoBinding], ) -> Result>> { let external = self.prepare_external_bindings(bindings)?; - self.run_scoped(inputs, &HashMap::new(), &external)? + let result = self.run_scoped(inputs, &HashMap::new(), &external); + self.release_step_workspace()?; + result? .into_iter() .map(|output| match output { None => Ok(None), @@ -142,7 +373,9 @@ impl Executor { bindings: &mut [DeviceIoBinding], ) -> Result { let external = self.prepare_external_bindings(bindings)?; - match self.run_scoped_mode(inputs, &HashMap::new(), &external, RunMode::Capture)? { + let result = self.run_scoped_mode(inputs, &HashMap::new(), &external, RunMode::Capture); + self.release_step_workspace()?; + match result? { ScopedRunResult::Executed(outputs) => { let mut tensors = Vec::with_capacity(outputs.len()); for output in outputs { @@ -195,7 +428,9 @@ impl Executor { self.ep.replay_device_graph()?; return Ok(true); } - match self.run_scoped_mode(&[], &HashMap::new(), &external, RunMode::Replay)? { + let result = self.run_scoped_mode(&[], &HashMap::new(), &external, RunMode::Replay); + self.release_step_workspace()?; + match result? { // `run_scoped_mode` clears `capture_schedule` when a branch flip // retired the graph this step; report that so the caller re-arms. ScopedRunResult::Executed(_) => Ok(self.capture_schedule.is_some()), diff --git a/crates/onnx-runtime-session/src/executor/build.rs b/crates/onnx-runtime-session/src/executor/build.rs index 7b5436ec56..47c91967f9 100644 --- a/crates/onnx-runtime-session/src/executor/build.rs +++ b/crates/onnx-runtime-session/src/executor/build.rs @@ -555,6 +555,10 @@ impl Executor { scan_inline_single_trip_enabled: scan_inline_single_trip_env_enabled(), scan_inline_single_trip_count: 0, kernel_bindings: vec![None; plan_len], + persistent_workspace: None, + step_workspace: None, + inherited_workspace: None, + workspace_preparation_required: false, }; // 5) Fully-static graphs are materialized eagerly (buffers + the whole diff --git a/crates/onnx-runtime-session/src/executor/control_flow.rs b/crates/onnx-runtime-session/src/executor/control_flow.rs index 5ecbba3c46..2b71991b6e 100644 --- a/crates/onnx-runtime-session/src/executor/control_flow.rs +++ b/crates/onnx-runtime-session/src/executor/control_flow.rs @@ -138,10 +138,20 @@ impl ChildExecutor { /// Execute the body with formal inputs in declared order and lexical values /// supplied by name. A cached plan is reused for matching dtype/shapes. + #[cfg(test)] pub(crate) fn run( &mut self, formal_inputs: &[&Tensor], outer_scope: &HashMap, + ) -> Result> { + self.run_with_workspace(formal_inputs, outer_scope, None) + } + + pub(crate) fn run_with_workspace( + &mut self, + formal_inputs: &[&Tensor], + outer_scope: &HashMap, + workspace: Option, ) -> Result> { if self.formal_names.len() != formal_inputs.len() { return Err(SessionError::Internal(format!( @@ -188,6 +198,11 @@ impl ChildExecutor { }; self.runs += 1; + self.compiled[cache_index].exec.inherited_workspace = + workspace.map(|view| (view.ptr().0 as usize, view.bytes())); + self.compiled[cache_index] + .exec + .workspace_preparation_required = workspace.is_some(); let inputs = self .input_names .iter() @@ -540,12 +555,23 @@ impl Executor { self.subgraph_execs.insert(prepared.key.clone(), child); } + let workspace = self + .persistent_workspace + .as_mut() + .or(self.step_workspace.as_mut()) + .map(|prepared| { + WorkspaceView::new(DevicePtrMut(prepared.buffer.as_mut_ptr()), prepared.bytes) + }) + .or_else(|| { + self.inherited_workspace + .map(|(ptr, bytes)| WorkspaceView::new(DevicePtrMut(ptr as *mut _), bytes)) + }); let child = self .subgraph_execs .get_mut(&prepared.key) .expect("child present"); let before = child.stats(); - let result = child.run(formal_inputs, &prepared.captures); + let result = child.run_with_workspace(formal_inputs, &prepared.captures, workspace); let after = child.stats(); self.control_flow_stats.subgraph_builds += after.builds - before.builds; self.control_flow_stats.subgraph_runs += after.runs - before.runs; diff --git a/crates/onnx-runtime-session/src/executor/dispatch.rs b/crates/onnx-runtime-session/src/executor/dispatch.rs index 68e0dde195..3f7590feca 100644 --- a/crates/onnx-runtime-session/src/executor/dispatch.rs +++ b/crates/onnx-runtime-session/src/executor/dispatch.rs @@ -422,6 +422,12 @@ impl Executor { shared_buffers: &mut self.shared_buffers, views_meta: &mut self.views, pinned: &mut self.pinned, + persistent_workspace: &mut self.persistent_workspace, + step_workspace: &mut self.step_workspace, + inherited_workspace: self + .inherited_workspace + .map(|(ptr, bytes)| WorkspaceView::new(DevicePtrMut(ptr as *mut _), bytes)), + workspace_preparation_required: self.workspace_preparation_required, }; // Build the (possibly strided) input views once; they feed both the @@ -1154,6 +1160,10 @@ struct KernelDispatchContext<'a> { shared_buffers: &'a mut HashMap>, views_meta: &'a mut HashMap, pinned: &'a mut HashSet, + persistent_workspace: &'a mut Option, + step_workspace: &'a mut Option, + inherited_workspace: Option, + workspace_preparation_required: bool, } impl KernelDispatchContext<'_> { @@ -1415,9 +1425,94 @@ impl KernelDispatchContext<'_> { }); let execution = { let _s = phase_span!("exec_kernel.compute"); + let metadata = views + .iter() + .map(|view| TensorMetadata::new(view.dtype, view.shape, !view.is_absent())) + .collect::>(); + let requirement = kernel.workspace_requirement(&metadata)?; + let prepared = match requirement.lifetime { + WorkspaceLifetime::SessionPersistent => &mut *self.persistent_workspace, + WorkspaceLifetime::StepScoped => &mut *self.step_workspace, + }; + let workspace = if requirement.bytes == 0 { + None + } else if let Some(inherited) = self.inherited_workspace { + let required = usize::try_from(requirement.bytes).map_err(|_| { + EpError::KernelFailed(format!( + "kernel workspace requirement {} does not fit usize", + requirement.bytes + )) + })?; + if required > inherited.bytes() { + Err(EpError::KernelFailed(format!( + "node {} (op '{}::{}') workspace invariant mismatch: execute requires {} bytes, enclosing preparation supplied {} bytes", + node.id.0, + node.domain, + node.op_type, + required, + inherited.bytes() + )))?; + } + Some(inherited) + } else { + let required = usize::try_from(requirement.bytes).map_err(|_| { + EpError::KernelFailed(format!( + "kernel workspace requirement {} does not fit usize", + requirement.bytes + )) + })?; + let needs_replacement = prepared.as_ref().is_none_or(|workspace| { + required > workspace.bytes || requirement.alignment > workspace.alignment + }); + if needs_replacement && !self.workspace_preparation_required { + // Sequential dispatch is the scratch hand-off boundary. An + // EP may enqueue asynchronous device work, so synchronize + // before retiring the old disposable workspace. Release it + // before acquiring the replacement to avoid charging both. + self.ep.sync()?; + if let Some(old) = prepared.take() { + self.ep.deallocate(old.buffer)?; + } + let lease = self + .ep + .reserve_workspace(requirement.bytes, requirement.role)?; + let buffer = self.ep.allocate(required, requirement.alignment)?; + *prepared = Some(PreparedWorkspace { + buffer, + _lease: lease, + bytes: required, + alignment: requirement.alignment, + }); + } + let prepared = prepared.as_mut().ok_or_else(|| { + EpError::KernelFailed(format!( + "node {} (op '{}::{}') reached execution without prepared {:?} workspace", + node.id.0, node.domain, node.op_type, requirement.lifetime + )) + })?; + if required > prepared.bytes || requirement.alignment > prepared.alignment { + Err(EpError::KernelFailed(format!( + "node {} (op '{}::{}') workspace invariant mismatch: execute requires {} bytes aligned to {}, prepared {} bytes aligned to {}", + node.id.0, + node.domain, + node.op_type, + required, + requirement.alignment, + prepared.bytes, + prepared.alignment + )))?; + } + Some(WorkspaceView::new( + DevicePtrMut(prepared.buffer.as_mut_ptr()), + prepared.bytes, + )) + }; match &kernel_inputs { - Some(inputs) => kernel.execute_with_inputs(inputs, outs), - None => kernel.execute(views, outs), + Some(inputs) if requirement.bytes == 0 => kernel.execute_with_inputs(inputs, outs), + Some(_) => Err(EpError::KernelFailed( + "workspace-bearing lazy-input kernels are not supported".into(), + )), + None => kernel.execute_with_workspace(views, outs, workspace), } }; execution.map_err(|error| { diff --git a/crates/onnx-runtime-session/src/executor/mod.rs b/crates/onnx-runtime-session/src/executor/mod.rs index 372acea2ba..ef49d17fa0 100644 --- a/crates/onnx-runtime-session/src/executor/mod.rs +++ b/crates/onnx-runtime-session/src/executor/mod.rs @@ -51,7 +51,8 @@ use std::time::{Duration, Instant}; use onnx_runtime_ep_api::{ CaptureRegionShapeStatus, DeviceBuffer, DevicePtr, DevicePtrMut, EpError, ExecutionProvider, ExternalMmapRegion, Kernel, KernelInput, KernelMatch, LazyWeight, LazyWeightBoundary, - ResidentWeight, StructuralCaptureDecline, TensorBacking, TensorMut, TensorView, WeightHandle, + ResidentWeight, StructuralCaptureDecline, TensorBacking, TensorMetadata, TensorMut, TensorView, + WeightHandle, WorkspaceLifetime, WorkspaceRequirement, WorkspaceView, }; type OptionalTensorSpecs = Vec)>>; @@ -607,6 +608,12 @@ impl Drop for Executor { for (_, buf) in self.buffers.drain() { let _ = self.ep.deallocate(buf); } + if let Some(workspace) = self.persistent_workspace.take() { + let _ = self.ep.deallocate(workspace.buffer); + } + if let Some(workspace) = self.step_workspace.take() { + let _ = self.ep.deallocate(workspace.buffer); + } self.shared_buffers.clear(); } } diff --git a/crates/onnx-runtime-session/src/executor/run.rs b/crates/onnx-runtime-session/src/executor/run.rs index b1c13830a8..3fde26db00 100644 --- a/crates/onnx-runtime-session/src/executor/run.rs +++ b/crates/onnx-runtime-session/src/executor/run.rs @@ -193,7 +193,7 @@ impl Executor { Ok(()) } - fn validate_required_inputs( + pub(super) fn validate_required_inputs( &self, inputs: &[(&str, &Tensor)], external: &ExternalBindings, diff --git a/crates/onnx-runtime-session/src/executor/state.rs b/crates/onnx-runtime-session/src/executor/state.rs index 98e55cf55b..94b70eac3f 100644 --- a/crates/onnx-runtime-session/src/executor/state.rs +++ b/crates/onnx-runtime-session/src/executor/state.rs @@ -270,6 +270,20 @@ pub(crate) struct Executor { /// Control-flow and sequence nodes always have `None` (they don't use the /// kernel cache). pub(super) kernel_bindings: Vec>, + pub(super) persistent_workspace: Option, + pub(super) step_workspace: Option, + /// Non-owning view of an enclosing executor's prepared workspace. Nested + /// control-flow executors run sequentially, so they may reuse the parent's + /// peak allocation without reserving or allocating a second buffer. + pub(super) inherited_workspace: Option<(usize, usize)>, + pub(super) workspace_preparation_required: bool, +} + +pub(super) struct PreparedWorkspace { + pub(super) buffer: DeviceBuffer, + pub(super) _lease: Option, + pub(super) bytes: usize, + pub(super) alignment: usize, } /// After this many consecutive buffer-identity signature mismatches, F5 Stage 2 diff --git a/crates/onnx-runtime-session/src/executor/tests.rs b/crates/onnx-runtime-session/src/executor/tests.rs index 2e353a0e62..2be7b82e80 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -1,4 +1,4 @@ -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use onnx_runtime_ep_api::{ CaptureSupport, Cost, EpConfig, EpError, ExecutionProviderCapabilities, Fence, Kernel, @@ -813,6 +813,8 @@ fn only_capacity_aware_inputs_keep_physical_capacity() { struct WeightDeliveryKernel { deliveries: Arc>>, + workspace_bytes: u64, + workspace_bytes_per_row: u64, } impl WeightDeliveryKernel { @@ -832,6 +834,30 @@ impl WeightDeliveryKernel { } impl Kernel for WeightDeliveryKernel { + fn workspace_requirement( + &self, + inputs: &[TensorMetadata<'_>], + ) -> onnx_runtime_ep_api::Result { + let rows = inputs + .first() + .and_then(|input| input.shape.first()) + .copied() + .unwrap_or(1) as u64; + let bytes = if self.workspace_bytes_per_row == 0 { + self.workspace_bytes + } else { + self.workspace_bytes_per_row + .checked_mul(rows) + .ok_or_else(|| EpError::KernelFailed("test workspace overflow".into()))? + }; + Ok(WorkspaceRequirement { + bytes, + alignment: if rows >= 4 { 512 } else { 256 }, + lifetime: WorkspaceLifetime::SessionPersistent, + role: onnx_runtime_memory_governor::MemoryRole::Workspace { step_scoped: false }, + }) + } + fn execute( &self, inputs: &[TensorView], @@ -875,6 +901,10 @@ struct WeightDeliveryEp { device: onnx_runtime_ir::DeviceId, allocations: Arc, host_uploads: Arc, + workspace_bytes: u64, + workspace_bytes_per_row: u64, + fail_next_allocation: Arc, + fail_allocation_size: Arc, } impl WeightDeliveryEp { @@ -920,6 +950,10 @@ impl WeightDeliveryEp { device, allocations, host_uploads, + workspace_bytes: 0, + workspace_bytes_per_row: 0, + fail_next_allocation: Arc::new(AtomicBool::new(false)), + fail_allocation_size: Arc::new(AtomicUsize::new(0)), } } @@ -1017,10 +1051,28 @@ impl ExecutionProvider for WeightDeliveryEp { ) -> onnx_runtime_ep_api::Result> { Ok(Box::new(WeightDeliveryKernel { deliveries: Arc::clone(&self.deliveries), + workspace_bytes: self.workspace_bytes, + workspace_bytes_per_row: self.workspace_bytes_per_row, })) } fn allocate(&self, size: usize, alignment: usize) -> onnx_runtime_ep_api::Result { + if self.fail_next_allocation.swap(false, Ordering::Relaxed) { + return Err(EpError::OutOfMemory { + requested: size, + available: 0, + }); + } + if self + .fail_allocation_size + .compare_exchange(size, 0, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + return Err(EpError::OutOfMemory { + requested: size, + available: 0, + }); + } self.allocations.fetch_add(1, Ordering::Relaxed); if self.device.is_host_accessible() { return self.cpu.allocate(size, alignment); @@ -1153,6 +1205,160 @@ fn weight_delivery_fixture() -> (Graph, Arc, std::path::PathBuf) { (graph, Arc::new(store), path) } +#[test] +fn prepare_reserves_static_nested_qmoe_workspace_and_child_reuses_it() { + fn branch(name: &str) -> Graph { + let mut graph = Graph::new(); + let input = graph.create_named_value( + format!("{name}_input"), + DataType::Float32, + static_shape([4]), + ); + graph.set_initializer( + input, + WeightRef::Inline(onnx_runtime_ir::TensorData::from_raw( + DataType::Float32, + vec![4], + [1.0f32, 2.0, 3.0, 4.0] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect(), + )), + ); + let output = graph.create_named_value( + format!("{name}_output"), + DataType::Float32, + static_shape([4]), + ); + let mut node = Node::new( + NodeId(0), + "BlockQuantizedMoE", + vec![Some(input)], + vec![output], + ); + node.domain = onnx_runtime_ir::RUNTIME_DOMAIN.into(); + graph.insert_node(node); + graph.add_output(output); + graph + } + + let mut graph = Graph::new(); + graph.opset_imports.insert(String::new(), 17); + graph + .opset_imports + .insert(onnx_runtime_ir::RUNTIME_DOMAIN.into(), 1); + let cond = graph.create_named_value("cond", DataType::Bool, static_shape([])); + graph.add_input(cond); + let output = graph.create_named_value("output", DataType::Float32, static_shape([4])); + graph.add_output(output); + let if_id = NodeId(0); + graph.insert_node(Node::new(if_id, "If", vec![Some(cond)], vec![output])); + graph + .subgraphs + .insert((if_id, "then_branch".into()), branch("then")); + graph + .subgraphs + .insert((if_id, "else_branch".into()), branch("else")); + + let deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut ep = WeightDeliveryEp::new(false, deliveries); + ep.workspace_bytes = 4096; + let mut executor = Executor::build(graph, Arc::new(WeightStore::new()), Arc::new(ep)).unwrap(); + let cond = Tensor::from_raw(DataType::Bool, vec![], &[1]).unwrap(); + let requirement = executor + .prepare_with_device_bindings(&[("cond", &cond)], &mut []) + .unwrap(); + assert_eq!(requirement.bytes, 4096); + assert_eq!(executor.persistent_workspace.as_ref().unwrap().bytes, 4096); + + let output = executor.run(&[("cond", &cond)]).unwrap(); + assert_eq!(output[0].to_vec_f32(), vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn inference_session_fallback_workspace_grows_retries_and_reuses() { + let mut graph = Graph::new(); + graph + .opset_imports + .insert(onnx_runtime_ir::RUNTIME_DOMAIN.into(), 1); + let rows = SymbolId(0); + let shape = vec![Dim::Symbolic(rows)]; + let input = graph.create_named_value("input", DataType::Float32, shape.clone()); + graph.add_input(input); + let output = graph.create_named_value("output", DataType::Float32, shape); + graph.add_output(output); + let mut node = Node::new( + NodeId(0), + "BlockQuantizedMoE", + vec![Some(input)], + vec![output], + ); + node.domain = onnx_runtime_ir::RUNTIME_DOMAIN.into(); + graph.insert_node(node); + + let inputs = crate::io_meta(&graph, &graph.inputs); + let outputs = crate::io_meta(&graph, &graph.outputs); + let deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut ep = WeightDeliveryEp::new(false, deliveries); + ep.workspace_bytes_per_row = 1024; + let fail_size = Arc::clone(&ep.fail_allocation_size); + let exec = Executor::build(graph, Arc::new(WeightStore::new()), Arc::new(ep)).unwrap(); + let mut session = crate::InferenceSession { + inputs, + outputs, + model_metadata: crate::ModelMetadata::default(), + exec, + decode_inline_exec: None, + ep_context_config: crate::EpContextDumpConfig::default(), + }; + + let small = Tensor::from_f32(&[2], &[1.0, 2.0]).unwrap(); + assert_eq!( + session.run(&[("input", &small)]).unwrap()[0].to_vec_f32(), + vec![1.0, 2.0] + ); + assert_eq!( + session.exec.persistent_workspace.as_ref().unwrap().bytes, + 2048 + ); + assert_eq!( + session + .exec + .persistent_workspace + .as_ref() + .unwrap() + .alignment, + 256 + ); + + let large = Tensor::from_f32(&[4], &[1.0, 2.0, 3.0, 4.0]).unwrap(); + fail_size.store(4096, Ordering::Relaxed); + let _error = session + .run(&[("input", &large)]) + .expect_err("workspace replacement allocation must fail once"); + assert!( + session.exec.persistent_workspace.is_none(), + "failed growth must leave a valid empty slot" + ); + + assert_eq!( + session.run(&[("input", &large)]).unwrap()[0].to_vec_f32(), + vec![1.0, 2.0, 3.0, 4.0] + ); + let grown = session.exec.persistent_workspace.as_ref().unwrap(); + assert_eq!(grown.bytes, 4096); + assert_eq!(grown.alignment, 512); + let ptr = grown.buffer.as_ptr(); + + assert_eq!( + session.run(&[("input", &small)]).unwrap()[0].to_vec_f32(), + vec![1.0, 2.0] + ); + let reused = session.exec.persistent_workspace.as_ref().unwrap(); + assert_eq!(reused.bytes, 4096); + assert_eq!(reused.buffer.as_ptr(), ptr); +} + fn two_node_weight_delivery_fixture() -> (Graph, Arc, std::path::PathBuf) { static NEXT_FILE: AtomicU64 = AtomicU64::new(0); let root = std::env::var_os("CARGO_TARGET_DIR") diff --git a/crates/onnx-runtime-session/src/lib.rs b/crates/onnx-runtime-session/src/lib.rs index 62182d05ab..f7918e52dc 100644 --- a/crates/onnx-runtime-session/src/lib.rs +++ b/crates/onnx-runtime-session/src/lib.rs @@ -37,6 +37,7 @@ pub use executor::{ enable_exec_phase_profile_for_process, exec_phase_stats, plan_double_buffer, print_exec_phase_profile, reset_dense_prefetch_gap_stats, reset_exec_phase_profile, }; +pub use onnx_runtime_ep_api::WorkspaceRequirement; pub use onnx_runtime_loader::{ EpContextDumpConfig, EpContextPartition, Model as EncoderModel, ModelMetadata, }; @@ -1135,6 +1136,20 @@ impl InferenceSession { self.exec.run_with_device_bindings(inputs, bindings) } + /// Prepare exact kernel workspace for a bound run without launching kernels. + pub fn prepare_with_device_bindings( + &mut self, + inputs: &[(&str, &Tensor)], + bindings: &mut [DeviceIoBinding], + ) -> Result { + self.exec.prepare_with_device_bindings(inputs, bindings) + } + + /// Locations of graph nodes that declare owned kernel workspace. + pub fn workspace_node_locations(&self) -> Vec { + self.exec.workspace_node_locations() + } + /// Lazily build the decode-specialized inlined-body sibling executor /// (Inc-1b PR-2, `cohaagen-27b-inc1b-design.md` §1). Idempotent: a second /// call is a no-op. Returns `true` when a decode-inline plan is now diff --git a/scripts/build_tiny_gemma4_vlm_cuda_routed.py b/scripts/build_tiny_gemma4_vlm_cuda_routed.py index d470425aa7..5e8e6844c2 100644 --- a/scripts/build_tiny_gemma4_vlm_cuda_routed.py +++ b/scripts/build_tiny_gemma4_vlm_cuda_routed.py @@ -126,8 +126,94 @@ def build_decoder_routed(path: Path) -> None: past_key = tensor_value("past_key_values.0.key", ir.DataType.FLOAT, [1, 1, "past_sequence", HIDDEN]) past_value = tensor_value("past_key_values.0.value", ir.DataType.FLOAT, [1, 1, "past_sequence", HIDDEN]) + # A zero-valued block-quantized MoE sits on the live inputs_embeds path. Its + # output is added residually, preserving the fixture's closed-form tokens + # while making workspace scale with the exact routed prefill sequence. + qmoe_hidden = 32 + qmoe_projection = initializer( + "qmoe_projection", np.zeros((HIDDEN, qmoe_hidden), dtype=np.float32) + ) + qmoe_projected = node( + "MatMul", [inputs_embeds, qmoe_projection], "qmoe_projected" + ) + qmoe_input_shape = constant( + "qmoe_input_shape", np.array([-1, qmoe_hidden], dtype=np.int64) + ) + qmoe_input = node( + "Reshape", + [qmoe_projected.outputs[0], qmoe_input_shape.outputs[0]], + "qmoe_input", + ) + router_logits_3d = node( + "ReduceMean", + [router_state], + "qmoe_router_logits_3d", + attributes=[ + ir.AttrInt64s("axes", [2]), + ir.AttrInt64("keepdims", 1), + ], + ) + router_shape = constant("qmoe_router_shape", np.array([-1, 1], dtype=np.int64)) + router_logits = node( + "Reshape", + [router_logits_3d.outputs[0], router_shape.outputs[0]], + "qmoe_router_logits", + ) + packed_zero = np.zeros((1, qmoe_hidden, 1, 17), dtype=np.uint8) + qmoe_fc1 = initializer("qmoe_fc1", packed_zero) + qmoe_fc2 = initializer("qmoe_fc2", packed_zero.copy()) + qmoe_fc1_bias = initializer( + "qmoe_fc1_bias", np.zeros((1, qmoe_hidden), dtype=np.float32) + ) + qmoe_fc2_bias = initializer( + "qmoe_fc2_bias", np.zeros((1, qmoe_hidden), dtype=np.float32) + ) + qmoe = ir.Node( + "pkg.nxrt", + "BlockQuantizedMoE", + [ + qmoe_input.outputs[0], + router_logits.outputs[0], + qmoe_fc1, + qmoe_fc1_bias, + qmoe_fc2, + qmoe_fc2_bias, + ], + { + "k": ir.AttrInt64("k", 1), + "activation_type": ir.AttrString("activation_type", "identity"), + "normalize_routing_weights": ir.AttrInt64("normalize_routing_weights", 0), + "swiglu_fusion": ir.AttrInt64("swiglu_fusion", 0), + "format": ir.AttrString("format", "mxfp4"), + "block_layout_version": ir.AttrInt64("block_layout_version", 1), + }, + outputs=[ir.Value(name="qmoe_output")], + ) + qmoe.outputs[0].type = ir.TensorType(ir.DataType.FLOAT) + qmoe.outputs[0].shape = ir.Shape(["rows", qmoe_hidden]) + qmoe_reduce = node( + "ReduceMean", + [qmoe.outputs[0]], + "qmoe_reduce", + attributes=[ + ir.AttrInt64s("axes", [1]), + ir.AttrInt64("keepdims", 1), + ], + ) + qmoe_restore_shape = constant( + "qmoe_restore_shape", np.array([1, -1, 1], dtype=np.int64) + ) + qmoe_restored = node( + "Reshape", + [qmoe_reduce.outputs[0], qmoe_restore_shape.outputs[0]], + "qmoe_restored", + ) + qmoe_residual = node( + "Add", [inputs_embeds, qmoe_restored.outputs[0]], "qmoe_residual" + ) + lm_head = initializer("lm_head", LM_HEAD) - matmul = node("MatMul", [inputs_embeds, lm_head], "logits_base") + matmul = node("MatMul", [qmoe_residual.outputs[0], lm_head], "logits_base") tie_bias = initializer("tie_bias", TIE_BIAS) logits_biased = node("Add", [matmul.outputs[0], tie_bias], "logits_biased") @@ -153,6 +239,17 @@ def build_decoder_routed(path: Path) -> None: [inputs_embeds, router_state, attention_mask, position_ids, past_key, past_value], [logits.outputs[0], present_key.outputs[0], present_value.outputs[0]], nodes=[ + qmoe_projected, + qmoe_input_shape, + qmoe_input, + router_logits_3d, + router_shape, + router_logits, + qmoe, + qmoe_reduce, + qmoe_restore_shape, + qmoe_restored, + qmoe_residual, matmul, logits_biased, router_contrib, @@ -163,8 +260,18 @@ def build_decoder_routed(path: Path) -> None: present_key, present_value, ], - initializers=[lm_head, tie_bias, router_zero, value_offset], - opset_imports={"": 13}, + initializers=[ + qmoe_fc1, + qmoe_fc2, + qmoe_fc1_bias, + qmoe_fc2_bias, + qmoe_projection, + lm_head, + tie_bias, + router_zero, + value_offset, + ], + opset_imports={"": 13, "pkg.nxrt": 1}, name="tiny_gemma4_decoder_routed", ) save_model( diff --git a/tests/fixtures/tiny-gemma4-vlm-cuda-routed/decoder.onnx.textproto b/tests/fixtures/tiny-gemma4-vlm-cuda-routed/decoder.onnx.textproto index cabf590700..10ca5e5ec9 100644 --- a/tests/fixtures/tiny-gemma4-vlm-cuda-routed/decoder.onnx.textproto +++ b/tests/fixtures/tiny-gemma4-vlm-cuda-routed/decoder.onnx.textproto @@ -3,35 +3,189 @@ producer_name: "onnx-genai tiny-gemma4-vlm-cuda-routed fixture" graph { node { input: "inputs_embeds" + input: "qmoe_projection" + output: "qmoe_projected" + name: "node_MatMul_0" + op_type: "MatMul" + } + node { + output: "qmoe_input_shape" + name: "node_Constant_1" + op_type: "Constant" + attribute { + name: "value" + t { + dims: 2 + data_type: 7 + name: "qmoe_input_shape_value" + raw_data: "\377\377\377\377\377\377\377\377 \000\000\000\000\000\000\000" + } + type: TENSOR + } + } + node { + input: "qmoe_projected" + input: "qmoe_input_shape" + output: "qmoe_input" + name: "node_Reshape_2" + op_type: "Reshape" + } + node { + input: "router_state" + output: "qmoe_router_logits_3d" + name: "node_ReduceMean_3" + op_type: "ReduceMean" + attribute { + name: "axes" + ints: 2 + type: INTS + } + attribute { + name: "keepdims" + i: 1 + type: INT + } + } + node { + output: "qmoe_router_shape" + name: "node_Constant_4" + op_type: "Constant" + attribute { + name: "value" + t { + dims: 2 + data_type: 7 + name: "qmoe_router_shape_value" + raw_data: "\377\377\377\377\377\377\377\377\001\000\000\000\000\000\000\000" + } + type: TENSOR + } + } + node { + input: "qmoe_router_logits_3d" + input: "qmoe_router_shape" + output: "qmoe_router_logits" + name: "node_Reshape_5" + op_type: "Reshape" + } + node { + input: "qmoe_input" + input: "qmoe_router_logits" + input: "qmoe_fc1" + input: "qmoe_fc1_bias" + input: "qmoe_fc2" + input: "qmoe_fc2_bias" + output: "qmoe_output" + name: "node_BlockQuantizedMoE_6" + op_type: "BlockQuantizedMoE" + attribute { + name: "k" + i: 1 + type: INT + } + attribute { + name: "activation_type" + s: "identity" + type: STRING + } + attribute { + name: "normalize_routing_weights" + i: 0 + type: INT + } + attribute { + name: "swiglu_fusion" + i: 0 + type: INT + } + attribute { + name: "format" + s: "mxfp4" + type: STRING + } + attribute { + name: "block_layout_version" + i: 1 + type: INT + } + domain: "pkg.nxrt" + } + node { + input: "qmoe_output" + output: "qmoe_reduce" + name: "node_ReduceMean_7" + op_type: "ReduceMean" + attribute { + name: "axes" + ints: 1 + type: INTS + } + attribute { + name: "keepdims" + i: 1 + type: INT + } + } + node { + output: "qmoe_restore_shape" + name: "node_Constant_8" + op_type: "Constant" + attribute { + name: "value" + t { + dims: 3 + data_type: 7 + name: "qmoe_restore_shape_value" + raw_data: "\001\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\001\000\000\000\000\000\000\000" + } + type: TENSOR + } + } + node { + input: "qmoe_reduce" + input: "qmoe_restore_shape" + output: "qmoe_restored" + name: "node_Reshape_9" + op_type: "Reshape" + } + node { + input: "inputs_embeds" + input: "qmoe_restored" + output: "qmoe_residual" + name: "node_Add_10" + op_type: "Add" + } + node { + input: "qmoe_residual" input: "lm_head" output: "logits_base" - name: "node_MatMul_0" + name: "node_MatMul_11" op_type: "MatMul" } node { input: "logits_base" input: "tie_bias" output: "logits_biased" - name: "node_Add_1" + name: "node_Add_12" op_type: "Add" } node { input: "router_state" input: "router_zero" output: "router_contrib" - name: "node_MatMul_2" + name: "node_MatMul_13" op_type: "MatMul" } node { input: "logits_biased" input: "router_contrib" output: "logits" - name: "node_Add_3" + name: "node_Add_14" op_type: "Add" } node { output: "kv_axis" - name: "node_Constant_4" + name: "node_Constant_15" op_type: "Constant" attribute { name: "value" @@ -48,21 +202,21 @@ graph { input: "inputs_embeds" input: "kv_axis" output: "current_key" - name: "node_Unsqueeze_5" + name: "node_Unsqueeze_16" op_type: "Unsqueeze" } node { input: "current_key" input: "value_offset" output: "current_value" - name: "node_Add_6" + name: "node_Add_17" op_type: "Add" } node { input: "past_key_values.0.key" input: "current_key" output: "present.0.key" - name: "node_Concat_7" + name: "node_Concat_18" op_type: "Concat" attribute { name: "axis" @@ -74,7 +228,7 @@ graph { input: "past_key_values.0.value" input: "current_value" output: "present.0.value" - name: "node_Concat_8" + name: "node_Concat_19" op_type: "Concat" attribute { name: "axis" @@ -83,6 +237,45 @@ graph { } } name: "tiny_gemma4_decoder_routed" + initializer { + dims: 1 + dims: 32 + dims: 1 + dims: 17 + data_type: 2 + name: "qmoe_fc1" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + initializer { + dims: 1 + dims: 32 + dims: 1 + dims: 17 + data_type: 2 + name: "qmoe_fc2" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + initializer { + dims: 1 + dims: 32 + data_type: 1 + name: "qmoe_fc1_bias" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + initializer { + dims: 1 + dims: 32 + data_type: 1 + name: "qmoe_fc2_bias" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + initializer { + dims: 4 + dims: 32 + data_type: 1 + name: "qmoe_projection" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } initializer { dims: 4 dims: 8 @@ -287,8 +480,28 @@ graph { } } } + value_info { + name: "qmoe_output" + type { + tensor_type { + elem_type: 1 + shape { + dim { + dim_param: "rows" + } + dim { + dim_value: 32 + } + } + } + } + } } opset_import { domain: "" version: 13 } +opset_import { + domain: "pkg.nxrt" + version: 1 +}