Avoid redundant vLLM /tokenize calls using native prompt token IDs

responses_api_models/vllm_model/app.py

@@ -470,11 +470,7 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict:
                 top_logprobs=0,
                 # Typically passed via OpenAI client extra_body.
                 return_tokens_as_token_ids=True,
-                # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids
-                # For prompt and generation token IDs
-                # return_token_ids=True,
-                # For prompt token IDs
-                # prompt_logprobs=0,
+                return_token_ids=True,
             )
 
         if self.config.uses_reasoning_parser and not self.config.preserve_reasoning_in_assistant_content:
@@ -744,48 +740,32 @@ async def chat_completions(
                 )
             log_probs = logprobs_block["content"]
             generation_log_probs = [log_prob["logprob"] for log_prob in log_probs]
-
-            """
-            START TODO remove this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids
-            """
-            # Looks like `"token_id:151667"`
             generation_token_ids = [log_prob["token"].removeprefix("token_id:") for log_prob in log_probs]
 
-            # The tokenize endpoint doesn't accept any sampling parameters
-            # The only relevant params are model, messages, and tools.
-            #
-            # IMPORTANT: pass through chat-template knobs (e.g. enable_thinking)
-            # when tokenizing, otherwise `prompt_token_ids` (and therefore logged
-            # `prompt_str`) can be built with different chat template settings than
-            # the actual generation request.
-            tokenize_body_dict = dict()
-            for key in ("model", "messages", "tools", "chat_template_kwargs"):
-                if key in body_dict:
-                    tokenize_body_dict[key] = body_dict[key]
-
-            # The base url has /v1 at the end but vLLM's tokenize endpoint does not have v1, hence the ..
-            tokenize_response = await client.create_tokenize(**tokenize_body_dict)
-            """
-            END
-            """
+            # Fast path: use the IDs from the original generation request so the
+            # prompt does not need to be rendered and tokenized a second time.
+            prompt_token_ids = chat_completion_dict.pop("prompt_token_ids", None)
+
+            # Older vLLM versions may accept return_token_ids without returning
+            # prompt_token_ids. Preserve /tokenize as a compatibility fallback.
+            if prompt_token_ids is None:
+                tokenize_body_dict = {
+                    key: body_dict[key]
+                    for key in ("model", "messages", "tools", "chat_template_kwargs")
+                    if key in body_dict
+                }
+                tokenize_response = await client.create_tokenize(**tokenize_body_dict)
+                prompt_token_ids = tokenize_response["tokens"]
 
-            message_dict = choice_dict["message"]
-            message_dict.update(
-                dict(
-                    # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids
-                    # prompt_token_ids=chat_completion_dict["prompt_token_ids"],
-                    prompt_token_ids=tokenize_response["tokens"],
-                    # generation_token_ids=choice_dict["token_ids"],
-                    generation_token_ids=generation_token_ids,
-                    generation_log_probs=generation_log_probs,
-                )
+            choice_dict["message"].update(
+                prompt_token_ids=prompt_token_ids,
+                generation_token_ids=generation_token_ids,
+                generation_log_probs=generation_log_probs,
             )
 
-            # Clean the duplicated information
+            # Remove vLLM-specific duplicate fields before validating the response.
             choice_dict.pop("logprobs")
-            # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids
-            # chat_completion_dict.pop("prompt_token_ids")
-            # choice_dict.pop("token_ids")
+            choice_dict.pop("token_ids", None)
 
         return NeMoGymChatCompletion.model_validate(chat_completion_dict)
 

responses_api_models/vllm_model/tests/test_app.py

@@ -4597,6 +4597,7 @@ def test_capture_path_pins_top_logprobs_to_zero(self) -> None:
         assert result["top_logprobs"] == 0
         assert result["logprobs"] is True
         assert result["return_tokens_as_token_ids"] is True
+        assert result["return_token_ids"] is True
 
         # Inbound non-zero value must also be overridden, not inherited.
         result = model._preprocess_chat_completion_create_params(
@@ -4661,13 +4662,51 @@ def _capture_chat_completion_dict(
             ],
         }
 
+    @mark.parametrize("prompt_token_ids", [[10, 20, 30], []])
+    def test_capture_path_uses_native_prompt_token_ids(self, prompt_token_ids: list[int]) -> None:
+        """Native prompt IDs, including an empty list, must bypass /tokenize."""
+        model = _make_top_logprobs_model(return_token_id_information=True)
+        app = model.setup_webserver()
+        captured_kwargs: dict[str, Any] = {}
+
+        async def mock_create_chat_completion(**kwargs):
+            captured_kwargs.update(kwargs)
+            response = self._capture_chat_completion_dict(
+                logprobs={
+                    "content": [
+                        {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []},
+                    ]
+                }
+            )
+            response["prompt_token_ids"] = prompt_token_ids
+            response["choices"][0]["token_ids"] = [123]
+            return response
+
+        mock_client = MagicMock(spec=NeMoGymAsyncOpenAI)
+        mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion)
+        mock_client.create_tokenize = AsyncMock()
+        model._clients = [mock_client]
+
+        response = TestClient(app).post(
+            "/v1/chat/completions",
+            json={"messages": [{"role": "user", "content": "hi"}]},
+        )
+
+        assert response.status_code == 200
+        assert captured_kwargs["return_token_ids"] is True
+        assert response.json()["choices"][0]["message"]["prompt_token_ids"] == prompt_token_ids
+        assert "prompt_token_ids" not in response.json()
+        assert "token_ids" not in response.json()["choices"][0]
+        mock_client.create_tokenize.assert_not_called()
+
     def test_capture_path_succeeds_with_inbound_null_top_logprobs(self) -> None:
         """End-to-end regression: a request with top_logprobs=null no longer empties capture;
         token ids and logprobs come back populated (and coerced to int)."""
         model = _make_top_logprobs_model(return_token_id_information=True)
         app = model.setup_webserver()
 
         captured_kwargs: dict[str, Any] = {}
