Skip to content

Conversation

@dongxuy04
Copy link
Collaborator

@dongxuy04 dongxuy04 commented Nov 3, 2025

Summary by CodeRabbit

  • New Features

    • Added load balancer support for multiple Mixture of Experts backends, enabling improved token routing and distribution.
    • Introduced auxiliary CUDA stream support for optimized load balancer operations.
    • Enhanced memory management with CPU tensor pageout hints for efficient memory handling.
    • Extended model architecture support for GptOssForCausalLM.
  • Tests

    • Expanded test coverage with parameterized MoE backend configurations and online load balancer scenarios.

Description

Enable EPLB for trtllm-gen and cutlass backend:

  • Refactor EPLB logic to base MoE class
  • Add EPLB weight loading for more quant methods of trtllm-gen backend like W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod.
  • Enable online EPLB for trtllm-gen and cutlass backend
  • Optimize memory usage by unmap some weights that might not used again
  • Explicitly mark EPLB unsupport and unverified quant methods
  • Add and update some tests.

Known issues:

For now, there might be some issues when using EPLB backend with TRTLLM-Gen backend and number of expert per group is larger than 32. The reason is that our current TRTLLM-Gen fused routing kernel doesn't support that as no models using that. But when using EPLB with redundant experts like EPLB288 for DeepSeek-R1, it seems that we have more than that. In thant case, ideally, we should not run the router inside TRTLLM-Gen, but since there are some issues with AutoTuner when using topk input in that case, it is not enabled now.

Test with GPT-OSS 120B on 4 GB200 nodes (16 GPU) with EP=16, BS=512, trtllm-bench IFB mode.

Without EPLB:
0: ===========================================================
0: = PERFORMANCE OVERVIEW
0: ===========================================================
0: Request Throughput (req/sec): 214.6129
0: Total Output Throughput (tokens/sec): 211175.4890
0: Total Token Throughput (tokens/sec): 441672.5608
0: Total Latency (ms): 38171.0446
0: Average request latency (ms): 35473.2092
0: Per User Output Throughput [w/ ctx] (tps/user): 27.7429
0: Per GPU Output Throughput (tps/gpu): 13198.4681

With 144 slot EPLB (16 redundant experts):
0: ===========================================================
0: = PERFORMANCE OVERVIEW
0: ===========================================================
0: Request Throughput (req/sec): 224.4187
0: Total Output Throughput (tokens/sec): 220824.1976
0: Total Token Throughput (tokens/sec): 461852.7904
0: Total Latency (ms): 36503.1961
0: Average request latency (ms): 33789.7656
0: Per User Output Throughput [w/ ctx] (tps/user): 29.1216
0: Per GPU Output Throughput (tps/gpu): 13801.5124

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

  • 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.

@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23395 [ run ] triggered by Bot. Commit: 6b6b9d1

@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23402 [ run ] triggered by Bot. Commit: f899b38

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23402 [ run ] completed with state SUCCESS. Commit: f899b38
/LLM/main/L0_MergeRequest_PR pipeline #17626 completed with status: 'FAILURE'

@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23429 [ run ] triggered by Bot. Commit: dfdc029

@dongxuy04 dongxuy04 force-pushed the user/dongxuy/trtllmgen_eplb branch from dfdc029 to abe56ca Compare November 4, 2025 03:12
@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23450 [ run ] triggered by Bot. Commit: abe56ca

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23429 [ run ] completed with state ABORTED. Commit: dfdc029
LLM/main/L0_MergeRequest_PR #17644 (Blue Ocean) completed with status: ABORTED

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23450 [ run ] completed with state SUCCESS. Commit: abe56ca
/LLM/main/L0_MergeRequest_PR pipeline #17659 completed with status: 'FAILURE'

@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@dongxuy04 dongxuy04 force-pushed the user/dongxuy/trtllmgen_eplb branch from 60d5d1e to 495e260 Compare November 6, 2025 05:41
@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23708 [ run ] triggered by Bot. Commit: 495e260

@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23723 [ run ] triggered by Bot. Commit: 65cccaa

@tensorrt-cicd
Copy link
Collaborator

PR_Github #23708 [ run ] completed with state ABORTED. Commit: 495e260
LLM/main/L0_MergeRequest_PR #17840 (Blue Ocean) completed with status: ABORTED

