feat: model when vLLM's async engine admits, preempts and returns outputs

THIRD_PARTY_NOTICES.md

@@ -364,6 +364,33 @@ in virtual time, which committed batch an arriving request can join while the
 engine core keeps a queue of `max_concurrent_batches`; it does not reproduce
 asynchronous worker execution. That revision has no root NOTICE file.
 
+The async-scheduling block frees, preemption rules and engine-loop timing in
+repository-root files `crates/core/src/engine/scheduler/vllm/core.rs` and
+`crates/core/src/engine/scheduler/rank.rs` are an independently written,
+modified Rust simulation adaptation of behavioral contracts at vLLM tag
+`v0.30.0` (commit `ced6857afa0ea7b2e3f0846a62e1394e90f15607`):
+
+- https://github.com/vllm-project/vllm/blob/ced6857afa0ea7b2e3f0846a62e1394e90f15607/vllm/v1/core/sched/scheduler.py
+- https://github.com/vllm-project/vllm/blob/ced6857afa0ea7b2e3f0846a62e1394e90f15607/vllm/v1/engine/core.py
+- https://github.com/vllm-project/vllm/blob/ced6857afa0ea7b2e3f0846a62e1394e90f15607/vllm/config/parallel.py
+- https://github.com/vllm-project/vllm/blob/ced6857afa0ea7b2e3f0846a62e1394e90f15607/vllm/v1/executor/uniproc_executor.py
+
+and of the CUDA-graph dispatch at vLLM tag `v0.32.0` (commit
+`5df63a7203cdb9328d824a27b682e77c7b9928be`):
+
+- https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/vllm/v1/worker/gpu/cudagraph_utils.py
+
+Copyright contributors to the vLLM project. Licensed under Apache-2.0:
+https://github.com/vllm-project/vllm/blob/ced6857afa0ea7b2e3f0846a62e1394e90f15607/LICENSE
+https://github.com/vllm-project/vllm/blob/5df63a7203cdb9328d824a27b682e77c7b9928be/LICENSE
+
+No upstream method bodies or test fixtures are copied. The adaptation models,
+in virtual time, when a finished request's blocks return to the pool, which
+running request may be preempted while a pass is in flight, when the engine
+loop decides a pass and returns the previous pass's outputs, and which passes
+run without a CUDA graph; it does not reproduce asynchronous worker execution.
+Neither revision has a root NOTICE file.
+
 The following files are derived from vLLM's attention test utilities at tag
 `v0.11.0` (commit `b8b302cde434df8c9289a2b465406b47ebab1c2d`):
 

crates/core/README.md

