From d2ddf3ce0f214ca35928cd4e967ad258bfc5b319 Mon Sep 17 00:00:00 2001 From: ethereal49 Date: Tue, 10 Nov 2026 06:27:13 -1000 Subject: [PATCH] [None][fix] Pass modeling_v2 target to bench KV-cache sizing and reporting Signed-off-by: ethereal49 --- tensorrt_llm/bench/benchmark/low_latency.py | 10 +- tensorrt_llm/bench/benchmark/throughput.py | 10 +- tensorrt_llm/bench/benchmark/utils/general.py | 42 +++++-- tensorrt_llm/bench/dataclasses/reporting.py | 7 +- tensorrt_llm/bench/tuning/settings.py | 101 ++++++++++++++++- tests/unittest/others/test_bench.py | 103 ++++++++++++++++++ 6 files changed, 253 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index ba85a97d246f..4eb36082e5d8 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/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) diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index a28d6e00a2b2..416575c4324a 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/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) diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index d2a133a9be40..e6e842cf0e49 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/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"), diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index 0f3e6954d96d..00eafe9bf6bc 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/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": diff --git a/tensorrt_llm/bench/tuning/settings.py b/tensorrt_llm/bench/tuning/settings.py index 9f47292975ac..23cb79513848 100644 --- a/tensorrt_llm/bench/tuning/settings.py +++ b/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, diff --git a/tests/unittest/others/test_bench.py b/tests/unittest/others/test_bench.py index c73e4e3c1cfa..7da9aa1008bd 100644 --- a/tests/unittest/others/test_bench.py +++ b/tests/unittest/others/test_bench.py @@ -13,12 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import io +import os from copy import deepcopy from pathlib import Path +from types import SimpleNamespace from typing import get_args from unittest import mock import pytest +import torch from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.reporting import PerfItemTuple, ReportUtility, StatsKeeper @@ -355,6 +358,106 @@ def _perf_item(rid, num_out, decoding_iteration): ) +@pytest.mark.parametrize("caller", ["sizing", "reporting"]) +@pytest.mark.parametrize( + ("mode", "attention_dp", "autotune", "expected"), + [ + ("auto", True, False, QuantAlgo.NVFP4), + ("require", True, False, QuantAlgo.NVFP4), + ("auto", False, False, QuantAlgo.FP8), + ("auto", True, True, QuantAlgo.FP8), + ("off", True, False, QuantAlgo.FP8), + ], +) +def test_bench_kv_cache_quant_uses_modeling_v2_target( + tmp_path, monkeypatch, caller, mode, attention_dp, autotune, expected +): + from tensorrt_llm._torch._experimental.modeling_v2 import _router_index + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.pyexecutor import model_loader + from tensorrt_llm.bench.benchmark.utils import general + from tensorrt_llm.bench.tuning import settings as tuning_settings + + class Nvfp4Target: + supports_nvfp4_kv_cache = True + + def route(config, llm_args, capability): + assert capability == (10, 7) + assert config.mapping.tp_size == config.mapping.moe_ep_size == 4 + assert config.moe_backend == "MEGAMOE_CUTEDSL" + assert llm_args.kv_cache_config.dtype == "nvfp4" + if not config.mapping.enable_attention_dp: + raise _router_index.NoTarget("requires attention DP") + if os.environ.get("MEGAMOE_TACTIC_AUTOTUNE") == "1": + raise _router_index.NoTarget("incompatible with tactic autotune") + return Nvfp4Target + + def load_config(model, trust_remote_code=True, **kwargs): + return ModelConfig( + pretrained_config=SimpleNamespace( + architectures=["BenchTestArchitecture"], torch_dtype=torch.bfloat16 + ), + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + **kwargs, + ) + + monkeypatch.setattr(ModelConfig, "from_pretrained", load_config) + monkeypatch.delenv("MEGAMOE_TACTIC_AUTOTUNE", raising=False) + monkeypatch.setattr( + _router_index, "registered_routes", lambda: {"BenchTestArchitecture": route} + ) + monkeypatch.setattr(_router_index, "_RESOLVED_MODES", {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 7)) + monkeypatch.setattr(tuning_settings, "get_sm_version", lambda: 107) + monkeypatch.setattr(model_loader, "get_sm_version", lambda: 107) + # Exercise the validator's SM107 capability gate without loading a CUDA target. + monkeypatch.setattr( + model_loader, + "supports_fp4_mla_attention", + lambda config: config.modeling_v2_target is Nvfp4Target, + ) + + metadata = mock.Mock(max_isl=8, avg_isl=8, avg_osl=4) + yaml_options = tmp_path / "extra.yaml" + yaml_options.write_text( + f"modeling_v2: '{mode}'\nenable_attention_dp: {str(attention_dp).lower()}\n" + "moe_config: {backend: MEGAMOE_CUTEDSL}\nkv_cache_config: {dtype: nvfp4}\n" + "tensor_parallel_size: 2\nmoe_expert_parallel_size: 2\n" + f"env_overrides: {{MEGAMOE_TACTIC_AUTOTUNE: '{int(autotune)}'}}\n" + ) + params = { + "tp": 4, + "ep": 4, + "pp": 1, + "gpus_per_node": 4, + "extra_llm_api_options": yaml_options, + } + cli_keys = {"tensor_parallel_size", "moe_expert_parallel_size"} + if caller == "sizing": + monkeypatch.setattr(tuning_settings, "get_model_config", mock.Mock()) + size = mock.Mock(return_value=(8, 128)) + monkeypatch.setattr(tuning_settings, "get_benchmark_engine_settings", size) + general.get_settings(params, metadata, "model", None, explicit_cli_keys=cli_keys) + actual = size.call_args.args[1].kv_cache_quant_algo + else: + params.update(max_batch_size=8, max_num_tokens=128) + settings = general.get_settings(params, metadata, "model", None) + runtime = RuntimeConfig( + **settings, + model="model", + extra_llm_api_options=str(yaml_options), + explicit_cli_keys=cli_keys, + ) + keeper = StatsKeeper() + keeper.register_request_perf_item(_perf_item(0, num_out=4, decoding_iteration=3)) + report = ReportUtility(keeper, metadata, runtime, mock.Mock(), runtime.get_llm_args()) + monkeypatch.setattr(report, "_query_gpu_info", lambda: {}) + actual = report.get_statistics_dict()["engine"]["kv_cache_dtype"] + assert actual == expected + assert "MEGAMOE_TACTIC_AUTOTUNE" not in os.environ + + def test_acceptance_length_weighted_by_iterations(): """AR/AL averages are weighted by per-request decoding iterations.