[None][chore] Refine AlltoAll benchmark scripts. by bobboli · Pull Request #11649 · NVIDIA/TensorRT-LLM

bobboli · 2026-02-23T09:34:42Z

Summary by CodeRabbit

New Features
- Added --perfect_router option for deterministic expert routing across ranks.
- Added --random_seed CLI argument to control reproducibility of inputs.
- Enhanced benchmark output with routing configuration and seed metadata.
- Improved kernel breakdown reporting with per-rank timing data for better performance insights.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

…t judgment. Add random seed to input. Signed-off-by: Bo Li <[email protected]>

Signed-off-by: Bo Li <[email protected]>

coderabbitai · 2026-02-23T09:40:27Z

📝 Walkthrough

Walkthrough

Introduces deterministic routing capability to the benchmark via a perfect_router option that balances token assignment across ranks and experts, implements refined GPU kernel scope detection with stricter containment logic, adds per-rank random seeding for reproducible inputs across MPI processes, and expands benchmark output to include kernel breakdown data with per-rank timings.

Changes

Cohort / File(s)	Summary
Input Generation & Routing `tests/microbenchmarks/bench_moe_comm.py` (_make_inputs function)	Adds `experts_per_rank` and `perfect_router` parameters; implements deterministic balanced routing schedule when `perfect_router=True` to distribute tokens across ranks/experts; preserves existing random routing when disabled.
Kernel Scope Detection `tests/microbenchmarks/bench_moe_comm.py` (_find_scope function)	Introduces strict containment-based logic for GPU kernel categorization (dispatch vs. combine); replaces previous scoping approach with explicit containment checks and non-overlap assertions to eliminate ambiguity.
Random Seeding & Reproducibility `tests/microbenchmarks/bench_moe_comm.py` (worker execution)	Seeds CPU and CUDA RNG per MPI rank using `_WORKER_ENV` environment variable propagation; ensures unique, reproducible inputs across distributed ranks.
Kernel Breakdown Aggregation `tests/microbenchmarks/bench_moe_comm.py` (kernel breakdown collection & merging)	Collects local kernel timing data, performs MPI allgather, merges per-rank times while maintaining deterministic kernel-name ordering; categorizes output by dispatch_kernels, combine_kernels, and other_kernels with per-rank timings.
CLI & Output Schema `tests/microbenchmarks/bench_moe_comm.py` (parse_args, output assembly)	Adds `--random_seed` and `--perfect_router` CLI options; expands benchmark metadata payload with `perfect_router`, `random_seed`, `experts_per_rank`, and per-rank kernel breakdown; ensures output directory creation and non-blocking MPI shutdown.
Synchronization Changes `tests/microbenchmarks/bench_moe_comm.py` (MPI barrier adjustments)	Removes several MPI barriers in error paths and control-flow adjustments; reflects revised synchronization strategy.

Sequence Diagram(s)

sequenceDiagram
    participant CLI as CLI / Main
    participant MPI as MPI Pool
    participant Worker as Worker Process
    participant GPU as GPU Kernel
    
    CLI->>CLI: parse_args (random_seed, perfect_router)
    CLI->>MPI: spawn workers with _WORKER_ENV
    
    MPI->>Worker: execute_benchmark_single
    Worker->>Worker: seed RNG per rank
    
    alt perfect_router enabled
        Worker->>Worker: _make_inputs (deterministic schedule)
        Note over Worker: Balance tokens across<br/>ranks & experts
    else
        Worker->>Worker: _make_inputs (random routing)
        Note over Worker: Random expert selection
    end
    
    Worker->>GPU: dispatch kernel (hidden_states)
    GPU-->>Worker: dispatch results
    
    Worker->>GPU: combine kernel (outputs)
    GPU-->>Worker: combine results
    
    Worker->>Worker: _find_scope (categorize kernels)
    Note over Worker: Containment-based<br/>scope detection
    
    Worker->>MPI: allgather kernel timings
    MPI->>Worker: merged per-rank data
    
    Worker->>Worker: aggregate & sort timings
    Worker-->>CLI: return benchmark metrics
    
    CLI->>CLI: assemble output (dispatch/combine/other kernels)
    CLI->>CLI: write JSON with per-rank breakdown

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	PR description contains only the template with all required sections left blank and unchecked.	Author must fill in the Description, Test Coverage sections and confirm the PR Checklist. A clear explanation of the changes and their testing is required.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[None][chore] Refine AlltoAll benchmark scripts' is directly related to the main changes in the PR, which refine the AlltoAll benchmark scripts with new routing and random seeding features.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings (stacked PR)
📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/microbenchmarks/bench_moe_comm.py (1)
651-653: Consider naming convention for module-level constant.

