Skip to content

perf: skip KVCR deadline queries for ranks without pulls - #611

Merged
jthomson04 merged 1 commit into
ai-dynamo:mainfrom
PeaBrane:codex/kvcr-idle-deadline
Oct 27, 2026
Merged

jthomson04 merged 1 commit into
ai-dynamo:mainfrom
PeaBrane:codex/kvcr-idle-deadline

Conversation

@PeaBrane

Copy link
Copy Markdown
Contributor

Why and what changed

Every vLLM rank queried the shared KVCR registry for its next deadline even when it had no incoming pull. At large worker/attention-DP counts, those repeated locks and lookups add substantial replay overhead.

Return early when the adapter's staged-pull map is empty. Detached pulls keep their stages until their blocks are consumed, so their wakeups remain intact. The lifecycle regression covers empty, parked, moving, detached, source-only, landed-but-unconsumed and completed states. Transfer advancement, bandwidth allocation, counters and event ordering are unchanged.

Review map

  • Risk level: low; the guard uses the existing staged-pull lifecycle, including detached pulls.
  • crates/core/src/engine/scheduler/vllm/kvcr.rs: the guard and its lifecycle regression.
  • crates/core/src/engine/kvcr.rs: measured baseline bookkeeping TODO; broader optimization remains deferred.
  • docs/replay/engine/kv-cache.md: describe the shortcut and remaining overhead.
  • Independent of the draining-source correctness fix in fix: reject KVCR staging from draining sources #600. Both PRs target main and can merge in either order. No public API, configuration or output-schema change.

Evidence

The same guard was measured in commit pair 52e131bc → 7376ce71, using 108 clean release trials with six interleaved repetitions per variant, one pinned CPU and an A/A control first. The conservative A/A floor was 21.9%.

Configuration Before µs/request After µs/request Replay time saved After/before, 95% CI
512 workers, DP1, hinted 122.17 82.55 32.4% 0.676 [0.582, 0.769]
2048 workers, DP1, hinted 573.42 386.53 32.6% 0.674 [0.649, 0.721]
512 workers, DP8, hinted 1062.61 501.58 52.8% 0.472 [0.437, 0.533]

All six pairs favor the guard in each listed case. Whole benchmark-process time, including setup and output serialization, falls 24.0%, 29.8% and 50.7%, respectively. At W512, KVCR-off is unchanged within noise (ratio 0.996 [0.905, 1.065]); idle KVCR improves 30.6%. Hinted gains at W8/32/128 are below the conservative noise floor. Every trial is retained.

Full before/after CanonicalV1 bytes, including KVCR statistics, match in all timed trials and 12 additional pairs: 20 distinct worker/arm/bandwidth/DP configurations. No fields were removed. The temporary benchmark and instrumentation are excluded from the PR.

The initial pre-optimization study measured hinted/G2-only ratios of 1.168/1.226/1.182/1.975/3.165 at W=8/32/128/512/2048; these are the baseline ratios retained in the TODO. At W2048 it counted 1.02 billion deadline queries versus 3.10 million take_landed calls. The proposed timestamp/concurrency/source cost model fit poorly (R²=0.58), so no asymptotic claim is made. The candidate profile retained 14,671 samples with zero reported loss; registry next_deadline still accounts for 25.7% of recognized replay stacks. Truncated/unresolved callers limit attribution.

  • Standalone branch on cd6efc5e: cargo test -p aisimulate-core — 2,488 passed, zero failed, one ignored, across 19 test binaries.
  • The lifecycle fixture uses a 1 ms first-byte delay and a one-block transfer taking 1 ms, yielding deadlines at 1 ms and 2 ms; a landed pull remains due until consumed.
  • cargo fmt --all -- --check and cargo clippy -p aisimulate-core --all-targets passed; Clippy has warnings on unchanged code.
  • The timing/parity evidence above isolates the identical guard on the earlier revision. Timings were not rerun on the newer standalone base. CI and an independent review of this extraction are not claimed.

Modeling or data provenance

Synthetic fixed-worker fixtures: fixed vLLM timing, private G2, 200 requests/worker/second, eight 11-turn sessions per worker, and 20% of follow-ups migrate with a prior-worker hint. Default KVCR timing is 1 ms / 50 GB/s. Measurements used release Rust 1.98.1 and one logical CPU on an Intel Xeon Platinum 8559C. Each before/after pair holds workload, DP size and cache capacity fixed. Native KVCR parity is not claimed.

Tracking

Performance follow-up to #454; companion correctness PR: #600. Further transfer-bookkeeping optimizations remain deferred and must preserve CanonicalV1 output bytes.

This pull request and its description were written by Codex.

