From b25ef1d7503a2889090bfe0298c74115e1bab882 Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 12 Jun 2026 00:26:33 -0700 Subject: [PATCH 01/24] [None][feat] support dynamic-tree MTP decoding Enable the MTP one-model path to reuse dynamic-tree speculative decoding while keeping the linear path unchanged. Stabilize the Blackwell CUDA graph warmup path with bounded capture shapes and explicit spec-dec generation autotuning. Signed-off-by: qgai --- examples/llm-api/quickstart_advanced.py | 3 + .../_torch/attention_backend/trtllm.py | 8 +- .../_torch/modules/mamba/mamba2_mixer.py | 80 +- .../_torch/pyexecutor/mamba_cache_manager.py | 55 +- .../_torch/pyexecutor/model_engine.py | 69 + tensorrt_llm/_torch/speculative/mtp.py | 59 +- .../_torch/speculative/mtp_dynamic_tree.py | 1199 +++++++++++++++++ .../_torch/speculative/spec_tree_manager.py | 103 +- tensorrt_llm/_torch/speculative/utils.py | 18 + tensorrt_llm/llmapi/llm_args.py | 49 +- .../test_dynamic_tree_slot_storage.py | 81 ++ 11 files changed, 1677 insertions(+), 47 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 718d4a410e2c..d2acd490d44a 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/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( diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 40d15970398b..df13bdf357a2 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1082,10 +1082,10 @@ def update_spec_dec_param( pos_src = torch.index_select(slot_storage.position_offsets, 0, slot_ids)[:, :n_dt] - pos_dst = self.spec_decoding_position_offsets[:num_gens * - n_dt].view( - num_gens, - n_dt) + compact_total = num_gens * n_dt + compact_offsets = self.spec_decoding_position_offsets[: + compact_total] + pos_dst = compact_offsets.view(num_gens, n_dt) pos_dst.copy_(pos_src, non_blocking=True) actual_mask_width = math.ceil(n_dt / 32) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index e807a904978f..99e7f053878e 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -430,16 +430,61 @@ 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: per-request tree links drive tree-aware + # conv/SSM recurrence. Off for linear MTP (path unchanged). + 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 fuses parent derivation: it reads next/sibling and + # writes the parent map into this buffer, which the SSM + # update then consumes to restore each token's parent state. + retrieve_parent_token = torch.empty( + (num_decodes, draft_token_num), + dtype=torch.int32, + device=state_indices_d.device) + + # Prefer the cache_manager's persistent arange tensor when + # available (allocated once at __init__, lives as long as the + # cache manager). The functools-cached fallback's storage can + # be co-allocated into the CUDA-graph private memory pool + # during the first warmup pass; the second warmup pass / capture + # then reads garbage from that recycled memory because no live + # tensor pins those exact bytes during the inter-pass + # allocator reset. The kv_cache_manager-owned tensor is + # outside the graph pool so its bytes stay valid. Linear + # MTP/non-dynamic-tree paths take the dense conv1d branch + # below and never reach here, so the linear path is unchanged. + _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) + # Reshape for batch processing. reshape (not view) because tree + # tokens may be non-contiguous; for linear (contiguous) tokens + # reshape is exactly view, so the linear path is unaffected. + 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 +494,15 @@ def conv1d(): conv_state_indices=state_indices_d[:num_decodes], intermediate_conv_window=intermediate_conv_states, intermediate_state_indices=intermediate_state_indices, + # Tree links (None for linear MTP -> dense conv path). + 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 +637,25 @@ 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's selective_state_update has no tree-parent + # restore; use the native Triton kernel (which does) for + # dynamic tree. Linear MTP keeps the configured func. + 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 +672,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: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index d02e19d5b95c..48a7ab7eeeda 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -784,14 +784,23 @@ def _drop(tensor): torch.cuda.empty_cache() @torch.compile(options={"max-autotune": True}) - def update_mamba_states(self, attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor): + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts num_accepted_draft_tokens = num_accepted_tokens[ num_contexts:num_contexts + num_gens] - 1 + # Linear MTP: accepted path is a contiguous chain (leaf at depth + # num_accepted-1). Dynamic tree: states are recorded in tree order, so + # the leaf's buffer position is its tree-node index (supplied by the + # worker). None -> linear (unchanged). + accepted_positions = (accepted_leaf_positions if accepted_leaf_positions + is not None else num_accepted_draft_tokens) state_indices_d = state_indices[num_contexts:num_contexts + num_gens] src_state_indices = self.intermediate_state_indices[:num_gens] @@ -828,7 +837,7 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", ssm_states = self.mamba_cache.temporal intermediate_ssm_cache = self.mamba_cache.intermediate_ssm accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, - num_accepted_draft_tokens] + accepted_positions] ssm_states[:, state_indices_d, :] = accepted_ssm_state # Conv: both paths save all intermediate conv windows, carry over the accepted one. @@ -836,7 +845,7 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window accepted_conv_state = intermediate_conv_window_cache[:, src_state_indices, - num_accepted_draft_tokens] + accepted_positions] conv_states[:, state_indices_d, :] = accepted_conv_state @@ -980,15 +989,18 @@ def mamba_layer_cache( def shutdown(self): self._impl.shutdown() - def update_mamba_states(self, attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor): + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): # Non-speculative configs don't allocate intermediate state; the # promotion is a clean no-op. if not self._impl.is_speculative(): return self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, - state_indices) + state_indices, accepted_leaf_positions) class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): @@ -1755,10 +1767,12 @@ def is_speculative(self) -> bool: return self.spec_config is not None @nvtx_range("hybrid_update_mamba_states") - def update_mamba_states(self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None): + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): if self.local_num_mamba_layers == 0: return batch_size = attn_metadata.num_seqs @@ -1767,6 +1781,14 @@ def update_mamba_states(self, num_accepted_draft_tokens = ( num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( torch.int32) + # Intermediate-buffer position of each request's accepted leaf state. + # Linear MTP: the accepted path is a contiguous chain, so the leaf sits + # at depth (num_accepted - 1). Dynamic tree: the 31 verified tokens are + # laid out in tree order, so the leaf's buffer position is its tree-node + # index (root/golden at 0), supplied by the worker. + accepted_positions = (accepted_leaf_positions.to(torch.int32) + if accepted_leaf_positions is not None else + num_accepted_draft_tokens) # Match the API of MambaCacheManager.update_mamba_states: callers # may pass per-request state slot indices explicitly (e.g. MTP via # attn_metadata.mamba_metadata.state_indices). Fall back to this @@ -1814,16 +1836,15 @@ def update_mamba_states(self, # Legacy: copy the accepted SSM state from the intermediate buffer. _promote_mamba_state_triton(self.all_ssm_states, self.intermediate_ssm_states, - src_state_indices, - num_accepted_draft_tokens, + src_state_indices, accepted_positions, state_indices_d) # Conv: both paths save all intermediate conv windows, carry over the # accepted one. _promote_mamba_state_triton(self.all_conv_states, self.intermediate_conv_states, - src_state_indices, - num_accepted_draft_tokens, state_indices_d) + src_state_indices, accepted_positions, + state_indices_d) @torch.inference_mode() def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d83f51a1bbe3..8d6c5a53ffbc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1384,6 +1384,65 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): clear_memory_buffers() torch.cuda.empty_cache() + def _need_spec_dec_gen_autotuner_warmup( + self, resource_manager: ResourceManager) -> bool: + """Whether dynamic-tree spec-dec generation needs explicit autotuning.""" + if not getattr(self.llm_args, "enable_autotuner", True): + return False + if not self.cuda_graph_runner.enabled: + return False + if self.spec_config is None or self.is_draft_model: + return False + if not self.spec_config.spec_dec_mode.use_one_engine(): + return False + if not getattr(self.spec_config, "use_dynamic_tree", False): + return False + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + return isinstance(kv_cache_manager, MambaHybridCacheManager) + + def _run_spec_dec_gen_autotuner_warmup( + self, resource_manager: ResourceManager) -> None: + """Profile the CUDA-graph generation MoE shape before capture.""" + if not self._need_spec_dec_gen_autotuner_warmup(resource_manager): + return + + draft_len = self.max_total_draft_tokens + effective_max_seq_len = self.max_seq_len + if self.mapping is not None and self.mapping.has_cp_helix(): + effective_max_seq_len = self.max_seq_len // self.mapping.cp_size + effective_max_seq_len = min(effective_max_seq_len, self.max_num_tokens) + + cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) + AutoTuner.get().setup_distributed_state(self.mapping, self.dist) + logger.info("Running spec-dec generation autotuner warmup...") + with self.no_cuda_graph(), autotune(cache_path=cache_path): + for bs in sorted(self._cuda_graph_batch_sizes, reverse=True): + if bs > self.batch_size: + continue + warmup_request = self._create_cuda_graph_warmup_request( + resource_manager, bs, draft_len, effective_max_seq_len) + with self._release_batch_context(warmup_request, + resource_manager) as batch: + if batch is None: + continue + logger.info( + f"Run pre-capture autotuner warmup at generation shape " + f"(bs={bs}, draft_len={draft_len}, " + f"max_seq_len={effective_max_seq_len})") + self.enable_spec_decode = True + self.runtime_draft_len = draft_len + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) + torch.cuda.synchronize() + + self.enable_spec_decode = self.is_spec_decode + self.runtime_draft_len = self.max_draft_len + logger.info( + f"[Autotuner] Cache size after spec-dec generation warmup is {len(AutoTuner.get().profiling_cache)}" + ) + def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. @@ -1521,6 +1580,16 @@ def _capture_generation_cuda_graphs(self, if self.mapping is not None and self.mapping.has_cp_helix(): effective_max_seq_len = self.max_seq_len // self.mapping.cp_size + # Blackwell trtllm-gen custom-mask spec-dec dynamic-tree generation + # FMHA reads out of bounds during CUDA-graph capture when the capture + # dummy uses the model's full context length. The benchmark uses a + # max-token budget, so capture a graph within that budget. + if (self.spec_config is not None + and self.spec_config.spec_dec_mode.use_one_engine() + and getattr(self.spec_config, 'use_dynamic_tree', False)): + effective_max_seq_len = min(effective_max_seq_len, + self.max_num_tokens) + sparse_config = self.sparse_attention_config if (isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) and sparse_config.needs_separate_short_long_cuda_graphs()): diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index fcf0603a04c1..69ab6a00b7a4 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -140,6 +140,20 @@ class MTPSpecMetadata(SpecMetadata): # CUDA graph, we use this tensor to store the number of input tokens for the # subsequent draft forward. subseq_all_rank_num_tokens: Optional[List[int]] = None + # Dynamic-tree drafting (one-model MTP-Eagle). When use_dynamic_tree is True, + # the target's multi-token verify forward uses a per-slot tree mask sourced + # from spec_tree_manager.slot_storage (see model_engine.update_spec_dec_param). + use_dynamic_tree: bool = False + dynamic_tree_max_topK: Optional[int] = None + spec_tree_manager: Optional[object] = None + # Dynamic-tree per-gen-request tree links for the Mamba tree-aware verify + # forward, gathered from spec_tree_manager.slot_storage in prepare(). + # Shapes [num_generations, max_total_draft_tokens + 1], int32, root at + # index 0. None for linear MTP. retrieve_parent_token stays None: the + # mixer's conv1d kernel derives + fills it from next_token/next_sibling. + retrieve_next_token: Optional[torch.Tensor] = None + retrieve_next_sibling: Optional[torch.Tensor] = None + retrieve_parent_token: Optional[torch.Tensor] = None def __post_init__(self) -> None: if self.mtp_hidden_states_manager is not None: @@ -169,6 +183,12 @@ def __post_init__(self) -> None: self.mtp_num_modules, device='cuda', ) + # Dynamic-tree drafting drives the target's multi-token verify forward + # through the spec-dec tree mask (see model_engine.update_spec_dec_param, + # which routes on these flags + spec_tree_manager.use_dynamic_tree). + if self.use_dynamic_tree: + self.is_spec_dec_tree = True + self.is_spec_dec_dynamic_tree = True @property def all_rank_num_seqs(self): @@ -194,7 +214,15 @@ def prepare(self): # while MTP Eagle worker uses (max_draft_len + 1) input tokens in the 1st draft # forward and only one input token in the following draft forward. # This num_tokens is used to set the all_rank_num_tokens for attention dp. - if not self.spec_dec_mode.is_mtp_eagle_one_model(): + if self.use_dynamic_tree: + # Target verify forward processes (max_total_draft_tokens + 1) tokens + # per gen request, but the draft step-0 forward (which consumes + # all_rank_num_tokens) processes only (max_draft_len + 1) per request + # after prepare_drafter_inputs repacks the accepted path. Correct + # num_tokens to the draft step-0 count for attention dp. + self.num_tokens -= self.num_generations * ( + self.max_total_draft_tokens - self.max_draft_len) + elif not self.spec_dec_mode.is_mtp_eagle_one_model(): self.num_tokens -= self.num_generations if self.mtp_hidden_states_manager is not None: # MTP vanilla or use relaxed acceptance @@ -238,6 +266,35 @@ def prepare(self): if gen_request_ids: sa_manager.prepare(gen_request_ids, self.runtime_draft_len) + # Dynamic tree: gather per-gen-request tree links for the Mamba + # tree-aware verify forward (the tree being verified this step was + # built last step into slot_storage). all_ids_buf is laid out + # [ctx | gen] and was filled by model_engine.fill_all_slot_ids before + # this prepare() runs (see _prepare_inputs). Token order matches the + # target forward's verify layout (root/golden at index 0). + self.retrieve_next_token = None + self.retrieve_next_sibling = None + self.retrieve_parent_token = None + if self.use_dynamic_tree and self.spec_tree_manager is not None: + num_gens = self.num_generations + if num_gens > 0: + num_contexts = num_seqs - num_gens + slot_storage = self.spec_tree_manager.slot_storage + gen_slot_ids = slot_storage.all_ids_buf[ + num_contexts:num_contexts + num_gens] + next_token, next_sibling = slot_storage.next_links_from_slots( + gen_slot_ids, num_gens) + # No-tree gen slots (CUDA-graph/warmup dummies, and a real + # slot's first decode before a tree exists) have sentinel + # links. The Mamba tree-aware verify conv1d/SSU reads these + # unconditionally, so substitute a valid linear chain for those + # rows; real-tree rows are untouched. + slot_storage.apply_no_tree_linear_chain(next_token, + next_sibling, + gen_slot_ids, num_gens) + self.retrieve_next_token = next_token + self.retrieve_next_sibling = next_sibling + class MTPSampler(SpecSamplerBase): """ diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py new file mode 100644 index 000000000000..708621080615 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -0,0 +1,1199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MTP-Eagle one-model dynamic tree speculative decoding (greedy only). + +This worker subclasses :class:`MTPEagleWorker` and reuses the drafter-agnostic +dynamic-tree bookkeeping/verification helpers (``DynamicTreeOpsConverter`` and +``SpecTreeManager``) that were originally written for +``Eagle3OneModelDynamicTreeWorker``. The only piece that is drafter-specific is +the draft loop: instead of running the eagle drafter +(``draft_model.model(...)`` with ``apply_eagle3_fc`` + hidden-state capture), +this worker runs ``draft_model.mtp_layers[0]`` repeatedly (the linear MTP-Eagle +drafter) but grows a topK tree at each layer instead of a linear chain. + +Scope (intentional simplifications): +- GREEDY only (temperature 0). The rejection-sampling path is not implemented. +- The new code path is only reached when ``spec_config.use_dynamic_tree`` is + True; the linear ``MTPEagleWorker`` path is unchanged. +""" + +import math +from typing import TYPE_CHECKING, List, Optional + +import torch +import triton + +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._utils import get_sm_version, nvtx_range +from tensorrt_llm.mapping import Mapping + +from ..distributed.ops import allgather +from ..model_config import ModelConfig +from ..pyexecutor.llm_request import LlmRequest +from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.scheduler import ScheduledRequests +from .eagle3 import MTPEagleWorker + +# Reuse the drafter-agnostic fused helpers from the eagle3 dynamic-tree worker. +# These operate purely on token/score/mask tensors and do not touch the eagle +# drafter, so they are safe to share with the MTP drafter. +from .eagle3_dynamic_tree import ( + _build_mask_and_position, + _gather_repack_step0_kernel, + _resample_final_tokens, + _select_topk_draft_tokens, +) +from .mtp import MTPHiddenStatesManager + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + +class MTPEagleDynamicTreeWorker(MTPEagleWorker): + """MTP-Eagle one-model worker with dynamic tree drafting + verification. + + Inherits the linear MTP-Eagle drafting/sampling primitives from + :class:`MTPEagleWorker` (``draft_sampler``, ``prepare_drafter_inputs``, + ``update_mtp_hidden_states``, ``_prepare_next_new_tokens``, ...) and adds the + dynamic-tree draft loop, greedy verification, and tree construction. + """ + + def __init__( + self, + spec_config: "MTPDecodingConfig", + model_config: Optional[ModelConfig] = None, + use_separate_draft_kv_cache: bool = False, + *, + mapping: Optional[Mapping] = None, + ): + super().__init__(spec_config, + model_config, + use_separate_draft_kv_cache, + mapping=mapping) + assert getattr(spec_config, "use_dynamic_tree", False), ( + "MTPEagleDynamicTreeWorker requires use_dynamic_tree=True" + ) + + from .dynamic_tree_ops import DynamicTreeOpsConverter + + self.K = spec_config.dynamic_tree_max_topK + self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 + self.tokens_per_gen_step = spec_config.tokens_per_gen_step + # _max_batch_size is auto-populated by py_executor_creator from the + # global max_batch_size (mirrors EagleDecodingConfig). It must be set by + # the time we get here. + assert spec_config._max_batch_size is not None, ( + "MTPDecodingConfig._max_batch_size was not populated; " + "py_executor_creator should have set it from the global max_batch_size." + ) + self._max_batch_size = spec_config._max_batch_size + + K = self.K + max_draft_len = spec_config.max_draft_len + max_batch_size = self._max_batch_size + loop_max_tokens = K * max_draft_len # draft loop working size + + # spec_tree_manager is lazily bound from the resource manager. + self.spec_tree_manager = None + self._d2t = None + + # === Pre-allocated draft-loop buffers (CUDA-graph safe) === + # Mirror Eagle3OneModelDynamicTreeWorker.__init__ buffer strategy. + self.draft_tokens_buffer = torch.zeros( + max_batch_size, loop_max_tokens, dtype=torch.int32, device="cuda" + ) + self.position_ids_buffer = torch.zeros( + max_batch_size, loop_max_tokens, dtype=torch.int32, device="cuda" + ) + self.history_draft_tokens_buffer = torch.zeros( + (max_batch_size, (K + K * K * (max_draft_len - 1))), dtype=torch.int32, device="cuda" + ) + self.history_score_buffer = torch.zeros( + (max_batch_size, K + K * K * (max_draft_len - 1)), dtype=torch.float32, device="cuda" + ) + self.history_draft_tokens_parent_buffer = torch.zeros( + (max_batch_size, max(K * (max_draft_len - 1) + 1, K + 1)), + dtype=torch.int64, + device="cuda", + ) + self.tree_mask_buffer = torch.zeros( + (max_batch_size * loop_max_tokens * loop_max_tokens), dtype=torch.int32, device="cuda" + ) + self.tree_mask_init_buffer = ( + torch.eye(K, dtype=torch.int32, device="cuda").unsqueeze(0).repeat(max_batch_size, 1, 1) + ) + self.tree_ops_converter = DynamicTreeOpsConverter( + dynamic_tree_max_topK=K, + max_draft_len=max_draft_len, + max_total_draft_tokens=self.max_total_draft_tokens, + max_batch_size=max_batch_size, + device=torch.device("cuda"), + ) + + self._max_path_len = max_draft_len + 1 + # Step-0 spec-dec reset buffers (mirror Eagle3OneModelDynamicTreeWorker): + # the target verify forward leaves a tokens_per_gen_step-wide tree mask, + # tree position offsets, and a kv_lens inflated by tokens_per_gen_step. + # The step-0 draft attends only to the max_path_len accepted-path tokens, + # so it resets to an 8-wide causal mask + causal offsets and rewinds + # kv_lens by (tokens_per_gen_step - max_path_len). Linear MTP needs none + # of this because there tokens_per_gen_step == max_path_len. + self._kv_correction = self.tokens_per_gen_step - self._max_path_len + self._step0_causal_mask = torch.tensor( + [(1 << (t + 1)) - 1 for t in range(self._max_path_len)], + dtype=torch.int32, + device="cuda", + ) + self._causal_offs = torch.arange(self._max_path_len, device="cuda", dtype=torch.int32) + self._last_selected_parents = None + self._parent_init_arange = torch.arange(-1, K, device="cuda", dtype=torch.int32) + + # Accepted-path bookkeeping for KV relocation and output. + self._accepted_draft_indices_tensor = torch.full( + (max_batch_size, max_draft_len), -1, dtype=torch.int32, device="cuda" + ) + self._kv_head_dim_bytes = None + + # === Verification buffers (greedy only) === + N = self.tokens_per_gen_step + self._accepted_tokens_buf = torch.zeros( + max_batch_size, self._max_path_len, dtype=torch.int32, device="cuda" + ) + self._num_accepted_tokens_buf = torch.ones(max_batch_size, dtype=torch.int32, device="cuda") + self._target_tokens_buf = torch.zeros(max_batch_size * N, dtype=torch.int64, device="cuda") + self._candidates_buf = torch.zeros(max_batch_size, N, dtype=torch.int32, device="cuda") + self._target_predict_buf = torch.zeros(max_batch_size, N, dtype=torch.int32, device="cuda") + + # === Hidden-state management for the growing-context draft loop === + # These mirror the eagle reference's _hs_write_buffer / _accumulated_hs + # but store the MTP layer's output hidden states (no eagle capture). + self._hs_write_buffer = None + self._accumulated_hs = None + self._hs_read_map = torch.zeros( + max_batch_size, loop_max_tokens, dtype=torch.long, device="cuda" + ) + self._step0_hs = None + self._hs_dim = None + + # === Step-0 repack scratch (graph-safe; mirrors eagle3 reference) === + # The target verify forward lays gen tokens out as + # tokens_per_gen_step rows per request, but the MTP draft step-0 only + # attends to the accepted-path (max_path_len) tokens. These buffers hold + # the repacked gen-only [num_gens * max_path_len] inputs; rows are sized + # by the static max (max_batch_size * tokens_per_gen_step) and the hidden + # buffer is lazily sized once the hidden dim is known. + max_total_tokens = max_batch_size * self.tokens_per_gen_step + self._step0_input_ids_buf = torch.zeros(max_total_tokens, dtype=torch.int32, device="cuda") + self._step0_position_ids_buf = torch.zeros( + max_total_tokens, dtype=torch.int32, device="cuda" + ) + self._step0_hidden_states_buf = None + self._gather_ids_buf = torch.zeros(max_total_tokens, dtype=torch.long, device="cuda") + + # Mask repack scratch (graph-safe; avoids .contiguous() in the loop). + buf_dim = max(self.max_total_draft_tokens + 1, loop_max_tokens) + mask_width = (buf_dim + 31) // 32 + self._mask_repack_buf = torch.zeros( + max_batch_size * buf_dim * mask_width, dtype=torch.int32, device="cuda" + ) + # sm>=100 (except 120/121): prepareCustomMask keeps padded 3D; no repack. + sm = get_sm_version() + self._needs_mask_repack = sm < 100 or sm in (120, 121) + + # ------------------------------------------------------------------ # + # Helpers (mirroring eagle3 dynamic-tree worker, drafter-agnostic) # + # ------------------------------------------------------------------ # + def _apply_spec_metadata(self, attn_metadata, batch_size, query_len): + """Set spec-dec gen lengths and refresh the C++ position-offset view.""" + attn_metadata.spec_decoding_generation_lengths[:batch_size] = query_len + attn_metadata.update_position_offsets_for_cpp(query_len) + + def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): + """Compact the padded [n_req, buf_dim, ceil(buf_dim/32)] mask into the + flat prefix XQA expects when n_tok < buf_dim. See the eagle reference + for the detailed rationale.""" + buf_dim = mask_buf.shape[1] + if n_tok >= buf_dim or n_req <= 1: + return + mask_width = math.ceil(n_tok / 32) + total_elems = n_req * n_tok * mask_width + scratch = self._mask_repack_buf[:total_elems].view(n_req, n_tok, mask_width) + scratch.copy_(mask_buf[:n_req, :n_tok, :mask_width]) + flat = mask_buf.view(-1) + flat[:total_elems] = scratch.view(-1) + + @nvtx_range("mtp_dyn._ensure_spec_tree_manager") + def _ensure_spec_tree_manager(self, resource_manager): + """Lazily bind spec_tree_manager and KV head metadata.""" + if self.spec_tree_manager is not None: + return + from ..pyexecutor.resource_manager import ResourceManagerType + + spec_rm = resource_manager.get_resource_manager(ResourceManagerType.SPEC_RESOURCE_MANAGER) + assert spec_rm is not None and hasattr(spec_rm, "spec_tree_manager"), ( + "Dynamic tree mode requires spec_tree_manager in resource_manager" + ) + self.spec_tree_manager = spec_rm.spec_tree_manager + + if self._kv_head_dim_bytes is None: + cache_mgr = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if cache_mgr is not None and hasattr(cache_mgr, "head_dim"): + from tensorrt_llm.bindings import DataType + + _dtype_bytes = { + DataType.HALF: 2, + DataType.BF16: 2, + DataType.FLOAT: 4, + DataType.FP8: 1, + DataType.INT8: 1, + DataType.NVFP4: 0.5, + } + self._kv_head_dim_bytes = int( + cache_mgr.head_dim * _dtype_bytes.get(cache_mgr.dtype, 0.5) + ) + + @nvtx_range("mtp_dyn.sample") + def sample( + self, logits: torch.Tensor, max_top_k: int, draft_model=None + ) -> tuple[torch.Tensor, torch.Tensor]: + """TopK sampling with softmax for the dynamic tree (greedy=topK). + + Returns (topk_indices [.., K], topk_values [.., K]). MTP shares the + target vocabulary through ``shared_head``/``lm_head``, so unlike EAGLE3 + there is no draft->target (d2t) token remap. + + TP correctness: ``DeepseekV3MTPHead`` forces ``gather_output=False`` on + the column-parallel lm_head in pure TP, so ``logits`` here is a per-rank + vocab shard ``[.., vocab/tp]``. Sampling top-K on a shard yields + DIFFERENT draft tokens/scores per rank, so the per-rank trees (and hence + ``num_accepted_tokens`` in the next verify) diverge across TP ranks and + the downstream attention/MoE collectives desync (hang at TP>1). The + linear ``MTPWorker.draft_sampler`` avoids this by all-gathering the + local argmax; the dynamic tree needs full top-K + probabilities, so we + all-gather the full sharded logits (stripping lm_head column padding) + and run softmax+top-K on the replicated full vocab, exactly matching the + TP=1 path. Gated on pure TP (tp_size>1, attention DP off) where the + lm_head output is actually sharded. + """ + mapping = ( + getattr(self.model_config, "mapping", None) if self.model_config is not None else None + ) + if mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp: + logits = allgather(logits, mapping, dim=-1) + if draft_model is not None: + vocab_size = draft_model.lm_head.num_embeddings + logits = logits[..., :vocab_size] + probs = torch.softmax(logits, dim=-1) + topk_values, topk_indices = torch.topk(probs, k=max_top_k, dim=-1) + return topk_indices, topk_values + + def update_draft_tokens_and_scores( + self, + cur_draft_idx, + new_draft_tokens, + new_draft_scores, + previous_draft_scores, + batch_size, + attn_metadata=None, + ): + """Grow the tree: write tokens/scores to history buffers + masks. + + Identical bookkeeping to the eagle reference (drafter-agnostic).""" + if cur_draft_idx == 0: + new_draft_scores = new_draft_scores.reshape(batch_size, self.K) + new_draft_tokens_2d = new_draft_tokens.reshape(batch_size, self.K) + self.draft_tokens_buffer[:batch_size, : self.K] = new_draft_tokens_2d + self.history_draft_tokens_buffer[:batch_size, : self.K] = new_draft_tokens_2d + self.history_score_buffer[:batch_size, : self.K] = new_draft_scores + # Parent buffer: -1 for root, 0..K-1 for first layer. + self.history_draft_tokens_parent_buffer[:batch_size, : self.K + 1] = ( + self._parent_init_arange + ) + self.prepare_tree_mask_and_position_offset(cur_draft_idx, attn_metadata, None) + return new_draft_scores + + ( + real_draft_tokens, + topk_values, + topk_indices, + selected_parents, + new_draft_tokens, + new_draft_scores, + ) = _select_topk_draft_tokens( + new_draft_tokens, new_draft_scores, previous_draft_scores, self.K + ) + + num_tokens_previous_layer = cur_draft_idx * self.K + num_tokens_current_layer = (cur_draft_idx + 1) * self.K + self.draft_tokens_buffer[ + :batch_size, num_tokens_previous_layer:num_tokens_current_layer + ] = real_draft_tokens + + write_start = self.K + (cur_draft_idx - 1) * self.K * self.K + write_end = write_start + self.K * self.K + self.history_draft_tokens_buffer[:batch_size, write_start:write_end] = new_draft_tokens + self.history_score_buffer[:batch_size, write_start:write_end] = new_draft_scores + + self._last_selected_parents = selected_parents + self.prepare_tree_mask_and_position_offset(cur_draft_idx, attn_metadata, selected_parents) + + if cur_draft_idx < self.max_draft_len - 1: + next_layer_start = cur_draft_idx * self.K + 1 + next_layer_end = next_layer_start + self.K + parents_relative_indices = topk_indices + self.K**2 * (cur_draft_idx - 1) + self.K + self.history_draft_tokens_parent_buffer[ + :batch_size, next_layer_start:next_layer_end + ] = parents_relative_indices + return topk_values + + def resampling_final_draft_tokens(self, batch_size: int): + """Reconstruct the final tree from history buffers.""" + return _resample_final_tokens( + self.history_score_buffer[:batch_size, :], + self.history_draft_tokens_buffer[:batch_size, :], + self.max_total_draft_tokens, + ) + + def prepare_tree_mask_and_position_offset( + self, cur_draft_idx, attn_metadata, selected_parents=None + ): + """Prepare mask + position offsets for the next draft layer. + + Drafter-agnostic; identical to the eagle reference.""" + if attn_metadata.spec_decoding_packed_mask is None: + return + spec_tree_manager = self.spec_tree_manager + batch_size = attn_metadata.num_seqs + num_tokens_current_layer = self.K * (cur_draft_idx + 1) + num_tokens_previous_layer = self.K * cur_draft_idx + packed_mask = attn_metadata.spec_decoding_packed_mask + if cur_draft_idx == 0: + spec_tree_manager.compute_spec_dec_packed_mask( + self.tree_mask_init_buffer[:batch_size], + packed_mask[:batch_size, :num_tokens_current_layer, :], + ) + self.tree_mask_buffer[ + : batch_size * num_tokens_current_layer * num_tokens_current_layer + ].copy_(self.tree_mask_init_buffer[:batch_size].view(-1)) + attn_metadata.spec_decoding_position_offsets.fill_(0) + self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) + else: + num_parent_mask = batch_size * cur_draft_idx * self.K * cur_draft_idx * self.K + parent_mask = self.tree_mask_buffer[:num_parent_mask].reshape( + batch_size, cur_draft_idx * self.K, cur_draft_idx * self.K + ) + + prev_total = batch_size * num_tokens_previous_layer + previous_position_offsets = attn_metadata.spec_decoding_position_offsets[ + :prev_total + ].view(batch_size, num_tokens_previous_layer) + + current_mask, new_positions = _build_mask_and_position( + parent_mask, + selected_parents, + self.tree_mask_init_buffer[:batch_size], + previous_position_offsets, + self.K, + ) + + spec_tree_manager.compute_spec_dec_packed_mask( + current_mask, packed_mask[:batch_size, :num_tokens_current_layer, :] + ) + self.tree_mask_buffer[ + : batch_size * num_tokens_current_layer * num_tokens_current_layer + ].copy_(current_mask.reshape(-1)) + + cur_total = batch_size * num_tokens_current_layer + attn_metadata.spec_decoding_position_offsets[:cur_total] = new_positions.reshape(-1) + self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) + + if self._needs_mask_repack: + self._repack_mask_padded_to_packed(packed_mask, batch_size, num_tokens_current_layer) + + def update_hidden_states( + self, + cur_draft_idx, + batch_size, + step0_hs=None, + hidden_states_to_save=None, + selected_parents=None, + ): + """Manage the growing-context hidden states for the MTP draft loop. + + Unlike the eagle reference (which saves eagle prenorm hidden states), + we save the MTP layer's OUTPUT hidden states. The gather/parent logic is + otherwise identical.""" + if cur_draft_idx == 0: + hs_dim = step0_hs.shape[-1] + self._hs_dim = hs_dim + if self._hs_write_buffer is None or self._hs_write_buffer.shape[2] != hs_dim: + self._hs_write_buffer = torch.zeros( + self._max_batch_size, + self.max_draft_len * self.K, + hs_dim, + device=step0_hs.device, + dtype=step0_hs.dtype, + ) + if self._accumulated_hs is None or self._accumulated_hs.shape[2] != hs_dim: + self._accumulated_hs = torch.zeros( + self._max_batch_size, + self.max_draft_len * self.K, + hs_dim, + device=step0_hs.device, + dtype=step0_hs.dtype, + ) + # All K depth-0 tokens share step0_hs (the parent hidden state). + self._accumulated_hs[:batch_size, : self.K] = step0_hs.unsqueeze(1).expand( + -1, self.K, -1 + ) + self._step0_hs = step0_hs + else: + num_tokens_per_req = cur_draft_idx * self.K + hs_to_save_reshaped = hidden_states_to_save.reshape(batch_size, num_tokens_per_req, -1) + self._hs_write_buffer[:batch_size, :num_tokens_per_req] = hs_to_save_reshaped + parent_offset = (cur_draft_idx - 1) * self.K + self._hs_read_map[ + :batch_size, cur_draft_idx * self.K : (cur_draft_idx + 1) * self.K + ] = parent_offset + selected_parents + num_tokens_next = (cur_draft_idx + 1) * self.K + read_idx = self._hs_read_map[:batch_size, self.K : num_tokens_next] + hs_dim = self._hs_write_buffer.shape[2] + self._accumulated_hs[:batch_size, self.K : num_tokens_next] = torch.gather( + self._hs_write_buffer[:batch_size], 1, read_idx.unsqueeze(-1).expand(-1, -1, hs_dim) + ) + + # ------------------------------------------------------------------ # + # Verification (greedy only) # + # ------------------------------------------------------------------ # + @nvtx_range("mtp_dyn.sample_and_accept_draft_tokens") + def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_metadata): + """Greedy dynamic-tree verification of the PREVIOUS step's tree. + + Overrides MTPWorker.sample_and_accept_draft_tokens. Returns + (accepted_tokens [bs, max_path_len], num_accepted_tokens [bs]).""" + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + N = self.tokens_per_gen_step + max_path_len = self._max_path_len + + if logits.dim() == 1: + logits = logits.unsqueeze(0) + + # Reset output buffers. + self._accepted_tokens_buf[:batch_size].zero_() + accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] + self._num_accepted_tokens_buf[:batch_size].fill_(1) + num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] + self._accepted_draft_indices_tensor[:batch_size].fill_(-1) + + num_flat_tokens = logits.shape[0] + torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) + target_tokens = self._target_tokens_buf[:num_flat_tokens] + + # Context requests: accept the sampled golden token only. + accepted_tokens[:num_contexts, 0].copy_(target_tokens[:num_contexts]) + + if num_gens > 0: + spec_tree_manager = self.spec_tree_manager + target_predict = self._target_predict_buf[:num_gens] + target_predict.copy_(target_tokens[num_contexts:].reshape(num_gens, N)) + + # First-step bootstrap / CUDA-graph warmup: no tree exists yet, so + # accept only the golden token (the first of the gen tokens). + if spec_tree_manager is None: + num_accepted_tokens[num_contexts:batch_size] = 1 + accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0] + self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 + return accepted_tokens, num_accepted_tokens + + # candidates[:, 0] = golden token, candidates[:, 1:] = draft tokens. + candidates = self._candidates_buf[:num_gens] + candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1) + candidates[:, 0] = target_predict[:, 0] + + slot_storage = spec_tree_manager.slot_storage + gen_slot_ids = slot_storage.all_ids_buf[num_contexts : num_contexts + num_gens] + tree_valid = slot_storage.has_tree[gen_slot_ids] + retrieve_packed = slot_storage.pack_retrieve_from_slots(gen_slot_ids, num_gens) + + accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_greedy_out_packed( + candidates, + retrieve_packed, + target_predict, + num_gens, + self._max_path_len, + tree_valid=tree_valid, + ) + ) + + accepted_draft_count = accept_token_num[:num_gens] + num_accepted_tokens[num_contexts:batch_size] = (accepted_draft_count + 1).to( + torch.int32 + ) + accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) + # accept_index stores root at slot 0; subtract 1 so root/padding 0 + # becomes the sentinel -1 (tree node index into the draft tokens). + self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( + accept_index[:num_gens, 1:max_path_len] - 1 + ).to(torch.int32) + + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, self.max_draft_len + ) + return accepted_tokens, num_accepted_tokens + + def _accepted_leaf_intermediate_positions(self, num_accepted_tokens, num_contexts, num_gens): + """Tree-node position of each gen request's accepted leaf in the mamba + intermediate-state buffer. + + The mixer records states in verify-token order: position 0 is the + golden/root token, positions 1.. are the draft nodes in tree order. + ``_accepted_draft_indices_tensor[r, c]`` holds the (c+1)-th accepted + draft node's tree index (= candidate position - 1), so the deepest + accepted node lives at column ``num_accepted - 2`` and its buffer + position is that index + 1. When only the golden token is accepted + (num_accepted == 1) the leaf is the root at position 0. + """ + accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens].to(torch.int64) + # Column of the deepest accepted draft node, clamped to >=0 for the + # golden-only case (its value is ignored via the mask below). + last_col = (accepted - 2).clamp_(min=0) + draft_idx = self._accepted_draft_indices_tensor[num_contexts : num_contexts + num_gens].to( + torch.int64 + ) + leaf = torch.gather(draft_idx, 1, last_col.unsqueeze(1)).squeeze(1) + 1 + # Golden-only requests (num_accepted == 1) take the root at position 0. + return torch.where(accepted > 1, leaf, torch.zeros_like(leaf)) + + @nvtx_range("mtp_dyn._relocate_kv_eagerly") + def _relocate_kv_eagerly(self, attn_metadata, batch_size): + """Move accepted draft tokens' KV from tree positions to the linear + prefix the next step expects. Mirrors the eagle reference. + + Mamba-2 hybrid handling: the parent KVCacheManager spans every global + layer, but Mamba layers carry recurrent state (``num_kv_heads == 0``) + and live in a *separate* C++ pool from the attention layers. The + ``update_kv_cache_draft_token_location_2d`` op addresses a single pool + with a uniform per-layer stride (``layerIdx * 2 * bytesPerBlock`` off + one base pointer, see ``updateKVBlockArrayDraftTokenLocation2D``), and + the stored block offsets are scaled by *that pool's* layer count + (``flat_index3(blockIdx, 0, fieldIdx, pool.numLayers, kvFactor)``). So + we must drive the op with the attention pool only: its compact layer + count, its uniform head count, and its slice of the pool-pointer and + block-offset tensors. For a pure-attention model there is exactly one + pool and this reduces to the original call.""" + cache_mgr = getattr(attn_metadata, "kv_cache_manager", None) + if cache_mgr is None or self._kv_head_dim_bytes is None: + return + if not hasattr(cache_mgr, "num_kv_heads_per_layer"): + return + + # Attention layers are those with KV heads (Mamba layers are zeroed). + # The set is static for a given model, so the layerCount / head count / + # pool index below are CUDA-graph-safe (no data-dependent control flow). + kv_heads = cache_mgr.num_kv_heads_per_layer + attn_heads = set(h for h in kv_heads if h > 0) + assert len(attn_heads) == 1, ( + "update_kv_cache_draft_token_location_2d requires uniform " + f"num_kv_heads across attention layers, got {list(kv_heads)}" + ) + attn_num_heads = attn_heads.pop() + attn_layer_offsets = [i for i, h in enumerate(kv_heads) if h > 0] + attn_num_layers = len(attn_layer_offsets) + + # All attention layers share one pool (same head count => same pool in + # the C++ WindowBlockManager). Resolve its index from the layer->pool + # mapping so we slice the correct pool's pointers/offsets; the op reads + # pool_pointers[0]/[1] as that pool's primary/secondary base. + pool_mapping = getattr(cache_mgr, "kv_cache_pool_mapping", None) + if pool_mapping is not None: + attn_pool_indices = set(int(pool_mapping[off][0]) for off in attn_layer_offsets) + assert len(attn_pool_indices) == 1, ( + "update_kv_cache_draft_token_location_2d requires all attention " + f"layers in one KV pool, got pools {sorted(attn_pool_indices)}" + ) + attn_pool_idx = attn_pool_indices.pop() + else: + attn_pool_idx = 0 + + pool_pointers = cache_mgr.kv_cache_pool_pointers[attn_pool_idx] + block_offsets = attn_metadata.kv_cache_block_offsets[attn_pool_idx] + + torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d( + self._accepted_draft_indices_tensor[:batch_size], + self._num_accepted_tokens_buf[:batch_size], + attn_metadata.kv_lens_cuda[:batch_size], + True, + attn_num_layers, + attn_num_heads, + self._kv_head_dim_bytes, + cache_mgr.max_total_draft_tokens, + cache_mgr.max_attention_window_vec[0], + pool_pointers, + block_offsets, + cache_mgr.max_blocks_per_seq, + cache_mgr.tokens_per_block, + None, + ) + + # ------------------------------------------------------------------ # + # Top-level forward # + # ------------------------------------------------------------------ # + @nvtx_range("mtp_dyn.forward") + def forward( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Dynamic-tree MTP-Eagle forward. + + Step-by-step (see module/report for divergence notes): + (a) verify the previous tree against target logits (greedy); + (b) update the Mamba hybrid cache with num_accepted_tokens, and + relocate accepted draft-token KV; + (c) run the MTP draft TREE loop and build the new tree into + spec_tree_manager.slot_storage; + (d) return the MTPEagleWorker output contract. + """ + if resource_manager is not None: + self._ensure_spec_tree_manager(resource_manager) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + raw_logits = logits + + self._execute_guided_decoder_if_present(logits) + + # (a) Verify previous tree (greedy). Also relocates accepted KV. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, spec_metadata, attn_metadata + ) + if num_gens > 0: + self._relocate_kv_eagerly(attn_metadata, batch_size) + + # (b) Update Mamba hybrid cache for accepted variable-length paths. + # The mixer records one intermediate state per verified token in + # tree order (root/golden at buffer position 0), so the accepted + # leaf's state lives at its tree-node position, not at linear depth + # num_accepted-1. Compute that per-gen-request leaf position from the + # accepted draft-node tree indices and pass it to the cache rollback. + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + accepted_leaf_positions = self._accepted_leaf_intermediate_positions( + num_accepted_tokens, num_contexts, num_gens + ) + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + accepted_leaf_positions=accepted_leaf_positions, + ) + + # Save attn/spec metadata before the draft loop mutates it. + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # (c) Run the MTP draft tree loop -> build + store the next tree. + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + next_draft_tokens = self._forward_draft_loop( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + draft_kv_cache_manager=draft_kv_cache_manager, + num_contexts=num_contexts, + num_gens=num_gens, + batch_size=batch_size, + ) + + # Restore attn metadata to support cuda graph. + self._restore_attn_metadata_from_spec_dec(attn_metadata) + attn_metadata.use_spec_decoding = True + + # (d) Prepare next_new_tokens for overlap scheduler. + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + "accepted_draft_tokens_indices": self._accepted_draft_indices_tensor[:batch_size], + } + + # ------------------------------------------------------------------ # + # Step-0 drafter-input repack (dynamic tree) # + # ------------------------------------------------------------------ # + @nvtx_range("mtp_dyn._prepare_step0_drafter_inputs") + def _prepare_step0_drafter_inputs( + self, + input_ids, + position_ids, + last_tokens_idx, + hidden_states, + accepted_tokens, + attn_metadata, + ): + """Repack step-0 drafter inputs to the accepted-path layout. + + ``MTPEagleWorker.prepare_drafter_inputs`` only repacks the gen + ``input_ids`` to ``num_gens * max_path_len`` rows while leaving + ``hidden_states`` / ``position_ids`` in the target verify layout + (``num_gens * tokens_per_gen_step`` rows). For the linear MTP path the + two layouts coincide (``tokens_per_gen_step == max_path_len``); for the + dynamic tree they differ, and the NemotronHMTP ``eh_proj`` fusion + ``cat([enorm(embed(input_ids)), hnorm(hidden_states)], dim=-1)`` then + sees mismatched row counts and crashes. + + This mirrors ``Eagle3OneModelDynamicTreeWorker.prepare_1st_drafter_inputs``: + a single fused Triton kernel gathers the TARGET ``hidden_states`` at the + accepted tree positions and repacks ``input_ids`` / ``position_ids`` + into the ``[ctx | gen (num_gens * max_path_len)]`` layout, and writes the + per-gen-request "last accepted token" gather id into ``_gather_ids_buf``. + + Unlike eagle3 there is no ``spec_metadata.hidden_states`` / + ``apply_eagle3_fc`` / ``layers_to_capture`` indirection: MTP shares the + target vocab and consumes the target hidden states directly. + + Returns the drafter ``inputs`` dict. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + num_ctx_tokens = attn_metadata.num_ctx_tokens + + # Context input_ids: shift-left + place golden token at last positions + # (identical to MTPEagleWorker.prepare_drafter_inputs context path). + input_ids_ctx = self._prepare_context_input_ids( + input_ids, num_ctx_tokens, last_tokens_idx, accepted_tokens, num_contexts + ) + + if num_gens > 0: + max_path_len = self._max_path_len + num_gen_tokens = num_gens * max_path_len + + hidden_dim = hidden_states.shape[-1] + if ( + self._step0_hidden_states_buf is None + or self._step0_hidden_states_buf.shape[-1] != hidden_dim + ): + self._step0_hidden_states_buf = torch.zeros( + self._step0_input_ids_buf.shape[0], + hidden_dim, + dtype=hidden_states.dtype, + device="cuda", + ) + + # accepted_tokens[num_contexts:] is the accepted path (incl golden + # at col 0); it is exactly the eagle reference's ``_accept_token``. + accept_token = accepted_tokens[num_contexts:batch_size] + + BLOCK_H = triton.next_power_of_2(hidden_dim) + _gather_repack_step0_kernel[(num_gens * max_path_len,)]( + hidden_states, + accept_token, + position_ids, + self._accepted_draft_indices_tensor[num_contexts:batch_size], + self._num_accepted_tokens_buf, + self._step0_hidden_states_buf, + self._step0_input_ids_buf, + self._step0_position_ids_buf, + self._gather_ids_buf, + num_ctx_tokens, + num_contexts, + self.tokens_per_gen_step, + max_path_len, + self.max_draft_len, + hidden_dim, + num_ctx_tokens, # gather_id references combined [ctx|gen] tensor + BLOCK_H=BLOCK_H, + ) + + input_ids = torch.cat( + [input_ids_ctx, self._step0_input_ids_buf[:num_gen_tokens]], dim=0 + ) + position_ids = torch.cat( + [position_ids[:num_ctx_tokens], self._step0_position_ids_buf[:num_gen_tokens]], + dim=0, + ) + hidden_states = torch.cat( + [hidden_states[:num_ctx_tokens], self._step0_hidden_states_buf[:num_gen_tokens]], + dim=0, + ) + + attn_metadata._seq_lens[num_contexts:batch_size].fill_(max_path_len) + attn_metadata._seq_lens_cuda[num_contexts:batch_size].fill_(max_path_len) + attn_metadata.on_update() + else: + # Context-only (warmup): no gen tokens to repack. + input_ids = input_ids_ctx + + return { + "input_ids": input_ids, + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + } + + # ------------------------------------------------------------------ # + # MTP draft tree loop # + # ------------------------------------------------------------------ # + def _forward_draft_loop( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + num_accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + ): + """MTP dynamic-tree draft loop with growing context. + + Step 0 runs mtp_layers[0] over the accepted-path tokens (max_path_len + per request, repacked by _prepare_step0_drafter_inputs), then expands + topK. Each subsequent layer re-runs mtp_layers[0] over ALL accumulated + tree tokens and expands topK per surviving parent. Finally the tree is + resampled and built into slot_storage so the NEXT target forward uses + the tree mask. + """ + spec_tree_manager = self.spec_tree_manager + + assert batch_size <= self._max_batch_size, ( + f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}" + ) + + # --- Step 0: one MTP forward over accepted golden tokens --- + # Repack input_ids/position_ids AND hidden_states to the accepted-path + # layout (num_gens * max_path_len gen rows) so the NemotronHMTP eh_proj + # fusion sees matching row counts. The fused kernel also writes the + # per-request last-accepted-token gather id into _gather_ids_buf. + position_ids, last_tokens_idx = self.prepare_position_ids_and_last_tokens( + position_ids, attn_metadata.seq_lens_cuda + ) + inputs = self._prepare_step0_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + last_tokens_idx=last_tokens_idx, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + ) + + # Step-0 causal spec-dec reset (mirrors Eagle3OneModelDynamicTreeWorker). + # The target verify forward left a tokens_per_gen_step-wide tree mask + + # position offsets and an inflated kv_lens; the step-0 draft attends only + # to the max_path_len accepted-path tokens, so reset to an 8-wide causal + # mask + causal offsets and rewind kv_lens. None in prefill-only warmup. + num_step0_tokens = self._max_path_len + if attn_metadata.spec_decoding_generation_lengths is not None: + total = num_gens * num_step0_tokens + dst = attn_metadata.spec_decoding_position_offsets[:total].view( + num_gens, num_step0_tokens + ) + dst.copy_(self._causal_offs[:num_step0_tokens].unsqueeze(0).expand(num_gens, -1)) + self._apply_spec_metadata(attn_metadata, num_gens, num_step0_tokens) + packed_mask = attn_metadata.spec_decoding_packed_mask + packed_mask[:num_gens].zero_() + packed_mask[:num_gens, :num_step0_tokens, 0] = self._step0_causal_mask[ + :num_step0_tokens + ] + if self._needs_mask_repack: + self._repack_mask_padded_to_packed(packed_mask, num_gens, num_step0_tokens) + attn_metadata.use_spec_decoding = num_gens > 0 + if num_gens > 0 and hasattr(attn_metadata, "kv_lens_cuda"): + attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= self._kv_correction + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states = draft_model.mtp_layers[0]( + embed_tokens=draft_model.embed_tokens, + all_rank_num_tokens=spec_metadata.all_rank_num_tokens, + **inputs, + ) + + # Gather the per-request "last accepted token" hidden state (the + # tree root for depth-0 expansion). The Triton kernel already wrote + # gen gather ids (num_ctx_tokens + gen_idx * max_path_len + + # num_accepted - 1) into _gather_ids_buf; prepend the context last- + # token ids. This indexes the post-repack [ctx | gen max_path_len] + # hidden_states layout. + self._gather_ids_buf[:num_contexts].copy_(last_tokens_idx[:num_contexts]) + gather_ids = self._gather_ids_buf[:batch_size] + + step0_hs = hidden_states[gather_ids] + logits = draft_model.mtp_layers[0].shared_head( + step0_hs, draft_model.lm_head, attn_metadata, True + ) + + new_draft_tokens, new_draft_scores = self.sample( + logits, self.K, draft_model=draft_model + ) + previous_draft_scores = self.update_draft_tokens_and_scores( + cur_draft_idx=0, + new_draft_tokens=new_draft_tokens, + new_draft_scores=new_draft_scores, + previous_draft_scores=None, + batch_size=batch_size, + attn_metadata=attn_metadata, + ) + self.update_hidden_states(cur_draft_idx=0, batch_size=batch_size, step0_hs=step0_hs) + self._prepare_draft_layer_metadata( + 0, + attn_metadata, + batch_size, + gather_ids, + num_contexts, + num_gens, + num_accepted_tokens, + inputs, + ) + + # --- Subsequent layers: grow the tree --- + for layer_idx in range(1, self.max_draft_len): + num_tokens_per_req = layer_idx * self.K + num_infer_tokens = batch_size * num_tokens_per_req + + inp_hs = self._accumulated_hs[:batch_size, :num_tokens_per_req, :].reshape( + num_infer_tokens, -1 + ) + inp_ids = self.draft_tokens_buffer[:batch_size, :num_tokens_per_req].reshape(-1) + inp_pos = self.position_ids_buffer[:batch_size, :num_tokens_per_req].reshape(-1) + layer_inputs = { + "input_ids": inp_ids, + "position_ids": inp_pos, + "hidden_states": inp_hs, + "attn_metadata": attn_metadata, + } + + hidden_states = draft_model.mtp_layers[0]( + embed_tokens=draft_model.embed_tokens, + all_rank_num_tokens=spec_metadata.subseq_all_rank_num_tokens, + **layer_inputs, + ) + + # Take the last K hidden states per request (the new leaves). + hs_reshaped = hidden_states.reshape(batch_size, num_tokens_per_req, -1) + selected_hs = hs_reshaped[:, -self.K :, :].reshape(batch_size * self.K, -1) + logits = draft_model.mtp_layers[0].shared_head( + selected_hs, draft_model.lm_head, attn_metadata, True + ) + + new_draft_tokens, new_draft_scores = self.sample( + logits, self.K, draft_model=draft_model + ) + new_draft_tokens = new_draft_tokens.reshape(batch_size, self.K, self.K) + new_draft_scores = new_draft_scores.reshape(batch_size, self.K, self.K) + + previous_draft_scores = self.update_draft_tokens_and_scores( + cur_draft_idx=layer_idx, + new_draft_tokens=new_draft_tokens, + new_draft_scores=new_draft_scores, + previous_draft_scores=previous_draft_scores, + batch_size=batch_size, + attn_metadata=attn_metadata, + ) + self.update_hidden_states( + cur_draft_idx=layer_idx, + batch_size=batch_size, + hidden_states_to_save=hidden_states, + selected_parents=self._last_selected_parents, + ) + self._prepare_draft_layer_metadata(layer_idx, attn_metadata, batch_size) + + # Resample the final tree and build it into slot_storage. + real_draft_tokens, topk_score_indices = self.resampling_final_draft_tokens(batch_size) + + if spec_tree_manager is not None and num_gens > 0: + self.tree_ops_converter.build_dynamic_tree( + history_draft_tokens_parent_buffer=self.history_draft_tokens_parent_buffer[ + num_contexts:batch_size + ], + topk_score_indices=topk_score_indices[num_contexts:], + tree_mask=spec_tree_manager.spec_dec_packed_mask[:num_gens], + positions=spec_tree_manager.spec_dec_position_offsets[:num_gens], + retrieve_index=spec_tree_manager.retrieve_index[:num_gens], + retrieve_next_token=spec_tree_manager.retrieve_next_token[:num_gens], + retrieve_next_sibling=spec_tree_manager.retrieve_next_sibling[:num_gens], + use_packed_mask=True, + ) + slot_storage = spec_tree_manager.slot_storage + gen_slots = slot_storage.all_ids_buf[num_contexts:batch_size] + spec_tree_manager.scatter_to_slot_storage(slot_storage, gen_slots, num_gens) + + return real_draft_tokens + + def _prepare_draft_layer_metadata( + self, + cur_draft_idx, + attn_metadata, + batch_size, + gather_ids=None, + num_contexts=0, + num_gens=0, + num_accepted_tokens=None, + inputs=None, + ): + """Set up attn_metadata seq_lens/kv_lens for the next drafter layer. + + Drafter-agnostic; mirrors the eagle reference's prepare_for_generation + (which is itself derived from MTPEagleWorker's i==0 / i>0 metadata + updates).""" + if cur_draft_idx == 0: + base_pos = inputs["position_ids"][gather_ids] + 1 + self.position_ids_buffer[:batch_size, : self.K] = base_pos.unsqueeze(1).expand( + -1, self.K + ) + + attn_metadata._seq_lens[:batch_size].fill_(self.K) + attn_metadata._seq_lens_cuda[:batch_size].fill_(self.K) + attn_metadata.on_update() + + if inputs["attn_metadata"].kv_cache_manager is not None: + attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + + if hasattr(attn_metadata, "kv_lens_cuda"): + # KV rewind: remove unaccepted draft-path tokens, then add the K + # new depth-0 tokens. + if num_gens > 0: + attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( + self._max_path_len + ) - num_accepted_tokens[num_contexts:batch_size] + attn_metadata.kv_lens_cuda[:batch_size] += self.K + attn_metadata.use_spec_decoding = True + attn_metadata.update_for_spec_dec() + else: + num_tokens_previous_layer = cur_draft_idx * self.K + num_tokens_current_layer = self.K * (cur_draft_idx + 1) + prev_pos = self.position_ids_buffer[:batch_size, :num_tokens_previous_layer] + self.position_ids_buffer[ + :batch_size, num_tokens_previous_layer:num_tokens_current_layer + ] = prev_pos[:, -self.K :] + 1 + attn_metadata._seq_lens[:batch_size].fill_(num_tokens_current_layer) + attn_metadata._seq_lens_cuda[:batch_size].fill_(num_tokens_current_layer) + attn_metadata.on_update() + if hasattr(attn_metadata, "kv_lens_cuda"): + attn_metadata.kv_lens_cuda[:batch_size] += self.K + attn_metadata.update_for_spec_dec() + + +class MTPEagleDynamicTreeResourceManager(BaseResourceManager): + """Resource manager for one-model MTP-Eagle dynamic tree mode. + + Composes: + - a ``SpecTreeManager`` (exposed as ``.spec_tree_manager``) so the model + engine / attention backend can wire the per-slot tree mask for the + target's multi-token verify forward, and the worker can build/store the + tree, and + - an ``MTPHiddenStatesManager`` so MTPEagleWorker's drafter-input + preparation (mtp_past_tokens / mtp_past_hidden_states slot pools and + slot_ids) keeps working. + """ + + hidden_states: Optional[torch.Tensor] = None + + def __init__( + self, + config: "MTPDecodingConfig", + dtype: torch.dtype, + hidden_size: int, + max_num_requests: int, + sa_manager=None, + ): + from .spec_tree_manager import SpecTreeManager + + self.max_num_requests = max_num_requests + self.spec_tree_manager = SpecTreeManager( + max_num_requests=max_num_requests, + use_dynamic_tree=True, + max_draft_len=config.max_draft_len, + max_total_draft_tokens=config.tokens_per_gen_step - 1, + eagle_choices=None, + dynamic_tree_max_topK=config.dynamic_tree_max_topK, + ) + # MTP hidden-state slot pools (needed by MTPEagleWorker drafter inputs). + self._mtp_hidden_states_manager = MTPHiddenStatesManager( + config, dtype, hidden_size, max_num_requests, sa_manager=sa_manager + ) + + # Expose the MTPHiddenStatesManager surface MTPSpecMetadata expects. + @property + def slot_manager(self): + return self._mtp_hidden_states_manager.slot_manager + + @property + def mtp_past_hidden_states_pool(self): + return self._mtp_hidden_states_manager.mtp_past_hidden_states_pool + + @property + def mtp_past_tokens_pool(self): + return self._mtp_hidden_states_manager.mtp_past_tokens_pool + + @property + def sa_manager(self): + return self._mtp_hidden_states_manager.sa_manager + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + self._mtp_hidden_states_manager.prepare_resources(scheduled_batch) + + def update_resources(self, scheduled_batch: ScheduledRequests): + self._mtp_hidden_states_manager.update_resources(scheduled_batch) + + def free_resources(self, request: LlmRequest): + # Clear tree validity for the freed slot, then free the MTP slot. + if request.py_seq_slot is not None: + self.spec_tree_manager.slot_storage.mark_invalid(request.py_seq_slot) + self._mtp_hidden_states_manager.free_resources(request) + + def add_dummy_requests(self, request_ids: List[int]): + # Dynamic-tree dummies use slot_storage.dummy_slot_id (no per-request + # slot), but the MTP hidden-state pool still needs a slot per dummy. + self._mtp_hidden_states_manager.add_dummy_requests(request_ids) + + def shutdown(self): + self._mtp_hidden_states_manager.shutdown() + + def get_max_resource_count(self) -> int: + return self.max_num_requests + + def get_needed_resource_to_completion(self, request: LlmRequest): + return 0 diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 74dd622826f1..6c5a5fb93bfb 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -16,24 +16,79 @@ class DynamicTreeSlotStorage: Buffers are [S, ...] where S = num_slots + 1 (+1 for CUDA graph dummy). """ - def __init__(self, num_slots: int, n_dt: int, mask_width: int): + def __init__(self, + num_slots: int, + n_dt: int, + mask_width: int, + topK: int = 1): S = num_slots + 1 self.dummy_slot_id = num_slots - # Slot buffers — C++ kernel writes directly via slotIds - self.packed_mask = torch.zeros((S, n_dt, mask_width), - dtype=torch.int32, - device='cuda') - self.position_offsets = torch.zeros((S, n_dt), - dtype=torch.int32, - device='cuda') + # Slot buffers — C++ kernel writes directly via slotIds. + # position_offsets / packed_mask init to a valid degenerate LINEAR chain + # (token i at depth i, attending to tokens 0..i incl. self), NOT zeros. + # The reserved CUDA-graph dummy slot (warmup/capture) is never written by + # the C++ tree scatter, so the spec-dec FMHA reads these rows as-is; a + # zeros packed_mask has no self-attention bit and OOBs the trtllm-gen + # kernel. Mirrors the retrieve_next_token chain below; real trees + # overwrite the row via the C++ scatter. bit (32*w + j) of packed_mask + # [i, w] set <=> token i attends to tree token (32*w + j); the causal + # value matches the static-tree mask formula 2^(i+1)-1. + _tok = torch.arange(n_dt, device='cuda') + self.position_offsets = _tok.to(torch.int32).unsqueeze(0).repeat( + S, 1).contiguous() + _bit = torch.arange(mask_width * 32, device='cuda') + _causal = ((_bit.unsqueeze(0) <= _tok.unsqueeze(1)) + & (_bit.unsqueeze(0) < n_dt)).view(n_dt, mask_width, 32) + _w = 2**torch.arange(32, dtype=torch.int64, device='cuda') + self.packed_mask = (_causal.to(torch.int64) * _w).sum(-1).to( + torch.int32).unsqueeze(0).repeat(S, 1, 1).contiguous() + + # Override ONLY the reserved CUDA-graph/warmup dummy slot with a + # bounded-depth K-ary tree (parent[i] = (i-1)//topK). Unlike a real + # slot's no-tree fallback — read only 1 token wide on its first decode — + # the dummy slot is read at the FULL n_dt-wide generation shape by the + # spec-dec verify forward during the CUDA-graph generation warmup. A + # depth-(n_dt-1) linear chain there presents a tree real requests never + # produce (real max depth = max_draft_len, sparse ancestor mask). The + # K-ary template mirrors the drafter's topK expansion so the warmup's + # dummy metadata matches what real dynamic-tree requests feed the + # trtllm-gen FMHA. Scoped to the dummy row so real slots (and the + # accepted eager path) keep the linear fallback above unchanged. + _k = max(int(topK), 1) + _depth = torch.zeros(n_dt, dtype=torch.int32) + _adj = torch.zeros(n_dt, n_dt, dtype=torch.bool) + for _i in range(n_dt): + _adj[_i, _i] = True # self + if _i > 0: + _p = (_i - 1) // _k # parent index < _i, its row already final + _depth[_i] = _depth[_p] + 1 + _adj[_i] |= _adj[_p] # inherit ancestors (incl. root) + self.position_offsets[self.dummy_slot_id] = _depth.to(device='cuda') + _adj_pad = torch.zeros(n_dt, mask_width * 32, dtype=torch.bool) + _adj_pad[:, :n_dt] = _adj + _dummy_mask = (_adj_pad.to(device='cuda').view(n_dt, mask_width, 32).to( + torch.int64) * _w).sum(-1).to(torch.int32) + self.packed_mask[self.dummy_slot_id] = _dummy_mask self.retrieve_index = torch.zeros((S, n_dt), dtype=torch.int32, device='cuda') - self.retrieve_next_token = torch.full((S, n_dt), - -1, - dtype=torch.int32, - device='cuda') + + # Degenerate linear-chain next-token links for no-tree slots (dummy + # CUDA-graph/warmup requests and a real slot's first decode before any + # tree is built). Token i's child is i+1 (parent i-1); the leaf has no + # child (-1). The Mamba tree-aware conv1d/SSU verify path indexes + # retrieve_next_token unconditionally (it does not gate on has_tree), so + # every no-tree row must describe a valid chain rather than sentinels. + # retrieve_next_token is initialized to the chain (not -1) so a + # never-built slot — notably the reserved dummy slot used by CUDA-graph + # capture/warmup, which never passes through the C++ scatter or + # prepare()'s has_tree substitution — still gathers valid, in-bounds + # parent links. Real trees overwrite the row via the C++ scatter. + chain = torch.arange(1, n_dt + 1, dtype=torch.int32, device='cuda') + chain[n_dt - 1] = -1 + self._no_tree_next_token = chain + self.retrieve_next_token = chain.unsqueeze(0).repeat(S, 1) self.retrieve_next_sibling = torch.full((S, n_dt), -1, dtype=torch.int32, @@ -111,6 +166,29 @@ def next_links_from_slots(self, slot_ids, count): torch.index_select(self.retrieve_next_sibling, 0, ids, out=next_sibling) return next_token, next_sibling + def apply_no_tree_linear_chain(self, next_token, next_sibling, slot_ids, + count): + """Overwrite no-tree rows' links with a valid degenerate linear chain. + + For gen slots whose tree was not built this step (has_tree False: + CUDA-graph/warmup dummies, and a real slot's first decode), the gathered + links are sentinels/uninitialized. The Mamba tree-aware verify path + reads them unconditionally, so replace those rows in-place with the + linear-chain template (next_token[i]=i+1, leaf -1; next_sibling -1) so + the conv1d/SSU parent traversal stays in-bounds and matches the + captured-graph op sequence a real-tree forward replays. Rows with a + real tree (has_tree True) are left untouched. + """ + if count == 0: + return next_token, next_sibling + ids = slot_ids[:count] + no_tree = ~self.has_tree[ids] # [count] + mask = no_tree.unsqueeze(1) # [count, 1] broadcasts over n_dt + next_token.copy_(torch.where(mask, self._no_tree_next_token, + next_token)) + next_sibling.masked_fill_(mask, -1) + return next_token, next_sibling + class SpecTreeManager: use_dynamic_tree: bool # Whether using dynamic tree @@ -284,6 +362,7 @@ def init_tree_info_for_dynamic_tree(self): num_slots=self.num_trees, n_dt=num_draft_with_root, mask_width=mask_width, + topK=self.dynamic_tree_max_topK, ) def scatter_to_slot_storage(self, ss, gen_slots, num_gens): diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 4f9c5af8846a..9b82ba52bf8d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -27,6 +27,8 @@ from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker from .model_drafter import ModelDrafter from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker +from .mtp_dynamic_tree import (MTPEagleDynamicTreeResourceManager, + MTPEagleDynamicTreeWorker) from .ngram import NGramDrafter, NGramPoolManager from .pard import PARDSpecMetadata, PARDWorker from .sa_worker import SASampler, SASpecMetadata, SAWorker @@ -112,6 +114,7 @@ def get_spec_metadata(spec_config, num_seq_slots=num_seq_slots, draft_vocab_size=draft_vocab_size, spec_resource_manager=spec_resource_manager, + use_dynamic_tree=getattr(spec_config, 'use_dynamic_tree', False), ) if spec_config.spec_dec_mode.is_mtp_vanilla(): return MTPSpecMetadata( @@ -273,6 +276,16 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): if sa_cfg is not None: sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, max_seq_len) + # Dynamic tree needs a SpecTreeManager (for the target's tree-mask verify + # forward) composed with the MTP hidden-state slot pools. + if getattr(spec_config, 'use_dynamic_tree', False): + return MTPEagleDynamicTreeResourceManager( + spec_config, + model_config.torch_dtype, + model_config.hidden_size, + max_num_requests, + sa_manager=sa_manager, + ) if spec_config.use_relaxed_acceptance_for_thinking or sa_manager is not None: # Unified resource manager: the unified worker reads # ``relaxed_delta_pool`` from ``Eagle3ResourceManager`` (mirrors the @@ -449,6 +462,11 @@ def get_spec_worker(spec_config, use_separate_draft_kv_cache, mapping=mapping) if spec_dec_mode.is_mtp_eagle_one_model(): + if getattr(spec_config, 'use_dynamic_tree', False): + return MTPEagleDynamicTreeWorker(spec_config, + model_config, + use_separate_draft_kv_cache, + mapping=mapping) return MTPEagleWorker(spec_config, model_config, use_separate_draft_kv_cache, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 57c3de25dcfe..43dc927f3c6b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2425,6 +2425,23 @@ class MTPDecodingConfig(DecodingBaseConfig): "When using EAGLE-style MTP, use faster one-model implementation (drafter as submodule) vs two-model." ) + use_dynamic_tree: bool = Field( + default=False, + description= + "Enable EAGLE-style dynamic-tree drafting for one-model MTP. When True, " + "each draft step expands dynamic_tree_max_topK candidates per node and the " + "tree is verified against the target, instead of a linear chain.") + dynamic_tree_max_topK: Optional[int] = Field( + default=None, + description= + "Top-K candidates expanded per node per draft layer when use_dynamic_tree " + "is enabled. Required when use_dynamic_tree is True.") + + # Backs the dynamic-tree worker's pre-allocated, batch-indexed CUDA buffers; + # MUST equal the global max_batch_size. Auto-populated by py_executor_creator + # (mirrors Eagle3DecodingConfig). PrivateAttr -- not a user-tunable knob. + _max_batch_size: Optional[int] = PrivateAttr(default=None) + sa_config: Optional[SAEnhancerConfig] = Field( default=None, status="beta", @@ -2471,13 +2488,37 @@ def _remap_deprecated_num_nextn_predict_layers(cls, data): def set_max_total_draft_tokens(self): # Leave max_draft_len as None ("use the model's num_nextn_predict_layers") # when the user doesn't set it; update_spec_config_from_model_config - # resolves it from the checkpoint before the model runs. When the user - # does set it, validate and mirror to max_total_draft_tokens (current MTP - # only supports a linear tree). + # resolves it from the checkpoint before the model runs. if self.max_draft_len is not None: if self.max_draft_len <= 0: raise ValueError("max_draft_len must be > 0 for MTP") - self.max_total_draft_tokens = self.max_draft_len + + # Dynamic-tree MTP: mirror EagleDecodingConfig. Honor an explicit + # max_total_draft_tokens within [max_draft_len, dynamic_tree_max_topK * max_draft_len]; + # otherwise default to dynamic_tree_max_topK * max_draft_len. + if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: + self.use_dynamic_tree = True + if self.max_draft_len is None: + raise ValueError( + "max_draft_len must be set when use_dynamic_tree is True") + if self.dynamic_tree_max_topK is None or self.dynamic_tree_max_topK <= 0: + raise ValueError( + "dynamic_tree_max_topK must be > 0 when use_dynamic_tree is True" + ) + default_max_total_draft_tokens = self.dynamic_tree_max_topK * self.max_draft_len + if self.max_total_draft_tokens is None: + self.max_total_draft_tokens = default_max_total_draft_tokens + elif self.max_total_draft_tokens < self.max_draft_len: + raise ValueError( + f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be >= " + f"max_draft_len ({self.max_draft_len})") + elif self.max_total_draft_tokens > default_max_total_draft_tokens: + raise ValueError( + f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be <= " + f"dynamic_tree_max_topK * max_draft_len ({default_max_total_draft_tokens})" + ) + elif self.max_draft_len is not None: + self.max_total_draft_tokens = self.max_draft_len # linear chain return self @model_validator(mode="after") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py new file mode 100644 index 000000000000..6566e13b47c3 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py @@ -0,0 +1,81 @@ +"""Unit tests for DynamicTreeSlotStorage no-tree link fallback. + +The Mamba tree-aware verify conv1d/SSU path reads retrieve_next_token / +retrieve_next_sibling unconditionally (it does not gate on has_tree). For +gen slots without a built tree (CUDA-graph/warmup dummies, and a real slot's +first decode), the gathered links are sentinels; apply_no_tree_linear_chain +must replace those rows with a valid degenerate linear chain so the kernel +indexing stays in-bounds. Real-tree rows must be left untouched. +""" + +import unittest + +import pytest +import torch + +from tensorrt_llm._torch.speculative.spec_tree_manager import DynamicTreeSlotStorage + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="DynamicTreeSlotStorage allocates CUDA buffers" +) +class TestNoTreeLinearChain(unittest.TestCase): + def _make(self, num_slots=4, n_dt=6): + # mask_width is unused by the link path; any positive value is fine. + return DynamicTreeSlotStorage(num_slots=num_slots, n_dt=n_dt, mask_width=1) + + def test_chain_template_is_valid(self): + n_dt = 6 + ss = self._make(n_dt=n_dt) + # next_token[i] = i+1, leaf has no child (-1). + expected = torch.tensor([1, 2, 3, 4, 5, -1], dtype=torch.int32, device="cuda") + self.assertTrue(torch.equal(ss._no_tree_next_token, expected)) + + def test_no_tree_rows_get_chain_real_rows_untouched(self): + num_slots, n_dt = 4, 6 + ss = self._make(num_slots=num_slots, n_dt=n_dt) + + # Slot 1 has a (synthetic) real tree; mark it valid and give it links + # distinct from the chain so we can detect any accidental overwrite. + real_next_token = torch.full((n_dt,), 3, dtype=torch.int32, device="cuda") + real_next_sibling = torch.full((n_dt,), 2, dtype=torch.int32, device="cuda") + ss.retrieve_next_token[1] = real_next_token + ss.retrieve_next_sibling[1] = real_next_sibling + ss.has_tree[1] = True + + # Gather slots [dummy, real, dummy] -> rows 0 and 2 are no-tree. + slot_ids = torch.tensor( + [ss.dummy_slot_id, 1, ss.dummy_slot_id], dtype=torch.long, device="cuda" + ) + count = 3 + next_token, next_sibling = ss.next_links_from_slots(slot_ids, count) + ss.apply_no_tree_linear_chain(next_token, next_sibling, slot_ids, count) + + chain = ss._no_tree_next_token + # No-tree rows -> linear chain, sibling all -1. + self.assertTrue(torch.equal(next_token[0], chain)) + self.assertTrue(torch.equal(next_token[2], chain)) + self.assertTrue( + torch.equal(next_sibling[0], torch.full((n_dt,), -1, dtype=torch.int32, device="cuda")) + ) + self.assertTrue( + torch.equal(next_sibling[2], torch.full((n_dt,), -1, dtype=torch.int32, device="cuda")) + ) + # Real-tree row -> original links preserved. + self.assertTrue(torch.equal(next_token[1], real_next_token)) + self.assertTrue(torch.equal(next_sibling[1], real_next_sibling)) + + def test_chain_links_are_in_bounds(self): + # Every link is either a valid token index in [0, n_dt) or the -1 stop + # sentinel; nothing points outside the per-request token range. + n_dt = 8 + ss = self._make(n_dt=n_dt) + slot_ids = torch.tensor([ss.dummy_slot_id], dtype=torch.long, device="cuda") + next_token, next_sibling = ss.next_links_from_slots(slot_ids, 1) + ss.apply_no_tree_linear_chain(next_token, next_sibling, slot_ids, 1) + valid = (next_token == -1) | ((next_token >= 0) & (next_token < n_dt)) + self.assertTrue(bool(valid.all().item())) + + +if __name__ == "__main__": + unittest.main() From 09b44c7364f3bd286be7f397a5e31e95cb81bec4 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 16 Jun 2026 00:23:00 -0700 Subject: [PATCH 02/24] [None][fix] Fix Nemotron MTP dynamic tree metadata Signed-off-by: qgai --- .../trtllmGenKernels/fmha/fmhaKernels.h | 6 +- .../_torch/models/modeling_nemotron_h.py | 2 + .../_torch/speculative/mtp_dynamic_tree.py | 105 +++++++++++++++++- 3 files changed, 106 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index d1681c39a56e..9f83bc7c8185 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/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.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; diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 5c2e60c4706a..3bb491fd1b91 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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 diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 708621080615..6a6989d376c1 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -212,6 +212,61 @@ def __init__( sm = get_sm_version() self._needs_mask_repack = sm < 100 or sm in (120, 121) + def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): + super()._prepare_attn_metadata_for_spec_dec(attn_metadata) + + batch_size = attn_metadata.num_seqs + if hasattr(attn_metadata, "kv_lens_cuda"): + # Keep kv_lens_cuda itself alive because TRTLLM attention holds a + # runtime view into it. + self._saved_kv_lens_cuda = attn_metadata.kv_lens_cuda[:batch_size].clone() + else: + self._saved_kv_lens_cuda = None + + # The draft loop overwrites the target verify tree mask/positions with + # draft-layer causal/growing-tree masks. Restore them before the next + # target forward. + if attn_metadata.spec_decoding_packed_mask is not None: + self._saved_packed_mask = attn_metadata.spec_decoding_packed_mask[:batch_size].clone() + else: + self._saved_packed_mask = None + if attn_metadata.spec_decoding_position_offsets is not None: + self._saved_position_offsets = attn_metadata.spec_decoding_position_offsets.clone() + self._saved_position_offsets_cpp = attn_metadata.spec_decoding_position_offsets_cpp + else: + self._saved_position_offsets = None + self._saved_position_offsets_cpp = None + if attn_metadata.spec_decoding_generation_lengths is not None: + self._saved_generation_lengths = attn_metadata.spec_decoding_generation_lengths[ + :batch_size + ].clone() + else: + self._saved_generation_lengths = None + + def _restore_attn_metadata_from_spec_dec(self, attn_metadata): + super()._restore_attn_metadata_from_spec_dec(attn_metadata) + + if self._saved_kv_lens_cuda is not None: + batch_size = self._saved_kv_lens_cuda.shape[0] + attn_metadata.kv_lens_cuda[:batch_size].copy_(self._saved_kv_lens_cuda) + self._saved_kv_lens_cuda = None + + if self._saved_packed_mask is not None: + batch_size = self._saved_packed_mask.shape[0] + attn_metadata.spec_decoding_packed_mask[:batch_size].copy_(self._saved_packed_mask) + self._saved_packed_mask = None + if self._saved_position_offsets is not None: + attn_metadata.spec_decoding_position_offsets.copy_(self._saved_position_offsets) + attn_metadata.spec_decoding_position_offsets_cpp = self._saved_position_offsets_cpp + self._saved_position_offsets = None + self._saved_position_offsets_cpp = None + if self._saved_generation_lengths is not None: + batch_size = self._saved_generation_lengths.shape[0] + attn_metadata.spec_decoding_generation_lengths[:batch_size].copy_( + self._saved_generation_lengths + ) + self._saved_generation_lengths = None + # ------------------------------------------------------------------ # # Helpers (mirroring eagle3 dynamic-tree worker, drafter-agnostic) # # ------------------------------------------------------------------ # @@ -220,6 +275,22 @@ def _apply_spec_metadata(self, attn_metadata, batch_size, query_len): attn_metadata.spec_decoding_generation_lengths[:batch_size] = query_len attn_metadata.update_position_offsets_for_cpp(query_len) + def _refresh_blackwell_tree_mask_metadata(self, attn_metadata): + if not getattr(attn_metadata, "use_spec_decoding", False): + return + if not getattr(attn_metadata, "is_spec_dec_dynamic_tree", False): + return + + first_sparse = getattr(attn_metadata, "spec_bl_tree_first_sparse_mask_offset_kv", None) + bl_tree_mask = getattr(attn_metadata, "spec_decoding_bl_tree_mask", None) + if first_sparse is None and bl_tree_mask is None: + return + + if bl_tree_mask is not None: + bl_tree_mask.zero_() + if first_sparse is not None: + attn_metadata.update_blackwell_first_sparse_mask_offset() + def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): """Compact the padded [n_req, buf_dim, ceil(buf_dim/32)] mask into the flat prefix XQA expects when n_tok < buf_dim. See the eagle reference @@ -539,7 +610,6 @@ def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_ tree_valid=tree_valid, ) ) - accepted_draft_count = accept_token_num[:num_gens] num_accepted_tokens[num_contexts:batch_size] = (accepted_draft_count + 1).to( torch.int32 @@ -714,6 +784,7 @@ def forward( ) # Save attn/spec metadata before the draft loop mutates it. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens self._prepare_attn_metadata_for_spec_dec(attn_metadata) # (c) Run the MTP draft tree loop -> build + store the next tree. @@ -735,6 +806,7 @@ def forward( # Restore attn metadata to support cuda graph. self._restore_attn_metadata_from_spec_dec(attn_metadata) + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens attn_metadata.use_spec_decoding = True # (d) Prepare next_new_tokens for overlap scheduler. @@ -942,6 +1014,12 @@ def _forward_draft_loop( attn_metadata.use_spec_decoding = num_gens > 0 if num_gens > 0 and hasattr(attn_metadata, "kv_lens_cuda"): attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= self._kv_correction + self._refresh_blackwell_tree_mask_metadata(attn_metadata) + if spec_metadata.all_rank_num_tokens is not None: + # Step-0 draft repacks gen requests to max_path_len tokens. Attention + # reads rank token counts from attn_metadata, while MoE also gets + # the same counts via the explicit all_rank_num_tokens argument. + attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_tokens with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states = draft_model.mtp_layers[0]( @@ -991,6 +1069,16 @@ def _forward_draft_loop( for layer_idx in range(1, self.max_draft_len): num_tokens_per_req = layer_idx * self.K num_infer_tokens = batch_size * num_tokens_per_req + subseq_all_rank_num_tokens = None + if spec_metadata.all_rank_num_seqs is not None: + # Subsequent dynamic-tree draft forwards process the full + # growing tree context, not one token per request. Attention + # DP/MoE communication sizes must therefore scale with the + # current per-request tree width. + subseq_all_rank_num_tokens = [ + n * num_tokens_per_req for n in spec_metadata.all_rank_num_seqs + ] + attn_metadata.all_rank_num_tokens = subseq_all_rank_num_tokens inp_hs = self._accumulated_hs[:batch_size, :num_tokens_per_req, :].reshape( num_infer_tokens, -1 @@ -1003,10 +1091,10 @@ def _forward_draft_loop( "hidden_states": inp_hs, "attn_metadata": attn_metadata, } - hidden_states = draft_model.mtp_layers[0]( embed_tokens=draft_model.embed_tokens, - all_rank_num_tokens=spec_metadata.subseq_all_rank_num_tokens, + all_rank_num_tokens=subseq_all_rank_num_tokens + or spec_metadata.subseq_all_rank_num_tokens, **layer_inputs, ) @@ -1092,15 +1180,19 @@ def _prepare_draft_layer_metadata( attn_metadata.num_contexts = 0 if hasattr(attn_metadata, "kv_lens_cuda"): - # KV rewind: remove unaccepted draft-path tokens, then add the K - # new depth-0 tokens. + # Match linear MTPEagleWorker's first-step cache-len semantics: + # generation rows only rewind unaccepted verify tokens here. The + # K depth-0 draft tokens are added after their forward writes KV; + # otherwise attention can read unwritten draft KV slots. if num_gens > 0: attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( self._max_path_len ) - num_accepted_tokens[num_contexts:batch_size] - attn_metadata.kv_lens_cuda[:batch_size] += self.K + if num_contexts > 0: + attn_metadata.kv_lens_cuda[:num_contexts] += self.K attn_metadata.use_spec_decoding = True attn_metadata.update_for_spec_dec() + self._refresh_blackwell_tree_mask_metadata(attn_metadata) else: num_tokens_previous_layer = cur_draft_idx * self.K num_tokens_current_layer = self.K * (cur_draft_idx + 1) @@ -1114,6 +1206,7 @@ def _prepare_draft_layer_metadata( if hasattr(attn_metadata, "kv_lens_cuda"): attn_metadata.kv_lens_cuda[:batch_size] += self.K attn_metadata.update_for_spec_dec() + self._refresh_blackwell_tree_mask_metadata(attn_metadata) class MTPEagleDynamicTreeResourceManager(BaseResourceManager): From 522f1153fd76ca0b0ceb4fb801ddb3104a616ac5 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 16 Jun 2026 03:54:31 -0700 Subject: [PATCH 03/24] [None][fix] Fix MTP dynamic tree slot reuse Signed-off-by: qgai --- .../_torch/speculative/mtp_dynamic_tree.py | 24 +++++++++++++++---- .../_torch/speculative/spec_tree_manager.py | 10 ++++---- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 6a6989d376c1..1e085a6f18b4 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -610,15 +610,29 @@ def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_ tree_valid=tree_valid, ) ) - accepted_draft_count = accept_token_num[:num_gens] + tree_valid_i = tree_valid[:num_gens] + accepted_draft_count = torch.where( + tree_valid_i, + accept_token_num[:num_gens], + torch.zeros_like(accept_token_num[:num_gens]), + ) num_accepted_tokens[num_contexts:batch_size] = (accepted_draft_count + 1).to( torch.int32 ) - accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) + + gen_accepted_tokens = accept_token[:num_gens].to(torch.int32) + bootstrap_accepted_tokens = torch.zeros_like(gen_accepted_tokens) + bootstrap_accepted_tokens[:, 0] = target_predict[:, 0] + accepted_tokens[num_contexts:batch_size] = torch.where( + tree_valid_i.unsqueeze(1), gen_accepted_tokens, bootstrap_accepted_tokens + ) # accept_index stores root at slot 0; subtract 1 so root/padding 0 # becomes the sentinel -1 (tree node index into the draft tokens). - self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( - accept_index[:num_gens, 1:max_path_len] - 1 + gen_accepted_indices = (accept_index[:num_gens, 1:max_path_len] - 1).to(torch.int32) + self._accepted_draft_indices_tensor[num_contexts:batch_size] = torch.where( + tree_valid_i.unsqueeze(1), + gen_accepted_indices, + torch.full_like(gen_accepted_indices, -1), ).to(torch.int32) num_accepted_tokens = self._apply_force_accepted_tokens( @@ -641,10 +655,10 @@ def _accepted_leaf_intermediate_positions(self, num_accepted_tokens, num_context accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens].to(torch.int64) # Column of the deepest accepted draft node, clamped to >=0 for the # golden-only case (its value is ignored via the mask below). - last_col = (accepted - 2).clamp_(min=0) draft_idx = self._accepted_draft_indices_tensor[num_contexts : num_contexts + num_gens].to( torch.int64 ) + last_col = (accepted - 2).clamp_(min=0, max=draft_idx.shape[1] - 1) leaf = torch.gather(draft_idx, 1, last_col.unsqueeze(1)).squeeze(1) + 1 # Golden-only requests (num_accepted == 1) take the root at position 0. return torch.where(accepted > 1, leaf, torch.zeros_like(leaf)) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 6c5a5fb93bfb..62f06b7a6761 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -43,6 +43,8 @@ def __init__(self, _w = 2**torch.arange(32, dtype=torch.int64, device='cuda') self.packed_mask = (_causal.to(torch.int64) * _w).sum(-1).to( torch.int32).unsqueeze(0).repeat(S, 1, 1).contiguous() + self._no_tree_position_offsets = self.position_offsets[0].clone() + self._no_tree_packed_mask = self.packed_mask[0].clone() # Override ONLY the reserved CUDA-graph/warmup dummy slot with a # bounded-depth K-ary tree (parent[i] = (i-1)//topK). Unlike a real @@ -136,12 +138,12 @@ def mark_valid(self, slot_ids, count): self.has_tree.narrow(0, self.dummy_slot_id, 1).fill_(False) def mark_invalid(self, slot_id): - """Clear validity and reset slot data.""" + """Clear validity and restore valid no-tree metadata.""" self.has_tree[slot_id] = False - self.packed_mask[slot_id] = 0 - self.position_offsets[slot_id] = 0 + self.packed_mask[slot_id] = self._no_tree_packed_mask + self.position_offsets[slot_id] = self._no_tree_position_offsets self.retrieve_index[slot_id] = 0 - self.retrieve_next_token[slot_id] = -1 + self.retrieve_next_token[slot_id] = self._no_tree_next_token self.retrieve_next_sibling[slot_id] = -1 def pack_retrieve_from_slots(self, slot_ids, count): From 8b3dee810988b3a5f1a0aad3a9c829000062c8c7 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 16 Jun 2026 05:34:21 -0700 Subject: [PATCH 04/24] [None][fix] Align MTP dynamic tree draft length Signed-off-by: qgai --- tensorrt_llm/_torch/attention_backend/trtllm.py | 8 ++++---- tensorrt_llm/_torch/speculative/utils.py | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index df13bdf357a2..40d15970398b 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1082,10 +1082,10 @@ def update_spec_dec_param( pos_src = torch.index_select(slot_storage.position_offsets, 0, slot_ids)[:, :n_dt] - compact_total = num_gens * n_dt - compact_offsets = self.spec_decoding_position_offsets[: - compact_total] - pos_dst = compact_offsets.view(num_gens, n_dt) + pos_dst = self.spec_decoding_position_offsets[:num_gens * + n_dt].view( + num_gens, + n_dt) pos_dst.copy_(pos_src, non_blocking=True) actual_mask_width = math.ceil(n_dt / 32) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 9b82ba52bf8d..71f890a24e85 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -374,7 +374,10 @@ def get_spec_decoder( # MTP Eagle one-model now uses the same sampler as Eagle3 one-model. return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) if spec_config.spec_dec_mode.is_mtp_vanilla(): - return MTPSampler(sampler_args, nextn=spec_config.max_draft_len) + nextn = spec_config.max_draft_len + if getattr(spec_config, "use_dynamic_tree", False): + nextn = spec_config.max_total_draft_tokens + return MTPSampler(sampler_args, nextn=nextn) if spec_config.spec_dec_mode.is_eagle3( ) or spec_config.spec_dec_mode.is_mtp_eagle(): # TorchSampler handles Eagle3 gracefully, by integrating d2t into the sampling process From d2928903eada2abeb8e7301ebeb0cc8f9c8e5f49 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 16 Jun 2026 21:16:34 -0700 Subject: [PATCH 05/24] [None][fix] Rebuild dynamic tree mask for MTP draft loop Signed-off-by: qgai --- cpp/tensorrt_llm/common/attentionOp.cpp | 1 + cpp/tensorrt_llm/common/attentionOp.h | 1 + .../kernels/decoderMaskedMultiheadAttention/xqaParams.h | 1 + cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h | 2 +- .../kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h | 1 + cpp/tensorrt_llm/kernels/xqaDispatcher.cpp | 1 + cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 3 ++- cpp/tensorrt_llm/thop/attentionOp.cpp | 4 +++- cpp/tensorrt_llm/thop/attentionOp.h | 3 ++- tensorrt_llm/_torch/attention_backend/fmha/fallback.py | 1 + tensorrt_llm/_torch/attention_backend/trtllm.py | 1 + tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py | 3 +++ 12 files changed, 18 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index b91c0ef98df6..3bec69f23ba2 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/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()) diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index f7337c9c9cb2..d716bc5d5a81 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/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; diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h index e421be0a6bd7..dc7794752e47 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h +++ b/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; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 9f83bc7c8185..8bcc29e0dbdb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -497,7 +497,7 @@ class TllmGenFmhaKernel // Prepare custom mask for spec-decoding generation kernels if needed. bool const prepareSpecDecTreeMask = params.mIsSpecDecTree - && (params.mLayerIdx == 0 + && (params.mForcePrepareSpecDecTreeMask || params.mLayerIdx == 0 || (params.mSpecDecodingTargetMaxGenLen > 0 && params.mMaxSeqLenQ != params.mSpecDecodingTargetMaxGenLen)); if (prepareSpecDecTreeMask) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h index 83b816c346f5..9252650ac67b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/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) diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp index 35fd02e7f127..8a5eeee91c6b 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/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; diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 23bd36d67b82..dd1ff0db5410 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/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::arg("force_prepare_spec_dec_tree_mask") = false, "Multi-head attention operation", + nb::call_guard()); m.def( "get_helix_workspace_size_per_rank", diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 848554c512d6..0bf5bc001206 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1084,7 +1084,8 @@ void attention(torch::Tensor q, std::optional k, std::optional compressed_kv_cache_pool_ptr, bool const is_cross, std::optional cross_kv, std::optional relative_attention_bias, int64_t relative_attention_max_distance, std::optional spec_decoding_target_max_draft_tokens, std::optional quant_scale_qkv, - std::optional dsv4_inv_rope_cos_sin_cache, bool enable_dsv4_epilogue_fusion) + std::optional 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 k, std::optionalmSpecDecodingTargetMaxGenLen = static_cast(spec_decoding_target_max_draft_tokens.value()) + 1; } + op->mForcePrepareSpecDecTreeMask = force_prepare_spec_dec_tree_mask; op->mUseSparseAttention = false; op->mUseTllmGenSparseAttentionPaged = false; diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index b4f0e90bc9a7..f28209166e0d 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -95,7 +95,8 @@ void attention(torch::Tensor q, std::optional k, std::optional relative_attention_bias = std::nullopt, int64_t relative_attention_max_distance = 0, std::optional spec_decoding_target_max_draft_tokens = std::nullopt, std::optional quant_scale_qkv = std::nullopt, - std::optional dsv4_inv_rope_cos_sin_cache = std::nullopt, bool enable_dsv4_epilogue_fusion = false); + std::optional 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 { diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 8f1da5dc36f9..9c081bf19dfc 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/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, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 40d15970398b..a3f4e89d8d85 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/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 diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 1e085a6f18b4..9838cc199a96 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -799,7 +799,9 @@ def forward( # Save attn/spec metadata before the draft loop mutates it. original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + original_force_prepare_spec_dec_tree_mask = attn_metadata.force_prepare_spec_dec_tree_mask self._prepare_attn_metadata_for_spec_dec(attn_metadata) + attn_metadata.force_prepare_spec_dec_tree_mask = True # (c) Run the MTP draft tree loop -> build + store the next tree. draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) @@ -821,6 +823,7 @@ def forward( # Restore attn metadata to support cuda graph. self._restore_attn_metadata_from_spec_dec(attn_metadata) attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + attn_metadata.force_prepare_spec_dec_tree_mask = original_force_prepare_spec_dec_tree_mask attn_metadata.use_spec_decoding = True # (d) Prepare next_new_tokens for overlap scheduler. From 214b1e36d548faa219e3b41f8e7af2e978100f8c Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 24 Jun 2026 03:51:29 -0700 Subject: [PATCH 06/24] [None][chore] Simplify dynamic tree comments Signed-off-by: qgai --- .../_torch/modules/mamba/mamba2_mixer.py | 30 +- .../_torch/pyexecutor/mamba_cache_manager.py | 11 +- .../_torch/pyexecutor/model_engine.py | 10 - tensorrt_llm/_torch/speculative/mtp.py | 34 +-- .../_torch/speculative/mtp_dynamic_tree.py | 262 +++--------------- .../_torch/speculative/spec_tree_manager.py | 50 +--- tensorrt_llm/_torch/speculative/utils.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 12 +- .../test_dynamic_tree_slot_storage.py | 16 +- 9 files changed, 68 insertions(+), 360 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 99e7f053878e..688c48b147de 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -430,8 +430,7 @@ def forward( f"{draft_token_num} must match fixed replay step " f"width {replay_step_width}.") - # Dynamic-tree verify: per-request tree links drive tree-aware - # conv/SSM recurrence. Off for linear MTP (path unchanged). + # 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 @@ -450,25 +449,14 @@ def forward( "tensors on spec_metadata.") retrieve_next_token = retrieve_next_token[:num_decodes] retrieve_next_sibling = retrieve_next_sibling[:num_decodes] - # conv1d fuses parent derivation: it reads next/sibling and - # writes the parent map into this buffer, which the SSM - # update then consumes to restore each token's parent state. + # 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's persistent arange tensor when - # available (allocated once at __init__, lives as long as the - # cache manager). The functools-cached fallback's storage can - # be co-allocated into the CUDA-graph private memory pool - # during the first warmup pass; the second warmup pass / capture - # then reads garbage from that recycled memory because no live - # tensor pins those exact bytes during the inter-pass - # allocator reset. The kv_cache_manager-owned tensor is - # outside the graph pool so its bytes stay valid. Linear - # MTP/non-dynamic-tree paths take the dense conv1d branch - # below and never reach here, so the linear path is unchanged. + # 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: @@ -478,9 +466,7 @@ def forward( attn_metadata.kv_cache_manager.get_max_resource_count(), state_indices_d.device)[:num_decodes] - # Reshape for batch processing. reshape (not view) because tree - # tokens may be non-contiguous; for linear (contiguous) tokens - # reshape is exactly view, so the linear path is unaffected. + # 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) @@ -494,7 +480,7 @@ def conv1d(): conv_state_indices=state_indices_d[:num_decodes], intermediate_conv_window=intermediate_conv_states, intermediate_state_indices=intermediate_state_indices, - # Tree links (None for linear MTP -> dense conv path). + # None on linear MTP. retrieve_next_token=retrieve_next_token, retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, @@ -646,9 +632,7 @@ def convert_dt(): B_d_4d = B_d_4d.contiguous() C_d_4d = C_d_4d.contiguous() if is_dyn_tree: - # flashinfer's selective_state_update has no tree-parent - # restore; use the native Triton kernel (which does) for - # dynamic tree. Linear MTP keeps the configured func. + # flashinfer SSU cannot restore tree-parent states. ssu_func = selective_state_update_native ssu_extra = dict( retrieve_parent_token=retrieve_parent_token) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 48a7ab7eeeda..0274ac32ceae 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -795,10 +795,7 @@ def update_mamba_states( num_gens = batch_size - num_contexts num_accepted_draft_tokens = num_accepted_tokens[ num_contexts:num_contexts + num_gens] - 1 - # Linear MTP: accepted path is a contiguous chain (leaf at depth - # num_accepted-1). Dynamic tree: states are recorded in tree order, so - # the leaf's buffer position is its tree-node index (supplied by the - # worker). None -> linear (unchanged). + # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. accepted_positions = (accepted_leaf_positions if accepted_leaf_positions is not None else num_accepted_draft_tokens) state_indices_d = state_indices[num_contexts:num_contexts + num_gens] @@ -1781,11 +1778,7 @@ def update_mamba_states( num_accepted_draft_tokens = ( num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( torch.int32) - # Intermediate-buffer position of each request's accepted leaf state. - # Linear MTP: the accepted path is a contiguous chain, so the leaf sits - # at depth (num_accepted - 1). Dynamic tree: the 31 verified tokens are - # laid out in tree order, so the leaf's buffer position is its tree-node - # index (root/golden at 0), supplied by the worker. + # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. accepted_positions = (accepted_leaf_positions.to(torch.int32) if accepted_leaf_positions is not None else num_accepted_draft_tokens) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8d6c5a53ffbc..278a7bf4a660 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1580,16 +1580,6 @@ def _capture_generation_cuda_graphs(self, if self.mapping is not None and self.mapping.has_cp_helix(): effective_max_seq_len = self.max_seq_len // self.mapping.cp_size - # Blackwell trtllm-gen custom-mask spec-dec dynamic-tree generation - # FMHA reads out of bounds during CUDA-graph capture when the capture - # dummy uses the model's full context length. The benchmark uses a - # max-token budget, so capture a graph within that budget. - if (self.spec_config is not None - and self.spec_config.spec_dec_mode.use_one_engine() - and getattr(self.spec_config, 'use_dynamic_tree', False)): - effective_max_seq_len = min(effective_max_seq_len, - self.max_num_tokens) - sparse_config = self.sparse_attention_config if (isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) and sparse_config.needs_separate_short_long_cuda_graphs()): diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 69ab6a00b7a4..b596b3c19ae3 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -140,17 +140,12 @@ class MTPSpecMetadata(SpecMetadata): # CUDA graph, we use this tensor to store the number of input tokens for the # subsequent draft forward. subseq_all_rank_num_tokens: Optional[List[int]] = None - # Dynamic-tree drafting (one-model MTP-Eagle). When use_dynamic_tree is True, - # the target's multi-token verify forward uses a per-slot tree mask sourced - # from spec_tree_manager.slot_storage (see model_engine.update_spec_dec_param). + # Dynamic-tree MTP-Eagle uses per-slot masks during target verify. use_dynamic_tree: bool = False dynamic_tree_max_topK: Optional[int] = None spec_tree_manager: Optional[object] = None - # Dynamic-tree per-gen-request tree links for the Mamba tree-aware verify - # forward, gathered from spec_tree_manager.slot_storage in prepare(). - # Shapes [num_generations, max_total_draft_tokens + 1], int32, root at - # index 0. None for linear MTP. retrieve_parent_token stays None: the - # mixer's conv1d kernel derives + fills it from next_token/next_sibling. + # Per-generation tree links for Mamba verify. Parent links are derived in + # conv1d. Shape: [num_generations, max_total_draft_tokens + 1]. retrieve_next_token: Optional[torch.Tensor] = None retrieve_next_sibling: Optional[torch.Tensor] = None retrieve_parent_token: Optional[torch.Tensor] = None @@ -183,9 +178,7 @@ def __post_init__(self) -> None: self.mtp_num_modules, device='cuda', ) - # Dynamic-tree drafting drives the target's multi-token verify forward - # through the spec-dec tree mask (see model_engine.update_spec_dec_param, - # which routes on these flags + spec_tree_manager.use_dynamic_tree). + # Enable target-side tree-mask routing for dynamic-tree MTP. if self.use_dynamic_tree: self.is_spec_dec_tree = True self.is_spec_dec_dynamic_tree = True @@ -215,11 +208,7 @@ def prepare(self): # forward and only one input token in the following draft forward. # This num_tokens is used to set the all_rank_num_tokens for attention dp. if self.use_dynamic_tree: - # Target verify forward processes (max_total_draft_tokens + 1) tokens - # per gen request, but the draft step-0 forward (which consumes - # all_rank_num_tokens) processes only (max_draft_len + 1) per request - # after prepare_drafter_inputs repacks the accepted path. Correct - # num_tokens to the draft step-0 count for attention dp. + # Step-0 draft uses max_draft_len + 1 tokens per generation. self.num_tokens -= self.num_generations * ( self.max_total_draft_tokens - self.max_draft_len) elif not self.spec_dec_mode.is_mtp_eagle_one_model(): @@ -266,12 +255,7 @@ def prepare(self): if gen_request_ids: sa_manager.prepare(gen_request_ids, self.runtime_draft_len) - # Dynamic tree: gather per-gen-request tree links for the Mamba - # tree-aware verify forward (the tree being verified this step was - # built last step into slot_storage). all_ids_buf is laid out - # [ctx | gen] and was filled by model_engine.fill_all_slot_ids before - # this prepare() runs (see _prepare_inputs). Token order matches the - # target forward's verify layout (root/golden at index 0). + # Gather per-generation tree links for Mamba verify. self.retrieve_next_token = None self.retrieve_next_sibling = None self.retrieve_parent_token = None @@ -284,11 +268,7 @@ def prepare(self): num_contexts:num_contexts + num_gens] next_token, next_sibling = slot_storage.next_links_from_slots( gen_slot_ids, num_gens) - # No-tree gen slots (CUDA-graph/warmup dummies, and a real - # slot's first decode before a tree exists) have sentinel - # links. The Mamba tree-aware verify conv1d/SSU reads these - # unconditionally, so substitute a valid linear chain for those - # rows; real-tree rows are untouched. + # No-tree rows need valid links; real-tree rows are unchanged. slot_storage.apply_no_tree_linear_chain(next_token, next_sibling, gen_slot_ids, num_gens) diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 9838cc199a96..ff945f0bdbcb 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -12,22 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""MTP-Eagle one-model dynamic tree speculative decoding (greedy only). - -This worker subclasses :class:`MTPEagleWorker` and reuses the drafter-agnostic -dynamic-tree bookkeeping/verification helpers (``DynamicTreeOpsConverter`` and -``SpecTreeManager``) that were originally written for -``Eagle3OneModelDynamicTreeWorker``. The only piece that is drafter-specific is -the draft loop: instead of running the eagle drafter -(``draft_model.model(...)`` with ``apply_eagle3_fc`` + hidden-state capture), -this worker runs ``draft_model.mtp_layers[0]`` repeatedly (the linear MTP-Eagle -drafter) but grows a topK tree at each layer instead of a linear chain. - -Scope (intentional simplifications): -- GREEDY only (temperature 0). The rejection-sampling path is not implemented. -- The new code path is only reached when ``spec_config.use_dynamic_tree`` is - True; the linear ``MTPEagleWorker`` path is unchanged. -""" +"""MTP-Eagle one-model dynamic tree speculative decoding (greedy only).""" import math from typing import TYPE_CHECKING, List, Optional @@ -46,9 +31,7 @@ from ..pyexecutor.scheduler import ScheduledRequests from .eagle3 import MTPEagleWorker -# Reuse the drafter-agnostic fused helpers from the eagle3 dynamic-tree worker. -# These operate purely on token/score/mask tensors and do not touch the eagle -# drafter, so they are safe to share with the MTP drafter. +# Reuse drafter-agnostic dynamic-tree helpers. from .eagle3_dynamic_tree import ( _build_mask_and_position, _gather_repack_step0_kernel, @@ -62,13 +45,7 @@ class MTPEagleDynamicTreeWorker(MTPEagleWorker): - """MTP-Eagle one-model worker with dynamic tree drafting + verification. - - Inherits the linear MTP-Eagle drafting/sampling primitives from - :class:`MTPEagleWorker` (``draft_sampler``, ``prepare_drafter_inputs``, - ``update_mtp_hidden_states``, ``_prepare_next_new_tokens``, ...) and adds the - dynamic-tree draft loop, greedy verification, and tree construction. - """ + """MTP-Eagle worker with dynamic-tree draft and greedy verify.""" def __init__( self, @@ -91,9 +68,7 @@ def __init__( self.K = spec_config.dynamic_tree_max_topK self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 self.tokens_per_gen_step = spec_config.tokens_per_gen_step - # _max_batch_size is auto-populated by py_executor_creator from the - # global max_batch_size (mirrors EagleDecodingConfig). It must be set by - # the time we get here. + # Set by py_executor_creator from the global max_batch_size. assert spec_config._max_batch_size is not None, ( "MTPDecodingConfig._max_batch_size was not populated; " "py_executor_creator should have set it from the global max_batch_size." @@ -109,8 +84,7 @@ def __init__( self.spec_tree_manager = None self._d2t = None - # === Pre-allocated draft-loop buffers (CUDA-graph safe) === - # Mirror Eagle3OneModelDynamicTreeWorker.__init__ buffer strategy. + # Pre-allocated draft-loop buffers (CUDA-graph safe). self.draft_tokens_buffer = torch.zeros( max_batch_size, loop_max_tokens, dtype=torch.int32, device="cuda" ) @@ -143,13 +117,7 @@ def __init__( ) self._max_path_len = max_draft_len + 1 - # Step-0 spec-dec reset buffers (mirror Eagle3OneModelDynamicTreeWorker): - # the target verify forward leaves a tokens_per_gen_step-wide tree mask, - # tree position offsets, and a kv_lens inflated by tokens_per_gen_step. - # The step-0 draft attends only to the max_path_len accepted-path tokens, - # so it resets to an 8-wide causal mask + causal offsets and rewinds - # kv_lens by (tokens_per_gen_step - max_path_len). Linear MTP needs none - # of this because there tokens_per_gen_step == max_path_len. + # Step-0 draft resets verify-time tree metadata to accepted-path width. self._kv_correction = self.tokens_per_gen_step - self._max_path_len self._step0_causal_mask = torch.tensor( [(1 << (t + 1)) - 1 for t in range(self._max_path_len)], @@ -176,9 +144,7 @@ def __init__( self._candidates_buf = torch.zeros(max_batch_size, N, dtype=torch.int32, device="cuda") self._target_predict_buf = torch.zeros(max_batch_size, N, dtype=torch.int32, device="cuda") - # === Hidden-state management for the growing-context draft loop === - # These mirror the eagle reference's _hs_write_buffer / _accumulated_hs - # but store the MTP layer's output hidden states (no eagle capture). + # Hidden states for the growing-context draft loop. self._hs_write_buffer = None self._accumulated_hs = None self._hs_read_map = torch.zeros( @@ -187,13 +153,7 @@ def __init__( self._step0_hs = None self._hs_dim = None - # === Step-0 repack scratch (graph-safe; mirrors eagle3 reference) === - # The target verify forward lays gen tokens out as - # tokens_per_gen_step rows per request, but the MTP draft step-0 only - # attends to the accepted-path (max_path_len) tokens. These buffers hold - # the repacked gen-only [num_gens * max_path_len] inputs; rows are sized - # by the static max (max_batch_size * tokens_per_gen_step) and the hidden - # buffer is lazily sized once the hidden dim is known. + # Step-0 repack scratch for accepted-path inputs. max_total_tokens = max_batch_size * self.tokens_per_gen_step self._step0_input_ids_buf = torch.zeros(max_total_tokens, dtype=torch.int32, device="cuda") self._step0_position_ids_buf = torch.zeros( @@ -223,9 +183,7 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): else: self._saved_kv_lens_cuda = None - # The draft loop overwrites the target verify tree mask/positions with - # draft-layer causal/growing-tree masks. Restore them before the next - # target forward. + # Restore verify metadata after the draft loop mutates it. if attn_metadata.spec_decoding_packed_mask is not None: self._saved_packed_mask = attn_metadata.spec_decoding_packed_mask[:batch_size].clone() else: @@ -268,7 +226,7 @@ def _restore_attn_metadata_from_spec_dec(self, attn_metadata): self._saved_generation_lengths = None # ------------------------------------------------------------------ # - # Helpers (mirroring eagle3 dynamic-tree worker, drafter-agnostic) # + # Helpers # # ------------------------------------------------------------------ # def _apply_spec_metadata(self, attn_metadata, batch_size, query_len): """Set spec-dec gen lengths and refresh the C++ position-offset view.""" @@ -292,9 +250,7 @@ def _refresh_blackwell_tree_mask_metadata(self, attn_metadata): attn_metadata.update_blackwell_first_sparse_mask_offset() def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): - """Compact the padded [n_req, buf_dim, ceil(buf_dim/32)] mask into the - flat prefix XQA expects when n_tok < buf_dim. See the eagle reference - for the detailed rationale.""" + """Compact padded masks into the flat prefix XQA expects.""" buf_dim = mask_buf.shape[1] if n_tok >= buf_dim or n_req <= 1: return @@ -339,25 +295,7 @@ def _ensure_spec_tree_manager(self, resource_manager): def sample( self, logits: torch.Tensor, max_top_k: int, draft_model=None ) -> tuple[torch.Tensor, torch.Tensor]: - """TopK sampling with softmax for the dynamic tree (greedy=topK). - - Returns (topk_indices [.., K], topk_values [.., K]). MTP shares the - target vocabulary through ``shared_head``/``lm_head``, so unlike EAGLE3 - there is no draft->target (d2t) token remap. - - TP correctness: ``DeepseekV3MTPHead`` forces ``gather_output=False`` on - the column-parallel lm_head in pure TP, so ``logits`` here is a per-rank - vocab shard ``[.., vocab/tp]``. Sampling top-K on a shard yields - DIFFERENT draft tokens/scores per rank, so the per-rank trees (and hence - ``num_accepted_tokens`` in the next verify) diverge across TP ranks and - the downstream attention/MoE collectives desync (hang at TP>1). The - linear ``MTPWorker.draft_sampler`` avoids this by all-gathering the - local argmax; the dynamic tree needs full top-K + probabilities, so we - all-gather the full sharded logits (stripping lm_head column padding) - and run softmax+top-K on the replicated full vocab, exactly matching the - TP=1 path. Gated on pure TP (tp_size>1, attention DP off) where the - lm_head output is actually sharded. - """ + """TopK sampling for dynamic tree; all-gather sharded TP logits.""" mapping = ( getattr(self.model_config, "mapping", None) if self.model_config is not None else None ) @@ -379,9 +317,7 @@ def update_draft_tokens_and_scores( batch_size, attn_metadata=None, ): - """Grow the tree: write tokens/scores to history buffers + masks. - - Identical bookkeeping to the eagle reference (drafter-agnostic).""" + """Grow the tree and update history buffers.""" if cur_draft_idx == 0: new_draft_scores = new_draft_scores.reshape(batch_size, self.K) new_draft_tokens_2d = new_draft_tokens.reshape(batch_size, self.K) @@ -440,9 +376,7 @@ def resampling_final_draft_tokens(self, batch_size: int): def prepare_tree_mask_and_position_offset( self, cur_draft_idx, attn_metadata, selected_parents=None ): - """Prepare mask + position offsets for the next draft layer. - - Drafter-agnostic; identical to the eagle reference.""" + """Prepare mask and position offsets for the next draft layer.""" if attn_metadata.spec_decoding_packed_mask is None: return spec_tree_manager = self.spec_tree_manager @@ -501,11 +435,7 @@ def update_hidden_states( hidden_states_to_save=None, selected_parents=None, ): - """Manage the growing-context hidden states for the MTP draft loop. - - Unlike the eagle reference (which saves eagle prenorm hidden states), - we save the MTP layer's OUTPUT hidden states. The gather/parent logic is - otherwise identical.""" + """Manage growing-context hidden states for the MTP draft loop.""" if cur_draft_idx == 0: hs_dim = step0_hs.shape[-1] self._hs_dim = hs_dim @@ -550,10 +480,7 @@ def update_hidden_states( # ------------------------------------------------------------------ # @nvtx_range("mtp_dyn.sample_and_accept_draft_tokens") def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_metadata): - """Greedy dynamic-tree verification of the PREVIOUS step's tree. - - Overrides MTPWorker.sample_and_accept_draft_tokens. Returns - (accepted_tokens [bs, max_path_len], num_accepted_tokens [bs]).""" + """Greedy verification of the previous dynamic tree.""" batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts @@ -582,8 +509,7 @@ def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_ target_predict = self._target_predict_buf[:num_gens] target_predict.copy_(target_tokens[num_contexts:].reshape(num_gens, N)) - # First-step bootstrap / CUDA-graph warmup: no tree exists yet, so - # accept only the golden token (the first of the gen tokens). + # No prior tree exists on bootstrap/warmup; accept the golden token. if spec_tree_manager is None: num_accepted_tokens[num_contexts:batch_size] = 1 accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0] @@ -626,8 +552,7 @@ def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_ accepted_tokens[num_contexts:batch_size] = torch.where( tree_valid_i.unsqueeze(1), gen_accepted_tokens, bootstrap_accepted_tokens ) - # accept_index stores root at slot 0; subtract 1 so root/padding 0 - # becomes the sentinel -1 (tree node index into the draft tokens). + # Convert root/padding index 0 to draft-node sentinel -1. gen_accepted_indices = (accept_index[:num_gens, 1:max_path_len] - 1).to(torch.int32) self._accepted_draft_indices_tensor[num_contexts:batch_size] = torch.where( tree_valid_i.unsqueeze(1), @@ -641,17 +566,7 @@ def sample_and_accept_draft_tokens(self, input_ids, logits, spec_metadata, attn_ return accepted_tokens, num_accepted_tokens def _accepted_leaf_intermediate_positions(self, num_accepted_tokens, num_contexts, num_gens): - """Tree-node position of each gen request's accepted leaf in the mamba - intermediate-state buffer. - - The mixer records states in verify-token order: position 0 is the - golden/root token, positions 1.. are the draft nodes in tree order. - ``_accepted_draft_indices_tensor[r, c]`` holds the (c+1)-th accepted - draft node's tree index (= candidate position - 1), so the deepest - accepted node lives at column ``num_accepted - 2`` and its buffer - position is that index + 1. When only the golden token is accepted - (num_accepted == 1) the leaf is the root at position 0. - """ + """Return each accepted leaf's position in the Mamba state buffer.""" accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens].to(torch.int64) # Column of the deepest accepted draft node, clamped to >=0 for the # golden-only case (its value is ignored via the mask below). @@ -665,30 +580,14 @@ def _accepted_leaf_intermediate_positions(self, num_accepted_tokens, num_context @nvtx_range("mtp_dyn._relocate_kv_eagerly") def _relocate_kv_eagerly(self, attn_metadata, batch_size): - """Move accepted draft tokens' KV from tree positions to the linear - prefix the next step expects. Mirrors the eagle reference. - - Mamba-2 hybrid handling: the parent KVCacheManager spans every global - layer, but Mamba layers carry recurrent state (``num_kv_heads == 0``) - and live in a *separate* C++ pool from the attention layers. The - ``update_kv_cache_draft_token_location_2d`` op addresses a single pool - with a uniform per-layer stride (``layerIdx * 2 * bytesPerBlock`` off - one base pointer, see ``updateKVBlockArrayDraftTokenLocation2D``), and - the stored block offsets are scaled by *that pool's* layer count - (``flat_index3(blockIdx, 0, fieldIdx, pool.numLayers, kvFactor)``). So - we must drive the op with the attention pool only: its compact layer - count, its uniform head count, and its slice of the pool-pointer and - block-offset tensors. For a pure-attention model there is exactly one - pool and this reduces to the original call.""" + """Move accepted draft KV from tree positions to the linear prefix.""" cache_mgr = getattr(attn_metadata, "kv_cache_manager", None) if cache_mgr is None or self._kv_head_dim_bytes is None: return if not hasattr(cache_mgr, "num_kv_heads_per_layer"): return - # Attention layers are those with KV heads (Mamba layers are zeroed). - # The set is static for a given model, so the layerCount / head count / - # pool index below are CUDA-graph-safe (no data-dependent control flow). + # Mamba layers have zero KV heads; relocate attention-layer KV only. kv_heads = cache_mgr.num_kv_heads_per_layer attn_heads = set(h for h in kv_heads if h > 0) assert len(attn_heads) == 1, ( @@ -699,10 +598,7 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): attn_layer_offsets = [i for i, h in enumerate(kv_heads) if h > 0] attn_num_layers = len(attn_layer_offsets) - # All attention layers share one pool (same head count => same pool in - # the C++ WindowBlockManager). Resolve its index from the layer->pool - # mapping so we slice the correct pool's pointers/offsets; the op reads - # pool_pointers[0]/[1] as that pool's primary/secondary base. + # Resolve the attention KV pool used by the relocation op. pool_mapping = getattr(cache_mgr, "kv_cache_pool_mapping", None) if pool_mapping is not None: attn_pool_indices = set(int(pool_mapping[off][0]) for off in attn_layer_offsets) @@ -749,16 +645,7 @@ def forward( draft_model, resource_manager=None, ): - """Dynamic-tree MTP-Eagle forward. - - Step-by-step (see module/report for divergence notes): - (a) verify the previous tree against target logits (greedy); - (b) update the Mamba hybrid cache with num_accepted_tokens, and - relocate accepted draft-token KV; - (c) run the MTP draft TREE loop and build the new tree into - spec_tree_manager.slot_storage; - (d) return the MTPEagleWorker output contract. - """ + """Run verify, cache promotion, and next-tree drafting.""" if resource_manager is not None: self._ensure_spec_tree_manager(resource_manager) @@ -776,12 +663,7 @@ def forward( if num_gens > 0: self._relocate_kv_eagerly(attn_metadata, batch_size) - # (b) Update Mamba hybrid cache for accepted variable-length paths. - # The mixer records one intermediate state per verified token in - # tree order (root/golden at buffer position 0), so the accepted - # leaf's state lives at its tree-node position, not at linear depth - # num_accepted-1. Compute that per-gen-request leaf position from the - # accepted draft-node tree indices and pass it to the cache rollback. + # Dynamic-tree Mamba states are stored by tree-node position. if self._is_mamba_hybrid_cache is None: self._is_mamba_hybrid_cache = isinstance( attn_metadata.kv_cache_manager, MambaHybridCacheManager @@ -857,36 +739,13 @@ def _prepare_step0_drafter_inputs( accepted_tokens, attn_metadata, ): - """Repack step-0 drafter inputs to the accepted-path layout. - - ``MTPEagleWorker.prepare_drafter_inputs`` only repacks the gen - ``input_ids`` to ``num_gens * max_path_len`` rows while leaving - ``hidden_states`` / ``position_ids`` in the target verify layout - (``num_gens * tokens_per_gen_step`` rows). For the linear MTP path the - two layouts coincide (``tokens_per_gen_step == max_path_len``); for the - dynamic tree they differ, and the NemotronHMTP ``eh_proj`` fusion - ``cat([enorm(embed(input_ids)), hnorm(hidden_states)], dim=-1)`` then - sees mismatched row counts and crashes. - - This mirrors ``Eagle3OneModelDynamicTreeWorker.prepare_1st_drafter_inputs``: - a single fused Triton kernel gathers the TARGET ``hidden_states`` at the - accepted tree positions and repacks ``input_ids`` / ``position_ids`` - into the ``[ctx | gen (num_gens * max_path_len)]`` layout, and writes the - per-gen-request "last accepted token" gather id into ``_gather_ids_buf``. - - Unlike eagle3 there is no ``spec_metadata.hidden_states`` / - ``apply_eagle3_fc`` / ``layers_to_capture`` indirection: MTP shares the - target vocab and consumes the target hidden states directly. - - Returns the drafter ``inputs`` dict. - """ + """Repack step-0 drafter inputs to accepted-path layout.""" num_contexts = attn_metadata.num_contexts batch_size = attn_metadata.num_seqs num_gens = batch_size - num_contexts num_ctx_tokens = attn_metadata.num_ctx_tokens - # Context input_ids: shift-left + place golden token at last positions - # (identical to MTPEagleWorker.prepare_drafter_inputs context path). + # Match MTPEagleWorker context input repack. input_ids_ctx = self._prepare_context_input_ids( input_ids, num_ctx_tokens, last_tokens_idx, accepted_tokens, num_contexts ) @@ -907,8 +766,7 @@ def _prepare_step0_drafter_inputs( device="cuda", ) - # accepted_tokens[num_contexts:] is the accepted path (incl golden - # at col 0); it is exactly the eagle reference's ``_accept_token``. + # Accepted path includes the golden token at column 0. accept_token = accepted_tokens[num_contexts:batch_size] BLOCK_H = triton.next_power_of_2(hidden_dim) @@ -976,26 +834,14 @@ def _forward_draft_loop( num_gens, batch_size, ): - """MTP dynamic-tree draft loop with growing context. - - Step 0 runs mtp_layers[0] over the accepted-path tokens (max_path_len - per request, repacked by _prepare_step0_drafter_inputs), then expands - topK. Each subsequent layer re-runs mtp_layers[0] over ALL accumulated - tree tokens and expands topK per surviving parent. Finally the tree is - resampled and built into slot_storage so the NEXT target forward uses - the tree mask. - """ + """Draft the next dynamic tree with growing context.""" spec_tree_manager = self.spec_tree_manager assert batch_size <= self._max_batch_size, ( f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}" ) - # --- Step 0: one MTP forward over accepted golden tokens --- - # Repack input_ids/position_ids AND hidden_states to the accepted-path - # layout (num_gens * max_path_len gen rows) so the NemotronHMTP eh_proj - # fusion sees matching row counts. The fused kernel also writes the - # per-request last-accepted-token gather id into _gather_ids_buf. + # Step 0: run MTP over accepted-path rows. position_ids, last_tokens_idx = self.prepare_position_ids_and_last_tokens( position_ids, attn_metadata.seq_lens_cuda ) @@ -1008,11 +854,7 @@ def _forward_draft_loop( attn_metadata=attn_metadata, ) - # Step-0 causal spec-dec reset (mirrors Eagle3OneModelDynamicTreeWorker). - # The target verify forward left a tokens_per_gen_step-wide tree mask + - # position offsets and an inflated kv_lens; the step-0 draft attends only - # to the max_path_len accepted-path tokens, so reset to an 8-wide causal - # mask + causal offsets and rewind kv_lens. None in prefill-only warmup. + # Reset verify-time tree metadata to accepted-path width. num_step0_tokens = self._max_path_len if attn_metadata.spec_decoding_generation_lengths is not None: total = num_gens * num_step0_tokens @@ -1033,9 +875,7 @@ def _forward_draft_loop( attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= self._kv_correction self._refresh_blackwell_tree_mask_metadata(attn_metadata) if spec_metadata.all_rank_num_tokens is not None: - # Step-0 draft repacks gen requests to max_path_len tokens. Attention - # reads rank token counts from attn_metadata, while MoE also gets - # the same counts via the explicit all_rank_num_tokens argument. + # Keep attention/MoE token counts aligned with step-0 repack. attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_tokens with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): @@ -1045,12 +885,7 @@ def _forward_draft_loop( **inputs, ) - # Gather the per-request "last accepted token" hidden state (the - # tree root for depth-0 expansion). The Triton kernel already wrote - # gen gather ids (num_ctx_tokens + gen_idx * max_path_len + - # num_accepted - 1) into _gather_ids_buf; prepend the context last- - # token ids. This indexes the post-repack [ctx | gen max_path_len] - # hidden_states layout. + # Gather each request's root hidden state for depth-0 expansion. self._gather_ids_buf[:num_contexts].copy_(last_tokens_idx[:num_contexts]) gather_ids = self._gather_ids_buf[:batch_size] @@ -1082,16 +917,13 @@ def _forward_draft_loop( inputs, ) - # --- Subsequent layers: grow the tree --- + # Subsequent layers grow the tree. for layer_idx in range(1, self.max_draft_len): num_tokens_per_req = layer_idx * self.K num_infer_tokens = batch_size * num_tokens_per_req subseq_all_rank_num_tokens = None if spec_metadata.all_rank_num_seqs is not None: - # Subsequent dynamic-tree draft forwards process the full - # growing tree context, not one token per request. Attention - # DP/MoE communication sizes must therefore scale with the - # current per-request tree width. + # Token counts scale with the current tree width. subseq_all_rank_num_tokens = [ n * num_tokens_per_req for n in spec_metadata.all_rank_num_seqs ] @@ -1177,11 +1009,7 @@ def _prepare_draft_layer_metadata( num_accepted_tokens=None, inputs=None, ): - """Set up attn_metadata seq_lens/kv_lens for the next drafter layer. - - Drafter-agnostic; mirrors the eagle reference's prepare_for_generation - (which is itself derived from MTPEagleWorker's i==0 / i>0 metadata - updates).""" + """Set attn_metadata seq_lens/kv_lens for the next draft layer.""" if cur_draft_idx == 0: base_pos = inputs["position_ids"][gather_ids] + 1 self.position_ids_buffer[:batch_size, : self.K] = base_pos.unsqueeze(1).expand( @@ -1197,10 +1025,7 @@ def _prepare_draft_layer_metadata( attn_metadata.num_contexts = 0 if hasattr(attn_metadata, "kv_lens_cuda"): - # Match linear MTPEagleWorker's first-step cache-len semantics: - # generation rows only rewind unaccepted verify tokens here. The - # K depth-0 draft tokens are added after their forward writes KV; - # otherwise attention can read unwritten draft KV slots. + # Rewind only unaccepted verify tokens; draft KV is added later. if num_gens > 0: attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( self._max_path_len @@ -1227,17 +1052,7 @@ def _prepare_draft_layer_metadata( class MTPEagleDynamicTreeResourceManager(BaseResourceManager): - """Resource manager for one-model MTP-Eagle dynamic tree mode. - - Composes: - - a ``SpecTreeManager`` (exposed as ``.spec_tree_manager``) so the model - engine / attention backend can wire the per-slot tree mask for the - target's multi-token verify forward, and the worker can build/store the - tree, and - - an ``MTPHiddenStatesManager`` so MTPEagleWorker's drafter-input - preparation (mtp_past_tokens / mtp_past_hidden_states slot pools and - slot_ids) keeps working. - """ + """Resource manager for MTP dynamic-tree mode.""" hidden_states: Optional[torch.Tensor] = None @@ -1295,8 +1110,7 @@ def free_resources(self, request: LlmRequest): self._mtp_hidden_states_manager.free_resources(request) def add_dummy_requests(self, request_ids: List[int]): - # Dynamic-tree dummies use slot_storage.dummy_slot_id (no per-request - # slot), but the MTP hidden-state pool still needs a slot per dummy. + # Dummies still need MTP hidden-state slots. self._mtp_hidden_states_manager.add_dummy_requests(request_ids) def shutdown(self): diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 62f06b7a6761..1a44a1bd064b 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -24,16 +24,8 @@ def __init__(self, S = num_slots + 1 self.dummy_slot_id = num_slots - # Slot buffers — C++ kernel writes directly via slotIds. - # position_offsets / packed_mask init to a valid degenerate LINEAR chain - # (token i at depth i, attending to tokens 0..i incl. self), NOT zeros. - # The reserved CUDA-graph dummy slot (warmup/capture) is never written by - # the C++ tree scatter, so the spec-dec FMHA reads these rows as-is; a - # zeros packed_mask has no self-attention bit and OOBs the trtllm-gen - # kernel. Mirrors the retrieve_next_token chain below; real trees - # overwrite the row via the C++ scatter. bit (32*w + j) of packed_mask - # [i, w] set <=> token i attends to tree token (32*w + j); the causal - # value matches the static-tree mask formula 2^(i+1)-1. + # Default no-tree rows to a valid linear chain; real trees overwrite + # these rows via scatter. _tok = torch.arange(n_dt, device='cuda') self.position_offsets = _tok.to(torch.int32).unsqueeze(0).repeat( S, 1).contiguous() @@ -46,17 +38,8 @@ def __init__(self, self._no_tree_position_offsets = self.position_offsets[0].clone() self._no_tree_packed_mask = self.packed_mask[0].clone() - # Override ONLY the reserved CUDA-graph/warmup dummy slot with a - # bounded-depth K-ary tree (parent[i] = (i-1)//topK). Unlike a real - # slot's no-tree fallback — read only 1 token wide on its first decode — - # the dummy slot is read at the FULL n_dt-wide generation shape by the - # spec-dec verify forward during the CUDA-graph generation warmup. A - # depth-(n_dt-1) linear chain there presents a tree real requests never - # produce (real max depth = max_draft_len, sparse ancestor mask). The - # K-ary template mirrors the drafter's topK expansion so the warmup's - # dummy metadata matches what real dynamic-tree requests feed the - # trtllm-gen FMHA. Scoped to the dummy row so real slots (and the - # accepted eager path) keep the linear fallback above unchanged. + # CUDA-graph dummy rows use a bounded K-ary tree to match real warmup + # verify shapes. _k = max(int(topK), 1) _depth = torch.zeros(n_dt, dtype=torch.int32) _adj = torch.zeros(n_dt, n_dt, dtype=torch.bool) @@ -76,17 +59,8 @@ def __init__(self, dtype=torch.int32, device='cuda') - # Degenerate linear-chain next-token links for no-tree slots (dummy - # CUDA-graph/warmup requests and a real slot's first decode before any - # tree is built). Token i's child is i+1 (parent i-1); the leaf has no - # child (-1). The Mamba tree-aware conv1d/SSU verify path indexes - # retrieve_next_token unconditionally (it does not gate on has_tree), so - # every no-tree row must describe a valid chain rather than sentinels. - # retrieve_next_token is initialized to the chain (not -1) so a - # never-built slot — notably the reserved dummy slot used by CUDA-graph - # capture/warmup, which never passes through the C++ scatter or - # prepare()'s has_tree substitution — still gathers valid, in-bounds - # parent links. Real trees overwrite the row via the C++ scatter. + # Mamba verify reads next links unconditionally, so no-tree rows must be + # valid linear chains instead of sentinels. chain = torch.arange(1, n_dt + 1, dtype=torch.int32, device='cuda') chain[n_dt - 1] = -1 self._no_tree_next_token = chain @@ -170,17 +144,7 @@ def next_links_from_slots(self, slot_ids, count): def apply_no_tree_linear_chain(self, next_token, next_sibling, slot_ids, count): - """Overwrite no-tree rows' links with a valid degenerate linear chain. - - For gen slots whose tree was not built this step (has_tree False: - CUDA-graph/warmup dummies, and a real slot's first decode), the gathered - links are sentinels/uninitialized. The Mamba tree-aware verify path - reads them unconditionally, so replace those rows in-place with the - linear-chain template (next_token[i]=i+1, leaf -1; next_sibling -1) so - the conv1d/SSU parent traversal stays in-bounds and matches the - captured-graph op sequence a real-tree forward replays. Rows with a - real tree (has_tree True) are left untouched. - """ + """Replace no-tree rows with linear-chain links in-place.""" if count == 0: return next_token, next_sibling ids = slot_ids[:count] diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 71f890a24e85..b5cdc5f30202 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -276,8 +276,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): if sa_cfg is not None: sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, max_seq_len) - # Dynamic tree needs a SpecTreeManager (for the target's tree-mask verify - # forward) composed with the MTP hidden-state slot pools. + # Dynamic tree combines SpecTreeManager with MTP hidden-state slots. if getattr(spec_config, 'use_dynamic_tree', False): return MTPEagleDynamicTreeResourceManager( spec_config, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 43dc927f3c6b..a99fe45b812f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2437,9 +2437,7 @@ class MTPDecodingConfig(DecodingBaseConfig): "Top-K candidates expanded per node per draft layer when use_dynamic_tree " "is enabled. Required when use_dynamic_tree is True.") - # Backs the dynamic-tree worker's pre-allocated, batch-indexed CUDA buffers; - # MUST equal the global max_batch_size. Auto-populated by py_executor_creator - # (mirrors Eagle3DecodingConfig). PrivateAttr -- not a user-tunable knob. + # Internal max batch size for dynamic-tree worker buffers. _max_batch_size: Optional[int] = PrivateAttr(default=None) sa_config: Optional[SAEnhancerConfig] = Field( @@ -2486,16 +2484,12 @@ def _remap_deprecated_num_nextn_predict_layers(cls, data): @model_validator(mode="after") def set_max_total_draft_tokens(self): - # Leave max_draft_len as None ("use the model's num_nextn_predict_layers") - # when the user doesn't set it; update_spec_config_from_model_config - # resolves it from the checkpoint before the model runs. + # None means update_spec_config_from_model_config resolves it from checkpoint. if self.max_draft_len is not None: if self.max_draft_len <= 0: raise ValueError("max_draft_len must be > 0 for MTP") - # Dynamic-tree MTP: mirror EagleDecodingConfig. Honor an explicit - # max_total_draft_tokens within [max_draft_len, dynamic_tree_max_topK * max_draft_len]; - # otherwise default to dynamic_tree_max_topK * max_draft_len. + # Dynamic tree defaults max_total_draft_tokens to topK * max_draft_len. if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: self.use_dynamic_tree = True if self.max_draft_len is None: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py index 6566e13b47c3..3cfc3ca58107 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py @@ -1,12 +1,4 @@ -"""Unit tests for DynamicTreeSlotStorage no-tree link fallback. - -The Mamba tree-aware verify conv1d/SSU path reads retrieve_next_token / -retrieve_next_sibling unconditionally (it does not gate on has_tree). For -gen slots without a built tree (CUDA-graph/warmup dummies, and a real slot's -first decode), the gathered links are sentinels; apply_no_tree_linear_chain -must replace those rows with a valid degenerate linear chain so the kernel -indexing stays in-bounds. Real-tree rows must be left untouched. -""" +"""Unit tests for DynamicTreeSlotStorage no-tree link fallback.""" import unittest @@ -35,8 +27,7 @@ def test_no_tree_rows_get_chain_real_rows_untouched(self): num_slots, n_dt = 4, 6 ss = self._make(num_slots=num_slots, n_dt=n_dt) - # Slot 1 has a (synthetic) real tree; mark it valid and give it links - # distinct from the chain so we can detect any accidental overwrite. + # Slot 1 has real-tree links distinct from the fallback chain. real_next_token = torch.full((n_dt,), 3, dtype=torch.int32, device="cuda") real_next_sibling = torch.full((n_dt,), 2, dtype=torch.int32, device="cuda") ss.retrieve_next_token[1] = real_next_token @@ -66,8 +57,7 @@ def test_no_tree_rows_get_chain_real_rows_untouched(self): self.assertTrue(torch.equal(next_sibling[1], real_next_sibling)) def test_chain_links_are_in_bounds(self): - # Every link is either a valid token index in [0, n_dt) or the -1 stop - # sentinel; nothing points outside the per-request token range. + # Links are valid token indices or the -1 stop sentinel. n_dt = 8 ss = self._make(n_dt=n_dt) slot_ids = torch.tensor([ss.dummy_slot_id], dtype=torch.long, device="cuda") From 095a98e7b0d5c0913344a53ec7b13de20e4122cd Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 24 Jun 2026 03:52:41 -0700 Subject: [PATCH 07/24] [None][test] Remove dynamic tree slot storage test Signed-off-by: qgai --- .../test_dynamic_tree_slot_storage.py | 71 ------------------- 1 file changed, 71 deletions(-) delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py deleted file mode 100644 index 3cfc3ca58107..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Unit tests for DynamicTreeSlotStorage no-tree link fallback.""" - -import unittest - -import pytest -import torch - -from tensorrt_llm._torch.speculative.spec_tree_manager import DynamicTreeSlotStorage - - -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="DynamicTreeSlotStorage allocates CUDA buffers" -) -class TestNoTreeLinearChain(unittest.TestCase): - def _make(self, num_slots=4, n_dt=6): - # mask_width is unused by the link path; any positive value is fine. - return DynamicTreeSlotStorage(num_slots=num_slots, n_dt=n_dt, mask_width=1) - - def test_chain_template_is_valid(self): - n_dt = 6 - ss = self._make(n_dt=n_dt) - # next_token[i] = i+1, leaf has no child (-1). - expected = torch.tensor([1, 2, 3, 4, 5, -1], dtype=torch.int32, device="cuda") - self.assertTrue(torch.equal(ss._no_tree_next_token, expected)) - - def test_no_tree_rows_get_chain_real_rows_untouched(self): - num_slots, n_dt = 4, 6 - ss = self._make(num_slots=num_slots, n_dt=n_dt) - - # Slot 1 has real-tree links distinct from the fallback chain. - real_next_token = torch.full((n_dt,), 3, dtype=torch.int32, device="cuda") - real_next_sibling = torch.full((n_dt,), 2, dtype=torch.int32, device="cuda") - ss.retrieve_next_token[1] = real_next_token - ss.retrieve_next_sibling[1] = real_next_sibling - ss.has_tree[1] = True - - # Gather slots [dummy, real, dummy] -> rows 0 and 2 are no-tree. - slot_ids = torch.tensor( - [ss.dummy_slot_id, 1, ss.dummy_slot_id], dtype=torch.long, device="cuda" - ) - count = 3 - next_token, next_sibling = ss.next_links_from_slots(slot_ids, count) - ss.apply_no_tree_linear_chain(next_token, next_sibling, slot_ids, count) - - chain = ss._no_tree_next_token - # No-tree rows -> linear chain, sibling all -1. - self.assertTrue(torch.equal(next_token[0], chain)) - self.assertTrue(torch.equal(next_token[2], chain)) - self.assertTrue( - torch.equal(next_sibling[0], torch.full((n_dt,), -1, dtype=torch.int32, device="cuda")) - ) - self.assertTrue( - torch.equal(next_sibling[2], torch.full((n_dt,), -1, dtype=torch.int32, device="cuda")) - ) - # Real-tree row -> original links preserved. - self.assertTrue(torch.equal(next_token[1], real_next_token)) - self.assertTrue(torch.equal(next_sibling[1], real_next_sibling)) - - def test_chain_links_are_in_bounds(self): - # Links are valid token indices or the -1 stop sentinel. - n_dt = 8 - ss = self._make(n_dt=n_dt) - slot_ids = torch.tensor([ss.dummy_slot_id], dtype=torch.long, device="cuda") - next_token, next_sibling = ss.next_links_from_slots(slot_ids, 1) - ss.apply_no_tree_linear_chain(next_token, next_sibling, slot_ids, 1) - valid = (next_token == -1) | ((next_token >= 0) & (next_token < n_dt)) - self.assertTrue(bool(valid.all().item())) - - -if __name__ == "__main__": - unittest.main() From 3b0f0ab10a1b3193ad6b00fc5c5024dac1f1a8e1 Mon Sep 17 00:00:00 2001 From: qgai Date: Sun, 28 Jun 2026 23:44:04 -0700 Subject: [PATCH 08/24] [None][fix] Fix MTP dynamic tree metadata Signed-off-by: qgai --- cpp/tensorrt_llm/common/attentionOp.h | 10 +-- .../_torch/pyexecutor/model_engine.py | 8 +- tensorrt_llm/_torch/speculative/eagle3.py | 31 +++++++- tensorrt_llm/_torch/speculative/mtp.py | 39 +--------- .../_torch/speculative/mtp_dynamic_tree.py | 10 ++- tensorrt_llm/_torch/speculative/utils.py | 3 +- .../speculative/hw_agnostic/test_mtp.py | 20 +++++ .../_torch/speculative/test_eagle3.py | 73 +++++++++++++++++++ 8 files changed, 142 insertions(+), 52 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index d716bc5d5a81..c76907803c69 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -572,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); diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 278a7bf4a660..ea9d0c349f66 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -620,7 +620,13 @@ def __init__( ) or self.model_is_wrapped self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 self.max_draft_len = spec_config.max_draft_len - self.runtime_draft_len = spec_config.max_draft_len + # Mutable per-iteration draft length (updated each iteration when + # dynamic draft length is enabled; otherwise stays fixed). Tree + # modes verify all tree nodes per step, which can be wider than the + # tree depth used by the drafter loop. + self.runtime_draft_len = (self.max_total_draft_tokens + if not spec_config.is_linear_tree else + self.max_draft_len) else: self.without_logits = False diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d9b5be1ecc1e..b7861c8c4ab9 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -398,6 +398,10 @@ class Eagle3OneModelSpecMetadata(SpecMetadata): # prepare() before self.num_tokens is decremented to the attention-DP subseq # shape; maybe_capture_hidden_states must bound by this, not self.num_tokens. num_capture_tokens: int = 0 + # Per-generation tree links for Mamba verify in dynamic-tree one-model paths. + retrieve_next_token: Optional[torch.Tensor] = None + retrieve_next_sibling: Optional[torch.Tensor] = None + retrieve_parent_token: Optional[torch.Tensor] = None def __post_init__(self): if self.layers_to_capture is None: @@ -443,9 +447,10 @@ def __post_init__(self): self.hidden_size * len(self.layers_to_capture)), dtype=self.dtype, device='cuda') - if (self.spec_resource_manager is not None - and self.spec_resource_manager.batch_indices_cuda is not None): - self.batch_indices_cuda = self.spec_resource_manager.batch_indices_cuda + batch_indices_cuda = getattr(self.spec_resource_manager, + "batch_indices_cuda", None) + if batch_indices_cuda is not None: + self.batch_indices_cuda = batch_indices_cuda assert self.batch_indices_cuda.shape[0] >= self.max_num_requests, ( f"batch_indices_cuda shape mismatch: " f"{type(self.spec_resource_manager).__name__} has " @@ -530,6 +535,26 @@ def prepare(self): if gen_request_ids: sa_manager.prepare(gen_request_ids, self.runtime_draft_len) + self.retrieve_next_token = None + self.retrieve_next_sibling = None + self.retrieve_parent_token = None + spec_tree_manager = getattr(self.spec_resource_manager, + 'spec_tree_manager', None) + if self.use_dynamic_tree and spec_tree_manager is not None: + num_gens = self.num_generations + if num_gens > 0: + num_contexts = num_seqs - num_gens + slot_storage = spec_tree_manager.slot_storage + gen_slot_ids = slot_storage.all_ids_buf[ + num_contexts:num_contexts + num_gens] + next_token, next_sibling = slot_storage.next_links_from_slots( + gen_slot_ids, num_gens) + slot_storage.apply_no_tree_linear_chain(next_token, + next_sibling, + gen_slot_ids, num_gens) + self.retrieve_next_token = next_token + self.retrieve_next_sibling = next_sibling + def maybe_capture_hidden_states( self, layer_id: int, diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index b596b3c19ae3..fcf0603a04c1 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -140,15 +140,6 @@ class MTPSpecMetadata(SpecMetadata): # CUDA graph, we use this tensor to store the number of input tokens for the # subsequent draft forward. subseq_all_rank_num_tokens: Optional[List[int]] = None - # Dynamic-tree MTP-Eagle uses per-slot masks during target verify. - use_dynamic_tree: bool = False - dynamic_tree_max_topK: Optional[int] = None - spec_tree_manager: Optional[object] = None - # Per-generation tree links for Mamba verify. Parent links are derived in - # conv1d. Shape: [num_generations, max_total_draft_tokens + 1]. - retrieve_next_token: Optional[torch.Tensor] = None - retrieve_next_sibling: Optional[torch.Tensor] = None - retrieve_parent_token: Optional[torch.Tensor] = None def __post_init__(self) -> None: if self.mtp_hidden_states_manager is not None: @@ -178,10 +169,6 @@ def __post_init__(self) -> None: self.mtp_num_modules, device='cuda', ) - # Enable target-side tree-mask routing for dynamic-tree MTP. - if self.use_dynamic_tree: - self.is_spec_dec_tree = True - self.is_spec_dec_dynamic_tree = True @property def all_rank_num_seqs(self): @@ -207,11 +194,7 @@ def prepare(self): # while MTP Eagle worker uses (max_draft_len + 1) input tokens in the 1st draft # forward and only one input token in the following draft forward. # This num_tokens is used to set the all_rank_num_tokens for attention dp. - if self.use_dynamic_tree: - # Step-0 draft uses max_draft_len + 1 tokens per generation. - self.num_tokens -= self.num_generations * ( - self.max_total_draft_tokens - self.max_draft_len) - elif not self.spec_dec_mode.is_mtp_eagle_one_model(): + if not self.spec_dec_mode.is_mtp_eagle_one_model(): self.num_tokens -= self.num_generations if self.mtp_hidden_states_manager is not None: # MTP vanilla or use relaxed acceptance @@ -255,26 +238,6 @@ def prepare(self): if gen_request_ids: sa_manager.prepare(gen_request_ids, self.runtime_draft_len) - # Gather per-generation tree links for Mamba verify. - self.retrieve_next_token = None - self.retrieve_next_sibling = None - self.retrieve_parent_token = None - if self.use_dynamic_tree and self.spec_tree_manager is not None: - num_gens = self.num_generations - if num_gens > 0: - num_contexts = num_seqs - num_gens - slot_storage = self.spec_tree_manager.slot_storage - gen_slot_ids = slot_storage.all_ids_buf[ - num_contexts:num_contexts + num_gens] - next_token, next_sibling = slot_storage.next_links_from_slots( - gen_slot_ids, num_gens) - # No-tree rows need valid links; real-tree rows are unchanged. - slot_storage.apply_no_tree_linear_chain(next_token, - next_sibling, - gen_slot_ids, num_gens) - self.retrieve_next_token = next_token - self.retrieve_next_sibling = next_sibling - class MTPSampler(SpecSamplerBase): """ diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index ff945f0bdbcb..14ae6972da8e 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -201,6 +201,11 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): else: self._saved_generation_lengths = None + def prepare_position_ids_and_last_tokens(self, position_ids, seq_lens_cuda): + position_ids = position_ids.squeeze(0) + last_tokens_idx = torch.cumsum(seq_lens_cuda, dim=0, dtype=torch.long) - 1 + return position_ids, last_tokens_idx + def _restore_attn_metadata_from_spec_dec(self, attn_metadata): super()._restore_attn_metadata_from_spec_dec(attn_metadata) @@ -1030,10 +1035,8 @@ def _prepare_draft_layer_metadata( attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( self._max_path_len ) - num_accepted_tokens[num_contexts:batch_size] - if num_contexts > 0: - attn_metadata.kv_lens_cuda[:num_contexts] += self.K + attn_metadata.kv_lens_cuda[:batch_size] += self.K attn_metadata.use_spec_decoding = True - attn_metadata.update_for_spec_dec() self._refresh_blackwell_tree_mask_metadata(attn_metadata) else: num_tokens_previous_layer = cur_draft_idx * self.K @@ -1047,7 +1050,6 @@ def _prepare_draft_layer_metadata( attn_metadata.on_update() if hasattr(attn_metadata, "kv_lens_cuda"): attn_metadata.kv_lens_cuda[:batch_size] += self.K - attn_metadata.update_for_spec_dec() self._refresh_blackwell_tree_mask_metadata(attn_metadata) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index b5cdc5f30202..d800751b209f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -556,7 +556,8 @@ def update_spec_config_from_model_config(spec_config, model_config): f"using max_draft_len={effective_draft_len} draft tokens.") spec_config.max_draft_len = effective_draft_len - spec_config.max_total_draft_tokens = spec_config.max_draft_len + if not spec_config.use_dynamic_tree: + spec_config.max_total_draft_tokens = spec_config.max_draft_len def update_spec_config_from_loaded_model(spec_config, model) -> None: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index a381dbd94907..a537e38b97df 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -11,12 +11,17 @@ from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker +from tensorrt_llm._torch.speculative.utils import update_spec_config_from_model_config from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) from utils.llm_data import llm_models_root +class _MTPPretrainedConfig: + num_nextn_predict_layers = 1 + + def unittest_name_func(testcase_func, param_num, param): name = param.args[0] return "%s_%s" % ( @@ -25,6 +30,21 @@ def unittest_name_func(testcase_func, param_num, param): ) +def test_mtp_dynamic_tree_preserves_max_total_draft_tokens(): + spec_config = MTPDecodingConfig( + max_draft_len=6, + max_total_draft_tokens=15, + use_dynamic_tree=True, + dynamic_tree_max_topK=4, + ) + + update_spec_config_from_model_config(spec_config, _MTPPretrainedConfig()) + + assert spec_config.max_draft_len == 6 + assert spec_config.max_total_draft_tokens == 15 + assert spec_config.tokens_per_gen_step == 16 + + class TestMTPSampleAndAcceptDraftTokens(unittest.TestCase): def setUp(self): tensorrt_llm.logger.set_level("warning") diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 4e7faecbbea2..452fd5b289dc 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -21,6 +21,8 @@ from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ _extend_full_attention_windows_for_spec_decode from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig) @@ -282,6 +284,77 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): assert draft_kwargs["max_blocks"] is None +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_mtp_eagle_one_model_dynamic_tree_metadata_prepares_mamba_links(): + max_num_requests = 3 + max_draft_len = 2 + max_total_draft_tokens = 2 + spec_tree_manager = SpecTreeManager( + max_num_requests=max_num_requests, + use_dynamic_tree=True, + max_total_draft_tokens=max_total_draft_tokens, + max_draft_len=max_draft_len, + eagle_choices=None, + dynamic_tree_max_topK=2, + ) + slot_storage = spec_tree_manager.slot_storage + slot_storage.all_ids_buf[:max_num_requests].copy_( + torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + slot_storage.has_tree[1] = True + slot_storage.retrieve_next_token[1] = torch.tensor([2, -1, -1], + dtype=torch.int32, + device="cuda") + slot_storage.retrieve_next_sibling[1] = torch.tensor([-1, -1, -1], + dtype=torch.int32, + device="cuda") + + class _ResourceManager: + hidden_states = None + slot_manager = None + sa_manager = None + + def __init__(self): + self.spec_tree_manager = spec_tree_manager + self.batch_indices_cuda = torch.empty(max_num_requests, + dtype=torch.int, + device="cuda") + + metadata = Eagle3OneModelSpecMetadata( + max_draft_len=max_draft_len, + max_total_draft_tokens=max_total_draft_tokens, + spec_dec_mode=SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL, + max_num_requests=max_num_requests, + num_layers=1, + hidden_size=1, + max_num_tokens=16, + spec_resource_manager=_ResourceManager(), + use_dynamic_tree=True, + ) + metadata.request_ids = [10, 11, 12] + metadata.seq_lens = [ + 1, max_total_draft_tokens + 1, max_total_draft_tokens + 1 + ] + metadata.num_generations = 2 + metadata.num_tokens = 1 + 2 * (max_total_draft_tokens + 1) + + metadata.prepare() + + assert metadata.retrieve_next_token is not None + assert metadata.retrieve_next_sibling is not None + assert metadata.retrieve_next_token.shape == (2, max_total_draft_tokens + 1) + assert torch.equal(metadata.retrieve_next_token[0], + slot_storage.retrieve_next_token[1]) + assert torch.equal( + metadata.retrieve_next_token[1], + torch.tensor([1, 2, -1], dtype=torch.int32, device="cuda")) + assert torch.equal( + metadata.retrieve_next_sibling[1], + torch.full((max_total_draft_tokens + 1, ), + -1, + dtype=torch.int32, + device="cuda")) + + @pytest.mark.parametrize( "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model", [ From 30cf2e4552bc6707cd566169f3159eb8449a6ad3 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 00:02:33 -0700 Subject: [PATCH 09/24] [None][chore] Remove unused spec-dec autotuner warmup Signed-off-by: qgai --- .../_torch/pyexecutor/model_engine.py | 59 ------------------- 1 file changed, 59 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ea9d0c349f66..6adfc7ec14cc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1390,65 +1390,6 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): clear_memory_buffers() torch.cuda.empty_cache() - def _need_spec_dec_gen_autotuner_warmup( - self, resource_manager: ResourceManager) -> bool: - """Whether dynamic-tree spec-dec generation needs explicit autotuning.""" - if not getattr(self.llm_args, "enable_autotuner", True): - return False - if not self.cuda_graph_runner.enabled: - return False - if self.spec_config is None or self.is_draft_model: - return False - if not self.spec_config.spec_dec_mode.use_one_engine(): - return False - if not getattr(self.spec_config, "use_dynamic_tree", False): - return False - kv_cache_manager = resource_manager.get_resource_manager( - self.kv_cache_manager_key) - return isinstance(kv_cache_manager, MambaHybridCacheManager) - - def _run_spec_dec_gen_autotuner_warmup( - self, resource_manager: ResourceManager) -> None: - """Profile the CUDA-graph generation MoE shape before capture.""" - if not self._need_spec_dec_gen_autotuner_warmup(resource_manager): - return - - draft_len = self.max_total_draft_tokens - effective_max_seq_len = self.max_seq_len - if self.mapping is not None and self.mapping.has_cp_helix(): - effective_max_seq_len = self.max_seq_len // self.mapping.cp_size - effective_max_seq_len = min(effective_max_seq_len, self.max_num_tokens) - - cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) - AutoTuner.get().setup_distributed_state(self.mapping, self.dist) - logger.info("Running spec-dec generation autotuner warmup...") - with self.no_cuda_graph(), autotune(cache_path=cache_path): - for bs in sorted(self._cuda_graph_batch_sizes, reverse=True): - if bs > self.batch_size: - continue - warmup_request = self._create_cuda_graph_warmup_request( - resource_manager, bs, draft_len, effective_max_seq_len) - with self._release_batch_context(warmup_request, - resource_manager) as batch: - if batch is None: - continue - logger.info( - f"Run pre-capture autotuner warmup at generation shape " - f"(bs={bs}, draft_len={draft_len}, " - f"max_seq_len={effective_max_seq_len})") - self.enable_spec_decode = True - self.runtime_draft_len = draft_len - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) - torch.cuda.synchronize() - - self.enable_spec_decode = self.is_spec_decode - self.runtime_draft_len = self.max_draft_len - logger.info( - f"[Autotuner] Cache size after spec-dec generation warmup is {len(AutoTuner.get().profiling_cache)}" - ) - def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. From 106b8082832e8f3765b7555b5ff03c626d799670 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 00:08:44 -0700 Subject: [PATCH 10/24] [None][chore] Clean up dynamic tree slot storage Signed-off-by: qgai --- tensorrt_llm/_torch/speculative/eagle3.py | 3 - .../_torch/speculative/spec_tree_manager.py | 107 +++++++++++------- 2 files changed, 65 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index b7861c8c4ab9..062b61c03be9 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -549,9 +549,6 @@ def prepare(self): num_contexts:num_contexts + num_gens] next_token, next_sibling = slot_storage.next_links_from_slots( gen_slot_ids, num_gens) - slot_storage.apply_no_tree_linear_chain(next_token, - next_sibling, - gen_slot_ids, num_gens) self.retrieve_next_token = next_token self.retrieve_next_sibling = next_sibling diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 1a44a1bd064b..119b073a52e7 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -20,51 +20,36 @@ def __init__(self, num_slots: int, n_dt: int, mask_width: int, - topK: int = 1): + top_k: int = 1): S = num_slots + 1 self.dummy_slot_id = num_slots # Default no-tree rows to a valid linear chain; real trees overwrite # these rows via scatter. - _tok = torch.arange(n_dt, device='cuda') - self.position_offsets = _tok.to(torch.int32).unsqueeze(0).repeat( + no_tree_position_offsets, no_tree_packed_mask = self._make_no_tree_metadata( + n_dt, mask_width) + self.position_offsets = no_tree_position_offsets.unsqueeze(0).repeat( S, 1).contiguous() - _bit = torch.arange(mask_width * 32, device='cuda') - _causal = ((_bit.unsqueeze(0) <= _tok.unsqueeze(1)) - & (_bit.unsqueeze(0) < n_dt)).view(n_dt, mask_width, 32) - _w = 2**torch.arange(32, dtype=torch.int64, device='cuda') - self.packed_mask = (_causal.to(torch.int64) * _w).sum(-1).to( - torch.int32).unsqueeze(0).repeat(S, 1, 1).contiguous() + self.packed_mask = no_tree_packed_mask.unsqueeze(0).repeat( + S, 1, 1).contiguous() self._no_tree_position_offsets = self.position_offsets[0].clone() self._no_tree_packed_mask = self.packed_mask[0].clone() # CUDA-graph dummy rows use a bounded K-ary tree to match real warmup # verify shapes. - _k = max(int(topK), 1) - _depth = torch.zeros(n_dt, dtype=torch.int32) - _adj = torch.zeros(n_dt, n_dt, dtype=torch.bool) - for _i in range(n_dt): - _adj[_i, _i] = True # self - if _i > 0: - _p = (_i - 1) // _k # parent index < _i, its row already final - _depth[_i] = _depth[_p] + 1 - _adj[_i] |= _adj[_p] # inherit ancestors (incl. root) - self.position_offsets[self.dummy_slot_id] = _depth.to(device='cuda') - _adj_pad = torch.zeros(n_dt, mask_width * 32, dtype=torch.bool) - _adj_pad[:, :n_dt] = _adj - _dummy_mask = (_adj_pad.to(device='cuda').view(n_dt, mask_width, 32).to( - torch.int64) * _w).sum(-1).to(torch.int32) - self.packed_mask[self.dummy_slot_id] = _dummy_mask + dummy_position_offsets, dummy_packed_mask = self._make_dummy_tree_metadata( + n_dt, mask_width, top_k) + self.position_offsets[self.dummy_slot_id] = dummy_position_offsets + self.packed_mask[self.dummy_slot_id] = dummy_packed_mask self.retrieve_index = torch.zeros((S, n_dt), dtype=torch.int32, device='cuda') # Mamba verify reads next links unconditionally, so no-tree rows must be # valid linear chains instead of sentinels. - chain = torch.arange(1, n_dt + 1, dtype=torch.int32, device='cuda') - chain[n_dt - 1] = -1 - self._no_tree_next_token = chain - self.retrieve_next_token = chain.unsqueeze(0).repeat(S, 1) + self._no_tree_next_token = self._make_no_tree_next_token(n_dt) + self.retrieve_next_token = self._no_tree_next_token.unsqueeze(0).repeat( + S, 1) self.retrieve_next_sibling = torch.full((S, n_dt), -1, dtype=torch.int32, @@ -88,6 +73,57 @@ def __init__(self, dtype=torch.int32, device='cuda') + @staticmethod + def _pack_bool_mask(mask: torch.Tensor, mask_width: int) -> torch.Tensor: + """Pack a bool attention mask into int32 words.""" + num_rows, num_bits = mask.shape + total_bits = mask_width * 32 + padded_mask = torch.zeros((num_rows, total_bits), + dtype=torch.bool, + device=mask.device) + padded_mask[:, :num_bits] = mask + weights = 2**torch.arange(32, dtype=torch.int64, device=mask.device) + return (padded_mask.view(num_rows, mask_width, 32).to(torch.int64) * + weights).sum(-1).to(torch.int32) + + @classmethod + def _make_no_tree_metadata( + cls, n_dt: int, + mask_width: int) -> tuple[torch.Tensor, torch.Tensor]: + token_ids = torch.arange(n_dt, device='cuda') + position_offsets = token_ids.to(torch.int32) + causal_mask = token_ids.unsqueeze(1) >= token_ids.unsqueeze(0) + return position_offsets, cls._pack_bool_mask(causal_mask, mask_width) + + @classmethod + def _make_dummy_tree_metadata( + cls, n_dt: int, mask_width: int, + top_k: int) -> tuple[torch.Tensor, torch.Tensor]: + top_k = max(int(top_k), 1) + position_offsets = [0] * n_dt + ancestor_mask = [[False] * n_dt for _ in range(n_dt)] + for token_idx in range(n_dt): + ancestor_mask[token_idx][token_idx] = True + if token_idx > 0: + parent_idx = (token_idx - 1) // top_k + position_offsets[token_idx] = position_offsets[parent_idx] + 1 + ancestor_mask[token_idx][:token_idx] = ancestor_mask[ + parent_idx][:token_idx] + + position_offsets = torch.tensor(position_offsets, + dtype=torch.int32, + device='cuda') + ancestor_mask = torch.tensor(ancestor_mask, + dtype=torch.bool, + device='cuda') + return position_offsets, cls._pack_bool_mask(ancestor_mask, mask_width) + + @staticmethod + def _make_no_tree_next_token(n_dt: int) -> torch.Tensor: + next_token = torch.arange(1, n_dt + 1, dtype=torch.int32, device='cuda') + next_token[n_dt - 1] = -1 + return next_token + def fill_all_slot_ids(self, context_requests, generation_requests): """Fill all_ids_buf for full batch [ctx | gen] via one HtoD copy.""" dummy_slot = self.dummy_slot_id @@ -142,19 +178,6 @@ def next_links_from_slots(self, slot_ids, count): torch.index_select(self.retrieve_next_sibling, 0, ids, out=next_sibling) return next_token, next_sibling - def apply_no_tree_linear_chain(self, next_token, next_sibling, slot_ids, - count): - """Replace no-tree rows with linear-chain links in-place.""" - if count == 0: - return next_token, next_sibling - ids = slot_ids[:count] - no_tree = ~self.has_tree[ids] # [count] - mask = no_tree.unsqueeze(1) # [count, 1] broadcasts over n_dt - next_token.copy_(torch.where(mask, self._no_tree_next_token, - next_token)) - next_sibling.masked_fill_(mask, -1) - return next_token, next_sibling - class SpecTreeManager: use_dynamic_tree: bool # Whether using dynamic tree @@ -328,7 +351,7 @@ def init_tree_info_for_dynamic_tree(self): num_slots=self.num_trees, n_dt=num_draft_with_root, mask_width=mask_width, - topK=self.dynamic_tree_max_topK, + top_k=self.dynamic_tree_max_topK, ) def scatter_to_slot_storage(self, ss, gen_slots, num_gens): From ca9c3a8433e491c2fca5faf9ccd3bac2f7991a3a Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 00:17:33 -0700 Subject: [PATCH 11/24] [None][chore] Optimize dynamic tree slot staging Signed-off-by: qgai --- .../_torch/speculative/spec_tree_manager.py | 63 ++++++++++++------- 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 119b073a52e7..7f6ec8d5869e 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -24,19 +24,19 @@ def __init__(self, S = num_slots + 1 self.dummy_slot_id = num_slots - # Default no-tree rows to a valid linear chain; real trees overwrite - # these rows via scatter. + # Bootstrap/reused slots may not have a tree yet; keep their metadata + # as a valid linear chain so verification kernels can read it directly. no_tree_position_offsets, no_tree_packed_mask = self._make_no_tree_metadata( n_dt, mask_width) self.position_offsets = no_tree_position_offsets.unsqueeze(0).repeat( S, 1).contiguous() self.packed_mask = no_tree_packed_mask.unsqueeze(0).repeat( S, 1, 1).contiguous() - self._no_tree_position_offsets = self.position_offsets[0].clone() - self._no_tree_packed_mask = self.packed_mask[0].clone() + self._no_tree_position_offsets = no_tree_position_offsets + self._no_tree_packed_mask = no_tree_packed_mask - # CUDA-graph dummy rows use a bounded K-ary tree to match real warmup - # verify shapes. + # CUDA-graph dummies use a deterministic K-ary tree, matching real + # dynamic-tree mask/position shapes without depending on request state. dummy_position_offsets, dummy_packed_mask = self._make_dummy_tree_metadata( n_dt, mask_width, top_k) self.position_offsets[self.dummy_slot_id] = dummy_position_offsets @@ -100,23 +100,31 @@ def _make_dummy_tree_metadata( cls, n_dt: int, mask_width: int, top_k: int) -> tuple[torch.Tensor, torch.Tensor]: top_k = max(int(top_k), 1) - position_offsets = [0] * n_dt - ancestor_mask = [[False] * n_dt for _ in range(n_dt)] - for token_idx in range(n_dt): - ancestor_mask[token_idx][token_idx] = True - if token_idx > 0: - parent_idx = (token_idx - 1) // top_k - position_offsets[token_idx] = position_offsets[parent_idx] + 1 - ancestor_mask[token_idx][:token_idx] = ancestor_mask[ - parent_idx][:token_idx] - - position_offsets = torch.tensor(position_offsets, - dtype=torch.int32, - device='cuda') - ancestor_mask = torch.tensor(ancestor_mask, + token_ids = torch.arange(n_dt, device='cuda') + parents = torch.where(token_ids > 0, (token_ids - 1) // top_k, + token_ids) + ancestor_chain = torch.empty((n_dt, n_dt), + dtype=torch.long, + device='cuda') + current = token_ids + for depth in range(n_dt): + ancestor_chain[:, depth] = current + current = parents[current] + + # Pack bits directly from the parent chain instead of materializing a + # dense bool mask and repacking it. + valid_ancestors = torch.ones((n_dt, n_dt), dtype=torch.bool, device='cuda') - return position_offsets, cls._pack_bool_mask(ancestor_mask, mask_width) + valid_ancestors[:, 1:] = ancestor_chain[:, 1:] != ancestor_chain[:, :-1] + bit_values = (1 << (ancestor_chain % 32)).to(torch.int32) + bit_values.masked_fill_(~valid_ancestors, 0) + packed_mask = torch.zeros((n_dt, mask_width), + dtype=torch.int32, + device='cuda') + packed_mask.scatter_add_(1, ancestor_chain // 32, bit_values) + position_offsets = valid_ancestors.sum(-1).to(torch.int32) - 1 + return position_offsets, packed_mask @staticmethod def _make_no_tree_next_token(n_dt: int) -> torch.Tensor: @@ -162,9 +170,16 @@ def pack_retrieve_from_slots(self, slot_ids, count): return self._verify_staging[:0] ids = slot_ids[:count] staging = self._verify_staging[:count] - staging[:, :, 0] = self.retrieve_index[ids] - staging[:, :, 1] = self.retrieve_next_token[ids] - staging[:, :, 2] = self.retrieve_next_sibling[ids] + # Avoid advanced-indexing temporaries on the verification path. + torch.index_select(self.retrieve_index, 0, ids, out=staging[:, :, 0]) + torch.index_select(self.retrieve_next_token, + 0, + ids, + out=staging[:, :, 1]) + torch.index_select(self.retrieve_next_sibling, + 0, + ids, + out=staging[:, :, 2]) return staging def next_links_from_slots(self, slot_ids, count): From a3c9307eb4cc805cb853448fdbf65b2836929682 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 00:40:10 -0700 Subject: [PATCH 12/24] [None][chore] Reduce dynamic tree cleanup diff Signed-off-by: qgai --- .../_torch/speculative/spec_tree_manager.py | 38 +++++++++---------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 7f6ec8d5869e..061086c19a02 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -1,6 +1,5 @@ import logging import math -from itertools import accumulate from typing import List import torch @@ -203,11 +202,11 @@ class SpecTreeManager: # Auxiliary buffers # The top k list for each draft layer. - top_k_list: list + top_k_list = [] # The user input eagle choices, only available when using static tree. - eagle_choices: List[List[int]] - # If dynamic tree, each request has their own tree. If static tree, all requests share the same tree. - num_trees: int + eagle_choices: List[List[int]] = None + # If dynamice tree, each request has their own tree. If static tree, all requests share the same tree. + num_trees: int = None # Convert the choice to a path. Each path is an array of indices from the root to other nodes in the tree. # shape: [num_trees, max_total_draft_tokens + 1, max_draft_len + 1] @@ -225,33 +224,37 @@ class SpecTreeManager: # shape: [num_trees, max_total_draft_tokens + 1], device tensor. spec_dec_position_offsets: torch.Tensor = None - ############################ Auxiliary buffers for the static tree. ############################ + # TODO: Optimized together with the subsequent dynamic tree. + # Auxiliary buffers for the static tree. # Considering that the static tree does not modify the tree structure during inference, we can calculate some buffers in advance. # NOTE: Most of these buffers are introduced due to limitations of XQA: # With tree attention, XQA cannot simply take the tokens to be processed in the next round as input. Instead, it needs to take ALL of their parent nodes as input. # This incurs additional computation, but it is unavoidable. # NOTE: The reason why most of these auxiliary buffers are with `len == max_draft_len - 1` is that: we do not need to prepare specific input data for the first draft layer. + # The top k value for each draft layer. Device tensor. top_k_list_cuda: list[torch.Tensor] = None # The max top k value for all draft layers. Which is used for torch.topk and cuda graph. max_top_k = -1 - # Gather the required draft tokens among the 'max_total_draft_tokens + 1' tokens. + # Gather the required draft tokens from all currently generated draft tokens as the input of the next draft layer. # Only the nodes has child(s) this layer and all their parents nodes will be gathered. + # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. Each element is a tensor with shape [num_tokens_for_next_layer]. tokens_gather_idx_for_drafter_model: list[torch.Tensor] = None + # Gather the required logits from all currently generated logits. + # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. + logits_gather_idx: list[torch.Tensor] = None + # The packed mask for the drafter model's attention (i.e., xqa). - # shape: [1, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)], device tensor. spec_dec_packed_mask_for_drafter_model: torch.Tensor = None # The read indices offset for the drafter model. - # shape: [max_total_draft_tokens + 1], device tensor. hidden_states_read_indices_offset_for_drafter_model: torch.Tensor = None # The write back start indices for the drafter tokens between different draft layers. - # shape: [max_draft_len + 1], device tensor. draft_tokens_indices_cumsum: torch.Tensor = None # Work buffers for dynamic tree build kernel output @@ -384,6 +387,7 @@ def scatter_to_slot_storage(self, ss, gen_slots, num_gens): 0, ids, self.retrieve_next_sibling[:num_gens]) ss.mark_valid(ids, num_gens) + # For the static tree def init_tree_info_for_static_tree(self): self.index_mapping_set = {} self.nodes_list_per_layer = [[] for _ in range(self.max_draft_len + 1)] @@ -432,7 +436,9 @@ def init_tree_info_for_static_tree(self): pin_memory=prefer_pinned())) # 6) Compute the spec decoding according to the eagle_paths for the target model - self.compute_spec_dec_mask_matrix(0) + for i, path in enumerate(self.eagle_paths[0]): + indices = path[path > -1] + self.spec_dec_mask_matrix[0][i, indices] = 1 self.compute_spec_dec_packed_mask(self.spec_dec_mask_matrix, self.spec_dec_packed_mask) @@ -471,6 +477,7 @@ def init_tree_info_for_static_tree(self): num_nodes_per_layer = [0] num_nodes_per_layer.extend( [len(node_list) for node_list in self.nodes_list_per_layer[1:]]) + from itertools import accumulate self.draft_tokens_indices_cumsum = torch.tensor(list( accumulate(num_nodes_per_layer)), dtype=torch.int32, @@ -521,15 +528,6 @@ def get_top_k_list(self, draft_layer_id): assert draft_layer_id >= 0 return self.top_k_list[draft_layer_id] - def compute_spec_dec_mask_matrix(self, tree_idx=0): - if self.eagle_paths is None: - raise RuntimeError( - "compute_spec_dec_mask_matrix() is not supported in dynamic tree mode" - ) - for i, path in enumerate(self.eagle_paths[0]): - indices = path[path > -1] - self.spec_dec_mask_matrix[0][i, indices] = 1 - def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask): bs, num_tokens, num_tokens_attend = mask_matrix.shape assert mask_matrix.ndim == 3, f"Expected 3D mask_matrix, got {mask_matrix.ndim}D" From 1852d05f3d9857149d4dcacf8aa09e74adb3cc88 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:00:08 -0700 Subject: [PATCH 13/24] [None][test] Use Nemotron Super dynamic tree test Signed-off-by: qgai --- .../_torch/speculative/test_eagle3.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 452fd5b289dc..16c797a69fa4 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -1150,6 +1150,99 @@ def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, assert len(results[0].outputs[0].token_ids) > 0 +@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.high_cuda_memory +@skip_pre_blackwell +def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( + use_cuda_graph: bool, disable_overlap_scheduler: bool): + if torch.cuda.device_count() < 8: + pytest.skip("Nemotron Super dynamic-tree MTP test requires 8 GPUs") + + models_path = llm_models_root() + model_path = f"{models_path}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" + + max_batch_size = 1 + max_draft_len = 6 + dynamic_tree_max_topK = 10 + max_total_draft_tokens = 31 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + mamba_ssm_cache_dtype="float16", + free_gpu_memory_fraction=0.8) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=model_path, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + pipeline_parallel_size=1, + moe_config=MoeConfig(backend="TRTLLM"), + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + ) + + spec_config = MTPDecodingConfig( + max_draft_len=max_draft_len, + mtp_eagle_one_model=True, + use_dynamic_tree=True, + dynamic_tree_max_topK=dynamic_tree_max_topK, + max_total_draft_tokens=max_total_draft_tokens, + ) + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + sampling_params = SamplingParams(max_tokens=128, temperature=0) + num_tokens = 0 + num_drafted = 0 + num_accepted = 0 + tok_ids = llm_spec.tokenizer.encode("The future of AI is") + + for output in llm_spec.generate_async(tok_ids, + sampling_params, + streaming=True): + new_tokens = output.outputs[0].token_ids + num_drafted += max_draft_len + num_accepted += len(new_tokens) - num_tokens - 1 + num_tokens = len(new_tokens) + + accept_rate = num_accepted / num_drafted + assert accept_rate > 0.20 + + raw_prompts = ["The capital of France is"] + prompts = [ + llm_spec.tokenizer.apply_chat_template( + [{ + "role": "user", + "content": p + }], + tokenize=False, + add_generation_prompt=True, + ) for p in raw_prompts + ] + sampling_params = SamplingParams(max_tokens=10, temperature=0) + + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + assert text_spec == text_ref + + +if __name__ == "__main__": + unittest.main() + + + @pytest.mark.parametrize("use_cuda_graph", [True, False]) def test_eagle3_lora(use_cuda_graph: bool): """Test LoRA with 3 requests and max_batch_size=4. From a648cff8aa90f7df277500421bbd3c24bed52681 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:20:23 -0700 Subject: [PATCH 14/24] [None][test] Restore MTP test coverage Signed-off-by: qgai --- .../_torch/speculative/test_eagle3.py | 47 ++++++++----------- 1 file changed, 19 insertions(+), 28 deletions(-) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 16c797a69fa4..1bd3ed592c9b 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -11,7 +11,8 @@ import torch from test_common.llm_data import with_mocked_hf_download_for_single_gpu from utils.llm_data import llm_models_root -from utils.util import skip_blackwell, skip_num_gpus_less_than +from utils.util import (skip_blackwell, skip_num_gpus_less_than, + skip_pre_blackwell) from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata @@ -25,7 +26,7 @@ from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, - KvCacheConfig) + KvCacheConfig, MoeConfig, MTPDecodingConfig) from tensorrt_llm.lora_helper import LoraConfig sys.path.append(os.path.join(os.path.dirname(__file__), '..')) @@ -1164,8 +1165,6 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( max_batch_size = 1 max_draft_len = 6 - dynamic_tree_max_topK = 10 - max_total_draft_tokens = 31 kv_cache_config = KvCacheConfig(enable_block_reuse=False, mamba_ssm_cache_dtype="float16", free_gpu_memory_fraction=0.8) @@ -1184,23 +1183,27 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( kv_cache_config=kv_cache_config, max_seq_len=8192, ) - - spec_config = MTPDecodingConfig( - max_draft_len=max_draft_len, - mtp_eagle_one_model=True, - use_dynamic_tree=True, - dynamic_tree_max_topK=dynamic_tree_max_topK, - max_total_draft_tokens=max_total_draft_tokens, - ) + spec_config = MTPDecodingConfig(max_draft_len=max_draft_len, + mtp_eagle_one_model=True, + use_dynamic_tree=True, + dynamic_tree_max_topK=10, + max_total_draft_tokens=31) llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + prompt = llm_spec.tokenizer.apply_chat_template( + [{ + "role": "user", + "content": "The future of AI is" + }], + tokenize=False, + add_generation_prompt=True, + ) + tok_ids = llm_spec.tokenizer.encode(prompt) sampling_params = SamplingParams(max_tokens=128, temperature=0) num_tokens = 0 num_drafted = 0 num_accepted = 0 - tok_ids = llm_spec.tokenizer.encode("The future of AI is") - for output in llm_spec.generate_async(tok_ids, sampling_params, streaming=True): @@ -1212,25 +1215,13 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( accept_rate = num_accepted / num_drafted assert accept_rate > 0.20 - raw_prompts = ["The capital of France is"] - prompts = [ - llm_spec.tokenizer.apply_chat_template( - [{ - "role": "user", - "content": p - }], - tokenize=False, - add_generation_prompt=True, - ) for p in raw_prompts - ] sampling_params = SamplingParams(max_tokens=10, temperature=0) - - results_spec = llm_spec.generate(prompts, sampling_params) + results_spec = llm_spec.generate([prompt], sampling_params) generated_text_spec = [result.outputs[0].text for result in results_spec] llm_spec.shutdown() llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) + results_ref = llm_ref.generate([prompt], sampling_params) generated_text_ref = [result.outputs[0].text for result in results_ref] llm_ref.shutdown() From 10cdcc45a47edb3c3db296234ebb503b60103509 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:22:53 -0700 Subject: [PATCH 15/24] [None][test] Keep dynamic tree test diff minimal Signed-off-by: qgai --- .../speculative/hw_agnostic/test_mtp.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index a537e38b97df..a381dbd94907 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -11,17 +11,12 @@ from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker -from tensorrt_llm._torch.speculative.utils import update_spec_config_from_model_config from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) from utils.llm_data import llm_models_root -class _MTPPretrainedConfig: - num_nextn_predict_layers = 1 - - def unittest_name_func(testcase_func, param_num, param): name = param.args[0] return "%s_%s" % ( @@ -30,21 +25,6 @@ def unittest_name_func(testcase_func, param_num, param): ) -def test_mtp_dynamic_tree_preserves_max_total_draft_tokens(): - spec_config = MTPDecodingConfig( - max_draft_len=6, - max_total_draft_tokens=15, - use_dynamic_tree=True, - dynamic_tree_max_topK=4, - ) - - update_spec_config_from_model_config(spec_config, _MTPPretrainedConfig()) - - assert spec_config.max_draft_len == 6 - assert spec_config.max_total_draft_tokens == 15 - assert spec_config.tokens_per_gen_step == 16 - - class TestMTPSampleAndAcceptDraftTokens(unittest.TestCase): def setUp(self): tensorrt_llm.logger.set_level("warning") From f4cda88979009c9f2dfc8481eab44f0f59a2d4ad Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:30:22 -0700 Subject: [PATCH 16/24] [None][chore] Trim spec tree manager diff Signed-off-by: qgai --- .../_torch/speculative/spec_tree_manager.py | 86 +++++++++---------- 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 061086c19a02..63c031b42a89 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -1,5 +1,6 @@ import logging import math +from itertools import accumulate from typing import List import torch @@ -202,11 +203,11 @@ class SpecTreeManager: # Auxiliary buffers # The top k list for each draft layer. - top_k_list = [] + top_k_list: list # The user input eagle choices, only available when using static tree. - eagle_choices: List[List[int]] = None - # If dynamice tree, each request has their own tree. If static tree, all requests share the same tree. - num_trees: int = None + eagle_choices: List[List[int]] + # If dynamic tree, each request has their own tree. If static tree, all requests share the same tree. + num_trees: int # Convert the choice to a path. Each path is an array of indices from the root to other nodes in the tree. # shape: [num_trees, max_total_draft_tokens + 1, max_draft_len + 1] @@ -224,40 +225,39 @@ class SpecTreeManager: # shape: [num_trees, max_total_draft_tokens + 1], device tensor. spec_dec_position_offsets: torch.Tensor = None - # TODO: Optimized together with the subsequent dynamic tree. - # Auxiliary buffers for the static tree. + ############################ Auxiliary buffers for the static tree. ############################ # Considering that the static tree does not modify the tree structure during inference, we can calculate some buffers in advance. # NOTE: Most of these buffers are introduced due to limitations of XQA: # With tree attention, XQA cannot simply take the tokens to be processed in the next round as input. Instead, it needs to take ALL of their parent nodes as input. # This incurs additional computation, but it is unavoidable. # NOTE: The reason why most of these auxiliary buffers are with `len == max_draft_len - 1` is that: we do not need to prepare specific input data for the first draft layer. - # The top k value for each draft layer. Device tensor. top_k_list_cuda: list[torch.Tensor] = None # The max top k value for all draft layers. Which is used for torch.topk and cuda graph. max_top_k = -1 - # Gather the required draft tokens from all currently generated draft tokens as the input of the next draft layer. + # Gather the required draft tokens among the 'max_total_draft_tokens + 1' tokens. # Only the nodes has child(s) this layer and all their parents nodes will be gathered. - # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. Each element is a tensor with shape [num_tokens_for_next_layer]. tokens_gather_idx_for_drafter_model: list[torch.Tensor] = None - # Gather the required logits from all currently generated logits. - # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. - logits_gather_idx: list[torch.Tensor] = None - # The packed mask for the drafter model's attention (i.e., xqa). + # shape: [1, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)], device tensor. spec_dec_packed_mask_for_drafter_model: torch.Tensor = None # The read indices offset for the drafter model. + # shape: [max_total_draft_tokens + 1], device tensor. hidden_states_read_indices_offset_for_drafter_model: torch.Tensor = None # The write back start indices for the drafter tokens between different draft layers. + # shape: [max_draft_len + 1], device tensor. draft_tokens_indices_cumsum: torch.Tensor = None - # Work buffers for dynamic tree build kernel output + ############################ Auxiliary buffers for the dynamic tree. ############################ + # CUDA kernel outputs for dynamic tree verification. + # These are produced by build_dynamic_tree CUDA kernel and used by verify_dynamic_tree_greedy. + # shape: [num_trees, max_total_draft_tokens + 1], int32, device tensor. retrieve_index: torch.Tensor = None retrieve_next_token: torch.Tensor = None retrieve_next_sibling: torch.Tensor = None @@ -305,14 +305,17 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, device='cuda', ).unsqueeze(0).repeat(self.num_trees, 1, 1) - n_dt = self.max_total_draft_tokens + 1 + # CUDA kernel facing — rows = max_total_draft_tokens + 1, + # columns widened to match attn_metadata mask_width so that the + # Hopper flat copy in update_spec_dec_param needs no per-row padding. self.spec_dec_packed_mask = torch.zeros( - (self.num_trees, n_dt, math.ceil(n_dt / 32)), + (self.num_trees, self.max_total_draft_tokens + 1, + math.ceil(self._internal_buf_dim / 32)), dtype=torch.int32, device='cuda', ) self.spec_dec_position_offsets = torch.zeros( - (self.num_trees, n_dt), + (self.num_trees, self.max_total_draft_tokens + 1), dtype=torch.int32, device='cuda', ) @@ -340,16 +343,8 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, self.init_tree_info_for_static_tree() def init_tree_info_for_dynamic_tree(self): + # Allocate retrieve buffers for CUDA kernel outputs num_draft_with_root = self.max_total_draft_tokens + 1 - - self.top_k_list = [ - torch.ones(self.dynamic_tree_max_topK, - dtype=torch.int32, - device='cpu', - pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK - ] - - # Work buffers for build_dynamic_tree kernel output self.retrieve_index = torch.zeros((self.num_trees, num_draft_with_root), dtype=torch.int32, device='cuda') @@ -364,11 +359,19 @@ def init_tree_info_for_dynamic_tree(self): dtype=torch.int32, device='cuda') - mask_width = math.ceil(num_draft_with_root / 32) + # For the dynamic tree + # To the internal layer, the number of nodes is the same as the dynamic_tree_max_topK. + self.top_k_list = [ + torch.ones(self.dynamic_tree_max_topK, + dtype=torch.int32, + device='cpu', + pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK + ] + self.slot_storage = DynamicTreeSlotStorage( num_slots=self.num_trees, n_dt=num_draft_with_root, - mask_width=mask_width, + mask_width=self.spec_dec_packed_mask.shape[-1], top_k=self.dynamic_tree_max_topK, ) @@ -387,7 +390,6 @@ def scatter_to_slot_storage(self, ss, gen_slots, num_gens): 0, ids, self.retrieve_next_sibling[:num_gens]) ss.mark_valid(ids, num_gens) - # For the static tree def init_tree_info_for_static_tree(self): self.index_mapping_set = {} self.nodes_list_per_layer = [[] for _ in range(self.max_draft_len + 1)] @@ -436,9 +438,7 @@ def init_tree_info_for_static_tree(self): pin_memory=prefer_pinned())) # 6) Compute the spec decoding according to the eagle_paths for the target model - for i, path in enumerate(self.eagle_paths[0]): - indices = path[path > -1] - self.spec_dec_mask_matrix[0][i, indices] = 1 + self.compute_spec_dec_mask_matrix(0) self.compute_spec_dec_packed_mask(self.spec_dec_mask_matrix, self.spec_dec_packed_mask) @@ -477,7 +477,6 @@ def init_tree_info_for_static_tree(self): num_nodes_per_layer = [0] num_nodes_per_layer.extend( [len(node_list) for node_list in self.nodes_list_per_layer[1:]]) - from itertools import accumulate self.draft_tokens_indices_cumsum = torch.tensor(list( accumulate(num_nodes_per_layer)), dtype=torch.int32, @@ -528,6 +527,15 @@ def get_top_k_list(self, draft_layer_id): assert draft_layer_id >= 0 return self.top_k_list[draft_layer_id] + def compute_spec_dec_mask_matrix(self, tree_idx=0): + if self.eagle_paths is None: + raise RuntimeError( + "compute_spec_dec_mask_matrix() is not supported in dynamic tree mode" + ) + for i, path in enumerate(self.eagle_paths[0]): + indices = path[path > -1] + self.spec_dec_mask_matrix[0][i, indices] = 1 + def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask): bs, num_tokens, num_tokens_attend = mask_matrix.shape assert mask_matrix.ndim == 3, f"Expected 3D mask_matrix, got {mask_matrix.ndim}D" @@ -538,21 +546,13 @@ def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask): # Use cached bit weights weights = self._pack_weights - src = mask_matrix if mask_matrix.dtype == torch.int32 else mask_matrix.to( - torch.int32) - - if num_blocks == 1 and num_tokens_attend <= 32: - result = self._pack_result_buf[:bs, :num_tokens, :1] - torch.sum(src * weights[:num_tokens_attend], - dim=-1, - out=result[:, :, 0]) - packed_mask[:, :num_tokens, :1] = result - return packed_mask # Pad into pre-allocated buffer total_bits = num_blocks * 32 padded_m = self._padded_mask_buf[:bs, :num_tokens, :total_bits] padded_m.zero_() + src = mask_matrix if mask_matrix.dtype == torch.int32 else mask_matrix.to( + torch.int32) padded_m[:, :, :num_tokens_attend].copy_(src) # Reshape last dim into [num_blocks, 32] for blocked packing From 6bedd4a5d2fd474035ea69afb856df47af129fe8 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:39:35 -0700 Subject: [PATCH 17/24] [None][chore] Align spec tree diff with latest main Signed-off-by: qgai --- .../_torch/speculative/spec_tree_manager.py | 61 +++++++++---------- 1 file changed, 28 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index 63c031b42a89..dc88d288251b 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -170,16 +170,9 @@ def pack_retrieve_from_slots(self, slot_ids, count): return self._verify_staging[:0] ids = slot_ids[:count] staging = self._verify_staging[:count] - # Avoid advanced-indexing temporaries on the verification path. - torch.index_select(self.retrieve_index, 0, ids, out=staging[:, :, 0]) - torch.index_select(self.retrieve_next_token, - 0, - ids, - out=staging[:, :, 1]) - torch.index_select(self.retrieve_next_sibling, - 0, - ids, - out=staging[:, :, 2]) + staging[:, :, 0] = self.retrieve_index[ids] + staging[:, :, 1] = self.retrieve_next_token[ids] + staging[:, :, 2] = self.retrieve_next_sibling[ids] return staging def next_links_from_slots(self, slot_ids, count): @@ -254,10 +247,7 @@ class SpecTreeManager: # shape: [max_draft_len + 1], device tensor. draft_tokens_indices_cumsum: torch.Tensor = None - ############################ Auxiliary buffers for the dynamic tree. ############################ - # CUDA kernel outputs for dynamic tree verification. - # These are produced by build_dynamic_tree CUDA kernel and used by verify_dynamic_tree_greedy. - # shape: [num_trees, max_total_draft_tokens + 1], int32, device tensor. + # Work buffers for dynamic tree build kernel output retrieve_index: torch.Tensor = None retrieve_next_token: torch.Tensor = None retrieve_next_sibling: torch.Tensor = None @@ -305,17 +295,14 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, device='cuda', ).unsqueeze(0).repeat(self.num_trees, 1, 1) - # CUDA kernel facing — rows = max_total_draft_tokens + 1, - # columns widened to match attn_metadata mask_width so that the - # Hopper flat copy in update_spec_dec_param needs no per-row padding. + n_dt = self.max_total_draft_tokens + 1 self.spec_dec_packed_mask = torch.zeros( - (self.num_trees, self.max_total_draft_tokens + 1, - math.ceil(self._internal_buf_dim / 32)), + (self.num_trees, n_dt, math.ceil(n_dt / 32)), dtype=torch.int32, device='cuda', ) self.spec_dec_position_offsets = torch.zeros( - (self.num_trees, self.max_total_draft_tokens + 1), + (self.num_trees, n_dt), dtype=torch.int32, device='cuda', ) @@ -343,8 +330,16 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, self.init_tree_info_for_static_tree() def init_tree_info_for_dynamic_tree(self): - # Allocate retrieve buffers for CUDA kernel outputs num_draft_with_root = self.max_total_draft_tokens + 1 + + self.top_k_list = [ + torch.ones(self.dynamic_tree_max_topK, + dtype=torch.int32, + device='cpu', + pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK + ] + + # Work buffers for build_dynamic_tree kernel output self.retrieve_index = torch.zeros((self.num_trees, num_draft_with_root), dtype=torch.int32, device='cuda') @@ -359,19 +354,11 @@ def init_tree_info_for_dynamic_tree(self): dtype=torch.int32, device='cuda') - # For the dynamic tree - # To the internal layer, the number of nodes is the same as the dynamic_tree_max_topK. - self.top_k_list = [ - torch.ones(self.dynamic_tree_max_topK, - dtype=torch.int32, - device='cpu', - pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK - ] - + mask_width = math.ceil(num_draft_with_root / 32) self.slot_storage = DynamicTreeSlotStorage( num_slots=self.num_trees, n_dt=num_draft_with_root, - mask_width=self.spec_dec_packed_mask.shape[-1], + mask_width=mask_width, top_k=self.dynamic_tree_max_topK, ) @@ -546,13 +533,21 @@ def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask): # Use cached bit weights weights = self._pack_weights + src = mask_matrix if mask_matrix.dtype == torch.int32 else mask_matrix.to( + torch.int32) + + if num_blocks == 1 and num_tokens_attend <= 32: + result = self._pack_result_buf[:bs, :num_tokens, :1] + torch.sum(src * weights[:num_tokens_attend], + dim=-1, + out=result[:, :, 0]) + packed_mask[:, :num_tokens, :1] = result + return packed_mask # Pad into pre-allocated buffer total_bits = num_blocks * 32 padded_m = self._padded_mask_buf[:bs, :num_tokens, :total_bits] padded_m.zero_() - src = mask_matrix if mask_matrix.dtype == torch.int32 else mask_matrix.to( - torch.int32) padded_m[:, :, :num_tokens_attend].copy_(src) # Reshape last dim into [num_blocks, 32] for blocked packing From 7239241eacbc5105568444989baf74e0f080eda9 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 29 Jun 2026 01:50:16 -0700 Subject: [PATCH 18/24] [None][refactor] Simplify dynamic tree metadata helper Signed-off-by: qgai --- .../_torch/speculative/spec_tree_manager.py | 32 +++---------------- 1 file changed, 5 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index dc88d288251b..545b5bc7bb06 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -26,8 +26,8 @@ def __init__(self, # Bootstrap/reused slots may not have a tree yet; keep their metadata # as a valid linear chain so verification kernels can read it directly. - no_tree_position_offsets, no_tree_packed_mask = self._make_no_tree_metadata( - n_dt, mask_width) + no_tree_position_offsets, no_tree_packed_mask = self._make_kary_tree_metadata( + n_dt, mask_width, top_k=1) self.position_offsets = no_tree_position_offsets.unsqueeze(0).repeat( S, 1).contiguous() self.packed_mask = no_tree_packed_mask.unsqueeze(0).repeat( @@ -37,7 +37,7 @@ def __init__(self, # CUDA-graph dummies use a deterministic K-ary tree, matching real # dynamic-tree mask/position shapes without depending on request state. - dummy_position_offsets, dummy_packed_mask = self._make_dummy_tree_metadata( + dummy_position_offsets, dummy_packed_mask = self._make_kary_tree_metadata( n_dt, mask_width, top_k) self.position_offsets[self.dummy_slot_id] = dummy_position_offsets self.packed_mask[self.dummy_slot_id] = dummy_packed_mask @@ -74,30 +74,8 @@ def __init__(self, device='cuda') @staticmethod - def _pack_bool_mask(mask: torch.Tensor, mask_width: int) -> torch.Tensor: - """Pack a bool attention mask into int32 words.""" - num_rows, num_bits = mask.shape - total_bits = mask_width * 32 - padded_mask = torch.zeros((num_rows, total_bits), - dtype=torch.bool, - device=mask.device) - padded_mask[:, :num_bits] = mask - weights = 2**torch.arange(32, dtype=torch.int64, device=mask.device) - return (padded_mask.view(num_rows, mask_width, 32).to(torch.int64) * - weights).sum(-1).to(torch.int32) - - @classmethod - def _make_no_tree_metadata( - cls, n_dt: int, - mask_width: int) -> tuple[torch.Tensor, torch.Tensor]: - token_ids = torch.arange(n_dt, device='cuda') - position_offsets = token_ids.to(torch.int32) - causal_mask = token_ids.unsqueeze(1) >= token_ids.unsqueeze(0) - return position_offsets, cls._pack_bool_mask(causal_mask, mask_width) - - @classmethod - def _make_dummy_tree_metadata( - cls, n_dt: int, mask_width: int, + def _make_kary_tree_metadata( + n_dt: int, mask_width: int, top_k: int) -> tuple[torch.Tensor, torch.Tensor]: top_k = max(int(top_k), 1) token_ids = torch.arange(n_dt, device='cuda') From d45200674f4c9bf3cca7179eeeeb2f91124e024a Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 9 Jul 2026 01:18:01 -0700 Subject: [PATCH 19/24] [None][chore] Remove unused imports Signed-off-by: qgai --- tests/unittest/_torch/speculative/test_eagle3.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 1bd3ed592c9b..61d1ed024cb0 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -22,8 +22,6 @@ from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ _extend_full_attention_windows_for_spec_decode from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig) From 1bd96b6f184be86a4b04e5c3b0bd49c8330c7d1d Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 19:15:51 -0700 Subject: [PATCH 20/24] [None][doc] Note MTP dynamic tree is greedy-only in feature matrix Signed-off-by: qgai --- docs/source/features/feature-combination-matrix.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/features/feature-combination-matrix.md b/docs/source/features/feature-combination-matrix.md index b56c1b219af9..d91322fc8c11 100644 --- a/docs/source/features/feature-combination-matrix.md +++ b/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 | --- | | | | From 383abc22a043c046a15f3700bfd9c32ff2ef9fa2 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 13 Jul 2026 19:41:05 -0700 Subject: [PATCH 21/24] [None][fix] Require explicit use_dynamic_tree opt-in for dynamic tree Signed-off-by: qgai --- tensorrt_llm/llmapi/llm_args.py | 33 ++++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a99fe45b812f..a61f9c4e16af 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1971,8 +1971,10 @@ class EagleDecodingConfig(DecodingBaseConfig): ) dynamic_tree_max_topK: Optional[int] = Field( default=None, - description="The topK value for each layer when dynamic tree is enabled." - ) + description= + "The topK value for each layer when dynamic tree is enabled. Required " + "when use_dynamic_tree is True; ignored (with a warning) when " + "use_dynamic_tree is False.") num_eagle_layers: Optional[int] = Field( default=None, description= @@ -2048,9 +2050,17 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': # So the number of choices also represents the number of max draft nodes. self.max_total_draft_tokens = len(self.eagle_choices) + # Dynamic tree is enabled only by an explicit use_dynamic_tree=True; + # dynamic_tree_max_topK alone does not turn it on. + if not self.use_dynamic_tree and self.dynamic_tree_max_topK is not None: + logger.warning( + "dynamic_tree_max_topK is set but use_dynamic_tree is False; " + "ignoring dynamic_tree_max_topK and using the linear draft path." + ) + self.dynamic_tree_max_topK = None + # Dynamic tree logic - if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: - self.use_dynamic_tree = True + if self.use_dynamic_tree: if self.eagle_choices is not None: raise ValueError( "If use_dynamic_tree is True, eagle_choices should be None") @@ -2435,7 +2445,8 @@ class MTPDecodingConfig(DecodingBaseConfig): default=None, description= "Top-K candidates expanded per node per draft layer when use_dynamic_tree " - "is enabled. Required when use_dynamic_tree is True.") + "is enabled. Required when use_dynamic_tree is True; ignored (with a " + "warning) when use_dynamic_tree is False.") # Internal max batch size for dynamic-tree worker buffers. _max_batch_size: Optional[int] = PrivateAttr(default=None) @@ -2489,9 +2500,17 @@ def set_max_total_draft_tokens(self): if self.max_draft_len <= 0: raise ValueError("max_draft_len must be > 0 for MTP") + # Dynamic tree is enabled only by an explicit use_dynamic_tree=True; + # dynamic_tree_max_topK alone does not turn it on. + if not self.use_dynamic_tree and self.dynamic_tree_max_topK is not None: + logger.warning( + "dynamic_tree_max_topK is set but use_dynamic_tree is False; " + "ignoring dynamic_tree_max_topK and using the linear draft path." + ) + self.dynamic_tree_max_topK = None + # Dynamic tree defaults max_total_draft_tokens to topK * max_draft_len. - if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: - self.use_dynamic_tree = True + if self.use_dynamic_tree: if self.max_draft_len is None: raise ValueError( "max_draft_len must be set when use_dynamic_tree is True") From dc165b1cf12bc86da74695252df1c07fd572de34 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 15 Jul 2026 00:37:39 -0700 Subject: [PATCH 22/24] [None][refactor] Collapse dynamic tree range checks into one condition Signed-off-by: qgai --- tensorrt_llm/llmapi/llm_args.py | 32 ++++++++++++++------------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a61f9c4e16af..4ec8383b5a51 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2080,16 +2080,12 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': logger.warning( f"max_total_draft_tokens is not provided, use the default value {default_max_total_draft_tokens} (default_max_total_draft_tokens = dynamic_tree_max_topK * max_draft_len)" ) - else: - if self.max_total_draft_tokens < self.max_draft_len: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) should be >= max_draft_len ({self.max_draft_len})" - ) - if self.max_total_draft_tokens > self.dynamic_tree_max_topK * self.max_draft_len: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) should be <= " - f"dynamic_tree_max_topK * max_draft_len ({self.dynamic_tree_max_topK * self.max_draft_len})" - ) + elif not (self.max_draft_len <= self.max_total_draft_tokens <= + default_max_total_draft_tokens): + raise ValueError( + f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be in " + f"[max_draft_len ({self.max_draft_len}), dynamic_tree_max_topK * " + f"max_draft_len ({default_max_total_draft_tokens})]") # Linear tree if self.max_total_draft_tokens is None: @@ -2521,15 +2517,15 @@ def set_max_total_draft_tokens(self): default_max_total_draft_tokens = self.dynamic_tree_max_topK * self.max_draft_len if self.max_total_draft_tokens is None: self.max_total_draft_tokens = default_max_total_draft_tokens - elif self.max_total_draft_tokens < self.max_draft_len: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be >= " - f"max_draft_len ({self.max_draft_len})") - elif self.max_total_draft_tokens > default_max_total_draft_tokens: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be <= " - f"dynamic_tree_max_topK * max_draft_len ({default_max_total_draft_tokens})" + logger.warning( + f"max_total_draft_tokens is not provided, use the default value {default_max_total_draft_tokens} (default_max_total_draft_tokens = dynamic_tree_max_topK * max_draft_len)" ) + elif not (self.max_draft_len <= self.max_total_draft_tokens <= + default_max_total_draft_tokens): + raise ValueError( + f"max_total_draft_tokens ({self.max_total_draft_tokens}) must be in " + f"[max_draft_len ({self.max_draft_len}), dynamic_tree_max_topK * " + f"max_draft_len ({default_max_total_draft_tokens})]") elif self.max_draft_len is not None: self.max_total_draft_tokens = self.max_draft_len # linear chain return self From 024d23754dc49fe1587e5567461a5c52aa1f1cee Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 16 Jul 2026 23:21:48 -0700 Subject: [PATCH 23/24] [None][chore] Drop churn leftovers in eagle3 test after rebase Signed-off-by: qgai --- .../_torch/speculative/test_eagle3.py | 76 ------------------- 1 file changed, 76 deletions(-) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 61d1ed024cb0..33cf3d635661 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -283,77 +283,6 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): assert draft_kwargs["max_blocks"] is None -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_mtp_eagle_one_model_dynamic_tree_metadata_prepares_mamba_links(): - max_num_requests = 3 - max_draft_len = 2 - max_total_draft_tokens = 2 - spec_tree_manager = SpecTreeManager( - max_num_requests=max_num_requests, - use_dynamic_tree=True, - max_total_draft_tokens=max_total_draft_tokens, - max_draft_len=max_draft_len, - eagle_choices=None, - dynamic_tree_max_topK=2, - ) - slot_storage = spec_tree_manager.slot_storage - slot_storage.all_ids_buf[:max_num_requests].copy_( - torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) - slot_storage.has_tree[1] = True - slot_storage.retrieve_next_token[1] = torch.tensor([2, -1, -1], - dtype=torch.int32, - device="cuda") - slot_storage.retrieve_next_sibling[1] = torch.tensor([-1, -1, -1], - dtype=torch.int32, - device="cuda") - - class _ResourceManager: - hidden_states = None - slot_manager = None - sa_manager = None - - def __init__(self): - self.spec_tree_manager = spec_tree_manager - self.batch_indices_cuda = torch.empty(max_num_requests, - dtype=torch.int, - device="cuda") - - metadata = Eagle3OneModelSpecMetadata( - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - spec_dec_mode=SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL, - max_num_requests=max_num_requests, - num_layers=1, - hidden_size=1, - max_num_tokens=16, - spec_resource_manager=_ResourceManager(), - use_dynamic_tree=True, - ) - metadata.request_ids = [10, 11, 12] - metadata.seq_lens = [ - 1, max_total_draft_tokens + 1, max_total_draft_tokens + 1 - ] - metadata.num_generations = 2 - metadata.num_tokens = 1 + 2 * (max_total_draft_tokens + 1) - - metadata.prepare() - - assert metadata.retrieve_next_token is not None - assert metadata.retrieve_next_sibling is not None - assert metadata.retrieve_next_token.shape == (2, max_total_draft_tokens + 1) - assert torch.equal(metadata.retrieve_next_token[0], - slot_storage.retrieve_next_token[1]) - assert torch.equal( - metadata.retrieve_next_token[1], - torch.tensor([1, 2, -1], dtype=torch.int32, device="cuda")) - assert torch.equal( - metadata.retrieve_next_sibling[1], - torch.full((max_total_draft_tokens + 1, ), - -1, - dtype=torch.int32, - device="cuda")) - - @pytest.mark.parametrize( "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model", [ @@ -1227,11 +1156,6 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( assert text_spec == text_ref -if __name__ == "__main__": - unittest.main() - - - @pytest.mark.parametrize("use_cuda_graph", [True, False]) def test_eagle3_lora(use_cuda_graph: bool): """Test LoRA with 3 requests and max_batch_size=4. From 9bd515a6bf140300b8d4dcb2122ef2f2c558b3d7 Mon Sep 17 00:00:00 2001 From: qgai Date: Sun, 19 Jul 2026 18:41:55 -0700 Subject: [PATCH 24/24] [None][chore] Apply ruff-format Signed-off-by: qgai --- tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 14ae6972da8e..5dca84abae8e 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -55,10 +55,7 @@ def __init__( *, mapping: Optional[Mapping] = None, ): - super().__init__(spec_config, - model_config, - use_separate_draft_kv_cache, - mapping=mapping) + super().__init__(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) assert getattr(spec_config, "use_dynamic_tree", False), ( "MTPEagleDynamicTreeWorker requires use_dynamic_tree=True" )