[None][feat] Add Gemma4 MTP assistant support#15833
Draft
lfr-0531 wants to merge 9 commits into
Draft
Conversation
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #56933 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #56933 [ run ] completed with state
|
lfr-0531
force-pushed
the
user/fanrongl/gemma4-mtp-support
branch
from
July 2, 2026 15:03
346825e to
49e054e
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57206 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57206 [ run ] completed with state
|
lfr-0531
force-pushed
the
user/fanrongl/gemma4-mtp-support
branch
2 times, most recently
from
July 7, 2026 01:56
ca4cd19 to
4b170de
Compare
lfr-0531
force-pushed
the
user/fanrongl/gemma4-mtp-support
branch
from
July 21, 2026 15:37
4b170de to
c250ca1
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60734 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60734 [ run ] completed with state
|
Add standalone Gemma4 assistant config, model, weight loading, and target KV-cache sharing for two-model MTP decoding. Support masked assistant logits and CUDA graph execution, and add structural coverage for assistant configuration and KV source mapping. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Preserve Gemma4 assistant hidden-state and position semantics across multi-token drafting, and capture target hidden states for verification. Refresh FlashInfer paged-prefill graph state after request turnover and add regression coverage for drafting and CUDA graph replay. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Normalize the Gemma4 assistant CUDA graph key across capture and runtime lookup so the full drafting loop replays its captured graph instead of silently falling back to eager execution. Add regression coverage for the first-draft state transition. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Reuse the existing linear drafting loop through model capabilities, enforce assistant configuration invariants, and reject unsupported shared-KV combinations.\n\nFix vocab-parallel ordered logits, shared target KV cache ownership and budgeting, CUDA graph capacity, and FlashInfer KV pool identity. Add focused regressions for configuration, TP logits, graph sizing, cache budgets, and same-head-dimension pools. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
lfr-0531
force-pushed
the
user/fanrongl/gemma4-mtp-support
branch
from
July 22, 2026 15:48
e132e69 to
925108f
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61013 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61013 [ run ] completed with state
|
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Gemma4 provides separate assistant checkpoints for multi-token prediction (MTP), but the TensorRT-LLM PyTorch backend did not recognize the assistant configuration or support running these checkpoints as a two-model speculative decoding setup.
This change adds a compatibility configuration and model implementation for
Gemma4AssistantForCausalLM, including target hidden-state and embedding fusion, target KV-cache sharing, ordered/centroid-masked logits, weight mapping, and CUDA Graph-safe position handling. It also updates the two-model MTP setup to load the requested draft checkpoint, preserves all four assistant layers, avoids allocating a duplicate full-size KV cache, and forwards speculative metadata through the Gemma4 multimodal wrapper.The changes span configuration, model execution, checkpoint loading, and executor resource setup because the assistant model needs coordinated support across those existing layers. There are no C++ changes, public API changes, or new dependencies. Existing Gemma4 execution is preserved; the new behavior is selected only for
gemma4_assistantcheckpoints.Related PRs: None.
Test Coverage
pre-commit run --fileson all ten changed files.pytest -q tests/unittest/_torch/modeling/test_modeling_gemma4.py -k 'TestGemma4ModelInstantiation or TestGemma4Assistant'(15 passed).MTP Results
Each value is
acceptance length / E2E speedup(batch size = 1).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.