Per the coding guidelines, global variables should use upper snake_case with prefix G_ (e.g., G_WORKER_ENV). However, the leading underscore indicates this is module-private, which is also a valid Python convention. This is a minor stylistic point.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@tests/microbenchmarks/bench_moe_comm.py` around lines 651 - 653, The
module-level constant _WORKER_ENV should follow the project's naming convention
for globals; rename _WORKER_ENV to G_WORKER_ENV and update every usage/reference
in the module (e.g., any tests or functions that read _WORKER_ENV) to the new
name, ensuring imports/exports and any module-private intent remain correct;
keep the dictionary contents unchanged and run tests to verify no references
were missed.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/microbenchmarks/bench_moe_comm.py`:
- Line 887: The code calls os.makedirs(os.path.dirname(args.output_file),
exist_ok=True) which fails when args.output_file has no directory component;
update the logic to first compute dirpath = os.path.dirname(args.output_file)
and only call os.makedirs(dirpath, exist_ok=True) if dirpath is non-empty (e.g.,
if dirpath: os.makedirs(...)); reference args.output_file and the os.makedirs
call to locate the change.

---

Nitpick comments:
In `@tests/microbenchmarks/bench_moe_comm.py`:
- Around line 651-653: The module-level constant _WORKER_ENV should follow the
project's naming convention for globals; rename _WORKER_ENV to G_WORKER_ENV and
update every usage/reference in the module (e.g., any tests or functions that
read _WORKER_ENV) to the new name, ensuring imports/exports and any
module-private intent remain correct; keep the dictionary contents unchanged and
run tests to verify no references were missed.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6943087 and 268bfa3.

📒 Files selected for processing (1)

tests/microbenchmarks/bench_moe_comm.py

Signed-off-by: Bo Li <[email protected]>

…beginning. Signed-off-by: Bo Li <[email protected]>

Signed-off-by: Bo Li <[email protected]>

…lose to timed iters. Signed-off-by: Bo Li <[email protected]>

Signed-off-by: Bo Li <[email protected]>

bobboli · 2026-03-05T17:12:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-05T17:42:11Z

PR_Github #37892 [ run ] triggered by Bot. Commit: dbc91e7 Link to invocation

tensorrt-cicd · 2026-03-05T21:11:32Z

PR_Github #37892 [ run ] completed with state SUCCESS. Commit: dbc91e7
/LLM/main/L0_MergeRequest_PR pipeline #29340 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

Signed-off-by: Bo Li <[email protected]>

bobboli · 2026-03-10T16:39:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-10T16:45:52Z

PR_Github #38465 [ run ] triggered by Bot. Commit: aef4d60 Link to invocation

bobboli · 2026-03-10T16:47:49Z

/bot kill

Signed-off-by: Bo Li <[email protected]>

tensorrt-cicd · 2026-03-10T16:54:31Z

PR_Github #38467 [ kill ] triggered by Bot. Commit: aef4d60 Link to invocation

Signed-off-by: Bo Li <[email protected]>

tensorrt-cicd · 2026-03-10T16:54:35Z

PR_Github #38465 [ run ] completed with state ABORTED. Commit: aef4d60

Link to invocation

bobboli · 2026-03-10T16:54:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-10T16:55:03Z

PR_Github #38467 [ kill ] completed with state SUCCESS. Commit: aef4d60
Successfully killed previous jobs for commit aef4d60

Link to invocation

tburt-nv

Licenses look good

cxxfilt==0.3.0
BSD-2 license: https://github.com/afq984/python-cxxfilt/blob/v0.3.0/LICENSE

NVIDIA dependencies:
cuda-bindings==13.1.1
cuda-pathfinder==1.4.1
cupti-python==13.0.1
nvidia-cuda-cupti==13.1.115

Co-authored-by: tburt-nv <[email protected]> Signed-off-by: Bo Li <[email protected]>

bobboli · 2026-03-10T17:00:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-10T17:05:06Z

PR_Github #38469 [ run ] triggered by Bot. Commit: 3872f9b Link to invocation

tensorrt-cicd · 2026-03-10T17:07:45Z

PR_Github #38471 [ run ] triggered by Bot. Commit: 3872f9b Link to invocation

tensorrt-cicd · 2026-03-10T17:07:48Z

PR_Github #38469 [ run ] completed with state ABORTED. Commit: 3872f9b

Link to invocation

tensorrt-cicd · 2026-03-10T21:59:47Z

PR_Github #38471 [ run ] completed with state SUCCESS. Commit: 3872f9b
/LLM/main/L0_MergeRequest_PR pipeline #29827 completed with status: 'FAILURE'

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

bobboli · 2026-03-11T03:30:53Z

/bot run --disable-fail-fast --reuse-test

