[None][feat] Enable MTP for Step-3.7 NVFP4 and port Step-3.7VL vision tower to TRT-LLM modules#14926
Conversation
📝 WalkthroughWalkthroughThis PR ports the Step3p7 vision stack to TRT‑LLM building blocks, adds optional ChangesStep3p7 Vision Porting and MTP Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py (1)
801-819:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAssert that the
mtp_nextn=3cases actually stay on the MTP path.Right now these tests only prove that MMMU still runs after passing
speculative_config. If Step-3.7 silently falls back to non-speculative decoding, the new[0, 3]matrix can still pass and we would miss the regression this PR is meant to cover. Please add one direct assertion for the nonzero branch in both tests (or via a shared helper) so coverage is explicitly sufficient for MTP enablement.Suggested assertion
def test_fp8_block_scales(self, mtp_nextn): with self._make_llm(f"{llm_models_root()}/Step-3.7-Flash-FP8", mtp_nextn=mtp_nextn) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + if mtp_nextn > 0: + assert llm.args.speculative_config is not None + assert llm.args.speculative_config.max_draft_len == mtp_nextn task = MMMU(self.MODEL_NAME) task.evaluate( llm, sampling_params=self.sampling_params, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS, ) @@ def test_nvfp4(self, mtp_nextn): with self._make_llm( f"{llm_models_root()}/Step-3.7-Flash-NVFP4", mtp_nextn=mtp_nextn ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + if mtp_nextn > 0: + assert llm.args.speculative_config is not None + assert llm.args.speculative_config.max_draft_len == mtp_nextn task = MMMU(self.MODEL_NAME) task.evaluate( llm, sampling_params=self.sampling_params, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS, )As per coding guidelines, for
tests/**: “suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR.”🤖 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 `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py` around lines 801 - 819, Add an explicit assertion in both test_fp8_block_scales and test_nvfp4 to ensure the nonzero mtp_nextn branch actually uses the MTP/speculative path: when mtp_nextn == 3 assert the LLM instance reports MTP/speculative enabled (e.g. check llm.args.mtp_nextn == 3 and/or llm.args.speculative_config is truthy and llm.args.speculative_config.mtp_nextn == 3) before running MMMU.evaluate so the test fails if Step-3.7 silently falls back to non-speculative decoding.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/modeling_step3p7vl.py (2)
936-937: 💤 Low valueType annotation mismatch with initialization.
The lists are typed as
List[torch.Tensor]but initialized withNonevalues. This works at runtime but is technically a type annotation inaccuracy.✨ Suggested fix for type accuracy
- full_feats: Dict[int, List[torch.Tensor]] = {r: [None] * len(full_imgs[r]) for r in order} - patch_feats: Dict[int, List[torch.Tensor]] = {r: [None] * len(patch_imgs[r]) for r in order} + full_feats: Dict[int, List[Optional[torch.Tensor]]] = {r: [None] * len(full_imgs[r]) for r in order} + patch_feats: Dict[int, List[Optional[torch.Tensor]]] = {r: [None] * len(patch_imgs[r]) for r in order}Note: This also requires importing
Optionalfor the inner list type if not already used.🤖 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_step3p7vl.py` around lines 936 - 937, The type annotations for full_feats and patch_feats are inaccurate: they declare List[torch.Tensor] but initialize lists with None; change their types to Dict[int, List[Optional[torch.Tensor]]] (import Optional from typing if not present) so the annotation matches initialization, e.g. update the annotations for full_feats and patch_feats in modeling_step3p7vl.py to use Optional[torch.Tensor] for the inner list element type.
318-318: 💤 Low valueOptional: Use tuple unpacking for cleaner shape construction.
The static analysis hint suggests using unpacking syntax which is slightly more idiomatic.
✨ Suggested style improvement
- pad_shape = q_real.shape[:-1] + (self.head_dim_pad,) + pad_shape = (*q_real.shape[:-1], self.head_dim_pad)🤖 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_step3p7vl.py` at line 318, The pad_shape assignment currently builds the tuple via concatenation; change it to use tuple unpacking for clarity: replace the expression assigning pad_shape (which references q_real and self.head_dim_pad) with an unpacking form like pad_shape = (*q_real.shape[:-1], self.head_dim_pad) so the shape construction is more idiomatic and easier to read.
🤖 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.
Outside diff comments:
In `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`:
- Around line 801-819: Add an explicit assertion in both test_fp8_block_scales
and test_nvfp4 to ensure the nonzero mtp_nextn branch actually uses the
MTP/speculative path: when mtp_nextn == 3 assert the LLM instance reports
MTP/speculative enabled (e.g. check llm.args.mtp_nextn == 3 and/or
llm.args.speculative_config is truthy and llm.args.speculative_config.mtp_nextn
== 3) before running MMMU.evaluate so the test fails if Step-3.7 silently falls
back to non-speculative decoding.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_step3p7vl.py`:
- Around line 936-937: The type annotations for full_feats and patch_feats are
inaccurate: they declare List[torch.Tensor] but initialize lists with None;
change their types to Dict[int, List[Optional[torch.Tensor]]] (import Optional
from typing if not present) so the annotation matches initialization, e.g.
update the annotations for full_feats and patch_feats in modeling_step3p7vl.py
to use Optional[torch.Tensor] for the inner list element type.
- Line 318: The pad_shape assignment currently builds the tuple via
concatenation; change it to use tuple unpacking for clarity: replace the
expression assigning pad_shape (which references q_real and self.head_dim_pad)
with an unpacking form like pad_shape = (*q_real.shape[:-1], self.head_dim_pad)
so the shape construction is more idiomatic and easier to read.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9a9d48ea-055c-438b-853e-1554da434630
📒 Files selected for processing (8)
tensorrt_llm/_torch/models/modeling_step3p7.pytensorrt_llm/_torch/models/modeling_step3p7vl.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/modeling/test_modeling_step3p7.pytests/unittest/_torch/modeling/test_modeling_step3p7vl.py
|
/bot --help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #52054 [ run ] triggered by Bot. Commit: |
|
/bot run --stage-list "DGX_B200-4_GPUs-PyTorch-3" |
|
PR_Github #52086 [ run ] triggered by Bot. Commit: |
|
PR_Github #52054 [ run ] completed with state |
|
PR_Github #52086 [ run ] completed with state
|
|
/bot run --stage-list "DGX_B200-4_GPUs-PyTorch-3" |
|
PR_Github #52201 [ run ] triggered by Bot. Commit: |
yechank-nvidia
left a comment
yechank-nvidia
left a comment
There was a problem hiding this comment.
LGTM. No major blocking components. Thanks for addressing previous PR's comments.
|
PR_Github #52201 [ run ] completed with state
|
d7d192e to
810d431
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-PyTorch-3" |
|
PR_Github #52290 [ run ] triggered by Bot. Commit: |
|
@coderabbitai why didn't you review the code as a QA engineer at the first commit? |
|
You can check and adjust these settings in your (ᵕ ᴗ ᵕ)∫ ✅ Action performedReview finished.
|
|
@coderabbitai review the pr as a QA engineer to check the test coverage. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
echo "=== modeling_step3p7.py (relevant forward method) ==="
rg -n "spec_input_ids|spec_worker|spec_token" tensorrt_llm/_torch/models/modeling_step3p7.py -C 5
echo ""
echo "=== modeling_step3p7vl.py key exports and new code ==="
rg -n "spec_input_ids|image_token_id|freqs_for_grid|Step3VisionRope2D|Step3VisionAttention|Step3VisionMLP|Step3VisionBlock|Step3VisionTransformer|Step3p7VisionEncoder|Step3p7VisionTower" tensorrt_llm/_torch/models/modeling_step3p7vl.py | head -80
echo ""
echo "=== test_modeling_step3p7.py full content ==="
cat tests/unittest/_torch/modeling/test_modeling_step3p7.py
echo ""
echo "=== test_modeling_step3p7vl.py full content ==="
cat tests/unittest/_torch/modeling/test_modeling_step3p7vl.pyLength of output: 78756 🏁 Script executed: #!/bin/bash
echo "=== integration test for pytorch ==="
rg -n "Step3_7|step3p7|mtp_nextn|MTPDecodingConfig" tests/integration/defs/accuracy/test_llm_api_pytorch.py -C 3
echo ""
echo "=== integration test for multimodal ==="
rg -n "Step3_7|step3p7|mtp_nextn|MTPDecodingConfig" tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py -C 3
echo ""
echo "=== l0_dgx_b200.yml new entries ==="
rg -n "step3|Step3|mtp_nextn" tests/integration/test_lists/test-db/l0_dgx_b200.yml -C 3
echo ""
echo "=== qa llm_function_core.txt ==="
rg -n "step3|Step3|mtp_nextn" tests/integration/test_lists/qa/llm_function_core.txt -C 2
echo ""
echo "=== New methods in modeling_step3p7vl.py ==="
rg -n "def " tensorrt_llm/_torch/models/modeling_step3p7vl.py
echo ""
echo "=== spec_input_ids usage in vl file ==="
rg -n "spec_input_ids" tensorrt_llm/_torch/models/modeling_step3p7vl.py -C 5
echo ""
echo "=== head_dim padding logic ==="
rg -n "pad|zero|128|96" tensorrt_llm/_torch/models/modeling_step3p7vl.py | head -40
echo ""
echo "=== update_spec_config_from_model_config ==="
rg -rn "update_spec_config_from_model_config|mtp_num_hidden_layers|num_nextn_predict_layers" tensorrt_llm/ | head -30Length of output: 108651
✅ Good coverage
|
7d348d5 to
df1ad38
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52978 [ run ] triggered by Bot. Commit: |
|
PR_Github #52907 [ run ] completed with state
|
|
PR_Github #52978 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53059 [ run ] triggered by Bot. Commit: |
|
PR_Github #53059 [ run ] completed with state
|
df1ad38 to
7798287
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53072 [ run ] triggered by Bot. Commit: |
|
PR_Github #53072 [ run ] completed with state
|
update_spec_config_from_model_config previously read the MTP layer count only from `num_nextn_predict_layers`, which DeepSeek-style configs expose. Qwen3Next-style configs (including Qwen3.5) instead expose it as `mtp_num_hidden_layers`, so MTP spec config resolution failed on those checkpoints. Read whichever field is present and fall back to a single shared MTP / EAGLE layer when neither exists. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> [None][feat] Enable MTP for Step3p7 NVFP4 and fix spec token ids on VLM path The NVFP4 modelopt export ships the 3 plain-path MTP layers (45..47) under the bare model.layers.* namespace, so it can now run speculative decoding like the FP8/BF16 checkpoints. On the multimodal path the main model consumes fused inputs_embeds and input_ids is None, so the VLM wrapper forwards the pre-fusion token ids via a new spec_input_ids argument (image tokens remapped into vocab range) so the MTP/spec worker always receives valid token ids. Update accuracy tests to parametrize mtp_nextn, refresh the checkpoint unit test for the new MTP layout, and adjust the L0/QA test lists accordingly. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> [None][feat] Port Step3p7VL vision tower to TRT-LLM Attention module Replace the raw-torch SDPA vision attention with TRT-LLM's Attention module so the vision encoder runs through the backend FMHA dispatch (context-only, FULL mask, varlen via per-image attn_metadata). The PerceptionEncoder head_dim (96) is not a trtllm-gen FMHA cubin size, so q/k/v/o head dims are zero-padded to 128 while 2D RoPE runs on the real channels; a compensating q_scaling preserves the softmax scale. Fused HF in_proj_*/out_proj weights are remapped onto qkv_proj/o_proj during weight loading with the head_dim padding. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> [None][feat] Port Step3p7VL vision MLP to TRT-LLM MLP module Dispatch the non-gated vision FFN through the base TRT-LLM MLP module (up_proj -> activation -> down_proj) instead of a bespoke c_fc/c_proj nn.Module. HF mlp.c_fc / mlp.c_proj weights are remapped onto up_proj / down_proj in Step3p7VisionTower._remap_vision_weights (renamed from _remap_vision_attention_weights). Update unit tests accordingly. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Precompute the vision tower's 2D-RoPE cos/sin tables once per forward and thread them through every layer instead of re-evaluating cos()/sin() inside each attention call. Restrict multimodal device copies to the tensors the vision tower consumes on GPU, and skip the spec_input_ids OOV-sentinel rewrite when speculative decoding is off. Update gsm8k/mmmu accuracy references for the NVFP4 MTP path. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
… update (NVIDIA#14997)" This reverts commit 316430f.
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
7798287 to
d86dc09
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53219 [ run ] triggered by Bot. Commit: |
|
PR_Github #53219 [ run ] completed with state
|
|
/bot run |
|
PR_Github #53274 [ run ] triggered by Bot. Commit: |
|
PR_Github #53274 [ run ] completed with state |
6 participants
Description
Enables Multi-Token Prediction (MTP) speculative decoding for the Step-3.7-Flash NVFP4 checkpoint and ports the Step-3.7VL vision tower and MLP onto TRT-LLM's native modules. Bundled changes:
Enable MTP for Step-3.7 NVFP4 + fix spec token ids on the VLM path. The NVFP4 ModelOpt export ships the 3 plain-path MTP layers (45..47) under the bare
model.layers.*namespace, so the NVFP4 checkpoint can now run speculative decoding like the FP8/BF16 checkpoints. On the multimodal path the main model consumes fusedinputs_embedsandinput_idsisNone, so the VLM wrapper now forwards the pre-fusion token ids via a newspec_input_idsargument (image tokens remapped into vocab range) so the MTP/spec worker always receives valid token ids.Port the Step-3.7VL vision tower to the TRT-LLM
Attentionmodule. Replaces the raw-torch SDPA vision attention so the vision encoder runs through the backend FMHA dispatch (context-only,FULLmask, varlen via per-imageattn_metadata). The PerceptionEncoderhead_dim(96) is not a trtllm-gen FMHA cubin size, so q/k/v/o head dims are zero-padded to 128 while 2D RoPE runs on the real channels; a compensatingq_scalingpreserves the softmax scale. Fused HFin_proj_*/out_projweights are remapped ontoqkv_proj/o_projduring weight loading.Port the Step-3.7VL vision MLP to the base TRT-LLM
MLPmodule. Dispatches the non-gated vision FFN throughup_proj -> activation -> down_projinstead of a bespokec_fc/c_projnn.Module. HFmlp.c_fc/mlp.c_projweights are remapped ontoup_proj/down_projduring weight loading.Resolve MTP layer count across checkpoint field names.
update_spec_config_from_model_configpreviously read the MTP layer count only fromnum_nextn_predict_layers(DeepSeek-style). Qwen3Next-style configs (including Qwen3.5) expose it asmtp_num_hidden_layers; the resolver now reads whichever field is present and falls back to a single shared MTP / EAGLE layer when neither exists.Optimize vision RoPE and multimodal H2D copies. Precomputes the vision tower's 2D-RoPE
cos/sintables once per forward pass and threads them through all 47 blocks instead of re-evaluatingcos()/sin()inside every attention call. Restricts multimodal device copies to only the tensors the vision tower actually consumes on GPU. Skips thespec_input_idsOOV-sentinel rewrite when speculative decoding is disabled.Test Coverage
Unit tests:
tests/unittest/_torch/modeling/test_modeling_step3p7.py— refreshed for the NVFP4 MTP layer layout; removedtest_config_and_weight_accounting_1_nvfp4waive.tests/unittest/_torch/modeling/test_modeling_step3p7vl.py— vision tower / MLP TRT-LLM module port and weight remapping, RoPE precomputation path.Accuracy / integration tests (parametrized over
mtp_nextn∈ {0, 3}):accuracy/test_llm_api_pytorch.py::TestStep3_7::test_nvfp4accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scalesaccuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_fp8_block_scalesaccuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4Accuracy reference updates:
references/gsm8k.yaml: added NVFP4 + MTP entry (accuracy 88, reuses non-spec baseline).references/mmmu.yaml: updated FP8 and NVFP4 baselines from 64 → 60; added FP8+MTP and NVFP4+MTP entries.Test list updates:
tests/integration/test_lists/test-db/l0_dgx_b200.yml: moved NVFP4 MTP case to QA.tests/integration/test_lists/qa/llm_function_core.txt: added Step-3.7 / Step-3.7VL NVFP4 + FP8 MTP accuracy cases (both text and multimodal, mtp_nextn=0 and 3).tests/integration/test_lists/waives.txt: removedtest_config_and_weight_accounting_1_nvfp4skip (NVFP4 MTP now works).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES to the best of my knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
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