Plan governed QMoE workspace before admission

.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.

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>,
+) -> 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<TokenId> {
     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,14 +165,29 @@ 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();
         }
 
         // 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<()> {

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<onnx_runtime_session::WorkspaceRequirement> {
+        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<onnx_runtime_session::WorkspaceRequirement> {
+        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::<HashMap<_, _>>();
+        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<_>>();
+                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::<Vec<_>>();
+        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.
     ///

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<onnx_runtime_session::WorkspaceRequirement> {
+        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<onnx_runtime_session::WorkspaceRequirement> {
+        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<onnx_runtime_session::WorkspaceRequirement> {
+        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<onnx_runtime_session::WorkspaceRequirement> {
+        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.

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<TokenId> = prompt_tokens
                 .iter()