[None][feat] Afmoe trinity support

docs/source/models/supported-models.md

@@ -5,6 +5,7 @@ The following is a table of supported models for the PyTorch backend:
 
 | Architecture                         | Model                              | HuggingFace Example                          |
 | ------------------------------------ | ---------------------------------- | -------------------------------------------- |
+| `AfmoeForCausalLM`                   | Arcee Foundation MoE (Trinity)     | `arcee-ai/Trinity-Mini`                      |
 | `BertForSequenceClassification`      | BERT-based                         | `textattack/bert-base-uncased-yelp-polarity` |
 | `Cohere2ForCausalLM`                 | Command A                          | `CohereLabs/c4ai-command-a-03-2025`          |
 | `DeciLMForCausalLM`                  | Nemotron                           | `nvidia/Llama-3_1-Nemotron-51B-Instruct`     |

tensorrt_llm/_torch/models/__init__.py

@@ -5,6 +5,7 @@
 # under transformers >= 5.5; see _torch/configs/__init__.py.
 import tensorrt_llm._torch.configs  # noqa: F401
 
+from .modeling_afmoe import AfmoeForCausalLM
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_bert import BertForSequenceClassification
 from .modeling_clip import CLIPVisionModel
@@ -51,6 +52,7 @@
 
 # Note: for better readiblity, this should have same order as imports above
 __all__ = [
+    "AfmoeForCausalLM",
     "AutoModelForCausalLM",
     "BertForSequenceClassification",
     "CLIPVisionModel",

tensorrt_llm/_torch/models/checkpoints/__init__.py

@@ -1,4 +1,5 @@
 from .base_checkpoint_loader import BaseCheckpointLoader
+from .hf.afmoe_weight_mapper import AfmoeHfWeightMapper
 from .hf.checkpoint_loader import HfCheckpointLoader
 from .hf.config_loader import HfConfigLoader
 from .hf.gemma3_weight_mapper import Gemma3HfWeightMapper
@@ -24,9 +25,9 @@
 from .mx.checkpoint_loader import MXCheckpointLoader
 
 __all__ = [
-    "HfConfigLoader", "HfWeightLoader", "HfWeightMapper", "MistralConfigLoader",
-    "MistralWeightMapper", "MistralCheckpointLoader", "BaseCheckpointLoader",
-    "HfCheckpointLoader", "NemotronHHfWeightMapper",
+    "AfmoeHfWeightMapper", "HfConfigLoader", "HfWeightLoader", "HfWeightMapper",
+    "MistralConfigLoader", "MistralWeightMapper", "MistralCheckpointLoader",
+    "BaseCheckpointLoader", "HfCheckpointLoader", "NemotronHHfWeightMapper",
     "NemotronNasHfWeightMapper", "Gemma3HfWeightMapper",
     "MixtralHfWeightMapper", "Llama4HfWeightMapper", "Qwen2MoeHfWeightMapper",
     "Qwen3MoeHfWeightMapper", "Qwen2VLHfWeightMapper",

tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py

@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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.
+
+import torch
+from torch import nn
+
+from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
+from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm._torch.modules.fused_moe.interface import MoE
+
+
+@register_mapper("HF", "AfmoeForCausalLM")
+class AfmoeHfWeightMapper(HfWeightMapper):
+    def __init__(self):
+        super().__init__()
+
+        self.params_map = {
+            # MoE expert weights: gate_proj->w1, up_proj->w3, down_proj->w2
+            r"(.*experts\.\d+\.)gate_proj(.*)": r"\1w1\2",
+            r"(.*experts\.\d+\.)up_proj(.*)": r"\1w3\2",
+            r"(.*experts\.\d+\.)down_proj(.*)": r"\1w2\2",
+            # HF router weight path -> TRT-LLM gate path
+            r"(.*)\.router\.gate\.(.*)": r"\1.gate.\2",
+            # expert_bias -> gate.e_score_correction_bias
+            r"(.*)\.mlp\.expert_bias(.*)": r"\1.mlp.gate.e_score_correction_bias\2",
+        }
+
+    def preprocess_weights(self, weights: dict) -> dict:
+        weights = self.rename_by_params_map(self.params_map, weights)
+        weights = self._fuse_attention_gate(weights)
+        return weights
+
+    def _fuse_attention_gate(self, weights: dict) -> dict:
+        """Fuse the separate attention ``gate_proj`` into ``q_proj``.
+
+        AfmoeAttention uses ``attn_output_gate=True``, so the gate weights are
+        interleaved with the query weights per head and loaded through the fused
+        QKV projection. The HF checkpoint stores ``q_proj`` and ``gate_proj`` as
+        two separate matrices of shape ``[num_heads * head_dim, hidden]``; the
+        fused QKV projection expects the query slot laid out per head as
+        ``[head0_q, head0_gate, head1_q, head1_gate, ...]`` (see
+        ``Attention.forward`` where ``q_gate`` is viewed as
+        ``[..., num_heads, 2 * head_dim]`` and chunked into q/gate).
+        """
+        marker = ".self_attn.gate_proj."
+        gate_keys = [k for k in weights if marker in k]
+        if not gate_keys:
+            return weights
+
+        num_heads = self.model.config.num_attention_heads
+        for gate_key in gate_keys:
+            prefix, suffix = gate_key.split(marker)
+            q_key = f"{prefix}.self_attn.q_proj.{suffix}"
+            if q_key not in weights:
+                continue
+            weights[q_key] = self._interleave_per_head(weights[q_key], weights[gate_key], num_heads)
+            del weights[gate_key]
+        return weights
+
+    @staticmethod
+    def _interleave_per_head(q: torch.Tensor, gate: torch.Tensor, num_heads: int) -> torch.Tensor:
+        """Interleave q and gate rows per head: ``[h0_q, h0_gate, h1_q, ...]``.
+
+        Works for 2D weights ``[num_heads * per_head, hidden]`` as well as 1D
+        biases and FP8 block scales, since the split is always taken along the
+        leading (output) dimension.
+        """
+        assert q.shape[0] % num_heads == 0, (
+            f"q_proj rows {q.shape[0]} not divisible by num_heads {num_heads}"
+        )
+        assert gate.shape == q.shape, f"gate_proj shape {gate.shape} != q_proj shape {q.shape}"
+        per_head = q.shape[0] // num_heads
+        tail = q.shape[1:]
+        q = q.reshape(num_heads, per_head, *tail)
+        gate = gate.reshape(num_heads, per_head, *tail)
+        fused = torch.stack([q, gate], dim=1)
+        return fused.reshape(num_heads * 2 * per_head, *tail).contiguous()
+
+    def is_special_instance_module(self, module: nn.Module) -> bool:
+        return isinstance(module, MoE)
+
+    def handle_special_instance_module(
+        self,
+        module: nn.Module,
+        module_name: str,
+        module_weights: dict,
+        allow_partial_loading: bool = False,
+    ) -> None:
+        if isinstance(module, MoE):
+            module.load_weights(
+                weights=[module_weights],
+                allow_partial_loading=allow_partial_loading,
+            )