feat(RL): add nvext Tokens-in-Tokens-Out and RL related response protocol and frontend support

[biswapanda]

Summary

• add Rust lib/llm protocol, preprocessor, and response plumbing for nvext Tokens-in-Tokens-Out
• carry nvext.token_data, cache_salt, extra_fields, token constraints, prompt_logprobs, detokenize, and stop token IDs through the Rust request/response path
• validate passthrough cache_salt and stop_token_ids; reject truncate_prompt_tokens until it is supported
• add Rust unit coverage for chat/completion request extraction and nvext response fields

Scope

This PR is the Rust lib/llm half of the nvext Tokens-in-Tokens-Out split. It intentionally excludes the vLLM worker implementation so reviewers can focus on the public/internal request protocol and response mapping first.

Where To Review

• lib/llm/src/protocols/openai/nvext.rs
• lib/llm/src/preprocessor.rs
• lib/llm/src/protocols/openai/completions.rs
• lib/llm/src/protocols/openai/chat_completions.rs
• lib/llm/src/protocols/openai/validate.rs

Local Tests

• cargo test -p dynamo-llm protocols::openai::completions::tests:: --lib
• cargo test -p dynamo-llm protocols::openai::chat_completions::tests:: --lib
• rustfmt --edition 2024 --check lib/llm/src/protocols/openai/validate.rs lib/llm/src/protocols/openai/completions.rs lib/llm/src/protocols/openai/chat_completions.rs
• git diff --check

Follow-up

