Repository navigation
[TRTLLM-12635][feat] add bench_moe microbenchmark - #14507
Merged
xxi-nv merged 2 commits intoMay 25, 2026
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR introduces a comprehensive MoE microbenchmark framework for TensorRT-LLM, including structured benchmark specifications, multi-dimensional search-space expansion, advanced routing-control planning with materialization, multiple timing implementations (eager, CUDA graph, CUPTI), per-candidate execution orchestration, MPI-based worker loops with resume and CUDA-poison recovery, dashboard-oriented JSON reporting, optional Excel workbook generation, and extensive design documentation. It also makes the CuTe DSL dependency optional to improve compatibility. ChangesComplete MoE Microbenchmark Implementation
Optional CuTe DSL Dependency
🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (12)
tests/microbenchmarks/bench_moe/__init__.py (1)
23-23: 💤 Low valueOptional:
from __future__ import annotationsis unnecessary.Python 3.10+ (the minimum version for this project per setup.py) supports PEP 585 generics and type hints natively. This import can be removed.
♻️ Proposed cleanup
-from __future__ import annotations - import importlibBased on learnings: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features throughout the codebase, and you do not need to add
from __future__ import annotations.🤖 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/microbenchmarks/bench_moe/__init__.py` at line 23, Remove the unnecessary future import: delete the line "from __future__ import annotations" from tests/microbenchmarks/bench_moe/__init__.py since the project targets Python 3.10+, which natively supports PEP 563/annotations behavior; no code changes elsewhere are required.tests/microbenchmarks/bench_moe/__main__.py (1)
33-33: 💤 Low valueOptional:
from __future__ import annotationsis unnecessary.Python 3.10+ (the minimum version for this project) supports type hints natively without this future import.
♻️ Proposed cleanup
-from __future__ import annotations - import sysBased on learnings: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features throughout the codebase, and you do not need to add
from __future__ import annotations.🤖 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/microbenchmarks/bench_moe/__main__.py` at line 33, Remove the unnecessary future import: delete the line "from __future__ import annotations" from tests/microbenchmarks/bench_moe/__main__.py (the import statement itself is the symbol to remove) since the project requires Python 3.10+, which natively supports postponed evaluation of annotations; after removal, run tests/lint to ensure no formatting or unused-import warnings remain.tests/microbenchmarks/test_bench_moe_routing_control.py (1)
719-719: 💤 Low valueConsider tuple unpacking for cleaner concatenation.
For consistency with modern Python style, consider using unpacking instead of concatenation.
✨ Optional style improvement
- assert search.comm_methods == bench_moe._FORCED_COMM_ENV_VALUES + ("AUTO",) + assert search.comm_methods == (*bench_moe._FORCED_COMM_ENV_VALUES, "AUTO")🤖 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/microbenchmarks/test_bench_moe_routing_control.py` at line 719, Replace the explicit tuple concatenation with tuple unpacking for cleaner style: update the assertion comparing search.comm_methods to use a tuple constructed via unpacking of bench_moe._FORCED_COMM_ENV_VALUES plus "AUTO" (i.e., build the expected tuple with (*bench_moe._FORCED_COMM_ENV_VALUES, "AUTO") and compare that to search.comm_methods).tests/microbenchmarks/bench_moe/cli.py (3)
654-655: ⚡ Quick winSpecify encoding when opening JSON config file.
open(args.config_file)uses the platform default encoding which may vary. Explicitly specifyencoding="utf-8"for consistent JSON parsing across platforms.Proposed fix
- with open(args.config_file) as f: + with open(args.config_file, encoding="utf-8") as f: cfg = json.load(f)🤖 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/microbenchmarks/bench_moe/cli.py` around lines 654 - 655, The JSON config is opened with open(args.config_file) which relies on platform default encoding; update the file open call that loads cfg (the json.load call using args.config_file) to explicitly pass encoding="utf-8" when opening the file so JSON parsing is consistent across platforms.
540-551: ⚡ Quick winDuplicate
_coerce_str_tuplefunction.This function is duplicated verbatim in
search.py(lines 228-239). Consider keeping a single definition in one module and importing it in the other to reduce maintenance burden.🤖 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/microbenchmarks/bench_moe/cli.py` around lines 540 - 551, There is a duplicate function _coerce_str_tuple defined in this module and in search.py; remove the duplicate here and import the single canonical definition instead (or move it to a shared util module and import from there). Update the top-level imports to import _coerce_str_tuple from the chosen module (or add the new util import), and ensure callers in this file still reference _coerce_str_tuple unchanged; run tests to confirm nothing else relies on the local definition.
125-479: 💤 Low valueAccessing private
parser._actionsfor CLI tracking.Lines 472-476 iterate over
parser._actionsto detect which flags were explicitly provided. This is a common pattern but relies on argparse internals. Consider documenting this as intentional or usingargparse.Action'sdefaultsentinel pattern as an alternative.Otherwise the argument grouping and help text are well-structured.
🤖 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/microbenchmarks/bench_moe/cli.py` around lines 125 - 479, The code uses the private parser._actions in the parse_args function to build args._cli_provided by scanning option_strings; either mark this as intentional in a comment/PR note or replace with a public-safe approach: attach a custom argparse.Action subclass (or use a unique default sentinel per argument) to record when an option is explicitly set, and then populate args._cli_provided from those Actions instead of iterating parser._actions; update references to parser._actions, the option_strings loop, and args._cli_provided accordingly and add a short comment explaining the chosen approach.tests/microbenchmarks/bench_moe/search.py (2)
118-134: ⚡ Quick winPrefix unused variable with underscore.
moe_epfrom line 120 is never used. Prefix it with_to signal intentional discard.Proposed fix
# Mapping layout gate. try: - moe_ep, moe_tp, enable_dp = _resolve_mapping_layout(config, world_size) + _moe_ep, moe_tp, enable_dp = _resolve_mapping_layout(config, world_size) except ValueError as exc: return False, str(exc)🤖 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/microbenchmarks/bench_moe/search.py` around lines 118 - 134, The variable moe_ep returned from _resolve_mapping_layout is unused; rename it to a throwaway name (e.g., _moe_ep or simply _) in the assignment to indicate intentional discard and satisfy linters: change the left-hand side of "moe_ep, moe_tp, enable_dp = _resolve_mapping_layout(config, world_size)" to use a prefixed-underscore variable for the first element while leaving moe_tp and enable_dp unchanged.
36-54: 💤 Low valueConsider documenting the broad exception catch.
Line 53-54 catches a bare
Exception, which static analysis flags. While this is defensive (anycan_implementfailure should skip the backend), adding a brief comment explaining why broad catching is intentional here would suppress reviewer/linter concerns.The specific exceptions from
can_implementmay vary by backend implementation, so the current approach is pragmatic.🤖 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/microbenchmarks/bench_moe/search.py` around lines 36 - 54, In _check_backend_can_implement, add a short explanatory comment immediately above the "except Exception as exc:" block explaining that we intentionally catch broad exceptions from backend_cls.can_implement because backend implementations may raise diverse, backend-specific errors and any failure should cause the backend to be skipped; also include a linter hint/suppression (e.g., note about deliberately broad catch or a noqa/B902) if your project linter requires it. This documents intent and silences static-analysis warnings while preserving the existing fallback behavior.tests/microbenchmarks/bench_moe/routing/builders.py (1)
254-255: ⚡ Quick winSpecify explicit encoding when opening JSON files.
Opening files without an explicit encoding parameter relies on the platform default, which can vary. For JSON files, UTF-8 is the standard encoding.
♻️ Proposed fix
- with open(path) as f: + with open(path, encoding="utf-8") as f: payload = json.load(f)🤖 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/microbenchmarks/bench_moe/routing/builders.py` around lines 254 - 255, The file opens a JSON file with open(path) and relies on platform default encoding; update the open call to specify encoding='utf-8' so json.load reads UTF-8 reliably (the lines around open(path) and payload = json.load(f) should be changed to open with an explicit encoding argument).tests/microbenchmarks/bench_moe/routing/__init__.py (1)
76-111: 💤 Low valueConsider sorting
__all__for consistency.The static analysis tool notes that
__all__is not sorted. While the current ordering loosely groups by submodule, alphabetical sorting improves maintainability and makes it easier to spot missing/duplicate entries.🤖 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/microbenchmarks/bench_moe/routing/__init__.py` around lines 76 - 111, The __all__ list is not alphabetically sorted; please reorder the entries in the __all__ list (the symbol named __all__) into ascending alphabetical order (by string) while preserving existing quoted names and commas and without changing any symbol names; ensure duplicates are removed if found and keep the final list formatting consistent with the current style so imports/export behavior is unchanged (e.g., entries like "_COMM_PATTERN_NAMES", "_EXPERT_PATTERN_NAMES", "RoutingPlan", "_build_dispatch_matrix", etc. should be reordered alphabetically).tests/microbenchmarks/bench_moe/timing/cuda_graph.py (1)
66-130: 💤 Low valueAdd
strict=Truetozip()for length mismatch detection.Both
cupti_timesandforward_times_msshould have exactlyiterselements. Addingstrict=Truewill surface bugs if their lengths ever diverge.♻️ Suggested fix
forward_times_ms = [ - ct if ct is not None else et for ct, et in zip(cupti_times, forward_times_ms) + ct if ct is not None else et for ct, et in zip(cupti_times, forward_times_ms, strict=True) ]🤖 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/microbenchmarks/bench_moe/timing/cuda_graph.py` around lines 66 - 130, The list comprehension that merges cupti_times with forward_times_ms should use zip(..., strict=True) to ensure both sequences have exactly iters elements; update the expression in the function that computes forward_times_ms (the list comprehension using cupti_times and forward_times_ms) to call zip(cupti_times, forward_times_ms, strict=True) so a length mismatch raises immediately (refer to symbols cupti_times, forward_times_ms, and the list comprehension where they are zipped). Ensure any test expectations or callers still produce exactly iters elements.tests/microbenchmarks/bench_moe/case_runner.py (1)
846-857: 💤 Low valueConsider logging the exception from
moe.destroy()instead of silently passing.The try-except-pass at line 852-853 silently swallows any exception during cleanup. While it's correct that cleanup failures shouldn't crash the benchmark, logging the exception would help diagnose resource leak or teardown issues during development.
♻️ Suggested improvement
if moe is not None: try: moe.destroy() - except Exception: - pass + except Exception as exc: + _maybe_print_rank0(f"[bench_moe] moe.destroy() failed: {type(exc).__name__}: {exc}")🤖 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/microbenchmarks/bench_moe/case_runner.py` around lines 846 - 857, The cleanup finally silently swallows exceptions from moe.destroy(); change the except block to catch the exception as a variable (e.g., except Exception as e:) and log it (using logging.exception or the module's existing logger) with a clear message like "error destroying moe" so teardown failures are recorded but not raised; keep the subsequent environment restoration for TRTLLM_FORCE_COMM_METHOD unchanged.
🤖 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 `@tests/microbenchmarks/bench_moe/routing/native_logits.py`:
- Around line 237-240: The try/except around calling apply_orig(router_logits)
is catching all Exception and silently swallowing errors; change this to catch
only the expected exception(s) (e.g., AttributeError, TypeError, or a custom
expected exception) or check for the callable before calling (e.g., if
callable(apply_orig): ...), and if you must catch an error, log the failure at
debug level using the existing test/logger facility instead of pass; reference
apply_orig and router_logits when making the change so you only handle the
specific failure modes.
In `@tests/microbenchmarks/test_bench_moe_routing_control.py`:
- Line 302: The test unpacks obs_hist from the call to
bench_moe._observe_routing_metrics but never uses it; rename the unused unpacked
variable to start with an underscore (e.g., _obs_hist) in the tuple assignment
where bench_moe._observe_routing_metrics is called so it follows Python
conventions for unused variables and avoids linter warnings.
---
Nitpick comments:
In `@tests/microbenchmarks/bench_moe/__init__.py`:
- Line 23: Remove the unnecessary future import: delete the line "from
__future__ import annotations" from tests/microbenchmarks/bench_moe/__init__.py
since the project targets Python 3.10+, which natively supports PEP
563/annotations behavior; no code changes elsewhere are required.
In `@tests/microbenchmarks/bench_moe/__main__.py`:
- Line 33: Remove the unnecessary future import: delete the line "from
__future__ import annotations" from tests/microbenchmarks/bench_moe/__main__.py
(the import statement itself is the symbol to remove) since the project requires
Python 3.10+, which natively supports postponed evaluation of annotations; after
removal, run tests/lint to ensure no formatting or unused-import warnings
remain.
In `@tests/microbenchmarks/bench_moe/case_runner.py`:
- Around line 846-857: The cleanup finally silently swallows exceptions from
moe.destroy(); change the except block to catch the exception as a variable
(e.g., except Exception as e:) and log it (using logging.exception or the
module's existing logger) with a clear message like "error destroying moe" so
teardown failures are recorded but not raised; keep the subsequent environment
restoration for TRTLLM_FORCE_COMM_METHOD unchanged.
In `@tests/microbenchmarks/bench_moe/cli.py`:
- Around line 654-655: The JSON config is opened with open(args.config_file)
which relies on platform default encoding; update the file open call that loads
cfg (the json.load call using args.config_file) to explicitly pass
encoding="utf-8" when opening the file so JSON parsing is consistent across
platforms.
- Around line 540-551: There is a duplicate function _coerce_str_tuple defined
in this module and in search.py; remove the duplicate here and import the single
canonical definition instead (or move it to a shared util module and import from
there). Update the top-level imports to import _coerce_str_tuple from the chosen
module (or add the new util import), and ensure callers in this file still
reference _coerce_str_tuple unchanged; run tests to confirm nothing else relies
on the local definition.
- Around line 125-479: The code uses the private parser._actions in the
parse_args function to build args._cli_provided by scanning option_strings;
either mark this as intentional in a comment/PR note or replace with a
public-safe approach: attach a custom argparse.Action subclass (or use a unique
default sentinel per argument) to record when an option is explicitly set, and
then populate args._cli_provided from those Actions instead of iterating
parser._actions; update references to parser._actions, the option_strings loop,
and args._cli_provided accordingly and add a short comment explaining the chosen
approach.
In `@tests/microbenchmarks/bench_moe/routing/__init__.py`:
- Around line 76-111: The __all__ list is not alphabetically sorted; please
reorder the entries in the __all__ list (the symbol named __all__) into
ascending alphabetical order (by string) while preserving existing quoted names
and commas and without changing any symbol names; ensure duplicates are removed
if found and keep the final list formatting consistent with the current style so
imports/export behavior is unchanged (e.g., entries like "_COMM_PATTERN_NAMES",
"_EXPERT_PATTERN_NAMES", "RoutingPlan", "_build_dispatch_matrix", etc. should be
reordered alphabetically).
In `@tests/microbenchmarks/bench_moe/routing/builders.py`:
- Around line 254-255: The file opens a JSON file with open(path) and relies on
platform default encoding; update the open call to specify encoding='utf-8' so
json.load reads UTF-8 reliably (the lines around open(path) and payload =
json.load(f) should be changed to open with an explicit encoding argument).
In `@tests/microbenchmarks/bench_moe/search.py`:
- Around line 118-134: The variable moe_ep returned from _resolve_mapping_layout
is unused; rename it to a throwaway name (e.g., _moe_ep or simply _) in the
assignment to indicate intentional discard and satisfy linters: change the
left-hand side of "moe_ep, moe_tp, enable_dp = _resolve_mapping_layout(config,
world_size)" to use a prefixed-underscore variable for the first element while
leaving moe_tp and enable_dp unchanged.
- Around line 36-54: In _check_backend_can_implement, add a short explanatory
comment immediately above the "except Exception as exc:" block explaining that
we intentionally catch broad exceptions from backend_cls.can_implement because
backend implementations may raise diverse, backend-specific errors and any
failure should cause the backend to be skipped; also include a linter
hint/suppression (e.g., note about deliberately broad catch or a noqa/B902) if
your project linter requires it. This documents intent and silences
static-analysis warnings while preserving the existing fallback behavior.
In `@tests/microbenchmarks/bench_moe/timing/cuda_graph.py`:
- Around line 66-130: The list comprehension that merges cupti_times with
forward_times_ms should use zip(..., strict=True) to ensure both sequences have
exactly iters elements; update the expression in the function that computes
forward_times_ms (the list comprehension using cupti_times and forward_times_ms)
to call zip(cupti_times, forward_times_ms, strict=True) so a length mismatch
raises immediately (refer to symbols cupti_times, forward_times_ms, and the list
comprehension where they are zipped). Ensure any test expectations or callers
still produce exactly iters elements.
In `@tests/microbenchmarks/test_bench_moe_routing_control.py`:
- Line 719: Replace the explicit tuple concatenation with tuple unpacking for
cleaner style: update the assertion comparing search.comm_methods to use a tuple
constructed via unpacking of bench_moe._FORCED_COMM_ENV_VALUES plus "AUTO"
(i.e., build the expected tuple with (*bench_moe._FORCED_COMM_ENV_VALUES,
"AUTO") and compare that to search.comm_methods).
🪄 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: 1e56b205-d1ef-4f7f-a53f-e36a04576328
📒 Files selected for processing (29)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytests/microbenchmarks/BENCH_MOE_DASHBOARD_DESIGN.mdtests/microbenchmarks/BENCH_MOE_DESIGN.mdtests/microbenchmarks/BENCH_MOE_ROUTING_CONTROL_DESIGN.mdtests/microbenchmarks/bench_moe/BENCH_MOE_USER_GUIDE.mdtests/microbenchmarks/bench_moe/__init__.pytests/microbenchmarks/bench_moe/__main__.pytests/microbenchmarks/bench_moe/build.pytests/microbenchmarks/bench_moe/case_runner.pytests/microbenchmarks/bench_moe/cli.pytests/microbenchmarks/bench_moe/mapping.pytests/microbenchmarks/bench_moe/reporting.pytests/microbenchmarks/bench_moe/results.pytests/microbenchmarks/bench_moe/routing/__init__.pytests/microbenchmarks/bench_moe/routing/builders.pytests/microbenchmarks/bench_moe/routing/materialize.pytests/microbenchmarks/bench_moe/routing/native_logits.pytests/microbenchmarks/bench_moe/routing/parsing.pytests/microbenchmarks/bench_moe/search.pytests/microbenchmarks/bench_moe/specs.pytests/microbenchmarks/bench_moe/timing/__init__.pytests/microbenchmarks/bench_moe/timing/autotune.pytests/microbenchmarks/bench_moe/timing/cuda_graph.pytests/microbenchmarks/bench_moe/timing/cupti.pytests/microbenchmarks/bench_moe/timing/eager.pytests/microbenchmarks/bench_moe/utils.pytests/microbenchmarks/bench_moe/worker.pytests/microbenchmarks/sweep_moe.pytests/microbenchmarks/test_bench_moe_routing_control.py
xxi-nv
force-pushed
the
xxi/bench-moe-refactor-20260520-rebased
branch
from
May 25, 2026 01:03
b0e4b2a to
2c4e704
Compare
Add a modular MoE microbenchmark with routing-control workloads, resume/reporting support, and documented MPI/Slurm launch modes so GB200 benchmark sweeps are reproducible. Signed-off-by: xxi <xxi@nvidia.com>
xxi-nv
force-pushed
the
xxi/bench-moe-refactor-20260520-rebased
branch
from
May 25, 2026 01:10
2c4e704 to
201fabc
Compare
xxi-nv
changed the title
leslie-fang25
approved these changes
May 25, 2026
leslie-fang25
left a comment
leslie-fang25
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, just left a few nits
Tighten routing-control patch cleanup and reuse shared helpers so review feedback is resolved without changing benchmark behavior. Signed-off-by: xxi <xxi@nvidia.com>
xxi-nv
enabled auto-merge (squash)
May 25, 2026 06:14
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50156 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50156 [ run ] completed with state |
KleinBlueC
pushed a commit
to KleinBlueC/TensorRT-LLM
that referenced
this pull request
May 26, 2026
Signed-off-by: xxi <xxi@nvidia.com>
bmarimuthu-nv
pushed a commit
to nv-auto-deploy/TensorRT-LLM
that referenced
this pull request
May 28, 2026
Signed-off-by: xxi <xxi@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test plan