Conversation
📝 WalkthroughWalkthroughThe AllReduce benchmark now supports configurable strategies, execution controls, rank-aware timing aggregation, p90 reporting, raw JSONL output, coordinated distributed initialization, and expanded result metadata. Shared utilities, documentation, and unit tests cover these behaviors. ChangesAllReduce benchmark controls and timing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant allreduce_comm
participant ProcessGroup
participant AllReduceBackend
participant JSONL
CLI->>allreduce_comm: provide strategy and execution controls
allreduce_comm->>ProcessGroup: validate or initialize distributed state
allreduce_comm->>AllReduceBackend: run AllReduce validation and iterations
AllReduceBackend-->>allreduce_comm: return per-rank timing samples
allreduce_comm->>JSONL: write aggregated timing metadata
JSONL-->>allreduce_comm: return write status
allreduce_comm-->>CLI: return benchmark results
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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 |
|
Local pre-commit, focused tests, and the 2xB300 integrated smoke are complete. Could a maintainer please approve the external CI for this PR? @flashinfer-bot run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
benchmarks/routines/allreduce_comm.py (1)
454-469: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winA rank-local kernel failure makes the ranks diverge at
comm.allgather.
bench_gpu_timecan raiseRuntimeErroron one rank only. That rank returnsNoneat Line 463. The other ranks continue tocomm.allgatherat Line 469 and block. The run then hangs instead of reporting an error.Convert the failure into a collective decision before the gather.
🛡️ Proposed fix
+ local_bench_error = None try: times = bench_gpu_time( ... ) except RuntimeError as e: - if rank == 0: - elem_size = torch.tensor([], dtype=input_dtype).element_size() - msg_size_mb = num_tokens * hidden_size * elem_size / (1024 * 1024) - print( - f"[ERROR] Kernel failed for shape=({num_tokens}, {hidden_size}) " - f"msg_size={msg_size_mb:.1f} MiB: {e}" - ) - return None + times = [] + local_bench_error = f"{type(e).__name__}: {e}" + + bench_error = gather_rank_errors(comm, "AllReduce benchmark", local_bench_error) + if bench_error is not None: + if rank == 0: + elem_size = torch.tensor([], dtype=input_dtype).element_size() + msg_size_mb = num_tokens * hidden_size * elem_size / (1024 * 1024) + print( + f"[ERROR] Kernel failed for shape=({num_tokens}, {hidden_size}) " + f"msg_size={msg_size_mb:.1f} MiB: {bench_error}" + ) + return None🤖 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 `@benchmarks/routines/allreduce_comm.py` around lines 454 - 469, Update the RuntimeError handling in the benchmark routine around bench_gpu_time so all ranks participate in a collective failure decision before returning. Communicate whether any rank failed, ensure every rank takes the same return path when a failure occurs, and only execute comm.allgather when no rank has failed; preserve the existing rank-0 error reporting.
🧹 Nitpick comments (2)
benchmarks/routines/allreduce_comm.py (1)
686-700: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the dead assignment to
needs_mnnvl.Line 696 sets
needs_mnnvl = False, but no later code readsneeds_mnnvl. The value has no effect.♻️ Proposed cleanup
if args.ar_backend == "auto": # Preserve one collective control flow: every rank falls back # to TRT-LLM instead of letting auto re-select failed MNNVL. backend_list = ["trtllm"] - needs_mnnvl = False🤖 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 `@benchmarks/routines/allreduce_comm.py` around lines 686 - 700, Remove the unused needs_mnnvl = False assignment from the MNNVL initialization failure branch in the allreduce backend selection flow. Preserve the backend_list fallback to ["trtllm"] and the surrounding control flow unchanged.benchmarks/routines/allreduce_comm_utils.py (1)
165-184: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
allow_nan=Falsecan abort the JSONL write for valid benchmark records.
bench_gpu_timecan producenanvalues, and the result dict already usestorch.nanfortflopselsewhere in the benchmark. If any timing value isNaN,json.dumpsraisesValueErrorand the whole record is dropped. The caller converts this into a cross-rankRuntimeError, so oneNaNsample aborts the benchmark run.Consider writing
nullfor non-finite values instead of failing.🤖 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 `@benchmarks/routines/allreduce_comm_utils.py` around lines 165 - 184, Update append_jsonl to serialize non-finite numeric values, including NaN and infinity, as JSON null rather than raising from json.dumps. Preserve existing JSON output for finite values and ensure complete benchmark records are still written.
🤖 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.
Outside diff comments:
In `@benchmarks/routines/allreduce_comm.py`:
- Around line 454-469: Update the RuntimeError handling in the benchmark routine
around bench_gpu_time so all ranks participate in a collective failure decision
before returning. Communicate whether any rank failed, ensure every rank takes
the same return path when a failure occurs, and only execute comm.allgather when
no rank has failed; preserve the existing rank-0 error reporting.
---
Nitpick comments:
In `@benchmarks/routines/allreduce_comm_utils.py`:
- Around line 165-184: Update append_jsonl to serialize non-finite numeric
values, including NaN and infinity, as JSON null rather than raising from
json.dumps. Preserve existing JSON output for finite values and ensure complete
benchmark records are still written.
In `@benchmarks/routines/allreduce_comm.py`:
- Around line 686-700: Remove the unused needs_mnnvl = False assignment from the
MNNVL initialization failure branch in the allreduce backend selection flow.
Preserve the backend_list fallback to ["trtllm"] and the surrounding control
flow unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e0dd8bc7-992e-4a34-91a0-3925388b411b
📒 Files selected for processing (5)
benchmarks/README.mdbenchmarks/routines/allreduce_comm.pybenchmarks/routines/allreduce_comm_utils.pybenchmarks/routines/flashinfer_benchmark_utils.pytests/test_allreduce_comm_utils.py
|
Addressed the review in commit 4de67ad: rank-local �ench_gpu_time failures are now gathered before any rank enters the raw-sample �llgather, so all ranks take the same return path; the dead |
|
Correction: the remote Linux hook invocation for commit 4de67ad did not complete because the SSH file transfer stalled; I am not counting it as a pass. The verified follow-up checks are 17 focused unit tests, |
|
Follow-up verification is now complete at |
1 participant
📌 Description
This PR makes the
allreduce_fusionbenchmark suitable for reproducing small, fixed-latency communication regressions without changing the AllReduce API or CUDA kernels.It:
--enable_pdlsetting instead of hardcoding PDL on in the timed path;max,rank0, ormeanaggregation exactly once;The historical routine benchmarked forced oneshot and twoshot strategies, and that remains the default. The one intentional default correction is PDL: the routine now follows the global CLI contract, so PDL is off unless
--enable_pdlis present.🔍 Related Issues
Related to #3825. This benchmark-only change does not claim to fix the reported latency regression.
It also follows the PDL CLI semantics introduced in #3435.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.🧪 Tests
unittest, etc.).Validation performed:
pre-commit run --all-files: all hooks passed on the initial five-file patch and again foraa056ac2in a clean WSL-native clone; the latter avoids Windows worktree line-ending conversion;max, 10 warmups/50 samples;rank0, 10 warmups/30 samples;mean, 10 warmups/30 samples.Reviewer Notes
Please focus on the request-vs-effective provenance wording and whether the optional raw JSONL surface is appropriately scoped.
strategy_request,trigger_completion_at_end_request, andtiming_mode_requestdeliberately do not claim backend-resolved behavior that the current APIs cannot observe.This contribution was developed with assistance from OpenAI Codex. I reviewed the code, test coverage, and the archived B300 evidence before submission.
Summary by CodeRabbit
New Features
Documentation
Tests