[None][refactor] Register modeling_v2 routes with a routes decorator

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -71,9 +71,8 @@ LLM(model=..., modeling_v2=...)  ->  ModelLoader
                                         +-- modeling_v2_resolve(config, llm_args)
                                               |
                     _router_index.py          +-- architectures[0] -> routing module
-                    models/<x>/routing.py     +-- route(config, mapping, compute_capability,
-                                              |         spec_config, kv_cache_quant_algo,
-                                              |         model_config, llm_args)
+                    models/<x>/routing.py     +-- route(model_config, llm_args,
+                                              |         compute_capability)
                                               |         -> target class, or NoTarget
                                         |
                                      config.modeling_v2_target = class

tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py

@@ -27,7 +27,6 @@
 """
 
 from ._router_index import (
-    MODELING_V2_ROUTERS,
     MODELING_V2_VALIDATE_ENV,
     ModelingV2Mode,
     NoTarget,
@@ -36,7 +35,6 @@
 )
 
 __all__ = [
-    "MODELING_V2_ROUTERS",
     "MODELING_V2_VALIDATE_ENV",
     "ModelingV2Mode",
     "NoTarget",

tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py

@@ -1,63 +1,78 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Architecture name -> routing module, and the resolver that reads the table.
+"""Architecture name -> ``route``, and the resolver that consults it.
 
 The checkpoint's own ``architectures[0]`` (``GptOssForCausalLM``,
-``DeepseekV3ForCausalLM``) selects a routing module, never a target: one
+``DeepseekV3ForCausalLM``) selects a family's ``route``, never a target: one
 architecture has many deployments and a target serves exactly one of them.
 ``modeling_v2_resolve`` returns the target *class*, the model loader records it
 on ``ModelConfig.modeling_v2_target``, and ``AutoModelForCausalLM._resolve_class``
 builds that class instead of the registry's.
 
 Two levels, on purpose:
 
-* This table maps ``architectures[0]`` -> routing module. Its keys are things
-  that already exist in every checkpoint's config.json, so adding a new
-  checkpoint, GPU arch or parallel topology never touches it -- only a new
-  architecture family does.
-* The routing module's ``route`` owns the rest of the decision as one
-  forward-reading tree over the facts that fix which code runs: the
-  checkpoint's config, the parallel mapping, the GPU's compute capability, the
-  speculative-decoding config and the KV-cache algorithm the engine will run
-  with. It also receives the ``ModelConfig`` and the LLM arguments, from which a
-  family derives any other engine choice its targets depend on. It returns the
-  target class, or raises :class:`NoTarget` with the reason in words. Reading
-  that one file tells you where any configuration lands.
+* Each family's ``routing.py`` declares the ``architectures[0]`` values its
+  ``route`` serves with ``@routes(...)``, and ``models/__init__.py`` imports
+  every family's ``routing.py``, which is what registers them. The names are
+  things that already exist in every checkpoint's config.json, so adding a
+  new checkpoint, GPU arch or parallel topology never touches either -- only
+  a new architecture family does.
+* That ``route`` owns the rest of the decision as one
+  forward-reading tree over the facts that fix which code runs. It receives
+  the ``ModelConfig``, the LLM arguments and the GPU's compute capability, and
+  reads from the first two the checkpoint's config, the parallel mapping, the
+  speculative-decoding config, the KV-cache algorithm the engine will run with
+  (:func:`engine_kv_cache_quant_algo`) and any other engine choice its targets
+  depend on. It returns the target class, or raises :class:`NoTarget` with the
+  reason in words. Reading that one file tells you where any configuration
+  lands.
 
 Tuning knobs -- ``max_batch_size``, ``max_num_tokens``, CUDA graph settings --
 are deliberately not routing inputs. They do not change which code runs, so a
 target claims a configuration regardless of them; ``require`` promises that the
 target ran, not that a particular knob setting was measured.
 
