feat(kv-router): expose prefix hashes, policy class, and worker capacity to selection policies by PeaBrane · Pull Request #15449 · ai-dynamo/dynamo

PeaBrane · 2026-09-30T22:18:12Z

Summary

Adds three read-only inputs to the public worker-selection plugin API
(plugins::worker_selection). Porting recently published routing policies onto the plugin API
showed they needed inputs a plugin could not get. The stacked follow-up PR ports those policies.

WorkerSelectionContext::prefix_hashes() borrows the prompt's chained per-block tracking
hashes. Entry i identifies blocks 0..=i, so a policy can key on a shared prefix for hot-spot
detection, consistent hashing, or two-choice hashing.
- It is None where the host tracks no active blocks, such as disaggregated prefill pools.
- It is unique to the request where the host assumes no KV reuse, such as disaggregated
  decode pools.
WorkerSelectionContext::policy_class() returns the request's policy class. A request
classifier's set_policy_class override is now recorded on the request after validation, so
selection sees the classifier's decision rather than only the caller's value. Nothing
downstream of classification read the field before this change.
WorkerSelectionContext::worker_capacity(worker) looks up the worker's advertised
total_kv_blocks and max_num_batched_tokens. It maps the zero that some backends advertise
for "unknown" to None.

worker_capacity is a lookup into the worker table the host already holds, not a new
WorkerInputs group. A first version carried capacity beside every candidate row and slowed
the default policy by about 2% at 1,024 workers, even though the default never reads capacity.
The lookup leaves candidate rows, views, and picker columns unchanged.

Reviewer guide:

Accessors: plugins/worker_selection/context.rs and inputs.rs (WorkerCapacityInput).
Host wiring: scheduling/selector/mod.rs.
Classifier write-back: scheduling/queue.rs.
Contract test: scheduling/selector/policy.rs.
Docs: custom-worker-selection.mdx.

Validation

Default selection cost is unchanged. Criterion benches compared origin/main against this
PR, interleaved as baseline, candidate, baseline, candidate on one pinned core. The builtin
default policy (default_selection/t0/plugin), baseline vs this PR, per selection:

Workers origin/main This PR

64 2.89–2.92 µs 2.91–2.92 µs

1,024 53.5–54.3 µs 54.0–54.6 µs

4,096 266–268 µs 259–275 µs

kv-router's composed-policy benches (custom_worker_selection, unused_preferred_taint_metadata)
are equal or faster.
Replay routing is byte-identical. Offline replay canonical reports match origin/main byte
for byte, in two fresh processes each, across four configurations: vLLM and SGLang semantics,
each aggregated and disaggregated. Each ran the first 5,000 Mooncake trace rows on the
replay-bench seeded router. The four digests are the same as the baseline's.
- Lifecycle coverage matches the baseline: 5,000 completions per row, 9 vLLM preemptions,
  and 10 SGLang retractions.
- These were measured on the top of the stack. Its kv-router code is identical to this PR's.
Tests:
- New contract test request_facts_and_worker_capacity_reach_policy_components exercises all
  three accessors through a scorer and a picker, including zero-as-unknown capacity and an
  unknown worker.
- classifier_overrides_are_validated_before_enqueue now asserts the write-back.
Suites and lint: cargo test -p dynamo-kv-router --lib passes (1,032 tests). This PR on
its own also passes the builtin and mocker lib tests. Workspace clippy is clean.

🤖 Generated with Claude Code

Summary by CodeRabbit

New Features
- Worker-selection policies can now access prompt-block prefix hashes, the request’s policy class, and advertised worker capacities, including KV-block and batched-token limits. Unavailable or unknown values are reported as unavailable.
- The selected policy class is now available to worker selection.

…policies Worker-selection plugins could not key on a shared prompt prefix, see the class a request classifier assigned, or read a worker's advertised capacity. - WorkerSelectionContext::prefix_hashes() borrows the request's chained per-block tracking hashes. - WorkerSelectionContext::policy_class() returns the requested class. A classifier's override is now recorded on the request before selection. - WorkerSelectionContext::worker_capacity() looks up a worker's advertised total_kv_blocks and max_num_batched_tokens, mapping the zero some backends advertise for "unknown" to None. Capacity is a lookup rather than a WorkerInputs group: carrying it alongside every candidate row slowed the default policy by about 2% at 1,024 workers, although the default never reads it. Signed-off-by: PeaBrane <yanrpei@gmail.com>

github-actions · 2026-09-30T22:23:34Z

🌿 Fern Docs Preview: https://nvidia-preview-01a0ff67-5e6c-737f-9ce0-8968b047e351.docs.buildwithfern.com/dynamo/dev

coderabbitai · 2026-09-30T22:24:34Z

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Note

Repository guideline files applied to this review (6)

