[None][feat] Support Nemotron dynamic-tree MTP decoding

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -216,6 +216,7 @@ bool AttentionOp::convertMMHAParamsToXQAParams(tensorrt_llm::kernels::XQAParams&
     // Medusa mode will have multiple query tokens.
     xqaParams.multi_query_tokens = mIsSpecDecodingEnabled && mUseSpecDecoding;
     xqaParams.is_spec_dec_tree = mIsSpecDecTree;
+    xqaParams.force_prepare_spec_dec_tree_mask = mForcePrepareSpecDecTreeMask;
     xqaParams.layer_idx = generationsParams.layer_idx;
 
     if (mKVCacheQuantMode.hasInt8KvCache())

cpp/tensorrt_llm/common/attentionOp.h

@@ -501,6 +501,7 @@ class AttentionOp
     int32_t mSpecDecodingMaxGenerationLength = 1;
     // Static spec-dec tree length used by FMHA autotuning.
     int32_t mSpecDecodingTargetMaxGenLen = 0;
+    bool mForcePrepareSpecDecTreeMask = false;
     bool mIsMLAEnabled = false;
     bool mIsGenerationMLA = false;
     bool mUseGenFlashMLA = false;
@@ -571,11 +572,11 @@ class AttentionOp
             mCrossAttention, mMaxDistance, mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput,
             mFP8ContextMLA, mFP8GenerationMLA, mChunkPrefillBufferBatchSize, mDenseContextFMHA, mHasFullAttentionMask,
             mIsSpecDecodingEnabled, mUseSpecDecoding, mIsSpecDecTree, mSpecDecodingIsGenerationLengthVariable,
-            mSpecDecodingMaxGenerationLength, mSpecDecodingTargetMaxGenLen, mIsMLAEnabled, mIsGenerationMLA,
-            mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttentionPaged, mUseTllmGenSparseAttention,
-            mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin,
-            mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA,
-            mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant,
+            mSpecDecodingMaxGenerationLength, mSpecDecodingTargetMaxGenLen, mForcePrepareSpecDecTreeMask, mIsMLAEnabled,
+            mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttentionPaged,
+            mUseTllmGenSparseAttention, mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads,
+            mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast,
+            mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant,
             mFusesDsv4InvRopeFp8Quant, mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1),
             mSkipSoftmaxThresholdScaleFactorPrefill, mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ,
             mSageAttnNumEltsPerBlkK, mSageAttnNumEltsPerBlkV, mSageAttnQkInt8);

cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h

@@ -63,6 +63,7 @@ struct XQAParams
     int64_t* spec_decoding_bl_tree_mask_offset;        // for blackwell spec-dec tree mask offset
     uint32_t* spec_decoding_bl_tree_mask;              // for blackwell spec-dec tree mask
     int32_t* spec_bl_tree_first_sparse_mask_offset_kv; // for blackwell spec-dec tree first sparse mask offset kv
+    bool force_prepare_spec_dec_tree_mask = false;
     int32_t const* mrope_position_deltas = nullptr;
     // Helix parallelism params.
     int32_t const* helix_position_offsets = nullptr;

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h

@@ -496,7 +496,11 @@ class TllmGenFmhaKernel
         tg::CudaRunner::Grid grid{numCtasX, numCtasY, numCtasZ};
 
         // Prepare custom mask for spec-decoding generation kernels if needed.
