[https://nvbugs/5934461][fix] Cherry-picks 11767 (logits support in disagg)

[pcastonguay]

Summary by CodeRabbit

• New Features

  • Added method to retrieve generation logits while bypassing the exclusion flag.
  • Extended disaggregated parameters to support first-generation logits propagation.
  • Improved logits handling for first-token responses in disaggregated execution.

• Tests

  • Added comprehensive tests for logits serialization and disaggregated generation flows including request cancellation support.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR introduces first-generation logits propagation through the disaggregated execution pipeline. It adds first_gen_logits field to DisaggregatedParams, implements prepending and detection logic in PyExecutor, extends PyResult with an unexcluded logits accessor, and implements serialization/deserialization with dtype handling in the OpenAI protocol layer. Integration tests validate the new cancellation and logits propagation workflows.

Changes

Cohort / File(s)	Summary
Logits Field & Propagation tensorrt_llm/disaggregated_params.py, tensorrt_llm/executor/result.py	Adds first_gen_logits field to DisaggregatedParams and propagates it from generation_logits in context-only responses.
PyExecutor Prepending Logic tensorrt_llm/_torch/pyexecutor/py_executor.py	Adds _maybe_prepend_logprobs_and_logits() and _has_prepended_logits() helper methods to prepend disaggregated logits and detect prepended state; integrates with disaggregated transmission and first-token response handling.
PyResult Accessor tensorrt_llm/_torch/pyexecutor/llm_request.py	Adds get_latest_logits_unexcluded() method to retrieve latest generation logits while bypassing exclusion flags.
Serialization & Deserialization tensorrt_llm/serve/openai_protocol.py	Adds _serialize_first_gen_logits() and _deserialize_first_gen_logits() with base64 encoding, shape/dtype metadata, and bf16→fp16 conversion; integrates into disaggregated parameter conversion functions.
Integration & Unit Tests tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py, tests/unittest/_torch/executor/test_chunked_logits.py, tests/unittest/disaggregated/test_openai_disagg_service.py	Adds new MPI signaling tags (MPI_CANCEL, MPI_STARTED), cancellation-aware worker orchestration with streaming/non-streaming request handling, tests for get_latest_logits_unexcluded() and serialization roundtrip.
Test Registry tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/test-db/l0_h100.yml	Registers new integration tests (test_disaggregated_cancel_gen_requests, test_disaggregated_logits) for TinyLlama model across multiple test suites.

Sequence Diagram

sequenceDiagram
    participant Executor as PyExecutor
    participant Result as PyResult
    participant Params as DisaggregatedParams
    participant Protocol as OpenAI Protocol

    Note over Executor,Protocol: First-Gen Logits Propagation Flow

    Executor->>Executor: _maybe_prepend_logprobs_and_logits()
    activate Executor
    Executor->>Result: append generation_logits
    activate Result
    Result->>Result: store logits chunks
    deactivate Result
    deactivate Executor

    Executor->>Result: get_latest_logits_unexcluded()
    activate Result
    Result-->>Executor: return unexcluded logits
    deactivate Result

    Executor->>Params: first_gen_logits assignment
    activate Params
    Params->>Params: store first_gen_logits
    deactivate Params

    Params->>Protocol: serialize first_gen_logits
    activate Protocol
    Protocol->>Protocol: base64 encode + dtype handling
    Protocol-->>Params: JSON-safe representation
    deactivate Protocol

    Protocol->>Params: deserialize first_gen_logits
    activate Params
    Params->>Params: restore torch.Tensors
    deactivate Params

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [https://nvbugs/5934461][fix] Propagate logits from prefill to decode in disagg #11767: Directly related—both modify disaggregated first_gen_logits feature including field addition, serialization in openai_protocol, and prepend/expose logic in pyexecutor and result handling.

Suggested reviewers

• hchings
• chuangz0
• schetlur-nv

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.64% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	PR description contains only the template placeholder with no actual explanation of changes, objectives, test coverage, or checklist completion details filled in.	Fill in the Description, Test Coverage, and PR Checklist sections with specific details about the logits propagation feature, related test cases, and verification steps.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title references a specific bug fix (NVBugs 5934461) and clearly indicates this is a cherry-pick of logits support in disaggregated generation, which is the primary focus of the changeset across multiple files.

_{✏️ 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: 3

Caution

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

⚠️ Outside diff range comments (1)

tensorrt_llm/executor/result.py (1)

305-312: ⚠️ Potential issue | 🔴 Critical

Add missing logger import

logger.warning() at line 305 references an undefined variable. The logger module is used but never imported, causing a NameError at runtime when disaggregated serving encounters a logprobs mismatch.

Proposed fix

+from ..logger import logger

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/executor/result.py` around lines 305 - 312, The warning call
uses an undefined symbol `logger`, causing a NameError; fix by adding a
module-level logger and importing the stdlib logging module (e.g., import
logging and initialize logger = logging.getLogger(__name__)) at the top of
tensorrt_llm/executor/result.py so the `logger.warning(...)` call (around the
logprobs mismatch message) has a defined logger object; place the
import/initialization near other imports to match file style.

🤖 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 1137-1176: The serializer and deserializer currently assume
well-formed inputs; update _serialize_first_gen_logits to validate the top-level
first_gen_logits is a list/sequence and that each element is a torch.Tensor
(raise ValueError with a clear message otherwise), and validate tensor
attributes (e.g., contiguous/cpu convertible) before encoding; similarly update
_deserialize_first_gen_logits to validate the top-level serialized is a list,
each item is a dict and that "data" is a string, "shape" is a list of ints, and
"dtype" is a valid numpy dtype string (raise ValueError for violations), and
check that the decoded bytes length is compatible with the inferred dtype/shape
before constructing the ndarray and torch.Tensor.

In `@tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py`:
- Around line 794-796: The test currently swallows exceptions in the
cancellation block by printing the exception in the except Exception as e
handler and continuing to finally, which can hide failures; update the except
block in the cancellation test to either re-raise the caught exception (raise)
after logging or call pytest.fail(str(e)) so the test fails on errors—modify the
except Exception as e handler that prints "Exception encountered: {e}" to raise
the exception (or invoke pytest.fail) to ensure real regressions surface.

In `@tests/unittest/_torch/executor/test_chunked_logits.py`:
- Around line 179-180: The test constructs an incorrect 4D expected tensor by
calling .unsqueeze(1) before transpose; remove .unsqueeze(1) so the expected
shape matches the 3D output of get_latest_logits_unexcluded(). Update the
assertion lines that build expected from logits_b (e.g., replace
logits_b.unsqueeze(1).transpose(0, 1) with logits_b.transpose(0, 1) or simply
logits_b.transpose as appropriate) and make the same change for the similar
block around lines 196-202 to ensure shapes match before torch.allclose;
reference get_latest_logits_unexcluded, logits_b, and unexcluded when locating
the fixes.

---

Outside diff comments:
In `@tensorrt_llm/executor/result.py`:
- Around line 305-312: The warning call uses an undefined symbol `logger`,
causing a NameError; fix by adding a module-level logger and importing the
stdlib logging module (e.g., import logging and initialize logger =
logging.getLogger(__name__)) at the top of tensorrt_llm/executor/result.py so
the `logger.warning(...)` call (around the logprobs mismatch message) has a
defined logger object; place the import/initialization near other imports to
match file style.

---

ℹ️ 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 e922a86 and 867f36e.

📒 Files selected for processing (10)

• 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/integration/test_lists/test-db/l0_h100.yml
• tests/unittest/_torch/executor/test_chunked_logits.py
• tests/unittest/disaggregated/test_openai_disagg_service.py

[pcastonguay]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37367 [ run ] triggered by Bot. Commit: 867f36e Link to invocation

[brb-nv]

LGTM.

[tensorrt-cicd]

PR_Github #37367 [ run ] completed with state SUCCESS. Commit: 867f36e
/LLM/release-1.3.0rc5.post1/L0_MergeRequest_PR pipeline #9 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 #37411 [ run ] triggered by Bot. Commit: 867f36e Link to invocation

[tensorrt-cicd]

PR_Github #37411 [ run ] completed with state SUCCESS. Commit: 867f36e
/LLM/release-1.3.0rc5.post1/L0_MergeRequest_PR pipeline #12 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 #37506 [ run ] triggered by Bot. Commit: 867f36e Link to invocation

[tensorrt-cicd]

PR_Github #37506 [ run ] completed with state SUCCESS. Commit: 867f36e
/LLM/release-1.3.0rc5.post1/L0_MergeRequest_PR pipeline #15 completed with status: 'SUCCESS'

Link to invocation