Skip to content

feat(router): pluggable KVHints policy interface with session awareness - #15679

Draft
karen-sy wants to merge 1 commit into
mainfrom
karenc/router-pluggable-kvhintpolicy
Draft

karen-sy wants to merge 1 commit into
mainfrom
karenc/router-pluggable-kvhintpolicy

Conversation

@karen-sy

@karen-sy karen-sy commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Overview:

This PR extends the existing RouterPluginRegistry with a configurable KvHintPolicy hook that formulates request-attached KV-hint actions after Dynamo chooses the worker for a request. This gives custom router policies read-only access to session-aware block lineage without changing worker selection or coupling SessionPrefixIndexer to any specific action or execution policy.

Details:

  • Adds a KvHintPolicy contract whose synchronous formulate() callback receives the request ID, canonical session context, and selected worker/DP rank.
  • Exposes read-only SessionPrefixIndexer queries through KvHintPolicyContext, including worker-qualified session frontiers and root-to-frontier block-hash paths.
  • Reuses the existing RouterPluginRegistry: statically linked policies register one constructor and are selected through router_policy_config using kv_hint_policy: {type, parameters}.
  • Runs the policy after worker selection and appends its actions to any envelope already produced during selection.
  • Fails open on policy errors or panics, preserving request dispatch without the policy's actions.

User example

Dynamo exposes KvHintPolicy as the policy interface and KvHintPolicyContext as its read-only view of the request, session, selected worker, and session lineage. A user implements the policy, registers it under a name with the router's existing RouterPluginRegistry, and returns the KvHintAction values to attach to matching requests. The existing router_policy_config surface was introduced in #10744, and #12536 established the registered-policy selection pattern reused here:

struct RetainSession;

impl KvHintPolicy for RetainSession {
    fn formulate(
        &self,
        context: &KvHintPolicyContext<'_>,
    ) -> Result<Vec<KvHintAction>, Box<KvHintPolicyError>> {
        let Some(session) = context.session_context() else {
            return Ok(Vec::new());
        };
        let Some(paths) = context.get_session_block_lineage(
            session.session_id(),
            context.selected_worker(),
            None,
        )? else {
            return Ok(Vec::new());
        };
        let block_hashes = paths.into_iter().flatten().collect::<Vec<_>>();

        Ok(vec![KvHintAction::new(
            format!("retain:{}", context.request_id()),
            "kv.retain",
            "1.0",
            BTreeMap::from([("block_hashes".into(), serde_json::json!(block_hashes))]),
        )])
    }
}

registry.register_kv_hint_policy(
    "retain-session",
    Arc::new(|_, _, _, _| Ok(Box::new(RetainSession))),
)?;

The registration makes retain-session available to the router. Select it in router_policy_config with:

kv_hint_policy:
  type: retain-session
  parameters: {}

The kv.retain action and payload above are illustrative; the engine consumer defines the supported action semantics.

System-level view

flowchart TB
    subgraph R["Dynamo router"]
        direction TB
        REQ["Session-tagged request"]
        IDX["G1/G2<br/>residency indexes"]
        SEL["Choose worker<br/>and DP rank"]
        SPI["SessionPrefixIndexer"]
        CTX["Policy context<br/>request + session + chosen worker"]
        POL["KvHintPolicy<br/>formulate actions"]
        ATTACH["Attach policy actions<br/>to routed request"]

        REQ --> SEL
        IDX --> SEL
        SEL -. "enqueue matched prefix" .-> SPI
        SEL --> CTX
        CTX --> POL
        SPI -->|"read committed lineage"| POL
        POL -->|"KvHintAction list"| ATTACH
    end

    ATTACH -->|"request + KvHint"| ENG["Chosen<br/>inference worker"]
    ENG -->|"canonical KV events"| IDX
    ENG -->|"ordered session updates"| SPI
Loading

Request ordering

  1. The existing router indexes and worker-selection policy choose a worker and DP rank.
  2. Existing route-time session tracking enqueues the selected worker's matched prefix for SessionPrefixIndexer.
  3. KvHintPolicy::formulate() runs with the selected worker and may query committed session lineage for that worker.
  4. compose_kv_hint() appends policy actions to any envelope already produced during selection. Otherwise, the policy actions form a new envelope.
  5. Dynamo dispatches the request and the combined envelope to the selected engine worker.

Session-index updates are ordered but asynchronous. A policy query observes committed state and must not assume that blocks first associated with or created by the current request are already indexed.

Boundaries

  • This PR does not implement a concrete KV policy or define new KV-hint action semantics.
  • KvHintPolicy does not participate in worker selection; it receives the already selected worker.
  • Hints remain request-attached. This does not add an asynchronous control channel.
  • Session queries return no index result when SessionPrefixIndexer is disabled.
  • Policy callbacks are synchronous and may execute concurrently across requests, so implementations must return promptly and synchronize their own mutable state.

Where should the reviewer start?

  1. lib/kv-router/src/plugins/kv_hint.rs — public policy contract and session-query API.
  2. lib/kv-router/src/plugins/kv_hint/registry.rs and lib/kv-router/src/plugins/registry.rs — integration with the existing router plugin registry and per-partition policy construction.
  3. lib/llm/src/kv_router.rs — policy installation, context construction, and fail-open invocation.
  4. lib/llm/src/kv_router/routing_host/kv.rs — post-selection execution, action precedence, and request dispatch.
  5. lib/kv-router/tests/plugin_api_compat.rs — external plugin registration and API compatibility coverage.

Related Issues

