[None][chore] Align LlmArgs with some Pydantic best practices

AGENTS.md

@@ -102,6 +102,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT)
 - **Avoid broad exception handling** — catch specific exceptions, not bare `except:` (see `CODING_GUIDELINES.md`).
 - **Python import style is enforced** — `from package.subpackage import module`, never `from module import Class`. Pre-commit will not catch this.
 - **One concern per PR** — avoid scope creep. If a PR touches unrelated areas, split it.
+- **User-facing configuration classes** - when editing or defining any user-facing configuration classes (particularly `LlmArgs` or any class used in its fields), you **MUST** follow the Pydantic guidelines in `CODING_GUIDELINES.md`.
 
 ## Development Workflow
 

CODING_GUIDELINES.md

@@ -404,6 +404,43 @@ 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.
 
+### 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 fail fast when users specify invalid 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.)
+- **Do not define `__init__` methods** - this bypasses Pydantic's validation and type coercion, and can cause subtle bugs with model inheritance. Instead:
+  - For validation logic, use `@field_validator` or `@model_validator`
+  - For post-validation initialization (e.g. setting up private fields and state), use [`model_post_init()`](https://docs.pydantic.dev/latest/api/base_model/#pydantic.BaseModel.model_post_init)
+  - For custom construction patterns, use classmethods (e.g. `from_yaml()`)
+  - See [Defining a custom `__init__()`](https://docs.pydantic.dev/latest/concepts/models/#defining-a-custom-__init__) for more details
+
+**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.
+  - Good: `MyModel.model_dump()`
+  - Bad: `MyModel.to_dict()`
+  - Note: you can override `model_dump()` to customize its behavior, but avoid doing so unless absolutely necessary.
+- Avoid defining `from_dict()` / `from_kwargs()` methods - prefer constructing the class directly from arguments.
+  - Good: `MyModel(**kwargs)`, `MyModel(**my_dict)`
+  - Bad: `MyModel.from_dict(kwargs)`, `MyModel.from_kwargs(kwargs)`
+
 #### Docstring Syntax
 ##### Classes and Functions
 Use the [Google style](https://google.github.io/styleguide/pyguide.html), which can be parsed by Sphinx.

examples/models/core/gemma/convert_checkpoint.py

@@ -172,8 +172,9 @@ def compute_quant_algo(args: argparse.Namespace) -> Optional[QuantAlgo]:
 def create_quant_config(args: argparse.Namespace) -> QuantConfig:
     quant_algo = compute_quant_algo(args)
     GemmaForCausalLM.assert_valid_quant_algo(quant_algo)
-    quant_config = QuantConfig(quant_algo=quant_algo,
-                               smoothquant_val=args.smoothquant)
+    quant_config = QuantConfig(quant_algo=quant_algo)
+    if args.smoothquant is not None:
+        quant_config.smoothquant_val = args.smoothquant
 
     if args.fp8_kv_cache:
         quant_config.kv_cache_quant_algo = QuantAlgo.FP8

examples/models/core/llama/convert_checkpoint.py

@@ -443,7 +443,7 @@ def from_cli_args(args):
             'moe_ep_size': args.moe_ep_size,
             'cp_size': args.cp_size,
         },
-        'quantization': args_to_quant_config(args).to_dict()
+        'quantization': args_to_quant_config(args).model_dump()
     }
     config.update(args_to_build_options(args))
     return config

examples/models/core/mllama/convert_checkpoint.py

@@ -350,7 +350,7 @@ def from_cli_args(args):
             'moe_tp_size': args.moe_tp_size,
             'moe_ep_size': args.moe_ep_size,
         },
-        'quantization': args_to_quant_config(args).to_dict()
+        'quantization': args_to_quant_config(args).model_dump()
     }
     config.update(args_to_build_options(args))
     return config

tensorrt_llm/_torch/auto_deploy/llm.py

