How captured chat responses skip building logprob objects that capture throws away

Explainer for NVIDIA-NeMo/RL #4634. RL permalinks pin the PR head 5b6fa89; vLLM links pin v0.29.0, the version RL uses.

Before
For every captured response, vLLM built one OpenAI logprob object per generated token, and RL then deleted all of them unread.
This PR
For captured, non-streaming requests, an override returns None instead of building those objects. Everything that is saved or sent stays the same.
Status
Proposed in this PR; the review found no correctness problem. One small follow-up is in the review: a test that catches a vLLM upgrade turning the fast path off.

Background. During Gym rollouts, RL's vLLM server records the exact token IDs and log-probabilities of each model call. This is token capture. It writes them to a shared store (it stages them), then removes them from the HTTP reply to Gym. Dotted words have a short explanation: hover or tap them.

vLLM's chat handler formats the reply in chat_completion_full_generator. For each output it calls self._create_chat_logprobs(...), which builds one object per generated token. For a 32k-token answer, that is 32k objects, built on the server's event loop. Capture never reads them. RL attaches the raw arrays from the engine output (utils.py:470), and Gym's capture code reads those message fields first (adapter L34-42). Then RL deletes the objects (vllm_worker_async.py:880).

One captured reply, both ways

Same request, one generated token (the PR's own test values: token 7, logprob −0.5) Before (and every request that isn't captured) With this PR (captured request) 1 engine output final_res.outputs[0]: token_ids [7], logprobs [{7: Logprob(-0.5)}] — same in both 2 format logprobs serving.py:1241 content: [{token: "token_id:7", logprob: -0.5, bytes: [120], top_logprobs: []}] ← 1 per token None the override returns before building anything captured_chat_logprobs.py:52-64 3 attach raw arrays worker:1270 message.generation_token_ids [7], message.generation_log_probs [-0.5] read from the engine output in step 1, not from step 2 — same in both 4 JSON payload choices[0].logprobs: {content: [ … ]} choices[0].logprobs: null 5 capture stages Gym adapter:34 reads message fields → stages ids [7], logprobs [-0.5] choice.logprobs is only a fallback when the message fields are missing — never here 6 reply to Gym worker:880 choice.pop("logprobs") → {message, ng_commit_coords} before: deletes the objects from step 2 unread · after: deletes a null — reply is identical the only steps that differ: step 2 skips the per-token objects, so step 4 carries null identical in both — staged data (5) and the reply to Gym (6) do not change Cost of step 2 on the left One object, plus a bytes list and a top-logprobs list, for every generated token. It runs on the server's event loop, so a long answer delays every other request on that server while it is built and serialized. In the PR's earlier CPU benchmark, about 300 ms of CPU per response went here.

The two paths join again at step 5. From there on, nothing differs: the same arrays are staged, and Gym gets the same reply. The only change a client could see is step 4. That payload never leaves the server, because step 6 removes logprobs before sending.

When the fast path turns on

Where it lands: RL only, one new module plus a gate in the worker. captured_chat_logprobs.py overrides vLLM's _create_chat_logprobs; vllm_worker_async.py decides when. No Gym or vLLM change: Gym already reads the message fields first, and vLLM already accepts a None there (protocol.py:106).

@@ vllm_worker_async.py base L1245-L1250 → head L1251-L1263 · NeMoRLOpenAIServingChatMixin.chat_completion_full_generator
-                response = await super().chat_completion_full_generator(
-                    request,
-                    capture_result_generator(),
-                    *args,
-                    **kwargs,
-                )
+                bypass_logprobs = bool(
+                    not request.stream
+                    and request.logprobs
+                    and return_as_token_id
+                    and id(request) in worker_self._capture_calls
+                )
+                with self._bypass_chat_logprobs(enabled=bypass_logprobs):
+                    response = await super().chat_completion_full_generator(
+                        request,
+                        capture_result_generator(),
+                        *args,
+                        **kwargs,
+                    )
                 ...
                 if request.logprobs and return_as_token_id:
                     response = attach_token_information_to_chat_response_choices(
ConditionWhy it is needed
not request.streamStreaming replies use a different vLLM function, so they are never touched.
request.logprobs and return_as_token_idThe same test as the attach step below it, so the raw arrays are always attached when the objects are skipped.
id(request) in _capture_callsOnly requests that capture accepted. Their logprobs are always deleted at step 6.

The on/off value is kept per server object, asyncio task and thread. So a second request on the same server, a child task, or a call from a worker thread still builds logprobs normally. The PR's tests check each of these.

So what. Captured replies stop paying for per-token objects that nobody reads. In the PR's earlier CPU benchmark (simulated engine, 768 active sequences, measured once), p95 time to build a reply fell from 1,420 ms to 17.7 ms, and finished replies in a 180 s window rose from 324 to 2,559. Saved data and replies to Gym do not change. The one risk is quiet: if a vLLM upgrade stops calling _create_chat_logprobs from the full generator, the speedup turns off with no error. The review asks for a one-line test to catch that.

💬 review comment — test_vllm_captured_chat_logprobs.py:122