[None][feat] Add Gemma4 MTP assistant support by lfr-0531 · Pull Request #15833 · NVIDIA/TensorRT-LLM

lfr-0531 · 2026-07-01T11:18:08Z

Dev Engineer Review

Adds Gemma4 MTP assistant support for the PyTorch backend.
Adds Gemma4AssistantConfig and registers model_type="gemma4_assistant".
Adds assistant checkpoint mapping and Gemma4AssistantForCausalLM.
Implements ordered and centroid-masked logits, hidden-state projections, recurrent drafting, and target KV-cache sharing.
Extends speculative decoding metadata, FlashInfer draft views, VSWA pool routing, and CUDA Graph handling.
Updates model loading to load draft weights for shared-KV configurations.
Updates documentation and examples for supported Gemma4 target/assistant pairs.
The assistant configuration rejects partial KV sharing, MoE blocks, double-wide MLPs, and unsupported input dimensions.
Review the reported CI failures before approval. The available CI summary does not include specific failure details.
Confirm that the PreTrainedConfig base-class changes preserve compatibility with existing Gemma4 unified configurations.
Confirm that trtllm-gen availability and shared/separate KV-cache paths are handled consistently.
No C++ changes, new dependencies, or test-list changes are described in the provided changes.

QA Engineer Review

Test changes

tests/unittest/_torch/attention/test_flashinfer_attention.py
- Added TestFlashInferAttention.test_separate_kv_draft_metadata_uses_draft_manager.
tests/unittest/_torch/modeling/test_modeling_gemma4.py
- Added TestGemma4Assistant.test_assistant_config.
- Added TestGemma4Assistant.test_assistant_rejects_partial_kv_sharing.
- Added TestGemma4Assistant.test_ordered_embedding_combines_vocab_parallel_shards.
- Added TestGemma4Assistant.test_assistant_uses_target_kv_sources.
- Added TestGemma4CUDAGraph.test_shared_kv_draft_view.
- Updated CUDA Graph expectations to use kv_pool_id.
tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
- Added test_mtp_shared_kv_config.
- Added test_mtp_shared_kv_draft_inputs.
tests/unittest/_torch/speculative/hw_agnostic/test_sa.py
- Updated documentation comments only.

Test-list coverage

No tests/integration/test_lists/, test-db/, qa/, or waives.txt changes are described.

Verdict

Needs follow-up. The tests cover configuration, embedding, target KV sources, draft metadata, and shared-KV input preparation. The CI summary reports repeated pipeline failures without specific test results. Confirm the failures and verify coverage for CUDA Graph replay, FlashInfer metadata refresh, and supported Gemma4 checkpoint pairs.

Description

Gemma4 provides dedicated MTP assistant checkpoints, but the TensorRT-LLM PyTorch backend previously did not recognize the gemma4_assistant configuration or support these checkpoints in the MTP drafting workflow.

This PR adds a compatibility configuration and model implementation for Gemma4AssistantForCausalLM, covering checkpoint registration and weight mapping, target embedding and hidden-state fusion, assistant projections, ordered and centroid-masked logits, recurrent drafting, and speculative metadata forwarding through the Gemma4 multimodal wrapper. It also extends the MTP runtime and FlashInfer metadata handling so the Q-only assistant layers can read the target model's KV cache, including correct KV-page views, position handling, request turnover, and CUDA Graph capture and replay. Documentation and examples are updated for all four supported Gemma4 target/assistant checkpoint pairs.

The changes span configuration, model execution, speculative decoding, attention metadata, checkpoint loading, and executor resource setup because KV-cache ownership and drafting state must remain consistent across these layers. The assistant reuses the target KV cache, avoiding allocation of another full-size GPU KV cache. Existing Gemma4 execution is unchanged unless a gemma4_assistant checkpoint is configured. There are no C++ changes, public API changes, or new dependencies.

The current scope supports FlashInfer with linear MTP drafting. Dynamic-tree decoding, draft-length scheduling, suffix-automaton enhancement, rejection sampling, and guided decoding are rejected explicitly. The main implementation risk is shared-KV metadata correctness across context and generation requests; focused unit coverage and end-to-end CUDA Graph testing are included for that path.

Related PRs: None.

Test Coverage

Relevant unit suites:
- tests/unittest/_torch/modeling/test_modeling_gemma4.py
- tests/unittest/_torch/modeling/test_gemma4_multimodal.py
- tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
Validated configuration loading and end-to-end generation for E2B, E4B, 26B-A4B, and 31B assistant checkpoints with MTP draft lengths 1 through 5.
Verified CUDA Graph enabled and disabled execution; no CUDA error, OOM, assertion failure, timeout, or executor error was observed.
Built and installed the native components and verified that runtime imports used this checkout's Python sources and the corresponding installed bindings.

Acceptance Length Results

Model	MTP=1	MTP=2	MTP=3	MTP=4	MTP=5
Gemma4 E2B	1.6628	2.0878	2.3535	2.5226	2.6221
Gemma4 E4B	1.6890	2.1483	2.4509	2.6467	2.7813
Gemma4 26B-A4B	1.7871	2.3962	2.8241	3.1284	3.3558
Gemma4 31B	1.8042	2.4270	2.8803	3.2130	3.4387

E2E Speedup Results

Model	MTP=1	MTP=2	MTP=3	MTP=4	MTP=5
Gemma4 E2B	1.7908x	1.7693x	1.5802x	2.2393x	2.1387x
Gemma4 E4B	1.7688x	1.4682x	2.2993x	2.4579x	2.4926x
Gemma4 26B-A4B	2.0339x	1.7078x	2.4778x	2.7620x	2.6781x
Gemma4 31B	1.6632x	2.4062x	2.7035x	2.8527x	3.1425x

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.

lfr-0531 · 2026-07-01T13:20:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T13:27:10Z

PR_Github #56933 [ run ] triggered by Bot. Commit: 346825e Link to invocation

tensorrt-cicd · 2026-07-01T21:04:11Z

PR_Github #56933 [ run ] completed with state SUCCESS. Commit: 346825e
/LLM/main/L0_MergeRequest_PR pipeline #45738 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

lfr-0531 · 2026-07-02T15:29:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T15:35:34Z

PR_Github #57206 [ run ] triggered by Bot. Commit: a49a1ba Link to invocation

tensorrt-cicd · 2026-07-03T08:46:25Z

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

lfr-0531 · 2026-07-21T16:34:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T16:40:48Z

PR_Github #60734 [ run ] triggered by Bot. Commit: e132e69 Link to invocation

tensorrt-cicd · 2026-07-21T20:38:41Z

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

lfr-0531 · 2026-07-22T15:49:13Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-22T15:56:01Z

PR_Github #61013 [ run ] triggered by Bot. Commit: 925108f Link to invocation

tensorrt-cicd · 2026-07-22T17:12:35Z

PR_Github #61013 [ run ] completed with state FAILURE. Commit: 925108f
/LLM/main/L0_MergeRequest_PR pipeline #49269 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 · 2026-08-01T01:25:33Z

PR_Github #63163 [ run ] triggered by Bot. Commit: 5dcaf06 Link to invocation

tensorrt-cicd · 2026-08-01T06:00:13Z

PR_Github #63163 [ run ] completed with state SUCCESS. Commit: 5dcaf06
/LLM/main/L0_MergeRequest_PR pipeline #51247 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

lfr-0531 · 2026-08-01T14:59:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-01T15:05:31Z

PR_Github #63218 [ run ] triggered by Bot. Commit: 5dcaf06 Link to invocation

tensorrt-cicd · 2026-08-01T15:05:31Z

PR_Github #63218 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

lfr-0531 · 2026-08-02T03:13:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-02T03:20:30Z

PR_Github #63254 [ run ] triggered by Bot. Commit: 5dcaf06 Link to invocation

tensorrt-cicd · 2026-08-02T03:20:31Z

PR_Github #63254 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

lfr-0531 · 2026-08-03T01:58:34Z

/bot run --disable-fail-fast

Allow only the validated shared-KV one-engine path while preserving the existing guard for other FlashInfer combinations. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

lfr-0531 · 2026-08-03T05:00:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-03T05:06:02Z

PR_Github #63356 [ run ] triggered by Bot. Commit: 8df01d7 Link to invocation

tensorrt-cicd · 2026-08-03T09:05:38Z

PR_Github #63356 [ run ] completed with state FAILURE. Commit: 8df01d7
/LLM/main/L0_MergeRequest_PR pipeline #51343 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

lfr-0531 · 2026-08-03T09:09:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-03T09:16:43Z

PR_Github #63405 [ run ] triggered by Bot. Commit: 8df01d7 Link to invocation

tensorrt-cicd · 2026-08-03T13:43:15Z

PR_Github #63405 [ run ] completed with state SUCCESS. Commit: 8df01d7
/LLM/main/L0_MergeRequest_PR pipeline #51384 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

lfr-0531 · 2026-08-03T15:37:14Z

/bot run

tensorrt-cicd · 2026-08-03T15:44:32Z

PR_Github #63481 [ run ] triggered by Bot. Commit: 8df01d7 Link to invocation

lfr-0531 · 2026-08-03T15:51:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-03T15:57:19Z

PR_Github #63484 [ run ] triggered by Bot. Commit: 4207c44 Link to invocation

tensorrt-cicd · 2026-08-03T16:00:48Z

PR_Github #63481 [ run ] completed with state ABORTED. Commit: 8df01d7

Link to invocation

tensorrt-cicd · 2026-08-03T19:43:20Z

PR_Github #63484 [ run ] completed with state SUCCESS. Commit: 4207c44
/LLM/main/L0_MergeRequest_PR pipeline #51454 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

lfr-0531 · 2026-08-04T00:42:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-04T00:49:19Z

PR_Github #63600 [ run ] triggered by Bot. Commit: 4207c44 Link to invocation

tensorrt-cicd · 2026-08-04T00:53:16Z

PR_Github #63599 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: run run run

Link to invocation

tensorrt-cicd · 2026-08-04T01:50:41Z

PR_Github #63600 [ run ] completed with state SUCCESS. Commit: 4207c44
/LLM/main/L0_MergeRequest_PR pipeline #51560 completed with status: 'SUCCESS'

CI Report

Link to invocation

yihwang-nv

Thanks, LGTM!

github-actions Bot assigned lfr-0531 Jul 1, 2026

lfr-0531 force-pushed the user/fanrongl/gemma4-mtp-support branch from 346825e to 49e054e Compare July 2, 2026 15:03

lfr-0531 force-pushed the user/fanrongl/gemma4-mtp-support branch 2 times, most recently from ca4cd19 to 4b170de Compare July 7, 2026 01:56

lfr-0531 force-pushed the user/fanrongl/gemma4-mtp-support branch from 4b170de to c250ca1 Compare July 21, 2026 15:37

lfr-0531 force-pushed the user/fanrongl/gemma4-mtp-support branch from e132e69 to 925108f Compare July 22, 2026 15:48

lfr-0531 force-pushed the user/fanrongl/gemma4-mtp-support branch from cbae5b4 to 224b5a3 Compare July 28, 2026 09:05

lfr-0531 added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 28, 2026

lfr-0531 marked this pull request as ready for review July 28, 2026 09:16

lfr-0531 requested review from a team as code owners July 28, 2026 09:16

lfr-0531 requested review from PerkzZheng, chang-l, kris1025, schetlur-nv and tongyuantongyu July 28, 2026 09:16

yihwang-nv reviewed Aug 2, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

[None][fix] scope FlashInfer one-engine guard

8df01d7

Allow only the validated shared-KV one-engine path while preserving the existing guard for other FlashInfer combinations. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Merge branch 'main' into user/fanrongl/gemma4-mtp-support

4207c44

yihwang-nv approved these changes Aug 4, 2026

View reviewed changes

lfr-0531 merged commit 048ae4a into NVIDIA:main Aug 4, 2026
7 checks passed

Uh oh!

Conversation

lfr-0531 commented Jul 1, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Test changes

Test-list coverage

Verdict

Description

Test Coverage

Acceptance Length Results

E2E Speedup Results

PR Checklist

GitHub Bot Help

Uh oh!

lfr-0531 commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

lfr-0531 commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

lfr-0531 commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

lfr-0531 commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Aug 1, 2026

Uh oh!

tensorrt-cicd commented Aug 1, 2026

Uh oh!

lfr-0531 commented Aug 1, 2026

Uh oh!

tensorrt-cicd commented Aug 1, 2026

Uh oh!

tensorrt-cicd commented Aug 1, 2026

Uh oh!

lfr-0531 commented Aug 2, 2026

Uh oh!

tensorrt-cicd commented Aug 2, 2026

Uh oh!

tensorrt-cicd commented Aug 2, 2026

Uh oh!

Uh oh!

lfr-0531 commented Aug 3, 2026

Uh oh!

lfr-0531 commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

lfr-0531 commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

lfr-0531 commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

lfr-0531 commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

lfr-0531 commented Jul 1, 2026 •

edited by coderabbitai Bot

Loading