[macandro96]

What does this PR do?

Captured chat responses currently build and serialize OpenAI logprob objects for every generated token, then discard those objects after staging the raw token IDs and selected-token logprobs to TQ. This PR skips that redundant object construction for admitted, non-streaming captures that use the raw-array response path.

Targets main directly. Main already stages captured raw arrays and strips standard OpenAI logprobs after terminal staging. This optimization requires no prefix-recovery or checkpoint-stack changes.

Changes

• Scope the bypass around full-response formatting only when the request is non-streaming, requests logprobs, resolves return_tokens_as_token_ids=true, and has an admitted capture state.
• Override vLLM's _create_chat_logprobs at the serving MRO boundary. Keep the original request and engine sampling parameters intact. Continue attaching and validating the engine-native raw arrays through the existing helper.
• Isolate the scope by context, serving instance, task, and thread, and reset it on success, error, or cancellation. Other requests delegate to vLLM normally.
• Keep vLLM's response parsing and the existing capture staging/stripping behavior. No checkpoint protocol or runtime configuration changes.

Observed performance gains

Historical CPU A/B from tq-chat-response-v2-768-20261007-085403, benchmark revision reused-engine-inputs-bounded-drain-v2:

Steady-state metric	Full OpenAI logprob construction	Bypass
Response construction wall time, mean	356.3 ms	6.43 ms
Response construction wall time, p95	1,420.4 ms	17.68 ms
Response serialization CPU time, mean	30.68 ms	1.90 ms
Terminal latency including gate wait, p95	31.80 s	1.93 s
Terminal executor queue wait, p95	26.60 s	1.45 s
Terminal rows completed in the 180 s load window	324	2,559

That is about 80× lower response-construction p95, 94% lower terminal p95, and 7.9× more steady terminal completions in this workload. The full path spent 307.4 ms of CPU time on average in _create_chat_logprobs; the bypass omits that work. Raw-array attachment remained about 2 ms CPU per response in both modes.

Setup: one simulated generation worker, 768 active sequences, token lengths 2,048 / 8,192 / 32,768 / 65,536, four Simple TQ storage actors, 32 terminal writers, pending-terminal cap 64, cached physical-core counts, audit disabled, and checkpoint attempts every 30 s. The read-only synthetic engine inputs were reused, avoiding repeated fixture construction as a confounder. Measurements above use steady_* distributions and exclude initialization and final drain.

The three fence-drain measurements were 4.57 / 8.32 / 1.55 s with full logprobs versus 0.58 / 0.26 / 0.48 s with the bypass. Prefix cut time did not improve consistently: the benchmark's cut_and_verify_seconds were 29.62 / 18.45 / 13.77 s versus 23.13 / 20.60 / 20.22 s. This change removes response-formatting contention; it does not remove all prefix-processing work.

These measurements come from the earlier prefix-checkpointing experiment; its fence/cut phases are context for the gains, not functionality added by this PR. These are historical results from the benchmark's full versus skip-logprobs modes, not a GPU measurement of this new commit or a 131k-rollout experiment. The benchmark exercised production cut/fence/terminal methods, Gym capture, and installed vLLM/RL response builders, with a simulated engine and shared serving loop. It excluded text/tool parsing, GPU generation, Gym HTTP, training, and checkpoint export. Aggregate terminal timings include queueing and other work; they are not isolated storage PUT latencies. Each mode was measured once.

Internal artifact root:
/lustre/fsw/portfolios/nemotron/projects/nemotron_sw_post/users/amahishi/projects/pipeline/artifacts/tq-chat-response-v2-768-20261007-085403

Results: full/seqs-768-audit-off/result.json and skip-logprobs/seqs-768-audit-off/result.json.

Usage

Automatic for eligible captured requests; no new flag is needed. Keep engine logprobs enabled. Ordinary API responses, streaming requests, unadmitted capture markers, and requests without the raw-array path retain their normal formatting behavior.

Testing

• 89 passed, 4 skipped across test_captured_chat_logprobs.py, test_vllm_chat_template_wiring.py, test_vllm_utils.py, and test_vllm_token_capture_hosting.py.
• Serving-path regressions cover admitted capture, server defaults, ordinary/unadmitted/streaming/non-raw requests, unchanged request fields, preserved raw arrays, and rejection of missing selected-token logprobs.
• Scope tests cover concurrent requests, nested scopes, other serving instances, inherited child-task/thread contexts, exceptions, and cancellation.
• Ruff lint/format checks, the new module's Pyrefly check, and git diff --check passed.

Local validation used a Python 3.12 CPU environment with the Gym revision pinned by main (17abf2ddc76f74e40d0ba5a15b0401710a46db26). The four skipped tests require vLLM, which is not installed locally; one is a new engine-free integration test against the real vLLM formatter. GPU generation and end-to-end Gym HTTP validation remain for CI or a container smoke run.

Before your PR is "Ready for review"

• Read and followed contributor guidelines.
• Added necessary regression tests.
• Ran focused unit tests and static checks.
• Run optional vLLM tests and GPU/container functional validation.
• Documented behavior and historical measurement scope above.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[macandro96]

/ok to test ff5c044

[macandro96]

/ok to test 1cfe1e3

[macandro96]

/ok to test 64de786

[macandro96]

/ok to test 5b6fa89

[terrykong]

Thanks for this. The change is small, well-scoped, and safe. We found no correctness issues.

When the bypass is on, nothing reads the objects it skips:

• the raw arrays come from the engine output (utils.py#L470);
• Gym's capture adapter reads the message fields first (adapters/vllm.py#L34-L42);
• _finish_request_capture strips choice.logprobs anyway (vllm_worker_async.py#L880);
• vLLM's ChatCompletionResponseChoice.logprobs is Optional, so leaving it unset is allowed (protocol.py#L106).

The scope design is good. The separate module imports without vLLM, so the scope tests run in the CPU lane. Each part of the (instance, task, thread) check has its own test: we planted 9 mutants and every one was caught. And when the bypass doesn't apply, it falls back to vLLM's normal path.

One low-severity inline comment: add a test that fails if a vLLM bump stops calling _create_chat_logprobs from the full generator. Today that would turn off the speedup without any test failing.

No ask on perf numbers. They are from an earlier run, but the reason for the speedup is clear: it stops building per-token objects that get thrown away.

We did not run anything on GPU or with vLLM installed. The CPU tests passed locally (18 passed, 1 skipped because vLLM is missing).

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4634-logprob-bypass.html

Generated by Claude Code

[macandro96]

/ok to test 7a3b6e1