From 7d764b685e4fbbf58be9bf64cd3065e0bd602cbf Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 15 Jan 2026 16:53:06 -0800 Subject: [PATCH 01/15] [None][chore] Standardize LlmArgs initialization and parsing around Pydantic Signed-off-by: Anish Shanbhag --- CODING_GUIDELINES.md | 2 + tensorrt_llm/_torch/auto_deploy/llm.py | 13 +- tensorrt_llm/bench/benchmark/__init__.py | 59 - tensorrt_llm/bench/benchmark/low_latency.py | 105 +- tensorrt_llm/bench/benchmark/throughput.py | 99 +- tensorrt_llm/bench/benchmark/utils/general.py | 232 ++-- .../bench/dataclasses/configuration.py | 194 +-- tensorrt_llm/bench/dataclasses/reporting.py | 70 +- tensorrt_llm/commands/eval.py | 59 +- tensorrt_llm/commands/serve.py | 268 +--- tensorrt_llm/llmapi/build_cache.py | 58 +- tensorrt_llm/llmapi/llm.py | 159 +-- tensorrt_llm/llmapi/llm_args.py | 1233 ++++++----------- tensorrt_llm/llmapi/llm_create.py | 85 ++ tensorrt_llm/llmapi/llm_utils.py | 6 +- tensorrt_llm/llmapi/mm_encoder.py | 3 +- tensorrt_llm/mapping.py | 24 +- .../accuracy/test_disaggregated_serving.py | 4 +- ...g_config_deepseek_v3_lite_empty_batch.yaml | 8 - .../unit/singlegpu/shim/test_llm_config.py | 2 +- .../api_stability/references/llm.yaml | 2 +- tests/unittest/llmapi/test_config_database.py | 6 +- tests/unittest/llmapi/test_llm.py | 14 +- tests/unittest/llmapi/test_llm_args.py | 140 +- .../all_models/llmapi/tensorrt_llm/1/model.py | 2 +- 25 files changed, 919 insertions(+), 1928 deletions(-) create mode 100644 tensorrt_llm/llmapi/llm_create.py diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 70f0c1bfbed6..41487c9e427b 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -408,6 +408,8 @@ foo.SomeClass() ##### Classes and Functions Use the [Google style](https://google.github.io/styleguide/pyguide.html), which can be parsed by Sphinx. +When defining any user-facing configuration classes (e.g. `LlmArgs` or any class used in its fields), always use Pydantic classes rather than dataclasses or vanilla classes. + ##### Attributes and Variables Attributes and variables can be documented inline. Attribute docstrings will be rendered under the docstring for the class. For example: ```python diff --git a/tensorrt_llm/_torch/auto_deploy/llm.py b/tensorrt_llm/_torch/auto_deploy/llm.py index 38c7524db239..c36feb71958f 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm.py +++ b/tensorrt_llm/_torch/auto_deploy/llm.py @@ -1,11 +1,12 @@ import types -from typing import Any, Dict, List, Optional, Tuple +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union import torch from ...executor.result import CompletionOutput from ...inputs.registry import DefaultInputProcessor, ExtraProcessedInputs -from ...llmapi.llm import RequestOutput, _TorchLLM +from ...llmapi.llm import BaseLLM, RequestOutput, _TorchLLM from ...llmapi.tokenizer import TokenizerBase, TransformersTokenizer, tokenizer_factory from ...sampling_params import SamplingParams from .distributed import common as dist_ad @@ -112,9 +113,9 @@ def factory(self) -> ModelFactory: self._factory = self.args.create_factory() return self._factory - def __init__(self, *args, **kwargs): - kwargs["backend"] = "_autodeploy" - super().__init__(*args, **kwargs) + def __init__(self, model: Union[str, Path], **kwargs: Any) -> None: + args = LlmArgs(model=model, **kwargs) + BaseLLM.__init__(args) def _try_load_tokenizer(self) -> Optional[TokenizerBase]: if self.args.skip_tokenizer_init: @@ -159,7 +160,7 @@ class DemoLLM(LLM): """ def __init__(self, **kwargs): - self.args: LlmArgs = LlmArgs.from_kwargs(**kwargs) + self.args: LlmArgs = LlmArgs(**kwargs) self.mpi_session = None self.runtime_context = None diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 4c07816a4059..63165db9768a 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -4,11 +4,8 @@ from pydantic import AliasChoices, BaseModel, Field -from tensorrt_llm import LLM as PyTorchLLM -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter from tensorrt_llm.bench.build.build import get_model_config -from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.logger import logger @@ -21,7 +18,6 @@ class GeneralExecSettings(BaseModel): backend: str = Field( default="pytorch", description="The backend to use when running benchmarking") - beam_width: int = Field(default=1, description="Number of search beams") model_path: Optional[Path] = Field(default=None, description="Path to model checkpoint") concurrency: int = Field( @@ -36,15 +32,8 @@ class GeneralExecSettings(BaseModel): default=-1, description="End-of-sequence token ID, -1 to disable EOS") iteration_log: Optional[Path] = Field( default=None, description="Path where iteration logging is written") - kv_cache_percent: float = Field( - default=0.90, - validation_alias=AliasChoices("kv_cache_percent", - "kv_cache_free_gpu_mem_fraction"), - description="Percentage of memory for KV Cache after model load") max_input_len: int = Field(default=4096, description="Maximum input sequence length") - max_seq_len: Optional[int] = Field(default=None, - description="Maximum sequence length") modality: Optional[str] = Field( default=None, description="Modality of multimodal requests") model: Optional[str] = Field(default=None, description="Model name or path") @@ -75,54 +64,6 @@ def checkpoint_path(self) -> Path: return self.model_path or self.model -def ignore_trt_only_args(kwargs: dict, backend: str): - """Ignore TensorRT-only arguments for non-TensorRT backends. - - Args: - kwargs: Dictionary of keyword arguments to be passed to the LLM constructor. - backend: The backend type (e.g., "pytorch", "_autodeploy"). - """ - trt_only_args = [ - "batching_type", - "normalize_log_probs", - "extended_runtime_perf_knob_config", - ] - for arg in trt_only_args: - if kwargs.pop(arg, None): - logger.warning(f"Ignore {arg} for {backend} backend.") - - -def get_llm(runtime_config: RuntimeConfig, kwargs: dict): - """Create and return an appropriate LLM instance based on the backend configuration. - - Args: - runtime_config: Runtime configuration containing backend selection and settings. - kwargs: Additional keyword arguments to pass to the LLM constructor. - - Returns: - An instance of the appropriate LLM class for the specified backend. - """ - llm_cls = LLM - - if runtime_config.backend != None: - ignore_trt_only_args(kwargs, runtime_config.backend) - - if runtime_config.iteration_log is not None: - kwargs["enable_iter_perf_stats"] = True - - if runtime_config.backend == 'pytorch': - llm_cls = PyTorchLLM - - elif runtime_config.backend == "_autodeploy": - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - - kwargs["world_size"] = kwargs.pop("tensor_parallel_size", None) - llm_cls = AutoDeployLLM - - llm = llm_cls(**kwargs) - return llm - - def get_general_cli_options( params: Dict, bench_env: BenchmarkEnvironment) -> GeneralExecSettings: """Get general execution settings from command line parameters. diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index c36a7dc43e49..5637ed7aad1d 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -8,27 +8,23 @@ import yaml from click_option_group import (MutuallyExclusiveOptionGroup, OptionGroup, optgroup) -from huggingface_hub import snapshot_download from tensorrt_llm.bench.benchmark import (generate_json_report, - get_general_cli_options, get_llm) + get_general_cli_options) from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark -from tensorrt_llm.bench.benchmark.utils.general import generate_warmup_dataset -from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig +from tensorrt_llm.bench.benchmark.utils.general import ( + ALL_SUPPORTED_BACKENDS, generate_warmup_dataset, + get_exec_settings_for_backend, update_sampler_args_with_extra_options) from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility -from tensorrt_llm.llmapi import CapacitySchedulerPolicy -from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode - -# isort: off -from tensorrt_llm.bench.benchmark.utils.general import ( - get_settings_from_engine, get_settings, - update_sampler_args_with_extra_options, ALL_SUPPORTED_BACKENDS) -# isort: on from tensorrt_llm.bench.utils.data import (create_dataset_from_stream, initialize_tokenizer, update_metadata_for_multimodal) +from tensorrt_llm.llmapi import CapacitySchedulerPolicy +from tensorrt_llm.llmapi.llm_args import CudaGraphConfig +from tensorrt_llm.llmapi.llm_create import create_llm_from_llm_args from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode from tensorrt_llm.sampling_params import SamplingParams @@ -222,67 +218,25 @@ def latency_command( # The accurate table for multimodal models will be logged after the benchmark is done. logger.info(metadata.get_summary_for_print()) - # Engine configuration parsing for PyTorch backend - kwargs = {} - if options.backend and options.backend.lower( - ) in ALL_SUPPORTED_BACKENDS and options.backend.lower() != "tensorrt": - 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) - 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 - elif options.backend.lower() == "tensorrt": - assert options.max_seq_len is None, ( - "max_seq_len is not a runtime parameter for C++ backend") - exec_settings, build_cfg = get_settings_from_engine(options.engine_dir) - engine_max_seq_len = build_cfg["max_seq_len"] + # Engine configuration parsing + exec_settings = get_exec_settings_for_backend(params, metadata, options, + bench_env) - if metadata.max_sequence_length > engine_max_seq_len: - raise RuntimeError( - f"Engine supports a max sequence of {engine_max_seq_len}. Provided " - "dataset contains a maximum sequence of " - f"{metadata.max_sequence_length}. Please rebuild a new engine to" - "support this dataset.") - else: - raise click.BadParameter( - f"{options.backend} is not a known backend, check help for available options.", - param_hint="backend") + exec_settings.revision = bench_env.revision - exec_settings["model"] = options.model - exec_settings["revision"] = bench_env.revision - engine_tokens = exec_settings["settings_config"]["max_num_tokens"] - - # Update configuration with runtime options - exec_settings["settings_config"][ - "kv_cache_percent"] = options.kv_cache_percent - exec_settings["settings_config"]["max_batch_size"] = 1 - exec_settings["settings_config"]["max_num_tokens"] = engine_tokens - exec_settings["settings_config"]["beam_width"] = options.beam_width - exec_settings["settings_config"]["chunking"] = False - exec_settings["settings_config"][ - "scheduler_policy"] = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT - - # Performance options - exec_settings["performance_options"]["cuda_graphs"] = True - exec_settings["performance_options"]["multi_block_mode"] = True - - exec_settings["extra_llm_api_options"] = params.get("extra_llm_api_options") + # Update llm_args with runtime options + llm_args = exec_settings.llm_args + llm_args.max_batch_size = 1 + llm_args.cuda_graph_config = CudaGraphConfig(max_batch_size=1) + llm_args.enable_chunked_prefill = False + llm_args.scheduler_config.capacity_scheduler_policy = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT # Decoding Options if medusa_choices is not None: with open(medusa_choices, "r") as medusa_yml: - exec_settings["decoding_config"]["medusa_choices"] = \ - yaml.load(medusa_yml, Loader=yaml.SafeLoader) - - # Construct the runtime configuration dataclass. - runtime_config = RuntimeConfig(**exec_settings) - - llm = None - kwargs = kwargs | runtime_config.get_llm_args() - kwargs['backend'] = options.backend + medusa_choices_data = yaml.load(medusa_yml, Loader=yaml.SafeLoader) + if llm_args.decoding_config is not None: + llm_args.decoding_config.medusa_choices = medusa_choices_data # Set environment variables for setting runtime options. default_env_overrides = { @@ -292,23 +246,22 @@ def latency_command( "TRTLLM_ENABLE_PDL": "1", } # Update defaults with existing overrides (user preference takes priority) - default_env_overrides.update(kwargs.get("env_overrides", {})) - kwargs["env_overrides"] = default_env_overrides + llm_args.env_overrides = default_env_overrides | llm_args.env_overrides + llm = None try: logger.info("Setting up latency benchmark.") + llm = create_llm_from_llm_args(llm_args) - llm = get_llm(runtime_config, kwargs) - - ignore_eos = True if runtime_config.decoding_config.decoding_mode == SpeculativeDecodingMode.NONE else False + ignore_eos = True if llm_args.decoding_config.decoding_mode == SpeculativeDecodingMode.NONE else False eos_id = tokenizer.eos_token_id if not ignore_eos else -1 pad_id = tokenizer.pad_token_id if not ignore_eos else -1 sampler_args = { "end_id": eos_id, "pad_id": pad_id, - "n": options.beam_width, - "use_beam_search": options.beam_width > 1 + "n": llm_args.max_beam_width, + "use_beam_search": llm_args.max_beam_width > 1 } sampler_args = update_sampler_args_with_extra_options( @@ -354,8 +307,8 @@ def latency_command( # For multimodal models, we need to update the metadata with the correct input lengths metadata = update_metadata_for_multimodal(metadata, statistics) - report_utility = ReportUtility(statistics, metadata, runtime_config, - logger, kwargs, True) + report_utility = ReportUtility(statistics, metadata, exec_settings, + logger, llm_args, True) # Generate reports for statistics, output tokens, and request info. generate_json_report(options.report_json, report_utility.get_statistics_dict) diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 6e9f6b87eccb..e109d3d4a131 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -8,29 +8,24 @@ import click from click_option_group import (MutuallyExclusiveOptionGroup, OptionGroup, optgroup) -from huggingface_hub import snapshot_download from tensorrt_llm.bench.benchmark import (GeneralExecSettings, generate_json_report, - get_general_cli_options, get_llm) + get_general_cli_options) from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark -from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir - -# isort: off -from tensorrt_llm.bench.benchmark.utils.general import ( - get_settings_from_engine, get_settings, ALL_SUPPORTED_BACKENDS) -# isort: on from tensorrt_llm.bench.benchmark.utils.general import ( - generate_warmup_dataset, update_sampler_args_with_extra_options) -from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig + ALL_SUPPORTED_BACKENDS, generate_warmup_dataset, + get_exec_settings_for_backend, update_sampler_args_with_extra_options) from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility from tensorrt_llm.bench.utils.data import (create_dataset_from_stream, initialize_tokenizer, update_metadata_for_multimodal) from tensorrt_llm.llmapi import CapacitySchedulerPolicy +from tensorrt_llm.llmapi.llm_create import create_llm_from_llm_args from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir @click.command(name="throughput") @@ -299,16 +294,12 @@ def throughput_command( # Parameters from CLI image_data_format: str = params.get("image_data_format", "pt") data_device: str = params.get("data_device", "cpu") - no_skip_tokenizer_init: bool = params.get("no_skip_tokenizer_init", False) # Get general CLI options using the centralized function options: GeneralExecSettings = get_general_cli_options(params, bench_env) tokenizer = initialize_tokenizer(options.checkpoint_path) # Extract throughput-specific options not handled by GeneralExecSettings - max_batch_size = params.get("max_batch_size") - max_num_tokens = params.get("max_num_tokens") - enable_chunked_context: bool = params.get("enable_chunked_context") scheduler_policy: str = params.get("scheduler_policy") custom_module_dirs: list[Path] = params.pop("custom_module_dirs", []) @@ -320,9 +311,6 @@ def throughput_command( f"Failed to import custom module from {custom_module_dir}: {e}") raise e - # Runtime kwargs and option tracking. - kwargs = {} - # Dataset Loading and Preparation with open(options.dataset_path, "r") as dataset: metadata, requests = create_dataset_from_stream( @@ -348,79 +336,26 @@ def throughput_command( logger.info(metadata.get_summary_for_print()) # Engine configuration parsing - if options.backend and options.backend.lower( - ) in ALL_SUPPORTED_BACKENDS and options.backend.lower() != "tensorrt": - # If we're dealing with a model name, perform a snapshot download to - # make sure we have a local copy of the model. - 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) - 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 - elif options.backend.lower() == "tensorrt": - assert options.max_seq_len is None, ( - "max_seq_len is not a runtime parameter for C++ backend") - exec_settings, build_cfg = get_settings_from_engine(options.engine_dir) - engine_max_seq_len = build_cfg["max_seq_len"] + exec_settings = get_exec_settings_for_backend(params, metadata, options, + bench_env) - # TODO: Verify that the engine can handle the max/min ISL/OSL. - if metadata.max_sequence_length > engine_max_seq_len: - raise RuntimeError( - f"Engine supports a max sequence of {engine_max_seq_len}. " - "Provided dataset contains a maximum sequence of " - f"{metadata.max_sequence_length}. Please rebuild a new engine " - "to support this dataset.") - else: - raise click.BadParameter( - f"{options.backend} is not a known backend, check help for available options.", - param_hint="backend") + exec_settings.revision = bench_env.revision + exec_settings.iteration_log = options.iteration_log - exec_settings["model"] = options.model - exec_settings["revision"] = bench_env.revision - engine_bs = exec_settings["settings_config"]["max_batch_size"] - engine_tokens = exec_settings["settings_config"]["max_num_tokens"] + # Update llm_args with runtime options + policy = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT if scheduler_policy == "guaranteed_no_evict" else CapacitySchedulerPolicy.MAX_UTILIZATION + exec_settings.llm_args.scheduler_config.capacity_scheduler_policy = policy - # Runtime Options - runtime_max_bs = max_batch_size or engine_bs - runtime_max_tokens = max_num_tokens or engine_tokens - - # Update configuration with runtime options - exec_settings["settings_config"][ - "kv_cache_percent"] = options.kv_cache_percent - exec_settings["settings_config"]["max_batch_size"] = runtime_max_bs - exec_settings["settings_config"]["max_num_tokens"] = runtime_max_tokens - exec_settings["settings_config"]["beam_width"] = options.beam_width - exec_settings["settings_config"][ - "scheduler_policy"] = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT if scheduler_policy == "guaranteed_no_evict" else CapacitySchedulerPolicy.MAX_UTILIZATION - exec_settings["settings_config"]["chunking"] = enable_chunked_context - - # Dynamic runtime features. - exec_settings["settings_config"]["dynamic_max_batch_size"] = True - - # LlmArgs - exec_settings["extra_llm_api_options"] = params.pop("extra_llm_api_options") - exec_settings["iteration_log"] = options.iteration_log - - # Construct the runtime configuration dataclass. - runtime_config = RuntimeConfig(**exec_settings) llm = None - try: logger.info("Setting up throughput benchmark.") - kwargs = kwargs | runtime_config.get_llm_args() - kwargs['skip_tokenizer_init'] = not no_skip_tokenizer_init - kwargs['backend'] = options.backend - - llm = get_llm(runtime_config, kwargs) + llm = create_llm_from_llm_args(exec_settings.llm_args) sampler_args = { "end_id": options.eos_id, "pad_id": options.eos_id, - "n": options.beam_width, - "use_beam_search": options.beam_width > 1 + "n": exec_settings.llm_args.max_beam_width, + "use_beam_search": exec_settings.llm_args.max_beam_width > 1 } sampler_args = update_sampler_args_with_extra_options( sampler_args, params.pop("sampler_options")) @@ -464,8 +399,8 @@ def throughput_command( # For multimodal models, we need to update the metadata with the correct input lengths metadata = update_metadata_for_multimodal(metadata, statistics) - report_utility = ReportUtility(statistics, metadata, runtime_config, - logger, kwargs, options.streaming) + report_utility = ReportUtility(statistics, metadata, exec_settings, + logger, options.streaming) # Generate reports for statistics, output tokens, and request info. generate_json_report(options.report_json, report_utility.get_statistics_dict) diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index 3a35008daba3..cdbd17cac592 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -1,19 +1,30 @@ from __future__ import annotations -import json +from importlib.metadata import version from pathlib import Path from random import choices, shuffle from typing import Dict, List, Tuple, Union import yaml +from huggingface_hub import snapshot_download from tensorrt_llm._torch.pyexecutor.model_loader import \ validate_and_set_kv_cache_quant +from tensorrt_llm.bench.benchmark import GeneralExecSettings from tensorrt_llm.bench.build.build import (get_benchmark_engine_settings, get_model_config) from tensorrt_llm.bench.build.dataclasses import NemotronHybridConfig -from tensorrt_llm.bench.dataclasses.general import (DatasetMetadata, +from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig +from tensorrt_llm.bench.dataclasses.general import (BenchmarkEnvironment, + DatasetMetadata, InferenceRequest) +from tensorrt_llm.builder import get_engine_version +from tensorrt_llm.llmapi import (BatchingType, CapacitySchedulerPolicy, + ContextChunkingPolicy, + ExtendedRuntimePerfKnobConfig) +from tensorrt_llm.llmapi.llm_args import (BaseLlmArgs, CudaGraphConfig, + TrtLlmArgs) +from tensorrt_llm.llmapi.llm_create import get_llm_args_from_cli_params from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo @@ -25,48 +36,63 @@ ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy", "tensorrt"] +def apply_general_llm_args_overrides(llm_args: BaseLlmArgs, max_batch_size: int, + max_num_tokens: int) -> BaseLlmArgs: + llm_args.max_batch_size = max_batch_size + llm_args.max_num_tokens = max_num_tokens + + llm_args.scheduler_config.capacity_scheduler_policy = CapacitySchedulerPolicy.MAX_UTILIZATION + if llm_args.enable_chunked_prefill: + llm_args.scheduler_config.context_chunking_policy = ContextChunkingPolicy.FIRST_COME_FIRST_SERVED + + # Disable KV reuse for benchmarking + llm_args.kv_cache_config.enable_block_reuse = False + + if isinstance(llm_args, TrtLlmArgs): + llm_args.extended_runtime_perf_knob_config = ExtendedRuntimePerfKnobConfig( + cuda_graph_mode=True, + multi_block_mode=True, + cuda_graph_cache_size=1000, + ) + llm_args.batching_type = BatchingType.INFLIGHT + + # Misc overrides + llm_args.skip_tokenizer_init = True + llm_args.trust_remote_code = True + llm_args.cuda_graph_config = CudaGraphConfig( + enable_padding=True, + max_batch_size=max_batch_size, + ) + + def get_settings_from_engine( - engine_path: Path -) -> Tuple[Dict[str, Union[str, int]], Dict[str, Union[str, int]]]: + engine_path: Path) -> Tuple[RuntimeConfig, Dict[str, Union[str, int]]]: """Retrieve basic engine information. Args: engine_path (Path): Path to a TRT-LLM engine directory. Returns: - Tuple[Dict[str, Union[str, int]], Dict[str, Union[str, int]]]: Engine - properties parsed from the engine at engine_path. + Tuple[RuntimeConfig, Dict[str, Union[str, int]]]: RuntimeConfig and engine properties parsed from the engine at engine_path. """ - config_path = engine_path / "config.json" - runtime_config = {} + # TrtLlmArgs automatically loads build_config and parallel_config from the engine + llm_args = TrtLlmArgs(model=str(engine_path.absolute())) + apply_general_llm_args_overrides(llm_args, + llm_args.build_config.max_batch_size, + llm_args.build_config.max_num_tokens) - with open(config_path, "r") as config_json: - config = json.load(config_json) + runtime_config = RuntimeConfig( + sw_version=get_engine_version(str(engine_path)), + engine_dir=engine_path.absolute(), + backend="tensorrt", + llm_args=llm_args, + ) - mapping = config["pretrained_config"]["mapping"] - engine_build_cfg = config["build_config"] - - executor_settings = { - "max_batch_size": engine_build_cfg["max_batch_size"], - "max_num_tokens": engine_build_cfg["max_num_tokens"], - } - - runtime_config.update({ - "sw_version": config["version"], - "engine_dir": str(engine_path.absolute()), - "settings_config": executor_settings, - "mapping": mapping, - }) - - runtime_config["performance_options"] = {} - runtime_config["decoding_config"] = { - "decoding_mode": engine_build_cfg["speculative_decoding_mode"] - } - return runtime_config, engine_build_cfg + return runtime_config, llm_args.build_config.model_dump() def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, - model_path: Union[Path, None]) -> Dict[str, Union[str, int]]: + model_path: Union[Path, None]) -> RuntimeConfig: """Retrieve basic runtime config for pytorch backend path Args: @@ -74,40 +100,39 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, model (str): Model name. model_path (Union[Path, None]): Path to the model. Returns: - Dict[str, Union[str, int]]: Properties for runtime config. + RuntimeConfig: Properties for runtime config. """ - extra_llm_api_options = params.get("extra_llm_api_options") - enable_chunked_prefill = params.get("enable_chunked_prefill", False) - - kv_cache_dtype = "auto" mamba_ssm_cache_dtype = params.get("mamba_ssm_cache_dtype", "auto") - kv_cache_config = {} - if extra_llm_api_options: - with open(extra_llm_api_options, 'r') as f: - llm_args_dict = yaml.safe_load(f) - kv_cache_config = llm_args_dict.get("kv_cache_config", { - "dtype": "auto", - }) - kv_cache_dtype = kv_cache_config.get("dtype", "auto") - mamba_ssm_cache_dtype = kv_cache_config.get("mamba_ssm_cache_dtype", - mamba_ssm_cache_dtype) - - enable_chunked_prefill = llm_args_dict.get("enable_chunked_prefill", - enable_chunked_prefill) - - mapping = { - "pp_size": params.get("pp"), - "tp_size": params.get("tp"), - "world_size": params.get("pp") * params.get("tp"), - "moe_ep_size": params.get("ep"), - "moe_cluster_size": params.get("cluster_size"), - "gpus_per_node": params.get("gpus_per_node"), - } - - if params.get("max_batch_size") and params.get("max_num_tokens"): - logger.info("Use user-provided max batch size and max num tokens.") - max_batch_size, max_num_tokens = params.get( - "max_batch_size"), params.get("max_num_tokens") + # TODO: unify CLI parameter naming across trtllm-serve / trtllm-bench to avoid the need to manually + # specify parameters here + llm_args = get_llm_args_from_cli_params( + model=model, + backend=params.get("backend"), + extra_llm_api_options=params.get("extra_llm_api_options"), + max_batch_size=params.get("max_batch_size"), + max_num_tokens=params.get("max_num_tokens"), + max_seq_len=params.get("max_seq_len"), + max_input_len=params.get("max_input_len"), + max_beam_width=params.get("beam_width"), + tensor_parallel_size=params.get("tp"), + pipeline_parallel_size=params.get("pp"), + moe_expert_parallel_size=params.get("ep"), + moe_cluster_parallel_size=params.get("cluster_size"), + enable_chunked_prefill=params.get("enable_chunked_context"), + skip_tokenizer_init=not params.get("no_skip_tokenizer_init", False), + ) + + tp_size = llm_args.tensor_parallel_size + pp_size = llm_args.pipeline_parallel_size + + if llm_args.max_seq_len is None: + llm_args.max_seq_len = dataset_metadata.max_sequence_length + + # Check if max_batch_size and max_num_tokens were explicitly set in either CLI params or + # config YAML file + if {"max_batch_size", "max_num_tokens"}.issubset(llm_args.model_fields_set): + logger.info("Using user-provided max batch size and max num tokens.") + max_batch_size, max_num_tokens = llm_args.max_batch_size, llm_args.max_num_tokens else: model_config = get_model_config(model, model_path) @@ -115,8 +140,8 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, model_config.set_mamba_ssm_cache_dtype(mamba_ssm_cache_dtype) from tensorrt_llm._torch.model_config import ModelConfig - model = model_path or model - tllm_model_config = ModelConfig.from_pretrained(model, + resolved_model = model_path or model + tllm_model_config = ModelConfig.from_pretrained(resolved_model, trust_remote_code=True) if (kv_cache_dtype is None @@ -129,8 +154,8 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, max_batch_size, max_num_tokens = get_benchmark_engine_settings( model_config, tllm_model_config.quant_config, - params.get("tp"), - params.get("pp"), + tp_size, + pp_size, dataset_metadata.avg_isl, dataset_metadata.avg_osl, params.get("kv_cache_free_gpu_mem_fraction"), @@ -142,7 +167,7 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, ) # If chunked prefill is disabled, we need to ensure that the max_num_tokens is at least the max_isl - if not enable_chunked_prefill: + if not llm_args.enable_chunked_prefill: logger.warning( f"Chunked prefill is disabled, but max_num_tokens ({max_num_tokens}) is less than the max ISL ({dataset_metadata.max_isl}). " f"Forcing max_num_tokens to {dataset_metadata.max_isl + max_batch_size}." @@ -154,36 +179,47 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, # Expecting this to be the max of chunk block and max_num_tokens. pass - cuda_graph_config = { - "enable_padding": True, - "max_batch_size": max_batch_size - } - - kv_cache_config["dtype"] = kv_cache_dtype - kv_cache_config["mamba_ssm_cache_dtype"] = mamba_ssm_cache_dtype - - pyt_options = { - "cuda_graph_config": cuda_graph_config, - "kv_cache_config": kv_cache_config, - } - - backend = params.get("backend", "pytorch") - return { - "sw_version": "1.2", - "model_path": model_path, - "settings_config": { - "max_batch_size": int(max_batch_size), - "max_num_tokens": int(max_num_tokens), - "chunking": enable_chunked_prefill, - }, - "mapping": mapping, - "backend": backend, - "decoding_config": {}, - "performance_options": { - "cuda_graphs": True, - "pytorch_config": pyt_options, - } - } + apply_general_llm_args_overrides(llm_args, max_batch_size, max_num_tokens) + + return RuntimeConfig( + sw_version=version("tensorrt_llm"), + model_path=model_path, + backend=params.get("backend", "pytorch"), + llm_args=llm_args, + ) + + +def get_exec_settings_for_backend( + params: dict, + metadata: DatasetMetadata, + options: GeneralExecSettings, + bench_env: "BenchmarkEnvironment", +) -> RuntimeConfig: + backend = options.backend.lower() + if backend in ALL_SUPPORTED_BACKENDS and backend != "tensorrt": + 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) + return exec_settings + elif backend == "tensorrt": + assert params.get("max_seq_len") is None, ( + "max_seq_len is not a runtime parameter for C++ backend") + exec_settings, build_cfg = get_settings_from_engine(options.engine_dir) + engine_max_seq_len = build_cfg["max_seq_len"] + + if metadata.max_sequence_length > engine_max_seq_len: + raise RuntimeError( + f"Engine supports a max sequence of {engine_max_seq_len}. " + f"Provided dataset contains a maximum sequence of " + f"{metadata.max_sequence_length}. Please rebuild a new engine " + "to support this dataset.") + return exec_settings + else: + raise ValueError( + f"{options.backend} is not a known backend, check help for available options.", + ) def generate_warmup_dataset(requests, steps) -> List[InferenceRequest]: diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index 45d27b557677..b154f0a4c23f 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -1,204 +1,18 @@ from __future__ import annotations -from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Literal, Optional -from pydantic import (BaseModel, Field, PositiveFloat, field_validator, - model_validator) +from pydantic import BaseModel -import tensorrt_llm.bindings.executor as trtllm -from tensorrt_llm.llmapi import (BatchingType, CapacitySchedulerPolicy, - ContextChunkingPolicy, DynamicBatchConfig, - ExtendedRuntimePerfKnobConfig, KvCacheConfig, - SchedulerConfig) -from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_options -from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode - -SPECULATIVE_MAP = { - SpeculativeDecodingMode.NONE: lambda *args: None, - SpeculativeDecodingMode.MEDUSA: trtllm.DecodingMode.Medusa, -} +from tensorrt_llm.llmapi.llm_args import BaseLlmArgs class RuntimeConfig(BaseModel): - model: str model_path: Optional[Path] = None engine_dir: Optional[Path] = None revision: Optional[str] = None sw_version: str - settings_config: ExecutorSettingsConfig - # TODO: this is a dict corresponding to the Mapping class, the type should be - # changed to Mapping after the Mapping class is migrated to a Pydantic model. - mapping: Dict[str, Any] - decoding_config: Optional[DecodingConfig] = None - performance_options: PerformanceOptions backend: Literal["pytorch", "_autodeploy", None] = None - extra_llm_api_options: Optional[str] = None iteration_log: Optional[Path] = None - - def get_llm_args(self) -> Dict: - model = self.engine_dir or self.model_path or self.model - - llm_args = { - "scheduler_config": - self.settings_config.get_scheduler_config(), - "model": - model, - "skip_tokenizer_init": - True, - "pipeline_parallel_size": - self.mapping["pp_size"], - "tensor_parallel_size": - self.mapping["tp_size"], - "gpus_per_node": - self.mapping["gpus_per_node"], - "moe_expert_parallel_size": - self.mapping["moe_ep_size"], - "moe_cluster_parallel_size": - self.mapping["moe_cluster_size"], - "trust_remote_code": - True, - "enable_chunked_prefill": - self.settings_config.chunking, - "extended_runtime_perf_knob_config": - self.performance_options.get_perf_config(), - "decoding_config": - self.decoding_config.get_decoding_config(), - "batching_type": - BatchingType.INFLIGHT, - "max_batch_size": - self.settings_config.max_batch_size, - "max_num_tokens": - self.settings_config.max_num_tokens, - } - - backend_config_map = { - "pytorch": self.performance_options.get_pytorch_perf_config, - "_autodeploy": self.performance_options.get_autodeploy_perf_config - } - - if self.backend in backend_config_map: - llm_args.update(backend_config_map[self.backend]()) - - kv_cache_config = self.settings_config.get_kvcache_config().__dict__ - backend_cache_config = llm_args.pop("kv_cache_config", {}) - llm_args["kv_cache_config"] = backend_cache_config | kv_cache_config - - updated_llm_args = update_llm_args_with_extra_options( - llm_args, self.extra_llm_api_options) - - if self.backend == "pytorch": - cuda_graph_config = updated_llm_args.pop( - "cuda_graph_config", llm_args["cuda_graph_config"]) - if cuda_graph_config: - # Use runtime max_batch_size as cuda_graph_config.max_batch_size - # if both max_batch_size and batch_sizes are not set. - batch_sizes_set = cuda_graph_config.get("batch_sizes", - None) is not None - max_batch_size_set = cuda_graph_config.get( - "max_batch_size", None) is not None - if not batch_sizes_set and not max_batch_size_set: - cuda_graph_config[ - "max_batch_size"] = self.settings_config.max_batch_size - updated_llm_args["cuda_graph_config"] = cuda_graph_config - - return updated_llm_args - - @model_validator(mode="after") - def validate_full_config(self) -> RuntimeConfig: - # TODO: Check engine to make sure it can support Medusa. - return self - - -@dataclass -class PerformanceOptions: - cuda_graphs: bool = False - multi_block_mode: bool = True - cuda_graph_cache_size: int = 1000 - pytorch_config: Dict[str, Any] = Field(default_factory=dict) - - def get_perf_config(self) -> ExtendedRuntimePerfKnobConfig: - config = ExtendedRuntimePerfKnobConfig() - config.cuda_graph_mode = self.cuda_graphs - config.multi_block_mode = self.multi_block_mode - config.cuda_graph_cache_size = self.cuda_graph_cache_size - - return config - - def get_pytorch_perf_config(self): - return self.pytorch_config - - def get_autodeploy_perf_config(self) -> Dict: - AutoDeployPerfConfig = dict - ad_config = AutoDeployPerfConfig() - return ad_config - - -class DecodingConfig(BaseModel): - medusa_choices: Optional[List[List[int]]] = None - decoding_mode: SpeculativeDecodingMode = SpeculativeDecodingMode.NONE - - @field_validator("decoding_mode") - @classmethod - def decoding_mode_validator( - cls, value: Union[str, int, - SpeculativeDecodingMode]) -> SpeculativeDecodingMode: - return SpeculativeDecodingMode(value) - - @model_validator(mode="after") - def validate_speculative_decoding(self) -> DecodingConfig: - if self.medusa_choices and self.decoding_mode != SpeculativeDecodingMode.MEDUSA: - raise RuntimeError( - "Attempting to use set Medusa choices with a non-Medusa engine." - " Verify that you are using a Medusa engine.") - - return self - - def get_decoding_config(self) -> trtllm.DecodingConfig: - """Create a populated TRT-LLM DecodingConfig.""" - kwargs = {"decoding_mode": SPECULATIVE_MAP[self.decoding_mode]()} - - if self.medusa_choices is not None: - kwargs["medusa_choices"] = self.medusa_choices - - return trtllm.DecodingConfig(**kwargs) - - -class ExecutorSettingsConfig(BaseModel): - chunking: bool = True - scheduler_policy: CapacitySchedulerPolicy = Field( - default=CapacitySchedulerPolicy.MAX_UTILIZATION) - max_batch_size: int - max_num_tokens: int - kv_cache_percent: PositiveFloat = Field(default=.90, le=1.0) - kv_cache_reuse: bool = False - dynamic_max_batch_size: bool = True - dynamic_max_num_tokens: bool = False # Will enable after more validation. - - def get_dynamic_config(self) -> DynamicBatchConfig: - window_size = 128 if self.dynamic_max_batch_size else 0 - return DynamicBatchConfig( - enable_batch_size_tuning=self.dynamic_max_batch_size, - enable_max_num_tokens_tuning=self.dynamic_max_num_tokens, - dynamic_batch_moving_average_window=window_size, - ) - - def get_kvcache_config(self) -> KvCacheConfig: - return KvCacheConfig( - free_gpu_memory_fraction=self.kv_cache_percent, - enable_block_reuse=False, - ) - - def get_scheduler_config(self) -> SchedulerConfig: - if self.chunking: - return SchedulerConfig( - capacity_scheduler_policy=self.scheduler_policy, - context_chunking_policy=ContextChunkingPolicy. - FIRST_COME_FIRST_SERVED, - dynamic_batch_config=self.get_dynamic_config(), - ) - else: - return SchedulerConfig( - capacity_scheduler_policy=self.scheduler_policy, - dynamic_batch_config=self.get_dynamic_config()) + llm_args: BaseLlmArgs diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index 3ce443755844..01c3161013c3 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/tensorrt_llm/bench/dataclasses/reporting.py @@ -19,7 +19,6 @@ from tensorrt_llm.bench.dataclasses.statistics import (BenchmarkStatistics, PercentileStats, RequestRecord) -from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.logger import Logger from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode @@ -186,7 +185,6 @@ def __init__(self, dataset_metadata: DatasetMetadata, rt_cfg: RuntimeConfig, logger: Logger, - kwargs: Dict[str, Any], streaming: bool = False) -> None: """Initialize the ReportingController. @@ -200,7 +198,7 @@ def __init__(self, self.dataset_metadata = dataset_metadata self.rt_cfg = rt_cfg self.logger = logger - self.kwargs = kwargs + self.llm_args = rt_cfg.llm_args self.raw_statistics = statistics self.statistics = statistics.generate_statistics_summary( self.get_max_draft_len()) @@ -311,7 +309,7 @@ def get_statistics_dict(self) -> Dict[str, Any]: """ stats_dict = { "engine": { - "model": self.rt_cfg.model, + "model": self.llm_args.model, "model_path": str(self.rt_cfg.model_path), "engine_dir": str(self.rt_cfg.engine_dir), "revision": self.rt_cfg.revision, @@ -323,20 +321,9 @@ def get_statistics_dict(self) -> Dict[str, Any]: stats_dict["machine"] = self._query_gpu_info() # Retrieve KV cache information. - kv_cache_config = self.kwargs.get("kv_cache_config", KvCacheConfig()) - if isinstance(kv_cache_config, KvCacheConfig): - kv_cache_dtype = kv_cache_config.dtype - kv_cache_mem_percent = kv_cache_config.free_gpu_memory_fraction - elif isinstance(kv_cache_config, dict): - kv_cache_dtype = kv_cache_config.get("dtype", "auto") - kv_cache_mem_percent = kv_cache_config.get( - "free_gpu_memory_fraction") - else: - raise ValueError( - f"Invalid kv_cache_config type: {type(kv_cache_config)}.") - - kv_cache_mem_percent = kv_cache_mem_percent \ - if kv_cache_mem_percent is not None else None + kv_cache_config = self.llm_args.kv_cache_config + kv_cache_dtype = kv_cache_config.dtype + kv_cache_mem_percent = kv_cache_config.free_gpu_memory_fraction # Engine/Backend details if self.rt_cfg.backend not in ('pytorch', '_autodeploy'): @@ -364,7 +351,7 @@ def get_statistics_dict(self) -> Dict[str, Any]: from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._utils import torch_dtype_to_str - model = self.rt_cfg.model_path or self.rt_cfg.model + model = self.rt_cfg.model_path or self.llm_args.model model_config = ModelConfig.from_pretrained(model, trust_remote_code=True) @@ -384,16 +371,26 @@ def get_statistics_dict(self) -> Dict[str, Any]: } # World and runtime info + world_size = self.llm_args.tensor_parallel_size * self.llm_args.pipeline_parallel_size stats_dict["world_info"] = { - "tp_size": self.rt_cfg.mapping["tp_size"], - "pp_size": self.rt_cfg.mapping["pp_size"], - "ep_size": self.rt_cfg.mapping["moe_ep_size"], - "world_size": self.rt_cfg.mapping["world_size"], - "max_batch_size": self.rt_cfg.settings_config.max_batch_size, - "max_num_tokens": self.rt_cfg.settings_config.max_num_tokens, - "scheduling_policy": self.rt_cfg.settings_config.scheduler_policy, - "kv_cache_percentage": kv_cache_mem_percent, - "issue_rate": self.convert_rate_to_s(self.statistics.issue_rate_ns) + "tp_size": + self.llm_args.tensor_parallel_size, + "pp_size": + self.llm_args.pipeline_parallel_size, + "ep_size": + self.llm_args.moe_expert_parallel_size, + "world_size": + world_size, + "max_batch_size": + self.llm_args.max_batch_size, + "max_num_tokens": + self.llm_args.max_num_tokens, + "scheduling_policy": + self.llm_args.scheduler_config.capacity_scheduler_policy.value, + "kv_cache_percentage": + kv_cache_mem_percent, + "issue_rate": + self.convert_rate_to_s(self.statistics.issue_rate_ns) } # Request details @@ -429,7 +426,7 @@ def get_statistics_dict(self) -> Dict[str, Any]: self.per_user_output_throughput_tok_s, # Output throughput per GPU (total throughput / world size) "output_throughput_per_gpu_tok_s": - self.output_throughput_tok_s / self.rt_cfg.mapping["world_size"], + self.output_throughput_tok_s / world_size, # Request latency percentiles "request_latency_percentiles_ms": self.statistics.request_latency_percentiles.model_dump( @@ -481,17 +478,16 @@ def get_statistics_dict(self) -> Dict[str, Any]: } spec_decoding, decoding_mode = False, None - if (self.rt_cfg.decoding_config - and self.rt_cfg.decoding_config.decoding_mode + if (self.llm_args.decoding_config + and self.llm_args.decoding_config.decoding_mode != SpeculativeDecodingMode.NONE): # cpp decoding spec_decoding = True decoding_mode = self.rt_cfg.decoding_config.decoding_mode.values[1] - elif ("speculative_config" in self.kwargs - and self.kwargs["speculative_config"] is not None): + elif (self.llm_args.speculative_config is not None): # pytorch speculative decoding spec_decoding = True - decoding_mode = self.kwargs["speculative_config"].decoding_type + decoding_mode = self.llm_args.speculative_config.decoding_type if (spec_decoding): stats_dict["decoding_stats"] = { "mode": @@ -741,9 +737,7 @@ def report_statistics(self) -> None: def get_max_draft_len(self) -> int: """Get max_draft_len from speculative_config.""" - # Try to get from speculative_config - if ("speculative_config" in self.kwargs - and self.kwargs["speculative_config"] is not None): - return self.kwargs["speculative_config"].max_draft_len or 0 + if self.llm_args.speculative_config is not None: + return self.llm_args.speculative_config.max_draft_len or 0 return 0 diff --git a/tensorrt_llm/commands/eval.py b/tensorrt_llm/commands/eval.py index 8bfd8bab0f72..1480bc4575bd 100644 --- a/tensorrt_llm/commands/eval.py +++ b/tensorrt_llm/commands/eval.py @@ -12,19 +12,18 @@ # 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 typing import Optional +from typing import Any import click import tensorrt_llm.profiler as profiler +from tensorrt_llm.llmapi.llm_create import (create_llm_from_llm_args, + get_llm_args_from_cli_params) -from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..evaluate import (GSM8K, MMLU, MMMU, CnnDailymail, GPQADiamond, GPQAExtended, GPQAMain, JsonModeEval, LongBenchV1, LongBenchV2) -from ..llmapi import BuildConfig, KvCacheConfig -from ..llmapi.llm_utils import update_llm_args_with_extra_options +from ..llmapi import BuildConfig from ..logger import logger, severity_map @@ -118,56 +117,12 @@ default=False, help="Flag for disabling KV cache reuse.") @click.pass_context -def main(ctx, model: str, tokenizer: Optional[str], - custom_tokenizer: Optional[str], log_level: str, backend: str, - max_beam_width: int, max_batch_size: int, max_num_tokens: int, - max_seq_len: int, tp_size: int, pp_size: int, ep_size: Optional[int], - gpus_per_node: Optional[int], kv_cache_free_gpu_memory_fraction: float, - trust_remote_code: bool, revision: Optional[str], - extra_llm_api_options: Optional[str], disable_kv_cache_reuse: bool): +def main(ctx, model: str, log_level: str, **params: dict[str, Any]): logger.set_level(log_level) - - kv_cache_config = KvCacheConfig( - free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction, - enable_block_reuse=not disable_kv_cache_reuse) - - llm_args = { - "model": model, - "tokenizer": tokenizer, - "custom_tokenizer": custom_tokenizer, - "tensor_parallel_size": tp_size, - "pipeline_parallel_size": pp_size, - "moe_expert_parallel_size": ep_size, - "gpus_per_node": gpus_per_node, - "trust_remote_code": trust_remote_code, - "revision": revision, - "kv_cache_config": kv_cache_config, - } - - if backend == 'pytorch': - llm_cls = PyTorchLLM - llm_args.update(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) - elif backend == 'tensorrt': - llm_cls = LLM - build_config = BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) - llm_args.update(build_config=build_config) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - - if extra_llm_api_options is not None: - llm_args = update_llm_args_with_extra_options(llm_args, - extra_llm_api_options) + llm_args = get_llm_args_from_cli_params(model=model, **params) profiler.start("trtllm init") - llm = llm_cls(**llm_args) + llm = create_llm_from_llm_args(llm_args) profiler.stop("trtllm init") elapsed_time = profiler.elapsed_time_in_sec("trtllm init") logger.info(f"TRTLLM initialization time: {elapsed_time:.3f} seconds.") diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 76cbde9646f8..a6814ad5c27f 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -11,29 +11,24 @@ import click import torch -import yaml from strenum import StrEnum from torch.cuda import device_count -from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm import MultimodalEncoder -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._utils import mpi_rank from tensorrt_llm.executor.utils import LlmLauncherEnvs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig -from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, - DynamicBatchConfig, KvCacheConfig, - SchedulerConfig) +from tensorrt_llm.llmapi import BuildConfig from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole, - extract_disagg_cluster_config, parse_disagg_config_file, parse_metadata_server_config_file) -from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict +from tensorrt_llm.llmapi.llm_args import BaseLlmArgs +from tensorrt_llm.llmapi.llm_create import (create_llm_from_llm_args, + get_llm_args_from_cli_params) from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.logger import logger, severity_map -from tensorrt_llm.mapping import CpType from tensorrt_llm.serve import OpenAIDisaggServer, OpenAIServer from tensorrt_llm.serve.tool_parser import ToolParserFactory from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir @@ -86,106 +81,17 @@ def _signal_handler_cleanup_child(signum, frame): sys.exit(128 + signum) -def get_llm_args( - model: str, - tokenizer: Optional[str] = None, - custom_tokenizer: Optional[str] = None, - backend: str = "pytorch", - max_beam_width: int = BuildConfig.model_fields["max_beam_width"]. - default, - max_batch_size: int = BuildConfig.model_fields["max_batch_size"]. - default, - max_num_tokens: int = BuildConfig.model_fields["max_num_tokens"]. - default, - max_seq_len: int = BuildConfig.model_fields["max_seq_len"].default, - tensor_parallel_size: int = 1, - pipeline_parallel_size: int = 1, - context_parallel_size: int = 1, - cp_config: Optional[dict] = None, - moe_expert_parallel_size: Optional[int] = None, - gpus_per_node: Optional[int] = None, - free_gpu_memory_fraction: float = 0.9, - num_postprocess_workers: int = 0, - trust_remote_code: bool = False, - revision: Optional[str] = None, - reasoning_parser: Optional[str] = None, - fail_fast_on_attention_window_too_large: bool = False, - otlp_traces_endpoint: Optional[str] = None, - enable_chunked_prefill: bool = False, - **llm_args_extra_dict: Any): - - if gpus_per_node is None: - gpus_per_node = device_count() - if gpus_per_node == 0: - raise ValueError("No GPU devices found on the node") - build_config = BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) - kv_cache_config = KvCacheConfig( - free_gpu_memory_fraction=free_gpu_memory_fraction, ) - - dynamic_batch_config = DynamicBatchConfig( - enable_batch_size_tuning=True, - enable_max_num_tokens_tuning=False, - dynamic_batch_moving_average_window=128) - scheduler_config = SchedulerConfig( - capacity_scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, - dynamic_batch_config=dynamic_batch_config, - ) - if cp_config is not None and "cp_type" in cp_config: - cp_config = cp_config.copy() - try: - cp_config["cp_type"] = CpType[cp_config["cp_type"].upper()] - except KeyError: - raise ValueError(f"Invalid cp_type: {cp_config['cp_type']}. " \ - f"Must be one of: {', '.join([t.name for t in CpType])}") - - llm_args = { - "model": model, - "scheduler_config": scheduler_config, - "tokenizer": tokenizer, - "custom_tokenizer": custom_tokenizer, - "tensor_parallel_size": tensor_parallel_size, - "pipeline_parallel_size": pipeline_parallel_size, - "context_parallel_size": context_parallel_size, - "cp_config": cp_config if cp_config is not None else {}, - "moe_expert_parallel_size": moe_expert_parallel_size, - "gpus_per_node": gpus_per_node, - "trust_remote_code": trust_remote_code, - "revision": revision, - "build_config": build_config, - "max_batch_size": max_batch_size, - "max_num_tokens": max_num_tokens, - "max_beam_width": max_beam_width, - "max_seq_len": max_seq_len, - "kv_cache_config": kv_cache_config, - "backend": backend, - "num_postprocess_workers": num_postprocess_workers, - "postprocess_tokenizer_dir": tokenizer or model, - "reasoning_parser": reasoning_parser, - "fail_fast_on_attention_window_too_large": - fail_fast_on_attention_window_too_large, - "otlp_traces_endpoint": otlp_traces_endpoint, - "enable_chunked_prefill": enable_chunked_prefill, - } - - return llm_args, llm_args_extra_dict - - def launch_server( host: str, port: int, - llm_args: dict, + llm_args: BaseLlmArgs, tool_parser: Optional[str] = None, chat_template: Optional[str] = None, metadata_server_cfg: Optional[MetadataServerConfig] = None, server_role: Optional[ServerRole] = None, disagg_cluster_config: Optional[DisaggClusterConfig] = None, multimodal_server_config: Optional[MultimodalServerConfig] = None): - - backend = llm_args["backend"] - model = llm_args["model"] + model = llm_args.model addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) address_family = socket.AF_INET6 if all( @@ -200,22 +106,7 @@ def launch_server( except OSError as e: raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}") - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == '_autodeploy': - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - - # AutoDeploy does not support build_config - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - elif backend == 'tensorrt' or backend == 'trt': - llm_args.pop("backend") - llm = LLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") + llm = create_llm_from_llm_args(llm_args) server = OpenAIServer(llm=llm, model=model, @@ -233,7 +124,7 @@ def launch_server( asyncio.run(server(host, port, sockets=[s])) -def launch_grpc_server(host: str, port: int, llm_args: dict): +def launch_grpc_server(host: str, port: int, llm_args: BaseLlmArgs): """ Launch a gRPC server for TensorRT-LLM. @@ -243,7 +134,7 @@ def launch_grpc_server(host: str, port: int, llm_args: dict): Args: host: Host to bind to port: Port to bind to - llm_args: Arguments for LLM initialization (from get_llm_args) + llm_args: Arguments for LLM initialization """ import grpc @@ -259,32 +150,15 @@ def launch_grpc_server(host: str, port: int, llm_args: dict): async def serve_grpc_async(): logger.info("Initializing TensorRT-LLM gRPC server...") - - backend = llm_args.get("backend") - model_path = llm_args.get("model", "") - - if backend == "pytorch": - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == "_autodeploy": - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - elif backend == "tensorrt" or backend == "trt": - llm_args.pop("backend") - llm = LLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - + + llm = create_llm_from_llm_args(llm_args) logger.info("Model loaded successfully") # Create request manager request_manager = GrpcRequestManager(llm) # Create servicer - servicer = TrtllmServiceServicer(request_manager, model_path=model_path) + servicer = TrtllmServiceServicer(request_manager, model_path=llm_args.model) # Create gRPC server server = grpc.aio.server( @@ -352,12 +226,12 @@ def signal_handler(): def launch_mm_encoder_server( + model: str, host: str, port: int, encoder_args: dict, metadata_server_cfg: Optional[MetadataServerConfig] = None, ): - model = encoder_args["model"] encoder_args.pop("build_config") mm_encoder = MultimodalEncoder(**encoder_args) @@ -600,24 +474,13 @@ def convert(self, value: Any, param: Optional["click.Parameter"], default=False, help="Run gRPC server instead of OpenAI HTTP server. " "gRPC server accepts pre-tokenized requests and returns raw token IDs.") -def serve( - model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], - host: str, port: int, log_level: str, backend: str, max_beam_width: int, - max_batch_size: int, max_num_tokens: int, max_seq_len: int, - tensor_parallel_size: int, pipeline_parallel_size: int, - context_parallel_size: int, moe_expert_parallel_size: Optional[int], - moe_cluster_parallel_size: Optional[int], gpus_per_node: Optional[int], - free_gpu_memory_fraction: float, num_postprocess_workers: int, - trust_remote_code: bool, revision: Optional[str], - extra_llm_api_options: Optional[str], reasoning_parser: Optional[str], - tool_parser: Optional[str], metadata_server_config_file: Optional[str], - server_role: Optional[str], - fail_fast_on_attention_window_too_large: bool, - otlp_traces_endpoint: Optional[str], enable_chunked_prefill: bool, - disagg_cluster_uri: Optional[str], media_io_kwargs: Optional[str], - custom_module_dirs: list[Path], chat_template: Optional[str], - grpc: bool): - """Running an OpenAI API compatible server (or gRPC server with --grpc flag) +def serve(model: str, host: str, port: int, log_level: str, + tool_parser: Optional[str], + metadata_server_config_file: Optional[str], + server_role: Optional[str], disagg_cluster_uri: Optional[str], + media_io_kwargs: Optional[str], custom_module_dirs: list[Path], + chat_template: Optional[str], grpc: bool, **params: dict[str, Any]): + """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ @@ -630,49 +493,24 @@ def serve( logger.error( f"Failed to import custom module from {custom_module_dir}: {e}") raise e - llm_args, _ = get_llm_args( - model=model, - tokenizer=tokenizer, - custom_tokenizer=custom_tokenizer, - backend=backend, - max_beam_width=max_beam_width, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_seq_len=max_seq_len, - tensor_parallel_size=tensor_parallel_size, - pipeline_parallel_size=pipeline_parallel_size, - context_parallel_size=context_parallel_size, - moe_expert_parallel_size=moe_expert_parallel_size, - moe_cluster_parallel_size=moe_cluster_parallel_size, - gpus_per_node=gpus_per_node, - free_gpu_memory_fraction=free_gpu_memory_fraction, - num_postprocess_workers=num_postprocess_workers, - trust_remote_code=trust_remote_code, - revision=revision, - reasoning_parser=reasoning_parser, - fail_fast_on_attention_window_too_large= - fail_fast_on_attention_window_too_large, - otlp_traces_endpoint=otlp_traces_endpoint, - enable_chunked_prefill=enable_chunked_prefill) - - llm_args_extra_dict = {} - if extra_llm_api_options is not None: - with open(extra_llm_api_options, 'r') as f: - llm_args_extra_dict = yaml.safe_load(f) - llm_args = update_llm_args_with_extra_dict(llm_args, llm_args_extra_dict) + + llm_args = get_llm_args_from_cli_params(model=model, **params) metadata_server_cfg = parse_metadata_server_config_file( metadata_server_config_file) # Specify disagg_cluster_config in config file or through command line "--disagg_cluster_uri", # but disagg_cluster_uri takes precedence over cluster uri in config file - disagg_cluster_config = llm_args.pop("disagg_cluster", None) - if disagg_cluster_config: - disagg_cluster_config = extract_disagg_cluster_config( - disagg_cluster_config, disagg_cluster_uri) - elif disagg_cluster_uri: - disagg_cluster_config = DisaggClusterConfig( - cluster_uri=disagg_cluster_uri) + + # TODO: disagg_cluster_config is broken + disagg_cluster_config = None + # disagg_cluster_config = llm_args.pop("disagg_cluster", None) + # if disagg_cluster_config: + # disagg_cluster_config = extract_disagg_cluster_config( + # disagg_cluster_config, disagg_cluster_uri) + # elif disagg_cluster_uri: + # disagg_cluster_config = DisaggClusterConfig( + # cluster_uri=disagg_cluster_uri) if metadata_server_cfg is not None or disagg_cluster_config is not None: assert ( @@ -762,34 +600,24 @@ def serve( default=None, help="Path to metadata server config file") def serve_encoder(model: str, host: str, port: int, log_level: str, - max_batch_size: int, max_num_tokens: int, - gpus_per_node: Optional[int], trust_remote_code: bool, - extra_encoder_options: Optional[str], - metadata_server_config_file: Optional[str]): + metadata_server_config_file: Optional[str], + **params: dict[str, Any]): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ logger.set_level(log_level) - # TODO: expose more argument progressivly - llm_args, _ = get_llm_args(model=model, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - gpus_per_node=gpus_per_node, - trust_remote_code=trust_remote_code) - - encoder_args_extra_dict = {} - if extra_encoder_options is not None: - with open(extra_encoder_options, 'r') as f: - encoder_args_extra_dict = yaml.safe_load(f) - encoder_args = update_llm_args_with_extra_dict(llm_args, - encoder_args_extra_dict) + # TODO: encoder args should have its own subclass of LlmArgs, which would eliminate the need to convert + # to a dict here + encoder_args = get_llm_args_from_cli_params(model=model, + **params).model_dump() metadata_server_cfg = parse_metadata_server_config_file( metadata_server_config_file) - launch_mm_encoder_server(host, port, encoder_args, metadata_server_cfg) + launch_mm_encoder_server(model, host, port, encoder_args, + metadata_server_cfg) @click.command("disaggregated") @@ -919,12 +747,8 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX) server_cfg = disagg_cfg.server_configs[int(instance_idx)] - llm_args, llm_args_extra_dict = get_llm_args(**server_cfg.other_args) - llm_args = update_llm_args_with_extra_dict(llm_args, - llm_args_extra_dict) - - # Ignore the non-LLM args - llm_args.pop("router", None) + # TODO: this might be a misnomer because other_args is not actually CLI params + llm_args = get_llm_args_from_cli_params(**server_cfg.other_args) _launch_disaggregated_server(config_file, llm_args) return @@ -944,9 +768,8 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): instance_idx) server_cfg = disagg_cfg.server_configs[instance_idx] - llm_args, llm_args_extra_dict = get_llm_args(**server_cfg.other_args) - llm_args = update_llm_args_with_extra_dict(llm_args, - llm_args_extra_dict) + # TODO: this might be a misnomer because other_args is not actually CLI params + llm_args = get_llm_args_from_cli_params(**server_cfg.other_args) _launch_disaggregated_leader(sub_comm, instance_idx, config_file, log_level) @@ -964,7 +787,8 @@ class DisaggLauncherEnvs(StrEnum): TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT = "TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT" -def _launch_disaggregated_server(disagg_config_file: str, llm_args: dict): +def _launch_disaggregated_server(disagg_config_file: str, + llm_args: BaseLlmArgs): # Launching the server instance_idx = os.environ.get(DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX) assert instance_idx is not None, f"{DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX} should be set by the launcher" diff --git a/tensorrt_llm/llmapi/build_cache.py b/tensorrt_llm/llmapi/build_cache.py index a666ab1fff4c..b22a80dc514b 100644 --- a/tensorrt_llm/llmapi/build_cache.py +++ b/tensorrt_llm/llmapi/build_cache.py @@ -10,6 +10,7 @@ from typing import Any, List, Optional import filelock +from pydantic import BaseModel, Field, model_validator import tensorrt_llm from tensorrt_llm.builder import BuildConfig @@ -28,41 +29,36 @@ def get_build_cache_config_from_env() -> tuple[bool, str]: return build_cache_enabled, build_cache_root -class BuildCacheConfig: +class BuildCacheConfig(BaseModel): """ Configuration for the build cache. - Attributes: - cache_root (str): The root directory for the build cache. - max_records (int): The maximum number of records to store in the cache. - max_cache_storage_gb (float): The maximum amount of storage (in GB) to use for the cache. - Note: The build-cache assumes the weights of the model are not changed during the execution. If the weights are changed, you should remove the caches manually. """ - def __init__(self, - cache_root: Optional[Path] = None, - max_records: int = 10, - max_cache_storage_gb: float = 256): - self._cache_root = cache_root - self._max_records = max_records - self._max_cache_storage_gb = max_cache_storage_gb - - @property - def cache_root(self) -> Path: - _build_cache_enabled, _build_cache_root = get_build_cache_config_from_env( - ) - return self._cache_root or Path(_build_cache_root) - - @property - def max_records(self) -> int: - return self._max_records - - @property - def max_cache_storage_gb(self) -> float: - return self._max_cache_storage_gb + cache_root: Optional[str] = Field( + default=None, + description= + "The root directory for the build cache. Falls back to env var if not provided." + ) + max_records: int = Field( + default=10, + gt=0, + description="The maximum number of records to store in the cache.") + max_cache_storage_gb: float = Field( + default=256.0, + description="The maximum amount of storage (in GB) to use for the cache." + ) + + @model_validator(mode="after") + def set_default_cache_root(self) -> "BuildCacheConfig": + """Set cache_root from environment variable if not provided.""" + if self.cache_root is None: + _, default_cache_root = get_build_cache_config_from_env() + self.cache_root = default_cache_root + return self class BuildCache: @@ -76,17 +72,11 @@ class BuildCache: CACHE_VERSION = 0 def __init__(self, config: Optional[BuildCacheConfig] = None): - - _, default_cache_root = get_build_cache_config_from_env() config = config or BuildCacheConfig() - - self.cache_root = config.cache_root or Path(default_cache_root) + self.cache_root = config.cache_root self.max_records = config.max_records self.max_cache_storage_gb = config.max_cache_storage_gb - if config.max_records < 1: - raise ValueError("max_records should be greater than 0") - def free_storage_in_gb(self) -> float: ''' Get the free storage capacity of the cache. ''' # measure the root directory diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 5ff6dc807494..ecfc05fad5b6 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -12,9 +12,7 @@ import transformers from tqdm import tqdm -from transformers import PreTrainedTokenizerBase -from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.inputs.data import TextPrompt from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, @@ -40,8 +38,8 @@ from ..sampling_params import SamplingParams from ..scheduling_params import SchedulingParams from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, - TRT_LLMARGS_EXPLICIT_DOCSTRING, PeftCacheConfig, - PybindMirror, TorchLlmArgs, TrtLlmArgs) + TRT_LLMARGS_EXPLICIT_DOCSTRING, BaseLlmArgs, + PeftCacheConfig, PybindMirror, TorchLlmArgs, TrtLlmArgs) from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig, LlmBuildStats, ModelLoader, _ModelRuntimeContext) from .mpi_session import MpiPoolSession, external_mpi_comm_available @@ -119,78 +117,23 @@ class BaseLLM: The base class for all LLM classes. """ - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - - self._executor_cls = kwargs.pop("executor_cls", GenerationExecutor) - self._orchestrator_type = kwargs.get("orchestrator_type", None) + def __init__(self, args: BaseLlmArgs) -> None: self._llm_id = None + self.args = args log_level = logger.level logger.set_level("info") # force display the backend - try: - env_overrides = kwargs.get("env_overrides", None) - self._process_env_overrides(env_overrides) - - backend = kwargs.get('backend', None) - if backend == "pytorch": - logger.info("Using LLM with PyTorch backend") - llm_args_cls = TorchLlmArgs - if self._orchestrator_type == "ray" or mpi_disabled(): - self._orchestrator_type = "ray" - os.environ["TLLM_DISABLE_MPI"] = "1" - # Propagate to args construction - kwargs["orchestrator_type"] = "ray" - - elif backend == '_autodeploy': - logger.info("Using LLM with AutoDeploy backend") - from .._torch.auto_deploy.llm_args import \ - LlmArgs as AutoDeployLlmArgs - llm_args_cls = AutoDeployLlmArgs - else: - logger.info("Using LLM with TensorRT backend") - llm_args_cls = TrtLlmArgs - - # check the kwargs and raise ValueError directly - valid_keys = set( - list(llm_args_cls.model_fields.keys()) + - ['_mpi_session', 'backend']) - for key in kwargs: - if key not in valid_keys: - raise ValueError( - f"{self.__class__.__name__} got invalid argument: {key}" - ) - - self.args = llm_args_cls.from_kwargs( - model=model, - tokenizer=tokenizer, - tokenizer_mode=tokenizer_mode, - skip_tokenizer_init=skip_tokenizer_init, - trust_remote_code=trust_remote_code, - tensor_parallel_size=tensor_parallel_size, - dtype=dtype, - revision=revision, - tokenizer_revision=tokenizer_revision, - **kwargs) + self._process_env_overrides(self.args.env_overrides) - except Exception as e: - logger.error( - f"Failed to parse the arguments for the LLM constructor: {e}") - raise e + if self.args.backend == "pytorch": + logger.info("Using LLM with PyTorch backend") + elif self.args.backend == "_autodeploy": + logger.info("Using LLM with AutoDeploy backend") + else: + logger.info("Using LLM with TensorRT backend") - finally: - logger.set_level(log_level) # restore the log level + logger.set_level(log_level) # restore the log level logger_debug(f"LLM.args.mpi_session: {self.args.mpi_session}\n", "yellow") @@ -767,9 +710,8 @@ def _try_load_tokenizer(self) -> Optional[TokenizerBase]: # TODO smor- need to refine what is the desired behavior if lora is enabled # in terms of the tokenizer initialization process - if hasattr(self.args, "backend") and self.args.backend in [ - "pytorch", "_autodeploy" - ] and self.args.lora_config is not None: + if self.args.backend in ["pytorch", "_autodeploy" + ] and self.args.lora_config is not None: num_lora_dirs = len(self.args.lora_config.lora_dir) if num_lora_dirs == 1: tokenizer_path = self.args.lora_config.lora_dir[0] @@ -865,23 +807,9 @@ class _TrtLLM(BaseLLM): Parameters: """ - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - # TODO: deprecate backend in LLM kwargs - - super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init, - trust_remote_code, tensor_parallel_size, dtype, - revision, tokenizer_revision, **kwargs) + def __init__(self, model: Union[str, Path], **kwargs: Any) -> None: + args = TrtLlmArgs(model=model, **kwargs) + super().__init__(args) @property def workspace(self) -> Path: @@ -1026,7 +954,7 @@ def _build_model(self): or (self.args.build_config and self.args.build_config.gather_context_logits)) - self._executor = self._executor_cls.create( + self._executor = GenerationExecutor.create( self._engine_dir, executor_config=self._executor_config, batched_logits_processor=self.args.batched_logits_processor, @@ -1049,36 +977,12 @@ class _TorchLLM(BaseLLM): Parameters: """ - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - - # TODO: deprecate backend in LLM kwargs - backend = kwargs.pop("backend", "pytorch") - + def __init__(self, model: Union[str, Path], **kwargs: Any) -> None: # Validate that users don't pass TrtLlmArgs-specific arguments self._validate_args_for_torch_backend(kwargs) - super().__init__(model, - tokenizer, - tokenizer_mode, - skip_tokenizer_init, - trust_remote_code, - tensor_parallel_size, - dtype, - revision, - tokenizer_revision, - backend=backend, - **kwargs) + args = TorchLlmArgs(model=model, **kwargs) + super().__init__(args) @set_api_status("prototype") def _collective_rpc(self, @@ -1129,7 +1033,7 @@ def _build_model(self): # TODO: revisit gather_context_logits return_logits = self.args.gather_generation_logits - self._executor = self._executor_cls.create( + self._executor = GenerationExecutor.create( self._engine_dir, executor_config=None, batched_logits_processor=self.args.batched_logits_processor, @@ -1168,23 +1072,10 @@ def _validate_args_for_torch_backend(self, kwargs: dict) -> None: ) +# NOTE: this is defined as a subclass instead of simply setting LLM = _TorchLLM so that +# LLM.__name__ is "LLM" class LLM(_TorchLLM): - - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init, - trust_remote_code, tensor_parallel_size, dtype, - revision, tokenizer_revision, **kwargs) + pass # sphinx will ignore the LLM's docstring if it is not explicitly set diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index eac238bb0a97..916cde35d854 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -8,14 +8,15 @@ from dataclasses import dataclass from enum import Enum, EnumMeta from pathlib import Path -from typing import (Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, +from typing import (Annotated, Any, Dict, List, Literal, Optional, Set, Tuple, Type, TypeAlias, TypeVar, Union, get_args, get_origin) import torch import yaml -from pydantic import AliasChoices, BaseModel +from pydantic import BaseModel, ConfigDict from pydantic import Field as PydanticField -from pydantic import PrivateAttr, field_validator, model_validator +from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, + PrivateAttr, field_validator, model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -27,7 +28,7 @@ from tensorrt_llm.lora_helper import (LoraConfig, get_default_trtllm_modules_to_hf_modules) -from .._utils import mpi_rank +from .._utils import mpi_disabled, mpi_rank # yapf: disable # isort: off @@ -51,7 +52,7 @@ # yapf: enable from ..builder import BuildConfig, EngineConfig from ..logger import logger -from ..mapping import Mapping +from ..mapping import CpType, Mapping from ..models.automodel import AutoConfig from ..models.modeling_utils import (PretrainedConfig, QuantAlgo, QuantConfig, SpeculativeDecodingMode) @@ -60,8 +61,6 @@ from .tokenizer import TokenizerBase, tokenizer_factory from .utils import generate_api_docs_as_docstring, get_type_repr -# TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import - TypeBaseModel = TypeVar("T", bound=BaseModel) @@ -113,7 +112,7 @@ class CudaGraphConfig(StrictBaseModel): default=None, description="List of batch sizes to create CUDA graphs for.") - max_batch_size: int = Field( + max_batch_size: NonNegativeInt = Field( default=0, description="Maximum batch size for CUDA graphs.") enable_padding: bool = Field( @@ -122,14 +121,36 @@ class CudaGraphConfig(StrictBaseModel): "If true, batches are rounded up to the nearest cuda_graph_batch_size. This is usually a net win for performance." ) - @field_validator('max_batch_size') - @classmethod - def validate_cuda_graph_max_batch_size(cls, v): - """Validate cuda_graph_config.max_batch_size is non-negative.""" - if v < 0: - raise ValueError( - "cuda_graph_config.max_batch_size must be non-negative") - return v + @model_validator(mode='after') + def validate_cuda_graph_config(self) -> 'CudaGraphConfig': + """Validate CUDA graph configuration. + + Ensures that: + 1. If batch_sizes is provided, max_batch_size must be 0 + 2. If batch_sizes is not provided, it is generated based on max_batch_size + 3. If both are provided, batch_sizes must match the generated values + """ + if self.batch_sizes: + self.batch_sizes = sorted(self.batch_sizes) + if self.max_batch_size != 0: + if self.batch_sizes != CudaGraphConfig._generate_cuda_graph_batch_sizes( + self.max_batch_size, self.enable_padding): + raise ValueError( + "Please don't set both cuda_graph_config.batch_sizes " + "and cuda_graph_config.max_batch_size.\n" + f"cuda_graph_config.batch_sizes: {self.batch_sizes}, " + f"cuda_graph_config.max_batch_size: {self.max_batch_size}" + ) + else: + self.max_batch_size = max(self.batch_sizes) + else: + max_batch_size = self.max_batch_size or 128 + generated_sizes = CudaGraphConfig._generate_cuda_graph_batch_sizes( + max_batch_size, self.enable_padding) + self.batch_sizes = generated_sizes + self.max_batch_size = max_batch_size + + return self @staticmethod def _generate_cuda_graph_batch_sizes(max_batch_size: int, @@ -188,41 +209,14 @@ class BaseSparseAttentionConfig(StrictBaseModel): """ Configuration for sparse attention. """ + algorithm: str + seq_len_threshold: Optional[int] = Field( default=None, description= "The sequence length threshold for separating short and long sequences." ) - @property - def algorithm(self) -> str: - raise NotImplementedError("Algorithm must be implemented in subclasses") - - @classmethod - def from_dict(cls, data: dict): - # dispatch to the correct sparse attention config - config_classes = { - "rocket": RocketSparseAttentionConfig, - "dsa": DeepSeekSparseAttentionConfig, - "skip_softmax": SkipSoftmaxAttentionConfig, - } - - algorithm = data.get("algorithm", None) - if algorithm is None: - raise ValueError(f"Sparse attention algorithm is required") - - config_class = config_classes.get(algorithm.lower()) - if config_class is None: - raise ValueError(f"Invalid algorithm: {algorithm}") - - # Remove 'algorithm' before passing to subclass constructor - # It's a ClassVar in subclasses, and used for dispatching to the correct subclass - data = {k: v for k, v in data.items() if k != 'algorithm'} - return config_class(**data) - - def _check_fields(self): - pass - def supports_backend(self, backend: str) -> bool: """ Override if the speculation algorithm does not support @@ -247,7 +241,7 @@ class RocketSparseAttentionConfig(BaseSparseAttentionConfig): """ Configuration for RocketKV sparse attention. """ - algorithm: ClassVar[str] = "rocket" + algorithm: Literal["rocket"] = "rocket" window_size: Optional[int] = Field( default=32, description="The window size for snap KV.") kernel_size: Optional[int] = Field( @@ -263,10 +257,6 @@ class RocketSparseAttentionConfig(BaseSparseAttentionConfig): description="KT cache dtype", ) - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -278,7 +268,7 @@ class DeepSeekSparseAttentionConfig(BaseSparseAttentionConfig): """ Configuration for DeepSeek Sparse Attention. """ - algorithm: ClassVar[str] = "dsa" + algorithm: Literal["dsa"] = "dsa" index_n_heads: Optional[int] = Field( default=None, description="The number of heads for the indexer.") index_head_dim: Optional[int] = Field( @@ -292,10 +282,6 @@ class DeepSeekSparseAttentionConfig(BaseSparseAttentionConfig): description= "Whether to skip the MQA and Top-K in the indexer for short sequences.") - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -312,15 +298,11 @@ class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """ Configuration for skip softmax attention. """ - algorithm: ClassVar[str] = "skip_softmax" + algorithm: Literal["skip_softmax"] = "skip_softmax" threshold_scale_factor: Optional[Union[float, Dict[str, float]]] = Field( default=None, description="The threshold scale factor for skip softmax attention.") - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -475,37 +457,21 @@ class MoeConfig(StrictBaseModel): "Use low precision combine in MoE operations (only for NVFP4 quantization). When enabled, uses lower precision for combining expert outputs to improve performance." ) - @classmethod - def from_dict(cls, data: dict): - return cls(**data) + +Nvfp4Backend = Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core'] class Nvfp4GemmConfig(StrictBaseModel): """ Configuration for NVFP4 GEMM backend selection. """ - allowed_backends: List[str] = Field( + allowed_backends: List[Nvfp4Backend] = Field( default=['cutlass', 'cublaslt', 'cuda_core'], + min_length=1, description="List of backends to consider for auto-selection. " "Default excludes 'cutedsl' for faster build time. " - "Add 'cutedsl' for extreme performance at the cost of longer server launch time. " - "Valid values: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'.") - - @model_validator(mode="after") - def validate_allowed_backends(self) -> 'Nvfp4GemmConfig': - valid_backends = {'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'} - invalid = set(self.allowed_backends) - valid_backends - if invalid: - raise ValueError( - f"Invalid backends in allowed_backends: {invalid}. " - f"Valid backends are: {sorted(valid_backends)}") - if not self.allowed_backends: - raise ValueError("allowed_backends cannot be empty.") - return self - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) + "Add 'cutedsl' for extreme performance at the cost of longer server launch time." + ) class AttentionDpConfig(StrictBaseModel): @@ -520,9 +486,54 @@ class AttentionDpConfig(StrictBaseModel): default=10, description="The number of iterations to wait for batching.") + @model_validator(mode='after') + def validate_attention_dp_config(self) -> 'AttentionDpConfig': + """Validate attention DP configuration. + + Ensures that: + 1. If enable_balance is true, batching_wait_iters must be >= 0 + 2. If enable_balance is true, timeout_iters must be >= 0 + """ + if self.enable_balance: + if self.batching_wait_iters < 0: + raise ValueError( + "attention_dp_config.batching_wait_iters must be greater or equal to 0 when enable_balance is true" + ) + if self.timeout_iters < 0: + raise ValueError( + "attention_dp_config.timeout_iters must be greater or equal to 0 when enable_balance is true" + ) + return self + + +class CpConfig(StrictBaseModel): + """ + Configuration for context parallelism. + """ + cp_type: CpType = Field(default=CpType.ULYSSES, + description="Context parallel type.") + tokens_per_block: Optional[int] = Field( + default=None, + description="Number of tokens per block. Used in HELIX parallelism.") + use_nccl_for_alltoall: Optional[bool] = Field( + default=None, + description= + "Whether to use NCCL for alltoall communication. Used in HELIX parallelism. Defaults to True." + ) + cp_anchor_size: Optional[int] = Field( + default=None, description="Anchor size for STAR attention.") + block_size: Optional[int] = Field( + default=None, description="Block size for STAR attention.") + + @field_validator("cp_type", mode="before") @classmethod - def from_dict(cls, data: dict): - return cls(**data) + def validate_cp_type(cls, v): + """Normalize cp_type string to uppercase.""" + if v is None: + return None + if isinstance(v, str): + return v.upper() + return v class _ParallelConfig(StrictBaseModel): @@ -535,7 +546,7 @@ class _ParallelConfig(StrictBaseModel): moe_cluster_size: int = -1 moe_tp_size: int = -1 moe_ep_size: int = -1 - cp_config: dict = Field(default_factory=dict) + cp_config: Optional[CpConfig] = Field(default=None) pp_partition: Optional[List[int]] = Field(default=None) enable_attention_dp: bool = False enable_lm_head_tp_in_adp: bool = False @@ -578,19 +589,21 @@ def is_multi_gpu(self) -> bool: return self.world_size > 1 def to_mapping(self) -> Mapping: - return Mapping(world_size=self.world_size, - rank=mpi_rank(), - gpus_per_node=self.gpus_per_node, - tp_size=self.tp_size, - pp_size=self.pp_size, - pp_partition=self.pp_partition, - cp_size=self.cp_size, - cp_config=self.cp_config, - enable_attention_dp=self.enable_attention_dp, - enable_lm_head_tp_in_adp=self.enable_lm_head_tp_in_adp, - moe_cluster_size=self.moe_cluster_size, - moe_tp_size=self.moe_tp_size, - moe_ep_size=self.moe_ep_size) + return Mapping( + world_size=self.world_size, + rank=mpi_rank(), + gpus_per_node=self.gpus_per_node, + tp_size=self.tp_size, + pp_size=self.pp_size, + pp_partition=self.pp_partition, + cp_size=self.cp_size, + # TODO: Mapping still uses cp_config as a dict; migrate to CpConfig + cp_config=self.cp_config.model_dump() if self.cp_config else {}, + enable_attention_dp=self.enable_attention_dp, + enable_lm_head_tp_in_adp=self.enable_lm_head_tp_in_adp, + moe_cluster_size=self.moe_cluster_size, + moe_tp_size=self.moe_tp_size, + moe_ep_size=self.moe_ep_size) class CalibConfig(StrictBaseModel): @@ -618,26 +631,6 @@ class CalibConfig(StrictBaseModel): description= "The maximum sequence length to initialize tokenizer for calibration.") - @classmethod - def from_dict(cls, config: dict) -> 'CalibConfig': - """Create a CalibConfig instance from a dict. - - Args: - config (dict): The dict used to create CalibConfig. - - Returns: - tensorrt_llm.llmapi.CalibConfig: The CalibConfig created from dict. - """ - return cls(**config) - - def to_dict(self) -> dict: - """Dump a CalibConfig instance to a dict. - - Returns: - dict: The dict dumped from CalibConfig. - """ - return self.model_dump() - class _ModelFormatKind(Enum): HF = 0 @@ -646,73 +639,65 @@ class _ModelFormatKind(Enum): class DecodingBaseConfig(StrictBaseModel): - # The number of the drafter layers. - max_draft_len: Optional[int] = None - # The number of draft tokens in the draft tokens tree. - # If it's a linear tree, each draft layer will only generate one draft token. - # In this case, max_draft_len == max_total_draft_tokens. - # If it's a static or dynamic tree, each draft layer may generate more than one draft token. - # In this case, max_total_draft_tokens >= max_draft_len. - max_total_draft_tokens: Optional[int] = None - # The speculative (draft) model. Accepts either: - # - A HuggingFace Hub model ID (str), e.g., "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" - # which will be automatically downloaded. - # - A local filesystem path to a downloaded model directory. + max_draft_len: Optional[NonNegativeInt] = Field( + default=None, description="The number of drafter layers.") + + max_total_draft_tokens: Optional[int] = Field( + default=None, + description= + "The number of draft tokens in the draft tokens tree. If it's a linear tree, each draft layer will " + "only generate one draft token. In this case, max_draft_len == max_total_draft_tokens. If it's a static or " + "dynamic tree, each draft layer may generate more than one draft token. In this case, " + "max_total_draft_tokens >= max_draft_len.") + speculative_model: Optional[Union[str, Path]] = Field( default=None, - validation_alias=AliasChoices("speculative_model", - "speculative_model_dir")) - - # PyTorch only. - # When specified, speculation will be disabled at batch sizes above - # this value. Otherwise, speculation will always be on. - max_concurrency: Optional[int] = None - - # Developer interface: dynamically adjust draft length based on active batch size in runtime. - # Maps batch size to draft lengths. For example: - # {1: 4, 4: 2, 8: 0} means: - # - batch_size >= 1: use draft_len=4 - # - batch_size >= 4: use draft_len=2 - # - batch_size >= 8: use draft_len=0 (disable speculation) - # draft_len_schedule is enforced to contain batch_size=1 and its according draft_len equals max_draft_len for consistency - # for example, if max_draft_len=4, the schedule must contain {1: 4} - draft_len_schedule: Optional[dict[int, int]] = None - - load_format: Optional[str] = None - # PyTorch only. - # Rolling average window size (N) for acceptance length across completed requests. - # If not set or set to 0, the feature is disabled. - acceptance_window: Optional[int] = None - # PyTorch only. - # Threshold for average acceptance length; speculation will be disabled - # permanently once the rolling average over the last N completed requests - # (N = acceptance_window) drops below this value. - acceptance_length_threshold: Optional[float] = None - - # Prototype. If true, allows non-greedy sampling when speculation is used. Only applicable - # to 1-model code paths; non-greedy sampling is always enabled on 2-model paths. - allow_advanced_sampling: bool = False - - # Validate acceptance controls at field level so they run on model creation - @field_validator('acceptance_window') - @classmethod - def _validate_acceptance_window(cls, v: Optional[int]): - if v is None: - return v - if v < 0: - raise ValueError( - f"acceptance_window must be >= 0 (0 disables), got {v}") - return v + description= + "The speculative (draft) model. Accepts either (1) a HuggingFace Hub model ID (e.g. 'yuhuili/EAGLE3-LLaMA3.1-Instruct-8B')," + "which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory." + ) - @field_validator('acceptance_length_threshold') - @classmethod - def _validate_acceptance_length_threshold(cls, v: Optional[float]): - if v is None: - return v - if v < 0: - raise ValueError( - f"acceptance_length_threshold must be >= 0, got {v}") - return v + max_concurrency: Optional[NonNegativeInt] = Field( + default=None, + description= + "When specified, speculation will be disabled at batch sizes above this value. Otherwise, " + "speculation will always be on. PyTorch backend only.") + + draft_len_schedule: Optional[dict[int, int]] = Field( + default=None, + description= + "Developer interface: dynamically adjust draft length based on active batch size in runtime. " + "Maps batch size to draft lengths. For example, {1: 4, 4: 2, 8: 0} means: " + "batch_size >= 1 uses draft_len=4, batch_size >= 4 uses draft_len=2, " + "batch_size >= 8 uses draft_len=0 (disable speculation). " + "draft_len_schedule is enforced to contain batch_size=1 and its according draft_len equals " + "max_draft_len for consistency; for example, if max_draft_len=4, the schedule must contain {1: 4}." + ) + + load_format: Optional[str] = Field( + default=None, description="The load format of the speculative model.") + + acceptance_window: Optional[NonNegativeInt] = Field( + default=None, + description= + "The rolling average window size (N) for acceptance length across completed requests. " + "If not set or set to 0, the feature is disabled. PyTorch backend only." + ) + + acceptance_length_threshold: Optional[NonNegativeFloat] = Field( + default=None, + description= + "The threshold for average acceptance length; speculation will be disabled permanently once the " + "rolling average over the last N completed requests (N = acceptance_window) drops below this value. " + "PyTorch backend only.") + + allow_advanced_sampling: bool = Field( + default=False, + status="prototype", + description= + "If true, allows non-greedy sampling when speculation is used. Only applicable " + "to 1-model code paths; non-greedy sampling is always enabled on 2-model paths." + ) # If set, drafting is allowed to use chain drafter. _allow_chain_drafter: bool = PrivateAttr(True) @@ -766,41 +751,6 @@ def validate_draft_len_schedule_and_sort(cls, v, info): return dict(sorted(v.items(), key=lambda x: x[0])) return v - @classmethod - def from_dict(cls, data: dict, backend: Optional[str] = None): - # dispatch to the correct decoding config - decoding_type = data.get("decoding_type") - config_classes = { - "MTP": MTPDecodingConfig, - "Medusa": MedusaDecodingConfig, - "Eagle": EagleDecodingConfig, - "Eagle3": Eagle3DecodingConfig, - "Lookahead": LookaheadDecodingConfig, - "NGram": NGramDecodingConfig, - "DraftTarget": DraftTargetDecodingConfig, - "SaveState": SaveHiddenStatesDecodingConfig, - "UserProvided": UserProvidedDecodingConfig, - "AUTO": AutoDecodingConfig, - } - - backend = backend.lower() if isinstance(backend, str) else backend - if decoding_type == "Eagle" and backend in ("pytorch", "_autodeploy"): - data = dict(data) - data.pop("decoding_type") - spec_cfg = Eagle3DecodingConfig(**data) - spec_cfg._decoding_type_alias = "Eagle" - return spec_cfg - - config_class = config_classes.get(decoding_type) - if config_class is None: - raise ValueError(f"Invalid decoding type: {decoding_type}") - data.pop("decoding_type") - - return config_class(**data) - - def _check_fields(self): - pass - def supports_backend(self, backend: str) -> bool: """ Override if the speculation algorithm does not support @@ -808,11 +758,6 @@ def supports_backend(self, backend: str) -> bool: """ return True - def validate(self) -> None: - """ - Do any additional error checking here. - """ - @functools.cached_property def spec_dec_mode(self): # spec_dec_mode has more functionality than the raw decoding_mode string. @@ -875,40 +820,65 @@ def validate_calibration_file_path(self) -> 'LayerwiseBenchmarksConfig': class MedusaDecodingConfig(DecodingBaseConfig): + decoding_type: Literal["Medusa"] = "Medusa" medusa_choices: Optional[List[List[int]]] = None num_medusa_heads: Optional[int] = None - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.max_total_draft_tokens = self.max_draft_len # Current Medusa only support linear tree - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - - decoding_type: ClassVar[str] = "Medusa" + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_total_draft_tokens = self.max_draft_len # Current Medusa only supports linear tree + return self def supports_backend(self, backend: str) -> bool: return backend not in ("pytorch", "_autodeploy") class EagleDecodingConfig(DecodingBaseConfig): - eagle_choices: Optional[List[List[int]]] = None - greedy_sampling: Optional[bool] = True - posterior_threshold: Optional[float] = None - # Whether to use dynamic tree. - use_dynamic_tree: Optional[bool] = False - # The topK value for each layer when enable dynamic tree. - dynamic_tree_max_topK: Optional[int] = None - # The number of eagle layer. will not be used in pytorch flow, just for compatibility with TRT flow - num_eagle_layers: Optional[int] = None - # The number of non-leaves in each layer. - max_non_leaves_per_layer: Optional[int] = None - eagle3_one_model: Optional[bool] = True - eagle3_layers_to_capture: Optional[Set[int]] = None - # The model architecture of the eagle3 model. - # choices: llama3, mistral_large3 - eagle3_model_arch: str = "llama3" + decoding_type: Literal["Eagle"] = "Eagle" + eagle_choices: Optional[List[List[int]]] = Field( + default=None, + description= + "Static tree structure for draft token generation. Each sublist represents a path in the tree. Mutually exclusive with use_dynamic_tree." + ) + greedy_sampling: Optional[bool] = Field( + default=True, + description= + "Whether to use greedy sampling (Top-1 with token equality acceptance) or typical acceptance with multinomial sampling." + ) + posterior_threshold: Optional[float] = Field( + default=None, + description= + "Minimum token probability threshold for typical acceptance. Corresponds to epsilon in https://arxiv.org/pdf/2401.10774." + ) + use_dynamic_tree: Optional[bool] = Field( + default=False, + description= + "Whether to use dynamic tree (Eagle-2 algorithm). Mutually exclusive with eagle_choices." + ) + dynamic_tree_max_topK: Optional[int] = Field( + default=None, + description="The topK value for each layer when dynamic tree is enabled." + ) + num_eagle_layers: Optional[int] = Field( + default=None, + description= + "The number of eagle layers. Will not be used in pytorch flow, just for compatibility with TRT flow." + ) + max_non_leaves_per_layer: Optional[int] = Field( + default=None, description="The number of non-leaves in each layer.") + eagle3_one_model: Optional[bool] = Field( + default=True, + description= + "Whether to use the faster one-model implementation (draft as submodule) or the two-model implementation." + ) + eagle3_layers_to_capture: Optional[Set[int]] = Field( + default=None, + description= + "Target model layer indices to capture hidden states from for the EAGLE3 draft model. Defaults to {1, num_layers//2-1, num_layers-4}." + ) + eagle3_model_arch: Literal["llama3", "mistral_large3"] = Field( + default="llama3", + description="The model architecture of the eagle3 model.") @field_validator('eagle_choices', mode='before') @classmethod @@ -928,8 +898,8 @@ def validate_eagle_choices(cls, v): @model_validator(mode='after') def validate_eagle_config(self) -> 'EagleDecodingConfig': - if self.max_draft_len is None: - raise ValueError("max_draft_len is required for Eagle") + if self.max_draft_len is None or self.max_draft_len <= 0: + raise ValueError("max_draft_len must be > 0 for Eagle") self.num_eagle_layers = self.max_draft_len self.max_total_draft_tokens = self.max_draft_len # If using linear-tree, the max_total_draft_tokens is the same as max_draft_len @@ -941,10 +911,9 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': # Checks whether the input eagle choices is valid # and reset the max_draft_len and num_eagle_layers if necessary if self.eagle_choices is not None: - # If eagle_choices is provided, use_dynamic_tree should not be used if self.use_dynamic_tree: raise ValueError( - "If eagle_choices is provided, use_dynamic_tree need to be False" + "If eagle_choices is provided, use_dynamic_tree should be False" ) # Get num_eagle_layers from eagle_choices @@ -980,15 +949,11 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': return self - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - - decoding_type: ClassVar[str] = "Eagle" - - def validate(self) -> None: + @model_validator(mode="after") + def validate_speculative_model(self) -> None: if self.speculative_model is None: raise ValueError("Draft model must be provided for EAGLE") + return self def check_eagle_choices(self): # 1) Check connectivity @@ -1036,37 +1001,28 @@ def is_linear_tree(self) -> bool: class Eagle3DecodingConfig(EagleDecodingConfig): - decoding_type: ClassVar[str] = "Eagle3" + decoding_type: Literal["Eagle3"] = "Eagle3" class SaveHiddenStatesDecodingConfig(DecodingBaseConfig): + decoding_type: Literal["SaveState"] = "SaveState" output_directory: str write_interval: int = 20 file_prefix: str = "data" - eagle3_layers_to_capture: Optional[Set[int]] = None + eagle3_layers_to_capture: Set[int] = Field(..., min_length=1) max_total_draft_tokens: Optional[int] = Field(default=1, init=False) eagle_choices: Optional[List[List[int]]] = Field(default=None, init=False) - def model_post_init(self, __context): + _last_hidden_in_save: bool = PrivateAttr(default=True) + + @model_validator(mode="after") + def set_last_hidden_in_save(self): self._last_hidden_in_save = True - if self.eagle3_layers_to_capture is None: - self._last_hidden_in_save = False - elif -1 not in self.eagle3_layers_to_capture: + if -1 not in self.eagle3_layers_to_capture: self._last_hidden_in_save = False self.eagle3_layers_to_capture.add(-1) - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - - decoding_type: ClassVar[str] = "SaveState" - - def validate(self) -> None: - if self.output_directory is None or not self.eagle3_layers_to_capture: - raise ValueError( - "Save directory and layers to capture must be provided") - @functools.cached_property def spec_dec_mode(self): from tensorrt_llm._torch.speculative.interface import \ @@ -1086,109 +1042,123 @@ def num_capture_layers(self): class UserProvidedDecodingConfig(DecodingBaseConfig): + decoding_type: Literal["User_Provided"] = "User_Provided" # Cannot use real type annotations due to circular imports drafter: object # Type is Drafter resource_manager: object = None # Type is Optional[ResourceManager] - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.max_total_draft_tokens = self.max_draft_len # Current UserProvided only support linear tree - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - - decoding_type: ClassVar[str] = "User_Provided" + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_total_draft_tokens = self.max_draft_len # Current UserProvided only supports linear tree + return self class NGramDecodingConfig(DecodingBaseConfig): """ Configuration for NGram drafter speculative decoding. - - Arguments: - max_draft_len: int - The length maximum of draft tokens (can be understood as length maximum of output draft tokens). - - max_matching_ngram_size: int - The length maximum of searching tokens (can be understood as length maximum of input tokens to search). - - is_keep_all: bool = True - Whether to keep all candidate pattern-matches pairs, only one match is kept for each pattern if False. - - is_use_oldest: bool = True - Whether to provide the oldest match when pattern is hit, the newest one is provided if False. - - is_public_pool: bool = True - Whether to use a common pool for all requests, or the pool is private for each request if False. """ - max_matching_ngram_size: int = 0 - is_keep_all: bool = True - is_use_oldest: bool = True - is_public_pool: bool = True - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.max_total_draft_tokens = self.max_draft_len # Current NGram only support linear tree - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) + decoding_type: Literal["NGram"] = "NGram" + max_matching_ngram_size: PositiveInt = Field( + default=2, + description= + "The length maximum of searching tokens (can be understood as length maximum of input tokens " + "to search).") + is_keep_all: bool = Field( + default=True, + description= + "Whether to keep all candidate pattern-matches pairs, only one " + "match is kept for each pattern if False.") + is_use_oldest: bool = Field( + default=True, + description="Whether to provide the oldest match when pattern is hit, " + "the newest one is provided if False.") + is_public_pool: bool = Field( + default=True, + description="Whether to use a common pool for all requests, or the pool " + "is private for each request if False.") - decoding_type: ClassVar[str] = "NGram" + @model_validator(mode="after") + def validate_ngram_config(self): + if self.max_draft_len is None or self.max_draft_len <= 0: + raise ValueError("max_draft_len must be > 0 for NGram") + self.max_total_draft_tokens = self.max_draft_len # Current NGram only supports linear tree + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" class DraftTargetDecodingConfig(DecodingBaseConfig): + decoding_type: Literal["Draft_Target"] = "Draft_Target" - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.max_total_draft_tokens = self.max_draft_len # Current DraftTarget only support linear tree - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - - decoding_type: ClassVar[str] = "Draft_Target" + @model_validator(mode="after") + def validate_draft_target_config(self): + if self.max_draft_len is None or self.max_draft_len <= 0: + raise ValueError("max_draft_len must be > 0 for DraftTarget") + if self.speculative_model is None: + raise ValueError( + "speculative_model must be specified for DraftTarget") + self.max_total_draft_tokens = self.max_draft_len # Current DraftTarget only supports linear tree + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" or backend == "_autodeploy" class MTPDecodingConfig(DecodingBaseConfig): - num_nextn_predict_layers: int = 1 - use_relaxed_acceptance_for_thinking: bool = False - relaxed_topk: int = 1 - relaxed_delta: float = 0. - use_mtp_vanilla: bool = False - mtp_eagle_one_model: bool = True + decoding_type: Literal["MTP"] = "MTP" + num_nextn_predict_layers: PositiveInt = Field( + default=1, + description= + "Number of MTP modules. Each module predicts the next token, so N modules produce N draft tokens." + ) + use_relaxed_acceptance_for_thinking: bool = Field( + default=False, + description= + "Enable relaxed acceptance during thinking phase for reasoning models. Accepts draft tokens matching any top-K candidate instead of exact top-1." + ) + relaxed_topk: int = Field( + default=1, + description= + "Number of top candidate tokens to consider for relaxed acceptance. Draft token is accepted if it matches any of these." + ) + relaxed_delta: float = Field( + default=0., + description= + "Probability threshold for relaxed acceptance. Only candidates with prob >= (top-1 prob - delta) are kept." + ) + use_mtp_vanilla: bool = Field( + default=False, + description= + "Force vanilla MTP mode (sequential MTP layers). When False, uses EAGLE-style MTP for single-layer checkpoints." + ) + mtp_eagle_one_model: bool = Field( + default=True, + description= + "When using EAGLE-style MTP, use faster one-model implementation (drafter as submodule) vs two-model." + ) # TODO: remove this after distinguishing `max_draft_len` and `num_nextn_predict_layers` # Now we need a flag when MTPDecodingConfig is updated by PyTorchModelEngine. - num_nextn_predict_layers_from_model_config: int = 1 - - # When encounter , start thinking phase. - # When encounter , end thinking phase. - # [thinking phase] [real output] - begin_thinking_phase_token: int = 128798 - end_thinking_phase_token: int = 128799 - - def __init__(self, **kwargs): - super().__init__(**kwargs) - if 'num_nextn_predict_layers' in kwargs: - self.max_draft_len = kwargs['num_nextn_predict_layers'] - self.max_total_draft_tokens = kwargs[ - 'num_nextn_predict_layers'] # Current MTP only support linear tree + num_nextn_predict_layers_from_model_config: int = Field(default=1, + init=False) - @classmethod - def from_dict(cls, data: dict): - out = cls(**data) - out.max_draft_len = out.num_nextn_predict_layers - out.max_total_draft_tokens = out.num_nextn_predict_layers # Current MTP only support linear tree - return out + begin_thinking_phase_token: int = Field( + default=128798, + description= + "Token ID marking start of thinking phase. Relaxed acceptance only applies within this phase." + ) + end_thinking_phase_token: int = Field( + default=128799, + description= + "Token ID marking end of thinking phase. Strict acceptance resumes after this." + ) - decoding_type: ClassVar[str] = "MTP" + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_draft_len = self.num_nextn_predict_layers + self.max_total_draft_tokens = self.num_nextn_predict_layers # Current MTP only supports linear tree + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -1220,15 +1190,12 @@ class AutoDecodingConfig(DecodingBaseConfig): Attributes that are inherited from the base class are ignored. """ - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.max_total_draft_tokens = self.max_draft_len # Current Auto only support linear tree - - @classmethod - def from_dict(cls, data: dict): - return cls(**data) + decoding_type: Literal["AUTO"] = "AUTO" - decoding_type: ClassVar[str] = "AUTO" + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_total_draft_tokens = self.max_draft_len # Current Auto only supports linear tree + return self def supports_backend(self, backend: str) -> bool: return backend == "pytorch" @@ -1492,13 +1459,16 @@ class DynamicBatchConfig(StrictBaseModel, PybindMirror): Controls how batch size and token limits are dynamically adjusted at runtime. """ enable_batch_size_tuning: bool = Field( + default=True, description="Controls if the batch size should be tuned dynamically") enable_max_num_tokens_tuning: bool = Field( + default=False, description="Controls if the max num tokens should be tuned dynamically" ) dynamic_batch_moving_average_window: int = Field( + default=128, description= "The window size for moving average of input and output length which is used to calculate dynamic batch size and max num tokens" ) @@ -1520,8 +1490,9 @@ class SchedulerConfig(StrictBaseModel, PybindMirror): context_chunking_policy: Optional[ContextChunkingPolicy] = Field( default=None, description="The context chunking policy to use") - dynamic_batch_config: Optional[DynamicBatchConfig] = Field( - default=None, description="The dynamic batch config to use") + dynamic_batch_config: DynamicBatchConfig = Field( + default_factory=DynamicBatchConfig, + description="The dynamic batch config to use") def _to_pybind(self): return _SchedulerConfig( @@ -1609,40 +1580,29 @@ class LookaheadDecodingConfig(DecodingBaseConfig, PybindMirror): Configuration for lookahead speculative decoding. """ - max_window_size: int = Field( + decoding_type: Literal["Lookahead"] = "Lookahead" + max_window_size: PositiveInt = Field( default=_LookaheadDecodingConfig.get_default_lookahead_decoding_window( ), description="Number of NGrams in lookahead branch per step.") - max_ngram_size: int = Field( + max_ngram_size: PositiveInt = Field( default=_LookaheadDecodingConfig.get_default_lookahead_decoding_ngram(), description="Number of tokens per NGram.") - max_verification_set_size: int = Field( + max_verification_set_size: PositiveInt = Field( default=_LookaheadDecodingConfig. get_default_lookahead_decoding_verification_set(), description="Number of NGrams in verification branch per step.") - @field_validator('max_window_size', 'max_ngram_size', - 'max_verification_set_size') - @classmethod - def validate_positive_values(cls, v): - if v <= 0: - raise ValueError(f"Value must be positive, got {v}") - return v - - def __init__(self, **data): - super().__init__(**data) - self.max_total_draft_tokens = self.max_draft_len # Current Lookahead only support linear tree - self._check_fields() + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_total_draft_tokens = self.max_draft_len # Current Lookahead only supports linear tree + return self def calculate_speculative_resource(self): return _LookaheadDecodingConfig.calculate_speculative_resource_tuple( self.max_window_size, self.max_ngram_size, self.max_verification_set_size) - @classmethod - def from_dict(cls, data: dict): - return cls(**data) - def _to_pybind(self): return _LookaheadDecodingConfig(self.max_window_size, self.max_ngram_size, @@ -1651,25 +1611,30 @@ def _to_pybind(self): def supports_backend(self, backend: str) -> bool: return backend not in ("pytorch", "_autodeploy") - decoding_type: ClassVar[str] = "Lookahead" - - -SpeculativeConfig: TypeAlias = Optional[Union[ - DraftTargetDecodingConfig, - EagleDecodingConfig, - LookaheadDecodingConfig, - MedusaDecodingConfig, - MTPDecodingConfig, - NGramDecodingConfig, - UserProvidedDecodingConfig, - SaveHiddenStatesDecodingConfig, - AutoDecodingConfig, -]] - -SparseAttentionConfig: TypeAlias = Union[ - RocketSparseAttentionConfig, - DeepSeekSparseAttentionConfig, - SkipSoftmaxAttentionConfig, + +SpeculativeConfig: TypeAlias = Annotated[ + Union[ + DraftTargetDecodingConfig, + EagleDecodingConfig, + Eagle3DecodingConfig, + LookaheadDecodingConfig, + MedusaDecodingConfig, + MTPDecodingConfig, + NGramDecodingConfig, + UserProvidedDecodingConfig, + SaveHiddenStatesDecodingConfig, + AutoDecodingConfig, + ], + Field(discriminator="decoding_type"), +] + +SparseAttentionConfig: TypeAlias = Annotated[ + Union[ + RocketSparseAttentionConfig, + DeepSeekSparseAttentionConfig, + SkipSoftmaxAttentionConfig, + ], + Field(discriminator="algorithm"), ] @@ -1687,8 +1652,9 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "The maximum number of tokens that should be stored in the KV cache. If both `max_tokens` and `free_gpu_memory_fraction` are specified, memory corresponding to the minimum will be used." ) - max_attention_window: Optional[List[int]] = Field( + max_attention_window: Optional[List[PositiveInt]] = Field( default=None, + min_length=1, description= "Size of the attention window for each sequence. Only the last tokens will be stored in the KV cache. If the number of elements in `max_attention_window` is less than the number of layers, `max_attention_window` will be repeated multiple times to the number of layers." ) @@ -1698,6 +1664,8 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): "Number of sink tokens (tokens to always keep in attention window).") free_gpu_memory_fraction: Optional[float] = Field( default=0.9, + ge=0, + le=1, description= "The fraction of GPU memory fraction that should be allocated for the KV cache. Default is 90%. If both `max_tokens` and `free_gpu_memory_fraction` are specified, memory corresponding to the minimum will be used." ) @@ -1739,7 +1707,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): ) use_uvm: bool = Field(default=False, description="Whether to use UVM for the KV cache.") - max_gpu_total_bytes: int = Field( + max_gpu_total_bytes: NonNegativeInt = Field( default=0, description= "The maximum size in bytes of GPU memory that can be allocated for the KV cache. If both `max_gpu_total_bytes` and `free_gpu_memory_fraction` are specified, memory corresponding to the minimum will be allocated." @@ -1779,47 +1747,6 @@ def _to_pybind(self): attention_dp_events_gather_period_ms, max_gpu_total_bytes=self.max_gpu_total_bytes) - @field_validator('free_gpu_memory_fraction') - @classmethod - def validate_free_gpu_memory_fraction(cls, v: float): - """Validates that the fraction is between 0.0 and 1.0.""" - if not 0 <= v <= 1: - raise ValueError( - "kv_cache_config.free_gpu_memory_fraction must be a float between 0 and 1" - ) - return v - - @field_validator('max_gpu_total_bytes') - @classmethod - def validate_max_gpu_total_bytes(cls, v: int): - if v < 0: - raise ValueError( - "kv_cache_config.max_gpu_total_bytes must be non-negative") - return v - - @field_validator('max_attention_window') - @classmethod - def validate_max_attention_window(cls, v: Optional[List[int]]): - # Allow unset - if v is None: - return v - - # Must be a non-empty list of positive integers - if not isinstance(v, list) or len(v) == 0: - raise ValueError( - "kv_cache_config.max_attention_window must be a non-empty list of positive integers" - ) - for i in v: - if not isinstance(i, int): - raise ValueError( - "kv_cache_config.max_attention_window must contain only integers" - ) - if i <= 0: - raise ValueError( - "kv_cache_config.max_attention_window values must be positive" - ) - return v - @PybindMirror.mirror_pybind_fields(_ExtendedRuntimePerfKnobConfig) class ExtendedRuntimePerfKnobConfig(StrictBaseModel, PybindMirror): @@ -1868,16 +1795,14 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): default=None, description="The max number of tokens the transfer buffer can fit.") - kv_transfer_timeout_ms: Optional[int] = Field( + kv_transfer_timeout_ms: Optional[PositiveInt] = Field( default=None, - gt=0, description= "Timeout in milliseconds for KV cache transfer. Requests exceeding this timeout will be cancelled." ) - kv_transfer_sender_future_timeout_ms: Optional[int] = Field( + kv_transfer_sender_future_timeout_ms: Optional[PositiveInt] = Field( default=1000, - gt=0, description= "Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms" ) @@ -1935,10 +1860,7 @@ class BaseLlmArgs(StrictBaseModel): """ Base class for both TorchLlmArgs and TrtLlmArgs. It contains all the arguments that are common to both. """ - model_config = { - "arbitrary_types_allowed": True, - "extra": "forbid", - } + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") # Explicit arguments model: Union[str, Path] = Field( @@ -2034,9 +1956,10 @@ class BaseLlmArgs(StrictBaseModel): "Pipeline parallel partition, a list of each rank's layer number.", status="prototype") - cp_config: Optional[dict] = Field(default_factory=dict, - description="Context parallel config.", - status="prototype") + cp_config: Optional[CpConfig] = Field( + default=None, + description="Context parallel config.", + status="prototype") load_format: Literal['auto', 'dummy'] = Field( default='auto', @@ -2055,6 +1978,9 @@ class BaseLlmArgs(StrictBaseModel): lora_config: Optional[LoraConfig] = Field( default=None, description="LoRA configuration for the model.") + quant_config: QuantConfig = Field(default_factory=QuantConfig, + description="Quantization config.") + # Several options from ExecutorConfig, expanded here for less hierarchy kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") @@ -2105,21 +2031,21 @@ class BaseLlmArgs(StrictBaseModel): status="prototype") # Speculative decoding parameters - speculative_config: SpeculativeConfig = Field( + speculative_config: Optional[SpeculativeConfig] = Field( default=None, description="Speculative decoding config.") - max_batch_size: Optional[int] = Field(default=None, + max_batch_size: Optional[int] = Field(default=2048, description="The maximum batch size.") # generation constraints max_input_len: Optional[int] = Field( - default=None, description="The maximum input length.") + default=1024, description="The maximum input length.") max_seq_len: Optional[int] = Field( default=None, description="The maximum sequence length.") - max_beam_width: Optional[int] = Field(default=None, - description="The maximum beam width.") + max_beam_width: int = Field(default=1, + description="The maximum beam width.") max_num_tokens: Optional[int] = Field( default=8192, description="The maximum number of tokens.") @@ -2219,36 +2145,10 @@ def speculative_model(self) -> Optional[Union[str, Path]]: return self.speculative_config.speculative_model if self.speculative_config is not None else None @classmethod - def from_kwargs(cls, **kwargs: Any) -> "BaseLlmArgs": - """Create `LlmArgs` instance from kwargs. - - Args: - kwargs (Any): Arguments passed to `LlmArgs` constructor. - - Returns: - tensorrt_llm.llmapi.llm_utils.BaseLlmArgs: The `BaseLlmArgs` instance. - """ - kwargs = BaseLlmArgs._check_consistency(dict(kwargs)) - ret = cls(**kwargs) - return ret - - @staticmethod - def _check_consistency(kwargs_dict: Dict[str, Any]) -> Dict[str, Any]: - # max_beam_width is not included since vague behavior due to lacking the support for dynamic beam width during - # generation - black_list = set(["max_beam_width"]) - executor_config_attrs = set( - attr for attr in dir(_ExecutorConfig) if not attr.startswith('_') - and callable(getattr(_ExecutorConfig, attr))) - executor_config_attrs -= black_list - llm_args_attr = set(BaseLlmArgs.model_fields.keys()) - # NOTE: When cpp ExecutorConfig add new options, please add the new options into `LlmArgs` with docs as well - # ASK chunweiy for help if you are not sure about the new options. - assert executor_config_attrs.issubset( - llm_args_attr - ), f"New options found in underlying ExecutorConfig: {llm_args_attr - executor_config_attrs}" - - return kwargs_dict + def from_yaml(cls, yaml_path: Union[str, Path]): + with open(yaml_path, "r") as f: + config_dict = yaml.safe_load(f) + return cls(**config_dict) @field_validator("dtype") @classmethod @@ -2271,13 +2171,6 @@ def validate_gpus_per_node(cls, v, info): v = torch.cuda.device_count() return v - @field_validator("model") - @classmethod - def validate_model(cls, v, info): - if not isinstance(v, (str, Path)): - raise ValueError(f"Invalid model: {v}") - return v - @model_validator(mode="after") def validate_parallel_config(self): if self.moe_cluster_parallel_size is None: @@ -2304,9 +2197,17 @@ def validate_parallel_config(self): return self @model_validator(mode="after") - def set_default_max_input_len(self): - if self.max_input_len is None: - self.max_input_len = 1024 + def set_postprocess_tokenizer_dir(self): + """Set postprocess_tokenizer_dir if not explicitly provided. + + This must run before validate_and_init_tokenizer since that validator + converts tokenizer from a path to an object. + """ + if self.postprocess_tokenizer_dir is None: + if isinstance(self.tokenizer, (str, Path)): + self.postprocess_tokenizer_dir = str(self.tokenizer) + else: + self.postprocess_tokenizer_dir = str(self.model) return self @model_validator(mode="after") @@ -2419,27 +2320,21 @@ class TrtLlmArgs(BaseLlmArgs): description="The workspace for the model.") # Once set, the model will reuse the build_cache - enable_build_cache: object = Field( - default=False, - description="Enable build cache.", - json_schema_extra={ - "type": f"Union[{get_type_repr(BuildCacheConfig)}, bool]" - }) + enable_build_cache: Union[BuildCacheConfig, + bool] = Field(default=False, + description="Enable build cache.") extended_runtime_perf_knob_config: Optional[ ExtendedRuntimePerfKnobConfig] = Field( default=None, description="Extended runtime perf knob config.") - calib_config: Optional[CalibConfig] = Field( - default=None, description="Calibration config.", validate_default=True) - - # Quantization and calibration configurations - quant_config: Optional[QuantConfig] = Field( - default=None, description="Quantization config.", validate_default=True) + calib_config: CalibConfig = Field(default_factory=CalibConfig, + description="Calibration config.") - embedding_parallel_mode: str = Field( - default='SHARDING_ALONG_VOCAB', - description="The embedding parallel mode.") + embedding_parallel_mode: Literal[ + 'NONE', 'SHARDING_ALONG_VOCAB', 'SHARDING_ALONG_HIDDEN'] = Field( + default='SHARDING_ALONG_VOCAB', + description="The embedding parallel mode.") fast_build: bool = Field(default=False, description="Enable fast build.") @@ -2614,6 +2509,14 @@ def validate_speculative_config(self): return self + @model_validator(mode="after") + def set_orchestrator_type_for_mpi_disabled(self): + if mpi_disabled(): + self.orchestrator_type = "ray" + elif self.orchestrator_type == "ray": + os.environ["TLLM_DISABLE_MPI"] = "1" + return self + def _load_config_from_engine(self, engine_dir: Path): engine_config = EngineConfig.from_json_file(engine_dir / "config.json") self._pretrained_config = engine_config.pretrained_config @@ -2713,20 +2616,6 @@ def validate_model_format_misc(self): self._model_format = model_format return self - @field_validator('calib_config', mode='before') - @classmethod - def init_calib_config(cls, v): - if v is None: - return CalibConfig() - return v - - @field_validator("quant_config", mode='before') - @classmethod - def validate_quant_config(cls, v, info): - if v is None: - v = QuantConfig() - return v - @model_validator(mode="after") def setup_embedding_parallel_mode(self): if self.embedding_parallel_mode == 'NONE': @@ -2746,9 +2635,6 @@ def validate_enable_build_cache(self): return self self.enable_build_cache = BuildCacheConfig() if isinstance( self.enable_build_cache, bool) else self.enable_build_cache - if not isinstance(self.enable_build_cache, BuildCacheConfig): - raise ValueError( - f"Invalid build_cache_config: {self.enable_build_cache}") return self @model_validator(mode="after") @@ -2787,7 +2673,7 @@ class TorchCompileConfig(StrictBaseModel): default=False, description="Enable piecewise CUDA graph in torch.compile.") - capture_num_tokens: Optional[List[int]] = Field( + capture_num_tokens: Optional[List[PositiveInt]] = Field( default=None, description= "List of num of tokens to capture the piecewise CUDA graph for. If not provided, the number of tokens will be the same as cuda_graph_config.batch_sizes." @@ -2798,8 +2684,6 @@ class TorchCompileConfig(StrictBaseModel): def validate_capture_num_tokens(cls, v): if v is None: return v - if any(t <= 0 for t in v): - raise ValueError("capture_num_tokens must contain positive ints.") return sorted(set(v), reverse=True) enable_userbuffers: bool = Field( @@ -2807,23 +2691,17 @@ def validate_capture_num_tokens(cls, v): description= "When torch compile is enabled, userbuffers is enabled by default.") - max_num_streams: int = Field( + max_num_streams: PositiveInt = Field( default=1, description= "The maximum number of CUDA streams to use for torch.compile.") - @field_validator('max_num_streams') - @classmethod - def validate_torch_compile_max_num_streams(cls, v): - """Validate torch_compile_config.max_num_streams >= 1.""" - if v < 1: - raise ValueError( - "torch_compile_config.max_num_streams must be >= 1") - return v - - @staticmethod - def _generate_capture_num_tokens() -> List[int]: - return [2**i for i in range(8)] + [i for i in range(256, 3073, 256)] + @model_validator(mode='after') + def set_default_capture_num_tokens(self) -> 'TorchCompileConfig': + if self.enable_piecewise_cuda_graph and self.capture_num_tokens is None: + self.capture_num_tokens = [2**i for i in range(8) + ] + [i for i in range(256, 3073, 256)] + return self class TorchLlmArgs(BaseLlmArgs): @@ -2837,7 +2715,7 @@ class TorchLlmArgs(BaseLlmArgs): cuda_graph_config: Optional[CudaGraphConfig] = Field( default_factory=CudaGraphConfig, - description="CUDA graph config.If true, use CUDA graphs for decoding. \ + description="CUDA graph config. If true, use CUDA graphs for decoding. \ CUDA graphs are only created for the batch sizes in cuda_graph_config.batch_sizes, \ and are enabled for batches that consist of decoding requests *only* \ (the reason is that it's hard to capture a single graph with prefill requests \ @@ -2904,13 +2782,13 @@ class TorchLlmArgs(BaseLlmArgs): "The maximum number of requests for perf metrics. Must also set request_perf_metrics to true to get perf metrics.", status="prototype") - batch_wait_timeout_ms: float = Field( + batch_wait_timeout_ms: NonNegativeFloat = Field( default=0, description= "If greater than 0, the request queue might wait up to batch_wait_timeout_ms to receive max_batch_size requests, if fewer than max_batch_size requests are currently available. If 0, no waiting occurs.", status="prototype") - batch_wait_timeout_iters: int = Field( + batch_wait_timeout_iters: NonNegativeInt = Field( default=0, description= "Maximum number of iterations the scheduler will wait to accumulate new coming requests for improved GPU utilization efficiency. If greater than 0, the scheduler will delay batch processing to gather more requests up to the specified iteration limit. If 0, disables timeout-iters-based batching delays.", @@ -2918,6 +2796,8 @@ class TorchLlmArgs(BaseLlmArgs): batch_wait_max_tokens_ratio: float = Field( default=0, + ge=0, + le=1, description= "Token accumulation threshold ratio for batch scheduling optimization. If greater than 0, the scheduler will accumulate requests locally until the total token count reaches batch_wait_max_tokens_ratio * max_num_tokens. This mechanism enhances GPU utilization efficiency by ensuring adequate batch sizes.If 0 disables token-based batching delays.", status="prototype") @@ -2950,7 +2830,7 @@ class TorchLlmArgs(BaseLlmArgs): ) # TODO: make this a per-request parameter - stream_interval: int = Field( + stream_interval: PositiveInt = Field( default=1, description= "The iteration interval to create responses under the streaming mode. " @@ -3030,9 +2910,6 @@ class TorchLlmArgs(BaseLlmArgs): "Only enable it if you intend to use this feature.", status="prototype") - # PrivateVars - _quant_config: Optional[QuantConfig] = PrivateAttr(default=None) - disable_flashinfer_sampling: bool = Field( default=False, description= @@ -3051,22 +2928,8 @@ class TorchLlmArgs(BaseLlmArgs): description="Configuration for layer-wise benchmarks calibration.", status="prototype") - @property - def quant_config(self) -> QuantConfig: - if self._quant_config is None: - self._quant_config = QuantConfig() - return self._quant_config - - @quant_config.setter - def quant_config(self, value: QuantConfig): - self._quant_config = value - # TODO: remove backend later - @field_validator('backend', mode='before') - def init_backend(cls, v): - if v is None: - return 'pytorch' - return v + backend: Literal["pytorch"] = "pytorch" @field_validator('load_format', mode='before') @classmethod @@ -3095,12 +2958,8 @@ def extra_resource_managers(self, value: Dict[str, object]) -> None: self._extra_resource_managers = value @model_validator(mode="after") - def validate_misc(self): + def set_model_format(self): self._model_format = _ModelFormatKind.HF - if self.max_beam_width is None: - self.max_beam_width = 1 - if self.max_batch_size is None: - self.max_batch_size = 2048 return self @model_validator(mode="after") @@ -3111,33 +2970,19 @@ def validate_speculative_config(self): f"Speculation type {self.speculative_config.decoding_type} does not " f"support backend {self.backend}") - if isinstance(self.speculative_config, EagleDecodingConfig): - if (getattr(self.speculative_config, "_decoding_type_alias", - None) == "Eagle" or type(self.speculative_config) - is EagleDecodingConfig): - logger.warning( - "speculative_config.decoding_type 'Eagle' is not supported on the PyTorch backend; only 'Eagle3' is supported. " - "'Eagle' is treated as 'Eagle3' for backward compatibility. " - "EAGLE (v1/v2) draft checkpoints are incompatible with Eagle3—use an Eagle3 draft model." - ) - assert self.speculative_config.max_draft_len > 0 - assert self.speculative_config.speculative_model is not None, "EAGLE3 draft model must be specified." - elif isinstance(self.speculative_config, NGramDecodingConfig): - assert self.speculative_config.max_draft_len > 0 and self.speculative_config.max_matching_ngram_size > 0 - elif isinstance(self.speculative_config, DraftTargetDecodingConfig): - assert self.speculative_config.max_draft_len > 0 - assert self.speculative_config.speculative_model is not None, "Draft model must be specified." - elif isinstance(self.speculative_config, MTPDecodingConfig): - assert self.speculative_config.num_nextn_predict_layers > 0 - self.speculative_config.max_draft_len = self.speculative_config.num_nextn_predict_layers - elif isinstance(self.speculative_config, - UserProvidedDecodingConfig): - pass - elif isinstance(self.speculative_config, AutoDecodingConfig): - pass - elif isinstance(self.speculative_config, - SaveHiddenStatesDecodingConfig): - assert self.backend in ['pytorch'] + # If user passed decoding_type: Eagle on pytorch, convert to Eagle3 with warning + if type(self.speculative_config) is EagleDecodingConfig: + logger.warning( + "speculative_config.decoding_type 'Eagle' is not supported on the PyTorch backend; only 'Eagle3' is supported. " + "'Eagle' is treated as 'Eagle3' for backward compatibility. " + "EAGLE (v1/v2) draft checkpoints are incompatible with Eagle3—use an Eagle3 draft model." + ) + # Convert EagleDecodingConfig to Eagle3DecodingConfig + eagle_data = self.speculative_config.model_dump(exclude={"decoding_type"}) + self.speculative_config = Eagle3DecodingConfig(**eagle_data) + + if isinstance(self.speculative_config, + SaveHiddenStatesDecodingConfig): logger.warning( "SaveHiddenStatesDecodingConfig is active, setting max_batch_size to 1, disabling overlap scheduler, and setting cuda_graph_config to None" ) @@ -3145,23 +2990,12 @@ def validate_speculative_config(self): self.disable_overlap_scheduler = True self.cuda_graph_config = None self.speculative_config.max_draft_len = 1 - else: - raise ValueError( - f"Unrecognized speculative config type {type(self.speculative_config)}" - ) else: self.decoding_config = None return self - @model_validator(mode="after") - def validate_stream_interval(self): - if self.stream_interval <= 0: - raise ValueError( - f"stream_interval must be positive, got {self.stream_interval}") - return self - @model_validator(mode="after") def validate_checkpoint_format(self): if self.checkpoint_format is not None and self.checkpoint_loader is not None: @@ -3204,54 +3038,6 @@ def validate_load_balancer(self) -> 'TorchLlmArgs': ) from e return self - @model_validator(mode='after') - def validate_cuda_graph_config(self) -> 'TorchLlmArgs': - """Validate CUDA graph configuration. - - Ensures that: - 1. If cuda_graph_config.batch_sizes is provided, cuda_graph_config.max_batch_size must be 0 - 2. If cuda_graph_config.batch_sizes is not provided, it is generated based on cuda_graph_config.max_batch_size - 3. If both are provided, cuda_graph_config.batch_sizes must match the generated values - """ - if self.cuda_graph_config is None: - return self - - config = self.cuda_graph_config - - if config.batch_sizes: - config.batch_sizes = sorted(config.batch_sizes) - if config.max_batch_size != 0: - if config.batch_sizes != CudaGraphConfig._generate_cuda_graph_batch_sizes( - config.max_batch_size, config.enable_padding): - raise ValueError( - "Please don't set both cuda_graph_config.batch_sizes " - "and cuda_graph_config.max_batch_size.\n" - f"cuda_graph_config.batch_sizes: {self.cuda_graph_config.batch_sizes}, " - f"cuda_graph_config.max_batch_size: {self.cuda_graph_config.max_batch_size}" - ) - else: - config.max_batch_size = max(config.batch_sizes) - else: - max_batch_size = config.max_batch_size or 128 - generated_sizes = CudaGraphConfig._generate_cuda_graph_batch_sizes( - max_batch_size, config.enable_padding) - config.batch_sizes = generated_sizes - config.max_batch_size = max_batch_size - - return self - - @model_validator(mode='after') - def validate_torch_compile_config(self) -> 'TorchLlmArgs': - if self.torch_compile_config is None: - return self - - config = self.torch_compile_config - if config.enable_piecewise_cuda_graph: - if config.capture_num_tokens is None: - config.capture_num_tokens = TorchCompileConfig._generate_capture_num_tokens( - ) - return self - @model_validator(mode='after') def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': if self.kv_cache_config is None: @@ -3272,12 +3058,12 @@ def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': @model_validator(mode='after') def validate_helix_tokens_per_block(self) -> 'TorchLlmArgs': """Validate that cp_config.tokens_per_block matches kv_cache_config.tokens_per_block when HELIX parallelism is active.""" - if self.context_parallel_size == 1 or self.cp_config is None or not self.cp_config: + if self.context_parallel_size == 1 or self.cp_config is None: return self - cp_type = self.cp_config.get('cp_type', None) - if cp_type is not None and str(cp_type).upper() == 'HELIX': - cp_tokens_per_block = self.cp_config.get('tokens_per_block', None) + cp_type = self.cp_config.cp_type + if cp_type == CpType.HELIX: + cp_tokens_per_block = self.cp_config.tokens_per_block if cp_tokens_per_block is not None: kv_tokens_per_block = self.kv_cache_config.tokens_per_block assert cp_tokens_per_block == kv_tokens_per_block, ( @@ -3307,52 +3093,6 @@ def warn_on_unstable_feature_usage(self) -> 'TorchLlmArgs': return self - @model_validator(mode='after') - def validate_attention_dp_config(self) -> 'TorchLlmArgs': - """Validate attention DP configuration. - - Ensures that: - 1. If attention_dp_config.enable_balance is true, attention_dp_config.batching_wait_iters must be greater or equal to 0 - 2. If attention_dp_config.enable_balance is true, attention_dp_config.timeout_iters must be greater or equal to 0 - """ - if self.attention_dp_config is None: - return self - - config = self.attention_dp_config - if config.enable_balance: - if config.batching_wait_iters < 0: - raise ValueError( - "attention_dp_config.batching_wait_iters must be greater or equal to 0 when enable_balance is true" - ) - if config.timeout_iters < 0: - raise ValueError( - "attention_dp_config.timeout_iters must be greater or equal to 0 when enable_balance is true" - ) - return self - - @model_validator(mode='after') - def validate_batch_wait_timeout_ms(self) -> 'TorchLlmArgs': - """Validate batch wait timeout.""" - if self.batch_wait_timeout_ms < 0: - raise ValueError("batch_wait_timeout_ms must be greater than 0") - return self - - @model_validator(mode='after') - def validate_batch_wait_timeout_iters(self) -> 'TorchLlmArgs': - if self.batch_wait_timeout_iters < 0: - raise ValueError( - f"batch_wait_timeout_iters must be >= 0, got {self.batch_wait_timeout_iters}" - ) - return self - - @model_validator(mode='after') - def validate_batch_wait_max_tokens_ratio(self) -> 'TorchLlmArgs': - if self.batch_wait_max_tokens_ratio < 0 or self.batch_wait_max_tokens_ratio > 1: - raise ValueError( - f"batch_wait_max_tokens_ratio must be in range [0, 1], got {self.batch_wait_max_tokens_ratio}" - ) - return self - @model_validator(mode='after') def validate_ray_worker_extension_cls(self) -> 'TorchLlmArgs': if self.ray_worker_extension_cls is not None and self.orchestrator_type != "ray": @@ -3379,85 +3119,6 @@ def get_executor_config( return executor_config -def update_llm_args_with_extra_dict( - llm_args: Dict, - llm_args_dict: Dict, - extra_llm_api_options: Optional[str] = None) -> Dict: - - # Deep merge kv_cache_config to prevent partial YAML kv_cache_config from replacing the complete kv_cache_config - if 'kv_cache_config' in llm_args and 'kv_cache_config' in llm_args_dict: - # Convert KvCacheConfig object to dict if necessary - base_kv_config = llm_args['kv_cache_config'] - if isinstance(base_kv_config, KvCacheConfig): - base_kv_config = base_kv_config.model_dump(exclude_unset=True) - llm_args_dict['kv_cache_config'] = base_kv_config | llm_args_dict[ - 'kv_cache_config'] - - field_mapping = { - "quant_config": QuantConfig, - "calib_config": CalibConfig, - "build_config": BuildConfig, - "decoding_config": DecodingConfig, - "enable_build_cache": BuildCacheConfig, - "speculative_config": DecodingBaseConfig, - "lora_config": LoraConfig, - "moe_config": MoeConfig, - "nvfp4_gemm_config": Nvfp4GemmConfig, - "attention_dp_config": AttentionDpConfig, - "sparse_attention_config": BaseSparseAttentionConfig, - "kv_cache_config": KvCacheConfig, - } - for field_name, field_type in field_mapping.items(): - if field_name in llm_args_dict: - # Some fields need to be converted manually. - if field_name in ["speculative_config", "sparse_attention_config"]: - if field_name == "speculative_config": - backend = llm_args_dict.get("backend") or llm_args.get( - "backend") - llm_args_dict[field_name] = field_type.from_dict( - llm_args_dict[field_name], backend=backend) - else: - llm_args_dict[field_name] = field_type.from_dict( - llm_args_dict[field_name]) - else: - llm_args_dict[field_name] = field_type( - **llm_args_dict[field_name]) - extra_llm_str = f"because it's specified in {extra_llm_api_options}" if extra_llm_api_options else "" - logger.warning(f"Overriding {field_name} {extra_llm_str}") - - llm_args = llm_args | llm_args_dict - - # build_config only works for TensorRT backend, it will be ignored in PyTorch backend - if "build_config" in llm_args: - # Ensure build_config is a BuildConfig object, not a dict - if isinstance(llm_args["build_config"], dict): - llm_args["build_config"] = BuildConfig(**llm_args["build_config"]) - - for key in [ - "max_batch_size", - "max_num_tokens", - "max_beam_width", - "max_seq_len", - ]: - if key in llm_args_dict: - logger.info( - f"Overriding {key} from build_config to {llm_args_dict[key]}" - ) - setattr(llm_args["build_config"], key, llm_args_dict[key]) - - return llm_args - - -def update_llm_args_with_extra_options(llm_args: Dict, - extra_llm_api_options: str) -> Dict: - if extra_llm_api_options is not None: - with open(extra_llm_api_options, 'r') as f: - llm_args_dict = yaml.safe_load(f) - llm_args = update_llm_args_with_extra_dict(llm_args, llm_args_dict, - extra_llm_api_options) - return llm_args - - def get_model_format(model_dir: str, trust_remote_code: bool = False) -> _ModelFormatKind: ''' Get the format of the model. ''' diff --git a/tensorrt_llm/llmapi/llm_create.py b/tensorrt_llm/llmapi/llm_create.py new file mode 100644 index 000000000000..b23625da395e --- /dev/null +++ b/tensorrt_llm/llmapi/llm_create.py @@ -0,0 +1,85 @@ +from typing import Any + +import yaml + +from tensorrt_llm._torch.auto_deploy.llm import LLM as AutoDeployLLM +from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs as AutoDeployLlmArgs +from tensorrt_llm.llmapi.llm import _TorchLLM, _TrtLLM +from tensorrt_llm.llmapi.llm_args import BaseLlmArgs, TorchLlmArgs, TrtLlmArgs +from tensorrt_llm.logger import logger + + +def _deep_merge(base: dict, override: dict) -> dict: + """Deep merge override into base. Override values take precedence.""" + result = base.copy() + for key, value in override.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = _deep_merge(result[key], value) + else: + result[key] = value + return result + + +def get_llm_args_from_cli_params(model: str, **params: dict[str, Any]) -> BaseLlmArgs: + """Construct an instance of LlmArgs from CLI params and optionally a config YAML file. + + This function constructs LlmArgs according to the following logic: + 1. CLI params are parsed into a dict that aligns with the structure of LlmArgs. + 2. If a config YAML file is provided, it is deep-merged on top. + 3. The LlmArgs class for the correct backend is instantiated with the merged dict. + + This means that the config file takes precedence over CLI params. + """ + extra_llm_api_options = params.pop("extra_llm_api_options", None) + llm_args_dict = {} + + # Handle CLI params that map to nested config fields + params.setdefault("kv_cache_config", {}) + if "free_gpu_memory_fraction" in params: + params["kv_cache_config"]["free_gpu_memory_fraction"] = params.pop( + "free_gpu_memory_fraction" + ) + # TODO: align CLI param naming with trtllm-serve / trtllm-bench to avoid duplication with above + if "kv_cache_free_gpu_memory_fraction" in params: + params["kv_cache_config"]["free_gpu_memory_fraction"] = params.pop( + "kv_cache_free_gpu_memory_fraction" + ) + if "disable_kv_cache_reuse" in params: + params["kv_cache_config"]["enable_block_reuse"] = not params.pop("disable_kv_cache_reuse") + + # Override with config file + config_dict = {} + if extra_llm_api_options is not None: + with open(extra_llm_api_options, "r") as f: + config_dict = yaml.safe_load(f) or {} + llm_args_dict = _deep_merge(params, config_dict) + + # Determine backend and the corresponding LlmArgs class + backend = llm_args_dict.get("backend", "pytorch") + if backend == "pytorch": + llm_args_cls = TorchLlmArgs + elif backend == "_autodeploy": + llm_args_cls = AutoDeployLlmArgs + elif backend == "tensorrt": + llm_args_cls = TrtLlmArgs + else: + raise ValueError(f"Invalid backend: {backend}") + + logger.info(f"Creating LLM args from dict: {llm_args_dict}") + return llm_args_cls(model=model, **llm_args_dict) + + +def create_llm_from_llm_args(llm_args: BaseLlmArgs): + """Create and return the appropriate LLM instance for the backend corresponding to llm_args.""" + if llm_args.backend == "pytorch": + llm_cls = _TorchLLM + elif llm_args.backend == "_autodeploy": + llm_cls = AutoDeployLLM + else: + llm_cls = _TrtLLM + + # TODO: llm_args is converted to a dict here and then immediately parsed back into an LlmArgs instance + # in the LLM constructor, which is a bit hacky. This is due to the fact that the LLM constructor needs to + # accept kwargs (rather than an LlmArgs instance) for the offline API. + llm = llm_cls(**llm_args.model_dump()) + return llm diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 83e26ce93e46..acbc68c756dd 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -34,9 +34,7 @@ LlmArgs, LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, NGramDecodingConfig, UserProvidedDecodingConfig, _ModelFormatKind, - _ModelWrapper, _ParallelConfig, - update_llm_args_with_extra_dict, - update_llm_args_with_extra_options) + _ModelWrapper, _ParallelConfig) from .mpi_session import MPINodeState, MpiSession from .tokenizer import TransformersTokenizer, load_hf_tokenizer # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import @@ -974,6 +972,4 @@ class LlmBuildStats: 'CachedModelLoader', 'EagleDecodingConfig', 'Eagle3DecodingConfig', - 'update_llm_args_with_extra_dict', - 'update_llm_args_with_extra_options', ] diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py index 3ff85dd42bef..4398d802b9b7 100644 --- a/tensorrt_llm/llmapi/mm_encoder.py +++ b/tensorrt_llm/llmapi/mm_encoder.py @@ -4,6 +4,7 @@ from tqdm import tqdm from tensorrt_llm._utils import nvtx_range_debug +from tensorrt_llm.executor.executor import GenerationExecutor from tensorrt_llm.inputs import create_input_processor, prompt_inputs from tensorrt_llm.inputs.data import PromptInputs from tensorrt_llm.sampling_params import SamplingParams @@ -60,7 +61,7 @@ def _build_model(self): assert isinstance(self.args, TorchLlmArgs) self.args.mm_encoder_only = True - self._executor = self._executor_cls.create( + self._executor = GenerationExecutor.create( self._engine_dir, executor_config=None, model_world_size=self.args.parallel_config.world_size, diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index a650e76122d7..74d3e02ee7e7 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -12,25 +12,25 @@ # 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 enum import IntEnum from typing import List import torch +from strenum import StrEnum from torch.distributed import ProcessGroup from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl from tensorrt_llm._utils import mpi_disabled -class CpType(IntEnum): +class CpType(StrEnum): # CP type for ulysses parallelism - ULYSSES = 0 + ULYSSES = "ULYSSES" # CP type for star attention - STAR = 1 + STAR = "STAR" # CP type for ring attention - RING = 2 + RING = "RING" # CP type for helix parallelism - HELIX = 3 + HELIX = "HELIX" class MappingBase: @@ -69,11 +69,13 @@ def __init__( # Convert cp_type to CpType enum if it is a string. if cp_config is not None: if "cp_type" in cp_config and isinstance(cp_config["cp_type"], str): - try: - cp_config["cp_type"] = CpType[cp_config["cp_type"].upper()] - except KeyError: - raise ValueError(f"Invalid cp_type: {cp_config['cp_type']}. " \ - f"Must be one of: {', '.join([t.name for t in CpType])}") + cp_type_str = cp_config["cp_type"].upper() + if cp_type_str not in CpType.__members__: + raise ValueError( + f"Invalid cp_type: {cp_config['cp_type']}. " + f"Must be one of: {', '.join([t.name for t in CpType])}" + ) + cp_config["cp_type"] = CpType(cp_type_str) cp_type = cp_config.get("cp_type", CpType.ULYSSES) moe_world_size = tp_size if cp_type == CpType.ULYSSES else tp_size * cp_size diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 4050f8a26ef8..f03e0b89c46b 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -182,8 +182,8 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): with open(gen_server_config_path, "w") as f: yaml.dump(gen_server_config, f) - args = LlmArgs.from_kwargs(model=model_name, - tensor_parallel_size=tensor_parallel_size) + args = LlmArgs(model=model_name, tensor_parallel_size=tensor_parallel_size) + if "FP4" in model_name: args.quant_config.quant_algo = "NVFP4" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml index 920fa0f05342..7fce3cfe6e64 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml @@ -4,10 +4,6 @@ model: DeepSeek-V3-Lite/bf16 backend: "pytorch" context_servers: num_instances: 1 - build_config: - max_batch_size: 10 - max_num_tokens: 512 - max_seq_len: 768 max_batch_size: 10 max_num_tokens: 512 max_seq_len: 768 @@ -27,10 +23,6 @@ context_servers: - "localhost:8001" generation_servers: num_instances: 1 - build_config: - max_batch_size: 1 - max_num_tokens: 2048 - max_seq_len: 2560 tensor_parallel_size: 1 enable_attention_dp: false enable_lm_head_tp_in_adp: false diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_llm_config.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_llm_config.py index ce4812dfe3b3..163d3e658575 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_llm_config.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_llm_config.py @@ -163,7 +163,7 @@ def test_non_registered_model_factory(model_factory: str): ("moe_tensor_parallel_size", 2), ("moe_expert_parallel_size", 2), ("enable_attention_dp", True), - ("cp_config", {"some_key": "some_value"}), + ("cp_config", {"cp_type": "HELIX"}), ], ) def test_parallel_config_validation(parallel_field, invalid_value): diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index fab224c41c42..d57add2fc0bf 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -15,7 +15,7 @@ methods: default: False status: beta cp_config: - annotation: Optional[dict] + annotation: Optional[tensorrt_llm.llmapi.llm_args.CpConfig] default: null status: prototype pp_partition: diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py index 4084481ef2cc..405b7977f160 100644 --- a/tests/unittest/llmapi/test_config_database.py +++ b/tests/unittest/llmapi/test_config_database.py @@ -20,7 +20,7 @@ import pytest import yaml -from tensorrt_llm.llmapi.llm_args import TorchLlmArgs, update_llm_args_with_extra_dict +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs CONFIG_ROOT = Path(__file__).parents[3] / "examples" / "configs" DATABASE_DIR = CONFIG_ROOT / "database" @@ -54,9 +54,7 @@ def test_config_validates_against_llm_args(config_path: Path): with open(config_path) as f: config_dict = yaml.safe_load(f) or {} - base_args = TorchLlmArgs(model="dummy/model", skip_tokenizer_init=True) - merged = update_llm_args_with_extra_dict(base_args.model_dump(), config_dict) - TorchLlmArgs(**merged) + TorchLlmArgs(model="dummy/model", skip_tokenizer_init=True, **config_dict) @pytest.mark.part0 diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index b80f600ce32f..370a2c543692 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -217,18 +217,18 @@ def test_llm_args_invalid_usage(): runtime_max_num_tokens = 2 # Update build_config with warning msg if runtime arguments are passed. - llm_args = LlmArgs.from_kwargs(model='test-model', - max_batch_size=runtime_max_batch_size, - max_num_tokens=runtime_max_num_tokens) + llm_args = LlmArgs(model='test-model', + max_batch_size=runtime_max_batch_size, + max_num_tokens=runtime_max_num_tokens) assert llm_args.build_config.max_batch_size == runtime_max_batch_size assert llm_args.build_config.max_num_tokens == runtime_max_num_tokens # Conflict between build_config and runtime_params build_config = BuildConfig(max_batch_size=5, max_num_tokens=7) - llm_args = LlmArgs.from_kwargs(model='test-model', - build_config=build_config, - max_batch_size=runtime_max_batch_size, - max_num_tokens=runtime_max_num_tokens) + llm_args = LlmArgs(model='test-model', + build_config=build_config, + max_batch_size=runtime_max_batch_size, + max_num_tokens=runtime_max_num_tokens) assert llm_args.build_config.max_batch_size == build_config.max_batch_size assert llm_args.build_config.max_num_tokens == build_config.max_num_tokens diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index a2b086feaa67..8d2e6297292e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -7,12 +7,10 @@ import tensorrt_llm.bindings.executor as tle from tensorrt_llm import LLM as TorchLLM from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.builder import LoraConfig from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, SchedulerConfig) from tensorrt_llm.llmapi.llm_args import * from tensorrt_llm.llmapi.utils import print_traceback_on_error -from tensorrt_llm.plugin import PluginConfig from .test_llm import llama_model_path @@ -54,19 +52,16 @@ def _yaml_to_dict(self, yaml_content: str) -> dict: dict_content = yaml.safe_load(f) return dict_content - def test_update_llm_args_with_extra_dict_with_speculative_config(self): + def test_llm_args_yaml_with_speculative_config(self): yaml_content = """ speculative_config: decoding_type: Lookahead max_window_size: 4 max_ngram_size: 3 """ - dict_content = self._yaml_to_dict(yaml_content) + llm_args_dict = self._yaml_to_dict(yaml_content) - llm_args = TrtLlmArgs(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) - llm_args = TrtLlmArgs(**llm_args_dict) + llm_args = TrtLlmArgs(model=llama_model_path, **llm_args_dict) assert llm_args.speculative_config.max_window_size == 4 assert llm_args.speculative_config.max_ngram_size == 3 assert llm_args.speculative_config.max_verification_set_size == 4 @@ -77,28 +72,21 @@ def test_llm_args_with_invalid_yaml(self): max_num_tokens: 1 max_seq_len: 1 """ - dict_content = self._yaml_to_dict(yaml_content) + llm_args_dict = self._yaml_to_dict(yaml_content) - llm_args = TrtLlmArgs(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) with pytest.raises(ValueError): - llm_args = TrtLlmArgs(**llm_args_dict) + llm_args = TrtLlmArgs(model=llama_model_path, **llm_args_dict) def test_llm_args_with_build_config(self): - # build_config isn't a Pydantic yaml_content = """ build_config: max_beam_width: 4 max_batch_size: 8 max_num_tokens: 256 """ - dict_content = self._yaml_to_dict(yaml_content) + llm_args_dict = self._yaml_to_dict(yaml_content) - llm_args = TrtLlmArgs(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) - llm_args = TrtLlmArgs(**llm_args_dict) + llm_args = TrtLlmArgs(model=llama_model_path, **llm_args_dict) assert llm_args.build_config.max_beam_width == 4 assert llm_args.build_config.max_batch_size == 8 assert llm_args.build_config.max_num_tokens == 256 @@ -110,12 +98,9 @@ def test_llm_args_with_kvcache_config(self): max_tokens: 1024 max_attention_window: [1024, 1024, 1024] """ - dict_content = self._yaml_to_dict(yaml_content) + llm_args_dict = self._yaml_to_dict(yaml_content) - llm_args = TrtLlmArgs(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) - llm_args = TrtLlmArgs(**llm_args_dict) + llm_args = TrtLlmArgs(model=llama_model_path, **llm_args_dict) assert llm_args.kv_cache_config.enable_block_reuse == True assert llm_args.kv_cache_config.max_tokens == 1024 assert llm_args.kv_cache_config.max_attention_window == [ @@ -128,12 +113,9 @@ def test_llm_args_with_pydantic_options(self): max_num_tokens: 256 max_seq_len: 128 """ - dict_content = self._yaml_to_dict(yaml_content) + llm_args_dict = self._yaml_to_dict(yaml_content) - llm_args = TrtLlmArgs(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) - llm_args = TrtLlmArgs(**llm_args_dict) + llm_args = TrtLlmArgs(model=llama_model_path, **llm_args_dict) assert llm_args.max_batch_size == 16 assert llm_args.max_num_tokens == 256 assert llm_args.max_seq_len == 128 @@ -365,54 +347,6 @@ def test_PeftCacheConfig_from_pybind_gets_python_only_default_values_when_none( assert config.lora_prefetch_dir == "." -def test_update_llm_args_with_extra_dict_with_nested_dict(): - llm_api_args_dict = { - "model": - "dummy-model", - "build_config": - None, # Will override later. - "extended_runtime_perf_knob_config": - ExtendedRuntimePerfKnobConfig(multi_block_mode=True), - "kv_cache_config": - KvCacheConfig(enable_block_reuse=False), - "peft_cache_config": - PeftCacheConfig(num_host_module_layer=0), - "scheduler_config": - SchedulerConfig(capacity_scheduler_policy=CapacitySchedulerPolicy. - GUARANTEED_NO_EVICT) - } - plugin_config = PluginConfig(dtype='float16', nccl_plugin=None) - build_config = BuildConfig(max_input_len=1024, - lora_config=LoraConfig(lora_ckpt_source='hf'), - plugin_config=plugin_config) - extra_llm_args_dict = { - "build_config": build_config.model_dump(mode="json"), - } - - llm_api_args_dict = update_llm_args_with_extra_dict(llm_api_args_dict, - extra_llm_args_dict, - "build_config") - initialized_llm_args = TrtLlmArgs(**llm_api_args_dict) - - def check_nested_dict_equality(dict1, dict2, path=""): - if not isinstance(dict1, dict) or not isinstance(dict2, dict): - if dict1 != dict2: - raise ValueError(f"Mismatch at {path}: {dict1} != {dict2}") - return True - if dict1.keys() != dict2.keys(): - raise ValueError(f"Different keys at {path}:") - for key in dict1: - new_path = f"{path}.{key}" if path else key - if not check_nested_dict_equality(dict1[key], dict2[key], new_path): - raise ValueError(f"Mismatch at {path}: {dict1} != {dict2}") - return True - - build_config_dict1 = build_config.model_dump(mode="json") - build_config_dict2 = initialized_llm_args.build_config.model_dump( - mode="json") - check_nested_dict_equality(build_config_dict1, build_config_dict2) - - class TestTorchLlmArgsCudaGraphSettings: def test_cuda_graph_batch_sizes_case_0(self): @@ -568,7 +502,7 @@ def test_to_dict_and_from_dict(self): args = TrtLlmArgs(model=llama_model_path, build_config=build_config) args_dict = args.model_dump() - new_args = TrtLlmArgs.from_kwargs(**args_dict) + new_args = TrtLlmArgs(**args_dict) assert new_args.model_dump() == args_dict @@ -590,36 +524,6 @@ def test_build_config_from_engine(self): assert args.max_num_tokens == 16 assert args.max_batch_size == 4 - def test_model_dump_does_not_mutate_original(self): - """Test that model_dump() and update_llm_args_with_extra_dict don't mutate the original.""" - # Create args with specific build_config values - build_config = BuildConfig( - max_batch_size=8, - max_num_tokens=256, - ) - args = TrtLlmArgs(model=llama_model_path, build_config=build_config) - - # Store original values - original_max_batch_size = args.build_config.max_batch_size - original_max_num_tokens = args.build_config.max_num_tokens - - # Convert to dict and pass through update_llm_args_with_extra_dict with overrides - args_dict = args.model_dump() - extra_dict = { - "max_batch_size": 128, - "max_num_tokens": 1024, - } - updated_dict = update_llm_args_with_extra_dict(args_dict, extra_dict) - - # Verify original args was NOT mutated - assert args.build_config.max_batch_size == original_max_batch_size - assert args.build_config.max_num_tokens == original_max_num_tokens - - # Verify updated dict has new values - new_args = TrtLlmArgs(**updated_dict) - assert new_args.build_config.max_batch_size == 128 - assert new_args.build_config.max_num_tokens == 1024 - class TestStrictBaseModelArbitraryArgs: """Test that StrictBaseModel prevents arbitrary arguments from being accepted.""" @@ -627,9 +531,9 @@ class TestStrictBaseModelArbitraryArgs: def test_cuda_graph_config_arbitrary_args(self): """Test that CudaGraphConfig rejects arbitrary arguments.""" # Valid arguments should work - config = CudaGraphConfig(batch_sizes=[1, 2, 4], max_batch_size=8) + config = CudaGraphConfig(batch_sizes=[1, 2, 4], max_batch_size=4) assert config.batch_sizes == [1, 2, 4] - assert config.max_batch_size == 8 + assert config.max_batch_size == 4 # Arbitrary arguments should be rejected with pytest.raises( @@ -845,3 +749,19 @@ class TestConfig(StrictBaseModel): pydantic_core._pydantic_core.ValidationError) as exc_info: TestConfig(field1="test", field2=100, extra_field="should_fail") assert "extra_field" in str(exc_info.value) + + +def test_executor_config_consistency(): + """Verify that BaseLlmArgs exposes all ExecutorConfig options.""" + # max_beam_width is not included since vague behavior due to lacking the support for dynamic beam width during + # generation + black_list = set(["max_beam_width"]) + executor_config_attrs = set(attr for attr in dir(tle.ExecutorConfig) + if not attr.startswith('_') + and callable(getattr(tle.ExecutorConfig, attr))) + executor_config_attrs -= black_list + llm_args_attr = set(BaseLlmArgs.model_fields.keys()) + # NOTE: When cpp ExecutorConfig add new options, please add the new options into `LlmArgs` with docs as well + # ASK chunweiy for help if you are not sure about the new options. + assert executor_config_attrs.issubset(llm_args_attr), \ + f"New options found in underlying ExecutorConfig: {executor_config_attrs - llm_args_attr}" diff --git a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py index 2594fd47760b..922b970f002c 100755 --- a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py +++ b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py @@ -49,7 +49,6 @@ from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._utils import global_mpi_rank, global_mpi_size from tensorrt_llm.llmapi.llm import RequestOutput -from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict @dataclass @@ -157,6 +156,7 @@ def initialize(self, args): text_output_config["data_type"]) if global_mpi_rank() == 0: # Initialize engine arguments + # TODO self.llm_engine_args = update_llm_args_with_extra_dict( {}, get_model_config(os.environ.get('LLM_CONFIG_PATH', From bb85a34da81c6018929ceae16e1d4cf37bed0fc4 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 16 Jan 2026 12:13:13 -0800 Subject: [PATCH 02/15] Add TODO Signed-off-by: Anish Shanbhag --- tensorrt_llm/llmapi/mm_encoder.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py index 4398d802b9b7..da31d9b6fef8 100644 --- a/tensorrt_llm/llmapi/mm_encoder.py +++ b/tensorrt_llm/llmapi/mm_encoder.py @@ -14,6 +14,7 @@ from .mpi_session import external_mpi_comm_available +# TODO: separate EncoderArgs ? class MultimodalEncoder(_TorchLLM): """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend. """ From 09ab80a27bbc5673293c600356c259407d7611fd Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 29 Jan 2026 22:37:01 -0800 Subject: [PATCH 03/15] Update coding guidelines Signed-off-by: Anish Shanbhag --- CODING_GUIDELINES.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 41487c9e427b..1cfe0e9fdd43 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -410,6 +410,21 @@ Use the [Google style](https://google.github.io/styleguide/pyguide.html), which When defining any user-facing configuration classes (e.g. `LlmArgs` or any class used in its fields), always use Pydantic classes rather than dataclasses or vanilla classes. +- avoid to_dict / from_dict +- avoid dicts / object as any part of a pydantic class +- use discriminated unions +- prefer (1) positiveint/nonnegativeint/nonnegativefloat/positivefloat, or (2) if not possible, gt/ge/le/lt +- use min_length to enforce minimum length of a list +- use pydantic field descriptions instead of comments +- use Literal instead of str when a field should only accept certain values. avoid defining a custom validator on a str field. +- raise ValueError instead of assertion +- co-locate model validation logic within the class itself rather than in a parent class, unless it depends on other fields in the parent class (e.g. don't validate cuda_graph_config in BaseLlmArgs; also add example of OK case) + +TODO: +- ask cursor for more antipatterns fixed in this pr +- unit tests: e.g. for get_llm_args_from_cli_params, other tests you removed, etc +- add validator to BaseLlmArgs.__init_subclass__ to check that all subfields are Pydantic + ##### Attributes and Variables Attributes and variables can be documented inline. Attribute docstrings will be rendered under the docstring for the class. For example: ```python From 796264d4401382d83c9cd288d8d44eb9d5a91220 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 29 Jan 2026 22:45:31 -0800 Subject: [PATCH 04/15] Add TODO Signed-off-by: Anish Shanbhag --- 3rdparty/DeepGEMM | 1 + 3rdparty/NVTX | 1 + 3rdparty/cppzmq | 1 + 3rdparty/cutlass | 1 + 3rdparty/cxxopts | 1 + 3rdparty/json | 1 + 3rdparty/nanobind | 1 + 3rdparty/pybind11 | 1 + 3rdparty/ucxx | 1 + 3rdparty/xgrammar | 1 + CODING_GUIDELINES.md | 1 + 11 files changed, 11 insertions(+) create mode 160000 3rdparty/DeepGEMM create mode 160000 3rdparty/NVTX create mode 160000 3rdparty/cppzmq create mode 160000 3rdparty/cutlass create mode 160000 3rdparty/cxxopts create mode 160000 3rdparty/json create mode 160000 3rdparty/nanobind create mode 160000 3rdparty/pybind11 create mode 160000 3rdparty/ucxx create mode 160000 3rdparty/xgrammar diff --git a/3rdparty/DeepGEMM b/3rdparty/DeepGEMM new file mode 160000 index 000000000000..0315934ce27c --- /dev/null +++ b/3rdparty/DeepGEMM @@ -0,0 +1 @@ +Subproject commit 0315934ce27c5c0b05dfff1a5eb101a5f8872cfe diff --git a/3rdparty/NVTX b/3rdparty/NVTX new file mode 160000 index 000000000000..a1ceb0677f67 --- /dev/null +++ b/3rdparty/NVTX @@ -0,0 +1 @@ +Subproject commit a1ceb0677f67371ed29a2b1c022794f077db5fe7 diff --git a/3rdparty/cppzmq b/3rdparty/cppzmq new file mode 160000 index 000000000000..c94c20743ed7 --- /dev/null +++ b/3rdparty/cppzmq @@ -0,0 +1 @@ +Subproject commit c94c20743ed7d4aa37835a5c46567ab0790d4acc diff --git a/3rdparty/cutlass b/3rdparty/cutlass new file mode 160000 index 000000000000..57e3cfb47a2d --- /dev/null +++ b/3rdparty/cutlass @@ -0,0 +1 @@ +Subproject commit 57e3cfb47a2d9e0d46eb6335c3dc411498efa198 diff --git a/3rdparty/cxxopts b/3rdparty/cxxopts new file mode 160000 index 000000000000..eb787304d67e --- /dev/null +++ b/3rdparty/cxxopts @@ -0,0 +1 @@ +Subproject commit eb787304d67ec22f7c3a184ee8b4c481d04357fd diff --git a/3rdparty/json b/3rdparty/json new file mode 160000 index 000000000000..55f93686c015 --- /dev/null +++ b/3rdparty/json @@ -0,0 +1 @@ +Subproject commit 55f93686c01528224f448c19128836e7df245f72 diff --git a/3rdparty/nanobind b/3rdparty/nanobind new file mode 160000 index 000000000000..a0ed2587f108 --- /dev/null +++ b/3rdparty/nanobind @@ -0,0 +1 @@ +Subproject commit a0ed2587f1089ef7657e2ed49ad6756b01c74e9f diff --git a/3rdparty/pybind11 b/3rdparty/pybind11 new file mode 160000 index 000000000000..f99ffd7e0300 --- /dev/null +++ b/3rdparty/pybind11 @@ -0,0 +1 @@ +Subproject commit f99ffd7e03001810a3e722bf48ad1a9e08415d7d diff --git a/3rdparty/ucxx b/3rdparty/ucxx new file mode 160000 index 000000000000..16eaa57c8d98 --- /dev/null +++ b/3rdparty/ucxx @@ -0,0 +1 @@ +Subproject commit 16eaa57c8d98c8ef54d666a2d2b11e76cfa565f5 diff --git a/3rdparty/xgrammar b/3rdparty/xgrammar new file mode 160000 index 000000000000..e4e816f5f0fe --- /dev/null +++ b/3rdparty/xgrammar @@ -0,0 +1 @@ +Subproject commit e4e816f5f0fe39f5b1601a17a4552307fa3b70ff diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 1cfe0e9fdd43..b618a980c360 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -423,6 +423,7 @@ When defining any user-facing configuration classes (e.g. `LlmArgs` or any class TODO: - ask cursor for more antipatterns fixed in this pr - unit tests: e.g. for get_llm_args_from_cli_params, other tests you removed, etc +- tests for all of the above Pydantic best practices - add validator to BaseLlmArgs.__init_subclass__ to check that all subfields are Pydantic ##### Attributes and Variables From a5c8152f46bb275bc192f2856cc867d16c155625 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 29 Jan 2026 23:16:05 -0800 Subject: [PATCH 05/15] Revert 3rdparty changes Signed-off-by: Anish Shanbhag --- 3rdparty/DeepGEMM | 1 - 3rdparty/NVTX | 1 - 3rdparty/cppzmq | 1 - 3rdparty/cutlass | 1 - 3rdparty/cxxopts | 1 - 3rdparty/json | 1 - 3rdparty/nanobind | 1 - 3rdparty/pybind11 | 1 - 3rdparty/ucxx | 1 - 3rdparty/xgrammar | 1 - 10 files changed, 10 deletions(-) delete mode 160000 3rdparty/DeepGEMM delete mode 160000 3rdparty/NVTX delete mode 160000 3rdparty/cppzmq delete mode 160000 3rdparty/cutlass delete mode 160000 3rdparty/cxxopts delete mode 160000 3rdparty/json delete mode 160000 3rdparty/nanobind delete mode 160000 3rdparty/pybind11 delete mode 160000 3rdparty/ucxx delete mode 160000 3rdparty/xgrammar diff --git a/3rdparty/DeepGEMM b/3rdparty/DeepGEMM deleted file mode 160000 index 0315934ce27c..000000000000 --- a/3rdparty/DeepGEMM +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 0315934ce27c5c0b05dfff1a5eb101a5f8872cfe diff --git a/3rdparty/NVTX b/3rdparty/NVTX deleted file mode 160000 index a1ceb0677f67..000000000000 --- a/3rdparty/NVTX +++ /dev/null @@ -1 +0,0 @@ -Subproject commit a1ceb0677f67371ed29a2b1c022794f077db5fe7 diff --git a/3rdparty/cppzmq b/3rdparty/cppzmq deleted file mode 160000 index c94c20743ed7..000000000000 --- a/3rdparty/cppzmq +++ /dev/null @@ -1 +0,0 @@ -Subproject commit c94c20743ed7d4aa37835a5c46567ab0790d4acc diff --git a/3rdparty/cutlass b/3rdparty/cutlass deleted file mode 160000 index 57e3cfb47a2d..000000000000 --- a/3rdparty/cutlass +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 57e3cfb47a2d9e0d46eb6335c3dc411498efa198 diff --git a/3rdparty/cxxopts b/3rdparty/cxxopts deleted file mode 160000 index eb787304d67e..000000000000 --- a/3rdparty/cxxopts +++ /dev/null @@ -1 +0,0 @@ -Subproject commit eb787304d67ec22f7c3a184ee8b4c481d04357fd diff --git a/3rdparty/json b/3rdparty/json deleted file mode 160000 index 55f93686c015..000000000000 --- a/3rdparty/json +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 55f93686c01528224f448c19128836e7df245f72 diff --git a/3rdparty/nanobind b/3rdparty/nanobind deleted file mode 160000 index a0ed2587f108..000000000000 --- a/3rdparty/nanobind +++ /dev/null @@ -1 +0,0 @@ -Subproject commit a0ed2587f1089ef7657e2ed49ad6756b01c74e9f diff --git a/3rdparty/pybind11 b/3rdparty/pybind11 deleted file mode 160000 index f99ffd7e0300..000000000000 --- a/3rdparty/pybind11 +++ /dev/null @@ -1 +0,0 @@ -Subproject commit f99ffd7e03001810a3e722bf48ad1a9e08415d7d diff --git a/3rdparty/ucxx b/3rdparty/ucxx deleted file mode 160000 index 16eaa57c8d98..000000000000 --- a/3rdparty/ucxx +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 16eaa57c8d98c8ef54d666a2d2b11e76cfa565f5 diff --git a/3rdparty/xgrammar b/3rdparty/xgrammar deleted file mode 160000 index e4e816f5f0fe..000000000000 --- a/3rdparty/xgrammar +++ /dev/null @@ -1 +0,0 @@ -Subproject commit e4e816f5f0fe39f5b1601a17a4552307fa3b70ff From 9dcca9bdfd89e75628fb71b531d00eb4c7dfb746 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Thu, 29 Jan 2026 23:40:06 -0800 Subject: [PATCH 06/15] Fix precommit Signed-off-by: Anish Shanbhag --- tensorrt_llm/commands/serve.py | 5 +++-- tensorrt_llm/llmapi/llm_args.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index a6814ad5c27f..169a2dcd4600 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -150,7 +150,7 @@ def launch_grpc_server(host: str, port: int, llm_args: BaseLlmArgs): async def serve_grpc_async(): logger.info("Initializing TensorRT-LLM gRPC server...") - + llm = create_llm_from_llm_args(llm_args) logger.info("Model loaded successfully") @@ -158,7 +158,8 @@ async def serve_grpc_async(): request_manager = GrpcRequestManager(llm) # Create servicer - servicer = TrtllmServiceServicer(request_manager, model_path=llm_args.model) + servicer = TrtllmServiceServicer(request_manager, + model_path=llm_args.model) # Create gRPC server server = grpc.aio.server( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 916cde35d854..afe05a4e7eee 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2978,7 +2978,8 @@ def validate_speculative_config(self): "EAGLE (v1/v2) draft checkpoints are incompatible with Eagle3—use an Eagle3 draft model." ) # Convert EagleDecodingConfig to Eagle3DecodingConfig - eagle_data = self.speculative_config.model_dump(exclude={"decoding_type"}) + eagle_data = self.speculative_config.model_dump( + exclude={"decoding_type"}) self.speculative_config = Eagle3DecodingConfig(**eagle_data) if isinstance(self.speculative_config, From ea6f07af4ad04adaa42a4ac74b4f5172520a3336 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 11:05:53 -0800 Subject: [PATCH 07/15] Fixes Signed-off-by: Anish Shanbhag --- CODING_GUIDELINES.md | 2 + tensorrt_llm/_torch/auto_deploy/llm_args.py | 20 ++++++-- tensorrt_llm/bench/build/build.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 51 +++++++++++-------- ...sagg_config_ctxtp2_gentp1_trt_backend.yaml | 2 +- .../disagg_config_gen_only_trt_backend.yaml | 2 +- .../disagg_config_trt_backend.yaml | 2 +- .../defs/stress_test/stress_test.py | 4 +- tests/integration/defs/test_e2e.py | 10 ++-- .../integration/test_lists/test-db/l0_a10.yml | 10 ++-- tests/integration/test_lists/waives.txt | 4 +- .../unittest/llmapi/apps/_test_openai_chat.py | 8 +-- .../llmapi/apps/_test_openai_completions.py | 12 ++--- .../unittest/llmapi/apps/_test_openai_misc.py | 4 +- .../llmapi/apps/_test_openai_multi_chat.py | 2 +- .../llmapi/apps/_test_openai_multi_gpu.py | 2 +- .../llmapi/apps/_test_openai_reasoning.py | 4 +- .../apps/_test_trtllm_serve_top_logprobs.py | 4 +- tests/unittest/llmapi/test_llm.py | 9 ++-- tests/unittest/llmapi/test_llm_args.py | 24 ++++----- tests/unittest/others/test_tracing.py | 2 +- 21 files changed, 100 insertions(+), 80 deletions(-) diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index b618a980c360..00129ee8bc12 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -419,11 +419,13 @@ When defining any user-facing configuration classes (e.g. `LlmArgs` or any class - use Literal instead of str when a field should only accept certain values. avoid defining a custom validator on a str field. - raise ValueError instead of assertion - co-locate model validation logic within the class itself rather than in a parent class, unless it depends on other fields in the parent class (e.g. don't validate cuda_graph_config in BaseLlmArgs; also add example of OK case) +- don't define manual validate() methods, prefer using model_validator / field_validator TODO: - ask cursor for more antipatterns fixed in this pr - unit tests: e.g. for get_llm_args_from_cli_params, other tests you removed, etc - tests for all of the above Pydantic best practices +- test to validate that all llmargs classes have field descriptions for all fields - add validator to BaseLlmArgs.__init_subclass__ to check that all subfields are Pydantic ##### Attributes and Variables diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index c44e9a939205..ed0e1db08802 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -1,3 +1,4 @@ +import functools from importlib.resources import files from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Type, Union @@ -7,6 +8,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict from ...llmapi.llm_args import ( + BaseLlmArgs, BuildConfig, EagleDecodingConfig, SamplerType, @@ -28,16 +30,26 @@ def _get_config_dict() -> SettingsConfigDict: ) +@functools.cache +def _get_dummy_llm_args(llm_args_cls: Type[BaseLlmArgs]) -> BaseLlmArgs: + """Create a cached dummy LlmArgs instance to retrieve post-validation default values.""" + return llm_args_cls(model="dummy") + + def _check_for_default_value_only( - cls: Type[BaseSettings], value: Any, info: ValidationInfo, msg: str + llm_args_cls: Type[BaseSettings], value: Any, info: ValidationInfo, msg: str ) -> Any: - """Check if the value is the default value for the field. + """Raise ValueError if value differs from the field's default. - If the value is not the default value, raise a ValueError. + Compares against both the raw field default and the post-validation default + (from a dummy instance) to handle cases where validators modify defaults. """ + dummy_llm_args = _get_dummy_llm_args(llm_args_cls) field_name = info.field_name assert field_name is not None, "field_name should be set for validated field." - if value != cls.model_fields[field_name].get_default(call_default_factory=True): + field_default = llm_args_cls.model_fields[field_name].get_default(call_default_factory=True) + post_validation_default = getattr(dummy_llm_args, field_name) + if value not in {field_default, post_validation_default}: raise ValueError(msg) return value diff --git a/tensorrt_llm/bench/build/build.py b/tensorrt_llm/bench/build/build.py index 10f21ba86c05..d5729b610372 100644 --- a/tensorrt_llm/bench/build/build.py +++ b/tensorrt_llm/bench/build/build.py @@ -324,7 +324,7 @@ def build_command( # Build the LLM engine with the LLMAPI. llm = LLM(checkpoint_path, - tokenizer, + tokenizer=tokenizer, dtype=model_config.dtype, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index afe05a4e7eee..e447a037c249 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -126,23 +126,23 @@ def validate_cuda_graph_config(self) -> 'CudaGraphConfig': """Validate CUDA graph configuration. Ensures that: - 1. If batch_sizes is provided, max_batch_size must be 0 + 1. If batch_sizes is provided, max_batch_size is derived from it 2. If batch_sizes is not provided, it is generated based on max_batch_size - 3. If both are provided, batch_sizes must match the generated values + 3. User should not set both with contradictory values """ if self.batch_sizes: self.batch_sizes = sorted(self.batch_sizes) - if self.max_batch_size != 0: - if self.batch_sizes != CudaGraphConfig._generate_cuda_graph_batch_sizes( - self.max_batch_size, self.enable_padding): - raise ValueError( - "Please don't set both cuda_graph_config.batch_sizes " - "and cuda_graph_config.max_batch_size.\n" - f"cuda_graph_config.batch_sizes: {self.batch_sizes}, " - f"cuda_graph_config.max_batch_size: {self.max_batch_size}" - ) - else: - self.max_batch_size = max(self.batch_sizes) + derived_max = max(self.batch_sizes) + if self.max_batch_size == 0: + # max_batch_size not set, derive from batch_sizes + self.max_batch_size = derived_max + elif self.max_batch_size != derived_max: + # User provided both with contradictory values + raise ValueError( + "Please don't set both cuda_graph_config.batch_sizes " + "and cuda_graph_config.max_batch_size.\n" + f"cuda_graph_config.batch_sizes: {self.batch_sizes}, " + f"cuda_graph_config.max_batch_size: {self.max_batch_size}") else: max_batch_size = self.max_batch_size or 128 generated_sizes = CudaGraphConfig._generate_cuda_graph_batch_sizes( @@ -1009,7 +1009,7 @@ class SaveHiddenStatesDecodingConfig(DecodingBaseConfig): output_directory: str write_interval: int = 20 file_prefix: str = "data" - eagle3_layers_to_capture: Set[int] = Field(..., min_length=1) + eagle3_layers_to_capture: Optional[Set[int]] = None max_total_draft_tokens: Optional[int] = Field(default=1, init=False) eagle_choices: Optional[List[List[int]]] = Field(default=None, init=False) @@ -1018,10 +1018,17 @@ class SaveHiddenStatesDecodingConfig(DecodingBaseConfig): @model_validator(mode="after") def set_last_hidden_in_save(self): + if self.eagle3_layers_to_capture is None: + self._last_hidden_in_save = False + return self + if len(self.eagle3_layers_to_capture) == 0: + raise ValueError( + "eagle3_layers_to_capture must be non-empty if provided") self._last_hidden_in_save = True if -1 not in self.eagle3_layers_to_capture: self._last_hidden_in_save = False self.eagle3_layers_to_capture.add(-1) + return self @functools.cached_property def spec_dec_mode(self): @@ -2509,14 +2516,6 @@ def validate_speculative_config(self): return self - @model_validator(mode="after") - def set_orchestrator_type_for_mpi_disabled(self): - if mpi_disabled(): - self.orchestrator_type = "ray" - elif self.orchestrator_type == "ray": - os.environ["TLLM_DISABLE_MPI"] = "1" - return self - def _load_config_from_engine(self, engine_dir: Path): engine_config = EngineConfig.from_json_file(engine_dir / "config.json") self._pretrained_config = engine_config.pretrained_config @@ -3110,6 +3109,14 @@ def validate_ray_placement_config(self) -> 'TorchLlmArgs': ) return self + @model_validator(mode='after') + def set_orchestrator_type_for_mpi_disabled(self) -> 'TorchLlmArgs': + if mpi_disabled(): + self.orchestrator_type = "ray" + elif self.orchestrator_type == "ray": + os.environ["TLLM_DISABLE_MPI"] = "1" + return self + def get_executor_config( self, _hf_model_dir: Optional[Path] = None, diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml index fa5dffa518b8..8d73c5e1c6cf 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml @@ -2,7 +2,7 @@ hostname: localhost port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "trt" +backend: "tensorrt" context_servers: num_instances: 1 tensor_parallel_size: 2 diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml index ad706f8bf1f4..65bc5cd9ede8 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml @@ -1,7 +1,7 @@ hostname: localhost port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 -backend: "trt" +backend: "tensorrt" context_servers: num_instances: 0 generation_servers: diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml index 3eb275c87e04..fe6982601204 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml @@ -2,7 +2,7 @@ hostname: localhost port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "trt" +backend: "tensorrt" context_servers: num_instances: 1 tensor_parallel_size: 1 diff --git a/tests/integration/defs/stress_test/stress_test.py b/tests/integration/defs/stress_test/stress_test.py index f5fed7b0ebbb..93386cdd4bc1 100644 --- a/tests/integration/defs/stress_test/stress_test.py +++ b/tests/integration/defs/stress_test/stress_test.py @@ -373,7 +373,7 @@ def is_port_available(port: int, "test_mode", ["stress-test", "stress-stage-alone", "stress-test-with-accuracy"], ids=lambda x: x) -@pytest.mark.parametrize("backend", ["trt", "pytorch"], ids=lambda x: x) +@pytest.mark.parametrize("backend", ["tensorrt", "pytorch"], ids=lambda x: x) @pytest.mark.parametrize("capacity_scheduler_policy", ["GUARANTEED_NO_EVICT", "MAX_UTILIZATION"], ids=lambda x: x) @@ -409,7 +409,7 @@ def test_run_stress_test(config, stress_time_timeout, backend, Args: config: Model configuration for the test (injected by pytest.mark.parametrize) stress_time_timeout: Tuple of (stress_time, stress_timeout) in seconds - backend: Backend to use ("trt" or "pytorch") + backend: Backend to use ("tensorrt" or "pytorch") capacity_scheduler_policy: Scheduler policy ("GUARANTEED_NO_EVICT", "MAX_UTILIZATION") test_mode: Test mode ("stress-test" or "stress-stage-alone") """ diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 918f135c2166..07f72adf34af 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -1579,7 +1579,7 @@ def test_trtllm_serve_lora_example(llm_root, llm_venv): str(test_root / "_test_trtllm_serve_lora.py")]) -@pytest.mark.parametrize("backend", ["pytorch", "trt"]) +@pytest.mark.parametrize("backend", ["pytorch", "tensorrt"]) def test_trtllm_serve_top_logprobs(llm_root, llm_venv, backend: str): example_root = Path(os.path.join(llm_root, "examples", "serve")) test_root = unittest_path() / "llmapi" / "apps" @@ -1593,7 +1593,7 @@ def test_trtllm_serve_top_logprobs(llm_root, llm_venv, backend: str): ]) -@pytest.mark.parametrize("backend", ["pytorch", "trt"]) +@pytest.mark.parametrize("backend", ["pytorch", "tensorrt"]) def test_openai_misc_example(llm_root, llm_venv, backend: str): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ @@ -1614,7 +1614,7 @@ def test_openai_cache_salt(llm_root, llm_venv): str(test_root / "_test_openai_cache_salt.py")]) -@pytest.mark.parametrize("backend", ["pytorch", "trt"]) +@pytest.mark.parametrize("backend", ["pytorch", "tensorrt"]) def test_openai_completions_example(llm_root, llm_venv, backend: str): test_root = unittest_path() / "llmapi" / "apps" filter_expr = f"{backend} and not sampler" @@ -1624,7 +1624,7 @@ def test_openai_completions_example(llm_root, llm_venv, backend: str): ]) -@pytest.mark.parametrize("backend", ["pytorch", "trt"]) +@pytest.mark.parametrize("backend", ["pytorch", "tensorrt"]) def test_openai_chat_example(llm_root, llm_venv, backend: str): test_root = unittest_path() / "llmapi" / "apps" filter_expr = f"{backend} and not sampler" @@ -1634,7 +1634,7 @@ def test_openai_chat_example(llm_root, llm_venv, backend: str): ]) -@pytest.mark.parametrize("backend", ["pytorch", "trt"]) +@pytest.mark.parametrize("backend", ["pytorch", "tensorrt"]) def test_openai_reasoning(llm_root, llm_venv, backend: str): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 4288f8ef18f1..5d0e94be744e 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -155,11 +155,11 @@ l0_a10: - llmapi/test_llm_examples.py::test_llmapi_kv_cache_connector[Qwen2-0.5B] - test_e2e.py::test_openai_health - test_e2e.py::test_trtllm_serve_example - - test_e2e.py::test_trtllm_serve_top_logprobs[trt] - - test_e2e.py::test_openai_misc_example[trt] - - test_e2e.py::test_openai_completions_example[trt] - - test_e2e.py::test_openai_chat_example[trt] - - test_e2e.py::test_openai_reasoning[trt] + - test_e2e.py::test_trtllm_serve_top_logprobs[tensorrt] + - test_e2e.py::test_openai_misc_example[tensorrt] + - test_e2e.py::test_openai_completions_example[tensorrt] + - test_e2e.py::test_openai_chat_example[tensorrt] + - test_e2e.py::test_openai_reasoning[tensorrt] - test_e2e.py::test_trtllm_bench_sanity[--non-streaming-FP16-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_latency_sanity[FP16-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - unittest/trt/quantization diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d4820523e3d3..a3dd979565f2 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -154,7 +154,7 @@ triton_server/test_triton.py::test_python_bls_unit_tests[python-bls-unit-tests] triton_server/test_triton.py::test_mistral_ib[mistral-ib] SKIP (https://nvbugs/5477399) triton_server/test_triton.py::test_eagle[eagle] SKIP (https://nvbugs/5477378) examples/test_mixtral.py::test_llm_mixtral_moe_plugin_lora_4gpus[Mixtral-8x7B-v0.1-chinese-mixtral-lora] SKIP (https://nvbugs/5477421) -test_e2e.py::test_openai_chat_example[trt] SKIP (https://nvbugs/5477444) +test_e2e.py::test_openai_chat_example[tensorrt] SKIP (https://nvbugs/5477444) test_e2e.py::test_trtllm_bench_iteration_log[TRT-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] SKIP (https://nvbugs/5448523) accuracy/test_llm_api_pytorch.py::TestLlama3_2_3B::test_auto_dtype SKIP (https://nvbugs/5520319) examples/test_llama.py::test_llm_llama_1gpu_fp8_kv_cache[llama-v2-7b-hf-bfloat16] SKIP (https://nvbugs/5527940) @@ -196,7 +196,7 @@ unittest/_torch/modules/test_fused_moe.py::test_fused_moe_alltoall[DeepEPLowLate unittest/_torch/auto_deploy/unit/multigpu/test_ad_build_small_multi.py::test_build_ad[meta-llama/Meta-Llama-3.1-8B-Instruct-llm_extra_args0-2] SKIP (https://nvbugs/5680755) full:H100_PCIe/unittest/llmapi/test_llm_pytorch.py::test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache SKIP (https://nvbugs/5682551) unittest/_torch/speculative/test_draft_len_schedule.py::test_correctness_across_batch_sizes[model_drafter-schedule1] SKIP (https://nvbugs/5680911) -test_e2e.py::test_openai_completions_example[trt] SKIP (https://nvbugs/5701450) +test_e2e.py::test_openai_completions_example[tensorrt] SKIP (https://nvbugs/5701450) triton_server/test_triton_llm.py::test_llmapi_backend[4-0-disableDecoupleMode-tensorrt_llm] SKIP (https://nvbugs/5701480) unittest/_torch/modules/tests_lora_modules/test_lora_attention_pytorch_flow_vs_trt.py::TestLoraAttentionPytorchFlowVsTRT::test_lora_attention SKIP (https://nvbugs/5701421) test_e2e.py::test_ptp_quickstart_multimodal_2gpu[mistral-small-3.1-24b-instruct-Mistral-Small-3.1-24B-Instruct-2503] SKIP (https://nvbugs/5648560) diff --git a/tests/unittest/llmapi/apps/_test_openai_chat.py b/tests/unittest/llmapi/apps/_test_openai_chat.py index d35935b1e03c..86083a90360a 100644 --- a/tests/unittest/llmapi/apps/_test_openai_chat.py +++ b/tests/unittest/llmapi/apps/_test_openai_chat.py @@ -22,7 +22,7 @@ def model_name(): return "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param @@ -70,7 +70,7 @@ def server(model_name: str, backend: str, extra_llm_api_options: bool, args = ["--backend", f"{backend}"] args.extend(["--kv_cache_free_gpu_memory_fraction", "0.2"]) # for co-existence with other servers - if backend == "trt": + if backend == "tensorrt": args.extend(["--max_beam_width", "4"]) if extra_llm_api_options: args.extend( @@ -542,7 +542,7 @@ async def test_chat_completion_with_invalid_logit_bias( def test_chat_cached_tokens(client: openai.OpenAI, model_name: str, backend: str, extra_llm_api_options: bool): - if backend == "trt": + if backend == "tensorrt": pytest.skip("Cached tokens is not supported in trt backend yet") messages = [{ @@ -583,7 +583,7 @@ def test_chat_cached_tokens(client: openai.OpenAI, model_name: str, async def test_chat_cached_tokens_stream(async_client: openai.AsyncOpenAI, model_name: str, backend: str, extra_llm_api_options: bool): - if backend == "trt": + if backend == "tensorrt": pytest.skip("Cached tokens is not supported in trt backend yet") messages = [{ diff --git a/tests/unittest/llmapi/apps/_test_openai_completions.py b/tests/unittest/llmapi/apps/_test_openai_completions.py index 03f08a5c6b8f..a8e8cc4b0e1f 100644 --- a/tests/unittest/llmapi/apps/_test_openai_completions.py +++ b/tests/unittest/llmapi/apps/_test_openai_completions.py @@ -19,7 +19,7 @@ def model_name(): return "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param @@ -56,7 +56,7 @@ def server(model_name: str, backend: str, num_postprocess_workers: int, args.extend(["--kv_cache_free_gpu_memory_fraction", "0.2"]) # for co-existence with other servers args.extend(["--num_postprocess_workers", f"{num_postprocess_workers}"]) - if backend == "trt": + if backend == "tensorrt": args.extend( ["--extra_llm_api_options", temp_extra_llm_api_options_file]) with RemoteOpenAIServer(model_path, args) as remote_server: @@ -459,7 +459,7 @@ async def test_completion_with_invalid_logit_bias( def test_completion_logprobs(client: openai.OpenAI, model_name: str, backend: str, num_postprocess_workers: int): """Test completion with logprobs enabled (non-streaming).""" - if backend == "trt" and num_postprocess_workers > 0: + if backend == "tensorrt" and num_postprocess_workers > 0: pytest.skip("Logprobs is not supported in TRT processpool mode") prompt = "Hello, my name is" @@ -505,7 +505,7 @@ async def test_completion_logprobs_streaming(async_client: openai.AsyncOpenAI, backend: str, model_name: str, num_postprocess_workers: int): """Test completion with logprobs enabled (streaming).""" - if backend == "trt" and num_postprocess_workers > 0: + if backend == "tensorrt" and num_postprocess_workers > 0: pytest.skip("Logprobs is not supported in TRT processpool mode") prompt = "Hello, my name is" @@ -559,7 +559,7 @@ async def test_completion_logprobs_streaming(async_client: openai.AsyncOpenAI, def test_completion_cached_tokens(client: openai.OpenAI, model_name: str, backend: str): - if backend == "trt": + if backend == "tensorrt": pytest.skip("Cached tokens is not supported in trt backend yet") prompt = "This is a test prompt" @@ -588,7 +588,7 @@ def test_completion_cached_tokens(client: openai.OpenAI, model_name: str, @pytest.mark.asyncio(loop_scope="module") async def test_completion_cached_tokens_stream(async_client: openai.AsyncOpenAI, model_name: str, backend: str): - if backend == "trt": + if backend == "tensorrt": pytest.skip("Cached tokens is not supported in trt backend yet") prompt = "This is a test prompt" diff --git a/tests/unittest/llmapi/apps/_test_openai_misc.py b/tests/unittest/llmapi/apps/_test_openai_misc.py index 9e1b1a8dbe9e..05c998de2ce7 100644 --- a/tests/unittest/llmapi/apps/_test_openai_misc.py +++ b/tests/unittest/llmapi/apps/_test_openai_misc.py @@ -10,7 +10,7 @@ from .openai_server import RemoteOpenAIServer -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param @@ -20,7 +20,7 @@ def model_name(backend): # Note: TRT backend does not support Qwen3-0.6B-Base, # and PyTorch backend does not support going over the limit of "max_position_embeddings" tokens # of TinyLlama. - if backend == "trt": + if backend == "tensorrt": return "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" else: return "Qwen3/Qwen3-0.6B-Base" diff --git a/tests/unittest/llmapi/apps/_test_openai_multi_chat.py b/tests/unittest/llmapi/apps/_test_openai_multi_chat.py index a628659ebe4c..4af1b9e7a80b 100644 --- a/tests/unittest/llmapi/apps/_test_openai_multi_chat.py +++ b/tests/unittest/llmapi/apps/_test_openai_multi_chat.py @@ -65,7 +65,7 @@ def engine_from_fp8_quantization(model_name): def server(model_name: str, engine_from_fp8_quantization: str): model_path = get_model_path(model_name) args = [ - "--tp_size", "2", "--tokenizer", model_path, "--backend", "trt", + "--tp_size", "2", "--tokenizer", model_path, "--backend", "tensorrt", "--max_num_tokens", "20480", "--max_batch_size", "128" ] with RemoteOpenAIServer(engine_from_fp8_quantization, diff --git a/tests/unittest/llmapi/apps/_test_openai_multi_gpu.py b/tests/unittest/llmapi/apps/_test_openai_multi_gpu.py index 6ac65c42b25e..ffb09fe6c249 100644 --- a/tests/unittest/llmapi/apps/_test_openai_multi_gpu.py +++ b/tests/unittest/llmapi/apps/_test_openai_multi_gpu.py @@ -15,7 +15,7 @@ def model_name(): return "llama-models-v3/llama-v3-8b-instruct-hf" -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param diff --git a/tests/unittest/llmapi/apps/_test_openai_reasoning.py b/tests/unittest/llmapi/apps/_test_openai_reasoning.py index 455108fc388c..421602d04a5e 100644 --- a/tests/unittest/llmapi/apps/_test_openai_reasoning.py +++ b/tests/unittest/llmapi/apps/_test_openai_reasoning.py @@ -18,7 +18,7 @@ def model_name(request) -> str: # yapf: enable -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param @@ -26,7 +26,7 @@ def backend(request): @pytest.fixture(scope="module") def server(model_name: str, backend: str): # Skip specific model/backend combinations - if model_name == "Qwen3/Qwen3-0.6B" and backend == "trt": + if model_name == "Qwen3/Qwen3-0.6B" and backend == "tensorrt": pytest.skip("Qwen3 model not supported with trt backend") model_path = get_model_path(model_name) diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py index aeb7be4a740c..a03bed239f23 100644 --- a/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py +++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py @@ -13,7 +13,7 @@ def model_name(): return "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" -@pytest.fixture(scope="module", params=["trt", "pytorch"]) +@pytest.fixture(scope="module", params=["tensorrt", "pytorch"]) def backend(request): return request.param @@ -39,7 +39,7 @@ def temp_extra_llm_api_options_file(tmp_path_factory): def server(model_name: str, backend: str, temp_extra_llm_api_options_file: str): model_path = get_model_path(model_name) args = ["--backend", f"{backend}"] - if backend == "trt": + if backend == "tensorrt": args += ["--extra_llm_api_options", temp_extra_llm_api_options_file] with RemoteOpenAIServer(model_path, args) as remote_server: yield remote_server diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 370a2c543692..a371b4e3d724 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -1894,11 +1894,12 @@ async def main(): "prompt_logprobs, logprobs, return_context_logits, return_generation_logits, backend", [ # TRT backend test cases - (2, None, True, False, "trt"), # prompt_logprobs with context_logits - (None, 2, False, False, "trt"), # generation logprobs only (top-2) + (2, None, True, False, "tensorrt" + ), # prompt_logprobs with context_logits + (None, 2, False, False, "tensorrt"), # generation logprobs only (top-2) (2, None, False, False, - "trt"), # prompt_logprobs without context_logits - (None, None, False, False, "trt"), # no logprobs at all + "tensorrt"), # prompt_logprobs without context_logits + (None, None, False, False, "tensorrt"), # no logprobs at all ]) def test_llm_return_logprobs(prompt_logprobs: Optional[int], logprobs: Optional[int], diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 8d2e6297292e..99114498ddee 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3,6 +3,7 @@ import pydantic_core import pytest import yaml +from pydantic import TypeAdapter import tensorrt_llm.bindings.executor as tle from tensorrt_llm import LLM as TorchLLM @@ -25,7 +26,7 @@ def test_LookaheadDecodingConfig(): assert config.max_verification_set_size == 4 # from dict - config = LookaheadDecodingConfig.from_dict({ + config = LookaheadDecodingConfig(**{ "max_window_size": 4, "max_ngram_size": 3, "max_verification_set_size": 4 @@ -135,10 +136,11 @@ def test_llm_args_with_model_kwargs(self, llm_args_cls): def test_decoding_type_eagle3_parses_to_eagle3_decoding_config(): - spec_cfg = DecodingBaseConfig.from_dict( + adapter = TypeAdapter(SpeculativeConfig) + spec_cfg = adapter.validate_python( dict(decoding_type="Eagle3", max_draft_len=3, - speculative_model_dir="/path/to/draft/model")) + speculative_model="/path/to/draft/model")) assert isinstance(spec_cfg, Eagle3DecodingConfig) @@ -152,12 +154,9 @@ def _capture_warning(msg, *args, **kwargs): monkeypatch.setattr(llm_args_mod.logger, "warning", _capture_warning) - spec_cfg = DecodingBaseConfig.from_dict(dict( - decoding_type="Eagle", - max_draft_len=3, - speculative_model_dir="/path/to/draft/model"), - backend="pytorch") - assert isinstance(spec_cfg, Eagle3DecodingConfig) + spec_cfg = EagleDecodingConfig(decoding_type="Eagle", + max_draft_len=3, + speculative_model="/path/to/draft/model") TorchLlmArgs(model=llama_model_path, speculative_config=spec_cfg) @@ -167,10 +166,9 @@ def _capture_warning(msg, *args, **kwargs): def test_decoding_type_eagle3_errors_on_tensorrt_backend(): - spec_cfg = DecodingBaseConfig.from_dict( - dict(decoding_type="Eagle3", - max_draft_len=3, - speculative_model_dir="/path/to/draft/model")) + spec_cfg = Eagle3DecodingConfig(decoding_type="Eagle3", + max_draft_len=3, + speculative_model="/path/to/draft/model") with pytest.raises(ValueError, match="only supported on the PyTorch backend"): TrtLlmArgs(model=llama_model_path, speculative_config=spec_cfg) diff --git a/tests/unittest/others/test_tracing.py b/tests/unittest/others/test_tracing.py index 01da3716d30a..a5f929c6177c 100644 --- a/tests/unittest/others/test_tracing.py +++ b/tests/unittest/others/test_tracing.py @@ -105,7 +105,7 @@ def server( ): model_path = get_model_path(model_name) args = ["--backend", f"{backend}"] - if backend == "trt": + if backend == "tensorrt": args.extend(["--max_beam_width", "4"]) args.extend(["--extra_llm_api_options", temp_extra_llm_api_options_file]) args.extend(["--num_postprocess_workers", f"{num_postprocess_workers}"]) From 2d2eeced7a9ef5b71d435bc5c007b705be404046 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 12:09:20 -0800 Subject: [PATCH 08/15] Fixes Signed-off-by: Anish Shanbhag --- tensorrt_llm/llmapi/llm_args.py | 45 ++++++++++--------- tensorrt_llm/llmapi/llm_create.py | 3 ++ tensorrt_llm/llmapi/llm_utils.py | 2 +- .../defs/accuracy/test_cli_flow.py | 12 +++-- .../references/calib_config.yaml | 9 ---- .../api_stability/references/llm.yaml | 4 -- .../references_committed/llm.yaml | 11 ++--- tests/unittest/llmapi/test_llm.py | 4 +- .../all_models/llmapi/tensorrt_llm/1/model.py | 12 ++--- 9 files changed, 43 insertions(+), 59 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e447a037c249..5865958e17bb 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1985,9 +1985,6 @@ class BaseLlmArgs(StrictBaseModel): lora_config: Optional[LoraConfig] = Field( default=None, description="LoRA configuration for the model.") - quant_config: QuantConfig = Field(default_factory=QuantConfig, - description="Quantization config.") - # Several options from ExecutorConfig, expanded here for less hierarchy kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") @@ -2051,8 +2048,8 @@ class BaseLlmArgs(StrictBaseModel): max_seq_len: Optional[int] = Field( default=None, description="The maximum sequence length.") - max_beam_width: int = Field(default=1, - description="The maximum beam width.") + max_beam_width: Optional[int] = Field(default=1, + description="The maximum beam width.") max_num_tokens: Optional[int] = Field( default=8192, description="The maximum number of tokens.") @@ -2178,6 +2175,24 @@ def validate_gpus_per_node(cls, v, info): v = torch.cuda.device_count() return v + @model_validator(mode="after") + def normalize_optional_fields_to_defaults(self): + """Normalize certain fields to their declared default values in case a user explicitly sets them to None. + + This is necessary because downstream code expects these fields to be non-None. + """ + for field_name in ( + "max_batch_size", + "max_input_len", + "max_beam_width", + "max_num_tokens", + ): + if getattr(self, field_name) is None: + field_info = self.__class__.model_fields.get(field_name) + if field_info is not None and field_info.default is not None: + setattr(self, field_name, field_info.default) + return self + @model_validator(mode="after") def validate_parallel_config(self): if self.moe_cluster_parallel_size is None: @@ -2335,6 +2350,9 @@ class TrtLlmArgs(BaseLlmArgs): ExtendedRuntimePerfKnobConfig] = Field( default=None, description="Extended runtime perf knob config.") + quant_config: QuantConfig = Field(default_factory=QuantConfig, + description="Quantization config.") + calib_config: CalibConfig = Field(default_factory=CalibConfig, description="Calibration config.") @@ -3038,23 +3056,6 @@ def validate_load_balancer(self) -> 'TorchLlmArgs': ) from e return self - @model_validator(mode='after') - def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': - if self.kv_cache_config is None: - return self - - assert self.quant_config is not None - if self.kv_cache_config.dtype == "auto": - return self - elif self.kv_cache_config.dtype == 'fp8': - self.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - else: - logger.warning( - f"Cannot sync quant_config.kv_cache_quant_algo with kv_cache_config.dtype of {self.kv_cache_config.dtype}, " - "please update the validator") - - return self - @model_validator(mode='after') def validate_helix_tokens_per_block(self) -> 'TorchLlmArgs': """Validate that cp_config.tokens_per_block matches kv_cache_config.tokens_per_block when HELIX parallelism is active.""" diff --git a/tensorrt_llm/llmapi/llm_create.py b/tensorrt_llm/llmapi/llm_create.py index b23625da395e..6bc5d916f59f 100644 --- a/tensorrt_llm/llmapi/llm_create.py +++ b/tensorrt_llm/llmapi/llm_create.py @@ -65,6 +65,9 @@ def get_llm_args_from_cli_params(model: str, **params: dict[str, Any]) -> BaseLl else: raise ValueError(f"Invalid backend: {backend}") + # Remove model from llm_args_dict if present, since we pass it explicitly as the first argument + del llm_args_dict["model"] + logger.info(f"Creating LLM args from dict: {llm_args_dict}") return llm_args_cls(model=model, **llm_args_dict) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index acbc68c756dd..9012f392fecb 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -498,7 +498,7 @@ def _load_model_from_hf(self): dtype=self.llm_args.dtype, mapping=self.mapping, quant_config=self.llm_args.quant_config, - **self.llm_args.calib_config.to_dict(), + **self.llm_args.calib_config.model_dump(), trust_remote_code=self.llm_args.trust_remote_code, ) if self.llm_args.parallel_config.is_multi_gpu: diff --git a/tests/integration/defs/accuracy/test_cli_flow.py b/tests/integration/defs/accuracy/test_cli_flow.py index 20258e88baa7..8b173d49fb81 100644 --- a/tests/integration/defs/accuracy/test_cli_flow.py +++ b/tests/integration/defs/accuracy/test_cli_flow.py @@ -14,8 +14,6 @@ # limitations under the License. import pytest -from tensorrt_llm.llmapi import (EagleDecodingConfig, LookaheadDecodingConfig, - MedusaDecodingConfig) from tensorrt_llm.quantization import QuantAlgo from ..conftest import (get_sm_version, llm_models_root, parametrize_with_ids, @@ -428,7 +426,7 @@ class TestVicuna7B(CliFlowAccuracyTestHarness): def test_lookahead(self, mocker): mocker.patch.object(CnnDailymail, "MAX_BATCH_SIZE", 8) - self.run(spec_dec_algo=LookaheadDecodingConfig.decoding_type, + self.run(spec_dec_algo="Lookahead", extra_build_args=[ "--max_draft_len=83", "--speculative_decoding_mode=lookahead_decoding" @@ -448,7 +446,7 @@ def test_medusa(self, cuda_graph, mocker): extra_summarize_args.append("--cuda_graph_mode") self.run(dtype="float16", - spec_dec_algo=MedusaDecodingConfig.decoding_type, + spec_dec_algo="Medusa", extra_convert_args=[ f"--medusa_model_dir={self.MEDUSA_MODEL_PATH}", "--num_medusa_heads=4" @@ -476,7 +474,7 @@ def test_eagle(self, cuda_graph, chunked_context, typical_acceptance, extra_summarize_args.extend( ["--eagle_posterior_threshold=0.09", "--temperature=0.7"]) - self.run(spec_dec_algo=EagleDecodingConfig.decoding_type, + self.run(spec_dec_algo="Eagle", extra_convert_args=[ f"--eagle_model_dir={self.EAGLE_MODEL_PATH}", "--max_draft_len=63", "--num_eagle_layers=4", @@ -503,7 +501,7 @@ def test_eagle_2(self, cuda_graph, chunked_context, mocker): if chunked_context: extra_summarize_args.append("--enable_chunked_context") - self.run(spec_dec_algo=EagleDecodingConfig.decoding_type, + self.run(spec_dec_algo="Eagle", extra_convert_args=[ f"--eagle_model_dir={self.EAGLE_MODEL_PATH}", "--max_draft_len=63", "--num_eagle_layers=4", @@ -848,7 +846,7 @@ def test_medusa_fp8_prequantized(self, mocker): "--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" ] self.run(dtype="float16", - spec_dec_algo=MedusaDecodingConfig.decoding_type, + spec_dec_algo="Medusa", extra_build_args=["--speculative_decoding_mode=medusa"], extra_summarize_args=extra_summarize_args) diff --git a/tests/unittest/api_stability/references/calib_config.yaml b/tests/unittest/api_stability/references/calib_config.yaml index 37f6ff1bee6a..833a2ecbc489 100644 --- a/tests/unittest/api_stability/references/calib_config.yaml +++ b/tests/unittest/api_stability/references/calib_config.yaml @@ -23,13 +23,4 @@ methods: annotation: int default: 2048 return_annotation: None - from_dict: - parameters: - config: - annotation: dict - default: inspect._empty - return_annotation: tensorrt_llm.llmapi.llm_utils.CalibConfig - to_dict: - parameters: {} - return_annotation: dict properties: {} diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d57add2fc0bf..d7c07cc088a2 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -83,10 +83,6 @@ methods: annotation: Optional[str] default: null status: deprecated - build_config: - annotation: Optional[tensorrt_llm.llmapi.llm_args.BuildConfig] - default: null - status: deprecated cuda_graph_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.CudaGraphConfig] default: null diff --git a/tests/unittest/api_stability/references_committed/llm.yaml b/tests/unittest/api_stability/references_committed/llm.yaml index a8662702dc6c..7f2cf1265aad 100644 --- a/tests/unittest/api_stability/references_committed/llm.yaml +++ b/tests/unittest/api_stability/references_committed/llm.yaml @@ -59,21 +59,21 @@ methods: default: null # Speculative decoding speculative_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.DraftTargetDecodingConfig, tensorrt_llm.llmapi.llm_args.EagleDecodingConfig,tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_args.MedusaDecodingConfig, tensorrt_llm.llmapi.llm_args.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, tensorrt_llm.llmapi.llm_args.UserProvidedDecodingConfig, tensorrt_llm.llmapi.llm_args.AutoDecodingConfig, tensorrt_llm.llmapi.llm_args.SaveHiddenStatesDecodingConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.DraftTargetDecodingConfig, tensorrt_llm.llmapi.llm_args.EagleDecodingConfig, tensorrt_llm.llmapi.llm_args.Eagle3DecodingConfig, tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_args.MedusaDecodingConfig, tensorrt_llm.llmapi.llm_args.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, tensorrt_llm.llmapi.llm_args.UserProvidedDecodingConfig, tensorrt_llm.llmapi.llm_args.AutoDecodingConfig, tensorrt_llm.llmapi.llm_args.SaveHiddenStatesDecodingConfig, NoneType] default: null # generation constraints max_batch_size: annotation: Optional[int] - default: null + default: 2048 max_input_len: annotation: Optional[int] - default: null + default: 1024 max_seq_len: annotation: Optional[int] default: null max_beam_width: annotation: Optional[int] - default: null + default: 1 max_num_tokens: annotation: Optional[int] default: 8192 @@ -81,9 +81,6 @@ methods: load_format: annotation: Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat] default: 0 - enable_tqdm: - annotation: bool - default: false enable_chunked_prefill: annotation: bool default: false diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index a371b4e3d724..32e2789610c8 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -131,7 +131,9 @@ def llm_test_harness(model_dir: str, llm_cls = LLM_torch if backend == "pytorch" else LLM - with assert_resource_freed(llm_cls, model_dir, tokenizer, + with assert_resource_freed(llm_cls, + model_dir, + tokenizer=tokenizer, **llm_kwargs) as llm: outputs = llm.generate(inputs, sampling_params=sampling_params) print(outputs) diff --git a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py index 922b970f002c..b74fe6d1dedd 100755 --- a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py +++ b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py @@ -155,14 +155,10 @@ def initialize(self, args): self.output_dtype = pb_utils.triton_string_to_numpy( text_output_config["data_type"]) if global_mpi_rank() == 0: - # Initialize engine arguments - # TODO - self.llm_engine_args = update_llm_args_with_extra_dict( - {}, - get_model_config(os.environ.get('LLM_CONFIG_PATH', - 'model.yaml'), - exclude_keys=["triton_config"]), - ) + # Initialize engine arguments from model.yaml (excluding triton_config) + self.llm_engine_args = get_model_config( + os.environ.get('LLM_CONFIG_PATH', 'model.yaml'), + exclude_keys=["triton_config"]) self.logger.log_info( f"[trtllm] rank{global_mpi_rank()} is starting trtllm engine with args: {self.llm_engine_args}" ) From 321d89c8e875fb746384a7a251304753318f0c8c Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 12:30:10 -0800 Subject: [PATCH 09/15] Fixes Signed-off-by: Anish Shanbhag --- examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py | 2 +- tensorrt_llm/bench/benchmark/utils/general.py | 1 + tensorrt_llm/llmapi/llm_create.py | 2 +- tests/unittest/llmapi/test_llm_args.py | 7 ++----- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py b/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py index 1738242267d9..1b26dc879be9 100644 --- a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py +++ b/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py @@ -56,7 +56,7 @@ def __init__( free_gpu_memory_fraction: float = 0.9, trust_remote_code: bool = False, use_cuda_graph: bool = True, - backend: str = 'trt', + backend: str = 'tensorrt', max_context_length: Optional[int] = None, moe_expert_parallel_size: Optional[int] = None, moe_backend: Optional[str] = "TRTLLM", diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index cdbd17cac592..21f08cbb82a9 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -124,6 +124,7 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, tp_size = llm_args.tensor_parallel_size pp_size = llm_args.pipeline_parallel_size + kv_cache_dtype = llm_args.kv_cache_config.dtype if llm_args.max_seq_len is None: llm_args.max_seq_len = dataset_metadata.max_sequence_length diff --git a/tensorrt_llm/llmapi/llm_create.py b/tensorrt_llm/llmapi/llm_create.py index 6bc5d916f59f..9b58dea5fcc1 100644 --- a/tensorrt_llm/llmapi/llm_create.py +++ b/tensorrt_llm/llmapi/llm_create.py @@ -66,7 +66,7 @@ def get_llm_args_from_cli_params(model: str, **params: dict[str, Any]) -> BaseLl raise ValueError(f"Invalid backend: {backend}") # Remove model from llm_args_dict if present, since we pass it explicitly as the first argument - del llm_args_dict["model"] + llm_args_dict.pop("model", None) logger.info(f"Creating LLM args from dict: {llm_args_dict}") return llm_args_cls(model=model, **llm_args_dict) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 99114498ddee..3aaa2ab19922 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -127,11 +127,8 @@ def test_llm_args_with_model_kwargs(self, llm_args_cls): model_kwargs: num_hidden_layers: 2 """ - dict_content = self._yaml_to_dict(yaml_content) - llm_args = llm_args_cls(model=llama_model_path) - llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(), - dict_content) - llm_args = llm_args_cls(**llm_args_dict) + llm_args_dict = self._yaml_to_dict(yaml_content) + llm_args = llm_args_cls(model=llama_model_path, **llm_args_dict) assert llm_args.model_kwargs['num_hidden_layers'] == 2 From 8de9b4c1d23064ecc856ecdb3f655ac0ce86ffab Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 13:05:43 -0800 Subject: [PATCH 10/15] Fixes Signed-off-by: Anish Shanbhag --- .../_torch/speculative/save_hidden_state.py | 2 +- tensorrt_llm/bench/benchmark/utils/general.py | 48 +++++++++++-------- .../bench/dataclasses/configuration.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 13 +++-- 4 files changed, 36 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/save_hidden_state.py b/tensorrt_llm/_torch/speculative/save_hidden_state.py index 936ca5d2cae8..e7a6b16e63e4 100644 --- a/tensorrt_llm/_torch/speculative/save_hidden_state.py +++ b/tensorrt_llm/_torch/speculative/save_hidden_state.py @@ -47,7 +47,7 @@ def _process_request(self, request: LlmRequest, resource_manager) -> None: "input_ids": input_ids, "hidden_state": hidden_states, } - if len(self.spec_config.eagle3_layers_to_capture) > 1: + if self.spec_config.num_capture_layers > 1: if self.spec_config._last_hidden_in_save: out_dict[ "aux_hidden_states"] = resource_manager.hidden_states[:num_tokens, :].cpu( diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index 21f08cbb82a9..49956820b07b 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -91,7 +91,8 @@ def get_settings_from_engine( return runtime_config, llm_args.build_config.model_dump() -def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, +def get_settings(params: dict, dataset_metadata: DatasetMetadata, + options: GeneralExecSettings, model: str, model_path: Union[Path, None]) -> RuntimeConfig: """Retrieve basic runtime config for pytorch backend path @@ -105,22 +106,27 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, mamba_ssm_cache_dtype = params.get("mamba_ssm_cache_dtype", "auto") # TODO: unify CLI parameter naming across trtllm-serve / trtllm-bench to avoid the need to manually # specify parameters here - llm_args = get_llm_args_from_cli_params( - model=model, - backend=params.get("backend"), - extra_llm_api_options=params.get("extra_llm_api_options"), - max_batch_size=params.get("max_batch_size"), - max_num_tokens=params.get("max_num_tokens"), - max_seq_len=params.get("max_seq_len"), - max_input_len=params.get("max_input_len"), - max_beam_width=params.get("beam_width"), - tensor_parallel_size=params.get("tp"), - pipeline_parallel_size=params.get("pp"), - moe_expert_parallel_size=params.get("ep"), - moe_cluster_parallel_size=params.get("cluster_size"), - enable_chunked_prefill=params.get("enable_chunked_context"), - skip_tokenizer_init=not params.get("no_skip_tokenizer_init", False), - ) + # Note: max_batch_size and max_num_tokens are only passed if explicitly set (not None), so that + # model_fields_set correctly reflects user intent. If not in model_fields_set, the benchmark code + # below will compute optimal values via heuristics. + llm_args_kwargs = { + "backend": params.get("backend"), + "extra_llm_api_options": params.get("extra_llm_api_options"), + "max_seq_len": params.get("max_seq_len"), + "max_input_len": params.get("max_input_len"), + "max_beam_width": params.get("beam_width"), + "tensor_parallel_size": params.get("tp"), + "pipeline_parallel_size": params.get("pp"), + "moe_expert_parallel_size": params.get("ep"), + "moe_cluster_parallel_size": params.get("cluster_size"), + "enable_chunked_prefill": params.get("enable_chunked_context"), + "skip_tokenizer_init": not params.get("no_skip_tokenizer_init", False), + } + for param in ("max_batch_size", "max_num_tokens"): + if params.get(param) is not None: + llm_args_kwargs[param] = params[param] + + llm_args = get_llm_args_from_cli_params(model=model, **llm_args_kwargs) tp_size = llm_args.tensor_parallel_size pp_size = llm_args.pipeline_parallel_size @@ -129,8 +135,10 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, if llm_args.max_seq_len is None: llm_args.max_seq_len = dataset_metadata.max_sequence_length - # Check if max_batch_size and max_num_tokens were explicitly set in either CLI params or - # config YAML file + if options.iteration_log is not None: + llm_args.enable_iter_perf_stats = True + + # Check if max_batch_size and max_num_tokens were explicitly set (via CLI or config YAML) if {"max_batch_size", "max_num_tokens"}.issubset(llm_args.model_fields_set): logger.info("Using user-provided max batch size and max num tokens.") max_batch_size, max_num_tokens = llm_args.max_batch_size, llm_args.max_num_tokens @@ -201,7 +209,7 @@ def get_exec_settings_for_backend( if bench_env.checkpoint_path is None: snapshot_download(options.model, revision=bench_env.revision) - exec_settings = get_settings(params, metadata, bench_env.model, + exec_settings = get_settings(params, metadata, options, bench_env.model, bench_env.checkpoint_path) return exec_settings elif backend == "tensorrt": diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index b154f0a4c23f..e55a27c83704 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -13,6 +13,6 @@ class RuntimeConfig(BaseModel): engine_dir: Optional[Path] = None revision: Optional[str] = None sw_version: str - backend: Literal["pytorch", "_autodeploy", None] = None + backend: Literal["pytorch", "_autodeploy", "tensorrt", None] = None iteration_log: Optional[Path] = None llm_args: BaseLlmArgs diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5865958e17bb..d5d629da0d72 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -13,7 +13,7 @@ import torch import yaml -from pydantic import BaseModel, ConfigDict +from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, PrivateAttr, field_validator, model_validator) @@ -652,6 +652,8 @@ class DecodingBaseConfig(StrictBaseModel): speculative_model: Optional[Union[str, Path]] = Field( default=None, + validation_alias=AliasChoices("speculative_model", + "speculative_model_dir"), description= "The speculative (draft) model. Accepts either (1) a HuggingFace Hub model ID (e.g. 'yuhuili/EAGLE3-LLaMA3.1-Instruct-8B')," "which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory." @@ -2430,16 +2432,19 @@ def validate_build_config_with_runtime_params(self): ) if self.max_seq_len is not None: if self.max_seq_len != self.build_config.max_seq_len: + self.max_seq_len = self.build_config.max_seq_len logger.warning( f"max_seq_len [{self.max_seq_len}] is overridden by build_config.max_seq_len [{self.build_config.max_seq_len}] in build_config" ) if self.max_beam_width is not None: if self.max_beam_width != self.build_config.max_beam_width: + self.max_beam_width = self.build_config.max_beam_width logger.warning( f"max_beam_width [{self.max_beam_width}] is overridden by build_config.max_beam_width [{self.build_config.max_beam_width}] in build_config" ) if self.max_input_len is not None: if self.max_input_len != self.build_config.max_input_len: + self.max_input_len = self.build_config.max_input_len logger.warning( f"max_input_len [{self.max_input_len}] is overridden by build_config.max_input_len [{self.build_config.max_input_len}] in build_config" ) @@ -2467,12 +2472,6 @@ def validate_build_config_remaining(self): 'enable_prompt_adapter') and self.enable_prompt_adapter: self.build_config.max_prompt_embedding_table_size = self.max_prompt_adapter_token * self.build_config.max_batch_size - if self.max_beam_width is None: - if self.build_config: - self.max_beam_width = self.build_config.max_beam_width - else: - self.max_beam_width = 1 - return self @model_validator(mode="after") From 1e9fdfa58f88abf49f0c948cd5e30652d03e6aa1 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 13:16:19 -0800 Subject: [PATCH 11/15] Fixes Signed-off-by: Anish Shanbhag --- tensorrt_llm/llmapi/build_cache.py | 4 ++-- tensorrt_llm/llmapi/llm_create.py | 16 +++++++++++++++- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/llmapi/build_cache.py b/tensorrt_llm/llmapi/build_cache.py index b22a80dc514b..4b391c9341d5 100644 --- a/tensorrt_llm/llmapi/build_cache.py +++ b/tensorrt_llm/llmapi/build_cache.py @@ -38,7 +38,7 @@ class BuildCacheConfig(BaseModel): changed, you should remove the caches manually. """ - cache_root: Optional[str] = Field( + cache_root: Optional[Path] = Field( default=None, description= "The root directory for the build cache. Falls back to env var if not provided." @@ -57,7 +57,7 @@ def set_default_cache_root(self) -> "BuildCacheConfig": """Set cache_root from environment variable if not provided.""" if self.cache_root is None: _, default_cache_root = get_build_cache_config_from_env() - self.cache_root = default_cache_root + self.cache_root = Path(default_cache_root) return self diff --git a/tensorrt_llm/llmapi/llm_create.py b/tensorrt_llm/llmapi/llm_create.py index 9b58dea5fcc1..935ea9ba28df 100644 --- a/tensorrt_llm/llmapi/llm_create.py +++ b/tensorrt_llm/llmapi/llm_create.py @@ -1,3 +1,17 @@ +# 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 typing import Any import yaml @@ -39,7 +53,7 @@ def get_llm_args_from_cli_params(model: str, **params: dict[str, Any]) -> BaseLl params["kv_cache_config"]["free_gpu_memory_fraction"] = params.pop( "free_gpu_memory_fraction" ) - # TODO: align CLI param naming with trtllm-serve / trtllm-bench to avoid duplication with above + # TODO: align CLI param naming across trtllm-serve / trtllm-bench to avoid duplication with above if "kv_cache_free_gpu_memory_fraction" in params: params["kv_cache_config"]["free_gpu_memory_fraction"] = params.pop( "kv_cache_free_gpu_memory_fraction" From e7e86f8692628e070a0d2a897a823b0da781754a Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 16:28:40 -0800 Subject: [PATCH 12/15] Add tests for best practices Signed-off-by: Anish Shanbhag --- tensorrt_llm/llmapi/llm_args.py | 14 +- tests/unittest/llmapi/test_llm_args.py | 279 ++++++++++++++++++++++++- 2 files changed, 283 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d5d629da0d72..715f2a5c2676 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -58,8 +58,9 @@ SpeculativeDecodingMode) from ..sampling_params import BatchedLogitsProcessor from .build_cache import BuildCacheConfig +from .mpi_session import MpiSession from .tokenizer import TokenizerBase, tokenizer_factory -from .utils import generate_api_docs_as_docstring, get_type_repr +from .utils import generate_api_docs_as_docstring TypeBaseModel = TypeVar("T", bound=BaseModel) @@ -2000,12 +2001,8 @@ class BaseLlmArgs(StrictBaseModel): "Guided decoding backend. llguidance is supported in PyTorch backend only." ) - batched_logits_processor: Optional[object] = Field( - default=None, - description="Batched logits processor.", - json_schema_extra={ - "type": f"Optional[{get_type_repr(BatchedLogitsProcessor)}]" - }) + batched_logits_processor: Optional[BatchedLogitsProcessor] = Field( + default=None, description="Batched logits processor.") iter_stats_max_iterations: Optional[int] = Field( default=None, @@ -2089,10 +2086,9 @@ class BaseLlmArgs(StrictBaseModel): deprecated="Use speculative_config instead.", ) - mpi_session: Optional[object] = Field( + mpi_session: Optional[MpiSession] = Field( default=None, description="The optional MPI session to use for this LLM instance.", - json_schema_extra={"type": "Optional[MpiSession]"}, exclude=True, alias="_mpi_session") diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 3aaa2ab19922..581bab780c0c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1,13 +1,19 @@ import tempfile +from dataclasses import is_dataclass +from enum import Enum +from pathlib import Path +from typing import Any, get_args, get_origin import pydantic_core import pytest import yaml -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter import tensorrt_llm.bindings.executor as tle from tensorrt_llm import LLM as TorchLLM from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm._torch.auto_deploy.llm_args import \ + LlmArgs as AutoDeployLlmArgs from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, SchedulerConfig) from tensorrt_llm.llmapi.llm_args import * @@ -760,3 +766,274 @@ def test_executor_config_consistency(): # ASK chunweiy for help if you are not sure about the new options. assert executor_config_attrs.issubset(llm_args_attr), \ f"New options found in underlying ExecutorConfig: {executor_config_attrs - llm_args_attr}" + + +def _get_all_llm_args_classes(): + """Get all subclasses of BaseLlmArgs by traversing the class hierarchy.""" + subclasses = [] + to_visit = [BaseLlmArgs] + while to_visit: + cls = to_visit.pop() + for subclass in cls.__subclasses__(): + subclasses.append(subclass) + to_visit.append(subclass) + return subclasses + + +def _get_all_pydantic_models_from_llm_args(): + """ + Get all Pydantic models referenced by BaseLlmArgs and its subclasses, + including nested models. + """ + from typing import get_args + + from pydantic import BaseModel + + visited = set() + models = [] + + def extract_types_from_annotation(annotation): + """Extract all concrete types from an annotation (handles Union, Optional, List, etc.).""" + args = get_args(annotation) + if args: + # Generic type - recurse into all type arguments + return [ + t for arg in args if arg is not type(None) + for t in extract_types_from_annotation(arg) + ] + elif isinstance(annotation, type): + return [annotation] + return [] + + def visit_model(model_cls): + if model_cls in visited: + return + if not isinstance(model_cls, type) or not issubclass( + model_cls, BaseModel): + return + + visited.add(model_cls) + models.append(model_cls) + + # Visit all field types + for field_name, field_info in model_cls.model_fields.items(): + annotation = field_info.annotation + if annotation is None: + continue + + # Extract all types from the annotation + types_in_annotation = extract_types_from_annotation(annotation) + for type_cls in types_in_annotation: + if isinstance(type_cls, type) and issubclass( + type_cls, BaseModel): + visit_model(type_cls) + + # Start with BaseLlmArgs and all its subclasses + for cls in [BaseLlmArgs] + _get_all_llm_args_classes(): + visit_model(cls) + + return models + + +# Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases. +# Avoid adding to this list unless absolutely necessary, especially if a field is user-facing. +# Keys are Pydantic model classes, values are lists of field names. +EXEMPT_FIELDS: dict[type, list[str]] = { + BaseLlmArgs: [ + "batched_logits_processor", # ABC type (BatchedLogitsProcessor) + "decoding_config", # deprecated field, typed as object + "model_kwargs", # typed as Dict[str, Any] for flexibility + "mpi_session", # ABC type (MpiSession) + ], + TorchLlmArgs: [ + "checkpoint_loader", # TODO: typed as object due to circular import + ], + AutoDeployLlmArgs: [ + "draft_checkpoint_loader", # object due to circular import + "transforms", # typed as Dict[str, Dict[str, Any]] for flexibility + "model_kwargs", # typed as Dict[str, Any] for flexibility + "tokenizer_kwargs", # typed as Dict[str, Any] for flexibility + ], + UserProvidedDecodingConfig: [ + "drafter", # TODO: typed as object due to circular import (actual type: Drafter) + "resource_manager", # TODO: typed as object due to circular import (actual type: ResourceManager) + ], + MoeConfig: ["load_balancer"], # allows multiple types including dict + RayPlacementConfig: ["placement_groups"], # contains Ray-specific types +} + + +class TestUserFacingPydanticModels: + """ + Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices. + """ + + def _is_allowed_type(self, annotation, model_cls: type, + field_name: str) -> tuple[bool, str]: + """ + Check if a type annotation is allowed according to the rules below: + + Allowed: + - Pydantic models + - Enums + - primitives (str, int, float, bool, bytes, Path, type(None)) + - Union of allowed types + - Optional of allowed types + Not allowed: + - object, Any, bare containers (dict, list, etc. without type parameters) + - dataclasses (should be converted to Pydantic models) + - regular classes that aren't Pydantic models, Enums, or primitives + - Pydantic models that do not inherit from StrictBaseModel + - Any other type that cannot be translated to a JSON schema + + Returns (is_compatible, reason) tuple. + """ + + # Check if this field is exempt (check class and all parent classes) + for cls in model_cls.__mro__: + if field_name in EXEMPT_FIELDS.get(cls, []): + return True, "exempt" + + # Check explicitly blocked types + if annotation is object: + return False, "bare 'object' type (use a specific type or Pydantic model)" + if annotation is Any: + return False, "'Any' type (use a specific type)" + + # Check for bare container types (missing type parameters) + if annotation in (dict, list, tuple, set, frozenset): + name = annotation.__name__ + return False, f"bare '{name}' type (use {name}[...] with specific type parameters)" + + # Check for dataclasses (should use Pydantic models instead) + if isinstance(annotation, type) and is_dataclass(annotation): + return False, f"dataclass '{annotation.__name__}' (convert to Pydantic BaseModel)" + + # Check for regular classes that aren't Pydantic models, enums, or primitives + _ALLOWED_CLASS_BASES = (BaseModel, Enum, str, int, float, bool, bytes, + Path, type(None)) + if isinstance(annotation, type) and not issubclass( + annotation, _ALLOWED_CLASS_BASES): + return False, f"class '{annotation.__name__}' is not a Pydantic model (convert to StrictBaseModel)" + + # Require user-facing Pydantic models to forbid arbitrary fields + if issubclass(annotation, BaseModel) and not issubclass( + annotation, StrictBaseModel): + return False, f"Pydantic model '{annotation.__name__}' is not a StrictBaseModel (convert to StrictBaseModel)" + + # Recursively check generic type arguments for disallowed types + origin = get_origin(annotation) + if origin is not None: + for arg in get_args(annotation): + if arg is type(None): + continue + ok, reason = self._is_allowed_type(arg, model_cls, field_name) + if not ok: + container_name = getattr(origin, "__name__", str(origin)) + return False, f"in {container_name}: {reason}" + + # Final validation: can Pydantic generate a JSON schema for this type? + try: + TypeAdapter(annotation).json_schema() + return True, "schema compatible" + except Exception as e: + return False, f"cannot derive JSON schema: {type(e).__name__}: {e}" + + def test_all_fields_have_descriptions(self): + """Test that all fields in LlmArgs classes (including subfields) have descriptions.""" + violations = [] + + for cls in _get_all_pydantic_models_from_llm_args(): + for field_name, field_info in cls.model_fields.items(): + if field_name.startswith("_"): + # Skip private / non user-facing fields + continue + if field_info.description is None or field_info.description.strip( + ) == "": + violations.append( + f"{cls.__name__}.{field_name}: missing description") + + if violations: + pytest.fail( + f"The following fields are missing descriptions:\n" + + "\n".join(violations) + + "\n\nPlease add a description to each by using Field(description=\"...\")." + ) + + def test_all_fields_have_allowed_types(self): + """Test that all fields in LlmArgs classes (including subfields) have types that are allowed + (i.e. are Pydantic-compatible) according to the logic in _is_allowed_type.""" + violations = [] + + for cls in _get_all_pydantic_models_from_llm_args(): + for field_name, field_info in cls.model_fields.items(): + if field_name.startswith("_"): + # Skip private / non user-facing fields + continue + is_compatible, reason = self._is_allowed_type( + field_info.annotation, cls, field_name) + if not is_compatible: + violations.append(f"{cls.__name__}.{field_name}: {reason}") + + if violations: + pytest.fail( + f"The following user-facing fields have types that are not allowed:\n" + + "\n".join(violations) + + "\n\nPlease use Pydantic-compatible types (primitives, Pydantic models that inherit from StrictBaseModel, " + "or other compatible types). If this is intentional, add the field " + "to EXEMPT_FIELDS.") + + # Methods that shouldn't be manually defined on Pydantic models + _FORBIDDEN_METHODS = { + "from_dict": + "Construct the class directly from the dict instead, i.e. MyModel(**my_dict).", + "to_dict": "Use Pydantic's model_dump() instead.", + "validate": + "Use Pydantic's @field_validator or @model_validator instead.", + "is_valid": + "Use Pydantic's @field_validator or @model_validator instead.", + } + + def test_no_manual_serialization_or_validation_methods(self): + """Test that no Pydantic models define manual serialization/validation methods.""" + violations = [] + + for cls in _get_all_pydantic_models_from_llm_args(): + for method_name, suggestion in self._FORBIDDEN_METHODS.items(): + if method_name in cls.__dict__: + violations.append( + f"{cls.__name__}.{method_name}(): {suggestion}") + + if violations: + pytest.fail( + f"The following models define forbidden methods:\n" + + "\n".join(violations) + + "\n\nPydantic models should use built-in methods instead. " + "See: https://docs.pydantic.dev/latest/concepts/serialization/") + + def test_no_mutable_default_values(self): + """Test that no fields use mutable default values directly. + + Mutable defaults like [] or {} are shared across instances and cause bugs. + Use Field(default_factory=list) or Field(default_factory=dict) instead. + """ + violations = [] + + for cls in _get_all_pydantic_models_from_llm_args(): + for field_name, field_info in cls.model_fields.items(): + if field_name.startswith("_"): + continue + # Check if the default is a mutable type instance + default = field_info.default + if isinstance(default, (list, dict, set)): + type_name = type(default).__name__ + violations.append( + f"{cls.__name__}.{field_name}: uses mutable default {type_name}. " + f"Use Field(default_factory={type_name}) instead.") + + if violations: + pytest.fail( + f"The following fields use mutable default values:\n" + + "\n".join(violations) + + "\n\nMutable defaults are shared across instances and cause bugs. " + "Use Field(default_factory=...) instead.") From 8295fbeb798ad4140613731b899d94fc7cbfa478 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 16:28:52 -0800 Subject: [PATCH 13/15] Add Pydantic coding guidelines Signed-off-by: Anish Shanbhag --- CODING_GUIDELINES.md | 55 ++++++++++++++++++++++++++++---------------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 00129ee8bc12..f23f7029946e 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -404,29 +404,44 @@ foo.SomeClass() 1. For interfaces that may be used outside a file, prefer docstrings over comments. 2. Comments should be reserved for code within a function, or interfaces that are local to a file. -#### Docstring Syntax -##### Classes and Functions -Use the [Google style](https://google.github.io/styleguide/pyguide.html), which can be parsed by Sphinx. - -When defining any user-facing configuration classes (e.g. `LlmArgs` or any class used in its fields), always use Pydantic classes rather than dataclasses or vanilla classes. - -- avoid to_dict / from_dict -- avoid dicts / object as any part of a pydantic class -- use discriminated unions -- prefer (1) positiveint/nonnegativeint/nonnegativefloat/positivefloat, or (2) if not possible, gt/ge/le/lt -- use min_length to enforce minimum length of a list -- use pydantic field descriptions instead of comments -- use Literal instead of str when a field should only accept certain values. avoid defining a custom validator on a str field. -- raise ValueError instead of assertion -- co-locate model validation logic within the class itself rather than in a parent class, unless it depends on other fields in the parent class (e.g. don't validate cuda_graph_config in BaseLlmArgs; also add example of OK case) -- don't define manual validate() methods, prefer using model_validator / field_validator +#### Pydantic Guidelines + +When defining any user-facing configuration classes (particularly `LlmArgs` or any class used in its fields), always use Pydantic classes rather than dataclasses or vanilla classes. + +**Model Structure:** +- Inherit from `StrictBaseModel` (which sets `extra="forbid"`) to catch typos in field names +- Use [discriminated unions](https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions) when a field needs to accept one of several possible config classes (e.g. `speculative_config` accepts any of `EagleDecodingConfig`, `MedusaDecodingConfig`, etc.) + +**Field Definitions:** +- Add descriptions to all user-facing fields via `Field(description="...")`. Avoid using comments for descriptions. +- Avoid `dict`, `object`, `Any` as field types - use properly typed alternatives +- Avoid defining mutable defaults directly; use `default_factory` instead. + - Good: `Field(default_factory=list)`, `Field(default_factory=dict)`, `Field(default_factory=MyClass)` + - Bad: `Field(default=[])`, `Field(default={})`, `Field(default=MyClass())` directly +- Use `Literal["value1", "value2"]` instead of `str` when a field should only accept certain values +- Prefer `PositiveInt`, `NonNegativeInt`, `NonNegativeFloat`, `PositiveFloat`, `Field(gt=0)`, `Field(ge=0)`, etc. for numeric constraints instead of defining custom validators +- Use `Field(min_length=1)` to enforce minimum length of a list + +**Validation:** +- Use `@field_validator` and `@model_validator` instead of manual `validate()` or `is_valid()` methods +- Raise `ValueError` instead of using assertions +- Co-locate validation logic within the class itself rather than in a parent class, unless it depends on fields in the parent + +**Serialization:** +- Avoid defining `to_dict()` methods - prefer Pydantic's built-in `model_dump()` to convert to a dictionary. + - Note: you can override `model_dump()` to customize its behavior, but avoid doing so unless absolutely necessary. + - Good: `MyModel.model_dump()` + - Bad: `MyModel.to_dict()` +- Avoid defining `from_dict()` methods - prefer constructing the class directly from arguments. + - Good: `MyModel(**kwargs)` + - Bad: `MyModel.from_dict(kwargs)` TODO: -- ask cursor for more antipatterns fixed in this pr - unit tests: e.g. for get_llm_args_from_cli_params, other tests you removed, etc -- tests for all of the above Pydantic best practices -- test to validate that all llmargs classes have field descriptions for all fields -- add validator to BaseLlmArgs.__init_subclass__ to check that all subfields are Pydantic + +#### Docstring Syntax +##### Classes and Functions +Use the [Google style](https://google.github.io/styleguide/pyguide.html), which can be parsed by Sphinx. ##### Attributes and Variables Attributes and variables can be documented inline. Attribute docstrings will be rendered under the docstring for the class. For example: From 7523ca5607e28b6cb2428e2e59a2e2947ea0ea7a Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Fri, 30 Jan 2026 16:33:59 -0800 Subject: [PATCH 14/15] Add TODO Signed-off-by: Anish Shanbhag --- tensorrt_llm/llmapi/llm_create.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/llmapi/llm_create.py b/tensorrt_llm/llmapi/llm_create.py index 935ea9ba28df..015fd16c5a45 100644 --- a/tensorrt_llm/llmapi/llm_create.py +++ b/tensorrt_llm/llmapi/llm_create.py @@ -34,6 +34,9 @@ def _deep_merge(base: dict, override: dict) -> dict: return result +# TODO: why does config need to take precedence over CLI params? you did this for a reason + + def get_llm_args_from_cli_params(model: str, **params: dict[str, Any]) -> BaseLlmArgs: """Construct an instance of LlmArgs from CLI params and optionally a config YAML file. From 22603530eda4545710e490e0a5ceaa3863f760ef Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Wed, 4 Feb 2026 15:23:49 -0800 Subject: [PATCH 15/15] Fix CompletionOutput pickling Signed-off-by: Anish Shanbhag --- tensorrt_llm/executor/result.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 0642c1aff1fe..3de70b7c7cfb 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -158,6 +158,18 @@ def token_ids_diff(self) -> List[int]: def logprobs_diff(self) -> TokenLogprobs | List[float]: return self.logprobs[self._last_logprobs_len:] + def __getstate__(self): + # Exclude _incremental_states which contains unpicklable DecodeStream objects + return { + slot: getattr(self, slot) + for slot in self.__slots__ if slot != "_incremental_states" + } + + def __setstate__(self, state): + for slot, value in state.items(): + setattr(self, slot, value) + self._incremental_states = None + class GenerationResultBase: ''' This holds the core logic of the GenerationResult class. '''