[TRTLLM-14511][feat] refactor per-request perf metrics for multi-process trtllm-server

examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh

@@ -92,11 +92,5 @@ for concurrency in ${concurrency_list}; do
     echo "Benchmark with concurrency ${concurrency} done"
 done
 
-# Fetch perf metrics from disagg server
-echo "Fetching perf metrics from http://${hostname}:${port}/perf_metrics ..."
-curl -s "http://${hostname}:${port}/perf_metrics" > ${log_path}/perf_metrics.json 2>&1 || true
-if [ -s "${log_path}/perf_metrics.json" ]; then
-    echo "Perf metrics saved to ${log_path}/perf_metrics.json"
-else
-    echo "Warning: perf_metrics response was empty or endpoint not available"
-fi
+# Configure perf_metrics_output_dir on each server to persist per-request JSONL.
+echo "Per-request metrics are available in the configured server-side JSONL output."

tensorrt_llm/llmapi/disagg_utils.py

@@ -92,6 +92,8 @@ class DisaggServerConfig():
     otlp_config: Optional[OtlpConfig] = None
     max_retries: int = 1
     perf_metrics_max_requests: int = 0
+    return_perf_metrics: bool = False
+    perf_metrics_output_dir: Optional[str] = None
     disagg_cluster_config: Optional[DisaggClusterConfig] = None
     node_id: int = uuid.getnode(
     ) % 256  # Assuming only one disagg-server is running on a machine, modulo 256.
