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

Introduce the first PR slice from the MX/GMS prototype: checkpoint_format="MX" support using upstream modelexpress MxLiveWeightLoader and publish_model_params, while intentionally excluding GMS/load_format changes.

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Made-with: Cursor
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Made-with: Cursor

Author: chienchunhung

setup.py

@@ -427,6 +427,21 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
     scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'],
     extras_require={
         "devel": devel_deps,
+        # NOTE: The MX (modelexpress) Python package used by
+        # tensorrt_llm._torch.models.checkpoints.mx is intentionally NOT
+        # declared as an ``[mx]`` extra while this integration is at
+        # prototype status. ``modelexpress`` v0.3.0 is on PyPI
+        # (Apache-2.0) but is brand-new (Beta status, single release) and
+        # still needs onboarding into NVIDIA's OSS package allowlist.
+        #
+        # Until allowlisting is complete, users who want to exercise the
+        # MX code path install the dependency manually:
+        #
+        #     pip install "modelexpress>=0.3.0,<0.4.0"
+        #
+        # Restoring one-line ``pip install tensorrt_llm[mx]`` ergonomics
+        # is a single revert of this hunk once the OSS-allowlist step is
+        # complete (tracked in §15 of the design doc as MX-7).
     },
     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_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,16 @@ 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. Generic HF
+                loaders ignore these; MX uses ``model`` for direct P2P writes.
 
         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("HF")
 class HfWeightLoader(BaseWeightLoader):
@@ -59,8 +60,8 @@ def _get_local_available_host_memory() -> int:
                                               op=_MPI.MIN)
         return available_host_memory
 
-    def load_weights(self, checkpoint_dir: str,
-                     mapping: Mapping) -> dict[str, Any]:
+    def load_weights(self, checkpoint_dir: str, mapping: Mapping,
+                     **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, as there really is no need

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

@@ -7,6 +7,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"]