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..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,33 +29,10 @@ 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} + apply_sharding_hints: 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" + # 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/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/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index b53170f46470..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", ) @@ -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/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() 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)