[https://nvbugs/6221450][fix] AutoDeploy: Qwen3.5 400B NVFP4 accuracy regression fix (#14667)

Signed-off-by: Taylor Yeonbok Lee <249374542+taylor-yb-lee@users.noreply.github.com>

Author: taylor-yb-lee

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

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]

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

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()

tests/integration/test_lists/waives.txt

@@ -19,7 +19,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)