diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py index 7a21c4961ca..62328ccf309 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py @@ -87,6 +87,10 @@ def _insert_fused_moe_ops(gm: GraphModule, backend: Literal["auto", "trtllm", "t graph.get_attr(new_key_w_up), graph.get_attr(new_key_w_down), ), + kwargs={ + "mlp_style": mlp_style_val, + "act_fn": act_fn_val, + }, ) node.replace_all_uses_with(new_node)