feat: Add logprobs support to TRTLLM backend

[esoba]

Overview:

Add logprobs extraction and propagation for TRTLLM backend

Details:

• Added _extract_logprobs() method to extract and format logprobs from TRTLLM engine output
• Propagates log_probs and top_logprobs to Dynamo's response format
• Handles top_logprobs=0 edge case using max(1, logprobs)
• Added ChatPayloadWithLogprobs and CompletionPayloadWithLogprobs test utilities
• Added e2e tests for aggregated and disaggregated logprobs configurations

Where should the reviewer start?

• components/src/dynamo/trtllm/request_handlers/handler_base.py - main logprobs extraction logic
• tests/serve/test_trtllm.py - new test configurations
• 'tests/utils/payloads.py' - Added dataclass for payload w/logits
• 'tests/utils/payload_builder.py' - Added arguments for logprobs to chat/completion payloads

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: [FEATURE]: Add logprobs support to TRTLLM Backend #4684
• Follows approach from PoC: logprobs support for vLLM backend #4678 + feat: Add logprobs support to vLLM backend (#4683) #4697

Summary by CodeRabbit

• New Features

  • Added log probability (logprobs) support to retrieve token-level probability information in responses
  • Users can now request logprobs and top_logprobs parameters to obtain probability data for generated tokens and alternative top choices

• Tests

  • Added test configurations and utilities for validating logprobs functionality across different parameter combinations

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[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.

[github-actions]

👋 Hi esoba! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

Walkthrough

The changes implement log-probability support across TRTLLM request handlers and test utilities. A new static helper method extracts logprob information from generated tokens, request option handling propagates logprobs configuration into sampling parameters, and test utilities add payload types and builders to validate logprobs in responses.

Changes

Cohort / File(s)	Summary
Handler Implementation components/src/dynamo/trtllm/request_handlers/handler_base.py	Added _extract_logprobs() static helper to extract logprob information for newly produced tokens. Extended generate_locally to propagate output_options.logprobs and output_options.prompt_logprobs into SamplingParams (clamping logprobs to at least 1), then invoke extraction and attach log_probs and top_logprobs to outgoing payloads.
Test Configuration tests/serve/test_trtllm.py	Added two new TRTLLMConfig test entries ("aggregated_logprobs" and "disaggregated_logprobs") with dedicated script names, model ports, and request payloads exercising logprobs and top_logprobs combinations. Extended completions test path with logprobs=3 parameter.
Test Payload Utilities tests/utils/payload_builder.py, tests/utils/payloads.py	Updated chat_payload() and completion_payload() to accept logprobs parameters and conditionally return new payload types (ChatPayloadWithLogprobs, CompletionPayloadWithLogprobs). Added three new payload validation classes that extend base payloads to assert and validate logprobs structure in responses (token, logprob, top_logprobs fields).

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• handler_base.py: Review the _extract_logprobs() extraction logic for token alignment and edge cases (None logprobs, list conversion).
• payload_builder.py and payloads.py: Verify conditional payload type selection and validation assertions match expected response structures.
• test_trtllm.py: Confirm new test configurations cover intended logprobs/top_logprobs combinations and request payload construction.

Poem

🐰 A token hops with logprobs bright,
Extracted now from TRTLLM's light,
Top choices ranked in every span,
Validated by our testing plan,
Probabilities bloom where logs once sat! 🌟

Pre-merge checks

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.55% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description follows the template structure with all required sections completed: Overview explains the change, Details lists specific modifications with file references, Where should the reviewer start section guides reviewers to key files, and Related Issues section references issue #4684.
Title check	✅ Passed	The title 'feat: Add logprobs support to TRTLLM backend' accurately and concisely summarizes the main change: adding log probability support to the TRTLLM backend, which is reflected across all modified files.

---

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.}

[rmccorm4]

/ok to test 111d0e8

[rmccorm4]

@esoba please fix the failing precommit check: https://github.com/ai-dynamo/dynamo/actions/runs/19946850289/job/57198198828?pr=4759 - just some linting

[tanmayv25]

The code changes look good to me. But there are still some code formatting errors showing up pre-commit runs.

[tanmayv25]

Can you fix the conflict in payload_builder.py and failing DCO test?
https://github.com/ai-dynamo/dynamo/pull/4759/checks?check_run_id=57495128760