[Refactor] Extract model specific logics in export lib

modelopt/torch/export/layer_utils.py

@@ -60,6 +60,13 @@
     RgLruConfig,
 )
 from .model_config_utils import pad_weights
+from .modeling import (
+    NULL_HOOKS,
+    ExportContext,
+    match_by_decoder_type,
+    match_mlp_block,
+    match_moe_block,
+)
 from .postprocess import view_as_float8_e4m3fn_if_needed, view_as_uint8_if_needed
 from .quant_utils import (
     get_activation_scaling_factor,
@@ -93,25 +100,14 @@ def get_experts_list(
     """
     experts_list = []
 
-    # Define linear layer names for different model types
-    if "mixtralforcausallm" in model_type:
-        linear_names = ["w1", "w2", "w3"]
-    elif any(
-        qwen_variant in model_type
-        for qwen_variant in [
-            "qwenmoeforcausallm",
-            "qwen2moeforcausallm",
-            "qwen3moeforcausallm",
-            "qwen3nextforcausallm",
-        ]
-    ):
-        linear_names = ["gate_proj", "down_proj", "up_proj"]
-    elif "nemotronhforcausallm" in model_type:
-        linear_names = ["up_proj", "down_proj"]
-    elif "gemma4" in model_type:
-        linear_names = ["gate_proj", "down_proj", "up_proj"]
-    else:
+    # Expert linear names live in modeling/families/*. The grouped export path only
+    # supports families whose experts are iterable per-expert sub-modules (Mixtral,
+    # Qwen MoE, NemotronH, Gemma4); stacked/fused layouts (DBRX, GptOss, ...) raise
+    # NotImplementedError here and are handled by other paths.
+    spec = match_moe_block(module)
+    if spec is None or not spec.has_iterable_experts or spec.expert_linear_names is None:
         raise NotImplementedError(f" {model_type} not supported")
+    linear_names = list(spec.expert_linear_names)
 
     # Common logic for all supported model types
     experts_list.extend(
@@ -770,26 +766,18 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer):
         "c_fc_1",  # exaone
     }
 
-    # Arctic (llama-based MoE, decoder_type is "llama") has MLP keyword conflicts with Qwen
-    # Arctic's residual MLP use w1 for fc, w2 for proj, w3 for gate
-    if type(module).__name__ in ["ArcticMLP", "InternLM2MLP"]:
-        fc_keywords.discard("w2")
-        gate_keywords.discard("w1")
-        fc_keywords.add("w1")
-        proj_keywords.add("w2")
-        gate_keywords.add("w3")
-
-    if decoder_type == "mpt":
-        fc_keywords.add("up_proj")
-        gate_keywords.discard("up_proj")
-
-    if type(module).__name__ in [
-        "TLGv4MLP",
-        "Phi3SmallMLP",
-        "NemotronMLP",
-    ]:  # for TLGv4ForCausalLM
-        fc_keywords.add("up_proj")
-        gate_keywords.discard("up_proj")
+    # Per-family MLP keyword role overrides live in modeling/families/* (e.g. Arctic/
+    # InternLM2 w1/w2/w3 remap; MPT/Phi3Small/TLGv4/Nemotron up_proj→fc). Resolve by the
+    # MLP module class and by decoder_type; apply on top of the shared keyword sets.
+    _role_sets = {"fc": fc_keywords, "gate": gate_keywords, "proj": proj_keywords}
+    for _spec in (match_mlp_block(module), match_by_decoder_type(decoder_type)):
+        if _spec is None or not _spec.mlp_keyword_roles:
+            continue
+        for _kw, _role in _spec.mlp_keyword_roles.items():
+            fc_keywords.discard(_kw)
+            gate_keywords.discard(_kw)
+            proj_keywords.discard(_kw)
+            _role_sets[_role].add(_kw)
 
     fc_linear: nn.Module = None
     gate_linear: nn.Module = None
@@ -865,11 +853,11 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer):
 
     assert config.proj is not None and config.fc is not None, "proj or fc can not be found"
 
-    # Override hidden_act based on decoder_type
-    if decoder_type in ["bloom", "glm"]:
-        hidden_act = "gelu"
-    if decoder_type == "phi3":
-        hidden_act = "swiglu"
+    # Per-family activation override lives in modeling/families/* (e.g. bloom/glm → gelu,
+    # phi3 → swiglu). Unmatched families fall through to activation detection below.
+    _spec = match_by_decoder_type(decoder_type)
+    if _spec is not None and _spec.forced_activation is not None:
+        hidden_act = _spec.forced_activation
 
     if hidden_act is None:
         if hasattr(module, "activation"):
@@ -992,6 +980,12 @@ def module_match_name_list(module, name_list):
         if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"):
             return [first_proj_attr, "down_proj"]
 
+    # Resolve expert names from the model-family registry; fall back to the mapping
+    # below when no family spec matches the MoE block.
+    spec = match_moe_block(module)
+    if spec is not None and spec.expert_linear_names is not None:
+        return list(spec.expert_linear_names)
+
     if module_match_name_list(
         module,
         [
@@ -1581,6 +1575,11 @@ def build_decoder_config(
     """Builds the full decoder config from the module."""
     quantization = get_quantization_format(module)
     config = DecoderLayerConfig(decoder_type=decoder_type)
+
+    # Per-model behavioral hooks (default no-op) and the context passed into them.
+    _spec = match_by_decoder_type(decoder_type)
+    hooks = _spec.hooks if (_spec is not None and _spec.hooks is not None) else NULL_HOOKS
+    hook_ctx = ExportContext(build_layernorm_config=build_layernorm_config)
     # Supporting different attention layer config in MCoreGPTModel. If per layer config
     # exists, override the global config.
     if hasattr(module, "self_attention") and hasattr(module.self_attention, "config"):
@@ -1641,26 +1640,17 @@ def __init__(self):
         # and residual_layernorm could be after post_layernorm
         if is_layernorm(layer):
             layernorm_config = build_layernorm_config(layer)
+            # Family-specific layernorm slots (e.g. Gemma2/3) are placed by the hook.
+            if hooks.place_submodule(name, layer, layernorm_config, config, hook_ctx):
+                pass
             # Special attributes naming for Encoder-Decoder models
-            if model_type_is_enc_dec(decoder_type):
+            elif model_type_is_enc_dec(decoder_type):
                 _update_encoder_decoder_layernorm_config(
                     model_metadata_config, config, layernorm_config
                 )
             # For all decoder only models
             elif name in ["ln_mlp"]:
                 config.mlp_layernorm = layernorm_config
-            elif (
-                config.decoder_type in ["gemma2", "gemma3"]
-            ) and "post_attention_layernorm" in name:
-                config.post_layernorm = layernorm_config
-            elif (
-                config.decoder_type in ["gemma2", "gemma3"]
-            ) and "pre_feedforward_layernorm" in name:
-                config.pre_feedforward_layernorm = layernorm_config
-            elif (
-                config.decoder_type in ["gemma2", "gemma3"]
-            ) and "post_feedforward_layernorm" in name:
-                config.post_feedforward_layernorm = layernorm_config
             elif config.input_layernorm is None:
                 config.input_layernorm = layernorm_config
             elif config.post_layernorm is None:
@@ -1681,8 +1671,12 @@ def __init__(self):
                 attention_config = build_attention_config(
                     layer, model_metadata_config, config, tp_size=tp_size
                 )
+                # Family-specific attention placement (self/cross routing, extra norms)
+                # is handled by the hook; otherwise use the default placement below.
+                if hooks.place_submodule(name, layer, attention_config, config, hook_ctx):
+                    pass
                 # For decoder of Encoder-Decoder model with self, cross attention
-                if (
+                elif (
                     model_type_is_enc_dec(decoder_type)
                     and model_metadata_config["enc_dec"] == "dec"
                 ):
@@ -1691,19 +1685,8 @@ def __init__(self):
                         config.self_attention = attention_config
                     else:
                         config.cross_attention = attention_config
-                elif decoder_type == "mllama":
-                    if "cross" in type(layer).__name__.lower():
-                        config.cross_attention = attention_config
-                    else:
-                        config.self_attention = attention_config
                 else:
                     config.attention = attention_config
-                    if decoder_type == "gemma3":
-                        # Gemma3 has q_norm and k_norm within self_attn
-                        q_layernorm_config = build_layernorm_config(layer.q_norm)
-                        k_layernorm_config = build_layernorm_config(layer.k_norm)
-                        config.attention.q_layernorm = q_layernorm_config
-                        config.attention.k_layernorm = k_layernorm_config
 
         elif is_moe(layer):
             if quantization not in [QUANTIZATION_NONE, QUANTIZATION_FP8]:

modelopt/torch/export/model_config_export.py

@@ -60,6 +60,7 @@
     pack_linear_weights,
     split_config_and_weights,
 )
+from .modeling import match_by_decoder_type
 from .postprocess import (
     check_weight_shape_valid,
     pad_embedding_lm_head,
@@ -312,8 +313,10 @@ def torch_to_tensorrt_llm_checkpoint(
                     # the model head weight is None so we just skip processing
                     config.share_embedding_table = True
                     continue
-                # TRT LLM forces the embedding table to be shared for the following models.
-                force_share_embedding_table = decoder_type in ["gemma", "gemma2", "gemma3"]
+                # TRT LLM forces the embedding table to be shared for some families
+                # (see modeling/families/*); the equality check below still gates it.
+                _spec = match_by_decoder_type(decoder_type)
+                force_share_embedding_table = bool(_spec and _spec.force_share_embedding_table)
                 if force_share_embedding_table and torch.equal(
                     module.weight, config.vocab_embedding.weight
                 ):

modelopt/torch/export/modeling/__init__.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model-family export descriptors.
+
+Holds per-model export data organized by model family. The export code resolves a
+``ModelSpec`` via the registry lookups and reads its fields; an unmatched lookup
+returns ``None`` so callers fall back to their default behavior.
+"""
+
+from . import families
+from .base import ModelSpec
+from .context import ExportContext
+from .hooks import NULL_HOOKS, ModelHooks
+from .registry import (
+    iter_pqs_fuse_rules,
+    match_by_architecture,
+    match_by_decoder_type,
+    match_mlp_block,
+    match_moe_block,
+    register,
+)
+
+__all__ = [
+    "NULL_HOOKS",
+    "ExportContext",
+    "ModelHooks",
+    "ModelSpec",
+    "iter_pqs_fuse_rules",
+    "match_by_architecture",
+    "match_by_decoder_type",
+    "match_mlp_block",
+    "match_moe_block",
+    "register",
+]

modelopt/torch/export/modeling/base.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Per-model-family export descriptor.
+
+A ``ModelSpec`` declares how one model family differs from the generic export path,
+so the export code can read these values instead of branching on model names. Each
+spec holds per-model data only, not export logic.
+"""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .hooks import ModelHooks
+
+__all__ = ["ModelSpec"]
+
+
+@dataclass
+class ModelSpec:
+    """Per-model-family export data.
+
+    A spec is resolved from a model (or one of its sub-modules) via any of its matching
+    keys:
+        architectures: exact HuggingFace architecture class names (e.g. "Qwen3MoeForCausalLM").
+        decoder_types: ModelOpt ``decoder_type`` strings (e.g. "bloom", "phi3"). Exact match.
+        moe_block_names: MoE block class-name substrings, matched case-insensitively
+            (e.g. "Qwen3MoeSparseMoeBlock").
+        mlp_block_names: MLP module class names, matched by exact equality
+            (e.g. "ArcticMLP", "Phi3SmallMLP").
+
+    Per-model fields:
+        expert_linear_names: expert linear projection names for this family's MoE block,
+            e.g. ("gate_proj", "down_proj", "up_proj"). ``None`` if not applicable.
+        has_iterable_experts: True when experts are stored as per-expert iterable
+            sub-modules (Mixtral, Qwen MoE, NemotronH, Gemma4) and can be grouped by
+            ``get_experts_list``. False for stacked or fused layouts (DBRX, GptOss).
+        forced_activation: activation that overrides MLP activation detection
+            (e.g. Bloom/GLM → "gelu", Phi3 → "swiglu"). ``None`` for no override.
+        force_share_embedding_table: True for families that share the embedding/output
+            table (Gemma/Gemma2/Gemma3); still gated by a weight-equality check.
+        mlp_keyword_roles: overrides mapping a child-module name to its MLP role, e.g.
+            {"up_proj": "fc"} (MPT/Phi3Small) or {"w1": "fc", "w2": "proj", "w3": "gate"}
+            (Arctic/InternLM2). Each keyword is removed from the default role sets and
+            added to its target role. ``None`` for no override.
+        pqs_fuse_rules: AWQ pre_quant_scale fusion rules, each a
+            ``(module_class_substrings, fuse_into, fuse_from)`` triple: for a module whose
+            class name contains one of the substrings, the pre_quant_scale on ``fuse_from``
+            is folded into ``fuse_into`` (e.g. attention o_proj → v_proj, MLP down_proj →
+            up_proj).
+        hooks: optional :class:`ModelHooks` for behavioral seams that no value can express
+            (custom config-slot placement, module-tree unwrap, MoE building). ``None`` uses
+            the engine defaults.
+    """
+
+    name: str
+
+    # Matching keys.
+    architectures: tuple[str, ...] = ()
+    decoder_types: tuple[str, ...] = ()
+    moe_block_names: tuple[str, ...] = ()
+    mlp_block_names: tuple[str, ...] = ()
+
+    # MoE expert layout.
+    expert_linear_names: tuple[str, ...] | None = None
+    has_iterable_experts: bool = False
+
+    # MLP and activation.
+    forced_activation: str | None = None
+    mlp_keyword_roles: dict[str, str] | None = None
+
+    # Embedding.
+    force_share_embedding_table: bool = False
+
+    # AWQ pre_quant_scale fusion.
+    pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = ()
+
+    # Behavioral overrides for structural export seams (None = use the engine default).
+    hooks: "ModelHooks | None" = None
+
+    # Free-form extension slot for future per-model fields.
+    _extra: dict = field(default_factory=dict, repr=False)

modelopt/torch/export/modeling/context.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Context passed from the export engine into per-model hooks.
+
+The engine fills this with the builder callables a hook may need and passes it in, so
+hooks never import the engine modules directly (keeping ``modeling/`` free of cycles).
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+__all__ = ["ExportContext"]
+
+
+@dataclass
+class ExportContext:
+    """Builders and metadata available to ``ModelHooks`` methods."""
+
+    build_layernorm_config: Callable