[None][fix] Pass modeling_v2 target to bench KV-cache sizing and reporting

tensorrt_llm/bench/benchmark/low_latency.py

@@ -248,8 +248,12 @@ def latency_command(
     if bench_env.checkpoint_path is None:
         snapshot_download(options.model, revision=bench_env.revision)
 
-    exec_settings = get_settings(params, metadata, bench_env.model,
-                                 bench_env.checkpoint_path)
+    explicit_cli_keys = collect_explicit_cli_keys()
+    exec_settings = get_settings(params,
+                                 metadata,
+                                 bench_env.model,
+                                 bench_env.checkpoint_path,
+                                 explicit_cli_keys=explicit_cli_keys)
     kwargs_max_sql = options.max_seq_len or metadata.max_sequence_length
     logger.info(f"Setting PyTorch max sequence length to {kwargs_max_sql}")
     kwargs["max_seq_len"] = kwargs_max_sql
@@ -270,7 +274,7 @@ def latency_command(
 
     exec_settings["extra_llm_api_options"] = params.get("extra_llm_api_options")
     exec_settings["iteration_log"] = options.iteration_log
-    exec_settings["explicit_cli_keys"] = collect_explicit_cli_keys()
+    exec_settings["explicit_cli_keys"] = explicit_cli_keys
 
     # Construct the runtime configuration dataclass.
     runtime_config = RuntimeConfig(**exec_settings)

tensorrt_llm/bench/benchmark/throughput.py

@@ -391,8 +391,12 @@ def throughput_command(
     if bench_env.checkpoint_path is None:
         snapshot_download(options.model, revision=bench_env.revision)
 
-    exec_settings = get_settings(params, metadata, bench_env.model,
-                                 bench_env.checkpoint_path)
+    explicit_cli_keys = collect_explicit_cli_keys()
+    exec_settings = get_settings(params,
+                                 metadata,
+                                 bench_env.model,
+                                 bench_env.checkpoint_path,
+                                 explicit_cli_keys=explicit_cli_keys)
     kwargs_max_sql = options.max_seq_len or metadata.max_sequence_length
     logger.info(f"Setting PyTorch max sequence length to {kwargs_max_sql}")
     kwargs["max_seq_len"] = kwargs_max_sql
@@ -423,7 +427,7 @@ def throughput_command(
     # LlmArgs
     exec_settings["extra_llm_api_options"] = params.pop("extra_llm_api_options")
     exec_settings["iteration_log"] = options.iteration_log
-    exec_settings["explicit_cli_keys"] = collect_explicit_cli_keys()
+    exec_settings["explicit_cli_keys"] = explicit_cli_keys
 
     # Construct the runtime configuration dataclass.
     runtime_config = RuntimeConfig(**exec_settings)

tensorrt_llm/bench/benchmark/utils/general.py

@@ -7,14 +7,11 @@
 import yaml
 
 from tensorrt_llm._errors import ConfigInputError
-from tensorrt_llm._torch.pyexecutor.model_loader import \
-    validate_and_set_kv_cache_quant
 from tensorrt_llm.bench.dataclasses.general import (DatasetMetadata,
                                                     InferenceRequest)
+from tensorrt_llm.bench.tuning import settings
 from tensorrt_llm.bench.tuning.dataclasses import (NemotronHybridConfig,
                                                    Qwen3HybridConfig)
-from tensorrt_llm.bench.tuning.settings import (get_benchmark_engine_settings,
-                                                get_model_config)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.quantization.mode import QuantAlgo
 from tensorrt_llm.version import __version__
@@ -27,14 +24,21 @@
 ALL_SUPPORTED_BACKENDS = ["pytorch"]
 
 
-def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str,
-                 model_path: Union[Path, None]) -> Dict[str, Union[str, int]]:
+def get_settings(
+        params: dict,
+        dataset_metadata: DatasetMetadata,
+        model: str,
+        model_path: Union[Path, None],
+        *,
+        explicit_cli_keys: set[str] | None = None
+) -> Dict[str, Union[str, int]]:
     """Retrieve basic runtime config for pytorch backend path
 
     Args:
         params (dict): Configuration parameters.
         model (str): Model name.
         model_path (Union[Path, None]): Path to the model.
+        explicit_cli_keys: LLM argument names set explicitly on the CLI.
     Returns:
         Dict[str, Union[str, int]]: Properties for runtime config.
     """
@@ -45,6 +49,7 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str,
     kv_cache_dtype = "auto"
     mamba_ssm_cache_dtype = params.get("mamba_ssm_cache_dtype", "auto")
     kv_cache_config = {}
+    llm_args_dict = {}
     if extra_llm_api_options:
         with open(extra_llm_api_options, 'r') as f:
             llm_args_dict = yaml.safe_load(f)
@@ -81,7 +86,7 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str,
             f"Initial settings: max_batch_size={max_batch_size} (user-provided), "
             f"max_num_tokens={max_num_tokens} (user-provided).")
     else:
-        model_config = get_model_config(model, model_path)
+        model_config = settings.get_model_config(model, model_path)
 
         if isinstance(model_config,
                       (NemotronHybridConfig, Qwen3HybridConfig
@@ -98,9 +103,28 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str,
             kv_cache_dtype = _KV_CACHE_MAP.get(
                 tllm_model_config.quant_config.quant_algo, "auto")
 
-        validate_and_set_kv_cache_quant(tllm_model_config, kv_cache_dtype)
+        tllm_model_config = settings.resolve_benchmark_kv_cache_quant(
+            tllm_model_config,
+            model,
+            kv_cache_dtype,
+            {
+                "model": model,
+                "trust_remote_code": True,
+                "skip_tokenizer_init": True,
+                "tensor_parallel_size": mapping["tp_size"],
+                "pipeline_parallel_size": mapping["pp_size"],
+                "moe_expert_parallel_size": mapping["moe_ep_size"],
+                "moe_cluster_parallel_size": mapping["moe_cluster_size"],
+                "gpus_per_node": mapping["gpus_per_node"],
+                "kv_cache_config": {
+                    "dtype": kv_cache_dtype
+                },
+            },
+            extra_options=llm_args_dict,
+            explicit_cli_keys=explicit_cli_keys,
+        )
 
-        max_batch_size, max_num_tokens = get_benchmark_engine_settings(
+        max_batch_size, max_num_tokens = settings.get_benchmark_engine_settings(
             model_config,
             tllm_model_config.quant_config,
             params.get("tp"),

tensorrt_llm/bench/dataclasses/reporting.py

@@ -25,13 +25,12 @@
 except ImportError:
     pynvml = None
 
-from tensorrt_llm._torch.pyexecutor.model_loader import \
-    validate_and_set_kv_cache_quant
 from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig
 from tensorrt_llm.bench.dataclasses.general import DatasetMetadata
 from tensorrt_llm.bench.dataclasses.statistics import (BenchmarkStatistics,
                                                        PercentileStats,
                                                        RequestRecord)
+from tensorrt_llm.bench.tuning.settings import resolve_benchmark_kv_cache_quant
 from tensorrt_llm.llmapi import KvCacheConfig
 from tensorrt_llm.logger import Logger
 from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode
@@ -474,7 +473,9 @@ def get_statistics_dict(self) -> Dict[str, Any]:
         model_config = ModelConfig.from_pretrained(model,
                                                    trust_remote_code=True)
 
-        validate_and_set_kv_cache_quant(model_config, kv_cache_dtype)
+        model_config = resolve_benchmark_kv_cache_quant(model_config, model,
+                                                        kv_cache_dtype,
+                                                        self.kwargs)
 
         stats_dict["engine"] |= {
             "backend":

tensorrt_llm/bench/tuning/settings.py

@@ -1,24 +1,40 @@
 from __future__ import annotations
 
+import os
+from contextlib import contextmanager
+from dataclasses import replace
 from pathlib import Path
-from typing import Tuple
+from typing import TYPE_CHECKING, Tuple
 
+import torch
+from pydantic import BaseModel
+
+from tensorrt_llm._torch._experimental.modeling_v2 import modeling_v2_resolve
 from tensorrt_llm._torch.pyexecutor.config_utils import (
     is_nemotron_hybrid,
     is_qwen3_hybrid,
     load_pretrained_config,
 )
+from tensorrt_llm._torch.pyexecutor.model_loader import (
+    ModelLoader,
+    construct_checkpoint_loader,
+    validate_and_set_kv_cache_quant,
+)
+from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.bench.tuning.dataclasses import (
     ModelConfig,
     NemotronHybridConfig,
     Qwen3HybridConfig,
 )
 from tensorrt_llm.bench.tuning.heuristics import calc_engine_setting
-from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+from tensorrt_llm.llmapi.llm_args import TorchLlmArgs, update_llm_args_with_extra_dict
 from tensorrt_llm.llmapi.llm_utils import QuantConfig
 from tensorrt_llm.logger import logger
 from tensorrt_llm.quantization.mode import QuantAlgo
 
+if TYPE_CHECKING:
+    from tensorrt_llm._torch.model_config import ModelConfig as TorchModelConfig
+
 TUNED_QUANTS = {
     QuantAlgo.NVFP4,
     QuantAlgo.FP8,
@@ -32,6 +48,87 @@
 DEFAULT_MAX_NUM_TOKENS = TorchLlmArgs.model_fields["max_num_tokens"].default
 
 
+@contextmanager
+def _routing_environment(overrides: dict[str, str] | None):
+    """Make the LLM's environment overrides visible to the bench router."""
+    overrides = overrides or {}
+    previous = {name: os.environ.get(name) for name in overrides}
+    try:
+        os.environ.update({name: str(value) for name, value in overrides.items()})
+        yield
+    finally:
+        for name, value in previous.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
+def resolve_benchmark_kv_cache_quant(
+    config: TorchModelConfig,
+    model: str | Path,
+    kv_cache_dtype: str,
+    llm_options: dict,
+    *,
+    extra_options: dict | None = None,
+    explicit_cli_keys: set[str] | None = None,
+) -> TorchModelConfig:
+    """Resolve target-dependent cache quantization without loading model weights."""
+    requests_nvfp4 = kv_cache_dtype == "nvfp4" or (
+        kv_cache_dtype == "auto" and config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4
+    )
+    # Only SM107's NVFP4 fallback depends on the target implementation.
+    if requests_nvfp4 and torch.cuda.is_available() and get_sm_version() == 107:
+        if extra_options:
+            llm_options = update_llm_args_with_extra_dict(
+                llm_options, extra_options, explicit_cli_keys=explicit_cli_keys
+            )
+        if llm_options.get("modeling_v2") != "off":
+            # Model defaults mutate nested configs; keep the caller's options intact.
+            llm_args = TorchLlmArgs(
+                **{
+                    name: value.model_copy(deep=True) if isinstance(value, BaseModel) else value
+                    for name, value in llm_options.items()
+                }
+            )
+            checkpoint_loader = construct_checkpoint_loader(
+                llm_args.checkpoint_loader,
+                llm_args.checkpoint_format,
+                mx_config=llm_args.mx_config,
+                checkpoint_io_policy=llm_args.checkpoint_io_policy,
+                load_format=llm_args.load_format,
+                partial_model_loading=llm_args.is_partial_model_loading,
+            )
+            checkpoint_dir = str(model)
+            try:
+                llm_args = ModelLoader.load_config_and_apply_defaults(
+                    checkpoint_dir, llm_args, checkpoint_loader
+                )
+                routing_config = checkpoint_loader.load_config(
+                    checkpoint_dir,
+                    trust_remote_code=llm_args.trust_remote_code,
+                    mapping=llm_args.parallel_config.to_mapping(),
+                    spec_config=llm_args.speculative_config,
+                    mm_encoder_only=llm_args.mm_encoder_only,
+                    attn_backend=llm_args.attn_backend,
+                    sparse_attention_config=llm_args.sparse_attention_config,
+                    moe_backend=llm_args.moe_config.backend,
+                    moe_load_balancer=llm_args.moe_config.load_balancer,
+                    model_kwargs=llm_args.model_kwargs,
+                )
+            finally:
+                checkpoint_loader.cleanup()
+            # Bench runs on the driver; the engine checks agreement across workers.
+            with _routing_environment(llm_args.env_overrides):
+                target = modeling_v2_resolve(routing_config, llm_args)
+            if target is not None:
+                config = replace(routing_config, modeling_v2_target=target)
+                kv_cache_dtype = llm_args.kv_cache_config.dtype
+
+    validate_and_set_kv_cache_quant(config, kv_cache_dtype)
+    return config
+
+
 def get_benchmark_engine_settings(
     model_config: ModelConfig,
     quant_config: QuantConfig,