[#11440] [feat] AutoDeploy : Support Qwen3.5 - #11394
Conversation
|
@CodeRabbit summary |
✅ Actions performedSummary regeneration triggered. |
📝 WalkthroughWalkthroughAdds cached Gated Delta Rule support (FLA and Torch backends), model export patches for Qwen3Next, MOE export/sharding enhancements (reduce/expand experts), a full Qwen3.5 MoE reference model, transforms to insert cached gated-delta and split fused MoE for sharding, config files, shim/small fixes, and extensive unit/integration tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Inference Client
participant Model as Patched Model
participant GDN as GatedDeltaNet Layer
participant CacheOp as Cached Gated Delta Op
participant KV as KV Cache State
Client->>Model: forward (prefill + decode)
Model->>GDN: call patched GDN forward
rect rgba(100,200,100,0.5)
Note over GDN,CacheOp: Prefill phase (per-sequence)
GDN->>CacheOp: q,k,v,g,meta (batch_info,cu_seqlen,slot_idx,use_initial_states)
CacheOp->>KV: allocate/load delta_cache
CacheOp->>CacheOp: run chunk_gated_delta_rule -> y_prefill, final_state
CacheOp->>KV: write final_state into delta_cache
CacheOp-->>GDN: return y_prefill
end
rect rgba(100,100,200,0.5)
Note over GDN,CacheOp: Decode phase (token-by-token)
loop per token
GDN->>CacheOp: single-token q,k,v,g + meta
CacheOp->>KV: read delta_cache for slot
CacheOp->>CacheOp: fused_recurrent_gated_delta_rule_fwd -> y_token, updated_state
CacheOp->>KV: update delta_cache
CacheOp-->>GDN: return y_token
end
end
GDN-->>Model: attention outputs
Model-->>Client: final outputs
sequenceDiagram
participant Transform as insert_cached_gated_delta_rule
participant IR as Graph / source_attn_node
participant Backend as AttentionDescriptor
participant CacheInit as Cache Initializer
Transform->>IR: find source attention nodes (torch_gated_delta_rule)
Transform->>Backend: query source/cached ops and metadata keys
Backend-->>Transform: provide cached op and initializer function
Transform->>CacheInit: call get_cache_initializers(node, cache_config)
CacheInit-->>Transform: StateResourceHandler (delta_cache) descriptor
Transform->>IR: replace source node with cached op, wire delta_cache and metadata
Transform-->>IR: emit updated graph with cached gated-delta operator
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (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 |
65515b0 to
2d406cd
Compare
|
@CodeRabbit summary |
✅ Actions performedSummary regeneration triggered. |
e4d957d to
4dbc52f
Compare
There was a problem hiding this comment.
Actionable comments posted: 9
🤖 Fix all issues with AI agents
Verify each finding against the current code and only fix it if needed.
In `@examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml`:
- Around line 15-18: Remove the leftover debug overrides by deleting the
commented-out blocks "text_config" and "vision_config" (the num_hidden_layers
and depth test overrides) from the YAML so the config contains only the intended
production settings; if you want to keep examples, move them to a clearly marked
"examples" or "notes" section with explanatory text instead of inline commented
overrides.
In `@examples/auto_deploy/model_registry/configs/qwen3Next.yaml`:
- Around line 1-20: The filename qwen3Next.yaml uses camelCase while other
configs use snake_case; rename the file to qwen3_next.yaml and update any
references to it (e.g., in deployment manifests, model registry entries, CI
jobs, or README) so consumers reference the new name; ensure references to keys
in this file such as runtime, model_factory, kv_cache_config, and
transforms.export_to_gm remain unchanged and run the test/CI to verify no broken
links.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/fla/torch_backend_gated_delta.py`:
- Line 156: The custom op registration currently declares mutates_args=() but
the implementation of the op (registered via
`@torch.library.custom_op`("auto_deploy::torch_cached_gated_delta_rule",
mutates_args=())) mutates delta_cache in-place (assignments like
delta_cache[slot] = final_state.squeeze(0).to(delta_cache.dtype) and
delta_cache[slot] = new_state.squeeze(0).to(delta_cache.dtype)); update the
decorator to declare that delta_cache is mutated by changing mutates_args to
("delta_cache",) to match the FLA backend (fla_backend_gated_delta.py) so Torch
compiler/FX know the op has side effects.
In `@tensorrt_llm/_torch/auto_deploy/export/export.py`:
- Around line 187-192: The code in list_arg_indices assumes list-of-node
arguments start at index 3 (list_arg_indices and node.args) and that MOE ops
follow a fixed op-set/layout; update the export logic to document and guard this
convention: add a clear inline comment describing the expected argument
convention (e.g., "args[0:3] = hidden_states, selected_experts, routing_weights;
args[3:] = per-expert weight/scale lists") next to the list_arg_indices
computation, and add a lightweight validation (e.g., an assertion or warning
using the node/op identifier) when the node's arg count or types don't match the
convention so new MOE ops with different layouts are flagged rather than
silently skipped; keep references to list_arg_indices and node.args in the added
comment and validation.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py`:
- Around line 4-12: Replace the direct symbol import from modeling_qwen3_5_moe
with a module-level import (e.g., import the modeling_qwen3_5_moe module) and
re-export the needed classes by referencing them off that module (so
Qwen3_5MoeForCausalLM and Qwen3_5MoeForConditionalGeneration are assigned from
modeling_qwen3_5_moe.Qwen3_5MoeForCausalGeneration, etc.), keeping the existing
__all__ entries; this preserves the module namespace per guidelines and ensures
consumers still get the same exported names.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py`:
- Around line 1-2: The file modeling_qwen3_5_moe.py is missing the required
NVIDIA Apache License 2.0 header; update the top of the file (before any imports
or code in modeling_qwen3_5_moe.py) to add the full NVIDIA copyright header in
Apache License 2.0 format with the latest modification year (2025), matching the
header style used in other new files such as fla_backend_gated_delta.py and
torch_backend_gated_delta.py; ensure the header includes the full license
boilerplate text and copyright notice for NVIDIA CORPORATION.
- Around line 1381-1382: The loop over total_input_ids shadows the loop variable
ids by reassigning it with a masked version (ids = ids[attention_mask[i] == 1]);
instead create a new variable name (e.g., masked_ids or filtered_ids) and use
that within the loop to hold the masked tensor, leaving the original loop
variable ids intact; update any subsequent references in the loop body from ids
to the new name and keep references to total_input_ids and attention_mask as-is.
In `@tensorrt_llm/_torch/auto_deploy/models/patches/qwen3_next.py`:
- Around line 18-27: The current direct imports of Qwen3NextDynamicCache,
Qwen3NextGatedDeltaNet, Qwen3NextModel, Qwen3NextSparseMoeBlock, and
apply_mask_to_padding_states violate the namespace rule; change the transformers
import to import the module (e.g., from transformers.models.qwen3_next import
modeling_qwen3_next) and update all usages to reference
modeling_qwen3_next.Qwen3NextModel, modeling_qwen3_next.Qwen3NextDynamicCache,
modeling_qwen3_next.Qwen3NextGatedDeltaNet,
modeling_qwen3_next.Qwen3NextSparseMoeBlock, and
modeling_qwen3_next.apply_mask_to_padding_states instead of the direct names;
keep the existing local export imports (BaseExportPatch, ExportPatchRegistry)
as-is if they already follow the guideline.
- Around line 1-12: Add the standard NVIDIA Apache-2.0 copyright header at the
top of the new module
tensorrt_llm/_torch/auto_deploy/models/patches/qwen3_next.py (above the module
docstring) with the current year and the required license boilerplate; ensure
the header follows the project's canonical NVIDIA header format used in other
source files so the file-level comment includes the copyright owner and the
Apache License 2.0 notice.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py`:
- Around line 190-219: The replacement currently hardcodes is_gated_mlp=True and
act_fn=ActivationType.Silu in _replace_torch_moe_fused_ops which can
misrepresent the original torch_moe_fused node; update the code to read these
options from the original node's kwargs (e.g., inspect node.kwargs for
"is_gated_mlp" and "act_fn"), convert/normalize act_fn to the expected type
(int(ActivationType) only if necessary), and pass those extracted values into
the graph.call_function for replacement_op, falling back to the current defaults
only if the keys are absent; ensure you reference the original node (node) and
replacement_op when copying args/kwargs so behavior is preserved for non-Qwen
models.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`:
- Around line 141-144: The change flips the default for the argument whose
description is "When True, apply simple shard (column split + all_gather) to
'leftover' linear nodes that are not part of any layer subgraph." which causes
unprocessed linear nodes to stop being sharded by default; if this was not
intentional, revert the default back to True for that parameter (restore
previous default), update any AutoDeploy presets that rely on implicit sharding
and the user-facing docs/changelog to call out the behavioral change, and add a
short note in the function/class docstring where that parameter is declared so
future reviewers see the compatibility impact.
- Around line 1888-1896: The code incorrectly treats conv1d_node.args[-1] as the
output-channel dim; change logic to derive conv_dim from the conv weight tensor
(for torch_causal_conv1d use weight.shape[0]) and use that conv_dim for the
assertion and later parameter modification; specifically, when computing
conv_split_sizes_original from split_node_after_conv and when updating the conv
parameter (the code around conv1d_node, split_node_after_conv, and the later
parameter modification lines), replace use of conv1d_node.args[-1] with the
derived weight.shape[0] (and ensure the division by world_size uses this integer
conv_dim).
- Around line 1882-1886: The code calls bfs returning two values
(split_node_after_conv, depth) but depth is unused and the predicate checks
len(list(n.users)) >= 3 then later indexes a single-element list; update the bfs
call to only capture the node (e.g., split_node_after_conv = bfs(...)[0] or
change bfs to return a single value) and replace any single-element list
indexing with direct extraction (avoid unnecessary list(...) wrapping) in the
block using conv1d_node and the is_op predicate referencing
torch.ops.aten.split_with_sizes and torch.ops.aten.split; apply the same cleanup
to the analogous occurrence around the second block that uses bfs with
conv1d_node and the split ops.
In `@tensorrt_llm/_torch/auto_deploy/utils/node_utils.py`:
- Around line 690-697: The loop that extends boundary_nodes is matching any
torch.ops.aten.add via is_op, which is too permissive; update the predicate used
in the bfs call (where boundary_nodes, bfs and is_op are referenced) to use an
is_residual_add (or equivalent guard) instead of raw is_op(...,
torch.ops.aten.add) so only true residual/add-for-skip connections are treated
as layer boundaries; keep include_root=False and preserve the append logic
(boundary_nodes.append(next_res_add)) but ensure the new predicate checks
residual semantics (bias/fused adds should not match).
- Around line 511-513: Replace the current two-line guard that checks the
presence and truthiness of the key with a single get-based check: use
gm.meta.get("_weight_mapping_computed") to decide early return, then assign
gm.meta["_weight_mapping_computed"] = True as before; this simplifies the logic
around the "_weight_mapping_computed" meta flag on gm.meta.
In `@tensorrt_llm/bench/benchmark/throughput.py`:
- Around line 323-327: Remove the unnecessary noqa suppression on the
side-effect import: in the conditional that checks options.backend ==
"_autodeploy" and imports tensorrt_llm._torch.auto_deploy (the block that
eagerly registers custom model configs before AutoConfig.from_pretrained),
delete the trailing " # noqa: F401" so the import remains documented by the
existing comment but without a no-op lint directive.
In
`@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py`:
- Around line 215-244: The PreStackedMoEModel class is duplicated; extract it
into a shared test utility (e.g., _model_test_utils) and parameterize
differences (input dimensionality, batch/sequence sizes, default
hidden_size/intermediate_size/num_experts/top_k) so both tests import the same
class; update the copies to import the shared PreStackedMoEModel and pass
specific parameters (or override get_input) instead of maintaining two
near-identical definitions (refer to class PreStackedMoEModel, its __init__,
forward, and get_input to locate and consolidate logic).
In
`@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py`:
- Around line 1079-1081: The loop over gm_transformed.named_parameters()
declares an unused variable name causing a lint warning; update the loop in the
test (the for loop that currently reads "for name, param in
gm_transformed.named_parameters():") to use an underscore for the unused
variable (e.g., "for _, param in gm_transformed.named_parameters():") so only
param is treated as used before calling gm_transformed.load_state_dict(...).
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/fla/test_fla_cached_gated_delta_rule.py`:
- Around line 33-44: Extract the duplicated helper function _random_inputs into
a shared test utility (e.g., a helpers module or conftest) and replace the local
definitions in both test_fla_cached_gated_delta_rule.py and
test_torch_cached_gated_delta_rule.py with an import; specifically, move the
function body exactly as-is, export it (or add as a pytest fixture if
preferred), update both test files to import _random_inputs from the shared
module, and remove the duplicate local implementations so tests use the single
canonical helper.
In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_5_moe.py`:
- Around line 414-420: The loop currently reassigns the loop variable expert_idx
(from expert_hit) with expert_idx = expert_idx[0], shadowing the original;
change this to use a new local name (e.g., idx or expert_idx_tuple) or unpack
the tensor result so you don't overwrite the loop variable: when iterating over
expert_hit in the for loop, extract the scalar index into a fresh variable and
use that (for example use idx = expert_idx[0] or unpack expert_idx_tuple) before
comparing to num_experts and calling torch.where on expert_mask to compute
top_k_pos and token_idx.
- Around line 1146-1150: The test unpacks two values from model.get_rope_index
but the second one is unused (deltas), so rename the unpacked variable to
_deltas to indicate intentional discard; update the two occurrences where you do
"position_ids, deltas = model.get_rope_index(...)" (including the call near the
assertion and the later unpack at line ~1172) to "position_ids, _deltas =
model.get_rope_index(...)" so static analysis no longer flags an unused
variable.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_gdn_patches.py`:
- Around line 17-20: The test currently imports Qwen3NextGatedDeltaNet and
torch_chunk_gated_delta_rule directly; change these to import the module(s)
instead and reference the symbols through the module namespace (e.g., import
transformers.models.qwen3_next.modeling_qwen3_next as qwen3_next) so you use
qwen3_next.Qwen3NextGatedDeltaNet and qwen3_next.torch_chunk_gated_delta_rule,
and update the internal patched forward import to use the module reference (use
qwen3_next._patched_gdn_forward) to comply with the "from package.subpackage
import module" guideline.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_patches.py`:
- Around line 64-66: The broad "except Exception as e" in the layer extraction
block swallows all errors and returns None, hiding test failures; remove the
broad try/except so exceptions propagate in the test (or if you must catch,
narrow it to expected exceptions like ImportError or ValueError and re-raise
others), and stop returning None so the test at the assertion (module is not
None) fails with the original error instead of a generic None result; locate the
try/except that prints "Error extracting layer" and update it accordingly.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gated_delta_rule_cache.py`:
- Around line 36-65: Duplicate test helper classes (DummyFactory and
GatedDeltaRuleModel) appear in test_gated_delta_rule_cache.py and
test_torch_gated_delta_rule_cache.py; consolidate them into a shared test
utility module (e.g., tests/unittest/_torch/auto_deploy/test_utils.py) and
update both test files to import DummyFactory and GatedDeltaRuleModel from that
module, then parameterize the tests by backend/dtype/tolerance (use
pytest.mark.parametrize) so each test passes the backend name, dtype and
tolerance values into the same test function instead of duplicating nearly
identical files; update references to DummyFactory.get_cache_config_updates,
DummyFactory.build_model, and GatedDeltaRuleModel to use the shared
implementations.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py`:
- Around line 1079-1114: The PreStackedMoEModel class is duplicated across test
files; extract it into a single shared test utility (e.g., a test helper module)
and have both tests import and instantiate it with different default parameters
as needed. Move the class definition (symbols: PreStackedMoEModel, methods
__init__, forward, get_input) into a shared test utility module, parameterize
hidden_size/num_experts/top_k/default dtype, adjust both test modules to import
PreStackedMoEModel and override any differing defaults or reshape logic in their
local test setup, and remove the duplicated class definitions from the
individual test files.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`:
- Around line 336-339: Replace the direct class imports with a module import and
update usages: change the import of Glm4MoeLiteConfig and Glm4MoeLiteForCausalLM
to import the module modeling_glm4_moe_lite (i.e., use from
tensorrt_llm._torch.auto_deploy.models.custom import modeling_glm4_moe_lite),
then update all references in this test (e.g., Glm4MoeLiteConfig and
Glm4MoeLiteForCausalLM) to modeling_glm4_moe_lite.Glm4MoeLiteConfig and
modeling_glm4_moe_lite.Glm4MoeLiteForCausalLM so the code follows the required
module-level import style.
- Line 249: Remove the unnecessary "# noqa: F401, E402" directives from the
import statements that include the module symbol import
tensorrt_llm._torch.auto_deploy.custom_ops (and the other import in this file
with the same noqa), i.e., delete the trailing noqa comment(s) so the import
lines are clean and rely on configured linters rather than disabled/no-op
directives.
- Around line 308-312: The helper function _count_moe_experts inside the test
duplicates the module-level _count_moe_experts_in_graph; remove the local
_count_moe_experts and call _count_moe_experts_in_graph instead, or move/reorder
the _count_moe_experts_in_graph definition so it is defined before the test;
update any call sites in the test to use _count_moe_experts_in_graph
(referencing that function name) and ensure imports/visibility remain correct.
🧹 Nitpick comments (19)
🤖 Fix all nitpicks with AI agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml`: - Around line 15-18: Remove the leftover debug overrides by deleting the commented-out blocks "text_config" and "vision_config" (the num_hidden_layers and depth test overrides) from the YAML so the config contains only the intended production settings; if you want to keep examples, move them to a clearly marked "examples" or "notes" section with explanatory text instead of inline commented overrides. In `@examples/auto_deploy/model_registry/configs/qwen3Next.yaml`: - Around line 1-20: The filename qwen3Next.yaml uses camelCase while other configs use snake_case; rename the file to qwen3_next.yaml and update any references to it (e.g., in deployment manifests, model registry entries, CI jobs, or README) so consumers reference the new name; ensure references to keys in this file such as runtime, model_factory, kv_cache_config, and transforms.export_to_gm remain unchanged and run the test/CI to verify no broken links. In `@tensorrt_llm/_torch/auto_deploy/export/export.py`: - Around line 187-192: The code in list_arg_indices assumes list-of-node arguments start at index 3 (list_arg_indices and node.args) and that MOE ops follow a fixed op-set/layout; update the export logic to document and guard this convention: add a clear inline comment describing the expected argument convention (e.g., "args[0:3] = hidden_states, selected_experts, routing_weights; args[3:] = per-expert weight/scale lists") next to the list_arg_indices computation, and add a lightweight validation (e.g., an assertion or warning using the node/op identifier) when the node's arg count or types don't match the convention so new MOE ops with different layouts are flagged rather than silently skipped; keep references to list_arg_indices and node.args in the added comment and validation. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py`: - Around line 1381-1382: The loop over total_input_ids shadows the loop variable ids by reassigning it with a masked version (ids = ids[attention_mask[i] == 1]); instead create a new variable name (e.g., masked_ids or filtered_ids) and use that within the loop to hold the masked tensor, leaving the original loop variable ids intact; update any subsequent references in the loop body from ids to the new name and keep references to total_input_ids and attention_mask as-is. In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py`: - Around line 190-219: The replacement currently hardcodes is_gated_mlp=True and act_fn=ActivationType.Silu in _replace_torch_moe_fused_ops which can misrepresent the original torch_moe_fused node; update the code to read these options from the original node's kwargs (e.g., inspect node.kwargs for "is_gated_mlp" and "act_fn"), convert/normalize act_fn to the expected type (int(ActivationType) only if necessary), and pass those extracted values into the graph.call_function for replacement_op, falling back to the current defaults only if the keys are absent; ensure you reference the original node (node) and replacement_op when copying args/kwargs so behavior is preserved for non-Qwen models. In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py`: - Around line 1882-1886: The code calls bfs returning two values (split_node_after_conv, depth) but depth is unused and the predicate checks len(list(n.users)) >= 3 then later indexes a single-element list; update the bfs call to only capture the node (e.g., split_node_after_conv = bfs(...)[0] or change bfs to return a single value) and replace any single-element list indexing with direct extraction (avoid unnecessary list(...) wrapping) in the block using conv1d_node and the is_op predicate referencing torch.ops.aten.split_with_sizes and torch.ops.aten.split; apply the same cleanup to the analogous occurrence around the second block that uses bfs with conv1d_node and the split ops. In `@tensorrt_llm/_torch/auto_deploy/utils/node_utils.py`: - Around line 511-513: Replace the current two-line guard that checks the presence and truthiness of the key with a single get-based check: use gm.meta.get("_weight_mapping_computed") to decide early return, then assign gm.meta["_weight_mapping_computed"] = True as before; this simplifies the logic around the "_weight_mapping_computed" meta flag on gm.meta. In `@tensorrt_llm/bench/benchmark/throughput.py`: - Around line 323-327: Remove the unnecessary noqa suppression on the side-effect import: in the conditional that checks options.backend == "_autodeploy" and imports tensorrt_llm._torch.auto_deploy (the block that eagerly registers custom model configs before AutoConfig.from_pretrained), delete the trailing " # noqa: F401" so the import remains documented by the existing comment but without a no-op lint directive. In `@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py`: - Around line 215-244: The PreStackedMoEModel class is duplicated; extract it into a shared test utility (e.g., _model_test_utils) and parameterize differences (input dimensionality, batch/sequence sizes, default hidden_size/intermediate_size/num_experts/top_k) so both tests import the same class; update the copies to import the shared PreStackedMoEModel and pass specific parameters (or override get_input) instead of maintaining two near-identical definitions (refer to class PreStackedMoEModel, its __init__, forward, and get_input to locate and consolidate logic). In `@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py`: - Around line 1079-1081: The loop over gm_transformed.named_parameters() declares an unused variable name causing a lint warning; update the loop in the test (the for loop that currently reads "for name, param in gm_transformed.named_parameters():") to use an underscore for the unused variable (e.g., "for _, param in gm_transformed.named_parameters():") so only param is treated as used before calling gm_transformed.load_state_dict(...). In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/fla/test_fla_cached_gated_delta_rule.py`: - Around line 33-44: Extract the duplicated helper function _random_inputs into a shared test utility (e.g., a helpers module or conftest) and replace the local definitions in both test_fla_cached_gated_delta_rule.py and test_torch_cached_gated_delta_rule.py with an import; specifically, move the function body exactly as-is, export it (or add as a pytest fixture if preferred), update both test files to import _random_inputs from the shared module, and remove the duplicate local implementations so tests use the single canonical helper. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_5_moe.py`: - Around line 414-420: The loop currently reassigns the loop variable expert_idx (from expert_hit) with expert_idx = expert_idx[0], shadowing the original; change this to use a new local name (e.g., idx or expert_idx_tuple) or unpack the tensor result so you don't overwrite the loop variable: when iterating over expert_hit in the for loop, extract the scalar index into a fresh variable and use that (for example use idx = expert_idx[0] or unpack expert_idx_tuple) before comparing to num_experts and calling torch.where on expert_mask to compute top_k_pos and token_idx. - Around line 1146-1150: The test unpacks two values from model.get_rope_index but the second one is unused (deltas), so rename the unpacked variable to _deltas to indicate intentional discard; update the two occurrences where you do "position_ids, deltas = model.get_rope_index(...)" (including the call near the assertion and the later unpack at line ~1172) to "position_ids, _deltas = model.get_rope_index(...)" so static analysis no longer flags an unused variable. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_gdn_patches.py`: - Around line 17-20: The test currently imports Qwen3NextGatedDeltaNet and torch_chunk_gated_delta_rule directly; change these to import the module(s) instead and reference the symbols through the module namespace (e.g., import transformers.models.qwen3_next.modeling_qwen3_next as qwen3_next) so you use qwen3_next.Qwen3NextGatedDeltaNet and qwen3_next.torch_chunk_gated_delta_rule, and update the internal patched forward import to use the module reference (use qwen3_next._patched_gdn_forward) to comply with the "from package.subpackage import module" guideline. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_patches.py`: - Around line 64-66: The broad "except Exception as e" in the layer extraction block swallows all errors and returns None, hiding test failures; remove the broad try/except so exceptions propagate in the test (or if you must catch, narrow it to expected exceptions like ImportError or ValueError and re-raise others), and stop returning None so the test at the assertion (module is not None) fails with the original error instead of a generic None result; locate the try/except that prints "Error extracting layer" and update it accordingly. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gated_delta_rule_cache.py`: - Around line 36-65: Duplicate test helper classes (DummyFactory and GatedDeltaRuleModel) appear in test_gated_delta_rule_cache.py and test_torch_gated_delta_rule_cache.py; consolidate them into a shared test utility module (e.g., tests/unittest/_torch/auto_deploy/test_utils.py) and update both test files to import DummyFactory and GatedDeltaRuleModel from that module, then parameterize the tests by backend/dtype/tolerance (use pytest.mark.parametrize) so each test passes the backend name, dtype and tolerance values into the same test function instead of duplicating nearly identical files; update references to DummyFactory.get_cache_config_updates, DummyFactory.build_model, and GatedDeltaRuleModel to use the shared implementations. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py`: - Around line 1079-1114: The PreStackedMoEModel class is duplicated across test files; extract it into a single shared test utility (e.g., a test helper module) and have both tests import and instantiate it with different default parameters as needed. Move the class definition (symbols: PreStackedMoEModel, methods __init__, forward, get_input) into a shared test utility module, parameterize hidden_size/num_experts/top_k/default dtype, adjust both test modules to import PreStackedMoEModel and override any differing defaults or reshape logic in their local test setup, and remove the duplicated class definitions from the individual test files. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py`: - Line 249: Remove the unnecessary "# noqa: F401, E402" directives from the import statements that include the module symbol import tensorrt_llm._torch.auto_deploy.custom_ops (and the other import in this file with the same noqa), i.e., delete the trailing noqa comment(s) so the import lines are clean and rely on configured linters rather than disabled/no-op directives. - Around line 308-312: The helper function _count_moe_experts inside the test duplicates the module-level _count_moe_experts_in_graph; remove the local _count_moe_experts and call _count_moe_experts_in_graph instead, or move/reorder the _count_moe_experts_in_graph definition so it is defined before the test; update any call sites in the test to use _count_moe_experts_in_graph (referencing that function name) and ensure imports/visibility remain correct.tensorrt_llm/bench/benchmark/throughput.py (1)
323-327: Remove unnecessary# noqa: F401directive.Ruff reports that
F401is not enabled, making thisnoqadirective a no-op. The side-effect import is intentional and well-documented by the comment above, so the suppression isn't needed.Suggested fix
- import tensorrt_llm._torch.auto_deploy # noqa: F401 + import tensorrt_llm._torch.auto_deploy🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/bench/benchmark/throughput.py` around lines 323 - 327, Remove the unnecessary noqa suppression on the side-effect import: in the conditional that checks options.backend == "_autodeploy" and imports tensorrt_llm._torch.auto_deploy (the block that eagerly registers custom model configs before AutoConfig.from_pretrained), delete the trailing " # noqa: F401" so the import remains documented by the existing comment but without a no-op lint directive.examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml (1)
15-18: Consider removing commented-out debug overrides before merging.The commented-out
text_configandvision_configblocks appear to be debugging/testing aids for running with fewer layers. If they're not intended as documentation for users, they could be cleaned up to avoid confusion.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml` around lines 15 - 18, Remove the leftover debug overrides by deleting the commented-out blocks "text_config" and "vision_config" (the num_hidden_layers and depth test overrides) from the YAML so the config contains only the intended production settings; if you want to keep examples, move them to a clearly marked "examples" or "notes" section with explanatory text instead of inline commented overrides.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py (2)
249-249: Remove unusednoqadirectives.Ruff reports that
F401andE402are not enabled, making thenoqacomment a no-op. If the linter configuration doesn't check these rules, the directive is unnecessary clutter. Same applies to line 336.-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401, E402 +import tensorrt_llm._torch.auto_deploy.custom_ops🤖 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` at line 249, Remove the unnecessary "# noqa: F401, E402" directives from the import statements that include the module symbol import tensorrt_llm._torch.auto_deploy.custom_ops (and the other import in this file with the same noqa), i.e., delete the trailing noqa comment(s) so the import lines are clean and rely on configured linters rather than disabled/no-op directives.
308-312: Duplicate of_count_moe_experts_in_graphdefined at line 371.The local
_count_moe_expertslambda inside the test is functionally identical to the module-level_count_moe_experts_in_graphhelper. Reuse the existing helper to avoid duplication.Suggested fix
# --- structural check: both graphs must have the right expert count --- - 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_expertsNote: Since
_count_moe_experts_in_graphis defined later in the file (line 371), you'd need to move it above this test or reorder the definitions.🤖 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 308 - 312, The helper function _count_moe_experts inside the test duplicates the module-level _count_moe_experts_in_graph; remove the local _count_moe_experts and call _count_moe_experts_in_graph instead, or move/reorder the _count_moe_experts_in_graph definition so it is defined before the test; update any call sites in the test to use _count_moe_experts_in_graph (referencing that function name) and ensure imports/visibility remain correct.examples/auto_deploy/model_registry/configs/qwen3Next.yaml (1)
1-20: Minor filename inconsistency:qwen3Next.yamluses camelCase.Other config files in this PR (e.g.,
qwen3.5_moe_400b.yaml,qwen3_5_moe_35b.yaml) follow snake_case. Consider renaming toqwen3_next.yamlfor consistency, unless the intent is to match the HuggingFace model name exactly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_deploy/model_registry/configs/qwen3Next.yaml` around lines 1 - 20, The filename qwen3Next.yaml uses camelCase while other configs use snake_case; rename the file to qwen3_next.yaml and update any references to it (e.g., in deployment manifests, model registry entries, CI jobs, or README) so consumers reference the new name; ensure references to keys in this file such as runtime, model_factory, kv_cache_config, and transforms.export_to_gm remain unchanged and run the test/CI to verify no broken links.tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (1)
190-219: Hardcodedis_gated_mlp=Trueandact_fn=Silulimits generality.
_replace_torch_moe_fused_opsunconditionally injectsis_gated_mlp=Trueandact_fn=int(ActivationType.Silu)for everytorch_moe_fusednode. This is correct for Qwen3.5 MoE but will silently produce wrong results if a future model usestorch_moe_fusedwith a non-gated MLP or different activation.Consider extracting these from the original node's kwargs (if available) or adding a comment clarifying the assumption.
Suggested approach
if is_op(node, torch.ops.auto_deploy.torch_moe_fused): + # torch_moe_fused currently only used by Qwen3.5 MoE (gated SiLU). + # TODO: extract from node kwargs when other variants appear. with graph.inserting_before(node): new_node = graph.call_function( replacement_op, args=node.args, kwargs={"is_gated_mlp": True, "act_fn": int(ActivationType.Silu)}, )🤖 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_moe.py` around lines 190 - 219, The replacement currently hardcodes is_gated_mlp=True and act_fn=ActivationType.Silu in _replace_torch_moe_fused_ops which can misrepresent the original torch_moe_fused node; update the code to read these options from the original node's kwargs (e.g., inspect node.kwargs for "is_gated_mlp" and "act_fn"), convert/normalize act_fn to the expected type (int(ActivationType) only if necessary), and pass those extracted values into the graph.call_function for replacement_op, falling back to the current defaults only if the keys are absent; ensure you reference the original node (node) and replacement_op when copying args/kwargs so behavior is preserved for non-Qwen models.tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py (1)
1079-1081: Unused loop variablename.The Ruff warning is valid —
nameis not used in the loop body. Use_to indicate the variable is intentionally unused.Suggested fix
- for name, param in gm_transformed.named_parameters(): + for _, param in gm_transformed.named_parameters(): param.data.fill_(0.0)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py` around lines 1079 - 1081, The loop over gm_transformed.named_parameters() declares an unused variable name causing a lint warning; update the loop in the test (the for loop that currently reads "for name, param in gm_transformed.named_parameters():") to use an underscore for the unused variable (e.g., "for _, param in gm_transformed.named_parameters():") so only param is treated as used before calling gm_transformed.load_state_dict(...).tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py (1)
215-244:PreStackedMoEModellargely duplicates the class intest_moe_fusion.py.A near-identical
PreStackedMoEModelexists intests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py(lines 1078–1113). Consider extracting this into a shared test utility (e.g.,_model_test_utils) to avoid drift between the two copies. The main differences are the input dimensionality (2D vs 3D) and default sizes, which could be parameterized.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py` around lines 215 - 244, The PreStackedMoEModel class is duplicated; extract it into a shared test utility (e.g., _model_test_utils) and parameterize differences (input dimensionality, batch/sequence sizes, default hidden_size/intermediate_size/num_experts/top_k) so both tests import the same class; update the copies to import the shared PreStackedMoEModel and pass specific parameters (or override get_input) instead of maintaining two near-identical definitions (refer to class PreStackedMoEModel, its __init__, forward, and get_input to locate and consolidate logic).tensorrt_llm/_torch/auto_deploy/utils/node_utils.py (1)
511-513: Small simplification for the mapping guard.Ruff flags the key check as unnecessary;
dict.getkeeps the logic simpler.Suggested tweak
- if "_weight_mapping_computed" in gm.meta and gm.meta["_weight_mapping_computed"]: + if gm.meta.get("_weight_mapping_computed"): return🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/utils/node_utils.py` around lines 511 - 513, Replace the current two-line guard that checks the presence and truthiness of the key with a single get-based check: use gm.meta.get("_weight_mapping_computed") to decide early return, then assign gm.meta["_weight_mapping_computed"] = True as before; this simplifies the logic around the "_weight_mapping_computed" meta flag on gm.meta.tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py (1)
1882-1886: Minor cleanup for unuseddepthand single-element list indexing.These trigger Ruff warnings and can be made a bit tighter.
Suggested cleanup
- split_node_after_conv, depth = bfs( + split_node_after_conv, _ = bfs( conv1d_node, lambda n: is_op(n, [torch.ops.aten.split_with_sizes, torch.ops.aten.split]) and len(list(n.users)) >= 3, ) @@ - delta_node = list(filtered_nodes(subgraph_nodes, is_any_delta_op))[0] + delta_node = next(filtered_nodes(subgraph_nodes, is_any_delta_op), None) + assert delta_node is not None, "Expected a delta op in the subgraph"Also applies to: 1976-1978
🤖 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/sharding.py` around lines 1882 - 1886, The code calls bfs returning two values (split_node_after_conv, depth) but depth is unused and the predicate checks len(list(n.users)) >= 3 then later indexes a single-element list; update the bfs call to only capture the node (e.g., split_node_after_conv = bfs(...)[0] or change bfs to return a single value) and replace any single-element list indexing with direct extraction (avoid unnecessary list(...) wrapping) in the block using conv1d_node and the is_op predicate referencing torch.ops.aten.split_with_sizes and torch.ops.aten.split; apply the same cleanup to the analogous occurrence around the second block that uses bfs with conv1d_node and the split ops.tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_gdn_patches.py (1)
17-20: Imports directly reference classes/functions from a module file.Per coding guidelines, imports should maintain the namespace:
from package.subpackage import modulerather than importing individual classes or functions. For the internal import on line 24, this would be:-from tensorrt_llm._torch.auto_deploy.models.patches.qwen3_next import _patched_gdn_forward +from tensorrt_llm._torch.auto_deploy.models.patches import qwen3_nextThen use
qwen3_next._patched_gdn_forward. The transformers imports (lines 17-20) are harder to restructure due to the deep module path, so those are less critical.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/models/test_qwen3_next_gdn_patches.py` around lines 17 - 20, The test currently imports Qwen3NextGatedDeltaNet and torch_chunk_gated_delta_rule directly; change these to import the module(s) instead and reference the symbols through the module namespace (e.g., import transformers.models.qwen3_next.modeling_qwen3_next as qwen3_next) so you use qwen3_next.Qwen3NextGatedDeltaNet and qwen3_next.torch_chunk_gated_delta_rule, and update the internal patched forward import to use the module reference (use qwen3_next._patched_gdn_forward) to comply with the "from package.subpackage import module" guideline.tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_next_patches.py (1)
64-66: Broadexcept Exceptionswallows all errors silently.This catches every exception (including unexpected ones like
TypeError,AttributeErrorfrom config issues) and returnsNone, making test debugging difficult. The test at line 79 will just assertmodule is not None, giving a generic failure message. Consider either letting exceptions propagate naturally (since this is a test) or catching specific exceptions you expect (e.g.,ImportError,ValueError).Suggested fix: let the test fail directly on unexpected errors
- try: - config = AutoConfig.for_model("qwen3_next") - ... - return module - except Exception as e: - print(f"Error extracting layer: {e}") - return None + config = AutoConfig.for_model("qwen3_next") + ... + if module is None: + print(f"Layer '{layer_name}' not found in the model.") + else: + print(f"Successfully extracted layer '{layer_name}'.") + return moduleAs per coding guidelines, "When using try-except blocks, limit the except to the smallest set of errors possible. Avoid bare
except:clauses."🤖 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/models/test_qwen3_next_patches.py` around lines 64 - 66, The broad "except Exception as e" in the layer extraction block swallows all errors and returns None, hiding test failures; remove the broad try/except so exceptions propagate in the test (or if you must catch, narrow it to expected exceptions like ImportError or ValueError and re-raise others), and stop returning None so the test at the assertion (module is not None) fails with the original error instead of a generic None result; locate the try/except that prints "Error extracting layer" and update it accordingly.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/fla/test_fla_cached_gated_delta_rule.py (1)
33-44: Duplicated_random_inputshelper across test files.This helper is identical to the one in
test_torch_cached_gated_delta_rule.py(lines 32-43). Consider extracting it to a shared test utility module (e.g., a conftest or helper module) to avoid duplication.🤖 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/fla/test_fla_cached_gated_delta_rule.py` around lines 33 - 44, Extract the duplicated helper function _random_inputs into a shared test utility (e.g., a helpers module or conftest) and replace the local definitions in both test_fla_cached_gated_delta_rule.py and test_torch_cached_gated_delta_rule.py with an import; specifically, move the function body exactly as-is, export it (or add as a pytest fixture if preferred), update both test files to import _random_inputs from the shared module, and remove the duplicate local implementations so tests use the single canonical helper.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py (1)
1079-1114: DuplicatePreStackedMoEModelclass across test files.An almost identical
PreStackedMoEModelclass exists intests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py(lines 214-243). The only differences are default parameter values (hidden_size=64vs32,num_experts=3vs4) and the 3D reshape in this version'sforward/get_input. Consider extracting a shared test utility to reduce duplication, though the different default dimensions may justify keeping them separate for now.🤖 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_moe_fusion.py` around lines 1079 - 1114, The PreStackedMoEModel class is duplicated across test files; extract it into a single shared test utility (e.g., a test helper module) and have both tests import and instantiate it with different default parameters as needed. Move the class definition (symbols: PreStackedMoEModel, methods __init__, forward, get_input) into a shared test utility module, parameterize hidden_size/num_experts/top_k/default dtype, adjust both test modules to import PreStackedMoEModel and override any differing defaults or reshape logic in their local test setup, and remove the duplicated class definitions from the individual test files.tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_qwen3_5_moe.py (2)
414-420: Loop variableexpert_idxis shadowed by reassignment.
expert_idxis reassigned on the next line, which shadows the loop variable. This is a minor code smell that violates the "avoid shadowing variables" guideline.Suggested fix
- for expert_idx in expert_hit: - expert_idx = expert_idx[0] - if expert_idx == num_experts: + for expert_entry in expert_hit: + expert_idx = expert_entry[0] + if expert_idx == num_experts:As per coding guidelines, "Avoid shadowing variables declared in an outer scope."
🤖 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/models/test_qwen3_5_moe.py` around lines 414 - 420, The loop currently reassigns the loop variable expert_idx (from expert_hit) with expert_idx = expert_idx[0], shadowing the original; change this to use a new local name (e.g., idx or expert_idx_tuple) or unpack the tensor result so you don't overwrite the loop variable: when iterating over expert_hit in the for loop, extract the scalar index into a fresh variable and use that (for example use idx = expert_idx[0] or unpack expert_idx_tuple) before comparing to num_experts and calling torch.where on expert_mask to compute top_k_pos and token_idx.
1146-1150: Unused unpacked variabledeltas.Static analysis flags
deltasas unused on lines 1148 and 1172. Prefix with underscore to indicate intentional discard.Suggested fix
- position_ids, deltas = model.get_rope_index(input_ids, image_grid_thw=image_grid_thw) + position_ids, _deltas = model.get_rope_index(input_ids, image_grid_thw=image_grid_thw)And similarly on line 1172:
- position_ids, deltas = model.get_rope_index(input_ids, attention_mask=attention_mask) + position_ids, _deltas = model.get_rope_index(input_ids, attention_mask=attention_mask)🤖 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/models/test_qwen3_5_moe.py` around lines 1146 - 1150, The test unpacks two values from model.get_rope_index but the second one is unused (deltas), so rename the unpacked variable to _deltas to indicate intentional discard; update the two occurrences where you do "position_ids, deltas = model.get_rope_index(...)" (including the call near the assertion and the later unpack at line ~1172) to "position_ids, _deltas = model.get_rope_index(...)" so static analysis no longer flags an unused variable.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gated_delta_rule_cache.py (1)
36-65: Significant code duplication withtest_torch_gated_delta_rule_cache.py.
DummyFactoryandGatedDeltaRuleModelare virtually identical between this file andtest_torch_gated_delta_rule_cache.py. The only meaningful differences between the two test files are the backend name, dtype, and tolerance values. Consider extracting the shared classes into a common test utility and parameterizing the test by backend.Also applies to: 64-123
🤖 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_gated_delta_rule_cache.py` around lines 36 - 65, Duplicate test helper classes (DummyFactory and GatedDeltaRuleModel) appear in test_gated_delta_rule_cache.py and test_torch_gated_delta_rule_cache.py; consolidate them into a shared test utility module (e.g., tests/unittest/_torch/auto_deploy/test_utils.py) and update both test files to import DummyFactory and GatedDeltaRuleModel from that module, then parameterize the tests by backend/dtype/tolerance (use pytest.mark.parametrize) so each test passes the backend name, dtype and tolerance values into the same test function instead of duplicating nearly identical files; update references to DummyFactory.get_cache_config_updates, DummyFactory.build_model, and GatedDeltaRuleModel to use the shared implementations.tensorrt_llm/_torch/auto_deploy/export/export.py (1)
187-192: Hardcoded arg index 3 and MOE op set — consider documenting the convention.The list-of-node arguments are assumed to start at index 3 (
for i in range(3, len(node.args))), and the MOE op set on Lines 174-178 is hardcoded. If a new MOE op is introduced with a different argument layout, this function will silently skip it. A brief comment explaining the assumed argument convention (e.g., "args[0:3] = hidden_states, selected_experts, routing_weights; args[3:] = per-expert weight lists") would help future maintainers.🤖 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 187 - 192, The code in list_arg_indices assumes list-of-node arguments start at index 3 (list_arg_indices and node.args) and that MOE ops follow a fixed op-set/layout; update the export logic to document and guard this convention: add a clear inline comment describing the expected argument convention (e.g., "args[0:3] = hidden_states, selected_experts, routing_weights; args[3:] = per-expert weight/scale lists") next to the list_arg_indices computation, and add a lightweight validation (e.g., an assertion or warning using the node/op identifier) when the node's arg count or types don't match the convention so new MOE ops with different layouts are flagged rather than silently skipped; keep references to list_arg_indices and node.args in the added comment and validation.tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py (1)
1381-1382: Variableidsis shadowed within the loop body.Line 1382 reassigns the loop variable
ids(fromfor i, ids in enumerate(total_input_ids)) to a masked version. While functionally correct, this shadows the outer variable and can cause confusion during debugging.Proposed fix
for i, ids in enumerate(total_input_ids): - ids = ids[attention_mask[i] == 1] + active_ids = ids[attention_mask[i] == 1] - vision_start_indices = torch.argwhere(ids == vision_start_token_id).squeeze(1) - vision_tokens = ids[vision_start_indices + 1] + vision_start_indices = torch.argwhere(active_ids == vision_start_token_id).squeeze(1) + vision_tokens = active_ids[vision_start_indices + 1] image_nums = int((vision_tokens == image_token_id).sum().item()) video_nums = int((vision_tokens == video_token_id).sum().item()) - input_tokens = ids.tolist() + input_tokens = active_ids.tolist()As per coding guidelines: "Avoid shadowing variables declared in an outer scope."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py` around lines 1381 - 1382, The loop over total_input_ids shadows the loop variable ids by reassigning it with a masked version (ids = ids[attention_mask[i] == 1]); instead create a new variable name (e.g., masked_ids or filtered_ids) and use that within the loop to hold the masked tensor, leaving the original loop variable ids intact; update any subsequent references in the loop body from ids to the new name and keep references to total_input_ids and attention_mask as-is.
26ef3d9 to
fcec2aa
Compare
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #36129 [ run ] triggered by Bot. Commit: |
|
PR_Github #36129 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36133 [ run ] triggered by Bot. Commit: |
|
PR_Github #36133 [ run ] completed with state
|
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
2881ac9 to
abba963
Compare
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36199 [ run ] triggered by Bot. Commit: |
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36206 [ run ] triggered by Bot. Commit: |
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36213 [ run ] triggered by Bot. Commit: |
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36215 [ run ] triggered by Bot. Commit: |
|
PR_Github #36215 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --reuse-test |
|
PR_Github #36270 [ run ] triggered by Bot. Commit: |
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #36272 [ run ] triggered by Bot. Commit: |
|
PR_Github #36272 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --reuse-test |
|
PR_Github #36300 [ run ] triggered by Bot. Commit: |
|
PR_Github #36300 [ run ] completed with state |
Summary by CodeRabbit
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.