Update TensorRT-LLM backend (triton-inference-server#290)

* Update TensorRT-LLM backend

Author: kaiyux

README.md

@@ -218,7 +218,7 @@ The following table shows the fields that may to be modified before deployment:
 | `max_beam_width` | Optional (default=1). The maximum beam width that any request may ask for when using beam search.|
 | `max_tokens_in_paged_kv_cache` | Optional (default=unspecified). The maximum size of the KV cache in number of tokens. If unspecified, value is interpreted as 'infinite'. KV cache allocation is the min of max_tokens_in_paged_kv_cache and value derived from kv_cache_free_gpu_mem_fraction below. |
 | `max_attention_window_size` | Optional (default=max_sequence_length). When using techniques like sliding window attention, the maximum number of tokens that are attended to generate one token. Defaults attends to all tokens in sequence. |
-| `kv_cache_free_gpu_mem_fraction` | Optional (default=0.85). Set to a number between 0 and 1 to indicate the maximum fraction of GPU memory (after loading the model) that may be used for KV cache.|
+| `kv_cache_free_gpu_mem_fraction` | Optional (default=0.9). Set to a number between 0 and 1 to indicate the maximum fraction of GPU memory (after loading the model) that may be used for KV cache.|
 | `max_num_sequences` | Optional (default=`max_batch_size` if `enable_trt_overlap` is `false` and to `2 * max_batch_size` if `enable_trt_overlap` is `true`, where `max_batch_size` is the TRT engine maximum batch size). Maximum number of sequences that the in-flight batching scheme can maintain state for.
 | `enable_trt_overlap` | Optional (default=`true`). Set to `true` to partition available requests into 2 'microbatches' that can be run concurrently to hide exposed CPU runtime |
 | `exclude_input_in_output` | Optional (default=`false`). Set to `true` to only return completion tokens in a response. Set to `false` to return the prompt tokens concatenated with the generated tokens  |
@@ -346,6 +346,7 @@ He was a member of the French Academy of Sciences and the French Academy of Arts
 Soyer was a member of the French Academy of Sciences and
 ```
 
+#### Early stopping
 You can also stop the generation process early by using the `--stop-after-ms`
 option to send a stop request after a few milliseconds:
 
@@ -357,6 +358,54 @@ You will find that the generation process is stopped early and therefore the
 number of generated tokens is lower than 200. You can have a look at the
 client code to see how early stopping is achieved.
 
+#### Return context logits and/or generation logits
+If you want to get context logits and/or generation logits, you need to enable `--gather_context_logits` and/or `--gather_generation_logits` when building the engine (or `--enable gather_all_token_logits` to enable both at the same time). For more setting details about these two flags, please refer to [build.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/gpt/build.py) or [gpt_runtime](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/gpt_runtime.md).
+
+After launching the server, you could get the output of logits by passing the corresponding parameters `--return-context-logits` and/or `--return-generation-logits` in the client scripts (`end_to_end_grpc_client.py` and `inflight_batcher_llm_client.py`). For example:
+```bash
+python3 inflight_batcher_llm/client/inflight_batcher_llm_client.py --request-output-len 20 --tokenizer-dir /path/to/tokenizer/ \
+--return-context-logits \
+--return-generation-logits
+```
+
+The result should be similar to the following:
+```
+Input sequence:  [28524, 287, 5093, 12, 23316, 4881, 11, 30022, 263, 8776, 355, 257]
+Got completed request
+Input: Born in north-east France, Soyer trained as a
+Output beam 0:  has since worked in restaurants in London,
+Output sequence:  [21221, 878, 3867, 284, 3576, 287, 262, 1903, 6303, 82, 13, 679, 468, 1201, 3111, 287, 10808, 287, 3576, 11]
+context_logits.shape: (1, 12, 50257)
+context_logits: [[[ -65.9822     -62.267445   -70.08991   ...  -76.16964    -78.8893
+    -65.90678  ]
+  [-103.40278   -102.55243   -106.119026  ... -108.925415  -109.408585
+   -101.37687  ]
+  [ -63.971176   -64.03466    -67.58809   ...  -72.141235   -71.16892
+    -64.23846  ]
+  ...
+  [ -80.776375   -79.1815     -85.50916   ...  -87.07368    -88.02817
+    -79.28435  ]
+  [ -10.551408    -7.786484   -14.524468  ...  -13.805856   -15.767286
+     -7.9322424]
+  [-106.33096   -105.58956   -111.44852   ... -111.04858   -111.994194
+   -105.40376  ]]]
+generation_logits.shape: (1, 1, 20, 50257)
+generation_logits: [[[[-106.33096  -105.58956  -111.44852  ... -111.04858  -111.994194
+    -105.40376 ]
+   [ -77.867424  -76.96638   -83.119095 ...  -87.82542   -88.53957
+     -75.64877 ]
+   [-136.92282  -135.02484  -140.96051  ... -141.78284  -141.55045
+    -136.01668 ]
+   ...
+   [-100.03721   -98.98237  -105.25507  ... -108.49254  -109.45882
+     -98.95136 ]
+   [-136.78777  -136.16165  -139.13437  ... -142.21495  -143.57468
+    -134.94667 ]
+   [  19.222942   19.127287   14.804495 ...   10.556551    9.685863
+      19.625107]]]]
+```
+
+
 ### Launch Triton server *within Slurm based clusters*
 
 #### Prepare some scripts

all_models/gpt/ensemble/config.pbtxt

@@ -76,6 +76,12 @@ input [
     dims: [ 1 ]
     optional: true
   },
+  {
+    name: "frequency_penalty"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    optional: true
+  },
   {
     name: "random_seed"
     data_type: TYPE_UINT64
@@ -187,6 +193,10 @@ ensemble_scheduling {
           key: "presence_penalty"
           value: "presence_penalty"
       }
+      input_map {
+          key: "frequency_penalty"
+          value: "frequency_penalty"
+      }
       input_map {
           key: "random_seed"
           value: "random_seed"

all_models/gpt/tensorrt_llm/1/model.py

@@ -173,6 +173,8 @@ def execute(self, requests):
                     request, 'min_length')
                 inputs['presence_penalty'] = get_input_scalar_by_name(
                     request, 'presence_penalty')
+                inputs['frequency_penalty'] = get_input_scalar_by_name(
+                    request, 'frequency_penalty')
                 inputs['random_seed'] = get_input_scalar_by_name(
                     request, 'random_seed')
                 inputs['output_log_probs'] = get_input_scalar_by_name(
@@ -203,6 +205,8 @@ def execute(self, requests):
                 sampling_config.min_length = inputs['min_length']
             if inputs['presence_penalty'] is not None:
                 sampling_config.presence_penalty = inputs['presence_penalty']
+            if inputs['frequency_penalty'] is not None:
+                sampling_config.frequency_penalty = inputs['frequency_penalty']
             sampling_config.random_seed = inputs['random_seed']
             sampling_config.output_log_probs = inputs['output_log_probs']
             if self.remove_input_padding:

all_models/gpt/tensorrt_llm/config.pbtxt

@@ -92,6 +92,13 @@ input [
     reshape: { shape: [ ] }
     optional: true
   },
+  {
+    name: "frequency_penalty"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    reshape: { shape: [ ] }
+    optional: true
+  },
   {
     name: "random_seed"
     data_type: TYPE_UINT64

all_models/inflight_batcher_llm/ensemble/config.pbtxt

@@ -104,6 +104,12 @@ input [
     dims: [ 1 ]
     optional: true
   },
+  {
+    name: "frequency_penalty"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    optional: true
+  },
   {
     name: "random_seed"
     data_type: TYPE_UINT64
@@ -116,6 +122,18 @@ input [
     dims: [ 1 ]
     optional: true
   },
+  {
+    name: "return_context_logits"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "return_generation_logits"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
   {
     name: "beam_width"
     data_type: TYPE_INT32
@@ -168,6 +186,16 @@ output [
     name: "output_log_probs"
     data_type: TYPE_FP32
     dims: [ -1, -1 ]
+  },
+  {
+    name: "context_logits"
+    data_type: TYPE_FP32
+    dims: [ -1, -1 ]
+  },
+  {
+    name: "generation_logits"
+    data_type: TYPE_FP32
+    dims: [ -1, -1, -1 ]
   }
 ]
 ensemble_scheduling {
@@ -199,6 +227,14 @@ ensemble_scheduling {
         key: "EMBEDDING_BIAS_WEIGHTS"
         value: "embedding_bias_weights"
       }
+      input_map {
+        key: "END_ID"
+        value: "end_id"
+      }
+      input_map {
+        key: "PAD_ID"
+        value: "pad_id"
+      }
       output_map {
         key: "REQUEST_INPUT_LEN"
         value: "_REQUEST_INPUT_LEN"
@@ -223,6 +259,14 @@ ensemble_scheduling {
         key: "EMBEDDING_BIAS"
         value: "_EMBEDDING_BIAS"
       }
+      output_map {
+        key: "OUT_END_ID"
+        value: "_PREPROCESSOR_END_ID"
+      }
+      output_map {
+        key: "OUT_PAD_ID"
+        value: "_PREPROCESSOR_PAD_ID"
+      }
     },
     {
       model_name: "tensorrt_llm"
@@ -241,11 +285,11 @@ ensemble_scheduling {
       }
       input_map {
           key: "end_id"
-          value: "end_id"
+          value: "_PREPROCESSOR_END_ID"
       }
       input_map {
           key: "pad_id"
-          value: "pad_id"
+          value: "_PREPROCESSOR_PAD_ID"
       }
       input_map {
           key: "embedding_bias"
@@ -279,6 +323,10 @@ ensemble_scheduling {
           key: "presence_penalty"
           value: "presence_penalty"
       }
+      input_map {
+          key: "frequency_penalty"
+          value: "frequency_penalty"
+      }
       input_map {
           key: "random_seed"
           value: "random_seed"
@@ -287,6 +335,14 @@ ensemble_scheduling {
           key: "return_log_probs"
           value: "return_log_probs"
       }
+      input_map {
+          key: "return_context_logits"
+          value: "return_context_logits"
+      }
+      input_map {
+          key: "return_generation_logits"
+          value: "return_generation_logits"
+      }
       input_map {
           key: "beam_width"
           value: "beam_width"
@@ -326,6 +382,14 @@ ensemble_scheduling {
       output_map {
         key: "output_log_probs"
         value: "_OUTPUT_LOG_PROBS"
+      },
+      output_map {
+        key: "context_logits"
+        value: "_CONTEXT_LOGITS"
+      },
+      output_map {
+        key: "generation_logits"
+        value: "_GENERATION_LOGITS"
       }
     },
     {
@@ -343,6 +407,14 @@ ensemble_scheduling {
         key: "OUTPUT_LOG_PROBS"
         value: "_OUTPUT_LOG_PROBS"
       }
+      input_map {
+        key: "CONTEXT_LOGITS"
+        value: "_CONTEXT_LOGITS"
+      }
+      input_map {
+        key: "GENERATION_LOGITS"
+        value: "_GENERATION_LOGITS"
+      }
       input_map {
         key: "SEQUENCE_LENGTH"
         value: "_SEQUENCE_LENGTH"
@@ -359,6 +431,14 @@ ensemble_scheduling {
         key: "OUT_CUM_LOG_PROBS"
         value: "cum_log_probs"
       }
+      output_map {
+        key: "OUT_CONTEXT_LOGITS"
+        value: "context_logits"
+      }
+      output_map {
+        key: "OUT_GENERATION_LOGITS"
+        value: "generation_logits"
+      }
     }
   ]
 }

all_models/inflight_batcher_llm/postprocessing/1/model.py

@@ -126,6 +126,14 @@ def execute(self, requests):
             output_log_probs = pb_utils.get_input_tensor_by_name(
                 request, 'OUTPUT_LOG_PROBS').as_numpy()
 
+            # Get context logits
+            context_logits = pb_utils.get_input_tensor_by_name(
+                request, 'CONTEXT_LOGITS').as_numpy()
+
+            # Get generation logits
+            generation_logits = pb_utils.get_input_tensor_by_name(
+                request, 'GENERATION_LOGITS').as_numpy()
+
             # Reshape Input
             # tokens_batch = tokens_batch.reshape([-1, tokens_batch.shape[0]])
             # tokens_batch = tokens_batch.T
@@ -145,6 +153,12 @@ def execute(self, requests):
             out_output_log_probs = pb_utils.Tensor('OUT_OUTPUT_LOG_PROBS',
                                                    output_log_probs)
 
+            out_context_logits = pb_utils.Tensor('OUT_CONTEXT_LOGITS',
+                                                 context_logits)
+
+            out_generation_logits = pb_utils.Tensor('OUT_GENERATION_LOGITS',
+                                                    generation_logits)
+
             # Create InferenceResponse. You can set an error here in case
             # there was a problem with handling this inference request.
             # Below is an example of how you can set errors in inference
@@ -153,7 +167,8 @@ def execute(self, requests):
             # pb_utils.InferenceResponse(
             #    output_tensors=..., TritonError("An error occurred"))
             inference_response = pb_utils.InferenceResponse(output_tensors=[
-                output_tensor, out_cum_log_probs, out_output_log_probs
+                output_tensor, out_cum_log_probs, out_output_log_probs,
+                out_context_logits, out_generation_logits
             ])
             responses.append(inference_response)
 

all_models/inflight_batcher_llm/postprocessing/config.pbtxt

@@ -47,6 +47,18 @@ input [
     name: "OUTPUT_LOG_PROBS"
     data_type: TYPE_FP32
     dims: [ -1, -1 ]
+  },
+  {
+    name: "CONTEXT_LOGITS"
+    data_type: TYPE_FP32
+    dims: [ -1, -1 ]
+    optional: true
+  },
+  {
+    name: "GENERATION_LOGITS"
+    data_type: TYPE_FP32
+    dims: [ -1, -1, -1 ]
+    optional: true
   }
 ]
 output [
@@ -64,6 +76,16 @@ output [
     name: "OUT_OUTPUT_LOG_PROBS"
     data_type: TYPE_FP32
     dims: [ -1, -1 ]
+  },
+  {
+    name: "OUT_CONTEXT_LOGITS"
+    data_type: TYPE_FP32
+    dims: [ -1, -1 ]
+  },
+  {
+    name: "OUT_GENERATION_LOGITS"
+    data_type: TYPE_FP32
+    dims: [ -1, -1, -1 ]
   }
 ]