Avoid shared registry locks and deadline lookups when a rank has no
incoming pull. Retain deadlines for detached pulls until their blocks
are consumed, with lifecycle regression coverage.

Record the measured baseline bookkeeping cost and defer the remaining
transfer optimizations. This change is independent of the draining-source
staging fix in PR ai-dynamo#600.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Rudy Pei <yanrpei@gmail.com>
@PeaBrane
PeaBrane requested review from a team as code owners October 27, 2026 16:36
@copy-pr-bot

copy-pr-bot Bot commented Oct 27, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Oct 27, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

No actionable comments were generated in the recent review.

Recent review info
Run configuration
Commits

Reviewing files that changed from the base of the PR and between cd6efc5 and 8304832.

Files selected for processing: 3
  • crates/core/src/engine/kvcr.rs
  • crates/core/src/engine/scheduler/vllm/kvcr.rs
  • docs/replay/engine/kv-cache.md
Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Recent review details
Context from checks skipped due to timeout: 1
  • GitHub Check: CodeRabbit
Additional context used
Code guidelines: 3
crates/core/src/engine/scheduler/vllm/CLAUDE.md — auto-discovered
AGENTS.md — configured
REVIEW.md — configured
Path-based instructions: 6
Review scheduler state transitions and event ordering, including cancel, preempt, drain, retry, duplicate, empty, and terminal paths.

CodeRabbit configuration file

Files:

  • crates/core/src/engine/scheduler/vllm/kvcr.rs
  • crates/core/src/engine/kvcr.rs
Check commands, defaults, supported runtimes, public names, and claims against executable behavior.

CodeRabbit configuration file

Files:

  • docs/replay/engine/kv-cache.md
Read REVIEW.md before commenting.

CodeRabbit configuration file

Files:

  • docs/replay/engine/kv-cache.md
  • crates/core/src/engine/scheduler/vllm/kvcr.rs
  • crates/core/src/engine/kvcr.rs
Applicability reference: crates/core/src/engine/scheduler/vllm/AGENTS.md: CLAUDE.md Source excerpt: vLLM admits based on the current known sequence and permits decode-time preemption.

CodeRabbit inference engine (crates/core/src/engine/scheduler/vllm/CLAUDE.md)

Files:

  • crates/core/src/engine/scheduler/vllm/kvcr.rs
Source excerpt: Keep run logs, investigation notes, and benchmark results out of `docs/`; measured evidence belongs under `benchmarks/evidence/`.

CodeRabbit inference engine (AGENTS.md)

Files:

  • docs/replay/engine/kv-cache.md
  • crates/core/src/engine/scheduler/vllm/kvcr.rs
  • crates/core/src/engine/kvcr.rs
Source excerpt: Binary parquet changes require a machine-readable summary of keys, shapes, units, row counts, coverage, and anomalies plus proof that a real consumer reaches the new data.

CodeRabbit inference engine (REVIEW.md)

Files:

  • docs/replay/engine/kv-cache.md
  • crates/core/src/engine/scheduler/vllm/kvcr.rs
  • crates/core/src/engine/kvcr.rs
Multi-repo context ai-dynamo/dynamo, ai-dynamo/aiconfigurator

Linked repositories findings

ai-dynamo/dynamo

  • Dynamo directly pins both the Python aisimulate package and Rust aisimulate-core to the exact same version, so release compatibility must be maintained across both bindings. [::ai-dynamo/dynamo::pyproject.toml:15-18] [::ai-dynamo/dynamo::Cargo.toml:53-56]
  • Offline virtual-clock replay is delegated to AISimulate, with Dynamo registering DynamoReplayRunnerFactory as an AISimulate runner plugin. [::ai-dynamo/dynamo::components/src/dynamo/mocker/main.py:4] [::ai-dynamo/dynamo::pyproject.toml:127-128]
  • Dynamo’s replay factory imports AISimulate’s replay contracts and advertises replay-spec API version 1; the changed internal KVCR scheduling behavior therefore runs beneath this downstream replay boundary without a declared public API change. [::ai-dynamo/dynamo::components/src/dynamo/replay/simulation.py:21-32] [::ai-dynamo/dynamo::components/src/dynamo/replay/simulation.py:52-75]
  • Rust replay entrypoints lower Dynamo configuration into AISimulate-owned ReplaySpec/Replayer contracts and use CURRENT_REPLAY_SPEC_VERSION. [::ai-dynamo/dynamo::lib/mocker/src/replay/offline/entrypoints.rs:4-15] [::ai-dynamo/dynamo::lib/mocker/src/replay/offline/entrypoints.rs:118-145]
  • Dynamo’s KV-event adapter consumes AISimulate engine events and replay artifacts, making preservation of event ordering and replay output compatibility important even though no direct next_deadline or CanonicalV1 reference exists downstream. [::ai-dynamo/dynamo::lib/mocker/src/replay/offline/extensions/kv_events/mod.rs:4-18] [::ai-dynamo/dynamo::lib/mocker/src/replay/offline/extensions/kv_events/mod.rs:181-205]

