Repository navigation
[None][fix] Propagate decoder extra_attrs in Gemma4 multimodal wrapper - #19511
pamelap-nvidia wants to merge 2 commits into
Conversation
Gemma4Attention registers itself into the inner ModelConfig.extra_attrs
that is passed to Gemma4ForCausalLM, but PyTorchModelEngine.model_forward
publishes the outer wrapper's model_config.extra_attrs to the thread-local
before running the model. The two are different objects, so attn_layers
never reaches the thread-local and extract_extra_attrs asserts
"Attention layer is not registered".
Eager execution is unaffected because the custom inplace attention op is
gated on is_torch_compiling() or is_in_breakable_cuda_graph(), so the
lookup only runs under CUDA graph capture. The visible symptom is that
both prefill_cuda_graph_backend=piecewise and =breakable fail to
initialize on Gemma4.
Propagating the inner config's extra_attrs to the outer one after the
decoder is constructed fixes both paths.
Measured on gemma-4-E4B-it (BF16, TP1, 1x RTX PRO 6000 Blackwell, SM120)
with prefill_cuda_graph_backend=breakable, ISL 25 / OSL 5:
concurrency 1: TTFT 48.3 -> 14.5 ms, request latency 83.0 -> 50.5 ms,
throughput 11.9 -> 19.4 req/s
concurrency 4: TTFT 93.3 -> 23.2 ms, request latency 129.8 -> 61.3 ms,
throughput 30.2 -> 63.5 req/s
Inter-token latency is unchanged, as expected: decode was already CUDA
graphed and is memory-bandwidth bound. The gain is entirely prefill
kernel launch overhead, which nsys showed at 1426 launches and 83% GPU
idle per context forward step.
Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
|
Understand this PR’s impact Explore downstream dependencies and potential security impact with Blast Radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. Walkthrough
ChangesGemma4MM configuration propagation
Priority: ➖ Normal Estimated code review effort: 1 (Trivial) | ~5 minutes Change: Bug fix Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
|
/bot run |
|
PR_Github #74865 [ run ] triggered by Bot. Commit: |
|
PR_Github #74865 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The measurements predate the 2026-09-21 merge of main into this branch, so they were taken on a head where this aliasing did not yet exist. Please re-run prefill_cuda_graph_backend=breakable and =piecewise on current main without this line — if initialization succeeds and the TTFT delta reproduces, the PR can be closed. If it still fails with "Attention layer is not registered", the real defect is elsewhere, and the fix should ship with a test asserting attn_layers is present in the published extra_attrs after construction.
| # PyTorchModelEngine.model_forward). Without propagating them, any CUDA | ||
| # graph capture path that takes the custom inplace attention op fails | ||
| # with "Attention layer is not registered". | ||
| self.model_config.extra_attrs.update(llm_model_config.extra_attrs) |
There was a problem hiding this comment.
This is a no-op at the current head: self.model_config.extra_attrs and llm_model_config.extra_attrs are already the same dict object, so this is d.update(d).
Chain on this branch: line 1078 rebinds model_config_cp.extra_attrs to the top-level model_config.extra_attrs, and get_sub_model_config line 909 rebinds sub_config.extra_attrs to that same dict, so Gemma4Attention registers straight into it. Both of those lines came from #17408 (merged 2026-09-16), which this branch picked up in the 2026-09-21 merge of main — after this commit was written on 2026-09-01.
As written, the line is dead code and the comment describes a mechanism that no longer holds. Please drop both.
Dev Engineer Review
The change propagates the decoder
ModelConfig.extra_attrsto the outer Gemma4 configuration after decoder construction. This restoresattn_layerslookup during CUDA graph capture for piecewise and breakable prefill backends. Eager execution and public declarations remain unchanged. Reported Gemma4 measurements show lower TTFT and request latency, higher throughput at concurrency 1 and 4, and unchanged inter-token latency.QA Engineer Review
No test changes.
Per-File QA Perspective
tensorrt_llm/_torch/models/modeling_gemma4mm.py: Verify Gemma4 multimodal CUDA graph initialization with both prefill backends. Verifyattn_layersregistration and confirm that eager execution remains unchanged.Description
Gemma4Attention registers itself into the inner ModelConfig.extra_attrs that is passed to Gemma4ForCausalLM, but PyTorchModelEngine.model_forward publishes the outer wrapper's model_config.extra_attrs to the thread-local before running the model. The two are different objects, so attn_layers never reaches the thread-local and extract_extra_attrs asserts "Attention layer is not registered".
Eager execution is unaffected because the custom inplace attention op is gated on is_torch_compiling() or is_in_breakable_cuda_graph(), so the lookup only runs under CUDA graph capture. The visible symptom is that both prefill_cuda_graph_backend=piecewise and =breakable fail to initialize on Gemma4.
Propagating the inner config's extra_attrs to the outer one after the decoder is constructed fixes both paths.
Measured on gemma-4-E4B-it (BF16, TP1, 1x RTX PRO 6000 Blackwell, SM120) with prefill_cuda_graph_backend=breakable, ISL 25 / OSL 5:
concurrency 1: TTFT 48.3 -> 14.5 ms, request latency 83.0 -> 50.5 ms,
throughput 11.9 -> 19.4 req/s
concurrency 4: TTFT 93.3 -> 23.2 ms, request latency 129.8 -> 61.3 ms,
throughput 30.2 -> 63.5 req/s
Inter-token latency is unchanged, as expected: decode was already CUDA graphed and is memory-bandwidth bound. The gain is entirely prefill kernel launch overhead, which nsys showed at 1426 launches and 83% GPU idle per context forward step.
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.