feat(router): orchestrate conditional disagg bypass

[karen-sy]

Important

This PR is [4/5] in the conditional-disaggregation stack.
Merge order is bottom-up: #11718 merges to main first, and #11723 merges last.
Do not merge this PR until every earlier PR in the checklist has merged and this branch has been rebased/retargeted as needed.

Link to the DEP: #11514

Stack merge order:

• feat(router): add conditional disagg policy kernel #11718 — feat(router): add conditional disagg policy kernel
• feat(backends): handle conditional disagg bypass annotation #11719 — feat(backends): handle conditional disagg bypass annotation
• feat(router): add decode affinity scoring #11720 — feat(router): add decode affinity scoring
• feat(router): orchestrate conditional disagg bypass #11725 — feat(router): orchestrate conditional disagg bypass
• feat(router): expose conditional disagg configuration #11723 — feat(router): expose conditional disagg configuration

Reference to original (closed) combined PR, with addressed review comments: ai-dynamo/dynamo#11357

Overview:

Details:

Where should the reviewer start?

Related Issues

⚠️ This section is required. Choose one path below and delete the other.

🔗 This PR is linked to an issue:

• Closes #XXXX

🚫 This PR is NOT linked to an issue:

• Confirmed — no related issue

Summary by CodeRabbit

• New Features

  • Added conditional routing that can bypass remote prefill and send requests directly to an appropriate decode worker.
  • Added session-affinity support to improve consistency when selecting workers across request phases.
  • Added advisory worker selection and load checks to make routing decisions without reserving capacity.
  • Added configurable safeguards based on prefill and decode worker workload.

• Bug Fixes

  • Improved propagation of routing affinity and configuration across prefill and decode paths.

[devin-ai-integration]

Devin Review found 2 potential issues.

[PeaBrane]

AI-assisted review of the conditional-disagg orchestration and stacked contracts. Posting these as neutral questions and edge cases for consideration.

[PeaBrane]

Approving on the condition that the AI-assisted review comments above are resolved.

[coderabbitai]

Walkthrough

Changes

The PR adds conditional-disaggregation support to the prefill router, including decode-worker selection without admission, session-affinity lookup, worker-busy gating, bypass stream generation, and updated PrefillRouter::new call sites.

Conditional disaggregation routing

