feat(router): conditional disaggregation - #11357
Open
karen-sy wants to merge 10 commits into
Open
Conversation
github-actions
Bot
added
feat
backend::vllm
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
|
Contributor
WalkthroughThis PR adds conditional disaggregation across router config, policy selection, busy-state checks, and decode-worker bypass handling. It also wires new Python/Rust bindings, CLI flags, and documentation for the new routing behavior. ChangesConditional Disaggregation Feature
Estimated code review effort: 4 (Complex) | ~75 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
karen-sy
marked this pull request as draft
July 7, 2026 22:45
Signed-off-by: Karen Chung <karenc@nvidia.com>
github-actions
Bot
added
the
documentation
Contributor
karen-sy
marked this pull request as ready for review
July 8, 2026 16:21
Signed-off-by: Karen Chung <karenc@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
components/src/dynamo/vllm/handlers.py (1)
2758-2787: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMirror the bypass annotation check in
_generate_text_mode— decode workers can run in text mode, butgenerate()routes those requests straight to_generate_text_mode, which still treats all decode workers as plain decode-only. That leaves conditional-disagg bypass requests using the wrong abort-defer path; match the_generate_token_modehandling here too.🤖 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 `@components/src/dynamo/vllm/handlers.py` around lines 2758 - 2787, Mirror the bypass-annotation handling from _generate_token_mode in _generate_text_mode so decode workers can switch to AGG when BYPASS_REMOTE_PREFILL_ANNOTATION is present. Update the request routing in _generate_text_mode to inspect request["annotations"] the same way, set the local decode-only flag accordingly, and ensure the abort-defer path uses the same conditional-disagg behavior as the token-mode flow.
🧹 Nitpick comments (5)
components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py (1)
653-664: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider adding coverage for the
_generate_locally_implbypass wiring.This test only validates
_setup_disaggregated_params_for_modein isolation. The end-to-end bypass behavior in_generate_locally_impl— annotation detection, overwritingrequest["disaggregated_params"], skipping the DECODE "params is None" error, and skipping_DeferredAbort— is not exercised by any test in this file.🤖 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 `@components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py` around lines 653 - 664, Add an end-to-end test around _generate_locally_impl in TRTLLMHandlerBase that uses BYPASS_REMOTE_PREFILL_ANNOTATION to drive the decode bypass path. Verify the method detects the annotation, rewrites request["disaggregated_params"], avoids the DECODE “params is None” error, and does not raise _DeferredAbort; keep the existing _setup_disaggregated_params_for_mode test if useful, but extend coverage through the actual _generate_locally_impl flow.components/src/dynamo/trtllm/request_handlers/handler_base.py (1)
690-699: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated bypass-condition logic.
The bypass condition
self.disaggregation_mode == DisaggregationMode.DECODE and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or [])is computed independently at Line 691-693 and again at Line 978-980. Since this logic determines execution mode (local prefill+decode vs. remote handoff), keeping it in one place reduces the risk of the two checks drifting out of sync in future edits.♻️ Proposed refactor to share the bypass check
+ def _has_bypass_remote_prefill(self, request: dict) -> bool: + return ( + self.disaggregation_mode == DisaggregationMode.DECODE + and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or []) + ) + def _setup_disaggregated_params_for_mode( self, request: dict, ep_disaggregated_params: Optional[Any], ) -> tuple[Any, Any, dict]: ... - use_request_disagg_params = request.get(HEALTH_CHECK_KEY) or ( - self.disaggregation_mode == DisaggregationMode.DECODE - and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or []) - ) + use_request_disagg_params = request.get( + HEALTH_CHECK_KEY + ) or self._has_bypass_remote_prefill(request)- bypass_remote_prefill = ( - self.disaggregation_mode == DisaggregationMode.DECODE - and BYPASS_REMOTE_PREFILL_ANNOTATION in (request.get("annotations") or []) - ) + bypass_remote_prefill = self._has_bypass_remote_prefill(request)Also applies to: 977-986
🤖 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 `@components/src/dynamo/trtllm/request_handlers/handler_base.py` around lines 690 - 699, The request bypass condition is duplicated in handler_base’s request handling flow, which risks the prefill/decode mode checks drifting apart over time. Factor the bypass logic into a single shared helper or boolean within the relevant request-processing path in HandlerBase, and reuse it both where disaggregated params are selected and where the later execution-mode decision is made so the decode-mode + BYPASS_REMOTE_PREFILL_ANNOTATION check stays consistent.components/src/dynamo/vllm/handlers.py (1)
2772-2787: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider asserting
is_decode_onlyflips correctly for the deferred-abort guard.The bypass path changes
is_decode_only, which feeds directly into_deferred_abort_guard's abort-safety semantics. The new tests only assert on chunks/multimodal extraction; a small assertion on the guard'sis_decode_onlyargument would pin down this safety-relevant behavior.🤖 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 `@components/src/dynamo/vllm/handlers.py` around lines 2772 - 2787, The bypass path in the handler logic updates is_decode_only and then passes it into _deferred_abort_guard, so add a test assertion that this flag becomes False when BYPASS_REMOTE_PREFILL_ANNOTATION is present and remains True otherwise. Use the existing handler flow around DisaggregationMode.DECODE, request.get("annotations"), and _deferred_abort_guard to verify the exact is_decode_only argument being sent, not just the chunking or multimodal outputs.lib/kv-router/src/conditional_disagg.rs (1)
109-118: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate policy-gating logic risks drift.
policy_needs_prefill_worker_busyre-implements the same enabled+policy-kind matching that each concrete policy'sneeds_prefill_worker_busy()(Lines 223-225, 273-275) already encodes viamake_conditional_disagg_policy. Two independent sources of truth for the same decision can silently diverge as policies are added/changed.♻️ Suggested consolidation
-pub fn policy_needs_prefill_worker_busy(config: Option<&KvRouterConfig>) -> bool { - let Some(config) = config else { return false }; - if !config.conditional_disagg_enabled { - return false; - } - matches!( - config.conditional_disagg_policy, - ConditionalDisaggPolicyKind::PrefillLoad | ConditionalDisaggPolicyKind::IslOrLoad, - ) -} +pub fn policy_needs_prefill_worker_busy(config: Option<&KvRouterConfig>) -> bool { + make_conditional_disagg_policy(config).needs_prefill_worker_busy() +}🤖 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/kv-router/src/conditional_disagg.rs` around lines 109 - 118, Consolidate the duplicated prefill-worker-busy gating logic by making policy_needs_prefill_worker_busy reuse the existing policy-specific needs_prefill_worker_busy behavior created by make_conditional_disagg_policy. The current function independently checks conditional_disagg_enabled plus ConditionalDisaggPolicyKind matches, which duplicates the decision already encoded in the concrete policy implementations; refactor it to delegate through the configured policy instance (or a shared helper used by both) so there is a single source of truth. Keep the existing symbols policy_needs_prefill_worker_busy, needs_prefill_worker_busy, and make_conditional_disagg_policy as the main anchor points for the change.components/src/dynamo/frontend/frontend_args.py (1)
162-184: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate validation logic vs.
router/args.py.The policy-choice,
eff_isl_threshold >= 0, andeff_isl_ratio_thresholdrange checks (Lines 162-176) are duplicated verbatim incomponents/src/dynamo/router/args.py(DynamoRouterConfig.validate, Lines 71-85). Consider extracting a shared validator (e.g., alongsidewarn_conditional_disagg_prefill_busy_threshold_resolutioninkv_router_args.py) to avoid the two implementations drifting.♻️ Suggested consolidation
+def validate_conditional_disagg_common(config: "KvRouterConfigBase") -> None: + if config.conditional_disagg_policy not in CONDITIONAL_DISAGG_POLICY_CHOICES: + raise ValueError( + "--router-conditional-disagg-policy must be one of " + + ", ".join(f"'{c}'" for c in CONDITIONAL_DISAGG_POLICY_CHOICES) + ) + if config.conditional_disagg_eff_isl_threshold < 0: + raise ValueError("--router-conditional-disagg-eff-isl-threshold must be >= 0") + if not 0.0 <= config.conditional_disagg_eff_isl_ratio_threshold <= 1.0: + raise ValueError( + "--router-conditional-disagg-eff-isl-ratio-threshold must be in [0.0, 1.0]" + )Then call
validate_conditional_disagg_common(self)from bothfrontend_args.pyandrouter/args.py.🤖 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 `@components/src/dynamo/frontend/frontend_args.py` around lines 162 - 184, The validation for conditional disaggregation settings is duplicated between frontend_args.py and DynamoRouterConfig.validate, so extract the shared checks for policy choice, eff_isl_threshold >= 0, and eff_isl_ratio_threshold range into a common helper near warn_conditional_disagg_prefill_busy_threshold_resolution in kv_router_args.py. Update both frontend_args.py and router/args.py to call the shared validator (for example, validate_conditional_disagg_common(self)) so the logic stays in sync and avoids drift.
🤖 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 `@components/src/dynamo/trtllm/request_handlers/handler_base.py`:
- Around line 690-699: The bypass branch in handler_base.py returns only
LlmDisaggregatedParams and drops ep_disaggregated_params, which prevents
_prepare_input_for_generation() from forwarding encode-worker embedding handles
into multimodal_processor.process_openai_request(). Update the
use_request_disagg_params handling in the request path to preserve and return
ep_disaggregated_params alongside the disaggregated params, and add coverage for
conditional-disagg bypass requests that include EPD/multimodal data.
In `@docs/components/router/router-configuration.md`:
- Line 24: The `--router-conditional-disagg` description in the router
configuration docs should not imply a hard prerequisite that is not enforced
here; either update the text to state `--router-mode kv` is an assumed
deployment requirement, or align it with the actual validation performed by the
router configuration code. Use the `--router-conditional-disagg` entry and the
related router mode wording to make the prerequisite clear without overclaiming
enforcement.
In `@lib/kv-router/src/scheduling/selector.rs`:
- Around line 212-225: The new decode-only early-return in
DefaultWorkerSelector::calculate_logit (the branch gated by self.worker_type ==
"decode" and !request.track_prefill_tokens) is not covered by existing tests.
Add a focused test that constructs DefaultWorkerSelector with worker_type set to
"decode" and a request with track_prefill_tokens disabled so this branch is
exercised directly, then assert the expected logit behavior; keep the existing
prefill-accounting test intact since it uses a non-decode worker_type and cannot
hit this path.
---
Outside diff comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 2758-2787: Mirror the bypass-annotation handling from
_generate_token_mode in _generate_text_mode so decode workers can switch to AGG
when BYPASS_REMOTE_PREFILL_ANNOTATION is present. Update the request routing in
_generate_text_mode to inspect request["annotations"] the same way, set the
local decode-only flag accordingly, and ensure the abort-defer path uses the
same conditional-disagg behavior as the token-mode flow.
---
Nitpick comments:
In `@components/src/dynamo/frontend/frontend_args.py`:
- Around line 162-184: The validation for conditional disaggregation settings is
duplicated between frontend_args.py and DynamoRouterConfig.validate, so extract
the shared checks for policy choice, eff_isl_threshold >= 0, and
eff_isl_ratio_threshold range into a common helper near
warn_conditional_disagg_prefill_busy_threshold_resolution in kv_router_args.py.
Update both frontend_args.py and router/args.py to call the shared validator
(for example, validate_conditional_disagg_common(self)) so the logic stays in
sync and avoids drift.
In `@components/src/dynamo/trtllm/request_handlers/handler_base.py`:
- Around line 690-699: The request bypass condition is duplicated in
handler_base’s request handling flow, which risks the prefill/decode mode checks
drifting apart over time. Factor the bypass logic into a single shared helper or
boolean within the relevant request-processing path in HandlerBase, and reuse it
both where disaggregated params are selected and where the later execution-mode
decision is made so the decode-mode + BYPASS_REMOTE_PREFILL_ANNOTATION check
stays consistent.
In `@components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py`:
- Around line 653-664: Add an end-to-end test around _generate_locally_impl in
TRTLLMHandlerBase that uses BYPASS_REMOTE_PREFILL_ANNOTATION to drive the decode
bypass path. Verify the method detects the annotation, rewrites
request["disaggregated_params"], avoids the DECODE “params is None” error, and
does not raise _DeferredAbort; keep the existing
_setup_disaggregated_params_for_mode test if useful, but extend coverage through
the actual _generate_locally_impl flow.
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 2772-2787: The bypass path in the handler logic updates
is_decode_only and then passes it into _deferred_abort_guard, so add a test
assertion that this flag becomes False when BYPASS_REMOTE_PREFILL_ANNOTATION is
present and remains True otherwise. Use the existing handler flow around
DisaggregationMode.DECODE, request.get("annotations"), and _deferred_abort_guard
to verify the exact is_decode_only argument being sent, not just the chunking or
multimodal outputs.
In `@lib/kv-router/src/conditional_disagg.rs`:
- Around line 109-118: Consolidate the duplicated prefill-worker-busy gating
logic by making policy_needs_prefill_worker_busy reuse the existing
policy-specific needs_prefill_worker_busy behavior created by
make_conditional_disagg_policy. The current function independently checks
conditional_disagg_enabled plus ConditionalDisaggPolicyKind matches, which
duplicates the decision already encoded in the concrete policy implementations;
refactor it to delegate through the configured policy instance (or a shared
helper used by both) so there is a single source of truth. Keep the existing
symbols policy_needs_prefill_worker_busy, needs_prefill_worker_busy, and
make_conditional_disagg_policy as the main anchor points for the change.
🪄 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: 00e55a81-b317-402d-bce3-f138e2247709
📒 Files selected for processing (27)
components/src/dynamo/common/configuration/groups/kv_router_args.pycomponents/src/dynamo/frontend/frontend_args.pycomponents/src/dynamo/router/args.pycomponents/src/dynamo/trtllm/request_handlers/handler_base.pycomponents/src/dynamo/trtllm/tests/test_trtllm_handler_base.pycomponents/src/dynamo/vllm/args.pycomponents/src/dynamo/vllm/backend_args.pycomponents/src/dynamo/vllm/handlers.pycomponents/src/dynamo/vllm/main.pycomponents/src/dynamo/vllm/tests/test_vllm_worker_handler.pydocs/components/router/router-configuration.mddocs/components/router/router-disaggregated-serving.mdlib/bindings/python/rust/llm/entrypoint.rslib/bindings/python/src/dynamo/_core.pyilib/kv-router/src/conditional_disagg.rslib/kv-router/src/lib.rslib/kv-router/src/scheduling/config.rslib/kv-router/src/scheduling/local.rslib/kv-router/src/scheduling/queue.rslib/kv-router/src/scheduling/selector.rslib/llm/src/discovery/watcher.rslib/llm/src/kv_router.rslib/llm/src/kv_router/prefill_router/activation.rslib/llm/src/kv_router/prefill_router/conditional_bypass.rslib/llm/src/kv_router/prefill_router/mod.rslib/llm/src/kv_router/push_router.rslib/llm/src/kv_router/scheduler.rs
…sagg-main Signed-off-by: Karen Chung <karenc@nvidia.com>
…main Signed-off-by: Karen Chung <karenc@nvidia.com>
Signed-off-by: Karen Chung <karenc@nvidia.com>
| | ----------------------------------------------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------- | | ||
| | `--router-conditional-disagg` | Disabled | Enables conditional disaggregation. Requires `--router-mode kv` and separate prefill/decode worker pools. | | ||
| | `--router-conditional-disagg-policy` | `isl_bounding` | Selects the bypass policy: `isl_bounding`, `prefill_load`, or `isl_or_load`. | | ||
| | `--router-conditional-disagg-eff-isl-threshold` | `2048` | Sets the effective ISL token threshold for `isl_bounding` and `isl_or_load`. | |
Contributor
There was a problem hiding this comment.
I don't understand difference between isl-threshold and isl-ratio from this
Contributor
Author
There was a problem hiding this comment.
added some more info, can you check again?
| | `--router-conditional-disagg-eff-isl-threshold` | `2048` | Sets the effective ISL token threshold for `isl_bounding` and `isl_or_load`. | | ||
| | `--router-conditional-disagg-eff-isl-ratio-threshold` | `0.7` | Sets the effective/raw ISL ratio threshold for `isl_bounding` and `isl_or_load`. | | ||
| | `--router-conditional-disagg-prefill-busy-threshold` | Unset | Sets the prefill busy threshold for `prefill_load` and `isl_or_load`. When unset, those policies inherit `--router-queue-threshold` if it is set. | | ||
| | `--router-conditional-disagg-decode-busy-threshold` | Unset | A decode loadedness threshold. When set, gates decode workers from local prefill work when decode-side KV pressure % is above the threshold. | |
Contributor
There was a problem hiding this comment.
Can you give an acceptable range? Is this like 90% or something?
Contributor
Author
There was a problem hiding this comment.
added tuning recommendation
| ); | ||
| } | ||
| (None, None) => { | ||
| tracing::warn!( |
Contributor
There was a problem hiding this comment.
Can this be more actionable? As a user I wouldn't know really what this means or what to do.
Contributor
Author
There was a problem hiding this comment.
improved the message, can you check again
| if config.conditional_disagg_enabled | ||
| && let Some(threshold) = config.conditional_disagg_decode_busy_threshold | ||
| { | ||
| tracing::info!( |
Contributor
There was a problem hiding this comment.
Same here, what does circuit breaker mean? How would a user understand this? We should have what this is in the docs if we are printing it.
Contributor
Author
There was a problem hiding this comment.
improved the message, can you check again
| ); | ||
| } | ||
| (None, Some(threshold)) => { | ||
| tracing::info!( |
Contributor
There was a problem hiding this comment.
Same here. I don't know what load gate means? Why does it inherit the router-queue-threshold?
Contributor
Author
There was a problem hiding this comment.
improved the message, can you check again
Contributor
|
This PR needs a boatload of reviewers. Do you think its worth more tightly scoping it and the changes? |
Contributor
|
I think I'm ok with this shape if you guys are |
Contributor
We would need 9 different code-owner groups to review to technically merge it without a force.
|
Contributor
Author
|
Mmm yeah I can convert it into smaller stacked PR's. Don't want all these codeowner groups to have to comb thru all the diffs. Let me address @PeaBrane and @alec-flowers comments so far first, break PR down into smaller ones, and link here |
Signed-off-by: Karen Chung <karenc@nvidia.com>
Signed-off-by: Karen Chung <karenc@nvidia.com>
karen-sy
force-pushed
the
karenc/conditional-disagg-main
branch
from
July 15, 2026 03:45
a42f5fe to
8add3b5
Compare
…sagg-main Signed-off-by: Karen Chung <karenc@nvidia.com>
This was referenced Jul 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview:
For motivation/design/prelim results, see DEP here: #11514
Adds conditional disaggregation support to the KV router so selected requests can bypass the remote prefill worker and run prefill+decode locally on the chosen decode worker when the effective ISL is small enough or when the configured load policy says prefill is busy. Enables decode KV-affinity routing for conditional-disagg mode.
Supports vLLM and TRTLLM.
Details:
Summary
isl_boundingprefill_loadisl_or_load--router-conditional-disagg-decode-busy-threshold.x-bypass-remote-prefill.context_and_generation.Validation
cargo fmt --allcargo test -p dynamo-kv-router conditional_disaggcargo test -p dynamo-llm prefill_routercargo checkfromlib/bindings/pythonpython -m py_compileon modified Python handler/config/test filesWhere should the reviewer start?
docs/components/routerlib/kv-router/src/conditional_disagg.rslib/llm/src/kv_router/prefill_router/conditional_bypass.rslib/llm/src/kv_router/prefill_router/mod.rslib/kv-router/src/scheduling/selector.rscomponents/src/dynamo/vllm/handlers.pycomponents/src/dynamo/trtllm/request_handlers/handler_base.pycomponents/src/dynamo/common/configuration/groups/kv_router_args.pySummary by CodeRabbit
Summary
New Features
Bug Fixes