[milesial]

Why and what changed

With aligned state caching, a prompt that ends exactly on a block boundary can skip its reusable prompt checkpoint. The scheduler stops at the prompt boundary, but the cache publishes the final full block strictly before that boundary. Subsequent decoding then leaves no reusable prompt state.

Calculate the scheduler boundary from prompt.saturating_sub(1) to match the cache publisher. For a 16-token prompt and four-token blocks, a repeated request can reuse 12 tokens after the first request completes a long decode.

Review map

• Risk level: medium.
• Start with VllmCore::align_state_prefill and the new state-cache regression in crates/core/src/engine/scheduler/vllm/core.rs.
• Public or serialized contract changed: no.
• Compatibility or rollback concern: aligned prefill gains the missing checkpoint split. Token-only replay is unchanged.

Evidence

The regression completes a 16-token prompt with 16 output tokens, then submits the same prompt. It requires 12 reused tokens: floor((16 - 1) / 4) * 4. This covers checkpoint retention beyond the first decode block.

Fast CI, Full CI, and independent reviews are pending. This draft does not claim merge readiness.

AgentX KV reuse impact

With this fix, a lower-load AgentX replay reaches 95.38% token-weighted KV prefix reuse on one TP1 aggregated worker.
The run uses C8, which means eight active play lanes, with no speculation.
All requests use the same replica, so this result needs no session-pinning change.

The workload uses a 16-play AgentX sample with 4,908 source requests, seed 42, and a 3,600-second profile after warmup.
The worker uses 64 GiB of modeled cache, 64-token blocks, and aligned state caching with prefix_match_unit=64.

Single-replica TP1 replay, C8	Prefix match unit	Reused input tokens	Total input tokens	KV prefix reuse
Earlier run, before the fix	Unset	45,356,608	112,180,416	40.43%
Lower-load run with the fix	64	144,463,552	151,455,744	95.38%

The ratio divides reused input tokens by total input tokens across completed requests.
It measures token reuse rather than the fraction of requests with a cache hit.
The fixed run completes 1,578 requests with no cancellations or readmissions.

The C8 comparison changes both the boundary fix and the prefix match unit.
Its 54.95-percentage-point difference does not isolate the effect of this PR.
These duration-limited runs also complete different request counts.

For attribution, the separate matched comparison keeps eight TP1 workers, C64, round-robin routing, and prefix_match_unit=64 unchanged.
That comparison measures 0.44% reuse before the boundary fix and 58.37% after it, a gain of 57.93 percentage points.

Modeling or data provenance

The scheduler regression is synthetic. The AgentX results describe simulator cache behavior, with functional_only qualification and no matched silicon calibration.
No performance tables, timing equations, or model metadata change.

Tracking

PR #542 separately allows full-block state-cache alignment with speculation. Both PRs target main independently.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: ai-dynamo/aisimulate/.coderabbit.yaml
• Review profile: ASSERTIVE
• Plan: Enterprise
• Run ID: 403dea22-28ad-42c1-ad06-7c23b840d54a

📥 Commits

Reviewing files that changed from the base of the PR and between a0a35e7 and 89e0d07.

📒 Files selected for processing: 1

• crates/core/src/engine/scheduler/vllm/core.rs

🔗 Linked repositories identified

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

• ai-dynamo/dynamo (manual)
• ai-dynamo/aiconfigurator (manual)

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

📜 Recent review details

🧰 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: 5

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/core.rs

Read REVIEW.md before commenting.

⚙️ CodeRabbit configuration file

Files:

• crates/core/src/engine/scheduler/vllm/core.rs

Applicability reference: crates/core/src/engine/scheduler/vllm/AGENTS.md: CLAUDE.md Source excerpt: `core.rs` owns common queue mechanics, KV allocation, state transitions, completion, and event emission.

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

Files:

• crates/core/src/engine/scheduler/vllm/core.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:

• crates/core/src/engine/scheduler/vllm/core.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:

• crates/core/src/engine/scheduler/vllm/core.rs

🔀 Multi-repo context ai-dynamo/dynamo, ai-dynamo/aiconfigurator

Linked repositories findings

ai-dynamo/dynamo

• No downstream align_state_prefill, VllmCore, or state-prefill API consumer was found. The only nearby protocol field is the generic optional sequence_length, so this PR’s internal scheduler fix does not appear to require a Dynamo API update. [::ai-dynamo/dynamo::lib/llm/src/protocols/common.rs:795]

ai-dynamo/aiconfigurator

• AIC models scheduler capacity using token-sized blocks: it validates scheduler_block_size as a positive integer and returns floor-divided KV capacity. This is consistent with the PR’s block-boundary semantics, but contains no state-checkpoint behavior. [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/memory.py:1075-1123] [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/memory.py:1155-1159]
• kv_block_size remains an optional EngineConfig/public API field; no checkpoint or replay serialization contract references it. [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/engine.py:251-282] [::ai-dynamo/aiconfigurator::tests/cross_package/test_core_public_api.py:74-78]

🔇 Additional comments: 2

crates/core/src/engine/scheduler/vllm/core.rs: 2

2601-2602: LGTM!

3727-3751: LGTM!

📝 Summary

Risk level: Low to moderate.

Human review should focus on:

1. prompt.saturating_sub(1) boundary handling.
2. Checkpoint reuse after completed decode.
3. Interaction with speculative decoding and state-cache alignment.

align_state_prefill now computes last_full from the prompt length. This retains the replay boundary when the prompt ends on a full block. Existing tests cover partial hits, exact boundaries, residual tails, decode, small budgets, and token-only behavior.