Layer / File(s)	Summary
Non-admission scheduler selection lib/kv-router/src/scheduling/*, lib/llm/src/kv_router/scheduler.rs	Schedulers add non-admission selection, shared worker-selection logic, worker-load queries, and response handling.
Admission-aware KV routing lib/llm/src/kv_router.rs, lib/llm/src/kv_router/push_router/selection.rs	KV routing distinguishes lifecycle-tracked admission from advisory selection and returns projected decode-block information.
Prefill router state and affinity wiring lib/llm/src/kv_router/prefill_router/*, lib/llm/src/kv_router/push_router.rs, lib/llm/src/entrypoint/input/common.rs, lib/llm/src/discovery/watcher.rs, deploy/inference-gateway/ext-proc/src/epp.rs, lib/bindings/c/src/lib.rs	PrefillRouter receives decode routing and conditional-disaggregation configuration, while initialization paths and session-affinity wiring are updated.
Conditional bypass execution lib/llm/src/kv_router/prefill_router/conditional_bypass.rs, lib/llm/src/kv_router/prefill_router/mod.rs	Conditional bypass evaluates routing pins, overlap, policy, and busy thresholds, then emits a bypass annotation and chains the decode stream when selected; tests cover policy and pin resolution.

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

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	Required sections are mostly empty, and the Related Issues section still has placeholders instead of a real issue reference or no-issue confirmation.	Fill in Overview, Details, and Where should the reviewer start, and replace the Related Issues placeholder with a real issue link or the no-issue confirmation.
Docstring Coverage	⚠️ Warning	Docstring coverage is 53.06% which is insufficient. The required threshold is 80.00%.	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 matches the main change: orchestrating conditional disaggregation bypass in router code.
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.

---

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

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

lib/llm/src/kv_router/prefill_router/conditional_bypass.rs (1)

231-243: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Drop the dead decode_chosen_worker_busy write
ConditionalDisaggDecisionInput is Copy, so this still compiles, but the updated value is never used: should_bypass_remote_prefill has already run, and the decision/log path reads decode_busy directly. Remove the assignment unless a policy will consume the field.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_router/prefill_router/conditional_bypass.rs` around lines 231
- 243, Remove the unused input reassignment via with_decode_chosen_worker_busy
in the conditional bypass flow. Keep the decode_busy computation intact, since
the existing decision and logging paths already use it directly, and do not add
a replacement field update unless a policy consumes that field.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/kv-router/src/scheduling/queue.rs`:
- Around line 693-716: The capacity checks in worker_is_prefill_busy and
projected_decode_load_exceeds currently return None when configuration capacity
is unavailable, unlike the queue path’s DEFAULT_MAX_BATCHED_TOKENS fallback.
Align both methods with that queue behavior by using the established fallback
when capacity is missing, while preserving their existing Option-based
worker/config lookup handling.

In `@lib/llm/src/kv_router/prefill_router/conditional_bypass.rs`:
- Around line 174-175: Update both conditional-disaggregation probes in
conditional_bypass.rs: at lines 174-175 and 352-353, pass the request’s
policy_class and session_id to find_best_match_details_without_admission instead
of None, reusing the values derived by the normal KvPushRouter::select_request
path so both advisory checks inspect the same class queue and session view as
real routing.
- Around line 125-156: Reorder the pin resolution in the conditional bypass flow
so resolve_request_decode_pin is evaluated before decode_affinity_target. If the
request pin resolves, use it without aborting when affinity lacks a DP rank;
only evaluate and fall back to affinity when no request pin exists, while
preserving unresolved request-pin behavior and the existing pinned_worker
precedence.

---

Nitpick comments:
In `@lib/llm/src/kv_router/prefill_router/conditional_bypass.rs`:
- Around line 231-243: Remove the unused input reassignment via
with_decode_chosen_worker_busy in the conditional bypass flow. Keep the
decode_busy computation intact, since the existing decision and logging paths
already use it directly, and do not add a replacement field update unless a
policy consumes that field.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4c2d1a8f-45f6-46b1-b4db-a3b8e29d0669

📥 Commits

Reviewing files that changed from the base of the PR and between cf495c8 and 3eb4288.

📒 Files selected for processing (13)

• deploy/inference-gateway/ext-proc/src/epp.rs
• lib/bindings/c/src/lib.rs
• lib/kv-router/src/scheduling/local.rs
• lib/kv-router/src/scheduling/queue.rs
• lib/llm/src/discovery/watcher.rs
• lib/llm/src/entrypoint/input/common.rs
• lib/llm/src/kv_router.rs
• lib/llm/src/kv_router/prefill_router/activation.rs
• lib/llm/src/kv_router/prefill_router/conditional_bypass.rs
• lib/llm/src/kv_router/prefill_router/mod.rs
• lib/llm/src/kv_router/push_router.rs
• lib/llm/src/kv_router/push_router/selection.rs
• lib/llm/src/kv_router/scheduler.rs

[MatejKosec]

One design question after tracing this path:

The scheduler already has the information needed here: which worker it selected, the projected load for that worker, and the worker’s capacity. At the moment, the advisory selection returns only part of that information. PrefillRouter then asks back down through KvRouter → KvScheduler → LocalScheduler → SchedulerQueue whether the selected worker exceeds a threshold. That is why methods such as projected_decode_load_exceeds and worker_is_prefill_busy appear at several layers even though the calculation itself happens only once.

Would it be cleaner for the advisory selection to return the relevant load measurements together with the selected worker? For example, a small result could contain the worker, cached tokens, projected decode blocks, and total KV capacity. The conditional-disaggregation code could then apply its configured threshold directly.

I think that would make the responsibilities easier to follow. The scheduler would report the measurements it already owns, while conditional disaggregation would decide what those measurements mean for the bypass policy. It would also let the worker selection and the associated load/capacity measurements come from the same observation, rather than selecting first and querying some of the state afterward.

I’m not suggesting a broad snapshot of everything in the scheduler—just a small, purpose-built advisory result. This would not remove the acknowledged race between advisory selection and real admission, but it would make this path more cohesive and avoid adding more one-off load queries through every layer as conditional disaggregation evolves.

I do not think this needs to block the PR, but I wanted to raise it as a possible cleanup direction.

[MatejKosec]

2. select_without_admission_inner now has its own copy of most of the code in select_worker_for_request. Both functions calculate the current worker loads, remove workers that cannot accept the request, choose the best remaining worker, and collect information about that worker. The new function needs some extra information, which explains why the code was copied. The problem is that there are now two places that define how a worker is chosen. They behave the same today, but a future change to worker selection could be added to one function and missed in the other. Then the advisory path and the normal admission path could make different choices for reasons that are very difficult to see while reading the code. I think both paths should use one shared selection function, with the advisory path adding the chosen worker’s load and capacity to the result.

[MatejKosec]

Approved, once comments resolved.