feat(sdk): model EPD (encoder disaggregation) for VL disagg sweep#1340
feat(sdk): model EPD (encoder disaggregation) for VL disagg sweep#1340wenpengw-nv wants to merge 3 commits into
Conversation
WalkthroughChangesEPD support adds dedicated encoder-worker selection for vision-language inference, including encoder-only benchmarking, task and CLI configuration, encoder-aware sweep matching, and report fields. Encoder model and backend contract
EPD candidate enumeration and matching
Task and CLI configuration flow
EPD output schema and reporting
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (5)
src/aiconfigurator/cli/main.py (1)
1178-1179: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring not updated for the two new params.
enable_epdandencoder_tpwere added to the signature, but theArgs:docstring block below (visible in context, lines ~1206-1215) doesn't document either one.As per path instructions, "For new or changed user-facing options, verify docs updates and generator/SDK wiring."
📝 Proposed docstring addition
image_height: int = 0, image_width: int = 0, num_images: int = 1, enable_epd: bool = False, encoder_tp: list[int] | None = None,enable_chunked_prefill: Whether to enable chunked prefill for finer context token sweep. + enable_epd: EPD (vision-language models): run the vision encoder on dedicated encode + workers instead of colocated with prefill. Requires an image workload. + encoder_tp: EPD encode-worker TP sizes to sweep (requires enable_epd). Default: [1, 2, 4, 8]. enable_wideep: Whether to enable Wide Expert Parallelism (WideEP) for MoE models.🤖 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/aiconfigurator/cli/main.py` around lines 1178 - 1179, Update the Args: docstring for the function containing enable_epd and encoder_tp to document both new parameters, including their purpose, accepted values, and defaults. Ensure the descriptions match the existing CLI option behavior and any generator/SDK wiring.Source: Path instructions
src/aiconfigurator/sdk/task_v2.py (1)
1138-1155: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winFail fast:
enable_epddoesn't validate an image workload is configured.The
--enable-epdCLI help text states EPD "requires an image workload (--image-height/--image-width)", and_get_encoder_worker_candidatesinsweep.pydoes raise a clearValueErrorwhen there's no image input — but only afterget_model, ViT-op construction, and per-(tp,batch)backend.run_encoder_staticcalls. Adding the check here fails fast before any of that model/backend work runs, consistent with the same-file philosophy of_check_prefix_discipline(explicit over silent-then-late failures).Based on coding guidelines, "Catch unsupported system/backend/version combinations early with fast validation to prevent wasting user time on late failures after profiling."
♻️ Proposed fix
if (self.encoder_tp_candidates or self.encoder_batch_candidates) and not self.enable_epd: raise ValueError("encoder_tp_candidates/encoder_batch_candidates require enable_epd=True.") + if self.enable_epd and not (self.image_height > 0 and self.image_width > 0): + raise ValueError( + "enable_epd (EPD) requires an image workload; set image_height/image_width." + )🤖 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/aiconfigurator/sdk/task_v2.py` around lines 1138 - 1155, Add an early validation in _validate_disagg for enable_epd that requires both image workload dimensions, image_height and image_width, to be configured; raise a clear ValueError when either is missing, before any model or backend work occurs. Keep the existing encoder candidate validation intact and align the message with the CLI requirement that EPD needs --image-height/--image-width.Source: Coding guidelines
src/aiconfigurator/cli/report_and_save.py (1)
221-226: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueMinor defensive nit:
int(row.get(...) or 0)would raise on NaN.
int(x or 0)only guardsNone/0/""; afloat('nan')(truthy in Python) would pass through andint(nan)raisesValueError. Current sweep semantics appear to always populate(e)tp/(e)workerswith concrete ints (0 for plain PD, real values for EPD) rather thanNaN, so this isn't reachable today — flagging only as a defensive hardening for future schema changes.🤖 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/aiconfigurator/cli/report_and_save.py` around lines 221 - 226, The encoder worker parsing in the report formatting logic is not defensive against NaN values. Update the `(e)workers` and `(e)tp` conversions in the `gpus_replica_str` construction to safely treat missing, invalid, or NaN values as zero before calling `int`, while preserving existing formatting for valid values.src/aiconfigurator/sdk/backends/base_backend.py (1)
303-327: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a direct unit test for
run_encoder_static.The EPD sweep test (
test_sweep_disagg_epd_composes_encoder_stage) mocks_get_encoder_worker_candidatesentirely, sorun_encoder_static's latency/power/memory composition (and thenot model.encoder_opsbranches inrun_static/run_static_latency_only) aren't exercised end-to-end anywhere in the provided diff. Worth a small backend-level test (e.g. a VL model withencoder_colocated=False) to pin the contract directly.🤖 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/aiconfigurator/sdk/backends/base_backend.py` around lines 303 - 327, The new run_encoder_static composition lacks direct unit coverage. Add a backend-level test using a vision-language model with encoder_colocated=False that mocks _run_encoder_phase and _get_encoder_component_memory_for_runtime, then verifies latency correction, phase-average power, and returned memory; also cover zero-latency/no-encoder behavior and the related run_static and run_static_latency_only branches.tests/unit/sdk/sweep/test_sweep.py (1)
209-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLGTM on the happy-path EPD composition test — arithmetic double-checked against
sweep.pyand it all reconciles.One gap: no test here covers the negative/edge EPD paths introduced in this PR —
enable_epd=True+autoscale=TrueraisingValueError, or the case where every encoder candidate's latency exceeds the TTFT budget (emptyencoder_choices). Worth adding if not covered elsewhere.As per path instructions,
tests/**: "Check that tests cover the changed behavior rather than only the happy path."🤖 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 `@tests/unit/sdk/sweep/test_sweep.py` around lines 209 - 346, Add tests alongside test_sweep_disagg_epd_composes_encoder_stage covering both negative EPD paths: assert sweep_disagg raises ValueError when enable_epd=True and autoscale=True, and mock encoder candidates whose encoder_latency exceeds the TTFT budget, then assert the resulting encoder_choices path is empty and handled as expected. Reuse the existing common_kwargs and candidate monkeypatches, referencing sweep_disagg and _get_encoder_worker_candidates.Source: Path instructions
🤖 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.
Nitpick comments:
In `@src/aiconfigurator/cli/main.py`:
- Around line 1178-1179: Update the Args: docstring for the function containing
enable_epd and encoder_tp to document both new parameters, including their
purpose, accepted values, and defaults. Ensure the descriptions match the
existing CLI option behavior and any generator/SDK wiring.
In `@src/aiconfigurator/cli/report_and_save.py`:
- Around line 221-226: The encoder worker parsing in the report formatting logic
is not defensive against NaN values. Update the `(e)workers` and `(e)tp`
conversions in the `gpus_replica_str` construction to safely treat missing,
invalid, or NaN values as zero before calling `int`, while preserving existing
formatting for valid values.
In `@src/aiconfigurator/sdk/backends/base_backend.py`:
- Around line 303-327: The new run_encoder_static composition lacks direct unit
coverage. Add a backend-level test using a vision-language model with
encoder_colocated=False that mocks _run_encoder_phase and
_get_encoder_component_memory_for_runtime, then verifies latency correction,
phase-average power, and returned memory; also cover zero-latency/no-encoder
behavior and the related run_static and run_static_latency_only branches.
In `@src/aiconfigurator/sdk/task_v2.py`:
- Around line 1138-1155: Add an early validation in _validate_disagg for
enable_epd that requires both image workload dimensions, image_height and
image_width, to be configured; raise a clear ValueError when either is missing,
before any model or backend work occurs. Keep the existing encoder candidate
validation intact and align the message with the CLI requirement that EPD needs
--image-height/--image-width.
In `@tests/unit/sdk/sweep/test_sweep.py`:
- Around line 209-346: Add tests alongside
test_sweep_disagg_epd_composes_encoder_stage covering both negative EPD paths:
assert sweep_disagg raises ValueError when enable_epd=True and autoscale=True,
and mock encoder candidates whose encoder_latency exceeds the TTFT budget, then
assert the resulting encoder_choices path is empty and handled as expected.
Reuse the existing common_kwargs and candidate monkeypatches, referencing
sweep_disagg and _get_encoder_worker_candidates.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8658d488-ec38-4f07-8ecc-09db5bd8d9e5
📒 Files selected for processing (10)
src/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/common.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/picking.pysrc/aiconfigurator/sdk/sweep.pysrc/aiconfigurator/sdk/task_v2.pytests/unit/sdk/sweep/test_sweep.py
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: CodeRabbit / Review
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build and Test (e2e)
- GitHub Check: Cargo Deny
- GitHub Check: Build and Test (unit)
- GitHub Check: AIC Accuracy Regression Testing
🧰 Additional context used
📓 Path-based instructions (5)
**/*.py
📄 CodeRabbit inference engine (.claude/rules/generator/guard_rails.md)
**/*.py: Catch unsupported system/backend/version combinations early with fast validation to prevent wasting user time on late failures after profiling.
MoE model detection must useconfig.json, not model name patterns, to correctly identify MoE models with non-standard naming.
Usemodel_familynotmodel_namefor model compatibility checks;model_nameincludes size/quantization info that changes across variants.
safe_model_namemust be generated BEFORE saving results to avoid race condition with rawmodel_pathcontaining slashes.
Quantization type must be inferred from model config, not name patterns, to handle inconsistent naming conventions.
Files:
src/aiconfigurator/sdk/picking.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/common.pytests/unit/sdk/sweep/test_sweep.pysrc/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/sweep.py
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
src/aiconfigurator/sdk/picking.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/common.pytests/unit/sdk/sweep/test_sweep.pysrc/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/sweep.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/picking.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/common.pysrc/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/sweep.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/sweep/test_sweep.py
src/aiconfigurator/cli/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/cli/**: - Check that CLI argument changes preserve backward compatibility, validation behavior, defaults, and plain-output expectations.
- For new or changed user-facing options, verify docs updates and generator/SDK wiring.
- Watch for non-TTY assumptions in tests or output formatting.
Files:
src/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.py
🧠 Learnings (3)
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/picking.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/common.pysrc/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/sweep.py
📚 Learning: 2026-06-26T03:59:40.944Z
Learnt from: YijiaZhao
Repo: ai-dynamo/aiconfigurator PR: 1247
File: src/aiconfigurator/sdk/models/hybrid_moe.py:0-0
Timestamp: 2026-06-26T03:59:40.944Z
Learning: When implementing SGLang prefill CP modeling paths (e.g., for SWA/windowed attention) in these model files, size any CP all-gather communication (such as the payload for `cp_allgather_and_save_kv_cache`) by the full new-token sequence length being prefetched, not by the sliding-window cap. Rationale (SGLang v0.5.13): `cp_allgather_and_save_kv_cache` gathers the full per-layer new-token KV and writes the full result to each rank’s KV pool; the sliding-window limit should constrain only how much KV is stored/kept, not the volume of the all-gather communication.
Applied to files:
src/aiconfigurator/sdk/models/qwen3vl.py
📚 Learning: 2026-03-17T07:41:41.843Z
Learnt from: YijiaZhao
Repo: ai-dynamo/aiconfigurator PR: 600
File: src/aiconfigurator/sdk/backends/vllm_backend.py:42-43
Timestamp: 2026-03-17T07:41:41.843Z
Learning: In src/aiconfigurator/sdk/backends/vllm_backend.py, do not include seq_imbalance_correction_scale or gen_seq_imbalance_correction_scale in the AGG cache key isl/osl/b/ctx_tokens. These correction scales are per-instance and constant for a given usage of VLLMBackend, so they do not affect cache hits. Update the cache key computation to exclude these fields and adjust any tests that validate the cache key content. This guidance targets this file; apply similar logic to other backends only if they share the same per-instance scale semantics.
Applied to files:
src/aiconfigurator/sdk/backends/base_backend.py
🔇 Additional comments (13)
src/aiconfigurator/sdk/config.py (1)
32-35: LGTM!src/aiconfigurator/sdk/task_v2.py (1)
447-456: LGTM!Also applies to: 1129-1130, 1456-1459
src/aiconfigurator/cli/main.py (1)
274-287: LGTM!Also applies to: 1380-1381, 2408-2409
src/aiconfigurator/cli/report_and_save.py (1)
148-152: LGTM!Also applies to: 175-176, 253-260
src/aiconfigurator/sdk/common.py (1)
784-790: LGTM!src/aiconfigurator/sdk/models/qwen3vl.py (1)
45-54: LGTM!Also applies to: 85-94
src/aiconfigurator/sdk/backends/base_backend.py (1)
303-327: LGTM!Also applies to: 525-527, 600-602
src/aiconfigurator/sdk/sweep.py (5)
36-36: LGTM!Also applies to: 62-64, 81-87, 207-212
657-773: LGTM! Verified the TTFT/power/GPU-dilution arithmetic in_overlay_encoder_stageagainsttest_sweep_disagg_epd_composes_encoder_stage's expected numbers — all reconcile correctly.
774-915: LGTM! Per-encoder-choice TTFT budgeting and theenc is Nonevs EPD branch in the matching loop check out against the test assertions.
918-1077: LGTM!Also applies to: 1182-1215
1142-1181: 🎯 Functional CorrectnessIndependent GPU caps are gated together here. If
max_prefill_gpusormax_decode_gpusis meant to work on its own, this should enforce each bound separately; otherwise a single cap is silently ignored. The loop/budget check also looks duplicated with_match_workers, so a small shared helper would keep the two paths aligned.src/aiconfigurator/sdk/picking.py (1)
155-161: LGTM!
AIC Accuracy Regression TestingWorkflow result: Old revision: Summary
Comparison summary table
|
Add encode/prefill/decode (EPD) modeling to sweep_disagg for vision-language models, mirroring how real EPD deployments work (e.g. SGLang --encoder-only / --language-only): - enable_epd (Task field / CLI --enable-epd) runs the vision encoder on dedicated encode workers instead of colocated with prefill. - Encode-worker candidates are enumerated over encoder tp (default 1/2/4/8, mirroring the encode instance's --tp-size) x batch (default 1..32, modeling cross-request image batching); per-worker throughput is batch / encode latency via the new BaseBackend.run_encoder_static. - Prefill workers become language-only: ModelConfig.encoder_colocated= False skips building the ViT ops, so vision tokens still extend the context (via _visual_context_tokens) but no encoder compute or memory lands on the prefill worker. - Encode -> prefill is sequential per request, so each encode choice spends its batch latency out of the TTFT budget before the prefill SLA filter, and the composed row reports ttft = encode latency + corrected prefill ttft. - Rate matching gains the third pool: for each (p, d) worker pair the encode pool is sized to the smallest count that does not bottleneck the rate-matched throughput, with encode GPUs counted in the per-GPU objective and the replica GPU budget. Results fill the existing (e)* columns plus a new (e)bs column; the CLI report renders them when an encode pool is present. The plain PD path is unchanged (single no-encoder choice with the full TTFT budget); parity between sweep._rate_match_dict and picking._build_disagg_summary_dict is preserved. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: wenpengw-nv <wenpengw@nvidia.com>
f419f72 to
a21d0c5
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/aiconfigurator/cli/report_and_save.py`:
- Around line 310-317: Include context parallelism in both CLI GPU breakdown
equations: at src/aiconfigurator/cli/report_and_save.py lines 310-317, multiply
the aggregate worker GPU term in the e_workers branch by row.get("cp", 1); at
lines 220-225, multiply the prefill and decode worker terms by their respective
CP fields. Preserve the existing formatting and worker-count calculations.
In `@src/aiconfigurator/sdk/sweep.py`:
- Around line 846-851: Update the E+agg cell-matching flow around the shown
worker loop to accept the deployment GPU budget from Task.total_gpus, and skip
candidates whose computed num_gpu exceeds that budget before evaluating best.
Wire the new budget through the Task-to-sweep call chain and add or update tests
covering oversized cells and valid selections.
- Around line 842-843: Update the filtering condition in the sweep configuration
flow around encoder_latency and ttft_target so configurations whose combined
latency exactly equals the target are allowed, while values exceeding the target
remain rejected. Add a focused regression test covering the exact-boundary case.
🪄 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: a6f2d206-40d7-40bb-bf4f-bf72d460c638
📒 Files selected for processing (5)
src/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/sweep.pysrc/aiconfigurator/sdk/task_v2.pytests/unit/sdk/sweep/test_sweep.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/aiconfigurator/cli/main.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (7)
**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
**/*: Treat this file as the only always-injected rule file; do not add new always-on rules here without human approval. New rule files must includepaths:frontmatter.
When reviewing changes in a governed area, read that area's rule files first, even if path-based auto-loading does not occur during a read-only review. Rule violations are review findings even when the code works.
A task must remain within its module: collector tasks may change onlycollector/and its tests, generator tasks onlysrc/aiconfigurator/generator/and its tests, and SDK tasks must not modify either. Cross-module contract changes require explicit human approval.
Files:
src/aiconfigurator/sdk/task_v2.pytests/unit/sdk/sweep/test_sweep.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/sweep.py
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
src/aiconfigurator/sdk/task_v2.pytests/unit/sdk/sweep/test_sweep.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/sweep.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
Run
ruff check .andruff format --check .to validate Python linting and formatting.
Files:
src/aiconfigurator/sdk/task_v2.pytests/unit/sdk/sweep/test_sweep.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/sweep.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/sweep.py
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/sweep/test_sweep.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; usepytest -m "unit or build"for the build-test subset when required data is available.
Files:
tests/unit/sdk/sweep/test_sweep.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/sweep/test_sweep.py
src/aiconfigurator/cli/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/cli/**: - Check that CLI argument changes preserve backward compatibility, validation behavior, defaults, and plain-output expectations.
- For new or changed user-facing options, verify docs updates and generator/SDK wiring.
- Watch for non-TTY assumptions in tests or output formatting.
Files:
src/aiconfigurator/cli/report_and_save.py
🧠 Learnings (1)
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/task_v2.pysrc/aiconfigurator/sdk/sweep.py
🔇 Additional comments (2)
src/aiconfigurator/sdk/task_v2.py (1)
449-460: LGTM!Also applies to: 1183-1186, 1436-1439, 1492-1495
src/aiconfigurator/cli/report_and_save.py (1)
148-175: LGTM!Also applies to: 252-281, 331-345
| if enc["encoder_latency"] + r["ttft"] >= ttft_target: | ||
| continue |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Allow configurations that exactly meet the TTFT target.
Line 842 rejects equality, although SLA filtering elsewhere treats the target as inclusive.
Proposed fix
- if enc["encoder_latency"] + r["ttft"] >= ttft_target:
+ if enc["encoder_latency"] + r["ttft"] > ttft_target:
continuePlease add an exact-boundary regression test. As per path instructions, prefer an inline fix when it is clear and limited to the hunk.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if enc["encoder_latency"] + r["ttft"] >= ttft_target: | |
| continue | |
| if enc["encoder_latency"] + r["ttft"] > ttft_target: | |
| continue |
🤖 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/aiconfigurator/sdk/sweep.py` around lines 842 - 843, Update the filtering
condition in the sweep configuration flow around encoder_latency and ttft_target
so configurations whose combined latency exactly equals the target are allowed,
while values exceeding the target remain rejected. Add a focused regression test
covering the exact-boundary case.
Source: Path instructions
| for a_num in range(1, _MAX_AGG_WORKERS_EPD + 1): | ||
| e_num = max(1, math.ceil(rate_one * a_num / encoder_capacity)) | ||
| num_gpu = gpus_one * a_num + enc["num_total_gpus"] * e_num | ||
| key = (rate_one * a_num / num_gpu, -num_gpu) | ||
| if best is None or key > best[0]: | ||
| best = (key, a_num, e_num) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Constrain E+agg cells by the deployment GPU budget.
This loop may select a cell larger than Task.total_gpus; downstream reporting then computes zero replicas. Pass the budget into this matcher and skip cells where num_gpu exceeds it. A one-click fix is unsafe because this requires Task-to-sweep contract wiring and tests.
🤖 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/aiconfigurator/sdk/sweep.py` around lines 846 - 851, Update the E+agg
cell-matching flow around the shown worker loop to accept the deployment GPU
budget from Task.total_gpus, and skip candidates whose computed num_gpu exceeds
that budget before evaluating best. Wire the new budget through the
Task-to-sweep call chain and add or update tests covering oversized cells and
valid selections.
enable_epd on serving_mode='agg' runs the vision encoder on dedicated encode workers while prefill+decode stay aggregated: - sweep_agg(enable_epd): agg workers become language-only (encoder_colocated=False, vision tokens stay in context); encode candidates reuse _get_encoder_worker_candidates (tp x batch) - _rate_match_agg_epd: per (agg row, encode choice) pick the integer cell of (a)workers agg + (e)workers encode workers (encode pool sized to not bottleneck, ties keep the smaller cell); TTFT/(e)* overlay reuses _overlay_encoder_stage - enable_epd stays the only EPD switch in both serving modes; the encoder candidate lists are pure search-space knobs (check moved to the shared validate()) - default (enable_epd=False) keeps the encoder inline in agg and disagg - CLI --enable-epd now turns the agg experiment into E+agg; the agg report table renders (a)workers/(e)workers/(e)tp/(e)bs and the gpus/replica cell breakdown Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: wenpengw-nv <wenpengw@nvidia.com>
- ModelConfig.encoder_colocated -> language_only: the field describes the deployed worker (mirrors SGLang --language-only), not the model; hosting the ViT is a weight-placement property like tp_size - image H/W resolve via smart resize (round to the nearest patch*merge multiple, as upstream VL processors) instead of floor - encoder batch schedule capped at 8 (SGLANG_ENCODER_MAX_BATCH_SIZE default); larger batches don't exist in the deployed system - vision tokens enter the batch boundaries: agg ctx-token grid and feasibility guards use the effective ISL (restores legacy find_best_agg parity), disagg prefill batch bound and the Task-side token budget divide/derive by the same effective ISL - symmetric TTFT queueing correction: the disagg EPD encode stage carries the same autoscale correction factor as prefill (the inline PD baseline corrects its whole E+P ttft); the agg path stays raw, matching run_agg's inline convention - hetero encoder placement: Task.encoder_system_name puts the encode pool on its own system (GPU type); backend/version follow the P/agg side - cli_default Python API exposes enable_epd / encoder_tp / encoder_system; new --encoder-system CLI flag Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: wenpengw-nv <wenpengw@nvidia.com>
234b286 to
7f72810
Compare
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/aiconfigurator/sdk/task_v2.py`:
- Around line 1199-1206: Update validate() near the encoder_knobs_set check to
reject enable_epd=True when the task is text-only (non-VL), raising a clear
ValueError before sweep or model enumeration. Preserve the existing validation
for encoder_* settings and allow enable_epd for VL tasks.
🪄 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: 1718f5c2-5eee-4bfa-9f59-71f8b24f83e2
📒 Files selected for processing (10)
src/aiconfigurator/cli/api.pysrc/aiconfigurator/cli/main.pysrc/aiconfigurator/cli/report_and_save.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/sweep.pysrc/aiconfigurator/sdk/task_v2.pytests/unit/sdk/backends/test_encoder.pytests/unit/sdk/sweep/test_sweep.py
🚧 Files skipped from review as they are similar to previous changes (4)
- src/aiconfigurator/cli/main.py
- src/aiconfigurator/cli/report_and_save.py
- tests/unit/sdk/sweep/test_sweep.py
- src/aiconfigurator/sdk/sweep.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (7)
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/backends/test_encoder.py
**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
**/*: Treat this file as the only always-injected rule file; do not add new always-on rules here without human approval. New rule files must includepaths:frontmatter.
When reviewing changes in a governed area, read that area's rule files first, even if path-based auto-loading does not occur during a read-only review. Rule violations are review findings even when the code works.
A task must remain within its module: collector tasks may change onlycollector/and its tests, generator tasks onlysrc/aiconfigurator/generator/and its tests, and SDK tasks must not modify either. Cross-module contract changes require explicit human approval.
Files:
tests/unit/sdk/backends/test_encoder.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/cli/api.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/task_v2.py
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
tests/unit/sdk/backends/test_encoder.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/cli/api.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/task_v2.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
Run
ruff check .andruff format --check .to validate Python linting and formatting.
Files:
tests/unit/sdk/backends/test_encoder.pysrc/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/cli/api.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/task_v2.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; usepytest -m "unit or build"for the build-test subset when required data is available.
Files:
tests/unit/sdk/backends/test_encoder.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/backends/test_encoder.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/task_v2.py
src/aiconfigurator/cli/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/cli/**: - Check that CLI argument changes preserve backward compatibility, validation behavior, defaults, and plain-output expectations.
- For new or changed user-facing options, verify docs updates and generator/SDK wiring.
- Watch for non-TTY assumptions in tests or output formatting.
Files:
src/aiconfigurator/cli/api.py
🧠 Learnings (3)
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/models/qwen3vl.pysrc/aiconfigurator/sdk/config.pysrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/task_v2.py
📚 Learning: 2026-06-26T03:59:40.944Z
Learnt from: YijiaZhao
Repo: ai-dynamo/aiconfigurator PR: 1247
File: src/aiconfigurator/sdk/models/hybrid_moe.py:0-0
Timestamp: 2026-06-26T03:59:40.944Z
Learning: When implementing SGLang prefill CP modeling paths (e.g., for SWA/windowed attention) in these model files, size any CP all-gather communication (such as the payload for `cp_allgather_and_save_kv_cache`) by the full new-token sequence length being prefetched, not by the sliding-window cap. Rationale (SGLang v0.5.13): `cp_allgather_and_save_kv_cache` gathers the full per-layer new-token KV and writes the full result to each rank’s KV pool; the sliding-window limit should constrain only how much KV is stored/kept, not the volume of the all-gather communication.
Applied to files:
src/aiconfigurator/sdk/models/qwen3vl.py
📚 Learning: 2026-03-17T07:41:41.843Z
Learnt from: YijiaZhao
Repo: ai-dynamo/aiconfigurator PR: 600
File: src/aiconfigurator/sdk/backends/vllm_backend.py:42-43
Timestamp: 2026-03-17T07:41:41.843Z
Learning: In src/aiconfigurator/sdk/backends/vllm_backend.py, do not include seq_imbalance_correction_scale or gen_seq_imbalance_correction_scale in the AGG cache key isl/osl/b/ctx_tokens. These correction scales are per-instance and constant for a given usage of VLLMBackend, so they do not affect cache hits. Update the cache key computation to exclude these fields and adjust any tests that validate the cache key content. This guidance targets this file; apply similar logic to other backends only if they share the same per-instance scale semantics.
Applied to files:
src/aiconfigurator/sdk/backends/base_backend.py
🔇 Additional comments (10)
src/aiconfigurator/sdk/config.py (1)
32-36: LGTM!src/aiconfigurator/sdk/models/qwen3vl.py (1)
51-54: LGTM!Also applies to: 93-94
src/aiconfigurator/sdk/backends/base_backend.py (4)
229-246: LGTM!
308-332: LGTM!
530-543: 🎯 Functional CorrectnessCorrect fix: language-only workers no longer lose vision context tokens.
Previously, when
model.encoder_opswas empty (EPD language-only worker) the old code path fell through to_run_encoder_phase, which early-returnsimg_ctx_tokens=0— silently dropping vision tokens from the LLM context length for language-only workers. Routing this case through_visual_context_tokens(keyed onencoder_config, notencoder_ops) fixes that.
605-626: LGTM! Mirrors the same language-only vision-context fix asrun_static_latency_onlyabove.tests/unit/sdk/backends/test_encoder.py (1)
545-572: LGTM!src/aiconfigurator/sdk/task_v2.py (2)
449-465: LGTM!Also applies to: 1092-1102, 1434-1461, 1463-1523, 1578-1602
1604-1611: 🗄️ Data Integrity & IntegrationNo separate encoder-version lookup needed
encoder_system_nameintentionally reuses the prefill/agg backend and version, so resolving a second version here would be redundant.> Likely an incorrect or invalid review comment.src/aiconfigurator/cli/api.py (1)
148-150: LGTM! Wiring tobuild_default_tasks/Taskis consistent; the "requires image workload" gap is flagged at its root cause intask_v2.py'svalidate().Also applies to: 198-203, 263-265
| # EPD is switched on explicitly and only by enable_epd (both serving | ||
| # modes); every other encoder_* field is a pure search-space / | ||
| # placement knob and must not act as an implicit switch. | ||
| encoder_knobs_set = ( | ||
| self.encoder_tp_candidates or self.encoder_batch_candidates or self.encoder_system_name | ||
| ) | ||
| if encoder_knobs_set and not self.enable_epd: | ||
| raise ValueError("encoder_* settings require enable_epd=True.") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== task_v2 relevant symbols ==\n'
rg -n "enable_epd|image_height|image_width|encoder_tp_candidates|encoder_batch_candidates|encoder_system_name" src/aiconfigurator/sdk/task_v2.py src/aiconfigurator/cli/api.py docs -S || true
printf '\n== all enable_epd references ==\n'
rg -n "enable_epd" src docs -S || trueRepository: ai-dynamo/aiconfigurator
Length of output: 6780
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== task_v2 validation window ==\n'
sed -n '1170,1225p' src/aiconfigurator/sdk/task_v2.py | cat -n
printf '\n== task_v2 encoder DB window ==\n'
sed -n '1598,1620p' src/aiconfigurator/sdk/task_v2.py | cat -n
printf '\n== sweep EPD windows ==\n'
sed -n '396,490p' src/aiconfigurator/sdk/sweep.py | cat -n
printf '\n---\n'
sed -n '1070,1185p' src/aiconfigurator/sdk/sweep.py | cat -nRepository: ai-dynamo/aiconfigurator
Length of output: 16513
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== encoder candidate helper ==\n'
rg -n "def _get_encoder_worker_candidates|visual|image_height|image_width|language_only|enable_epd" src/aiconfigurator/sdk/sweep.py src/aiconfigurator/sdk/task_v2.py -S
printf '\n== helper definition window ==\n'
start=$(rg -n "def _get_encoder_worker_candidates" src/aiconfigurator/sdk/sweep.py | head -n1 | cut -d: -f1)
if [ -n "${start:-}" ]; then
sed -n "${start},$((start+140))p" src/aiconfigurator/sdk/sweep.py | cat -n
fi
printf '\n== visual model helpers ==\n'
rg -n "visual_context|image.*workload|vl|vision encoder|encoder_config|extra_params" src/aiconfigurator/sdk -SRepository: ai-dynamo/aiconfigurator
Length of output: 44986
Reject enable_epd for text-only tasks in validate()
_get_encoder_worker_candidates() already fails later with a less direct error, but this should be caught here so enable_epd=True on a non-VL task surfaces a clear ValueError before sweep/model enumeration.
🛡️ Suggested Change
if encoder_knobs_set and not self.enable_epd:
raise ValueError("encoder_* settings require enable_epd=True.")
+ if self.enable_epd and not (self.image_height and self.image_width):
+ raise ValueError("enable_epd=True requires an image workload (image_height/image_width).")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # EPD is switched on explicitly and only by enable_epd (both serving | |
| # modes); every other encoder_* field is a pure search-space / | |
| # placement knob and must not act as an implicit switch. | |
| encoder_knobs_set = ( | |
| self.encoder_tp_candidates or self.encoder_batch_candidates or self.encoder_system_name | |
| ) | |
| if encoder_knobs_set and not self.enable_epd: | |
| raise ValueError("encoder_* settings require enable_epd=True.") | |
| # EPD is switched on explicitly and only by enable_epd (both serving | |
| # modes); every other encoder_* field is a pure search-space / | |
| # placement knob and must not act as an implicit switch. | |
| encoder_knobs_set = ( | |
| self.encoder_tp_candidates or self.encoder_batch_candidates or self.encoder_system_name | |
| ) | |
| if encoder_knobs_set and not self.enable_epd: | |
| raise ValueError("encoder_* settings require enable_epd=True.") | |
| if self.enable_epd and not (self.image_height and self.image_width): | |
| raise ValueError("enable_epd=True requires an image workload (image_height/image_width).") |
🤖 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/aiconfigurator/sdk/task_v2.py` around lines 1199 - 1206, Update
validate() near the encoder_knobs_set check to reject enable_epd=True when the
task is text-only (non-VL), raising a clear ValueError before sweep or model
enumeration. Preserve the existing validation for encoder_* settings and allow
enable_epd for VL tasks.
1 participant
Overview:
Adds encode/prefill/decode (EPD) modeling for vision-language models to the disagg sweep, mirroring how real EPD deployments work.
Details:
--enable-epd(Task fieldenable_epd): the vision encoder runs on dedicated encode workers instead of colocated with prefill.1/2/4/8) × batch (default1..32, modeling cross-request image batching). Per-worker throughput = batch / encode latency.ModelConfig.encoder_colocated=False).(e)*columns (+ new(e)bs).The plain PD path is unchanged when EPD is off.
Usage
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
Bug Fixes
Tests