[None][perf] MTP draft-step glue trim (TRTLLM_MTP_TAIL_TRIM, default on; "0" = original path)

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

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

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()

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
 

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,

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]