@@ -159,7 +159,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

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -187,7 +187,12 @@ def validate_and_init_tokenizer(self):
         " no processes are spawned and the model is run on a single GPU (only for ``demollm``).",
     )
 
-    runtime: Literal["demollm", "trtllm"] = Field(default="trtllm")
+    runtime: Literal["demollm", "trtllm"] = Field(
+        default="trtllm",
+        description="The runtime backend to use. 'trtllm' is a production-grade runtime optimized for "
+        "high-performance inference. 'demollm' is a lightweight runtime for development and testing "
+        "with a simplified scheduler and KV-cache manager for easier debugging.",
+    )
 
     device: str = Field(default="cuda", description="The device to use for the model.", frozen=True)
 

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -3416,7 +3416,7 @@ def _prepare_inputs(
                 return self._prepare_star_attention_inputs(
                     scheduled_requests, kv_cache_manager, attn_metadata,
                     resource_manager)
-            elif CpType.HELIX == cp_type:
+            elif cp_type in (CpType.HELIX, CpType.ULYSSES):
                 # Take the usual route of _prepare_tp_inputs.
                 pass
             else:

tensorrt_llm/builder.py

@@ -22,14 +22,15 @@
 
 import numpy as np
 import tensorrt as trt
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from ._common import _is_building, check_max_num_tokens, serialize_engine
 from ._utils import (get_sm_version, np_bfloat16, np_float8, str_dtype_to_trt,
                      to_json_file, trt_gte)
 from .functional import PositionEmbeddingType
 from .graph_rewriting import optimize
 from .llmapi.kv_cache_type import KVCacheType
+from .llmapi.utils import StrictBaseModel
 from .logger import logger
 from .lora_helper import LoraConfig
 from .models import PretrainedConfig, PretrainedModel
@@ -450,7 +451,7 @@ def save_config(builder_config: BuilderConfig, config_path: str):
         logger.info(f'Config saved to {config_path}.')
 
 
-class BuildConfig(BaseModel):
+class BuildConfig(StrictBaseModel):
     """Configuration class for TensorRT LLM engine building parameters.
 
     This class contains all the configuration parameters needed to build a TensorRT LLM engine,

tensorrt_llm/commands/build.py

@@ -458,7 +458,6 @@ def main():
         )
 
     plugin_config = PluginConfig.from_arguments(args)
-    plugin_config.validate()
     if args.fast_build:
         plugin_config.manage_weights = True
 

tensorrt_llm/commands/serve.py

@@ -189,7 +189,7 @@ def get_llm_args(
         KvCacheConfig(free_gpu_memory_fraction=free_gpu_memory_fraction)
         if free_gpu_memory_fraction != kv_cache_default_fraction else None,
         "cp_config":
-        cp_config if cp_config else {},
+        cp_config,
         "build_config":
         BuildConfig(max_batch_size=max_batch_size,
                     max_num_tokens=max_num_tokens,

tensorrt_llm/llmapi/build_cache.py

@@ -10,10 +10,12 @@
 from typing import Any, List, Optional
 
 import filelock
+from pydantic import Field, model_validator
 
 import tensorrt_llm
 from tensorrt_llm.builder import BuildConfig
-from tensorrt_llm.llmapi.utils import enable_llm_debug, print_colored
+from tensorrt_llm.llmapi.utils import (StrictBaseModel, enable_llm_debug,
+                                       print_colored)
 from tensorrt_llm.logger import logger
 
 
@@ -28,41 +30,36 @@ def get_build_cache_config_from_env() -> tuple[bool, str]:
     return build_cache_enabled, build_cache_root
 
 
-class BuildCacheConfig:
+class BuildCacheConfig(StrictBaseModel):
     """
     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[Path] = 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 = Path(default_cache_root)
+        return self
 
 
 class BuildCache:
@@ -76,17 +73,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

tensorrt_llm/llmapi/llm.py

@@ -188,17 +188,16 @@ def __init__(self,
                         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.args = llm_args_cls(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)
 
         except Exception as e:
             logger.error(