@@ -189,6 +191,8 @@ def extract_disagg_cfg(hostname: str = 'localhost',
                        port: int = 8000,
                        max_retries: int = 1,
                        perf_metrics_max_requests: int = 0,
+                       return_perf_metrics: bool = False,
+                       perf_metrics_output_dir: Optional[str] = None,
                        context_servers: Optional[dict] = None,
                        generation_servers: Optional[dict] = None,
                        conditional_disagg_config: Optional[dict] = None,
@@ -207,10 +211,12 @@ def extract_disagg_cfg(hostname: str = 'localhost',
     context_servers = context_servers or {}
     generation_servers = generation_servers or {}
 
+    inherited_args = dict(kwargs)
+
     # If parameters are specified outside the context_severs and generation_servers sections,
     # make sure they match
     # Also inherit the values from the top-level
-    for key, value in kwargs.items():
+    for key, value in inherited_args.items():
         for server_type, servers in [("context_servers", context_servers),
                                      ("generation_servers", generation_servers)
                                      ]:
@@ -241,11 +247,19 @@ def extract_disagg_cfg(hostname: str = 'localhost',
 
     otlp_config = OtlpConfig(**otlp_config) if otlp_config else None
 
-    config = DisaggServerConfig(server_configs, hostname, port,
-                                ctx_router_config, gen_router_config,
-                                conditional_disagg_config, otlp_config,
-                                max_retries, perf_metrics_max_requests,
-                                disagg_cluster_config)
+    config = DisaggServerConfig(
+        server_configs=server_configs,
+        hostname=hostname,
+        port=port,
+        ctx_router_config=ctx_router_config,
+        gen_router_config=gen_router_config,
+        conditional_disagg_config=conditional_disagg_config,
+        otlp_config=otlp_config,
+        max_retries=max_retries,
+        perf_metrics_max_requests=perf_metrics_max_requests,
+        return_perf_metrics=return_perf_metrics,
+        perf_metrics_output_dir=perf_metrics_output_dir,
+        disagg_cluster_config=disagg_cluster_config)
     if node_id is not None:
         node_id_space = 1 << DISAGG_NODE_ID_BITS
         if not 0 <= node_id < node_id_space:

tensorrt_llm/llmapi/llm_args.py

@@ -4390,15 +4390,27 @@ class BaseLlmArgs(StrictBaseModel):
         status="deprecated",
         telemetry=TelemetryField.categorical('pytorch', '_autodeploy'))
 
-    return_perf_metrics: bool = Field(default=False,
-                                      description="Return perf metrics.",
-                                      status="prototype")
+    return_perf_metrics: bool = Field(
+        default=False,
+        description=
+        "Allow serving responses to include per-request performance metrics when "
+        "the request sets X-TRTLLM-return-metrics: 1.",
+        status="prototype")
+
+    perf_metrics_output_dir: Optional[str] = Field(
+        default=None,
+        description="Directory for per-process performance metrics JSONL "
+        "files. Setting this enables collection even when "
+        "return_perf_metrics is false.",
+        status="prototype",
+        telemetry=False)
 
     perf_metrics_max_requests: NonNegativeInt = Field(
         default=0,
         description=
-        "The maximum number of requests for perf metrics. Must also set return_perf_metrics to true to get perf metrics.",
-        status="prototype")
+        "Deprecated compatibility field. Completed per-request metrics are no "
+        "longer retained in memory.",
+        status="deprecated")
 
     prometheus_metrics_config: Optional[PrometheusMetricsConfig] = Field(
         default=None,

tensorrt_llm/scaffolding/task.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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 __future__ import annotations
 
 import json
@@ -116,8 +130,8 @@ class GenerationTask(Task):
 
     # Server-side request id captured from the OpenAI client's streaming
     # chunks (every chunk's ``chunk.id`` field).  Trace-replay clients use
-    # this to look up the matching record in trtllm-serve's
-    # ``/perf_metrics`` drain and attach per-request KV-cache statistics
+    # this to correlate response-carried per-request metrics and attach
+    # KV-cache statistics
     # (``num_reused_blocks`` / ``num_missed_blocks`` / ``free_num_blocks``)
     # to the per-LLM-call row in the step JSON.  ``None`` until the worker
     # observes the first chunk (or when not running against an

tensorrt_llm/scaffolding/trace_replay/replay.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
 import asyncio
 import hashlib
 import random
@@ -256,7 +270,7 @@ class DropPathStats:
     POSTs to the server, one immediately after — for every conv id this
     branch owns.  Each probe's ``request_id`` (server-assigned, captured
     by the worker from the OpenAI streaming chunk) is recorded here and
-    later joined with trtllm-serve's ``/perf_metrics`` drain to recover
+    later joined with response-carried metrics to recover
     ``num_reused_blocks`` (the per-request KV-cache hit count) and
     ``free_num_blocks`` (the post-call snapshot of the free-block pool).
 
@@ -399,7 +413,7 @@ async def _run(self):
                     if event is None:  # sentinel
                         # Fire end-of-branch retention probes BEFORE returning,
                         # so the parent ``wait_all_done`` covers them and the
-                        # client's downstream /perf_metrics drain sees the
+                        # client receives the response-carried metrics for the
                         # probe records.  For a child branch this is at
                         # parallel_end (immediately after any drop_kv_cache);
                         # for the root branch it is at end-of-session.
@@ -503,7 +517,7 @@ async def _handle_drop_kv_cache(self, event: TraceEvent):
         # ``max_tokens=1, ignore_eos=True`` probe per (conv_id, phase) pair
         # immediately before and after the truncate.  Each probe's
         # request_id is recorded here; the per-request KV-cache hit
-        # accounting is joined in later via /perf_metrics.  The probes
+        # accounting is joined later from response-carried metrics. The probes
         # serialize the truncate against a real before/after measurement
         # so downstream verification can mechanically prove the truncate
         # actually freed the blocks the engine claimed it did (P0.6.a).

tensorrt_llm/scaffolding/worker.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
 import asyncio
 import copy
 import json
@@ -341,8 +355,8 @@ async def generation_handler(self, task: GenerationTask) -> TaskStatus:
                 now = time.perf_counter()
                 # Every chunk carries the server-assigned request id. Capture
                 # it once (the first non-None value) so callers can later
-                # correlate per-request perf metrics drained from
-                # ``/perf_metrics`` with the GenerationTask that issued them.
+                # correlate response-carried per-request metrics with the
+                # GenerationTask that issued them.
                 if request_id is None:
                     cid = getattr(chunk, "id", None)
                     if cid is not None:

tensorrt_llm/serve/_perf_metrics_schema.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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 List, Optional, Union
+
+from typing_extensions import NotRequired, TypedDict
+
+
+class StepMetrics(TypedDict):
+    forward_start_time: float
+    forward_end_time: float
+    sample_start_time: float
+    sample_end_time: float
+    gpu_forward_time: float
+    gpu_sample_time: float
+    token_time: NotRequired[float]
+    scheduled_time: NotRequired[float]
+    prev_batch_token_time: NotRequired[float]
+    iter: NotRequired[int]
+
+
+class TimeBreakdownMetrics(TypedDict):
+    step_metrics: NotRequired[List[StepMetrics]]
+    ctx_chunk_metrics: NotRequired[List[StepMetrics]]
+    ctx_gpu_forward_time: NotRequired[float]
+    ctx_gpu_sample_time: NotRequired[float]
+
+
+class TimingMetrics(TypedDict):
+    arrival_time: Optional[float]
+    first_scheduled_time: NotRequired[Optional[float]]
+    first_token_time: NotRequired[Optional[float]]
+    last_token_time: Optional[float]
+    server_arrival_time: NotRequired[Optional[float]]
+    server_first_token_time: NotRequired[Optional[float]]
+    kv_cache_size: NotRequired[int]
+    kv_cache_transfer_start: NotRequired[Optional[float]]
+    kv_cache_transfer_end: NotRequired[Optional[float]]
+
+
+class KvCacheMetrics(TypedDict):
+    num_total_allocated_blocks: int
+    num_new_allocated_blocks: int
+    num_reused_blocks: int
+    num_missed_blocks: int
+
+
+class SpeculativeDecodingMetrics(TypedDict):
+    acceptance_rate: float
+    total_accepted_draft_tokens: int
+    total_draft_tokens: int
+
+
+class PerfMetrics(TypedDict):
+    timing_metrics: TimingMetrics
+    first_iter: NotRequired[int]
+    last_iter: NotRequired[int]
+    kv_cache_metrics: NotRequired[KvCacheMetrics]
+    speculative_decoding: NotRequired[SpeculativeDecodingMetrics]
+
+
+class WorkerPerfMetrics(TypedDict):
+    request_id: Union[int, str]
+    perf_metrics: PerfMetrics
+    ctx_request_id: NotRequired[int]
+    time_breakdown_metrics: NotRequired[TimeBreakdownMetrics]
+
+
+class WorkerPerfMetricsRecord(WorkerPerfMetrics):
+    status: str
+    disagg_request_id: NotRequired[int]
+
+
+class DisaggPerfMetricsRecord(TypedDict):
+    ctx_server: str
+    gen_server: str
+    disagg_server_arrival_time: float
+    disagg_ctx_dispatch_time: Optional[float]
+    disagg_server_first_token_time: Optional[float]
+    status: str
+    disagg_request_id: NotRequired[int]
+    ctx_perf_metrics: NotRequired[WorkerPerfMetrics]
+    gen_perf_metrics: NotRequired[WorkerPerfMetrics]
+
+
+PerfMetricsRecord = Union[WorkerPerfMetricsRecord, DisaggPerfMetricsRecord]