Add num_input_tokens additional output

Author: kthui

ci/L0_additional_outputs_vllm/additional_outputs_test.py

@@ -44,6 +44,7 @@ def _get_inputs(
         sampling_parameters=None,
         return_finish_reason=None,
         return_cumulative_logprob=None,
+        return_num_input_tokens=None,
         return_num_output_tokens=None,
     ):
         inputs = []
@@ -76,6 +77,12 @@ def _get_inputs(
                 np.array([return_cumulative_logprob], dtype=bool)
             )
 
+        if return_num_input_tokens is not None:
+            inputs.append(grpcclient.InferInput("return_num_input_tokens", [1], "BOOL"))
+            inputs[-1].set_data_from_numpy(
+                np.array([return_num_input_tokens], dtype=bool)
+            )
+
         if return_num_output_tokens is not None:
             inputs.append(
                 grpcclient.InferInput("return_num_output_tokens", [1], "BOOL")
@@ -135,6 +142,18 @@ def _assert_cumulative_logprob(self, return_cumulative_logprob):
             assert cumulative_logprob != prev_cumulative_logprob
             prev_cumulative_logprob = cumulative_logprob
 
+    def _assert_num_input_tokens(self, return_num_input_tokens):
+        for response in self._responses:
+            result, error = response["result"], response["error"]
+            assert error is None
+            num_input_tokens_np = result.as_numpy(name="num_input_tokens")
+            if return_num_input_tokens is None or return_num_input_tokens == False:
+                assert num_input_tokens_np is None
+                continue
+            num_input_tokens = num_input_tokens_np.astype(int)
+            assert num_input_tokens > 0
+            assert num_input_tokens <= len(self._prompt)
+
     def _assert_num_output_tokens(self, return_num_output_tokens):
         for response in self._responses:
             result, error = response["result"], response["error"]
@@ -166,12 +185,14 @@ def _assert_num_output_tokens(self, return_num_output_tokens):
     @pytest.mark.parametrize("stream", [True, False])
     @pytest.mark.parametrize("return_finish_reason", [None, True, False])
     @pytest.mark.parametrize("return_cumulative_logprob", [None, True, False])
+    @pytest.mark.parametrize("return_num_input_tokens", [None, True, False])
     @pytest.mark.parametrize("return_num_output_tokens", [None, True, False])
     def test_additional_outputs(
         self,
         stream,
         return_finish_reason,
         return_cumulative_logprob,
+        return_num_input_tokens,
         return_num_output_tokens,
     ):
         inputs = self._get_inputs(
@@ -180,10 +201,12 @@ def test_additional_outputs(
             sampling_parameters=self._sampling_parameters,
             return_finish_reason=return_finish_reason,
             return_cumulative_logprob=return_cumulative_logprob,
+            return_num_input_tokens=return_num_input_tokens,
             return_num_output_tokens=return_num_output_tokens,
         )
         self._llm_infer(inputs)
         self._assert_text_output_valid()
         self._assert_finish_reason(return_finish_reason)
         self._assert_cumulative_logprob(return_cumulative_logprob)
+        self._assert_num_input_tokens(return_num_input_tokens)
         self._assert_num_output_tokens(return_num_output_tokens)

ci/L0_additional_outputs_vllm/test.sh

@@ -40,7 +40,7 @@ sed -i 's/"gpu_memory_utilization": 0.5/"gpu_memory_utilization": 0.3/' models/v
 RET=0
 
 # Test
-SERVER_LOG="vllm_opt.server.log"
+SERVER_LOG="additional_outputs_test.server.log"
 SERVER_ARGS="--model-repository=models"
 run_server
 if [ "$SERVER_PID" == "0" ]; then

src/model.py

@@ -104,6 +104,12 @@ def _auto_complete_inputs_and_outputs(auto_complete_model_config):
                 "dims": [1],
                 "optional": True,
             },
+            {
+                "name": "return_num_input_tokens",
+                "data_type": "TYPE_BOOL",
+                "dims": [1],
+                "optional": True,
+            },
             {
                 "name": "return_num_output_tokens",
                 "data_type": "TYPE_BOOL",
@@ -125,6 +131,7 @@ def _auto_complete_inputs_and_outputs(auto_complete_model_config):
             {"name": "text_output", "data_type": "TYPE_STRING", "dims": [-1]},
             {"name": "finish_reason", "data_type": "TYPE_STRING", "dims": [-1]},
             {"name": "cumulative_logprob", "data_type": "TYPE_FP32", "dims": [-1]},
+            {"name": "num_input_tokens", "data_type": "TYPE_UINT32", "dims": [1]},
             {"name": "num_output_tokens", "data_type": "TYPE_UINT32", "dims": [-1]},
         ]
 
@@ -377,10 +384,11 @@ def _get_input_tensors(self, request):
         else:
             parameters = request.parameters()
 
-        # return_finish_reason, return_cumulative_logprob, return_num_output_tokens
+        # additional outputs
         additional_outputs = {
             "return_finish_reason": None,
             "return_cumulative_logprob": None,
+            "return_num_input_tokens": None,
             "return_num_output_tokens": None,
         }
         for tensor_name in additional_outputs.keys():
@@ -496,6 +504,15 @@ def _create_response(
                 )
             )
 
+        # num_input_tokens
+        if additional_outputs["return_num_input_tokens"]:
+            num_input_tokens = len(request_output.prompt_token_ids)
+            output_tensors.append(
+                pb_utils.Tensor(
+                    "num_input_tokens", np.asarray(num_input_tokens, dtype=np.uint32)
+                )
+            )
+
         # num_output_tokens
         if additional_outputs["return_num_output_tokens"]:
             if prev_request_output is None: