perf(router): reduce scheduler and dispatch overhead

[jthomson04]

Summary

Reduce scheduler and dispatch work for requests served by large worker pools, with admission queueing disabled, enabled, or mixed across policy classes. Reuse load-projection storage, apply lifecycle state updates directly, and clone only the selected discovered instance during transport lookup.

• Reuse admission projection storage. Refresh values for each selection and retain actor-owned storage through booking and response capture, including failures. Apply attempt-checked output-block updates directly for frontend and standalone callers; output updates do not drain queues.
• Apply lifecycle updates directly and combine capacity notifications. Apply normal prefill completion and explicit release through the existing sequence-state methods. For queued schedulers, record affected workers in a bounded 256-entry buffer; overflow requests a full recheck. The actor combines notifications into one drain pass and services them alongside admission commands and deadlines, including when the admission channel is full. Before an enqueue drains pending admissions, it consumes capacity hints and rechecks worker lanes so an older blocked pinned request can compete with a later shared arrival. Capacity wakes retain the cleanup gate.
• Clone only the selected instance during dispatch. Incorporate the final change from perf(runtime): clone only the selected instance during dispatch #14545 as a separate commit. Client::instance_by_id() borrows the current discovery snapshot and clones only the matching instance. PushRouter::resolve_transport() uses that lookup and preserves worker-removal, fallback, and error behavior.

Selection, booking, queue order, fairness, and overlap refresh remain inside the actor. Attempt checks, replica publication, reservation cleanup, and booking ownership remain intact. Destructor cleanup, expiry, admission handoff, and failed-admission rollback retain their actor paths. Successful lifecycle calls confirm the state change and notification. LocalScheduler::update_queue() waits for an explicit recheck when queueing is enabled and the actor is running; with queueing disabled it returns immediately without an admission barrier.

No new dependencies, tasks, runtime settings, HTTP/Python interfaces, or wire changes. The current base already includes the numeric-map optimization and builtin worker-selection refactor. This PR preserves that implementation; it does not retain changes that only tuned the old reference selector.

Reviewer entry points

• lib/kv-router/src/scheduling/local.rs: direct lifecycle updates and update_queue() timing.
• lib/kv-router/src/scheduling/queue.rs: serialized selection/booking, reusable projection storage, and capacity rechecks before draining admissions.
• lib/kv-router/src/scheduling/queue/capacity.rs: bounded worker hints, overflow fallback, and retained wakeups.
• lib/runtime/src/component/client.rs and lib/runtime/src/pipeline/network/egress/push_router.rs: selected-instance lookup and transport resolution.
• Existing tests cover capacity ordering, overflow, cleanup progress, projection reuse, stale attempts, queue-off updates, and discovery lookup updates/removal.

Validation

Rebased onto main 0fac8e85d0c1a27bbd6b6d75d386b757e7f6f322. Current combined head: 3cce7cb547.

• KV-router library tests with standalone-selection,bench: 1,286 passed, 2 ignored.
• Runtime selected-instance lookup test: 1 passed; runtime dispatch tests: 40 passed, including worker removal and fallback behavior.
• Frontend router tests (dynamo-llm --no-default-features --lib kv_router::): 389 passed.
• Online replay tests: 32 passed.
• KV-router all-target Clippy and runtime library Clippy passed with warnings denied. Formatting and patch whitespace checks passed.
• All four commits have DCO sign-offs. The new runtime commit has identical added/deleted lines to the final perf(runtime): clone only the selected instance during dispatch #14545 patch.

No new performance benchmark was run for this rebase or the combined change. The results below describe earlier scheduler revisions, before the instance-lookup change and the current main refactor. They do not measure the current combined head.

Review-fix benchmark on the same base

Control 7c21112254 versus candidate 55d56abf0c, both on ae5ca99fba. Five Grace nodes, 16 workers with actual backlog, AgentX, concurrency 6,144, 120-second load, identical warmup/drain/policy and frozen mock-worker/client binaries. Separate frontend builds, verified live binary identities, and fresh processes each run.

One control/candidate pair, followed by the permitted reverse-order pair to check the initial CPU increase:

Run order	Version	Completions/s	TTFT p95 (ms)	Post-first-token p95 (s)	CPU ms/completion
1	control	865.1	354.8	35.19	21.433
2	candidate	965.6	325.1	29.43	21.865
3	candidate-repeat	862.1	414.3	30.92	21.465
4	control-repeat	641.3	1092.2	44.78	20.821

The candidate had higher throughput and lower TTFT and post-first-token latency in both pairs. CPU per completion was higher in both pairs: +2.0% and +3.1%. This is a measured tradeoff, not a CPU saving.

The size of a sustained speedup and the cause of the CPU increase remain inconclusive. The two control runs differed by 25.9% in throughput. Request size and worker distribution also changed. Those observations do not prove the cause of the differences, and no slowdown is assigned to an external cause.

All 426,472 records have complete client coverage, with zero errors. All 16 workers were used in every run; long-queue peaks ranged from 1,365 to 1,503 per frontend. Cleanup passed on all five nodes after each run. Output tracking was disabled; no cancellation-heavy benchmark was added. The setup attempt that stopped before load is excluded from these results.

Earlier optimization benchmark evidence

The benchmark used base 4067d08f, step 1 2c483e0059, and both steps da44ae41e1, before the rebases onto #14149 and main. These are measurements of the original implementation, not a new benchmark of the rebased head.

Matched isolated builds and fresh processes on five Grace nodes: two frontend processes, AgentX workload, 2,048 workers with queueing off, and 16 workers with actual backlog for queueing on/mixed. One clean run per build/mode, separate diagnostic runs, and one extra queue-on baseline/full pair to check the initial regression. Binary identities and complete client-record coverage were verified; rare transport errors were accepted and retained in the artifacts.

The diagnostic runs show 11–27% less actor service time per enqueue on both frontends across all three modes, and roughly half the admission-channel traffic. Actor service is elapsed processing time, including overlap-refresh waits, rather than CPU time.

Queue mode	Base requests/s	Step 1 requests/s	Both steps requests/s
Off	5,197.9	5,475.4	5,681.8
On, first comparison	650.3	780.8	472.9
Mixed	175.2	539.5	962.5
On, additional control pair	729.0	—	943.2

With queueing off, both steps improved throughput by 9.3%, reduced TTFT p95 by 62.7%, and reduced CPU per completed request by 1.9%; post-first-token p95 increased by 10.5%. The first queue-on slowdown did not repeat in the control pair. The results support the reduction in scheduler work, but run variation prevents a claim of a stable throughput gain in all modes or a confirmed external cause for the first slowdown. Output tracking was disabled in these workloads, so its direct-update performance was not measured.

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-01a0b691-fabb-75cb-b29b-8109ea2bd9f6.docs.buildwithfern.com/dynamo/dev

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: ai-dynamo/dynamo/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 309c07fe-bd7a-4a90-9584-65382a4fc1e7

📥 Commits

Reviewing files that changed from the base of the PR and between ae5ca99 and 7c21112.

📒 Files selected for processing (13)

• lib/kv-router/benches/worker_selection.rs
• lib/kv-router/src/scheduling/local.rs
• lib/kv-router/src/scheduling/overlap.rs
• lib/kv-router/src/scheduling/overlap_refresh.rs
• lib/kv-router/src/scheduling/queue.rs
• lib/kv-router/src/scheduling/queue/capacity.rs
• lib/kv-router/src/scheduling/selector/default.rs
• lib/kv-router/src/scheduling/selector/mod.rs
• lib/kv-router/src/scheduling/types.rs
• lib/kv-router/src/sequences/multi_worker.rs
• lib/kv-router/src/sequences/prompt_registry.rs
• lib/kv-router/src/services/selection/core/tests.rs
• lib/llm/src/kv_router.rs

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

---

Walkthrough

The scheduler now applies lifecycle changes synchronously, coalesces capacity rechecks, and reuses worker-load projection maps. Overlap data uses FxHashMap. Tests and benchmarks cover direct release, projection reuse, larger worker sets, and scoring results.

Changes

Scheduler capacity and projection

Layer / File(s)	Summary
Projection and overlap contracts lib/kv-router/src/scheduling/overlap.rs, lib/kv-router/src/sequences/..., lib/llm/src/kv_router.rs, lib/kv-router/src/scheduling/local.rs	Overlap signals use FxHashMap. Worker-load projection and cached-token data now flow through caller-provided maps.
Direct lifecycle updates and capacity notifications lib/kv-router/src/scheduling/queue/capacity.rs, lib/kv-router/src/scheduling/queue.rs, lib/kv-router/src/scheduling/local.rs, lib/kv-router/src/services/selection/core/tests.rs	Lifecycle mutations and booking release update capacity directly. CapacityUpdates coalesces targeted or full queue rechecks.
Reusable projected-load selection lib/kv-router/src/scheduling/queue.rs	Selection and admission reuse actor-owned projected-load storage across success and error paths.
Performance fixtures and scheduler validation lib/kv-router/benches/worker_selection.rs, lib/kv-router/src/scheduling/selector/..., lib/kv-router/src/scheduling/types.rs, lib/kv-router/src/scheduling/overlap_refresh.rs	Benchmarks cover projection and 2,048 workers. Fixtures use FxHashMap, and scoring assertions include worker logits and selection metadata.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Merge Risk: ⚪ Minimal · up to 7c211

No actionable correctness, availability, or integration risk was established, so the PR is ready for normal merge checks.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.61% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 92 functions across 13 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The description provides detailed overview, implementation details, reviewer entry points, and validation results. It does not include the required Related Issues section or confirm that no related is…	Add the required Related Issues section and select one path: link the applicable issue with a Closes or Relates to entry, or confirm that this PR is not linked to an issue.