ai-dynamo/aiconfigurator

  • AIconfigurator is a migration/compatibility reference: it directs users to AISimulate and shows EngineReplayRunnerFactory as the replacement replay integration. [::ai-dynamo/aiconfigurator::docs/aisimulate_migration.md:48-65]
  • Its transition policy explicitly limits the repository to bug, security, and migration-blocking fixes, so this PR’s runtime behavior should be reviewed against AISimulate/Dynamo compatibility rather than treated as an AIC implementation change. [::ai-dynamo/aiconfigurator::tests/unit/test_repository_transition_policy.py:10-15] [::ai-dynamo/aiconfigurator::tests/unit/test_repository_transition_policy.py:49-59]
📝 Summary

Risk: Low to medium. The change only skips deadline lookups when an adapter has no staged pulls. Human attention should focus on: (1) detached pulls retaining their stages and wakeups, (2) deadline behavior when pulls land or complete, and (3) confirming the reported replay-byte equivalence and performance measurements.

Behavior and contracts

  • VllmKvcrAdapter::next_deadline returns None when stages is empty. Otherwise, it queries the registry for this rank’s next deadline.
  • Detached pulls remain staged until their blocks land, so they continue to contribute deadlines.
  • The change does not alter public APIs or configuration, according to the supplied summary. The inspected method is pub(super).
  • Comments and documentation record the optimization and its performance context. No changes to transfer behavior are reported.

Evidence and readiness

  • The commit adds a regression test for detached pulls and tests deadline behavior as stages change. Source inspection confirms the empty-stage guard.
  • The author reports passing core tests, formatting, and Clippy. Performance and CanonicalV1 byte comparisons were measured on an earlier revision, not rerun on this standalone base.
  • CI results and an independent review are not supplied. Review-finding counts are unavailable.
  • Technical quality: The guard is localized, and the supplied tests target detached-pull wakeups.
  • Merge readiness: Not established by the available evidence. Confirm current-base validation and CI before relying on the earlier measurements.

VllmKvcrAdapter::next_deadline now skips registry lookups when no pulls are staged. A test covers detached-pull deadlines and cross-rank registry advancement. Comments and documentation describe the optimization and replay-output constraint.

Changes: KVCR Deadline Lookup

Layer / File(s) Summary
Skip deadline lookups without staged pulls
crates/core/src/engine/scheduler/vllm/kvcr.rs, crates/core/src/engine/kvcr.rs, docs/replay/engine/kv-cache.md
The adapter returns None when no pulls are staged and continues to use registry deadlines for staged pulls. The test covers detached pulls and another rank advancing the registry. Comments and documentation describe KVCR bookkeeping and the requirement to preserve CanonicalV1 replay output bytes.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

  • ai-dynamo/aisimulate#600: Includes the same deadline early return and detached-pull lifecycle test, plus changes to draining-source staging.
Merge Risk: Minimal · up to 83048

No merge-blocking issue was identified in the deadline lookup optimization.

Pre-merge checks: Passed 7 · Inconclusive 1

Failed Failed checks: 1 inconclusive

Check name Status Explanation Resolution
Review Evidence Inconclusive The visible PR description names relevant commands and results (cargo test -p aisimulate-core, cargo fmt --all -- --check, and cargo clippy -p aisimulate-core --all-targets), reports negative an… Provide the complete PR description, including the omitted modeling/data-provenance section, and rerun this check.
Passed Passed checks: 7 passed
Check name Status Explanation
Title check Passed The title precisely states the behavioral change: KVCR deadline queries are skipped for ranks without pulls.
Description check Passed The description explains the problem, implementation, lifecycle coverage, risks, evidence, modeling assumptions, compatibility impact, and tracking. It does not provide explicit CodeRabbit or reviewed…
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Cross-Layer Contract Passed No public input, schema field, default, or serialized output declaration changed. The internal next_deadline guard is traced from VllmKvcrAdapter through HostOffloadAdapter and the scheduler ret…
Modeling And Data Evidence Passed The PR supplies provenance and reproducible before/after evidence for the selection-logic change. The diff adds the stages.is_empty() guard in VllmKvcrAdapter::next_deadline and a lifecycle regres…
Compatibility Boundaries Passed No compatibility-boundary failure is introduced. The PR changes only two internal Rust implementation/test files and one documentation file; it changes no Python surface, schema, binding, version, dep…
Full details: Description check