@@ -61,8 +61,15 @@ The replay engine models vLLM's attention-DP prefill cadence through
 [configuration and validation notes](../../docs/replay/engine/workers.md#vllm-prefill-schedule-interval).
 
 `EngineConfig::async_scheduling` models vLLM's async-scheduling batch queue,
-which delays requests that arrive at a busy engine; see
+which delays requests that arrive at a busy engine. It also frees a finished
+request's blocks for the pass after next and, with a CPU tier, does not preempt
+a request that the pass in flight runs; see
 [vLLM async scheduling](../../docs/replay/engine/workers.md#vllm-async-scheduling).
+`EngineConfig::launch` adds the engine loop's launch costs. A rank then holds
+each pass's outputs until the driver releases them through
+`GeneralizedMockerEngine::next_emission_ms` and `take_due_emissions`, as
+offline replay does; see
+[vLLM pass launch](../../docs/replay/engine/workers.md#vllm-pass-launch).
 
 SGLang uses `EngineConfig::prefill_decode_interval` to block prefill for N
 scheduler rounds after a group-wide EXTEND. Its default is zero and its semantics

crates/core/src/engine/common/protocols.rs

@@ -200,6 +200,9 @@ pub(crate) struct MockEngineArgs {
     /// `EngineConfig::max_concurrent_batches`.
     #[builder(default = "1")]
     pub max_concurrent_batches: usize,
+    /// `EngineConfig::launch`.
+    #[builder(default = "None")]
+    pub launch: Option<crate::engine::PassLaunchConfig>,
     #[builder(default = "true")]
     pub enable_prefix_caching: bool,
     #[builder(default = "true")]

crates/core/src/engine/config.rs

@@ -673,6 +673,53 @@ impl NativeHostOffloadConfig {
     }
 }
 
+/// CPU launch costs of vLLM's async engine loop (`scheduler.launch`).
+///
+/// vLLM decides and launches pass k+1 before it returns pass k's outputs, so a
+/// pass's outputs leave once the next pass is launched, and a request joins the
+/// first pass decided after it arrives.
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct PassLaunchConfig {
+    /// From deciding a pass without prefill and with at most
+    /// `max_cudagraph_capture_size` scheduled tokens to the end of its launch,
+    /// in ms.
+    pub decode_ms: f64,
+    /// The same for such a pass with prefill (piecewise CUDA graphs).
+    pub piecewise_ms: f64,
+    /// The same for a larger pass, which runs eager with or without prefill.
+    pub eager_ms: f64,
+    /// vLLM's `max_cudagraph_capture_size`.
+    pub max_cudagraph_capture_size: usize,
+    /// The part of an eager pass the GPU has run before its launch can end:
+    /// the CPU waits once the CUDA launch queue is full.
+    pub eager_queue_fraction: f64,
+}
+
+impl PassLaunchConfig {
+    fn validate(&self) -> Result<()> {
+        for (field, ms) in [
+            ("decode_ms", self.decode_ms),
+            ("piecewise_ms", self.piecewise_ms),
+            ("eager_ms", self.eager_ms),
+        ] {
+            ensure!(
+                ms.is_finite() && ms >= 0.0,
+                "launch.{field} must be finite and non-negative"
+            );
+        }
+        ensure!(
+            self.max_cudagraph_capture_size > 0,
+            "launch.max_cudagraph_capture_size must be positive"
+        );
+        ensure!(
+            (0.0..1.0).contains(&self.eager_queue_fraction),
+            "launch.eager_queue_fraction must lie in [0, 1)"
+        );
+        Ok(())
+    }
+}
+
 /// Recurrent-state allocation size for one simulated rank/GPU.
 /// Token block geometry and pool capacity use the existing EngineConfig fields.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -756,6 +803,14 @@ pub struct EngineConfig {
     /// misses. `None` selects vLLM's default: on, except with ngram drafts.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub async_scheduling: Option<bool>,
+    /// vLLM async scheduling's CPU launch costs: a pass's outputs leave once
+    /// the next pass is launched, and an arrival joins the first pass decided
+    /// after it. `None` keeps outputs at the pass end and the pass-count rule.
+    /// A rank holds the outputs until the driver releases them through
+    /// `GeneralizedMockerEngine::take_due_emissions`, as offline replay does;
+    /// other drivers must leave this unset.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub launch: Option<PassLaunchConfig>,
     /// Whether complete blocks remain reusable after request release.
     #[serde(default = "default_true")]
     pub enable_prefix_caching: bool,
@@ -864,6 +919,8 @@ struct EngineConfigWire {
     prefill_decode_interval: usize,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     async_scheduling: Option<bool>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    launch: Option<PassLaunchConfig>,
     #[serde(default = "default_true")]
     enable_prefix_caching: bool,
     #[serde(default = "default_true")]
@@ -949,6 +1006,7 @@ impl<'de> Deserialize<'de> for EngineConfig {
             prefill_schedule_interval: wire.prefill_schedule_interval,
             prefill_decode_interval: wire.prefill_decode_interval,
             async_scheduling: wire.async_scheduling,
+            launch: wire.launch,
             enable_prefix_caching: wire.enable_prefix_caching,
             enable_chunked_prefill: wire.enable_chunked_prefill,
             speedup_ratio: wire.speedup_ratio,
@@ -996,6 +1054,7 @@ impl Default for EngineConfig {
             prefill_schedule_interval: default_prefill_schedule_interval(),
             prefill_decode_interval: 0,
             async_scheduling: None,
+            launch: None,
             enable_prefix_caching: true,
             enable_chunked_prefill: true,
             speedup_ratio: 1.0,
@@ -1132,6 +1191,17 @@ impl EngineConfig {
             self.backend == Backend::Vllm || self.async_scheduling.is_none(),
             "async_scheduling is supported only for backend=vllm"
         );
+        if let Some(launch) = &self.launch {
+            launch.validate()?;
+            ensure!(
+                self.max_concurrent_batches() > 1,
+                "launch requires vLLM async_scheduling"
+            );
+            ensure!(
+                self.worker_type == WorkerType::Aggregated,
+                "launch is supported only on aggregated workers"
+            );
+        }
         ensure!(
             self.backend == Backend::Sglang || !self.sglang.host_loop,
             "sglang.host_loop is supported only for backend=sglang"
@@ -2261,6 +2331,50 @@ mod tests {
         }
     }
 
+    #[test]
+    fn launch_needs_async_scheduling_on_an_aggregated_vllm_worker() {
+        let launch = PassLaunchConfig {
+            decode_ms: 5.0,
+            piecewise_ms: 50.0,
+            eager_ms: 150.0,
+            max_cudagraph_capture_size: 512,
+            eager_queue_fraction: 0.4,
+        };
+        let config = EngineConfig {
+            launch: Some(launch),
+            ..EngineConfig::default()
+        };
+        config.validate().unwrap();
+        let decoded: EngineConfig =
+            serde_json::from_str(&serde_json::to_string(&config).unwrap()).unwrap();
+        assert_eq!(decoded.launch, Some(launch));
+        for (edit, message) in [
+            (
+                (|config: &mut EngineConfig| config.async_scheduling = Some(false))
+                    as fn(&mut EngineConfig),
+                "launch requires vLLM async_scheduling",
+            ),
+            (
+                |config| config.worker_type = WorkerType::Decode,
+                "launch is supported only on aggregated workers",
+            ),
+            (
+                |config| config.launch.as_mut().unwrap().eager_queue_fraction = 1.0,
+                "launch.eager_queue_fraction must lie in [0, 1)",
+            ),
+        ] {
+            let mut invalid = config.clone();
+            edit(&mut invalid);
+            assert!(
+                invalid
+                    .validate()
+                    .unwrap_err()
+                    .to_string()
+                    .contains(message)
+            );
+        }
+    }
+
     #[test]
     fn deserialization_rejects_negative_prefill_decode_interval() {
         let error = serde_json::from_value::<EngineConfig>(serde_json::json!({

crates/core/src/engine/generalized/contracts.rs

@@ -413,4 +413,18 @@ pub trait RankEngine: Sized {
 
     /// Whether this rank owns no request, pass, or internal work.
     fn is_drained(&self) -> bool;
+
+    /// Earliest time at which outputs held past their pass's completion leave
+    /// the rank. Unlike internal deadlines it is not masked by a pass in flight:
+    /// vLLM returns a pass's outputs while the next pass runs. The default
+    /// holds none.
+    fn next_emission_ms(&self) -> Option<f64> {
+        None
+    }
+
+    /// Release the held outputs due at `now_ms` as completion effects that
+    /// carry only outputs.
+    fn take_due_emissions(&mut self, _now_ms: f64) -> Result<Option<Self::PassCompletionEffects>> {
+        Ok(None)
+    }
 }

crates/core/src/engine/generalized/engine.rs

@@ -411,6 +411,41 @@ impl<C: RankEngine> GeneralizedMockerEngine<C> {
         })
     }
 
+    /// Earliest held-output release across ranks (see
+    /// [`RankEngine::next_emission_ms`]).
+    pub fn next_emission_ms(&self) -> Option<f64> {
+        if self.poisoned.is_some() {
+            return None;
+        }
+        self.ranks
+            .iter()
+            .filter_map(RankEngine::next_emission_ms)
+            .min_by(f64::total_cmp)
+    }
+
+    /// Release every rank's held outputs due at `now_ms`, in rank order, even
+    /// while a pass is in flight.
+    pub fn take_due_emissions(
+        &mut self,
+        now_ms: f64,
+    ) -> Result<EngineEffects<C::PassCompletionEffects>> {
+        self.ensure_healthy()?;
+        validate_time(now_ms, "emission time")?;
+        let mut effects = Vec::new();
+        for dp_rank in 0..self.ranks.len() {
+            let due = self.ranks[dp_rank]
+                .take_due_emissions(now_ms)
+                .with_context(|| format!("releasing outputs of attention-DP rank {dp_rank}"));
+            if let Some(due) = due.map_err(|error| self.poison(error))? {
+                effects.push(RankEffects {
+                    dp_rank: dp_rank as u32,
+                    effects: due,
+                });
+            }
+        }
+        Ok(EngineEffects { by_rank: effects })
+    }
+
     /// Earliest valid internal-work deadline across all ranks.
     pub fn next_internal_deadline_ms(&self) -> Option<f64> {
         if self.poisoned.is_some() || self.pending_pass.is_some() {

crates/core/src/engine/mod.rs

@@ -37,7 +37,7 @@ pub use common::running_mean::RunningMean;
 pub use common::speculative::normalize_conditional_accept_rates;
 pub use config::{
     Backend, EngineConfig, FrontendConfig, FrontendStage, G2Connector, G2Scope, G3OffloadConfig,
-    G3Scope, KvcrConfig, NativeHostOffloadConfig, PreemptionMode, SglangConfig,
+    G3Scope, KvcrConfig, NativeHostOffloadConfig, PassLaunchConfig, PreemptionMode, SglangConfig,
     SglangSchedulePolicy, StateCacheConfig, TrtllmCapacityPolicy, TrtllmConfig, WorkerType,
 };
 pub use g3_offload::{G3IoStats, G3Stats};

crates/core/src/engine/runtime.rs

@@ -87,6 +87,12 @@ impl EngineFactory {
         self.config.native_host_offload.is_some() || self.config.frontend.is_some()
     }
 
+    /// Whether ranks built by this factory can hold outputs past their pass
+    /// (vLLM `launch`).
+    pub(crate) fn can_hold_outputs(&self) -> bool {
+        self.config.launch.is_some()
+    }
+
     /// Build one scheduler/KV/timing rank with an explicit identity.
     pub fn build_rank(&self, identity: RankIdentity) -> Result<SchedulerRank> {
         let seed_offset = engine_seed_offset(identity)?;

crates/core/src/engine/scheduler/mod.rs

@@ -292,6 +292,10 @@ pub(crate) struct EnginePassResult {
     /// Forward pass metrics snapshot for this iteration.
     pub(crate) fpm: Option<ForwardPassSnapshot>,
     pub(crate) decode_acceptance: DecodeAcceptance,
+    /// vLLM `launch`: when the engine loop returns the previous pass's
+    /// outputs, after launching this pass. `None` when the previous pass's
+    /// outputs left at its end.
+    pub(crate) previous_returned_ms: Option<f64>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]