No public or serialized contract changes are present. The change affects scheduler behavior only.

The supplied source supports the intended fix and focused regression coverage. Test execution, CI results, review findings, and merge readiness remain unavailable. The code appears technically focused, but merge readiness is not established.

align_state_prefill now derives the final full-block boundary from the prompt length. A regression test checks that an identical 16-token prompt reuses 12 input tokens after the first request completes.

Changes: Prefill Reuse

Layer / File(s)	Summary
Preserve the final full-block boundary crates/core/src/engine/scheduler/vllm/core.rs	align_state_prefill uses the prompt length to derive the final full-block boundary. The regression test checks that a second identical 16-token request reuses 12 input tokens.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~8 minutes

Merge Risk: ⚪ Minimal · up to 89e0d

The change preserves block-aligned prompt reuse without introducing an actionable merge-blocking risk.

Pre-merge checks: 7 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Review Evidence		The description gives boundary details and an expected result, but it does not name an exact test command or report a command result. It also does not distinguish local, fake-runner, parity, and produ…	Update the Evidence section with the exact local test command and result. Label any fake-runner, parity, and production evidence separately. Report Fast CI or Full CI only with the tested commit and actual result; otherwise mark them unavai…

Passed checks: 7 passed

Check name	Status	Explanation
Title check		The title precisely describes the behavioral change: retaining state checkpoints for prompts aligned to block boundaries.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Cross-Layer Contract		PASS. The pull request changes only the private VllmCore::align_state_prefill method and adds a Rust regression test in core.rs. The method is called only by internal scheduling, and VllmCore is…
Modeling And Data Evidence		The PR changes selection logic with documented provenance and reproducible numeric evidence. align_state_prefill now uses prompt.saturating_sub(1), matching the state-cache publisher’s replay boun…
Compatibility Boundaries		PASS: The PR changes only crates/core/src/engine/scheduler/vllm/core.rs. It updates the private align_state_prefill method and adds a #[cfg(test)] regression test. No Python files, schemas, bind…
Description check		The description follows the required sections and clearly states the problem, implementation, risk, compatibility impact, regression evidence, expected-value derivation, modeling provenance, and relat…

Full details: Cross-Layer Contract

Explanation

PASS. The pull request changes only the private VllmCore::align_state_prefill method and adds a Rust regression test in core.rs. The method is called only by internal scheduling, and VllmCore is pub(crate) and re-exported internally. No public input, schema field, default, serialized field, CLI contract, or output definition changes. The new aligned-prompt path is tested with a 16-token prompt and a required 12-token reuse assertion. No stale cross-layer consumer is introduced.

Full details: Modeling And Data Evidence

Explanation

The PR changes selection logic with documented provenance and reproducible numeric evidence. align_state_prefill now uses prompt.saturating_sub(1), matching the state-cache publisher’s replay boundary in vllm_backend.rs. The added regression submits the same 16-token prompt after a long decode and asserts the explained golden value of 12 reused tokens (floor((16 - 1) / 4) * 4). This is concrete validation, not binary-only reporting.

Full details: Compatibility Boundaries

Explanation

PASS: The PR changes only crates/core/src/engine/scheduler/vllm/core.rs. It updates the private align_state_prefill method and adds a #[cfg(test)] regression test. No Python files, schemas, bindings, version metadata, packaging files, or documented defaults changed. The one-wheel/one-crate artifact contract and AISimulate separation are therefore unaffected.

Full details: Description check

Explanation

The description follows the required sections and clearly states the problem, implementation, risk, compatibility impact, regression evidence, expected-value derivation, modeling provenance, and related work. It also identifies pending CI and review status. Exact test commands, commit identifiers, and before/after output are not provided, but these omissions do not make the description incomplete or off-topic.

Full details: Review Evidence

Explanation

The description gives boundary details and an expected result, but it does not name an exact test command or report a command result. It also does not distinguish local, fake-runner, parity, and production evidence. Fast CI and Full CI are marked pending, not reported as results. The repository confirms the PR template requires exact commands and results ('.github/pull_request_template.md:23-37'). No inactive nested workflow is cited as hosted CI evidence.

Resolution

Update the Evidence section with the exact local test command and result. Label any fake-runner, parity, and production evidence separately. Report Fast CI or Full CI only with the tested commit and actual result; otherwise mark them unavailable or pending without presenting them as evidence.

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

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

[tedzhouhk]

Reviewed 89e0d074. The focused native state-cache suite passed (68 tests). Approving with the following backend-fidelity finding still outstanding; this approval does not establish matched-vLLM accuracy.

[P1] At core.rs:2602, prompt=16/block=4 without SD now splits prefill into 12+4 and retains state12. The pinned vLLM splitter executes 16; its allocator does not materialize state12, even with internal-checkpoint capability. Executing those upstream methods gives zero reuse for an identical prompt, versus 12 here. This changes both checkpoint availability and prefill iterations, not only bookkeeping.

Please distinguish this simulator checkpoint policy from the pinned serving behavior, or tie it to a validated serving-side fix/version. Preserve the functional-only qualification. The upstream check used actual methods with minimal metadata stubs, not a full vLLM/GPU run.

[dreamtalen]

Checked the fidelity change: upstream vLLM fixed this boundary in vllm-project/vllm#60533. The Mamba align split now stops at num_tokens - 1 rounded down to the block, so a block-aligned prompt materializes its replay-boundary checkpoint. This PR follows the fixed upstream behavior. LGTM.