fix(mocker): include outputs in mooncake delta prompts (#11060)

Signed-off-by: Yongming Ding <yongmingd@nvidia.com>

Author: dreamtalen

lib/mocker/src/loadgen/driver.rs

@@ -5,6 +5,8 @@ use std::cmp::Ordering;
 use std::collections::BinaryHeap;
 
 use anyhow::{Context, Result, anyhow, bail};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
 use rustc_hash::FxHashMap;
 use uuid::Uuid;
 
@@ -38,9 +40,12 @@ enum PromptMode {
     DeltaCumulative,
 }
 
+const SYNTHETIC_DELTA_OUTPUT_SEED: u64 = 0xD37A_0A7E_5EED;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum TurnOutcome {
     Completed,
+    Rejected,
     Cancelled,
 }
 
@@ -78,6 +83,7 @@ struct TurnRuntime {
 struct InFlightTurn {
     session_index: usize,
     turn_index: usize,
+    emitted_output_tokens: usize,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -384,6 +390,7 @@ impl WorkloadDriver {
             u32::try_from(engine_block_size).context("engine_block_size does not fit in u32")?;
         let trace_block_size = trace.block_size;
         let is_concurrency = matches!(&policy, SchedulingPolicy::Concurrency(_));
+        let mut output_rng = StdRng::seed_from_u64(SYNTHETIC_DELTA_OUTPUT_SEED);
         let sessions: Vec<SessionRuntime> = trace
             .sessions
             .into_iter()
@@ -402,12 +409,22 @@ impl WorkloadDriver {
                         } else {
                             None
                         };
+                        let tokens = turn.synthesize_tokens(trace_block_size)?;
+                        let output_token_ids = match turn.output_token_ids {
+                            Some(output_token_ids) => Some(output_token_ids),
+                            None if prompt_mode == PromptMode::DeltaCumulative => Some(
+                                (0..turn.max_output_tokens)
+                                    .map(|_| output_rng.random::<u32>())
+                                    .collect(),
+                            ),
+                            None => None,
+                        };
                         Ok(TurnRuntime {
                             request_id: None,
-                            tokens: turn.synthesize_tokens(trace_block_size)?,
+                            tokens,
                             replay_key: turn.replay_key,
                             max_output_tokens: turn.max_output_tokens,
-                            output_token_ids: turn.output_token_ids,
+                            output_token_ids,
                             delay_after_previous_ms: turn.delay_after_previous_ms,
                             priority: turn.priority,
                             strict_priority: turn.strict_priority,
@@ -417,7 +434,16 @@ impl WorkloadDriver {
                     })
                     .collect::<Result<Vec<_>>>()?;
                 let cumulative_capacity = if prompt_mode == PromptMode::DeltaCumulative {
-                    turns.iter().map(|turn| turn.tokens.len()).sum()
+                    turns
+                        .iter()
+                        .map(|turn| {
+                            turn.tokens.len()
+                                + turn
+                                    .output_token_ids
+                                    .as_ref()
+                                    .map_or(0, |output| output.len())
+                        })
+                        .sum()
                 } else {
                     0
                 };
@@ -538,6 +564,7 @@ impl WorkloadDriver {
                 InFlightTurn {
                     session_index,
                     turn_index,
+                    emitted_output_tokens: 0,
                 },
             );
             emitted.push(ReadyTurn {
@@ -553,9 +580,60 @@ impl WorkloadDriver {
         emitted
     }
 
+    pub fn on_output_token(&mut self, request_uuid: Uuid, token_id: u32) -> Result<()> {
+        if self.prompt_mode == PromptMode::Full {
+            return Ok(());
+        }
+        let in_flight = self
+            .in_flight
+            .get(&request_uuid)
+            .copied()
+            .ok_or_else(|| anyhow!("unknown workload request output for {request_uuid}"))?;
+
+        let turn = &self.sessions[in_flight.session_index].turns[in_flight.turn_index];
+        let planned_output_tokens = turn
+            .output_token_ids
+            .as_ref()
+            .expect("delta turns must have planned output tokens");
+        let expected_token = planned_output_tokens
+            .get(in_flight.emitted_output_tokens)
+            .ok_or_else(|| {
+                anyhow!(
+                    "workload request {request_uuid} emitted more than {} planned output tokens",
+                    planned_output_tokens.len()
+                )
+            })?;
+        if token_id != *expected_token {
+            bail!(
+                "workload request {request_uuid} emitted token {token_id} at position {}, expected {}",
+                in_flight.emitted_output_tokens,
+                expected_token
+            );
+        }
+
+        let in_flight = self
+            .in_flight
+            .get_mut(&request_uuid)
+            .expect("validated in-flight request must still exist");
+        in_flight.emitted_output_tokens = in_flight
+            .emitted_output_tokens
+            .checked_add(1)
+            .context("workload emitted output token count overflow")?;
+        Ok(())
+    }
+
     pub fn on_complete(&mut self, request_uuid: Uuid, now_ms: f64) -> Result<()> {
+        self.on_terminal(request_uuid, now_ms, false)
+    }
+
+    pub fn on_terminal(&mut self, request_uuid: Uuid, now_ms: f64, rejected: bool) -> Result<()> {
+        let outcome = if rejected {
+            TurnOutcome::Rejected
+        } else {
+            TurnOutcome::Completed
+        };
         let resolution = self
-            .resolve_turn(request_uuid, now_ms, TurnOutcome::Completed)?
+            .resolve_turn(request_uuid, now_ms, outcome)?
             .expect("completed turns require an in-flight request");
         self.apply_resolution(resolution, now_ms);
         Ok(())
@@ -569,7 +647,7 @@ impl WorkloadDriver {
     ) -> Result<Option<TurnResolution>> {
         let Some(in_flight) = self.in_flight.get(&request_uuid).copied() else {
             return match outcome {
-                TurnOutcome::Completed => Err(anyhow!(
+                TurnOutcome::Completed | TurnOutcome::Rejected => Err(anyhow!(
                     "unknown workload request completion for {request_uuid}"
                 )),
                 TurnOutcome::Cancelled => Ok(None),
@@ -604,8 +682,23 @@ impl WorkloadDriver {
         }
 
         let request_id = turn.request_id.clone();
+        if outcome == TurnOutcome::Rejected && in_flight.emitted_output_tokens != 0 {
+            bail!(
+                "rejected workload request {request_uuid} emitted {} output tokens",
+                in_flight.emitted_output_tokens
+            );
+        }
+        let completed_output_tokens = (outcome == TurnOutcome::Completed
+            && self.prompt_mode == PromptMode::DeltaCumulative)
+            .then(|| {
+                let planned_output_tokens = turn
+                    .output_token_ids
+                    .as_ref()
+                    .expect("delta turns must have planned output tokens");
+                planned_output_tokens[..in_flight.emitted_output_tokens].to_vec()
+            });
         let (next_turn_index, next_ready_at_ms, session_ended) = match outcome {
-            TurnOutcome::Completed => {
+            TurnOutcome::Completed | TurnOutcome::Rejected => {
                 let next_turn_index = in_flight
                     .turn_index
                     .checked_add(1)
@@ -626,6 +719,11 @@ impl WorkloadDriver {
         session.in_flight = None;
         session.next_turn_index = next_turn_index;
         session.next_ready_at_ms = next_ready_at_ms;
+        if next_ready_at_ms.is_some()
+            && let Some(output_tokens) = completed_output_tokens
+        {
+            session.cumulative_tokens.extend(output_tokens);
+        }
         if let Some(ready_at_ms) = next_ready_at_ms {
             self.ready_sessions.push(ReadySession {
                 ready_at_ms,
@@ -1057,7 +1155,7 @@ mod tests {
     }
 
     #[test]
-    fn accumulating_delta_mode_emits_cumulative_input_tokens() {
+    fn accumulating_delta_mode_includes_previous_output_tokens() {
         let trace = Trace {
             block_size: 4,
             sessions: vec![SessionTrace {
@@ -1066,8 +1164,8 @@ mod tests {
                 turns: vec![
                     TurnTrace {
                         input_length: 6,
-                        max_output_tokens: 1,
-                        output_token_ids: None,
+                        max_output_tokens: 2,
+                        output_token_ids: Some(vec![20, 21]),
                         replay_key: None,
                         hash_ids: vec![10, 11],
                         delay_after_previous_ms: 0.0,
@@ -1094,20 +1192,145 @@ mod tests {
         let first = driver.pop_ready(0.0, usize::MAX);
         assert_eq!(first.len(), 1);
         assert_eq!(first[0].request.tokens, vec![10, 10, 10, 10, 11, 11]);
+        assert_eq!(first[0].request.output_token_ids, Some(vec![20, 21]));
         assert_eq!(first[0].request.priority, 3);
         assert_eq!(first[0].request.strict_priority, 4);
+        driver.on_output_token(first[0].request_uuid, 20).unwrap();
+        driver.on_output_token(first[0].request_uuid, 21).unwrap();
         driver.on_complete(first[0].request_uuid, 10.0).unwrap();
 
         let second = driver.pop_ready(15.0, usize::MAX);
         assert_eq!(second.len(), 1);
         assert_eq!(
             second[0].request.tokens,
-            vec![10, 10, 10, 10, 11, 11, 12, 12, 12]
+            vec![10, 10, 10, 10, 11, 11, 20, 21, 12, 12, 12]
         );
         assert_eq!(second[0].request.priority, -2);
         assert_eq!(second[0].request.strict_priority, 7);
     }
 
+    #[test]
+    fn accumulating_delta_mode_plans_missing_output_token_ids() {
+        let trace = Trace {
+            block_size: 1,
+            sessions: vec![SessionTrace {
+                session_id: "a".into(),
+                first_arrival_timestamp_ms: Some(0.0),
+                turns: vec![
+                    TurnTrace {
+                        input_length: 2,
+                        max_output_tokens: 3,
+                        hash_ids: vec![10, 11],
+                        ..Default::default()
+                    },
+                    TurnTrace {
+                        input_length: 1,
+                        max_output_tokens: 1,
+                        hash_ids: vec![12],
+                        ..Default::default()
+                    },
+                ],
+            }],
+        };
+        let mut driver = WorkloadDriver::new_concurrency_accumulating_deltas(trace, 1, 1).unwrap();
+
+        let first = driver.pop_ready(0.0, usize::MAX);
+        assert_eq!(first.len(), 1);
+        let planned_output = first[0]
+            .request
+            .output_token_ids
+            .clone()
+            .expect("delta replay should plan synthetic outputs");
+        assert_eq!(planned_output.len(), 3);
+        for &token_id in &planned_output {
+            driver
+                .on_output_token(first[0].request_uuid, token_id)
+                .unwrap();
+        }
+        driver.on_complete(first[0].request_uuid, 1.0).unwrap();
+
+        let second = driver.pop_ready(1.0, usize::MAX);
+        assert_eq!(second.len(), 1);
+        let mut expected = vec![10, 11];
+        expected.extend(planned_output);
+        expected.push(12);
+        assert_eq!(second[0].request.tokens, expected);
+        assert_eq!(
+            second[0].request.output_token_ids.as_ref().map(Vec::len),
+            Some(1)
+        );
+    }
+
+    #[test]
+    fn accumulating_delta_mode_appends_only_emitted_output_tokens() {
+        let trace = Trace {
+            block_size: 1,
+            sessions: vec![SessionTrace {
+                session_id: "a".into(),
+                first_arrival_timestamp_ms: Some(0.0),
+                turns: vec![
+                    TurnTrace {
+                        input_length: 1,
+                        max_output_tokens: 3,
+                        output_token_ids: Some(vec![20, 21, 22]),
+                        hash_ids: vec![10],
+                        ..Default::default()
+                    },
+                    TurnTrace {
+                        input_length: 1,
+                        max_output_tokens: 1,
+                        hash_ids: vec![12],
+                        ..Default::default()
+                    },
+                ],
+            }],
+        };
+        let mut driver = WorkloadDriver::new_concurrency_accumulating_deltas(trace, 1, 1).unwrap();
+
+        let first = driver.pop_ready(0.0, usize::MAX);
+        driver.on_output_token(first[0].request_uuid, 20).unwrap();
+        driver.on_output_token(first[0].request_uuid, 21).unwrap();
+        driver.on_complete(first[0].request_uuid, 1.0).unwrap();
+
+        let second = driver.pop_ready(1.0, usize::MAX);
+        assert_eq!(second[0].request.tokens, vec![10, 20, 21, 12]);
+    }
+
+    #[test]
+    fn accumulating_delta_mode_does_not_append_rejected_output_tokens() {
+        let trace = Trace {
+            block_size: 1,
+            sessions: vec![SessionTrace {
+                session_id: "a".into(),
+                first_arrival_timestamp_ms: Some(0.0),
+                turns: vec![
+                    TurnTrace {
+                        input_length: 1,
+                        max_output_tokens: 2,
+                        output_token_ids: Some(vec![20, 21]),
+                        hash_ids: vec![10],
+                        ..Default::default()
+                    },
+                    TurnTrace {
+                        input_length: 1,
+                        max_output_tokens: 1,
+                        hash_ids: vec![12],
+                        ..Default::default()
+                    },
+                ],
+            }],
+        };
+        let mut driver = WorkloadDriver::new_concurrency_accumulating_deltas(trace, 1, 1).unwrap();
+
+        let first = driver.pop_ready(0.0, usize::MAX);
+        driver
+            .on_terminal(first[0].request_uuid, 1.0, true)
+            .unwrap();
+
+        let second = driver.pop_ready(1.0, usize::MAX);
+        assert_eq!(second[0].request.tokens, vec![10, 12]);
+    }
+
     #[test]
     fn agentic_mode_releases_turn_after_dependency_completion_plus_delay() {
         let trace = AgenticTrace {

lib/mocker/src/loadgen/types.rs

@@ -54,11 +54,11 @@ pub enum DynamoRequestTrace {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum TraceFileFormat {
     Mooncake,
-    /// Mooncake-shaped rows where follow-up turns contain input deltas.
-    /// Offline replay accumulates those deltas per session in token space before
-    /// computing engine block hashes. Use this only for delta traces: it expands
-    /// compact session deltas into cumulative prompts and can use much more
-    /// memory than `Mooncake`.
+    /// Mooncake-shaped rows where follow-up turns contain new input deltas.
+    /// Offline replay accumulates each generated output and the next input delta
+    /// per session before computing engine block hashes. Use this only for delta
+    /// traces: it expands compact session turns into cumulative prompts and can
+    /// use much more memory than `Mooncake`.
     MooncakeDelta,
     /// Mooncake request/cache rows plus explicit request-level workflow
     /// dependencies. Each row dispatches after `wait_for` completions plus its