feat(mocker): model prefill→decode KV stranding + abort for disagg cascade tests

components/src/dynamo/mocker/args.py

@@ -544,6 +544,19 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         "Default: 64.0 (inter-node InfiniBand). Set to 0 to disable KV transfer delay. "
         "For intra-node NVLink, typical value is ~450.",
     )
+    parser.add_argument(
+        "--kv-transfer-abort-timeout-ms",
+        type=int,
+        default=None,
+        help="Prefill-to-decode handshake timeout in milliseconds. When set, "
+        "the prefill worker holds KV cache after compute until decode connects to its "
+        "bootstrap server, up to this timeout. If decode does not arrive in time, the "
+        "prefill aborts the request (KV released, abort error surfaced); late-arriving "
+        "decodes for the same room get a clean ABORT response. Mirrors production "
+        "VLLM_NIXL_ABORT_REQUEST_TIMEOUT. Decode-side: must wait for local KV "
+        "availability before connecting to prefill. Default: None (legacy behavior, "
+        "prefill ACKs immediately).",
+    )
     parser.add_argument(
         "--kv-cache-dtype",
         type=str,

components/src/dynamo/mocker/config.py

@@ -306,6 +306,9 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs:
         bandwidth_g3_to_g2_gbps=getattr(args, "bandwidth_g3_to_g2_gbps", None),
         bandwidth_g2_to_g4_gbps=getattr(args, "bandwidth_g2_to_g4_gbps", None),
         bandwidth_g4_to_g2_gbps=getattr(args, "bandwidth_g4_to_g2_gbps", None),
+        kv_transfer_abort_timeout_ms=getattr(
+            args, "kv_transfer_abort_timeout_ms", None
+        ),
         reasoning=_parse_reasoning_config(getattr(args, "reasoning", None)),
         sglang=_build_sglang_args(args),
         trtllm=_build_trtllm_args(args),

lib/bindings/python/rust/llm/replay.rs

@@ -171,7 +171,7 @@ impl MockEngineArgs {
 #[pymethods]
 impl MockEngineArgs {
     #[new]
-    #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None))]
+    #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, kv_transfer_abort_timeout_ms=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None))]
     #[allow(clippy::too_many_arguments)]
     fn new(
         engine_type: &str,
@@ -205,6 +205,7 @@ impl MockEngineArgs {
         bootstrap_port: Option<u16>,
         kv_bytes_per_token: Option<usize>,
         kv_transfer_bandwidth: Option<f64>,
+        kv_transfer_abort_timeout_ms: Option<u64>,
         reasoning: Option<ReasoningConfig>,
         zmq_kv_events_port: Option<u16>,
         zmq_replay_port: Option<u16>,
@@ -275,6 +276,7 @@ impl MockEngineArgs {
             .bandwidth_g3_to_g2_gbps(bandwidth_g3_to_g2_gbps)
             .bandwidth_g2_to_g4_gbps(bandwidth_g2_to_g4_gbps)
             .bandwidth_g4_to_g2_gbps(bandwidth_g4_to_g2_gbps)
+            .kv_transfer_abort_timeout_ms(kv_transfer_abort_timeout_ms)
             .reasoning(reasoning.map(|config| config.inner()))
             .zmq_kv_events_port(zmq_kv_events_port)
             .zmq_replay_port(zmq_replay_port)

lib/bindings/python/src/dynamo/_core.pyi

@@ -1893,6 +1893,7 @@ class MockEngineArgs:
         bootstrap_port: Optional[int] = None,
         kv_bytes_per_token: Optional[int] = None,
         kv_transfer_bandwidth: Optional[float] = None,
+        kv_transfer_abort_timeout_ms: Optional[int] = None,
         reasoning: Optional[ReasoningConfig] = None,
         zmq_kv_events_port: Optional[int] = None,
         zmq_replay_port: Optional[int] = None,

lib/llm/src/mocker.rs

@@ -258,6 +258,67 @@ impl MockEngine {
         });
     }
 
+    /// Wait until the scheduler at `dp_rank` reports both **block** and **sequence**
+    /// headroom for admitting a new request, up to `timeout`. Used by the decode side
+    /// of disagg before connecting to a prefill bootstrap server.
+    ///
+    /// Models real NIXL admission behavior in vLLM v1 and sglang — both engines gate
+    /// the disaggregated KV recv on **two independent budgets**: a block/token budget
+    /// (enough KV-cache blocks to absorb the projected sequence) and a sequence-slot
+    /// budget (a free slot in the per-request scheduling state — vLLM `max_num_seqs` /
+    /// `len(self.running)`; sglang's `DecodeReqToTokenPool`). Both must have headroom;
+    /// either alone blocks the NIXL recv from being armed. Without the seq-slot check
+    /// the prefill-side abort path is unreachable from realistic seq-bound load shapes.
+    ///
+    /// Returns Ok(()) immediately if schedulers haven't reported metrics yet
+    /// (total_blocks == 0) so warmup is graceful. Also treats `max_num_seqs == 0`
+    /// as "no seq cap configured" (mirrors `MockEngineArgs.max_num_seqs == None`).
+    ///
+    /// **No timeout on this side.** Real vLLM and sglang decode workers do
+    /// not run an independent "give up waiting for capacity" timer — a
+    /// WAITING_FOR_REMOTE_KVS request stays WAITING until the scheduler admits it, the
+    /// request context is dropped, or the prefill side's `kv_transfer_abort_timeout_ms`
+    /// fires and closes the bootstrap room (which surfaces to the decode side later when
+    /// it attempts `connect_to_prefill(room_id)` — that connect then fails with a
+    /// closed-room error). Cancellation is delegated to the surrounding request context.
+    pub async fn wait_for_decode_kv_capacity(&self, dp_rank: u32) -> Result<()> {
+        let schedulers = self
+            ._schedulers
+            .get()
+            .ok_or_else(|| anyhow::anyhow!("schedulers not initialized"))?;
+        let scheduler_idx = dp_rank as usize;
+        if scheduler_idx >= schedulers.len() {
+            return Err(anyhow::anyhow!(
+                "dp_rank {dp_rank} out of bounds (have {} schedulers)",
+                schedulers.len()
+            ));
+        }
+        let mut rx = schedulers[scheduler_idx].metrics_receiver();
+
+        loop {
+            let metrics = rx.borrow().clone();
+            // total_blocks == 0 means the scheduler hasn't published yet (still warming up).
+            // Don't block on warmup — the scheduler's own queue will handle admission
+            // once metrics start flowing.
+            let has_block_capacity =
+                metrics.total_blocks == 0 || metrics.active_decode_blocks < metrics.total_blocks;
+            // max_num_seqs == 0 means "no seq cap configured" — mirror vLLM/sglang's
+            // treatment of unset max_running_requests as effectively unbounded.
+            // running_requests is the live sequence count (vLLM len(self.running) /
+            // sglang req_to_token_pool occupancy).
+            let has_seq_capacity =
+                metrics.max_num_seqs == 0 || metrics.running_requests < metrics.max_num_seqs;
+            if has_block_capacity && has_seq_capacity {
+                return Ok(());
+            }
+            // Wait indefinitely for the next metrics update — no decode-side timer.
+            // Cancellation lands via the request context being dropped.
+            rx.changed()
+                .await
+                .map_err(|_| anyhow::anyhow!("scheduler metrics channel closed"))?;
+        }
+    }
+
     /// Send a request to the appropriate scheduler, waiting for initialization if needed.
     pub async fn direct(&self, request: DirectRequest, dp_rank: usize) {
         let sender = self.request_sender(dp_rank).await;
@@ -523,12 +584,37 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
             .await;
 
         // Bootstrap rendezvous for disaggregated serving
-        // - Decode: send receiver metadata to prefill, then wait for prefill completion
-        // - Prefill: wait for decode metadata before emitting output, then complete_room()
+        // - Decode: wait for own KV capacity, then send receiver metadata to prefill
+        //   and wait for prefill completion or abort
+        // - Prefill: wait for decode metadata (bounded by abort_timeout when set) before emitting
+        //   output, then complete_room(); on abort-timeout abort_room()
         let bootstrap_room = request.bootstrap_info.as_ref().map(|b| b.bootstrap_room);
+        let abort_timeout = self
+            .engine_args
+            .kv_transfer_abort_timeout_ms
+            .map(Duration::from_millis);
+
         if let Some(bootstrap_info) = &request.bootstrap_info
             && self.engine_args.is_decode()
         {
+            // Gate decode on local KV capacity before connecting (see
+            // wait_for_decode_kv_capacity). Only when abort_timeout is configured, to
+            // preserve the legacy "skip the wait" path for older DGDs.
+            if abort_timeout.is_some() {
+                self.wait_for_decode_kv_capacity(dp_rank)
+                    .await
+                    .map_err(|e| Error::msg(format!("Decode KV wait failed: {e}")))?;
+            }
+            // Forensic logging: emit decode-side wait-start event so post-hoc
+            // analysis can reconstruct (decode_worker, target_prefill, room) stranding graphs.
+            tracing::info!(
+                target: "mocker::kv_abort",
+                decode_dp_rank = dp_rank,
+                room_id = bootstrap_info.bootstrap_room,
+                target_host = %bootstrap_info.bootstrap_host,
+                target_port = bootstrap_info.bootstrap_port,
+                "decode_kv_wait_start"
+            );
             connect_to_prefill(
                 &bootstrap_info.bootstrap_host,
                 bootstrap_info.bootstrap_port,
@@ -579,15 +665,49 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
         // Spawn a task to handle the complex async logic
         tokio::spawn(async move {
             if let Some((server, room_id, sender, direct_request)) = delayed_prefill_submission {
+                // Forensic logging: pin-start event with room_id + prefill dp_rank.
+                // Lets post-hoc analysis join with decode logs on room_id to build the
+                // (prefill -> decode) stranding graph. While this prefill waits for decode to
+                // arrive, its scheduler request stays active so KV stays pinned.
+                let pin_start = std::time::Instant::now();
+                tracing::info!(
+                    target: "mocker::kv_abort",
+                    prefill_dp_rank = dp_rank,
+                    room_id,
+                    "prefill_kv_pin_start"
+                );
                 tokio::select! {
-                    result = server.wait_for_decode_ready(room_id) => {
+                    // Bound the decode-arrival wait by abort_timeout (when set). On
+                    // timeout, abort the room so waiting/late decodes get a clean ABORT instead
+                    // of hanging, surface the abort to the client, and end the stream.
+                    result = server.wait_for_decode_ready(room_id, abort_timeout) => {
                         if let Err(e) = result {
+                            tracing::warn!(
+                                "Prefill aborting transfer for room {room_id}: {e}"
+                            );
+                            server.abort_room(room_id);
+                            tracing::info!(
+                                target: "mocker::kv_abort",
+                                prefill_dp_rank = dp_rank,
+                                room_id,
+                                outcome = "aborted",
+                                duration_ms = pin_start.elapsed().as_millis() as u64,
+                                "prefill_kv_pin_end"
+                            );
                             let _ = stream_tx.send(LLMEngineOutput::error(format!(
-                                "Bootstrap wait for decode metadata failed: {e}"
+                                "NIXL transfer aborted: {e}"
                             )));
                             active_requests.remove(&request_uuid);
                             return;
                         }
+                        tracing::info!(
+                            target: "mocker::kv_abort",
+                            prefill_dp_rank = dp_rank,
+                            room_id,
+                            outcome = "completed",
+                            duration_ms = pin_start.elapsed().as_millis() as u64,
+                            "prefill_kv_pin_end"
+                        );
                     }
                     _ = async_context.stopped() => {
                         let _ = stream_tx.send(LLMEngineOutput::cancelled());
@@ -653,6 +773,9 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
                         }
 
                         if signal.completed {
+                            // By this point a disagg prefill has already waited for
+                            // decode to arrive (or aborted) in wait_for_decode_ready below, so
+                            // KV stayed pinned during the wait — modeling real NIXL behavior.
                             if stream_tx.send(output).is_err() {
                                 tracing::error!("Output stream receiver closed.");
                                 break;

lib/mocker/src/common/protocols.rs

@@ -518,6 +518,7 @@ struct MockEngineArgsSerde {
     bootstrap_port: OptionalConfigValue<u16>,
     kv_bytes_per_token: OptionalConfigValue<usize>,
     kv_transfer_bandwidth: OptionalConfigValue<f64>,
+    kv_transfer_abort_timeout_ms: OptionalConfigValue<u64>,
     num_g2_blocks: OptionalConfigValue<usize>,
     num_g3_blocks: OptionalConfigValue<usize>,
     enable_g4_storage: OptionalConfigValue<bool>,
@@ -729,6 +730,17 @@ pub struct MockEngineArgs {
     #[validate(range(min = 0.0))]
     pub kv_transfer_bandwidth: Option<f64>,
 
+    /// Timeout (milliseconds) for the prefill→decode NIXL handshake. Mirrors the
+    /// production `VLLM_NIXL_ABORT_REQUEST_TIMEOUT` env var. When a disagg prefill
+    /// completes compute but no decode arrives within this window, the prefill aborts
+    /// the request (KV released, abort error surfaced to the client); late-arriving
+    /// decodes for the same room receive a clean ABORT response.
+    ///
+    /// `None` (default) disables the prefill-side wait — preserves the legacy
+    /// behavior where prefill ACKs immediately and never holds KV waiting for decode.
+    #[builder(default = "None")]
+    pub kv_transfer_abort_timeout_ms: Option<u64>,
+
     /// KVBM G2 (host DRAM) block capacity. When the `kvbm-offload`
     /// feature is enabled, setting this explicitly opts the mocker into
     /// G2 offload simulation. When unset or set to 0, no G2 offload engine
@@ -1061,6 +1073,11 @@ impl TryFrom<MockEngineArgsSerde> for MockEngineArgs {
         if let Some(kv_transfer_bandwidth) = compat.kv_transfer_bandwidth.into_nullable() {
             builder = builder.kv_transfer_bandwidth(kv_transfer_bandwidth);
         }
+        if let Some(kv_transfer_abort_timeout_ms) =
+            compat.kv_transfer_abort_timeout_ms.into_nullable()
+        {
+            builder = builder.kv_transfer_abort_timeout_ms(kv_transfer_abort_timeout_ms);
+        }
         if let Some(num_g2_blocks) = compat.num_g2_blocks.into_nullable() {
             builder = builder.num_g2_blocks(num_g2_blocks);
         }
@@ -1341,6 +1358,7 @@ mod tests {
             "bandwidth_g3_to_g2_gbps": args.bandwidth_g3_to_g2_gbps,
             "bandwidth_g2_to_g4_gbps": args.bandwidth_g2_to_g4_gbps,
             "bandwidth_g4_to_g2_gbps": args.bandwidth_g4_to_g2_gbps,
+            "kv_transfer_abort_timeout_ms": args.kv_transfer_abort_timeout_ms,
             "reasoning": args.reasoning,
             "zmq_kv_events_port": args.zmq_kv_events_port,
             "zmq_replay_port": args.zmq_replay_port,

lib/mocker/src/scheduler/sglang/config.rs

@@ -40,6 +40,9 @@ pub(super) struct SglangConfig {
     pub(super) kv_bytes_per_token: Option<usize>,
     pub(super) kv_transfer_bandwidth: Option<f64>,
     pub(super) speculative_max_tokens: Option<usize>,
+    /// Sequence-slot cap (MockEngineArgs.max_num_seqs). 0 == unlimited. Used by the
+    /// decode-side admission wait to model the seq-slot budget.
+    pub(super) max_num_seqs: u64,
 }
 
 impl SglangConfig {
@@ -89,6 +92,7 @@ impl SglangConfig {
             kv_bytes_per_token: args.kv_bytes_per_token,
             kv_transfer_bandwidth: args.kv_transfer_bandwidth,
             speculative_max_tokens: args.aic_nextn.map(|nextn| nextn + 1),
+            max_num_seqs: args.max_num_seqs.map(|v| v as u64).unwrap_or(0),
         }
     }
 }

lib/mocker/src/scheduler/sglang/core.rs

@@ -249,16 +249,21 @@ impl SglangCore {
                 .count(),
             output_signals: decode.output_signals,
             admissions: admit.admissions,
-            mocker_metrics: MockerMetrics::from_parts(
-                self.dp_rank,
-                active_decode_blocks,
-                self.config.total_kv_tokens.div_ceil(self.config.block_size) as u64,
-                self.running.len() as u64,
-                self.waiting.len() as u64,
-                0,
-                sglang_cache_hit_tokens,
-                sglang_cache_total_tokens,
-            ),
+            mocker_metrics: {
+                let mut metrics = MockerMetrics::from_parts(
+                    self.dp_rank,
+                    active_decode_blocks,
+                    self.config.total_kv_tokens.div_ceil(self.config.block_size) as u64,
+                    self.running.len() as u64,
+                    self.waiting.len() as u64,
+                    0,
+                    sglang_cache_hit_tokens,
+                    sglang_cache_total_tokens,
+                );
+                // Expose the seq-slot cap for the decode-side admission wait.
+                metrics.max_num_seqs = self.config.max_num_seqs;
+                metrics
+            },
             router_event_visibility: RouterEventVisibility::PassEnd,
             kv_events: self
                 .kv_event_buffer