[https://nvbugs/6195110][fix] Restore DeepSeek shared-weights vanilla MTP path by zhaoyangwang-nvidia · Pull Request #14457 · NVIDIA/TensorRT-LLM

zhaoyangwang-nvidia · 2026-05-22T09:18:00Z

The decoupling refactor (#12341) made num_nextn_predict_layers a checkpoint-only property and clamped max_draft_len to it for vanilla MTP. This silently dropped the shared-weights hack that lets vanilla run on single-MTP-layer checkpoints (e.g., DeepSeek-V3-Lite) with max_draft_len > ckpt MTP count, breaking the lone CI coverage of the vanilla path (TestDeepSeekV3Lite::test_fp8_block_scales[mtp=vanilla]).

Restore the hack by expanding pretrained_config.num_nextn_predict_layers to max_draft_len before model construction when use_mtp_vanilla=True, preserving the original checkpoint count as _ckpt_num_nextn_predict_layers for weight loader mod-indexing and FP8 exclude_modules duplication. Updates the affected modeling files (deepseekv3) to read ckpt_nextn from the preserved field, with a fallback to num_nextn_predict_layers when no expansion happened.

Remove the waive for the failing test.

Summary by CodeRabbit

Release Notes

Bug Fixes
- Improved Multi-Token Prediction (MTP) layer handling across supported models with enhanced checkpoint compatibility and accurate weight mapping.
- Fixed MTP layer expansion logic for vanilla MTP decoding configurations.
Tests
- Re-enabled test coverage for DeepSeekV3Lite MTP functionality.

Description

The decoupling refactor (#12341) made num_nextn_predict_layers a
checkpoint-only property and clamped max_draft_len to it for vanilla MTP.
This silently dropped the "shared-weights" hack that lets a single-MTP-layer
checkpoint run vanilla MTP with max_draft_len > ckpt MTP count, breaking the
only CI coverage of the path
(TestDeepSeekV3Lite::test_fp8_block_scales[mtp=vanilla]).

This shared-weights vanilla mode is a special MTP configuration: one
checkpoint MTP layer, run for multiple draft tokens, keeping a separate KV
cache per draft position — by sharing that single layer's weights across the
expanded layer instances. It is not used in any real-world workload today, but
DeepSeek indicated they want to keep the path alive for their model, so we
restore it and guard it with one CI test.

Test Coverage

TestDeepSeekV3Lite::test_fp8_block_scales[mtp=vanilla] (verified on H200,
sm_90 + CUTLASS MoE, same as CI H100).

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.

zhaoyangwang-nvidia · 2026-05-22T09:30:01Z

/bot run

coderabbitai · 2026-05-22T09:34:03Z

📝 Walkthrough

Walkthrough

This PR implements checkpoint layer count preservation for vanilla MTP decoding expansion. When expanding MTP layers beyond the checkpoint's original capacity, the code now saves the original checkpoint layer count and uses it consistently for layer index remapping across all model architectures (DeepseekV3, ExaoneMoe, GLM, NemotronH).

Changes

MTP Vanilla Decoding Layer Preservation

Layer / File(s)	Summary
Vanilla MTP config expansion and checkpoint preservation `tensorrt_llm/_torch/pyexecutor/model_loader.py`	When vanilla MTP decoding with `max_draft_len` expands `num_nextn_predict_layers` beyond the checkpoint's original count, the original checkpoint layer count is preserved in `_ckpt_num_nextn_predict_layers` for downstream use.
DeepseekV3 checkpoint-aware MTP layer mapping `tensorrt_llm/_torch/models/modeling_deepseekv3.py`	Weight loader uses the preserved checkpoint layer count to detect shared MTP weights and compute correct modulo indices for layer remapping; initialization derives `ckpt_nextn` from the preserved checkpoint count for layer duplication.
Multi-model initialization with checkpoint-aware ckpt_nextn `tensorrt_llm/_torch/models/modeling_exaone_moe.py`, `tensorrt_llm/_torch/models/modeling_glm.py`, `tensorrt_llm/_torch/models/modeling_nemotron_h.py`	ExaoneMoe, GLM, and NemotronH models consistently derive `ckpt_nextn` from the preserved `_ckpt_num_nextn_predict_layers` attribute (fallback to `num_nextn_predict_layers`) during initialization and weight loading to correctly map expanded MTP layers.
Test waive removal `tests/integration/test_lists/waives.txt`	Removes the waived test entry for DeepSeekV3Lite fp8 block scales with vanilla MTP parameterization, confirming the fix resolves the underlying issue.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

ZhanruiSunCh
yuxianq
govind-ramnarayan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and specifically summarizes the main change: restoring vanilla MTP shared-weights path for DeepSeek.
Description check	✅ Passed	The PR description is comprehensive and clearly explains the issue, solution, test coverage, and checklist items.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

tensorrt-cicd · 2026-05-22T09:35:52Z

PR_Github #49912 [ run ] triggered by Bot. Commit: f5c6957 Link to invocation

zhaoyangwang-nvidia · 2026-05-22T09:38:33Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-05-22T09:45:39Z

PR_Github #49913 [ run ] triggered by Bot. Commit: a8ffd3e Link to invocation

tensorrt-cicd · 2026-05-22T09:49:17Z

PR_Github #49912 [ run ] completed with state ABORTED. Commit: f5c6957

Link to invocation

tensorrt-cicd · 2026-05-22T13:16:06Z

PR_Github #49913 [ run ] completed with state SUCCESS. Commit: a8ffd3e
/LLM/main/L0_MergeRequest_PR pipeline #39492 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

zhaoyangwang-nvidia · 2026-05-25T03:19:32Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-05-25T03:25:49Z

PR_Github #50133 [ run ] triggered by Bot. Commit: a8ffd3e Link to invocation

tensorrt-cicd · 2026-05-25T08:10:36Z

PR_Github #50133 [ run ] completed with state SUCCESS. Commit: a8ffd3e
/LLM/main/L0_MergeRequest_PR pipeline #39685 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

zhaoyangwang-nvidia · 2026-05-25T08:20:29Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-05-25T08:27:18Z

PR_Github #50182 [ run ] triggered by Bot. Commit: a8ffd3e Link to invocation

tensorrt-cicd · 2026-05-25T13:07:43Z

PR_Github #50182 [ run ] completed with state SUCCESS. Commit: a8ffd3e
/LLM/main/L0_MergeRequest_PR pipeline #39725 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

zhaoyangwang-nvidia · 2026-05-25T14:21:14Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-05-25T14:27:22Z

PR_Github #50223 [ run ] triggered by Bot. Commit: 1687d2e Link to invocation

tensorrt-cicd · 2026-05-25T19:35:38Z

PR_Github #50223 [ run ] completed with state FAILURE. Commit: 1687d2e
/LLM/main/L0_MergeRequest_PR pipeline #39760 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

zhaoyangwang-nvidia · 2026-05-26T02:47:40Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-05-26T02:54:23Z

PR_Github #50259 [ run ] triggered by Bot. Commit: 1687d2e Link to invocation

tensorrt-cicd · 2026-06-01T02:25:22Z

PR_Github #51283 [ run ] triggered by Bot. Commit: dc3b17f Link to invocation

tensorrt-cicd · 2026-06-01T03:37:58Z

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

zhaoyangwang-nvidia · 2026-06-01T03:52:24Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-06-01T03:57:46Z

PR_Github #51300 [ run ] triggered by Bot. Commit: dc3b17f Link to invocation

…ding for GLM/Exaone/Nemotron Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

tensorrt-cicd · 2026-06-01T08:36:46Z

PR_Github #51300 [ run ] completed with state SUCCESS. Commit: dc3b17f
/LLM/main/L0_MergeRequest_PR pipeline #40717 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

…M/Exaone/Nemotron-H The MTPWorker vanilla runtime path (per-iter in-place shifts and shared attn_metadata across distinct MTP layers) hits a latent CUDA Graph IMA in warmup on GLM/Exaone/Nemotron-H when ModelLoader expands a single checkpoint MTP layer to multiple model MTP layer instances. The IMA is independent of the load-side fixes and reproduces even without expansion, so it lives in the worker itself, not in this PR. Reject the shared-weights expansion path upfront in Glm4MoeForCausalLM/ExaoneMoeForCausalLM/NemotronHForCausalLM with a clear NotImplementedError when `_ckpt_num_nextn_predict_layers` is set on the pretrained config (i.e. ModelLoader actually expanded). Eagle MTP and regular vanilla MTP (max_draft_len <= ckpt num_nextn_predict_layers) on these models, plus DeepSeek-V3 vanilla expansion, are unaffected. Also fix a pre-existing ConsumableWeightsDict incompatibility in ExaoneMoeWeightMapper.preprocess_weights: replace weights.pop(name) with explicit weights[new]=weights[name]+del weights[name] so HF loading of K-EXAONE no longer fails before reaching the model. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

…MTP only Drop the GLM/Exaone/Nemotron-H load-side extensions and runtime guards added in previous commits on this branch. The vanilla shared-weights runtime path (single ckpt MTP layer, max_draft_len > 1) was originally requested for DeepSeek-V3 and the test_fp8_block_scales[mtp=vanilla] CI case is the only end-to-end coverage we keep. The runtime path is not exercised in production on the other three model families, and local validation showed a latent IMA in MTPWorker's CUDA-graph warmup that is independent of the load-side fixes; making those models work end-to-end requires a separate runtime fix. Reverts (back to the Restore commit baseline): - modeling_glm.py: drop mark_consumed skip + has_shared_mtp_weights - modeling_exaone_moe.py: drop shared-weights vanilla guard - modeling_nemotron_h.py: drop shared-weights vanilla guard - exaone_moe_weight_mapper.py: drop modulo alias + ConsumableWeightsDict pop/setdefault compat - nemotron_h_weight_mapper.py: drop modulo alias - model_loader.py: drop 'or 1' fallback; restore strict gate on config-provided num_nextn_predict_layers Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

…ers plumbing on GLM/Exaone/Nemotron-H These files only had Restore-commit plumbing (`getattr(_ckpt_num_nextn_predict_layers, None) or num_nextn_predict_layers`) that supports the shared-weights vanilla MTP expansion path. The PR scope is now DeepSeek-V3 only, so this plumbing is dead code on the other three model families -- and worse, if it accidentally activates (user sets use_mtp_vanilla=True + max_draft_len > 1) it would silently load wrong weights or hit the known MTPWorker runtime IMA. Revert these three files fully to main to keep the PR strictly DeepSeek-scoped. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

… model_types only Gate the ModelLoader vanilla MTP expansion (max_draft_len > ckpt num_nextn_predict_layers -> rewrite num_nextn_predict_layers + record _ckpt_num_nextn_predict_layers) to {deepseek_v3, deepseek_v32, glm_moe_dsa}. These are exactly the model_types whose MTP path goes through DeepseekV3MTP + DeepseekV3WeightLoader, which is the only weight loader that correctly handles shared-weights expansion. For other MTP-capable models the expansion no longer fires; the natural min(max_draft_len, ckpt_nextn) clamp in MTPForCausalLM keeps the build size at ckpt count, so users effectively get up to ckpt_nextn draft tokens. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

…sion whitelist glm_moe_dsa shares DeepseekV3's code path (registered as the same ForCausalLM class), but it has not been end-to-end validated with the shared-weights vanilla MTP expansion path. Keep the whitelist to the DeepSeek model names only ({deepseek_v3, deepseek_v32}) to avoid implying broader support than the CI case (DeepSeek-V3-Lite fp8 vanilla) actually exercises. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

…ghts expansion Update the comment block in ModelLoader to explain that the shared-weights vanilla MTP path is a very special mode that has not been used in any real-world workload to date and is only kept alive because DeepSeek indicated they want to support it for their model. That is why the expansion is whitelisted to DeepSeek model_types only; other MTP-capable model families do not need this mode and their vanilla MTP usage falls through to the natural clamp inside MTPForCausalLM. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

zhaoyangwang-nvidia · 2026-06-01T11:59:28Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

zhaoyangwang-nvidia · 2026-06-02T03:21:52Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-06-02T03:27:31Z

PR_Github #51487 [ run ] triggered by Bot. Commit: 5f697c2 Link to invocation

tensorrt-cicd · 2026-06-02T03:56:59Z

PR_Github #51487 [ run ] completed with state FAILURE. Commit: 5f697c2
/LLM/main/L0_MergeRequest_PR pipeline #40892 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

zhaoyangwang-nvidia · 2026-06-02T05:50:51Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-06-02T05:56:22Z

PR_Github #51517 [ run ] triggered by Bot. Commit: 5f697c2 Link to invocation

tensorrt-cicd · 2026-06-02T09:22:25Z

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

zhaoyangwang-nvidia · 2026-06-02T12:31:41Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

zhaoyangwang-nvidia · 2026-06-03T01:36:46Z

/bot run --extra-stage "DGX_H100-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-06-03T01:42:41Z

PR_Github #51719 [ run ] triggered by Bot. Commit: 5f697c2 Link to invocation

tensorrt-cicd · 2026-06-03T06:43:32Z

PR_Github #51719 [ run ] completed with state SUCCESS. Commit: 5f697c2
/LLM/main/L0_MergeRequest_PR pipeline #41095 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

github-actions Bot assigned zhaoyangwang-nvidia May 22, 2026

zhaoyangwang-nvidia marked this pull request as ready for review May 22, 2026 09:29

zhaoyangwang-nvidia requested review from a team as code owners May 22, 2026 09:29

zhaoyangwang-nvidia requested review from danielafrimi, hlu1, joyang-nv and yechank-nvidia May 22, 2026 09:29

zhaoyangwang-nvidia force-pushed the fix/nvbugs-6195110-mtp-vanilla-shared-weights branch from f5c6957 to a8ffd3e Compare May 22, 2026 09:38

zhaoyangwang-nvidia force-pushed the fix/nvbugs-6195110-mtp-vanilla-shared-weights branch from 834f267 to 1687d2e Compare May 25, 2026 14:21

yechank-nvidia reviewed Jun 1, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/models/modeling_glm.py Outdated

Comment thread tensorrt_llm/_torch/models/modeling_exaone_moe.py Outdated

Comment thread tensorrt_llm/_torch/models/modeling_nemotron_h.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/model_loader.py

[https://nvbugs/6195110][fix] Complete MTP vanilla shared-weights loa…

79a998f

…ding for GLM/Exaone/Nemotron Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

zhaoyangwang-nvidia added 6 commits June 1, 2026 04:09

nv-guomingz approved these changes Jun 2, 2026

View reviewed changes

zhaoyangwang-nvidia changed the title ~~[https://nvbugs/6195110][fix] Restore MTP vanilla shared-weights path~~ [https://nvbugs/6195110][fix] Restore DeepSeek shared-weights vanilla MTP path Jun 2, 2026

Funatiq approved these changes Jun 2, 2026

View reviewed changes

lfr-0531 approved these changes Jun 2, 2026

View reviewed changes

sunnyqgg merged commit c938efa into NVIDIA:main Jun 3, 2026
19 checks passed

coderabbitai Bot mentioned this pull request Jun 4, 2026

[https://nvbugs/6248827][fix] Wrap both accesses in getattr(self.config, 'num_nextn_predict_layers', 0)… #14932

Closed

2 tasks

Uh oh!

Conversation

zhaoyangwang-nvidia commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

zhaoyangwang-nvidia commented May 22, 2026

Uh oh!

coderabbitai Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

zhaoyangwang-nvidia commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

zhaoyangwang-nvidia commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

zhaoyangwang-nvidia commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

zhaoyangwang-nvidia commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

zhaoyangwang-nvidia commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 1, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

zhaoyangwang-nvidia commented May 22, 2026 •

edited

Loading

coderabbitai Bot commented May 22, 2026 •

edited

Loading