[https://nvbugs/5926823][fix] Cherry-pick: Propagate logprobs from prefill to decode in disagg (#11727)

[pcastonguay]

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for log probabilities in disaggregated generation, enabling token probability tracking and propagation across generation phases.
  • Enhanced disaggregated parameter handling to capture and serialize token-level probability data.

• Tests

  • Added comprehensive test coverage for log probability serialization, deserialization, and disaggregated generation workflows with multiple configurations.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[pcastonguay]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds support for propagating log probabilities of the first generated token across disaggregated workers in multi-phase generation pipelines. It introduces a new first_gen_log_probs field to track these probabilities, implements serialization and deserialization logic, adds validation checks in the execution pipeline, and includes comprehensive tests.

Changes

Cohort / File(s)	Summary
Core Disaggregated Parameters tensorrt_llm/disaggregated_params.py	Added first_gen_log_probs: Optional[List] field to store log probabilities of the first generated token in disaggregated generation scenarios.
Request and Executor Processing tensorrt_llm/_torch/pyexecutor/llm_request.py, tensorrt_llm/_torch/pyexecutor/py_executor.py	Enhanced log_probs property with explicit None checks; added post-processing logic to append first_gen_log_probs to request results during disaggregated transmission completion with beam_width validation.
Result Processing and Validation tensorrt_llm/executor/result.py	Implemented conditional logprobs length assertions distinguishing between generation_only and non-generation_only disaggregated requests; added logic to propagate first-token log probabilities to disaggregated_params during context-only phases.
OpenAI Protocol Serialization tensorrt_llm/serve/openai_protocol.py	Added first_gen_log_probs field to DisaggregatedParams; introduced _serialize_first_gen_log_probs and _deserialize_first_gen_log_probs utilities to handle JSON-serializable conversion of token log-probability mappings.
Integration and Unit Tests tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py, tests/unittest/disaggregated/test_openai_disagg_service.py, tests/integration/test_lists/qa/llm_function_core.txt	Added end-to-end validation tests for logprobs propagation across disaggregated workers with context and generation phases; introduced TestFirstGenLogProbsSerializeRoundtrip test class covering serialization edge cases and roundtrip integrity.

Sequence Diagram

sequenceDiagram
    participant Client
    participant ContextWorker as Context Worker
    participant GenerationWorker as Generation Worker
    participant ResultHandler as Result Handler

    Client->>ContextWorker: Send context_only request
    ContextWorker->>ContextWorker: Process prefill tokens
    ContextWorker->>ContextWorker: Generate first_gen_log_probs
    ContextWorker->>ResultHandler: Return context phase result
    ResultHandler->>ResultHandler: Propagate first_gen_log_probs<br/>to disaggregated_params
    ResultHandler->>Client: Return context result with<br/>first_gen_log_probs

    Client->>GenerationWorker: Send generation_only request<br/>with first_gen_log_probs
    GenerationWorker->>GenerationWorker: Receive and validate<br/>first_gen_log_probs
    GenerationWorker->>GenerationWorker: Process decode phase
    GenerationWorker->>ResultHandler: Return generation phase result
    ResultHandler->>ResultHandler: Append first_gen_log_probs<br/>to final output
    ResultHandler->>Client: Return generation result<br/>with complete logprobs

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is largely incomplete, containing only the template structure with placeholder comments and no actual implementation details, objectives, or test coverage explanations.	Fill in the Description section explaining the issue and solution, and the Test Coverage section listing relevant tests that validate the logprobs propagation feature.
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.83% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly references the NVBugs ticket and describes the fix: propagating logprobs from prefill to decode in disaggregated mode, which aligns with the code changes.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/serve/openai_protocol.py`:
- Around line 1090-1133: Update _deserialize_first_gen_log_probs to strictly
validate types: ensure the top-level "serialized" is a list and each position is
a list (already present), but additionally enforce that each item dict contains
"token_id" as an int, "logprob" as a numeric (int/float) and "rank" if present
is either None or an int; convert numeric logprob to float and rank to int (or
None) before constructing Logprob. In practice, inside
_deserialize_first_gen_log_probs (the loop that iterates pos and item) add
explicit isinstance checks for token_id, logprob and rank, raise ValueError with
clear messages on bad types, cast logprob via float(item["logprob"]) and rank
via (None if item.get("rank") is None else int(item["rank"])), then call
Logprob(logprob=..., rank=...). Ensure error messages reference
first_gen_log_probs indices for easier debugging.

---

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 74f56e8 and 47b8940.

📒 Files selected for processing (8)

• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/disaggregated_params.py
• tensorrt_llm/executor/result.py
• tensorrt_llm/serve/openai_protocol.py
• tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/unittest/disaggregated/test_openai_disagg_service.py

[tensorrt-cicd]

PR_Github #37106 [ run ] triggered by Bot. Commit: 47b8940 Link to invocation

[brb-nv]

LGTM.

[tensorrt-cicd]

PR_Github #37106 [ run ] completed with state SUCCESS. Commit: 47b8940
/LLM/release-1.3.0rc5.post1/L0_MergeRequest_PR pipeline #3 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[pcastonguay]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37175 [ run ] triggered by Bot. Commit: 47b8940 Link to invocation

[tensorrt-cicd]

PR_Github #37175 [ run ] completed with state SUCCESS. Commit: 47b8940
/LLM/release-1.3.0rc5.post1/L0_MergeRequest_PR pipeline #5 completed with status: 'SUCCESS'

Link to invocation

[Superjomn]

LGTM on the llmapi changes.