-Routing modules are imported lazily: resolving a GptOss config never imports
-the DeepSeek tree, and a target's modeling code is imported only on the
-``route`` branch that returns it.
+Routing modules are imported on the first resolve, all of them, and hold no
+modeling code: a target's modeling code is imported only on the ``route``
+branch that returns it, so resolving a GptOss config never imports a DeepSeek
+target.
 """
 
 from __future__ import annotations
 
 import enum
 import os
-from importlib import import_module
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple
 
 from tensorrt_llm.logger import logger
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.model_config import ModelConfig
     from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
-_PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2"
+# architectures[0] -> the ``route`` that decides where it lands, filled by ``@routes``.
+_ROUTES: Dict[str, Callable[..., type]] = {}
 
-# architectures[0] -> routing module, relative to this package.
-MODELING_V2_ROUTERS = {
-    "GptOssForCausalLM": "models.gpt_oss.routing",
-    "DeepseekV3ForCausalLM": "models.deepseek_v3.routing",
-    # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the
-    # config loader's composite branch writes the same name back.
-    "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing",
-    "GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing",
-}
+
+def routes(*architectures: str):
+    """Declare that the decorated ``route`` decides where these ``architectures[0]`` values land."""
+
+    def register(route):
+        for arch in architectures:
+            assert arch not in _ROUTES, (
+                f"{arch} routed by both {_ROUTES[arch].__module__} and {route.__module__}"
+            )
+            _ROUTES[arch] = route
+        return route
+
+    return register
+
+
+def registered_routes() -> Dict[str, Callable[..., type]]:
+    """architectures[0] -> ``route``, for every family ``models/__init__.py`` imports."""
+    from . import models  # noqa: F401
+
+    return _ROUTES
 
 
 #: Turns the step contract check on. Off by default, so a served engine runs
@@ -99,7 +114,7 @@ def of(cls, llm_args: "TorchLlmArgs") -> "ModelingV2Mode":
 
 
 class NoTarget(Exception):
-    """Raised by a routing module's ``route`` when no target fits.
+    """Raised by a family's ``route`` when no target fits.
 
     The message says, in words, which fact ruled the targets out. Under
     ``auto`` the resolver logs it and falls back to the built-in implementation;
@@ -130,14 +145,6 @@ def tp_reaches_across_nodes(mapping) -> bool:
     return MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16)
 
 
-def routing_module(arch: str):
-    """Import and return the routing module for ``arch``, or None."""
-    name = MODELING_V2_ROUTERS.get(arch)
-    if name is None:
-        return None
-    return import_module(f"{_PACKAGE}.{name}")
-
-
 def _device_compute_capability() -> Tuple[int, int]:
     import torch
 
@@ -148,7 +155,7 @@ def _device_compute_capability() -> Tuple[int, int]:
     return torch.cuda.get_device_capability()
 
 
-def _kv_cache_quant_algo(config: "ModelConfig", llm_args: "TorchLlmArgs"):
+def engine_kv_cache_quant_algo(config: "ModelConfig", llm_args: "TorchLlmArgs"):
     """The KV-cache algorithm the engine will run with, which routing runs too early to read.
 
     The loader resolves the target before `validate_and_set_kv_cache_quant`
@@ -172,8 +179,8 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti
     """Decide which modeling_v2 target, if any, this deployment builds.
 
     Returns the target class, or None for the built-in implementation:
-    when ``llm_args.modeling_v2`` is ``off``, when no routing module claims the
-    architecture, or when the routing module's ``route`` raises ``NoTarget``.
+    when ``llm_args.modeling_v2`` is ``off``, when no ``route`` is registered
+    for the architecture, or when that ``route`` raises ``NoTarget``.
     In ``require`` every one of those but ``off`` raises ``ValueError`` instead.
     """
     mode = ModelingV2Mode.of(llm_args)
@@ -186,21 +193,11 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti
     arch = pretrained_config.architectures[0]
 
     try:
-        routing = routing_module(arch)
-        if routing is None:
-            known = ", ".join(sorted(MODELING_V2_ROUTERS)) or "(none)"
-            raise NoTarget(
-                f"no routing module for this architecture; routed architectures: {known}"
-            )
-        return routing.route(
-            pretrained_config,
-            config.mapping,
-            _device_compute_capability(),
-            config.spec_config,
-            _kv_cache_quant_algo(config, llm_args),
-            model_config=config,
-            llm_args=llm_args,
-        )
+        route = registered_routes().get(arch)
+        if route is None:
+            known = ", ".join(sorted(registered_routes())) or "(none)"
+            raise NoTarget(f"no route for this architecture; routed architectures: {known}")
+        return route(config, llm_args, _device_compute_capability())
     except NoTarget as why:
         if mode is ModelingV2Mode.REQUIRE:
             raise ValueError(f"modeling_v2='require', but no target fits {arch!r}: {why}") from None

tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py

@@ -5,4 +5,13 @@
 ``models/<x>/`` corresponds one-to-one with the built-in zoo's
 ``modeling_<x>.py``: the same architecture, split into one self-contained
 codebase per deployment target instead of one class serving every