@xxi-nv xxi-nv self-requested a review November 6, 2025 08:18
@dongxuy04 dongxuy04 marked this pull request as ready for review November 6, 2025 11:15
@dongxuy04 dongxuy04 requested review from a team as code owners November 6, 2025 11:15
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Nov 6, 2025

📝 Walkthrough

Walkthrough

Integrates load-balancer support across MoE backends with slot-based token routing, GPU/CPU stage synchronization, and expert statistics tracking. Refactors weight loading for shared-weight scenarios and expands test coverage to exercise multiple MoE backends with online expert-parallel load balancing.

Changes

Cohort / File(s) Summary
MoE Core Infrastructure
tensorrt_llm/_torch/modules/fused_moe/interface.py, tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
Adds aux_stream_dict parameter and load-balancer initialization scaffolding to MoE base class. Introduces private methods for load-balancer lifecycle management (_init_load_balancer, _load_balancer_route, _load_balancer_update_statistic, etc.). Implements advise_tensor_pageout for CPU memory management of shared weights; tracks loaded_shared_weights for deferred unmapping.
MoE Backend: Cutlass
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
Extends forward_chunk with repeating_info parameter. Integrates load-balancer routing into forward path with per-call GPU/CPU stage synchronization. Routes tokens to slots instead of experts when load-balancer is active. Wires ExpertStatistic tracking and alltoall/reducescatter logic to handle load-balancer statistics.
MoE Backend: TRTLLMGen
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
Adds aux_stream_dict parameter to constructor. Implements _supports_load_balancer() and integrates load-balancer routing, GPU/CPU stages, and statistics collection into forward paths. Extends alltoall preparation to pass and retrieve load-balancer statistics on final call.
MoE Backend: WideEP
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
Adds _supports_load_balancer() method. Replaces direct load-balancer integration with internal method delegation (load_balancer* methods). Removes external AllReduce import usage, centralizing load-balancer orchestration.
MoE Backend: CuteDSL & Factory
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py, tensorrt_llm/_torch/modules/fused_moe/create_moe.py
Adds repeating_info parameter to CuteDslFusedMoE.forward_chunk. Widens allowed MoE backends for load-balanced operation; threads aux_stream_dict to TRTLLMGenFusedMoE constructor.
Model: GPT-OSS
tensorrt_llm/_torch/models/modeling_gpt_oss.py
Adds moe_load_balancer_config to MLPBlock. Derives num_slots from config.moe_load_balancer.num_slots (defaulting to num_experts). Adjusts swiglu tensor shapes to use num_slots instead of num_experts.
Quantization
tensorrt_llm/_torch/modules/fused_moe/quantization.py
Adds EPLB verification checks (_online_eplb_not_supported, _online_eplb_not_verified). Extends weight-loading paths to support CPU/GPU device-agnostic operations with non_blocking transfers. Registers raw shared weights for unmapping via _add_raw_shared_weights_for_unmap. Adds ignore_weight_scale parameter to load_all_fp4_weight_scales_and_alphas.
Tests: Definitions & Metadata
tests/integration/defs/.test_durations, tests/integration/defs/accuracy/test_llm_api_pytorch.py
Updates test duration entries with moe_backend parameter. Adds moe_backend parametrization to existing online EPLB tests (bfloat16, nvfp4); introduces new test_w4_4gpus_online_eplb with kv_cache_dtype parameter.
Tests: Test Lists
tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/qa/llm_function_core_sanity.txt, tests/integration/test_lists/qa/llm_function_nim.txt, tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml, tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml, tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
Replaces online_eplb test entries without moe_backend to variants with moe_backend=WIDEEP, CUTLASS, or TRTLLM. Adds new w4_4gpus_online_eplb[fp8] test. Expands coverage across multiple hardware targets.

Sequence Diagram(s)

