[None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4 - #17396
[None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4#17396erictsai-nv wants to merge 2 commits into
Conversation
|
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 (5)
WalkthroughGemma3 and Gemma4 now enable KV cache manager v2 and select the Python transceiver runtime. Gemma4 defaults to the FLASHINFER backend. Gemma4 multimodal defaults also enable KV cache manager v2. Unit tests cover these defaults. ChangesGemma runtime defaults
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/modeling_gemma3.py (2)
296-301: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType and document the transceiver runtime hook.
get_preferred_transceiver_runtimehas no parameter or return annotations and no docstring. Match the base hook signature, returnstr, and document why Gemma3 requires the"PYTHON"runtime.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 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 296 - 301, Type and document the get_preferred_transceiver_runtime classmethod to match the base hook signature, including annotations for cls, pretrained_config, and the str return value. Add a Google-style docstring explaining why Gemma3 requires the "PYTHON" transceiver runtime, while preserving the existing return value.Source: Coding guidelines
288-294: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the required annotations and docstring to this hook.
get_model_defaultshas an untypedllm_argsparameter, a baredictreturn annotation, and no Google-style docstring. Use the type aliases from the base hook and document that this method enables the V2 KV-cache manager.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 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 288 - 294, Update Gemma3ForCausalLM.get_model_defaults to use the base hook’s type aliases for the llm_args parameter and return value, and add a Google-style docstring describing that it enables the V2 KV-cache manager. Preserve the existing defaults unchanged.Source: Coding guidelines
tensorrt_llm/_torch/models/modeling_gemma4.py (1)
1268-1274: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType and document the transceiver runtime hook.
get_preferred_transceiver_runtimehas no parameter or return annotations and no docstring. Match the base hook signature, returnstr, and document why Gemma4 requires the"PYTHON"runtime.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 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_gemma4.py` around lines 1268 - 1274, Update the Gemma4 classmethod get_preferred_transceiver_runtime to match the base hook’s annotated signature, including the pretrained_config parameter type and a str return annotation. Add a Google-style docstring documenting that Gemma4 requires the "PYTHON" transceiver runtime.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.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 1263-1266: Update
Gemma4ForConditionalGeneration.get_model_defaults to delegate to
Gemma4ForCausalLM.get_model_defaults(llm_args), preserving the multimodal
wrapper’s existing defaults while inheriting use_kv_cache_manager_v2=True.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma3.py`:
- Around line 296-301: Type and document the get_preferred_transceiver_runtime
classmethod to match the base hook signature, including annotations for cls,
pretrained_config, and the str return value. Add a Google-style docstring
explaining why Gemma3 requires the "PYTHON" transceiver runtime, while
preserving the existing return value.
- Around line 288-294: Update Gemma3ForCausalLM.get_model_defaults to use the
base hook’s type aliases for the llm_args parameter and return value, and add a
Google-style docstring describing that it enables the V2 KV-cache manager.
Preserve the existing defaults unchanged.
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 1268-1274: Update the Gemma4 classmethod
get_preferred_transceiver_runtime to match the base hook’s annotated signature,
including the pretrained_config parameter type and a str return annotation. Add
a Google-style docstring documenting that Gemma4 requires the "PYTHON"
transceiver runtime.
🪄 Autofix
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: e6c1195c-2109-49d9-a4a6-8f2b78c7fc8c
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_gemma4.pytests/unittest/_torch/modeling/test_modeling_gemma3.pytests/unittest/_torch/modeling/test_modeling_gemma4.py
Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…t_model_defaults Signed-off-by: Eric Tsai <ertsai@nvidia.com>
5b8bd57 to
5bfdca3
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #64567 [ run ] triggered by Bot. Commit: |
|
PR_Github #64567 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The mechanism and the [None] tag are fine for a default flip, but the coverage doesn't match what the description claims.
Model defaults are keyed by the checkpoint's architectures[0] (model_loader.py:415, :452). Gemma3 4B/12B/27B ship as Gemma3ForConditionalGeneration, which resolves to Gemma3VLM (modeling_gemma3vl.py:178) — that class has neither hook, so it keeps V1. The description says "all Gemma3 variants" and lists 27B as validated on the v2cpp arm; either the validation used a text-only checkpoint, or the resolution went somewhere I'm not seeing. Please reconcile — add the hook to Gemma3VLM if VLM V2 was actually what you validated, otherwise narrow the description to Gemma3ForCausalLM.
docs/source/features/kvcache.md:117 currently documents only hybrid Mamba as selecting V2 under auto. Since this changes OOTB behavior for two shipped model families, worth a sentence there.
Also worth confirming: [None] is right for a default flip that changes runtime behavior for released models — if there's a Gemma-V2-enablement JIRA, cite it so the rollback story is traceable.
| """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" | ||
| return { | ||
| "attn_backend": "FLASHINFER", | ||
| "kv_cache_config": { |
There was a problem hiding this comment.
The multimodal classes get the V2 default but not get_preferred_transceiver_runtime(). In disagg with NIXL, _resolve_transceiver_runtime_auto leaves the runtime at CPP, and _resolve_kv_cache_manager_v2_auto (llm_utils.py:610-616) then downgrades this default back to V1 with only an INFO log — so Gemma4 MM disagg silently keeps V1 while Gemma4 text gets V2. If MM disagg is out of scope here ([modeling_gemma4mm.py:968](https://github.com/NVIDIA/TensorRT-LLM/pull/17396/files#diff-c199d40bc9ba6c34e1a2ed072172254df5e14e0d33790fae775615591b4a744fR968) rejects the mm-disagg env path), say so in a comment; otherwise add the same -> "PYTHON" override on Gemma4MultimodalModelBase so both paths agree.
| vocab_size=model_config.pretrained_config.vocab_size) | ||
|
|
||
| @classmethod | ||
| def get_model_defaults(cls, llm_args) -> dict: |
There was a problem hiding this comment.
This only covers Gemma3ForCausalLM. AutoModelForCausalLM._resolve_class keys off architectures[0], so google/gemma-3-{4b,12b,27b}-it (Gemma3ForConditionalGeneration → Gemma3VLM, modeling_gemma3vl.py:178) hits neither get_model_defaults nor get_preferred_transceiver_runtime and stays on V1. That contradicts "all Gemma3 variants" plus the 27B validation in the description. If the VLM path is deliberately excluded (e.g. V2 not validated with the bidirectional image mask), please state that; otherwise add both hooks to Gemma3VLM.
| torch.testing.assert_close(attention_mask, expected_attention_mask) | ||
|
|
||
|
|
||
| def test_gemma3_model_defaults_select_v2(): |
There was a problem hiding this comment.
These four new tests just re-assert the dict literal on the line above them — they'll pass whether or not the default is ever honored. The behavior that can actually regress is the resolution path: given use_kv_cache_manager_v2="auto" and a Gemma3/Gemma4 config, does load_config_and_apply_defaults land on V2, and does an explicit user False still win? apply_model_defaults_to_llm_args / _resolve_kv_cache_manager_v2_auto are importable and need no GPU — one test through them would be worth all four of these.
3 participants
Dev Engineer Review
KVCacheManagerV2throughget_model_defaults().FLASHINFERbackend default."PYTHON"transceiver runtime for disaggregated NIXL serving.KVCacheManagerV2.QA Engineer Review
test_gemma3_model_defaults_select_v2().test_gemma3_prefers_python_transceiver().test_gemma4_model_defaults_select_v2().test_gemma4_prefers_python_transceiver().test-db/orqa/.Description
Enables
KVCacheManagerV2by default for Gemma3 and Gemma4.The change adds
get_model_defaults()toGemma3ForCausalLMand extends the existing one inGemma4ForCausalLMreturningkv_cache_config.use_kv_cache_manager_v2=True. Both classes also declareget_preferred_transceiver_runtime() -> "PYTHON"so that disaggregated serving over NIXL retains V2 instead of falling back to V1. This plugs into the per-model auto-selection infrastructure added in #15823, following the pattern established by GPT-OSS (#16942).Note: Gemma4 hybrid attention was already unconditionally routed to V2 via
_non_hybrid_kv_cache_manager_cls(_util.py:83-88); this change makes the selection consistent at thellm_argslevel and extends the default to non-hybrid Gemma4 and all Gemma3 variants.Since #14047 landed,
TLLM_KV_CACHE_MANAGER_V2_BACKENDdefaults tocpp, so this PR moves Gemma onto the C++ V2 core.Validation
H100 (agg + disagg) and B200 (NVFP4 MMMU): Gemma3 1B/27B and Gemma4 26B-A4B all passed across v1/v2py/v2cpp arms, accuracy within run noise.
Test Coverage
New:
test_gemma3_model_defaults_select_v2,test_gemma3_prefers_python_transceiver,test_gemma4_model_defaults_select_v2,test_gemma4_prefers_python_transceiver.Existing Gemma3/4 integration tests that do not set
use_kv_cache_manager_v2explicitly now run through the"auto"resolution path and exercise V2 by default.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.