feat(mocker): model prefill→decode KV stranding + abort for disagg cascade tests - #10557
feat(mocker): model prefill→decode KV stranding + abort for disagg cascade tests#10557nnshah1 wants to merge 15 commits into
Conversation
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
WalkthroughThis PR implements DIS-2147: an optional abort timeout for disaggregated prefill-to-decode KV cache transfers. A new ChangesDIS-2147 Disaggregated Serving Abort Timeout
🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
lib/mocker/src/scheduler/vllm/core.rs (1)
510-525:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReplace Linear ticket reference with GitHub issue number.
Comment references "DIS-2147". Replace with the corresponding GitHub issue number (e.g.,
GH-NNNNor#NNNN) per coding guidelines.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/mocker/src/scheduler/vllm/core.rs` around lines 510 - 525, The inline comment in function mocker_metrics() references a Linear ticket "DIS-2147"; update that comment to use the project's GitHub issue format (e.g., "GH-<number>" or "#<number>") per guidelines by replacing "DIS-2147" with the correct GitHub issue identifier in the comment above metrics.max_num_seqs so the comment reads e.g. "// GH-XXXX: expose the seq-slot cap..." instead of the Linear ticket; ensure only the comment text is changed and code/logic (mocker_metrics, metrics.max_num_seqs) remains untouched.Source: Coding guidelines
lib/mocker/src/services/bootstrap.rs (1)
190-216:⚠️ Potential issue | 🟠 Major | ⚡ Quick winThe decode side still times out independently after 30s.
Both the server-side wait and the client-side response read are still capped by
RENDEZVOUS_TIMEOUT. Ifkv_transfer_abort_timeout_msis larger than 30s, or prefill legitimately finishes later than that, decode fails early andhandle_connection()removes the room before prefill can complete/abort.⏱️ Suggested direction
- ImmediateOrWait::Wait(rx) => match tokio::time::timeout(RENDEZVOUS_TIMEOUT, rx).await { - Ok(Ok(RoomOutcome::Completed)) => { + ImmediateOrWait::Wait(rx) => match rx.await { + Ok(RoomOutcome::Completed) => { tracing::debug!("Bootstrap: room {room_id} prefill completed, sending ACK"); ACK_BYTE } - Ok(Ok(RoomOutcome::Aborted)) => { + Ok(RoomOutcome::Aborted) => { tracing::warn!( "Bootstrap: room {room_id} prefill aborted while decode waited, \ sending ABORT" ); ABORT_BYTE } - Ok(Ok(RoomOutcome::Pending)) => { + Ok(RoomOutcome::Pending) => { bail!("Bootstrap: room {room_id} sender fired with Pending outcome"); } - Ok(Err(_)) => { + Err(_) => { bail!("Bootstrap: room {room_id} sender dropped"); } - Err(_) => { - rooms.remove(&room_id); - bail!("Bootstrap: room {room_id} timeout waiting for prefill"); - } }, @@ - tokio::time::timeout(RENDEZVOUS_TIMEOUT, stream.read_exact(&mut response)) - .await - .map_err(|_| anyhow::anyhow!("Bootstrap: response timeout for room {room_id}"))? - .map_err(|e| anyhow::anyhow!("Bootstrap: read response failed: {e}"))?; + stream + .read_exact(&mut response) + .await + .map_err(|e| anyhow::anyhow!("Bootstrap: read response failed: {e}"))?;Based on supplied downstream context, prefill-side
abort_timeoutis intended to be authoritative, so these fixed 30s decode-side waits can now fail first.Also applies to: 365-370
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/mocker/src/services/bootstrap.rs` around lines 190 - 216, The decode-side wait uses the global RENDEZVOUS_TIMEOUT, causing premature timeouts; change the tokio::time::timeout calls (e.g. where ImmediateOrWait::Wait(rx) matches) to use the room-specific abort timeout (the prefill/kv_transfer_abort_timeout_ms value) instead of RENDEZVOUS_TIMEOUT, obtaining that timeout from the Room or prefill params available in bootstrap.rs/handle_connection() and applying the same change to the other occurrences around lines 365-370; ensure you pass the correctly converted Duration to tokio::time::timeout and only remove the room from rooms.remove(&room_id) if the room-specific timeout elapses.lib/bindings/python/rust/llm/replay.rs (1)
174-209:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate the Python typing stub for the new constructor keyword.
This runtime signature now accepts
kv_transfer_abort_timeout_ms, but the supplied_core.pyistub still stops atkv_transfer_bandwidth. Typed callers will get a bogus unexpected-keyword error until the stub is updated too.🧩 Follow-up stub change
--- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ 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,Based on supplied typing stub context,
lib/bindings/python/src/dynamo/_core.pyihas not been updated alongside this runtime signature.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/bindings/python/rust/llm/replay.rs` around lines 174 - 209, The Python typing stub for the constructor is out of sync: the Rust/pyO3 constructor function new now accepts the kv_transfer_abort_timeout_ms keyword but the typed stub (_core.pyi) still ends at kv_transfer_bandwidth; update the stub to add kv_transfer_abort_timeout_ms: Optional[int] (or appropriate integer type) to the constructor signature so typed callers no longer get unexpected-keyword errors, ensuring the parameter name and type match the runtime signature of new and any related overloads in the same stub.
🧹 Nitpick comments (2)
lib/llm/src/mocker.rs (1)
684-686: ⚡ Quick winUse lazy formatting in tracing::warn!.
The warning message uses eager string interpolation. Prefer lazy formatting with structured fields or positional arguments to avoid unnecessary allocations when the log level is disabled.
Suggested refactor
- tracing::warn!( - "Prefill aborting transfer for room {room_id}: {e}" - ); + tracing::warn!( + room_id = room_id, + error = %e, + "Prefill aborting transfer" + );🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/llm/src/mocker.rs` around lines 684 - 686, The tracing::warn! call is doing eager string interpolation for room_id and e; replace it with lazy formatting by passing room_id and e as structured fields or positional arguments to tracing::warn! instead of embedding them in the format string—locate the tracing::warn! invocation that references room_id and e in mocker.rs (the prefill aborting transfer log) and change it to use structured fields like room_id = %room_id and error = %e or use positional placeholders so the values are only formatted when the warn level is enabled.lib/mocker/src/common/protocols.rs (1)
1251-1260: ⚡ Quick winMake the round-trip test use a concrete timeout value.
This still serializes
kv_transfer_abort_timeout_msasNone, so the test won't fail if the new field stops round-tripping. Set a non-default value and assert it after restore.🧪 Suggested test tightening
let args = MockEngineArgs::builder() .worker_type(WorkerType::Decode) + .kv_transfer_abort_timeout_ms(Some(123)) .max_num_seqs(None) .max_num_batched_tokens(None) .reasoning(None) .sglang(None) .build() @@ let restored = MockEngineArgs::from_json_str(&payload.to_string()).unwrap(); assert_eq!(restored.worker_type, WorkerType::Decode); + assert_eq!(restored.kv_transfer_abort_timeout_ms, Some(123)); assert_eq!(restored.max_num_seqs, None); assert_eq!(restored.max_num_batched_tokens, None);Also applies to: 1295-1307
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/mocker/src/common/protocols.rs` around lines 1251 - 1260, The round-trip test builds MockEngineArgs without setting kv_transfer_abort_timeout_ms so it remains None and won't catch serialization regressions; update the MockEngineArgs::builder() call used in the round-trip test (the builder chain that sets worker_type, max_num_seqs, max_num_batched_tokens, reasoning, sglang, etc.) to explicitly set kv_transfer_abort_timeout_ms to a concrete non-default value (e.g., Some(1234)), then after restore/assertion verify the restored args.kv_transfer_abort_timeout_ms equals that value; apply the same change to the other builder usage in the file that mirrors this test.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@components/src/dynamo/mocker/args.py`:
- Around line 551-558: The help text string in
components/src/dynamo/mocker/args.py contains an internal ticket ID "DIS-2147";
update the help for the prefill-to-decode handshake timeout (the argument help
string around the prefill worker timeout) to remove "DIS-2147" and replace it
with either a public reference (e.g., "GH-<number>" or "#<number>") or plain
prose like "see public issue tracker" or simply omit the ticket reference,
preserving the existing explanatory text about behavior, defaults, and
decode-side requirements; edit the help string used when defining the timeout
argument so no internal Linear ID remains.
In `@lib/llm/src/mocker.rs`:
- Around line 600-607: The comment references an internal Linear ticket ID
"DIS-2147"; update that comment to reference the public GitHub issue format
(e.g., "GH-XXXX" or "`#XXXX`") per guidelines. Locate the comment block that
mentions DIS-2147 immediately above the abort_timeout check in the
function/method containing abort_timeout, and replace "DIS-2147" with the
correct GitHub issue number while preserving the rest of the comment and the
surrounding logic that calls self.wait_for_decode_kv_capacity(dp_rank).await and
maps errors with Error::msg.
- Around line 667-677: Update the hardcoded DIS-2147 reference in the comment
block and the tracing target to the project-standard GitHub issue format (e.g.,
"GH-NNNN" or "`#NNNN`"); specifically edit the comment starting with "// DIS-2147
forensic logging..." and the tracing::info! call target "mocker::dis2147" (and
any inline string "prefill_kv_pin_start" context if it embeds the ticket) to use
the GitHub issue identifier instead (e.g., replace DIS-2147 with GH-XXXX or
`#XXXX`) so both the comment and the tracing target reflect the new issue
numbering.
- Around line 775-777: Update the inline comment referencing "DIS-2147" to use
the GitHub issue format (e.g., "GH-NNNN" or "`#NNNN`") per guidelines: locate the
comment near the wait_for_decode_ready reference in mocker.rs (the block
mentioning "DIS-2147" and "KV stayed pinned") and replace "DIS-2147" with the
appropriate GitHub issue identifier (e.g., "GH-<number>" or "#<number>") while
keeping the rest of the comment text unchanged.
- Around line 679-709: The inline ticket reference "DIS-2147" in the comment and
tracing targets should be replaced with the repository issue format (e.g.,
"GH-XXXX" or "`#XXXX`"); update the top comment above the wait_for_decode_ready
call and the tracing target string "mocker::dis2147" to the chosen GitHub issue
identifier, and search within the same function/block (including tracing::warn!,
tracing::info! calls and any other occurrences around
server.wait_for_decode_ready, pin_start, abort_room, stream_tx, active_requests,
request_uuid) to replace all remaining "DIS-2147" references so they
consistently use the GitHub issue number.
In `@lib/mocker/src/common/protocols.rs`:
- Around line 711-718: Doc comment describing the "Timeout (milliseconds) for
the prefill→decode NIXL handshake" contains an internal ticket ID "DIS-2147";
remove that internal Linear reference and replace it with either a public GitHub
issue reference (e.g., GH-NNNN or `#NNNN`) or plain prose describing the behavior
change (e.g., "see related issue on repository" or "behavior introduced to
prevent resource leakage during late decodes") so the comment no longer includes
internal ticket identifiers; update the same doc block that mentions prefill
abort behavior to reflect the public reference or plain prose.
In `@lib/mocker/src/scheduler/sglang/config.rs`:
- Around line 43-45: Update the doc comment on the max_num_seqs field to replace
the internal ticket reference "DIS-2147" with the repository's GitHub issue
notation (e.g., "GH-NNNN" or "`#NNNN`") per guidelines; locate the comment above
the pub(super) max_num_seqs: u64 field and substitute the ticket string so the
comment reads something like "Used by the GH-<issue> decode-side admission wait
to model the seq-slot budget" (preserve surrounding text and formatting).
In `@lib/mocker/src/scheduler/sglang/core.rs`:
- Around line 252-266: The inline comment referencing "DIS-2147" should be
replaced with the approved GitHub issue format; locate the mocker_metrics block
where MockerMetrics is created (the block that sets metrics.max_num_seqs =
self.config.max_num_seqs) and update the comment "// DIS-2147: expose the
seq-slot cap for the decode-side admission wait." to use the GitHub issue
identifier (e.g., "// GH-<number>: expose the seq-slot cap for the decode-side
admission wait." or "// #<number>: ...") following the project's coding
guidelines so the reference is consistent and searchable.
In `@lib/mocker/src/scheduler/vllm/live.rs`:
- Around line 32-35: Doc comment for the struct field max_num_seqs contains an
internal Linear ticket reference "DIS-2147"; update that reference to the
corresponding GitHub issue format (e.g., replace "DIS-2147" with "`#2147`" or
"GH-2147") in the doc comment above pub max_num_seqs so it follows the project's
GitHub issue convention.
In `@lib/mocker/src/services/bootstrap.rs`:
- Around line 10-15: Replace all internal Linear ticket IDs "DIS-2147" in the
comments around the bootstrap behavior descriptions with a public reference or
plain prose; specifically update the comment lines that mention
wait_for_decode_ready(room_id, timeout), abort_room(room_id), and
complete_room(room_id) (and other occurrences noted in the comment) to either
reference a GitHub issue (e.g., GH-2147 or `#2147`) or describe it as "a
configured abort timeout" / "related capacity bug" instead of the internal
ticket ID so no Linear IDs remain in source comments and test labels.
---
Outside diff comments:
In `@lib/bindings/python/rust/llm/replay.rs`:
- Around line 174-209: The Python typing stub for the constructor is out of
sync: the Rust/pyO3 constructor function new now accepts the
kv_transfer_abort_timeout_ms keyword but the typed stub (_core.pyi) still ends
at kv_transfer_bandwidth; update the stub to add kv_transfer_abort_timeout_ms:
Optional[int] (or appropriate integer type) to the constructor signature so
typed callers no longer get unexpected-keyword errors, ensuring the parameter
name and type match the runtime signature of new and any related overloads in
the same stub.
In `@lib/mocker/src/scheduler/vllm/core.rs`:
- Around line 510-525: The inline comment in function mocker_metrics()
references a Linear ticket "DIS-2147"; update that comment to use the project's
GitHub issue format (e.g., "GH-<number>" or "#<number>") per guidelines by
replacing "DIS-2147" with the correct GitHub issue identifier in the comment
above metrics.max_num_seqs so the comment reads e.g. "// GH-XXXX: expose the
seq-slot cap..." instead of the Linear ticket; ensure only the comment text is
changed and code/logic (mocker_metrics, metrics.max_num_seqs) remains untouched.
In `@lib/mocker/src/services/bootstrap.rs`:
- Around line 190-216: The decode-side wait uses the global RENDEZVOUS_TIMEOUT,
causing premature timeouts; change the tokio::time::timeout calls (e.g. where
ImmediateOrWait::Wait(rx) matches) to use the room-specific abort timeout (the
prefill/kv_transfer_abort_timeout_ms value) instead of RENDEZVOUS_TIMEOUT,
obtaining that timeout from the Room or prefill params available in
bootstrap.rs/handle_connection() and applying the same change to the other
occurrences around lines 365-370; ensure you pass the correctly converted
Duration to tokio::time::timeout and only remove the room from
rooms.remove(&room_id) if the room-specific timeout elapses.
---
Nitpick comments:
In `@lib/llm/src/mocker.rs`:
- Around line 684-686: The tracing::warn! call is doing eager string
interpolation for room_id and e; replace it with lazy formatting by passing
room_id and e as structured fields or positional arguments to tracing::warn!
instead of embedding them in the format string—locate the tracing::warn!
invocation that references room_id and e in mocker.rs (the prefill aborting
transfer log) and change it to use structured fields like room_id = %room_id and
error = %e or use positional placeholders so the values are only formatted when
the warn level is enabled.
In `@lib/mocker/src/common/protocols.rs`:
- Around line 1251-1260: The round-trip test builds MockEngineArgs without
setting kv_transfer_abort_timeout_ms so it remains None and won't catch
serialization regressions; update the MockEngineArgs::builder() call used in the
round-trip test (the builder chain that sets worker_type, max_num_seqs,
max_num_batched_tokens, reasoning, sglang, etc.) to explicitly set
kv_transfer_abort_timeout_ms to a concrete non-default value (e.g., Some(1234)),
then after restore/assertion verify the restored
args.kv_transfer_abort_timeout_ms equals that value; apply the same change to
the other builder usage in the file that mirrors this test.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 50e5f1b3-74a2-4ade-8460-66c44daeb240
📒 Files selected for processing (10)
components/src/dynamo/mocker/args.pycomponents/src/dynamo/mocker/config.pylib/bindings/python/rust/llm/replay.rslib/llm/src/mocker.rslib/mocker/src/common/protocols.rslib/mocker/src/scheduler/sglang/config.rslib/mocker/src/scheduler/sglang/core.rslib/mocker/src/scheduler/vllm/core.rslib/mocker/src/scheduler/vllm/live.rslib/mocker/src/services/bootstrap.rs
…t timeout Adds optional NIXL-style backpressure to the online mocker's disagg bootstrap rendezvous so cascade Steps 2 and 3 become reproducible without GPU compute. Gated entirely on the new MockEngineArgs::kv_transfer_abort_timeout_ms field. When None (default), behavior is unchanged from pre-DIS-2147. Mechanism: - Prefill side: after signal.completed and before complete_room, call server.wait_for_decode_arrival(room_id, timeout). The scheduler-side request stays active during the wait, so KV is pinned (modeling real disagg). On Ok -> normal handoff. On Err (timeout) -> server.abort_room() + surface error to client. - Decode side: before connect_to_prefill, poll scheduler.metrics_receiver until active_decode_blocks < total_blocks (free KV available). The act of connecting is decode's 'ready' signal. - Bootstrap server: new RoomOutcome enum, new prefill_waiting oneshot, new abort_room() with TTL cleanup, new ABORT_BYTE (0x02). Late decode on aborted room -> ABORT_BYTE -> clean error. All 5 existing bootstrap tests pass; 5 new DIS-2147 scenario tests pass; 253 mocker lib tests pass; cargo check / fmt / clippy clean. Linear: DIS-2147 Signed-off-by: nnshah1 <neelays@nvidia.com> (cherry picked from commit aee86e0995113e1016f5cb081522c0fdd71d3773)
…EngineArgs The DIS-2147 prefill-stranding extension added kv_transfer_abort_timeout_ms to MockEngineArgs / MockEngineArgsBuilder in lib/mocker/src/common/protocols.rs but never updated the pyo3 binding at lib/bindings/python/rust/llm/replay.rs. Result: MockEngineArgs.__new__() rejects the kwarg with "got an unexpected keyword argument 'kv_transfer_abort_timeout_ms'" when the test framework tries to set it from YAML, which blocks the mocker prefill-stranding sanity test that validates the cascade simulation fidelity vs real GPU behavior. Add the field to the #[pyo3(signature)] decorator, the fn arg list, and the builder chain — mirrors kv_transfer_bandwidth which sits next to it in the struct. Signed-off-by: nnshah1 <neelays@nvidia.com> (cherry picked from commit 8f36a982d1e135a52d2ddfaec676c7455e389aa1)
The original DIS-2147 implementation (aee86e09951) gave wait_for_decode_kv_capacity its own timeout, sourced from the same --kv-transfer-abort-timeout-ms value as the prefill-side abort. That created a redundant decode-side error path that doesn't exist in real engines: - vLLM: WAITING_FOR_REMOTE_KVS requests stay WAITING until the scheduler admits them — no decode-side timer. - sglang: PreallocQueue blocks on req_to_token_pool + token budget — no decode-side timer either. In real disagg the only authoritative timer is the PREFILL side's kv_transfer_abort_timeout_ms. When it fires, the prefill bootstrap room closes. The decode side eventually attempts connect_to_prefill on that room_id, finds it closed, and surfaces the failure via the bootstrap connection — not via its own clock. Empirical confirmation (this branch, 20-burst test before F1): 18 prefill_kv_pin_end outcome="aborted" ✓ authoritative 36 "Decode KV wait timed out after 2s" ← redundant; mocker-only Remove the decode-side timeout. wait_for_decode_kv_capacity now blocks on the watch channel indefinitely until both block-budget AND seq-slot gates have headroom. Cancellation is delegated to the surrounding request context (the future is dropped when the FE drops the request). Signature change: wait_for_decode_kv_capacity drops the `timeout: Duration` parameter. Only one caller in lib/llm/src/mocker.rs:777, updated. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com> (cherry picked from commit 8d1990b0e3b8c16a3db3d94de94b7b340858fecb)
Trim redundant inline comment in decode bootstrap path; the design rationale already lives on wait_for_decode_kv_capacity. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
79d9d08 to
b874918
Compare
PeaBrane
left a comment
PeaBrane
left a comment
There was a problem hiding this comment.
Requesting changes because the current ordering does not create the KV-stranding behavior described by the PR: it waits before prefill is submitted, so no prefill KV is held and the overload cascade cannot be reproduced. The shared pinned-prefill scheduler state described inline should work cleanly for both the live mocker and replay harness.
| // 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) => { |
There was a problem hiding this comment.
Blocking: this wait happens before the later sender.send(direct_request), so the prefill request has not entered the scheduler or allocated any KV during the claimed pin interval. That models rendezvous delay, but not KV stranding or prefill-capacity exhaustion.
Please represent this as a shared scheduler lifecycle instead:
RunningPrefill -> PrefillPinned -> Released/Aborted
The scheduler should run prefill normally, retain the completed request and its KV allocation in PrefillPinned, include that KV in occupancy metrics, and expose an explicit release_pinned_prefill(request_id, outcome) transition. The live mocker can trigger that transition from decoder admission/transfer or timeout; replay can schedule the same transition as a simulation event. Keeping TCP/timeout orchestration outside the scheduler gives both paths identical resource accounting while allowing different timing drivers.
… state Addresses @PeaBrane's review: the prior ordering submitted the disagg prefill to the scheduler only AFTER wait_for_decode_ready resolved, so no KV was held during the wait — the cascade could not form (the pin logs were aspirational). - Ordering fix (mocker.rs): the disagg prefill is submitted FIRST so its KV is really allocated; wait_for_decode_ready now governs only the *release* of the already-pinned KV. Strand duration is event-driven and load-dependent (not a fixed time), matching the modeled cascade. - Shared pinned scheduler state (vllm + sglang cores, kv managers): on prefill completion of a disagg stranding-candidate (request carries a bootstrap room), the request's blocks are moved into a pinned set instead of freed, and stay counted in active blocks so the pool fills. release_pinned() runs the deferred free. - Release trigger by context: live → decode transfer-completion via bootstrap (release-pin control message), abort_timeout frees on no-show; replay (offline) → modeled time-based release (decode_end + handoff_delay), since offline can't observe cross-instance decode pickup. Driver stall-advances virtual time to the earliest pending pin deadline. Tests: new test_disagg_prefill_strands_kv_until_release (vllm) + sglang analogue prove the pool stays full while pinned, a second prefill is rejected (cascade), and release drops occupancy to 0; non-bootstrap prefill does not strand (no regression). cargo test -p dynamo-mocker (363) + dynamo-llm kv_router (19) pass; fmt + clippy -D warnings clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
|
@PeaBrane — addressed in 4a44959. You were right: the prior ordering submitted the prefill only after Implemented the shared pinned-prefill scheduler state you described:
This is the "works cleanly for both the live mocker and replay harness" split you asked for. Verification: new One open design question for you: in replay the time-based release uses |
CI (test_disagg_background_prefill_sticky) caught that pinned prefill KV leaked into the router-facing per-worker load metric (active_decode_blocks), inflating the stranding worker's reported load and diverting even the no-sticky control request away from normal KV-overlap routing. The pinned-block count has two distinct consumers: (1) pool capacity / admission — must include pinned KV so the pool fills and the cascade forms; (2) the router's per-worker load signal (ActiveLoad.active_decode_blocks) used for worker selection — must NOT include a stranded prefill awaiting handoff, or it mis-diverts unrelated sequences. Separate them: track a pinned_block_footprint (vllm) / pinned_token_footprint (sglang) and subtract it from the router-facing metric only; the admission path keeps reading num_active_blocks()/active_kv_blocks() directly (cascade preserved). sglang stranding test repointed its pool/cascade assertions to a capacity-view helper and now also asserts the router metric IS discounted. Verified: test_disagg_background_prefill_sticky passes; stranding unit tests (vllm+sglang) + non-bootstrap-no-strand pass; cargo test -p dynamo-mocker (363) + dynamo-llm kv_router (167) pass; fmt + clippy -D warnings clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
Captures intent + mechanism for modeling disagg prefill-KV stranding via the engine-opaque disaggregated_params (vLLM NIXL kv_transfer_params) instead of the sglang-shaped bootstrap rendezvous: prefill emits a transfer handle + pins KV; decode carries it and pulls → event-driven release (live via the reframed channel, replay via in-process correlation). sglang behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
github-actions
Bot
added
the
documentation
Re-architect disagg prefill->decode KV stranding to the vLLM/NIXL model
(engine-opaque disaggregated_params) instead of the sglang-shaped bootstrap
rendezvous. Mocker-side only — no prefill_router changes. Spec:
docs/superpowers/specs/2026-06-15-mocker-vllm-disagg-params-stranding-design.md
vLLM path:
- Prefill emits real disaggregated_params {transfer_id, prefill_host,
prefill_port} (replacing the dummy) and its request stream COMPLETES normally
— no stream-holding — so it flows through the existing prefill_router drain
without deadlock. Matches real vLLM, where the NixlConnector holds KV
independent of the request lifecycle.
- The scheduler pin (pin_completed) is decoupled from the request and survives
its completion; pinned KV counts toward pool capacity (cascade) but is
discounted from the router-facing load metric (so stranding doesn't mis-divert
overlap routing).
- Release is channel-server-driven and event-driven, keyed by transfer_id:
the prefill registers a pin with the BootstrapServer; a decode connect fires
the release (then ACK); kv_transfer_abort_timeout_ms with no decode releases +
ABORT. Strand duration is load-dependent (until that request's decode pulls),
never a fixed time.
- Replay: in-process transfer_id -> uuid correlator releases the pin when the
matching decode is scheduled (time-based release removed).
Channel: key generalized room_id -> transfer_id with RoomId/TransferId (=u64)
aliases so each engine keeps native vocabulary; connect/wait/ACK/ABORT behavior
unchanged. sglang keeps its bootstrap-triple disaggregated_params format and
behavior (path-gated on !vllm_disagg_handoff) — unchanged.
Verified in-container: dynamo-mocker 366 + dynamo-llm mocker 6 + kv_router 167
pass; new bootstrap pin tests (decode-connect release, abort-timeout release) +
replay event-driven release test; fmt + clippy -D warnings clean. prefill_router
untouched; sglang tests unchanged. The multi-process live e2e
(test_disagg_background_prefill_sticky) is covered by CI (not runnable locally —
needs frontend+NATS+etcd+bindings).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: nnshah1 <neelays@nvidia.com>
Address review (graham): - vLLM disagg transfer_id was rand::random() into a long-lived channel/pin registry — a collision would silently mis-release/leak a pinned KV. Use a process-monotonic AtomicU64 (collision-free), matching the replay path. - Build the emitted disaggregated_params via serde_json::json! (infallible) instead of to_value(..).expect() on the live prefill path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: nnshah1 <neelays@nvidia.com>
…-worker busy-spin
Make the vLLM disaggregated_params (NIXL-style) KV-stranding path engage in a live, multi-process mocker deployment with no bootstrap rendezvous and no router changes:
- config.py: a vLLM disagg prefill registers set_disaggregated_endpoint(None, None) so the unchanged prefill_router takes its existing NoBootstrapEndpoint -> output-disaggregated_params branch. sglang keeps advertising host/port (bootstrap path) and is byte-for-byte unchanged (vLLM-gated). main.py threads engine_type into build_runtime_config (the pyclass exposes no getter).
- mocker.rs: start the KV-transfer server for a vLLM prefill on an OS-assigned port even without --bootstrap-ports; emit {transfer_id, host, port} in disaggregated_params; decode reads it and connects. Add is_vllm() helper.
- live.rs: receive_requests now blocks when there is no *runnable* work (not when fully empty). A pinned-only worker previously hot-spun (is_empty() counts pinned), starving the decode workers whose admission releases the pin -> runaway prefill->decode handoff latency. A pinned-only worker can only advance via request_rx (new request or release_pin), so blocking on recv() is correct and deadlock-free.
- bootstrap.rs: register_pin always installs the reclaim timer (default RENDEZVOUS_TIMEOUT when no explicit abort timeout) for parity with sglang's wait_for_decode_ready, and stores the timer's AbortHandle so a decode-connect release cancels it (no lingering task on the happy path).
Signed-off-by: nnshah1 <neelays@nvidia.com>
…ss, abort parity Signed-off-by: nnshah1 <neelays@nvidia.com>
End-to-end check that Least Loaded frontend routing drives vLLM mocker prefill/decode workers launched WITHOUT --bootstrap-ports over the disaggregated_params (NIXL) path and strands/releases KV. The pin lifecycle (prefill_kv_pin_start / decode_kv_wait_start / prefill_kv_pin_end) is observable in worker logs under DYN_LOG=mocker::kv_abort=info. Signed-off-by: nnshah1 <neelays@nvidia.com>
…he PR Signed-off-by: nnshah1 <neelays@nvidia.com>
PeaBrane
left a comment
PeaBrane
left a comment
There was a problem hiding this comment.
Please broaden the new end-to-end test across the relevant router modes.
| enforce_disagg=True, | ||
| request_plane=request_plane, | ||
| event_plane="nats", | ||
| router_mode="least-loaded", |
There was a problem hiding this comment.
There’s no reason to limit this coverage to Least Loaded, even though it motivated the test. Please parameterize router_mode and cover at least round-robin, kv, and least-loaded so the bootstrap-free disaggregated_params path is validated independently of routing mode.
PeaBrane
left a comment
PeaBrane
left a comment
There was a problem hiding this comment.
The scheduler-level stranding model is useful, but the current implementation needs clearer core boundaries, extraction of the live disaggregation orchestration, and a faithful shared lifecycle for SGLang and vLLM before merging.
| /// the matching decode pulls the transfer. `None` disables stranding (the | ||
| /// request frees normally on completion). | ||
| #[serde(default, skip_serializing_if = "Option::is_none")] | ||
| pub bootstrap_room: Option<u64>, |
There was a problem hiding this comment.
Pinning belongs in the core scheduler, but bootstrap_room and release_pin expose live transport concepts here and make DirectRequest double as a control command. Could we model this explicitly, for example:
enum CompletionPolicy {
Release,
PinUntilReleased,
}
enum SchedulerInput {
Request {
request: DirectRequest,
completion_policy: CompletionPolicy,
},
ReleasePinned {
uuid: Uuid,
},
}Live and replay can retain their own SGLang room / vLLM transfer-ID to scheduler-UUID correlation outside core. This keeps the reusable core responsible for scheduler lifecycle rather than transport vocabulary.
| // `transfer_id`. The server fires the registered `release_pin` when the | ||
| // matching decode connects (then ACKs it) or on abort-timeout. | ||
| #[allow(clippy::type_complexity)] | ||
| let (prefill_transfer_id, emitted_disagg_params, vllm_disagg_handoff): ( |
There was a problem hiding this comment.
This has grown into a substantial disaggregated-transfer subsystem inside generate(): engine-specific wire parsing, transfer-ID generation, decode-capacity waiting, TCP coordination, pin callback registration, timeout handling, and stream ordering. Please extract it into mocker/disagg.rs behind typed structures, for example a DisaggProtocol plus a DisaggHandoff/state object. generate() should mainly resolve the handoff, submit the scheduler request, run the required lifecycle gates, and forward scheduler output.
| // The TWO disagg paths differ in WHO drives the pin release, and the | ||
| // difference is intentional — vLLM is the NIXL-faithful reframe: | ||
| // | ||
| // - **sglang bootstrap path** (`bootstrap_info` present): the decode |
There was a problem hiding this comment.
Cross-checking current SGLang, this is not actually unchanged. With the default optimistic_prefill_retries=0, decode physically preallocates KV and publishes destination metadata before prefill leaves its bootstrap queue. Only optimistic-prefill mode computes early enough to strand on decode capacity. SGLang also retains source KV until transfer success, not merely until decode connects.
I still think both engines should use a shared typed lifecycle, but it needs distinct states/events such as SourcePinned, DestinationReserved, Transferring, Completed, and Aborted. Please either preserve baseline SGLang ordering or expose optimistic-prefill behavior explicitly, and release the source pin only after the modeled transfer completes.
PeaBrane
left a comment
PeaBrane
left a comment
There was a problem hiding this comment.
There are a lot of independently risky moving parts bundled here: scheduler pin ownership for two engines, live channel rendezvous/reclaim, bootstrap-free endpoint registration, decoder admission, replay correlation, configuration/API changes, and the router E2E. It may be worth staging this as a small dependent PR series so each invariant is easier to review and validate, for example: (1) scheduler pin/release primitive plus core tests, (2) live vLLM handoff, endpoint advertising, and reclaim, (3) scheduler-owned decode admission/reservation, then (4) replay/SGLang parity and end-to-end coverage. The existing changes-requested verdict remains appropriate while the correctness points below and the earlier architecture comments are addressed.
| // 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 { |
There was a problem hiding this comment.
This is a metrics snapshot, not decoder admission or a capacity reservation. It proves only that at least one block and sequence slot are currently free; it does not account for this request's KV footprint, and multiple waiting decodes can all pass the same snapshot. connect_to_prefill then releases the source pin before the DirectRequest reaches the scheduler. Replay has the same issue when it releases at dispatch_decode. The pin release should be driven by a scheduler-owned reservation/admission event for this specific request.
| "prefill_kv_pin_start" | ||
| ); | ||
| let release_sender = sender.clone(); | ||
| server.register_pin( |
There was a problem hiding this comment.
The abort clock starts here immediately after scheduler submission, not when the scheduler actually transitions the completed prefill into pinned state. If prefill compute exceeds the timeout, release_pin removes the pin candidate before completion, so the request never strands and the configured transfer timeout incorrectly includes compute time. Please register/start this lifecycle from the scheduler's pin transition, or add an explicit pin-ready notification.
| pub(crate) fn is_empty(&self) -> bool { | ||
| self.waiting.is_empty() && self.running.is_empty() | ||
| // a worker holding pinned (stranded) KV is not idle. | ||
| self.waiting.is_empty() && self.running.is_empty() && self.pinned.is_empty() |
There was a problem hiding this comment.
This change makes a pinned-only SGLang core non-empty, but the unchanged live SGLang loop still blocks only on core.is_empty(). While a pin is held it will therefore repeatedly execute zero-work, zero-duration passes instead of waiting on request_rx. vLLM was changed to gate on !has_runnable_work() for exactly this reason; SGLang needs the same live-loop change and coverage.
| // The mocker prefill/decode run as local processes; the | ||
| // bootstrap server binds 0.0.0.0, so loopback is the | ||
| // reachable pull address for the modeled transfer. | ||
| "prefill_host": "127.0.0.1", |
There was a problem hiding this comment.
This is not reachable in a normal distributed mocker deployment: the repository's disaggregated mocker examples run prefill and decode as separate pods, so 127.0.0.1 on decode refers to the decode pod, not prefill. Please emit the configured advertised address, such as DYN_HTTP_RPC_HOST/the runtime registration host, together with the bound port.
| request | ||
| .prefill_result | ||
| .as_ref() | ||
| .and_then(|r| VllmDisaggParams::parse(&r.disaggregated_params)) |
There was a problem hiding this comment.
A present but malformed prefill_result.disaggregated_params is silently treated as no handoff, so this decode proceeds as aggregated while the source remains pinned until timeout. Absence can represent a genuinely aggregated request, but present-invalid handoff data should return a structured error rather than weakening the contract.
| discovery_backend = "etcd" | ||
| request_plane = "tcp" | ||
| shared_namespace = f"test-namespace-{generate_random_suffix()}" | ||
| mocker_args = {"speedup_ratio": SPEEDUP_RATIO, "block_size": BLOCK_SIZE} |
There was a problem hiding this comment.
This test does not currently distinguish the stranding implementation from an immediate-free implementation: requests are sequential, the KV pool is not constrained, kv_transfer_abort_timeout_ms remains None so the capacity wait is disabled, and the only assertion is request success. Please create actual decoder pressure/concurrency and assert an observable pin/admission/reclaim outcome. The test command helper will also need to forward the timeout argument if this path is enabled through it.
…replay Offline disagg replay released the stranded prefill's pin at decode *routing* (dispatch_decode = the router's worker pick), so the strand only modeled rendezvous latency, not slot-limited stranding — a decode-saturated run produced byte-for-byte the same metrics as no stranding. Move the release to the decode engine's *slot admission*: dispatch_decode is now pure routing (enqueue into the decode engine); release_prefill_pin_for fires from handle_decode_engine_effects for each request the decode engine admits (secures a seq-slot + KV blocks and pulls the transfer). A completion-path safety net releases pins for decodes that complete/reject without ever being admitted. Under decode saturation the admission — and thus the release — is delayed, so the prefill's KV stays pinned and pressures the prefill pool (the cascade), matching the live model. Signed-off-by: nnshah1 <neelays@nvidia.com>
|
will be done separately - |
2 participants
Summary
Models disaggregated prefill→decode KV-block stranding in the mocker the way real vLLM does it — via the engine-opaque
disaggregated_params(NIXLkv_transfer_params) — so cascade/overload tests reproduce the failure mode where a prefill worker holds KV for a decode that arrives late (or never), stranding the blocks until the transfer completes or times out.The vLLM path is additive and engine-gated: sglang keeps its bootstrap-triple
disaggregated_paramsformat and behavior unchanged, and the router is not modified — only which of its existing branches is taken.How it works
Stranding (scheduler core). A vLLM prefill runs, pins its KV on completion (the
ActiveSequenceand its block handles are retained), and emitsdisaggregated_params = {transfer_id, prefill_host, prefill_port}. Pinned KV still counts toward pool capacity (the cascade) but is discounted from the router-facing load metric, so a strand doesn't mis-divert overlap/load routing of unrelated requests. Release is event-driven and load-dependent (the strand lasts until that request's decode is admitted and pulls): live → the decode connects to the prefill bytransfer_idover the cross-process channel and the pull fires the pin release; replay → an in-processtransfer_id → pinned-uuidcorrelator releases when the matching decode is scheduled.Bootstrap-free live engagement (mocker-only). A vLLM disagg prefill registers
set_disaggregated_endpoint(None, None)— it advertises as a disagg/prefill worker with no bootstrap host/port — so the unchanged prefill_router takes its existingNoBootstrapEndpoint→ output-disaggregated_paramsbranch (execution.rs). The KV-transfer (channel) server starts for a vLLM prefill on an OS-assigned port even without--bootstrap-ports; the real port is emitted indisaggregated_paramsand the decode reads it. sglang always setsbootstrap_port, so it never reaches this branch — its behavior is byte-for-byte unchanged.Scheduler liveness fix. The live scheduler loop blocks when there is no runnable work (
!has_runnable_work()), not when it is fully empty (is_empty(), which counts pinned KV). Previously a pinned-only worker hot-spun — never blocking while holding a pin — which starved the decode workers whose admission releases the pin, compounding the prefill→decode handoff latency. A pinned-only worker can only advance via arequest_rxmessage (a new request or arelease_pin), so blocking onrecv()is correct and deadlock-free.Abort / reclaim.
register_pinalways installs the reclaim timer — an explicitkv_transfer_abort_timeout_mswhen set, elseRENDEZVOUS_TIMEOUT— for parity with sglang'swait_for_decode_ready, so a never-pulled pin is always reclaimed (and a cleanABORT_BYTEgoes to waiting/late decodes). The timer'sAbortHandleis stored and cancelled when a decode-connect releases the pin, so the happy path leaves no lingering task. Decode-side admission gates on both the block budget and the sequence-slot budget (max_num_seqs), matching real vLLM/sglang; block-only gating made the abort path unreachable under seq-saturation.Forensic logging at
target: "mocker::kv_abort"—prefill_kv_pin_start,decode_kv_wait_start,prefill_kv_pin_end {outcome, duration_ms}— reconstructs the stranding graph (which prefill held KV for which decode, how long, completed vs aborted).Channel key / naming. The cross-process channel uses a single engine-neutral
u64key,transfer_id. Type aliasesRoomIdandTransferId(both= u64) let each engine read in its native vocabulary (sglang bootstrap room vs vLLMkv_transfer_paramstransfer id) with no conversion — the same value hits the same map entry.Changed files
lib/llm/src/mocker.rsdisaggregated_params; decode parses it → live channel connect bytransfer_id. Starts the KV-transfer server for a vLLM prefill on an OS-assigned port (no--bootstrap-ports). sglang: existing bootstrap-room path, unchanged.components/src/dynamo/mocker/config.py,main.pyset_disaggregated_endpoint(None, None)(disagg, no bootstrap host/port) → the unchanged prefill_router takesNoBootstrapEndpoint. sglang keeps host/port.lib/mocker/src/services/bootstrap.rstransfer_id+RoomId/TransferIdaliases; pin registry with always-on reclaim timer (parity defaultRENDEZVOUS_TIMEOUT) whoseAbortHandleis cancelled on decode-connect release. connect/wait/ACK/ABORT machinery unchanged.lib/mocker/src/scheduler/vllm/{core.rs,live.rs}pin_completed/release_pinned/take_pinned); pinned footprint discounted from the router metric. Live loop blocks on no-runnable-work (not no-work) so a pinned-only worker parks instead of hot-spinning.lib/mocker/src/replay/offline/*transfer_id → pinned-uuidcorrelator releases when the matching decode is scheduled.lib/bindings/python/rust/llm/replay.rs,_core.pyikv_transfer_abort_timeout_msparam + stub.Edge cases
num_active_blocks→ new prefills are rejected/queued; the router-metric discount keeps stranding from mis-diverting overlap routing.release_pinned(uuid)is idempotent; the channel release fires at-most-once.disaggregated_params(aggregated): no pin, normal free.Non-goals
Test
cargo fmt/cargo clippyclean ondynamo-mocker+dynamo-llm;cargo test -p dynamo-mocker --lib= 366 passed (incl. pin/release/abort channel tests, replay event-driven release, and JSON round-trip). Live multi-process e2e (tests/router/test_router_e2e_with_mockers.py) validates the bootstrap-free disagg_params path (no_bootstrap), the unchanged bootstrap path (with_bootstrap), and background-prefill stickiness — green under load, with flat pin durations (no handoff-latency compounding). A newtests/router/test_ll_strand_check.pyexercises the full path under Least Loaded frontend routing (vLLM disagg, no--bootstrap-ports) end-to-end. The pinned-prefill liveness and abort/no-decode concurrency paths were reviewed for locks-across-await, missed wakeups, and deadlock.🤖 Generated with Claude Code