sequenceDiagram
    participant App as Application
    participant MoE as MoE Module
    participant LB as Load Balancer
    participant Expert as Expert Slot
    participant GPU as GPU Stage
    participant CPU as CPU Stage

    App->>MoE: forward_chunk(x, router_logits, repeating_info=(is_first, is_last))
    
    rect rgb(200, 220, 240)
    Note over MoE,GPU: GPU Wait Stage (is_first_call)
    MoE->>LB: _load_balancer_start_wait_gpu_stage(is_first)
    LB->>GPU: wait for prior stage completion
    GPU-->>LB: ready
    end
    
    rect rgb(220, 240, 220)
    Note over MoE,LB: Routing & Statistics
    MoE->>LB: _load_balancer_update_statistic(token_selected_experts, is_first, is_last)
    MoE->>LB: _load_balancer_route(token_selected_experts)
    LB-->>MoE: token_selected_slots
    end
    
    rect rgb(240, 220, 220)
    Note over MoE,Expert: Expert Computation
    MoE->>Expert: fused_moe(x, token_selected_slots, weights/biases)
    Expert-->>MoE: output
    end
    
    rect rgb(240, 240, 200)
    Note over MoE,GPU: GPU Done Stage
    MoE->>LB: _load_balancer_done_wait_gpu_stage(is_first)
    LB->>GPU: signal completion
    end
    
    rect rgb(220, 240, 240)
    Note over MoE,CPU: CPU Set Stage (is_last_call)
    MoE->>LB: _load_balancer_start_set_cpu_stage(is_last)
    LB->>CPU: update statistics, prepare next layer
    CPU-->>LB: ready
    MoE->>LB: _load_balancer_done_set_cpu_stage(is_last)
    end
    
    MoE-->>App: output
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Areas requiring extra attention:

  • interface.py: Comprehensive load-balancer scaffolding with multiple new private methods (_init_load_balancer, _load_balancer_route, _load_balancer_update_statistic, etc.); verify correctness of slot boundary calculations, stream handling, and synchronization logic across all entry/exit points.
  • fused_moe_cutlass.py: Dense integration of load-balancer routing, GPU/CPU stage coordination, and token-to-slot remapping; ensure token_selected_slots is consistently used throughout alltoall/reducescatter paths and that is_first_call/is_last_call flags correctly trigger synchronization barriers.
  • fused_moe_trtllm_gen.py: Similar load-balancer integration with additional aux_stream_dict parameter threading; verify aux stream lifecycle and non-blocking operations align with GPU synchronization assumptions.
  • quantization.py: Multiple weight-loading paths now support GPU-resident tensors with non_blocking transfers; verify device context consistency, especially for shared-weight scenarios with CPU/GPU staging.
  • moe_load_balancer.py: New advise_tensor_pageout function using OS-level madvise; validate page-aligned computation and error handling on various platforms; verify loaded_shared_weights tracking does not cause memory leaks.
  • Cross-backend consistency: WideEPMoE, CutlassFusedMoE, and TRTLLMGenFusedMoE now implement load-balancer support differently; ensure token routing semantics (slots vs. experts) are consistent and compatible with alltoall prepare/dispatch logic.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 29.79% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
Description check ⚠️ Warning PR description is largely incomplete and does not follow the required template structure. Critical sections are missing or only partially filled. Complete the PR description by filling out the 'Description' and 'Test Coverage' sections. Ensure the description clearly explains the issue, solution, and lists relevant test cases that validate the changes.
✅ Passed checks (1 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the main feature being added: enabling EPLB (Expert Parallel Load Balancing) for the trtllm-gen and cutlass backends.
✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

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.

❤️ Share

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

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (7)
tensorrt_llm/_torch/models/modeling_gpt_oss.py (1)

143-175: Guard self.num_slots // moe_ep_size against non-divisible slot counts.

self.num_slots now comes from the load balancer config and may not be divisible by config.mapping.moe_ep_size. When that's the case, the integer floor division silently drops remainder slots and produces the wrong tensor length for swiglu_*; when num_slots < moe_ep_size it also returns 0, leading to empty tensors and downstream crashes. Please either validate divisibility (e.g., assert num_slots % moe_ep_size == 0) or adjust the construction to reflect the actual per-rank slot distribution.

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (2)

432-486: Requeue statistic update for non-last MNNVL chunks.

In the MNNVL path, you call _load_balancer_get_local_statistic_tensor() only when is_last_call is true, so chunked executions skip exporting intermediate stats. However, _load_balancer_update_statistic still runs for every chunk with ignore_allreduce true. If chunking happens, we never push the accumulated local tensor to the balancer, causing outdated global stats. Please either gather stats on every chunk (when tensor is non-None) or defer ignore_allreduce=True to the final chunk only.


649-692: Ensure CPU-stage sync wraps alltoall_combine.

_load_balancer_start_set_cpu_stage is invoked before alltoall_combine, but for chunked MNNVL you may return early (when not using alltoall) without calling _load_balancer_done_set_cpu_stage. Please move both start/done calls into a try/finally or otherwise guarantee they execute around every return path so the balancer doesn’t get stuck waiting.

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (4)

81-186: Restore SM==120 escape hatch.

Prior code short-circuited on sm_version >= 120 before load-balancer routing. After extending aux_stream_dict and load balancer support, we still abort execution on new hardware. If TRTLLM-Gen remains unsupported on SM120+, that’s fine, but we now reject the load balancer at creation time without fallback. Please document or downgrade this to a runtime feature flag so other backends (e.g., Cutlass) can still progress under EPLB.


333-380: Missing load-balancer routing in non-post-comm branch.

When post_quant_comm is false (no alltoall/allgather needed), we skip token routing entirely and never call _load_balancer_route, so the balancer’s decisions are ignored. Please invoke the load balancer prior to _quantize_for_post_quant_comm or ensure routing always happens before data leaves the rank.


398-417: Guard gathered statistic reshape by tensor size.

gathered_loadbalancer_local_statistic_info.view((self.mapping.moe_ep_size, self.num_experts)) assumes a dense tensor. When some ranks return None (e.g., static routing), the gather op hands back an empty tuple, and .view raises. Add a size check before reshaping or ensure the op always returns a filled tensor.


800-820: Keep repeat index increment outside load-balancer guard.

self.repeat_idx is only advanced when self.layer_load_balancer is set. Without the balancer, repeat_idx remains 0, but upstream code (e.g., chunked execution guards) still relies on it. Move the increment out of the guard so the original behavior is preserved.

🧹 Nitpick comments (2)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py (1)

143-151: Remove unused load-balancer parameters.

forward_chunk signature now accepts repeating_info, all_rank_num_tokens, and use_dp_padding but doesn’t use them or interact with the load balancer. That makes the interface inconsistent and silences actual load-balancer wiring for CuteDSL. Either plumb the parameters through (start/done GPU waits, route slots, statistics update) or drop them to avoid misleading call sites.

tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)

3926-3961: Consider parameterizing over moe_backend for broader online EPLB coverage.

The test hardcodes backend="TRTLLM" (line 3947), unlike the DeepSeekV3Lite online EPLB tests above which are parameterized over multiple backends. Given the PR enables EPLB for both TRTLLM and CUTLASS, consider adding parameterization to test both backends:

+    @pytest.mark.parametrize("moe_backend", ["CUTLASS", "TRTLLM"],
+                             ids=["cutlass", "trtllm"])
     def test_w4_4gpus_online_eplb(self, kv_cache_dtype, mocker):
+    def test_w4_4gpus_online_eplb(self, kv_cache_dtype, moe_backend, mocker):
         ...
         pytorch_config = dict(disable_overlap_scheduler=False,
                               cuda_graph_config=CudaGraphConfig(),
-                              moe_config=MoeConfig(backend="TRTLLM",
+                              moe_config=MoeConfig(backend=moe_backend,
                                                    load_balancer=eplb_config))

If TRTLLM-only testing is intentional, please document the rationale in the test docstring.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e822184 and 65cccaa.

📒 Files selected for processing (17)
  • tensorrt_llm/_torch/models/modeling_gpt_oss.py (2 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py (2 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py (1 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (16 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (11 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (9 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/interface.py (4 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py (7 hunks)
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py (27 hunks)
  • tests/integration/defs/.test_durations (2 hunks)
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py (3 hunks)
  • tests/integration/test_lists/qa/llm_function_core.txt (1 hunks)
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt (1 hunks)
  • tests/integration/test_lists/qa/llm_function_nim.txt (1 hunks)
  • tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml (1 hunks)
  • tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml (1 hunks)
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml (1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

  • tensorrt_llm/_torch/models/modeling_gpt_oss.py
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
  • tensorrt_llm/_torch/modules/fused_moe/interface.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

  • tensorrt_llm/_torch/models/modeling_gpt_oss.py
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
  • tensorrt_llm/_torch/modules/fused_moe/interface.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

  • tensorrt_llm/_torch/models/modeling_gpt_oss.py
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
  • tensorrt_llm/_torch/modules/fused_moe/interface.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
🧠 Learnings (16)
📓 Common learnings
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.
📚 Learning: 2025-09-09T09:40:45.658Z
Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
📚 Learning: 2025-08-21T02:39:12.009Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
📚 Learning: 2025-09-17T02:48:52.732Z
Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7781
File: tests/integration/test_lists/waives.txt:313-313
Timestamp: 2025-09-17T02:48:52.732Z
Learning: In TensorRT-LLM, `tests/integration/test_lists/waives.txt` is specifically for waiving/skipping tests, while other test list files like those in `test-db/` and `qa/` directories are for different test execution contexts (pre-merge, post-merge, QA tests). The same test appearing in both waives.txt and execution list files is intentional - the test is part of test suites but will be skipped due to the waiver.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
📚 Learning: 2025-08-14T23:23:27.449Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tensorrt_llm/_torch/models/modeling_gpt_oss.py
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
  • tensorrt_llm/_torch/modules/fused_moe/interface.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
📚 Learning: 2025-08-19T03:35:20.866Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4616-4626
Timestamp: 2025-08-19T03:35:20.866Z
Learning: In the MOE profiler TMA workspace preparation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu), the overlapping of TMA WS regions for NONE and FINALIZE variants is deliberate design to save memory space, as confirmed by djns99. The comment "reuse the same pointers to save space" reflects this intentional behavior.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py
📚 Learning: 2025-07-28T17:06:08.621Z
Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_nim.txt
  • tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
📚 Learning: 2025-08-26T09:49:04.956Z
Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Applied to files:

  • tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
  • tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
📚 Learning: 2025-08-14T06:36:40.701Z
Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

Applied to files:

  • tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
📚 Learning: 2025-08-20T07:43:36.447Z
Learnt from: ChristinaZ
Repo: NVIDIA/TensorRT-LLM PR: 7068
File: cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh:169-172
Timestamp: 2025-08-20T07:43:36.447Z
Learning: In TensorRT-LLM MOE kernels, when processing up to 128 experts across 32 threads, each thread handles at most 4 experts (N < 5 constraint), where N represents candidates per thread rather than total system capacity.

Applied to files:

  • tensorrt_llm/_torch/models/modeling_gpt_oss.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
📚 Learning: 2025-08-06T13:58:07.506Z
Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.

Applied to files:

  • tests/integration/test_lists/qa/llm_function_core_sanity.txt
  • tests/integration/test_lists/qa/llm_function_core.txt
📚 Learning: 2025-08-08T22:03:40.707Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Applied to files:

  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
📚 Learning: 2025-08-09T20:57:04.084Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

Applied to files:

  • tensorrt_llm/_torch/modules/fused_moe/create_moe.py
  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
📚 Learning: 2025-08-19T12:45:11.997Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Applied to files:

  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
📚 Learning: 2025-09-19T21:28:13.751Z
Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Applied to files:

  • tensorrt_llm/_torch/modules/fused_moe/quantization.py
📚 Learning: 2025-08-21T21:48:35.135Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:399-417
Timestamp: 2025-08-21T21:48:35.135Z
Learning: CUTLASS extensions in TensorRT-LLM (located under cpp/tensorrt_llm/cutlass_extensions/) are designed to integrate with and extend functionality in the external CUTLASS repository. When analyzing these extensions, their consumers and functionality wiring may exist in the CUTLASS codebase rather than within TensorRT-LLM itself.

Applied to files:

  • tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
🧬 Code graph analysis (9)
tensorrt_llm/_torch/modules/fused_moe/create_moe.py (3)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (1)
  • WideEPMoE (26-945)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)
  • CutlassFusedMoE (28-782)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)
  • TRTLLMGenFusedMoE (29-853)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py (1)
tensorrt_llm/_torch/utils.py (1)
  • Fp4QuantizedTensor (110-117)
tensorrt_llm/_torch/modules/fused_moe/quantization.py (4)
tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py (2)
  • need_load_shared_weights (462-463)
  • get_load_expert_ids (452-454)
tensorrt_llm/logger.py (1)
  • warning (132-133)
tensorrt_llm/_torch/modules/fused_moe/interface.py (2)
  • _add_raw_shared_weights_for_unmap (295-299)
  • register_all_parameter_slot_and_to_fix_weight_fns (424-456)
tensorrt_llm/_torch/modules/linear.py (2)
  • load_weight_shard (65-118)
  • TensorParallelMode (50-62)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (5)
tensorrt_llm/_torch/utils.py (1)
  • Fp4QuantizedTensor (110-117)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (3)
  • _supports_load_balancer (254-256)
  • enable_alltoall (243-246)
  • moe_alltoall_backend (249-252)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (2)
  • _supports_load_balancer (933-935)
  • enable_alltoall (243-246)
tensorrt_llm/_torch/modules/fused_moe/interface.py (11)
  • _supports_load_balancer (301-306)
  • _load_balancer_start_wait_gpu_stage (317-320)
  • _load_balancer_done_wait_gpu_stage (322-325)
  • enable_alltoall (610-613)
  • AlltoallMethodType (26-34)
  • _load_balancer_update_statistic (327-352)
  • _load_balancer_route (354-361)
  • _load_balancer_get_local_statistic_tensor (373-378)
  • _load_balancer_update_statistic_with_gathered_statistic (380-385)
  • _load_balancer_start_set_cpu_stage (363-366)
  • _load_balancer_done_set_cpu_stage (368-371)
tensorrt_llm/_mnnvl_utils.py (1)
  • mnnvl_moe_alltoallv_prepare_without_allgather (402-446)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (3)
tensorrt_llm/_torch/distributed/ops.py (2)
  • allgather (265-298)
  • reducescatter (363-428)
tensorrt_llm/_torch/modules/fused_moe/interface.py (11)
  • _load_balancer_start_wait_gpu_stage (317-320)
  • _load_balancer_done_wait_gpu_stage (322-325)
  • enable_alltoall (610-613)
  • AlltoallMethodType (26-34)
  • _load_balancer_update_statistic (327-352)
  • _load_balancer_route (354-361)
  • _load_balancer_get_local_statistic_tensor (373-378)
  • _load_balancer_update_statistic_with_gathered_statistic (380-385)
  • _load_balancer_start_set_cpu_stage (363-366)
  • _load_balancer_done_set_cpu_stage (368-371)
  • _supports_load_balancer (301-306)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (3)
  • enable_alltoall (243-246)
  • moe_alltoall_backend (249-252)
  • _supports_load_balancer (254-256)
tests/integration/defs/accuracy/test_llm_api_pytorch.py (4)
tests/integration/defs/conftest.py (1)
  • parametrize_with_ids (1818-1844)
tensorrt_llm/llmapi/llm_args.py (2)
  • KvCacheConfig (1265-1409)
  • MoeConfig (268-302)
tensorrt_llm/_torch/model_config.py (1)
  • MoeLoadBalancerConfig (30-64)
tests/integration/defs/accuracy/accuracy_core.py (3)
  • GSM8K (334-349)
  • evaluate (184-247)
  • evaluate (766-776)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (4)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)
  • _supports_load_balancer (183-185)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (2)
  • _supports_load_balancer (933-935)
  • forward_chunk (382-694)
tensorrt_llm/_torch/modules/fused_moe/interface.py (9)
  • _supports_load_balancer (301-306)
  • _load_balancer_start_wait_gpu_stage (317-320)
  • _load_balancer_done_wait_gpu_stage (322-325)
  • _load_balancer_update_statistic (327-352)
  • _load_balancer_route (354-361)
  • _load_balancer_get_local_statistic_tensor (373-378)
  • _load_balancer_update_statistic_with_gathered_statistic (380-385)
  • _load_balancer_start_set_cpu_stage (363-366)
  • _load_balancer_done_set_cpu_stage (368-371)
tensorrt_llm/_mnnvl_utils.py (3)
  • MnnvlMoe (352-624)
  • mnnvl_moe_alltoallv_prepare_without_allgather (402-446)
  • mnnvl_moe_alltoallv (531-592)
tensorrt_llm/_torch/modules/fused_moe/interface.py (6)
tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py (23)
  • get_moe_load_balancer (1146-1156)
  • add_layer (907-943)
  • get_repeat_count (456-457)
  • set_initial_weight_assignments (534-543)
  • is_static_routing (459-460)
  • is_static_routing (880-882)
  • _add_raw_host_weight_for_unmap (502-511)
  • start_wait_gpu_stage (588-604)
  • done_wait_gpu_stage (606-615)
  • update_local_statistic (648-681)
  • update_statistic_with_local_ids (729-763)
  • route (797-814)
  • start_set_cpu_stage (617-634)
  • done_set_cpu_stage (636-646)
  • get_local_statistic_tensor (683-697)
  • update_statistic_with_gathered_statistic (699-727)
  • register_weight_slot (489-500)
  • make_tensor_host_accessible (561-563)
  • add_register_weight_fn (551-559)
  • add_to_migrate_weight_fn (545-549)
  • get_load_expert_ids (452-454)
  • share_host_tensor_with_shape (212-237)
  • pre_register_host_tensor_with_shape (187-210)
tensorrt_llm/_torch/model_config.py (4)
  • num_local_slots (46-47)
  • get_layer_initial_global_assignments (57-64)
  • slot_start (50-51)
  • slot_end (54-55)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)
  • _supports_load_balancer (254-256)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)
  • _supports_load_balancer (183-185)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (1)
  • _supports_load_balancer (933-935)
tensorrt_llm/_torch/distributed/ops.py (1)
  • AllReduce (554-710)
tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py (1)
tensorrt_llm/_utils.py (2)
  • data_ptr (974-975)
  • numel (1002-1003)
🪛 Ruff (0.14.3)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py

148-148: Unused method argument: all_rank_num_tokens

(ARG002)


149-149: Unused method argument: use_dp_padding

(ARG002)


150-150: Unused method argument: repeating_info

(ARG002)

tensorrt_llm/_torch/modules/fused_moe/quantization.py

2134-2135: Consider iterable unpacking instead of concatenation

Replace with iterable unpacking

(RUF005)


2139-2140: Consider iterable unpacking instead of concatenation

Replace with iterable unpacking

(RUF005)

tensorrt_llm/_torch/modules/fused_moe/interface.py

431-431: Loop control variable expert_id not used within loop body

(B007)

tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py

59-59: Avoid specifying long messages outside the exception class

(TRY003)


89-89: Avoid specifying long messages outside the exception class

(TRY003)

🔇 Additional comments (11)
tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml (1)

96-99: Inconsistency between AI summary and code state.

The AI-generated summary claims these tests are "Added" but the current code shows them commented out with a GDRCopy infrastructure dependency note. Additionally, the summary states that test_guided_decoding[xgrammar-mtp_nextn=0] is deleted, yet line 102 shows it remains in the file.

Please clarify: are these test entries intentionally disabled pending infrastructure setup, or should they be enabled as part of this PR? The presence of moe_backend parameters (WIDEEP, TRTLLM, CUTLASS) suggests these are new test variants designed to validate the PR's EPLB support across multiple MoE backends.

tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml (3)

45-45: No issues found. Test method and fp8 parameter both exist and are correctly supported.

Verification confirms that test_w4_4gpus_online_eplb exists in TestGPTOSS class at line 3931 of tests/integration/defs/accuracy/test_llm_api_pytorch.py with the parameterization ["auto", pytest.param("fp8", marks=skip_pre_blackwell)]. The fp8 variant is explicitly supported and properly marked for Blackwell devices. The YAML entry is valid.


37-40: All test methods and moe_backend parameters are valid and supported.

The test methods test_nvfp4_4gpus_online_eplb and test_bfloat16_4gpus_online_eplb exist in the test suite and are correctly parameterized with moe_backend values. The YAML entries use only supported backend combinations:

  • test_nvfp4_4gpus_online_eplb supports moe_backend values: WIDEEP, TRTLLM
  • test_bfloat16_4gpus_online_eplb supports moe_backend values: WIDEEP, CUTLASS, TRTLLM

All parameter combinations in lines 37-40 match the test definitions.


37-45: Test matrix alignment verified—all entries are valid.

All test methods and parameterized variants referenced in the yml file exist in tests/integration/defs/accuracy/test_llm_api_pytorch.py and correctly match the test definitions:

  • test_bfloat16_4gpus_online_eplb supports moe_backend variants (WIDEEP, CUTLASS, TRTLLM) and mtp_nextn=[0, 2]
  • test_nvfp4_4gpus_online_eplb supports moe_backend variants (WIDEEP, TRTLLM) and fp8kv=[True, False]
  • test_w4_4gpus correctly parameterized with backend and tp/dp variants
  • test_w4_4gpus_online_eplb supports kv_cache_dtype=[auto, fp8]
  • test_fp8_tp4 and test_nvfp4_tp4 both exist

The yml file lists a selected subset of generated test variants, which is standard practice for targeted test execution across different environments.

tests/integration/test_lists/qa/llm_function_core.txt (1)

467-470: Quiet add/remove; no issues spotted.

The test list entries simply switch backend parameters to WIDEEP and remain well-formed. Nothing to change.

tests/integration/defs/.test_durations (1)

228-229: LGTM! Test duration metadata updated consistently.

The test identifiers now include the moe_backend=WIDEEP parameter, which aligns with the broader parameterization of online EPLB tests across MOE backends introduced in this PR.

Also applies to: 255-256

tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml (1)

42-45: Backend coverage expanded appropriately for GB300 hardware.

The changes expand online EPLB test coverage across multiple MOE backends (WIDEEP, TRTLLM, CUTLASS), which is appropriate for comprehensive multi-GPU testing on GB300 hardware. The test variants cover key combinations:

  • nvfp4 with both WIDEEP and TRTLLM backends (fp8kv=True)
  • bfloat16 with both WIDEEP and CUTLASS backends (mtp_nextn=2)

Verify whether any of these online EPLB tests require the ISOLATION marker (like Line 41) to prevent resource conflicts or ensure deterministic behavior, especially given the multi-backend and multi-GPU nature of these tests.

tests/integration/test_lists/qa/llm_function_nim.txt (1)

280-283: LGTM! Online EPLB tests now explicitly specify WIDEEP backend.

The test identifiers are updated consistently to include moe_backend=WIDEEP, covering the key parameter combinations for online EPLB scenarios (bfloat16 with mtp_nextn variations and nvfp4 with fp8kv variations).

tests/integration/test_lists/qa/llm_function_core_sanity.txt (1)

55-56: LGTM! Sanity test list updated with explicit backend parameterization.

The online EPLB test entries now consistently include moe_backend=WIDEEP, maintaining appropriate coverage for sanity testing with the key parameter variations (mtp_nextn and fp8kv).

Also applies to: 60-61

tests/integration/defs/accuracy/test_llm_api_pytorch.py (2)

1620-1641: LGTM! MoE backend parameterization enables comprehensive online EPLB testing.

The parameterization correctly extends test coverage to verify online expert parallel load balancing across WIDEEP, CUTLASS, and TRTLLM backends. The implementation properly integrates the moe_backend parameter into the test logic and follows established patterns.


1645-1667: Verify CUTLASS exclusion from NVFP4 online EPLB testing.

The parameterization for this test includes only ["WIDEEP", "TRTLLM"], excluding CUTLASS which was present in the bfloat16 variant above (line 1620). Please confirm this is intentional—likely due to NVFP4 quantization not supporting CUTLASS with online EPLB, or a known issue on specific SM versions.

Run the following to check if there are any related comments or constraints in the MoE backend implementations:

Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
@dongxuy04 dongxuy04 force-pushed the user/dongxuy/trtllmgen_eplb branch from aaaf4bb to b94a262 Compare November 12, 2025 07:39
@dongxuy04 dongxuy04 requested a review from a team as a code owner November 12, 2025 07:39
@dongxuy04 dongxuy04 requested a review from liji-nv November 12, 2025 07:39
@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@dongxuy04
Copy link
Collaborator Author

/bot kill

@dongxuy04 dongxuy04 force-pushed the user/dongxuy/trtllmgen_eplb branch from b94a262 to 8378276 Compare November 12, 2025 07:43
@dongxuy04 dongxuy04 removed the request for review from liji-nv November 12, 2025 07:44
@dongxuy04
Copy link
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24269 [ run ] triggered by Bot. Commit: 8378276

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24226 [ run ] completed with state ABORTED. Commit: aaaf4bb
LLM/main/L0_MergeRequest_PR #18272 (Blue Ocean) completed with status: ABORTED

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24272 [ kill ] triggered by Bot. Commit: 8378276

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24269 [ run ] completed with state ABORTED. Commit: 8378276

@tensorrt-cicd
Copy link
Collaborator

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

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24273 [ run ] triggered by Bot. Commit: 8378276

@tensorrt-cicd
Copy link
Collaborator

PR_Github #24273 [ run ] completed with state SUCCESS. Commit: 8378276
/LLM/main/L0_MergeRequest_PR pipeline #18310 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

@dongxuy04 dongxuy04 merged commit 9241cca into NVIDIA:main Nov 12, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants