[https://nvbugs/6196391][fix] Carryover disagg TTFT improvements

[brb-nv]

Description

This MR carries over following fixes to main:

• [None][fix] avoid duplicate harmony tokenization & populate perf_metrics #14506
  a) Avoid retokenization on gen server in disagg - just pass them from ctx -> gen.
  b) Enable populating perf metrics.
• [https://nvbugs/6196391][fix] Eliminate O(N^2) memcpy in findBlocksInReuseTreeByBlockKey #14499
  a) Avoid O(N^2) memcpy in kv cache reuse function used by disagg codepath.

Test Coverage

$ pytest tests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py -s -v
$ pytest tests/unittest/llmapi/apps/_test_openai_chat_harmony.py::test_prompt_token_ids_skips_retokenization -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)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• 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

• New Features

  • Enabled performance metrics collection for requests when appropriate environment conditions are set.

• Improvements

  • Enhanced streaming response handling with improved timing and metrics extraction.
  • Optimized token processing to reuse pre-tokenized inputs when available, with expanded error logging for troubleshooting.

[brb-nv]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR optimizes BlockKey construction in the C++ KVCache manager and enhances the OpenAI server with conditional performance metrics collection and improved streaming response handling with token preparation caching.

Changes

BlockKey Construction Optimization

Layer / File(s)	Summary
BlockKey reuse tree optimization cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp	findBlocksInReuseTreeByBlockKey now emplaces BlockKey objects per token block with selective field copying instead of repeatedly copying the entire original blockKey, reducing per-iteration work.

OpenAI Server Metrics and Streaming

Layer / File(s)	Summary
Perf metrics collection enablement tensorrt_llm/serve/openai_server.py	Both openai_chat and chat_harmony enable per-request perf metric collection when TRTLLM_KVCACHE_TIME_OUTPUT_PATH is set by setting sampling_params.return_perf_metrics = True.
Token preparation with pre-tokenization caching tensorrt_llm/serve/openai_server.py	chat_harmony now reuses pre-tokenized request.prompt_token_ids when available, falling back to harmony adapter conversion with expanded logging on failure.
Streaming response flow with metrics extraction tensorrt_llm/serve/openai_server.py	Streaming generator in chat_harmony stamps first response with steady clock timestamp, yields post-processed chunks for all responses, emits [DONE] terminator, and calls _extract_metrics after completion.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change as carrying over TTFT improvements with specific issue reference, directly matching the commit descriptions.
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.
Description check	✅ Passed	The PR description references upstream PRs (14506, 14499) with clear objectives, includes specific test commands, and completes the provided checklist comprehensively.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 2

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

Inline comments:
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 1638-1640: The code currently logs full request payloads
(request.messages and request) and uses a bare except; change the logger.error
calls that reference request, request.messages and tools_dict to only emit safe
identifiers and a redacted summary (e.g., log request.id, request.model, and a
short redacted messages summary or list of tool names instead of full contents)
and replace the bare exception handling with a specific exception catch for
tokenization errors (e.g., catch TokenizationError or ValueError as e in the
tokenization block), log the error with request id/model and redacted summary,
and avoid printing any sensitive message/tool arguments; update the logger.error
calls and exception clause where the symbols logger, request, request.messages,
and tools_dict appear.
- Around line 1606-1608: In openai_server.py, replace the two bare "except:"
blocks with explicit handling: catch asyncio.CancelledError and re-raise it
(except asyncio.CancelledError: raise), then catch other errors with "except
Exception as e:" and log the traceback (logger.error(traceback.format_exc()) or
logger.exception(...)) before re-raising; also change the existing "raise e"
usage to a plain "raise" so the original traceback is preserved. Target the bare
except blocks and the "raise e" instance in the file and apply these exact
exception patterns to preserve cancellation semantics and original tracebacks.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4d26cc62-5cdd-48c8-aff4-6fa8bfa156d5

📥 Commits

Reviewing files that changed from the base of the PR and between f6ba936 and 99d48dc.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• tensorrt_llm/serve/openai_server.py

[tensorrt-cicd]

PR_Github #50929 [ run ] triggered by Bot. Commit: 99d48dc Link to invocation

[tensorrt-cicd]

PR_Github #50929 [ run ] completed with state SUCCESS. Commit: 99d48dc
/LLM/main/L0_MergeRequest_PR pipeline #40390 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51113 [ run ] triggered by Bot. Commit: c017564 Link to invocation

[SimengLiu-nv]

LGTM.

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51123 [ run ] triggered by Bot. Commit: dbd2a4c Link to invocation

[tensorrt-cicd]

PR_Github #51113 [ run ] completed with state ABORTED. Commit: c017564

Link to invocation

[tensorrt-cicd]

PR_Github #51123 [ run ] completed with state SUCCESS. Commit: dbd2a4c
/LLM/main/L0_MergeRequest_PR pipeline #40561 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51204 [ run ] triggered by Bot. Commit: dbd2a4c Link to invocation

[tensorrt-cicd]

PR_Github #51204 [ run ] completed with state FAILURE. Commit: dbd2a4c
/LLM/main/L0_MergeRequest_PR pipeline #40631 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51212 [ run ] triggered by Bot. Commit: dbd2a4c Link to invocation

[tensorrt-cicd]

PR_Github #51212 [ run ] completed with state SUCCESS. Commit: dbd2a4c
/LLM/main/L0_MergeRequest_PR pipeline #40636 completed with status: 'SUCCESS'

CI Report

Link to invocation