Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions tensorrt_llm/bench/benchmark/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
23 changes: 19 additions & 4 deletions tensorrt_llm/bench/benchmark/throughput.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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
}
Expand Down
1 change: 1 addition & 0 deletions tensorrt_llm/serve/openai_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
47 changes: 32 additions & 15 deletions tensorrt_llm/serve/openai_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
46 changes: 43 additions & 3 deletions tensorrt_llm/serve/scripts/benchmark_serving.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment on lines +747 to +750

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

rg -n -C 8 \
  'AIOHTTP_TIMEOUT|ClientTimeout|is_server_spec_dec_enabled' \
  . --glob '*.py'

Repository: NVIDIA/TensorRT-LLM

Length of output: 40033


🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
import ast

backend = Path("tensorrt_llm/serve/scripts/backend_request_func.py")
benchmark = Path("tensorrt_llm/serve/scripts/benchmark_serving.py")

backend_tree = ast.parse(backend.read_text())
benchmark_tree = ast.parse(benchmark.read_text())

timeout_values = []
for node in ast.walk(backend_tree):
    if isinstance(node, ast.Assign) and any(
        isinstance(target, ast.Name) and target.id == "AIOHTTP_TIMEOUT"
        for target in node.targets
    ):
        timeout_values.append(ast.unparse(node.value))

print("AIOHTTP_TIMEOUT assignments:", timeout_values)

for node in benchmark_tree.body:
    if isinstance(node, ast.AsyncFunctionDef) and node.name == "is_server_spec_dec_enabled":
        print(
            "is_server_spec_dec_enabled session timeout:",
            [
                ast.unparse(keyword.value)
                for stmt in ast.walk(node)
                if isinstance(stmt, ast.Call)
                and isinstance(stmt.func, ast.Attribute)
                and stmt.func.attr == "ClientSession"
                for keyword in stmt.keywords
                if keyword.arg == "timeout"
            ],
        )
        print(
            "is_server_spec_dec_enabled GET calls:",
            [
                ast.unparse(stmt)
                for stmt in ast.walk(node)
                if isinstance(stmt, ast.Call)
                and isinstance(stmt.func, ast.Attribute)
                and stmt.func.attr == "get"
            ],
        )

for node in benchmark_tree.body:
    if isinstance(node, ast.AsyncFunctionDef) and node.name == "is_server_spec_dec_enabled":
        print("helper source lines:", node.lineno, node.end_lineno)

for node in ast.walk(benchmark_tree):
    if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
        if node.func.attr == "run" and any(
            isinstance(arg, ast.Call)
            and isinstance(arg.func, ast.Name)
            and arg.func.id == "is_server_spec_dec_enabled"
            for arg in node.args
        ):
            print("main invocation:", ast.unparse(node))
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 457


Use a dedicated short timeout for the metadata probe.

AIOHTTP_TIMEOUT is six hours and also covers /v1/models, so an unresponsive server can delay fallback for up to six hours.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/serve/scripts/benchmark_serving.py` around lines 747 - 750,
Update the metadata probe in the session/get flow around models_url to use a
dedicated short timeout instead of the six-hour AIOHTTP_TIMEOUT. Keep the
existing long timeout for benchmark requests, but ensure the /v1/models request
fails promptly so fallback is not delayed.

if response.status != 200:
return False
payload = await response.json()
except Exception:
return False
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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():
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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,
Expand Down
Loading