From 89daace0b029975b685ea990bd447db3bec61486 Mon Sep 17 00:00:00 2001 From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Date: Thu, 4 Jun 2026 04:07:38 -0700 Subject: [PATCH] [None][feat] AutoDeploy: propagate layer_type hint across pattern rewrites AutoDeploy pattern matchers that collapse N fine-grained ops into one (e.g. the SwiGLU matchers) dropped the per-op layer_type sharding hint. The hint-driven sharder (apply_sharding_hints / shard_layers) then read layer_type=None on the fused op and, because the whitelist is fail-open, TP-sharded nodes meant to be replicated -- the root cause of the Qwen3.5-MoE NVFP4 shared-expert accuracy regression fixed point-wise in #14667. Generalize the fix into the matcher infra: add an optional layer_type arg to the intermediate SwiGLU ops, add collect_classification_hints / stamp_hints helpers in node_utils, and propagate the consensus layer_type from the matched nodes onto the replacement op in ADReplacementPatternEntry. Per-weight mechanics (tp_mode, output_sizes, ...) are intentionally not propagated: the fused op's ShardableNode re-derives them structurally and they legitimately differ across constituents (an MLA layer mixes none/colwise/rowwise yet shares one layer_type), so generic propagation of the sharding decision is ill-posed. Tested: GPU-free unit tests for the helpers (consensus / conflict-drop / schema-guarded stamping), an NVFP4 end-to-end regression mirroring the Qwen3.5-MoE bug, and the full transformations/library suite (329 passed). Fixes #14834 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> --- .../auto_deploy/custom_ops/linear/swiglu.py | 9 +++ .../_torch/auto_deploy/utils/node_utils.py | 73 +++++++++++++++++++ .../auto_deploy/utils/pattern_matcher.py | 17 +++++ .../library/test_fuse_swiglu.py | 72 +++++++++++++++++- .../library/test_nvfp4_swiglu.py | 44 ++++++++++- 5 files changed, 213 insertions(+), 2 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 92bdd51e6aba..a13ddd1de8c4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py +++ b/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]. diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index 286242ce5228..f80d55f18870 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/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, diff --git a/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py b/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py index fbd12dc7b147..f508c1a324d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py +++ b/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,6 +128,18 @@ 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, @@ -134,6 +147,10 @@ def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> Non 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 diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py index f3377e4b2e08..f272d05d6e55 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py +++ b/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" diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py index d89de5051688..765d7101aec5 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py +++ b/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}"