lib/kv-router/src/scheduling/AGENTS.md — auto-discovered
docs/fern/AGENTS.md — auto-discovered
AGENTS.md — auto-discovered
docs/fern/pages/AGENTS.md — auto-discovered
docs/fern/pages/community/contributing/documentation/documentation-style-guide.md — configured
lib/kv-router/AGENTS.md — auto-discovered

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: c2259dde-58b9-4059-89a6-f6c020022d3c

📥 Commits

Reviewing files that changed from the base of the PR and between 9dc2102 and db0d3c5.

📒 Files selected for processing (8)

docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx
lib/kv-router/src/plugins/worker_selection.rs
lib/kv-router/src/plugins/worker_selection/context.rs
lib/kv-router/src/plugins/worker_selection/inputs.rs
lib/kv-router/src/scheduling/queue.rs
lib/kv-router/src/scheduling/selector/mod.rs
lib/kv-router/src/scheduling/selector/policy.rs
lib/kv-router/src/scheduling/selector/reference.rs

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

Walkthrough

Worker-selection contexts now expose request prefix hashes, policy class, and worker capacity. Selector materialization supplies capacity lookups from worker configuration, and classifier validation stores the selected policy class on the request.

Changes

Worker Selection Context

Layer / File(s)	Summary
Expose request and worker capacity data `lib/kv-router/src/plugins/worker_selection/inputs.rs`, `lib/kv-router/src/plugins/worker_selection.rs`, `lib/kv-router/src/plugins/worker_selection/context.rs`, `docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx`	Adds `WorkerCapacityInput` and exports it. `WorkerSelectionContext` exposes prefix hashes, policy class, and worker capacity. The documentation describes these accessors.
Connect classifier and selector data `lib/kv-router/src/scheduling/queue.rs`, `lib/kv-router/src/scheduling/selector/mod.rs`, `lib/kv-router/src/scheduling/selector/policy.rs`, `lib/kv-router/src/scheduling/selector/reference.rs`	Classification stores the selected policy class on the request. Selector materialization provides a capacity lookup from configured workers. Tests check the exposed request data, worker capacity, and updated selector inputs.

Priority: ⬇️ Low

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

Merge Risk: ⚪ Minimal · up to db0d3

