[TRTLLM-16725][refactor] Simplify checkpoint I/O policy resolution

tensorrt_llm/_torch/models/checkpoints/auto_mapper.py

@@ -1,22 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from typing import Optional
 
 from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPER_MAPPING
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
 class AutoCheckpointMapper():
 
     @staticmethod
-    def get(format: str, name: Optional[str] = None) -> "BaseWeightMapper":
+    def get(format: CheckpointFormat | str,
+            name: Optional[str] = None) -> "BaseWeightMapper":
         if name is not None:
             try:
                 return MODEL_CLASS_MAPPER_MAPPING[f'{name}_{format}']()
             except KeyError:  # no mapper for this model architecture, resort to default
-                if format == "MX":
+                if format == CheckpointFormat.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']()
+                        return MODEL_CLASS_MAPPER_MAPPING[
+                            f'{name}_{CheckpointFormat.HF}']()
                     except KeyError:
                         pass
                 # TODO smor- a potential bug here, if the class isn't added to __init__, it will return the default mapper

tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py

@@ -21,6 +21,7 @@
     CheckpointCatalog
 from tensorrt_llm._torch.models.modeling_utils import \
     CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
@@ -56,7 +57,7 @@ def config_loader(self) -> BaseConfigLoader:
 
     @property
     @abstractmethod
-    def checkpoint_format(self) -> str:
+    def checkpoint_format(self) -> CheckpointFormat | str:
         ...
 
     def load_config(self, checkpoint_dir: str, **kwargs) -> ModelConfig:
@@ -111,7 +112,8 @@ def post_load_publish(self,
         """Publish format-specific loaded weights after the load path."""
 
     @classmethod
-    def get(cls, checkpoint_format: str, **kwargs) -> "BaseCheckpointLoader":
+    def get(cls, checkpoint_format: CheckpointFormat | str,
+            **kwargs) -> "BaseCheckpointLoader":
         try:
             return CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING[checkpoint_format](
                 **kwargs)

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

@@ -19,9 +19,10 @@
 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.moe.fused_moe.weight_owner import is_moe_weight_owner
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "AfmoeForCausalLM")
+@register_mapper(CheckpointFormat.HF, "AfmoeForCausalLM")
 class AfmoeHfWeightMapper(HfWeightMapper):
     def __init__(self):
         super().__init__()

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

@@ -20,10 +20,12 @@
 from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \
     HfWeightLoader
 from tensorrt_llm._torch.models.modeling_utils import register_checkpoint_loader
+from tensorrt_llm.checkpoint_format import CheckpointFormat
+from tensorrt_llm.llmapi.llm_args import CheckpointIOPolicy
 from tensorrt_llm.mapping import Mapping
 
 
-@register_checkpoint_loader("HF")
+@register_checkpoint_loader(CheckpointFormat.HF)
 class HfCheckpointLoader(BaseCheckpointLoader):
 
     def __init__(self,
@@ -40,7 +42,7 @@ def __init__(self,
         else:
             self._config_loader = config_loader
         self._weight_mapper = weight_mapper
-        self._checkpoint_format = "HF"
+        self._checkpoint_format = CheckpointFormat.HF
 
     def cleanup(self) -> None:
         # Clean up weight mapper first as it may hold model references
@@ -74,14 +76,14 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping,
             loader = self.weight_loader
             loader._reset_checkpoint_io_status()
             status = loader._last_checkpoint_io_status
-            status.selected = "native"
-            if loader.checkpoint_io_policy != "native":
+            status.selected = CheckpointIOPolicy.NATIVE
+            if loader.checkpoint_io_policy != CheckpointIOPolicy.NATIVE:
                 status.fallback_reason = (
                     "context-only model loading requires lazy safetensors "
                     "without full-checkpoint read-ahead")
             weights = loader._load_lazy_safetensors(
                 checkpoint_dir, kwargs.get("use_consolidated", False))
-            status.effective = "native"
+            status.effective = CheckpointIOPolicy.NATIVE
             loader._log_checkpoint_io_status()
             return weights
         return super().load_weights(checkpoint_dir, mapping=mapping, **kwargs)
@@ -92,7 +94,8 @@ def open_weight_session(self, checkpoint_dir: str, mapping: Mapping,
         """Delegate the optimized session only for the built-in HF path."""
         if (type(self) is HfCheckpointLoader
                 and type(self.weight_loader) is HfWeightLoader
-                and self.weight_loader.checkpoint_io_policy != "native" and
+                and self.weight_loader.checkpoint_io_policy
+                != CheckpointIOPolicy.NATIVE and
                 not self._uses_context_only_lazy_loading(kwargs.get("model"))):
             with self.weight_loader.open_weight_session(checkpoint_dir,
                                                         mapping=mapping,
@@ -138,5 +141,5 @@ def config_loader(self) -> Optional[BaseConfigLoader]:
         return self._config_loader
 
     @property
-    def checkpoint_format(self) -> str:
+    def checkpoint_format(self) -> CheckpointFormat:
         return self._checkpoint_format

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

@@ -5,10 +5,11 @@
 from tensorrt_llm._torch.models.checkpoints.base_config_loader import \
     BaseConfigLoader
 from tensorrt_llm._torch.models.modeling_utils import register_config_loader
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_config_loader("MX")
-@register_config_loader("HF")
+@register_config_loader(CheckpointFormat.MX)
+@register_config_loader(CheckpointFormat.HF)
 class HfConfigLoader(BaseConfigLoader):
 
     def load(self, checkpoint_dir: str, **kwargs) -> ModelConfig:

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

@@ -5,9 +5,10 @@
 
 from tensorrt_llm._torch.models.checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "Cosmos3ForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Cosmos3ForConditionalGeneration")
 class Cosmos3HfWeightMapper(Qwen3VLHfWeightMapper):
     """
     Cosmos3 unified checkpoints store the Reasoner LLM with the old `model.` prefix

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

@@ -16,9 +16,10 @@
 from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "Exaone4_5_ForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Exaone4_5_ForConditionalGeneration")
 class Exaone4_5HfWeightMapper(HfWeightMapper):
     def preprocess_weights(self, weights: dict):
         """Rename HF checkpoint prefixes; supports plain dict and ConsumableWeightsDict."""

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

@@ -18,9 +18,10 @@
 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.moe.fused_moe.weight_owner import is_moe_weight_owner
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "ExaoneMoeForCausalLM")
+@register_mapper(CheckpointFormat.HF, "ExaoneMoeForCausalLM")
 class ExaoneMoeWeightMapper(HfWeightMapper):
     def __init__(self):
         super().__init__()

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

@@ -1,12 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 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.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "Gemma3ForCausalLM")
-@register_mapper("HF", "Gemma3ForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Gemma3ForCausalLM")
+@register_mapper(CheckpointFormat.HF, "Gemma3ForConditionalGeneration")
 class Gemma3HfWeightMapper(HfWeightMapper):
 
     def should_skip_module(self, module_name: str) -> bool:

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

@@ -19,6 +19,7 @@
 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.linear import W4A16_AWQ_LinearMethod
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 _LANG_PREFIX = "model.language_model."
 _MODEL_PREFIX = "model."
@@ -29,10 +30,10 @@
 )
 
 
-@register_mapper("HF", "Gemma4ForCausalLM")
-@register_mapper("HF", "Gemma4ForConditionalGeneration")
-@register_mapper("HF", "Gemma4UnifiedForConditionalGeneration")
-@register_mapper("HF", "Gemma4AssistantForCausalLM")
+@register_mapper(CheckpointFormat.HF, "Gemma4ForCausalLM")
+@register_mapper(CheckpointFormat.HF, "Gemma4ForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Gemma4UnifiedForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Gemma4AssistantForCausalLM")
 class Gemma4HfWeightMapper(HfWeightMapper):
     @property
     def _is_vlm(self) -> bool:

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

@@ -30,6 +30,7 @@
 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.pyexecutor.config_utils import unwrap_glm5_next_text_config
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 from tensorrt_llm.quantization.mode import QuantAlgo
 
 # Match dotted namespaces exactly so unrelated decoder weights are not ignored.
@@ -155,8 +156,8 @@ def glm5_next_weight_owner(dest: str, num_layers: int) -> Any:
     return None
 
 
-@register_mapper("HF", "Glm5NextForConditionalGeneration")
-@register_mapper("HF", "Glm5NextForCausalLM")
+@register_mapper(CheckpointFormat.HF, "Glm5NextForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Glm5NextForCausalLM")
 class Glm5NextHfWeightMapper(HfWeightMapper):
     """Resolve checkpoint destinations and owners for the model's custom loader.
 

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

@@ -1,9 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \
     HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "Llama4ForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Llama4ForConditionalGeneration")
 class Llama4HfWeightMapper(HfWeightMapper):
     """
     Weight mapper for Llama4ForConditionalGeneration that handles the

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

@@ -1,8 +1,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "LlavaNextForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "LlavaNextForConditionalGeneration")
 class LlavaNextHfWeightMapper(HfWeightMapper):
     def preprocess_weights(self, weights: dict) -> dict:
         transformed_weights = {}

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

@@ -17,14 +17,15 @@
 
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 MINIMAX_M3_PARAMS_MAP = {
     r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": r"\1.gate.e_score_correction_bias",
 }
 
 
-@register_mapper("HF", "MiniMaxM3SparseForCausalLM")
-@register_mapper("HF", "MiniMaxM3SparseForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "MiniMaxM3SparseForCausalLM")
+@register_mapper(CheckpointFormat.HF, "MiniMaxM3SparseForConditionalGeneration")
 class MiniMaxM3HfWeightMapper(HfWeightMapper):
     """Handle M3 gate naming and MXFP8 GQA duplication for loader v2."""
 

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

@@ -1,11 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from torch import nn
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \
     BaseWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "MixtralForCausalLM")
+@register_mapper(CheckpointFormat.HF, "MixtralForCausalLM")
 class MixtralHfWeightMapper(BaseWeightMapper):
 
     def map_weights(self) -> None:

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

@@ -11,10 +11,11 @@
     HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
 from tensorrt_llm._torch.utils import split
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "NemotronHPuzzleForCausalLM")
-@register_mapper("HF", "NemotronHForCausalLM")
+@register_mapper(CheckpointFormat.HF, "NemotronHPuzzleForCausalLM")
+@register_mapper(CheckpointFormat.HF, "NemotronHForCausalLM")
 class NemotronHHfWeightMapper(HfWeightMapper):
 
     @staticmethod

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

@@ -1,10 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 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.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "DeciLMForCausalLM")
+@register_mapper(CheckpointFormat.HF, "DeciLMForCausalLM")
 class NemotronNasHfWeightMapper(HfWeightMapper):
     """Weight mapper for Nemotron-NAS / DeciLM.
 

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

@@ -1,10 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 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.moe.fused_moe.weight_owner import is_moe_weight_owner
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
 def _unfuse_moe_expert_weights(weights: dict) -> dict:
@@ -57,7 +61,7 @@ def _unfuse_moe_expert_weights(weights: dict) -> dict:
     return updated
 
 
-@register_mapper("HF", "Qwen2MoeForCausalLM")
+@register_mapper(CheckpointFormat.HF, "Qwen2MoeForCausalLM")
 class Qwen2MoeHfWeightMapper(HfWeightMapper):
 
     def is_special_instance_module(self, module: nn.Module) -> bool:

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

@@ -1,9 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \
     HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.checkpoint_format import CheckpointFormat
 
 
-@register_mapper("HF", "Qwen2VLForConditionalGeneration")
+@register_mapper(CheckpointFormat.HF, "Qwen2VLForConditionalGeneration")
 class Qwen2VLHfWeightMapper(HfWeightMapper):
     """
     Weight mapper for Qwen2VLForConditionalGeneration that handles the