fix(vllm): validate replaced prompt length

Author: joyang-nv

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -753,6 +753,23 @@ def model_post_init(self, context):
                 return super().model_post_init(context)
 
         class NeMoRLOpenAIServingMixin:
+            def _validate_chat_completion_prompt_token_ids(
+                self, request, prompt_token_ids: list[int]
+            ) -> None:
+                """Validate the final prompt after authoritative token replacement.
+
+                vLLM validates the chat-template rendering before this mixin replaces
+                its prefix with authoritative trajectory tokens.  The replacement can
+                make the actual engine prompt longer than the rendering that vLLM
+                checked, so every chat-completion return path must validate the final
+                token IDs again.  Tokenize requests intentionally remain unrestricted.
+                """
+                if isinstance(request, NeMoRLChatCompletionRequest):
+                    validate_prompt_token_count(
+                        prompt_token_count=len(prompt_token_ids),
+                        max_prompt_tokens=self.model_config.max_model_len - 1,
+                    )
+
             @staticmethod
             def _set_max_tokens(request, max_tokens: int) -> None:
                 """Set the request's max output tokens.
@@ -950,12 +967,16 @@ async def preprocess_chat(
                     not hasattr(request, "required_prefix_token_ids")
                     or request.required_prefix_token_ids is None
                 ):
+                    final_prompt_token_ids = res[1][0]["prompt_token_ids"]
+                    self._validate_chat_completion_prompt_token_ids(
+                        request, final_prompt_token_ids
+                    )
                     # Clamp the request's max output tokens so that input + output <= max_model_len.
                     if actual_request_max_tokens is not None:
                         self._clamp_max_tokens(
                             request,
                             actual_request_max_tokens,
-                            res[1][0]["prompt_token_ids"],
+                            final_prompt_token_ids,
                         )
                     return res
 
@@ -1006,6 +1027,10 @@ async def preprocess_chat(
 
                 engine_prompt["prompt_token_ids"] = final_prompt_token_ids
 
+                self._validate_chat_completion_prompt_token_ids(
+                    request, final_prompt_token_ids
+                )
+
                 # Clamp after prefix replacement since the prompt length may have changed.
                 if actual_request_max_tokens is not None:
                     self._clamp_max_tokens(

tests/unit/models/generation/test_vllm_generation.py

@@ -2637,6 +2637,26 @@ async def test_vllm_http_server_correct_merged_tokens_matches_baseline(
         vllm_http_server_generated_token["token"].removeprefix("token_id:")
     )
 
+    # vLLM validates the shorter chat-template rendering before NeMo RL
+    # replaces it with the authoritative trajectory prefix.  Reject an
+    # oversized replacement before it reaches the engine's fixed token buffer.
+    oversized_authoritative_prefix = [
+        *initial_tokenized_query_ids_prefix,
+        *([initial_tokenized_ids[0]] * 1024),
+    ]
+    oversized_response = requests.post(
+        url=f"{base_urls[0]}/chat/completions",
+        json=body | {"required_prefix_token_ids": oversized_authoritative_prefix},
+    )
+    assert oversized_response.status_code == 400
+    assert "maximum context length" in oversized_response.json()["error"]["message"]
+
+    # The rejection must leave EngineCore healthy for subsequent requests.
+    healthy_response = requests.post(
+        url=f"{base_urls[0]}/chat/completions", json=body_with_reference_token_ids
+    )
+    assert healthy_response.status_code == 200
+
     # A prompt tokenized by this server can be reused directly by the next
     # generation request. The chat conversation is still rendered for response
     # and tool parsing, but prompt encoding is bypassed.