[None][feat] AutoDeploy: propagate layer_type hint across pattern-matcher rewrites

tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py

@@ -68,6 +68,9 @@ def torch_swiglu_mlp(
         gate_bias: Optional gate projection bias of shape [intermediate_size].
         up_bias: Optional up projection bias of shape [intermediate_size].
         down_bias: Optional down projection bias of shape [hidden_size].
+        layer_type: Layer-classification sharding hint (e.g. "mlp"/"moe"/"shared_expert"),
+            propagated from the matched linears by the pattern matcher and consumed by
+            ``apply_sharding_hints`` (``shard_layers``). Does not affect the numeric result.
 
     Returns:
         Output tensor of shape [..., hidden_size].
@@ -184,6 +187,9 @@ def torch_nvfp4_swiglu_mlp(
         down_input_scale: Input scale for down projection.
         down_weight_scale: Per-block weight scale for down projection.
         down_alpha: Alpha (combined scale) for down projection.
+        layer_type: Layer-classification sharding hint (e.g. "mlp"/"moe"/"shared_expert"),
+            propagated from the matched linears by the pattern matcher and consumed by
+            ``apply_sharding_hints`` (``shard_layers``). Does not affect the numeric result.
 
     Returns:
         Output tensor of shape [..., hidden_size].
@@ -344,6 +350,9 @@ def torch_finegrained_fp8_swiglu_mlp(
         gate_weight_scale: Per-block weight scale for gate [N/128, K/128] float32.
         up_weight_scale: Per-block weight scale for up [N/128, K/128] float32.
         down_weight_scale: Per-block weight scale for down [N/128, K/128] float32.
+        layer_type: Layer-classification sharding hint (e.g. "mlp"/"moe"/"shared_expert"),
+            propagated from the matched linears by the pattern matcher and consumed by
+            ``apply_sharding_hints`` (``shard_layers``). Does not affect the numeric result.
 
     Returns:
         Output tensor of shape [..., hidden_size].

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -1339,6 +1339,79 @@ def set_op_args(node: Node, **name_value_pairs) -> None:
     node.kwargs = kwargs
 
 
+# Classification hints are layer-level: they are invariant across the fine-grained ops that
+# make up one logical layer (e.g. all projections of a SwiGLU MLP share the same
+# ``layer_type``) and are consumed by policy filters such as ``shard_layers`` -- NOT by
+# per-weight sharding mechanics. They are therefore the only sharding-related kwargs that are
+# well-defined to carry onto a fused/replacement op produced by an N->1 pattern rewrite (by
+# consensus). Per-weight mechanics (``tp_mode``, ``output_sizes``, ``tp_min_local_shape``,
+# ``tp_scaled_dim``, ``enable_sharding``) are intentionally NOT propagated: a fused op's
+# ShardableNode re-derives those structurally, so copying them across a rewrite is ill-defined
+# (the constituents legitimately disagree -- e.g. an MLA layer mixes ``tp_mode`` none/colwise/
+# rowwise while sharing a single ``layer_type``).
+CLASSIFICATION_HINT_NAMES = frozenset({"layer_type"})
+
+
+def _op_schema_arg_names(node: Node) -> set:
+    """Return the argument names declared by a call_function node's op schema.
+
+    Returns an empty set for non-call_function nodes or ops without an introspectable
+    schema, so callers can use it as a safe membership test.
+    """
+    if not isinstance(node, Node) or node.op != "call_function":
+        return set()
+    try:
+        return {a.name for a in _get_op_schema(node).arguments}
+    except (ValueError, RuntimeError):
+        return set()
+
+
+def collect_classification_hints(nodes: Iterable[Node]) -> dict:
+    """Return a consensus value for each classification hint across ``nodes``.
+
+    For every name in :data:`CLASSIFICATION_HINT_NAMES`, scan the call_function nodes that
+    declare it and collect the distinct *meaningful* values (ignoring ``None`` and the
+    ``"unknown"`` default). A name is included in the result only when exactly one such
+    value is observed; conflicting values are dropped with a warning, since a conflict
+    means the caller grouped nodes that belong to different logical layers.
+    """
+    result: dict = {}
+    for name in CLASSIFICATION_HINT_NAMES:
+        values = set()
+        for n in nodes:
+            if name not in _op_schema_arg_names(n):
+                continue
+            [value] = extract_op_args(n, name)
+            if value is not None and value != "unknown":
+                values.add(value)
+        if len(values) == 1:
+            result[name] = next(iter(values))
+        elif len(values) > 1:
+            ad_logger.warning(
+                f"Conflicting '{name}' hints {sorted(values)} among matched nodes; "
+                "not propagating to the replacement op (matched nodes may span layers)."
+            )
+    return result
+
+
+def stamp_hints(nodes: Iterable[Node], hints: dict) -> int:
+    """Set ``hints`` on every node in ``nodes`` whose op schema declares them.
+
+    Returns the number of nodes updated. Each hint is applied only to nodes whose op
+    actually declares that argument, so passing a heterogeneous node list is safe.
+    """
+    if not hints:
+        return 0
+    count = 0
+    for n in nodes:
+        names = _op_schema_arg_names(n)
+        to_set = {k: v for k, v in hints.items() if k in names}
+        if to_set:
+            set_op_args(n, **to_set)
+            count += 1
+    return count
+
+
 def predecessors(
     node: Node,
     depth: int = 1,

tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py

@@ -55,6 +55,7 @@
 from torch.fx import GraphModule
 
 from ..export import torch_export_to_gm
+from .node_utils import collect_classification_hints, stamp_hints
 
 
 @contextlib.contextmanager
@@ -127,13 +128,29 @@ def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> Non
         del node
         assert match.replacement_graph is not None
         output_nodes = match.output_nodes()
+
+        # Carry layer-level classification hints (currently ``layer_type``) from the matched
+        # nodes onto the replacement op(s). This is the only sharding-related metadata that is
+        # well-defined to propagate across an N->1 rewrite: it is invariant across the
+        # fine-grained ops of one logical layer and is consumed by the downstream hint-driven
+        # sharder (``apply_sharding_hints`` / ``shard_layers``). Per-weight mechanics (tp_mode,
+        # output_sizes, ...) are intentionally NOT carried -- the replacement op's ShardableNode
+        # re-derives them structurally. Collect before the rewrite (cheap; reads only the
+        # matched nodes) so the node-set diff below is paid only when there is a hint to carry.
+        class_hints = collect_classification_hints(match.nodes)
+        nodes_before = set(graph.nodes) if class_hints else None
+
         self.replace_with_graph(
             match,
             graph,
             match.replacement_graph,
             self.normalize_args(*match.args, **match.kwargs),
         )
 
+        if class_hints:
+            inserted_nodes = [n for n in graph.nodes if n not in nodes_before]
+            stamp_hints(inserted_nodes, class_hints)
+
         if len(output_nodes) > 1:
             # Torch's generic replacement path inserts the copied replacement graph relative to the
             # earliest matched output node. That is usually fine for single-output rewrites, but it

tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py

@@ -19,7 +19,12 @@
 from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import *  # noqa
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
-from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
+from tensorrt_llm._torch.auto_deploy.utils.node_utils import (
+    collect_classification_hints,
+    extract_op_args,
+    is_op,
+    stamp_hints,
+)
 
 
 class SwiGLUMLP(torch.nn.Module):
@@ -200,3 +205,68 @@ def test_swiglu_pattern_match_only():
     y_matched = gm_matched(x)
     y_model = model(x)
     torch.testing.assert_close(y_matched, y_model, atol=1e-3, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# layer_type (classification-hint) propagation across pattern rewrites
+#
+# These exercise the generic helpers used by the matcher infra to carry the
+# layer-level ``layer_type`` hint from the matched fine-grained ops onto a fused
+# replacement op (so downstream ``apply_sharding_hints`` / ``shard_layers`` still
+# sees it). They are graph-only and need no GPU. The end-to-end behaviour through
+# an actual matcher is covered on the NVFP4 path in ``test_nvfp4_swiglu.py``
+# (the BF16 SwiGLU matcher never fuses hint-carrying linears, since ``torch.export``
+# materializes the hints positionally, so there is no hint to lose there).
+# ---------------------------------------------------------------------------
+
+
+def test_collect_classification_hints_consensus_over_mixed_mechanics():
+    """Consensus recovers the shared layer_type even when constituents disagree on tp_mode."""
+    g = torch.fx.Graph()
+    x = g.placeholder("x")
+    nodes = []
+    for tp in ("colwise", "colwise", "rowwise"):  # gate/up colwise, down rowwise
+        nodes.append(
+            g.call_function(
+                torch.ops.auto_deploy.torch_linear_simple.default,
+                args=(x, x, None),
+                kwargs={"tp_mode": tp, "layer_type": "shared_expert"},
+            )
+        )
+    # A node left at the default ("unknown") must not contribute to the consensus.
+    nodes.append(
+        g.call_function(torch.ops.auto_deploy.torch_linear_simple.default, args=(x, x, None))
+    )
+    assert collect_classification_hints(nodes) == {"layer_type": "shared_expert"}
+
+
+def test_collect_classification_hints_conflict_is_dropped():
+    """Conflicting layer_type values (a rewrite spanning layers) are dropped, not guessed."""
+    g = torch.fx.Graph()
+    x = g.placeholder("x")
+    a = g.call_function(
+        torch.ops.auto_deploy.torch_linear_simple.default,
+        args=(x, x, None),
+        kwargs={"layer_type": "moe"},
+    )
+    b = g.call_function(
+        torch.ops.auto_deploy.torch_linear_simple.default,
+        args=(x, x, None),
+        kwargs={"layer_type": "mla"},
+    )
+    assert collect_classification_hints([a, b]) == {}
+
+
+def test_stamp_hints_only_on_declaring_ops():
+    """stamp_hints sets a hint only on ops whose schema declares it (aten.silu is skipped)."""
+    g = torch.fx.Graph()
+    x = g.placeholder("x")
+    swiglu = g.call_function(
+        torch.ops.auto_deploy.torch_swiglu_mlp.default,
+        args=(x, x, x, x, None, None, None),
+    )
+    silu = g.call_function(torch.ops.aten.silu.default, args=(x,))
+
+    assert stamp_hints([swiglu, silu], {"layer_type": "shared_expert"}) == 1
+    [lt] = extract_op_args(swiglu, "layer_type")
+    assert lt == "shared_expert"

tests/unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py

@@ -28,7 +28,7 @@
 import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
-from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
+from tensorrt_llm._torch.auto_deploy.utils.node_utils import extract_op_args, is_op, set_op_args
 from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale
 
 _skip_reason = "Requires NVFP4 (Blackwell+) and TRT-LLM ops"
@@ -390,3 +390,45 @@ def forward(self, x):
     assert _count_ops(gm_result, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 2, (
         "Original NVFP4 linear ops should be unchanged"
     )
+
+
+@pytest.mark.skipif(_skip_condition, reason=_skip_reason)
+def test_nvfp4_swiglu_pattern_propagates_layer_type():
+    """Regression guard for the Qwen3.5-MoE NVFP4 accuracy bug.
+
+    The SwiGLU matcher collapses the three NVFP4 linears into one
+    ``torch_nvfp4_swiglu_mlp`` op. It must carry the constituents' ``layer_type`` hint onto
+    that fused op so the downstream hint-driven sharder (``apply_sharding_hints`` /
+    ``shard_layers``) can still exclude (replicate) the shared expert. Before the fix the
+    fused op carried no ``layer_type`` and the fail-open whitelist TP-sharded it, corrupting
+    the shared-expert output.
+
+    ``layer_type`` is attached as a kwarg here exactly as ``quantize_nvfp4_linear_from_config``
+    does in production (the matcher ignores it for matching but it must survive the rewrite).
+    """
+    torch.manual_seed(0)
+    model = NVFP4SwiGLUMLP().to("cuda")
+    x = torch.randn(2, 128, device="cuda", dtype=torch.float16)
+
+    gm = torch_export_to_gm(model, args=(x,), clone=True)
+
+    # Tag every NVFP4 linear with a layer_type hint (as the quantization transform does).
+    for n in gm.graph.nodes:
+        if is_op(n, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear):
+            set_op_args(n, layer_type="shared_expert")
+    gm.recompile()
+
+    gm_matched = InferenceOptimizer(
+        None,
+        {"match_nvfp4_swiglu_pattern": {"stage": "pattern_matcher"}},
+    )(None, gm)
+
+    swiglu_nodes = [
+        n
+        for n in gm_matched.graph.nodes
+        if is_op(n, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default)
+    ]
+    assert len(swiglu_nodes) == 1, f"expected 1 torch_nvfp4_swiglu_mlp, got {len(swiglu_nodes)}"
+
+    [lt] = extract_op_args(swiglu_nodes[0], "layer_type")
+    assert lt == "shared_expert", f"layer_type was not propagated onto the fused op: got {lt!r}"