tensorrt-cicd · 2026-03-11T03:38:19Z

PR_Github #38532 [ run ] triggered by Bot. Commit: 3872f9b Link to invocation

tensorrt-cicd · 2026-03-11T05:38:33Z

PR_Github #38532 [ run ] completed with state SUCCESS. Commit: 3872f9b
/LLM/main/L0_MergeRequest_PR pipeline #29879 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Bo Li <[email protected]>

bobboli added 2 commits February 13, 2026 14:21

Handle the dispatch-combine range overlap via strict range containmen…

d58bf1d

…t judgment. Add random seed to input. Signed-off-by: Bo Li <[email protected]>

Support perfect router.

268bfa3

Signed-off-by: Bo Li <[email protected]>

bobboli requested review from rosenrodt and xxi-nv February 23, 2026 09:34

coderabbitai Bot reviewed Feb 23, 2026

View reviewed changes

Comment thread tests/microbenchmarks/bench_moe_comm.py Outdated

bobboli added 8 commits February 23, 2026 12:26

Support multi-node multi-rank.

7ca11c6

Signed-off-by: Bo Li <[email protected]>

Print metadata in output.

7af78de

Signed-off-by: Bo Li <[email protected]>

Address comment.

d2cee6a

Signed-off-by: Bo Li <[email protected]>

Refactor benchmark report generation to include metadata once at the …

83965ee

…beginning. Signed-off-by: Bo Li <[email protected]>

Requires strict GPU-rank matching.

22a9a18

Signed-off-by: Bo Li <[email protected]>

_sync may cause high variance at first timed iter. Put warmup iters c…

88e89eb

…lose to timed iters. Signed-off-by: Bo Li <[email protected]>

CUDA graph support.

6c7e67b

Signed-off-by: Bo Li <[email protected]>

PDL toggl.e

dbc91e7

Signed-off-by: Bo Li <[email protected]>

xxi-nv approved these changes Mar 6, 2026

View reviewed changes

bobboli added 4 commits March 8, 2026 19:23

Encapsulate warmup iters into cuda graph as well.

f19ebae

Signed-off-by: Bo Li <[email protected]>

Use cupti + cuda event to profile kernel time inside cuda graph.

1e2e50c

Signed-off-by: Bo Li <[email protected]>

Add prepare_dispatch to support NVLINK_TWO_SIDED.

d1a4cfd

Signed-off-by: Bo Li <[email protected]>

Refactor comments.

b61895f

Signed-off-by: Bo Li <[email protected]>

bobboli requested a review from a team as a code owner March 10, 2026 16:36

bobboli enabled auto-merge (squash) March 10, 2026 16:38

Use cuda graph by default.

d3b9192

Signed-off-by: Bo Li <[email protected]>

bobboli force-pushed the alltoall_bench_refine branch from aef4d60 to d3b9192 Compare March 10, 2026 16:54

Merge branch 'main' into alltoall_bench_refine

e805b82

Signed-off-by: Bo Li <[email protected]>

tburt-nv reviewed Mar 10, 2026

View reviewed changes

Comment thread requirements-dev.txt Outdated

Update requirements-dev.txt

3872f9b

Co-authored-by: tburt-nv <[email protected]> Signed-off-by: Bo Li <[email protected]>

tburt-nv approved these changes Mar 10, 2026

View reviewed changes

bobboli merged commit ecf0057 into NVIDIA:main Mar 11, 2026
4 of 5 checks passed

limin2021 pushed a commit to limin2021/TensorRT-LLM that referenced this pull request Mar 19, 2026

[None][chore] Refine AlltoAll benchmark scripts. (NVIDIA#11649)

9263527

Signed-off-by: Bo Li <[email protected]>

longcheng-nv pushed a commit to longcheng-nv/TensorRT-LLM that referenced this pull request Mar 31, 2026

[None][chore] Refine AlltoAll benchmark scripts. (NVIDIA#11649)

465742b

Signed-off-by: Bo Li <[email protected]>

Conversation

bobboli commented Feb 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

coderabbitai Bot commented Feb 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bobboli commented Mar 5, 2026

Uh oh!

tensorrt-cicd commented Mar 5, 2026

Uh oh!

tensorrt-cicd commented Mar 5, 2026

Uh oh!

bobboli commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

bobboli commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

bobboli commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

tburt-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bobboli commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

tensorrt-cicd commented Mar 10, 2026

Uh oh!

bobboli commented Mar 11, 2026

Uh oh!

tensorrt-cicd commented Mar 11, 2026

Uh oh!

tensorrt-cicd commented Mar 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bobboli commented Feb 23, 2026 •

edited

Loading

coderabbitai Bot commented Feb 23, 2026 •

edited

Loading