+        captured_tokenize_kwargs: dict[str, Any] = {}
 
         async def mock_create_chat_completion(**kwargs):
             captured_kwargs.update(kwargs)
@@ -4681,6 +4720,7 @@ async def mock_create_chat_completion(**kwargs):
             )
 
         async def mock_create_tokenize(**kwargs):
+            captured_tokenize_kwargs.update(kwargs)
             return {"tokens": [10, 20, 30]}
 
         mock_client = MagicMock(spec=NeMoGymAsyncOpenAI)
@@ -4691,18 +4731,92 @@ async def mock_create_tokenize(**kwargs):
         client = TestClient(app)
         response = client.post(
             "/v1/chat/completions",
-            json={"messages": [{"role": "user", "content": "hi"}], "top_logprobs": None},
+            json={
+                "messages": [{"role": "user", "content": "hi"}],
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "description": "Look something up",
+                            "parameters": {"type": "object", "properties": {}},
+                        },
+                    }
+                ],
+                "metadata": {"chat_template_kwargs": '{"enable_thinking": false}'},
+                "top_logprobs": None,
+            },
         )
         assert response.status_code == 200
 
         # The request forwarded to vLLM had top_logprobs pinned to 0, not the inbound null.
         assert captured_kwargs["top_logprobs"] == 0
+        assert captured_kwargs["return_token_ids"] is True
+        assert captured_tokenize_kwargs == {
+            key: captured_kwargs[key] for key in ("model", "messages", "tools", "chat_template_kwargs")
+        }
 
         message = response.json()["choices"][0]["message"]
         assert message["generation_token_ids"] == [123, 456]
         assert message["generation_log_probs"] == [-0.1, -0.2]
         assert message["prompt_token_ids"] == [10, 20, 30]
 
+    def test_capture_path_falls_back_for_null_prompt_token_ids(self) -> None:
+        model = _make_top_logprobs_model(return_token_id_information=True)
+        app = model.setup_webserver()
+
+        async def mock_create_chat_completion(**kwargs):
+            response = self._capture_chat_completion_dict(
+                logprobs={
+                    "content": [
+                        {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []},
+                    ]
+                }
+            )
+            response["prompt_token_ids"] = None
+            return response
+
+        mock_client = MagicMock(spec=NeMoGymAsyncOpenAI)
+        mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion)
+        mock_client.create_tokenize = AsyncMock(return_value={"tokens": [10, 20]})
+        model._clients = [mock_client]
+
+        response = TestClient(app).post(
+            "/v1/chat/completions",
+            json={"messages": [{"role": "user", "content": "hi"}]},
+        )
+
+        assert response.status_code == 200
+        assert response.json()["choices"][0]["message"]["prompt_token_ids"] == [10, 20]
+        mock_client.create_tokenize.assert_awaited_once()
+
+    def test_capture_path_propagates_tokenize_failure(self) -> None:
+        model = _make_top_logprobs_model(return_token_id_information=True)
+        app = model.setup_webserver()
+        tokenize_error = RuntimeError("tokenize failed")
+
+        async def mock_create_chat_completion(**kwargs):
+            return self._capture_chat_completion_dict(
+                logprobs={
+                    "content": [
+                        {"token": "token_id:123", "logprob": -0.1, "bytes": None, "top_logprobs": []},
+                    ]
+                }
+            )
+
+        mock_client = MagicMock(spec=NeMoGymAsyncOpenAI)
+        mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion)
+        mock_client.create_tokenize = AsyncMock(side_effect=tokenize_error)
+        model._clients = [mock_client]
+
+        with raises(RuntimeError, match="tokenize failed") as exc_info:
+            TestClient(app).post(
+                "/v1/chat/completions",
+                json={"messages": [{"role": "user", "content": "hi"}]},
+            )
+
+        assert exc_info.value is tokenize_error
+
     def test_capture_path_preserves_routed_experts(self) -> None:
         model = _make_top_logprobs_model(return_token_id_information=True)
         app = model.setup_webserver()