[https://nvbugs/5934461][fix] Propagate logits from prefill to decode in disagg

[brb-nv]

Description

This is to address the following bug by propagating logits from prefill to decode in disagg mode.
https://nvbugspro.nvidia.com/bug/5934461

Test Coverage

$ pytest tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[True-TinyLlama-1.1B-Chat-v1.0] -s -v
$ pytest tests/unittest/disaggregated/test_openai_disagg_service.py::TestFirstGenLogitsSerializeRoundtrip -s -v
$ pytest tests/unittest/_torch/executor/test_chunked_logits.py::TestGetLatestLogitsUnexcluded -s -v

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.

Summary by CodeRabbit

Release Notes

• New Features

  • Disaggregated generation logits are now available and properly propagated through prefill and generation phases
  • Added streaming support for disaggregated generation

• Bug Fixes

  • Fixed generation logits handling in context-only disaggregated responses

• Tests

  • Added integration tests validating disaggregated logits for streaming and non-streaming scenarios

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces disaggregated generation logits propagation, enabling capture and transmission of first-token logits from prefill to decode phases. It adds serialization/deserialization helpers, updates executor and result-handling logic, implements a repro workflow with an HTTP-based prefill server, and includes new integration tests with streaming support.

Changes

Cohort / File(s)	Summary
Core data model tensorrt_llm/disaggregated_params.py	Added first_gen_logits optional field to store per-beam, per-token logits from prefill phase.
Serialization and protocol tensorrt_llm/serve/openai_protocol.py	Implemented _serialize_first_gen_logits and _deserialize_first_gen_logits helpers with base64 encoding; integrated into to_disaggregated_params and to_llm_disaggregated_params conversion functions.
Executor integration tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/_torch/pyexecutor/llm_request.py	Added logic to prepend first_gen_logits to generation results when disaggregated transmission completes; enforced beam_width == 1 constraint; captured logits snapshot in first token response.
Result propagation tensorrt_llm/executor/result.py	Implemented collection of first generation logits from context-phase outputs and assignment to disaggregated_params for downstream consumption.
Repro workflow disagg_logits_repro.py	New module implementing PrefillEngine (HTTP server with /prefill endpoint), Engine client (remote prefill orchestration), serialization helpers, and pytest-based end-to-end test with CLI support.
Integration testing tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py	Added test_disaggregated_logits function with streaming support; worker-side logic to handle per-chunk logits collection; end-to-end validation across aggregated and disaggregated paths.
Test configuration tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/test-db/l0_a10.yml, tests/integration/test_lists/test-db/l0_h100.yml	Registered new test_disaggregated_logits parameterized test cases for TinyLlama-1.1B-Chat-v1.0 with streaming variants.
Unit tests tests/unittest/disaggregated/test_openai_disagg_service.py	Added TestFirstGenLogitsSerializeRoundtrip suite covering serialization/deserialization of first_gen_logits across dtypes (float32, float16, bfloat16) and multi-beam scenarios.

Sequence Diagram(s)

sequenceDiagram
    participant Client as Client/Engine
    participant PrefillSrv as Prefill Server
    participant LLM as LLM Model
    participant Executor as Executor
    
    Client->>PrefillSrv: HTTP POST /prefill (prompt)
    PrefillSrv->>LLM: generate_async(sampling_params)
    LLM-->>PrefillSrv: first_gen_tokens + first_gen_logits
    PrefillSrv->>PrefillSrv: serialize_logits(base64)
    PrefillSrv-->>Client: PrefillResponse (logits serialized)
    Client->>Client: deserialize_logits()
    Client->>Executor: generate_async(disaggregated_params)
    Executor->>Executor: prepend_first_gen_logits()
    Executor->>LLM: generate_async(decode phase)
    LLM-->>Executor: generation_logits
    Executor-->>Client: streaming response with logits

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• [TRTLLM-8921][feat] implement gen-first disagg_service #11020: Shares the same first_gen_logits disaggregated propagation flow with overlapping serialization, parameter, and executor modifications.

Suggested reviewers

• chuangz0
• schetlur-nv
• pcastonguay

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 8.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	The PR description is incomplete. While it identifies the bug being fixed and provides test coverage commands, it lacks key template sections like a detailed explanation of the solution and clarity on design changes.	Expand the description to explain the technical approach for propagating logits, any architectural changes, and clarify whether the tava architecture diagram needs updating given the design modifications.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: propagating logits from prefill to decode in disaggregated mode, which aligns with the primary objective of the pull request.

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

🧹 Nitpick comments (5)

tensorrt_llm/executor/result.py (1)

27-27: Use module-namespace import for logger.

Please import the logger module instead of importing logger directly to match repo import conventions.

Suggested change

-from ..logger import logger
+from .. import logger

As per coding guidelines "When importing in Python, always maintain the namespace. Import the module, not individual classes or functions".

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/executor/result.py` at line 27, Replace the direct object import
"from ..logger import logger" with a module-namespace import (e.g., "from ..
import logger") and update any usages in this file that reference the logger
object (references to logger.*) to access the logger via the module
(logger.logger.*) so the repo import convention is preserved; ensure all calls
in the file that used the imported name now reference the module member.

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

3280-3282: Make temporary exclude_last_generation_logits mutation exception-safe.

At Line 3280-3282, if generation_logits access throws, the flag can remain in the wrong state. Wrap restore in finally.

Suggested hardening

                 if has_prepended_logits:
                     req.set_exclude_last_generation_logits(False)
-                    logits_snapshot = req.py_result.generation_logits
-                    req.set_exclude_last_generation_logits(True)
+                    try:
+                        logits_snapshot = req.py_result.generation_logits
+                    finally:
+                        req.set_exclude_last_generation_logits(True)

🤖 Prompt for AI Agents

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

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3280 - 3282, The
temporary mutation of req.set_exclude_last_generation_logits must be
exception-safe: wrap the access to req.py_result.generation_logits in a
try/finally so that req.set_exclude_last_generation_logits(True) is always
called to restore state even if reading generation_logits throws; specifically,
set exclude to False, assign logits_snapshot from
req.py_result.generation_logits inside the try block, and restore exclude to
True in the finally block to guarantee consistent state for subsequent
operations.

tests/unittest/disaggregated/test_openai_disagg_service.py (2)

16-17: Use module-namespace access for the newly added helper imports.

For these newly added helpers, prefer module import + qualified access instead of direct symbol import.

Example refactor

 from tensorrt_llm.serve.openai_protocol import (
     CompletionRequest,
     CompletionResponse,
     CompletionResponseChoice,
     DisaggregatedParams,
     DisaggScheduleStyle,
     UsageInfo,
-    _deserialize_first_gen_logits,
-    _serialize_first_gen_logits,
 )
+from tensorrt_llm.serve import openai_protocol

-        assert _serialize_first_gen_logits(None) is None
-        assert _deserialize_first_gen_logits(None) is None
+        assert openai_protocol._serialize_first_gen_logits(None) is None
+        assert openai_protocol._deserialize_first_gen_logits(None) is None

As per coding guidelines "When importing in Python, always maintain the namespace. Import the module, not individual classes or functions (e.g., use from package.subpackage import foo then foo.SomeClass())."

🤖 Prompt for AI Agents

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

In `@tests/unittest/disaggregated/test_openai_disagg_service.py` around lines 16 -
17, Replace the direct symbol imports of _deserialize_first_gen_logits and
_serialize_first_gen_logits in the test_openai_disagg_service.py file with a
module-level import (import the module that defines them) and update all call
sites to use qualified access (e.g.,
module_name._deserialize_first_gen_logits(...) and
module_name._serialize_first_gen_logits(...)); keep the original function names
but remove them from the from ... import ... list so the test follows the
module-namespace import guideline.

---

193-194: Add strict=True to zip in the roundtrip comparison loop.

At Line 193, this avoids silent truncation if lengths diverge.

Suggested change

-        for orig, rest in zip(original, restored):
+        for orig, rest in zip(original, restored, strict=True):
             torch.testing.assert_close(rest, orig.cpu())

🤖 Prompt for AI Agents

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

In `@tests/unittest/disaggregated/test_openai_disagg_service.py` around lines 193
- 194, The roundtrip comparison loop uses zip(original, restored) which can
silently truncate if the two iterables differ in length; update that call to
zip(original, restored, strict=True) in the test loop (the one iterating over
orig and rest and calling torch.testing.assert_close(rest, orig.cpu())) so the
test will raise on length mismatches instead of silently ignoring extra
elements.

tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py (1)

655-657: Avoid broad catch-and-reraise in this test block.

Line 655 catches Exception and rethrows immediately. Prefer removing this block (rely on finally) or narrowing to specific exception types.

As per coding guidelines When using try-except blocks in Python, limit the except to the smallest set of errors possible.

🤖 Prompt for AI Agents

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

In `@tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py` around
lines 655 - 657, The test currently catches a broad Exception and immediately
re-raises it (the try/except around the teardown block in
test_disaggregated_single_gpu), which is unnecessary; remove the except
Exception: ... raise e block and rely on the existing finally for cleanup, or if
you need to handle known failure modes, replace it with explicit handlers for
the specific exception types you expect (e.g., TimeoutError, AssertionError) and
log contextual info before re-raising. Ensure you remove the print-and-reraise
pattern around the try so failures propagate naturally to the test runner and
cleanup in finally still executes.

🤖 Prompt for all review comments with AI agents

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

Inline comments:
In `@disagg_logits_repro.py`:
- Around line 29-32: The hardcoded local path in the MODEL_PATH getenv default
should be removed; update the os.getenv("TRTLLM_MODEL_PATH", ...) call in
disagg_logits_repro.py to use no machine-specific fallback (e.g.,
os.getenv("TRTLLM_MODEL_PATH") or an empty string) and add a short validation
after the assignment that raises a clear error (or exits) if MODEL_PATH is not
set, referencing MODEL_PATH so runs/CI must supply TRTLLM_MODEL_PATH rather than
relying on a local filesystem path.
- Around line 111-115: The thread starts uvicorn with host "0.0.0.0" which
exposes the endpoint externally; change the default host in the uvicorn.run
invocation used inside threading.Thread (the args/kwargs passed to uvicorn.run
in the Thread creation) from "0.0.0.0" to a loopback address like "127.0.0.1" or
"localhost" (or make it configurable via an environment variable) so the server
binds to localhost by default rather than all interfaces.

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2830-2834: The code currently raises a ValueError when beam_width
!= 1 in the "first_gen_logits transfer" block (variable beam_width in
py_executor.py), which can unwind the executor loop and kill unrelated requests;
instead handle this as a request-level rejection: detect beam_width != 1, log a
warning noting the unsupported per-request option, mark or return an error
status for that specific request (or skip disaggregated logits propagation for
it) and continue the executor loop rather than raising an exception. Locate the
check around beam_width and replace the raise with per-request error handling
(e.g., set the request result/error, emit a warning via the executor logger, and
continue) so only the offending request is affected.

In `@tensorrt_llm/serve/openai_protocol.py`:
- Around line 1125-1129: Validate and harden deserialization of request-facing
tensors by first ensuring item["data"] is valid base64 and decoding succeeds,
verifying item["dtype"] is a recognized numpy dtype, and confirming
item["shape"] is a sequence of positive ints whose product times dtype.itemsize
equals the decoded byte length before calling np.frombuffer and reshape; catch
and raise a clear ValueError on any mismatch or decode/reshape error instead of
trusting inputs, and only then append torch.from_numpy(np_array.copy()) to
result (refer to variables/functions: item["data"], item["dtype"],
item["shape"], np.frombuffer, reshape, result.append(torch.from_numpy(...)) for
where to implement these checks).

In `@tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py`:
- Around line 587-590: The test currently only checks
agg_output.generation_logits for non-None and shape but not correctness; update
the test to assert that agg_logits (from agg_output.generation_logits) matches
the expected propagated logits by comparing specific token logits (e.g.,
first-token logits or a small slice) against the corresponding logits from the
per-shard/single-request results (e.g., single_outputs[i].generation_logits or
the concatenated shard logits used to build the aggregate) or against a
precomputed expected tensor; compute the expected slice from the individual
outputs, then use an equality or close assertion (e.g., exact equality for ints
or approximate for floats) to ensure values were propagated correctly. Ensure
you update all places where agg_logits is checked (the
agg_output.generation_logits usage and the similar checks around the other test
blocks) to perform these value comparisons rather than only shape/non-None
checks.

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3280-3282: The temporary mutation of
req.set_exclude_last_generation_logits must be exception-safe: wrap the access
to req.py_result.generation_logits in a try/finally so that
req.set_exclude_last_generation_logits(True) is always called to restore state
even if reading generation_logits throws; specifically, set exclude to False,
assign logits_snapshot from req.py_result.generation_logits inside the try
block, and restore exclude to True in the finally block to guarantee consistent
state for subsequent operations.

In `@tensorrt_llm/executor/result.py`:
- Line 27: Replace the direct object import "from ..logger import logger" with a
module-namespace import (e.g., "from .. import logger") and update any usages in
this file that reference the logger object (references to logger.*) to access
the logger via the module (logger.logger.*) so the repo import convention is
preserved; ensure all calls in the file that used the imported name now
reference the module member.

In `@tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py`:
- Around line 655-657: The test currently catches a broad Exception and
immediately re-raises it (the try/except around the teardown block in
test_disaggregated_single_gpu), which is unnecessary; remove the except
Exception: ... raise e block and rely on the existing finally for cleanup, or if
you need to handle known failure modes, replace it with explicit handlers for
the specific exception types you expect (e.g., TimeoutError, AssertionError) and
log contextual info before re-raising. Ensure you remove the print-and-reraise
pattern around the try so failures propagate naturally to the test runner and
cleanup in finally still executes.

In `@tests/unittest/disaggregated/test_openai_disagg_service.py`:
- Around line 16-17: Replace the direct symbol imports of
_deserialize_first_gen_logits and _serialize_first_gen_logits in the
test_openai_disagg_service.py file with a module-level import (import the module
that defines them) and update all call sites to use qualified access (e.g.,
module_name._deserialize_first_gen_logits(...) and
module_name._serialize_first_gen_logits(...)); keep the original function names
but remove them from the from ... import ... list so the test follows the
module-namespace import guideline.
- Around line 193-194: The roundtrip comparison loop uses zip(original,
restored) which can silently truncate if the two iterables differ in length;
update that call to zip(original, restored, strict=True) in the test loop (the
one iterating over orig and rest and calling torch.testing.assert_close(rest,
orig.cpu())) so the test will raise on length mismatches instead of silently
ignoring extra elements.

---

ℹ️ 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 3e12071 and 8dd7a6d.

📒 Files selected for processing (11)

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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #36999 [ run ] triggered by Bot. Commit: 7189126 Link to invocation

[tensorrt-cicd]

PR_Github #36999 [ run ] completed with state SUCCESS. Commit: 7189126
/LLM/main/L0_MergeRequest_PR pipeline #28647 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37112 [ run ] triggered by Bot. Commit: dfeca65 Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37130 [ run ] triggered by Bot. Commit: 677f18d Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37217 [ run ] triggered by Bot. Commit: 677f18d Link to invocation

[tensorrt-cicd]

PR_Github #37217 [ run ] completed with state SUCCESS. Commit: 677f18d
/LLM/main/L0_MergeRequest_PR pipeline #28805 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]

Looks like some user buffer tests are failing. Could you take a closer look? Thanks.

[brb-nv]

Looks like some user buffer tests are failing. Could you take a closer look? Thanks.

These test failures are unrelated and occur on latest post-merge on main also, Patrice.
https://nvidia.slack.com/archives/C06BP0WMHEW/p1772388064551419

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37257 [ run ] triggered by Bot. Commit: 78db5ae Link to invocation

[tensorrt-cicd]

PR_Github #37257 [ run ] completed with state SUCCESS. Commit: 78db5ae
/LLM/main/L0_MergeRequest_PR pipeline #28836 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37290 [ run ] triggered by Bot. Commit: 78db5ae Link to invocation

[tensorrt-cicd]

PR_Github #37290 [ run ] completed with state SUCCESS. Commit: 78db5ae
/LLM/main/L0_MergeRequest_PR pipeline #28859 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37344 [ run ] triggered by Bot. Commit: 78db5ae Link to invocation

[tensorrt-cicd]

PR_Github #37344 [ run ] completed with state SUCCESS. Commit: 78db5ae
/LLM/main/L0_MergeRequest_PR pipeline #28905 completed with status: 'SUCCESS'

Link to invocation