-        if (params.mLayerIdx == 0 && params.mIsSpecDecTree)
+        bool const prepareSpecDecTreeMask = params.mIsSpecDecTree
+            && (params.mForcePrepareSpecDecTreeMask || params.mLayerIdx == 0
+                || (params.mSpecDecodingTargetMaxGenLen > 0
+                    && params.mMaxSeqLenQ != params.mSpecDecodingTargetMaxGenLen));
+        if (prepareSpecDecTreeMask)
         {
             int32_t stepQ = options.mTileSizeQ * options.mNumInstsQ;
             int32_t stepKv = options.mTileSizeKv * options.mNumInstsKv;

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h

@@ -369,6 +369,7 @@ struct TllmGenFmhaRunnerParams
     // row stride ceilDiv(mPackedMaskMaxSeqLenQ, 32) rather than ceilDiv(seqLenQ, 32).
     int32_t mPackedMaskMaxSeqLenQ = 0;
     int32_t mSpecDecodingTargetMaxGenLen = 0;
+    bool mForcePrepareSpecDecTreeMask = false;
 
     // set the attention mask type
     TllmGenFmhaRunnerParams& setAttentionMaskType(std::int8_t maskType)

cpp/tensorrt_llm/kernels/xqaDispatcher.cpp

@@ -534,6 +534,7 @@ void XqaDispatcher::runImpl(
         tllmRunnerParams.generalPackedCustoMaskPtr = params.spec_decoding_packed_mask;
         tllmRunnerParams.mPackedMaskMaxSeqLenQ = params.spec_decoding_max_generation_length;
         tllmRunnerParams.mSpecDecodingTargetMaxGenLen = mFixedParams.specDecodingTargetMaxGenLen;
+        tllmRunnerParams.mForcePrepareSpecDecTreeMask = params.force_prepare_spec_dec_tree_mask;
         tllmRunnerParams.customMaskPtr = params.spec_decoding_bl_tree_mask;
         tllmRunnerParams.customMaskOffsetsPtr = params.spec_decoding_bl_tree_mask_offset;
         tllmRunnerParams.firstSparseMaskOffsetsKvPtr = params.spec_bl_tree_first_sparse_mask_offset_kv;

cpp/tensorrt_llm/nanobind/thop/bindings.cpp

@@ -179,7 +179,8 @@ void initBindings(nb::module_& m)
         nb::arg("relative_attention_bias") = std::nullopt, nb::arg("relative_attention_max_distance") = 0,
         nb::arg("spec_decoding_target_max_draft_tokens") = std::nullopt, nb::arg("quant_scale_qkv") = std::nullopt,
         nb::arg("dsv4_inv_rope_cos_sin_cache") = std::nullopt, nb::arg("enable_dsv4_epilogue_fusion") = false,
-        "Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
+        nb::arg("force_prepare_spec_dec_tree_mask") = false, "Multi-head attention operation",
+        nb::call_guard<nb::gil_scoped_release>());
 
     m.def(
         "get_helix_workspace_size_per_rank",

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -1084,7 +1084,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<int64_t> compressed_kv_cache_pool_ptr, bool const is_cross, std::optional<torch::Tensor> cross_kv,
     std::optional<torch::Tensor> relative_attention_bias, int64_t relative_attention_max_distance,
     std::optional<int64_t> spec_decoding_target_max_draft_tokens, std::optional<torch::Tensor> quant_scale_qkv,
-    std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache, bool enable_dsv4_epilogue_fusion)
+    std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache, bool enable_dsv4_epilogue_fusion,
+    bool const force_prepare_spec_dec_tree_mask)
 {
     TLLM_LOG_TRACE("Attention op starts at layer %d", local_layer_idx);
     // Use these tensors to infer if the attention is using KV cache
@@ -1237,6 +1238,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     {
         op->mSpecDecodingTargetMaxGenLen = static_cast<int32_t>(spec_decoding_target_max_draft_tokens.value()) + 1;
     }
+    op->mForcePrepareSpecDecTreeMask = force_prepare_spec_dec_tree_mask;
 
     op->mUseSparseAttention = false;
     op->mUseTllmGenSparseAttentionPaged = false;

cpp/tensorrt_llm/thop/attentionOp.h

@@ -95,7 +95,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> relative_attention_bias = std::nullopt, int64_t relative_attention_max_distance = 0,
     std::optional<int64_t> spec_decoding_target_max_draft_tokens = std::nullopt,
     std::optional<torch::Tensor> quant_scale_qkv = std::nullopt,
-    std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache = std::nullopt, bool enable_dsv4_epilogue_fusion = false);
+    std::optional<torch::Tensor> dsv4_inv_rope_cos_sin_cache = std::nullopt, bool enable_dsv4_epilogue_fusion = false,
+    bool const force_prepare_spec_dec_tree_mask = false);
 
 struct KvCachePoolPointers
 {

docs/source/features/feature-combination-matrix.md

@@ -14,7 +14,7 @@
 | Speculative Decoding — Linear                               | Yes               | Yes        | Yes                | No                   | Yes                | No                | Yes                        | Yes                   | Yes             | ---                           |                                      |                                                           |               |                  |                |                          |                       |                 |          |
 | Speculative Decoding — Dynamic Trees                        | Yes               | Yes        | Yes                | No                   | Yes                | No                | Yes                        | Yes                   | Yes             | No                            | ---                                  |                                                           |               |                  |                |                          |                       |                 |          |
 | Speculative Decoding — Legacy Path (NGram, user-provided)   | Yes               | Yes        | Yes                | No                   | Yes                | No                | Yes                        | Yes                   | Yes             | No                            | No                                   | ---                                                       |               |                  |                |                          |                       |                 |          |
-| Torch Sampler                                               | Yes               | Yes        | Yes                | Yes                  | Yes                | Yes               | Yes                        | Yes                   | Yes             | Yes                           | Yes                                  | Yes                                                       | ---           |                  |                |                          |                       |                 |          |
+| Torch Sampler                                               | Yes               | Yes        | Yes                | Yes                  | Yes                | Yes               | Yes                        | Yes                   | Yes             | Yes                           | Yes (MTP dynamic tree: greedy only) | Yes                                                       | ---           |                  |                |                          |                       |                 |          |
 | TLLM C++ Sampler                                            | Yes               | Yes        | Yes                | Yes                  | Yes                | Yes               | Yes                        | Yes                   | Yes             | No                            | No                                   | No                                                        | No            | ---              |                |                          |                       |                 |          |
 | KV Cache Reuse                                              | Yes               | Yes        | Yes                | Yes                  | Yes                | Yes               | Yes                        | Yes                   | Yes             | Yes                           | Yes                                  | Yes                                                       | Yes           | Yes              | ---            |                          |                       |                 |          |
 | Sliding Window Attention                                    | Yes               | Yes        | Yes                | Yes                  | Yes                | Untested          | Yes                        | Yes                   | Yes             | Yes                           | Yes                                  | Yes                                                       | Yes           | Yes              | Yes            | ---                      |                       |                 |          |

examples/llm-api/quickstart_advanced.py

@@ -310,6 +310,9 @@ def setup_llm(args, **kwargs):
             relaxed_topk=args.relaxed_topk,
             relaxed_delta=args.relaxed_delta,
             mtp_eagle_one_model=args.use_one_model,
+            use_dynamic_tree=args.use_dynamic_tree,
+            dynamic_tree_max_topK=args.dynamic_tree_max_topK,
+            max_total_draft_tokens=args.max_total_draft_tokens,
             speculative_model=args.model_dir)
     elif spec_decode_algo == "EAGLE3":
         spec_config = Eagle3DecodingConfig(

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -98,6 +98,7 @@ def forward(
             spec_decoding_bl_tree_mask_offset=metadata.spec_decoding_bl_tree_mask_offset,
             spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask,
             spec_decoding_target_max_draft_tokens=metadata.max_total_draft_tokens,
+            force_prepare_spec_dec_tree_mask=metadata.force_prepare_spec_dec_tree_mask,
             spec_bl_tree_first_sparse_mask_offset_kv=metadata.spec_bl_tree_first_sparse_mask_offset_kv,
             num_sparse_topk=metadata.num_sparse_topk,
             flash_mla_tile_scheduler_metadata=metadata.flash_mla_tile_scheduler_metadata,

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -117,6 +117,7 @@ def effective_beam_width(self) -> int:
     is_spec_dec_tree: bool = False
     # if spec-dec tree wouldn't be changed at all, the mask won't be computed every step.
     is_spec_dec_dynamic_tree: bool = False
+    force_prepare_spec_dec_tree_mask: bool = False
 
     # parameters required for spec-dec mode
     max_total_draft_tokens: Optional[int] = None

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -1275,6 +1275,8 @@ def forward(
                 residual=residual,
                 attn_metadata=attn_metadata,
                 all_rank_num_tokens=all_rank_num_tokens,
+                spec_metadata=spec_metadata,
+                mamba_metadata=attn_metadata.mamba_metadata,
                 lora_params=lora_params,
             )
         return hidden_states

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -430,16 +430,47 @@ def forward(
                         f"{draft_token_num} must match fixed replay step "
                         f"width {replay_step_width}.")
 
-                intermediate_state_indices = _cached_arange(
-                    attn_metadata.kv_cache_manager.get_max_resource_count(),
-                    state_indices_d.device)[:num_decodes]
+                # Dynamic-tree verify uses per-request links; linear MTP skips it.
+                is_dyn_tree = getattr(spec_metadata, 'is_spec_dec_dynamic_tree',
+                                      False)
+                retrieve_next_token = retrieve_next_sibling = None
+                retrieve_parent_token = None
+                if is_dyn_tree:
+                    if use_replay:
+                        raise NotImplementedError(
+                            "Dynamic-tree Mamba verify is not supported with "
+                            "the replay SSM-cache path (TRTLLM_USE_MAMBA_REPLAY)."
+                        )
+                    retrieve_next_token = spec_metadata.retrieve_next_token
+                    retrieve_next_sibling = spec_metadata.retrieve_next_sibling
+                    assert (retrieve_next_token is not None
+                            and retrieve_next_sibling is not None), (
+                                "Dynamic-tree verify requires retrieve link "
+                                "tensors on spec_metadata.")
+                    retrieve_next_token = retrieve_next_token[:num_decodes]
+                    retrieve_next_sibling = retrieve_next_sibling[:num_decodes]
+                    # conv1d fills parent links used by tree-aware SSM restore.
+                    retrieve_parent_token = torch.empty(
+                        (num_decodes, draft_token_num),
+                        dtype=torch.int32,
+                        device=state_indices_d.device)
+
+                # Prefer the cache_manager-owned arange; cached fallback storage
+                # can be recycled by CUDA graph warmup.
+                _km_isi = getattr(attn_metadata.kv_cache_manager,
+                                  'intermediate_state_indices', None)
+                if _km_isi is not None:
+                    intermediate_state_indices = _km_isi[:num_decodes]
+                else:
+                    intermediate_state_indices = _cached_arange(
+                        attn_metadata.kv_cache_manager.get_max_resource_count(),
+                        state_indices_d.device)[:num_decodes]
 
-                # Reshape for batch processing
-                xbc_d_reshaped = xbc_d.view(num_decodes, draft_token_num,
-                                            -1).transpose(1, 2)
+                # Use reshape because dynamic-tree tokens may be non-contiguous.
+                xbc_d_reshaped = xbc_d.reshape(num_decodes, draft_token_num,
+                                               -1).transpose(1, 2)
 
                 def conv1d():
-                    # TODO:support tree structure [TRTLLM-10320]
                     xbc_d_processed = causal_conv1d_update_triton(
                         xbc_d_reshaped,
                         conv_states,
@@ -449,11 +480,15 @@ def conv1d():
                         conv_state_indices=state_indices_d[:num_decodes],
                         intermediate_conv_window=intermediate_conv_states,
                         intermediate_state_indices=intermediate_state_indices,
+                        # None on linear MTP.
+                        retrieve_next_token=retrieve_next_token,
+                        retrieve_next_sibling=retrieve_next_sibling,
+                        retrieve_parent_token=retrieve_parent_token,
                         # PDL chain: conv1d → precompute → main (replay only)
                         launch_dependent_kernels=use_replay,
                     )
 
-                    return xbc_d_processed.transpose(1, 2).view(
+                    return xbc_d_processed.transpose(1, 2).reshape(
                         num_decode_tokens, -1)
 
             else:
@@ -588,13 +623,23 @@ def convert_dt():
                         state_batch_indices=state_batch_indices,
                         disable_state_update=True,
                         intermediate_state_indices=intermediate_state_indices,
+                        # None for linear MTP; tree parent map for dynamic tree.
+                        retrieve_parent_token=retrieve_parent_token,
                     )
                 else:
                     # Triton kernel + flashinfer need contiguous for alignment.
                     x_d_4d = x_d_4d.contiguous()
                     B_d_4d = B_d_4d.contiguous()
                     C_d_4d = C_d_4d.contiguous()
-                    self.selective_state_update_func(
+                    if is_dyn_tree:
+                        # flashinfer SSU cannot restore tree-parent states.
+                        ssu_func = selective_state_update_native
+                        ssu_extra = dict(
+                            retrieve_parent_token=retrieve_parent_token)
+                    else:
+                        ssu_func = self.selective_state_update_func
+                        ssu_extra = {}
+                    ssu_func(
                         ssm_states,
                         x_d_4d,
                         dt_d_4d,
@@ -611,6 +656,7 @@ def convert_dt():
                         intermediate_states_buffer=intermediate_ssm_states,
                         cache_steps=draft_token_num,
                         intermediate_state_indices=intermediate_state_indices,
+                        **ssu_extra,
                         **philox_kwargs,
                     )
             else: