diff --git a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py index 2130b4e255c3..bd175dc86350 100644 --- a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py @@ -1,6 +1,6 @@ import torch -from ..flashinfer_utils import ENABLE_PDL, IS_FLASHINFER_AVAILABLE +from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl if IS_FLASHINFER_AVAILABLE: from flashinfer.activation import silu_and_mul @@ -11,7 +11,7 @@ # Warp this into custom op since flashinfer didn't warp it properly and we want to avoid graph break between mlp layer for user buffer optimization @torch.library.custom_op("trtllm::flashinfer_silu_and_mul", mutates_args=()) def flashinfer_silu_and_mul(x: torch.Tensor) -> torch.Tensor: - return silu_and_mul(x, enable_pdl=ENABLE_PDL) + return silu_and_mul(x, enable_pdl=get_env_enable_pdl()) @flashinfer_silu_and_mul.register_fake def _(x: torch.Tensor) -> torch.Tensor: @@ -21,7 +21,7 @@ def _(x: torch.Tensor) -> torch.Tensor: @torch.library.custom_op("trtllm::flashinfer_rmsnorm", mutates_args=()) def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - return rmsnorm(input, weight, eps, enable_pdl=ENABLE_PDL) + return rmsnorm(input, weight, eps, enable_pdl=get_env_enable_pdl()) @flashinfer_rmsnorm.register_fake def _(input: torch.Tensor, weight: torch.Tensor, @@ -32,7 +32,10 @@ def _(input: torch.Tensor, weight: torch.Tensor, mutates_args=()) def flashinfer_gemma_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - return gemma_rmsnorm(input, weight, eps, enable_pdl=ENABLE_PDL) + return gemma_rmsnorm(input, + weight, + eps, + enable_pdl=get_env_enable_pdl()) @flashinfer_gemma_rmsnorm.register_fake def _(input: torch.Tensor, weight: torch.Tensor, @@ -44,7 +47,11 @@ def _(input: torch.Tensor, weight: torch.Tensor, def flashinfer_fused_add_rmsnorm(input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float) -> None: - fused_add_rmsnorm(input, residual, weight, eps, enable_pdl=ENABLE_PDL) + fused_add_rmsnorm(input, + residual, + weight, + eps, + enable_pdl=get_env_enable_pdl()) @torch.library.custom_op("trtllm::flashinfer_gemma_fused_add_rmsnorm", mutates_args=("input", "residual")) @@ -56,7 +63,7 @@ def flashinfer_gemma_fused_add_rmsnorm(input: torch.Tensor, residual, weight, eps, - enable_pdl=ENABLE_PDL) + enable_pdl=get_env_enable_pdl()) @torch.library.custom_op( "trtllm::flashinfer_apply_rope_with_cos_sin_cache_inplace", diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 5b150665b5b9..8a58065c7930 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -8,13 +8,13 @@ def get_env_enable_pdl(): - return os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1" + enabled = os.environ.get("TRTLLM_ENABLE_PDL", "0") == "1" + if enabled and not getattr(get_env_enable_pdl, "_printed", False): + logger.info("PDL enabled") + setattr(get_env_enable_pdl, "_printed", True) + return enabled -ENABLE_PDL = get_env_enable_pdl() -if ENABLE_PDL: - logger.info("PDL is enabled") - if platform.system() != "Windows": try: import flashinfer diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py index a45efc8d8e50..a37f3798ec75 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py @@ -30,7 +30,7 @@ else: from typing_extensions import override -from ..flashinfer_utils import ENABLE_PDL +from ..flashinfer_utils import get_env_enable_pdl from .sampling_utils import ( GREEDY, GroupedStrategySampler, @@ -112,7 +112,7 @@ def _prepare_probs_with_temperature( probs = flashinfer.sampling.softmax( logits, temperature, - enable_pdl=ENABLE_PDL, + enable_pdl=get_env_enable_pdl(), ) return probs diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 63313427d4d6..c36a7dc43e49 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -import os from functools import partial from pathlib import Path @@ -46,12 +45,14 @@ help="Path to a serialized TRT-LLM engine.", ) @optgroup.option( + "--config", "--extra_llm_api_options", + "extra_llm_api_options", type=str, default=None, help= - "Path to a YAML file that overwrites the parameters specified by trtllm-bench." -) + "Path to a YAML file that overwrites the parameters specified by trtllm-bench. " + "Can be specified as either --config or --extra_llm_api_options.") @optgroup.option( "--backend", type=click.Choice(ALL_SUPPORTED_BACKENDS), @@ -192,6 +193,7 @@ def latency_command( ) -> None: """Run a latency test on a TRT-LLM engine.""" logger.info("Preparing to run latency benchmark...") + # Parameters from CLI # Model, experiment, and engine params options = get_general_cli_options(params, bench_env) @@ -263,14 +265,6 @@ def latency_command( exec_settings["settings_config"][ "scheduler_policy"] = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT - # Set environment variables for setting runtime options. - # TODO: Once passing of variables is fixed, these should work - # when using MPI in C++ runtime. - os.environ["TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG"] = "1" - os.environ["TRTLLM_MMHA_KERNEL_BLOCK_SIZE"] = "256" - os.environ["FORCE_MULTI_BLOCK_MODE"] = "1" - os.environ["TRTLLM_ENABLE_PDL"] = "1" - # Performance options exec_settings["performance_options"]["cuda_graphs"] = True exec_settings["performance_options"]["multi_block_mode"] = True @@ -290,6 +284,17 @@ def latency_command( kwargs = kwargs | runtime_config.get_llm_args() kwargs['backend'] = options.backend + # Set environment variables for setting runtime options. + default_env_overrides = { + "TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG": "1", + "TRTLLM_MMHA_KERNEL_BLOCK_SIZE": "256", + "FORCE_MULTI_BLOCK_MODE": "1", + "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 + try: logger.info("Setting up latency benchmark.") diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 171b14ada760..6e9f6b87eccb 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -61,12 +61,14 @@ help="Paths to custom module directories to import.", ) @optgroup.option( + "--config", "--extra_llm_api_options", + "extra_llm_api_options", type=str, default=None, help= - "Path to a YAML file that overwrites the parameters specified by trtllm-bench." -) + "Path to a YAML file that overwrites the parameters specified by trtllm-bench. " + "Can be specified as either --config or --extra_llm_api_options.") @optgroup.option("--sampler_options", type=click.Path(exists=True, readable=True, @@ -293,6 +295,7 @@ def throughput_command( ) -> None: """Run a throughput test on a TRT-LLM engine.""" logger.info("Preparing to run throughput benchmark...") + # Parameters from CLI image_data_format: str = params.get("image_data_format", "pt") data_device: str = params.get("data_device", "cpu") diff --git a/tensorrt_llm/commands/eval.py b/tensorrt_llm/commands/eval.py index 5668ab3452ae..2ac2c95bd842 100644 --- a/tensorrt_llm/commands/eval.py +++ b/tensorrt_llm/commands/eval.py @@ -97,10 +97,13 @@ default=None, help="The revision to use for the HuggingFace model " "(branch name, tag name, or commit id).") -@click.option("--extra_llm_api_options", +@click.option("--config", + "--extra_llm_api_options", + "extra_llm_api_options", type=str, default=None, - help="Path to a YAML file that overwrites the parameters") + help="Path to a YAML file that overwrites the parameters. " + "Can be specified as either --config or --extra_llm_api_options.") @click.option("--disable_kv_cache_reuse", is_flag=True, default=False, diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 90fe07ceb5a3..cabbc4753958 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -325,12 +325,14 @@ def convert(self, value: Any, param: Optional["click.Parameter"], help="The revision to use for the HuggingFace model " "(branch name, tag name, or commit id).") @click.option( + "--config", "--extra_llm_api_options", + "extra_llm_api_options", type=str, default=None, help= - "Path to a YAML file that overwrites the parameters specified by trtllm-serve." -) + "Path to a YAML file that overwrites the parameters specified by trtllm-serve. " + "Can be specified as either --config or --extra_llm_api_options.") @click.option( "--reasoning_parser", type=click.Choice(ReasoningParserFactory.parsers.keys()), diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 4ed3f56cbd4b..2199bee74af3 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -241,8 +241,17 @@ def worker_main( tokenizer: Optional[TokenizerBase] = None, llm_args: Optional[BaseLlmArgs] = None, ) -> None: + mpi_comm().barrier() + if llm_args is not None and llm_args.env_overrides: + # this is needed because MPI_Init seems to cache the env at import time. + # The cached env snapshot is used to spawn workers. + # Any env overrides to the main process after tensorrt_llm import + # may not get reflected in the spawned worker process, no matter how early, + # unless we update it explicitly here. + os.environ.update(llm_args.env_overrides) + if llm_args is not None and llm_args.trust_remote_code: _init_hf_modules() diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 0ca965d0e08c..e4b060ee48f4 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -135,6 +135,9 @@ def __init__(self, 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") @@ -587,6 +590,25 @@ def get_kv_cache_events_async(self, ''' return self._executor.aget_kv_events(timeout=timeout) + def _process_env_overrides(self, + env_overrides: Optional[dict[str, str]]) -> None: + if env_overrides is None: + return + logger.info("Processing LLM API environment variable overrides") + # TODO: If an env var is cached at import-time in code, overriding os.environ will + # unfortunately not update wherever the var is used. + # This is a known issue and only way to fix it is at every such usage to access it + # from os.environ on-demand. + for key, value in env_overrides.items(): + str_value = str(value) + if key in os.environ: + old_value = os.environ[key] + os.environ[key] = str_value + logger.info(f"Overriding {key}: '{old_value}' -> '{str_value}'") + else: + os.environ[key] = str_value + logger.info(f"Setting {key}='{str_value}'") + def _prepare_sampling_params( self, sampling_params: Optional[SamplingParams] = None) -> SamplingParams: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 49a4c3486d4b..f719f03fd52f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1926,6 +1926,12 @@ class BaseLlmArgs(StrictBaseModel): status="prototype", ) + env_overrides: Optional[Dict[str, str]] = Field( + default=None, + description= + "[EXPERIMENTAL] Environment variable overrides. NOTE: import-time-cached env vars in the code won’t update unless the code fetches them from os.environ on demand.", + status="prototype") + _parallel_config: Optional[_ParallelConfig] = PrivateAttr(default=None) _model_format: Optional[_ModelFormatKind] = PrivateAttr(default=None) _speculative_model: Optional[str] = PrivateAttr(default=None) diff --git a/tests/integration/defs/examples/serve/test_serve.py b/tests/integration/defs/examples/serve/test_serve.py old mode 100755 new mode 100644 index c861d525a265..d5bfb03f3093 --- a/tests/integration/defs/examples/serve/test_serve.py +++ b/tests/integration/defs/examples/serve/test_serve.py @@ -1,7 +1,12 @@ import os +import queue +import subprocess +import threading import time +import pytest import requests +import yaml from defs.conftest import llm_models_root, skip_no_hopper from defs.trt_test_alternative import popen, print_error, print_info from openai import OpenAI @@ -33,6 +38,17 @@ def check_server_ready(http_port="8000", timeout_timer=600, sleep_interval=5): ) +def wait_for_log(log_queue, expected_log, timeout=10): + """Waits for a specific log message to appear in the queue.""" + start_time = time.time() + while time.time() - start_time < timeout: + current_logs = "".join(list(log_queue.queue)) + if expected_log in current_logs: + return True + time.sleep(0.1) + return False + + def check_openai_chat_completion(http_port="8000", model_name="TinyLlama-1.1B-Chat-v1.0"): """ @@ -92,8 +108,10 @@ def check_openai_chat_completion(http_port="8000", raise +@pytest.mark.parametrize("config_flag", ["--extra_llm_api_options", "--config"]) @skip_no_hopper -def test_extra_llm_api_options(serve_test_root): +def test_config_file_loading(serve_test_root, config_flag): + """Test config file loading via both --extra_llm_api_options and --config flags.""" test_configs_root = f"{serve_test_root}/test_configs" # moe backend = CUTLASS which only supports fp8 blockscale on Hopper @@ -119,14 +137,84 @@ def test_extra_llm_api_options(serve_test_root): "8000", "--backend", "pytorch", - "--extra_llm_api_options", + config_flag, config_file, ] - print_info("Launching trtllm-serve...") + print_info(f"Launching trtllm-serve with {config_flag}...") with popen(cmd): check_server_ready() # Extract model name from the model path for consistency model_name = model_path.split('/')[-1] # "Qwen3-30B-A3B-FP8" # Test the server with OpenAI chat completion check_openai_chat_completion(model_name=model_name) + + +@skip_no_hopper +def test_env_overrides_pdl(tmp_path): + """ + This test ensures that the `env_overrides` configuration option effectively propagates + environment variables to the server workers. Specifically, it sets `TRTLLM_ENABLE_PDL=1` + (Programmatic Dependent Launch) via config and verifies it overrides the env var initially set to 0. + + 1. This model (TinyLlama-1.1B-Chat-v1.0) architecture uses RMSNorm, which triggers 'flashinfer' kernels that use PDL when `TRTLLM_ENABLE_PDL=1`. + 2. When `TRTLLM_ENABLE_PDL=1` is actually propagated into worker env, flashinfer custom ops log "PDL enabled" to stdout/stderr. + """ + pdl_enabled = "1" + config_file = tmp_path / "config.yml" + config_file.write_text( + yaml.dump({ + "backend": "pytorch", + "env_overrides": { + "TRTLLM_ENABLE_PDL": pdl_enabled + } + })) + + port = 8001 + env = os.environ.copy() + # Pre-load with 0 to verify override works + env.update({ + "TLLM_LOG_LEVEL": "INFO", + "LLM_MODELS_ROOT": llm_models_root(), + "TRTLLM_ENABLE_PDL": "0" + }) + + cmd = [ + "trtllm-serve", "serve", + f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0", + "--host", "0.0.0.0", "--port", + str(port), "--backend", "pytorch", "--config", + str(config_file) + ] + + with popen(cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True) as proc: + output_queue = queue.Queue() + threading.Thread( + target=lambda: + [output_queue.put(line) for line in iter(proc.stdout.readline, '')], + daemon=True).start() + + check_server_ready(http_port=str(port), timeout_timer=300) + response = OpenAI(base_url=f"http://localhost:{port}/v1", + api_key="tensorrt_llm").chat.completions.create( + model="TinyLlama-1.1B-Chat-v1.0", + messages=[{ + "role": "user", + "content": "Test" + }], + max_tokens=10) + assert response and response.choices[0].message.content + + # Directly wait for the PDL log we expect + if not wait_for_log(output_queue, "PDL enabled", timeout=10): + logs = "".join(list(output_queue.queue)) + print_error( + f"Timeout waiting for 'PDL enabled'. Captured logs:\n{logs}") + assert False + + logs = ''.join(output_queue.queue) + assert "Overriding TRTLLM_ENABLE_PDL: '0' -> '1'" in logs diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 69c24f0f5ee3..504f48f2d265 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -754,7 +754,9 @@ examples/test_eagle.py::test_phi_eagle_1gpu[Phi-3-mini-128k-instruct-eagle2] examples/test_eagle.py::test_phi_eagle_1gpu[Phi-3-small-128k-instruct-eagle2] examples/test_eagle.py::test_phi_eagle_1gpu[Phi-3.5-mini-instruct-eagle2] -examples/serve/test_serve.py::test_extra_llm_api_options +examples/serve/test_serve.py::test_config_file_loading[--extra_llm_api_options] +examples/serve/test_serve.py::test_config_file_loading[--config] +examples/serve/test_serve.py::test_env_overrides_pdl examples/serve/test_serve_negative.py::test_invalid_max_tokens examples/serve/test_serve_negative.py::test_invalid_temperature examples/serve/test_serve_negative.py::test_invalid_top_p[-0.1] diff --git a/tests/integration/test_lists/qa/llm_function_nim.txt b/tests/integration/test_lists/qa/llm_function_nim.txt index 05cec91517cd..77f0563016c2 100644 --- a/tests/integration/test_lists/qa/llm_function_nim.txt +++ b/tests/integration/test_lists/qa/llm_function_nim.txt @@ -403,7 +403,8 @@ test_e2e.py::test_llmapi_generation_logits[llama-3.3-models/Llama-3.3-70B-Instru test_e2e.py::test_ptp_quickstart_multimodal_kv_cache_reuse[mistral-small-3.1-24b-instruct-Mistral-Small-3.1-24B-Instruct-2503-0.6-image] test_e2e.py::test_ptp_quickstart_multimodal_chunked_prefill[mistral-small-3.1-24b-instruct-Mistral-Small-3.1-24B-Instruct-2503-0.6-image] -examples/serve/test_serve.py::test_extra_llm_api_options +examples/serve/test_serve.py::test_config_file_loading[--extra_llm_api_options] +examples/serve/test_serve.py::test_config_file_loading[--config] examples/serve/test_serve_negative.py::test_invalid_max_tokens examples/serve/test_serve_negative.py::test_invalid_temperature examples/serve/test_serve_negative.py::test_invalid_top_p[-0.1] @@ -465,7 +466,6 @@ test_e2e.py::test_ptp_star_attention_example[Llama3.1-8B-BF16-llama-3.1-model/Me test_e2e.py::test_ptp_scaffolding[DeepSeek-R1-Distill-Qwen-7B-DeepSeek-R1/DeepSeek-R1-Distill-Qwen-7B] test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] unittest/llmapi/test_llm_pytorch.py::test_gemma3_1b_instruct_multi_lora -examples/serve/test_serve.py::test_extra_llm_api_options # llm-api promote pytorch to default llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_logging llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_type_tensorrt diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index c08ec37c1197..eb8a5d09d78d 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -203,6 +203,10 @@ methods: annotation: bool default: False status: prototype + env_overrides: + annotation: Optional[Dict[str, str]] + default: null + status: prototype return_annotation: None generate: parameters: