Add lossless PCIe serving calibration probe - #81
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a distributed PCIe calibration probe for TP backend selection, CKV overlap, query splitting, and prefetch recommendations. It includes policy generation, correctness checks, CLI execution, documentation, calibration evidence, and unit tests. ChangesPCIe calibration probe
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant Torchrun
participant Probe
participant NCCL
participant DMA
participant QueryIndexer
participant Rank0
Torchrun->>Probe: initialize distributed calibration
Probe->>NCCL: measure TP and CKV collectives
Probe->>DMA: measure TP DMA all-reduce
Probe->>QueryIndexer: measure full and split indexing
Probe-->>Probe: validate outputs and aggregate timings
Probe->>Rank0: emit policy and calibration results
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 7
🧹 Nitpick comments (6)
tests/comm/test_pcie_overlap_probe.py (1)
178-188: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueResolve Ruff RUF005.
Line 185 can use iterable unpacking instead of list concatenation:
Proposed fix
- late_loss = points + [(16384, _backend("nccl"))] + late_loss = [*points, (16384, _backend("nccl"))]🤖 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/comm/test_pcie_overlap_probe.py` around lines 178 - 188, Update the late_loss construction in test_dma_crossover_requires_winning_tail to use iterable unpacking instead of list concatenation, while preserving the existing points followed by the final nccl entry.Source: Linters/SAST tools
sparkinfer/comm/pcie/overlap_probe.py (5)
1232-1266: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNineteen defaults are duplicated between
ProbeConfigand the parser.Every default here restates a
ProbeConfigfield default. They currently agree, but nothing enforces that, and a drift means the CLI and the programmatic API calibrate differently. Deriving them from the dataclass removes the class of bug entirely.♻️ Derive parser defaults from the dataclass
def _build_parser() -> argparse.ArgumentParser: + defaults = ProbeConfig() parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--tp-size", type=int, default=8) - parser.add_argument("--dcp-size", type=int, default=4) - parser.add_argument("--hidden-size", type=int, default=6144) - parser.add_argument("--tp-rows", type=int, default=8192) + parser.add_argument("--tp-size", type=int, default=defaults.tp_size) + parser.add_argument("--dcp-size", type=int, default=defaults.dcp_size) + parser.add_argument("--hidden-size", type=int, default=defaults.hidden_size) + parser.add_argument("--tp-rows", type=int, default=defaults.tp_rows) parser.add_argument( "--allreduce-rows", type=_parse_positive_ints, - default=(1, 8, 32, 128, 512, 2048, 8192), + default=defaults.allreduce_rows, )Apply the same pattern to the remaining options.
🤖 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 `@sparkinfer/comm/pcie/overlap_probe.py` around lines 1232 - 1266, Update _build_parser to derive every CLI option default from the corresponding ProbeConfig dataclass field instead of duplicating literal defaults. Apply this consistently to all options currently mirroring ProbeConfig, while preserving parser types, names, and existing CLI behavior.
1330-1352: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConfig validation errors surface as a traceback instead of a CLI error.
ProbeConfig.validate()only runs insideCollectiveOverlapProbe.__init__, i.e. afterdist.init_process_group. A bad--dma-wire-modeor--dcp-sizetherefore aborts mid-rendezvous on every rank with a rawValueError. Callingconfig.validate()right after construction lets the process exit cleanly with a usage message before any distributed setup — which matters given the PR goal that "failures are reportable without loading a model."♻️ Validate before distributed init
dma_wire_mode=args.dma_wire_mode, ) + try: + config.validate() + except ValueError as exc: + print(f"invalid probe configuration: {exc}", file=sys.stderr) + return 2 result = run_probe(config)🤖 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 `@sparkinfer/comm/pcie/overlap_probe.py` around lines 1330 - 1352, Call config.validate() immediately after constructing the ProbeConfig in the CLI flow, before run_probe(config) can initialize distributed state. Preserve the existing CollectiveOverlapProbe validation while ensuring invalid options such as dma_wire_mode or dcp_size produce the CLI’s normal usage error instead of a raw traceback.
683-685: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueEvery context binding shares one pair of scratch buffers.
full_plan/full_scratchandsplit_plan/split_scratchare allocated once and reused across allcontext_tokensbindings. That is correct for the current sequential measurement loop, but it is an implicit invariant: no two bindings may be in flight at once.measure_overlapalready runs work onside_streamconcurrently withmain_stream, so if the indexer phase is ever folded into the overlap path this becomes a silent data race on shared scratch.A short comment stating the "one binding active at a time" contract would make the constraint explicit.
Also applies to: 719-722
🤖 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 `@sparkinfer/comm/pcie/overlap_probe.py` around lines 683 - 685, Add a concise comment near the shared full_plan/full_scratch and split_plan/split_scratch allocations documenting that all context bindings reuse these buffers and only one binding may be active at a time; preserve the current sequential measurement behavior and apply the same contract to the corresponding later allocation.
1047-1047: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
samplesis re-annotated with a conflicting type in the same scope.Line 1047 declares
dict[str, list[float]], Line 1075 redeclares the same name asdict[str, list[tuple[float, ...]]]. Type checkers reject the second annotation as a redefinition. Rename one (e.g.backend_samples/overlap_samples).Also applies to: 1075-1075
🤖 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 `@sparkinfer/comm/pcie/overlap_probe.py` at line 1047, Rename one of the conflicting annotations for samples in the surrounding probe logic, such as using backend_samples for the dict[str, list[float]] collection and overlap_samples for the dict[str, list[tuple[float, ...]]] collection, then update all references so each variable retains its intended data and the type checker sees no same-scope redefinition.
699-722: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueB023 hints here are benign, but the closure is fragile.
Ruff flags
base_page_table/local_tokensas unbound loop variables at Lines 700 and 702; in practicemake_bindingis invoked immediately at Lines 719-722, so late binding never triggers. Binding them as default arguments silences the lint and protects against someone later deferring the call.♻️ Bind loop variables explicitly
- def make_binding(plan, scratch, rows: int): - page_table = base_page_table.expand(rows, -1) + def make_binding( + plan, + scratch, + rows: int, + base_page_table: torch.Tensor = base_page_table, + local_tokens: int = local_tokens, + ): + page_table = base_page_table.expand(rows, -1)🤖 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 `@sparkinfer/comm/pcie/overlap_probe.py` around lines 699 - 722, Update the nested make_binding function to capture the loop-scoped base_page_table and local_tokens values through default arguments, then use those captured parameters when expanding the page table and creating seqlens. Preserve the existing immediate binding calls and all other binding behavior.Source: Linters/SAST tools
🤖 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/pcie_calibration_probe.md`:
- Around line 63-76: Expand the “Validation snapshot” evidence around the
documented benchmark claims to include the exact command, commit and worktree
identity, GPU mode, correctness result, and raw timings used to derive the
percentages and crossover; alternatively, link a generated JSON artifact
containing all of that information. Ensure the evidence covers both adjacent and
interleaved topologies and clearly records ratio direction so the results are
reproducible and auditable.
In `@sparkinfer/comm/pcie/overlap_probe.py`:
- Around line 944-962: Update the overlap probe around first() and second() to
synchronize their launch with a shared CUDA start event: record the event after
setup, make both collective streams wait on it, then enqueue the two launches.
Preserve the existing completion synchronization and elapsed-time calculations
using the appropriate stream events.
- Around line 1092-1108: Update the overlap decision call around decide_overlap
to use the pessimistic tp_first measurements alongside the existing side_first
values, selecting the worst-case ordering for concurrent throughput, collective,
and wall medians. Keep both orderings in the result payload and ensure the
decision no longer relies solely on side_first.
- Around line 87-88: Update validate() to enforce all divisibility requirements
used by the owner-merge path: require tp_rows to divide evenly by tp_size, and
require query_split_rows to divide evenly by indexer_shards before computing
owner_rows. Raise clear ValueError messages for each invalid condition while
preserving the existing query_split_size validation.
- Around line 1185-1194: Update the calculation of maximum_safe_context_tokens
near recommend_prefetch_depth to report the largest contiguous enabled prefix,
rather than max(enabled). Sort or otherwise traverse the context-token decisions
in ascending order and stop at the first rejected entry, preserving 0 when no
prefix is enabled.
- Around line 980-984: Strengthen the CKV correctness gate in the validation
block around ckv_ok to verify the entire gathered ckv_output buffer, not just
one byte at each rank-chunk offset. Compare all elements against the expected
rank-chunk pattern using an efficient device-side O(bytes) check, while
preserving the existing ckv_ok boolean contract before timings are interpreted.
- Around line 77-92: Extend the configuration validation around indexer_topk so
it cannot exceed the smallest per-shard local index pool derived from
context_tokens and the configured sharding. Reject invalid values with a clear
ValueError before probe execution, ensuring paged top-k and run_row_topk never
receive a top-k larger than available local candidates.
---
Nitpick comments:
In `@sparkinfer/comm/pcie/overlap_probe.py`:
- Around line 1232-1266: Update _build_parser to derive every CLI option default
from the corresponding ProbeConfig dataclass field instead of duplicating
literal defaults. Apply this consistently to all options currently mirroring
ProbeConfig, while preserving parser types, names, and existing CLI behavior.
- Around line 1330-1352: Call config.validate() immediately after constructing
the ProbeConfig in the CLI flow, before run_probe(config) can initialize
distributed state. Preserve the existing CollectiveOverlapProbe validation while
ensuring invalid options such as dma_wire_mode or dcp_size produce the CLI’s
normal usage error instead of a raw traceback.
- Around line 683-685: Add a concise comment near the shared
full_plan/full_scratch and split_plan/split_scratch allocations documenting that
all context bindings reuse these buffers and only one binding may be active at a
time; preserve the current sequential measurement behavior and apply the same
contract to the corresponding later allocation.
- Line 1047: Rename one of the conflicting annotations for samples in the
surrounding probe logic, such as using backend_samples for the dict[str,
list[float]] collection and overlap_samples for the dict[str, list[tuple[float,
...]]] collection, then update all references so each variable retains its
intended data and the type checker sees no same-scope redefinition.
- Around line 699-722: Update the nested make_binding function to capture the
loop-scoped base_page_table and local_tokens values through default arguments,
then use those captured parameters when expanding the page table and creating
seqlens. Preserve the existing immediate binding calls and all other binding
behavior.
In `@tests/comm/test_pcie_overlap_probe.py`:
- Around line 178-188: Update the late_loss construction in
test_dma_crossover_requires_winning_tail to use iterable unpacking instead of
list concatenation, while preserving the existing points followed by the final
nccl entry.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3a23e4df-4f30-4403-997f-8cb46f8d88b4
📒 Files selected for processing (3)
docs/pcie_calibration_probe.mdsparkinfer/comm/pcie/overlap_probe.pytests/comm/test_pcie_overlap_probe.py
|
Final release-candidate integration validation passed. Image: The image contains SparkInfer integration E2E gates using the generated policy:
The DCP4 calibration selected query split from 8k, CKV prefetch depth 1, and a 24 MiB lossless DMA crossover. Compressed DMA remained explicit-only. |
2 participants
Summary
Add a standalone multi-process probe that measures the communication decisions used by the GLM DCP prefill stack on the exact selected GPU order:
Design constraints
Validation
13 passedintests/comm/test_pcie_overlap_probe.pygit diff --check master...HEADMeasured values are deployment data, not hard-coded policy in this PR.
Summary by CodeRabbit