-checkpoint."""
+checkpoint.
+
+Importing this package imports every family's ``routing`` module, and each one
+registers its ``route`` through ``@routes``. A new family adds one line here.
+"""
+
+from .deepseek_v3 import routing as deepseek_v3_routing  # noqa: F401
+from .glm_moe_dsa import routing as glm_moe_dsa_routing  # noqa: F401
+from .gpt_oss import routing as gpt_oss_routing  # noqa: F401
+from .qwen4_exp import routing as qwen4_exp_routing  # noqa: F401

tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py

@@ -11,12 +11,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Tuple
 
 from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
 from tensorrt_llm.quantization.mode import QuantAlgo
 
-from ..._router_index import NoTarget, tp_reaches_across_nodes
+from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes, tp_reaches_across_nodes
 
 # The GPUs the target is certified on: sm_100 (B200, GB200) and sm_103 (B300,
 # GB300). They run the same kernels.
@@ -74,24 +74,19 @@
     from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 
+@routes("DeepseekV3ForCausalLM")
 def route(
-    pretrained_config: Any,
-    mapping: Any,
-    compute_capability: Tuple[int, int],
-    spec_config: Optional[Any],
-    kv_cache_quant_algo: Optional[Any],
-    model_config: Optional[ModelConfig] = None,
-    llm_args: Optional[TorchLlmArgs] = None,
+    model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int]
 ) -> type:
     """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
 
-    ``kv_cache_quant_algo`` is the KV-cache algorithm the engine will run with
-    (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone.
-
     Each target is imported where it is returned, so declining costs no import
-    of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes
-    which of these targets runs, so neither is read.
+    of modeling code.
     """
+    pretrained_config = model_config.pretrained_config
+    mapping = model_config.mapping
+    spec_config = model_config.spec_config
+    kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args)
     # The checkpoint is recognized by its config shape; see the gpt_oss routing
     # module for what shape-sniffing does and does not pin. q_lora_rank is in
     # the shape because it changes the attention weight layout: a checkpoint

tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py

@@ -11,30 +11,35 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, Tuple
+from typing import TYPE_CHECKING, Tuple
 
 from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
 from tensorrt_llm.quantization.mode import QuantAlgo
 
-from ..._router_index import NoTarget
+from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes
+
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 # (layers, hidden, routed experts, q_lora_rank). Shape-sniffing also matches
 # fine-tunes of the same shape; the accuracy gate pins identity.
 _SHAPE = (78, 6144, 256, 2048)
 
 
+@routes("GlmMoeDsaForCausalLM")
 def route(
-    pretrained_config: Any,
-    mapping: Any,
-    compute_capability: Tuple[int, int],
-    spec_config: Optional[Any],
-    kv_cache_quant_algo: Optional[Any],
+    model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int]
 ) -> type:
     """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
 
     Each target is imported where it is returned, so declining costs no import
     of modeling code.
     """
+    pretrained_config = model_config.pretrained_config
+    mapping = model_config.mapping
+    spec_config = model_config.spec_config
+    kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args)
     shape = (
         pretrained_config.num_hidden_layers,
         pretrained_config.hidden_size,

tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py

@@ -10,9 +10,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Tuple
 
-from ..._router_index import NoTarget, tp_reaches_across_nodes
+from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes, tp_reaches_across_nodes
 
 # The GPUs the target is certified on: sm_100 (B200, GB200) and sm_103 (B300,
 # GB300). They run the same kernels.
@@ -47,24 +47,19 @@
     from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 
+@routes("GptOssForCausalLM")
 def route(
-    pretrained_config: Any,
-    mapping: Any,
-    compute_capability: Tuple[int, int],
-    spec_config: Optional[Any],
-    kv_cache_quant_algo: Optional[Any],
-    model_config: Optional[ModelConfig] = None,
-    llm_args: Optional[TorchLlmArgs] = None,
+    model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int]
 ) -> type:
     """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
 
-    ``kv_cache_quant_algo`` is the KV-cache algorithm the engine will run with
-    (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone.
-
     Each target is imported where it is returned, so declining costs no import
-    of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes
-    which of these targets runs, so neither is read.
+    of modeling code.
     """
+    pretrained_config = model_config.pretrained_config
+    mapping = model_config.mapping
+    spec_config = model_config.spec_config
+    kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args)
     # The checkpoint is recognized by its config shape, the upstream idiom
     # (``is_mla``, ``is_nemotron_hybrid`` do the same). A fine-tune with the
     # same shape routes here too; the accuracy gate is what pins the identity.