Validation

  • cargo fmt --all -- --check
  • cargo check -p dynamo-kv-router -p dynamo-llm --all-targets
  • cargo clippy -p dynamo-kv-router -p dynamo-llm --all-targets -- -D warnings
  • cargo test -p dynamo-kv-router — 1,047 passed; 2 ignored
  • cargo test -p dynamo-llm kv_hint_policy — passed

Signed-off-by: Karen Chung <karenc@nvidia.com>
@karen-sy
karen-sy requested a review from a team as a code owner October 5, 2026 23:08
@github-actions github-actions Bot added feat router Relates to routing, KV-aware routing, etc. labels Oct 5, 2026
@coderabbitai

coderabbitai Bot commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

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: 0aec729d-28c3-4714-bce9-6766f0f5c44f
📥 Commits

Reviewing files that changed from the base of the PR and between def3b79 and 940d96f.

📒 Files selected for processing (12)
  • lib/kv-router/src/plugins.rs
  • lib/kv-router/src/plugins/kv_hint.rs
  • lib/kv-router/src/plugins/kv_hint/config.rs
  • lib/kv-router/src/plugins/kv_hint/registry.rs
  • lib/kv-router/src/plugins/registry.rs
  • lib/kv-router/src/scheduling/config.rs
  • lib/kv-router/src/scheduling/policy_config.rs
  • lib/kv-router/src/scheduling/types.rs
  • lib/kv-router/tests/plugin_api_compat.rs
  • lib/llm/src/kv_router.rs
  • lib/llm/src/kv_router/plugins.rs
  • lib/llm/src/kv_router/routing_host/kv.rs

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


Walkthrough

The KV-router plugin system now supports configured KV-hint policies. The LLM router can install a policy, formulate actions from request context, and combine them with request and selection hints during dispatch.

Changes

KV-hint policy flow

Layer / File(s) Summary
Policy contract and configuration
lib/kv-router/src/plugins/kv_hint.rs, lib/kv-router/src/plugins/kv_hint/*, lib/kv-router/src/scheduling/policy_config.rs, lib/kv-router/src/scheduling/config.rs, lib/kv-router/src/scheduling/types.rs
Adds the policy context and callback contract, validates policy configuration and constructor registration, and exposes the configured policy through router settings.
Plugin registry and bundle resolution
lib/kv-router/src/plugins.rs, lib/kv-router/src/plugins/registry.rs, lib/kv-router/tests/plugin_api_compat.rs
Adds KV-hint policies to plugin registration and resolution. Tests register and construct a configured policy.
Router policy installation and formulation
lib/llm/src/kv_router.rs, lib/llm/src/kv_router/plugins.rs
The LLM router installs an optional policy and formulates actions using request, session, worker, and session-prefix index context. Policy errors and panics yield no actions.
Hint composition during dispatch
lib/llm/src/kv_router/routing_host/kv.rs
Dispatch combines request, selection, and policy actions in that order. It retains the request hint message ID when available, otherwise uses the selection hint ID or the request ID for policy-only actions.

Priority: ➖ Normal

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

Merge Risk: ⚪ Minimal · up to 940d9

This adds optional KV-hint policy support that is inactive unless configured. Policy failures fall back to no extra hints, so the risk of merging is low.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 57.53% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 73 functions across 12 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 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 identifies the new pluggable KV-hint policy interface and its session awareness.
Description check ✅ Passed The description includes the required overview, details, reviewer starting points, and a related-issue reference. It also explains the design, boundaries, and validation performed.
  • Fix all pre-merge checks with AI
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

@karen-sy
karen-sy marked this pull request as draft October 5, 2026 23:39
@devin-ai-integration

Copy link
Copy Markdown
Contributor

⏹️ Dynamo PR CI cancelled — run 37386801701 (attempt 1) on 940d96ff4e

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

Other Jobs
dynamo-runtime ⏹️ 1
frontend ⏹️ 1
planner ⏹️ 1
triton-runtime ⏹️ 1

Failure details

No job failed. All 4 non-passing jobs are image builds that GitHub cancelled during Build and Push Image after hitting their job time limit. This was not fail-fast and not a manual or concurrency cancel. The red gate checks only mirror these cancels. A slow or stuck image build or push is the likely cause, so rerunning the failed jobs is probably enough.

For agents
{"pr": 15679, "run_id": 37386801701, "run_attempt": 1, "head_sha": "940d96ff4e16e940834640363d2f9f971fe1e86f",
 "failures": [],
 "cancelled": [
  {"job": "dynamo-runtime / image / Build multi-arch cuda13.0", "job_id": 112022092446, "failed_step": "Build and Push Image", "signature": "exceeded maximum execution time of 1h0m0s", "log_cmd": "gh api repos/ai-dynamo/dynamo/actions/jobs/112022092446/logs"},
  {"job": "planner / Build multi-arch cpu", "job_id": 112022092176, "failed_step": "Build and Push Image", "signature": "exceeded maximum execution time of 1h0m0s", "log_cmd": "gh api repos/ai-dynamo/dynamo/actions/jobs/112022092176/logs"},
  {"job": "triton-runtime / Build multi-arch cpu", "job_id": 112022092058, "failed_step": "Build and Push Image", "signature": "exceeded maximum execution time of 1h0m0s", "log_cmd": "gh api repos/ai-dynamo/dynamo/actions/jobs/112022092058/logs"},
  {"job": "frontend / Build multi-arch cpu", "job_id": 112022092256, "failed_step": "Build and Push Image", "signature": "exceeded maximum execution time of 1h30m0s", "log_cmd": "gh api repos/ai-dynamo/dynamo/actions/jobs/112022092256/logs"}
 ]}

Posted automatically by Devin for run 37386801701. Updated on every full-CI run of this PR.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feat router Relates to routing, KV-aware routing, etc. size/XL

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant