[None][feat] Layer-wise benchmarks: make model init more general and support weights loading #10562
Conversation
📝 WalkthroughWalkthroughMajor refactoring of layer-wise benchmarks tooling that consolidates model-specific runners (DeepSeekV3, Qwen3-Next) into a single generic Runner class with expanded public APIs for KV cache dtype configuration, MoE backend selection, and load format parameters. Supporting infrastructure is reorganized accordingly. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
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_qwen3_next.py (2)
1254-1265: Add a docstring topost_load_weights()explaining its purpose.The method wires
layer.next_layer_layernormfor each layer, which is critical for fusion paths that expect this attribute. While the hook is properly invoked by the standard model loader after weight loading completes, the method lacks documentation. Add a docstring consistent with other implementations in the codebase (e.g.,modeling_deepseekv3.py).
836-984: Fix forward method return type annotation: should beTuple[torch.Tensor, torch.Tensor]nottorch.Tensor.The
forward()method declares return type-> torch.Tensor(line 864) but actually returnshidden_states, residual(a tuple at line 984), violating theDecoderLayerabstract contract which requiresUnion[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]. Correct the annotation to match the actual return value.Additionally, add an NVIDIA copyright header at the top of the file per coding guidelines. The file currently has only the Apache license header.
examples/layer_wise_benchmarks/parse.py (1)
1-14: Add NVIDIA copyright header (latest modification year).This
.pysource file should include the NVIDIA copyright header per repo guidelines.examples/layer_wise_benchmarks/run.py (1)
1-21: Add NVIDIA copyright header (latest modification year).This
.pysource file should include the NVIDIA copyright header per repo guidelines.tensorrt_llm/tools/layer_wise_benchmarks/runner.py (3)
413-502: Global monkeypatch restore may be incorrect forload_pretrained_config.You patch
tensorrt_llm._torch.model_config.load_pretrained_configbut restore it to the importedload_pretrained_configsymbol, not necessarily the module’s original attribute value. This can leak state if another caller patched it earlier, or if the module attribute isn’t the same function object.Proposed fix
def scaled_from_ctx(scaled_from, mapping): @@ - tensorrt_llm._torch.model_config.load_pretrained_config = make_load_pretrained_config( - mapping, load_pretrained_config - ) + load_pretrained_config_attr_orig = tensorrt_llm._torch.model_config.load_pretrained_config + tensorrt_llm._torch.model_config.load_pretrained_config = make_load_pretrained_config( + mapping, load_pretrained_config_attr_orig + ) @@ try: yield finally: - tensorrt_llm._torch.model_config.load_pretrained_config = load_pretrained_config + tensorrt_llm._torch.model_config.load_pretrained_config = load_pretrained_config_attr_orig CutlassFusedMoE.select_alltoall_method_type = select_alltoall_method_type_cutlass TRTLLMGenFusedMoE.select_alltoall_method_type = select_alltoall_method_type_trtllm_gen WideEPMoE.select_alltoall_method_type = select_alltoall_method_type_wide_ep
611-632: Add NVIDIA copyright header and fix residual fusion logic to handle non-contiguous layer selections.The code has two issues:
Missing copyright header: The file lacks the required NVIDIA copyright header (per coding guidelines for all source files).
Correctness risk with residual fusion: When
layer_indicesselects non-contiguous layers from a model where some layers havenext_layer_layernormand others don't, you can passNoneas the residual to a fused layer. While layers gracefully handleresidual=Noneby treating it ashidden_states, this breaks residual fusion—the actual residual connection is lost. For example, if a layer without fusion setsoutput = (layer(...), None)and the next selected layer expects fusion, it receivesNoneinstead of the true residual.Ensure residual fusion is either uniform across all selected layers or keep a valid residual representation throughout the iteration, even for non-fused layers.
714-743: Missing case for INT8 in kv-cache dtype dictionary.The dictionary at line 738-742 maps
model_config.quant_config.kv_cache_quant_algousing keys"FP8","NVFP4", andNone, butINT8is a valid member ofKV_CACHE_QUANT_ALGO_LISTthat can appear from pre-quantized checkpoints. Whencreate_kv_cache_manager()is called withkv_cache_dtype="auto"(the default),validate_and_set_kv_cache_quant()returns early without modification, leaving checkpoint values likeQuantAlgo.INT8unchanged. This will cause aKeyErrorwhen attempting dictionary lookup. Add support for"INT8"key or handle it explicitly.
🤖 Fix all issues with AI agents
In @examples/layer_wise_benchmarks/parse.py:
- Around line 241-256: The ambiguous-module counting is off-by-one because you
use bisect.bisect(problem["runs"], range_start) (right-side insertion) and then
treat that index as the run id; change to compute a consistent run_id by using
bisect.bisect_left(problem["runs"], range_start) or subtracting 1 from
bisect(...) and clamping to valid range indices, then verify the range_start
actually falls before that run's end (compare against problem["ranges"] or
problem["runs"]+run length) before incrementing num_matches_per_run[run_id];
also make the error message use the same run indexing convention used elsewhere
(either 0-based run_id or format as human-friendly 1-based) so the reported
“N-th run” matches the computed run_id and update references to args.module,
problem["ranges"], problem["runs"], bisect.bisect, num_matches_per_run, and
run_id accordingly.
In @examples/layer_wise_benchmarks/README.md:
- Around line 191-197: Fix minor grammar and backend naming in the README:
change the phrase "only support" to "only supports" in the FP8 error line, and
standardize CLI backend flags to the exact option names used elsewhere (use
"--moe-backend DEEPGEMM" and "--moe-backend-for-prefilling DEEPGEMM" or
"WIDEEP") instead of alternate spellings like "DEEPGEMM-prefilling"; ensure the
troubleshooting text references the exact flag strings shown in examples so
readers can copy/paste them directly.
- Around line 25-27: The README lines invoking mpi_launch.sh assume the
environment variable LLM_MODELS_ROOT is set and that the model path contains no
spaces; update the examples to note that LLM_MODELS_ROOT must be exported and to
quote the model path when calling ./run.sh (e.g., wrap $LLM_MODELS_ROOT/... in
quotes) and add a brief sentence above the examples clarifying this requirement
so users running mpi_launch.sh or run.sh with config_ctx.yaml/config_gen.yaml
won’t hit copy/paste failures on paths with spaces.
In @tensorrt_llm/tools/layer_wise_benchmarks/__init__.py:
- Around line 1-4: Add the required NVIDIA copyright header (with latest
meaningful modification year 2026) at the top of the module containing the
exports for mark_ranges, BalanceMethod, and Runner; insert the standard
multi-line header comment before the existing imports (before the line "from
.mark_utils import mark_ranges") so the file begins with the repository's NVIDIA
header followed by the existing import and __all__ lines.
In @tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py:
- Around line 1-13: Add the required NVIDIA copyright header (latest meaningful
modification year = 2026) at the very top of the module (above the first import
such as nvtx) for mark_utils.py; include the full copyright owner line and
appropriate license/SPDX identifier consistent with other files in tensorrt_llm,
ensuring the header matches the project's standard multi-line header and
references 2026 as the year.
In @tensorrt_llm/tools/layer_wise_benchmarks/runner.py:
- Around line 378-412: The tuple returned from model_loader.load(...) assigns an
unused moe_load_balancer; rename that variable to _moe_load_balancer to mark it
intentionally unused. Update the assignment "model, moe_load_balancer =
model_loader.load(...)" to "model, _moe_load_balancer = model_loader.load(...)"
and ensure no other code later references moe_load_balancer (use model and
model.model.layers as-is).
- Around line 1-39: Add the required NVIDIA copyright header (with the latest
modification year) to the top of this source file so it precedes the existing
imports; update the year to the current/latest modification year and ensure any
repo-standard SPDX/license line used elsewhere in the project is included. Place
the header above the first token (currently the import block referencing
tensorrt_llm and torch) so files like runner.py, which define symbols such as
get_attention_backend, ModelConfig, ModelLoader, and KVCacheManager, have the
standard NVIDIA license block at file start.
🧹 Nitpick comments (6)
tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py (2)
3-12: Follow “keep module namespace” import guideline (avoidfrom ... import MLPLayer, ...).Current
from ... import ...imports violate the repo guideline “Always maintain the namespace when importing Python modules”. Consider importing the modules and referencing symbols off the module.Proposed refactor
import nvtx -from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate, Deepseekv3MoE -from tensorrt_llm._torch.models.modeling_nemotron_h import MLPLayer, NemotronHMOE +import tensorrt_llm._torch.models.modeling_deepseekv3 as modeling_deepseekv3 +import tensorrt_llm._torch.models.modeling_nemotron_h as modeling_nemotron_h from tensorrt_llm._torch.models.modeling_qwen3_next import ( Qwen3NextGatedDeltaNet, Qwen3NextSparseMoeBlock, ) from tensorrt_llm._torch.modules.attention import MLA, Attention from tensorrt_llm._torch.modules.fused_moe.interface import MoE from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.mamba.mamba2_mixer import Mamba2Mixer +import tensorrt_llm._torch.modules.mamba.mamba2_mixer as mamba2_mixerdef mark_ranges(): - DeepseekV3Gate.forward = nvtx.annotate("DeepseekV3Gate")(DeepseekV3Gate.forward) - Deepseekv3MoE.forward = nvtx.annotate("Deepseekv3MoE")(Deepseekv3MoE.forward) - MLPLayer.forward = nvtx.annotate("MLPLayer")(MLPLayer.forward) - NemotronHMOE.forward = nvtx.annotate("NemotronHMOE")(NemotronHMOE.forward) + modeling_deepseekv3.DeepseekV3Gate.forward = nvtx.annotate("DeepseekV3Gate")( + modeling_deepseekv3.DeepseekV3Gate.forward + ) + modeling_deepseekv3.Deepseekv3MoE.forward = nvtx.annotate("Deepseekv3MoE")( + modeling_deepseekv3.Deepseekv3MoE.forward + ) + modeling_nemotron_h.MLPLayer.forward = nvtx.annotate("MLPLayer")( + modeling_nemotron_h.MLPLayer.forward + ) + modeling_nemotron_h.NemotronHMOE.forward = nvtx.annotate("NemotronHMOE")( + modeling_nemotron_h.NemotronHMOE.forward + ) Qwen3NextGatedDeltaNet.forward = nvtx.annotate("Qwen3NextGatedDeltaNet")( Qwen3NextGatedDeltaNet.forward ) @@ - Mamba2Mixer.forward = nvtx.annotate("Mamba2Mixer")(Mamba2Mixer.forward) + mamba2_mixer.Mamba2Mixer.forward = nvtx.annotate("Mamba2Mixer")( + mamba2_mixer.Mamba2Mixer.forward + )
15-30: Makemark_ranges()idempotent to avoid double-wrappingforward.If
mark_ranges()is called twice (e.g., in multiple benchmark entrypoints / tests), you’ll nest NVTX wrappers and distort ranges. Consider a module-level guard (and name globals withG_...per guidelines).Proposed change
+G_MARK_RANGES_DONE = False + def mark_ranges(): + global G_MARK_RANGES_DONE + if G_MARK_RANGES_DONE: + return DeepseekV3Gate.forward = nvtx.annotate("DeepseekV3Gate")(DeepseekV3Gate.forward) @@ GatedMLP.forward = nvtx.annotate("GatedMLP")(GatedMLP.forward) Mamba2Mixer.forward = nvtx.annotate("Mamba2Mixer")(Mamba2Mixer.forward) + G_MARK_RANGES_DONE = Truetests/unittest/tools/test_layer_wise_benchmarks.py (1)
124-151: Subprocess execution in tests: prefersys.executableand keep inputs non-shell.The commands are list-form (good), but
["python3", ...]can be brittle (S607). Considersys.executable. Also, given S603, it’d be good to ensurellm_root/model_rootare always trusted fixtures (no user-provided paths).Minimal tweak for hermetic python
import os +import sys from subprocess import check_call @@ check_call( - ["python3", "parse.py", "--profile-dir", profile_dir, f"--world-size={world_size}"], + [sys.executable, "parse.py", "--profile-dir", profile_dir, f"--world-size={world_size}"], cwd=llm_root / "examples" / "layer_wise_benchmarks", )examples/layer_wise_benchmarks/parse.py (2)
141-220: Potential behavioral change: module filtering now requires exact range-name match.
args.module in range_namesmeans the user must pass the exact NVTX range text, and it won’t match partial/module-prefix cases. If that’s intended, consider documenting it inREADME.md; otherwise, consider allowing substring/prefix matching (or a--module-regex).
313-337:--error-on-unknown-kernelcurrently only triggers on a narrow subset of “unknowns”.Right now
warned_namesis only populated forat::native::.*elementwise_kernel<...misses; other unparsed kernels will silently pass even with--error-on-unknown-kernel. If the intent is “fail on any kernel name not mapped byparser_keywords”, you likely want to add all unmatched names towarned_names(or a separate counter) and include the first N offenders in the exception for debuggability.Also applies to: 371-372
examples/layer_wise_benchmarks/run.py (1)
12-21: Import style: keep module namespace on imports (per repo Python guidelines).Repo guideline says to maintain namespace when importing modules; consider switching to module imports (e.g.,
import tensorrt_llm.tools.layer_wise_benchmarks as layer_wise_benchmarks) and usinglayer_wise_benchmarks.Runner, etc.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
examples/layer_wise_benchmarks/README.mdexamples/layer_wise_benchmarks/parse.pyexamples/layer_wise_benchmarks/run.pytensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/tools/layer_wise_benchmarks/__init__.pytensorrt_llm/tools/layer_wise_benchmarks/deepseekv3_runner.pytensorrt_llm/tools/layer_wise_benchmarks/mark_utils.pytensorrt_llm/tools/layer_wise_benchmarks/qwen3_next_runner.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytensorrt_llm/tools/layer_wise_benchmarks/runner_factory.pytensorrt_llm/tools/layer_wise_benchmarks/runner_interface.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/tools/test_layer_wise_benchmarks.py
💤 Files with no reviewable changes (4)
- tensorrt_llm/tools/layer_wise_benchmarks/runner_interface.py
- tensorrt_llm/tools/layer_wise_benchmarks/deepseekv3_runner.py
- tensorrt_llm/tools/layer_wise_benchmarks/runner_factory.py
- tensorrt_llm/tools/layer_wise_benchmarks/qwen3_next_runner.py
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces. Do not use tabs
Always maintain the namespace when importing Python modules, even if only one class or function from a module is used
Python filenames should use snake_case (e.g.,some_file.py)
Python classes should use PascalCase (e.g.,class SomeClass)
Python functions and methods should use snake_case (e.g.,def my_awesome_function():)
Python local variables should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile)
Python global variables should use upper snake_case with prefixG(e.g.,G_MY_GLOBAL)
Python constants should use upper snake_case (e.g.,MY_CONSTANT)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Use comments in Python for code within a function, or interfaces that are local to a file
Use Google-style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with the format"""<type>: Description"""
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of errors possible
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block for the main logic
Files:
tensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/tools/layer_wise_benchmarks/__init__.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytensorrt_llm/tools/layer_wise_benchmarks/mark_utils.pytests/unittest/tools/test_layer_wise_benchmarks.pyexamples/layer_wise_benchmarks/parse.pyexamples/layer_wise_benchmarks/run.py
**/*.{cpp,cc,cxx,h,hpp,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification
Files:
tensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/tools/layer_wise_benchmarks/__init__.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytensorrt_llm/tools/layer_wise_benchmarks/mark_utils.pytests/unittest/tools/test_layer_wise_benchmarks.pyexamples/layer_wise_benchmarks/parse.pyexamples/layer_wise_benchmarks/run.py
**/*.{md,rst}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
When documenting CLI commands for TensorRT-LLM tools like
trtllm-serve,trtllm-bench, ortrtllm-eval, prefer using--configover--extra_llm_api_optionsfor specifying configuration files
Files:
examples/layer_wise_benchmarks/README.md
🧠 Learnings (16)
📚 Learning: 2025-09-03T13:16:06.824Z
Learnt from: nvpohanh
Repo: NVIDIA/TensorRT-LLM PR: 7478
File: tensorrt_llm/_torch/models/modeling_llama.py:1315-1315
Timestamp: 2025-09-03T13:16:06.824Z
Learning: The Llama4VisionEncoder.load_weights method signature is `def load_weights(self, weights: Dict)` and should not be confused with Llama4ForConditionalGeneration.load_weights which has a different signature including weight_mapper parameter.
Applied to files:
tensorrt_llm/_torch/models/modeling_qwen3_next.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:
examples/layer_wise_benchmarks/README.mdtests/integration/test_lists/test-db/l0_b200.ymltests/unittest/tools/test_layer_wise_benchmarks.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:
examples/layer_wise_benchmarks/README.mdtensorrt_llm/tools/layer_wise_benchmarks/mark_utils.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:
examples/layer_wise_benchmarks/README.md
📚 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:
examples/layer_wise_benchmarks/README.md
📚 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/test-db/l0_b200.ymltests/unittest/tools/test_layer_wise_benchmarks.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/test-db/l0_b200.yml
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-08-14T21:04:50.248Z
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-08-15T06:46:54.897Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-09-29T15:14:28.503Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-09-23T14:58:05.372Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
📚 Learning: 2025-10-20T16:54:09.824Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:6-6
Timestamp: 2025-10-20T16:54:09.824Z
Learning: In tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py, the import `from ...modules.mamba.layernorm_gated import _layer_norm_fwd` is correct and should not be changed to modules.fla.layernorm_gated. The _layer_norm_fwd function exists in both modules/mamba/layernorm_gated.py and modules/fla/layernorm_gated.py, but the mamba version is the intended implementation for this use case.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py
📚 Learning: 2025-10-20T17:07:18.745Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py:98-116
Timestamp: 2025-10-20T17:07:18.745Z
Learning: In NemotronH models (tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py), the gate (self.gate) returns topk_indices and topk_weights that are already in the correct shape to be passed directly to torch_ops.auto_deploy.torch_moe without needing to reshape them when hidden_states is flattened.
Applied to files:
tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.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:
tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.pytests/unittest/tools/test_layer_wise_benchmarks.py
🧬 Code graph analysis (5)
tensorrt_llm/_torch/models/modeling_qwen3_next.py (3)
tensorrt_llm/_torch/models/modeling_speculative.py (1)
post_load_weights(560-561)tensorrt_llm/_torch/modules/attention.py (1)
post_load_weights(2278-2290)tensorrt_llm/_torch/models/modeling_qwen3_moe.py (1)
post_load_weights(419-426)
tensorrt_llm/tools/layer_wise_benchmarks/__init__.py (1)
tensorrt_llm/tools/layer_wise_benchmarks/runner.py (2)
BalanceMethod(42-46)Runner(358-863)
tensorrt_llm/tools/layer_wise_benchmarks/runner.py (3)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
get_kv_cache_manager_cls(48-56)tensorrt_llm/_torch/pyexecutor/config_utils.py (3)
is_mla(12-16)is_nemotron_hybrid(4-9)is_qwen3_next(19-23)tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
_construct_checkpoint_loader(157-179)
tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py (2)
tensorrt_llm/_torch/models/modeling_nemotron_h.py (1)
MLPLayer(49-79)tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py (1)
Mamba2Mixer(33-311)
tests/unittest/tools/test_layer_wise_benchmarks.py (2)
tests/integration/defs/conftest.py (2)
llm_root(195-196)llm_models_root(83-97)tests/integration/defs/trt_test_alternative.py (1)
check_call(250-258)
🪛 Ruff (0.14.10)
tensorrt_llm/tools/layer_wise_benchmarks/runner.py
144-146: Avoid specifying long messages outside the exception class
(TRY003)
405-405: Unpacked variable moe_load_balancer is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
705-707: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
tests/unittest/tools/test_layer_wise_benchmarks.py
130-130: subprocess call: check for execution of untrusted input
(S603)
147-147: subprocess call: check for execution of untrusted input
(S603)
148-148: Starting a process with a partial executable path
(S607)
examples/layer_wise_benchmarks/parse.py
245-245: Loop control variable i not used within loop body
(B007)
246-246: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
252-255: Avoid specifying long messages outside the exception class
(TRY003)
372-372: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (7)
tests/integration/test_lists/test-db/l0_b200.yml (1)
85-88: LGTM: new Nemotron layer-wise benchmark unit test is added to L0 B200 pre-merge list.tensorrt_llm/tools/layer_wise_benchmarks/__init__.py (1)
1-4: Export change looks consistent (Runner/Balancing now come fromrunner).tests/unittest/tools/test_layer_wise_benchmarks.py (1)
124-150: The--mamba-ssm-cache-dtypeflag is properly implemented and wired through the codebase.The flag is defined in the argument parser at
examples/layer_wise_benchmarks/run.py:52with valid choices (auto,float16,bfloat16,float32), defaults to"auto", and is correctly passed to bothRunner.create_kv_cache_manager()and theRunnerconstructor. The parameter is validated viavalidate_and_set_mamba_ssm_cache_dtype()inrunner.py.examples/layer_wise_benchmarks/README.md (1)
33-36: Verify the 32k GEN example has the intendedseq_len_q(and that--seq-len-kv-cachealone is sufficient).For the “32k context length” GEN command, you set
--seq-len-kv-cache 32769but do not explicitly set--seq-len-q; that’s fine only ifconfig_gen.yamlalready constrainsseq_len_qas intended for this scenario.tensorrt_llm/tools/layer_wise_benchmarks/runner.py (1)
358-375: No action required. The code is compliant with the project's actual Python version requirement.The
list[int]annotation conforms to the project's Python 3.10+ requirement (as declared insetup.py), which is the actual minimum version used throughout the codebase. This pattern is already widely adopted across the codebase, including in other type hints in the same module and beyond. The coding guidelines mentioning Python 3.8+ are outdated; the discrepancy should be addressed by updating the guidelines, not this code.Likely an incorrect or invalid review comment.
examples/layer_wise_benchmarks/run.py (2)
212-226: [Your rewritten review comment text here]
[Exactly ONE classification tag]
50-60: No functional issue with dtype/format casing—each parameter's casing is internally consistent and validated correctly.The CLI argument choices use different casing, but this is by design and does not cause problems:
--kv-cache-dtype(lowercase:fp8,nvfp4,auto) matches the validatorvalidate_and_set_kv_cache_quant, which expects exactly_VALID_KV_CACHE_DTYPES = ("fp8", "nvfp4", "auto").--load-format(uppercase:AUTO,DUMMY) is normalized viaTorchLlmArgs.convert_load_format(), which calls.upper()and maps to theLoadFormatenum.--moe-backend-for-prefilling(uppercase:CUTLASS,DEEPGEMM,WIDEEP) matchesMoeConfig.backend, which expects uppercase values.--mamba-ssm-cache-dtype(lowercase:auto,float16,bfloat16,float32) is handled byvalidate_and_set_mamba_ssm_cache_dtype(), which processes lowercase strings or converts viastr_dtype_to_torch().No KeyError or silent mismatch will occur; each parameter's casing is intentional and correctly validated downstream.
yuantailing
changed the title
yuantailing
force-pushed
the
layer_wise_benchmarks
branch
from
f501adb to
bf19329
Compare
|
/bot run --disable-fail-fast |
yuantailing
force-pushed
the
layer_wise_benchmarks
branch
from
bf19329 to
477ddc8
Compare
|
PR_Github #31205 [ run ] triggered by Bot. Commit: |
|
PR_Github #31205 [ run ] completed with state
|
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
Signed-off-by: Tailing Yuan <[email protected]>
yuantailing
force-pushed
the
layer_wise_benchmarks
branch
from
c911327 to
cc5ac0a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #31322 [ run ] triggered by Bot. Commit: |
Signed-off-by: Tailing Yuan <[email protected]>
|
PR_Github #31322 [ run ] completed with state
|
| if args.enable_attention_dp is None: | ||
| args.enable_attention_dp = False | ||
| if args.kv_cache_dtype is None: | ||
| args.kv_cache_dtype = "auto" |
There was a problem hiding this comment.
Why not set defaults when parser.add_argument?
There was a problem hiding this comment.
@nv-guomingz It LGTM but can you also help review this part please?
3 participants
Summary by CodeRabbit
New Features
Improvements
Tests
✏️ Tip: You can customize this high-level summary in your review settings.
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 thestage-listparameter 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.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.