[None][perf] Run the draft layer after attention only on drafting rows

docs/source/features/speculative-decoding.md

@@ -368,3 +368,15 @@ accuracy and acceptance-length validation.
 planner for packed FP4 MLA decode. The default is disabled (0); the
 rule-based split plan remains active. Validate end-to-end latency on the
 serving workload before enabling wave simulation.
+
+### Draft layer row picking
+
+On the first draft step of MTP (vanilla and MTP-Eagle) and single-layer EAGLE 3
+one-model drafting, a draft layer that supports it runs its attention on every
+input row (writing the draft KV cache of each) but the rest of the layer
+(norms, MLP or MoE) only on the one row per request the drafter reads: the last
+prompt row of a context request, the last accepted row of a generation
+request. `TLLM_SPEC_DISABLE_DRAFT_ROW_PICK=1` turns this off. The picked rows
+are numerically equivalent to the full layer's but not bitwise identical for
+MoE draft layers, whose kernels may round differently at another row count, so
+an occasional draft token can differ.

tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/layers.py

@@ -900,6 +900,9 @@ def forward_mlp(
 class GlmMtpLayer(_AdoptedLayerMixin, _PostAttentionHopMixin, DeepseekV3MTP):
     """The draft layer: the built-in MTP layer, with the target's decode composition."""
 
+    # Step 0 keeps only the live rows through its own hop (arm_live_rows below), not the
+    # built-in pick_rows argument.
+    supports_pick_rows = False
     # The draft front packs this rank's chunk of the projection input; None keeps the
     # built-in norms and concatenation.
     draft_front = None

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1742,6 +1742,8 @@ def forward_mlp(
 
 
 class DeepseekV3MTP(DeepseekV3DecoderLayer):
+    # forward takes ``pick_rows`` (see speculative.interface.draft_layer_picks_rows).
+    supports_pick_rows = True
 
     def __init__(self,
                  model_config: ModelConfig[PretrainedConfig],
@@ -1807,8 +1809,13 @@ def forward(
         attn_metadata: AttentionMetadata,
         all_rank_num_tokens: Optional[List[int]] = None,
         spec_metadata: Optional[SpecMetadata] = None,
+        pick_rows: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
+        """With ``pick_rows`` (indices into the input rows) the attention
+        runs on every input row and the rest of the layer, which is per-token,
+        on ``pick_rows`` only; the output then has ``len(pick_rows)`` rows and
+        ``all_rank_num_tokens`` counts those rows."""
 
         def norm_embeds():
             return self.enorm(embed_tokens(input_ids))  #emdedding
@@ -1846,6 +1853,11 @@ def norm_hidden():
                 enable_allreduce=not (self.disable_attn_allreduce)),
             **kwargs,
         )
+        if pick_rows is not None:
+            # Row-wise: a TP partial sum not yet all-reduced (PRE_MOE_FUSION)
+            # picks the same as the reduced rows.
+            hidden_states = hidden_states[pick_rows]
+            residual = residual[pick_rows]
 
         # MTP Layer Must have sparse MOE
         if self.fusion_config.PRE_MOE_FUSION:

tensorrt_llm/_torch/models/modeling_glm.py

@@ -847,6 +847,9 @@ def forward_mlp(
 
 
 class Glm4MTP(Glm4DecoderLayer):
+    # forward takes ``pick_rows`` (see speculative.interface.draft_layer_picks_rows).
+    supports_pick_rows = True
+
     def __init__(
         self,
         model_config: ModelConfig[PretrainedConfig],
@@ -903,8 +906,16 @@ def forward(
         attn_metadata: AttentionMetadata,
         all_rank_num_tokens: Optional[List[int]] = None,
         lora_params: Optional[dict] = None,
+        pick_rows: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
+        """With ``pick_rows`` (indices into the input rows) the attention runs on
+        every input row and the rest of the layer, which is per-token, on
+        ``pick_rows`` only; the output then has ``len(pick_rows)`` rows and
+        ``all_rank_num_tokens`` counts those rows."""
+        # LoRA params describe every input row; the draft workers pass none to MTP layers.
+        assert pick_rows is None or lora_params is None, "pick_rows does not support LoRA"
+
         def norm_embeds():
             return self.enorm(embed_tokens(input_ids))  # emdedding
 
@@ -941,6 +952,11 @@ def norm_hidden():
             lora_params=lora_params,
             **kwargs,
         )
+        if pick_rows is not None:
+            # Row-wise: a TP partial sum not yet all-reduced (PRE_MOE_FUSION)
+            # picks the same as the reduced rows.
+            hidden_states = hidden_states[pick_rows]
+            residual = residual[pick_rows]
 
         # MTP Layer Must have sparse MOE
         if self.fusion_config.PRE_MOE_FUSION:

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -774,7 +774,10 @@ def forward(
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         spec_metadata: SpecMetadata,
+        pick_rows: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        """With ``pick_rows`` the attention runs on every input row and the
+        rest of the layer on ``pick_rows`` only."""
         residual = hidden_states
 
         hidden_states = self.hidden_norm(hidden_states)
@@ -787,6 +790,9 @@ def forward(
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
         )
+        if pick_rows is not None:
+            hidden_states = hidden_states[pick_rows]
+            residual = residual[pick_rows]
 
         hidden_states, residual = self.post_attention_layernorm(
             hidden_states, residual)
@@ -887,6 +893,10 @@ def __init__(
                 use_mla=use_mla,
                 aux_stream=self.aux_stream,
             )
+        # forward takes ``pick_rows`` (see speculative.interface.
+        # draft_layer_picks_rows) with one decoder layer: a later layer's
+        # attention would need every row.
+        self.supports_pick_rows = self.num_layers == 1
 
         self.norm = RMSNorm(
             hidden_size=config.hidden_size,
@@ -943,7 +953,16 @@ def forward(
         spec_metadata: Optional[SpecMetadata] = None,
         hidden_states: Optional[torch.Tensor] = None,
         all_rank_num_tokens: Optional[List[int]] = None,
+        pick_rows: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        """``pick_rows`` (single decoder layer only, see
+        ``supports_pick_rows``): the layer's attention runs on every input row,
+        everything after it on ``pick_rows`` only, so both outputs have
+        ``len(pick_rows)`` rows."""
+        # The worker passes pick_rows only when supports_pick_rows: a second
+        # decoder layer's attention would need every row.
+        assert pick_rows is None or self.supports_pick_rows, (
+            "pick_rows needs a single draft decoder layer")
         # When ``all_rank_num_tokens`` is supplied the caller wants this draft
         # forward to run with a different attention-DP token distribution
         # (e.g. the worker's per-step value); restore the original on exit so
@@ -991,6 +1010,7 @@ def forward(
                     hidden_states=hidden_states,
                     attn_metadata=attn_metadata,
                     spec_metadata=spec_metadata,
+                    pick_rows=pick_rows,
                 )
 
             hidden_states, hidden_states_to_save = self.norm(

tensorrt_llm/_torch/speculative/eagle3.py

@@ -12,6 +12,7 @@
 from tensorrt_llm._device_copy import DeviceCopyBuffer
 from tensorrt_llm._torch.custom_ops import inplace_slice_copy
 from tensorrt_llm._utils import prefer_pinned
+from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
 from ..attention.backends import AttentionMetadata
@@ -22,9 +23,10 @@
 from ..pyexecutor.llm_request import LlmRequest
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.scheduler import ScheduledRequests
-from .fused_accept import step0_epilogue
-from .interface import (INVALID_PROMPT_LOOKAHEAD_TOKEN, SpecMetadata,
-                        SpecWorkerBase, supports_sharded_target_sampling)
+from .fused_accept import step0_epilogue, step0_gather_ids
+from .interface import (INVALID_PROMPT_LOOKAHEAD_TOKEN, DraftRowPick,
+                        SpecMetadata, SpecWorkerBase,
+                        supports_sharded_target_sampling)
 from .mtp import _select_mtp_position_ids
 from .sa_enhancer import SADraftEnhancer
 from .spec_tree_manager import SpecTreeManager
@@ -498,6 +500,9 @@ def __init__(self,
         # Set by set_draft_model: the MTP layer applies the draft loop's
         # deferred kv_lens increment itself.
         self._draft_folds_kv_lens_increment = False
+        # Whether the draft layer runs its post-attention part on step 0 only
+        # on the rows the loop reads (see draft_layer_picks_rows).
+        self._step0_row_pick = DraftRowPick("Draft step 0 layer")
         # MXFP8 copy of the LM head for greedy draft tokens (MTP Eagle), set
         # by the model once its weights are loaded.
         self.draft_head = None
@@ -513,6 +518,16 @@ def set_draft_model(self, draft_model) -> None:
             self.is_mtp_eagle and isinstance(draft_model, MTPForCausalLM)
             and draft_model.folds_kv_lens_increment)
 
+    def _step0_draft_layer_picks_rows(self, draft_model) -> bool:
+        """Whether the draft layer of step 0 takes ``pick_rows`` (see
+        draft_layer_picks_rows)."""
+        if self.is_mtp_eagle:
+            mtp_layers = getattr(draft_model, "mtp_layers", None)
+            draft_layer = mtp_layers[0] if mtp_layers else None
+        else:
+            draft_layer = getattr(draft_model, "model", None)
+        return self._step0_row_pick.resolve([draft_layer], self.mapping)
+
     @property
     def max_draft_len(self) -> int:
         return self.spec_config.max_draft_len
@@ -906,6 +921,33 @@ def _forward_linear_draft_loop(self,
             for i in range(runtime_draft_len):
                 if uses_mtp_index_share:
                     attn_metadata.set_skip_topk(i > 0)
+                # Step 0 with a row-picking draft layer: the rows the loop
+                # reads (step 0's gather ids) are known before the forward, and
+                # the layer runs everything after its attention on them only.
+                # Under attention DP every rank makes the same decision (it
+                # depends on the model only), and the post-attention
+                # collectives get the per-request counts the later steps use.
+                pick_rows = None
+                if i == 0 and self._step0_draft_layer_picks_rows(draft_model):
+                    # set_spec_metadata_all_rank_num_tokens sets both.
+                    assert (spec_metadata.all_rank_num_tokens is None
+                            or spec_metadata.subseq_all_rank_num_tokens
+                            is not None)
+                    if use_epilogue:
+                        pick_rows = step0_gather_ids(
+                            pre_draft_seq_lens, num_accepted_tokens,
+                            num_contexts, batch_size, runtime_draft_len,
+                            attn_metadata.num_ctx_tokens)
+                    if pick_rows is None:
+                        pick_rows = self._step0_gather_ids(
+                            step0_seq_lens, num_accepted_tokens, spec_metadata,
+                            num_contexts, num_gens, runtime_draft_len,
+                            attn_metadata.num_ctx_tokens)
+                    inputs["pick_rows"] = pick_rows
+                    logger.info_once(
+                        f"Draft step 0 picks {batch_size} of "
+                        f"{inputs['input_ids'].shape[0]} rows after attention",
+                        key="spec_draft_step0_pick_rows")
                 # Run draft model (mode-specific via helper). The helper
                 # passes ``all_rank_num_tokens`` as a kwarg so the draft model
                 # handles save/restore internally (Eagle3DraftModel.forward
@@ -920,7 +962,8 @@ def _forward_linear_draft_loop(self,
                 # request, in request order, so every per-request row
                 # selection is the prefix slice [0, batch_size); each
                 # ``[gather_ids]`` site below takes a slice as readily as an
-                # index tensor.
+                # index tensor. A picked step-0 output already holds just
+                # those rows, in request order.
                 if i == 0 and use_epilogue:
                     step0 = step0_epilogue(pre_draft_seq_lens,
                                            num_accepted_tokens,
@@ -935,30 +978,30 @@ def _forward_linear_draft_loop(self,
                                            accepted_tokens=accepted_tokens if
                                            token_buffer is not None else None,
                                            last_out=token_buffer[0] if
-                                           token_buffer is not None else None)
+                                           token_buffer is not None else None,
+                                           picked_gather_ids=pick_rows)
                     if step0 is None:
                         token_buffer = None  # needs the epilogue's row 0
                 if i == 0 and step0 is not None:
                     gather_ids = step0.gather_ids
+                elif i == 0 and pick_rows is not None:
+                    gather_ids = pick_rows
                 elif i == 0:
-                    last_tokens_idx = torch.cumsum(
-                        step0_seq_lens, dim=0, dtype=torch.long) - 1
-                    start_ids_gen = (
-                        spec_metadata.batch_indices_cuda[:num_gens] *
-                        (runtime_draft_len + 1)).long()
-                    gather_ids_gen = (start_ids_gen +
-                                      num_accepted_tokens[num_contexts:] - 1 +
-                                      attn_metadata.num_ctx_tokens)
-                    gather_ids = torch.concat(
-                        [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0)
+                    gather_ids = self._step0_gather_ids(
+                        step0_seq_lens, num_accepted_tokens, spec_metadata,
+                        num_contexts, num_gens, runtime_draft_len,
+                        attn_metadata.num_ctx_tokens)
                 else:
                     gather_ids = slice(0, batch_size)
+                # The output rows of this step's requests, in request order.
+                output_rows = (slice(0, batch_size)
+                               if pick_rows is not None else gather_ids)
                 # MTP Eagle: one row per request, the LM-head input on this
                 # step and the draft input on the next.
                 if i == 0 and step0 is not None:
                     request_hidden_states = step0.hidden
                 elif self.is_mtp_eagle:
-                    request_hidden_states = hidden_states[gather_ids]
+                    request_hidden_states = hidden_states[output_rows]
                 else:
                     request_hidden_states = None
 
@@ -1065,7 +1108,7 @@ def _forward_linear_draft_loop(self,
                             attn_metadata, True)
                 else:
                     logits = draft_model.logits_processor(
-                        hidden_states[gather_ids], draft_model.lm_head,
+                        hidden_states[output_rows], draft_model.lm_head,
                         attn_metadata, True)
 
                 if draft_head is None:
@@ -1115,7 +1158,7 @@ def _forward_linear_draft_loop(self,
                     if self.is_mtp_eagle:
                         hidden_states = request_hidden_states
                     else:
-                        hidden_states = hidden_states_to_save[gather_ids]
+                        hidden_states = hidden_states_to_save[output_rows]
                     if step0 is not None:
                         position_ids = step0.positions(i + 1)
                     else:
@@ -1188,15 +1231,37 @@ def _forward_linear_draft_loop(self,
 
         return next_draft_tokens, next_new_tokens
 
-    def _get_step_all_rank_num_tokens(self, spec_metadata, step_idx: int):
+    @staticmethod
+    def _step0_gather_ids(seq_lens, num_accepted_tokens, spec_metadata,
+                          num_contexts: int, num_gens: int,
+                          runtime_draft_len: int,
+                          num_ctx_tokens: int) -> torch.Tensor:
+        """Step 0's row of each request: a context request's last prompt row,
+        a generation request's last accepted row."""
+        last_tokens_idx = torch.cumsum(seq_lens, dim=0, dtype=torch.long) - 1
+        start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] *
+                         (runtime_draft_len + 1)).long()
+        gather_ids_gen = (start_ids_gen + num_accepted_tokens[num_contexts:] -
+                          1 + num_ctx_tokens)
+        return torch.concat([last_tokens_idx[:num_contexts], gather_ids_gen],
+                            dim=0)
+
+    def _get_step_all_rank_num_tokens(self,
+                                      spec_metadata,
+                                      step_idx: int,
+                                      picks_rows: bool = False):
         """Pick the right ``all_rank_num_tokens`` for this draft iteration.
 
         Step 0 uses ``spec_metadata.all_rank_num_tokens``; subsequent steps
         use ``spec_metadata.subseq_all_rank_num_tokens`` since every sequence
-        contributes a single token per iteration.
+        contributes a single token per iteration. A step-0 forward that picks
+        one row per request after attention (``picks_rows``) also uses the
+        per-sequence counts: they are the token counts of its post-attention
+        part, the only part that reads them.
         """
         return (spec_metadata.all_rank_num_tokens
-                if step_idx == 0 else spec_metadata.subseq_all_rank_num_tokens)
+                if step_idx == 0 and not picks_rows else
+                spec_metadata.subseq_all_rank_num_tokens)
 
     def _run_draft_forward(self, draft_model, inputs, spec_metadata,
                            step_idx: int):
@@ -1205,10 +1270,11 @@ def _run_draft_forward(self, draft_model, inputs, spec_metadata,
         ``all_rank_num_tokens`` is passed as a kwarg in both modes. For MTP
         Eagle it goes directly to ``mtp_layers[0]``; for Eagle3 it goes to
         ``Eagle3DraftModel.forward`` which guards it with a try/finally so
-        attn_metadata sees the original value on return.
+        attn_metadata sees the original value on return. ``inputs`` carries
+        ``pick_rows`` only on a step 0 whose draft layer supports it.
         """
         all_rank_num_tokens = self._get_step_all_rank_num_tokens(
-            spec_metadata, step_idx)
+            spec_metadata, step_idx, picks_rows="pick_rows" in inputs)
 
         if self.is_mtp_eagle:
             hidden_states = draft_model.mtp_layers[0](