From 45ef690884c6cfe9fd0f80e9a4bbae6205e0d75b Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Tue, 26 May 2026 19:23:10 -0700 Subject: [PATCH] [nvbugs/6226287][fix] auto-handle EOS for spec-dec across benchmark entry points trtllm-bench latency already auto-detects speculative decoding and chooses end-of-sequence handling accordingly (low_latency.py:332). The throughput sub-command and trtllm-serve's benchmark_serving.py never copied this behavior, so spec-dec runs through those entry points silently produced unstable acceptance-rate numbers unless callers/test frameworks hand-picked --eos_id / --ignore-eos. This change makes all three entry points consistent: * bench/benchmark/__init__.py: GeneralExecSettings.eos_id is now Optional[int] = None (sentinel "auto-resolve") instead of -1. * bench/benchmark/throughput.py: --eos_id click default flips to None; after RuntimeConfig is built, eos_id is resolved to tokenizer.eos_token_id when speculative decoding is active and -1 otherwise (mirrors low_latency). Explicit --eos_id from the user still wins. * serve/openai_protocol.py + serve/openai_server.py: ModelCard gains an optional metadata dict, and /v1/models now exposes the engine's speculative_config there. The existing speculative_config extraction in _log_config_info_metrics is factored into _extract_speculative_config and reused. * serve/scripts/benchmark_serving.py: --ignore-eos becomes a tri-state BooleanOptionalAction (default None). When unset, the client queries /v1/models and defaults ignore_eos = False when spec-dec is enabled, True otherwise. --ignore-eos / --no-ignore-eos still override. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tensorrt_llm/bench/benchmark/__init__.py | 7 ++- tensorrt_llm/bench/benchmark/throughput.py | 23 +++++++-- tensorrt_llm/serve/openai_protocol.py | 1 + tensorrt_llm/serve/openai_server.py | 47 +++++++++++++------ .../serve/scripts/benchmark_serving.py | 46 ++++++++++++++++-- 5 files changed, 100 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index fc59fe7c6ed8..556d4387cc47 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -57,8 +57,11 @@ class GeneralExecSettings(BaseModel): description="Path to dataset file") duration: Optional[int] = Field(default=None, description="Maximum run time in seconds") - eos_id: int = Field( - default=-1, description="End-of-sequence token ID, -1 to disable EOS") + eos_id: Optional[int] = Field( + default=None, + description="End-of-sequence token ID. -1 disables EOS, " + "None auto-resolves based on the speculative decoding mode " + "(matches low_latency.py behavior).") iteration_log: Optional[Path] = Field( default=None, description="Path where iteration logging is written") kv_cache_percent: float = Field( diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index a28d6e00a2b2..415753443e2d 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -47,6 +47,7 @@ from tensorrt_llm.commands._telemetry import TelemetryCommand from tensorrt_llm.llmapi import CapacitySchedulerPolicy from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode from tensorrt_llm.sampling_params import SamplingParams @@ -149,10 +150,13 @@ @optgroup.option( "--eos_id", type=int, - default=-1, + default=None, required=False, help= - "Set the end-of-sequence token for the benchmark. Set to -1 to disable EOS.", + "Set the end-of-sequence token for the benchmark. Set to -1 to disable EOS. " + "If unset, the value is auto-resolved from the runtime decoding_mode: " + "tokenizer.eos_token_id when speculative decoding is enabled, otherwise -1 " + "(mirrors trtllm-bench latency).", ) @optgroup.option( "--modality", @@ -455,9 +459,20 @@ def throughput_command( llm = get_llm(runtime_config, kwargs) startup_metrics = llm.startup_metrics + # Auto-resolve eos_id when --eos_id was not passed (mirrors low_latency.py): + # honor EOS for spec-dec runs (clean acceptance-rate measurement), else -1 + # for stable throughput numbers. + eos_id = options.eos_id + if eos_id is None: + decoding_config = runtime_config.decoding_config + spec_active = (decoding_config is not None + and decoding_config.decoding_mode + != SpeculativeDecodingMode.NONE) + eos_id = tokenizer.eos_token_id if spec_active else -1 + sampler_args = { - "end_id": options.eos_id, - "pad_id": options.eos_id, + "end_id": eos_id, + "pad_id": eos_id, "n": options.beam_width, "use_beam_search": options.beam_width > 1 } diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index bfc6f1755cbd..778caf753577 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -240,6 +240,7 @@ class ModelCard(OpenAIBaseModel): object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "tensorrt_llm" + metadata: Optional[Dict[str, Any]] = None class ModelList(OpenAIBaseModel): diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 4d0a9c22b315..40fb603f6721 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1075,6 +1075,31 @@ def _has_cache_transceiver_config(self) -> bool: return (cache_transceiver_config is not None and getattr( cache_transceiver_config, "backend", None) is not None) + def _extract_speculative_config(self) -> Optional[Dict[str, str]]: + """Build a serializable speculative-decoding summary from generator args. + + Returns None when the engine has no speculative/decoding config attached. + Used for both Prometheus config-info metrics and the /v1/models metadata + channel that benchmark clients (e.g. benchmark_serving.py) consult to + auto-resolve --ignore-eos for spec-dec runs. + """ + args = self.generator.args + spec_config_obj = getattr(args, "speculative_config", None) or getattr( + args, "decoding_config", None) + if spec_config_obj is None: + return None + speculative_config: Dict[str, str] = {"spec_enabled": "true"} + decoding_type = getattr(spec_config_obj, "decoding_type", None) + if decoding_type is not None: + speculative_config["spec_method"] = str(decoding_type) + max_draft_len = getattr(spec_config_obj, "max_draft_len", None) + if max_draft_len is not None: + speculative_config["spec_num_tokens"] = str(max_draft_len) + draft_model = getattr(spec_config_obj, "speculative_model", None) + if draft_model is not None: + speculative_config["spec_draft_model"] = str(draft_model) + return speculative_config + def _log_config_info_metrics(self) -> None: """Extract configuration from generator args and log as Prometheus info gauges.""" args = self.generator.args @@ -1128,20 +1153,7 @@ def _log_config_info_metrics(self) -> None: parallel_config["expert_parallel_size"] = str(ep_size) # Speculative decoding config - spec_config_obj = getattr(args, "speculative_config", None) or getattr( - args, "decoding_config", None) - speculative_config = None - if spec_config_obj is not None: - speculative_config = {"spec_enabled": "true"} - decoding_type = getattr(spec_config_obj, "decoding_type", None) - if decoding_type is not None: - speculative_config["spec_method"] = str(decoding_type) - max_draft_len = getattr(spec_config_obj, "max_draft_len", None) - if max_draft_len is not None: - speculative_config["spec_num_tokens"] = str(max_draft_len) - draft_model = getattr(spec_config_obj, "speculative_model", None) - if draft_model is not None: - speculative_config["spec_draft_model"] = str(draft_model) + speculative_config = self._extract_speculative_config() # KV cache config kv_cache_config_obj = getattr(args, "kv_cache_config", None) @@ -1803,7 +1815,12 @@ async def version(self) -> JSONResponse: return JSONResponse(content=ver) async def get_model(self) -> JSONResponse: - model_list = ModelList(data=[ModelCard(id=self.model)]) + speculative_config = self._extract_speculative_config() + metadata = ({ + "speculative_config": speculative_config + } if speculative_config else None) + model_list = ModelList( + data=[ModelCard(id=self.model, metadata=metadata)]) return JSONResponse(content=model_list.model_dump()) async def get_iteration_stats(self) -> JSONResponse: diff --git a/tensorrt_llm/serve/scripts/benchmark_serving.py b/tensorrt_llm/serve/scripts/benchmark_serving.py index f5018a24b469..b6e7cd76a454 100644 --- a/tensorrt_llm/serve/scripts/benchmark_serving.py +++ b/tensorrt_llm/serve/scripts/benchmark_serving.py @@ -737,6 +737,30 @@ async def fetch_energy_metrics(base_url: str) -> Optional[dict]: return None +async def is_server_spec_dec_enabled(base_url: str) -> bool: + """Return True iff /v1/models advertises an enabled speculative_config. + + Any network/parse failure is swallowed and reported as False so callers + can fall back to safe defaults. + """ + models_url = f"{base_url}/v1/models" + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + try: + async with session.get(models_url) as response: + if response.status != 200: + return False + payload = await response.json() + except Exception: + return False + try: + spec_enabled = (payload["data"][0]["metadata"]["speculative_config"] + ["spec_enabled"]) + except (KeyError, IndexError, TypeError): + return False + return str(spec_enabled).strip().lower() == "true" + + def _snapshot_perf_metrics(output_dir: str) -> dict[Path, int]: directory = Path(output_dir) if not directory.exists(): @@ -991,6 +1015,15 @@ def create_dataset_and_sample(dataset_name: str): if "temperature" not in sampling_params: sampling_params["temperature"] = 0.0 # Default to greedy decoding. + # Auto-resolve --ignore-eos when the user did not pass an explicit flag. + # Polarity matches trtllm-bench latency (low_latency.py:332): when the server + # advertises speculative_config.spec_enabled, requests honor EOS so the + # acceptance rate is measured accurately; otherwise EOS is ignored for + # stable throughput. A server that does not expose speculative_config (older + # build, or non-spec-dec model) is treated as "no spec-dec" → ignore_eos=True. + if args.ignore_eos is None: + args.ignore_eos = not asyncio.run(is_server_spec_dec_enabled(base_url)) + # Avoid GC - reduce pause times. gc.disable() @@ -1335,9 +1368,16 @@ def create_dataset_and_sample(dataset_name: str): ) parser.add_argument( "--ignore-eos", - action="store_true", - help="Set ignore_eos flag when sending the benchmark request." - "Warning: ignore_eos is not supported in deepspeed_mii and tgi.") + action=argparse.BooleanOptionalAction, + default=None, + help="Whether to ignore EOS when sending the benchmark request. " + "If unset, the default is auto-resolved from the server's " + "/v1/models speculative_config metadata: when speculative decoding " + "is enabled the request honors EOS (so the acceptance rate is " + "measured cleanly), otherwise EOS is ignored for stable throughput " + "(mirrors trtllm-bench latency). Pass --ignore-eos / --no-ignore-eos " + "to override. Warning: ignore_eos is not supported in deepspeed_mii " + "and tgi.") parser.add_argument( "--percentile-metrics", type=str,