Skip to content

[None][feat] Prefer V2 transceiver backend for Gemma and Llama#16787

Open
moraxu wants to merge 2 commits into
NVIDIA:mainfrom
moraxu:dev-mguzek-use-transceiver-v2-for-llama-and-gemma
Open

[None][feat] Prefer V2 transceiver backend for Gemma and Llama#16787
moraxu wants to merge 2 commits into
NVIDIA:mainfrom
moraxu:dev-mguzek-use-transceiver-v2-for-llama-and-gemma

Conversation

@moraxu

@moraxu moraxu commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • Added model-level transceiver runtime preference hooks for Gemma 3 (Gemma3ForCausalLM) and Llama (LlamaForCausalLM) via new get_preferred_transceiver_runtime(...) classmethods, both returning "PYTHON" with return type Optional[Literal["CPP", "PYTHON"]].
  • Updated typing imports/annotations for consistent API signatures (Any, Literal).
  • Extended verification via unit + integration tests to ensure the preferred backend is selected when transceiver runtime is set to "auto", and adjusted an integration disaggregated KV-cache timing test to exercise the “legacy C++ DataTransceiver CSV format” path by forcing UCX cache transceiver selection.

QA Engineer Review

  • Test code added/modified:
    • tests/unittest/llmapi/test_llm_args.py
      • Added TestTransceiverRuntimeAutoResolution.test_llama_and_gemma_model_preferences_adopted(self, model_cls) to assert _resolve_transceiver_runtime_auto(...) maps cache_transceiver_config.transceiver_runtime == "auto" to "PYTHON" for both LlamaForCausalLM and Gemma3ForCausalLM.
    • tests/integration/defs/disaggregated/test_disaggregated.py
      • Modified test_disaggregated_kv_cache_time_output to force UCX cache transceiver selection (sets TRTLLM_USE_NIXL_KVCACHE=0 and TRTLLM_USE_UCX_KVCACHE=1) while validating the legacy C++ DataTransceiver CSV format behavior.
  • Coverage in tests/integration/test_lists/ (CI/QA lists): Not applicable—no test-list files were modified in this PR (no direct evidence available here).
  • Verdict: needs follow-up.

Description

Prefer V2 transceiver backend for Gemma and Llama

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)

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

@moraxu
moraxu requested review from a team as code owners July 23, 2026 07:46
@moraxu
moraxu requested review from QiJune and lori-ren July 23, 2026 07:46
@coderabbitai

coderabbitai Bot commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 59833560-27d3-4e74-b290-6b36ea4eb087

📥 Commits

Reviewing files that changed from the base of the PR and between 3b8aeaa and df03c3a.

📒 Files selected for processing (1)
  • tests/integration/defs/disaggregated/test_disaggregated.py

Walkthrough

Llama and Gemma causal language model classes now declare Python as their preferred transceiver runtime. Tests cover automatic resolution, and the disaggregated KV-cache timing test explicitly selects UCX for metrics collection.

Changes

Transceiver runtime preferences

Layer / File(s) Summary
Model runtime preference APIs
tensorrt_llm/_torch/models/modeling_gemma3.py, tensorrt_llm/_torch/models/modeling_llama.py
Adds typed get_preferred_transceiver_runtime classmethods to both model classes, returning "PYTHON".
Automatic runtime resolution coverage
tests/unittest/llmapi/test_llm_args.py
Adds parameterized coverage confirming "auto" resolves to "PYTHON" for Llama and Gemma models.
Disaggregated metrics environment
tests/integration/defs/disaggregated/test_disaggregated.py
Forces UCX cache transceiver usage while running the KV-cache timing metrics test.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: reasonsolo

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Description check ⚠️ Warning The description states the goal, but the required Test Coverage section is empty and does not list any safeguards for the change. Add the relevant tests under Test Coverage and briefly explain how they validate the Gemma and Llama backend preference behavior.
Docstring Coverage ⚠️ Warning Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title is concise and accurately summarizes the main change: preferring the V2 transceiver backend for Gemma and Llama.
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.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_gemma3.py (1)

279-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document both resolver-facing runtime-preference hooks.

Add matching Google-style docstrings documenting pretrained_config and the returned runtime sentinel:

  • tensorrt_llm/_torch/models/modeling_gemma3.py#L279-L284: document Gemma3ForCausalLM.get_preferred_transceiver_runtime.
  • tensorrt_llm/_torch/models/modeling_llama.py#L1137-L1142: document LlamaForCausalLM.get_preferred_transceiver_runtime.
🤖 Prompt for 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.

In `@tensorrt_llm/_torch/models/modeling_gemma3.py` around lines 279 - 284, Add
matching Google-style docstrings to
Gemma3ForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_gemma3.py lines 279-284 and
LlamaForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_llama.py lines 1137-1142, documenting the
pretrained_config parameter and the returned "PYTHON" runtime sentinel while
preserving the existing behavior.

Source: Coding guidelines

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

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma3.py`:
- Around line 279-284: Add matching Google-style docstrings to
Gemma3ForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_gemma3.py lines 279-284 and
LlamaForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_llama.py lines 1137-1142, documenting the
pretrained_config parameter and the returned "PYTHON" runtime sentinel while
preserving the existing behavior.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 71444888-d1ad-47d4-8c6b-dcf66b86a774

📥 Commits

Reviewing files that changed from the base of the PR and between 93efd33 and 97231af.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/models/modeling_gemma3.py
  • tensorrt_llm/_torch/models/modeling_llama.py
  • tests/unittest/llmapi/test_llm_args.py

@Shixiaowei02
Shixiaowei02 requested a review from nv-xtf July 23, 2026 08:03
@Shixiaowei02

Copy link
Copy Markdown
Collaborator

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

@nv-xtf

nv-xtf commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

I checked this, and the situation is slightly different from PR 16482. The existing Kimi test used DEFAULT, which caused the test harness to fall back to UCX/CPP, so it had to be explicitly changed to NIXL + PYTHON. Llama and Gemma already have corresponding NIXL + PYTHON disaggregated accuracy tests, so I don’t think similar test changes are needed here.
The existing test_auto_dtype cases will continue to exercise UCX/CPP and can be retained as legacy coverage.

Signed-off-by: Michal Guzek <mguzek@nvidia.com>
@Shixiaowei02
Shixiaowei02 force-pushed the dev-mguzek-use-transceiver-v2-for-llama-and-gemma branch from 97231af to 3b8aeaa Compare July 23, 2026 13:02
@Shixiaowei02

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61293 [ run ] triggered by Bot. Commit: 3b8aeaa Link to invocation

Signed-off-by: Michal Guzek <mguzek@nvidia.com>
@moraxu
moraxu requested review from a team as code owners July 23, 2026 17:01
@moraxu

moraxu commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

I only updated one test using TinyLlama in tests/integration/defs/disaggregated/test_disaggregated.py

@moraxu

moraxu commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61343 [ run ] triggered by Bot. Commit: df03c3a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61293 [ run ] completed with state ABORTED. Commit: 3b8aeaa
/LLM/main/L0_MergeRequest_PR pipeline #49526 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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61343 [ run ] completed with state FAILURE. Commit: df03c3a
/LLM/main/L0_MergeRequest_PR pipeline #49570 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

@Shixiaowei02

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61484 [ run ] triggered by Bot. Commit: df03c3a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61484 [ run ] completed with state FAILURE. Commit: df03c3a
/LLM/main/L0_MergeRequest_PR pipeline #49704 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

@moraxu

moraxu commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator Author

I think my modified test errors here: https://prod.blsm.nvidia.com/sw-tensorrt-top-1/job/LLM/job/main/job/L0_MergeRequest_PR/49570/testReport/junit/A10-PyTorch-3.disaggregated/test_disaggregated/test_disaggregated_perf_metrics_TinyLlama_1_1B_Chat_v1_0_/ - I will fix it shortly

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants