Draft - DO NOT REVIEW - AD host time improvement for large prefill (and more changes that will be removed - #11565
Draft - DO NOT REVIEW - AD host time improvement for large prefill (and more changes that will be removed#11565MrGeva wants to merge 7 commits into
Conversation
Signed-off-by: Eran Geva <egeva@nvl72126-T15.cm.cluster>
Signed-off-by: Eran Geva <egeva@nvl72094-T16.cm.cluster>
Signed-off-by: Eran Geva <egeva@nvl72126-T15.cm.cluster>
Signed-off-by: Eran Geva <egeva@nvl72064-T17.cm.cluster>
Signed-off-by: Eran Geva <egeva@oci-hsg-cs-001-login-01.cm.cluster>
MrGeva
changed the title
📝 WalkthroughWalkthroughThis PR introduces comprehensive Auto-Deploy enhancements: a new TRT-LLM attention backend with host metadata management, SwiGLU MLP operations (standard and NVFP4-quantized variants), multi-stream execution for MLA attention and shared-expert MoE, graph transformation pipelines for pattern matching and fusion, and updated KV cache and sequence interface handling. Configuration files and extensive tests validate the new features across multiple models and backends. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ❌ 3❌ Failed checks (3 warnings)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py (1)
1-1:⚠️ Potential issue | 🟡 MinorMissing NVIDIA copyright header
The file has no copyright header. Per coding guidelines, all source files must contain an NVIDIA copyright header using the Apache License 2.0 format, and the year must be updated on any modified file. As per coding guidelines, "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification."
📄 Proposed header to add at the top of the file
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Optional🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py` at line 1, Add the required NVIDIA copyright header (Apache License 2.0) to the top of this source file including the year of latest meaningful modification; update the year when you modify the file in the future. Place the header above all imports (before the existing "from typing import Optional" line) so it appears as the first lines of the file, and ensure it matches the project's standard NVIDIA Apache-2.0 boilerplate.tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py (1)
1-1:⚠️ Potential issue | 🟡 MinorUpdate copyright year to 2026.
The file has been meaningfully modified in this 2026 PR, but the header still says
2025.🔧 Proposed fix
-# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines: "Include NVIDIA copyright header on ALL new files and update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py` at line 1, Update the copyright header year from 2025 to 2026 in the file tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py by editing the top-of-file copyright comment to read 2026 so the header matches the current PR modifications (update the existing copyright comment line).tests/integration/defs/accuracy/test_llm_api_autodeploy.py (2)
1-1:⚠️ Potential issue | 🟡 MinorUpdate copyright year to 2026.
The header says
Copyright (c) 2025but the file is being modified in 2026. As per coding guidelines: "update year on modified files."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py` at line 1, Update the file header copyright year from 2025 to 2026 by editing the top comment line that currently reads "Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved." to use 2026 instead so the header reflects the current modification year.
347-368:⚠️ Potential issue | 🟠 Major
multi_stream_moeis unconditionally enabled despite running withworld_sizein[1, 4, 8].
TestNemotronMOE.get_default_kwargsconditionally enablesmulti_stream_moewith"enabled": world_size == 1and includes a comment stating it "currently does not work for world_size > 1". In contrast,TestNemotronSuperV3.get_default_kwargsunconditionally sets"enabled": Trueand does not accept aworld_sizeparameter. SinceTestNemotronSuperV3.test_fp8runs withworld_sizein[1, 4, 8], multi-stream MoE will be enabled for multi-GPU configurations. Either add a conditional enable based onworld_size(if the limitation applies to SuperV3), or add a comment explaining why this limitation does not apply to SuperV3.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py` around lines 347 - 368, The get_default_kwargs in TestNemotronSuperV3 currently enables the multi_stream_moe transform unconditionally; update it to either conditionally set "enabled": (world_size == 1) based on the test's world_size parameter (matching the behavior in TestNemotronMOE) or add a clear comment in get_default_kwargs explaining why multi_stream_moe can safely be enabled for multi-GPU runs in TestNemotronSuperV3 (referencing TestNemotronMOE, multi_stream_moe, get_default_kwargs, test_fp8, and world_size so reviewers can verify the rationale).
🧹 Nitpick comments (20)
tensorrt_llm/_torch/auto_deploy/export/export.py (1)
280-280: Constant re-created on every loop iteration.
_tensor_list_typesis invariant — hoist it above thefor node in list(graph.nodes)loop (e.g., next tomoe_ops).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/export/export.py` at line 280, The tuple _tensor_list_types is being recreated inside the loop; move its definition out of the loop so it is created once (hoist it to the module scope near the existing moe_ops declaration) and then reference _tensor_list_types inside the for node in list(graph.nodes) loop to avoid repeated allocations.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py (2)
331-335: Nested_count_moe_expertsduplicates the module-level_count_moe_experts_in_graph.These two functions are identical. Now that
_count_moe_experts_in_graphlives at module scope (lines 394-399), the nested copy can be eliminated.♻️ Proposed fix
- def _count_moe_experts(gm): - for node in gm.graph.nodes: - if node.op == "call_function" and "torch_moe" in str(node.target): - return len(node.args[3]) # w1_weight list length - return 0 - - assert _count_moe_experts(gm_full) == num_experts - assert _count_moe_experts(gm_reduced) == num_experts + assert _count_moe_experts_in_graph(gm_full) == num_experts + assert _count_moe_experts_in_graph(gm_reduced) == num_experts🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py` around lines 331 - 335, Remove the duplicate nested function _count_moe_experts inside the test (the local definition that iterates gm.graph.nodes and checks for "torch_moe" nodes), and replace its usages with the module-level _count_moe_experts_in_graph; i.e., delete the nested def _count_moe_experts(gm) and update any calls in the test to call _count_moe_experts_in_graph(gm) instead so the single shared implementation is used.
432-453: Consider adding numerical correctness checks to match theSimpleMoEForExporttest pattern.The test only validates graph structure and state-dict keys; it never calls the exported modules and asserts output equality. The corresponding
test_moe_export_with_reduced_expertstest (lines 337–349) asserts bothall_close(y_full, y_reduced)andall_close(y_loaded, y_full). Without equivalent checks here, the GLM4 path could silently produce wrong outputs while still passing all structural assertions.♻️ Proposed addition
gm_reduced.load_state_dict(model.state_dict(), strict=False) + + # Numerical: reduced export must produce the same output as full export + with torch.no_grad(): + y_full = gm_full(**sample_kwargs) + y_reduced_after_load = gm_reduced(**sample_kwargs) + assert all_close(y_full.logits, y_reduced_after_load.logits), ( + "Output after state-dict reload into reduced export differs from full export" + ) # Source model must be fully restored🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py` around lines 432 - 453, Add numerical output equality checks like in the SimpleMoEForExport test: run a fixed input through gm_full and gm_reduced (before and after calling gm_reduced.load_state_dict(model.state_dict(), strict=False)), capture outputs (y_full, y_reduced, y_loaded) and assert they match using the existing all_close or torch.allclose helper; keep the existing structural/state-dict assertions but append asserts all_close(y_full, y_reduced) and all_close(y_loaded, y_full) to ensure the GLM4 export path produces numerically equivalent outputs. Use the existing symbols gm_full, gm_reduced, model.state_dict(), and the helper all_close/torch.allclose to locate where to add these checks.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py (4)
174-213: First two tests use>= 1for fused op count; consider exact count for consistency.
test_fuse_add_cast_rms_normandtest_fuse_add_rms_norm_no_castassert_count_fused_ops(gm_t) >= 1, while the later tests assert exact counts (== 1,== 2). Since each of these models has exactly one fusible pair,== 1would be a tighter and more informative assertion.Proposed fix
- assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph" + assert _count_fused_ops(gm_t) == 1, "Expected exactly 1 fused op in graph"(Apply to both
test_fuse_add_cast_rms_normandtest_fuse_add_rms_norm_no_cast.)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py` around lines 174 - 213, Both tests use a loose assertion "_count_fused_ops(gm_t) >= 1"; change this to an exact check "== 1" in test_fuse_add_cast_rms_norm and test_fuse_add_rms_norm_no_cast to match later tests and reflect that each model produces exactly one fused op; update the assertion strings if desired (e.g., keep "fused op not found in graph") so the checks in those functions (_count_fused_ops and _count_rms_norm_ops called after _apply_transform on gm) are consistent and stricter.
11-12: Import style: prefer importing the module, not the class directly.Per coding guidelines, use
from package.subpackage import modulerather thanfrom package.subpackage.module import ClassName. Consider:from tensorrt_llm._torch.auto_deploy.transform import interface from tensorrt_llm._torch.auto_deploy.transform.library import fused_add_rms_normThen reference
interface.TransformConfigandfused_add_rms_norm.FuseAddRMSNorm. Note: the same pattern exists on unchanged lines (9–10) so this may warrant a broader cleanup. As per coding guidelines: "Python imports must usefrom package.subpackage import modulestyle; never usefrom module import Class."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py` around lines 11 - 12, The tests import classes directly; change to module-level imports to follow project style: replace direct class imports of TransformConfig and FuseAddRMSNorm with module imports (import the transform.interface module and the transform.library.fused_add_rms_norm module) and update usages to reference interface.TransformConfig and fused_add_rms_norm.FuseAddRMSNorm so all references to TransformConfig and FuseAddRMSNorm use the module-qualified names.
138-140:_count_add_opsis defined but never used.This helper is not called anywhere in the file. Consider removing it or adding a test assertion that uses it (e.g., verifying add ops are eliminated after fusion).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py` around lines 138 - 140, The helper function _count_add_ops is defined but unused; either remove the function or use it to assert fusion behavior—update the test (e.g., in test_fused_add_rms_norm) to call _count_add_ops(gm) before and after running the transformation pass and add an assertion that the number of torch.ops.aten.add.Tensor nodes decreases (or becomes zero) after fusion, or simply delete the _count_add_ops definition if no such assertion is desired; reference the function name _count_add_ops and the graph module variable gm when making the change.
6-8: Remove unused# noqadirective.Ruff (RUF100) flags the blanket
# noqaon line 6 as unused. The import is actively used (in_count_fused_ops), so no suppression is needed.Proposed fix
-from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( # noqa - flashinfer_fused_add_rms_norm, -) +from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( + flashinfer_fused_add_rms_norm, +)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py` around lines 6 - 8, Remove the unnecessary blanket "# noqa" on the import line that brings in flashinfer_fused_add_rms_norm; the import is used (e.g., by _count_fused_ops) so simply delete the "# noqa" directive from the import statement importing flashinfer_fused_add_rms_norm to satisfy Ruff RUF100.tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py (1)
14-15: Consider using built-in generic types instead oftypingimports.
List,Tuple, andOptionalfromtypingcan be replaced with the built-in equivalents (list,tuple,X | None) since the codebase targets Python 3.10+.♻️ Proposed refactor
-from typing import List, Optional, Tuple- matches: List[Tuple[Node, Optional[Node], Node]] = [] + matches: list[tuple[Node, Node | None, Node]] = []- cast_node: Optional[Node] = None + cast_node: Node | None = NoneAnd the return type annotation on
_apply:- ) -> Tuple[GraphModule, TransformInfo]: + ) -> tuple[GraphModule, TransformInfo]:Based on learnings: TensorRT-LLM requires Python ≥ 3.10, so PEP 585 built-in generics and
X | Noneunion syntax are available throughout the codebase.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py` around lines 14 - 15, Replace typing generics with Python 3.10+ built-ins: remove the imports of List, Optional, Tuple and use list, tuple, and X | None in annotations instead. Update the module-level import statement (remove List/Optional/Tuple) and change any function/type hints that use List[T], Tuple[...], or Optional[X] to list[T], tuple[...], and X | None respectively; also update the return type annotation on the method/function named _apply to use the built-in generics/union syntax.tensorrt_llm/_torch/auto_deploy/utils/_graph.py (1)
23-33: Module-level caches for derived op infrastructure look good.The lazy-creation pattern with
_derived_op_libsand_derived_op_registrykeeps registrations persistent and avoids redundant work. One minor naming nit per coding guidelines: module-level mutable globals are conventionallyG_-prefixed with upper snake_case (e.g.,G_DERIVED_OP_LIBS), though the leading underscore already signals these are internal.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py` around lines 23 - 33, Rename the module-level mutable globals to follow the project's naming convention: change _derived_op_libs to G_DERIVED_OP_LIBS and _derived_op_registry to G_DERIVED_OP_REGISTRY, update all references/usages throughout the module (including any initialization and lookups) to the new names, and preserve their types and initial values (Dict[str, torch.library.Library] and Dict[str, Callable]) so behavior is unchanged; ensure any import or external references in this file are also updated to the new identifiers.tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py (1)
1-7: Comprehensive test suite forcreate_derived_custom_op.Good coverage of correctness, registration visibility, caching semantics, fake implementation (default and custom), schema preservation with defaults, and
OpOverloadinput handling.Minor: the import on Line 6 imports
create_derived_custom_opdirectly from_graph. Per coding guidelines, preferfrom tensorrt_llm._torch.auto_deploy.utils import _graphand reference as_graph.create_derived_custom_op. This is a nit though.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py` around lines 1 - 7, The test imports create_derived_custom_op directly from the _graph module; change the import to bring in the module instead and reference the function via the module to follow the coding guideline: replace the direct import with importing tensorrt_llm._torch.auto_deploy.utils import _graph and update all usages of create_derived_custom_op in the test to call _graph.create_derived_custom_op (ensure any other references in this test file are updated accordingly).tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py (2)
27-30: Import style violates coding guidelines.The guideline requires
from package.subpackage import modulerather than importing individual symbols from a module. This should import the module and access its members via the module namespace.Proposed fix
-from tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention import ( - _GlobalTrtllmPlanner, - prepare_trtllm_metadata_host, -) +from tensorrt_llm._torch.auto_deploy.custom_ops.attention import trtllm_attentionThen use
trtllm_attention._GlobalTrtllmPlannerandtrtllm_attention.prepare_trtllm_metadata_hostthroughout.As per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py` around lines 27 - 30, Replace the symbol-level import of _GlobalTrtllmPlanner and prepare_trtllm_metadata_host with a module import so code follows the "from package.subpackage import module" guideline: import the module tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention as trtllm_attention, then update all usages to call trtllm_attention._GlobalTrtllmPlanner and trtllm_attention.prepare_trtllm_metadata_host so references resolve via the module namespace.
462-581: Unused loop variables and minor linting issues in the paged KV-cache test.Lines 513 and 564:
sis captured infor i, s in enumerate(seq_lengths)but never used in the loop body. Rename to_to signal intent.Proposed fix
- for i, s in enumerate(seq_lengths): + for i, _ in enumerate(seq_lengths):(Apply at both line 513 and line 564.)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py` around lines 462 - 581, In test_trtllm_attention_with_paged_kvcache, there are two for-loops using unused loop variable s (for i, s in enumerate(seq_lengths)) around the prefill page assignment and the decode page-update sections; rename s to _ in both loops so it signals an intentionally unused variable and fixes the lint warnings (update the two occurrences of "for i, s in enumerate(seq_lengths)" to "for i, _ in enumerate(seq_lengths)" while leaving surrounding logic, page_assignments, pages_per_seq_list, and cache_locs handling unchanged).tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py (2)
14-31: Import style: several imports pull specific symbols from modules.Per the coding guidelines, Python imports should use
from package.subpackage import modulestyle. Multiple imports here import specific classes/functions directly (e.g.,from torch.fx import GraphModule, Node,from ...utils._graph import create_derived_custom_op). Consider importing the modules and using qualified access.As per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py` around lines 14 - 31, Replace all direct symbol imports with module-level imports and update usages to qualified names: import torch.fx as fx and replace GraphModule/Node with fx.GraphModule and fx.Node; import ...utils._graph as _graph and use _graph.create_derived_custom_op; import ...utils.multi_stream_utils as multi_stream_utils and use multi_stream_utils._make_aux_stream_impl, multi_stream_utils.cuda_stream_manager, and multi_stream_utils.record_event_passthrough; import ...utils.node_utils as node_utils and use node_utils.is_op; import ...shim.interface as shim_interface and use shim_interface.CachedSequenceInterface; import ...models.factory as factory and use factory.ModelFactory; import ..interface as interface and use interface.BaseTransform, interface.SharedConfig, interface.TransformInfo, interface.TransformRegistry. Update all references in this file accordingly so no `from ... import Symbol` remains.
126-185: Graph mutation during iteration over pre-computedpairs— verify safety.
pairsis computed before the loop (line 144), and each iteration modifies the graph by insertingrec_node, replacingkv_linearwithnew_node, and erasingkv_linear. While the current logic appears safe (fork_point remains valid, Q-linears are correctly re-queried from livefork_point.users, andnode_orderis only used for original Q-linear nodes), this pattern is subtle.One edge case: if two fork points share a Q-linear user (unlikely in MLA but possible in unusual graphs), the
node_orderdict built at line 149 could be stale for the second fork point'searliest_qlookup. Consider adding a brief comment about the safety invariant (e.g., "pairs are independent: no fork_point shares Q-linears with another").🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py` around lines 126 - 185, The loop in _execute_kv_proj_in_aux_stream builds node_order once from graph.nodes but then mutates the graph (inserting rec_node, replacing/erasing kv_linear), which can make the precomputed node_order stale when locating earliest_q for later fork_point iterations; fix by either recomputing node_order (i.e., node_order = {n: i for i, n in enumerate(graph.nodes)}) immediately before computing q_linears/earliest_q inside the for fork_point, kv_linear in pairs loop, or explicitly document and assert the safety invariant (e.g., add a clear comment and an assert in _execute_kv_proj_in_aux_stream that pairs are independent: no two fork_point entries share Q-linear users) so future maintainers know why the single precomputed node_order is safe.tensorrt_llm/_torch/auto_deploy/config/default.yaml (1)
142-143: Inconsistentenableddefaults betweenfuse_nvfp4_swigluandfuse_swiglu.
fuse_nvfp4_swiglu(line 142) is implicitly enabled (noenabled: false), whilefuse_swiglu(line 167) is explicitly disabled. Both depend on their respective pattern matchers (match_nvfp4_swiglu_patternandmatch_swiglu_pattern) which are disabled by default, so this is safe in practice — but the inconsistency is worth noting. If both fusors are intended to be no-ops when their matchers haven't run, consider aligning their enabled defaults for clarity.Also applies to: 167-169
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/config/default.yaml` around lines 142 - 143, The two fusor entries are inconsistent: fuse_nvfp4_swiglu lacks an explicit enabled flag while fuse_swiglu is explicitly disabled; both depend on match_nvfp4_swiglu_pattern and match_swiglu_pattern which are disabled by default. Make the defaults explicit and consistent—e.g., add enabled: false to fuse_nvfp4_swiglu (or alternatively set enabled: true on fuse_swiglu if you intend both to be active) so fuse_nvfp4_swiglu and fuse_swiglu share the same default behavior and avoid confusion.tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py (3)
30-31: Usecollections.dequefor BFS queues instead oflist.pop(0).
list.pop(0)is O(n) because it shifts every remaining element. In_find_merge_add(line 31) and_execute_shared_expert_in_aux_stream(line 114) the queues are used for BFS. Usingcollections.dequewithpopleft()gives O(1) dequeue.Proposed fix for _find_merge_add
+from collections import deque + def _find_merge_add(moe_node: Node) -> Optional[Node]: visited: Set[Node] = set() - queue = list(moe_node.users.keys()) + queue = deque(moe_node.users.keys()) while queue: - n = queue.pop(0) + n = queue.popleft() if n in visited: continue visited.add(n) if is_op(n, torch.ops.aten.add.Tensor): return n queue.extend(n.users.keys()) return NoneApply the same change to the BFS in
_execute_shared_expert_in_aux_stream(line 111–131).Also applies to: 113-114
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around lines 30 - 31, The BFS loops in _find_merge_add and _execute_shared_expert_in_aux_stream use a Python list as "queue" and call pop(0), which is O(n); change the queue to a collections.deque and replace pop(0) with popleft() for O(1) dequeues. Import deque from collections (or use collections.deque), initialize queue = deque(initial_items) where the list was created, and update all usages in those functions (including any while queue: n = queue.pop(0)) to n = queue.popleft(); ensure any code that extends the queue uses append/extend or appendleft as appropriate. Make no other logic changes.
95-102: Potential misclassification when bothaddinputs descend from the MoE node.If both
arg0andarg1are inmoe_node's ancestor set (e.g., they share an upstream path through the MoE node), the routed/shared classification would be wrong. The current check only testsarg0_ancestors; ifmoe_nodeis not inarg0_ancestorsit defaults to assigningarg1as routed, without verifyingarg1actually descends frommoe_node. In practice the warning at line 133 would catch the failure later, but a symmetric check would be more robust.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around lines 95 - 102, The current logic only checks _get_ancestors(arg0) to decide which add input is routed vs shared, which can misclassify when both arg0 and arg1 descend from moe_node; compute arg1_ancestors = _get_ancestors(arg1) and then choose routed_output/shared_output with a symmetric decision: if moe_node in arg0_ancestors or arg0 is moe_node set routed=arg0/shared=arg1; elif moe_node in arg1_ancestors or arg1 is moe_node set routed=arg1/shared=arg0; otherwise keep the existing fallback and warning behavior; update the block around add_node, arg0, arg1, _get_ancestors, routed_output, and shared_output accordingly.
153-168: Usereplace_input_with()to safely handle both args and kwargs when rewiring fork_point.The rewiring at lines 166–168 only substitutes
fork_pointinfirst_shared.args. Since line 146 checksall_input_nodes(which includes both args and kwargs), iffork_pointappears in kwargs, the substitution will be missed, silently leaving a broken dependency. Usenode.replace_input_with(fork_point, begin_aux_node)instead—it's available since PyTorch 1.11+ and already used throughout the codebase in similar transforms (e.g.,sharding.py,gather_last_token_ids.py). Apply the same fix to lines 178–180 foradd_node.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py` around lines 153 - 168, The current rewiring replaces fork_point only inside first_shared.args (and similarly for add_node.args), which misses occurrences in kwargs; replace that manual tuple rebuild with the PyTorch helper node.replace_input_with(fork_point, begin_aux_node) so both args and kwargs are updated safely. Locate where you create begin_aux_node via graph.call_function(begin_aux_stream_passthrough, args=(fork_point,)) and then call first_shared.replace_input_with(fork_point, begin_aux_node); do the same for add_node (replace its fork_point inputs with begin_aux_node) to ensure all inputs/kwargs are patched.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py (1)
5-5: Avoid wildcard imports; import the module instead.
from ... import *violates the project's import convention and makes it unclear which names are being imported. Since this import is used for its side effect (registering custom ops), import the module itself:Proposed fix
-from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import * # noqa +from tensorrt_llm._torch.auto_deploy.custom_ops.linear import swiglu as _swiglu_ops # noqa: F401 — side-effect import registers opsAs per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py` at line 5, Replace the wildcard import with a module import so names are not leaked; change the line that currently does "from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import *" to "from tensorrt_llm._torch.auto_deploy.custom_ops.linear import swiglu" (or import the module under an alias) so the swiglu module is imported for its side effects while avoiding a star import.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py (1)
96-102: Consider usingtorch.testing.assert_closefor consistent assertion style.
_assert_numerical_correctnessand_assert_cuda_graph_correctnessusetorch.allclosewith a manual error message.torch.testing.assert_closeprovides richer diagnostics (element-wise diff summaries) and is used elsewhere in the test codebase (e.g.,test_fuse_swiglu.py). A minor consistency nit.Also applies to: 105-126
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py` around lines 96 - 102, Replace the custom torch.allclose assertion in _assert_numerical_correctness (and similarly in _assert_cuda_graph_correctness) with torch.testing.assert_close to get richer, consistent diagnostics; call torch.testing.assert_close(out, ref, atol=atol) (or include rtol if needed) instead of the current assert torch.allclose(...) and remove the manual max-diff formatting, keeping the same inputs gm, model, and test_x so the outputs being compared remain identical.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py`:
- Around line 208-236: The store method currently calls .numpy() on tensors
which will fail for CUDA tensors; update InputBuffer.store to ensure the source
tensor is on CPU before calling .numpy() (either assert not .is_cuda or call
data = data.detach().to("cpu") ), then perform dtype conversion and
.contiguous() on data before taking .numpy(); ensure the host_view slice is used
as the dst and keep the existing numpy.copyto(dst, src) flow so memcpy into
pinned memory still works (refer to the store method, variables data, dtype,
host_view, dst, src).
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py`:
- Line 196: Rename the global instance _GlobalTrtllmPlanner to follow the upper
snake_case G_ convention: create G_TRTLLM_PLANNER = _TrtllmPlanner() and replace
every usage of _GlobalTrtllmPlanner (e.g., calls like
_GlobalTrtllmPlanner.reset(...), attribute accesses, imports or references) with
G_TRTLLM_PLANNER so all code references (including in functions/methods that
call planner methods on _TrtllmPlanner) are updated to the new symbol.
- Around line 275-415: The function misses using the incoming scale arg: compute
q_scaling = (1.0 / math.sqrt(head_dim)) if scale is None else scale (import math
if not already) inside the trtllm attention function (the one that builds
qkv_fused and calls thop.attention) and replace the hardcoded 1.0 argument for
q_scaling in the thop.attention call with this q_scaling variable so attention
uses the provided or default scaling; reference symbols: scale (parameter),
head_dim (inferred above), q_scaling (new local variable), and the
thop.attention call site.
In `@tensorrt_llm/_torch/auto_deploy/export/export.py`:
- Line 276: Replace the ambiguous EN DASH in the comment that reads "Collect
indices of List[Tensor] arguments from the op schema – these" with a plain
hyphen so it becomes "Collect indices of List[Tensor] arguments from the op
schema - these" to satisfy Ruff RUF003; locate that comment in export.py (the
comment text above) and update the punctuation only.
- Around line 278-279: The current access of op schema (variables op,
node.target, schema) can raise AttributeError/StopIteration if op lacks _schema
and _schemas; update the logic in the export logic around where op = node.target
to check attributes defensively: first check hasattr(op, "_schemas") and use
next(iter(op._schemas.values())), then elif hasattr(op, "_schema") use
op._schema, otherwise log a clear warning via ad_logger (including node.target)
and skip/continue processing the node (mirroring the pattern used in
node_utils.py) so failures are explicit and non-fatal.
In `@tensorrt_llm/_torch/auto_deploy/shim/interface.py`:
- Around line 541-545: The call is passing the bound method
self._kv_cache_manager.blocks_in_primary_pool instead of its integer result;
change the argument to pass the method return value by invoking it (i.e., call
blocks_in_primary_pool()) when calling self.info.update_cache_information so
update_cache_information receives an int; locate this in the block using
self.info.update_cache_information and
self._kv_cache_manager.blocks_in_primary_pool and replace the method reference
with its invocation.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py`:
- Around line 311-313: Move the mid-file import of torch_nvfp4_swiglu_mlp out of
the bottom of the module and place it with the other top-level imports; remove
the unnecessary "# noqa: E402". Change the import style to "from
...custom_ops.linear import swiglu" and update any usages (e.g.,
torch_nvfp4_swiglu_mlp) to reference swiglu.torch_nvfp4_swiglu_mlp (and
similarly swiglu.torch_swiglu_mlp) so references inside the FuseSwiGLU class and
elsewhere use the new module-qualified name.
- Around line 572-587: The fused node creation silently assumes gate and up
projections share input scale and alpha; before creating fused_nvfp4_swiglu_mlp,
validate that gate_input_scale_node and gate_alpha_node equal the corresponding
up projection values (e.g., up_input_scale_node and up_alpha_node) and raise a
clear error or at least log a debug/warning if they differ; place this check
immediately before the graph.inserting_after(node) block so the mismatch is
caught early and include the node names in the message for easier debugging.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py`:
- Around line 131-140: The insertion of new_cast before list(cast_node.users)[0]
can create a topological violation because new_cast depends on add_out; instead
insert new_cast after add_out using graph.inserting_after(add_out) to guarantee
ordering, and avoid building a full list of users by using
next(iter(cast_node.users), None) if you ever need a user—replace the
graph.inserting_before(list(cast_node.users)[0]) usage with
graph.inserting_after(add_out) and remove the list(...) allocation; keep
cast_node.replace_all_uses_with(new_cast), graph.erase_node(cast_node), and
erased.add(id(cast_node)) as-is.
- Around line 89-92: The loop over matches currently guards only add_node and
norm_node against the erased set, but not cast_node, which can lead to reading a
consumed cast_node (e.g., when later accessing cast_node.users); update the
stale-match guard in the for-loop that iterates "for add_node, cast_node,
norm_node in matches" to also skip when id(cast_node) is in erased, i.e.,
include id(cast_node) in the conditional that continues, so any match
referencing a previously erased cast_node is safely skipped before further
inspection or modification.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`:
- Around line 1-12: The new module multi_stream_attn.py is missing the required
NVIDIA copyright header and Apache-2.0 license notice; add the standard NVIDIA
copyright header block (with the year of latest meaningful modification) and the
Apache License, Version 2.0 boilerplate at the top of the file before the module
docstring so the file complies with project licensing guidelines.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`:
- Around line 1-18: Add the required NVIDIA Apache-2.0 copyright header block at
the very top of the file before the module docstring in multi_stream_moe.py;
ensure you insert the SPDX lines and the Apache License 2.0 boilerplate exactly
as in project guidelines so the file starting with the triple-quoted module
docstring now sits after the header, leaving all imports and symbols (e.g.,
GraphModule, ModelFactory, CachedSequenceInterface,
begin_aux_stream_passthrough, cuda_stream_manager, BaseTransform) unchanged.
In `@tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py`:
- Around line 1-21: This file is missing the required NVIDIA copyright/header
with the Apache License 2.0; add the standard NVIDIA copyright header block at
the very top of the module (above the module docstring) so the file containing
CudaStreamManager, the custom ops record_event/wait_event, and
_make_aux_stream_impl includes the required license and attribution text per
project guidelines.
- Around line 72-85: The add_device method has a TOCTOU race on self.devices;
initialize a threading.Lock (self._lock) in the class __init__ and then wrap the
critical section in add_device with that lock so the check-and-mutate is atomic:
acquire self._lock, check "if device not in self.devices", append to
self.devices and set up events/streams in the guarded block (or set a boolean in
the lock and perform non-thread-sensitive CUDA operations outside if needed),
and log the warning in the else branch — ensure you reference and use self._lock
consistently in add_device and __init__ to prevent concurrent duplicate
creation.
- Around line 52-70: The __init__ guard using hasattr(self, "streams") is
ineffective because streams is a class attribute, making the guard always true
and leaving instance attributes like self._lock and device registration
uninitialized; fix by removing the guard and turning the mutable class-level
attributes into instance attributes inside CudaStreamManager.__init__ (e.g., set
self.devices = [], self.events = {}, self.streams = {}, self._caller_streams =
{}), then initialize self._lock = RLock() and call
self.add_device(torch.cuda.current_device()) so the instance is properly
initialized on construction; keep AUX_STREAM_NAME and MAIN_STREAM_NAME as class
constants and retain the _Singleton metaclass behavior.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`:
- Around line 37-39: The helper _reset_trtllm_planner currently calls
_GlobalTrtllmPlanner.__init__() which bypasses guard logic and is fragile;
change it to either call the existing reset API with sensible defaults (e.g.
_GlobalTrtllmPlanner.reset(torch.device("cpu"), 32, 2048)) or add a dedicated
parameterless clear() method on _GlobalTrtllmPlanner that zeroes internal state
and call that from _reset_trtllm_planner instead; update references to use
_GlobalTrtllmPlanner.reset or the new clear() and remove the direct __init__
invocation.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py`:
- Line 8: The import line "from tensorrt_llm._torch.auto_deploy.custom_ops.rope
import triton_rope" has an unnecessary "# noqa: F401" that triggers Ruff RUF100;
remove the trailing " # noqa: F401" from that import to clear the false
suppression, or if you must preserve the suppression for another tool keep it as
"## noqa: F401" (double-hash) or declare F401 as an external rule in
pyproject.toml under lint.external = ["F401"] so Ruff will not remove it.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py`:
- Around line 14-15: Replace the top-level import that currently reads "import
tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401" with a
namespace-preserving import: "from tensorrt_llm._torch.auto_deploy import
custom_ops" and remove the unused "# noqa: F401"; keep the existing "from
tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm" line as-is so
the module registration side-effect (custom_ops) is preserved while following
the namespace import guideline.
- Around line 20-21: Rename the module-level globals _skip_reason and
_skip_condition to follow the G_ upper snake_case convention (e.g.,
G_SKIP_REASON and G_SKIP_CONDITION) and update all references to them in this
test file (test_nvfp4_swiglu.py) accordingly; ensure the values remain the same
(G_SKIP_REASON = "Requires NVFP4 (Blackwell+) and TRT-LLM ops" and
G_SKIP_CONDITION = not (fp4_compatible() and trtllm_ops_available())) and adjust
any usages in decorators or test logic that referenced
_skip_reason/_skip_condition.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`:
- Around line 359-362: Remove the unused "# noqa: E402" and change the import
style to keep the module namespace: replace the current "from
tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import
(Glm4MoeLiteConfig, Glm4MoeLiteForCausalLM)" with "from
tensorrt_llm._torch.auto_deploy.models.custom import modeling_glm4_moe_lite" and
then update all usages in this test (references to Glm4MoeLiteConfig and
Glm4MoeLiteForCausalLM) to use qualified names
modeling_glm4_moe_lite.Glm4MoeLiteConfig and
modeling_glm4_moe_lite.Glm4MoeLiteForCausalLM.
---
Outside diff comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py`:
- Line 1: Update the copyright header year from 2025 to 2026 in the file
tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py by
editing the top-of-file copyright comment to read 2026 so the header matches the
current PR modifications (update the existing copyright comment line).
In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py`:
- Line 1: Update the file header copyright year from 2025 to 2026 by editing the
top comment line that currently reads "Copyright (c) 2025 NVIDIA CORPORATION &
AFFILIATES. All rights reserved." to use 2026 instead so the header reflects the
current modification year.
- Around line 347-368: The get_default_kwargs in TestNemotronSuperV3 currently
enables the multi_stream_moe transform unconditionally; update it to either
conditionally set "enabled": (world_size == 1) based on the test's world_size
parameter (matching the behavior in TestNemotronMOE) or add a clear comment in
get_default_kwargs explaining why multi_stream_moe can safely be enabled for
multi-GPU runs in TestNemotronSuperV3 (referencing TestNemotronMOE,
multi_stream_moe, get_default_kwargs, test_fp8, and world_size so reviewers can
verify the rationale).
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py`:
- Line 1: Add the required NVIDIA copyright header (Apache License 2.0) to the
top of this source file including the year of latest meaningful modification;
update the year when you modify the file in the future. Place the header above
all imports (before the existing "from typing import Optional" line) so it
appears as the first lines of the file, and ensure it matches the project's
standard NVIDIA Apache-2.0 boilerplate.
---
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/config/default.yaml`:
- Around line 142-143: The two fusor entries are inconsistent: fuse_nvfp4_swiglu
lacks an explicit enabled flag while fuse_swiglu is explicitly disabled; both
depend on match_nvfp4_swiglu_pattern and match_swiglu_pattern which are disabled
by default. Make the defaults explicit and consistent—e.g., add enabled: false
to fuse_nvfp4_swiglu (or alternatively set enabled: true on fuse_swiglu if you
intend both to be active) so fuse_nvfp4_swiglu and fuse_swiglu share the same
default behavior and avoid confusion.
In `@tensorrt_llm/_torch/auto_deploy/export/export.py`:
- Line 280: The tuple _tensor_list_types is being recreated inside the loop;
move its definition out of the loop so it is created once (hoist it to the
module scope near the existing moe_ops declaration) and then reference
_tensor_list_types inside the for node in list(graph.nodes) loop to avoid
repeated allocations.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py`:
- Around line 14-15: Replace typing generics with Python 3.10+ built-ins: remove
the imports of List, Optional, Tuple and use list, tuple, and X | None in
annotations instead. Update the module-level import statement (remove
List/Optional/Tuple) and change any function/type hints that use List[T],
Tuple[...], or Optional[X] to list[T], tuple[...], and X | None respectively;
also update the return type annotation on the method/function named _apply to
use the built-in generics/union syntax.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`:
- Around line 14-31: Replace all direct symbol imports with module-level imports
and update usages to qualified names: import torch.fx as fx and replace
GraphModule/Node with fx.GraphModule and fx.Node; import ...utils._graph as
_graph and use _graph.create_derived_custom_op; import
...utils.multi_stream_utils as multi_stream_utils and use
multi_stream_utils._make_aux_stream_impl,
multi_stream_utils.cuda_stream_manager, and
multi_stream_utils.record_event_passthrough; import ...utils.node_utils as
node_utils and use node_utils.is_op; import ...shim.interface as shim_interface
and use shim_interface.CachedSequenceInterface; import ...models.factory as
factory and use factory.ModelFactory; import ..interface as interface and use
interface.BaseTransform, interface.SharedConfig, interface.TransformInfo,
interface.TransformRegistry. Update all references in this file accordingly so
no `from ... import Symbol` remains.
- Around line 126-185: The loop in _execute_kv_proj_in_aux_stream builds
node_order once from graph.nodes but then mutates the graph (inserting rec_node,
replacing/erasing kv_linear), which can make the precomputed node_order stale
when locating earliest_q for later fork_point iterations; fix by either
recomputing node_order (i.e., node_order = {n: i for i, n in
enumerate(graph.nodes)}) immediately before computing q_linears/earliest_q
inside the for fork_point, kv_linear in pairs loop, or explicitly document and
assert the safety invariant (e.g., add a clear comment and an assert in
_execute_kv_proj_in_aux_stream that pairs are independent: no two fork_point
entries share Q-linear users) so future maintainers know why the single
precomputed node_order is safe.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`:
- Around line 30-31: The BFS loops in _find_merge_add and
_execute_shared_expert_in_aux_stream use a Python list as "queue" and call
pop(0), which is O(n); change the queue to a collections.deque and replace
pop(0) with popleft() for O(1) dequeues. Import deque from collections (or use
collections.deque), initialize queue = deque(initial_items) where the list was
created, and update all usages in those functions (including any while queue: n
= queue.pop(0)) to n = queue.popleft(); ensure any code that extends the queue
uses append/extend or appendleft as appropriate. Make no other logic changes.
- Around line 95-102: The current logic only checks _get_ancestors(arg0) to
decide which add input is routed vs shared, which can misclassify when both arg0
and arg1 descend from moe_node; compute arg1_ancestors = _get_ancestors(arg1)
and then choose routed_output/shared_output with a symmetric decision: if
moe_node in arg0_ancestors or arg0 is moe_node set routed=arg0/shared=arg1; elif
moe_node in arg1_ancestors or arg1 is moe_node set routed=arg1/shared=arg0;
otherwise keep the existing fallback and warning behavior; update the block
around add_node, arg0, arg1, _get_ancestors, routed_output, and shared_output
accordingly.
- Around line 153-168: The current rewiring replaces fork_point only inside
first_shared.args (and similarly for add_node.args), which misses occurrences in
kwargs; replace that manual tuple rebuild with the PyTorch helper
node.replace_input_with(fork_point, begin_aux_node) so both args and kwargs are
updated safely. Locate where you create begin_aux_node via
graph.call_function(begin_aux_stream_passthrough, args=(fork_point,)) and then
call first_shared.replace_input_with(fork_point, begin_aux_node); do the same
for add_node (replace its fork_point inputs with begin_aux_node) to ensure all
inputs/kwargs are patched.
In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py`:
- Around line 23-33: Rename the module-level mutable globals to follow the
project's naming convention: change _derived_op_libs to G_DERIVED_OP_LIBS and
_derived_op_registry to G_DERIVED_OP_REGISTRY, update all references/usages
throughout the module (including any initialization and lookups) to the new
names, and preserve their types and initial values (Dict[str,
torch.library.Library] and Dict[str, Callable]) so behavior is unchanged; ensure
any import or external references in this file are also updated to the new
identifiers.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`:
- Around line 27-30: Replace the symbol-level import of _GlobalTrtllmPlanner and
prepare_trtllm_metadata_host with a module import so code follows the "from
package.subpackage import module" guideline: import the module
tensorrt_llm._torch.auto_deploy.custom_ops.attention.trtllm_attention as
trtllm_attention, then update all usages to call
trtllm_attention._GlobalTrtllmPlanner and
trtllm_attention.prepare_trtllm_metadata_host so references resolve via the
module namespace.
- Around line 462-581: In test_trtllm_attention_with_paged_kvcache, there are
two for-loops using unused loop variable s (for i, s in enumerate(seq_lengths))
around the prefill page assignment and the decode page-update sections; rename s
to _ in both loops so it signals an intentionally unused variable and fixes the
lint warnings (update the two occurrences of "for i, s in
enumerate(seq_lengths)" to "for i, _ in enumerate(seq_lengths)" while leaving
surrounding logic, page_assignments, pages_per_seq_list, and cache_locs handling
unchanged).
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_moe.py`:
- Around line 96-102: Replace the custom torch.allclose assertion in
_assert_numerical_correctness (and similarly in _assert_cuda_graph_correctness)
with torch.testing.assert_close to get richer, consistent diagnostics; call
torch.testing.assert_close(out, ref, atol=atol) (or include rtol if needed)
instead of the current assert torch.allclose(...) and remove the manual max-diff
formatting, keeping the same inputs gm, model, and test_x so the outputs being
compared remain identical.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py`:
- Line 5: Replace the wildcard import with a module import so names are not
leaked; change the line that currently does "from
tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import *" to "from
tensorrt_llm._torch.auto_deploy.custom_ops.linear import swiglu" (or import the
module under an alias) so the swiglu module is imported for its side effects
while avoiding a star import.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py`:
- Around line 174-213: Both tests use a loose assertion "_count_fused_ops(gm_t)
>= 1"; change this to an exact check "== 1" in test_fuse_add_cast_rms_norm and
test_fuse_add_rms_norm_no_cast to match later tests and reflect that each model
produces exactly one fused op; update the assertion strings if desired (e.g.,
keep "fused op not found in graph") so the checks in those functions
(_count_fused_ops and _count_rms_norm_ops called after _apply_transform on gm)
are consistent and stricter.
- Around line 11-12: The tests import classes directly; change to module-level
imports to follow project style: replace direct class imports of TransformConfig
and FuseAddRMSNorm with module imports (import the transform.interface module
and the transform.library.fused_add_rms_norm module) and update usages to
reference interface.TransformConfig and fused_add_rms_norm.FuseAddRMSNorm so all
references to TransformConfig and FuseAddRMSNorm use the module-qualified names.
- Around line 138-140: The helper function _count_add_ops is defined but unused;
either remove the function or use it to assert fusion behavior—update the test
(e.g., in test_fused_add_rms_norm) to call _count_add_ops(gm) before and after
running the transformation pass and add an assertion that the number of
torch.ops.aten.add.Tensor nodes decreases (or becomes zero) after fusion, or
simply delete the _count_add_ops definition if no such assertion is desired;
reference the function name _count_add_ops and the graph module variable gm when
making the change.
- Around line 6-8: Remove the unnecessary blanket "# noqa" on the import line
that brings in flashinfer_fused_add_rms_norm; the import is used (e.g., by
_count_fused_ops) so simply delete the "# noqa" directive from the import
statement importing flashinfer_fused_add_rms_norm to satisfy Ruff RUF100.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`:
- Around line 331-335: Remove the duplicate nested function _count_moe_experts
inside the test (the local definition that iterates gm.graph.nodes and checks
for "torch_moe" nodes), and replace its usages with the module-level
_count_moe_experts_in_graph; i.e., delete the nested def _count_moe_experts(gm)
and update any calls in the test to call _count_moe_experts_in_graph(gm) instead
so the single shared implementation is used.
- Around line 432-453: Add numerical output equality checks like in the
SimpleMoEForExport test: run a fixed input through gm_full and gm_reduced
(before and after calling gm_reduced.load_state_dict(model.state_dict(),
strict=False)), capture outputs (y_full, y_reduced, y_loaded) and assert they
match using the existing all_close or torch.allclose helper; keep the existing
structural/state-dict assertions but append asserts all_close(y_full, y_reduced)
and all_close(y_loaded, y_full) to ensure the GLM4 export path produces
numerically equivalent outputs. Use the existing symbols gm_full, gm_reduced,
model.state_dict(), and the helper all_close/torch.allclose to locate where to
add these checks.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py`:
- Around line 1-7: The test imports create_derived_custom_op directly from the
_graph module; change the import to bring in the module instead and reference
the function via the module to follow the coding guideline: replace the direct
import with importing tensorrt_llm._torch.auto_deploy.utils import _graph and
update all usages of create_derived_custom_op in the test to call
_graph.create_derived_custom_op (ensure any other references in this test file
are updated accordingly).
| def store( | ||
| self, | ||
| name: str, | ||
| data: List[Number], | ||
| data: torch.Tensor, | ||
| fill_value: Optional[Number] = None, | ||
| ) -> int: | ||
| """Store data into the host buffer. | ||
| """Store a tensor into the pinned host buffer. | ||
|
|
||
| Args: | ||
| name: Name of the tensor to store to. | ||
| data: List of values to store. | ||
| fill_value: Optional value to fill the entire tensor with before storing. | ||
| If None, only the provided data is written. | ||
| data: 1-D torch.Tensor to store. | ||
| fill_value: Optional value to fill the entire buffer with before storing. | ||
|
|
||
| Returns: | ||
| Number of elements stored. | ||
| """ | ||
| numel, dtype = self._tensor_specs[name] | ||
| host_view = self.get_host_view(name) | ||
|
|
||
| # Fill with default value if specified | ||
| if fill_value is not None: | ||
| host_view.fill_(fill_value) | ||
|
|
||
| # Convert list to tensor and copy to host buffer | ||
| length = len(data) | ||
| length = data.numel() | ||
| assert length <= numel, f"Data too large for buffer '{name}': {length} > {numel}" | ||
|
|
||
| temp_tensor = torch.tensor(data, dtype=dtype) | ||
| host_view[:length].copy_(temp_tensor) | ||
| # Use numpy for the memcpy into pinned memory — avoids torch dispatcher overhead | ||
| dst = host_view[:length].numpy() | ||
| src = (data if data.dtype == dtype else data.to(dtype)).numpy() | ||
| np.copyto(dst, src) | ||
|
|
There was a problem hiding this comment.
Guard against CUDA tensors in InputBuffer.store (numpy() requires CPU).
store() now accepts torch.Tensor but uses .numpy(), which will throw if the tensor is on CUDA. Add a CPU guard (or assert) to avoid runtime errors.
🛠️ Suggested fix
- src = (data if data.dtype == dtype else data.to(dtype)).numpy()
+ if data.device.type != "cpu":
+ data = data.to("cpu", non_blocking=True)
+ src = (data if data.dtype == dtype else data.to(dtype)).numpy()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py` around
lines 208 - 236, The store method currently calls .numpy() on tensors which will
fail for CUDA tensors; update InputBuffer.store to ensure the source tensor is
on CPU before calling .numpy() (either assert not .is_cuda or call data =
data.detach().to("cpu") ), then perform dtype conversion and .contiguous() on
data before taking .numpy(); ensure the host_view slice is used as the dst and
keep the existing numpy.copyto(dst, src) flow so memcpy into pinned memory still
works (refer to the store method, variables data, dtype, host_view, dst, src).
| return t | ||
|
|
||
|
|
||
| _GlobalTrtllmPlanner = _TrtllmPlanner() |
There was a problem hiding this comment.
Rename the global planner to the G_ upper snake_case convention.
♻️ Suggested change
-_GlobalTrtllmPlanner = _TrtllmPlanner()
+G_TRTLLM_PLANNER = _TrtllmPlanner()Update all references accordingly (e.g., G_TRTLLM_PLANNER.reset(...), etc.).
As per coding guidelines: Use upper snake_case with prefix 'G' for global variables (e.g., G_MY_GLOBAL = ...).
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| _GlobalTrtllmPlanner = _TrtllmPlanner() | |
| G_TRTLLM_PLANNER = _TrtllmPlanner() |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py` at
line 196, Rename the global instance _GlobalTrtllmPlanner to follow the upper
snake_case G_ convention: create G_TRTLLM_PLANNER = _TrtllmPlanner() and replace
every usage of _GlobalTrtllmPlanner (e.g., calls like
_GlobalTrtllmPlanner.reset(...), attribute accesses, imports or references) with
G_TRTLLM_PLANNER so all code references (including in functions/methods that
call planner methods on _TrtllmPlanner) are updated to the new symbol.
| scale: Optional[float], | ||
| sliding_window: Optional[int] = None, | ||
| kv_scale_orig_quant: float = 1.0, | ||
| kv_scale_quant_orig: float = 1.0, | ||
| ) -> torch.Tensor: | ||
| """TRT-LLM attention with paged KV cache for Auto-Deploy. | ||
|
|
||
| Infers num_heads, num_kv_heads, head_dim, and tokens_per_block from tensor shapes. | ||
| All max-size constants (max_num_requests, max_context_length) are read from | ||
| ``max_seq_info_host`` which is set once via ``SequenceInfo.update_cache_information()``. | ||
|
|
||
| Note: ``prepare_trtllm_metadata_host`` is guaranteed to be called before this op, | ||
| so all persistent planner buffers are already initialized. | ||
|
|
||
| Note: layer_idx is always passed as 0 to thop.attention because | ||
| the kv_cache tensor is already a strided view for the correct layer, | ||
| pool_pointers encodes kv_cache.data_ptr() (layer-specific), and | ||
| pool_mapping is all zeros. See module docstring for details. | ||
| """ | ||
| # Infer dimensions from tensor shapes (bsnd layout) | ||
| num_heads = q.shape[2] | ||
| num_kv_heads = k.shape[2] | ||
| head_dim = q.shape[3] | ||
| tokens_per_block = kv_cache.shape[3] # HND: [blocks, 2, heads, tpb, head_dim] | ||
|
|
||
| # Get batch dimensions and model-level constants from host tensors (no device sync) | ||
| num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() | ||
| num_seq = num_prefill + num_decode | ||
| num_tokens = num_prefill_tokens + num_decode | ||
| max_context_length = int(max_seq_info_host[0]) | ||
| max_num_requests = int(max_seq_info_host[3]) | ||
| # Use sliding_window for attention_window_size if provided, else full context length | ||
| attention_window_size = ( | ||
| sliding_window | ||
| if isinstance(sliding_window, int) and sliding_window > 0 | ||
| else max_context_length | ||
| ) | ||
|
|
||
| # Get per-layer pool pointer tensor (stable address for CUDA graph replay) | ||
| host_kv_cache_pool_pointers = _GlobalTrtllmPlanner.get_pool_pointers_for_layer(kv_cache) | ||
|
|
||
| # FP8 KV cache: lazily create scale tensors from float constants on first use | ||
| if kv_cache.dtype == torch.float8_e4m3fn: | ||
| if _GlobalTrtllmPlanner.kv_scale_orig_quant is None: | ||
| _GlobalTrtllmPlanner.kv_scale_orig_quant = torch.tensor( | ||
| [kv_scale_orig_quant], dtype=torch.float32, device=q.device | ||
| ) | ||
| _GlobalTrtllmPlanner.kv_scale_quant_orig = torch.tensor( | ||
| [kv_scale_quant_orig], dtype=torch.float32, device=q.device | ||
| ) | ||
| quant_mode = int(QuantMode.FP8_KV_CACHE) | ||
| else: | ||
| quant_mode = 0 | ||
|
|
||
| # Reshape Q, K, V to [num_tokens, num_heads * head_dim] and fuse | ||
| # Input is always [bs, 1] (generate-only) or [1, total_seq_len] (prefill/mixed), | ||
| # so b * s == num_tokens always holds. | ||
| q_shape_og = q.shape | ||
| q_flat = q.reshape(num_tokens, num_heads * head_dim) | ||
| k_flat = k.reshape(num_tokens, num_kv_heads * head_dim) | ||
| v_flat = v.reshape(num_tokens, num_kv_heads * head_dim) | ||
| qkv_fused = torch.cat([q_flat, k_flat, v_flat], dim=-1).contiguous() | ||
|
|
||
| # Prepare output | ||
| output = torch.empty(num_tokens, num_heads * head_dim, dtype=q.dtype, device=q.device) | ||
|
|
||
| # Map SequenceInfo fields to thop.attention args | ||
| sequence_length = seq_len_with_cache[:num_seq] # device | ||
| context_lengths = seq_len[:num_seq] # device | ||
| host_past_kv_lengths = _GlobalTrtllmPlanner.host_past_kv_lengths[:num_seq] # host (pinned) | ||
| host_context_lengths = _GlobalTrtllmPlanner.host_context_lengths[:num_seq] # host (pinned) | ||
|
|
||
| # thop-specific metadata from _GlobalTrtllmPlanner | ||
| host_request_types = _GlobalTrtllmPlanner.host_request_types[:num_seq] | ||
| host_total_kv_lens = _GlobalTrtllmPlanner.host_total_kv_lens | ||
|
|
||
| # Block offsets from host_prepare | ||
| kv_cache_block_offsets = _GlobalTrtllmPlanner.block_offsets | ||
|
|
||
| # Pool mapping (shared, always zeros since layer offset is in pool_pointers) | ||
| host_kv_cache_pool_mapping = _GlobalTrtllmPlanner.host_pool_mapping | ||
|
|
||
| # Pack parameters for thop.attention | ||
| rotary_embedding_scales = [1.0, 1.0, 1.0] | ||
| rotary_embedding_max_position_info = [max_context_length, max_context_length] | ||
| spec_decoding_bool_params = [False, False, False] | ||
| spec_decoding_tensor_params = [None, None, None] | ||
|
|
||
| sm_version = get_sm_version() | ||
| if sm_version >= 89: # Ada/Hopper | ||
| spec_decoding_tensor_params.extend([None, None, None]) | ||
|
|
||
| mla_tensor_params = [None, None] | ||
|
|
||
| thop.attention( | ||
| qkv_fused, # q (actually fused QKV) | ||
| None, # k (None when using fused QKV) | ||
| None, # v (None when using fused QKV) | ||
| output, # output | ||
| None, # output_sf (NVFP4) | ||
| _GlobalTrtllmPlanner.workspace, # workspace (module-level, like flashinfer) | ||
| sequence_length, # sequence_length | ||
| host_past_kv_lengths, # host_past_key_value_lengths | ||
| host_total_kv_lens, # host_total_kv_lens | ||
| context_lengths, # context_lengths | ||
| host_context_lengths, # host_context_lengths | ||
| host_request_types, # host_request_types | ||
| kv_cache_block_offsets, # kv_cache_block_offsets | ||
| host_kv_cache_pool_pointers, # host_kv_cache_pool_pointers | ||
| host_kv_cache_pool_mapping, # host_kv_cache_pool_mapping | ||
| None, # cache_indirection (beam search) | ||
| _GlobalTrtllmPlanner.kv_scale_orig_quant, # kv_scale_orig_quant | ||
| _GlobalTrtllmPlanner.kv_scale_quant_orig, # kv_scale_quant_orig | ||
| None, # out_scale | ||
| None, # rotary_inv_freq | ||
| None, # rotary_cos_sin | ||
| None, # latent_cache (MLA) | ||
| None, # q_pe (MLA) | ||
| None, # block_ids_per_seq | ||
| None, # attention_sinks | ||
| True, # is_fused_qkv | ||
| True, # update_kv_cache | ||
| 1, # predicted_tokens_per_seq | ||
| 0, # layer_idx (always 0; pool_pointers already encodes the layer offset) | ||
| num_heads, # num_heads | ||
| num_kv_heads, # num_kv_heads | ||
| head_dim, # head_size | ||
| tokens_per_block, # tokens_per_block | ||
| max_num_requests, # max_num_requests | ||
| max_context_length, # max_context_length | ||
| attention_window_size, # attention_window_size | ||
| 0, # sink_token_length | ||
| 1, # beam_width | ||
| int(AttentionMaskType.causal), # mask_type | ||
| quant_mode, # quant_mode | ||
| 1.0, # q_scaling | ||
| 0, # position_embedding_type | ||
| 0, # rotary_embedding_dim | ||
| 10000.0, # rotary_embedding_base | ||
| 0, # rotary_embedding_scale_type | ||
| rotary_embedding_scales, # rotary_embedding_scales |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -name "trtllm_attention.py" -type fRepository: NVIDIA/TensorRT-LLM
Length of output: 137
🏁 Script executed:
wc -l tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 139
🏁 Script executed:
sed -n '275,415p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | head -150Repository: NVIDIA/TensorRT-LLM
Length of output: 6504
🏁 Script executed:
sed -n '260,280p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 760
🏁 Script executed:
sed -n '400,430p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1353
🏁 Script executed:
rg "scale.*thop\.attention|thop\.attention.*scale" -A 5 -B 5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
rg "def.*attention.*scale" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
grep -n "scale" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1327
🏁 Script executed:
rg "q_scaling" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 612
🏁 Script executed:
rg "def.*attention" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -A 15 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 10853
🏁 Script executed:
fd "attention" tensorrt_llm/_torch/auto_deploy/custom_ops/ --type fRepository: NVIDIA/TensorRT-LLM
Length of output: 661
🏁 Script executed:
sed -n '275,450p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | grep -i "scale"Repository: NVIDIA/TensorRT-LLM
Length of output: 1105
🏁 Script executed:
rg "scale.*\*|scale.*\+" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | grep -A 50 "def _decode_attention"Repository: NVIDIA/TensorRT-LLM
Length of output: 1647
🏁 Script executed:
sed -n '1,50p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 1666
🏁 Script executed:
rg "scale \*|scale\)" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py -B 3 -A 3 | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
grep -n "q_scaling\|scale" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 858
🏁 Script executed:
# Check how torch_attention (the source op) uses scale
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 3996
🏁 Script executed:
# Check how q_scaling is used or documented in thop
rg "q_scaling" tensorrt_llm/ -B 5 -A 5 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 8132
🏁 Script executed:
# Check if there are other calls to thop.attention to see how q_scaling is typically set
rg "thop\.attention" tensorrt_llm/ -B 10 -A 10 | grep -A 5 -B 5 "q_scaling" | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check the function documentation or any comments about scale parameter
sed -n '260,290p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1313
🏁 Script executed:
# Check how scale is computed when None in similar functions
rg "scale.*sqrt|sqrt.*scale" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 1615
🏁 Script executed:
# Look for any other uses of q_scaling to understand its semantics
rg "q_scaling.*=" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if there's documentation about what q_scaling expects
rg "thop" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 389
🏁 Script executed:
# Verify the complete usage by checking lines around where the scale is passed to triton attention
sed -n '245,260p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 660
🏁 Script executed:
# Let's verify head_dim is available in the function scope
sed -n '290,310p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1090
🏁 Script executed:
# Check if scale parameter is documented anywhere in function docstring
sed -n '275,302p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1408
🏁 Script executed:
# Check if there are any callers of trtllm_mha_with_cache to understand how scale is passed
rg "trtllm_mha_with_cache" tensorrt_llm/ -B 3 -A 3 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 1523
🏁 Script executed:
# Double-check: verify that head_dim is definitely available where we need it
sed -n '275,360p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | grep -n "head_dim"Repository: NVIDIA/TensorRT-LLM
Length of output: 618
🏁 Script executed:
# Final verification: search for any use of 'scale' variable after it's received in the function
sed -n '275,450p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py | grep -n "scale" | grep -v "kv_scale\|rotary_embedding_scale\|out_scale\|softmax"Repository: NVIDIA/TensorRT-LLM
Length of output: 249
🏁 Script executed:
# Verify the exact line where q_scaling is hardcoded
sed -n '405,420p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 756
Wire scale into thop.attention to apply attention scaling correctly.
The scale parameter is accepted but never used in the function body, causing q_scaling to be hardcoded to 1.0 in the thop.attention call. This produces silent numerical drift when a non-default scale is provided. The pattern used in other attention implementations (triton_attention.py, torch_backend_attention.py) shows the correct approach: compute q_scaling = 1.0 / math.sqrt(head_dim) if scale is None else scale.
🛠️ Suggested fix
attention_window_size = (
sliding_window
if isinstance(sliding_window, int) and sliding_window > 0
else max_context_length
)
+ q_scaling = scale if scale is not None else head_dim ** -0.5
# Get per-layer pool pointer tensor (stable address for CUDA graph replay)
host_kv_cache_pool_pointers = _GlobalTrtllmPlanner.get_pool_pointers_for_layer(kv_cache)Then update the thop.attention call:
- 1.0, # q_scaling
+ q_scaling, # q_scaling🧰 Tools
🪛 Ruff (0.15.1)
[warning] 275-275: Unused function argument: scale
(ARG001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py`
around lines 275 - 415, The function misses using the incoming scale arg:
compute q_scaling = (1.0 / math.sqrt(head_dim)) if scale is None else scale
(import math if not already) inside the trtllm attention function (the one that
builds qkv_fused and calls thop.attention) and replace the hardcoded 1.0
argument for q_scaling in the thop.attention call with this q_scaling variable
so attention uses the provided or default scaling; reference symbols: scale
(parameter), head_dim (inferred above), q_scaling (new local variable), and the
thop.attention call site.
| continue | ||
|
|
||
| # Collect indices of list-of-node arguments (expert weight/scale lists) | ||
| # Collect indices of List[Tensor] arguments from the op schema – these |
There was a problem hiding this comment.
Ambiguous EN DASH (–) in comment — use a plain hyphen (-).
Ruff RUF003 flags this. Replace – (U+2013) with - (U+002D) to avoid lint warnings.
- # Collect indices of List[Tensor] arguments from the op schema – these
+ # Collect indices of List[Tensor] arguments from the op schema - these📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Collect indices of List[Tensor] arguments from the op schema – these | |
| # Collect indices of List[Tensor] arguments from the op schema - these |
🧰 Tools
🪛 Ruff (0.15.1)
[warning] 276-276: Comment contains ambiguous – (EN DASH). Did you mean - (HYPHEN-MINUS)?
(RUF003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/export/export.py` at line 276, Replace the
ambiguous EN DASH in the comment that reads "Collect indices of List[Tensor]
arguments from the op schema – these" with a plain hyphen so it becomes "Collect
indices of List[Tensor] arguments from the op schema - these" to satisfy Ruff
RUF003; locate that comment in export.py (the comment text above) and update the
punctuation only.
| op = node.target | ||
| schema = op._schema if hasattr(op, "_schema") else next(iter(op._schemas.values())) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's read the file and see the actual code around lines 278-279
cat -n tensorrt_llm/_torch/auto_deploy/export/export.py | sed -n '270,290p'Repository: NVIDIA/TensorRT-LLM
Length of output: 997
🏁 Script executed:
# Let's see a broader context around the problematic line to understand exception handling
cat -n tensorrt_llm/_torch/auto_deploy/export/export.py | sed -n '260,300p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1829
🏁 Script executed:
# Search for where this function is defined and what wraps it
rg -B 20 "op = node.target" tensorrt_llm/_torch/auto_deploy/export/export.py | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 855
🏁 Script executed:
# Check if there's exception handling around this loop
cat -n tensorrt_llm/_torch/auto_deploy/export/export.py | sed -n '240,350p'Repository: NVIDIA/TensorRT-LLM
Length of output: 5171
🏁 Script executed:
# Search for how these custom ops are registered and whether they guarantee _schema
fd -t f "\.py" tensorrt_llm/_torch/auto_deploy | xargs rg -l "torch_moe\|torch_quant_fp8_moe" | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look for the custom op registration code
rg -B 5 -A 10 "torch_moe.*\.define\|register.*torch_moe" tensorrt_llm/_torch/auto_deployRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check what PyTorch version is used and its op schema conventions
rg "pytorch|torch.*version" tensorrt_llm/_torch/auto_deploy/export/export.py setup.py pyproject.toml CMakeLists.txt 2>/dev/null | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 630
🏁 Script executed:
# Search for where these custom ops are registered
rg "torch_moe\|torch_quant_fp8_moe\|torch_quant_nvfp4_moe" tensorrt_llm/_torch --type py -B 3 -A 3 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look for the op definition in the auto_deploy module
fd -t f "\.py" tensorrt_llm/_torch/auto_deploy | xargs rg -l "define\|register" | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if there's a __init__.py or ops.py in auto_deploy
fd "(__init__|ops)" tensorrt_llm/_torch/auto_deployRepository: NVIDIA/TensorRT-LLM
Length of output: 1699
🏁 Script executed:
# Look for any cpp/cu files that register these ops
fd -t f "\.(cpp|cu|h)" tensorrt_llm/_torch/auto_deployRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check PyTorch OpOverload documentation and behavior in this codebase
rg "_schema\|_schemas" tensorrt_llm/_torch --type py -B 2 -A 2 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Search for "torch_moe" definition in export library
fd -t f "\.py" tensorrt_llm/_torch/auto_deploy/export/library && cat tensorrt_llm/_torch/auto_deploy/export/library/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1279
🏁 Script executed:
# Check the main __init__.py of auto_deploy
cat tensorrt_llm/_torch/auto_deploy/__init__.py | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 417
🏁 Script executed:
# Search more broadly for custom op definitions
rg "torch\.library\|define.*torch_moe\|library\.define" tensorrt_llm/_torch --type py -B 2 -A 5 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look for how these ops are accessed through torch.ops.auto_deploy
rg "auto_deploy\.\|torch\.ops" tensorrt_llm/_torch/auto_deploy --type py | grep -E "torch_moe|torch_quant" | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if there's C++ registration that defines the schema
find tensorrt_llm/_torch/auto_deploy -name "*.cpp" -o -name "*.cu" -o -name "*.h" | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look in the fused_moe directory for op definitions
cat tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 746
🏁 Script executed:
# Search for torch_moe definition in fused_moe
fd -t f "\.py" tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe && ls -la tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/Repository: NVIDIA/TensorRT-LLM
Length of output: 1046
🏁 Script executed:
# Search for where these specific ops are created
rg "torch\.ops\.auto_deploy\.(torch_moe|torch_quant)" tensorrt_llm/_torch --type py -B 3 -A 3 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 6898
🏁 Script executed:
# Look at the actual call site in export.py to understand what node.target is
rg -B 15 "if not is_op\(node, moe_ops\)" tensorrt_llm/_torch/auto_deploy/export/export.py | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 678
🏁 Script executed:
# Check PyTorch documentation about OpOverload and _schema/_schemas by looking at actual usage patterns
python3 << 'EOF'
import torch
# Test if standard PyTorch ops have _schema or _schemas
try:
op = torch.ops.aten.add
print(f"torch.ops.aten.add type: {type(op)}")
print(f"Has _schema: {hasattr(op, '_schema')}")
print(f"Has _schemas: {hasattr(op, '_schemas')}")
if hasattr(op, '_schema'):
print(f"_schema: {op._schema}")
if hasattr(op, '_schemas'):
print(f"_schemas keys: {list(op._schemas.keys())[:3]}")
except Exception as e:
print(f"Error: {e}")
# Test with a custom library op
try:
torch.library.define("my_lib::my_op(Tensor x) -> Tensor")
op = torch.ops.my_lib.my_op
print(f"\nCustom op type: {type(op)}")
print(f"Has _schema: {hasattr(op, '_schema')}")
print(f"Has _schemas: {hasattr(op, '_schemas')}")
except Exception as e:
print(f"Custom op error: {e}")
EOFRepository: NVIDIA/TensorRT-LLM
Length of output: 180
🏁 Script executed:
# Look at how torch_moe is registered
cat tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | head -150Repository: NVIDIA/TensorRT-LLM
Length of output: 6887
🏁 Script executed:
# Search for the schema definition in torch_moe.py
rg "library\.define\|_schema" tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py -B 2 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if there's a .define call that registers these ops with a schema
rg "library\.define" tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/ -B 2 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check what torch.library API is being used
cat tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | grep -A 100 "library" | head -120Repository: NVIDIA/TensorRT-LLM
Length of output: 5304
🏁 Script executed:
# Look for how the ops are exported/created at module level
tail -100 tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3692
🏁 Script executed:
# Let's verify how PyTorch's custom_op decorator creates op objects and their schema attributes
cat tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | grep -A 50 "@torch.library.custom_op"Repository: NVIDIA/TensorRT-LLM
Length of output: 11620
🏁 Script executed:
# Check if there's any op registration that creates multiple schemas (_schemas dict)
rg "register_fake\|impl\|CompositeImplicitAutograd" tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Verify the schema access by looking at whether the fallback path is ever exercised
rg "_schemas" tensorrt_llm/_torch/auto_deploy/export/export.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 154
🏁 Script executed:
# Verify that custom ops always have _schema by checking the decorator and API version
rg "torch.library" tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 440
🏁 Script executed:
# Check if there's any test that exercises the schema access
fd "test.*export" tensorrt_llm --type f | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Verify the imports and torch version constraints
head -50 tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/torch_moe.py | grep -E "import|torch"Repository: NVIDIA/TensorRT-LLM
Length of output: 636
🏁 Script executed:
# Let's verify whether modern torch.library.custom_op OpOverload objects actually have _schema or _schemas
# by checking PyTorch documentation and any reference code in the repo
rg "node.target" tensorrt_llm/_torch/auto_deploy --type py -B 5 -A 5 | grep -E "schema|target|OpOverload" | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 3361
🏁 Script executed:
# Check if there are other places where op schema is accessed, to see the pattern
rg "\._schema\b" tensorrt_llm/_torch --type py -B 2 -A 2 | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 1927
🏁 Script executed:
# Look for how is_op works to understand what node.target is
rg -A 10 "def is_op" tensorrt_llm/_torch/auto_deploy/utils/node_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 486
🏁 Script executed:
# Look at the get_op_schema function in node_utils.py to see the better pattern
cat -n tensorrt_llm/_torch/auto_deploy/utils/node_utils.py | sed -n '200,250p'Repository: NVIDIA/TensorRT-LLM
Length of output: 2118
🏁 Script executed:
# Find the exact location of get_op_schema function
rg -B 5 "def get_op_schema" tensorrt_llm/_torch/auto_deploy/utils/node_utils.py -A 15Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Search more broadly for get_op_schema
rg "get_op_schema" tensorrt_llm/_torch/auto_deploy/utils/node_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look at the broader context where the schema check appears
cat -n tensorrt_llm/_torch/auto_deploy/utils/node_utils.py | grep -B 10 -A 10 "if hasattr.*_schemas"Repository: NVIDIA/TensorRT-LLM
Length of output: 1016
🏁 Script executed:
# Now let's see if there's exception handling around the problematic code in export.py
cat -n tensorrt_llm/_torch/auto_deploy/export/export.py | sed -n '270,295p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1243
🏁 Script executed:
# Check if ad_logger is used anywhere else with warning
rg "ad_logger\.(warning|error)" tensorrt_llm/_torch/auto_deploy/export/export.py | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 45
Add defensive error handling when accessing the op schema.
Line 279 will raise an unhandled AttributeError or StopIteration if op lacks both _schema and _schemas (or if _schemas is empty). While these custom ops are guaranteed to have _schema via the @torch.library.custom_op() decorator, the lack of explicit error handling makes failures opaque. The codebase already uses a better pattern elsewhere (see node_utils.py lines 778–783) that checks both attributes and raises a clear RuntimeError if neither exists.
Better defensive approach
if hasattr(op, "_schemas"):
schema = next(iter(op._schemas.values()))
elif hasattr(op, "_schema"):
schema = op._schema
else:
ad_logger.warning(f"Cannot resolve schema for op '{node.target}'; skipping expansion")
continue🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/auto_deploy/export/export.py` around lines 278 - 279, The
current access of op schema (variables op, node.target, schema) can raise
AttributeError/StopIteration if op lacks _schema and _schemas; update the logic
in the export logic around where op = node.target to check attributes
defensively: first check hasattr(op, "_schemas") and use
next(iter(op._schemas.values())), then elif hasattr(op, "_schema") use
op._schema, otherwise log a clear warning via ad_logger (including node.target)
and skip/continue processing the node (mirroring the pattern used in
node_utils.py) so failures are explicit and non-fatal.
| def _reset_trtllm_planner(): | ||
| """Force a full reset of the global TRT-LLM planner so buffers are re-allocated.""" | ||
| _GlobalTrtllmPlanner.__init__() |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
rg -nP --type=py -C5 'class _GlobalTrtllmPlanner'Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
cat tests/unittest/_torch/auto_deploy/unit/singulgpu/custom_ops/attention/test_trtllm_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 192
🏁 Script executed:
rg -nP --type=py 'GlobalTrtllmPlanner'Repository: NVIDIA/TensorRT-LLM
Length of output: 2837
🏁 Script executed:
rg -nP --type=py -B5 -A30 'class _TrtllmPlanner'Repository: NVIDIA/TensorRT-LLM
Length of output: 4808
🏁 Script executed:
rg -nP --type=py -A50 'def reset\(self' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2535
🏁 Script executed:
head -50 tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1595
🏁 Script executed:
sed -n '40,80p' tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1499
🏁 Script executed:
rg -nP --type=py -B3 -A3 '_reset_trtllm_planner' tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2842
🏁 Script executed:
rg -nP --type=py 'def prepare_trtllm_metadata_host' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 100
🏁 Script executed:
sed -n '204,250p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2170
🏁 Script executed:
rg -nP --type=py -B3 -A10 'class _FlashInferPlanner' tensorrt_llm/Repository: NVIDIA/TensorRT-LLM
Length of output: 1787
🏁 Script executed:
rg -nP --type=py 'def __init__|def reset|def clear' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 144
🏁 Script executed:
sed -n '76,130p' tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2306
🏁 Script executed:
rg -nP --type=py -A5 'def reset.*device' tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_flashinfer_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
rg -nP --type=py '_reset.*planner|_GlobalFlashInferPlanner' tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_flashinfer_attention_op.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 582
Call reset() instead of __init__(), or add a dedicated parameterless clear() method for testing.
_GlobalTrtllmPlanner already has a reset(device, max_batch, max_blocks_per_seq) method designed for this purpose. Calling __init__() directly bypasses its guard logic ("if workspace is not None: return") and is fragile to future changes in the __init__ signature or cleanup requirements.
The test can either:
- Add a parameterless
clear()method that just resets state to None (mirroring the initialization without requiring parameters), similar to__init__()but explicitly intended for testing, or - Call
reset(torch.device("cpu"), 32, 2048)with reasonable defaults before each test, aligning with the FlashInfer pattern used elsewhere in the codebase.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`
around lines 37 - 39, The helper _reset_trtllm_planner currently calls
_GlobalTrtllmPlanner.__init__() which bypasses guard logic and is fragile;
change it to either call the existing reset API with sensible defaults (e.g.
_GlobalTrtllmPlanner.reset(torch.device("cpu"), 32, 2048)) or add a dedicated
parameterless clear() method on _GlobalTrtllmPlanner that zeroes internal state
and call that from _reset_trtllm_planner instead; update references to use
_GlobalTrtllmPlanner.reset or the new clear() and remove the direct __init__
invocation.
| from _custom_op_utils import torch_rope_reference | ||
|
|
||
| # Import after we've imported torch (to ensure custom ops are registered) | ||
| from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqa: F401 |
There was a problem hiding this comment.
# noqa: F401 is ineffective — Ruff fires RUF100
Ruff's RUF100 enforces that noqa suppressions are "valid" — i.e., the violations they suppress are actually being triggered. Since F401 is not enabled in this project's Ruff configuration, a noqa directive that no longer matches any diagnostic violations is likely included by mistake and should be removed to avoid confusion.
If F401 suppression must be preserved for another linter (e.g., flake8), the Ruff-recommended approach is to separate multi-tool directives with a second # character so --fix won't remove it, or alternatively declare F401 as an external rule in pyproject.toml via lint.external = ["F401"].
🔧 Proposed fix (simplest — remove the unused noqa)
-from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqa: F401
+from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqaOr, if only Ruff is used and F401 is never enabled:
-from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqa: F401
+from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # type: ignore[import] # side-effect importOr the cleanest option — declare F401 as an external rule so Ruff respects the directive for other tools:
# In pyproject.toml under [tool.ruff.lint]:
+external = ["F401"]🧰 Tools
🪛 Ruff (0.15.1)
[warning] 8-8: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py`
at line 8, The import line "from tensorrt_llm._torch.auto_deploy.custom_ops.rope
import triton_rope" has an unnecessary "# noqa: F401" that triggers Ruff RUF100;
remove the trailing " # noqa: F401" from that import to clear the false
suppression, or if you must preserve the suppression for another tool keep it as
"## noqa: F401" (double-hash) or declare F401 as an external rule in
pyproject.toml under lint.external = ["F401"] so Ruff will not remove it.
| import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 | ||
| from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm |
There was a problem hiding this comment.
Use a namespace-preserving import and remove the unused noqa.
Ruff flags the # noqa: F401 as unused, and the import style should follow the namespace rule.
♻️ Suggested change
-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401
+from tensorrt_llm._torch.auto_deploy import custom_ops # side-effect importAs per coding guidelines: Always maintain the namespace when importing. Use from package.subpackage import foo instead of from package.subpackage.foo import SomeClass or import package.
🧰 Tools
🪛 Ruff (0.15.1)
[warning] 14-14: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py`
around lines 14 - 15, Replace the top-level import that currently reads "import
tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401" with a
namespace-preserving import: "from tensorrt_llm._torch.auto_deploy import
custom_ops" and remove the unused "# noqa: F401"; keep the existing "from
tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm" line as-is so
the module registration side-effect (custom_ops) is preserved while following
the namespace import guideline.
| _skip_reason = "Requires NVFP4 (Blackwell+) and TRT-LLM ops" | ||
| _skip_condition = not (fp4_compatible() and trtllm_ops_available()) |
There was a problem hiding this comment.
Rename module-level globals to the G_ upper snake_case convention.
♻️ Suggested change
-_skip_reason = "Requires NVFP4 (Blackwell+) and TRT-LLM ops"
-_skip_condition = not (fp4_compatible() and trtllm_ops_available())
+G_SKIP_REASON = "Requires NVFP4 (Blackwell+) and TRT-LLM ops"
+G_SKIP_CONDITION = not (fp4_compatible() and trtllm_ops_available())
-@pytest.mark.skipif(_skip_condition, reason=_skip_reason)
+@pytest.mark.skipif(G_SKIP_CONDITION, reason=G_SKIP_REASON)As per coding guidelines: Use upper snake_case with prefix 'G' for global variables (e.g., G_MY_GLOBAL = ...).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py`
around lines 20 - 21, Rename the module-level globals _skip_reason and
_skip_condition to follow the G_ upper snake_case convention (e.g.,
G_SKIP_REASON and G_SKIP_CONDITION) and update all references to them in this
test file (test_nvfp4_swiglu.py) accordingly; ensure the values remain the same
(G_SKIP_REASON = "Requires NVFP4 (Blackwell+) and TRT-LLM ops" and
G_SKIP_CONDITION = not (fp4_compatible() and trtllm_ops_available())) and adjust
any usages in decorators or test logic that referenced
_skip_reason/_skip_condition.
| from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( # noqa: E402 | ||
| Glm4MoeLiteConfig, | ||
| Glm4MoeLiteForCausalLM, | ||
| ) |
There was a problem hiding this comment.
Remove unused # noqa: E402 and fix import style per coding guidelines.
Two issues:
- Ruff flags the
# noqa: E402comment as a non-enabled directive — it has no effect and should be dropped. - The import uses the
from package.subpackage.module import SomeClassform, which is explicitly disallowed. The guideline requiresfrom package.subpackage import module, keeping the namespace.
♻️ Proposed fix
-from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( # noqa: E402
- Glm4MoeLiteConfig,
- Glm4MoeLiteForCausalLM,
-)
+from tensorrt_llm._torch.auto_deploy.models.custom import modeling_glm4_moe_liteThen replace every usage below with the qualified names, e.g. modeling_glm4_moe_lite.Glm4MoeLiteConfig and modeling_glm4_moe_lite.Glm4MoeLiteForCausalLM.
As per coding guidelines: "Use from package.subpackage import foo instead of from package.subpackage.foo import SomeClass."
🧰 Tools
🪛 Ruff (0.15.1)
[warning] 359-359: Unused noqa directive (non-enabled: E402)
Remove unused noqa directive
(RUF100)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`
around lines 359 - 362, Remove the unused "# noqa: E402" and change the import
style to keep the module namespace: replace the current "from
tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import
(Glm4MoeLiteConfig, Glm4MoeLiteForCausalLM)" with "from
tensorrt_llm._torch.auto_deploy.models.custom import modeling_glm4_moe_lite" and
then update all usages in this test (references to Glm4MoeLiteConfig and
Glm4MoeLiteForCausalLM) to use qualified names
modeling_glm4_moe_lite.Glm4MoeLiteConfig and
modeling_glm4_moe_lite.Glm4MoeLiteForCausalLM.
…, so not working Signed-off-by: Eran Geva <egeva@nvl72036-T08.cm.cluster>
Signed-off-by: Eran Geva <egeva@nvl72036-T08.cm.cluster>
1 participant
Summary by CodeRabbit
New Features
Improvements
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.