The vLLM worker-side implementation should be reviewed separately so its diff remains scoped to components/src/dynamo/vllm/*.

Notes

supersedes #9382

Summary by CodeRabbit

• New Features
  • Added prompt log-probability data to API responses, enabling detailed token probability analysis at the prompt level
  • Introduced new generation control options: detokenization toggle, allowed token IDs filtering, and bad word token IDs blocking for finer-grained output control
  • Added completion token ID tracking in response payloads for improved token-level visibility

[coderabbitai]

Walkthrough

This PR extends the LLM backend and OpenAI protocol to surface prompt log-probabilities and additional sampling controls through request parsing, validation, preprocessing, response building, and delta emission across both chat and completion endpoints.

Changes

Prompt Log-Probabilities and Extended Sampling Support

Layer / File(s)	Summary
Core protocol types for prompt logprobs and sampling lib/llm/src/protocols/common/llm_backend.rs, lib/llm/src/protocols/common.rs	Introduces PromptLogprobEntry and PromptLogprobs types; adds optional prompt_logprobs field to BackendOutput and LLMEngineOutput with serde skip-serialization; extends SamplingOptions with detokenize, allowed_token_ids, and bad_words_token_ids fields.
OpenAI request protocol - CommonExt and trait implementations lib/llm/src/protocols/openai/common_ext.rs, lib/llm/src/protocols/openai/chat_completions.rs, lib/llm/src/protocols/openai/completions.rs, lib/llm/src/protocols/openai.rs	Extends CommonExt with new optional fields and implements CommonExtProvider trait methods on request types; implements NvExtProvider to expose unsupported_fields; updates get_stop_token_ids() to parse passthrough stop_token_ids from unsupported_fields; changes get_prompt_logprobs() to read from CommonExt or infer from echo setting; wires sampling options extraction into request processing.
Request validation for passthrough fields lib/llm/src/protocols/openai/validate.rs	Introduces PASSTHROUGH_EXTRA_FIELDS allowlist and updates validation to permit listed fields while type-checking cache_salt and stop_token_ids.
NvExt response protocol extensions lib/llm/src/protocols/openai/nvext.rs	Adds alternate DP rank header support; extends NvExtProvider trait and NvExtResponse with completion_token_ids and prompt_logprobs fields; extends response field selection with gating flags; updates build_response_nvext signature to accept backend inputs and conditionally populate fields based on selection and finish-reason; modifies aggregation to concatenate completion token IDs across chunks and overwrite prompt logprobs; adds cache_salt field to NvExt.
Request preprocessing and backend wiring lib/llm/src/preprocessor.rs, lib/llm/src/backend.rs, lib/llm/src/protocols/unified.rs	Adds nvext_passthrough_args helper for building nvext JSON; integrates into preprocessor request building and multimodal handling; populates max_tokens from remaining context length; wires prompt_logprobs from engine output into BackendOutput; adds CommonExtProvider delegation methods to UnifiedRequest.
Response emission through delta lib/llm/src/protocols/openai/chat_completions/delta.rs, lib/llm/src/protocols/openai/completions/delta.rs	Updates choice_from_postprocessor to capture and pass prompt_logprobs and completion_token_ids into build_response_nvext; updates debug logging to reference completion_token_ids field.
Test coverage lib/llm/src/protocols/openai/chat_completions.rs, lib/llm/src/protocols/openai/completions.rs, lib/llm/src/protocols/openai/nvext.rs, lib/llm/src/migration.rs, lib/llm/tests/test_streaming_usage.rs, lib/llm/tests/tool_choice.rs, lib/llm/tests/tool_choice_finish_reasons.rs	Adds tests for stop token ID passthrough, sampling token constraints, and truncate_prompt_tokens rejection; adds tito parity tests for completion_token_ids pass-through and prompt_logprobs final-chunk gating; updates all test fixtures to include prompt_logprobs: None in BackendOutput constructions.

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title mentions nvext tokens-in-tokens-out protocol and RL support, which directly aligns with the extensive changes across the LLM backend, protocol implementations, and response handling layers shown in the changeset.
Description check	✅ Passed	The pull request description is well-structured and complete, following the template with all required sections properly filled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/llm/src/protocols/unified.rs (1)

279-287: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Delegate unsupported_fields() for UnifiedRequest.

OpenAIPreprocessor::nvext_passthrough_args relies on NvExtProvider::unsupported_fields() fallback (e.g., cache_salt). With the current impl, UnifiedRequest always uses the trait default None, so those fields are silently dropped on this path.

Suggested fix

 impl NvExtProvider for UnifiedRequest {
     fn nvext(&self) -> Option<&NvExt> {
         self.inner.nvext.as_ref()
     }

     fn raw_prompt(&self) -> Option<String> {
         None
     }
+
+    fn unsupported_fields(
+        &self,
+    ) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
+        NvExtProvider::unsupported_fields(&self.inner)
+    }
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/protocols/unified.rs` around lines 279 - 287, The UnifiedRequest
NvExtProvider impl currently only implements nvext() and raw_prompt(), so
unsupported_fields() falls back to the trait default and drops passthrough
fields; add an unsupported_fields(&self) implementation that delegates through
the inner nvext (e.g., return self.inner.nvext.as_ref().and_then(|nv|
nv.unsupported_fields())) so OpenAIPreprocessor::nvext_passthrough_args sees
fields like cache_salt; keep the implementation inside the impl NvExtProvider
for UnifiedRequest and reference the existing nvext() and raw_prompt() methods
for placement.

🧹 Nitpick comments (1)

lib/llm/src/protocols/openai/nvext.rs (1)

45-49: ⚡ Quick win

Add explicit test coverage for the DP-rank header alias path.

Line 48 introduces x-data-parallel-rank fallback, but there’s no direct test for alias-only input (or both headers present). Add a focused test to prevent regressions in routing behavior.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/protocols/openai/nvext.rs` around lines 45 - 49, Add unit tests
for the DP-rank header alias path to ensure routing uses HEADER_DP_RANK_ALIAS
when HEADER_DP_RANK is absent and that precedence is correct when both are
present: create tests that call the header-parsing logic in nvext.rs (exercising
the dp_rank extraction that uses HEADER_DP_RANK .or_else(HEADER_DP_RANK_ALIAS)
and the dp_rank variable) with (a) only HEADER_DP_RANK_ALIAS set and (b) both
HEADER_DP_RANK and HEADER_DP_RANK_ALIAS set, asserting the returned dp_rank
matches the expected value and that no panics or unexpected routing occur.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@lib/llm/src/protocols/unified.rs`:
- Around line 279-287: The UnifiedRequest NvExtProvider impl currently only
implements nvext() and raw_prompt(), so unsupported_fields() falls back to the
trait default and drops passthrough fields; add an unsupported_fields(&self)
implementation that delegates through the inner nvext (e.g., return
self.inner.nvext.as_ref().and_then(|nv| nv.unsupported_fields())) so
OpenAIPreprocessor::nvext_passthrough_args sees fields like cache_salt; keep the
implementation inside the impl NvExtProvider for UnifiedRequest and reference
the existing nvext() and raw_prompt() methods for placement.

---

Nitpick comments:
In `@lib/llm/src/protocols/openai/nvext.rs`:
- Around line 45-49: Add unit tests for the DP-rank header alias path to ensure
routing uses HEADER_DP_RANK_ALIAS when HEADER_DP_RANK is absent and that
precedence is correct when both are present: create tests that call the
header-parsing logic in nvext.rs (exercising the dp_rank extraction that uses
HEADER_DP_RANK .or_else(HEADER_DP_RANK_ALIAS) and the dp_rank variable) with (a)
only HEADER_DP_RANK_ALIAS set and (b) both HEADER_DP_RANK and
HEADER_DP_RANK_ALIAS set, asserting the returned dp_rank matches the expected
value and that no panics or unexpected routing occur.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a6eb0b0e-a231-4196-8b4e-f258b6705f6f

📥 Commits

Reviewing files that changed from the base of the PR and between 7eb03b9 and a711882.

📒 Files selected for processing (17)

• lib/llm/src/backend.rs
• lib/llm/src/migration.rs
• lib/llm/src/preprocessor.rs
• lib/llm/src/protocols/common.rs
• lib/llm/src/protocols/common/llm_backend.rs
• lib/llm/src/protocols/openai.rs
• lib/llm/src/protocols/openai/chat_completions.rs
• lib/llm/src/protocols/openai/chat_completions/delta.rs
• lib/llm/src/protocols/openai/common_ext.rs
• lib/llm/src/protocols/openai/completions.rs
• lib/llm/src/protocols/openai/completions/delta.rs
• lib/llm/src/protocols/openai/nvext.rs
• lib/llm/src/protocols/openai/validate.rs
• lib/llm/src/protocols/unified.rs
• lib/llm/tests/test_streaming_usage.rs
• lib/llm/tests/tool_choice.rs
• lib/llm/tests/tool_choice_finish_reasons.rs

[biswapanda]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[richardhuo-nv]

LGTM!

[biswapanda]

Rebased on top of main.