[TRTLLM-11851][feat] Add MX-only P2P checkpoint loading support for TRTLLM

setup.py

@@ -427,6 +427,10 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
     scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'],
     extras_require={
         "devel": devel_deps,
+        # MX remains prototype-only and is intentionally not declared as an
+        # optional package extra until its external dependency completes OSS
+        # allowlist onboarding. Keep install instructions in docs/PR text
+        # rather than packaging metadata.
     },
     zip_safe=True,
     install_requires=required_deps,

tensorrt_llm/_torch/models/checkpoints/__init__.py

@@ -19,6 +19,7 @@
 from .mistral.config_loader import MistralConfigLoader
 from .mistral.weight_mapper import (MistralLarge3WeightMapper,
                                     MistralWeightMapper)
+from .mx.checkpoint_loader import MXCheckpointLoader
 
 __all__ = [
     "HfConfigLoader", "HfWeightLoader", "HfWeightMapper", "MistralConfigLoader",
@@ -28,5 +29,5 @@
     "Qwen3MoeHfWeightMapper", "Qwen2VLHfWeightMapper",
     "Qwen3_5MoeHfWeightMapper", "Qwen3NextHfWeightMapper",
     "LlavaNextHfWeightMapper", "MistralLarge3CheckpointLoader",
-    "MistralLarge3WeightMapper", "Qwen3VLHfWeightMapper"
+    "MistralLarge3WeightMapper", "MXCheckpointLoader", "Qwen3VLHfWeightMapper"
 ]

tensorrt_llm/_torch/models/checkpoints/auto_mapper.py

@@ -11,6 +11,14 @@ def get(format: str, name: Optional[str] = None) -> "BaseWeightMapper":
             try:
                 return MODEL_CLASS_MAPPER_MAPPING[f'{name}_{format}']()
             except KeyError:  # no mapper for this model architecture, resort to default
+                if format == "MX":
+                    # MX uses HF on-disk checkpoint format for fallback, so
+                    # an architecture-specific HF mapper is closer than the
+                    # generic MX/HF default mapper.
+                    try:
+                        return MODEL_CLASS_MAPPER_MAPPING[f'{name}_HF']()
+                    except KeyError:
+                        pass
                 # TODO smor- a potential bug here, if the class isn't added to __init__, it will return the default mapper
                 return MODEL_CLASS_MAPPER_MAPPING[format]()
         else:

tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py

@@ -65,6 +65,23 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping,
                                                mapping=mapping,
                                                **kwargs)
 
+    def is_weights_preloaded(self) -> bool:
+        """Whether the last load wrote weights directly into the model."""
+        return False
+
+    def post_load_apply(self,
+                        model: nn.Module,
+                        *,
+                        weights_preloaded: bool = False) -> None:
+        """Apply format-specific state after weights have been loaded."""
+
+    def post_load_publish(self,
+                          model: nn.Module,
+                          *,
+                          checkpoint_dir: str,
+                          weights_preloaded: bool = False) -> None:
+        """Publish format-specific loaded weights after the load path."""
+
     @classmethod
     def get(cls, checkpoint_format: str, **kwargs) -> "BaseCheckpointLoader":
         try:

tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import threading
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterator, Tuple, Union
@@ -91,15 +94,15 @@ def mark_consumed(self, prefix: str) -> int:
 class BaseWeightLoader(ABC):
 
     @abstractmethod
-    def load_weights(
-            self, checkpoint_dir: str,
-            mapping: Mapping) -> Union[Dict[str, Any], ConsumableWeightsDict]:
+    def load_weights(self, checkpoint_dir: str, mapping: Mapping,
+                     **kwargs) -> Union[Dict[str, Any], ConsumableWeightsDict]:
         """
         Loads weights from a checkpoint directory.
 
         Args:
             checkpoint_dir: A path to the checkpoint directory.
             mapping: A mapping object containing the distributed configuration.
+            **kwargs: Optional format-specific loader arguments.
 
         Returns:
             A dictionary (or ConsumableWeightsDict) where keys are tensor names

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

@@ -1,9 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.checkpoints.base_config_loader import \
     BaseConfigLoader
 from tensorrt_llm._torch.models.modeling_utils import register_config_loader
 
 
+@register_config_loader("MX")
 @register_config_loader("HF")
 class HfConfigLoader(BaseConfigLoader):
 

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

@@ -34,6 +34,7 @@
 from tensorrt_llm.mapping import Mapping
 
 
+@register_checkpoint_weight_loader("MX")
 @register_checkpoint_weight_loader("mistral")
 @register_checkpoint_weight_loader("mistral_large_3")
 @register_checkpoint_weight_loader("HF")
@@ -63,7 +64,8 @@ def _get_local_available_host_memory() -> int:
     def load_weights(self,
                      checkpoint_dir: str,
                      mapping: Mapping,
-                     use_consolidated: bool = False) -> dict[str, Any]:
+                     use_consolidated: bool = False,
+                     **kwargs) -> dict[str, Any]:
         weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors")
         # Some model checkpoint directories contain not only the sharded safetensors, but one
         # consolidated tensor. In the presence of both, we favor the former unless specified explicitly, as there really is no need

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

@@ -6,6 +6,7 @@
 from ..base_weight_mapper import BaseWeightMapper
 
 
+@register_mapper("MX")
 @register_mapper("HF")
 class HfWeightMapper(BaseWeightMapper):
 

tensorrt_llm/_torch/models/checkpoints/mx/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 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.
+
+from .checkpoint_loader import MXCheckpointLoader
+
+__all__ = ["MXCheckpointLoader"]