✅ Passed checks (3 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly and concisely summarizes the primary change: reducing scheduler and dispatch overhead.

Full details: Description check

Explanation

The description provides detailed overview, implementation details, reviewer entry points, and validation results. It does not include the required Related Issues section or confirm that no related issue exists.

• Fix all pre-merge checks with AI

---

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

[jthomson04]

Code review

Found 15 issues (each expanded in an inline comment with its failure scenario and suggested fix). Reviewed at head 7c211122549949e84f0742a29f7ed1519a945f17 against base ae5ca99fba. Comments 11 and 12 are anchored to the nearest changed line because the flagged line is outside the diff.

1. Admission ordering: the invariant that let handle_enqueued skip the blocked-lane recheck no longer holds

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 983 to 987 in 7c21112

	} else {
	tokio::select! {
	command = rx.recv() => command,
	_ = async {
	if let Some(updates) = &capacity_updates {

2. ensure_running() gates a mutation that no longer needs the actor, making lifecycle accounting inconsistent if the actor dies

dynamo/lib/kv-router/src/scheduling/local.rs

Lines 642 to 644 in 7c21112

	) -> Result<LifecycleMutationOutcome, KvSchedulerError> {
	self.queue.ensure_running()?;
	let outcome = self

3. Output-block updates now run the registry write path on every request task at every block boundary, and the PR states this was not benchmarked

dynamo/lib/kv-router/src/scheduling/local.rs

Lines 708 to 710 in 7c21112

	) -> Result<(), KvSchedulerError> {
	self.queue.ensure_running()?;
	self.add_output_block_if_booking_sync(booking, decay_fraction)

4. handle_capacity_updates has no has_ready() early return, unlike handle_update/handle_enqueued

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 1313 to 1316 in 7c21112

	async fn handle_capacity_updates(&mut self) {
	self.reject_expired(Instant::now());
	let Some(updates) = &self.capacity_updates else {

5. drain_cleanup() runs unconditionally on every capacity wake, bypassing the burst gate the command path uses

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 1319 to 1321 in 7c21112

	let mut recheck_all = updates.drain_into(&mut self.capacity_workers);
	recheck_all |= self.drain_cleanup();
	{

6. Unknown-worker -> recheck_all_workers is a symptom patch; worker removal is not handled at its source

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 1323 to 1331 in 7c21112

	for &worker in &self.capacity_workers {
	let known = configs.get(&worker.worker_id).is_some_and(|config| {
	let start = config.data_parallel_start_rank();
	(start..start.saturating_add(config.data_parallel_size()))
	.contains(&worker.dp_rank)
	});
	if !known {
	recheck_all = true;
	break;

7. Reuse worker_config_for_rank instead of re-implementing the dp-rank range check

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 1323 to 1328 in 7c21112

	for &worker in &self.capacity_workers {
	let known = configs.get(&worker.worker_id).is_some_and(|config| {
	let start = config.data_parallel_start_rank();
	(start..start.saturating_add(config.data_parallel_size()))
	.contains(&worker.dp_rank)
	});

8. lib/kv-router/src/sequences/AGENTS.md requires before/after benchmark numbers for prompt-registry projection changes

dynamo/lib/kv-router/src/sequences/prompt_registry.rs

Lines 301 to 308 in 7c21112

	pub(super) fn project_worker_loads_into(
	&self,
	token_sequence: Option<&[SequenceHash]>,
	decay_now: Instant,
	projections: &mut FxHashMap<WorkerWithDpRank, WorkerLoadProjection>,
	) {
	projections.clear();

9. The legacy schedule* path now allocates two FxHashMaps per call where it previously moved the caller's maps

dynamo/lib/kv-router/src/scheduling/local.rs

Lines 484 to 486 in 7c21112

	tier_overlap_blocks,
	effective_overlap_blocks: effective_overlap_blocks.into_iter().collect(),
	effective_cached_tokens: effective_cached_tokens.into_iter().collect(),

10. Dead indirection: three names now do the identical synchronous thing

dynamo/lib/kv-router/src/scheduling/local.rs

Lines 732 to 741 in 7c21112

	#[doc(hidden)]
	pub async fn enqueue_output_block_if_booking(
	&self,
	booking: &SchedulerBookingDescriptor,
	decay_fraction: Option<f64>,
	) -> Result<(), KvSchedulerError> {
	self.add_output_block_if_booking(booking, decay_fraction)
	.await
	}

11. Redundant state: capacity_updates.is_some() and queueing_enabled (line 519) encode the same fact

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 506 to 508 in 7c21112

	cleanup: Arc<AdmissionCleanup>,
	capacity_updates: Option<Arc<CapacityUpdates>>,
	direct_release: DirectBookingRelease,

12. Update { worker: Some(..) } is now only reachable from #[cfg(test)] update_worker

dynamo/lib/kv-router/src/scheduling/queue.rs

Lines 155 to 158 in 7c21112

	Update {
	worker: Option<WorkerWithDpRank>,
	ack_tx: oneshot::Sender<()>,
	},

13. update_queue() is documented as waiting for a recheck but is a no-op when queueing is disabled, and has no production caller

dynamo/lib/kv-router/src/scheduling/local.rs

Lines 668 to 672 in 7c21112

	/// Wait for an actor recheck of pending admission. This does not wait for
	/// requests blocked on worker capacity to finish or for the queue to empty.
	pub async fn update_queue(&self) {
	self.queue.update().await;
	}

14. #[inline(always)] added to already-inlinable private/generic fns with no selection-bench delta cited

dynamo/lib/kv-router/src/scheduling/selector/default.rs

Lines 253 to 255 in 7c21112

	#[inline(always)]
	fn default_row(

15. The reused projection table never shrinks; clear() is O(bucket count)

dynamo/lib/kv-router/src/sequences/prompt_registry.rs

Lines 307 to 310 in 7c21112

	) {
	projections.clear();
	let query_len = token_sequence.map_or(0, |query| query.len());
	let matched_depth = self.membership.compute_overlap_depths(token_sequence);

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[copy-pr-bot]

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.

[jthomson04]

/ok to test f1ffe12