Explanation

The description explains the problem, implementation, lifecycle coverage, risks, evidence, modeling assumptions, compatibility impact, and tracking. It does not provide explicit CodeRabbit or reviewed-commit entries, but the required information is otherwise mostly complete and the low-risk change does not require a Codex review commit.

Full details: Cross-Layer Contract

Explanation

No public input, schema field, default, or serialized output declaration changed. The internal next_deadline guard is traced from VllmKvcrAdapter through HostOffloadAdapter and the scheduler retry path. The added lifecycle test covers empty, parked, detached, landed-but-unconsumed, completed, and source-only states. Existing replay tests cover CanonicalV1 serialization and KVCR report fields. The PR description also supplies before/after CanonicalV1 byte-parity evidence.

Full details: Modeling And Data Evidence

Explanation

The PR supplies provenance and reproducible before/after evidence for the selection-logic change. The diff adds the stages.is_empty() guard in VllmKvcrAdapter::next_deadline and a lifecycle regression covering detached pulls and deadline preservation. The description identifies the measured commit pair, workload, hardware, toolchain, 108 release trials, interleaving, A/A control, confidence intervals, and three before/after configurations. It also reports CanonicalV1 byte parity across timed and additional configurations, unchanged KVCR-off timing, and retained trials. This is not binary-only validation and does not contain an unexplained output refresh.

Full details: Compatibility Boundaries

Explanation

No compatibility-boundary failure is introduced. The PR changes only two internal Rust implementation/test files and one documentation file; it changes no Python surface, schema, binding, version, dependency, or artifact metadata. VllmKvcrAdapter::next_deadline remains pub(super) and only adds an existing-lifecycle guard; its caller still consumes the same Option&lt;f64&gt; scheduling hint. KVCR defaults remain 1.0 ms and 50.0 GB/s in crates/core/src/engine/config.rs:392-399, and the documented defaults remain synchronized in docs/replay/engine/kv-cache.md:275-280. The report field and kvcr serialization remain unchanged in crates/core/src/replay/report.rs:49-52,481-483. The crate and Python package remain version 0.13.0 in crates/core/Cargo.toml:4-8 and python/aisimulate/pyproject.toml:4-7; no Dynamo dependency or import was added.

Full details: Review Evidence

Explanation

The visible PR description names relevant commands and results (cargo test -p aisimulate-core, cargo fmt --all -- --check, and cargo clippy -p aisimulate-core --all-targets), reports negative and boundary coverage, reports CanonicalV1 parity, and explicitly says CI is not claimed. The changed test also covers detached, landed-but-unconsumed, and completed states (crates/core/src/engine/scheduler/vllm/kvcr.rs:515-535). No inactive nested workflow is cited in the visible text. However, the supplied description is truncated during “Modeling or data provenance”. The complete description is required to verify the evidence distinctions and to rule out an inactive nested-workflow claim in the omitted text.

Resolution

Provide the complete PR description, including the omitted modeling/data-provenance section, and rerun this check.

  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

PeaBrane added a commit to PeaBrane/aisimulate-contrib that referenced this pull request Oct 27, 2026
Keep this PR focused on rejecting newly staged pulls from draining
sources and documenting the remaining modeled boundaries.

Remove the deadline shortcut, its lifecycle regression, the bookkeeping
TODO and performance documentation now owned by independent PR ai-dynamo#611.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Rudy Pei <yanrpei@gmail.com>
@jthomson04

Copy link
Copy Markdown
Contributor

/ok to test b5ecd52

@copy-pr-bot

copy-pr-bot Bot commented Oct 27, 2026

Copy link
Copy Markdown

/ok to test b5ecd52

@jthomson04, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

@jthomson04
jthomson04 enabled auto-merge (squash) October 27, 2026 16:59
@jthomson04

Copy link
Copy Markdown
Contributor

/ok to test 8304832

@jthomson04
jthomson04 merged commit ddbeb86 into ai-dynamo:main Oct 27, 2026
72 checks passed
carlory pushed a commit to carlory/aisimulate that referenced this pull request Oct 30, 2026
Avoid shared registry locks and deadline lookups when a rank has no
incoming pull. Retain deadlines for detached pulls until their blocks
are consumed, with lifecycle regression coverage.

Record the measured baseline bookkeeping cost and defer the remaining
transfer optimizations. This change is independent of the draining-source
staging fix in PR ai-dynamo#600.

Signed-off-by: Rudy Pei <yanrpei@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants