Add Qwen3.5 MTP AutoDeploy support by govind-ramnarayan · Pull Request #252 · nv-auto-deploy/TensorRT-LLM

govind-ramnarayan · 2026-05-18T21:32:57Z

@coderabbitai summary

Description

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-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.

… to sharding-IR canonical Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

…py staging Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

…Mixer pattern test The IR-fication of MambaRMSNormGated.forward() in 1777ad5 replaced the inlined gated_rms_norm_ref(...) Python helper with a single registered op torch.ops.auto_deploy.torch_rmsnorm_gated(...). The legacy sharder pattern detector correctly emits a column-wise WeightShardingInfo for the new node, but the test's hand-constructed expected_transformations had no branch for it and the existing "norm_weight or a_log" check did not match (the IR op call has hidden_states as args[0], not the weight). The detected_set then contained one extra entry vs expected_set, breaking test_sharding_pattern_detection[NemotronHMamba2Mixer-...] for all 8 parametrizations. Add an explicit elif branch matching torch_rmsnorm_gated, mirroring the existing torch_ssm branch (column-wise split, layer_type=SSM, no fused_weight_dims). Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

…l detector After the IR-fication of modeling_nemotron_h.py / modeling_deepseek.py / modeling_qwen3*.py in 1777ad5, the closing linear of each attention/MoE/ MLP block emits ``torch.ops.auto_deploy.all_reduce`` before the residual aten.add. DetectHiddenStatesForCapture.collect_residual_add_nodes walks ``closing_linear.users -> aten.add`` and bails on the first non-add user, so on the IR-fied graph it now stops at the all_reduce and finds zero residual adds, raising: AssertionError: Unable to find residual add nodes for layers. Expected: {-1}, Found: dict_keys([]) This breaks every MTP/Eagle one-model speculative decoding path on the IR-fied models (test_super_mtp_smoke on A30/DGX_B200/H100_PCIe single-GPU stages, and TestNemotronSuperV3::test_mtp[4] on DGX_B200-4_GPUs). Make the residual-walker transparent across torch.ops.auto_deploy.all_reduce (identity at world_size=1; sharding-injected collective otherwise) without recording it as the residual add. Track ``last_add`` separately from the walker position so we still report the correct residual node. Verified locally: test_super_mtp_smoke passes end-to-end (84s) on a patched-BF16 NemotronSuperV3 model. Cluster C umbrella failures (test_run_unit_tests / test_unittests_v2[singlegpu/smoke,standalone, multigpu/transformations]) cascade from the same root cause and should clear with this fix. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

…ort-discipline check Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

…match registry search path The model registry resolver in build_and_run_ad.py searches yaml_extra fragments under model_registry/configs/ (line 32: _REGISTRY_CONFIGS_DIR). enable_sharder_ir.yaml was sitting at model_registry/ instead, so the two IR model entries that reference it (nano_v3__enable_sharder_ir, num_hidden_layers_5__enable_sharder_ir) would fail path resolution. Drop the now-redundant explicit allow-list entry in .gitignore — the existing !model_registry/configs/*.yaml wildcard already covers it. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

…g the 4 IR-aware canonical files Append enable_sharder_ir.yaml to yaml_extra for the 20 entries in models.yaml whose model architecture is handled by one of the four modeling files promoted to IR-canonical in this PR. Identified via huggingface_hub.HfApi.model_info().config.architectures matching: - DeepseekV3ForCausalLM -> modeling_deepseek.py (4 entries) - NemotronHForCausalLM -> modeling_nemotron_h.py (5 entries) - Qwen3ForCausalLM -> modeling_qwen3.py (9 entries) - Qwen3_5MoeForConditional -> modeling_qwen3_5_moe.py (2 entries) After this change, running build_and_run_ad.py against any of these 20 HF model names with --use-registry will resolve to a yaml_extra list ending in enable_sharder_ir.yaml, which disables legacy detect_sharding and enables apply_sharding_hints. The IR-aware modeling code is now exercised end-to-end without requiring users to opt in explicitly. The two pre-existing -IR virtual entries (DeepSeek-R1-IR, NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-IR) are kept as smoke-test entries: they pair the IR overlay with the trimmed num_hidden_layers_5.yaml or nano_v3.yaml configs, useful for fast validation without going through the production sharding plan. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

…ormers 5.x Two latent bugs pre-existed in the file (originally modeling_qwen3_ir.py on upstream; renamed to canonical modeling_qwen3.py by this PR's switch commit) and only surfaced once IR sharding became the default for Qwen3 text models (this PR's models.yaml flip): 1. config.rope_theta -> config.rope_scaling['rope_theta'] In transformers 4.x, Qwen3Config.rope_theta sits at the top of the config object. In transformers 5.x, rope settings moved under config.rope_scaling (alongside rope_type). Reading config.rope_theta directly hits AttributeError when loading any Qwen3 model with the installed transformers 5.x runtime. Replaced with a defensive lookup that handles both layouts. 2. _tied_weights_keys = list -> dict transformers 5.x's PreTrainedModel.get_expanded_tied_weights_keys calls .keys() on _tied_weights_keys, which crashes when the class attribute is a list. Replaced with the dict form mapping {'lm_head.weight': 'model.embed_tokens.weight'} that matches what the analogous upstream commit (NVIDIA#12829) did for the other three IR-aware modeling files (deepseek, nemotron_h, qwen3_5_moe). Validated end-to-end with build_and_run_ad.py against Qwen/Qwen3-4B (ws=2, IR sharding default): 10/10 coherent prompts, apply_sharding_hints matches=468 per rank. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

…kill Extract the "Sharding-aware IR model porting" section from ad-model-onboard into a dedicated ad-sharding-ir-port skill. The two workflows serve different purposes (onboarding a new model vs mechanically adding sharding hints to an existing one) and benefit from isolation. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

Sharding IR is now the default path — hints are added directly to the canonical modeling_*.py instead of creating a separate _ir.py copy. The legacy dual-file pattern is deprecated. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> # Conflicts: # examples/auto_deploy/model_registry/models.yaml # tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py

…onboard SKILL.md to match main This bullet was modified during the PR NVIDIA#13478 work to add a detailed explanation about keeping the trtllm router kernels verbatim. That verbose content belongs in a sharding-IR-specific doc, not in the general ad-model-onboard skill. Restore the short upstream/main form. Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

govind-ramnarayan · 2026-05-18T21:34:03Z

+transforms:
+  insert_cached_causal_conv:
+    backend: triton_causal_conv
+  multi_stream_gemm:


Let's see if we can enable the multi stream ops and see if the accuracy test still consistently passes?

govind-ramnarayan · 2026-05-18T21:34:44Z

 - name: Qwen/Qwen3.5-397B-A17B
  config_id: qwen3_5_moe_400b
  yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'qwen3.5_moe_400b.yaml']
+- name: Qwen/Qwen3.5-397B-A17B-FP8


nit: No need to add the non-mtp FP8 version in this PR

govind-ramnarayan · 2026-05-18T21:36:06Z

+        a_ext = a_flat[extend_start:extend_end].view(num_extend, tokens_per_extend, HV)
+        b_ext = b_flat[extend_start:extend_end].view(num_extend, tokens_per_extend, HV)
+
+        if batch_info.get_max_draft_len() > 0 and intermediate_delta_cache is not None:


nit: I don't think we need this check on max_draft_len - extend requests are by definition speculative requests right now. Double-check with the triton SSM backends.

The issue with this is the following. Please check the code path to verify what I am describing.

In the currrent code path, intermediate_delta_cache is always emitted as a cache argument by this op, whether speculation is on or off. Depending on whether speculation is on, the cache is managed (in which case it is bound to the KV Cache Manager's intermediate state cache) or it is not (in which case it is left unmanaged and is a torch.empty()). It is never actually "None"

(Btw, this is also true in triton_backend_mamba, the cache is just declared as Optional when it probably should not be - the Optional is just used for disabling in unit tests).

So effectively this if statement is useless - we just currently make the assumption that num_extend > 0 => speculation is on (so max draft len > 0) and that speculation always involves using the intermediate delta cache.

Please verify this logic and make the edit - this should eliminate the else statement below.

govind-ramnarayan · 2026-05-18T21:55:36Z

+            "layers_handle_final_norm": True,
+            "_checkpoint_conversion_mapping": {
+                r"^mtp\.layers\.0\.": "model.layers.",
+                r"^mtp\.fc\.": "model.layers.fc.",


Wait, is there really an "fc" in the weights? Do we just silently drop this due to this code (which only initializes the "fc" module if the number of capture layers is more than 1)? https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py#L526-L535

In this case, we need some sort of unit test to make sure we load all the weights and are missing nothing in our module definition...

govind-ramnarayan · 2026-05-18T21:57:03Z

@@ -62,6 +62,9 @@ def _build_model(self, device: DeviceLikeType) -> nn.Module:

        # Get model type for config
        model_type = model_config.model_type
+        if model_type == "qwen3_5_moe" and hasattr(model_config, "text_config"):


Can you clarify this with a comment? It's not good to have model-specific code in this factory file. I would prefer putting something like this into the modeling code file and maybe calling a utility if we need. But first please clarify and also reply to this comment to explain to me.

govind-ramnarayan · 2026-05-18T22:00:28Z

+  tokens_per_block: 32
+speculative_config:
+  decoding_type: MTP
+  max_draft_len: 1


Let's try increasing this draft_len. Is there any test for Qwen3Next in PyTorch backend that we can use for inspiration? Otherwise can try... 6?

govind-ramnarayan · 2026-05-18T22:01:00Z

+# capture.
+enable_chunked_prefill: false
+max_num_tokens: 8192
+max_batch_size: 32


Note: May need to reduce the batch size to fit on H100s when we increase speculative length. Prioritize increasing speculative length as long as we can fit it in batch_size=8 or greater.

govind-ramnarayan · 2026-05-18T22:01:13Z

@@ -74,10 +75,12 @@ def get_eagle_layers(config, model_type: str) -> Union[nn.ModuleList, nn.Module]
            layers = build_llama_eagle_layers(config)
        case "nemotron_h":
            layers = build_nemotron_eagle_layers(config)
+        case "qwen3_5_moe" | "qwen3_5_moe_text":


nit: What is the difference between "qwen3_5_moe" and "qwen3_5_moe_text". Can you walk me through the code path that creates this difference?

govind-ramnarayan · 2026-05-18T22:02:48Z

+        self.pre_fc_norm_hidden = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.fc = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+        self.self_attn = Qwen3_5MoeAttention(config, layer_idx)
+        self.mlp = Qwen3_5MoeSparseMoeBlock(config)


a bit weird to call this MLP when it is an MoE. I guess this is needed for the checkpoint? Thoughts about changing it? I'm not sure if it is just customary to call it MLP when it is an MoE.

govind-ramnarayan · 2026-05-18T22:03:57Z

-AutoModelForCausalLMFactory.register_custom_model_cls(
-    "Qwen3_5MoeConfig", Qwen3_5MoeForConditionalGeneration
-)
+AutoModelForCausalLMFactory.register_custom_model_cls("Qwen3_5MoeConfig", Qwen3_5MoeForCausalLM)


I see - so this is how we are redirecting Qwen3_5 MoE to the text model? Let's avoid this - let's try to just run our test with a model that has Qwen3_5MoeTextConfig. Maybe we can just overwrite the model type or something in our test for now?

govind-ramnarayan · 2026-05-18T22:17:03Z

+        if get_device_count() < world_size:
+            pytest.skip("Not enough devices for world size, skipping test")
+        kwargs = self.get_default_kwargs()
+        monkeypatch.setenv("TRTLLM_ACCURACY_NO_REFERENCE", "1")


No need for this, just use this env var locally. Better yet - add in the appropriate accuracy for this test.

govind-ramnarayan · 2026-05-18T22:17:16Z

-        sampling_params = self.get_default_sampling_params()
-        model_path = hf_id_to_local_model_dir(self.MODEL_NAME)
+        kwargs.update(self.get_fp8_gsm8k_runtime_overrides())
+        monkeypatch.setenv("TRTLLM_ACCURACY_NO_REFERENCE", "1")


Ditto to below - no monkeypatch.

govind-ramnarayan · 2026-05-18T22:18:09Z

    @pytest.mark.skip_less_device_memory(80000)
    @pytest.mark.parametrize("world_size", [8])
-    def test_bf16(self, world_size, mocker):
+    def test_fp8_gsm8k(self, world_size, mocker, monkeypatch):


Let's reinstate the bf16 test - we just changed to fp8 for a quick comparison. Let's revert this test completely.

govind-ramnarayan · 2026-05-18T22:20:35Z

+    @skip_pre_hopper
+    @pytest.mark.skip_less_device_memory(80000)
+    @pytest.mark.parametrize("world_size", [8])
+    def test_fp8_mtp_gsm8k(self, world_size, mocker, monkeypatch):


Since we are doing bf16 tests for no MTP, maybe let's try just doing BF16 + MTP and see if we have enough memory for that - otherwise we can make this test FP8 and possibly make the no-MTP test FP8 for comparison so that we can see that MTP does not decrease accuracy.

govind-ramnarayan · 2026-05-18T22:26:55Z

@@ -84,6 +73,73 @@ def _preprocess_for_reference(q, k, a, b_proj, A_log, dt_bias, num_v_heads=None)
    return q_norm, k_norm, g, beta


+def _make_extend_case(device, dtype, *, num_k_heads=2, num_v_heads=2):


This name is super vague - let's name it "make_extend_kernel_inputs()"

govind-ramnarayan · 2026-05-18T22:34:03Z

+
+        y_ref[:, start:end] = y_seq.to(dtype)
+        for step_idx in range(tokens_per_extend):
+            _, prefix_final_state = fused_recurrent_gated_delta_rule_fwd(


It makes more sense here to use the prefill kernel for each step, use the extend kernel for the "maximum" step size, and then compare the states after prefill kernels to the extend kernel. I don't think there's any reason to dig into the kernels like this either. Just do something like:

Set up prefill inputs. For i=1, ... num_extend_tokens, run a prefill with i tokens and collect the state afterwards.

Set up extend inputs. Run a single extend request for num_extend_tokens.

Do both of these directly with the autodeploy custom op - not the individual kernel calls.

We can just set this up so num_extend_tokens is something small - like 2 or 3. And the sizes should be small enough that this test runs fast.

govind-ramnarayan · 2026-05-18T22:37:15Z

+)
+from tensorrt_llm.llmapi import MTPDecodingConfig
+
+


Note: Throughout this file remove "qwen_3" strings in the test titles - it's not needed in this file.

Let's actually append these tests to the existing qwen3_5_moe modeling tests, since they are included in that file. We can put them at the end and separate from the others with a note that they are for testing MTP modeling definitions and utilities.

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

This reverts commit 74b4f34. Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

greg-kwasniewski1 and others added 14 commits May 10, 2026 12:45

[None][feat][autodeploy] Switch deepseek/nemotron_h/qwen3/qwen3_5_moe…

486aef3

… to sharding-IR canonical Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

[None][fix][autodeploy] Restore AD_USE_IR_MODELS hook for future _ir.…

15da900

…py staging Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Made-with: Cursor

[None][fix][autodeploy] Use relative custom_ops import to satisfy imp…

3448e02

…ort-discipline check Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

Merge branch 'main' into gk/switch_to_ir_models

b1a4d00

Merge branch 'main' into gk/switch_to_ir_models

cf4e0cb

Merge remote-tracking branch 'upstream/main' into gk/switch_to_ir_models

43ce243

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> # Conflicts: # examples/auto_deploy/model_registry/models.yaml # tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py

github-actions Bot assigned govind-ramnarayan May 18, 2026

govind-ramnarayan commented May 18, 2026

View reviewed changes

govind-ramnarayan force-pushed the gramnarayan/qwen3-mtp branch from 562a273 to bd3dd4b Compare May 19, 2026 01:20

govind-ramnarayan added 6 commits May 20, 2026 17:29

Add Qwen3.5 MTP AutoDeploy support

b4ec174

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

worklog changes

45fdf85

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

address comments

18d8c24

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Fix hidden-state capture seed selection

636fee8

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Revert "Fix hidden-state capture seed selection"

0c83447

This reverts commit 74b4f34. Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

rebase on switch_to_ir_models + cleanup sharding utility changes

f480ae9

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan force-pushed the gramnarayan/qwen3-mtp branch from bd3dd4b to f480ae9 Compare May 21, 2026 00:47

govind-ramnarayan changed the base branch from main to gk/switch_to_ir_models May 21, 2026 00:54

govind-ramnarayan added 2 commits May 20, 2026 19:40

Fix Qwen3.5 MTP hidden-state capture

971736d

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Remove Qwen3.5 MTP debug artifacts

0878760

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

greg-kwasniewski1 force-pushed the gk/switch_to_ir_models branch from 0730798 to 84471f6 Compare May 24, 2026 15:41

		@@ -84,6 +73,73 @@ def _preprocess_for_reference(q, k, a, b_proj, A_log, dt_bias, num_v_heads=None)
		return q_norm, k_norm, g, beta


		def _make_extend_case(device, dtype, *, num_k_heads=2, num_v_heads=2):

Uh oh!

Conversation

govind-ramnarayan commented May 18, 2026

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

govind-ramnarayan May 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

govind-ramnarayan May 18, 2026 •

edited

Loading