feat(sdk): model PP stage imbalance and pipe fill instead of ideal linear speedup - #1459
feat(sdk): model PP stage imbalance and pipe fill instead of ideal linear speedup#1459tianhaox wants to merge 5 commits into
Conversation
…near speedup PP was modeled as an ideal linear speedup: a step's latency was the whole model's (`_num_layers` is never divided by `pp_size`) and throughput was simply multiplied by `pp_size`. That is exact only when every stage costs the same and the pipe is always full. A pipeline advances at `max_i(t_i)`, not the average. The un-sharded `lm_head` sits alone on the last stage, so for a decode step it sets the cycle: Qwen3-32B on H20 at pp=8 runs at 83.2% of the ideal cycle (6.65x, not 8x). Prefill steps stay near-ideal at 99.9% because per-layer work dominates there. Uneven layer splits previously only produced a warning saying "we're nothing to correct this", and a starved pipe still scaled linearly. Add `PipelineSchedule`, which folds the step model's existing per-op breakdown into per-stage times and returns: balance_factor = (step_total / pp_size) / cycle_time # latency effect fill_factor = min(1, num_microbatches / pp_size) # throughput effect `run_agg` inflates the mix/genonly step latencies by 1/balance_factor to get the real traversal time, so TTFT, TPOT and throughput all follow consistently downstream, and charges fill_factor against throughput only. This is orchestration, not op query math, so the Rust engine-step parity surface is untouched and both engines benefit -- the same layering AFD's pipeline model already uses. `pp_size == 1` returns exactly 1.0 from both factors, so single-stage results are bit-identical. Known gaps are documented in docs/PIPELINE_PARALLEL_MODELING.md: fill_factor needs silicon calibration, chunked-prefill x PP is not yet modeled, the P2P op still hardcodes inter_node_bw, and PP remains excluded from the automatic search. Signed-off-by: tianhaox <tianhaox@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📜 Recent review details⏰ Context from checks skipped due to timeout. (10)
🧰 Additional context used📓 Path-based instructions (2)**/*⚙️ CodeRabbit configuration file
Files:
docs/**⚙️ CodeRabbit configuration file
Files:
🔇 Additional comments (4)
WalkthroughPipeline parallel modeling now includes stage placement, layer partitioning, steady-state scheduling, P2P costs, balance and fill factors, backend integration, diagnostics, tests, and documentation. ChangesPipeline Parallel Modeling
Estimated code review effort: 4 (Complex) | ~45 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.
Actionable comments posted: 6
🤖 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 `@aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py`:
- Around line 1533-1537: The PP correction currently derives from Rust aggregate
synthetic operations, losing placement-specific embedding, logits, and P2P
costs. Update the schedule calculation around _pipeline_schedule, run_mixed, and
_get_genonly_step_latency to use a placement-preserving per-operation breakdown,
or an equivalent Python-only breakdown solely for balance computation, while
preserving Rust aggregate results for execution. Add integration coverage with
pp_size > 1 for both engine-step backends and verify equivalent corrections.
- Around line 1522-1537: Separate pipeline traversal latency from steady-state
cycle time: in base_backend.py lines 1522-1537, update the backend latency
calculations so TTFT and request TPOT use stage traversal time, while throughput
alone uses the balance-adjusted steady-state cycle. In
aic-core/src/aiconfigurator_core/sdk/pipeline.py lines 161-165, apply P2P
overhead only across actual inter-stage edges and omit the final-stage hop. In
tests/unit/sdk/test_pipeline_schedule.py lines 169-176, add coverage for a P2P
schedule with a last-stage bottleneck and assert distinct traversal and
throughput-cycle values. In docs/PIPELINE_PARALLEL_MODELING.md lines 110-117,
document the corrected traversal-versus-cycle semantics instead of universally
equating traversal with pp times cycle.
- Around line 126-128: Update the warning logic in the scheduling method around
warn_on_unclassified_ops so it inspects both model.context_ops and
model.generation_ops when mixed-step scheduling is active. Ensure unclassified
operations from either pipeline are included in the placement warning while
preserving the existing layer-count and schedule behavior.
In `@aic-core/src/aiconfigurator_core/sdk/pipeline.py`:
- Around line 120-121: Update layer_partition to validate an explicitly
configured self.partition before returning it: reject negative layer counts and
partitions whose sum differs from num_layers, while preserving
even_partition(num_layers, self.pp_size) for unset partitions. Ensure validation
remains compatible with generator-provided partition inputs.
In `@docs/PIPELINE_PARALLEL_MODELING.md`:
- Line 41: Add the appropriate language identifier, such as text, to both fenced
code blocks in the pipeline parallel modeling documentation, including the
blocks containing stage timing output and the balance_factor expression, while
preserving their contents.
In `@tests/unit/sdk/test_pipeline_schedule.py`:
- Around line 16-26: Update the import block in test_pipeline_schedule.py using
the repository-configured Ruff ordering, equivalent to running ruff check --fix
on the file. Preserve the imported symbols and commit the resulting
section/order changes.
🪄 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: 733fa9f1-2b8d-4a11-abc3-0017562323f6
📒 Files selected for processing (4)
aic-core/src/aiconfigurator_core/sdk/backends/base_backend.pyaic-core/src/aiconfigurator_core/sdk/pipeline.pydocs/PIPELINE_PARALLEL_MODELING.mdtests/unit/sdk/test_pipeline_schedule.py
📜 Review details
⏰ Context from checks skipped due to timeout. (10)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Build wheels (manylinux_2_28_x86_64)
- GitHub Check: Build wheels (macosx_arm64)
- GitHub Check: Build wheels (manylinux_2_28_aarch64)
- GitHub Check: Build and Test (unit)
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Cargo Deny
- GitHub Check: Build and Test (e2e)
- GitHub Check: aic-core public API contract
🧰 Additional context used
📓 Path-based instructions (4)
**/*
⚙️ 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:
docs/PIPELINE_PARALLEL_MODELING.mdaic-core/src/aiconfigurator_core/sdk/pipeline.pytests/unit/sdk/test_pipeline_schedule.pyaic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
docs/**
⚙️ CodeRabbit configuration file
docs/**: - Check that docs match changed CLI, SDK, generator, backend, and support-matrix behavior.
- Flag docs that describe unsupported runtimes, stale command names, or behavior not covered by tests or support-matrix evidence.
Files:
docs/PIPELINE_PARALLEL_MODELING.md
aic-core/src/aiconfigurator_core/sdk/**
⚙️ CodeRabbit configuration file
aic-core/src/aiconfigurator_core/sdk/**: - Verify core SDK API changes remain compatible with legacy aiconfigurator.sdk imports, generator inputs, profiler data flow, and documented examples.
- Flag upper-layer dependencies or silent schema drift introduced into the minimal core distribution.
Files:
aic-core/src/aiconfigurator_core/sdk/pipeline.pyaic-core/src/aiconfigurator_core/sdk/backends/base_backend.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/test_pipeline_schedule.py
🪛 GitHub Actions: Lint and Format / 1_Lint and Format (Ruff).txt
tests/unit/sdk/test_pipeline_schedule.py
[error] 16-26: Ruff I001: Import block is unsorted or unformatted. Run 'ruff check --fix .' to organize imports.
🪛 GitHub Actions: Lint and Format / Lint and Format (Ruff)
tests/unit/sdk/test_pipeline_schedule.py
[error] 16-26: Ruff I001: Import block is unsorted or unformatted. Organize the imports or run 'ruff check --fix'. Command failed with exit code 1.
🪛 markdownlint-cli2 (0.23.1)
docs/PIPELINE_PARALLEL_MODELING.md
[warning] 41-41: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
[warning] 71-71: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (1)
aic-core/src/aiconfigurator_core/sdk/pipeline.py (1)
85-105: 🗄️ Data Integrity & IntegrationKeep
PipelineSchedulecore-only.The legacy backend aliases the core backend, and documentation and tests use
aiconfigurator_core.sdk.pipeline.PipelineSchedule. No legacy re-export is required.
| if pp_size > 1: | ||
| warn_on_unclassified_ops(model.generation_ops, model._num_layers) | ||
| return schedule |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Inspect context operations for unclassified placement.
Mixed-step scheduling consumes context and generation operation latency. This warning only inspects model.generation_ops, so a low-scale unmarked context operation can still be silently distributed across all stages.
Proposed fix
if pp_size > 1:
- warn_on_unclassified_ops(model.generation_ops, model._num_layers)
+ warn_on_unclassified_ops([*model.context_ops, *model.generation_ops], model._num_layers)As per path instructions, “Model operation pipelines are defined through context_ops and generation_ops.”
📝 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 pp_size > 1: | |
| warn_on_unclassified_ops(model.generation_ops, model._num_layers) | |
| return schedule | |
| if pp_size > 1: | |
| warn_on_unclassified_ops([*model.context_ops, *model.generation_ops], model._num_layers) | |
| return schedule |
🤖 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 `@aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py` around lines
126 - 128, Update the warning logic in the scheduling method around
warn_on_unclassified_ops so it inspects both model.context_ops and
model.generation_ops when mixed-step scheduling is active. Ensure unclassified
operations from either pipeline are included in the placement warning while
preserving the existing layer-count and schedule behavior.
Source: Path instructions
| def layer_partition(self, num_layers: int) -> tuple[int, ...]: | ||
| return self.partition if self.partition is not None else even_partition(num_layers, self.pp_size) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate explicit partition contents.
Line 121 accepts negative partitions and partitions whose sum differs from num_layers. These values produce invalid stage times and incorrect balance factors.
Proposed fix
def layer_partition(self, num_layers: int) -> tuple[int, ...]:
- return self.partition if self.partition is not None else even_partition(num_layers, self.pp_size)
+ if self.partition is None:
+ return even_partition(num_layers, self.pp_size)
+ if any(layers < 0 for layers in self.partition) or sum(self.partition) != num_layers:
+ raise ValueError(f"partition must contain non-negative entries totaling {num_layers}")
+ return self.partitionAs per path instructions, “Verify core SDK API changes remain compatible with … generator inputs.”
📝 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.
| def layer_partition(self, num_layers: int) -> tuple[int, ...]: | |
| return self.partition if self.partition is not None else even_partition(num_layers, self.pp_size) | |
| def layer_partition(self, num_layers: int) -> tuple[int, ...]: | |
| if self.partition is None: | |
| return even_partition(num_layers, self.pp_size) | |
| if any(layers < 0 for layers in self.partition) or sum(self.partition) != num_layers: | |
| raise ValueError(f"partition must contain non-negative entries totaling {num_layers}") | |
| return self.partition |
🤖 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 `@aic-core/src/aiconfigurator_core/sdk/pipeline.py` around lines 120 - 121,
Update layer_partition to validate an explicitly configured self.partition
before returning it: reject negative layer counts and partitions whose sum
differs from num_layers, while preserving even_partition(num_layers,
self.pp_size) for unset partitions. Ensure validation remains compatible with
generator-provided partition inputs.
Source: Path instructions
| its cost is a large fraction of a single stage's layer work, so it sets the | ||
| cycle. Qwen3-32B (64 layers, vocab 151936) on H20 at `pp=8`: | ||
|
|
||
| ``` |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add language identifiers to fenced code blocks.
markdownlint reports MD040 for both fences.
Proposed fix
-```
+```text
stage times (ms): [3.153, 3.150, 3.150, 3.150, 3.150, 3.150, 3.150, 3.901]
@@
-```
+```text
balance_factor = (step_total / pp_size) / cycle_timeAlso applies to: 71-71
🧰 Tools
🪛 markdownlint-cli2 (0.23.1)
[warning] 41-41: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 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 `@docs/PIPELINE_PARALLEL_MODELING.md` at line 41, Add the appropriate language
identifier, such as text, to both fenced code blocks in the pipeline parallel
modeling documentation, including the blocks containing stage timing output and
the balance_factor expression, while preserving their contents.
Source: Linters/SAST tools
…yers `PipelineSchedule` conflated two things with different consumers: - **where the work lives** -- layer partition, op placement, per-stage compute times, per-hop link cost. Model and hardware knowledge. - **how the pipe runs** -- microbatch count, P2P overlap, and the closed-form `balance_factor` / `fill_factor` that AIC's mean-field step model needs. Those scalars are surrogates for what an event-driven simulator computes natively. The Dynamo Mocker embeds the compiled engine for per-iteration cost estimates and models request-level scheduling itself; if it consumed the schedule wholesale it would apply AIC's closed-form bubble penalty on top of its own simulated occupancy and charge for the same bubble twice. It needs the first half and must not inherit the second. Split into `PipelineLayout` (shared, no scheduling policy) and `PipelineSteadyState` (wraps a layout, adds occupancy). The backend hook splits to match: `_pipeline_layout` for the partition, `_pipeline_steady_state` for the microbatch policy, the latter composing the former so a backend that overrides only the partition affects both. Pure refactor -- no behavior change. Verified bit-identical end to end on H20 (tp8/pp1: TTFT 650.722 / TPOT 18.195 / 207.72 tok/s/gpu; tp1/pp8: 3717.314 / 89.919 / 333.21). Tests restructured along the same boundary, plus one that pins it: the layout must not expose occupancy-derived attributes. Doing this before the Mocker integration lands, while the only caller is `base_backend`, keeps the seam cheap to shape. Signed-off-by: tianhaox <tianhaox@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@aic-core/src/aiconfigurator_core/sdk/pipeline.py`:
- Around line 108-124: Add the legacy sdk pipeline compatibility module
corresponding to the core PipelineLayout API, and register "pipeline" in
CORE_SDK_LEAF_MODULES so it is included in cross-package compatibility checks.
Keep the wrapper’s public exports aligned with aiconfigurator_core.sdk.pipeline.
In `@docs/PIPELINE_PARALLEL_MODELING.md`:
- Around line 75-79: Update the two fenced code blocks in the documentation
hunk, including the blocks near stage_times and cycle_time, by adding the text
language identifier to each opening fence. Leave the block contents and closing
fences unchanged.
🪄 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: 5d1ba94b-ac66-464b-b357-8c0bbcbeea67
📒 Files selected for processing (4)
aic-core/src/aiconfigurator_core/sdk/backends/base_backend.pyaic-core/src/aiconfigurator_core/sdk/pipeline.pydocs/PIPELINE_PARALLEL_MODELING.mdtests/unit/sdk/test_pipeline_modeling.py
🚧 Files skipped from review as they are similar to previous changes (1)
- aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
📜 Review details
⏰ Context from checks skipped due to timeout. (10)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Build wheels (manylinux_2_28_x86_64)
- GitHub Check: Build wheels (manylinux_2_28_aarch64)
- GitHub Check: Cargo Deny
- GitHub Check: Build and Test (e2e)
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build wheels (macosx_arm64)
- GitHub Check: Build and Test (unit)
- GitHub Check: aic-core public API contract
🧰 Additional context used
📓 Path-based instructions (4)
**/*
⚙️ 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:
aic-core/src/aiconfigurator_core/sdk/pipeline.pytests/unit/sdk/test_pipeline_modeling.pydocs/PIPELINE_PARALLEL_MODELING.md
aic-core/src/aiconfigurator_core/sdk/**
⚙️ CodeRabbit configuration file
aic-core/src/aiconfigurator_core/sdk/**: - Verify core SDK API changes remain compatible with legacy aiconfigurator.sdk imports, generator inputs, profiler data flow, and documented examples.
- Flag upper-layer dependencies or silent schema drift introduced into the minimal core distribution.
Files:
aic-core/src/aiconfigurator_core/sdk/pipeline.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/test_pipeline_modeling.py
docs/**
⚙️ CodeRabbit configuration file
docs/**: - Check that docs match changed CLI, SDK, generator, backend, and support-matrix behavior.
- Flag docs that describe unsupported runtimes, stale command names, or behavior not covered by tests or support-matrix evidence.
Files:
docs/PIPELINE_PARALLEL_MODELING.md
🪛 markdownlint-cli2 (0.23.1)
docs/PIPELINE_PARALLEL_MODELING.md
[warning] 75-75: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
[warning] 86-86: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (2)
docs/PIPELINE_PARALLEL_MODELING.md (1)
65-68: LGTM!Also applies to: 81-84, 147-153, 197-204
aic-core/src/aiconfigurator_core/sdk/pipeline.py (1)
126-133: Validate explicit partition entries.
layer_partition()accepts negative entries and totals that differ fromnum_layers. This can create negative stage times or drop or duplicate layer work. Validate the entries before returning the explicit partition. Add tests for negative entries and invalid totals.Proposed change
def layer_partition(self, num_layers: int) -> tuple[int, ...]: - return self.partition if self.partition is not None else even_partition(num_layers, self.pp_size) + if self.partition is None: + return even_partition(num_layers, self.pp_size) + if any(layers < 0 for layers in self.partition) or sum(self.partition) != num_layers: + raise ValueError(f"partition must contain non-negative entries totaling {num_layers}") + return self.partitionAs per path instructions, “Verify core SDK API changes remain compatible with legacy aiconfigurator.sdk imports, generator inputs, profiler data flow, and documented examples.”
Source: Path instructions
| @dataclass(frozen=True) | ||
| class PipelineLayout: | ||
| """Which ops land on which stage, and what each stage costs. | ||
|
|
||
| Carries no scheduling policy: it answers "what does stage *i* cost for | ||
| this batch" and "what does one hop cost", nothing about occupancy. That | ||
| keeps it usable by both AIC's closed-form model and an event-driven | ||
| simulator that derives bubbles from occupancy itself. | ||
|
|
||
| Args: | ||
| pp_size: Number of pipeline stages. | ||
| partition: Explicit layers-per-stage. ``None`` uses | ||
| :func:`even_partition`. | ||
| """ | ||
|
|
||
| pp_size: int = 1 | ||
| partition: tuple[int, ...] | None = None |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline aic-core/src/aiconfigurator_core/sdk/pipeline.py --items all
# Locate legacy and core SDK import paths and public re-exports.
rg -n -C3 --glob '*.py' \
'from aiconfigurator(_core)?\.sdk|import aiconfigurator(_core)?\.sdk|PipelineLayout|PipelineSteadyState' .
# Trace configuration and profiler inputs that can construct or influence layouts.
rg -n -C3 --glob '*.py' \
'pipeline_microbatches|pp_size|partition|layer_partition|per_op_ms' aic-core tests
# Check documented examples for obsolete imports or incompatible PP claims.
rg -n -C3 \
'aiconfigurator\.sdk|PipelineLayout|PipelineSteadyState|pipeline_microbatches|pp_size' docs README.md 2>/dev/null || trueRepository: ai-dynamo/aiconfigurator
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- pipeline module outline ---'
ast-grep outline aic-core/src/aiconfigurator_core/sdk/pipeline.py --items all
echo '--- pipeline module ---'
sed -n '1,280p' aic-core/src/aiconfigurator_core/sdk/pipeline.py
echo '--- exact layout and pipeline symbols ---'
rg -n -C4 --glob '*.py' \
'\b(PipelineLayout|PipelineSteadyState|even_partition|pipeline_microbatches|per_op_ms)\b' \
aic-core/src aic-core/tests tests 2>/dev/null | head -n 1200
echo '--- legacy package and compatibility files ---'
git ls-files | rg '(^|/)(aiconfigurator|aic-core)/.*(sdk|compat|init).*\.py$' | head -n 400
echo '--- legacy-to-core imports and re-exports ---'
rg -n -C3 --glob '*.py' \
'aiconfigurator_core|aiconfigurator\.sdk|__all__|import_module|sys\.modules' \
aiconfigurator aic-core/src 2>/dev/null | head -n 1600
echo '--- docs and examples ---'
rg -n -C3 \
'aiconfigurator\.sdk|aiconfigurator_core\.sdk|PipelineLayout|PipelineSteadyState|pipeline_microbatches|per_op_ms|partition' \
docs README.md examples 2>/dev/null | head -n 1000Repository: ai-dynamo/aiconfigurator
Length of output: 46207
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- all tracked files containing pipeline API names ---'
rg -l --glob '*.py' \
'\b(PipelineLayout|PipelineSteadyState|even_partition|pipeline_microbatches|per_op_ms)\b' . \
| sort
echo '--- configuration, generator, and profiler definitions/usages ---'
rg -n -C5 --glob '*.py' \
'(class .*Config|def .*generate|def .*profile|profiler|generator|pipeline|pp_size|partition)' \
aic-core/src/aiconfigurator_core/sdk aiconfigurator/sdk tests/cross_package tests/unit/sdk 2>/dev/null \
| rg 'PipelineLayout|PipelineSteadyState|even_partition|pipeline_microbatches|per_op_ms|partition|pp_size|profiler|generator|Config' \
| head -n 1800
echo '--- package init/export surfaces ---'
for f in $(git ls-files '*__init__.py' | rg '(^|/)(aiconfigurator|aic-core).*(sdk|__init__)' | head -n 200); do
matches=$(rg -n 'pipeline|Pipeline|aiconfigurator_core' "$f" || true)
if [ -n "$matches" ]; then
echo "### $f"
printf '%s\n' "$matches"
fi
done
echo '--- docs/readmes mentioning pipeline API ---'
rg -n -C4 \
'PipelineLayout|PipelineSteadyState|even_partition|pipeline_microbatches|per_op_ms|partition' \
docs README.md aic-core/src/aiconfigurator_core/sdk/models/README.md 2>/dev/null | head -n 1000Repository: ai-dynamo/aiconfigurator
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- package layout ---'
git ls-files | rg '(^|/)(aiconfigurator|aiconfigurator_core)(/|$)' | head -n 500
echo '--- public package initializers ---'
for f in $(git ls-files | rg '(^|/)(__init__\.py)$' | rg 'aiconfigurator'); do
echo "### $f"
cat -n "$f" | head -n 220
done
echo '--- legacy SDK pipeline-related references only ---'
rg -n -C5 --glob '*.py' \
'pipeline|pp_size|per_op_ms|generation_ops|context_ops|profiler|profile|generator' \
aiconfigurator aic-core/src/aiconfigurator 2>/dev/null \
| rg 'pipeline|pp_size|per_op_ms|Pipeline|profile|generator|generation_ops|context_ops' \
| head -n 1800
echo '--- backend integration around pipeline calls ---'
sed -n '90,150p' aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
sed -n '1500,1585p' aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
sed -n '1420,1465p' aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
echo '--- setup/package mapping ---'
rg -n -C5 --glob 'pyproject.toml' --glob 'setup.py' --glob 'setup.cfg' --glob '*.toml' \
'aiconfigurator|aic-core|package|packages|src' . | head -n 1200Repository: ai-dynamo/aiconfigurator
Length of output: 46792
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- every pipeline method call and definition ---'
rg -n -C6 --glob '*.py' \
'(_pipeline_layout|_pipeline_steady_state|pipeline_microbatches|PipelineLayout|PipelineSteadyState)' \
aic-core/src tests | head -n 2200
echo '--- public estimate/generator signatures and forwarding ---'
rg -n -C8 --glob '*.py' \
'def (estimate|generate|compile|profile|query|run)|\*\*kwargs|pipeline_microbatches' \
aic-core/src/aiconfigurator_core/sdk aic-core/src/aiconfigurator 2>/dev/null \
| head -n 2400
echo '--- model operation construction and profiler breakdown shapes ---'
rg -n -C5 --glob '*.py' \
'(context_ops|generation_ops|per_op|breakdown|latency_breakdown|query\(|profiler)' \
aic-core/src/aiconfigurator_core/sdk/models \
aic-core/src/aiconfigurator_core/sdk/backends \
aic-core/src/aiconfigurator_core/sdk/perf_interp \
| head -n 2600
echo '--- documentation PP references ---'
rg -n -C5 \
'PP|pp_size|pipeline|pipeline parallel|chunked prefill|microbatch' \
docs README.md aic-core/src/aiconfigurator_core/sdk/models/README.md 2>/dev/null \
| head -n 1800Repository: ai-dynamo/aiconfigurator
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- project metadata and package mappings ---'
for f in $(git ls-files | rg '(^|/)(pyproject\.toml|setup\.py|setup\.cfg|tox\.ini|pytest\.ini)$'); do
echo "### $f"
rg -n -C8 'aiconfigurator|aic-core|package|packages|src|build-system|dependencies' "$f" || true
done
echo '--- exact pipeline call sites ---'
rg -n -C12 --glob '*.py' \
'(_pipeline_layout|_pipeline_steady_state|pipeline_microbatches|PipelineLayout|PipelineSteadyState)' \
aic-core/src tests/unit/sdk tests/cross_package \
| head -n 1600
echo '--- relevant backend signatures and forwarding ---'
sed -n '1360,1480p' aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
sed -n '1500,1580p' aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py
echo '--- docs and examples with PP or pipeline terms ---'
rg -n -C4 \
'PP|pp_size|pipeline|pipeline parallel|chunked prefill|microbatch' \
docs README.md aic-core/src/aiconfigurator_core/sdk/models/README.md 2>/dev/null \
| head -n 1600Repository: ai-dynamo/aiconfigurator
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
roots = [Path("aic-core/src"), Path("tests"), Path("docs")]
hits = []
for root in roots:
if not root.exists():
continue
for path in root.rglob("*.py"):
try:
tree = ast.parse(path.read_text(), filename=str(path))
except (OSError, SyntaxError):
continue
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom):
module = node.module or ""
if "aiconfigurator" in module or "pipeline" in module:
names = [a.name for a in node.names]
hits.append((str(path), node.lineno, f"from {module} import {', '.join(names)}"))
elif isinstance(node, ast.Import):
for alias in node.names:
if "aiconfigurator" in alias.name:
hits.append((str(path), node.lineno, f"import {alias.name}"))
print("--- parsed import references ---")
for item in sorted(hits):
print(f"{item[0]}:{item[1]}: {item[2]}")
print("--- pipeline API definitions and call keywords ---")
for path in [Path("aic-core/src/aiconfigurator_core/sdk/backends/base_backend.py")]:
tree = ast.parse(path.read_text(), filename=str(path))
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.name in {"_pipeline_layout", "_pipeline_steady_state"}:
args = [a.arg for a in node.args.args + node.args.kwonlyargs]
print(f"{path}:{node.lineno}: {node.name}({', '.join(args)})")
for call in ast.walk(node):
if isinstance(call, ast.Call):
text = ast.unparse(call.func)
if "_pipeline_" in text or any(
isinstance(k, ast.keyword) and k.arg == "pipeline_microbatches"
for k in call.keywords
):
print(f"{path}:{call.lineno}: {ast.unparse(call)}")
PYRepository: ai-dynamo/aiconfigurator
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- legacy SDK files ---'
git ls-files 'src/aiconfigurator/sdk/**' 'src/aiconfigurator/**' | head -n 500
echo '--- legacy SDK package initializers ---'
for f in $(git ls-files 'src/aiconfigurator/**/__init__.py' 'src/aiconfigurator/__init__.py'); do
echo "### $f"
cat -n "$f" | head -n 260
done
echo '--- legacy SDK modules that import or alias core modules ---'
rg -n -C5 --glob '*.py' \
'aiconfigurator_core|from \.|\.\.sdk|ModelConfig|RuntimeConfig|pipeline|pp_size|per_op_ms|generation_ops|context_ops' \
src/aiconfigurator/sdk src/aiconfigurator/generator 2>/dev/null \
| head -n 2400
echo '--- compatibility tests ---'
sed -n '1,220p' tests/cross_package/test_import_contract.py
rg -n -C8 --glob '*.py' \
'aiconfigurator\.sdk|aiconfigurator_core\.sdk|pipeline_microbatches|pp_size|partition|profiler|generator' \
tests/cross_package tests/unit tests/integration tests/e2e tests/golden \
| head -n 2200Repository: ai-dynamo/aiconfigurator
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- pipeline compatibility files ---'
for f in \
src/aiconfigurator/sdk/pipeline.py \
src/aiconfigurator/sdk/_compat.py \
src/aiconfigurator/sdk/backends/base_backend.py \
aic-core/src/aiconfigurator_core/sdk/__init__.py \
tests/cross_package/test_import_contract.py; do
if [ -f "$f" ]; then
echo "### $f"
cat -n "$f" | sed -n '1,240p'
else
echo "MISSING: $f"
fi
done
echo '--- exact pipeline paths and references ---'
git ls-files | rg '(^|/)pipeline\.py$|pipeline_microbatches|CORE_SDK_LEAF_MODULES' || true
python3 - <<'PY'
from pathlib import Path
import ast
core = Path("aic-core/src/aiconfigurator_core/sdk")
legacy = Path("src/aiconfigurator/sdk")
contract = Path("tests/cross_package/test_import_contract.py")
core_leaves = sorted(
str(p.relative_to(core)).removesuffix(".py").replace("/", ".")
for p in core.rglob("*.py")
if p.name != "__init__.py"
)
legacy_leaves = sorted(
str(p.relative_to(legacy)).removesuffix(".py").replace("/", ".")
for p in legacy.rglob("*.py")
if p.name != "__init__.py"
)
tree = ast.parse(contract.read_text())
declared = []
for node in tree.body:
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name) and target.id == "CORE_SDK_LEAF_MODULES":
declared = sorted(ast.literal_eval(node.value))
print("--- leaf-set comparison ---")
print("core-only:", sorted(set(core_leaves) - set(legacy_leaves)))
print("legacy-only:", sorted(set(legacy_leaves) - set(core_leaves)))
print("declared-missing:", sorted(set(core_leaves) - set(declared)))
print("declared-extra:", sorted(set(declared) - set(core_leaves)))
PYRepository: ai-dynamo/aiconfigurator
Length of output: 13821
Add the legacy pipeline compatibility wrapper. Add src/aiconfigurator/sdk/pipeline.py and include "pipeline" in CORE_SDK_LEAF_MODULES; the cross-package contract currently fails because aiconfigurator_core.sdk.pipeline has no legacy counterpart.
🤖 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 `@aic-core/src/aiconfigurator_core/sdk/pipeline.py` around lines 108 - 124, Add
the legacy sdk pipeline compatibility module corresponding to the core
PipelineLayout API, and register "pipeline" in CORE_SDK_LEAF_MODULES so it is
included in cross-package compatibility checks. Keep the wrapper’s public
exports aligned with aiconfigurator_core.sdk.pipeline.
Source: Path instructions
| ``` | ||
| stage_times(per_op_ms, num_layers) -> [t_0, t_1, ... t_{pp-1}] | ||
| per_hop_latency(per_op_ms) -> P2P cost of one stage-to-stage hop | ||
| layer_partition(num_layers) -> layers per stage | ||
| ``` |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add language identifiers to the new fenced code blocks.
The fences at Line 75 and Line 86 have no language identifier and trigger MD040. Add text to both openings.
This is the same unresolved issue noted in the previous review.
As per path instructions, the fix is small and limited to this documentation hunk, so a one-click suggested change is safe.
Proposed fix
-```
+```text
stage_times(per_op_ms, num_layers) -> [t_0, t_1, ... t_{pp-1}]
per_hop_latency(per_op_ms) -> P2P cost of one stage-to-stage hop
layer_partition(num_layers) -> layers per stage
-```
+```
@@
-```
+```text
cycle_time = max_i(stage_time_i) + p2p_per_hop
balance_factor = (step_total / pp_size) / cycle_time # per-microbatch latency effect
fill_factor = min(1, num_microbatches / pp_size) # throughput-only effect
efficiency = balance_factor * fill_factor
-```
+```Also applies to: 86-91
🧰 Tools
🪛 markdownlint-cli2 (0.23.1)
[warning] 75-75: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 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 `@docs/PIPELINE_PARALLEL_MODELING.md` around lines 75 - 79, Update the two
fenced code blocks in the documentation hunk, including the blocks near
stage_times and cycle_time, by adding the text language identifier to each
opening fence. Leave the block contents and closing fences unchanged.
Sources: Path instructions, Linters/SAST tools
…dder path The previous framing justified the layout/steady-state split with a scenario that is not the current contract, and understated the scope limit. Both are corrected here; no code behavior changes. What the embedder contract actually is: `ForwardPassPerfModel::estimate_forward_pass_time_ms` takes per-ATTENTION-DP- rank `ForwardPassMetrics` and reduces with `max` over ranks. There is no stage axis in it. A PP worker reports one forward pass and the caller expects the estimate to already include intra-worker parallelism, so PP is AIC's job to model, not the caller's. That also makes the placement obvious by symmetry: `forward_pass_time_ms` maxes over DP ranks, a pipeline cycle maxes over stages -- same reduction, different axis. So the split is justified by keeping the placement/partition RULES in one place (so a second derivation cannot drift from `stage_times`), not by a double-counting risk in a simulator that models occupancy itself. Rewrote the module docstring, the design doc, and the boundary test's rationale. Two gaps documented that the earlier text missed: 1. This change covers `run_agg` only. The planner/Mocker path (`ForwardPassPerfModel` -> `forward_pass_time_ms` -> `rank_latency_ms`) never enters `run_agg` and is still ideal-PP. Fixing it means the same max-over-stages inside `rank_latency_ms` -- a Rust change against a literal port, so its own PR under the rust-core parity rule. 2. The FPM online correction silently absorbs the PP error. `tune_with_fpms` learns `median(observed_ms / native_ms)` where `native_ms` is the whole-model estimate; the only guard is finite-and-positive, so the factor is UNBOUNDED in magnitude, and the region key is built from scheduler counts only, so `pp_size` is not a feature. Depending on what the emitter reports as `wall_time`, either PP's throughput benefit is absent (whole-iteration wall time) or the correction learns ~1/pp_size and stands in for a structural effect (per-stage wall time). Either way it is laundered through a fitted factor with no warning. Replaces the earlier "stage-indexed cost oracle / FFI is future work" note, which was written on the wrong premise. Signed-off-by: tianhaox <tianhaox@nvidia.com>
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 `@docs/PIPELINE_PARALLEL_MODELING.md`:
- Line 101: Update the documentation phrase in the steady-state collapse
discussion from “one step shape” to the hyphenated “one-step shape,” without
changing surrounding text.
🪄 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: 8f2faf40-04d9-49d4-90fd-5a4aec7698e2
📒 Files selected for processing (3)
aic-core/src/aiconfigurator_core/sdk/pipeline.pydocs/PIPELINE_PARALLEL_MODELING.mdtests/unit/sdk/test_pipeline_modeling.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/sdk/test_pipeline_modeling.py
- aic-core/src/aiconfigurator_core/sdk/pipeline.py
📜 Review details
⏰ Context from checks skipped due to timeout. (10)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Build wheels (macosx_arm64)
- GitHub Check: Build wheels (manylinux_2_28_x86_64)
- GitHub Check: Build wheels (manylinux_2_28_aarch64)
- GitHub Check: aic-core public API contract
- GitHub Check: Cargo Deny
- GitHub Check: Build and Test (e2e)
- GitHub Check: Build and Test (unit)
- GitHub Check: Rust/Python engine-step parity
🧰 Additional context used
📓 Path-based instructions (2)
**/*
⚙️ 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:
docs/PIPELINE_PARALLEL_MODELING.md
docs/**
⚙️ CodeRabbit configuration file
docs/**: - Check that docs match changed CLI, SDK, generator, backend, and support-matrix behavior.
- Flag docs that describe unsupported runtimes, stale command names, or behavior not covered by tests or support-matrix evidence.
Files:
docs/PIPELINE_PARALLEL_MODELING.md
🪛 LanguageTool
docs/PIPELINE_PARALLEL_MODELING.md
[grammar] ~101-~101: Use a hyphen to join words.
Context: ...e*, valid only because AIC evaluates one step shape and assumes every stage sees ...
(QB_NEW_EN_HYPHEN)
🪛 markdownlint-cli2 (0.23.1)
docs/PIPELINE_PARALLEL_MODELING.md
[warning] 115-115: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (2)
docs/PIPELINE_PARALLEL_MODELING.md (2)
115-121: 📐 Maintainability & Code QualityAdd a language identifier to this fenced block.
MD040 flags the fence at Line 115. The block contains pseudo-code, so use
textand preserve its contents.As per path instructions, this fix is small and limited to this documentation hunk, so a one-click suggested change is safe.
Suggested change
-``` +```textSources: Path instructions, Linters/SAST tools
95-100: LGTM!Also applies to: 102-114, 116-121, 161-163, 207-215, 217-236
| is what stops two derivations from drifting. | ||
|
|
||
| `balance_factor` and `fill_factor` are a different kind of thing: a | ||
| steady-state *collapse*, valid only because AIC evaluates one step shape and |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use a hyphen in “one-step shape.”
Change one step shape to one-step shape.
Suggested change
- steady-state *collapse*, valid only because AIC evaluates one step shape and
+ steady-state *collapse*, valid only because AIC evaluates one-step shape andAs per path instructions, this fix is small and limited to this documentation hunk, so a one-click suggested change is safe.
📝 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.
| steady-state *collapse*, valid only because AIC evaluates one step shape and | |
| steady-state *collapse*, valid only because AIC evaluates one-step shape and |
🧰 Tools
🪛 LanguageTool
[grammar] ~101-~101: Use a hyphen to join words.
Context: ...e*, valid only because AIC evaluates one step shape and assumes every stage sees ...
(QB_NEW_EN_HYPHEN)
🤖 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 `@docs/PIPELINE_PARALLEL_MODELING.md` at line 101, Update the documentation
phrase in the steady-state collapse discussion from “one step shape” to the
hyphenated “one-step shape,” without changing surrounding text.
Sources: Path instructions, Linters/SAST tools
Qwen3-32B bf16 dummy weights, 8x H20-3e, TRT-LLM 1.3.0rc20 against the h20e_sxm perf DB collected on the same stack and version, so absolute numbers are comparable. Decode-dominated (isl=256, osl=512), 8 GPUs held constant across every (tp, pp), chunked prefill off, aiperf closed-loop. Three findings, documented in a new section 6. 1. PP is a pure loss at constant GPU count for this dense model: no (tp, pp) beats tp8/pp1 at any concurrency (best 0.62x, pp=8 lands at 0.28-0.40x). Expected for weight-bound decode -- tp*pp fixed means per-GPU weight bytes are unchanged, so pipelining buys nothing while P2P hops and the smaller per-stage batch cost real time. Evidence for leaving PP out of the automatic search for now. 2. The "always full pipe" default is the largest single error. At most min(pp, C) microbatches can be in flight; the default assumes pp. At C=1 this over-predicts by up to 645% (pp=8: 357.6 vs 48.0 measured). Driving num_microbatches from real concurrency cuts that to -7%. The fill_factor mechanism added here is the right shape -- but nothing sets it, because run_agg derives concurrency as batch_size * pp_size rather than taking it as a target. Recorded as gap 1. 3. A concurrency-dependent gap remains unexplained. For C >= pp the fill correction is inert and AIC still over-predicts, growing from +7..44% at C=8 to +73..173% at C=128. Stage imbalance moves this by ~1 point, so it is not the cause. Measured behaviour is non-monotonic in pp (at C=128, pp=2 ITL 49.8ms vs pp=4 32.8ms) and pp=2 throughput falls between C=64 and C=128 while ITL doubles -- a saturation signature pointing at engine-side scheduling. Flagged explicitly as not-to-be-fitted until understood. The pp=1 baseline predicts within -8%..+11% throughput and 8% TPOT throughout, so all of the above is PP-specific rather than a perf-database problem. This reframes the change: the stage-imbalance model in commit 1 is correct but second-order against silicon. The dominant errors are the fill default and the unexplained saturation. Gap 1 is updated accordingly. Signed-off-by: tianhaox <tianhaox@nvidia.com>
…oss-node
The previous revision drew a conclusion the experiment could not support. The
silicon run is single-node, and PP cannot win inside one NVLink domain at
constant GPU count: per-GPU weight bytes are unchanged (tp*pp fixed), so
pipelining buys nothing while the P2P hops and the smaller per-stage batch cost
real time. Reporting that as "PP is a pure loss", and using it to argue PP
belongs out of the search, generalised from the one regime where PP is
structurally incapable of paying off.
PP exists for the case where TP would otherwise span the slow fabric. On this
spec (num_gpus_per_node=8, inter_node 50 GB/s vs intra_node 450 GB/s) TP
all-reduces across that 9x cliff every layer while PP crosses it once per stage
boundary. AIC models this, and the effect is large -- Qwen3-32B, isl=256/
osl=512, C=64:
16 GPUs (2 nodes): tp16/pp1 1714 tok/s vs tp8/pp2 7073 -> 4.1x
32 GPUs (4 nodes): tp32/pp1 1735 tok/s vs tp8/pp4 9275 -> 5.4x
So section 6.2 (fill default) and 6.3 (saturation gap) are restated as error-
mode stress tests: single-node PP exposes the model's failure modes precisely
because no real gain can hide them. Which findings generalise is now stated --
the fill default and stage imbalance are topology-independent, while the
saturation effect may be TRT-LLM single-node scheduling and is untested
cross-node.
Also corrects the search-space gap. Leaving PP out is not a conservative
default: at 16 GPUs the search can only offer tp16/pp1 when tp8/pp2 is worth
4x, which is a wrong answer on a real sizing decision, not a cautious one. It
stays gated on the gaps only because both over-predict, so enabling the
dimension now would select pipelines that are too deep.
Signed-off-by: tianhaox <tianhaox@nvidia.com>
Multi-backend PP validation data — 5 models, 3 backends, 8xH200I recently ran an independent PP validation campaign that may be useful context for this PR. The data covers the regime the silicon validation in §6 of Methodology
All three backends target the same models on the same hardware with the same workload, so throughput ratios are directly comparable. Observed PP efficiency
How much of the PP penalty does
|
| Model | PP | Total penalty | balance_factor explains | Scheduling residual |
|---|---|---|---|---|
| Qwen3-8B | 2 | 20.2% | 14.7% (73%) | 5.5% |
| Qwen3-8B | 8 | 68.5% | 40.7% (59%) | 27.8% |
| Qwen3-14B | 8 | 69.2% | 33.1% (48%) | 36.1% |
| Qwen3-32B-FP8 | 8 | 69.2% | 23.6% (34%) | 45.6% |
| Llama-70B-FP8 | 8 | 77.9% | 11.5% (15%) | 66.4% |
The balance_factor accounts for 15-73% of the total PP penalty, with the fraction decreasing for larger models (where per-layer compute dominates lm_head cost and the stages are more balanced). The majority of the penalty — especially for larger models — comes from the scheduling residual flagged in §6.3 of docs/PIPELINE_PARALLEL_MODELING.md.
The scheduling residual has structure
The residual (observed_eff / balance_factor) fits r = 1 - s * log2(pp) reasonably well for vLLM and SGLang. The coefficient s represents the scheduling/synchronization overhead component:
| Model | vLLM s | SGLang s | TRT-LLM s |
|---|---|---|---|
| Qwen3-8B | 0.158 | 0.150 | 0.283 |
| Qwen3-14B | 0.195 | 0.192 | 0.288 |
| Qwen3-32B-FP8 | 0.212 | 0.250 | 0.256 |
| Llama-70B-FP8 | 0.239 | 0.274 | -- |
Two patterns:
sincreases with model size (0.15 -> 0.24 for vLLM). This is counter-intuitive if the overhead were a fixed per-hop ms cost — it suggests the scheduling overhead scales with something beyond simple synchronization (memory bandwidth contention, KV cache management, activation tensor size).- vLLM and SGLang
svalues are close (within 0.03 for 8B-14B), diverging for 32B+ models. TRT-LLMsis consistently higher (~0.27).
Combined model accuracy
I compared three modeling approaches for vLLM (models 8B-70B, PP=2/4/8):
| Approach | Mean abs error | Max abs error |
|---|---|---|
| Flat empirical coefficient (c=0.228) | 2.7 pp | 9.5 pp |
| balance_factor * (1 - s_per_model * log2(pp)) | 1.8 pp | 4.7 pp |
| balance_factor * (1 - s_mean * log2(pp)) | 4.6 pp | 14.2 pp |
The combined model with per-model s is the most accurate, halving max error. But using a single mean s per backend is actually worse than the flat coefficient — s varies enough with model size that collapsing it into one number introduces more error than the balance_factor removes.
For SGLang: flat c gives 5.2/11.7 pp mean/max error; combined per-model gives 3.7/10.7 pp.
For TRT-LLM: all approaches struggle (6-8 pp mean, 16-21 pp max) due to the non-monotonic PP=4 behavior described below.
TRT-LLM PP=4 cliff
TRT-LLM shows a non-monotonic PP efficiency pattern across all model sizes tested:
| Model | PP=2 c | PP=4 c | PP=8 c |
|---|---|---|---|
| Qwen3-0.6B | 0.225 | 0.432 | 0.298 |
| Qwen3-8B | 0.241 | 0.404 | 0.298 |
| Qwen3-14B | 0.311 | 0.352 | 0.281 |
| Qwen3-32B-FP8 | 0.321 | 0.218 | 0.272 |
(Per-PP-level coefficient c from eff = 1 - c * log2(pp))
For small/medium models (0.6B-14B), the PP=4 per-level coefficient spikes to 0.35-0.43 — meaning PP=4 is disproportionately worse than either PP=2 or PP=8 would predict from a smooth curve. The per-concurrency ITL data confirms this is consistent across all concurrency levels (not a measurement artifact):
TRT-LLM Qwen3-8B conc=1 ITL:
TP=8/PP=1: 2.9 ms
TP=4/PP=2: 5.7 ms (2.0x)
TP=2/PP=4: 17.1 ms (5.9x) <-- step function
TP=1/PP=8: 18.7 ms (6.4x)
This pattern is absent from vLLM and SGLang (smooth degradation) and may relate to TRT-LLM's MPI-based pipeline implementation handling the TP=2 per-stage configuration differently. The effect diminishes for 32B-FP8 where per-stage compute is larger relative to the overhead.
This appears related to the non-monotonic saturation behavior identified in §6.3 of docs/PIPELINE_PARALLEL_MODELING.md — the scheduling overhead may depend on the interaction between TP and PP configuration within each stage, not just PP alone.
Implications
The balance_factor is structurally correct — it captures the right mechanism (lm_head imbalance). The data shows:
- The scheduling residual is the dominant component for models >= 14B (52-85% of total penalty). Any calibration effort should prioritize this over refining balance further.
- The residual is backend-specific. vLLM and SGLang share similar PP scheduling behavior; TRT-LLM is materially different. This has implications for whether the correction belongs in
BaseBackendvs per-backend overrides. - The residual scales with model size in a way a single per-backend constant doesn't capture. The
fill_factorhook may be the right place to absorb this, but the relationship needs characterization. - TRT-LLM's PP=4 cliff suggests the overhead depends on the TP*PP interaction, not just PP alone.
The raw benchmark data (224 result files across all backends) and analysis scripts are available if useful for calibration.
2 participants
Overview:
PP was modeled as an ideal linear speedup. A step's latency was the whole
model's (
_num_layersis never divided bypp_size) and throughput was simplymultiplied by
pp_size. That is exact only when every stage costs the same andthe pipe is always full — neither holds.
This adds
PipelineSchedule, which computes the two factors that were missing:stage imbalance and pipe fill.
pp_size == 1is bit-identical to before.Details:
What was wrong
A pipeline advances at
max_i(t_i), not the average. The un-shardedlm_headsits alone on the last stage, so for a decode step it sets the cycle.
Qwen3-32B (64 layers, vocab 151936) on H20 at
pp=8:Prefill steps stay at 99.9% — per-layer work dominates there, so this is a
decode-side effect.
Two more gaps:
num_layers % pp_size != 0only logged "we're nothing tocorrect this" while the fattest stage actually gates the cycle (64 layers over
6 stages costs
11/(64/6)), and a starved pipe still scaled linearly.The model
Op placement reuses the naming contract the model classes already require
(
GPTModeldocstring: "Same for logits_gemm"):*embedding*→ first stage,*logits_gemm*→ last,*p2p*→ link (charged once per hop, not as awhole-step total), everything else per-layer.
warn_on_unclassified_opsflags anew head-like op that would otherwise be silently smeared across every stage.
Layer partition defaults to an even split with the remainder on the leading
stages (vLLM
get_pp_indices/ TRT-LLM); an explicit partition is accepted.Where it plugs in — two touch points in
run_agg:Inflating the step latency to the real traversal time (
pp * cycle) is theload-bearing choice: TTFT, TPOT,
_total_step_latency_ms,_step_throughputand
_throughput_capall derive from it, so they stay consistent withoutindividual patches.
Why orchestration and not the op layer — per
.claude/rules/rust-core/parity.md,changing op query math obligates a mirrored Rust implementation plus oracle
anchors. This consumes the per-op breakdown after the step is evaluated, so
both engine-step paths benefit with no parity surface touched. Same layering
AFD's pipeline model already uses.
Measured effect (Qwen3-32B, 8x H20, isl=4000/osl=500/bs=32 — prefill-dominated,
so the correction is mild here; decode-heavy workloads move much more):
Where should the reviewer start?
aic-core/src/aiconfigurator_core/sdk/pipeline.py— the whole model, ~200 linesaic-core/src/aiconfigurator_core/sdk/backends/base_backend.py— two touchpoints in
run_aggplus the_pipeline_schedulebackend hookdocs/PIPELINE_PARALLEL_MODELING.md— derivation and the known-gaps listTesting
tests/unit/sdk/test_pipeline_schedule.py: 40 cases (placement, partitionincl. uneven splits, max-not-average cycle, P2P hop amortization + overlap,
starvation, the
pp_size == 1identity, validation).tests/unit: 2974 passed. The 3 failures intests/unit/collector/sglang/test_collect_mla_module.py(GLM-5 checkpointnaming) are pre-existing — reproduced on a clean baseline.
pp_size == 1verified bit-identical end to end on H20: TTFT 650.722 /TPOT 18.195 / 207.72 tokens/s/gpu, unchanged.
Known gaps (documented in the design doc, deliberately not in this PR)
fill_factoris linear — the real curve is likely steeper once schedulerand sync overhead are counted. Needs silicon calibration.
a RAW dependency on the KV they write, so they cannot occupy the pipe
simultaneously.
fill_factoris the intended hook; the coefficient needsmeasurement.
inter_node_bw(~9x pessimistic intra-node).SystemSpec.get_p2p_bandwidth(num_gpus)already implements the topologytiers on both sides, but the correct selector is
tp * pp, which the op isnot constructed with — fixing it touches every model's
P2Pconstructionplus the Rust mirror, so it belongs in its own change. Impact at
pp=8:0.67% of a decode step, 2.2% of a prefill step, conservative direction.
build_disagg_parallel_liststakes
should_enable_pp(defaultFalse) and no caller passesTrue.Turning it on should wait until 1 and 2 are calibrated.
Related Issues:
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation