[None][feat] Qwen3-Next MTP (#11370)

Signed-off-by: Izzy Putterman <iputterman@nvidia.com>

Author: IzzyPutterman

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

@@ -56,6 +56,7 @@ def preprocess_weights(self, weights: dict) -> dict:
         config = self.config.pretrained_config
         tp_size = self.config.mapping.tp_size
         tp_rank = self.config.mapping.tp_rank
+        mtp_layer_offset = config.num_hidden_layers
 
         if self.config.mapping.enable_attention_dp:
             tp_size = 1
@@ -66,10 +67,28 @@ def preprocess_weights(self, weights: dict) -> dict:
         linear_key_dim = config.linear_key_head_dim * config.linear_num_key_heads  # 16 * 128
         linear_value_dim = config.linear_value_head_dim * config.linear_num_value_heads  # 32 * 128
 
+        mtp_mapping = {
+            "mtp.fc": "fc",
+            "mtp.norm": "shared_head.norm",
+            "mtp.pre_fc_norm_embedding": "pre_fc_norm_embedding",
+            "mtp.pre_fc_norm_hidden": "pre_fc_norm_hidden",
+        }
+
         new_weights = {}
         for name, _ in weights.items():
             key = name
 
+            if key.startswith("mtp.layers."):
+                _, _, mtp_layer_idx, module_name = key.split(".", 3)
+                key = (f"model.layers.{mtp_layer_offset + int(mtp_layer_idx)}."
+                       f"{module_name}")
+            elif key.startswith("mtp."):
+                for mtp_prefix, trtllm_name in mtp_mapping.items():
+                    if key.startswith(mtp_prefix):
+                        suffix = key[len(mtp_prefix):]
+                        key = f"model.layers.{mtp_layer_offset}.{trtllm_name}{suffix}"
+                        break
+
             if "A_log" in key:
                 w = split(weights[name], tp_size, tp_rank)
                 w = w.to(torch.float32)