From 6fbc7380a5021fc2adc502eccf1460c769dd4d75 Mon Sep 17 00:00:00 2001 From: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Date: Wed, 27 May 2026 21:04:50 -0700 Subject: [PATCH 1/5] Try unwaive Qwen3.5 Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 30da247dcaa2..f66281984f97 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -17,7 +17,6 @@ accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_ accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] SKIP (https://nvbugs/6200112) accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-8] SKIP (https://nvbugs/6248757) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397) -accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6211441) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6191524) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/6084775) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_mtp] SKIP (https://nvbugs/6029882) From 49bba741d94141d92ded2157af2e5122166ec64e Mon Sep 17 00:00:00 2001 From: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Date: Thu, 28 May 2026 22:25:10 -0700 Subject: [PATCH 2/5] [None][fix] AutoDeploy: add shard_exclude_filter to replicate selected nodes under IR sharding Add a shard_exclude_filter config so nodes whose weight keys match (e.g. the Qwen3.5-MoE shared expert) are replicated instead of TP-sharded, and reorder the MoE all-reduce so the replicated shared output is added after it (not scaled by world size). Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> --- .../models/custom/modeling_qwen3_5_moe.py | 5 ++++- .../transform/library/sharding_ir.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index b53170f46470..5d63dfb0a952 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py @@ -765,8 +765,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: layer_type="moe", ) - expert_output = expert_output + shared_expert_output + # The shared expert is replicated (excluded from TP sharding), so all-reduce + # the sharded routed-expert output first, then add the replicated shared + # output; adding before would scale it by the TP world size. expert_output = torch.ops.auto_deploy.all_reduce(expert_output, layer_type="moe") + expert_output = expert_output + shared_expert_output expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim) return expert_output diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py index 0e49da0b52cf..78a7af72ee3f 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py @@ -889,6 +889,11 @@ class IRShardingConfig(TransformConfig): default=None, description="When set, only shard nodes whose layer_type hint is in this list.", ) + shard_exclude_filter: Optional[List[str]] = Field( + default=None, + description="Substrings matched against a node's weight parameter keys; matching " + "nodes are left replicated (not TP-sharded), e.g. ['shared_expert'].", + ) enable_attention_dp: bool = Field(default=False) dist_mapping: dict[str, int] = Field(default_factory=dict) dist_config: DistConfig = Field(default_factory=DistConfig) @@ -921,6 +926,13 @@ def _init_dist_config(self, rank: int, world_size: int): # ============================================================================= +def _node_weight_keys_match(node: Node, patterns: List[str]) -> bool: + """True if any of *node*'s weight/bias/scale param keys contains a pattern substring.""" + wn = extract_weight_nodes(node) + keys = [n.node_key for n in (*wn.weights, *wn.biases, *wn.scales)] + return any(pat in key for key in keys for pat in patterns) + + def _log_sharding_prelude(dc: DistConfig) -> None: """Log the sharding configuration before apply_sharding_hints runs.""" skip = " (skipping)" if dc.tp_size <= 1 else "" @@ -1092,6 +1104,7 @@ def _apply( _log_sharding_result(dc, num_updates) else: shard_layers = self.config.shard_layers + exclude_filter = self.config.shard_exclude_filter num_skipped = 0 all_dead_nodes = [] @@ -1109,6 +1122,11 @@ def _apply( if lt is not None and lt not in shard_layers: num_skipped += 1 continue + # Replicate (skip sharding) nodes excluded by config, e.g. the + # shared expert, which is cheaper replicated than TP-sharded. + if exclude_filter and _node_weight_keys_match(node, exclude_filter): + num_skipped += 1 + continue num_updates += shardable_node.apply(gm, dc, max_num_tokens) From f3f596682ed01137bf2d22456e9909664a9e6a5e Mon Sep 17 00:00:00 2001 From: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Date: Thu, 28 May 2026 22:25:36 -0700 Subject: [PATCH 3/5] [None][fix] AutoDeploy: replicate Qwen3.5-MoE shared expert + SYMM_MEM all-reduce Exclude the shared expert from TP sharding via shard_exclude_filter (cheaper replicated at this size), set SYMM_MEM all-reduce on the IR path, and drop the now-dead detect_sharding block. test_nvfp4[8] MMLU 86.4% (was ~27%). Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> --- .../configs/qwen3.5_moe_400b.yaml | 33 ++++--------------- 1 file changed, 7 insertions(+), 26 deletions(-) diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml index 70d0d7f41a0f..e69278a32c0d 100644 --- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml +++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml @@ -29,33 +29,14 @@ transforms: enabled: true fuse_nvfp4_moe: backend: trtllm_gen - detect_sharding: - # for long input, tp8ep1 gives better performance - # dist_mapping: {moe_tp: 8, moe_ep: 1} + # Sharding is driven by the hint-based IR path (enable_sharder_ir.yaml); TP/EP + # plans come from the in-model sharding hints, so no detect_sharding config is + # needed. Only runtime distribution options are set here. + apply_sharding_hints: + # Symmetric-memory all-reduce (MULTIMEM); matches the pre-IR sharding perf. allreduce_strategy: SYMM_MEM - shard_all_unprocessed: true - simple_shard_filter: "lm_head" - sharding_dims: ['tp','ep', 'bmm'] - # use only manual config for TP sharding - sharding_source: ['manual'] - manual_config: - tp_plan: - # GDN layer - "in_proj_qkv": "delta" - # attention layer - "q_proj": "colwise" - "k_proj": "colwise" - "v_proj": "colwise" - "o_proj": "rowwise" - # lm_head: "gather" = column split + all_gather (not "colwise" which - # requires a LayerSubgraph and crashes for standalone unprocessed nodes) - "lm_head": "gather" - # replicating shared experts (keep them commented out) - # "shared_expert_gate_proj": "colwise" - # "shared_expert_up_proj": "colwise" - # "shared_expert_down_proj": "rowwise" - # gating layer should be replicated as well - # "gate": "gather" + # Replicate the shared expert instead of TP-sharding it (cheaper at this size). + shard_exclude_filter: ["shared_expert"] multi_stream_moe: stage: compile enabled: true From dda311fe97e6cf8b80a836834db4d84facb10860 Mon Sep 17 00:00:00 2001 From: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Date: Fri, 29 May 2026 20:58:13 -0700 Subject: [PATCH 4/5] [None][fix] AutoDeploy: replicate Qwen3.5-MoE shared expert via shard_layers whitelist - Tag shared-expert linears layer_type="shared_expert" (was "moe"). - Drop shard_exclude_filter blacklist; reuse existing shard_layers whitelist. - yaml: shard_layers=["moe","delta","mha"] omits shared_expert -> replicated. Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> --- .../configs/qwen3.5_moe_400b.yaml | 8 ++------ .../models/custom/modeling_qwen3_5_moe.py | 6 +++--- .../transform/library/sharding_ir.py | 18 ------------------ 3 files changed, 5 insertions(+), 27 deletions(-) diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml index e69278a32c0d..adef5dc45183 100644 --- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml +++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml @@ -29,14 +29,10 @@ transforms: enabled: true fuse_nvfp4_moe: backend: trtllm_gen - # Sharding is driven by the hint-based IR path (enable_sharder_ir.yaml); TP/EP - # plans come from the in-model sharding hints, so no detect_sharding config is - # needed. Only runtime distribution options are set here. apply_sharding_hints: - # Symmetric-memory all-reduce (MULTIMEM); matches the pre-IR sharding perf. allreduce_strategy: SYMM_MEM - # Replicate the shared expert instead of TP-sharding it (cheaper at this size). - shard_exclude_filter: ["shared_expert"] + # Shared expert is excluded from sharding for performance purpose + shard_layers: ["moe", "delta", "mha"] multi_stream_moe: stage: compile enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index 5d63dfb0a952..dcb530f0f433 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py @@ -630,21 +630,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.gate_proj.weight, self.gate_proj.bias, tp_mode="colwise", - layer_type="moe", + layer_type="shared_expert", ) up = torch.ops.auto_deploy.torch_linear_simple( x, self.up_proj.weight, self.up_proj.bias, tp_mode="colwise", - layer_type="moe", + layer_type="shared_expert", ) return torch.ops.auto_deploy.torch_linear_simple( self.act_fn(gate) * up, self.down_proj.weight, self.down_proj.bias, tp_mode="rowwise", - layer_type="moe", + layer_type="shared_expert", ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py index 78a7af72ee3f..0e49da0b52cf 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py @@ -889,11 +889,6 @@ class IRShardingConfig(TransformConfig): default=None, description="When set, only shard nodes whose layer_type hint is in this list.", ) - shard_exclude_filter: Optional[List[str]] = Field( - default=None, - description="Substrings matched against a node's weight parameter keys; matching " - "nodes are left replicated (not TP-sharded), e.g. ['shared_expert'].", - ) enable_attention_dp: bool = Field(default=False) dist_mapping: dict[str, int] = Field(default_factory=dict) dist_config: DistConfig = Field(default_factory=DistConfig) @@ -926,13 +921,6 @@ def _init_dist_config(self, rank: int, world_size: int): # ============================================================================= -def _node_weight_keys_match(node: Node, patterns: List[str]) -> bool: - """True if any of *node*'s weight/bias/scale param keys contains a pattern substring.""" - wn = extract_weight_nodes(node) - keys = [n.node_key for n in (*wn.weights, *wn.biases, *wn.scales)] - return any(pat in key for key in keys for pat in patterns) - - def _log_sharding_prelude(dc: DistConfig) -> None: """Log the sharding configuration before apply_sharding_hints runs.""" skip = " (skipping)" if dc.tp_size <= 1 else "" @@ -1104,7 +1092,6 @@ def _apply( _log_sharding_result(dc, num_updates) else: shard_layers = self.config.shard_layers - exclude_filter = self.config.shard_exclude_filter num_skipped = 0 all_dead_nodes = [] @@ -1122,11 +1109,6 @@ def _apply( if lt is not None and lt not in shard_layers: num_skipped += 1 continue - # Replicate (skip sharding) nodes excluded by config, e.g. the - # shared expert, which is cheaper replicated than TP-sharded. - if exclude_filter and _node_weight_keys_match(node, exclude_filter): - num_skipped += 1 - continue num_updates += shardable_node.apply(gm, dc, max_num_tokens) From 4fa4379d9cd7ae36e04f231d26f74e8fa516f147 Mon Sep 17 00:00:00 2001 From: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> Date: Sat, 30 May 2026 13:55:24 -0700 Subject: [PATCH 5/5] [None][fix] AutoDeploy: preserve layer_type hint across SwiGLU fusion - SwiGLU fusion dropped the layer_type op-arg, so hint-driven sharding (shard_layers) saw lt=None and TP-sharded nodes meant to be replicated. - Add optional layer_type to the unfused SwiGLU ops and propagate the source linears' layer_type onto the fused node in the matchers. gate/up/down linear(layer_type=X) --fuse--> swiglu_mlp() [hint lost] gate/up/down linear(layer_type=X) --fuse--> swiglu_mlp(layer_type=X) [kept] Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com> --- .../auto_deploy/custom_ops/linear/swiglu.py | 6 ++ .../transform/library/fuse_swiglu.py | 60 +++++++++++++++++-- 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py index d28845d2e30b..92bdd51e6aba 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py @@ -51,6 +51,7 @@ def torch_swiglu_mlp( gate_bias: Optional[torch.Tensor], up_bias: Optional[torch.Tensor], down_bias: Optional[torch.Tensor], + layer_type: str = "unknown", ) -> torch.Tensor: """Standardized SwiGLU MLP operation. @@ -86,6 +87,7 @@ def _( gate_bias: Optional[torch.Tensor], up_bias: Optional[torch.Tensor], down_bias: Optional[torch.Tensor], + layer_type: str = "unknown", ) -> torch.Tensor: """Fake implementation for tracing.""" # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] @@ -159,6 +161,7 @@ def torch_nvfp4_swiglu_mlp( down_input_scale: torch.Tensor, down_weight_scale: torch.Tensor, down_alpha: torch.Tensor, + layer_type: str = "unknown", ) -> torch.Tensor: """NVFP4 quantized SwiGLU MLP operation (intermediate representation). @@ -230,6 +233,7 @@ def _( down_input_scale: torch.Tensor, down_weight_scale: torch.Tensor, down_alpha: torch.Tensor, + layer_type: str = "unknown", ) -> torch.Tensor: """Fake implementation for tracing.""" # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] @@ -323,6 +327,7 @@ def torch_finegrained_fp8_swiglu_mlp( gate_weight_scale: torch.Tensor, up_weight_scale: torch.Tensor, down_weight_scale: torch.Tensor, + layer_type: str = "unknown", ) -> torch.Tensor: """FineGrained FP8 quantized SwiGLU MLP operation (intermediate representation). @@ -382,6 +387,7 @@ def _( gate_weight_scale: torch.Tensor, up_weight_scale: torch.Tensor, down_weight_scale: torch.Tensor, + layer_type: str = "unknown", ) -> torch.Tensor: """Fake implementation for tracing.""" # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py index 700cb3693b27..781910fa8928 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py @@ -22,6 +22,7 @@ The SwiGLU pattern is: silu(x @ gate.T) * (x @ up.T) @ down.T """ +from contextlib import contextmanager from typing import Tuple, Type import torch @@ -47,7 +48,7 @@ eliminate_dead_code, get_attr_by_name, ) -from ...utils.node_utils import is_op +from ...utils.node_utils import extract_op_args, is_op, set_op_args from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern from ...utils.quantization_utils import ensure_tma_col_major from ..interface import ( @@ -59,6 +60,42 @@ ) +def _weight_key(node: Node): + """Stable key for a weight arg: the get_attr target FQN (survives node re-creation + during pattern replacement), falling back to node identity. ``args[1]`` is the + weight for both linear and SwiGLU ops.""" + w = node.args[1] if len(node.args) > 1 else None + if not isinstance(w, Node): + return None + return w.target if w.op == "get_attr" else w + + +@contextmanager +def preserve_layer_types(gm: GraphModule, linear_op, fused_op): + """Carry the ``layer_type`` hint across a fusion that consumes ``linear_op`` nodes + and emits ``fused_op`` nodes (which would otherwise drop the hint). + + Snapshots each source weight's ``layer_type`` before the rewrite, then re-applies + it to the fused node keyed by weight, so hint-driven sharding (``shard_layers``) + can still classify the fused node. Wrap the matcher's ``patterns.apply`` call. + """ + wmap = {} + for n in gm.graph.nodes: + if is_op(n, linear_op): + [lt] = extract_op_args(n, "layer_type") + key = _weight_key(n) + if lt is not None and key is not None: + wmap[key] = lt + yield + if not wmap: + return + for n in gm.graph.nodes: + if is_op(n, fused_op): + key = _weight_key(n) + if key is not None and key in wmap: + set_op_args(n, layer_type=wmap[key]) + + def _maybe_to_deepgemm_layout( weight: torch.Tensor, scale: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -242,7 +279,12 @@ def _apply( dummy_args=dummy_args_with_bias, ) - num_matches = patterns.apply(gm.graph) + with preserve_layer_types( + gm, + torch.ops.auto_deploy.torch_linear_simple.default, + torch.ops.auto_deploy.torch_swiglu_mlp.default, + ): + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile() @@ -554,7 +596,12 @@ def _apply( dummy_args=dummy_args, ) - num_matches = patterns.apply(gm.graph) + with preserve_layer_types( + gm, + torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default, + torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default, + ): + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile() @@ -835,7 +882,12 @@ def _apply( dummy_args=dummy_args, ) - num_matches = patterns.apply(gm.graph) + with preserve_layer_types( + gm, + torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default, + torch.ops.auto_deploy.torch_finegrained_fp8_swiglu_mlp.default, + ): + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile()