[None][feat] Prefer V2 transceiver backend for Gemma and Llama - #16787
[None][feat] Prefer V2 transceiver backend for Gemma and Llama#16787moraxu wants to merge 4 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughLlama and Gemma causal language model classes now declare Python as their preferred transceiver runtime. Tests cover automatic resolution, while disaggregated tests explicitly select UCX or the C++ serialization protocol. ChangesTransceiver runtime preferences
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_gemma3.py (1)
279-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument both resolver-facing runtime-preference hooks.
Add matching Google-style docstrings documenting
pretrained_configand the returned runtime sentinel:
tensorrt_llm/_torch/models/modeling_gemma3.py#L279-L284: documentGemma3ForCausalLM.get_preferred_transceiver_runtime.tensorrt_llm/_torch/models/modeling_llama.py#L1137-L1142: documentLlamaForCausalLM.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
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_llama.pytests/unittest/llmapi/test_llm_args.py
|
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. |
97231af to
3b8aeaa
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61293 [ run ] triggered by Bot. Commit: |
I only updated one test using |
|
/bot run --disable-fail-fast |
|
PR_Github #61343 [ run ] triggered by Bot. Commit: |
|
PR_Github #61293 [ run ] completed with state
|
|
PR_Github #61343 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61484 [ run ] triggered by Bot. Commit: |
|
PR_Github #61484 [ run ] completed with state
|
|
I think my modified test errors here: |
df03c3a to
763bec4
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/disaggregated/test_disaggregated.py (1)
1573-1580: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPair send and receive files by their shared prefix.
Selecting
send_files[0]andrecv_files[0]independently can select different<instanceId>_<rank>files when multiple outputs exist or one side is missing. Require a common prefix and select both files from that pair.Proposed fix
- send_files = sorted(f for f in os.listdir(output_path) - if f.endswith("_send.csv")) - recv_files = sorted(f for f in os.listdir(output_path) - if f.endswith("_recv.csv")) - assert send_files, f"no *_send.csv in {output_path}: {os.listdir(output_path)}" - assert recv_files, f"no *_recv.csv in {output_path}: {os.listdir(output_path)}" - send_file = os.path.join(output_path, send_files[0]) - recv_file = os.path.join(output_path, recv_files[0]) + send_files = { + f.rsplit("_", 1)[0]: f + for f in os.listdir(output_path) + if f.endswith("_send.csv") + } + recv_files = { + f.rsplit("_", 1)[0]: f + for f in os.listdir(output_path) + if f.endswith("_recv.csv") + } + common_prefixes = sorted(send_files.keys() & recv_files.keys()) + assert common_prefixes, ( + f"no matching send/recv CSV pair in {output_path}: " + f"{os.listdir(output_path)}" + ) + prefix = common_prefixes[0] + send_file = os.path.join(output_path, send_files[prefix]) + recv_file = os.path.join(output_path, recv_files[prefix])🤖 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 `@tests/integration/defs/disaggregated/test_disaggregated.py` around lines 1573 - 1580, Update the file-selection logic near the send_file and recv_file assignments to pair files by their shared <instanceId>_<rank> prefix. Derive prefixes from send_files and recv_files, require at least one prefix present in both sets, and select the send and receive filenames from the same common prefix instead of independently choosing index zero.
🤖 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.
Outside diff comments:
In `@tests/integration/defs/disaggregated/test_disaggregated.py`:
- Around line 1573-1580: Update the file-selection logic near the send_file and
recv_file assignments to pair files by their shared <instanceId>_<rank> prefix.
Derive prefixes from send_files and recv_files, require at least one prefix
present in both sets, and select the send and receive filenames from the same
common prefix instead of independently choosing index zero.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 481cddc7-3577-4c1c-8819-184f0f236c2e
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_llama.pytests/integration/defs/disaggregated/test_disaggregated.pytests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/models/modeling_gemma3.py
- tensorrt_llm/_torch/models/modeling_llama.py
- tests/unittest/llmapi/test_llm_args.py
|
PR_Github #61614 [ run ] triggered by Bot. Commit: |
|
PR_Github #61614 [ run ] completed with state
|
Signed-off-by: Michal Guzek <mguzek@nvidia.com>
Signed-off-by: Michal Guzek <mguzek@nvidia.com>
Signed-off-by: Michal Guzek <mguzek@nvidia.com>
763bec4 to
8b0a39a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61818 [ run ] triggered by Bot. Commit: |
|
PR_Github #61818 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61961 [ run ] triggered by Bot. Commit: |
|
PR_Github #61961 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61973 [ run ] triggered by Bot. Commit: |
|
PR_Github #61973 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62014 [ run ] triggered by Bot. Commit: |
|
PR_Github #62014 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62102 [ run ] triggered by Bot. Commit: |
|
PR_Github #62102 [ run ] completed with state
|
5 participants
Dev Engineer Review
Gemma3ForCausalLM) and Llama (LlamaForCausalLM) now exposeget_preferred_transceiver_runtime(pretrained_config: Any = None) -> Optional[Literal["CPP", "PYTHON"]]and always resolve to"PYTHON".Any,Literal) in the affected model files.test_disaggregated_perf_metricsandtest_disaggregated_kv_cache_time_outputnow force UCX over NIXL viaTRTLLM_USE_NIXL_KVCACHE=0andTRTLLM_USE_UCX_KVCACHE=1.test_disaggregated_single_gpuarbitrary KV transfer cases now force C++-serialized protocol by settingCacheTransceiverConfig(transceiver_runtime="CPP")for the relevant tests.QA Engineer Review
tests/unittest/llmapi/test_llm_args.pyTestTransceiverRuntimeAutoResolution.test_llama_and_gemma_model_preferences_adopted(self, model_cls)to verify_resolve_transceiver_runtime_auto(...)resolves"auto"to"PYTHON"for bothLlamaForCausalLMandGemma3ForCausalLM.tests/integration/defs/disaggregated/test_disaggregated.pytest_disaggregated_perf_metrics(perf_metrics mode) to force UCX-over-NIXL for the timing-metrics assertions.test_disaggregated_kv_cache_time_outputto also force UCX-over-NIXL while validatingTRTLLM_KVCACHE_TIME_OUTPUT_PATH.tests/integration/defs/disaggregated/test_disaggregated_single_gpu.pytest_arbitrary_kv_cache_transferto settransceiver_runtime="CPP".test_arbitrary_kv_cache_transfer_missing_blocksto settransceiver_runtime="CPP".tests/integration/test_lists/(CI/QA lists): Not verified in provided context.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.