Add aggregate mode support to AIDynamo workload - #982
Open
saivishal1999 wants to merge 1 commit into
Open
Conversation
In aggregate mode a single vLLM worker handles both prefill and decode on the same node(s), as opposed to disaggregated mode where separate prefill and decode workers run on distinct nodes. Changes: - ai_dynamo.py: add mode field (aggregate|disaggregate, default disaggregate) to AIDynamoArgs; skip decode worker model population and constraint checks in aggregate mode - slurm_command_gen_strategy.py: zero out decode node count and omit --decode-* CLI args when mode=aggregate - ai_dynamo.sh: guard decode TP/PP validation and workers-per-node calculation when no decode nodes are configured; fall back to prefill node as frontend-node in aggregate mode Tested on Lyris GB200: aggregate (1 node, prefill 1/1 decode 0/0) and disaggregated (2 nodes, prefill 1/1 decode 1/1) both pass.
saivishal1999
requested review from
jeffnvidia,
podkidyshev and
srivatsankrishnan
as code owners
July 28, 2026 17:15
Contributor
📝 WalkthroughWalkthroughChangesAI Dynamo aggregate mode
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Comment |
![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: 1
🤖 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 `@src/cloudai/workloads/ai_dynamo/ai_dynamo.py`:
- Around line 583-586: Update the aggregate-mode branch in the constraint-check
function so it skips only decode/split-specific validation while still executing
prefill GPU-capacity validation, including tp_times_pp_le_gpus_per_node. Remove
the unconditional early return and preserve the existing aggregate-mode handling
for checks that do not apply.
🪄 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: ASSERTIVE
Plan: Enterprise
Run ID: 4b622366-c12f-4a19-aaee-a617e29f1f49
📒 Files selected for processing (3)
src/cloudai/workloads/ai_dynamo/ai_dynamo.pysrc/cloudai/workloads/ai_dynamo/ai_dynamo.shsrc/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py
Comment on lines
+583
to
+586
| if tr.test.cmd_args.dynamo.mode == "aggregate": | ||
| logging.info("constraint_check skipped: aggregate mode has no prefill/decode split") | ||
| return True | ||
|
|
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Retain prefill GPU-capacity validation in aggregate mode.
This return also skips tp_times_pp_le_gpus_per_node for the prefill worker. An aggregate run with TP × PP exceeding available GPUs is accepted, then the runtime allocation can compute zero workers per node. Skip only decode/split-specific checks; continue validating prefill capacity.
This is evidenced by the allocation path in src/cloudai/workloads/ai_dynamo/ai_dynamo.sh Lines 376-390.
🤖 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 `@src/cloudai/workloads/ai_dynamo/ai_dynamo.py` around lines 583 - 586, Update
the aggregate-mode branch in the constraint-check function so it skips only
decode/split-specific validation while still executing prefill GPU-capacity
validation, including tp_times_pp_le_gpus_per_node. Remove the unconditional
early return and preserve the existing aggregate-mode handling for checks that
do not apply.
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
1 participant
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.
In aggregate mode a single vLLM worker handles both prefill and decode on the same node(s), as opposed to disaggregated mode where separate prefill and decode workers run on distinct nodes.
Changes:
Tested on Lyris GB200: aggregate (1 node, prefill 1/1 decode 0/0) and disaggregated (2 nodes, prefill 1/1 decode 1/1) both pass.