No merge-blocking issue is established for these read-only policy inputs. The change appears mergeable subject to normal checks.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description provides detailed scope, implementation guidance, validation results, and reviewer guidance. However, it omits the required Related Issues section and does not confirm whether the pull…	Add the required Related Issues section. Either link the relevant issue with a closing or reference statement, or check the confirmation that this pull request is not linked to an issue.
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 7 files. (1 skipped: …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: exposing prefix hashes, policy class, and worker capacity to selection policies.
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.

Full details: Description check

Explanation

The description provides detailed scope, implementation guidance, validation results, and reviewer guidance. However, it omits the required Related Issues section and does not confirm whether the pull request has no related issue.

Full details: Docstring Coverage

Explanation

Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 7 files. (1 skipped: 1 unsupported.)

Fix all pre-merge checks with AI

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

Autopilot is currently an internal CodeRabbit preview.

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

…acity Address review on the policy-input API: - Bound the capacity lookup by Sync so WorkerSelectionContext stays Send and Sync for plugins that score candidates in parallel. Every host config is already Sync. - Drop WorkerCapacityInput::max_num_batched_tokens until a policy needs it. - Document that the standalone selection service passes hash-only prefix hashes through unchecked and that policy_class is not validated. - Assert the contract test's scorer actually runs. Signed-off-by: PeaBrane <yanrpei@gmail.com>

devin-ai-integration · 2026-10-01T03:14:37Z

⏹️ Dynamo PR CI cancelled — run 37086387586 (attempt 1) on `2c41b73ce6`

Gate checks: ✅ backend-status-check · ✅ deploy-status-check · ❌ dynamo-status-check

Other	Jobs
dynamo-runtime	⏹️ 1 ✅ 7

Failure details

No jobs failed. The 1 cancelled job, dynamo-runtime / rust-gpu, was not a fail-fast or manual cancel: it hit the job's 30-minute execution limit (The job has exceeded the maximum execution time of 30m0s) during step Run Rust checks (block-manager + media-ffmpeg + MM routing + integration tests). When it was stopped, it was still running the kvbm_physical unit tests and every test so far had passed. Rerunning the failed jobs is likely enough. If it times out again, the Rust GPU test step is too slow for its 30-minute limit.

For agents

{"pr": 15449, "run_id": 37086387586, "run_attempt": 1, "head_sha": "2c41b73ce6f0a97e69afddc9f39cd7140ecfe935",
 "failures": [{"job": "dynamo-runtime / rust-gpu", "job_id": 111103285326, "conclusion": "cancelled", "failed_step": "Run Rust checks (block-manager + media-ffmpeg + MM routing + integration tests)", "signature": "job exceeded maximum execution time of 30m0s", "tests": [], "log_cmd": "gh api repos/ai-dynamo/dynamo/actions/jobs/111103285326/logs"}]}

_{Posted automatically by Devin for run 37086387586. Updated on every full-CI run of this PR.}

Signed-off-by: PeaBrane <yanrpei@gmail.com>

ishandhanani · 2026-10-02T20:32:20Z

 | `session_context()` | Optional session metadata described below |
+| `prefix_hashes()` | Optional chained hashes of the prompt's complete KV blocks; entry `i` identifies blocks `0..=i`, so equal entries mean an exactly shared prefix. `None` when the frontend does not track active blocks, as in disaggregated prefill pools; unique to the request when the host does not assume KV reuse, as in disaggregated decode pools. The standalone selection service passes a hash-only request's precomputed hashes through unchecked |
+| `policy_class()` | Optional policy class: the caller's value, replaced by a request classifier's override. A caller's value is not validated against the policy profile |
+| `worker_capacity(worker)` | Capacity advertised in a worker's runtime config: `total_kv_blocks()` per data-parallel rank, `None` when not advertised or advertised as zero. `None` for a worker the host does not know. Each call looks the worker up, so policies that never call it pay nothing |


I think there should be a cleaner way to access the MDC of each worker if a cost function wants it

PeaBrane requested review from a team as code owners September 30, 2026 22:18

pull-request-size Bot added the size/L label Sep 30, 2026

github-actions Bot added feat documentation Improvements or additions to documentation labels Sep 30, 2026

PeaBrane added this pull request to stack #15451 September 30, 2026 22:18

dynamo-review-agent Bot reviewed Sep 30, 2026

View reviewed changes

Comment thread lib/kv-router/src/scheduling/selector/mod.rs Resolved

dynamo-review-agent Bot reviewed Sep 30, 2026

View reviewed changes

Comment thread lib/kv-router/src/plugins/worker_selection/inputs.rs Resolved

Comment thread lib/kv-router/src/plugins/worker_selection/inputs.rs Resolved

Comment thread lib/kv-router/src/scheduling/selector/policy.rs Resolved

Merge branch 'main' into rupei/router-policy-api

2c41b73

Signed-off-by: PeaBrane <yanrpei@gmail.com>

ishandhanani reviewed Oct 3, 2026

View reviewed changes

ishandhanani mentioned this pull request Oct 3, 2026

feat(router): add experimental KV fetch-hint policy hook #15599

Open

aylonsk mentioned this pull request Oct 10, 2026

feat(kv-router): expose per-rank resident block counts to selection plugins #16078

Open

ydaiming mentioned this pull request Oct 15, 2026

[FEATURE]: KV router should not select a worker that is at its in-flight request limit #16532

Open

biswapanda mentioned this pull request Oct 15, 2026

fix(router): exclude ranks at their request capacity #16605

Open

kthui mentioned this pull request Oct 30, 2026

refactor(router): flatten policy class configuration #17723

Merged

1 task done

john-liu-dev mentioned this pull request Nov 5, 2026

feat(router): add a KV-capacity-aware decode worker-selection policy #18224

Draft

1 task

Navigation Menu

Repository navigation

feat(kv-router): expose prefix hashes, policy class, and worker capacity to selection policies - #15449

feat(kv-router): expose prefix hashes, policy class, and worker capacity to selection policies#15449
PeaBrane wants to merge 3 commits into
mainfrom
rupei/router-policy-api

PeaBrane commented Sep 30, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

github-actions Bot commented Sep 30, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Sep 30, 2026

❌ Failed checks (2 warnings)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

devin-ai-integration Bot commented Oct 1, 2026 •

edited

Loading

Uh oh!

ishandhanani Oct 2, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Workers	`origin/main`	This PR
64	2.89–2.92 µs	2.91–2.92 µs
1,024	53.5–54.3 µs	54.0–54.6 µs
4,096	266–268 µs	259–275 µs

Conversation

PeaBrane commented Sep 30, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Summary by CodeRabbit

Uh oh!

github-actions Bot commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Sep 30, 2026

Walkthrough

Changes

❌ Failed checks (2 warnings)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

devin-ai-integration Bot commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

⏹️ Dynamo PR CI cancelled — run 37086387586 (attempt 1) on 2c41b73ce6

Failure details

Uh oh!

ishandhanani Oct 2, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

PeaBrane commented Sep 30, 2026 •

edited by coderabbitai Bot

Loading

github-actions Bot commented Sep 30, 2026 •

edited

Loading

devin-ai-integration Bot commented Oct 1, 2026 •

edited

Loading

⏹️ Dynamo PR CI cancelled — run 37086387586 (attempt 1) on `2c41b73ce6`