[TRTLLM-6104] feat: add request_perf_metrics to triton LLMAPI backend (#5554)

Signed-off-by: Vivian Chen <[email protected]>

Author: xuanzic

triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py

@@ -24,15 +24,24 @@ def get_sampling_params_from_request(request, batch_size=1, batch_index=0):
     Used in llmapi/tensorrt_llm
     """
     sampling_params_args = [
-        'best_of', 'temperature', 'top_k', 'top_p', 'frequency_penalty',
-        'presence_penalty', 'max_tokens', 'seed', 'exclude_input_from_output'
+        'best_of',
+        'temperature',
+        'top_k',
+        'top_p',
+        'frequency_penalty',
+        'presence_penalty',
+        'max_tokens',
+        'seed',
+        'exclude_input_from_output',
+        'return_perf_metrics',
     ]
     param_mappings = {}
     for arg in sampling_params_args:
         param_mappings[f"sampling_param_{arg}"] = arg
     default_values = {
         'sampling_param_best_of': 1,
         'sampling_param_exclude_input_from_output': False,
+        'sampling_param_return_perf_metrics': False,
     }
     kwargs = convert_request_input_to_dict(request, param_mappings,
                                            default_values, batch_size,

triton_backend/all_models/llmapi/tensorrt_llm/1/model.py

@@ -33,6 +33,7 @@
 from contextlib import asynccontextmanager
 
 import numpy as np
+import pandas as pd
 import triton_python_backend_utils as pb_utils
 import yaml
 from helpers import (get_input_tensor_by_name, get_output_config_from_request,
@@ -450,6 +451,92 @@ def _create_response(self, request_output, output_config):
                     pb_utils.Tensor(output_name,
                                     np.asarray(tensor_data, dtype=np.object_)))
 
+        if hasattr(request_output.outputs[0], 'request_perf_metrics'
+                   ) and request_output.outputs[0].request_perf_metrics:
+
+            perf_metrics = request_output.outputs[0].request_perf_metrics
+
+            # kv cache perf metrics per request
+            kv_metrics = perf_metrics.kv_cache_metrics
+
+            response.append(
+                pb_utils.Tensor(
+                    "kv_cache_reused_block",
+                    np.asarray([kv_metrics.num_reused_blocks],
+                               dtype=self.output_dtype)))
+            response.append(
+                pb_utils.Tensor(
+                    "kv_cache_hit_rate",
+                    np.asarray([kv_metrics.kv_cache_hit_rate],
+                               dtype=self.output_dtype)))
+            response.append(
+                pb_utils.Tensor(
+                    "kv_cache_alloc_new_blocks",
+                    np.asarray([kv_metrics.num_new_allocated_blocks],
+                               dtype=self.output_dtype)))
+            response.append(
+                pb_utils.Tensor(
+                    "kv_cache_alloc_total_blocks",
+                    np.asarray([kv_metrics.num_total_allocated_blocks],
+                               dtype=self.output_dtype)))
+            response.append(
+                pb_utils.Tensor(
+                    "kv_cache_missed_block",
+                    np.asarray([kv_metrics.num_missed_blocks],
+                               dtype=self.output_dtype)))
+
+            # timing perf metrics per request
+            timing_metrics = perf_metrics.timing_metrics
+            response.append(
+                pb_utils.Tensor(
+                    "arrival_time_ns",
+                    np.asarray(
+                        [pd.Timedelta(timing_metrics.arrival_time).value],
+                        dtype=self.output_dtype)))
+
+            response.append(
+                pb_utils.Tensor(
+                    "first_scheduled_time_ns",
+                    np.asarray([
+                        pd.Timedelta(timing_metrics.first_scheduled_time).value
+                    ],
+                               dtype=self.output_dtype)))
+
+            response.append(
+                pb_utils.Tensor(
+                    "first_token_time_ns",
+                    np.asarray(
+                        [pd.Timedelta(timing_metrics.first_token_time).value],
+                        dtype=self.output_dtype)))
+
+            response.append(
+                pb_utils.Tensor(
+                    "last_token_time_ns",
+                    np.asarray(
+                        [pd.Timedelta(timing_metrics.last_token_time).value],
+                        dtype=self.output_dtype)))
+
+            #spec dec perf metrics per request
+            spec_dec_metrics = perf_metrics.speculative_decoding
+
+            response.append(
+                pb_utils.Tensor(
+                    "acceptance_rate",
+                    np.asarray([spec_dec_metrics.acceptance_rate],
+                               dtype=self.output_dtype)))
+
+            response.append(
+                pb_utils.Tensor(
+                    "total_accepted_draft_tokens",
+                    np.asarray([spec_dec_metrics.total_accepted_draft_tokens],
+                               dtype=self.output_dtype)))
+
+            response.append(
+                pb_utils.Tensor(
+                    "total_draft_tokens",
+                    np.asarray([spec_dec_metrics.total_draft_tokens],
+                               dtype=self.output_dtype)))
+
         return pb_utils.InferenceResponse(output_tensors=response)
 
     def finalize(self):

triton_backend/all_models/llmapi/tensorrt_llm/config.pbtxt

@@ -118,6 +118,12 @@ input [
     dims: [ 1 ]
     optional: true
   },
+  {
+    name: "sampling_param_return_perf_metrics"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
   ## Arguments for Controlling Response Output Fields ##
   {
     name: "return_finish_reason"
@@ -161,5 +167,65 @@ output [
     name: "cumulative_logprob"
     data_type: TYPE_FP32
     dims: [-1]
+  },
+  {
+    name: "kv_cache_reused_block"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "kv_cache_missed_block"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "kv_cache_alloc_new_blocks"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "kv_cache_alloc_total_blocks"
+    data_type: TYPE_INT32
+    dims: [-1]
+  },
+  {
+    name: "kv_cache_hit_rate"
+    data_type: TYPE_FP32
+    dims: [-1]
+  },
+  {
+    name: "arrival_time_ns"
+    data_type: TYPE_INT64
+    dims: [ 1 ]
+  },
+  {
+    name: "first_scheduled_time_ns"
+    data_type: TYPE_INT64
+    dims: [ 1 ]
+  },
+  {
+    name: "first_token_time_ns"
+    data_type: TYPE_INT64
+    dims: [ 1 ]
+  },
+  {
+    name: "last_token_time_ns"
+    data_type: TYPE_INT64
+    dims: [ 1 ]
+  },
+  {
+    name: "acceptance_rate"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+  },
+  {
+    name: "total_accepted_draft_tokens"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+  },
+  {
+    name: "total_draft_tokens"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
   }
 ]

triton_backend/all_models/tests/test_llmapi_python_backend.py

@@ -143,6 +143,7 @@ def inputs(streaming=False):
         "sampling_param_seed": [2],
         "return_finish_reason": [True],
         "return_stop_reason": [True],
+        "sampling_param_return_perf_metrics": [True]
     }
 
 
@@ -164,6 +165,7 @@ def test_get_sampling_params_from_request():
     assert config["frequency_penalty"] == 0.0
     assert config["presence_penalty"] == 0.0
     assert config["seed"] == 2
+    assert config["return_perf_metrics"] == True
     assert np.array_equal(config["stop"], np.array(['\n', 'stop']))