diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 55b88face9d4..134689192b27 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -6,8 +6,8 @@ from collections import namedtuple from dataclasses import dataclass, field from enum import Enum, IntEnum -from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Literal, Optional, - Protocol, Tuple, Type, TypeVar, Union) +from typing import (TYPE_CHECKING, Any, ClassVar, Dict, Generic, List, Literal, + Optional, Protocol, Tuple, Type, TypeVar, Union) import torch from typing_extensions import Self @@ -92,6 +92,11 @@ class AttentionMetadata: # Whether CUDA graph is enabled. is_cuda_graph: bool = field(default=False, repr=False) + # Whether on_update_kv_lens() is a pure function of the current seq/kv lengths and block + # offsets, so a second call fully overwrites the first. Backends opt in per class after + # auditing their hook; the model engine then skips a call that is repeated right after. + kv_lens_hook_idempotent: ClassVar[bool] = False + # The length of each sequence in the batch for query. # The shape is (batch_size), and located on CPU memory. # For sub metadata of cross attention, it's automatically diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 75c85356ef00..057fbd17f8a2 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -5,7 +5,7 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple +from typing import TYPE_CHECKING, ClassVar, Dict, Optional, Set, Tuple import torch @@ -28,6 +28,8 @@ class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): + # on_update_kv_lens() is overridden below and has not been audited for idempotence. + kv_lens_hook_idempotent: ClassVar[bool] = False # The set of compress ratios for the layers compress_ratio_set: Set[int] # The set of (compress ratio, attention type) for the layers diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index f16100b9f3bf..6653f457c289 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1179,6 +1179,7 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): """ num_contexts = metadata.num_contexts num_generations = metadata.num_generations + metadata.mtp_sched_metadata_stale = False gen_seq_lens = metadata.get_indexer_kv_lens( metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations] ) @@ -1683,6 +1684,18 @@ def sparse_attn_indexer( metadata.shared_topk_indices[:num_generations, :] ) elif has_decode and not metadata.skip_indexer_for_gen_reqs: + # getattr: test doubles for the metadata may predate this flag. + if getattr(metadata, "mtp_sched_metadata_stale", False): + # The MTP tail trim skipped the schedule rebuild (see + # DSAtrtllmAttentionMetadata._skip_indexer_sched_metadata) but + # this pass would read it: fail loudly. + raise RuntimeError( + "DSA indexer: MQA-logits scheduler metadata is stale " + "(skipped by the MTP tail trim) but this pass computes TopK " + f"(in_mtp_draft_loop={metadata.in_mtp_draft_loop}, " + f"indexer_skip_topk={metadata.indexer_skip_topk}, " + f"mtp_index_share={self.mtp_index_share})" + ) # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations] max_decode_len = gen_seq_lens.max().item() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index fb435e686e44..08d57f90bb67 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -7,7 +7,7 @@ import functools import os from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, ClassVar, List, Optional import torch @@ -72,6 +72,10 @@ def build_req_idx_per_token(seq_lens: torch.Tensor, num_tokens: int) -> torch.Te class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" + # on_update_kv_lens() rebuilds everything from the current seq/kv lengths and block offsets: + # idempotent, fully overwritten by the second call. + kv_lens_hook_idempotent: ClassVar[bool] = True + sparse_metadata_params: Optional[DSAMetadataParams] = None use_fp8_ds_mla: bool = field(default=False, init=False) # Store reference to indexer for preparation stage @@ -100,6 +104,16 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): indexer_skip_topk: bool = False in_mtp_draft_loop: bool = False mtp_num_accepted: Optional[torch.Tensor] = None + # MTP tail trim (TRTLLM_MTP_TAIL_TRIM). While set, on_update_kv_lens() leaves the DeepGEMM + # MQA-logits schedule, the 2D/expanded indexer KV lengths and the GVR row order untouched. + # Only the indexer's MQA-logits + TopK pass reads them, and the draft loop sets this exactly + # when that pass does not run before the next full rebuild: before an index-shared draft + # step (indexer_skip_topk, reuse-TopK path) and after the last draft step (the next target + # forward rebuilds in _preprocess_inputs). Slot mappings / indptrs are still refreshed. + _skip_indexer_sched_metadata: bool = False + # Records the skip until the next full rebuild so the indexer can assert it never reads the + # stale buffers. + mtp_sched_metadata_stale: bool = False # Whether skip the indexer for context requests skip_indexer_for_ctx_reqs: bool = False # Whether skip the indexer for generation requests @@ -500,6 +514,9 @@ def on_update_kv_lens(self) -> None: # transform_local_topk_and_prepare_pool_view() call. self._invalidate_pool_view_cache() + # MTP tail trim: leave the indexer's MQA-logits schedule as is (see the field). + skip_sched = self._skip_indexer_sched_metadata + # Default fused path: collapse the eager DSA decode-metadata chain # (req_idx_per_token + slot mappings + the two gen indptr cumsums) into # one fused Triton launch. Only for the pure-decode/generation step @@ -577,6 +594,8 @@ def on_update_kv_lens(self) -> None: dtype=torch.int64, out=self.gen_cached_token_indptr[1 : self.num_generations + 1], ) + self.mtp_sched_metadata_stale = skip_sched + if self.num_generations > 0 and not skip_sched: gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens) self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens @@ -636,7 +655,8 @@ def on_update_kv_lens(self) -> None: self._fused_dsa_meta_armed = True self._run_fused_dsa_decode_metadata() - self._compute_kv_lens_row_reorder() + if not skip_sched: + self._compute_kv_lens_row_reorder() self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) def _run_fused_dsa_decode_metadata(self): @@ -835,6 +855,8 @@ def create_buffers_for_indexer(self, capture_graph=False): self.indexer_skip_topk = False self.in_mtp_draft_loop = False self.mtp_num_accepted = None + self._skip_indexer_sched_metadata = False + self.mtp_sched_metadata_stale = False # Indexer metadata # Separate slot mappings for non-interleaved layout (flat byte indices) @@ -1171,6 +1193,19 @@ def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: def set_skip_topk(self, skip: bool) -> None: self.indexer_skip_topk = skip + def set_skip_indexer_sched_metadata(self, skip: bool) -> None: + self._skip_indexer_sched_metadata = skip + + def mtp_next_step_reuses_topk(self) -> bool: + """True when a following MTP draft step (which runs with + indexer_skip_topk) takes the indexer's reuse-TopK path instead of the + MQA-logits + TopK path.""" + return ( + self.in_mtp_draft_loop + and self.sparse_metadata_params.mtp_index_share + and self.shared_topk_indices is not None + ) + def set_in_mtp_draft_loop(self, active: bool) -> None: self.in_mtp_draft_loop = active diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 2c87709bdc5a..69f3c9cc6614 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -61,6 +61,22 @@ def set_skip_topk(self, skip: bool) -> None: """Ask the indexer to reuse its captured selection instead of scoring.""" +@runtime_checkable +class MTPIndexerScheduleTrimMetadata(Protocol): + """Index-sharing sparse metadata that can leave its indexer schedule untouched. + + Used with ``MTPIndexShareMetadata`` by the MTP draft loop (``TRTLLM_MTP_TAIL_TRIM``): + before a draft step whose indexer pass reuses the captured selection, and after + the last one, the schedule the scoring pass would need is not rebuilt. + """ + + def set_skip_indexer_sched_metadata(self, skip: bool) -> None: + """Make the next kv-lens update leave the indexer's scheduler metadata as is.""" + + def mtp_next_step_reuses_topk(self) -> bool: + """Whether the following draft step reuses the captured selection instead of scoring.""" + + def use_self_sampling_gvr( *, enable_heuristic_topk: bool, diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 6d20a9aa6394..d111694a594a 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -47,6 +47,9 @@ _MNNVL_ONE_SHOT_THRESHOLD_BYTES = 64 * 1024 * 8 * 2 +# Lamport buffers per MNNVL all-reduce workspace. +NUM_LAMPORT_BUFFERS = 3 + _thread_local = threading.local() # Mirrored by the autotune context so Dynamo can guard this branch without @@ -264,13 +267,66 @@ def get_or_scale_allreduce_mnnvl_workspace( buffer_size_bytes) +def _build_mnnvl_workspace(comm, mapping: Mapping, + buffer_size_bytes: int) -> _MnnvlWorkspace: + """Allocate and initialise one MNNVL all-reduce workspace over ``comm``. + + The workspace holds NUM_LAMPORT_BUFFERS Lamport buffers of ``buffer_size_bytes`` each. The + outcome is converged across ``comm``: every rank either returns a ready workspace or raises, + so no rank is left waiting on the handle exchange. ``comm`` stays owned by the caller. + """ + force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + use_fabric_handle = force_mn or mapping.is_multi_node() + candidate_workspace: Optional[_MnnvlWorkspace] = None + candidate_error: Optional[Exception] = None + try: + # Each workspace contains NUM_LAMPORT_BUFFERS buffers. + workspace_size_bytes = NUM_LAMPORT_BUFFERS * buffer_size_bytes + mcast_buf_handle = _make_mnnvl_mcast_buffer(comm, workspace_size_bytes, + mapping, use_fabric_handle) + # We use per FP32 element in the buffer for lamport sync + buffer = mcast_buf_handle.get_uc_buffer( + mapping.tp_rank, + (workspace_size_bytes // torch.float32.itemsize, ), + torch.float32, + 0, + ) + # Layout: [cur idx, dirty idx, bytes per buffer, dirty num stages, + # numBytesToClear[4], access count ptr]. Filled in by + # _initialize_allreduce_mnnvl_protocol once every rank has its buffers. + buffer_flags = torch.tensor( + [0] * 9, + dtype=torch.uint32, + device=torch.device("cuda", _mnnvl_device_index(mapping)), + ) + candidate_workspace = { + "handle": mcast_buf_handle, + "uc_buffer": buffer, + "buffer_flags": buffer_flags, + "buffer_size_bytes": buffer_size_bytes, + "comm": comm, + } + except Exception as error: + candidate_error = error + + if not _mnnvl_workspace_all_succeeded(comm, candidate_error is None): + raise RuntimeError( + "MNNVL workspace construction failed on at least one rank" + ) from candidate_error + if candidate_error is not None: + raise candidate_error + assert candidate_workspace is not None + # Also fences the handle exchange: every rank leaves this call knowing its peers have + # their buffers, so nobody signals through memory another rank has not initialised. + _initialize_allreduce_mnnvl_protocol(candidate_workspace) + return candidate_workspace + + def _get_or_scale_allreduce_mnnvl_workspace( mapping: Mapping, dtype: torch.dtype, buffer_size_bytes: Optional[int] = None) -> _MnnvlWorkspace: - NUM_LAMPORT_BUFFERS = 3 - # Use MNNVLAllReduce class to share across threads allreduce_mnnvl_workspaces = MNNVLAllReduce.allreduce_mnnvl_workspaces pending_comms = MNNVLAllReduce.allreduce_mnnvl_pending_comms @@ -282,8 +338,6 @@ def _get_or_scale_allreduce_mnnvl_workspace( # A safe method to get the element size of the dtype elem_size = torch.tensor([], dtype=dtype).element_size() - force_mn = os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" - use_fabric_handle = force_mn or mapping.is_multi_node() if mapping not in allreduce_mnnvl_workspaces or allreduce_mnnvl_workspaces[ mapping]["buffer_size_bytes"] < (buffer_size_bytes or 0): @@ -317,56 +371,41 @@ def _get_or_scale_allreduce_mnnvl_workspace( logger.debug( f"[MNNVL] Requested {req_buffer_size_bytes} bytes, is larger than the current workspace size. Scaling workspace for pp_rank {mapping.pp_rank}, tp_size {mapping.tp_size} from {allreduce_mnnvl_workspaces[mapping]['buffer_size_bytes']} to {buffer_size_bytes} bytes" ) - candidate_workspace: Optional[_MnnvlWorkspace] = None - candidate_error: Optional[Exception] = None - try: - # Each workspace contains NUM_LAMPORT_BUFFERS buffers. - workspace_size_bytes = NUM_LAMPORT_BUFFERS * buffer_size_bytes - mcast_buf_handle = _make_mnnvl_mcast_buffer(comm, - workspace_size_bytes, - mapping, - use_fabric_handle) - # We use per FP32 element in the buffer for lamport sync - buffer = mcast_buf_handle.get_uc_buffer( - mapping.tp_rank, - (workspace_size_bytes // torch.float32.itemsize, ), - torch.float32, - 0, - ) - # Layout: [cur idx, dirty idx, bytes per buffer, dirty num stages, - # numBytesToClear[4], access count ptr]. Filled in by - # _initialize_allreduce_mnnvl_protocol once every rank has its buffers. - buffer_flags = torch.tensor( - [0] * 9, - dtype=torch.uint32, - device=torch.device("cuda", _mnnvl_device_index(mapping)), - ) - candidate_workspace = { - "handle": mcast_buf_handle, - "uc_buffer": buffer, - "buffer_flags": buffer_flags, - "buffer_size_bytes": buffer_size_bytes, - "comm": comm, - } - except Exception as error: - candidate_error = error - - if not _mnnvl_workspace_all_succeeded(comm, candidate_error is None): - raise RuntimeError( - "MNNVL workspace construction failed on at least one rank" - ) from candidate_error - if candidate_error is not None: - raise candidate_error - assert candidate_workspace is not None - # Also fences the handle exchange: every rank leaves this call knowing its peers have - # their buffers, so nobody signals through memory another rank has not initialised. - _initialize_allreduce_mnnvl_protocol(candidate_workspace) + candidate_workspace = _build_mnnvl_workspace(comm, mapping, + buffer_size_bytes) # Hand ownership of the communicator to the workspace. pending_comms.pop(mapping, None) allreduce_mnnvl_workspaces[mapping] = candidate_workspace return allreduce_mnnvl_workspaces[mapping] +def _launch_mnnvl_allreduce(input: torch.Tensor, workspace: _MnnvlWorkspace, + dtype: torch.dtype, + params: AllReduceParams) -> List[torch.Tensor]: + """Launch the MNNVL (fusion) all-reduce kernel on ``input`` over ``workspace``. + + ``params.fusion_op`` NONE is a plain SUM; the fused variants apply the RMSNorm epilogue. + Returns the kernel's output list (one tensor for a plain SUM). + """ + # We don't expect the buffer to be directly used in this level. The tensor is only used for passing the pointer to the kernel + buffer_base = workspace["uc_buffer"].view(dtype).view( + NUM_LAMPORT_BUFFERS, -1) + # The buffer flags is tied to the buffer and used to save the state of the buffer + buffer_flags = workspace["buffer_flags"] + fusion_op = params.fusion_op + return torch.ops.trtllm.mnnvl_fusion_allreduce( + input, + params.norm_weight, # gamma + params.residual, # residual + params.eps, # epsilon + buffer_base, # comm_buffer + buffer_flags, # buffer_flags + fusion_op != AllReduceFusionOp.NONE, # rmsnorm_fusion + params.scale, # scale + int(fusion_op), + ) + + def userbuffers_allreduce_finalize( input: torch.Tensor, force_applying_finalize: bool = False) -> torch.Tensor: @@ -947,26 +986,12 @@ def forward( buffer_size_bytes=workspace_size_bytes, ) - # We don't expect the buffer to be directly used in this level. The tensor is only used for passing the pointer to the kernel - buffer_base = workspace["uc_buffer"].view(self.dtype).view(3, -1) - # The buffer flags is tied to the buffer and used to save the state of the buffer - buffer_flags = workspace["buffer_flags"] - is_fusion = fusion_op != AllReduceFusionOp.NONE if is_fusion and fusion_op not in MNNVLAllReduce.SUPPORTED_FUSION_OPS: return None - outputs = torch.ops.trtllm.mnnvl_fusion_allreduce( - input, - all_reduce_params.norm_weight, # gamma - all_reduce_params.residual, # residual - all_reduce_params.eps, # epsilon - buffer_base, # comm_buffer - buffer_flags, # buffer_flags - is_fusion, # rmsnorm_fusion - all_reduce_params.scale, # scale - int(fusion_op), - ) + outputs = _launch_mnnvl_allreduce(input, workspace, self.dtype, + all_reduce_params) return tuple(outputs) if is_fusion else outputs[0] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2e5ce22d0296..f318310f42af 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -70,6 +70,7 @@ restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) from ..speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN +from ..speculative.mtp_tail_trim import mtp_tail_trim_enabled from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_static_draft_len, update_draft_len from ..utils import (get_model_extra_attrs, @@ -3807,8 +3808,20 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]): Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get('attn_metadata') + # Overlap scheduling corrects kv_lens_cuda below from the runtime + # accepted-token counts and calls on_update_kv_lens() again after it. + spec_kv_correction = (self.enable_spec_decode + and not self._disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, 'kv_lens_cuda')) # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + # MTP tail trim: when the hook is idempotent and the correction below + # repeats it with nothing reading its results in between, the first + # call is fully overwritten by the second call and skipped. + if attn_meta is not None and not (spec_kv_correction + and mtp_tail_trim_enabled() and + attn_meta.kv_lens_hook_idempotent): attn_meta.on_update_kv_lens() if self.enable_spec_decode and not self._disable_overlap_scheduler: @@ -3836,7 +3849,7 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]): self. previous_pos_id_offsets_cuda[:previous_batch_tokens]) - if hasattr(inputs['attn_metadata'], 'kv_lens_cuda'): + if spec_kv_correction: if num_ctx_requests >= num_chunked_ctx_requests and num_chunked_ctx_requests > 0: # The generation requests with draft_tokens are treated as chunked context requests when extend_ctx returns True. inputs['attn_metadata'].kv_lens_cuda[ diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 55721b44aceb..e19bf7868298 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -15,7 +15,8 @@ from ..attention.backends import AttentionMetadata from ..attention.backends.flashinfer import FlashInferAttentionMetadata -from ..attention.backends.sparse.params import MTPIndexShareMetadata +from ..attention.backends.sparse.params import (MTPIndexerScheduleTrimMetadata, + MTPIndexShareMetadata) from ..model_config import ModelConfig from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest @@ -24,6 +25,7 @@ from .interface import (INVALID_PROMPT_LOOKAHEAD_TOKEN, SpecMetadata, SpecWorkerBase) from .mtp import _select_mtp_position_ids +from .mtp_tail_trim import mtp_tail_trim_enabled from .sa_enhancer import SADraftEnhancer from .spec_tree_manager import SpecTreeManager @@ -38,6 +40,24 @@ def _reset_mtp_index_share(attn_metadata: MTPIndexShareMetadata) -> None: attn_metadata.set_mtp_num_accepted(None) +def _update_for_next_draft_step(attn_metadata, + skip_indexer_sched: bool) -> None: + """``update_for_spec_dec`` before the next draft step. + + ``skip_indexer_sched`` (MTP tail trim) leaves the indexer's scheduler + metadata untouched for a step nothing reads it in; the invariant is + documented on ``DSAtrtllmAttentionMetadata._skip_indexer_sched_metadata``. + """ + if not skip_indexer_sched: + attn_metadata.update_for_spec_dec() + return + attn_metadata.set_skip_indexer_sched_metadata(True) + try: + attn_metadata.update_for_spec_dec() + finally: + attn_metadata.set_skip_indexer_sched_metadata(False) + + class Eagle3ResourceManager(BaseResourceManager): """ Eagle3 needs to save the hidden states for the draft model. When using @@ -781,6 +801,19 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # config turns the reuse on. uses_mtp_index_share = self.is_mtp_eagle and isinstance( attn_metadata, MTPIndexShareMetadata) + # TRTLLM_MTP_TAIL_TRIM (see mtp_tail_trim.py): numerics-preserving + # draft-loop glue trims. From draft step 1 on every sequence has one + # token and ``gather_ids`` is ``batch_indices_cuda[:batch_size]`` = + # arange(batch_size), so its row gathers are prefix slices. + tail_trim = self.is_mtp_eagle and mtp_tail_trim_enabled() + # Index-sharing metadata that can also skip its indexer schedule + # rebuild for a draft step nothing reads it in (see the Protocol). + trims_indexer_sched = (tail_trim and uses_mtp_index_share + and isinstance(attn_metadata, + MTPIndexerScheduleTrimMetadata)) + # Whether draft steps >= 1 take the indexer's reuse-TopK path; fixed + # once step 0 has run. + reuses_topk = False with contextlib.ExitStack() as draft_scope: if uses_mtp_index_share: attn_metadata.set_in_mtp_draft_loop(True) @@ -815,6 +848,12 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # uses try/finally); attn_metadata is left untouched here. hidden_states, hidden_states_to_save = self._run_draft_forward( draft_model, inputs, spec_metadata, i) + if i == 0: + # Constant for the rest of the loop: steps >= 1 run with + # indexer_skip_topk, so whether they reuse this step's + # TopK (and never read the MQA-logits schedule) is known. + reuses_topk = (trims_indexer_sched and + attn_metadata.mtp_next_step_reuses_topk()) # Compute gather_ids: on the first draft step each generation # request may have accepted multiple tokens, so we index into @@ -832,6 +871,15 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) else: gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + # MTP tail trim: from step 1 on gather_ids == arange(batch_size), + # so the row selections are prefix slices, and the hidden-state + # rows are selected once for the LM head and the next step. + # With the trim off every site keeps its own + # hidden_states[gather_ids] gather, exactly as upstream. + row_sel = slice(0, batch_size) if (tail_trim + and i > 0) else gather_ids + if tail_trim: + hidden_states_sel = hidden_states[row_sel] if self.guided_decoder is not None: new_tokens = inputs["input_ids"][gather_ids] @@ -871,8 +919,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, use_lm_head_tp_in_adp = (lm_head_tp_in_adp_configured and not advanced_draft_sampling) if self.is_mtp_eagle: + head_rows = (hidden_states_sel + if tail_trim else hidden_states[gather_ids]) if use_lm_head_tp_in_adp: - hidden_states_gathered = hidden_states[gather_ids] + hidden_states_gathered = head_rows token_count = hidden_states_gathered.view( -1, hidden_states_gathered.shape[-1]).shape[0] max_num_requests = spec_metadata.max_num_requests @@ -906,19 +956,17 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, local_full_vocab_forward = getattr( shared_head, "forward_local_full_vocab", None) if local_full_vocab_forward is None: - logits = draft_model.lm_head( - hidden_states[gather_ids]) + logits = draft_model.lm_head(head_rows) else: logits = local_full_vocab_forward( - hidden_states[gather_ids], + head_rows, draft_model.lm_head, attn_metadata, True, ) else: logits = draft_model.mtp_layers[0].shared_head( - hidden_states[gather_ids], draft_model.lm_head, - attn_metadata, True) + head_rows, draft_model.lm_head, attn_metadata, True) else: logits = draft_model.logits_processor( hidden_states[gather_ids], draft_model.lm_head, @@ -964,14 +1012,24 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # gather_ids to get one hidden state per request. # Eagle3: the EAGLE draft model returns a secondary # ``hidden_states_to_save`` specifically for this purpose. - if self.is_mtp_eagle: - hidden_states = hidden_states[gather_ids] + is_last_step = i == runtime_draft_len - 1 + if tail_trim and is_last_step: + # No next draft step consumes these. + hidden_states = position_ids = None else: - hidden_states = hidden_states_to_save[gather_ids] - position_ids = (_select_mtp_position_ids( - inputs["position_ids"], gather_ids) + 1) - - # Update attn_metadata for the next iteration. + if self.is_mtp_eagle: + hidden_states = (hidden_states_sel if tail_trim else + hidden_states[gather_ids]) + else: + hidden_states = hidden_states_to_save[gather_ids] + position_ids = (_select_mtp_position_ids( + inputs["position_ids"], row_sel) + 1) + + # Update attn_metadata for the next iteration. MTP tail trim: + # skip the indexer schedule rebuild when its next consumer is + # a full rebuild (index-shared draft step / next target step). + skip_indexer_sched = trims_indexer_sched and (is_last_step + or reuses_topk) if i == 0: attn_metadata._seq_lens[:batch_size].fill_(1) attn_metadata._seq_lens_cuda[:batch_size].fill_(1) @@ -994,7 +1052,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, self._prepare_flash_mla_generation_layout( attn_metadata, num_contexts, batch_size) if hasattr(attn_metadata, 'kv_lens_cuda'): - attn_metadata.update_for_spec_dec() + _update_for_next_draft_step(attn_metadata, + skip_indexer_sched) # Both Eagle3 and MTP Eagle drafters take ``draft_len + 1`` # tokens in the first draft step (attention runs in spec-dec @@ -1005,7 +1064,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, else: if hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[:batch_size] += 1 - attn_metadata.update_for_spec_dec() + _update_for_next_draft_step(attn_metadata, + skip_indexer_sched) inputs = { "input_ids": new_draft_token, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 37992e2f8835..ac3d2cb59509 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -33,6 +33,7 @@ TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE from ..pyexecutor.resource_manager import ResourceManagerType +from .mtp_tail_trim import gather_draft_argmax_pairs, mtp_tail_trim_enabled if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder @@ -2691,6 +2692,12 @@ def greedy_sample_draft_with_tp_gather(self, if (sharded and mapping is not None and getattr(mapping, "tp_size", 1) > 1 and not mapping.enable_attention_dp): + if mtp_tail_trim_enabled() and logits.stride(-1) == 1: + # One Triton launch for the local (idx, max) pair + the + # dedicated MNNVL mailbox exchange (NCCL fallback); the + # gathered tensor is bit-identical to the path below. + gathered = gather_draft_argmax_pairs(logits, mapping) + return self._get_draft_tokens_from_gathered(gathered) from ..distributed.ops import allgather combined = self._get_local_max_and_combined(logits) gathered = allgather(combined, mapping, dim=-1) diff --git a/tensorrt_llm/_torch/speculative/mtp_tail_trim.py b/tensorrt_llm/_torch/speculative/mtp_tail_trim.py new file mode 100644 index 000000000000..d5aeb7649342 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/mtp_tail_trim.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 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 draft-step glue trim (``TRTLLM_MTP_TAIL_TRIM``, on unless set to "0"). + +Numerics-preserving host/launch-level trims of the one-model MTP draft loop: + +* ``local_argmax_pack``: one Triton launch replaces ``torch.max(dim=-1)`` (a + single-CTA ``reduce_kernel`` over each rank's vocab shard) plus the int32 + cast / rank offset / two float casts / stack that pack each rank's (global + argmax index, max value) pair, and writes the pair straight into the + exchange buffer (the MNNVL mailbox row or the NCCL ``[rows, 2]`` input). + Same tie-break (lowest index) and NaN rule (NaN wins, lowest NaN index) as + ``at::native`` ``GreaterOrNan``; bf16 -> fp32 is exact, so the packed values + are bit-identical. +* ``DraftArgmaxMailbox``: exchanges the packed pairs through a one-shot MNNVL + SUM allreduce over a DEDICATED Lamport workspace instead of the NCCL + allgather. Each rank writes its pair into its own slot of a zero row, so the + fp32 sum is exact. The workspace is private because the one-shot kernel + clears the dirty Lamport buffer left by the PREVIOUS call on the same + workspace with the CURRENT launch's threads: a tiny payload must never + follow a large one on a shared workspace, or its few threads end up + clearing the model allreduce's payload. +""" + +import functools +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm.functional import AllReduceParams +from tensorrt_llm.logger import logger + +from ..distributed.ops import ( + _MNNVL_ONE_SHOT_THRESHOLD_BYTES, + MNNVLAllReduce, + _build_mnnvl_workspace, + _get_mnnvl_workspace_comm, + _launch_mnnvl_allreduce, + allgather, +) + +MTP_TAIL_TRIM_ENV = "TRTLLM_MTP_TAIL_TRIM" + + +@functools.lru_cache(maxsize=1) +def mtp_tail_trim_enabled() -> bool: + """Process-constant gate: unset / "1" = trimmed path, "0" = original path. + + Cached so the choice is fixed before CUDA-graph capture. + """ + enabled = os.environ.get(MTP_TAIL_TRIM_ENV, "1").strip() != "0" + logger.info( + f"MTP tail trim {'ON' if enabled else 'OFF'} " + f"({MTP_TAIL_TRIM_ENV}={os.environ.get(MTP_TAIL_TRIM_ENV)})" + ) + return enabled + + +# int32 max: "no valid column seen yet" in the per-lane running argmax. +# Instantiated (not annotated) so @triton.jit kernels may read it. +_SENTINEL_INDEX = tl.constexpr(2147483647) + + +@triton.jit +def _argmax_first_nan_combine(v1, i1, v2, i2): + # Total order of at::native GreaterOrNan: NaN beats everything, among + # NaNs or equal values the lower index wins, otherwise the larger value. + n1 = v1 != v1 + n2 = v2 != v2 + lower = i1 < i2 + take1 = (n1 & ((~n2) | lower)) | ((~n1) & (~n2) & ((v1 > v2) | ((v1 == v2) & lower))) + return tl.where(take1, v1, v2), tl.where(take1, i1, i2) + + +@triton.jit(do_not_specialize=["n_cols", "stride_out", "idx_offset", "slot", "out_width"]) +def _draft_local_argmax_pack_kernel( + logits_ptr, + out_ptr, + n_cols, + stride_row, + stride_out, + idx_offset, + slot, + out_width, + BLOCK: tl.constexpr, + OUT_BLOCK: tl.constexpr, + BF16_BITS: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + base = logits_ptr + row.to(tl.int64) * stride_row + best_v = tl.full([BLOCK], float("-inf"), tl.float32) + best_i = tl.full([BLOCK], _SENTINEL_INDEX, tl.int32) + for start in range(0, n_cols, BLOCK): + cols = start + offs + m = cols < n_cols + if BF16_BITS: + # Widen by moving the bf16 bits into the fp32 high half, exactly + # like at::BFloat16 -> float: keeps NaN payloads, which a cvt + # would canonicalize. + raw = tl.load(base + cols, mask=m, other=float("-inf")) + bits = raw.to(tl.int16, bitcast=True).to(tl.int32) << 16 + v = bits.to(tl.float32, bitcast=True) + else: + v = tl.load(base + cols, mask=m, other=float("-inf")).to(tl.float32) + # Chunks are visited in increasing column order, so within a lane + # the new element always has the larger index: replace only on a + # strictly better value (or NaN over non-NaN), or the first valid. + nv = v != v + nb = best_v != best_v + take = m & ((best_i == _SENTINEL_INDEX) | (nv & (~nb)) | ((~nv) & (~nb) & (v > best_v))) + best_v = tl.where(take, v, best_v) + best_i = tl.where(take, cols, best_i) + vmax, imax = tl.reduce((best_v, best_i), 0, _argmax_first_nan_combine) + if BF16_BITS: + # torch's bf16 max reduces in fp32 and rounds back through + # c10::BFloat16, which maps every NaN to 0x7FFF: match its bits. + vmax = tl.where( + vmax != vmax, tl.full([], 0x7FFF0000, tl.int32).to(tl.float32, bitcast=True), vmax + ) + idx_f = (imax + idx_offset).to(tl.float32) + ocols = tl.arange(0, OUT_BLOCK) + zeros = tl.zeros([OUT_BLOCK], tl.float32) + packed = tl.where(ocols == slot, idx_f, tl.where(ocols == slot + 1, vmax, zeros)) + tl.store(out_ptr + row.to(tl.int64) * stride_out + ocols, packed, mask=ocols < out_width) + + +def local_argmax_pack( + logits: torch.Tensor, idx_offset: int, out: torch.Tensor, slot: int = 0 +) -> torch.Tensor: + """Write (float(argmax + idx_offset), float(max)) of each logits row into + ``out[row, slot:slot + 2]`` and zeros into the rest of ``out[row]``. + + Bit-identical to ``SpecWorkerBase._get_local_max_and_combined`` (which is + ``torch.max(logits, -1)`` + int32 offset + float casts + stack) placed at + ``slot``. ``logits`` is 2D ``[rows, vocab_shard]`` with unit column stride; + ``out`` is a contiguous fp32 ``[rows, width]``. + """ + assert logits.dim() == 2 and logits.stride(-1) == 1 + assert out.dtype == torch.float32 and out.is_contiguous() + rows, n_cols = logits.shape + width = out.shape[1] + assert out.shape[0] == rows and slot + 2 <= width + if rows == 0: + return out + block = min(triton.next_power_of_2(n_cols), 8192) + num_warps = 16 if block >= 8192 else (8 if block >= 2048 else 4) + _draft_local_argmax_pack_kernel[(rows,)]( + logits, + out, + n_cols, + logits.stride(0), + out.stride(0), + idx_offset, + slot, + width, + BLOCK=block, + OUT_BLOCK=triton.next_power_of_2(width), + BF16_BITS=logits.dtype == torch.bfloat16, + num_warps=num_warps, + ) + return out + + +# Plain SUM: no fusion epilogue. +_SUM_ONLY = AllReduceParams() + + +class DraftArgmaxMailbox: + """Dedicated one-shot MNNVL SUM allreduce used as an allgather of the + per-rank draft argmax pairs (see the module docstring).""" + + WIDTH = 128 # fp32 per row: one warp of float4 accesses in the kernel + _instances: dict = {} + # Communicators split for a mapping whose mailbox construction failed + # (McastGPUBuffer only borrows them): kept for reuse instead of leaking one + # world-wide split per attempt, like allreduce_mnnvl_pending_comms does + # for the model workspace. + _pending_comms: dict = {} + + def __init__(self, mapping): + self.mapping = mapping + tp = mapping.tp_size + # Size the workspace so every exchange stays under the one-shot / + # two-shot switch in allreduceOp.cpp: one Lamport buffer holds the + # whole gathered payload (rows x WIDTH fp32 from each of the tp ranks). + self.max_rows = _MNNVL_ONE_SHOT_THRESHOLD_BYTES // (self.WIDTH * 4 * tp) + buffer_size_bytes = self.max_rows * self.WIDTH * 4 * tp + # Same communicator flavour as the model's MNNVL workspace (an MPI + # split under MPI, the TP process group under a non-MPI orchestrator). + comm = self._pending_comms.get(mapping) + if comm is None: + comm = _get_mnnvl_workspace_comm(mapping) + self._pending_comms[mapping] = comm + # Converged across the group: every rank gets a workspace or raises. + self._workspace = _build_mnnvl_workspace(comm, mapping, buffer_size_bytes) + # Hand ownership of the communicator to the workspace. + self._pending_comms.pop(mapping, None) + # Persistent pack-kernel target; the kernel fully overwrites the rows + # it is given on every exchange. + self._staging = torch.empty( + (self.max_rows, self.WIDTH), + dtype=torch.float32, + device=self._workspace["buffer_flags"].device, + ) + + @classmethod + def get(cls, mapping) -> Optional["DraftArgmaxMailbox"]: + """Mailbox for ``mapping`` or None (NCCL allgather) when unsupported. + + Built on the first eager call (CUDA-graph warm-up runs eagerly before + capture); every TP rank reaches it in the same step, and the success + is converged across the group so all ranks take the same path. + Requires a pure-TP world (the communicator split is collective over + the world) on an MNNVL fabric, the same condition under which the + model's own allreduces use the MNNVL kernels: a group that spans nodes + qualifies on its own, a single-node group only once the model has + built its own MNNVL workspace for this mapping (i.e. MNNVL was + requested explicitly). + """ + if mapping in cls._instances: + return cls._instances[mapping] + if torch.cuda.is_current_stream_capturing(): + return None # never allocate under capture; NCCL for this graph + box = None + try: + tp = getattr(mapping, "tp_size", 1) + model_uses_mnnvl = mapping in MNNVLAllReduce.allreduce_mnnvl_workspaces + eligible = ( + tp > 1 + and 2 * tp <= cls.WIDTH + and mapping.world_size == tp + and not mapping.has_cp() + and not mapping.enable_attention_dp + and MNNVLAllReduce.is_mnnvl( + mapping, torch.float32, explicitly_requested=model_uses_mnnvl + ) + ) + if eligible: + box = cls(mapping) + logger.info( + "MTP tail trim: draft argmax pairs exchanged through a " + f"dedicated MNNVL one-shot mailbox (tp={tp}, rows<=" + f"{box.max_rows})" + ) + except Exception as e: # noqa: BLE001 - fall back to NCCL + logger.warning( + f"MTP tail trim: MNNVL draft-argmax mailbox " + f"unavailable ({e}); using the NCCL allgather" + ) + box = None + cls._instances[mapping] = box + return box + + def staging_rows(self, rows: int) -> torch.Tensor: + """Pack-kernel target for an exchange of ``rows`` rows: a contiguous + ``[rows, WIDTH]`` prefix of the persistent staging buffer.""" + return self._staging[:rows] + + def exchange(self, mailbox_rows: torch.Tensor) -> torch.Tensor: + """SUM-allreduce of ``[rows, WIDTH]`` fp32 rows (one pair per rank).""" + return _launch_mnnvl_allreduce(mailbox_rows, self._workspace, torch.float32, _SUM_ONLY)[0] + + +def gather_draft_argmax_pairs(logits: torch.Tensor, mapping) -> torch.Tensor: + """Trimmed ``_get_local_max_and_combined`` + TP allgather along dim -1. + + Returns ``[rows, 2 * tp_size]`` fp32, rank-major (idx, val) pairs, exactly + the tensor the original NCCL path produced (the one exception: a NaN max + comes back with the canonical NaN payload from the mailbox's fp32 sum; + the argmax that consumes it treats every NaN alike, so tokens match). + """ + rows, vocab_per_rank = logits.shape + tp, rank = mapping.tp_size, mapping.tp_rank + idx_offset = rank * vocab_per_rank + box = DraftArgmaxMailbox.get(mapping) + if box is not None and rows <= box.max_rows: + mailbox = box.staging_rows(rows) + local_argmax_pack(logits, idx_offset, mailbox, slot=2 * rank) + return box.exchange(mailbox)[:, : 2 * tp] + combined = torch.empty((rows, 2), dtype=torch.float32, device=logits.device) + local_argmax_pack(logits, idx_offset, combined, slot=0) + return allgather(combined, mapping, dim=-1) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 456697cc3676..7352f18816c2 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -228,6 +228,7 @@ l0_b200: # function's leaves on Blackwell, keep the rest of the sampler dir. - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py + - unittest/_torch/speculative/test_mtp_tail_trim.py -m "not gpu2" - unittest/_torch/speculative/test_dspark_cute_dsl_attention.py - unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py # On-device suffix-automaton spec-dec kernel tests (CUDA graph capture; diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 25a83b2d12f5..f20f5760f31d 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -49,6 +49,7 @@ l0_gb200_multi_gpus: - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses + - unittest/_torch/speculative/test_mtp_tail_trim.py -m "gpu2" - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 31659d64e83c..727e71fb9712 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -62,6 +62,7 @@ l0_h100: # function's leaves on Hopper, keep the rest of the sampler dir. - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py + - unittest/_torch/speculative/test_mtp_tail_trim.py -m "not gpu2" - unittest/_torch/speculative/test_fused_sampling_op.py - unittest/_torch/speculative/test_rejection_buffers_guard.py - unittest/_torch/speculative/test_sa_hybrid_state_promotion.py diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index a668eec0bb87..af0fee5e25c5 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -3920,14 +3920,20 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper +@pytest.mark.parametrize("stale_sched", [False, True]) @pytest.mark.parametrize("step0_mode", ["gen", "context"]) @pytest.mark.parametrize("batch_size", [1, 2]) -def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size): +def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size, stale_sched): """Verify mtp_index_share: draft step 0 stashes each request's last-token Top-K and draft steps > 0 reuse it verbatim (bit-exact), across > 1 draft step. Covers both step-0 shapes: "gen" (steady state, next_n > 1) and "context" (the first gen round after prefill, which runs the context path). Step-0 Top-K correctness is covered by test_indexer_decode_custom_vs_fallback. + + stale_sched: the reuse steps run with the MQA-logits schedule marked stale + and poisoned, as TRTLLM_MTP_TAIL_TRIM leaves it (the draft loop skips its + rebuild before index-shared steps); reuse must not read it, and a step that + computes Top-K on a stale schedule must raise instead. """ torch.manual_seed(7) heads, head_dim, block_size = 32, 128, 64 @@ -4068,12 +4074,21 @@ def make_inputs(n_tokens): meta.in_mtp_draft_loop = True meta.shared_topk_indices = stash meta.indexer_skip_topk = True + if stale_sched: + meta.mtp_sched_metadata_stale = True + meta.scheduler_metadata_buffer.fill_(-1) + meta.scheduler_metadata_buffer_full_next_n.fill_(-1) + meta.kv_lens_cuda_2d.fill_(-1) hs, qs, ks_fp8, ks_scale, ws = make_inputs(batch_size) indexer._update_k_cache(ks_fp8, ks_scale, meta) topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws) assert torch.equal(topk, stash[:batch_size, :]), ( f"{step0_mode} draft reuse step {step} should copy the stash 1:1 (next_n=1)" ) + if stale_sched: + meta.indexer_skip_topk = False + with pytest.raises(RuntimeError, match="stale"): + indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws) @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") diff --git a/tests/unittest/_torch/speculative/test_mtp_tail_trim.py b/tests/unittest/_torch/speculative/test_mtp_tail_trim.py new file mode 100644 index 000000000000..d4e4aa7fa4ac --- /dev/null +++ b/tests/unittest/_torch/speculative/test_mtp_tail_trim.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. +"""TRTLLM_MTP_TAIL_TRIM: the trimmed MTP draft-step glue is bit-identical. + +Single GPU: the fused local-argmax pack kernel against the original +``_get_local_max_and_combined`` and a mocked MTP-Eagle draft loop with the trim +on vs off. Two GPUs (``mpi_pool_executor``): the dedicated MNNVL mailbox +exchange against the NCCL allgather, eager and under CUDA-graph replay, once +right after a model-sized MNNVL allreduce on the shared workspace. +""" + +import gc +import os +import pickle +import sys +import traceback +from contextlib import contextmanager +from types import SimpleNamespace + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from utils.util import skip_num_gpus_less_than + +import tensorrt_llm +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce, allgather +from tensorrt_llm._torch.speculative import eagle3 as eagle3_mod +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import SpecWorkerBase +from tensorrt_llm._torch.speculative.mtp_tail_trim import ( + DraftArgmaxMailbox, + gather_draft_argmax_pairs, + local_argmax_pack, +) +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +# needed since we reuse the mpi executor pool, first test running will leak a thread +pytestmark = pytest.mark.threadleak(enabled=False) + + +def _ref_combined(logits, tp_rank): + """The original path: SpecWorkerBase._get_local_max_and_combined.""" + fake = SimpleNamespace(mapping=SimpleNamespace(tp_rank=tp_rank)) + return SpecWorkerBase._get_local_max_and_combined(fake, logits) + + +def _draft_tokens(gathered): + """The original consumer: SpecWorkerBase._get_draft_tokens_from_gathered.""" + return SpecWorkerBase._get_draft_tokens_from_gathered(SimpleNamespace(), gathered) + + +def _make_logits(rows, vocab, dtype, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(rows, vocab, generator=g, device="cuda").to(dtype) + if rows > 1: + x[1, :] = x[1, 0] # all equal -> index 0 + if rows > 2: + x[2, 5] = x[2, vocab - 3] = x[2].max() + 1 # tie -> lower index + if rows > 3: + x[3, :] = float("-inf") # all -inf -> index 0 + if rows > 4: + x[4, 7] = float("nan") # NaN wins + x[4, 9] = float("nan") + if rows > 5: + x[5, vocab - 1] = float("inf") + return x + + +def _bits(t): + return t.contiguous().view(torch.int32) + + +def _assert_same(actual, expected): + """Equal values, NaN == NaN (the fp32 sum in the allreduce may + canonicalize a NaN payload; argmax treats every NaN alike).""" + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize( + "rows,vocab", [(1, 19360), (6, 19360), (7, 77), (9, 38720), (3, 8192), (2, 151936)] +) +@pytest.mark.parametrize("tp_rank", [0, 5]) +def test_local_argmax_pack_matches_torch(dtype, rows, vocab, tp_rank): + logits = _make_logits(rows, vocab, dtype, seed=rows * 1000 + vocab) + ref = _ref_combined(logits, tp_rank) + out = torch.empty(rows, 2, dtype=torch.float32, device="cuda") + local_argmax_pack(logits, tp_rank * vocab, out, slot=0) + assert torch.equal(_bits(out), _bits(ref)), (out, ref) + + # Mailbox layout: the pair at slot 2 * rank, exact zeros elsewhere. + box = torch.full((rows, 128), 7.0, dtype=torch.float32, device="cuda") + local_argmax_pack(logits, tp_rank * vocab, box, slot=2 * tp_rank) + expect = torch.zeros_like(box) + expect[:, 2 * tp_rank : 2 * tp_rank + 2] = ref + assert torch.equal(_bits(box), _bits(expect)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_local_argmax_pack_strided_rows_and_cuda_graph(): + base = _make_logits(8, 19360 + 64, torch.bfloat16, seed=3) + logits = base[:, :19360] # row stride != vocab + ref = _ref_combined(logits, 2) + out = torch.empty(8, 2, dtype=torch.float32, device="cuda") + local_argmax_pack(logits, 2 * 19360, out) + assert torch.equal(_bits(out), _bits(ref)) + + static_in = logits.clone() + static_out = torch.empty(8, 2, dtype=torch.float32, device="cuda") + local_argmax_pack(static_in, 2 * 19360, static_out) # warm up / compile + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + local_argmax_pack(static_in, 2 * 19360, static_out) + for seed in (11, 12): + new = _make_logits(8, 19360, torch.bfloat16, seed=seed) + static_in.copy_(new) + graph.replay() + assert torch.equal(_bits(static_out), _bits(_ref_combined(new, 2))) + + +# -------------------------------------------------------------------------- +# Mocked MTP-Eagle draft loop: trim on vs off must feed every draft step the +# same inputs and produce the same draft tokens. + + +class _FakeMetadata: + def __init__(self, seq_lens, kv_lens, num_contexts, device): + self._seq_lens = seq_lens.clone() + self._seq_lens_cuda = seq_lens.clone().to(device) + self.kv_lens_cuda = kv_lens.clone().to(device) + self.num_contexts = num_contexts + self.num_ctx_tokens = int(seq_lens[:num_contexts].sum()) + self.kv_cache_manager = None + self.use_spec_decoding = True + self.kv_snapshots = [] + + @property + def seq_lens_cuda(self): + return self._seq_lens_cuda + + def on_update(self): + pass + + def update_for_spec_dec(self): + self.kv_snapshots.append(self.kv_lens_cuda.clone()) + + +class _FakeMTPLayer: + def __init__(self, hidden, vocab, device): + g = torch.Generator(device="cpu").manual_seed(0) + self.w_h = (torch.randn(hidden, hidden, generator=g) * 0.2).to(device) + self.embed = torch.randn(vocab, hidden, generator=g).to(device) + self.head = torch.randn(hidden, vocab, generator=g).to(device) + self.calls = [] + + def __call__( + self, + embed_tokens, + all_rank_num_tokens, + input_ids, + position_ids, + hidden_states, + attn_metadata, + spec_metadata, + ): + self.calls.append((input_ids.clone(), position_ids.clone(), hidden_states.clone())) + out = torch.tanh( + hidden_states @ self.w_h + + self.embed[input_ids.long()] + + position_ids.float().unsqueeze(-1) * 0.01 + ) + return out + + def shared_head(self, hidden_states, lm_head, attn_metadata, return_context_logits): + return hidden_states @ self.head + + +def _run_mocked_draft_loop(monkeypatch, trim, draft_len=5): + device = "cuda" if torch.cuda.is_available() else "cpu" + hidden, vocab = 32, 97 + num_contexts, batch_size = 1, 3 + ctx_len = 4 + seq_lens = torch.tensor( + [ctx_len] + [draft_len + 1] * (batch_size - num_contexts), dtype=torch.int32 + ) + num_tokens = int(seq_lens.sum()) + g = torch.Generator(device="cpu").manual_seed(1) + input_ids = torch.randint(0, vocab, (num_tokens,), generator=g, dtype=torch.int32).to(device) + position_ids = torch.arange(num_tokens, dtype=torch.int32).to(device) + 40 + hidden_states = torch.randn(num_tokens, hidden, generator=g).to(device) + num_accepted = torch.tensor([1, 3, 1], dtype=torch.int32).to(device) + kv_lens = torch.tensor([ctx_len, 50, 60], dtype=torch.int32) + + md = _FakeMetadata(seq_lens, kv_lens, num_contexts, device) + layer = _FakeMTPLayer(hidden, vocab, device) + draft_model = SimpleNamespace(mtp_layers=[layer], lm_head=None, embed_tokens=None) + spec_metadata = SimpleNamespace( + runtime_draft_len=draft_len, + batch_indices_cuda=torch.arange(8, dtype=torch.int32).to(device), + wants_advanced_draft_sampling=False, + all_rank_num_tokens=None, + subseq_all_rank_num_tokens=None, + ) + + @contextmanager + def draft_context(attn_metadata, manager): + yield attn_metadata + + def run_draft_forward(dm, inputs, sm, i): + return dm.mtp_layers[0](embed_tokens=None, all_rank_num_tokens=None, **inputs), None + + def sample_draft_tokens(logits, sm, bs, draft_step=None, mapping_lm_head_tp=None): + return torch.argmax(logits, dim=-1).type(torch.int32) + + worker = SimpleNamespace( + is_mtp_eagle=True, + model_config=None, + guided_decoder=None, + sa_enhancer=None, + draft_kv_cache_context=draft_context, + _run_draft_forward=run_draft_forward, + sample_draft_tokens=sample_draft_tokens, + ) + monkeypatch.setattr(eagle3_mod, "mtp_tail_trim_enabled", lambda: trim) + inputs = { + "input_ids": input_ids, + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": md, + "spec_metadata": spec_metadata, + } + tokens = Eagle3OneModelWorker._forward_linear_draft_loop( + worker, + inputs, + md, + spec_metadata, + draft_model, + None, + num_contexts, + batch_size, + num_accepted, + None, + ) + return tokens, layer.calls, md.kv_snapshots + + +@pytest.mark.parametrize("draft_len", [1, 2, 5]) +def test_draft_loop_trim_is_bit_identical(monkeypatch, draft_len): + ref_tokens, ref_calls, ref_kv = _run_mocked_draft_loop(monkeypatch, False, draft_len) + out_tokens, out_calls, out_kv = _run_mocked_draft_loop(monkeypatch, True, draft_len) + assert torch.equal(ref_tokens, out_tokens) + assert len(ref_calls) == len(out_calls) == draft_len + for (ri, rp, rh), (oi, op, oh) in zip(ref_calls, out_calls): + assert torch.equal(ri, oi) + assert torch.equal(rp, op) + assert rh.dtype == oh.dtype and torch.equal(rh, oh) + assert len(ref_kv) == len(out_kv) + for a, b in zip(ref_kv, out_kv): + assert torch.equal(a, b) + + +# -------------------------------------------------------------------------- +# Multi-rank: the dedicated MNNVL mailbox exchange vs the NCCL allgather. + +_MNNVL_TEST_ENV = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR") + + +def _check_mailbox_matches_nccl(mapping: Mapping, rank: int, world: int) -> None: + model_ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.MNNVL, dtype=torch.bfloat16) + assert model_ar.mnnvl_allreduce is not None + big = torch.randn(6, 6144, device="cuda").to(torch.bfloat16) # model-sized payload + model_ar(big) # the model workspace exists before the mailbox eligibility check + box = DraftArgmaxMailbox.get(mapping) + assert box is not None, "dedicated mailbox did not initialise" + + def ref_path(logits): + return allgather(_ref_combined(logits, rank), mapping, dim=-1) + + def check(out, ref): + _assert_same(out, ref) + torch.testing.assert_close(_draft_tokens(out), _draft_tokens(ref), rtol=0, atol=0) + + # rows > max_rows takes the NCCL fallback with the pack kernel. + for rows in (1, 2, 6, 64, box.max_rows, box.max_rows + 1): + for vocab in (19360, 77): + logits = _make_logits(rows, vocab, torch.bfloat16, seed=100 * rows + vocab + 7 * rank) + if rows > 6: # same max on every rank -> lowest rank must win + logits[6, :] = 0 + logits[6, 3] = 5 + check(gather_draft_argmax_pairs(logits, mapping), ref_path(logits)) + + # A model-sized allreduce right before a one-row exchange: the mailbox's + # own workspace keeps the model's dirty Lamport buffer out of its launch. + logits = _make_logits(1, 19360, torch.bfloat16, seed=41 + rank) + model_ar(big) + check(gather_draft_argmax_pairs(logits, mapping), ref_path(logits)) + + # CUDA graph: model allreduce + pack + mailbox, replayed on fresh logits. + static_logits = _make_logits(1, 19360, torch.bfloat16, seed=rank) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + model_ar(big) + gather_draft_argmax_pairs(static_logits, mapping) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + model_ar(big) + graph_out = gather_draft_argmax_pairs(static_logits, mapping) + for it in range(3): + new = _make_logits(1, 19360, torch.bfloat16, seed=1000 * it + rank) + static_logits.copy_(new) + graph.replay() + torch.cuda.synchronize() + check(graph_out, ref_path(new)) + + +def _mailbox_worker(world: int) -> bool: + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + previous_env = {name: (name in os.environ, os.environ.get(name)) for name in _MNNVL_TEST_ENV} + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + mapping = Mapping(world_size=world, tp_size=world, rank=rank) + try: + MPI.COMM_WORLD.barrier() + MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) + DraftArgmaxMailbox._instances.pop(mapping, None) + gc.collect() + MPI.COMM_WORLD.barrier() + _check_mailbox_matches_nccl(mapping, rank, world) + except Exception: + traceback.print_exc() + raise + finally: + DraftArgmaxMailbox._instances.pop(mapping, None) + MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) + gc.collect() + for name, (was_present, value) in previous_env.items(): + if was_present: + os.environ[name] = value + else: + os.environ.pop(name, None) + return True + + +@skip_num_gpus_less_than(2) +@pytest.mark.gpu2 +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_mailbox_matches_nccl_allgather(mpi_pool_executor): + world = mpi_pool_executor.num_workers + results = mpi_pool_executor.map(_mailbox_worker, [world] * world) + for r in results: + assert r is True